diff --git a/python/benchmarks/bench_arrow.py b/python/benchmarks/bench_arrow.py index b59a43d42072..4ed1c91d579c 100644 --- a/python/benchmarks/bench_arrow.py +++ b/python/benchmarks/bench_arrow.py @@ -170,3 +170,46 @@ def peakmem_nested_ints_with_nulls_to_rows(self, n_rows, method): def peakmem_array_of_structs_to_rows(self, n_rows, method): self.convert(self.array_of_structs) + + +class ArrowStructMapColumnToRowsBenchmark: + """ + Benchmark for converting Arrow struct and map columns to Python rows. + + ``baseline`` measures plain ``column.to_pylist()``; ``bulk`` measures + ``ArrowTableToRowsConversion._to_pylist`` with the struct/map bulk paths. + """ + + params = [ + [100000, 1000000], + ["baseline", "bulk"], + ] + param_names = ["n_rows", "method"] + + def setup(self, n_rows, method): + from pyspark.sql.conversion import ArrowTableToRowsConversion + + self.structs = pa.array( + [{"a": i, "b": f"s{i}"} if i % 10 != 0 else None for i in range(n_rows)], + type=pa.struct([("a", pa.int64()), ("b", pa.string())]), + ) + self.maps = pa.array( + [ + [(f"k{i % 3}", i), (f"q{i % 5}", i + 1)] if i % 10 != 0 else None + for i in range(n_rows) + ], + type=pa.map_(pa.string(), pa.int64()), + ) + if method == "bulk": + self.convert = ArrowTableToRowsConversion._to_pylist + else: + self.convert = lambda column: column.to_pylist() + + def time_structs_to_rows(self, n_rows, method): + self.convert(self.structs) + + def time_maps_to_rows(self, n_rows, method): + self.convert(self.maps) + + def peakmem_structs_to_rows(self, n_rows, method): + self.convert(self.structs) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index fddb42c63d8c..e660ddc18e12 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -1011,8 +1011,11 @@ def _should_manual_bulk() -> bool: @staticmethod def _to_pylist(column: Union["pa.Array", "pa.ChunkedArray"]) -> List[Any]: """ - Equivalent to ``column.to_pylist()``, but converts (nested) list columns in bulk - instead of one scalar at a time. + Equivalent to ``column.to_pylist()``, but converts (nested) list, struct and map + columns in bulk instead of one scalar at a time. Structs become dicts (with + a fallback to ``to_pylist`` for duplicate field names, which raise ``ValueError`` + there) and maps become lists of ``(key, value)`` tuples, matching + ``StructScalar.as_py`` and ``MapScalar.as_py`` exactly. Internal helper for the worker and ``convert`` call sites; do not use externally. @@ -1043,9 +1046,44 @@ def _to_pylist(column: Union["pa.Array", "pa.ChunkedArray"]) -> List[Any]: result.extend(ArrowTableToRowsConversion._to_pylist(chunk)) return result - if (pa.types.is_list(column.type) or pa.types.is_large_list(column.type)) and len( - column - ) > 0: + if len(column) == 0: + return [] + + if pa.types.is_map(column.type): + # Maps have the same offsets layout as lists; each row becomes a + # list of (key, value) tuples, matching MapScalar.as_py. + n = len(column) + offsets = column.offsets.to_numpy(zero_copy_only=True).tolist() + start = offsets[0] + length = offsets[-1] - start + keys = ArrowTableToRowsConversion._to_pylist(column.keys.slice(start, length)) + items = ArrowTableToRowsConversion._to_pylist(column.items.slice(start, length)) + if column.null_count == 0: + return [ + list( + zip( + keys[offsets[i] - start : offsets[i + 1] - start], + items[offsets[i] - start : offsets[i + 1] - start], + ) + ) + for i in range(n) + ] + valid = column.is_valid().to_numpy(zero_copy_only=False).tolist() + return [ + ( + list( + zip( + keys[offsets[i] - start : offsets[i + 1] - start], + items[offsets[i] - start : offsets[i + 1] - start], + ) + ) + if valid[i] + else None + ) + for i in range(n) + ] + + elif pa.types.is_list(column.type) or pa.types.is_large_list(column.type): n = len(column) # List offset buffers never carry a validity bitmap, so this conversion is # always zero-copy; zero_copy_only=True asserts that invariant and would @@ -1063,6 +1101,26 @@ def _to_pylist(column: Union["pa.Array", "pa.ChunkedArray"]) -> List[Any]: for i in range(n) ] + elif pa.types.is_struct(column.type): + n = len(column) + names = [column.type.field(i).name for i in range(column.type.num_fields)] + if len(set(names)) != len(names): + # StructScalar.as_py raises ValueError on duplicate field names; + # let the generic path surface the same error. + return column.to_pylist() + fields = [ + ArrowTableToRowsConversion._to_pylist(column.field(i)) + for i in range(column.type.num_fields) + ] + if column.null_count == 0: + if not names: + return [{} for _ in range(n)] + return [dict(zip(names, row)) for row in zip(*fields)] + valid = column.is_valid().to_numpy(zero_copy_only=False).tolist() + if not names: + return [{} if m else None for m in valid] + return [dict(zip(names, row)) if m else None for row, m in zip(zip(*fields), valid)] + return column.to_pylist() @staticmethod diff --git a/python/pyspark/sql/tests/test_conversion.py b/python/pyspark/sql/tests/test_conversion.py index 7ce2dfe01411..d02dd8e73497 100644 --- a/python/pyspark/sql/tests/test_conversion.py +++ b/python/pyspark/sql/tests/test_conversion.py @@ -15,6 +15,7 @@ # limitations under the License. # import datetime +import decimal import unittest import unittest.mock from zoneinfo import ZoneInfo @@ -899,6 +900,37 @@ def test_matches_to_pylist(self): pa.array([], type=pa.list_(pa.int32())), pa.array([None, None], type=pa.list_(pa.string())), pa.array([[1, 2], None], type=pa.list_(pa.int64(), 2)), + # non-list leaves keep as_py semantics (native to_pylist) + pa.array([b"", None, b"\x00\xff"], type=pa.binary()), + pa.array([datetime.date(2020, 1, 2), None], type=pa.date32()), + pa.array([decimal.Decimal("1.23"), None], type=pa.decimal128(10, 2)), + pa.array([[b"x", None], None, [b""]], type=pa.list_(pa.binary())), + pa.array([[True, None], [False]], type=pa.list_(pa.bool_())), + # struct and map bulk paths + pa.array( + [{"a": 1, "b": "x"}, None, {"a": None, "b": None}], + type=pa.struct([("a", pa.int64()), ("b", pa.string())]), + ), + pa.array( + [{"s": {"a": 1}, "l": [1, None]}, None], + type=pa.struct( + [("s", pa.struct([("a", pa.int32())])), ("l", pa.list_(pa.int64()))] + ), + ), + pa.array([{}, None, {}], type=pa.struct([])), + pa.array([None] * 4, type=pa.struct([("a", pa.int32())])), + pa.array( + [[("k1", [1, None]), ("k2", None)], None, []], + type=pa.map_(pa.string(), pa.list_(pa.int32())), + ), + pa.array( + [{"m": [("k", 1)]}, None], + type=pa.struct([("m", pa.map_(pa.string(), pa.int64()))]), + ), + pa.array( + [[{"a": 1}, None], None], + type=pa.list_(pa.struct([("a", pa.int64())])), + ), ] for column in columns: views = [column, column.slice(1), column.slice(0, max(len(column) - 1, 0))] @@ -922,6 +954,20 @@ def test_int_list_with_nulls_stays_int(self): self.assertEqual(result, [[1, None, 3]]) self.assertEqual([type(v) for v in result[0]], [int, type(None), int]) + def test_struct_duplicate_field_names_still_raises(self): + import pyarrow as pa + + dup = pa.StructArray.from_arrays([pa.array([1, 2]), pa.array(["a", "b"])], names=["x", "x"]) + with self.assertRaises(ValueError): + ArrowTableToRowsConversion._to_pylist(dup) + + def test_struct_rows_are_distinct_dicts(self): + import pyarrow as pa + + result = ArrowTableToRowsConversion._to_pylist(pa.array([{}, {}], type=pa.struct([]))) + self.assertEqual(result, [{}, {}]) + self.assertIsNot(result[0], result[1]) + def test_convert_table_with_list_columns(self): import pyarrow as pa