NoahKusaba commented on code in PR #20:
URL: https://github.com/apache/datafusion-iceberg/pull/20#discussion_r4208941331
##########
crates/datafusion/tests/integration_datafusion_test.rs:
##########
@@ -977,3 +992,395 @@ async fn test_insert_into_partitioned() -> Result<(),
Box<dyn Error>> {
Ok(())
}
+
+/// Executes `plan`, which must have a single partition, and returns its rows.
+async fn run_batches(
+ plan: &dyn ExecutionPlan,
+ ctx: &SessionContext,
+) -> Result<Vec<RecordBatch>, Box<dyn Error>> {
+ assert_eq!(plan.properties().partitioning.partition_count(), 1);
+ let stream = plan.execute(0, ctx.task_ctx())?;
+ Ok(collect(stream).await?)
+}
+
+/// Executes `plan`, which must have a single partition, and renders its rows
+/// as a table.
+async fn run(
+ plan: &dyn ExecutionPlan,
+ ctx: &SessionContext,
+) -> Result<String, Box<dyn Error>> {
+ Ok(pretty_format_batches(&run_batches(plan, ctx).await?)?.to_string())
+}
+
+/// Rebuilds `scan` from its accessors alone, as a codec would.
+fn rebuild_scan(scan: &IcebergTableScan) -> IcebergTableScan {
+ IcebergTableScan::new_with_predicate(
+ scan.table().clone(),
+ scan.snapshot_id(),
+ scan.schema(),
+ scan.predicates().cloned(),
+ scan.limit(),
+ )
+}
+
+/// 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)),
+ );
+ expect![[r#"
+ +-------+
+ | count |
+ +-------+
+ | 2 |
+ +-------+"#]]
+ .assert_eq(&run(&rebuilt_commit, &ctx).await?);
+
+ // 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()))?;
+ // A later write, so a scan reading the current snapshot rather than the
+ // pinned one would return its row too.
+ ctx.sql("INSERT INTO catalog.test_plan_nodes.my_table VALUES (3,
'hopper')")
+ .await?
+ .collect()
+ .await?;
+ let latest_snapshot_id = client
+ .load_table(&ident)
+ .await?
+ .metadata()
+ .current_snapshot_id()
+ .unwrap();
+ assert_ne!(latest_snapshot_id, snapshot_id);
+ 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_eq!(
+ scan.predicates().map(ToString::to_string).as_deref(),
+ Some("foo1 = 1")
+ );
+ let rebuilt = rebuild_scan(scan);
+ assert_eq!(rebuilt.schema(), scan.schema());
+ assert_eq!(rebuilt.projection(), scan.projection());
+ let expected = run(scan, &ctx).await?;
+ expect![[r#"
+ +------+------+
+ | foo1 | foo2 |
+ +------+------+
+ | 1 | alan |
+ +------+------+"#]]
+ .assert_eq(&expected);
+ assert_eq!(run(&rebuilt, &ctx).await?, expected);
+
+ // Without a projection the scan reads every column by name, 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(), ["foo1".to_string(), "foo2".to_string()]);
+ let rebuilt = rebuild_scan(scan);
+ assert_eq!(rebuilt.limit(), Some(1));
+ let expected = run(scan, &ctx).await?;
+ expect![[r#"
+ +------+------+
+ | foo1 | foo2 |
+ +------+------+
+ | 1 | alan |
+ +------+------+"#]]
+ .assert_eq(&expected);
Review Comment:
Yeah I generally agree with you on this, using raw assertions in tests is
something i saw on either the Ballista or iceberg-rust repo and assumed it was
the standard? I'll make sure to follow your test standards that you're
establishing in future PR's.
I'll admit I haven't done a great job post-validating the readability and
elegance of ai generated tests. (sorry mbutrovich).
I'll leave this as my summary comment on your test critiques.
--
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]
