diff --git a/crates/datafusion/src/read/mod.rs b/crates/datafusion/src/read/mod.rs index 26e2fb4..4c0157d 100644 --- a/crates/datafusion/src/read/mod.rs +++ b/crates/datafusion/src/read/mod.rs @@ -19,4 +19,4 @@ pub(crate) mod predicate; pub(crate) mod scan; pub use predicate::convert_filters_to_predicate; -pub use scan::IcebergDataSource; +pub use scan::{IcebergDataSource, IcebergDataSourceBuilder}; diff --git a/crates/datafusion/src/read/scan.rs b/crates/datafusion/src/read/scan.rs index 4432f0b..afcf3ee 100644 --- a/crates/datafusion/src/read/scan.rs +++ b/crates/datafusion/src/read/scan.rs @@ -49,59 +49,24 @@ pub struct IcebergDataSource { schema: ArrowSchemaRef, /// Projection column names in the output schema. projection: Vec, - /// Filters to apply to the table scan - predicates: Option, + /// Filter to apply to the table scan + predicate: Option, /// Optional limit on the number of rows to return fetch: Option, } impl IcebergDataSource { - /// Creates a new [`IcebergDataSource`] object. - pub(crate) fn new( - table: Table, - snapshot_id: Option, - schema: ArrowSchemaRef, - projection: Option<&Vec>, - filters: &[Expr], - limit: Option, - ) -> Result { - let output_schema = match projection { - None => schema, - Some(projection) => Arc::new(schema.project(projection)?), - }; - Ok(Self::new_with_predicate( + /// Starts building a source that scans `table` and returns the columns of + /// `schema`, all of them unless [`IcebergDataSourceBuilder::with_projection`] + /// selects some. Columns are read from Iceberg by name. + pub fn builder(table: Table, schema: ArrowSchemaRef) -> IcebergDataSourceBuilder { + IcebergDataSourceBuilder { table, - snapshot_id, - output_schema, - convert_filters_to_predicate(filters), - limit, - )) - } - - /// Rebuilds a source from the parts needed to scan an Iceberg table. - /// - /// The supplied schema determines both the source's output schema and the - /// columns read from Iceberg. Callers can recreate a source from its - /// accessors after transporting those values to another process. - pub fn new_with_predicate( - table: Table, - snapshot_id: Option, - schema: ArrowSchemaRef, - predicates: Option, - fetch: Option, - ) -> Self { - let projection = schema - .fields() - .iter() - .map(|field| field.name().clone()) - .collect(); - Self { - table, - snapshot_id, schema, - projection, - predicates, - fetch, + projection: None, + snapshot_id: None, + predicate: None, + fetch: None, } } @@ -121,8 +86,8 @@ impl IcebergDataSource { &self.projection } - pub fn predicates(&self) -> Option<&Predicate> { - self.predicates.as_ref() + pub fn predicate(&self) -> Option<&Predicate> { + self.predicate.as_ref() } pub fn fetch(&self) -> Option { @@ -130,6 +95,82 @@ impl IcebergDataSource { } } +/// Builds an [`IcebergDataSource`]. Created by [`IcebergDataSource::builder`]. +/// +/// A table provider passes the table's schema and selects columns with +/// [`Self::with_projection`]. To rebuild a source, such as in another process, +/// pass its [`IcebergDataSource::schema`], which is already projected, and no +/// projection. +#[derive(Debug, Clone)] +pub struct IcebergDataSourceBuilder { + table: Table, + schema: ArrowSchemaRef, + projection: Option>, + snapshot_id: Option, + predicate: Option, + fetch: Option, +} + +impl IcebergDataSourceBuilder { + /// Returns only the columns of the schema at these indices, or all of them + /// if `None` (the default). Indices are into the schema passed to + /// [`IcebergDataSource::builder`]. + pub fn with_projection(mut self, projection: Option<&Vec>) -> Self { + self.projection = projection.cloned(); + self + } + + /// Scans this snapshot, or the table's current one if `None` (the default). + pub fn with_snapshot_id(mut self, snapshot_id: Option) -> Self { + self.snapshot_id = snapshot_id; + self + } + + /// Filters the scan with `predicate`, or not at all if `None` (the default). + pub fn with_predicate(mut self, predicate: Option) -> Self { + self.predicate = predicate; + self + } + + /// Filters the scan with the parts of `filters` that convert to an Iceberg + /// predicate. The rest are dropped, so the scan may return rows that + /// `filters` reject. + pub fn with_filters(self, filters: &[Expr]) -> Self { + self.with_predicate(convert_filters_to_predicate(filters)) + } + + /// Stops after `fetch` rows, or reads every row if `None` (the default). + pub fn with_fetch(mut self, fetch: Option) -> Self { + self.fetch = fetch; + self + } + + /// Builds the source, which reads the columns of its projected schema by name. + /// + /// # Errors + /// + /// Returns an error if the projection is out of bounds for the schema. + pub fn build(self) -> Result { + let schema = match &self.projection { + None => self.schema, + Some(projection) => Arc::new(self.schema.project(projection)?), + }; + let projection = schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(); + Ok(IcebergDataSource { + table: self.table, + snapshot_id: self.snapshot_id, + schema, + projection, + predicate: self.predicate, + fetch: self.fetch, + }) + } +} + impl DataSource for IcebergDataSource { fn apply_expressions( &self, @@ -147,7 +188,7 @@ impl DataSource for IcebergDataSource { self.table.clone(), self.snapshot_id, Some(self.projection.clone()), - self.predicates.clone(), + self.predicate.clone(), ); let stream = futures::stream::once(fut).try_flatten(); @@ -184,7 +225,7 @@ impl DataSource for IcebergDataSource { ) -> std::fmt::Result { write!(f, "format=iceberg",)?; write!(f, ", projection=[{}]", self.projection.join(", "))?; - if let Some(predicate) = &self.predicates { + if let Some(predicate) = &self.predicate { write!(f, ", predicate={predicate}")?; } if let Some(fetch) = self.fetch { @@ -232,7 +273,7 @@ async fn get_batch_stream( table: Table, snapshot_id: Option, column_names: Option>, - predicates: Option, + predicate: Option, ) -> Result> + Send>>> { let scan_builder = match snapshot_id { Some(snapshot_id) => table.scan().snapshot_id(snapshot_id), @@ -243,8 +284,8 @@ async fn get_batch_stream( Some(column_names) => scan_builder.select(column_names), None => scan_builder.select_all(), }; - if let Some(pred) = predicates { - scan_builder = scan_builder.with_filter(pred); + if let Some(predicate) = predicate { + scan_builder = scan_builder.with_filter(predicate); } let table_scan = scan_builder.build().map_err(to_datafusion_error)?; diff --git a/crates/datafusion/src/table_provider/catalog_table_provider.rs b/crates/datafusion/src/table_provider/catalog_table_provider.rs index 5145b42..6dbb79d 100644 --- a/crates/datafusion/src/table_provider/catalog_table_provider.rs +++ b/crates/datafusion/src/table_provider/catalog_table_provider.rs @@ -139,14 +139,12 @@ impl TableProvider for IcebergCatalogTableProvider { .map_err(to_datafusion_error)?; // Create scan with fresh metadata (always use current snapshot) - Ok(DataSourceExec::from_data_source(IcebergDataSource::new( - table, - None, // Always use current snapshot for catalog-backed provider - self.schema.clone(), - projection, - filters, - limit, - )?)) + let source = IcebergDataSource::builder(table, self.schema.clone()) + .with_projection(projection) + .with_filters(filters) + .with_fetch(limit) + .build()?; + Ok(DataSourceExec::from_data_source(source)) } fn supports_filters_pushdown( @@ -447,7 +445,7 @@ mod tests { } #[tokio::test] - async fn test_scan_rejects_out_of_range_projection() { + async fn test_scan_projection() { let (catalog, namespace, table_name, _temp_dir) = get_test_catalog_and_table().await; let table_ident = TableIdent::new(namespace, table_name); @@ -456,15 +454,22 @@ mod tests { Arc::new(schema_to_arrow_schema(table.metadata().current_schema()).unwrap()); let out_of_range = schema.fields().len(); - let err = IcebergDataSource::new( - table, - None, - schema, - Some(&vec![0, out_of_range]), - &[], - None, - ) - .unwrap_err(); + // Columns come back in the projection's order, not the table's. + let source = IcebergDataSource::builder(table.clone(), schema.clone()) + .with_projection(Some(&vec![1, 0])) + .build() + .unwrap(); + insta::assert_debug_snapshot!(source.projection(), @r#" + [ + "name", + "id", + ] + "#); + + let err = IcebergDataSource::builder(table, schema) + .with_projection(Some(&vec![0, out_of_range])) + .build() + .unwrap_err(); assert_eq!( err.to_string(), format!( diff --git a/crates/datafusion/src/table_provider/static_table_provider.rs b/crates/datafusion/src/table_provider/static_table_provider.rs index 2a435cc..5ebabc8 100644 --- a/crates/datafusion/src/table_provider/static_table_provider.rs +++ b/crates/datafusion/src/table_provider/static_table_provider.rs @@ -126,14 +126,13 @@ impl TableProvider for IcebergStaticTableProvider { limit: Option, ) -> Result> { // Use cached table (no refresh) - Ok(DataSourceExec::from_data_source(IcebergDataSource::new( - self.table.clone(), - self.snapshot_id, - self.schema.clone(), - projection, - filters, - limit, - )?)) + let source = IcebergDataSource::builder(self.table.clone(), self.schema.clone()) + .with_snapshot_id(self.snapshot_id) + .with_projection(projection) + .with_filters(filters) + .with_fetch(limit) + .build()?; + Ok(DataSourceExec::from_data_source(source)) } fn supports_filters_pushdown( diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 9013890..8d300a3 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -1019,13 +1019,12 @@ async fn run( /// Rebuilds `scan` from its accessors alone, as a codec would. fn rebuild_scan(scan: &IcebergDataSource) -> IcebergDataSource { - IcebergDataSource::new_with_predicate( - scan.table().clone(), - scan.snapshot_id(), - scan.schema(), - scan.predicates().cloned(), - scan.fetch(), - ) + IcebergDataSource::builder(scan.table().clone(), scan.schema()) + .with_snapshot_id(scan.snapshot_id()) + .with_predicate(scan.predicate().cloned()) + .with_fetch(scan.fetch()) + .build() + .unwrap() } async fn run_source( @@ -1141,7 +1140,7 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { .downcast_ref::() .expect("an Iceberg data source"); assert_eq!( - scan.predicates().map(ToString::to_string).as_deref(), + scan.predicate().map(ToString::to_string).as_deref(), Some("foo1 = 1") ); let rebuilt = rebuild_scan(scan); @@ -1183,13 +1182,9 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { // A scan reads the columns of its schema and no others, so one built over // part of the table returns only those columns, from the pinned snapshot. let foo2_only = Arc::new(pinned.schema().project(&[1])?); - let partial = IcebergDataSource::new_with_predicate( - pinned.table().clone(), - Some(snapshot_id), - foo2_only, - None, - None, - ); + let partial = IcebergDataSource::builder(pinned.table().clone(), foo2_only) + .with_snapshot_id(Some(snapshot_id)) + .build()?; expect![[r#" +--------+ | foo2 |