Feat(prql): Handle DESC with sort (#3299)

* update sort * add todo * lint * update
tobymao · Apr 11, 2024 · 4790414 · 4790414
1 parent 75e0c69
commit 4790414
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 21 deletions.
diff --git a/sqlglot/dialects/prql.py b/sqlglot/dialects/prql.py
@@ -104,6 +104,17 @@ def _parse_take(self, query: exp.Query) -> t.Optional[exp.Query]:
             num = self._parse_number()  # TODO: TAKE for ranges a..b
             return query.limit(num) if num else None
 
+        def _parse_ordered(
+            self, parse_method: t.Optional[t.Callable] = None
+        ) -> t.Optional[exp.Ordered]:
+            asc = self._match(TokenType.PLUS)
+            desc = self._match(TokenType.DASH) or (asc and False)
+            term = term = super()._parse_ordered(parse_method=parse_method)
+            if term and desc:
+                term.set("desc", True)
+                term.set("nulls_first", False)
+            return term
+
         def _parse_order_by(self, query: exp.Select) -> t.Optional[exp.Query]:
             l_brace = self._match(TokenType.L_BRACE)
             expressions = self._parse_csv(self._parse_ordered)

diff --git a/tests/dialects/test_prql.py b/tests/dialects/test_prql.py
@@ -5,48 +5,56 @@ class TestPRQL(Validator):
     dialect = "prql"
 
     def test_prql(self):
-        self.validate_identity("FROM x", "SELECT * FROM x")
-        self.validate_identity("FROM x DERIVE a + 1", "SELECT *, a + 1 FROM x")
-        self.validate_identity("FROM x DERIVE x = a + 1", "SELECT *, a + 1 AS x FROM x")
-        self.validate_identity("FROM x DERIVE {a + 1}", "SELECT *, a + 1 FROM x")
-        self.validate_identity("FROM x DERIVE {x = a + 1, b}", "SELECT *, a + 1 AS x, b FROM x")
+        self.validate_identity("from x", "SELECT * FROM x")
+        self.validate_identity("from x derive a + 1", "SELECT *, a + 1 FROM x")
+        self.validate_identity("from x derive x = a + 1", "SELECT *, a + 1 AS x FROM x")
+        self.validate_identity("from x derive {a + 1}", "SELECT *, a + 1 FROM x")
+        self.validate_identity("from x derive {x = a + 1, b}", "SELECT *, a + 1 AS x, b FROM x")
         self.validate_identity(
-            "FROM x DERIVE {x = a + 1, b} SELECT {y = x, 2}", "SELECT a + 1 AS y, 2 FROM x"
+            "from x derive {x = a + 1, b} select {y = x, 2}", "SELECT a + 1 AS y, 2 FROM x"
         )
-        self.validate_identity("FROM x TAKE 10", "SELECT * FROM x LIMIT 10")
-        self.validate_identity("FROM x TAKE 10 TAKE 5", "SELECT * FROM x LIMIT 5")
-        self.validate_identity("FROM x FILTER age > 25", "SELECT * FROM x WHERE age > 25")
+        self.validate_identity("from x take 10", "SELECT * FROM x LIMIT 10")
+        self.validate_identity("from x take 10 take 5", "SELECT * FROM x LIMIT 5")
+        self.validate_identity("from x filter age > 25", "SELECT * FROM x WHERE age > 25")
         self.validate_identity(
-            "FROM x DERIVE {x = a + 1, b} FILTER age > 25",
+            "from x derive {x = a + 1, b} filter age > 25",
             "SELECT *, a + 1 AS x, b FROM x WHERE age > 25",
         )
-        self.validate_identity("FROM x FILTER dept != 'IT'", "SELECT * FROM x WHERE dept <> 'IT'")
+        self.validate_identity("from x filter dept != 'IT'", "SELECT * FROM x WHERE dept <> 'IT'")
         self.validate_identity(
-            "FROM x FILTER p == 'product' SELECT { a, b }", "SELECT a, b FROM x WHERE p = 'product'"
+            "from x filter p == 'product' select { a, b }", "SELECT a, b FROM x WHERE p = 'product'"
         )
         self.validate_identity(
-            "FROM x FILTER age > 25 FILTER age < 27", "SELECT * FROM x WHERE age > 25 AND age < 27"
+            "from x filter age > 25 filter age < 27", "SELECT * FROM x WHERE age > 25 AND age < 27"
         )
         self.validate_identity(
-            "FROM x FILTER (age > 25 && age < 27)", "SELECT * FROM x WHERE (age > 25 AND age < 27)"
+            "from x filter (age > 25 && age < 27)", "SELECT * FROM x WHERE (age > 25 AND age < 27)"
         )
         self.validate_identity(
-            "FROM x FILTER (age > 25 || age < 27)", "SELECT * FROM x WHERE (age > 25 OR age < 27)"
+            "from x filter (age > 25 || age < 27)", "SELECT * FROM x WHERE (age > 25 OR age < 27)"
         )
         self.validate_identity(
-            "FROM x FILTER (age > 25 || age < 22) FILTER age > 26 FILTER age < 27",
+            "from x filter (age > 25 || age < 22) filter age > 26 filter age < 27",
             "SELECT * FROM x WHERE ((age > 25 OR age < 22) AND age > 26) AND age < 27",
         )
         self.validate_identity(
-            "FROM x SORT age",
+            "from x sort age",
             "SELECT * FROM x ORDER BY age",
         )
         self.validate_identity(
-            "FROM x SORT {age, name}",
+            "from x sort {-age}",
+            "SELECT * FROM x ORDER BY age DESC",
+        )
+        self.validate_identity(
+            "from x sort {age, name}",
             "SELECT * FROM x ORDER BY age, name",
         )
-        self.validate_identity("FROM x APPEND y", "SELECT * FROM x UNION ALL SELECT * FROM y")
-        self.validate_identity("FROM x REMOVE y", "SELECT * FROM x EXCEPT ALL SELECT * FROM y")
         self.validate_identity(
-            "FROM x INTERSECT y", "SELECT * FROM x INTERSECT ALL SELECT * FROM y"
+            "from x sort {-age, +name}",
+            "SELECT * FROM x ORDER BY age DESC, name",
+        )
+        self.validate_identity("from x append y", "SELECT * FROM x UNION ALL SELECT * FROM y")
+        self.validate_identity("from x remove y", "SELECT * FROM x EXCEPT ALL SELECT * FROM y")
+        self.validate_identity(
+            "from x intersect y", "SELECT * FROM x INTERSECT ALL SELECT * FROM y"
         )