ensure we sort on post-REPLACE values

pola-rs · Jun 25, 2024 · 64c8447 · 64c8447
1 parent 55da6c8
commit 64c8447
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 26 deletions.
diff --git a/crates/polars-sql/src/context.rs b/crates/polars-sql/src/context.rs
@@ -724,7 +724,7 @@ impl SQLContext {
         };
 
         lf = if group_by_keys.is_empty() {
-            // Establish final/selected cols, accounting for 'SELECT *' modifiers
+            // Final/selected cols, accounting for 'SELECT *' modifiers
             let mut retained_cols = Vec::with_capacity(projections.len());
             let have_order_by = !query.order_by.is_empty();
 
@@ -747,14 +747,29 @@ impl SQLContext {
                     });
                 }
             }
+
+            // Apply the remaining modifiers and establish the final projection
             if have_order_by {
                 lf = lf.with_columns(projections);
-                if !select_modifiers.rename.is_empty() {
-                    lf = lf.with_columns(select_modifiers.renamed_cols());
-                }
+            }
+            if !select_modifiers.replace.is_empty() {
+                lf = lf.with_columns(&select_modifiers.replace);
+            }
+            if !select_modifiers.rename.is_empty() {
+                lf = lf.with_columns(select_modifiers.renamed_cols());
+            }
+            if have_order_by {
                 lf = self.process_order_by(lf, &query.order_by, Some(&retained_cols))?
             }
-            lf.select(retained_cols)
+            lf = lf.select(retained_cols);
+
+            if !select_modifiers.rename.is_empty() {
+                lf = lf.rename(
+                    select_modifiers.rename.keys(),
+                    select_modifiers.rename.values(),
+                );
+            };
+            lf
         } else {
             lf = self.process_group_by(lf, &group_by_keys, &projections)?;
             lf = self.process_order_by(lf, &query.order_by, None)?;
@@ -793,17 +808,6 @@ impl SQLContext {
             },
             None => lf,
         };
-
-        // Apply final 'SELECT *' REPLACE modifier
-        if !select_modifiers.replace.is_empty() {
-            lf = lf.with_columns(select_modifiers.replace);
-        }
-        if !select_modifiers.rename.is_empty() {
-            lf = lf.rename(
-                select_modifiers.rename.keys(),
-                select_modifiers.rename.values(),
-            );
-        }
         Ok(lf)
     }
 

diff --git a/py-polars/tests/unit/sql/test_wildcard_opts.py b/py-polars/tests/unit/sql/test_wildcard_opts.py
@@ -121,34 +121,31 @@ def test_select_rename_exclude_sort(order_by: str, df: pl.DataFrame) -> None:
 
 
 @pytest.mark.parametrize(
-    ("replacements", "order_by", "check_cols", "expected"),
+    ("replacements", "check_cols", "expected"),
     [
         (
             "(ID // 3 AS ID)",
-            "",
             ["ID"],
-            [(111,), (222,), (333,)],
+            [(333,), (222,), (111,)],
         ),
         (
-            "((City || ':' || City) AS City, ID // 3 AS ID)",
-            "ORDER BY ID DESC",
+            "((City || ':' || City) AS City, ID // -3 AS ID)",
             ["City", "ID"],
             [
-                ("Metropolis:Metropolis", 333),
-                ("Themyscira:Themyscira", 222),
-                ("Gotham:Gotham", 111),
+                ("Gotham:Gotham", -111),
+                ("Themyscira:Themyscira", -222),
+                ("Metropolis:Metropolis", -333),
             ],
         ),
     ],
 )
 def test_select_replace(
     replacements: str,
-    order_by: str,
     check_cols: list[str],
     expected: list[tuple[Any]],
     df: pl.DataFrame,
 ) -> None:
-    res = df.sql(f"SELECT * REPLACE {replacements} FROM self {order_by}")
+    res = df.sql(f"SELECT * REPLACE {replacements} FROM self ORDER BY ID DESC")
 
     assert res.select(check_cols).rows() == expected
     assert res.columns == df.columns