(arrow-datafusion) branch main updated: Change remove from swap to shift in index map (#9049)

[alamb]

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

alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-datafusion.git


The following commit(s) were added to refs/heads/main by this push:
     new 92104a5446 Change remove from swap to shift in index map (#9049)
92104a5446 is described below

commit 92104a54469c8401b799fb3fa855f3f1cb9d8745
Author: Mustafa Akur <106137913+mustafasr...@users.noreply.github.com>
AuthorDate: Mon Jan 29 15:29:47 2024 +0300

    Change remove from swap to shift in index map (#9049)
    
    * Change remove from swap to shift in index map
    
    * Add new tests
---
 .../physical-expr/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 6f3f2aa99b..cd0ae09a92 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<_>>();