Push sort requirements through simple projections

datafusion/core/tests/core_integration.rs

@@ -63,6 +63,9 @@ mod tracing;
 /// Run all tests that are found in the `extension_types` directory
 mod extension_types;
 
+/// Helper functions for tests.
+mod helper;
+
 #[cfg(test)]
 #[ctor::ctor(unsafe)]
 fn init() {

datafusion/core/tests/dataframe/mod.rs

@@ -3007,22 +3007,22 @@ async fn test_count_wildcard_on_sort() -> Result<()> {
     assert_snapshot!(
         pretty_format_batches(&sql_results).unwrap(),
         @r"
-    +---------------+------------------------------------------------------------------------------------+
-    | plan_type     | plan                                                                               |
-    +---------------+------------------------------------------------------------------------------------+
-    | logical_plan  | Sort: count(*) ASC NULLS LAST                                                      |
-    |               |   Projection: t1.b, count(Int64(1)) AS count(*)                                    |
-    |               |     Aggregate: groupBy=[[t1.b]], aggr=[[count(Int64(1))]]                          |
-    |               |       TableScan: t1 projection=[b]                                                 |
-    | physical_plan | SortPreservingMergeExec: [count(*)@1 ASC NULLS LAST]                               |
-    |               |   SortExec: expr=[count(*)@1 ASC NULLS LAST], preserve_partitioning=[true]         |
-    |               |     ProjectionExec: expr=[b@0 as b, count(Int64(1))@1 as count(*)]                 |
-    |               |       AggregateExec: mode=FinalPartitioned, gby=[b@0 as b], aggr=[count(Int64(1))] |
-    |               |         RepartitionExec: partitioning=Hash([b@0], 4), input_partitions=1           |
-    |               |           AggregateExec: mode=Partial, gby=[b@0 as b], aggr=[count(Int64(1))]      |
-    |               |             DataSourceExec: partitions=1, partition_sizes=[1]                      |
-    |               |                                                                                    |
-    +---------------+------------------------------------------------------------------------------------+
+    +---------------+-------------------------------------------------------------------------------------+
+    | plan_type     | plan                                                                                |
+    +---------------+-------------------------------------------------------------------------------------+
+    | logical_plan  | Sort: count(*) ASC NULLS LAST                                                       |
+    |               |   Projection: t1.b, count(Int64(1)) AS count(*)                                     |
+    |               |     Aggregate: groupBy=[[t1.b]], aggr=[[count(Int64(1))]]                           |
+    |               |       TableScan: t1 projection=[b]                                                  |
+    | physical_plan | SortPreservingMergeExec: [count(*)@1 ASC NULLS LAST]                                |
+    |               |   ProjectionExec: expr=[b@0 as b, count(Int64(1))@1 as count(*)]                    |
+    |               |     SortExec: expr=[count(Int64(1))@1 ASC NULLS LAST], preserve_partitioning=[true] |
+    |               |       AggregateExec: mode=FinalPartitioned, gby=[b@0 as b], aggr=[count(Int64(1))]  |
+    |               |         RepartitionExec: partitioning=Hash([b@0], 4), input_partitions=1            |
+    |               |           AggregateExec: mode=Partial, gby=[b@0 as b], aggr=[count(Int64(1))]       |
+    |               |             DataSourceExec: partitions=1, partition_sizes=[1]                       |
+    |               |                                                                                     |
+    +---------------+-------------------------------------------------------------------------------------+
     "
     );
 

datafusion/core/tests/helper/mod.rs

@@ -0,0 +1,23 @@
+// 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.
+
+//! Shared helpers for the `core_integration` test crate.
+//!
+//! Keep cross-cutting test utilities here when they are used by multiple test
+//! modules under `core/tests`. Placing them in this submodule avoids creating
+//! an additional Cargo integration test target for each helper file.
+pub(crate) mod plan_metrics;

datafusion/core/tests/helper/plan_metrics.rs

@@ -0,0 +1,54 @@
+// 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.
+
+//! Helpers for aggregating execution metrics across a physical plan tree.
+//!
+//! `ExecutionPlan::metrics()` returns metrics for a single plan node only; it
+//! does not include metrics from child operators. These helpers recursively walk
+//! the plan tree so tests can assert on metrics that may move between operators
+//! after optimizer rewrites, such as pushing a `SortExec` below a
+//! `ProjectionExec`.
+
+use datafusion_physical_plan::ExecutionPlan;
+
+/// Returns the total number of spill events recorded by `plan` and all of its
+/// descendants.
+///
+/// Missing `spill_count` metrics are treated as zero.
+pub fn plan_spill_count(plan: &dyn ExecutionPlan) -> usize {
+    let own = plan.metrics().and_then(|m| m.spill_count()).unwrap_or(0);
+
+    own + plan
+        .children()
+        .into_iter()
+        .map(|child| plan_spill_count(child.as_ref()))
+        .sum::<usize>()
+}
+
+/// Returns the total number of spilled bytes recorded by `plan` and all of its
+/// descendants.
+///
+/// Missing `spilled_bytes` metrics are treated as zero.
+pub fn plan_spilled_bytes(plan: &dyn ExecutionPlan) -> usize {
+    let own = plan.metrics().and_then(|m| m.spilled_bytes()).unwrap_or(0);
+
+    own + plan
+        .children()
+        .into_iter()
+        .map(|child| plan_spilled_bytes(child.as_ref()))
+        .sum::<usize>()
+}

datafusion/core/tests/memory_limit/mod.rs

@@ -62,6 +62,8 @@ use async_trait::async_trait;
 use futures::StreamExt;
 use tokio::fs::File;
 
+use crate::helper::plan_metrics::{plan_spill_count, plan_spilled_bytes};
+
 #[cfg(test)]
 #[ctor::ctor(unsafe)]
 fn init() {
@@ -546,8 +548,7 @@ async fn test_external_sort_zero_merge_reservation() {
     let _result = collect(stream).await;
 
     // Ensures the query spilled during execution
-    let metrics = physical_plan.metrics().unwrap();
-    let spill_count = metrics.spill_count().unwrap();
+    let spill_count = plan_spill_count(physical_plan.as_ref());
     assert!(spill_count > 0);
 }
 
@@ -603,9 +604,8 @@ async fn test_sort_skewed_batches_spill() {
 
     // The query must actually spill, otherwise it never reaches the merge path
     // this test is meant to cover.
-    let metrics = physical_plan.metrics().unwrap();
     assert!(
-        metrics.spill_count().unwrap() > 0,
+        plan_spill_count(physical_plan.as_ref()) > 0,
         "expected the sort to spill to disk"
     );
 }
@@ -696,8 +696,8 @@ async fn test_disk_spill_limit_not_reached() -> Result<()> {
         .await
         .expect("Query execution failed");
 
-    let spill_count = plan.metrics().unwrap().spill_count().unwrap();
-    let spilled_bytes = plan.metrics().unwrap().spilled_bytes().unwrap();
+    let spill_count = plan_spill_count(plan.as_ref());
+    let spilled_bytes = plan_spilled_bytes(plan.as_ref());
 
     println!("spill count {spill_count}, spill bytes {spilled_bytes}");
     assert!(spill_count > 0);
@@ -732,8 +732,8 @@ async fn test_spill_file_compressed_with_zstd() -> Result<()> {
         .await
         .expect("Query execution failed");
 
-    let spill_count = plan.metrics().unwrap().spill_count().unwrap();
-    let spilled_bytes = plan.metrics().unwrap().spilled_bytes().unwrap();
+    let spill_count = plan_spill_count(plan.as_ref());
+    let spilled_bytes = plan_spilled_bytes(plan.as_ref());
 
     println!("spill count {spill_count}");
     assert!(spill_count > 0);
@@ -768,8 +768,8 @@ async fn test_spill_file_compressed_with_lz4_frame() -> Result<()> {
         .await
         .expect("Query execution failed");
 
-    let spill_count = plan.metrics().unwrap().spill_count().unwrap();
-    let spilled_bytes = plan.metrics().unwrap().spilled_bytes().unwrap();
+    let spill_count = plan_spill_count(plan.as_ref());
+    let spilled_bytes = plan_spilled_bytes(plan.as_ref());
 
     println!("spill count {spill_count}");
     assert!(spill_count > 0);