fix: preserve sorts before limits when projections hide sort keys

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<dyn ExecutionPlan>;
+    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()?;

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);

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;