fix: pass the session's batch_size to the Iceberg table scan

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<Vec<String>>,
     predicate: Option<Predicate>,
     enable_page_index: bool,
+    batch_size: usize,
 ) -> Result<Pin<Box<dyn Stream<Item = Result<RecordBatch>> + 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)?;
 

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<RecordBatch>)> {
         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))
     }
 }
 

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(())
+}