Feat(prql): handle NULL (#3331)

fool1280 · georgesittas · web-flow · commit 31744b26ed97 · 2024-04-17T09:07:56.000+03:00
* handle null in prql

* lint

* update

* Simplify

* elif -&gt; if

* Duplicate test

---------

Co-authored-by: Jo &lt;46752250+georgesittas@users.noreply.github.com&gt;
diff --git a/sqlglot/dialects/prql.py b/sqlglot/dialects/prql.py
@@ -55,6 +55,20 @@ class Parser(parser.Parser):
             "SORT": lambda self, query: self._parse_order_by(query),
         }
 
+        def _parse_equality(self) -> t.Optional[exp.Expression]:
+            eq = self._parse_tokens(self._parse_comparison, self.EQUALITY)
+            if not isinstance(eq, (exp.EQ, exp.NEQ)):
+                return eq
+
+            # https://prql-lang.org/book/reference/spec/null.html
+            if isinstance(eq.expression, exp.Null):
+                is_exp = exp.Is(this=eq.this, expression=eq.expression)
+                return is_exp if isinstance(eq, exp.EQ) else exp.Not(this=is_exp)
+            if isinstance(eq.this, exp.Null):
+                is_exp = exp.Is(this=eq.expression, expression=eq.this)
+                return is_exp if isinstance(eq, exp.EQ) else exp.Not(this=is_exp)
+            return eq
+
         def _parse_statement(self) -> t.Optional[exp.Expression]:
             expression = self._parse_expression()
             expression = expression if expression else self._parse_query()
diff --git a/tests/dialects/test_prql.py b/tests/dialects/test_prql.py
@@ -58,3 +58,11 @@ def test_prql(self):
         self.validate_identity(
             "from x intersect y", "SELECT * FROM x INTERSECT ALL SELECT * FROM y"
         )
+        self.validate_identity(
+            "from x filter a == null filter null != b",
+            "SELECT * FROM x WHERE a IS NULL AND NOT b IS NULL",
+        )
+        self.validate_identity(
+            "from x filter (a > 1 || null != b || c != null)",
+            "SELECT * FROM x WHERE (a > 1 OR NOT b IS NULL OR NOT c IS NULL)",
+        )
