Fix Partial Sort Get Slice Point Between Batches (#16881)

* Update partial_sort.rs

* Update partial_sort.rs

* Update partial_sort.rs

* add sql test

Author: berkaysynnada

datafusion/catalog/src/streaming.rs

@@ -20,15 +20,17 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use arrow::datatypes::SchemaRef;
-use async_trait::async_trait;
-
 use crate::Session;
 use crate::TableProvider;
-use datafusion_common::{plan_err, Result};
-use datafusion_expr::{Expr, TableType};
+
+use arrow::datatypes::SchemaRef;
+use datafusion_common::{plan_err, DFSchema, Result};
+use datafusion_expr::{Expr, SortExpr, TableType};
+use datafusion_physical_expr::{create_physical_sort_exprs, LexOrdering};
 use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec};
 use datafusion_physical_plan::ExecutionPlan;
+
+use async_trait::async_trait;
 use log::debug;
 
 /// A [`TableProvider`] that streams a set of [`PartitionStream`]
@@ -37,6 +39,7 @@ pub struct StreamingTable {
     schema: SchemaRef,
     partitions: Vec<Arc<dyn PartitionStream>>,
     infinite: bool,
+    sort_order: Vec<SortExpr>,
 }
 
 impl StreamingTable {
@@ -60,13 +63,21 @@ impl StreamingTable {
             schema,
             partitions,
             infinite: false,
+            sort_order: vec![],
         })
     }
+
     /// Sets streaming table can be infinite.
     pub fn with_infinite_table(mut self, infinite: bool) -> Self {
         self.infinite = infinite;
         self
     }
+
+    /// Sets the existing ordering of streaming table.
+    pub fn with_sort_order(mut self, sort_order: Vec<SortExpr>) -> Self {
+        self.sort_order = sort_order;
+        self
+    }
 }
 
 #[async_trait]
