perf: Add dynamic predicates for TopK

crates/polars-expr/src/dispatch/misc.rs

@@ -16,7 +16,7 @@ use polars_plan::dsl::ReshapeDimension;
 use polars_plan::plans::FusedOperator;
 #[cfg(feature = "cov")]
 use polars_plan::plans::IRCorrelationMethod;
-use polars_plan::plans::RowEncodingVariant;
+use polars_plan::plans::{DynamicPred, RowEncodingVariant};
 use polars_row::RowEncodingOptions;
 use polars_utils::IdxSize;
 use polars_utils::pl_str::PlSmallStr;
@@ -1054,3 +1054,7 @@ pub fn repeat(args: &[Column]) -> PolarsResult<Column> {
 
     Ok(c.new_from_index(0, n))
 }
+
+pub fn dynamic_pred(columns: &[Column], pred: &DynamicPred) -> PolarsResult<Column> {
+    pred.evaluate(columns)
+}

crates/polars-expr/src/dispatch/mod.rs

@@ -527,6 +527,9 @@ pub fn function_expr_to_udf(func: IRFunctionExpr) -> SpecialEq<Arc<dyn ColumnsUd
         F::RowDecode(fs, variants) => {
             map_as_slice!(misc::row_decode, fs.clone(), variants.clone())
         },
+        F::DynamicPred { pred } => {
+            map_as_slice!(misc::dynamic_pred, &pred)
+        },
     }
 }
 

crates/polars-lazy/src/tests/optimization_checks.rs

@@ -231,7 +231,7 @@ pub fn test_slice_pushdown_sort() -> PolarsResult<()> {
     assert!(lp_arena.iter(lp).all(|(_, lp)| {
         use IR::*;
         match lp {
-            Sort { slice, .. } => *slice == Some((1, 3)),
+            Sort { slice, .. } => matches!(slice, Some((1, 3, _))),
             Slice { .. } => false,
             _ => true,
         }

crates/polars-mem-engine/src/planner/lp.rs

@@ -451,7 +451,7 @@ fn create_physical_plan_impl(
             Ok(Box::new(executors::SortExec {
                 input,
                 by_column,
-                slice,
+                slice: slice.map(|t| (t.0, t.1)),
                 sort_options,
             }))
         },

crates/polars-plan/src/plans/aexpr/function_expr/mod.rs

@@ -87,6 +87,7 @@ pub use self::struct_::IRStructFunction;
 #[cfg(feature = "trigonometry")]
 pub use self::trigonometry::IRTrigonometricFunction;
 use super::*;
+use crate::plans::optimizer::DynamicPred;
 
 #[cfg_attr(feature = "ir_serde", derive(serde::Serialize, serde::Deserialize))]
 #[derive(Clone, PartialEq, Debug)]
@@ -387,6 +388,9 @@ pub enum IRFunctionExpr {
     RowEncode(Vec<DataType>, RowEncodingVariant),
     #[cfg(feature = "dtype-struct")]
     RowDecode(Vec<Field>, RowEncodingVariant),
+    DynamicPred {
+        pred: DynamicPred,
+    },
 }
 
 impl Hash for IRFunctionExpr {
@@ -690,6 +694,9 @@ impl Hash for IRFunctionExpr {
                 fs.hash(state);
                 variants.hash(state);
             },
+            DynamicPred { pred } => {
+                pred.id().hash(state);
+            },
         }
     }
 }
@@ -904,6 +911,7 @@ impl Display for IRFunctionExpr {
             RowEncode(..) => "row_encode",
             #[cfg(feature = "dtype-struct")]
             RowDecode(..) => "row_decode",
+            DynamicPred { .. } => "dynamic_predicate",
         };
         write!(f, "{s}")
     }
@@ -1233,6 +1241,7 @@ impl IRFunctionExpr {
             F::RowEncode(..) => FunctionOptions::elementwise(),
             #[cfg(feature = "dtype-struct")]
             F::RowDecode(..) => FunctionOptions::elementwise(),
+            F::DynamicPred { .. } => FunctionOptions::elementwise(),
         }
     }
 }

