(auron) branch master updated: [AURON #2497] Fix SortExec projected schema and sort key mapping (#2498)

[sl...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

slfan1989 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/auron.git


The following commit(s) were added to refs/heads/master by this push:
     new 04eeb90d2 [AURON #2497] Fix SortExec projected schema and sort key mapping (#2498)
04eeb90d2 is described below

commit 04eeb90d2911465eb350253eacc281a16bb5f37f
Author: linfeng <33...@users.noreply.github.com>
AuthorDate: Sat Sep 5 15:48:29 2026 +0800

    [AURON #2497] Fix SortExec projected schema and sort key mapping (#2498)
    
    # Which issue does this PR close?
    
    Closes #2497
    
    # Rationale for this change
    
    `SortExec` used the full input schema during projected execution and
    consumed original sort-key indices against the projected schema. Subset
    or reordered projections could therefore produce an incorrect schema or
    incorrectly mapped columns.
    
    # What changes are included in this PR?
    
    - Use the projected output schema in both projected execution paths.
    - Remap sort-key indices against every matching position in the
    projection.
    - Add regression tests for reordered SortMergeJoin projections and
    strict single-column SortExec projections.
    
    # Are there any user-facing changes?
    
    No API changes. Affected projected sort and sort-merge join plans now
    return columns with the correct schema
    
    # How was this patch tested?
    
    Passed unit tests
    
    # Was this patch authored or co-authored using generative AI tooling?
    - [x] Yes
    - [ ] No
    
    If yes, include: `Generated-by: gpt5`
    
    ASF guidance: https://www.apache.org/legal/generative-tooling.html
    
    Co-authored-by: Shilun Fan <sl...@apache.org>
    Signed-off-by: Shilun Fan <sl...@apache.org>
---
 .../datafusion-ext-plans/src/joins/test.rs         | 68 +++++++++++++++++++++-
 .../datafusion-ext-plans/src/sort_exec.rs          | 53 ++++++++++++++++-
 2 files changed, 117 insertions(+), 4 deletions(-)

diff --git a/native-engine/datafusion-ext-plans/src/joins/test.rs b/native-engine/datafusion-ext-plans/src/joins/test.rs
index 70d0d2fca..9b7e56db1 100644
--- a/native-engine/datafusion-ext-plans/src/joins/test.rs
+++ b/native-engine/datafusion-ext-plans/src/joins/test.rs
@@ -30,7 +30,10 @@ mod tests {
         assert_batches_sorted_eq,
         common::{JoinSide, Result},
         logical_expr::Operator,
-        physical_expr::expressions::{BinaryExpr, Column},
+        physical_expr::{
+            PhysicalSortExpr,
+            expressions::{BinaryExpr, Column},
+        },
         physical_plan::{ExecutionPlan, common, joins::utils::*, test::TestMemoryExec},
         prelude::{SessionConfig, SessionContext},
     };
@@ -38,10 +41,12 @@ mod tests {
     use crate::{
         broadcast_join_build_hash_map_exec::BroadcastJoinBuildHashMapExec,
         broadcast_join_exec::BroadcastJoinExec,
+        common::column_pruning::ExecuteWithColumnPruning,
         joins::{
             ColumnIndex, JoinFilter,
             join_utils::{JoinType, JoinType::*},
         },
+        sort_exec::SortExec,
         sort_merge_join_exec::SortMergeJoinExec,
     };
 
@@ -483,6 +488,67 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+    async fn sort_merge_join_projects_non_key_columns_from_sorted_inputs() -> Result<()> {
+        MemManager::init(1000000);
+        let session_ctx = SessionContext::new();
+        let task_ctx = session_ctx.task_ctx();
+        let left_input = build_table(
+            ("left_key", &vec![2, 1]),
+            ("left_unused", &vec![200, 100]),
+            ("left_value", &vec![20, 10]),
+        )?;
+        let right_input = build_table(
+            ("right_key", &vec![1, 2]),
+            ("right_unused", &vec![1000, 2000]),
+            ("right_value", &vec![100, 200]),
+        )?;
+        let left_key = Arc::new(Column::new("left_key", 0));
+        let right_key = Arc::new(Column::new("right_key", 0));
+        let left: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(
+            left_input,
+            vec![PhysicalSortExpr {
+                expr: left_key.clone(),
+                options: SortOptions::default(),
+            }],
+            None,
+            0,
+        ));
+        let right: Arc<dyn ExecutionPlan> = Arc::new(SortExec::new(
+            right_input,
+            vec![PhysicalSortExpr {
+                expr: right_key.clone(),
+                options: SortOptions::default(),
+            }],
+            None,
+            0,
+        ));
+        let schema = build_join_schema_for_test(&left.schema(), &right.schema(), Inner)?;
+        let join = SortMergeJoinExec::try_new(
+            schema,
+            left,
+            right,
+            vec![(left_key, right_key)],
+            Inner,
+            None,
+            vec![SortOptions::default()],
+        )?;
+
+        let stream = join.execute_projected(0, task_ctx, &[2, 1, 0, 5, 4, 3])?;
+        let batches = common::collect(stream).await?;
+        let expected = vec![
+            "+------------+-------------+----------+-------------+--------------+-----------+",
+            "| left_value | left_unused | left_key | right_value | right_unused | right_key |",
+            "+------------+-------------+----------+-------------+--------------+-----------+",
+            "| 10         | 100         | 1        | 100         | 1000         | 1         |",
+            "| 20         | 200         | 2        | 200         | 2000         | 2         |",
+            "+------------+-------------+----------+-------------+--------------+-----------+",
+        ];
+        assert_batches_sorted_eq!(expected, &batches);
+
+        Ok(())
+    }
+
     #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
     async fn join_inner_two() -> Result<()> {
         for test_type in ALL_TEST_TYPE {
diff --git a/native-engine/datafusion-ext-plans/src/sort_exec.rs b/native-engine/datafusion-ext-plans/src/sort_exec.rs
index 48a8b9237..6b7bb2b80 100644
--- a/native-engine/datafusion-ext-plans/src/sort_exec.rs
+++ b/native-engine/datafusion-ext-plans/src/sort_exec.rs
@@ -261,8 +261,9 @@ impl SortExec {
         context: Arc<TaskContext>,
         projection: &[usize],
     ) -> Result<SendableRecordBatchStream> {
+        let output_schema = Arc::new(self.schema().project(projection)?);
         let exec_ctx =
-            ExecutionContext::new(context.clone(), partition, self.schema(), &self.metrics);
+            ExecutionContext::new(context.clone(), partition, output_schema, &self.metrics);
         let sorter = self.init_sorter(exec_ctx.clone(), projection)?;
 
         let elapsed_compute = exec_ctx.baseline_metrics().elapsed_compute().clone();
@@ -307,8 +308,9 @@ impl SortExec {
         context: Arc<TaskContext>,
         projection: &[usize],
     ) -> Result<SendableRecordBatchWithKeyRowsStream> {
+        let output_schema = Arc::new(self.schema().project(projection)?);
         let exec_ctx =
-            ExecutionContext::new(context.clone(), partition, self.schema(), &self.metrics);
+            ExecutionContext::new(context.clone(), partition, output_schema, &self.metrics);
         let sorter = self.init_sorter(exec_ctx.clone(), projection)?;
 
         let elapsed_compute = exec_ctx.baseline_metrics().elapsed_compute().clone();
@@ -1212,7 +1214,12 @@ impl PruneSortKeysFromBatch {
         let mut relation = vec![];
         for (expr_idx, expr) in exprs.iter().enumerate() {
             if let Some(col) = expr.expr.as_any().downcast_ref::<Column>() {
-                relation.push((expr_idx, col.index()));
+                // Relation indices are consumed against the projected schema below.
+                relation.extend(input_projection.iter().enumerate().filter_map(
+                    |(projected_idx, &input_idx)| {
+                        (input_idx == col.index()).then_some((expr_idx, projected_idx))
+                    },
+                ));
             }
         }
 
@@ -1574,6 +1581,7 @@ mod test {
         physical_plan::{ExecutionPlan, common, test::TestMemoryExec},
         prelude::SessionContext,
     };
+    use futures::TryStreamExt;
 
     use super::common_prefix_len;
     use crate::sort_exec::SortExec;
@@ -1674,6 +1682,45 @@ mod test {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_sort_projected_with_key_rows_excludes_sort_key() -> Result<()> {
+        MemManager::init(100);
+        let task_ctx = SessionContext::new().task_ctx();
+        let input = build_table(
+            ("key", &vec![3, 1, 2]),
+            ("value", &vec![30, 10, 20]),
+            ("unused", &vec![300, 100, 200]),
+        )?;
+        let sort = SortExec::new(
+            input,
+            vec![PhysicalSortExpr {
+                expr: Arc::new(Column::new("key", 0)),
+                options: SortOptions::default(),
+            }],
+            None,
+            0,
+        );
+
+        let output = sort.execute_projected_with_key_rows(0, task_ctx, &[1])?;
+        assert_eq!(output.schema().fields().len(), 1);
+        assert_eq!(output.schema().field(0).name(), "value");
+        let batches = output
+            .map_ok(|batch| batch.batch)
+            .try_collect::<Vec<_>>()
+            .await?;
+        let expected = vec![
+            "+-------+",
+            "| value |",
+            "+-------+",
+            "| 10    |",
+            "| 20    |",
+            "| 30    |",
+            "+-------+",
+        ];
+        assert_batches_eq!(expected, &batches);
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_top_k_with_only_sort_column() -> Result<()> {
         MemManager::init(100);