Minor changes

alamb · Jan 7, 2025 · 16ef66c · 16ef66c
1 parent d2fe8a4
commit 16ef66c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/datafusion/physical-expr-common/Cargo.toml b/datafusion/physical-expr-common/Cargo.toml
@@ -42,3 +42,4 @@ datafusion-common = { workspace = true, default-features = true }
 datafusion-expr-common = { workspace = true }
 hashbrown = { workspace = true }
 itertools = { workspace = true }
+indexmap = { workspace = true }
diff --git a/datafusion/physical-expr-common/src/sort_expr.rs b/datafusion/physical-expr-common/src/sort_expr.rs
@@ -30,7 +30,8 @@ use arrow::datatypes::Schema;
 use arrow::record_batch::RecordBatch;
 use datafusion_common::Result;
 use datafusion_expr_common::columnar_value::ColumnarValue;
-use itertools::Itertools;
+use itertools::{izip, Itertools};
+use indexmap::IndexSet;
 
 /// Represents Sort operation for a column in a RecordBatch
 ///
@@ -568,11 +569,19 @@ impl LexRequirement {
     /// `vec![a Some(ASC)]`.
     pub fn collapse(self) -> Self {
         let mut output = Vec::<PhysicalSortRequirement>::new();
+        let mut exprs = IndexSet::new();
+        let mut reqs = vec![];
         for item in self {
-            if !output.iter().any(|req| req.expr.eq(&item.expr)) {
-                output.push(item);
+            let PhysicalSortRequirement { expr, options: req } = item;
+            // new insertion
+            if exprs.insert(expr) {
+                reqs.push(req);
             }
         }
+        debug_assert_eq!(reqs.len(), exprs.len());
+        for (expr, req) in izip!(exprs, reqs) {
+            output.push(PhysicalSortRequirement::new(expr, req));
+        }
         LexRequirement::new(output)
     }
 }