Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions python/datafusion/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand All @@ -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))
Expand Down
18 changes: 18 additions & 0 deletions python/tests/test_expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think you can just update the test_relational_expr to have some null values and incorporate this into the existing test.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated the tests to fold the None comparison coverage into test_relational_expr and pushed the change. Thanks for the suggestion.



def test_expr_to_variant():
# Taken from https://github.com/apache/datafusion-python/issues/781
from datafusion import SessionContext
Expand Down
Loading