From be27a4a51ddeed023eb8471a9ff9666ba6352839 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Thu, 8 Oct 2026 17:47:52 +0000 Subject: [PATCH 1/2] fix(python): translate Polars struct field access Signed-off-by: Daniel King --- vortex-python/python/vortex/polars_.py | 2 + vortex-python/test/test_polars_.py | 78 ++++++++++++++++++++++++++ 2 files changed, 80 insertions(+) diff --git a/vortex-python/python/vortex/polars_.py b/vortex-python/python/vortex/polars_.py index b7536aa0173..bfa703fa8ae 100644 --- a/vortex-python/python/vortex/polars_.py +++ b/vortex-python/python/vortex/polars_.py @@ -158,6 +158,8 @@ def _polars_to_vortex(expr: dict[str, Any]) -> ve.Expr: return ve.fill_null(_inputs[0], _inputs[1]) return ve.zip_(ve.is_null(_inputs[0]), _inputs[1], _inputs[0]) fn = expr["function"] + if isinstance(fn, dict) and "FieldByName" in fn.get("StructExpr", {}): + return ve.get_item(fn["StructExpr"]["FieldByName"], _inputs[0]) if isinstance(fn, dict) and "ReplaceTimeZone" in fn.get("TemporalExpr", {}): time_zone, non_existent = fn["TemporalExpr"]["ReplaceTimeZone"] if isinstance(time_zone, dict): diff --git a/vortex-python/test/test_polars_.py b/vortex-python/test/test_polars_.py index dab406729cf..55c0cda5326 100644 --- a/vortex-python/test/test_polars_.py +++ b/vortex-python/test/test_polars_.py @@ -5,6 +5,7 @@ import os from datetime import UTC, date, datetime, time from decimal import Decimal +from typing import Any, cast from zoneinfo import ZoneInfo import polars as pl @@ -369,3 +370,80 @@ def _time_zone_scan(tmp_path, values, policy=None) -> tuple[pl.LazyFrame, pl.Laz table = pa.table(columns) vx.io.write(vx.array(table), str(path)) return pl.DataFrame(table).lazy(), vx.open(str(path)).to_polars() + + +def test_polars_struct_field(tmp_path): + frame = pl.DataFrame({"id": [0, 1, 2, 3], "x": [{"a": 1}, None, {"a": 3}, {"a": None}]}) + expr = pl.col("x").struct.field("a") >= 2 + path = tmp_path / "struct_field.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() == [2] + + +def _unmasked_field(array: pa.Array[Any], *names: str) -> pa.Array[Any]: + for name in names: + array = cast(pa.StructArray, array).field(name) + return array + + +def _assert_struct_field_scan(tmp_path, array, predicate, expected): + table = pa.table({"id": list(range(len(array))), "x": array}) + polars_result = pl.DataFrame(table).lazy().filter(predicate).collect() + path = tmp_path / "struct_parent_nulls.vortex" + vx.io.write(vx.array(table), str(path)) + vortex_result = vx.open(str(path)).to_polars().filter(predicate).collect() + assert_frame_equal(polars_result, expected) + assert_frame_equal(vortex_result, expected) + + +def test_polars_struct_field_null_parent(tmp_path): + leaf = pa.array([10, 20]) + parent = pa.StructArray.from_arrays([leaf], names=["value"], mask=pa.array([True, False])) + assert parent.field("value").to_pylist() == [10, 20] + predicate = pl.col("x").struct.field("value") >= 0 + expected = pl.DataFrame({"id": [1], "x": [{"value": 20}]}) + _assert_struct_field_scan(tmp_path, parent, predicate, expected) + + +def test_polars_struct_field_null_inner_parent(tmp_path): + leaf = pa.array([10, 20]) + inner = pa.StructArray.from_arrays([leaf], names=["value"], mask=pa.array([True, False])) + outer = pa.StructArray.from_arrays([inner], names=["child"]) + assert _unmasked_field(outer, "child", "value").to_pylist() == [10, 20] + predicate = pl.col("x").struct.field("child").struct.field("value") >= 0 + expected = pl.DataFrame({"id": [1], "x": [{"child": {"value": 20}}]}) + _assert_struct_field_scan(tmp_path, outer, predicate, expected) + + +def test_polars_struct_field_null_outer_parent(tmp_path): + leaf = pa.array([10, 20]) + inner = pa.StructArray.from_arrays([leaf], names=["value"]) + outer = pa.StructArray.from_arrays([inner], names=["child"], mask=pa.array([True, False])) + assert _unmasked_field(outer, "child", "value").to_pylist() == [10, 20] + predicate = pl.col("x").struct.field("child").struct.field("value") >= 0 + expected = pl.DataFrame({"id": [1], "x": [{"child": {"value": 20}}]}) + _assert_struct_field_scan(tmp_path, outer, predicate, expected) + + +def test_polars_struct_field_null_middle_parent(tmp_path): + leaf = pa.array([10, 20]) + inner = pa.StructArray.from_arrays([leaf], names=["value"]) + middle = pa.StructArray.from_arrays([inner], names=["child"], mask=pa.array([True, False])) + outer = pa.StructArray.from_arrays([middle], names=["child"]) + assert _unmasked_field(outer, "child", "child", "value").to_pylist() == [10, 20] + predicate = pl.col("x").struct.field("child").struct.field("child").struct.field("value") >= 0 + expected = pl.DataFrame({"id": [1], "x": [{"child": {"child": {"value": 20}}}]}) + _assert_struct_field_scan(tmp_path, outer, predicate, expected) + + +def test_polars_struct_field_null_parents_and_leaf(tmp_path): + leaf = pa.array([10, 20, 30, None, 50]) + inner = pa.StructArray.from_arrays([leaf], names=["value"], mask=pa.array([True, False, True, False, False])) + outer = pa.StructArray.from_arrays([inner], names=["child"], mask=pa.array([False, True, True, False, False])) + assert _unmasked_field(outer, "child", "value").to_pylist() == [10, 20, 30, None, 50] + predicate = pl.col("x").struct.field("child").struct.field("value") >= 0 + expected = pl.DataFrame({"id": [4], "x": [{"child": {"value": 50}}]}) + _assert_struct_field_scan(tmp_path, outer, predicate, expected) From 2b93c82f7e2ab9685b9d87bfdbbbe6a35654b852 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Mon, 12 Oct 2026 15:31:21 -0400 Subject: [PATCH 2/2] fix(python): postpone PyArrow test annotations Signed-off-by: Daniel King --- vortex-python/test/test_polars_.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vortex-python/test/test_polars_.py b/vortex-python/test/test_polars_.py index 55c0cda5326..1bbf36cf2f0 100644 --- a/vortex-python/test/test_polars_.py +++ b/vortex-python/test/test_polars_.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors +from __future__ import annotations + import math import os from datetime import UTC, date, datetime, time