diff --git a/vortex-python/python/vortex/polars_.py b/vortex-python/python/vortex/polars_.py index 43d362fa202..e976583fdd0 100644 --- a/vortex-python/python/vortex/polars_.py +++ b/vortex-python/python/vortex/polars_.py @@ -159,6 +159,10 @@ def _polars_to_vortex(expr: dict[str, Any]) -> ve.Expr: expr = expr["Function"] _inputs = [_polars_to_vortex(e) for e in expr["input"]] + if expr["function"] == "FillNull": + if "Literal" in expr["input"][1]: + return ve.fill_null(_inputs[0], _inputs[1]) + return ve.zip_(ve.is_null(_inputs[0]), _inputs[1], _inputs[0]) fn = expr["function"] if "Boolean" in fn: fn = fn["Boolean"] diff --git a/vortex-python/test/test_polars_.py b/vortex-python/test/test_polars_.py index 1f42bb2160e..c89713a5c25 100644 --- a/vortex-python/test/test_polars_.py +++ b/vortex-python/test/test_polars_.py @@ -204,3 +204,41 @@ def test_polars_boolean_not(tmp_path): actual = vx.open(str(path)).to_polars().filter(expr).collect() assert_frame_equal(actual, expected_frame) assert actual["id"].to_list() == [2] + + +def test_polars_fill_null(tmp_path): + frame = pl.DataFrame({"id": [0, 1, 2, 3], "x": [1, None, None, 5], "y": [4, 5, None, 6]}) + expr = pl.col("x").fill_null(pl.col("y")) == 5 + path = tmp_path / "fill_null.vortex" + vx.io.write(vx.array(frame.to_arrow()), str(path)) + expected_frame = frame.lazy().filter(expr).collect() + actual = vx.open(str(path)).to_polars().filter(expr).collect() + assert_frame_equal(actual, expected_frame) + assert actual["id"].to_list() == [1, 3] + + +@pytest.mark.parametrize("fill_value, expected_ids", [(5, [1, 2, 3]), (0, [1, 2]), (-1, [1, 2])]) +def test_polars_fill_null_literal(tmp_path, fill_value, expected_ids): + frame = pl.DataFrame({"id": [0, 1, 2, 3], "x": [1, None, None, 5]}) + filled = pl.col("x").fill_null(fill_value) + assert polars_to_vortex(filled).serialize() == ve.fill_null(ve.column("x"), fill_value).serialize() + predicate = filled == fill_value + path = tmp_path / "fill_null_literal.vortex" + vx.io.write(vx.array(frame.to_arrow()), str(path)) + expected_frame = frame.lazy().filter(predicate).collect() + actual = vx.open(str(path)).to_polars().filter(predicate).collect() + assert_frame_equal(actual, expected_frame) + assert actual["id"].to_list() == expected_ids + + +def test_polars_fill_null_string_literal(tmp_path): + frame = pl.DataFrame({"id": [0, 1, 2], "x": ["a", None, "c"]}) + filled = pl.col("x").fill_null("z") + assert polars_to_vortex(filled).serialize() == ve.fill_null(ve.column("x"), "z").serialize() + expr = filled == "z" + path = tmp_path / "fill_null_string.vortex" + vx.io.write(vx.array(frame.to_arrow()), str(path)) + expected_frame = frame.lazy().filter(expr).collect() + actual = vx.open(str(path)).to_polars().filter(expr).collect() + assert_frame_equal(actual, expected_frame) + assert actual["id"].to_list() == [1]