refactor(datafusion): add table scan builder

Author: toutane

Cargo.lock

@@ -73,6 +73,7 @@ derive_builder = "0.20"
 dirs = "6"
 enum-ordinalize = "4.3.0"
 env_logger = "0.11.8"
+either = "1"
 expect-test = "1"
 faststr = "0.2.31"
 flate2 = "1.1.5"

Cargo.toml

@@ -33,6 +33,7 @@ anyhow = { workspace = true }
 async-trait = { workspace = true }
 dashmap = { workspace = true }
 datafusion = { workspace = true }
+either = { workspace = true }
 futures = { workspace = true }
 iceberg = { workspace = true }
 parquet = { workspace = true }

crates/integrations/datafusion/Cargo.toml

@@ -28,6 +28,7 @@ use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType};
 use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
 use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning, PlanProperties};
 use datafusion::prelude::Expr;
+use either::Either;
 use futures::{Stream, TryStreamExt};
 use iceberg::expr::Predicate;
 use iceberg::scan::{FileScanTask, TableScan};
@@ -39,11 +40,9 @@ use crate::to_datafusion_error;
 
 /// Iceberg [`Table`] scan as a DataFusion [`ExecutionPlan`].
 ///
-/// Has three construction modes: [`new`][Self::new] for a lazy
-/// single-partition scan, [`new_with_tasks`][Self::new_with_tasks] for an
-/// eager scan over pre-planned [`FileScanTask`] buckets, and
-/// [`new_with_tasks_from_predicate`][Self::new_with_tasks_from_predicate] for
-/// rebuilding an eager scan from a [`Predicate`].
+/// Prefer [`IcebergTableScanBuilder`] when composing optional scan inputs. The
+/// constructor helpers are retained for direct lazy scans and eager scans over
+/// pre-planned [`FileScanTask`] buckets.
 ///
 /// Note: in eager mode the underlying `TableScan` is rebuilt on every
 /// `execute(partition)` call. The per-build cost is bounded (no I/O) and