crates/polars-plan/src/plans/aexpr/function_expr/schema.rs

@@ -451,6 +451,7 @@ impl IRFunctionExpr {
             }),
             #[cfg(feature = "dtype-struct")]
             RowDecode(fields, _) => mapper.with_dtype(DataType::Struct(fields.to_vec())),
+            DynamicPred { .. } => mapper.with_dtype(DataType::Boolean),
         }
     }
 

crates/polars-plan/src/plans/builder_ir.rs

@@ -184,7 +184,7 @@ impl<'a> IRBuilder<'a> {
         let ir = IR::Sort {
             input: self.root,
             by_column,
-            slice,
+            slice: slice.map(|t| (t.0, t.1, None)),
             sort_options,
         };
         let node = self.lp_arena.add(ir);

crates/polars-plan/src/plans/conversion/dsl_to_ir/mod.rs

@@ -506,7 +506,7 @@ pub fn to_alp_impl(lp: DslPlan, ctxt: &mut DslConversionContext) -> PolarsResult
             let lp = IR::Sort {
                 input,
                 by_column,
-                slice,
+                slice: slice.map(|t| (t.0, t.1, None)),
                 sort_options,
             };
 

crates/polars-plan/src/plans/conversion/ir_to_dsl.rs

@@ -1,3 +1,5 @@
+use polars_utils::format_pl_smallstr;
+
 use super::*;
 
 /// converts a node from the AExpr arena to Expr
@@ -1175,6 +1177,12 @@ pub fn ir_function_to_dsl(input: Vec<Expr>, function: IRFunctionExpr) -> Expr {
             fs.into_iter().map(|f| (f.name, f.dtype.into())).collect(),
             v,
         ),
+        IF::DynamicPred { pred } => {
+            return Expr::Display {
+                inputs: input,
+                fmt_str: Box::new(format_pl_smallstr!("{pred:?}")),
+            };
+        },
     };
 
     Expr::Function { input, function }

crates/polars-plan/src/plans/ir/format.rs