@@ -85,16 +96,25 @@ impl TableProvider for StreamingTable {
 
     async fn scan(
         &self,
-        _state: &dyn Session,
+        state: &dyn Session,
         projection: Option<&Vec<usize>>,
         _filters: &[Expr],
         limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
+        let physical_sort = if !self.sort_order.is_empty() {
+            let df_schema = DFSchema::try_from(self.schema.as_ref().clone())?;
+            let eqp = state.execution_props();
+
+            create_physical_sort_exprs(&self.sort_order, &df_schema, eqp)?
+        } else {
+            vec![]
+        };
+
         Ok(Arc::new(StreamingTableExec::try_new(
             Arc::clone(&self.schema),
             self.partitions.clone(),
             projection,
-            None,
+            LexOrdering::new(physical_sort),
             self.infinite,
             limit,
         )?))

datafusion/core/tests/memory_limit/mod.rs

@@ -1084,9 +1084,9 @@ fn batches_byte_size(batches: &[RecordBatch]) -> usize {
 }
 
 #[derive(Debug)]
-struct DummyStreamPartition {
-    schema: SchemaRef,
-    batches: Vec<RecordBatch>,
+pub(crate) struct DummyStreamPartition {
+    pub(crate) schema: SchemaRef,
+    pub(crate) batches: Vec<RecordBatch>,
 }
 
 impl PartitionStream for DummyStreamPartition {

datafusion/core/tests/physical_optimizer/enforce_sorting.rs

@@ -17,6 +17,7 @@
 
 use std::sync::Arc;
 
+use crate::memory_limit::DummyStreamPartition;
 use crate::physical_optimizer::test_utils::{
     aggregate_exec, bounded_window_exec, bounded_window_exec_with_partition,
     check_integrity, coalesce_batches_exec, coalesce_partitions_exec, create_test_schema,
@@ -32,11 +33,11 @@ use arrow::compute::SortOptions;
 use arrow::datatypes::{DataType, SchemaRef};
 use datafusion_common::config::ConfigOptions;
 use datafusion_common::tree_node::{TreeNode, TransformedResult};
-use datafusion_common::{Result, ScalarValue};
+use datafusion_common::{Result, ScalarValue, TableReference};
 use datafusion_datasource::file_scan_config::FileScanConfigBuilder;
 use datafusion_datasource::source::DataSourceExec;
 use datafusion_expr_common::operator::Operator;
-use datafusion_expr::{JoinType, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition};
+use datafusion_expr::{JoinType, SortExpr, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition};
 use datafusion_execution::object_store::ObjectStoreUrl;
 use datafusion_functions_aggregate::average::avg_udaf;
 use datafusion_functions_aggregate::count::count_udaf;
@@ -61,7 +62,14 @@ use datafusion_physical_optimizer::enforce_sorting::sort_pushdown::{SortPushDown
 use datafusion_physical_optimizer::enforce_distribution::EnforceDistribution;
 use datafusion_physical_optimizer::output_requirements::OutputRequirementExec;
 use datafusion_physical_optimizer::PhysicalOptimizerRule;
-
+use datafusion::prelude::*;
+use arrow::array::{Int32Array, RecordBatch};
+use arrow::datatypes::{Field};
+use arrow_schema::Schema;
+use datafusion_execution::TaskContext;
+use datafusion_catalog::streaming::StreamingTable;
+
+use futures::StreamExt;
 use rstest::rstest;
 
 /// Create a sorted Csv exec
@@ -879,6 +887,7 @@ async fn test_soft_hard_requirements_multiple_soft_requirements() -> Result<()>
     assert_optimized!(expected_input, expected_optimized, physical_plan, true);
     Ok(())
 }
+
 #[tokio::test]
 async fn test_soft_hard_requirements_multiple_sorts() -> Result<()> {
     let schema = create_test_schema()?;
@@ -3842,3 +3851,124 @@ fn test_parallelize_sort_preserves_fetch() -> Result<()> {
     );
     Ok(())
 }
+
+#[tokio::test]
+async fn test_partial_sort_with_homogeneous_batches() -> Result<()> {
+    // Create schema for the table
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("a", DataType::Int32, false),
+        Field::new("b", DataType::Int32, false),
+        Field::new("c", DataType::Int32, false),
+    ]));
+
+    // Create homogeneous batches - each batch has the same values for columns a and b
+    let batch1 = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(Int32Array::from(vec![1, 1, 1])),
+            Arc::new(Int32Array::from(vec![1, 1, 1])),
+            Arc::new(Int32Array::from(vec![3, 2, 1])),
+        ],
+    )?;
+    let batch2 = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(Int32Array::from(vec![2, 2, 2])),
+            Arc::new(Int32Array::from(vec![2, 2, 2])),
+            Arc::new(Int32Array::from(vec![4, 6, 5])),
+        ],
+    )?;
+    let batch3 = RecordBatch::try_new(
+        schema.clone(),
+        vec![
+            Arc::new(Int32Array::from(vec![3, 3, 3])),
+            Arc::new(Int32Array::from(vec![3, 3, 3])),
+            Arc::new(Int32Array::from(vec![9, 7, 8])),
+        ],
+    )?;
+
+    // Create session with batch size of 3 to match our homogeneous batch pattern
+    let session_config = SessionConfig::new()
+        .with_batch_size(3)
+        .with_target_partitions(1);
+    let ctx = SessionContext::new_with_config(session_config);
+
+    let sort_order = vec![
+        SortExpr::new(
+            Expr::Column(datafusion_common::Column::new(
+                Option::<TableReference>::None,
+                "a",
+            )),
+            true,
+            false,
+        ),
+        SortExpr::new(
+            Expr::Column(datafusion_common::Column::new(
+                Option::<TableReference>::None,
+                "b",
+            )),
+            true,
+            false,
+        ),
+    ];
+    let batches = Arc::new(DummyStreamPartition {
+        schema: schema.clone(),
+        batches: vec![batch1, batch2, batch3],
+    }) as _;
+    let provider = StreamingTable::try_new(schema.clone(), vec![batches])?
+        .with_sort_order(sort_order)
+        .with_infinite_table(true);
+    ctx.register_table("test_table", Arc::new(provider))?;
+
+    let sql = "SELECT * FROM test_table ORDER BY a ASC, c ASC";
+    let df = ctx.sql(sql).await?;
+
+    let physical_plan = df.create_physical_plan().await?;
+
+    // Verify that PartialSortExec is used
+    let plan_str = displayable(physical_plan.as_ref()).indent(true).to_string();
+    assert!(
+        plan_str.contains("PartialSortExec"),
+        "Expected PartialSortExec in plan:\n{plan_str}",
+    );
+
+    let task_ctx = Arc::new(TaskContext::default());
+    let mut stream = physical_plan.execute(0, task_ctx.clone())?;
+
+    let mut collected_batches = Vec::new();
+    while let Some(batch) = stream.next().await {
+        let batch = batch?;
+        if batch.num_rows() > 0 {
+            collected_batches.push(batch);
+        }
+    }
+
+    // Assert we got 3 separate batches (not concatenated into fewer)
+    assert_eq!(
+        collected_batches.len(),
+        3,
+        "Expected 3 separate batches, got {}",
+        collected_batches.len()
+    );
+
+    // Verify each batch has been sorted within itself
+    let expected_values = [vec![1, 2, 3], vec![4, 5, 6], vec![7, 8, 9]];
+
+    for (i, batch) in collected_batches.iter().enumerate() {
+        let c_array = batch
+            .column(2)
+            .as_any()
+            .downcast_ref::<Int32Array>()
+            .unwrap();
+        let actual = c_array.values().iter().copied().collect::<Vec<i32>>();
+        assert_eq!(actual, expected_values[i], "Batch {i} not sorted correctly",);
+    }
+
+    assert_eq!(
+        task_ctx.runtime_env().memory_pool.reserved(),
+        0,
+        "Memory should be released after execution"
+    );
+
+    Ok(())
+}