From 92104a54469c8401b799fb3fa855f3f1cb9d8745 Mon Sep 17 00:00:00 2001
From: Mustafa Akur <106137913+mustafasrepo@users.noreply.github.com>
Date: Mon, 29 Jan 2024 15:29:47 +0300
Subject: [PATCH] Change remove from swap to shift in index map (#9049)

* Change remove from swap to shift in index map

* Add new tests
---
 .../src/equivalence/properties.rs             | 38 ++++++++++++++++++-
 1 file changed, 36 insertions(+), 2 deletions(-)

diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs
index 6f3f2aa99b27..cd0ae09a92bb 100644
--- a/datafusion/physical-expr/src/equivalence/properties.rs
+++ b/datafusion/physical-expr/src/equivalence/properties.rs
@@ -730,7 +730,7 @@ impl EquivalenceProperties {
             for (PhysicalSortExpr { expr, .. }, idx) in &ordered_exprs {
                 eq_properties =
                     eq_properties.add_constants(std::iter::once(expr.clone()));
-                search_indices.swap_remove(idx);
+                search_indices.shift_remove(idx);
             }
             // Add new ordered section to the state.
             result.extend(ordered_exprs);
@@ -1779,6 +1779,7 @@ mod tests {
         let col_c = &col("c", &test_schema)?;
         let col_d = &col("d", &test_schema)?;
         let col_e = &col("e", &test_schema)?;
+        let col_f = &col("f", &test_schema)?;
         let col_h = &col("h", &test_schema)?;
         // a + d
         let a_plus_d = Arc::new(BinaryExpr::new(
@@ -1795,7 +1796,7 @@ mod tests {
             descending: true,
             nulls_first: true,
         };
-        // [d ASC, h ASC] also satisfies schema.
+        // [d ASC, h DESC] also satisfies schema.
         eq_properties.add_new_orderings([vec![
             PhysicalSortExpr {
                 expr: col_d.clone(),
@@ -1836,6 +1837,39 @@ mod tests {
                 vec![col_c, col_e],
                 vec![(col_c, option_asc), (col_e, option_desc)],
             ),
+            // TEST CASE 7
+            (
+                vec![col_d, col_h, col_e, col_f, col_b],
+                vec![
+                    (col_d, option_asc),
+                    (col_e, option_desc),
+                    (col_h, option_desc),
+                    (col_f, option_asc),
+                    (col_b, option_asc),
+                ],
+            ),
+            // TEST CASE 8
+            (
+                vec![col_e, col_d, col_h, col_f, col_b],
+                vec![
+                    (col_e, option_desc),
+                    (col_d, option_asc),
+                    (col_h, option_desc),
+                    (col_f, option_asc),
+                    (col_b, option_asc),
+                ],
+            ),
+            // TEST CASE 9
+            (
+                vec![col_e, col_d, col_b, col_h, col_f],
+                vec![
+                    (col_e, option_desc),
+                    (col_d, option_asc),
+                    (col_b, option_asc),
+                    (col_h, option_desc),
+                    (col_f, option_asc),
+                ],
+            ),
         ];
         for (exprs, expected) in test_cases {
             let exprs = exprs.into_iter().cloned().collect::<Vec<_>>();