From 4ed089db3e5ac0d904725621e01a635911a5df60 Mon Sep 17 00:00:00 2001 From: nikkhiil-hadoop Date: Fri, 6 Nov 2026 15:22:32 +0530 Subject: [PATCH] fix: pass the session's batch_size to the Iceberg table scan IcebergDataSource::open did not pass execution.batch_size to get_batch_stream, which never called TableScanBuilder::with_batch_size, so scans always read arrow-rs's default 1024-row batches. Add a regression test that sets batch_size above 1024 (smaller sizes pass without the fix because DataFusion splits oversized batches after the scan), and an IcebergTestHarness::query_batches helper that returns unformatted batches. Closes #100 --- crates/datafusion/src/read/scan.rs | 3 ++ crates/datafusion/src/test_utils/harness.rs | 13 ++++-- crates/datafusion/tests/batch_size.rs | 44 +++++++++++++++++++++ 3 files changed, 56 insertions(+), 4 deletions(-) create mode 100644 crates/datafusion/tests/batch_size.rs diff --git a/crates/datafusion/src/read/scan.rs b/crates/datafusion/src/read/scan.rs index 3d9f752..a3259fd 100644 --- a/crates/datafusion/src/read/scan.rs +++ b/crates/datafusion/src/read/scan.rs @@ -195,6 +195,7 @@ impl DataSource for IcebergDataSource { .execution .parquet .enable_page_index, + context.session_config().batch_size(), ); let stream = futures::stream::once(fut).try_flatten(); @@ -281,6 +282,7 @@ async fn get_batch_stream( column_names: Option>, predicate: Option, enable_page_index: bool, + batch_size: usize, ) -> Result> + Send>>> { let scan_builder = match snapshot_id { Some(snapshot_id) => table.scan().snapshot_id(snapshot_id), @@ -296,6 +298,7 @@ async fn get_batch_stream( } let table_scan = scan_builder .with_row_selection_enabled(enable_page_index) + .with_batch_size(Some(batch_size)) .build() .map_err(to_datafusion_error)?; diff --git a/crates/datafusion/src/test_utils/harness.rs b/crates/datafusion/src/test_utils/harness.rs index d106467..2afbdab 100644 --- a/crates/datafusion/src/test_utils/harness.rs +++ b/crates/datafusion/src/test_utils/harness.rs @@ -21,6 +21,7 @@ use std::sync::Arc; use async_trait::async_trait; use bytes::Bytes; +use datafusion::arrow::array::RecordBatch; use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::dataframe::DataFrame; use datafusion::error::Result; @@ -100,14 +101,18 @@ impl IcebergTestHarness { } pub async fn query(&self, sql: &str) -> Result<(String, String)> { + let (plan, batches) = self.query_batches(sql).await?; + Ok((plan, pretty_format_batches(&batches)?.to_string())) + } + + /// Like [`Self::query`], but returns the batches unformatted, for tests + /// that inspect batch boundaries. + pub async fn query_batches(&self, sql: &str) -> Result<(String, Vec)> { let dataframe: DataFrame = self.ctx.sql(sql).await?; let plan = dataframe.create_physical_plan().await?; let batches = dataframe.collect().await?; - Ok(( - displayable(plan.as_ref()).indent(true).to_string(), - pretty_format_batches(&batches)?.to_string(), - )) + Ok((displayable(plan.as_ref()).indent(true).to_string(), batches)) } } diff --git a/crates/datafusion/tests/batch_size.rs b/crates/datafusion/tests/batch_size.rs new file mode 100644 index 0000000..67acf68 --- /dev/null +++ b/crates/datafusion/tests/batch_size.rs @@ -0,0 +1,44 @@ +// 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::error::Result; +use iceberg_datafusion::test_utils::IcebergTestHarness; + +#[tokio::test] +async fn scan_honors_session_batch_size() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + + // Batches above the session's batch_size are split after the scan, so a + // size below arrow-rs's default of 1024 would pass without the reader + // seeing it. Use a larger one. + harness + .query("SET datafusion.execution.batch_size = 2000") + .await?; + let (plan, batches) = harness + .query_batches("SELECT pickup_date FROM taxi") + .await?; + + // Nothing above the scan may merge batches, or the assertion below would + // not be about the scan. + insta::assert_snapshot!(plan, @r" + CooperativeExec + DataSourceExec: format=iceberg, projection=[pickup_date] + "); + let max_rows = batches.iter().map(|batch| batch.num_rows()).max(); + assert_eq!(max_rows, Some(2000)); + Ok(()) +}