viirya opened a new pull request, #59423:
URL: https://github.com/apache/spark/pull/59423
### What changes were proposed in this pull request?
In `ApplyInPandasWithStateSerializer.load_stream`, explicitly convert
`startOffset` / `numRows` read from the state struct to `int` before passing
them to `RecordBatch.slice`.
### Why are the changes needed?
PyArrow 26.0.0 added `operator.index(...)` validation to `offset` / `length`
in `RecordBatch.slice` (and `Table.slice` / `ChunkedArray.slice`). Before that,
a float was silently truncated by the Cython `int64_t` conversion.
On this branch the state struct column is converted to pandas as a Series of
dicts. When an Arrow batch contains padding state rows
(`ApplyInPandasWithStateWriter.EMPTY_STATE_METADATA_ROW`, written whenever a
batch has more data rows than state rows, i.e. a group with more than one row,
or chunking), the struct has nulls, so its int32 fields are promoted to float64
and `startOffset` / `numRows` become Python floats. With PyArrow 26 this fails
with:
```
File ".../python/pyspark/sql/pandas/serializers.py", line 1596, in
gen_data_and_state
data_batch_for_group = data_batch.slice(data_start_offset, num_data_rows)
...
TypeError: 'float' object cannot be interpreted as an integer
```
This breaks `applyInPandasWithState` for users on PyArrow 26, and
`FlatMapGroupsInPandasWithStateSuite` (5 tests) in the "sql - slow tests" job
of the branch-4.0 / branch-4.1 scheduled builds since PyArrow 26.0.0 was
released (that job installs an unpinned `pyarrow`; the last green branch-4.0
run had 25.0.1, the first failing one 26.0.0, with identical numpy / pandas).
branch-4.2 and later are not affected because SPARK-54962 converts nullable
integers with `integer_object_nulls=True`, so the values stay ints there.
### Does this PR introduce _any_ user-facing change?
Yes. `applyInPandasWithState` works again with PyArrow 26. No change with
earlier PyArrow versions.
### How was this patch tested?
Existing tests: `FlatMapGroupsInPandasWithStateSuite` covers this path and
fails with PyArrow 26 in CI.
Locally:
- `build/sbt 'sql/testOnly *FlatMapGroupsInPandasWithStateSuite'` passes on
JDK 17.
- Fed an Arrow IPC stream to `ApplyInPandasWithStateSerializer.load_stream`
with and without a padding state row and recorded the argument types passed to
`slice`: `float` before this change when padding is present, `int` after.
`operator.index` (what PyArrow 26 calls) raises the exact error above on the
former.
The change applies cleanly to branch-4.0 and branch-3.5, which have the same
code.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (claude-opus-5-5)
This pull request and its description were written by Isaac.
--
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]
