feat: wrap FilterExec in LanceFilterExec (#3939)

This wraps usage of Datafusion FilterExec in a LanceFilterExec node that
carries the original logical Expr for the filter alongside the wrapped
FilterExec.

The reason for this is to enable a datafusion optimizer rule to more
easily serialize the execution logic of the filter (using substrait) and
send the node over the wire.

Author: wkalt

rust/lance/src/dataset/scanner.rs

@@ -12,10 +12,10 @@ use arrow_array::{Array, Float32Array, Int64Array, RecordBatch};
 use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef, SortOptions};
 use arrow_select::concat::concat_batches;
 use async_recursion::async_recursion;
-use datafusion::common::SchemaExt;
+use datafusion::common::{DFSchema, SchemaExt};
 use datafusion::functions_aggregate;
 use datafusion::functions_aggregate::count::count_udaf;
-use datafusion::logical_expr::Expr;
+use datafusion::logical_expr::{col, lit, Expr};
 use datafusion::physical_expr::PhysicalSortExpr;
 use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec;
 use datafusion::physical_plan::expressions;
@@ -25,16 +25,15 @@ use datafusion::physical_plan::{
     aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy},
     display::DisplayableExecutionPlan,
     expressions::Literal,
-    filter::FilterExec,
     limit::GlobalLimitExec,
     repartition::RepartitionExec,
     union::UnionExec,
     ExecutionPlan, SendableRecordBatchStream,
 };
 use datafusion::scalar::ScalarValue;
-use datafusion_expr::Operator;
+use datafusion_expr::execution_props::ExecutionProps;
 use datafusion_physical_expr::{aggregate::AggregateExprBuilder, expressions::Column};
-use datafusion_physical_expr::{LexOrdering, Partitioning, PhysicalExpr};
+use datafusion_physical_expr::{create_physical_expr, LexOrdering, Partitioning, PhysicalExpr};
 use datafusion_physical_plan::{empty::EmptyExec, joins::HashJoinExec};
 use futures::future::BoxFuture;
 use futures::stream::{Stream, StreamExt};
@@ -69,7 +68,7 @@ use crate::index::DatasetIndexInternalExt;
 use crate::io::exec::fts::{BoostQueryExec, FlatMatchQueryExec, MatchQueryExec, PhraseQueryExec};
 use crate::io::exec::knn::MultivectorScoringExec;
 use crate::io::exec::scalar_index::{MaterializeIndexExec, ScalarIndexExec};
