Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 19 additions & 2 deletions crates/datafusion/src/table_provider/catalog_table_provider.rs
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,11 @@ impl IcebergCatalogTableProvider {
///
/// Loads the table once to get the initial schema, then stores the catalog
/// reference for future metadata refreshes on each operation.
///
/// # Errors
///
/// Returns an error if the catalog cannot load the table, or its schema has
/// no Arrow equivalent.
pub async fn try_new(
catalog: Arc<dyn Catalog>,
namespace: NamespaceIdent,
Expand All @@ -83,11 +88,23 @@ impl IcebergCatalogTableProvider {
.map_err(to_datafusion_error)?,
);

Ok(IcebergCatalogTableProvider {
Ok(Self::new_with_schema(catalog, table_ident, schema))
}

/// Creates a provider for `table_ident` that declares `schema`, without loading
/// the table. For rebuilding a provider in another process from [`Self::catalog`],
/// [`Self::table_ident`] and [`TableProvider::schema`]. `schema` is not checked
/// against the table; scans read its columns by name from the snapshot they scan.
pub fn new_with_schema(
catalog: Arc<dyn Catalog>,
table_ident: TableIdent,
schema: ArrowSchemaRef,
) -> Self {
IcebergCatalogTableProvider {
Comment on lines +98 to +103

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't love the fact that the public API of this project makes incorrect state representable.

People here can choose to shoot themselves in the foot and pass a schema that has nothing to do with the actual table. This problem goes a bit beyond just this new constructor method, and it's rooted into the IcebergCatalogTableProvider struct fields themselves.

One thing I notice is that IcebergCatalogTableProvider only needs the Catalog itself because it's required for constructing an IcebergCommitExec. It's a bit unfortunate that Catalog is dangerously close to a god object in charge of everything, some interface segregation there could have prevented from needing to thread it around everywhere.

Given that Catalog contains so many methods and not necessarily related functionality, I wonder if it'd make sense to have it be part of the DataFusion's SessionConfig.extensions and thread it around as part of the context, rather than manually requiring it to be stored in every struct.

catalog,
table_ident,
schema,
})
}
}

pub(crate) async fn metadata_table(
Expand Down
52 changes: 30 additions & 22 deletions crates/datafusion/tests/integration_datafusion_test.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1065,12 +1065,11 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box<dyn Error>> {
.expect("a catalog-backed provider");
assert_eq!(provider.table_ident(), &ident);
assert!(Arc::ptr_eq(provider.catalog(), &client));
let rebuilt = IcebergCatalogTableProvider::try_new(
let rebuilt = IcebergCatalogTableProvider::new_with_schema(
provider.catalog().clone(),
provider.table_ident().namespace().clone(),
provider.table_ident().name(),
)
.await?;
provider.table_ident().clone(),
provider.schema(),
);
assert_eq!(rebuilt.table_ident(), &ident);
assert_eq!(rebuilt.schema(), provider.schema());

Expand Down Expand Up @@ -1395,23 +1394,32 @@ async fn test_scan_after_schema_evolution_reads_provider_columns()
.collect()
.await?;

let plan = provider.scan(&ctx.state(), None, &[], None).await?;
assert_eq!(plan.schema(), provider.schema());
// The rows come from two data files, which may be read in either order.
// They are joined under their own schema, not the plan's, so that a column
// the plan does not report would show.
let batches = run_batches(plan.as_ref(), &ctx).await?;
let batch = concat_batches(&batches[0].schema(), &batches)?;
let order = sort_to_indices(batch.column_by_name("foo1").unwrap(), None, None)?;
let sorted = take_record_batch(&batch, &order)?;
expect![[r#"
+------+--------+
| foo1 | foo2 |
+------+--------+
| 1 | alan |
| 2 | turing |
+------+--------+"#]]
.assert_eq(&pretty_format_batches(&[sorted])?.to_string());
// A provider rebuilt from the original's parts, as a codec does, scans the
// same columns as the original, not the table's current ones.
let rebuilt = IcebergCatalogTableProvider::new_with_schema(
provider.catalog().clone(),
provider.table_ident().clone(),
provider.schema(),
);
for provider in [provider.as_ref(), &rebuilt] {
let plan = provider.scan(&ctx.state(), None, &[], None).await?;
assert_eq!(plan.schema(), provider.schema());
// The rows come from two data files, which may be read in either order.
// They are joined under their own schema, not the plan's, so that a
// column the plan does not report would show.
let batches = run_batches(plan.as_ref(), &ctx).await?;
let batch = concat_batches(&batches[0].schema(), &batches)?;
let order = sort_to_indices(batch.column_by_name("foo1").unwrap(), None, None)?;
let sorted = take_record_batch(&batch, &order)?;
expect![[r#"
+------+--------+
| foo1 | foo2 |
+------+--------+
| 1 | alan |
| 2 | turing |
+------+--------+"#]]
.assert_eq(&pretty_format_batches(&[sorted])?.to_string());
}

Ok(())
}
Loading