Automatically split large single RecordBatches in MemorySource into smaller batches (#16734)

- **Config Options**
  - Added `batch_split` execution option (enabled by default).
  - Made batch size configurable with default value `8192`.

- **Execution Plan Enhancements**
  - Introduced `BatchSplitStream` to split large record batches before passing them downstream.
  - Integrated `BatchSplitStream` within `DataSourceExec`.

- **Metrics**
  - Added `SplitMetrics` with `batches_splitted` counter.
  - Exposed split metrics via the execution plan's metrics system.

- **Testing**
  - Added new integration tests for batch splitting scenarios in `datasource_split.rs`:
    - Large batches split into multiple parts.
    - Exact batch size yields no split.
    - Small batches pass through unchanged.
    - Empty batches handled correctly.
    - Multiple empty batches processed safely.
  - Added unit tests for `BatchSplitStream` behavior.
  - Updated SQL logic tests (`.slt` files) to reflect changed execution plans and batch sizes.

- **SQL Logic Test Adjustments**
  - Updated expected physical plans and outputs across various `.slt` files.
  - Changed default execution batch size in tests for consistency.

Author: kosiew

datafusion/core/src/datasource/listing/table.rs

@@ -2026,7 +2026,6 @@ mod tests {
     #[tokio::test]
     async fn test_insert_into_append_new_parquet_files_session_overrides() -> Result<()> {
         let mut config_map: HashMap<String, String> = HashMap::new();
-        config_map.insert("datafusion.execution.batch_size".into(), "10".into());
         config_map.insert(
             "datafusion.execution.soft_max_rows_per_output_file".into(),
             "10".into(),
@@ -2091,7 +2090,7 @@ mod tests {
             "datafusion.execution.parquet.write_batch_size".into(),
             "5".into(),
         );
-        config_map.insert("datafusion.execution.batch_size".into(), "1".into());
+        config_map.insert("datafusion.execution.batch_size".into(), "10".into());
         helper_test_append_new_files_to_table(
             ParquetFormat::default().get_ext(),
             FileCompressionType::UNCOMPRESSED,

datafusion/core/tests/execution/datasource_split.rs

@@ -0,0 +1,123 @@
+// 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 arrow::{
+    array::{ArrayRef, Int32Array},
+    datatypes::{DataType, Field, Schema},
+    record_batch::RecordBatch,
+};
+use datafusion_datasource::memory::MemorySourceConfig;
+use datafusion_execution::TaskContext;
+use datafusion_physical_plan::{common::collect, ExecutionPlan};
+use std::sync::Arc;
+
+/// Helper function to create a memory source with the given batch size and collect all batches
+async fn create_and_collect_batches(
+    batch_size: usize,
+) -> datafusion_common::Result<Vec<RecordBatch>> {
+    let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+    let array = Int32Array::from_iter_values(0..batch_size as i32);
+    let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array) as ArrayRef])?;
+    let exec = MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None)?;
+    let ctx = Arc::new(TaskContext::default());
+    let stream = exec.execute(0, ctx)?;
+    collect(stream).await
+}
+
+/// Helper function to create a memory source with multiple batches and collect all results
+async fn create_and_collect_multiple_batches(
+    input_batches: Vec<RecordBatch>,
+) -> datafusion_common::Result<Vec<RecordBatch>> {
+    let schema = input_batches[0].schema();
+    let exec = MemorySourceConfig::try_new_exec(&[input_batches], schema, None)?;
+    let ctx = Arc::new(TaskContext::default());
+    let stream = exec.execute(0, ctx)?;
+    collect(stream).await
+}
+
+#[tokio::test]
+async fn datasource_splits_large_batches() -> datafusion_common::Result<()> {
+    let batch_size = 20000;
+    let batches = create_and_collect_batches(batch_size).await?;
+
+    assert!(batches.len() > 1);
+    let max = batches.iter().map(|b| b.num_rows()).max().unwrap();
+    assert!(
+        max <= datafusion_execution::config::SessionConfig::new()
+            .options()
+            .execution
+            .batch_size
+    );
+    let total: usize = batches.iter().map(|b| b.num_rows()).sum();
+    assert_eq!(total, batch_size);
+    Ok(())
+}
+
+#[tokio::test]
+async fn datasource_exact_batch_size_no_split() -> datafusion_common::Result<()> {
+    let session_config = datafusion_execution::config::SessionConfig::new();
+    let configured_batch_size = session_config.options().execution.batch_size;
+
+    let batches = create_and_collect_batches(configured_batch_size).await?;
+
+    // Should not split when exactly equal to batch_size
+    assert_eq!(batches.len(), 1);
+    assert_eq!(batches[0].num_rows(), configured_batch_size);
+    Ok(())
+}
+
+#[tokio::test]
+async fn datasource_small_batch_no_split() -> datafusion_common::Result<()> {
+    // Test with batch smaller than the batch size (8192)
+    let small_batch_size = 512; // Less than 8192
+
+    let batches = create_and_collect_batches(small_batch_size).await?;
+
+    // Should not split small batches below the batch size
+    assert_eq!(batches.len(), 1);
+    assert_eq!(batches[0].num_rows(), small_batch_size);
+    Ok(())
+}
+
+#[tokio::test]
+async fn datasource_empty_batch_clean_termination() -> datafusion_common::Result<()> {
+    let batches = create_and_collect_batches(0).await?;
+
+    // Empty batch should result in one empty batch
+    assert_eq!(batches.len(), 1);
+    assert_eq!(batches[0].num_rows(), 0);
+    Ok(())
+}
+
+#[tokio::test]
+async fn datasource_multiple_empty_batches() -> datafusion_common::Result<()> {
+    let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
+    let empty_array = Int32Array::from_iter_values(std::iter::empty::<i32>());
+    let empty_batch =
+        RecordBatch::try_new(schema.clone(), vec![Arc::new(empty_array) as ArrayRef])?;
+
+    // Create multiple empty batches
+    let input_batches = vec![empty_batch.clone(), empty_batch.clone(), empty_batch];
+    let batches = create_and_collect_multiple_batches(input_batches).await?;
+
+    // Should preserve empty batches without issues
+    assert_eq!(batches.len(), 3);
+    for batch in &batches {
+        assert_eq!(batch.num_rows(), 0);
+    }
+    Ok(())
+}

