mbutrovich commented on code in PR #2671:
URL: https://github.com/apache/iceberg-rust/pull/2671#discussion_r3605669905
##########
crates/integrations/datafusion/src/physical_plan/scan.rs:
##########
Review Comment:
`IcebergTableScan` is a leaf node (`children()` returns `vec![]` at
`:117-119`), so any non-empty `children` here is a caller error. Right now it
is dropped instead of surfaced. This is the same issue @timsaucer raised on the
new scan node in #2298, which you fixed there with a
`DataFusionError::Internal` matching `IcebergCommitExec`, and you noted
`IcebergTableScan::with_new_children` had the same latent bug. Since this PR is
the one that keeps and extends this node, it seems like the right place to
close it:
```rust
fn with_new_children(
self: Arc<Self>,
children: Vec<Arc<dyn ExecutionPlan>>,
) -> DFResult<Arc<dyn ExecutionPlan>> {
if !children.is_empty() {
return Err(DataFusionError::Internal(
"IcebergTableScan is a leaf node and cannot have
children".to_string(),
));
}
Ok(self)
}
```
##########
crates/integrations/datafusion/src/physical_plan/scan.rs:
##########
Review Comment:
Each `execute(partition)` call gets its own `remaining = limit`, so with N
output partitions the scan can emit up to `N * limit` rows before DataFusion's
top-level limit trims to `limit`. Results stay correct because a
`GlobalLimitExec` sits above the scan, but each partition reads up to `limit`
rows independently, so this is looser than the single-partition lazy path.
Either clamp the per-partition read, or add a one-line comment noting the
top-level limit makes the loose per-partition bound acceptable, so the intent
is explicit.
##########
crates/integrations/datafusion/tests/integration_datafusion_test.rs:
##########
@@ -95,6 +99,237 @@ fn get_table_creation(
Ok(creation)
}
+async fn get_multi_file_table_context(
+ namespace_name: &str,
+ table_name: &str,
+ data_file_count: usize,
+) -> Result<(Arc<MemoryCatalog>, NamespaceIdent, String)> {
+ let iceberg_catalog = Arc::new(get_iceberg_catalog().await);
+ let namespace = NamespaceIdent::new(namespace_name.to_string());
+ set_test_namespace(&iceberg_catalog, &namespace).await?;
+
+ let creation = get_table_creation(temp_path(), table_name, None)?;
+ iceberg_catalog.create_table(&namespace, creation).await?;
+
+ let write_ctx = SessionContext::new_with_config(
+ SessionConfig::new().with_target_partitions(data_file_count),
+ );
+ let arrow_schema = Arc::new(ArrowSchema::new(vec![
+ Field::new("foo1", DataType::Int32, false),
+ Field::new("foo2", DataType::Utf8, false),
+ ]));
+
+ let batches: Vec<RecordBatch> = (1..=data_file_count as i32)
+ .map(|idx| {
+ RecordBatch::try_new(arrow_schema.clone(), vec![
+ Arc::new(Int32Array::from(vec![idx])) as ArrayRef,
+ Arc::new(StringArray::from(vec![format!("row-{idx}")])) as
ArrayRef,
+ ])
+ })
+ .collect::<std::result::Result<_, _>>()?;
+
+ let partitions = batches.into_iter().map(|batch| vec![batch]).collect();
+ let source_table = Arc::new(MemTable::try_new(arrow_schema,
partitions).unwrap());
+ write_ctx
+ .register_table("source_table", source_table)
+ .unwrap();
+
+ let catalog =
Arc::new(IcebergCatalogProvider::try_new(iceberg_catalog.clone()).await?);
+ write_ctx.register_catalog("catalog", catalog);
+
+ let insert_sql =
+ format!("INSERT INTO catalog.{namespace_name}.{table_name} SELECT *
FROM source_table");
+ let batches = write_ctx
+ .sql(&insert_sql)
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+ assert_eq!(batches.len(), 1);
+
+ let rows_inserted = batches[0]
+ .column(0)
+ .as_any()
+ .downcast_ref::<UInt64Array>()
+ .unwrap();
+ assert_eq!(rows_inserted.value(0), data_file_count as u64);
+
+ Ok((iceberg_catalog, namespace, table_name.to_string()))
+}
+
+async fn get_read_context(
+ catalog: Arc<MemoryCatalog>,
+ target_partitions: usize,
+ enable_eager_scan_planning: Option<bool>,
+) -> Result<SessionContext> {
+ let ctx = SessionContext::new_with_config(read_session_config(
+ target_partitions,
+ enable_eager_scan_planning,
+ ));
+ let catalog = Arc::new(IcebergCatalogProvider::try_new(catalog).await?);
+ ctx.register_catalog("catalog", catalog);
+ Ok(ctx)
+}
+
+fn read_session_config(
+ target_partitions: usize,
+ enable_eager_scan_planning: Option<bool>,
+) -> SessionConfig {
+ let config =
SessionConfig::new().with_target_partitions(target_partitions);
+
+ match enable_eager_scan_planning {
+ Some(enabled) => {
+ let mut iceberg_config = IcebergDataFusionConfig::default();
+ iceberg_config.enable_eager_scan_planning = enabled;
+ config.with_option_extension(iceberg_config)
+ }
+ None => config,
+ }
+}
+
+async fn scan_partition_count(
+ ctx: &SessionContext,
+ namespace: &NamespaceIdent,
+ table_name: &str,
+) -> usize {
+ let provider = ctx.catalog("catalog").unwrap();
+ let namespace_name = &namespace[0];
+ let schema = provider.schema(namespace_name).unwrap();
+ let table = schema.table(table_name).await.unwrap().unwrap();
+
+ let state = ctx.state();
+ let plan = table.scan(&state, None, &[], None).await.unwrap();
+ plan.downcast_ref::<IcebergTableScan>()
+ .expect("Expected IcebergTableScan");
+ plan.properties().output_partitioning().partition_count()
+}
+
+#[tokio::test]
+async fn test_multi_file_scan_produces_multiple_partitions() -> Result<()> {
+ let data_file_count = 3;
+ // Ask for more partitions than files to verify scan planning does not
expose empty partitions.
+ let target_partitions = data_file_count + 1;
+ let (iceberg_catalog, namespace, table_name) =
get_multi_file_table_context(
+ "test_multi_file_scan_partitions",
+ "my_table",
+ data_file_count,
+ )
+ .await?;
+ let ctx = get_read_context(iceberg_catalog, target_partitions,
Some(true)).await?;
+ let actual_partition_count = scan_partition_count(&ctx, &namespace,
&table_name).await;
+
+ assert_eq!(actual_partition_count, data_file_count);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_multi_file_scan_defaults_to_single_lazy_partition() ->
Result<()> {
+ let data_file_count = 3;
+ let target_partitions = data_file_count + 1;
+ let (iceberg_catalog, namespace, table_name) =
get_multi_file_table_context(
+ "test_multi_file_scan_default_lazy",
+ "my_table",
+ data_file_count,
+ )
+ .await?;
+ let ctx = get_read_context(iceberg_catalog, target_partitions,
None).await?;
+
+ let actual_partition_count = scan_partition_count(&ctx, &namespace,
&table_name).await;
+
+ assert_eq!(actual_partition_count, 1);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_set_enable_eager_scan_planning() -> Result<()> {
+ let data_file_count = 3;
+ let target_partitions = data_file_count + 1;
+ let (iceberg_catalog, namespace, table_name) =
+ get_multi_file_table_context("test_set_eager_scan_planning",
"my_table", data_file_count)
+ .await?;
+ let ctx = get_read_context(iceberg_catalog, target_partitions,
Some(false)).await?;
+
+ ctx.sql("SET iceberg.enable_eager_scan_planning = true")
+ .await
+ .unwrap()
+ .collect()
+ .await
+ .unwrap();
+
+ let actual_partition_count = scan_partition_count(&ctx, &namespace,
&table_name).await;
+
+ assert_eq!(actual_partition_count, data_file_count);
+
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_multi_partition_scan_matches_single_partition_results() ->
Result<()> {
Review Comment:
`test_multi_partition_scan_matches_single_partition_results` is a good check
that single- and multi-partition scans return identical rows. But
`get_multi_file_table_context` writes one row per file (`:125`,
`Int32Array::from(vec![idx])`), so every file is a single row group with no
filtering opportunity. A divergence in `row_group_filtering_enabled` or
`batch_size` between the eager and lazy paths (comment 2) would not change the
result here, so the test would still pass. A file with multiple row groups (and
ideally a predicate that prunes some of them) would turn this into a real guard
on the eager/lazy read equivalence.
##########
crates/integrations/datafusion/src/physical_plan/scan.rs:
##########
@@ -138,16 +131,57 @@ impl ExecutionPlan for IcebergTableScan {
fn execute(
&self,
- _partition: usize,
+ partition: usize,
_context: Arc<TaskContext>,
) -> DFResult<SendableRecordBatchStream> {
- let fut = get_batch_stream(
- self.table.clone(),
- self.snapshot_id,
- self.projection.clone(),
- self.predicates.clone(),
- );
- let stream = futures::stream::once(fut).try_flatten();
+ let stream: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>>
= match &self
+ .file_task_groups
+ {
+ Some(file_task_groups) => {
+ let Some(file_task_group) =
file_task_groups.get(partition).cloned() else {
+ return
Err(datafusion::common::DataFusionError::Internal(format!(
+ "IcebergTableScan partition {partition} does not
exist; scan has {} partitions",
+ file_task_groups.len()
+ )));
+ };
+
+ let tasks: FileScanTaskStream = Box::pin(futures::stream::iter(
+ (0..file_task_group.len()).map(move |idx|
Ok(file_task_group[idx].clone())),
+ ));
+ let stream = self
+ .table
+ .reader_builder()
+ // Eager planning lets DataFusion drive scan concurrency
via output
+ // partitions. Match DataFusion's FileStream model, where
each
+ // output partition owns one ScanState; keep one data file
in
+ // flight per output partition here.
+ //
https://github.com/apache/datafusion/blob/ad8e7b7f2babe3fcddc3a4f9b5cd1ac0d1b16ad9/datafusion/datasource/src/file_stream/scan_state.rs#L42-L43
+ .with_data_file_concurrency_limit(1)
+ .build()
+ // TODO: Avoid cloning FileScanTasks here once ArrowReader
can accept shared tasks.
+ .read(tasks)
+ .map_err(to_datafusion_error)?
+ .stream()
+ .map_err(to_datafusion_error);
Review Comment:
The eager arm builds its reader from `self.table.reader_builder()` and
overrides only `.with_data_file_concurrency_limit(1)`:
```rust
let stream = self
.table
.reader_builder()
.with_data_file_concurrency_limit(1)
.build()
.read(tasks)
...
```
That inherits `ArrowReaderBuilder::new` defaults
(`crates/iceberg/src/arrow/reader/mod.rs:65-73`): `row_group_filtering_enabled
= true`, `row_selection_enabled = false`, `batch_size = None`. The lazy arm at
`:174-178` reaches the reader through `TableScan::to_arrow`, which sources
those same three fields from `TableScanBuilder::new` defaults
(`crates/iceberg/src/scan/mod.rs:77-81`) and applies them at
`crates/iceberg/src/scan/mod.rs:487-492`.
So one conceptual "how to read Iceberg data files" decision is now fed by
two default sources in two different crates, with nothing pinning their
equivalence. They match today, but if a reader default is flipped in one
struct, the eager and lazy paths diverge silently in pushdown/read behavior.
This split is also what introduced the divergence: #2298's eager path went
through `TableScan::read` and shared one config source, and dropping that
(consistent with the #2298 discussion about not making `read` public) moved the
eager arm onto `ArrowReaderBuilder` defaults directly.
Two options, in preference order:
- Thread the scan's own reader config into the eager arm so `concurrency =
1` is the only intentional eager/lazy difference. This keeps a single source of
truth for the read settings.
- If that is out of scope, at minimum add a comment here stating the eager
path intentionally mirrors `to_arrow`'s reader defaults, and back it with an
equivalence test over a multi-row-group file.
--
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]