@@ -851,8 +851,12 @@ pub fn write_ir_non_recursive(
                 f.write_char('[')?;
 
                 let mut comma = false;
-                if let Some((o, l)) = slice {
-                    write!(f, "slice: ({o}, {l})")?;
+                if let Some((o, l, dyn_pred)) = slice {
+                    if let Some(dyn_pred) = &dyn_pred {
+                        write!(f, "slice: ({o}, {l}, {dyn_pred:?})")?;
+                    } else {
+                        write!(f, "slice: ({o}, {l})")?;
+                    }
                     comma = true;
                 }
                 if sort_options.maintain_order {

crates/polars-plan/src/plans/ir/mod.rs

@@ -89,7 +89,7 @@ pub enum IR {
     Sort {
         input: Node,
         by_column: Vec<ExprIR>,
-        slice: Option<(i64, usize)>,
+        slice: Option<(i64, usize, Option<DynamicPred>)>,
         sort_options: SortMultipleOptions,
     },
     Cache {

crates/polars-plan/src/plans/optimizer/mod.rs

@@ -34,7 +34,7 @@ pub use cse::NaiveExprMerger;
 use delay_rechunk::DelayRechunk;
 pub use expand_datasets::ExpandedDataset;
 use polars_core::config::verbose;
-pub use predicate_pushdown::PredicatePushDown;
+pub use predicate_pushdown::{DynamicPred, PredicateExpr, PredicatePushDown, TrivialPredicateExpr};
 pub use projection_pushdown::ProjectionPushDown;
 pub use simplify_expr::{SimplifyBooleanRule, SimplifyExprRule};
 use slice_pushdown_lp::SlicePushDown;
@@ -196,16 +196,6 @@ pub fn optimize(
         true
     };
 
-    if run_pushdowns {
-        run_projection_predicate_pushdown(
-            root,
-            ir_arena,
-            expr_arena,
-            pushdown_maintain_errors,
-            &opt_flags,
-        )?;
-    }
-
     if opt_flags.slice_pushdown() {
         let mut slice_pushdown_opt = SlicePushDown::new(
             // We don't maintain errors on slice as the behavior is much more predictable that way.
@@ -223,6 +213,16 @@ pub fn optimize(
         rules.push(Box::new(slice_pushdown_opt));
     }
 
+    if run_pushdowns {
+        run_projection_predicate_pushdown(
+            root,
+            ir_arena,
+            expr_arena,
+            pushdown_maintain_errors,
+            &opt_flags,
+        )?;
+    }
+
     if opt_flags.fast_projection() {
         rules.push(Box::new(SimpleProjectionAndCollapse::new(
             opt_flags.eager(),

crates/polars-plan/src/plans/optimizer/predicate_pushdown/dynamic.rs

@@ -0,0 +1,114 @@
+use std::any::Any;
+use std::fmt::{Debug, Formatter};
+use std::hash::Hash;
+use std::sync::RwLock;
+use std::sync::atomic::{AtomicBool, Ordering};
+
+use polars_core::frame::column::ScalarColumn;
+use polars_utils::unique_id::UniqueId;
+#[cfg(feature = "ir_serde")]
+use serde::{Deserialize, Serialize};
+
+use super::*;
+
+pub trait PredicateExpr: Send + Sync + Any {
+    // Invariant: output column must be of type `Boolean`. If true a value is
+    // included, if false it is filtered out. If None is returned it is assumed
+    // all values are needed.
+    fn evaluate(&self, columns: &[Column]) -> PolarsResult<Option<Column>>;
+}
+
+pub struct TrivialPredicateExpr;
+
+impl PredicateExpr for TrivialPredicateExpr {
+    fn evaluate(&self, _columns: &[Column]) -> PolarsResult<Option<Column>> {
+        Ok(None)
+    }
+}
+
+#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
+struct Inner {
+    #[cfg_attr(feature = "ir_serde", serde(skip))]
+    pred: RwLock<Option<Arc<dyn PredicateExpr>>>,
+    #[cfg_attr(feature = "ir_serde", serde(skip))]
+    is_set: AtomicBool,
+    id: UniqueId,
+}
+
+#[derive(Clone)]
+#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
+pub struct DynamicPred {
+    inner: Arc<Inner>,
+}
+
+impl Debug for DynamicPred {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "dynamic_pred: {:}", self.id())
+    }
+}
+
+impl PartialEq for DynamicPred {
+    fn eq(&self, other: &Self) -> bool {
+        self.id() == other.id()
+    }
+}
+
+impl Hash for DynamicPred {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.inner.id.hash(state);
+    }
+}
+
+impl DynamicPred {
+    fn new() -> Self {
+        Self {
+            inner: Arc::new(Inner {
+                pred: Default::default(),
+                is_set: Default::default(),
+                id: UniqueId::new(),
+            }),
+        }
+    }
+
+    pub fn id(&self) -> &UniqueId {
+        &self.inner.id
+    }
+
+    pub fn set(&self, pred: Arc<dyn PredicateExpr>) {
+        {
+            let mut guard = self.inner.pred.write().unwrap();
+            *guard = Some(pred);
+        }
+        self.inner.is_set.store(true, Ordering::Release);
+    }
+
+    pub fn evaluate(&self, columns: &[Column]) -> PolarsResult<Column> {
+        if self.inner.is_set.load(Ordering::Acquire) {
+            let guard = self.inner.pred.read().unwrap();
+            let dyn_func = guard.as_ref().unwrap();
+            if let Some(pred) = dyn_func.evaluate(columns)? {
+                return Ok(pred);
+            }
+        }
+
+        let s = Scalar::new(DataType::Boolean, AnyValue::Boolean(true));
+        Ok(Column::Scalar(ScalarColumn::new(
+            columns[0].name().clone(),
+            s,
+            columns[0].len(),
+        )))
+    }
+}
+
+pub fn new_dynamic_pred(node: Node, arena: &mut Arena<AExpr>) -> (Node, DynamicPred) {
+    let pred = DynamicPred::new();
+    let function = IRFunctionExpr::DynamicPred { pred: pred.clone() };
+    let options = function.function_options();
+    let aexpr = AExpr::Function {
+        input: vec![ExprIR::from_node(node, arena)],
+        function,
+        options,
+    };
+
+    (arena.add(aexpr), pred)
+}

crates/polars-plan/src/plans/optimizer/predicate_pushdown/mod.rs

@@ -1,15 +1,18 @@
+mod dynamic;
 mod group_by;
 mod join;
 mod keys;
 mod utils;
 
+pub use dynamic::{DynamicPred, PredicateExpr, TrivialPredicateExpr};
 use polars_core::datatypes::PlHashMap;
 use polars_core::prelude::*;
 use polars_utils::idx_vec::UnitVec;
 use recursive::recursive;
 use utils::*;
 
 use super::*;
+use crate::plans::optimizer::predicate_pushdown::dynamic::new_dynamic_pred;
 use crate::prelude::optimizer::predicate_pushdown::group_by::process_group_by;
 use crate::prelude::optimizer::predicate_pushdown::join::process_join;
 use crate::utils::{check_input_node, has_aexpr};
@@ -572,7 +575,31 @@ impl PredicatePushDown {
             lp @ Union { .. } => {
                 self.pushdown_and_continue(lp, acc_predicates, lp_arena, expr_arena, false)
             },
-            lp @ Sort { .. } => {
+            Sort {
+                input,
+                by_column,
+                mut slice,
+                sort_options,
+            } => {
+                if let Some((offset, len, None)) = slice
+                    && by_column.len() == 1
+                {
+                    let n = by_column[0].node();
+                    if let AExpr::Column(_) = expr_arena.get(n) {
+                        let (dyn_pred_node, pred) = new_dynamic_pred(n, expr_arena);
+                        slice = Some((offset, len, Some(pred)));
+
+                        let predicate = ExprIR::from_node(dyn_pred_node, expr_arena);
+                        insert_predicate_dedup(&mut acc_predicates, &predicate, expr_arena);
+                    }
+                }
+
+                let lp = Sort {
+                    input,
+                    by_column,
+                    slice,
+                    sort_options,
+                };
                 self.pushdown_and_continue(lp, acc_predicates, lp_arena, expr_arena, true)
             },
             lp @ Sink { .. } | lp @ SinkMultiple { .. } => {

crates/polars-plan/src/plans/optimizer/slice_pushdown_lp.rs

@@ -347,7 +347,7 @@ impl SlicePushDown {
             (Sort {input, by_column, slice, sort_options}, Some(state)) => {
                 // The slice argument on Sort should be inserted by slice pushdown,
                 // so it shouldn't exist yet (or be idempotently the same).
-                let new_slice = Some((state.offset, state.len as usize));
+                let new_slice = Some((state.offset, state.len as usize, None));
                 assert!(slice.is_none() || slice == new_slice);
 
                 // first restart optimization in inputs and get the updated LP

crates/polars-python/src/lazyframe/visitor/expr_nodes.rs

@@ -1429,6 +1429,9 @@ pub(crate) fn into_py(py: Python<'_>, expr: &AExpr) -> PyResult<Py<PyAny>> {
                 IRFunctionExpr::RowDecode(..) => {
                     return Err(PyNotImplementedError::new_err("row_decode"));
                 },
+                IRFunctionExpr::DynamicPred { .. } => {
+                    return Err(PyNotImplementedError::new_err("dynamic_pred"));
+                },
             }?,
             options: py.None(),
         }

crates/polars-python/src/lazyframe/visitor/nodes.rs

@@ -492,7 +492,7 @@ pub(crate) fn into_py(py: Python<'_>, plan: &IR) -> PyResult<Py<PyAny>> {
                 sort_options.nulls_last.clone(),
                 sort_options.descending.clone(),
             ),
-            slice: *slice,
+            slice: slice.as_ref().map(|t| (t.0, t.1)),
         }
         .into_py_any(py),
         IR::Cache { input, id } => Cache {