Yicong-Huang opened a new pull request, #56157: URL: https://github.com/apache/spark/pull/56157
### What changes were proposed in this pull request? Make `SparkSession.createDataFrame` accept a pandas `DataFrame` whose columns are backed by a multi-chunk `pyarrow.ChunkedArray`. `pa.Array.from_pandas(series)` can return a `pa.ChunkedArray` instead of a `pa.Array` when: - the input pandas `Series` is backed by a chunked Arrow array (e.g. the pyarrow-backed string extension dtype `string[pyarrow]`), or - the data exceeds the maximum size of a single Arrow array (e.g. string data larger than 2 GB). `pa.RecordBatch.from_arrays` does not accept `ChunkedArray`, so `createDataFrame` previously failed with: ``` TypeError: Cannot convert pyarrow.lib.ChunkedArray to pyarrow.lib.Array ``` This change replaces `create_arrow_batch_from_pandas` with `create_arrow_batches_from_pandas`, which materializes the columns into a `pa.Table` (which accepts both `Array` and `ChunkedArray`) and then splits it back into one or more `RecordBatch`es aligned on a common chunk boundary via `Table.to_batches()`. Both the classic and Spark Connect `createDataFrame` paths are updated to consume the list. The rebatch is zero-copy: `pa.Table.from_arrays` wraps single `Array` columns in single-chunk `ChunkedArray`s, and `to_batches()` aligns chunk boundaries across columns using zero-copy slices. ### Why are the changes needed? `createDataFrame` from a pandas `DataFrame` is broken whenever any column is backed by a `ChunkedArray`, which is increasingly common with pyarrow-backed pandas extension dtypes and with large datasets. See [SPARK-46776](https://issues.apache.org/jira/browse/SPARK-46776). ### Does this PR introduce _any_ user-facing change? Yes. `SparkSession.createDataFrame(pandas_df)` no longer raises `TypeError: Cannot convert pyarrow.lib.ChunkedArray to pyarrow.lib.Array` when a column is backed by a multi-chunk pyarrow array. Behavior for inputs without `ChunkedArray` columns is unchanged. ### How was this patch tested? - Added `ArrowTestsMixin.test_createDataFrame_pandas_chunked_array_backed`, which constructs a `pandas.DataFrame` with a `string[pyarrow]` column whose storage is a 2-chunk `pa.ChunkedArray` and verifies `createDataFrame` returns the expected row count and values. The test is exercised both with Arrow execution enabled and disabled, and runs against both the classic and Spark Connect `createDataFrame` implementations through the existing parity test infrastructure. - Existing `test_createDataFrame_toggle`, `test_createDataFrame_arrow_pandas`, and `test_createDataFrame_pandas_respect_session_timezone` still pass. ### Was this patch authored or co-authored using generative AI tooling? No -- 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]
