[SPARK-55294][PYTHON] Add support for PyArrow-backed dtypes in pandas API

[adith-os]

What changes were proposed in this pull request?

This PR adds support for PyArrow-backed dtypes (e.g., bool[pyarrow]) in PySpark's pandas API. The changes include:

• Type detection and conversion: Added is_pyarrow_backed_dtype() function to detect PyArrow-backed dtypes and enhanced as_spark_type() to convert them to appropriate Spark types.
• Dtype preservation: Enhanced spark_type_to_pandas_dtype() with a use_arrow_dtypes parameter to preserve PyArrow dtypes when converting from Spark types back to pandas dtypes.
• Propagation through operations: Updated comparison operators in data_type_ops/base.py to preserve PyArrow dtypes in results, and propagated the use_arrow_dtypes parameter through InternalField.from_struct_field() and related code paths.

Why are the changes needed?

Starting with pandas 3.0, when PyArrow is installed, pandas automatically uses PyArrow-backed dtypes for string columns. Without this PR, PySpark's pandas API loses dtype information when working with PyArrow-backed dtypes

Does this PR introduce any user-facing change?

Yes. Users working with PyArrow-backed dtypes will now see their dtypes preserved through PySpark operations
Before the PR:

>>> 26/05/09 19:25:26 WARN GarbageCollectionMetrics: To enable non-built-in garbage collector(s) List(scavenge), users should configure it(them) to spark.eventLog.gcMetrics.youngGenerationGarbageCollectors or spark.eventLog.gcMetrics.oldGenerationGarbageCollectors
26/05/09 19:25:26 WARN GarbageCollectionMetrics: To enable non-built-in garbage collector(s) List(global, scavenge), users should configure it(them) to spark.eventLog.gcMetrics.youngGenerationGarbageCollectors or spark.eventLog.gcMetrics.oldGenerationGarbageCollectors
import warnings
>>> warnings.filterwarnings('ignore')
>>> 
>>> import pandas as pd
>>> print(f"pandas: {pd.__version__}")
pandas: 3.0.2
>>> 
>>> import pyspark.pandas as ps
>>> s1 = ps.Series(['a', 'b', 'c'], dtype='string')
>>> s2 = ps.Series(['a', 'x', 'c'], dtype='string')
>>> result = s1 == s2
>>> print(f"result dtype: {result.dtype}")
result dtype: boolean
>>> print(result)
0     True                                                                      
1    False
2     True
dtype: boolean
>>> 

After the PR:

>>> 26/05/09 19:19:05 WARN GarbageCollectionMetrics: To enable non-built-in garbage collector(s) List(scavenge), users should configure it(them) to spark.eventLog.gcMetrics.youngGenerationGarbageCollectors or spark.eventLog.gcMetrics.oldGenerationGarbageCollectors
26/05/09 19:19:05 WARN GarbageCollectionMetrics: To enable non-built-in garbage collector(s) List(global, scavenge), users should configure it(them) to spark.eventLog.gcMetrics.youngGenerationGarbageCollectors or spark.eventLog.gcMetrics.oldGenerationGarbageCollectors

>>> import warnings
>>> warnings.filterwarnings('ignore')
>>> 
>>> import pandas as pd
>>> print(f"pandas: {pd.__version__}")
pandas: 3.0.2
>>> 
>>> import pyspark.pandas as ps
>>> s1 = ps.Series(['a', 'b', 'c'], dtype='string')
s2 = ps.Series(['a', 'x', 'c'], dtype='string')
result = s1 == s2
print(f"result dtype: {result.dtype}")
print(result)

>>> s2 = ps.Series(['a', 'x', 'c'], dtype='string')
>>> result = s1 == s2
>>> print(f"result dtype: {result.dtype}")
result dtype: bool[pyarrow]
>>> print(result)
0     True
1    False
2     True
dtype: bool[pyarrow]
>>> 

How was this patch tested?

New tests in python/pyspark/pandas/tests/test_typedef.py:

• test_as_spark_type_pyarrow_dtypes() - Tests conversion to Spark types
• test_spark_type_to_pandas_dtype_with_arrow_flag() - Tests conversion back to pandas dtypes
• test_is_str_dtype_with_pyarrow() - Tests string dtype detection
• test_is_pyarrow_backed_dtype() - Tests PyArrow dtype detection

Integration test in python/pyspark/pandas/tests/data_type_ops/test_string_ops.py:

• test_pyarrow_backed_string_comparisons() - Tests dtype preservation in comparison operations

Was this patch authored or co-authored using generative AI tooling?

Yes. Generated-by: GPT 5.5

[adith-os]

Hi @gaogaotiantian, Could you please review the PR when you get a chance and let me know whether this approach matches your expectations, or if you had a different direction in mind? Thanks!

[HyukjinKwon]

cc @devin-petersohn too

[devin-petersohn]

This feels both incomplete and too verbose. ne and eq were the only APIs updated. There's some copy/pasted code as well. I left some high level comments.

I also haven't fully checked, but I think this will be a problem for pandas 3 and the string dtype.

[devin-petersohn]

Can we just reuse existing functionality?

spark/python/pyspark/sql/pandas/types.py

Lines 347 to 350 in 1360e5c

	def from_arrow_type(
	at: "pa.DataType",
	prefer_timestamp_ntz: bool = True,
	) -> DataType:

[devin-petersohn]

StringDtype.storage is part of the API. This feels overly defensive.

[devin-petersohn]

This looks identical to the code block above. Should be a helper if it's this many lines IMO

[gaogaotiantian]

Honestly, I'm a bit conservative when a new contributor making complex changes with LLM to support pandas 3. Our CI does not check pandas 3 (the nightly CI simply fails and can not be used against a certain PR). I've had a few cases where people just feed my description to LLM and generate some patch that of course passes the CI because we do not test pandas 3, while the code may not even work properly.

At this point, it's very difficult and time consuming to validate the patch and the approach. We have very limited resources on many different matters. I would say that we deprioritize supporting pandas 3, especially from new contributors with LLM, so that we can keep everything else going. In the future, when we have enough resources to deal with it, we can put more effort into it.