Support None comparisons for null expressions

Author: zeel2104

python/datafusion/expr.py

@@ -483,6 +483,8 @@ def __eq__(self, rhs: object) -> Expr:
 
         Accepts either an expression or any valid PyArrow scalar literal value.
         """
+        if rhs is None:
+            return self.is_null()
         if not isinstance(rhs, Expr):
             rhs = Expr.literal(rhs)
         return Expr(self.expr.__eq__(rhs.expr))
@@ -492,6 +494,8 @@ def __ne__(self, rhs: object) -> Expr:
 
         Accepts either an expression or any valid PyArrow scalar literal value.
         """
+        if rhs is None:
+            return self.is_not_null()
         if not isinstance(rhs, Expr):
             rhs = Expr.literal(rhs)
         return Expr(self.expr.__ne__(rhs.expr))

python/tests/test_expr.py

@@ -173,6 +173,24 @@ def test_relational_expr(test_ctx):
     assert df.filter(col("a") == "beta").count() == 0
 
 
+def test_relational_expr_none_uses_null_predicates():
+    ctx = SessionContext()
+
+    batch = pa.RecordBatch.from_arrays(
+        [
+            pa.array([1, 2, None]),
+            pa.array(["alpha", None, "gamma"], type=pa.string_view()),
+        ],
+        names=["a", "b"],
+    )
+    df = ctx.create_dataframe([[batch]], name="batch_with_nulls")
+
+    assert df.filter(col("a") == None).count() == 1  # noqa: E711
+    assert df.filter(col("a") != None).count() == 2  # noqa: E711
+    assert df.filter(col("b") == None).count() == 1  # noqa: E711
+    assert df.filter(col("b") != None).count() == 2  # noqa: E711
+
+
 def test_expr_to_variant():
     # Taken from https://github.com/apache/datafusion-python/issues/781
     from datafusion import SessionContext