Spenserrrr opened a new pull request, #57829:
URL: https://github.com/apache/spark/pull/57829
### What changes were proposed in this pull request?
`PandasToArrowConversion.convert` assembles its result with
`pa.RecordBatch.from_arrays`, which rejects a `pa.ChunkedArray`. The per-column
conversion it feeds from calls `pa.Array.from_pandas`, which returns a
`ChunkedArray` when the input pandas Series is backed by a chunked Arrow array.
This PR flattens those columns with `combine_chunks()` 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, so an unguarded call would break every non-chunked column. Its cost
on the common path is ~81 ns.
**Relation to SPARK-46776 / #56157.** The sibling function
`create_arrow_table_from_pandas` (`python/pyspark/sql/pandas/conversion.py`)
already handles this, and its docstring documents the rule in general terms:
`pa.Array.from_pandas` may return a `ChunkedArray`,
`pa.RecordBatch.from_arrays` does not accept one, but `pa.Table.from_arrays`
does. That fix covered the `createDataFrame` path. `PandasToArrowConversion` --
the pandas UDF return path -- was not covered: it is a separate assembler in a
different module (`python/pyspark/sql/conversion.py`), extracted from
`serializers.py` by SPARK-55349 shortly before, and its `convert_column` is
annotated `-> "pa.Array"`, so the assembly trusts an annotation that
`from_pandas` can violate.
**Alternative considered: `pa.Table.from_arrays(...).to_batches()`, i.e.
exactly what the sibling does.** Not taken here, for three reasons:
1. `convert` is annotated `-> "pa.RecordBatch"` and every one of its 14 call
sites in `worker.py` consumes exactly one batch (12 `yield`, 2 `return`). The
consumer is `ArrowStreamSerializer.dump_stream`, which calls
`writer.write_batch(batch)` -- that rejects a `Table`. So the return-type
change would ripple to all 14 sites.
2. More importantly, one of those sites cannot be split into N batches at
all. `applyInPandasWithState` (`worker.py`, `construct_record_batch`) builds a
batch whose `_0` column records the true unpadded row counts for the `_1` data
and `_2` state columns, which are padded to a common length. Splitting that
batch at an arbitrary chunk boundary would separate the count header from the
rows it describes -- silent data corruption rather than a type error. Making
that path safe is a design question that should not ride along with a bug fix.
3. `combine_chunks()` copies into one contiguous buffer, which
`Table.from_arrays` avoids. But since a `RecordBatch` structurally cannot hold
chunked data, the only alternatives are to copy once or to emit multiple
batches; for a bug fix, copying only when the input is genuinely chunked is the
smaller change.
Happy to switch to the multi-batch approach in a follow-up if reviewers
prefer it, with the `applyInPandasWithState` case handled explicitly.
**Out of scope: string data over 2 GB.** That also makes `from_pandas`
return a `ChunkedArray` (an int32-offset buffer holds at most 2 GiB - 1), but
it cannot be fixed while returning a single `pa.RecordBatch`: the combined data
would not fit one buffer. On that input `combine_chunks()` raises
`ArrowInvalid: offset overflow while concatenating arrays, consider casting
input from string to large_string first`, which at least names the remedy --
`spark.sql.execution.arrow.useLargeVarTypes=true` uses int64 offsets and avoids
the chunking entirely (verified). Previously this case produced the same opaque
`TypeError` as every other chunked input.
### Why are the changes needed?
A pandas UDF whose returned Series is backed by a chunked Arrow array 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, a chunked-backed Series 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. A user cannot defensively
avoid it without knowing to call `combine_chunks()` or `.astype(object)`.
Two things make this worse than the `createDataFrame` case that SPARK-46776
fixed:
- **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 -- the user
still gets correct results. A pandas UDF running in a Python worker has no such
net, so this is a hard task failure.
- **No error classification.** The assembly call sits outside the
`try/except` that wraps the per-column conversion and turns pyarrow errors into
`PySparkTypeError` / `PySparkValueError`, so this `TypeError` escapes
unclassified from a function that otherwise classifies everything.
### 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 that 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`,
and to pass with it.
`python/pyspark/sql/tests/test_conversion.py` passes in full (42 tests, 114
subtests) on pyarrow 24.0.0 / pandas 2.3.3. Also verified manually, beyond what
the test asserts: the nested case (a `pd.DataFrame` column, which re-enters the
same assembly through the struct recursion -- this is 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)
--
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]