-use crate::io::exec::{get_physical_optimizer, LanceScanConfig};
+use crate::io::exec::{get_physical_optimizer, LanceFilterExec, LanceScanConfig};
 use crate::io::exec::{
     knn::new_knn_exec, project, AddRowAddrExec, FilterPlan, KNNVectorDistanceExec,
     LancePushdownScanExec, LanceScanExec, Planner, PreFilterSource, ScanConfig, TakeExec,
@@ -1591,10 +1590,7 @@ impl Scanner {
         if let Some(refine_expr) = filter_plan.refine_expr {
             // We create a new planner specific to the node's schema, since
             // physical expressions reference column by index rather than by name.
-            let planner = Planner::new(plan.schema());
-            let physical_refine_expr = planner.create_physical_expr(&refine_expr)?;
-
-            plan = Arc::new(FilterExec::try_new(physical_refine_expr, plan)?);
+            plan = Arc::new(LanceFilterExec::try_new(refine_expr, plan)?);
         }
 
         // Stage 3: sort
@@ -2066,9 +2062,7 @@ impl Scanner {
 
             if let Some(expr) = filter_plan.full_expr.as_ref() {
                 // If there is a prefilter we need to manually apply it to the new data
-                let planner = Planner::new(scan_node.schema());
-                let physical_refine_expr = planner.create_physical_expr(expr)?;
-                scan_node = Arc::new(FilterExec::try_new(physical_refine_expr, scan_node)?);
+                scan_node = Arc::new(LanceFilterExec::try_new(expr.clone(), scan_node)?);
             }
 
             let flat_match_plan = Arc::new(FlatMatchQueryExec::new(
@@ -2186,10 +2180,7 @@ impl Scanner {
                 )
             };
             if let Some(refine_expr) = &filter_plan.refine_expr {
-                let planner = Planner::new(plan.schema());
-                let physical_refine_expr = planner.create_physical_expr(refine_expr)?;
-
-                plan = Arc::new(FilterExec::try_new(physical_refine_expr, plan)?);
+                plan = Arc::new(LanceFilterExec::try_new(refine_expr.clone(), plan)?);
             }
             Ok(self.flat_knn(plan, q)?)
         }
@@ -2254,9 +2245,7 @@ impl Scanner {
 
             if let Some(expr) = filter_plan.full_expr.as_ref() {
                 // If there is a prefilter we need to manually apply it to the new data
-                let planner = Planner::new(scan_node.schema());
-                let physical_refine_expr = planner.create_physical_expr(expr)?;
-                scan_node = Arc::new(FilterExec::try_new(physical_refine_expr, scan_node)?);
+                scan_node = Arc::new(LanceFilterExec::try_new(expr.clone(), scan_node)?);
             }
             // first we do flat search on just the new data
             let topk_appended = self.flat_knn(scan_node, &q)?;
@@ -2429,9 +2418,7 @@ impl Scanner {
         if let Some(post_take_filter) = post_take_filter {
             let planner = Planner::new(plan.schema());
             let optimized_filter = planner.optimize_expr(post_take_filter)?;
-            let physical_refine_expr = planner.create_physical_expr(&optimized_filter)?;
-
-            plan = Arc::new(FilterExec::try_new(physical_refine_expr, plan)?);
+            plan = Arc::new(LanceFilterExec::try_new(optimized_filter, plan)?);
         }
 
         if self.with_row_address {
@@ -2459,7 +2446,6 @@ impl Scanner {
             let scan_arrow_schema = Arc::new(scan_schema.as_ref().into());
             let planner = Planner::new(scan_arrow_schema);
             let optimized_filter = planner.optimize_expr(filter.clone())?;
-            let physical_refine_expr = planner.create_physical_expr(&optimized_filter)?;
 
             let new_data_scan = self.scan_fragments(
                 true,
@@ -2471,7 +2457,7 @@ impl Scanner {
                 None,
                 false,
             );
-            let filtered = Arc::new(FilterExec::try_new(physical_refine_expr, new_data_scan)?);
+            let filtered = Arc::new(LanceFilterExec::try_new(optimized_filter, new_data_scan)?);
             Some(Arc::new(project(filtered, plan.schema().as_ref())?))
         } else {
             None
@@ -2599,45 +2585,43 @@ impl Scanner {
             q.metric_type,
         )?);
 
-        // filter out elements out of distance range
-        let lower_bound_expr = q
+        let lower: Option<(Expr, Arc<dyn PhysicalExpr>)> = q
             .lower_bound
-            .map(|v| {
-                let lower_bound = expressions::lit(v);
-                expressions::binary(
-                    expressions::col(DIST_COL, flat_dist.schema().as_ref())?,
-                    Operator::GtEq,
-                    lower_bound,
-                    flat_dist.schema().as_ref(),
-                )
+            .map(|v| -> Result<(Expr, Arc<dyn PhysicalExpr>)> {
+                let logical = col(DIST_COL).gt_eq(lit(v));
+                let schema = flat_dist.schema();
+                let df_schema = DFSchema::try_from(schema)?;
+                let physical = create_physical_expr(&logical, &df_schema, &ExecutionProps::new())?;
+                Ok::<(Expr, Arc<dyn PhysicalExpr>), _>((logical, physical))
             })
             .transpose()?;
-        let upper_bound_expr = q
+
+        let upper = q
             .upper_bound
-            .map(|v| {
-                let upper_bound = expressions::lit(v);
-                expressions::binary(
-                    expressions::col(DIST_COL, flat_dist.schema().as_ref())?,
-                    Operator::Lt,
-                    upper_bound,
-                    flat_dist.schema().as_ref(),
-                )
+            .map(|v| -> Result<(Expr, Arc<dyn PhysicalExpr>)> {
+                let logical = col(DIST_COL).lt(lit(v));
+                let schema = flat_dist.schema();
+                let df_schema = DFSchema::try_from(schema)?;
+                let physical = create_physical_expr(&logical, &df_schema, &ExecutionProps::new())?;
+                Ok::<(Expr, Arc<dyn PhysicalExpr>), _>((logical, physical))
             })
             .transpose()?;
-        let filter_expr = match (lower_bound_expr, upper_bound_expr) {
-            (Some(lower), Some(upper)) => Some(expressions::binary(
-                lower,
-                Operator::And,
-                upper,
-                flat_dist.schema().as_ref(),
-            )?),
-            (Some(lower), None) => Some(lower),
-            (None, Some(upper)) => Some(upper),
+
+        let filter_expr = match (lower, upper) {
+            (Some((llog, _)), Some((ulog, _))) => {
+                let logical = llog.and(ulog);
+                let schema = flat_dist.schema();
+                let df_schema = DFSchema::try_from(schema)?;
+                let physical = create_physical_expr(&logical, &df_schema, &ExecutionProps::new())?;
+                Some((logical, physical))
+            }
+            (Some((llog, lphys)), None) => Some((llog, lphys)),
+            (None, Some((ulog, uphys))) => Some((ulog, uphys)),
             (None, None) => None,
         };
 
         let knn_plan: Arc<dyn ExecutionPlan> = if let Some(filter_expr) = filter_expr {
-            Arc::new(FilterExec::try_new(filter_expr, flat_dist)?)
+            Arc::new(LanceFilterExec::try_new(filter_expr.0, flat_dist)?)
         } else {
             flat_dist
         };
@@ -2655,12 +2639,10 @@ impl Scanner {
         )
         .with_fetch(Some(q.k));
 
-        let not_nulls = FilterExec::try_new(
-            expressions::is_not_null(expressions::col(DIST_COL, sort.schema().as_ref())?)?,
-            Arc::new(sort),
-        )?;
+        let logical_not_null = col(DIST_COL).is_not_null();
+        let not_nulls = Arc::new(LanceFilterExec::try_new(logical_not_null, Arc::new(sort))?);
 
-        Ok(Arc::new(not_nulls))
+        Ok(not_nulls)
     }
 
     fn get_fragments_as_bitmap(&self) -> RoaringBitmap {
@@ -2826,10 +2808,8 @@ impl Scanner {
                 let columns_in_filter = Planner::column_names_in_expr(refine_expr);
                 let filter_schema = Arc::new(self.dataset.schema().project(&columns_in_filter)?);
                 let filter_input = self.scan(true, false, true, None, filter_schema);
-                let planner = Planner::new(filter_input.schema());
-                let physical_refine_expr = planner.create_physical_expr(refine_expr)?;
                 let filtered_row_ids =
-                    Arc::new(FilterExec::try_new(physical_refine_expr, filter_input)?);
+                    Arc::new(LanceFilterExec::try_new(refine_expr.clone(), filter_input)?);
                 PreFilterSource::FilteredRowIds(filtered_row_ids)
             }
             // No prefilter

rust/lance/src/io/exec.rs

@@ -5,6 +5,7 @@
 //!
 //! WARNING: Internal API with no stability guarantees.
 
+mod filter;
 pub mod filtered_read;
 pub mod fts;
 pub(crate) mod knn;
@@ -19,6 +20,7 @@ mod take;
 pub mod testing;
 pub mod utils;
 
+pub use filter::LanceFilterExec;
 pub use knn::{ANNIvfPartitionExec, ANNIvfSubIndexExec, KNNVectorDistanceExec};
 pub use lance_datafusion::planner::Planner;
 pub use lance_index::scalar::expression::FilterPlan;

rust/lance/src/io/exec/filter.rs

@@ -0,0 +1,100 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright The Lance Authors
+
+use std::sync::Arc;
+
+use datafusion::{execution::TaskContext, logical_expr::Expr};
+use datafusion_physical_plan::{
+    filter::FilterExec, metrics::MetricsSet, DisplayAs, DisplayFormatType, ExecutionPlan,
+    PlanProperties, SendableRecordBatchStream, Statistics,
+};
+use lance_core::{error::DataFusionResult, Result};
+use lance_datafusion::planner::Planner;
+
+#[derive(Debug)]
+// LanceFilterExec is a wrapper around FilterExec that includes the original
+// expression for the filter node. In comparison to a FilterExec, this makes it
+// possible for an optimization rule to serialize the filter to substrait and
+// send it to a remote worker.
+pub struct LanceFilterExec {
+    expr: Expr,
+    pub filter: Arc<FilterExec>,
+}
+
+impl DisplayAs for LanceFilterExec {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        self.filter.fmt_as(t, f)
+    }
+}
+
+impl LanceFilterExec {
+    pub fn try_new(expr: Expr, input: Arc<dyn ExecutionPlan>) -> Result<Self> {
+        let planner = Planner::new(input.schema());
+        let predicate = planner.create_physical_expr(&expr)?;
+        let filter_exec = FilterExec::try_new(predicate.clone(), input)?;
+        Ok(Self {
+            expr,
+            filter: Arc::new(filter_exec),
+        })
+    }
+
+    pub fn expr(&self) -> &Expr {
+        &self.expr
+    }
+}
+
+impl ExecutionPlan for LanceFilterExec {
+    fn name(&self) -> &str {
+        "LanceFilterExec"
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+
+    fn properties(&self) -> &PlanProperties {
+        self.filter.properties()
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+        self.filter.children()
+    }
+
+    fn maintains_input_order(&self) -> Vec<bool> {
+        self.filter.maintains_input_order()
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
+    ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
+        self.filter.clone().with_new_children(children)
+    }
+
+    fn execute(
+        &self,
+        partition: usize,
+        context: Arc<TaskContext>,
+    ) -> DataFusionResult<SendableRecordBatchStream> {
+        self.filter.execute(partition, context)
+    }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        self.filter.metrics()
+    }
+
+    fn statistics(&self) -> DataFusionResult<Statistics> {
+        self.filter.statistics()
+    }
+
+    fn cardinality_effect(&self) -> datafusion_physical_plan::execution_plan::CardinalityEffect {
+        self.filter.cardinality_effect()
+    }
+
+    fn try_swapping_with_projection(
+        &self,
+        projection: &datafusion_physical_plan::projection::ProjectionExec,
+    ) -> datafusion::error::Result<Option<Arc<dyn ExecutionPlan>>> {
+        self.filter.try_swapping_with_projection(projection)
+    }
+}

rust/lance/src/io/exec/knn.rs

@@ -111,8 +111,8 @@ pub struct KNNVectorDistanceExec {
 
     /// The vector query to execute.
     pub query: ArrayRef,
-    column: String,
-    distance_type: DistanceType,
+    pub column: String,
+    pub distance_type: DistanceType,
 
     output_schema: SchemaRef,
     properties: PlanProperties,