fix(datafusion): use caller session for filter planning

python/python/tests/test_table_provider.py

@@ -91,3 +91,31 @@ def make_ctx():
     result = normalize(ctx.table("ffi_lance_table").limit(1, offset=1).collect())
     assert len(result) == 1
     assert result["col1"][0].as_py() == 1
+
+
+def test_custom_udf_filter(tmp_path):
+    pytest.importorskip("datafusion")
+    from datafusion import SessionContext, udf
+
+    def is_even(values: pa.Array) -> pa.Array:
+        return pa.array([value.as_py() % 2 == 0 for value in values], type=pa.bool_())
+
+    is_even_udf = udf(
+        is_even,
+        input_fields=[pa.int64()],
+        return_field=pa.bool_(),
+        volatility="stable",
+        name="is_even",
+    )
+
+    dataset = lance.write_dataset(pa.table({"i": [1, 2, 3, 4]}), str(tmp_path))
+    provider = FFILanceTableProvider(dataset, with_row_id=True, with_row_addr=True)
+
+    ctx = SessionContext()
+    ctx.register_table("numbers", provider)
+    ctx.register_udf(is_even_udf)
+
+    result = normalize(
+        ctx.sql("SELECT i FROM numbers WHERE i = 2 AND is_even(i)").collect()
+    )
+    assert result["i"].to_pylist() == [2]

rust/lance-datafusion/src/planner.rs

@@ -17,6 +17,7 @@ use arrow_buffer::OffsetBuffer;
 use arrow_cast::cast_with_options;
 use arrow_schema::{DataType as ArrowDataType, Field, SchemaRef, TimeUnit};
 use arrow_select::concat::concat;
+use datafusion::catalog::Session;
 use datafusion::common::DFSchema;
 use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor};
 use datafusion::config::ConfigOptions;
@@ -1035,6 +1036,16 @@ impl Planner {
         )?)
     }
 
+    /// Create a [`PhysicalExpr`] using the caller's DataFusion session.
+    pub fn create_physical_expr_with_session(
+        &self,
+        expr: &Expr,
+        session: &dyn Session,
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        let df_schema = DFSchema::try_from(self.schema.as_ref().clone())?;
+        Ok(session.create_physical_expr(expr.clone(), &df_schema)?)
+    }
+
     /// Collect the columns in the expression.
     ///
     /// The columns are returned in sorted order.

rust/lance/src/datafusion/dataframe.rs

@@ -112,7 +112,7 @@ impl TableProvider for LanceTableProvider {
 
     async fn scan(
         &self,
-        _state: &dyn Session,
+        state: &dyn Session,
         projection: Option<&Vec<usize>>,
         filters: &[Expr],
         limit: Option<usize>,
@@ -161,7 +161,9 @@ impl TableProvider for LanceTableProvider {
         scan.limit(limit.map(|l| l as i64), None)?;
         scan.scan_in_order(self.ordered);
 
-        scan.create_plan().await.map_err(DataFusionError::from)
+        scan.create_plan_with_session(state)
+            .await
+            .map_err(DataFusionError::from)
     }
 
     // Since we are using datafusion itself to apply the filters it should

rust/lance/src/dataset/scanner.rs

@@ -17,6 +17,7 @@ use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaR
 use arrow_select::concat::concat_batches;
 use async_recursion::async_recursion;
 use chrono::Utc;
+use datafusion::catalog::Session;
 use datafusion::common::{DFSchema, JoinType, NullEquality, exec_datafusion_err};
 use datafusion::functions_aggregate;
 use datafusion::logical_expr::{Expr, ScalarUDF, col, lit};
@@ -522,6 +523,7 @@ impl FilterPlan {
         &self,
         input: Arc<dyn ExecutionPlan>,
         scanner: &Scanner,
+        session: Option<&dyn Session>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         let mut plan = input;
 
@@ -538,9 +540,12 @@ impl FilterPlan {
         }
 
         if let Some(refine_expr) = &self.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.
-            plan = Arc::new(LanceFilterExec::try_new(refine_expr.clone(), plan)?);
+            plan = Arc::new(match session {
+                Some(session) => {
+                    LanceFilterExec::try_new_with_session(refine_expr.clone(), plan, session)?
+                }
+                None => LanceFilterExec::try_new(refine_expr.clone(), plan)?,
+            });
         }
 
         Ok(plan)
@@ -2788,8 +2793,22 @@ impl Scanner {
     /// 3. Sort
     /// 4. Limit / Offset
     /// 5. Take remaining columns / Projection
+    pub fn create_plan(&self) -> BoxFuture<'_, Result<Arc<dyn ExecutionPlan>>> {
+        Box::pin(self.create_plan_impl(None))
+    }
+
+    pub(crate) fn create_plan_with_session<'a>(
+        &'a self,
+        session: &'a dyn Session,
+    ) -> BoxFuture<'a, Result<Arc<dyn ExecutionPlan>>> {
+        Box::pin(self.create_plan_impl(Some(session)))
+    }
+
     #[instrument(level = "debug", skip_all)]
