diff --git a/datafusion/physical-optimizer/src/pushdown_sort.rs b/datafusion/physical-optimizer/src/pushdown_sort.rs index 5dfe221ed24c0..d6c8ab25e4e60 100644 --- a/datafusion/physical-optimizer/src/pushdown_sort.rs +++ b/datafusion/physical-optimizer/src/pushdown_sort.rs @@ -167,11 +167,15 @@ impl PhysicalOptimizerRule for PushdownSort { // // If the SortExec carried a fetch (LIMIT), we must preserve it. // First try pushing the limit into the source via `with_fetch()`. - // If the source doesn't support `with_fetch`, fall back to - // wrapping with GlobalLimitExec. + // If the source doesn't support `with_fetch`, preserve whether + // the limit applies to each partition or to the entire input. if let Some(fetch) = sort_exec.fetch() { let inner = inner.with_fetch(Some(fetch)).unwrap_or_else(|| { - Arc::new(GlobalLimitExec::new(inner, 0, Some(fetch))) + if sort_exec.preserve_partitioning() { + Arc::new(LocalLimitExec::new(inner, fetch)) + } else { + Arc::new(GlobalLimitExec::new(inner, 0, Some(fetch))) + } }); Ok(Transformed::yes(inner)) } else { diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 0c8c2c5add363..bb98dc8501109 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -73,6 +73,7 @@ use log::info; use sqlparser::ast; use tempfile::TempDir; +mod exact_sort_pushdown; mod range_partitioning; /// Context for running tests @@ -182,6 +183,10 @@ impl TestContext { let file_name = relative_path.file_name().unwrap().to_str().unwrap(); match file_name { + "sort_pushdown.slt" => { + exact_sort_pushdown::register_exact_sort_source(test_ctx.session_ctx()) + .expect("register exact sort source"); + } "parquet_missing_bounds.slt" => { register_parquet_missing_bounds(&mut test_ctx).await; } diff --git a/datafusion/sqllogictest/src/test_context/exact_sort_pushdown.rs b/datafusion/sqllogictest/src/test_context/exact_sort_pushdown.rs new file mode 100644 index 0000000000000..62c59b14869cb --- /dev/null +++ b/datafusion/sqllogictest/src/test_context/exact_sort_pushdown.rs @@ -0,0 +1,141 @@ +// 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 std::fmt::{self, Formatter}; +use std::sync::Arc; + +use arrow::array::record_batch; +use arrow::datatypes::{self as arrow_schema, SchemaRef}; +use async_trait::async_trait; +use datafusion::catalog::{Session, TableProvider}; +use datafusion::common::Result; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::{Expr, TableType}; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr}; +use datafusion::physical_plan::test::TestMemoryExec; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SortOrderPushdownResult, +}; +use datafusion::prelude::SessionContext; + +pub(super) fn register_exact_sort_source(ctx: &SessionContext) -> Result<()> { + let partitions = (0..4) + .map(|i| { + Ok(vec![record_batch!( + ("group_id", Int32, [0, 0, 1, 1]), + ("value", Int32, [i, i + 4, i + 8, i + 12]) + )?]) + }) + .collect::>>()?; + let input = TestMemoryExec::try_new(&partitions, partitions[0][0].schema(), None)?; + ctx.register_table("exact_sort_source", Arc::new(ExactSortSource(input)))?; + Ok(()) +} + +fn ordering() -> LexOrdering { + [ + PhysicalSortExpr::new_default(Arc::new(Column::new("group_id", 0))).nulls_last(), + PhysicalSortExpr::new_default(Arc::new(Column::new("value", 1))).nulls_last(), + ] + .into() +} + +/// Accepts exact sort pushdown without supporting fetch pushdown. Its rows are +/// already sorted within each partition; the hook makes that ordering explicit. +#[derive(Debug, Clone)] +struct ExactSortSource(TestMemoryExec); + +#[async_trait] +impl TableProvider for ExactSortSource { + fn schema(&self) -> SchemaRef { + self.0.schema() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _: &dyn Session, + projection: Option<&[usize]>, + _: &[Expr], + _: Option, + ) -> Result> { + assert!(projection.is_none_or(|indices| indices == [0, 1])); + Ok(Arc::new(self.clone())) + } +} + +impl DisplayAs for ExactSortSource { + fn fmt_as(&self, _: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "ExactSortSource") + } +} + +impl ExecutionPlan for ExactSortSource { + fn name(&self) -> &str { + "ExactSortSource" + } + + fn properties(&self) -> &Arc { + self.0.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + assert!(children.is_empty()); + Ok(self) + } + + fn apply_expressions( + &self, + _: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.0.execute(partition, context) + } + + fn try_pushdown_sort( + &self, + order: &[PhysicalSortExpr], + ) -> Result>> { + if order != ordering().as_ref() { + return Ok(SortOrderPushdownResult::Unsupported); + } + Ok(SortOrderPushdownResult::Exact { + inner: Arc::new(Self( + self.0.clone().try_with_sort_information(vec![ordering()])?, + )), + }) + } +} diff --git a/datafusion/sqllogictest/test_files/sort_pushdown.slt b/datafusion/sqllogictest/test_files/sort_pushdown.slt index da583f127c2b7..b24bfec0e55d4 100644 --- a/datafusion/sqllogictest/test_files/sort_pushdown.slt +++ b/datafusion/sqllogictest/test_files/sort_pushdown.slt @@ -22,6 +22,54 @@ SET datafusion.execution.parquet.pushdown_filters = true; statement ok SET datafusion.optimizer.enable_sort_pushdown = true; +# exact_sort_source has four partitions with interleaved values 0..15. +# Values 0..7 belong to group 0; values 8..15 belong to group 1. +# It supports exact (group_id, value) ordering, but cannot apply a limit itself. +# +# The partitioned window requires redistributing this source's rows by group. +# The outer LIMIT lets the optimizer insert a local TopK before redistribution, +# exercising the standalone sort-pushdown path. Without the window, this query +# uses the already-correct merge/sort path. +# Removing the source-side sort must retain its per-partition limit. +query III +SELECT group_id, value, + SUM(value) OVER ( + PARTITION BY group_id ORDER BY value ROWS UNBOUNDED PRECEDING + ) AS running_sum +FROM exact_sort_source +ORDER BY group_id, value +LIMIT 3; +---- +0 0 0 +0 1 1 +0 2 3 + +# The pushed-down sort must leave a LocalLimitExec below the repartition. +statement ok +SET datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN +SELECT group_id, value, + SUM(value) OVER ( + PARTITION BY group_id ORDER BY value ROWS UNBOUNDED PRECEDING + ) AS running_sum +FROM exact_sort_source +ORDER BY group_id, value +LIMIT 3; +---- +physical_plan +01)SortPreservingMergeExec: [group_id@0 ASC NULLS LAST, value@1 ASC NULLS LAST], fetch=3 +02)--ProjectionExec: expr=[group_id@0 as group_id, value@1 as value, sum(exact_sort_source.value) PARTITION BY [exact_sort_source.group_id] ORDER BY [exact_sort_source.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as running_sum] +03)----BoundedWindowAggExec: wdw=[sum(exact_sort_source.value) PARTITION BY [exact_sort_source.group_id] ORDER BY [exact_sort_source.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(exact_sort_source.value) PARTITION BY [exact_sort_source.group_id] ORDER BY [exact_sort_source.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +04)------SortExec: TopK(fetch=3), expr=[group_id@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +05)--------RepartitionExec: partitioning=Hash([group_id@0], 4), input_partitions=4 +06)----------LocalLimitExec: fetch=3 +07)------------ExactSortSource + +statement ok +RESET datafusion.explain.physical_plan_only; + # Test 1: Sort Pushdown for ordered Parquet files # Create a sorted dataset statement ok