Repository: prql/prql
Pull request #853 (2022-07-25 15:20:48)

Unique row number aliases
Uses the row_number_id to add a unique number to _rn definition to prevent clashes in take 

max-sixty: Perfect @charlie-sanders ! Thanks!

Patch:
diff --git a/prql-compiler/src/semantic/name_resolver.rs b/prql-compiler/src/semantic/name_resolver.rs
index ca7e20bc..9584438b 100644
--- a/prql-compiler/src/semantic/name_resolver.rs
+++ b/prql-compiler/src/semantic/name_resolver.rs
@@ -769,7 +769,7 @@ mod tests {
               PARTITION BY last_name
               ORDER BY
                 first_name
-            ) AS _rn
+            ) AS _rn_82
           FROM
             employees
         )
@@ -778,7 +778,7 @@ mod tests {
         FROM
           table_0
         WHERE
-          _rn <= 1
+          _rn_82 <= 1
         "###);
 
         let res = parse(
diff --git a/prql-compiler/src/sql/distinct.rs b/prql-compiler/src/sql/distinct.rs
index 1df02e89..5f78334a 100644
--- a/prql-compiler/src/sql/distinct.rs
+++ b/prql-compiler/src/sql/distinct.rs
@@ -83,8 +83,8 @@ impl<'a> DistinctMaker<'a> {
         let decl = Declaration::Expression(Box::new(Item::Windowed(windowed).into()));
         let row_number_id = self.context.declarations.push(decl, None);
 
-        // name it _rn
-        let mut ident = Node::from(Item::Ident("_rn".to_string()));
+        // name it _rn_X where X is the row_number_id
+        let mut ident = Node::from(Item::Ident(format!("_rn_{}", row_number_id)));
         ident.declared_at = Some(row_number_id);
 
         // add the two transforms
diff --git a/prql-compiler/src/sql/translator.rs b/prql-compiler/src/sql/translator.rs
index 70f3fb44..61e07d97 100644
--- a/prql-compiler/src/sql/translator.rs
+++ b/prql-compiler/src/sql/translator.rs
@@ -2197,7 +2197,7 @@ take 20
         WITH table_0 AS (
           SELECT
             employees.*,
-            ROW_NUMBER() OVER (PARTITION BY department) AS _rn
+            ROW_NUMBER() OVER (PARTITION BY department) AS _rn_81
           FROM
             employees
         )
@@ -2206,7 +2206,7 @@ take 20
         FROM
           table_0
         WHERE
-          _rn <= 3
+          _rn_81 <= 3
         "###);
 
         assert_display_snapshot!((resolve_and_translate(parse(r###"
@@ -2221,7 +2221,7 @@ take 20
               PARTITION BY department
               ORDER BY
                 salary
-            ) AS _rn
+            ) AS _rn_82
           FROM
             employees
         )
@@ -2230,7 +2230,7 @@ take 20
         FROM
           table_0
         WHERE
-          _rn BETWEEN 2
+          _rn_82 BETWEEN 2
           AND 3
         "###);
     }
@@ -2311,6 +2311,43 @@ take 20
         "###);
     }
 
+    #[test]
+    fn test_rn_ids_are_unique() {
+        assert_display_snapshot!((resolve_and_translate(parse(r###"
+        from y_orig
+        group [y_id] (
+          take 2 # take 1 uses `distinct` instead of partitioning, which might be a separate bug
+        )
+        group [x_id] (
+          take 3
+        )
+        "###,
+        ).unwrap()).unwrap()), @r###"
+        WITH table_0 AS (
+          SELECT
+            y_orig.*,
+            ROW_NUMBER() OVER (PARTITION BY y_id) AS _rn_82
+          FROM
+            y_orig
+        ),
+        table_1 AS (
+          SELECT
+            table_0.*,
+            ROW_NUMBER() OVER (PARTITION BY x_id) AS _rn_83
+          FROM
+            table_0
+          WHERE
+            _rn_82 <= 2
+        )
+        SELECT
+          table_1.*
+        FROM
+          table_1
+        WHERE
+          _rn_83 <= 3
+        "###);
+    }
+
     #[test]
     fn test_quoting() -> Result<()> {
         // GH-#822


Test patch:
diff --git a/book/tests/snapshots/snapshot__run_reference_prql@distinct-2.prql.snap b/book/tests/snapshots/snapshot__run_reference_prql@distinct-2.prql.snap
index 42528663..485dd450 100644
--- a/book/tests/snapshots/snapshot__run_reference_prql@distinct-2.prql.snap
+++ b/book/tests/snapshots/snapshot__run_reference_prql@distinct-2.prql.snap
@@ -1,7 +1,6 @@
 ---
 source: book/tests/snapshot.rs
 expression: sql
-input_file: book/tests/prql/language-features/distinct-2.prql
 ---
 WITH table_0 AS (
   SELECT
@@ -10,7 +9,7 @@ WITH table_0 AS (
       PARTITION BY department
       ORDER BY
         age
-    ) AS _rn
+    ) AS _rn_82
   FROM
     employees
 )
@@ -19,4 +18,4 @@ SELECT
 FROM
   table_0
 WHERE
-  _rn <= 1
+  _rn_82 <= 1
diff --git a/book/tests/snapshots/snapshot__run_reference_prql@group-2.prql.snap b/book/tests/snapshots/snapshot__run_reference_prql@group-2.prql.snap
index 960823cd..f4e996bf 100644
--- a/book/tests/snapshots/snapshot__run_reference_prql@group-2.prql.snap
+++ b/book/tests/snapshots/snapshot__run_reference_prql@group-2.prql.snap
@@ -1,7 +1,6 @@
 ---
 source: book/tests/snapshot.rs
 expression: sql
-input_file: book/tests/prql/transforms/group-2.prql
 ---
 WITH table_0 AS (
   SELECT
@@ -10,7 +9,7 @@ WITH table_0 AS (
       PARTITION BY role
       ORDER BY
         join_date
-    ) AS _rn
+    ) AS _rn_82
   FROM
     employees
 )
@@ -19,4 +18,4 @@ SELECT
 FROM
   table_0
 WHERE
-  _rn <= 1
+  _rn_82 <= 1