dbtsai commented on code in PR #57099:
URL: https://github.com/apache/spark/pull/57099#discussion_r3540541738
##########
python/pyspark/sql/conversion.py:
##########
@@ -980,6 +980,60 @@ class ArrowTableToRowsConversion:
Conversion from Arrow Table to Rows.
"""
+ @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.
+
+ ``Array.to_pylist()`` materializes one Scalar per element; for list
types each row
+ additionally allocates a C++ scalar, a Python Scalar wrapper and a
Python Array
+ wrapper for the row's values before converting elements one by one,
which is
+ several times slower than converting the flattened child values in a
single pass
+ and slicing the resulting Python list per row (see
apache/arrow#50326). The values
+ themselves are still converted by Arrow's own ``to_pylist``, so
results are exactly
+ identical: ``None`` stays ``None`` and values inside numeric lists
stay Python ints,
+ unlike a pandas round trip which would coerce them to floats/NaN.
NumPy is used
+ only for the offsets (non-null integers) and the validity bitmap
(booleans), so no
+ value coercion can occur.
+
+ This can be removed once the minimum supported PyArrow version
includes the fix
+ for apache/arrow#50326.
+ """
+ import pyarrow as pa
+ import pyarrow.types as pa_types
+
+ try:
+ import numpy # noqa: F401
+ except ImportError:
+ return column.to_pylist()
+
+ if isinstance(column, pa.ChunkedArray):
+ result: List[Any] = []
+ for chunk in column.chunks:
+ result.extend(ArrowTableToRowsConversion._to_pylist(chunk))
+ return result
+
+ column_type = column.type
+ if (pa_types.is_list(column_type) or
pa_types.is_large_list(column_type)) and len(
+ column
+ ) > 0:
+ n = len(column)
+ offsets = column.offsets.to_numpy(zero_copy_only=True).tolist()
+ start = offsets[0]
+ flat = ArrowTableToRowsConversion._to_pylist(
+ column.values.slice(start, offsets[-1] - start)
+ )
+ if column.null_count == 0:
+ return [flat[offsets[i] - start : offsets[i + 1] - start] for
i in range(n)]
+ valid = column.is_valid().to_numpy(zero_copy_only=False).tolist()
+ return [
Review Comment:
Minor (no change needed): the per-row slicing here is still O(rows) Python
(`offsets.tolist()` + a comprehension slicing `flat`), so the win is fewer
object allocations, not vectorization -- consistent with the benchmark numbers.
Fine as the documented interim tradeoff; noting it so the "bulk" framing isn't
read as fully vectorized.
##########
python/pyspark/sql/conversion.py:
##########
@@ -980,6 +980,60 @@ class ArrowTableToRowsConversion:
Conversion from Arrow Table to Rows.
"""
+ @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.
+
+ ``Array.to_pylist()`` materializes one Scalar per element; for list
types each row
+ additionally allocates a C++ scalar, a Python Scalar wrapper and a
Python Array
+ wrapper for the row's values before converting elements one by one,
which is
+ several times slower than converting the flattened child values in a
single pass
+ and slicing the resulting Python list per row (see
apache/arrow#50326). The values
+ themselves are still converted by Arrow's own ``to_pylist``, so
results are exactly
+ identical: ``None`` stays ``None`` and values inside numeric lists
stay Python ints,
+ unlike a pandas round trip which would coerce them to floats/NaN.
NumPy is used
+ only for the offsets (non-null integers) and the validity bitmap
(booleans), so no
+ value coercion can occur.
+
+ This can be removed once the minimum supported PyArrow version
includes the fix
+ for apache/arrow#50326.
+ """
+ import pyarrow as pa
+ import pyarrow.types as pa_types
+
+ try:
+ import numpy # noqa: F401
+ except ImportError:
+ return column.to_pylist()
+
+ if isinstance(column, pa.ChunkedArray):
+ result: List[Any] = []
+ for chunk in column.chunks:
+ result.extend(ArrowTableToRowsConversion._to_pylist(chunk))
+ return result
+
+ column_type = column.type
+ if (pa_types.is_list(column_type) or
pa_types.is_large_list(column_type)) and len(
+ column
+ ) > 0:
+ n = len(column)
+ offsets = column.offsets.to_numpy(zero_copy_only=True).tolist()
Review Comment:
Minor: `zero_copy_only=True` here is load-bearing and works only because
list offset buffers never carry a validity bitmap. That invariant is correct,
but non-obvious -- a one-line comment would stop a future reader from "fixing"
it to `zero_copy_only=False` (and would flag the assumption if some new Arrow
list variant ever violated it).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
