diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 084ae59a76c23..45dabaac90f59 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -43,7 +43,7 @@ use datafusion_physical_expr_common::sort_expr::{ LexOrdering, PhysicalSortExpr, PhysicalSortRequirement, OrderingRequirements }; use datafusion_physical_expr::{Distribution, Partitioning, PhysicalExpr}; -use datafusion_physical_expr::expressions::{col, BinaryExpr, Column, NotExpr}; +use datafusion_physical_expr::expressions::{col, is_not_null, BinaryExpr, Column, NotExpr}; use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::repartition::RepartitionExec; @@ -359,6 +359,150 @@ async fn test_do_not_remove_sort_with_limit() -> Result<()> { Ok(()) } +#[test] +fn test_do_not_remove_sort_with_limit_after_projection() -> Result<()> { + let schema = create_test_schema()?; + let source = parquet_exec(Arc::clone(&schema)); + let sort = sort_exec([sort_expr("non_nullable_col", &schema)].into(), source); + // Dropping the sort key hides the output ordering, but preserves the row + // sequence that determines which rows a limit skips or returns. + let projection = projection_exec( + vec![(col("nullable_col", &schema)?, "nullable_col".to_string())], + sort, + )?; + assert!(projection.output_ordering().is_none()); + + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 1; + let filter = filter_exec( + is_not_null(col("nullable_col", &projection.schema())?)?, + Arc::clone(&projection), + ); + for input in [projection, filter] { + for plan in [ + global_limit_exec(Arc::clone(&input), 1, None), + global_limit_exec(Arc::clone(&input), 1, Some(1)), + global_limit_exec(Arc::clone(&input), 0, Some(1)), + local_limit_exec(input, 1), + ] { + let optimized = + EnsureRequirements::new().optimize(Arc::clone(&plan), &config)?; + let optimized = SanityCheckPlan::new().optimize(optimized, &config)?; + assert_eq!(get_plan_string(&optimized), get_plan_string(&plan)); + } + } + Ok(()) +} + +#[test] +fn test_required_sort_after_projection() -> Result<()> { + let schema = create_test_schema()?; + let source = parquet_exec(Arc::clone(&schema)); + let sort = sort_exec([sort_expr("non_nullable_col", &schema)].into(), source); + let projection = projection_exec( + vec![(col("nullable_col", &schema)?, "nullable_col".to_string())], + sort, + )?; + let filter = filter_exec( + is_not_null(col("nullable_col", &projection.schema())?)?, + projection, + ); + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 1; + + for with_limit in [false, true] { + let input = if with_limit { + global_limit_exec(Arc::clone(&filter), 1, None) + } else { + Arc::clone(&filter) + }; + let window = bounded_window_exec( + "nullable_col", + [sort_expr("nullable_col", &input.schema())], + input, + ); + let optimized = EnsureRequirements::new().optimize(window, &config)?; + let optimized = SanityCheckPlan::new().optimize(optimized, &config)?; + let twice = + EnsureRequirements::new().optimize(Arc::clone(&optimized), &config)?; + assert_eq!(get_plan_string(&twice), get_plan_string(&optimized)); + + if with_limit { + // The window needs its own sort, but OFFSET still depends on the + // hidden sort to select rows. + assert_snapshot!(displayable(optimized.as_ref()).indent(true), @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] + GlobalLimitExec: skip=1, fetch=None + FilterExec: nullable_col@0 IS NOT NULL + ProjectionExec: expr=[nullable_col@0 as nullable_col] + SortExec: expr=[non_nullable_col@1 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + "#); + } else { + // Without a limit, the window's ordering supersedes the hidden sort. + assert_snapshot!(displayable(optimized.as_ref()).indent(true), @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] + FilterExec: nullable_col@0 IS NOT NULL + ProjectionExec: expr=[nullable_col@0 as nullable_col] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + "#); + } + } + Ok(()) +} + +#[tokio::test] +async fn test_do_not_remove_sort_below_fetch_after_projection() -> Result<()> { + let batch = + record_batch!(("a", Int32, [30, 10, 20, 40]), ("b", Int32, [2, 4, 1, 3]))?; + let schema = batch.schema(); + let source = Arc::new(DataSourceExec::new(Arc::new( + datafusion::datasource::memory::MemorySourceConfig::try_new( + &[vec![batch]], + Arc::clone(&schema), + None, + )?, + ))) as Arc; + let sort = sort_exec([sort_expr("b", &schema)].into(), source); + let filter = filter_exec(is_not_null(col("a", &schema)?)?, sort) + .with_fetch(Some(2)) + .unwrap(); + let projection = + projection_exec(vec![(col("a", &schema)?, "a".to_string())], filter)?; + assert!(projection.output_ordering().is_none()); + + // The window needs a new ordering, but the filter's embedded limit still + // depends on the hidden sort to select a=20 and a=30. + let window = + bounded_window_exec("a", [sort_expr("a", &projection.schema())], projection); + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 1; + let optimized = EnsureRequirements::new().optimize(window, &config)?; + let optimized = SanityCheckPlan::new().optimize(optimized, &config)?; + let twice = EnsureRequirements::new().optimize(Arc::clone(&optimized), &config)?; + assert_eq!(get_plan_string(&twice), get_plan_string(&optimized)); + + for plan in [optimized, twice] { + let results = + datafusion_physical_plan::collect(plan, Arc::new(TaskContext::default())) + .await?; + datafusion_common::assert_batches_eq!( + [ + "+----+-------+", + "| a | count |", + "+----+-------+", + "| 20 | 1 |", + "| 30 | 2 |", + "+----+-------+", + ], + &results + ); + } + Ok(()) +} + #[tokio::test] async fn test_union_inputs_sorted() -> Result<()> { let schema = create_test_schema()?; diff --git a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs index 8a4905365b7af..4d35061b7fa5c 100644 --- a/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs @@ -102,8 +102,9 @@ fn update_sort_ctx_children_data( child_node.data = if is_sort(child_plan) { // child is sort true - } else if is_limit(child_plan) { - // There is no sort linkage for this path, it starts at a limit. + } else if is_limit(child_plan) || child_plan.fetch().is_some() { + // Removing a sort below a row limit can change which rows survive, + // including limits embedded in operators such as FilterExec. false } else { // If a descendent is a sort, and the child maintains the sort. @@ -427,16 +428,14 @@ pub fn ensure_sorting( .zip(requirements.children) .enumerate() { - let physical_ordering = child.plan.output_ordering(); - if let Some(required) = required_ordering { let eq_properties = child.plan.equivalence_properties(); let req = required.into_single(); if !eq_properties.ordering_satisfy_requirement(req.clone())? { - // Make sure we preserve the ordering requirements: - if physical_ordering.is_some() { - child = update_child_to_remove_unnecessary_sort(idx, child, plan)?; - } + // The new sort supersedes linked sorts even when a projection + // hides their keys. Explicit and embedded limits stop sort + // linkage, preserving sorts that determine which rows survive. + child = update_child_to_remove_unnecessary_sort(idx, child, plan)?; child = add_sort_above( child, req, @@ -446,9 +445,11 @@ pub fn ensure_sorting( ); child = update_sort_ctx_children_data(child, true)?; } - } else if physical_ordering.is_none() || !plan.maintains_input_order()[idx] { - // We have a `SortExec` whose effect may be neutralized by another - // order-imposing operator, remove this sort: + } else if !plan.maintains_input_order()[idx] { + // Remove linked sorts only when this operator does not preserve + // their row order. A projection can hide the sort key, making + // output_ordering() None, while a later LIMIT or OFFSET still + // depends on that order to select rows. child = update_child_to_remove_unnecessary_sort(idx, child, plan)?; } updated_children.push(child); diff --git a/datafusion/sqllogictest/test_files/limit.slt b/datafusion/sqllogictest/test_files/limit.slt index 06d1c85f5f514..d5f1cdbc7be40 100644 --- a/datafusion/sqllogictest/test_files/limit.slt +++ b/datafusion/sqllogictest/test_files/limit.slt @@ -1257,3 +1257,58 @@ physical_plan statement ok DROP TABLE t22489; + +# Regression test for https://github.com/apache/datafusion/issues/23534 +# The sort below OFFSET selects rows even when a projection hides its sort key. +statement ok +CREATE TABLE t23534 (grp VARCHAR NOT NULL, sort_key BIGINT NOT NULL) +AS VALUES ('physical_second', 2), ('sorted_first', 1); + +query TI +SELECT grp, sort_key FROM t23534 ORDER BY sort_key ASC NULLS LAST OFFSET 1; +---- +physical_second 2 + +query TI +SELECT grp, COUNT(*) AS n +FROM ( + SELECT grp, sort_key FROM t23534 + ORDER BY sort_key ASC NULLS LAST OFFSET 1 +) q +GROUP BY grp +ORDER BY grp; +---- +physical_second 1 + +# One pass leaves a projection between the limit and sort. Correctness must +# not depend on a later logical rewrite moving that projection. +statement ok +SET datafusion.optimizer.max_passes = 1; + +query TI +SELECT grp, COUNT(*) AS n +FROM ( + SELECT grp, sort_key FROM t23534 + ORDER BY sort_key ASC NULLS LAST OFFSET 1 +) q +GROUP BY grp +ORDER BY grp; +---- +physical_second 1 + +query TI +SELECT grp, COUNT(*) AS n +FROM ( + SELECT grp, sort_key FROM t23534 + ORDER BY sort_key ASC NULLS LAST LIMIT 1 OFFSET 1 +) q +GROUP BY grp +ORDER BY grp; +---- +physical_second 1 + +statement ok +SET datafusion.optimizer.max_passes = 3; + +statement ok +DROP TABLE t23534;