NoahKusaba commented on code in PR #20:
URL: https://github.com/apache/datafusion-iceberg/pull/20#discussion_r4136478601
##########
crates/datafusion/src/physical_plan/scan.rs:
##########
@@ -64,23 +65,141 @@ impl IcebergTableScan {
projection: Option<&Vec<usize>>,
filters: &[Expr],
limit: Option<usize>,
- ) -> Self {
- let output_schema = match projection {
- None => schema.clone(),
- Some(projection) => Arc::new(schema.project(projection).unwrap()),
+ ) -> Result<Self> {
+ let (output_schema, projection) = match projection {
+ None => (schema, None),
+ Some(projection) => {
+ let output_schema = Arc::new(schema.project(projection)?);
+ let names = output_schema
+ .fields()
+ .iter()
+ .map(|field| field.name().clone())
+ .collect();
+ (output_schema, Some(names))
+ }
};
- let plan_properties = Self::compute_properties(output_schema.clone());
- let projection = get_column_names(schema.clone(), projection);
- let predicates = convert_filters_to_predicate(filters);
+ Self::new_with_predicate(
+ table,
+ snapshot_id,
+ output_schema,
+ projection,
+ convert_filters_to_predicate(filters),
+ limit,
+ )
+ }
+
+ /// Creates a scan of `table` from an already-converted Iceberg
+ /// [`Predicate`] rather than DataFusion filters, for rebuilding a scan
from
+ /// its parts, such as after sending them to another process. A predicate
+ /// cannot be converted back to the filters it came from.
+ ///
+ /// Each argument takes what the matching accessor returns (`predicate`
+ /// what [`Self::predicates`] does), so a scan is rebuilt from
+ /// [`ExecutionPlan::schema`] and [`Self::projection`]:
+ ///
+ /// - `snapshot_id`: the snapshot to read, or `None` for the table's
current
+ /// snapshot.
+ /// - `schema`: the Arrow schema the scan outputs. With `projection` set,
+ /// its fields must be the projected columns, in order. Without one, it
+ /// must be the full schema of the table as read, which is not checked.
+ /// - `projection`: the names of the columns to read from the table, or
+ /// `None` for all of them.
+ /// - `predicate`: pushed down to Iceberg to skip data files and rows. The
+ /// table providers report their filters as
+ ///
[`Inexact`](datafusion::logical_expr::TableProviderFilterPushDown::Inexact),
+ /// so DataFusion still applies them above the scan.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if `projection` does not name the fields of `schema`,
+ /// in order.
+ ///
+ /// # Example
+ ///
+ /// ```
+ /// use std::collections::HashMap;
+ ///
+ /// use datafusion::catalog::TableProvider;
+ /// use datafusion::physical_plan::ExecutionPlan;
+ /// use datafusion::prelude::{SessionContext, col, lit};
+ /// use datafusion_iceberg::IcebergStaticTableProvider;
+ /// use datafusion_iceberg::physical_plan::IcebergTableScan;
+ /// use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder};
+ /// use iceberg::spec::{NestedField, PrimitiveType, Schema, Type};
+ /// use iceberg::{Catalog, CatalogBuilder, NamespaceIdent, TableCreation};
+ ///
+ /// # tokio::runtime::Runtime::new()?.block_on(async {
+ /// # let warehouse = tempfile::tempdir()?;
+ /// # let props = HashMap::from([(
+ /// # MEMORY_CATALOG_WAREHOUSE.to_string(),
+ /// # warehouse.path().display().to_string(),
+ /// # )]);
+ /// # let catalog = MemoryCatalogBuilder::default().load("memory",
props).await?;
+ /// # let namespace = NamespaceIdent::new("ns".to_string());
+ /// # catalog.create_namespace(&namespace, HashMap::new()).await?;
+ /// # let schema = Schema::builder()
+ /// # .with_fields(vec![
+ /// # NestedField::required(1, "id",
Type::Primitive(PrimitiveType::Int)).into(),
+ /// # NestedField::optional(2, "name",
Type::Primitive(PrimitiveType::String))
+ /// # .into(),
+ /// # ])
+ /// # .build()?;
+ /// # let creation =
TableCreation::builder().name("t".to_string()).schema(schema).build();
+ /// # let table = catalog.create_table(&namespace, creation).await?;
+ /// let provider =
IcebergStaticTableProvider::try_new_from_table(table).await?;
+ /// let ctx = SessionContext::new();
+ /// let filters = [col("id").gt(lit(1))];
+ /// let plan = provider
+ /// .scan(&ctx.state(), Some(&vec![1]), &filters, None)
+ /// .await?;
+ /// let scan = plan.downcast_ref::<IcebergTableScan>().unwrap();
+ ///
+ /// // Rebuild an equivalent scan from the original's accessors alone.
+ /// let rebuilt = IcebergTableScan::new_with_predicate(
+ /// scan.table().clone(),
+ /// scan.snapshot_id(),
+ /// scan.schema(),
+ /// scan.projection().map(<[String]>::to_vec),
+ /// scan.predicates().cloned(),
+ /// scan.limit(),
+ /// )?;
+ /// assert_eq!(rebuilt.schema(), scan.schema());
+ /// assert_eq!(rebuilt.projection(), scan.projection());
+ /// assert_eq!(rebuilt.predicates(), scan.predicates());
+ /// # Ok::<(), Box<dyn std::error::Error>>(())
+ /// # })?;
+ /// # Ok::<(), Box<dyn std::error::Error>>(())
+ /// ```
+ pub fn new_with_predicate(
+ table: Table,
+ snapshot_id: Option<i64>,
+ schema: ArrowSchemaRef,
+ projection: Option<Vec<String>>,
+ predicate: Option<Predicate>,
+ limit: Option<usize>,
+ ) -> Result<Self> {
+ // The scan reads the named columns and reports `schema`, so the two
+ // must agree, or its batches would not match its schema.
+ if let Some(projection) = &projection {
+ let fields: Vec<&String> =
+ schema.fields().iter().map(|field| field.name()).collect();
+ if !projection.iter().eq(fields.iter().copied()) {
+ return plan_err!(
+ "IcebergTableScan projection {projection:?} does not match
the \
+ fields of its schema {fields:?}"
+ );
+ }
Review Comment:
Agreed, and I checked `collect_scan_field_ids` at that rev and see the same:
`None` and a `select` of every top-level name resolve to the same field ids. In
14aaa24, `new_with_predicate` no longer takes a projection: it always selects
`schema`'s field names, so the batches match the schema by construction, and it
can't fail. The mismatch check and its test are gone. Nothing relied on
`projection()` being `None`. It's only read by the display and the tests, and
the EXPLAIN tests already show named projections. The display gets clearer too,
since a full scan used to print `projection:[]`, the same as an empty
projection. `projection()` keeps its signature and returns the schema's names.
Selecting by name also fixes the same mismatch through the catalog-backed
provider, which pairs its cached schema with a freshly loaded table. After a
column is added and written to, a scan with no projection used to return that
column too. `test_scan_after_schema_evolution_reads_provider_columns` covers
it, and fails on the previous commit.
##########
crates/datafusion/tests/integration_datafusion_test.rs:
##########
@@ -977,3 +987,150 @@ async fn test_insert_into_partitioned() -> Result<(),
Box<dyn Error>> {
Ok(())
}
+
+/// Executes the single partition of `plan` and renders its rows as a table.
+async fn run(
+ plan: &dyn ExecutionPlan,
+ ctx: &SessionContext,
+) -> Result<String, Box<dyn Error>> {
+ let stream = plan.execute(0, ctx.task_ctx())?;
+ let batches = datafusion::physical_plan::common::collect(stream).await?;
+ Ok(pretty_format_batches(&batches)?.to_string())
+}
+
+/// Returns the first node of type `T` in `plan`, depth first.
+fn find_node<T: ExecutionPlan + 'static>(plan: &Arc<dyn ExecutionPlan>) ->
Option<&T> {
+ plan.downcast_ref::<T>()
+ .or_else(|| plan.children().into_iter().find_map(find_node::<T>))
+}
+
+/// The plan nodes and providers can be named and inspected from outside this
+/// crate, and rebuilt from their parts, as a codec that serializes them does.
+#[tokio::test]
+async fn test_plan_nodes_are_inspectable() -> Result<(), Box<dyn Error>> {
+ let iceberg_catalog = get_iceberg_catalog().await;
+ let namespace = NamespaceIdent::new("test_plan_nodes".to_string());
+ set_test_namespace(&iceberg_catalog, &namespace).await?;
+ let creation = get_table_creation(temp_path(), "my_table", None)?;
+ iceberg_catalog.create_table(&namespace, creation).await?;
+ let ident = TableIdent::new(namespace.clone(), "my_table".to_string());
+ let client: Arc<dyn Catalog> = Arc::new(iceberg_catalog);
+
+ let ctx = SessionContext::new();
+ let catalog = IcebergCatalogProvider::try_new(client.clone()).await?;
+ ctx.register_catalog("catalog", Arc::new(catalog));
+ let provider = ctx
+ .table_provider("catalog.test_plan_nodes.my_table")
+ .await?;
+ let provider = provider
+ .downcast_ref::<IcebergTableProvider>()
+ .expect("a catalog-backed provider");
+ assert_eq!(provider.table_ident(), &ident);
+ assert!(Arc::ptr_eq(provider.catalog(), &client));
+ let rebuilt = IcebergTableProvider::try_new(
+ provider.catalog().clone(),
+ provider.table_ident().namespace().clone(),
+ provider.table_ident().name(),
+ )
+ .await?;
+ assert_eq!(rebuilt.table_ident(), &ident);
+ assert_eq!(rebuilt.schema(), provider.schema());
+
+ // Write path: a commit above a write, both holding the table, and the
+ // commit going through the provider's catalog. The plan that runs is
+ // rebuilt from their accessors and children alone. The optimizer drops
+ // the coalesce above a single-partition write, so the rebuilt commit
+ // always gets one, as a codec would.
+ let insert = ctx
+ .sql("INSERT INTO catalog.test_plan_nodes.my_table VALUES (1, 'alan'),
(2, 'turing')")
+ .await?
+ .create_physical_plan()
+ .await?;
+ let commit = insert
+ .downcast_ref::<IcebergCommitExec>()
+ .expect("the insert plan is rooted at a commit");
+ assert_eq!(commit.table().identifier(), &ident);
+ assert!(Arc::ptr_eq(commit.catalog(), &client));
+ let write = find_node::<IcebergWriteExec>(&insert).expect("a write below
the commit");
+ assert_eq!(write.table().identifier(), &ident);
+ let rebuilt_write: Arc<dyn ExecutionPlan> = Arc::new(IcebergWriteExec::new(
+ write.table().clone(),
+ write.children()[0].clone(),
+ ));
+ let rebuilt_commit = IcebergCommitExec::new(
+ commit.table().clone(),
+ commit.catalog().clone(),
+ Arc::new(CoalescePartitionsExec::new(rebuilt_write)),
+ );
+ let inserted = run(&rebuilt_commit, &ctx).await?;
+ assert!(inserted.contains("| 2 |"), "{inserted}");
+
+ // Read path: a scan pinned to a snapshot, rebuilt from its accessors,
+ // returns the same rows.
+ let table = client.load_table(&ident).await?;
+ let snapshot_id = table.metadata().current_snapshot_id().unwrap();
+ let pinned =
+ IcebergStaticTableProvider::try_new_from_table_snapshot(table,
snapshot_id)
+ .await?;
+ assert_eq!(pinned.snapshot_id(), Some(snapshot_id));
+ ctx.register_table("pinned", Arc::new(pinned.clone()))?;
+ let plan = ctx
+ .sql("SELECT foo2 FROM pinned WHERE foo1 = 1")
+ .await?
+ .create_physical_plan()
+ .await?;
+ let scan = find_node::<IcebergTableScan>(&plan).expect("a scan");
+ assert!(scan.predicates().is_some(), "the filter is pushed down");
+ let rebuilt = IcebergTableScan::new_with_predicate(
+ scan.table().clone(),
+ scan.snapshot_id(),
+ scan.schema(),
+ scan.projection().map(<[String]>::to_vec),
+ scan.predicates().cloned(),
+ scan.limit(),
+ )?;
+ assert_eq!(rebuilt.schema(), scan.schema());
+ let expected = run(scan, &ctx).await?;
+ assert!(
+ expected.contains("alan") && !expected.contains("turing"),
+ "{expected}"
+ );
+ assert_eq!(run(&rebuilt, &ctx).await?, expected);
+
+ // Without a projection the scan reads every column, and its limit is kept.
+ let plan = pinned.scan(&ctx.state(), None, &[], Some(1)).await?;
+ let scan = plan.downcast_ref::<IcebergTableScan>().expect("a scan");
+ assert_eq!(scan.projection(), None);
+ let rebuilt = IcebergTableScan::new_with_predicate(
+ scan.table().clone(),
+ scan.snapshot_id(),
+ scan.schema(),
+ scan.projection().map(<[String]>::to_vec),
+ scan.predicates().cloned(),
+ scan.limit(),
+ )?;
+ assert_eq!(rebuilt.limit(), Some(1));
+ let expected = run(scan, &ctx).await?;
+ assert_eq!(expected.lines().count(), 5, "one row:\n{expected}");
Review Comment:
Done in 14aaa24: the tests now assert full tables with `expect!`, and the
pinned scan is checked against a later snapshot, so reading the current
snapshot instead would show. `$snapshots` is the exception: its ids, times and
paths change every run, so the test checks each row against the snapshot's
metadata and compares the summaries in full.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