@@ -63,8 +62,8 @@ pub struct IcebergTableScan {
     table_schema: ArrowSchemaRef,
     /// Projection indices as received by `scan`.
     projection_indices: Option<Vec<usize>>,
-    /// Filters to apply to the table scan.
-    predicates: Option<Predicate>,
+    /// Predicate to apply to the table scan.
+    predicate: Option<Predicate>,
     /// Pre-planned file scan tasks per partition (eager mode), or `None` (lazy mode).
     buckets: Option<Vec<Vec<FileScanTask>>>,
     /// `None` when partitioning is `UnknownPartitioning`.
@@ -73,7 +72,117 @@ pub struct IcebergTableScan {
     limit: Option<usize>,
 }
 
+/// Builder for [`IcebergTableScan`].
+#[derive(Debug)]
+pub struct IcebergTableScanBuilder {
+    table: Table,
+    snapshot_id: Option<i64>,
+    schema: ArrowSchemaRef,
+    projection: Option<Vec<usize>>,
+    filter_source: Option<Either<Vec<Expr>, Option<Predicate>>>,
+    limit: Option<usize>,
+    buckets: Option<Vec<Vec<FileScanTask>>>,
+    partitioning: Partitioning,
+    partition_keys_kind: Option<PartitionKeysKind>,
+}
+
+impl IcebergTableScanBuilder {
+    pub fn new(table: Table, schema: ArrowSchemaRef) -> Self {
+        Self {
+            table,
+            snapshot_id: None,
+            schema,
+            projection: None,
+            filter_source: None,
+            limit: None,
+            buckets: None,
+            partitioning: Partitioning::UnknownPartitioning(1),
+            partition_keys_kind: None,
+        }
+    }
+
+    pub fn with_snapshot_id(mut self, snapshot_id: Option<i64>) -> Self {
+        self.snapshot_id = snapshot_id;
+        self
+    }
+
+    pub fn with_projection(mut self, projection: Option<&[usize]>) -> Self {
+        self.projection = projection.map(<[usize]>::to_vec);
+        self
+    }
+
+    /// Sets the filter source for the scan.
+    ///
+    /// Use `Either::Left(filters)` for DataFusion filters, or
+    /// `Either::Right(predicate)` when rebuilding from an Iceberg predicate.
+    pub fn with_filter_source(mut self, filter_source: Either<&[Expr], Option<Predicate>>) -> Self {
+        self.filter_source = Some(match filter_source {
+            Either::Left(filters) => Either::Left(filters.to_vec()),
+            Either::Right(predicate) => Either::Right(predicate),
+        });
+        self
+    }
+
+    pub fn with_limit(mut self, limit: Option<usize>) -> Self {
+        self.limit = limit;
+        self
+    }
+
+    pub fn with_tasks(
+        mut self,
+        buckets: Vec<Vec<FileScanTask>>,
+        partitioning: Partitioning,
+    ) -> Self {
+        self.buckets = Some(buckets);
+        self.partitioning = partitioning;
+        self
+    }
+
+    pub(crate) fn with_partition_keys_kind(
+        mut self,
+        partition_keys_kind: Option<PartitionKeysKind>,
+    ) -> Self {
+        self.partition_keys_kind = partition_keys_kind;
+        self
+    }
+
+    pub fn build(self) -> IcebergTableScan {
+        let Self {
+            table,
+            snapshot_id,
+            schema,
+            projection,
+            filter_source,
+            limit,
+            buckets,
+            partitioning,
+            partition_keys_kind,
+        } = self;
+        let predicate = match filter_source {
+            Some(Either::Left(filters)) => convert_filters_to_predicate(&filters),
+            Some(Either::Right(predicate)) => predicate,
+            None => None,
+        };
+
+        IcebergTableScan::new_inner(
+            table,
+            snapshot_id,
+            schema,
+            projection.as_deref(),
+            predicate,
+            limit,
+            partitioning,
+            buckets,
+        )
+        .with_partition_keys_kind(partition_keys_kind)
+    }
+}
+
 impl IcebergTableScan {
+    pub fn builder(table: Table, schema: ArrowSchemaRef) -> IcebergTableScanBuilder {
+        IcebergTableScanBuilder::new(table, schema)
+    }
+
     /// Creates a lazy single-partition scan that plans and reads all tasks
     /// inside `execute(0)`. Used by
     /// [`IcebergStaticTableProvider`][crate::table::IcebergStaticTableProvider].
@@ -85,16 +194,12 @@ impl IcebergTableScan {
         filters: &[Expr],
         limit: Option<usize>,
     ) -> Self {
-        Self::new_inner(
-            table,
-            snapshot_id,
-            schema,
-            projection,
-            convert_filters_to_predicate(filters),
-            limit,
-            Partitioning::UnknownPartitioning(1),
-            None,
-        )
+        Self::builder(table, schema)
+            .with_snapshot_id(snapshot_id)
+            .with_projection(projection.map(Vec::as_slice))
+            .with_filter_source(Either::Left(filters))
+            .with_limit(limit)
+            .build()
     }
 
     /// Creates an eager multi-partition scan over pre-planned task buckets.
@@ -112,48 +217,21 @@ impl IcebergTableScan {
         buckets: Vec<Vec<FileScanTask>>,
         partitioning: Partitioning,
     ) -> Self {
-        Self::new_inner(
-            table,
-            snapshot_id,
-            schema,
-            projection,
-            convert_filters_to_predicate(filters),
-            limit,
-            partitioning,
-            Some(buckets),
-        )
-    }
-
-    /// Creates an eager multi-partition scan from a [`Predicate`].
-    #[allow(clippy::too_many_arguments)]
-    pub fn new_with_tasks_from_predicate(
-        table: Table,
-        snapshot_id: Option<i64>,
-        schema: ArrowSchemaRef,
-        projection: Option<&Vec<usize>>,
-        predicate: Option<Predicate>,
-        limit: Option<usize>,
-        buckets: Vec<Vec<FileScanTask>>,
-        partitioning: Partitioning,
-    ) -> Self {
-        Self::new_inner(
-            table,
-            snapshot_id,
-            schema,
-            projection,
-            predicate,
-            limit,
-            partitioning,
-            Some(buckets),
-        )
+        Self::builder(table, schema)
+            .with_snapshot_id(snapshot_id)
+            .with_projection(projection.map(Vec::as_slice))
+            .with_filter_source(Either::Left(filters))
+            .with_limit(limit)
+            .with_tasks(buckets, partitioning)
+            .build()
     }
 
     #[allow(clippy::too_many_arguments)]
     fn new_inner(
         table: Table,
         snapshot_id: Option<i64>,
         schema: ArrowSchemaRef,
-        projection: Option<&Vec<usize>>,
+        projection: Option<&[usize]>,
         predicate: Option<Predicate>,
         limit: Option<usize>,
         partitioning: Partitioning,
@@ -170,7 +248,7 @@ impl IcebergTableScan {
             Boundedness::Bounded,
         ));
         let table_schema = Arc::clone(&schema);
-        let projection_indices = projection.cloned();
+        let projection_indices = projection.map(<[usize]>::to_vec);
         let projection = get_column_names(schema, projection);
 
         Self {
@@ -180,7 +258,7 @@ impl IcebergTableScan {
             projection,
             table_schema,
             projection_indices,
-            predicates: predicate,
+            predicate,
             buckets,
             partition_keys_kind: None,
             limit,
@@ -215,8 +293,8 @@ impl IcebergTableScan {
         self.projection_indices.as_deref()
     }
 
-    pub fn predicates(&self) -> Option<&Predicate> {
-        self.predicates.as_ref()
+    pub fn predicate(&self) -> Option<&Predicate> {
+        self.predicate.as_ref()
     }
 
     /// Returns the pre-planned file task buckets, or an empty slice in lazy mode.
@@ -290,7 +368,7 @@ impl ExecutionPlan for IcebergTableScan {
             self.table.clone(),
             self.snapshot_id,
             self.projection.clone(),
-            self.predicates.clone(),
+            self.predicate.clone(),
             bucket,
         );
         let stream = Box::pin(futures::stream::once(fut).try_flatten())
@@ -319,7 +397,7 @@ impl DisplayAs for IcebergTableScan {
             .as_deref()
             .map_or(String::new(), |v| v.join(","));
         let predicate = self
-            .predicates
+            .predicate
             .as_ref()
             .map_or(String::new(), |p| p.to_string());
 
@@ -344,7 +422,7 @@ fn build_table_scan(
     table: Table,
     snapshot_id: Option<i64>,
     column_names: Option<Vec<String>>,
-    predicates: Option<Predicate>,
+    predicate: Option<Predicate>,
 ) -> DFResult<TableScan> {
     let scan_builder = match snapshot_id {
         Some(id) => table.scan().snapshot_id(id),
@@ -354,7 +432,7 @@ fn build_table_scan(
         Some(names) => scan_builder.select(names),
         None => scan_builder.select_all(),
     };
-    if let Some(pred) = predicates {
+    if let Some(pred) = predicate {
         scan_builder = scan_builder.with_filter(pred);
     }
     scan_builder.build().map_err(to_datafusion_error)
@@ -367,10 +445,10 @@ async fn build_record_batch_stream(
     table: Table,
     snapshot_id: Option<i64>,
     column_names: Option<Vec<String>>,
-    predicates: Option<Predicate>,
+    predicate: Option<Predicate>,
     bucket: Option<Vec<FileScanTask>>,
 ) -> DFResult<Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>>> {
-    let table_scan = build_table_scan(table, snapshot_id, column_names, predicates)?;
+    let table_scan = build_table_scan(table, snapshot_id, column_names, predicate)?;
     let stream: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>> = match bucket {
         Some(bucket) => {
             let task_stream = Box::pin(futures::stream::iter(
@@ -416,7 +494,7 @@ fn apply_limit(
 
 pub(super) fn get_column_names(
     schema: ArrowSchemaRef,
-    projection: Option<&Vec<usize>>,
+    projection: Option<&[usize]>,
 ) -> Option<Vec<String>> {
     projection.map(|v| {
         v.iter()
@@ -462,22 +540,19 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn test_predicate_constructor_exposes_rebuild_inputs() {
+    async fn test_builder_exposes_rebuild_inputs() {
         let schema = create_test_arrow_schema();
         let projection = vec![0usize, 2];
         let predicate = Reference::new("x").greater_than(Datum::long(5));
-        let scan = IcebergTableScan::new_with_tasks_from_predicate(
-            get_test_table_from_metadata_file().await,
-            None,
-            schema.clone(),
-            Some(&projection),
-            Some(predicate.clone()),
-            Some(100),
-            vec![vec![], vec![]],
-            Partitioning::UnknownPartitioning(2),
-        );
-
-        assert_eq!(scan.predicates(), Some(&predicate));
+        let scan =
+            IcebergTableScan::builder(get_test_table_from_metadata_file().await, schema.clone())
+                .with_projection(Some(&projection))
+                .with_filter_source(Either::Right(Some(predicate.clone())))
+                .with_limit(Some(100))
+                .with_tasks(vec![vec![], vec![]], Partitioning::UnknownPartitioning(2))
+                .build();
+
+        assert_eq!(scan.predicate(), Some(&predicate));
         assert_eq!(scan.table_schema().fields(), schema.fields());
         assert_eq!(scan.table_schema().fields().len(), 3);
         assert_eq!(scan.schema().fields().len(), 2);