-    pub async fn create_plan(&self) -> Result<Arc<dyn ExecutionPlan>> {
+    async fn create_plan_impl(
+        &self,
+        session: Option<&dyn Session>,
+    ) -> Result<Arc<dyn ExecutionPlan>> {
         log::trace!("creating scanner plan");
         self.validate_options()?;
 
@@ -2852,7 +2871,7 @@ impl Scanner {
                     self.take_source(take_op).await?
                 } else {
                     let planned_read = self
-                        .filtered_read_source(&mut filter_plan.expr_filter_plan)
+                        .filtered_read_source(&mut filter_plan.expr_filter_plan, session)
                         .await?;
                     if planned_read.limit_pushed_down {
                         use_limit_node = false;
@@ -2897,7 +2916,7 @@ impl Scanner {
         plan = self.take(plan, pre_filter_projection)?;
 
         // Filter
-        plan = filter_plan.refine_filter(plan, self).await?;
+        plan = filter_plan.refine_filter(plan, self, session).await?;
 
         // Aggregate (if set, applies aggregate and returns early)
         if let Some(agg) = &self.aggregate {
@@ -3124,6 +3143,7 @@ impl Scanner {
         make_deletions_null: bool,
         fragments: Option<Arc<Vec<Fragment>>>,
         scan_range: Option<Range<u64>>,
+        session: Option<&dyn Session>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         // Kept for the overlay stale-Take path below, which re-evaluates blocked stale rows.
         let user_projection = projection.clone();
@@ -3168,6 +3188,10 @@ impl Scanner {
             read_options = read_options.with_only_indexed_fragments();
         }
 
+        if let Some(session) = session {
+            read_options = read_options.with_physical_filters(session)?;
+        }
+
         // Mask data overlay files: a row with an overlay committed after an index it relies on
         // touched an indexed field can no longer be trusted to that index. Block just those rows
         // from the index result (their fragments stay indexed, so non-stale rows keep the index)
@@ -3217,7 +3241,12 @@ impl Scanner {
             .await?;
         let planner = Planner::new(stale_node.schema());
         let optimized_filter = planner.optimize_expr(filter.clone())?;
-        let filtered = Arc::new(LanceFilterExec::try_new(optimized_filter, stale_node)?);
+        let filtered = Arc::new(match session {
+            Some(session) => {
+                LanceFilterExec::try_new_with_session(optimized_filter, stale_node, session)?
+            }
+            None => LanceFilterExec::try_new(optimized_filter, stale_node)?,
+        });
         let stale_path: Arc<dyn ExecutionPlan> =
             Arc::new(project(filtered, plan.schema().as_ref())?);
 
@@ -3231,6 +3260,7 @@ impl Scanner {
     // Helper function for filtered read
     //
     // Delegates to legacy or new filtered read based on dataset storage version
+    #[allow(clippy::too_many_arguments)]
     async fn filtered_read(
         &self,
         filter_plan: &ExprFilterPlan,
@@ -3239,6 +3269,7 @@ impl Scanner {
         fragments: Option<Arc<Vec<Fragment>>>,
         scan_range: Option<Range<u64>>,
         is_prefilter: bool,
+        session: Option<&dyn Session>,
     ) -> Result<PlannedFilteredScan> {
         // Use legacy path if dataset uses legacy storage format
         if self.dataset.is_legacy_storage() {
@@ -3260,6 +3291,7 @@ impl Scanner {
                     make_deletions_null,
                     fragments,
                     scan_range,
+                    session,
                 )
                 .await?;
             Ok(PlannedFilteredScan {
@@ -3322,6 +3354,7 @@ impl Scanner {
     async fn filtered_read_source(
         &self,
         filter_plan: &mut ExprFilterPlan,
+        session: Option<&dyn Session>,
     ) -> Result<PlannedFilteredScan> {
         log::trace!("source is a filtered read");
 
@@ -3375,6 +3408,7 @@ impl Scanner {
             self.fragments.clone().map(Arc::new),
             scan_range,
             /*is_prefilter= */ false,
+            session,
         )
         .await
     }
@@ -4606,6 +4640,7 @@ impl Scanner {
                     Some(Arc::new(fragments)),
                     None,
                     /*is_prefilter=*/ true,
+                    None,
                 )
                 .await?;
             if let Some(refine_expr) = filter_plan.refine_expr.as_ref() {
@@ -4896,6 +4931,7 @@ impl Scanner {
                     self.fragments.clone().map(Arc::new),
                     None,
                     /*is_prefilter= */ true,
+                    None,
                 )
                 .await?;
 
@@ -6191,6 +6227,7 @@ impl Scanner {
                 Some(fragments),
                 None,
                 /*is_prefilter= */ true,
+                None,
             )
             .await?;
         Ok(PreFilterSource::FilteredRowIds(plan))

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

@@ -3,7 +3,7 @@
 
 use std::sync::Arc;
 
-use datafusion::{execution::TaskContext, logical_expr::Expr};
+use datafusion::{catalog::Session, execution::TaskContext, logical_expr::Expr};
 use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SendableRecordBatchStream,
     Statistics, filter::FilterExec, metrics::MetricsSet,
@@ -31,6 +31,24 @@ 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)?;
+        Self::try_new_with_predicate(expr, predicate, input)
+    }
+
+    pub fn try_new_with_session(
+        expr: Expr,
+        input: Arc<dyn ExecutionPlan>,
+        session: &dyn Session,
+    ) -> Result<Self> {
+        let planner = Planner::new(input.schema());
+        let predicate = planner.create_physical_expr_with_session(&expr, session)?;
+        Self::try_new_with_predicate(expr, predicate, input)
+    }
+
+    fn try_new_with_predicate(
+        expr: Expr,
+        predicate: Arc<dyn datafusion_physical_plan::PhysicalExpr>,
+        input: Arc<dyn ExecutionPlan>,
+    ) -> Result<Self> {
         let filter_exec = FilterExec::try_new(predicate.clone(), input)?;
         Ok(Self {
             expr,