datafusion/core/tests/execution/mod.rs

@@ -16,4 +16,5 @@
 // under the License.
 
 mod coop;
+mod datasource_split;
 mod logical_plan;

datafusion/datasource/src/source.rs

@@ -25,8 +25,10 @@ use std::sync::Arc;
 use datafusion_physical_plan::execution_plan::{
     Boundedness, EmissionType, SchedulingType,
 };
+use datafusion_physical_plan::metrics::SplitMetrics;
 use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet};
 use datafusion_physical_plan::projection::ProjectionExec;
+use datafusion_physical_plan::stream::BatchSplitStream;
 use datafusion_physical_plan::{
     DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
 };
@@ -267,7 +269,18 @@ impl ExecutionPlan for DataSourceExec {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
-        self.data_source.open(partition, Arc::clone(&context))
+        let stream = self.data_source.open(partition, Arc::clone(&context))?;
+        let batch_size = context.session_config().batch_size();
+        log::debug!(
+            "Batch splitting enabled for partition {partition}: batch_size={batch_size}"
+        );
+        let metrics = self.data_source.metrics();
+        let split_metrics = SplitMetrics::new(&metrics, partition);
+        Ok(Box::pin(BatchSplitStream::new(
+            stream,
+            batch_size,
+            split_metrics,
+        )))
     }
 
     fn metrics(&self) -> Option<MetricsSet> {

datafusion/physical-plan/src/metrics/baseline.rs

@@ -169,6 +169,23 @@ impl SpillMetrics {
     }
 }
 
+/// Metrics for tracking [`crate::stream::BatchSplitStream`] activity
+#[derive(Debug, Clone)]
+pub struct SplitMetrics {
+    /// Number of times an input [`RecordBatch`] was split
+    pub batches_splitted: Count,
+}
+
+impl SplitMetrics {
+    /// Create a new [`SplitMetrics`]
+    pub fn new(metrics: &ExecutionPlanMetricsSet, partition: usize) -> Self {
+        Self {
+            batches_splitted: MetricBuilder::new(metrics)
+                .counter("batches_splitted", partition),
+        }
+    }
+}
+
 /// Trait for things that produce output rows as a result of execution.
 pub trait RecordOutput {
     /// Record that some number of output rows have been produced

datafusion/physical-plan/src/metrics/mod.rs

@@ -32,7 +32,7 @@ use std::{
 use datafusion_common::HashMap;
 
 // public exports
-pub use baseline::{BaselineMetrics, RecordOutput, SpillMetrics};
+pub use baseline::{BaselineMetrics, RecordOutput, SpillMetrics, SplitMetrics};
 pub use builder::MetricBuilder;
 pub use custom::CustomMetricValue;
 pub use value::{Count, Gauge, MetricValue, ScopedTimerGuard, Time, Timestamp};