This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 225b74a2d307 [SPARK-58625][PYTHON] Support pa.ChunkedArray columns in
PandasToArrowConversion
225b74a2d307 is described below
commit 225b74a2d307ba12dab912e37e25f7423926d75d
Author: Spenser Sun <[email protected]>
AuthorDate: Fri Aug 7 13:44:43 2026 +0800
[SPARK-58625][PYTHON] Support pa.ChunkedArray columns in
PandasToArrowConversion
### What changes were proposed in this pull request?
`PandasToArrowConversion.convert` assembles its result with
`pa.RecordBatch.from_arrays`, which rejects a `pa.ChunkedArray` -- and the
per-column conversion it feeds from calls `pa.Array.from_pandas`, which returns
one when the input Series is backed by a chunked Arrow array. Flatten those
columns before assembling the batch:
```python
converted = [convert_column(col, field) for col, field in zip(columns,
schema.fields)]
# pa.Array.from_pandas returns a pa.ChunkedArray for a chunked arrow-backed
Series
# (e.g. a pyarrow-backed extension dtype), which pa.RecordBatch.from_arrays
rejects.
arrays = [a.combine_chunks() if isinstance(a, pa.ChunkedArray) else a for a
in converted]
return pa.RecordBatch.from_arrays(arrays, schema.names)
```
The `isinstance` guard is required -- `pa.Array` has no `combine_chunks`
attribute -- and costs ~81 ns on the common path.
**Relation to SPARK-46776 / #56157.** The sibling
`create_arrow_table_from_pandas` (`python/pyspark/sql/pandas/conversion.py`)
already handles this, and its docstring states the rule generally:
`pa.Array.from_pandas` may return a `ChunkedArray`,
`pa.RecordBatch.from_arrays` does not accept one, `pa.Table.from_arrays` does.
That covered `createDataFrame`. `PandasToArrowConversion` -- the pandas UDF
return path -- is a separate assembler in a different module, extracted from
`serializer [...]
**Alternative considered: `pa.Table.from_arrays(...).to_batches()`, i.e.
what the sibling does.** It also covers the >2 GB case below, but is not taken
here:
- `convert` is annotated `-> "pa.RecordBatch"`, and the consumer of all 14
`worker.py` call sites is `dump_stream`'s `writer.write_batch(batch)`, which
rejects a `Table`. The return-type change ripples to every site.
- `to_batches()` emits one batch per chunk boundary (unioned across
columns), so the batch grouping becomes data-dependent. That is harmless for 13
sites, but **`applyInPandasWithState` cannot be split**: its
`construct_record_batch` pads the `_1` data and `_2` state columns to a common
length and stores the true unpadded counts in `_0`. Splitting separates that
header from the rows it describes -- silent data corruption rather than an
error. Making it safe is a design question, not a [...]
- `combine_chunks()` copies, which `Table.from_arrays` avoids. But a
`RecordBatch` structurally cannot hold chunked data, so the choice is copy-once
or emit-N; copying only genuinely chunked input is the smaller change.
Happy to switch to the multi-batch approach in a follow-up if reviewers
prefer, with `applyInPandasWithState` handled explicitly.
**Out of scope: string data over 2 GB**, which also makes `from_pandas`
chunk (an int32-offset buffer holds at most 2 GiB - 1). It cannot be fixed
while returning a single `RecordBatch`, since the combined data would not fit
one buffer. `combine_chunks()` raises `ArrowInvalid: offset overflow while
concatenating arrays, consider casting input from string to large_string
first`, which names the remedy:
`spark.sql.execution.arrow.useLargeVarTypes=true` uses int64 offsets and avoids
chun [...]
### Why are the changes needed?
A pandas UDF returning a chunked arrow-backed Series fails with a raw
pyarrow error -- `TypeError: Cannot convert pyarrow.lib.ChunkedArray to
pyarrow.lib.Array`. Reproducible without a Spark session:
```python
import pandas as pd, pyarrow as pa
from pyspark.sql.conversion import PandasToArrowConversion
from pyspark.sql.types import StructType, StructField, StringType
chunked = pa.chunked_array([pa.array(["a", "b"]), pa.array(["c", "d",
"e"])])
series = pd.Series(chunked, dtype="string[pyarrow]")
schema = StructType([StructField("s", StringType())])
PandasToArrowConversion.convert([series], schema, arrow_cast=True)
```
Inside a UDF this arises from ordinary pandas operations: `pd.concat` of
two pyarrow-backed Series produces two chunks, and `.copy()` / `.reset_index()`
preserve the chunking -- so whether the UDF crashes depends on which operation
it happened to end with, and a user cannot avoid it without knowing to call
`combine_chunks()` or `.astype(object)`.
Two things make this worse than the `createDataFrame` case SPARK-46776
fixed. There is **no fallback**: `createDataFrame` is wrapped in a `try/except`
governed by `spark.sql.execution.arrow.pyspark.fallback.enabled` (default
`true`), which downgrades the failure to a warning plus a non-Arrow slow path,
so results are still correct; a UDF in a Python worker has no such net and the
task fails. And there is **no error classification**: the assembly call sits
outside the `try/except` that [...]
This was found while writing the `pa.Array.from_pandas` golden tests in
#57830 (SPARK-58602): pinning that function's return type showed it can hand
back a `pa.ChunkedArray`, which led here. That PR records the upstream
behavior; this one fixes the consumer.
### Does this PR introduce _any_ user-facing change?
Yes. A pandas UDF returning a chunked arrow-backed Series no longer fails
with `TypeError: Cannot convert pyarrow.lib.ChunkedArray to pyarrow.lib.Array`;
the column is flattened and the UDF succeeds. No change for non-chunked
columns, which is the common path.
### How was this patch tested?
New test `PandasToArrowConversionTests.test_convert_chunked_array_backed`
in `python/pyspark/sql/tests/test_conversion.py`, asserting the resulting
column is a `pa.Array` (not merely that the call does not raise) and that the
values survive. Verified to fail without the fix with the original `TypeError`.
`test_conversion.py` passes in full (42 tests, 114 subtests) on pyarrow
24.0.0 / pandas 2.3.3. Also checked manually, beyond what the test asserts: the
nested case (a `pd.DataFrame` column, which re-enters the same assembly through
the struct recursion -- the shape `mapInPandas` uses), a chunked
`int32[pyarrow]` column with nulls in both chunks, the `pd.concat` idiom, and
the >2 GB case discussed above.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.5)
Closes #57829 from Spenserrrr/chunkedarray-recordbatch-fix.
Authored-by: Spenser Sun <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 86d4fe8b3b792e3426d39b6de80ac21baa573d9c)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/sql/conversion.py | 5 ++++-
python/pyspark/sql/tests/test_conversion.py | 15 +++++++++++++++
2 files changed, 19 insertions(+), 1 deletion(-)
diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py
index e472c7304e2e..b326df555f79 100644
--- a/python/pyspark/sql/conversion.py
+++ b/python/pyspark/sql/conversion.py
@@ -503,7 +503,10 @@ class PandasToArrowConversion:
)
raise PySparkValueError(error_msg) from e
- arrays = [convert_column(col, field) for col, field in zip(columns,
schema.fields)]
+ converted = [convert_column(col, field) for col, field in zip(columns,
schema.fields)]
+ # pa.Array.from_pandas returns a pa.ChunkedArray for a chunked
arrow-backed Series
+ # (e.g. a pyarrow-backed extension dtype), which
pa.RecordBatch.from_arrays rejects.
+ arrays = [a.combine_chunks() if isinstance(a, pa.ChunkedArray) else a
for a in converted]
return pa.RecordBatch.from_arrays(arrays, schema.names)
diff --git a/python/pyspark/sql/tests/test_conversion.py
b/python/pyspark/sql/tests/test_conversion.py
index d02dd8e73497..3a5c9d520b75 100644
--- a/python/pyspark/sql/tests/test_conversion.py
+++ b/python/pyspark/sql/tests/test_conversion.py
@@ -461,6 +461,21 @@ class PandasToArrowConversionTests(unittest.TestCase):
result = PandasToArrowConversion.convert([cat_series], schema)
self.assertEqual(result.column(0).to_pylist(), ["a", "b", "a", "c"])
+ def test_convert_chunked_array_backed(self):
+ """Test a chunked arrow-backed series is converted to a single
Array."""
+ import pandas as pd
+ import pyarrow as pa
+
+ # pa.Array.from_pandas returns a ChunkedArray here, which
+ # pa.RecordBatch.from_arrays rejects.
+ chunked = pa.chunked_array([pa.array(["a", "b"]), pa.array(["c", "d",
"e"])])
+ series = pd.Series(chunked, dtype="string[pyarrow]")
+ schema = StructType([StructField("s", StringType())])
+
+ result = PandasToArrowConversion.convert([series], schema,
arrow_cast=True)
+ self.assertIsInstance(result.column(0), pa.Array)
+ self.assertEqual(result.column(0).to_pylist(), ["a", "b", "c", "d",
"e"])
+
@unittest.skipIf(not have_pyarrow, pyarrow_requirement_message)
class ConversionTests(unittest.TestCase):
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
