mbutrovich commented on code in PR #2671:
URL: https://github.com/apache/iceberg-rust/pull/2671#discussion_r3667578012
##########
crates/integrations/datafusion/src/physical_plan/scan.rs:
##########
@@ -138,18 +137,60 @@ 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();
-
- // Apply limit if specified
+ 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 = build_table_scan(&self.table, &self.scan_config)?
+ .arrow_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);
+
+ Box::pin(stream)
+ }
+ None => {
+ let table = self.table.clone();
+ let scan_config = self.scan_config.clone();
+ let fut = async move {
+ let table_scan = build_table_scan(&table, &scan_config)?;
+ let stream = table_scan
+ .to_arrow()
+ .await
+ .map_err(to_datafusion_error)?
+ .map_err(to_datafusion_error);
+ Ok::<_, datafusion::common::DataFusionError>(stream)
+ };
+
+ Box::pin(futures::stream::once(fut).try_flatten())
+ }
+ };
+
+ // Apply a scan-partition bound if specified. In eager planning this
is only
Review Comment:
This explains why the per-partition limit looseness is safe (a
`GlobalLimitExec` above the scan enforces the real bound), which is a fair
resolution of the earlier comment. But nothing exercises it: no test runs a
`SELECT ... LIMIT k` query against an eager multi-partition scan and checks the
row count and rows. Since this PR is what introduces the N-way over-read in the
first place (the single-partition lazy path couldn't over-read), a targeted
test here would confirm the `GlobalLimitExec` assumption holds end to end
rather than resting on the comment alone.
##########
crates/integrations/datafusion/src/config.rs:
##########
@@ -0,0 +1,31 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use datafusion::common::config::ConfigExtension;
+use datafusion::common::extensions_options;
+
+extensions_options! {
+ /// Configuration options for Iceberg's DataFusion integration.
+ pub struct IcebergDataFusionConfig {
Review Comment:
`SET iceberg.enable_eager_scan_planning = true` will not work on a plain
`SessionContext`. Nothing in this crate registers `IcebergDataFusionConfig` as
an extension anywhere except the test helper (`read_session_config` in
`integration_datafusion_test.rs:287-297`, via `.with_option_extension(...)`).
Checked `datafusion-common` 54.1.0 (pinned in this repo's `Cargo.lock`):
`ConfigOptions::set` looks up the prefix in `self.extensions.0`, and if nothing
has registered `"iceberg"`, it returns `Could not find config namespace
"iceberg"` rather than setting the option.
`test_set_enable_eager_scan_planning` (`integration_datafusion_test.rs:359`)
only passes because it pre-registers the extension with `Some(false)` before
running the `SET` statement, so it doesn't cover the case of a user who never
touches the extension in Rust first.
So today, using this feature via SQL `SET` requires the embedding
application to already know to call
`.with_option_extension(IcebergDataFusionConfig::default())` before `SET` can
target it at all, and that's not documented anywhere. Worth a line on
`IcebergDataFusionConfig`'s doc comment spelling out the registration
requirement, or registering it somewhere reachable (e.g.
`IcebergCatalogProvider::try_new`) so `SET` works without that precondition.
--
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]
