zahed1994 opened a new pull request, #58405:
URL: https://github.com/apache/spark/pull/58405
### What changes were proposed in this pull request?
This PR adds support for pandas `ArrowDtype` and PyArrow-backed string
dtypes in PySpark's pandas API (`pyspark.pandas`).
Key changes:
- `python/pyspark/pandas/typedef/typehints.py`:
- Added safe import check for `pandas.ArrowDtype`
(`extension_arrow_dtypes_available`).
- Updated `as_spark_type()` to delegate `ArrowDtype` directly to PySpark's
existing `from_arrow_type(tpe.pyarrow_dtype, prefer_timestamp_ntz)`.
- Updated `spark_type_to_pandas_dtype()` to return
`ArrowDtype(to_arrow_type(spark_type))` when `use_arrow_dtypes=True`.
- Added `is_pyarrow_backed_dtype(tpe)` to detect `ArrowDtype` and
PyArrow-backed `StringDtype` instances.
- Extended `is_str_dtype()` and `handle_dtype_as_extension_dtype()` to
recognize `ArrowDtype`.
- `python/pyspark/pandas/internal.py`:
- Added `use_arrow_dtypes: bool = False` parameter to
`InternalField.from_struct_field()` and forwarded it to
`spark_type_to_pandas_dtype()`.
- `python/pyspark/pandas/base.py`:
- Updated `column_op()` decorator to inspect input operands via
`is_pyarrow_backed_dtype(col.dtype)` and pass `use_arrow_dtypes` to
`InternalField.from_struct_field()`.
- `python/pyspark/pandas/utils.py` & `indexing.py`:
- Propagated `use_arrow_dtypes` through `combine_frames()` and
`__setitem__` struct field recreation.
---
### Why are the changes needed?
In pandas 2.0+ and pandas 3.0, PyArrow-backed dtypes (e.g. `int64[pyarrow]`,
`string[pyarrow]`, `bool[pyarrow]`) are increasingly used for Series and
DataFrames. Without this change, `pyspark.pandas` fails to recognize
`ArrowDtype` during `as_spark_type()` schema conversion and drops
PyArrow-backed return dtypes back to standard NumPy dtypes during operations
like comparisons (`s1 == s2`).
---
### Does this PR introduce _any_ user-facing change?
Yes. Users operating on PyArrow-backed pandas Series/DataFrames will now see
their dtypes recognized during type inference and preserved through PySpark
pandas operations:
#### 1. Schema Inference & Type Mapping
```python
import pandas as pd
import pyarrow as pa
from pyspark.pandas.typedef import as_spark_type, spark_type_to_pandas_dtype
# Converting PyArrow-backed dtypes to Spark DataType:
as_spark_type(pd.ArrowDtype(pa.int64())) # -> LongType()
as_spark_type(pd.ArrowDtype(pa.string())) # -> StringType()
as_spark_type(pd.ArrowDtype(pa.bool_())) # -> BooleanType()
# Converting Spark DataType back to PyArrow-backed dtypes:
spark_type_to_pandas_dtype(types.LongType(), use_arrow_dtypes=True) # ->
int64[pyarrow]
```
#### 2. Dtype Preservation in Operations
```python
import pandas as pd
import pyarrow as pa
import pyspark.pandas as ps
s1 = ps.Series(["a", "b", "c"], dtype=pd.ArrowDtype(pa.string()))
s2 = ps.Series(["a", "x", "c"], dtype=pd.ArrowDtype(pa.string()))
res = s1 == s2
# Before: res.dtype -> boolean (NumPy bool)
# After: res.dtype -> bool[pyarrow]
```
---
### How was this patch tested?
#### Unit & Integration Tests Added
- `python/pyspark/pandas/tests/test_typedef.py`:
- `test_as_spark_type_pyarrow_dtypes()`: Validates bidirectional mapping
for `bool[pyarrow]`, `int8[pyarrow]`, `int16[pyarrow]`, `int32[pyarrow]`,
`int64[pyarrow]`, `float[pyarrow]`, `double[pyarrow]`, `string[pyarrow]`.
- `test_is_pyarrow_backed_dtype()`: Validates detection of `ArrowDtype`
and `StringDtype(storage="pyarrow")`.
- `python/pyspark/pandas/tests/data_type_ops/test_string_ops.py`:
- `test_pyarrow_backed_string_ops()`: Validates that comparison operations
(`s1 == s2`) preserve `bool[pyarrow]` dtypes.
#### Local Test Execution & Verification
1. **Bidirectional Type Mapper Tests**:
```text
PASS: bool[pyarrow] <-> BooleanType()
PASS: int8[pyarrow] <-> ByteType()
PASS: int16[pyarrow] <-> ShortType()
PASS: int32[pyarrow] <-> IntegerType()
PASS: int64[pyarrow] <-> LongType()
PASS: float[pyarrow] <-> FloatType()
PASS: double[pyarrow] <-> DoubleType()
PASS: string[pyarrow] <-> StringType()
```
2. **Dtype Detection & InternalField Propagation**:
```text
PASS: is_pyarrow_backed_dtype(ArrowDtype(pa.int64())) == True
PASS: is_pyarrow_backed_dtype(ArrowDtype(pa.string())) == True
PASS: is_pyarrow_backed_dtype(StringDtype(storage='pyarrow')) == True
PASS: is_pyarrow_backed_dtype(StringDtype(storage='python')) == False
PASS: handle_dtype_as_extension_dtype(ArrowDtype(pa.int64())) == True
Arrow struct field dtype: bool[pyarrow]
SUCCESS: InternalField with use_arrow_dtypes=True returns bool[pyarrow]!
```
---
### 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]
