Skip to content

[SPARK-54966][PYTHON] Factor out legacy pandas-to-Arrow column conversion #58751

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Spenserrrr wants to merge 1 commit into
base: master
Choose a base branch
from
+201 −154
Open

[SPARK-54966][PYTHON] Factor out legacy pandas-to-Arrow column conversion #58751

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
257 changes: 151 additions & 106 deletions python/pyspark/sql/conversion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -309,7 +309,7 @@ class PandasToArrowConversion:
"""

@classmethod
def convert(
def from_pandas(
cls,
data: Union["pd.DataFrame", Sequence[Union["pd.Series", "pd.DataFrame"]]],
schema: StructType,
Expand All @@ -321,7 +321,7 @@ def convert(
assign_cols_by_name: bool = False,
int_to_decimal_coercion_enabled: bool = False,
ignore_unexpected_complex_type_values: bool = False,
is_legacy: bool = False,
use_legacy_error_handling: bool = False,
) -> "pa.RecordBatch":
"""
Convert a pandas DataFrame or list of Series/DataFrames to an Arrow RecordBatch.
Expand All @@ -348,13 +348,11 @@ def convert(
Whether to enable int to decimal coercion (default False)
ignore_unexpected_complex_type_values : bool
Whether to ignore unexpected complex type values in converter (default False)
is_legacy : bool
Whether to use the legacy pandas-to-Arrow conversion path. The legacy
path uses broader Arrow exception handling (ArrowException) to allow
more implicit type coercions (e.g., int->boolean, dict->struct via
ArrowTypeError). The non-legacy path only catches ArrowInvalid for
the cast fallback, so type mismatches like string->decimal raise
immediately. (default False)
use_legacy_error_handling : bool
Whether to use legacy error handling and error messages. Legacy handling
catches ArrowException (including ArrowTypeError) for the cast fallback.
The new error handling only catches ArrowInvalid for the fallback;
ArrowTypeError is raised without retrying the conversion. (default False)

Returns
-------
Expand All @@ -363,8 +361,7 @@ def convert(
import pandas as pd
import pyarrow as pa

from pyspark.errors import PySparkTypeError, PySparkValueError
from pyspark.sql.pandas.types import _create_converter_from_pandas, to_arrow_type
from pyspark.sql.pandas.types import to_arrow_type

# Handle empty schema (0 columns)
# Use dummy column + select([]) to preserve row count (PyArrow limitation workaround)
Expand All @@ -390,17 +387,11 @@ def convert(
else:
columns = list(data)

def convert_column(
col: Union["pd.Series", "pd.DataFrame"], field: StructField
) -> "pa.Array":
"""Convert a single column (Series or DataFrame) to an Arrow Array.

Uses field.name for error messages instead of series.name to avoid
copying the Series via rename() - a ~20% overhead on the hot path.
"""
converted: List[Union["pa.Array", "pa.ChunkedArray"]] = []
for col, field in zip(columns, schema.fields):
if isinstance(col, pd.DataFrame):
assert isinstance(field.dataType, StructType)
nested_batch = cls.convert(
nested_batch = cls.from_pandas(
col,
field.dataType,
timezone=timezone,
Expand All @@ -410,105 +401,159 @@ def convert_column(
assign_cols_by_name=assign_cols_by_name,
int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled,
ignore_unexpected_complex_type_values=ignore_unexpected_complex_type_values,
is_legacy=is_legacy,
use_legacy_error_handling=use_legacy_error_handling,
)
# Wrap the nested RecordBatch as a single StructArray column
return ArrowBatchTransformer.wrap_struct(nested_batch).column(0)
converted.append(ArrowBatchTransformer.wrap_struct(nested_batch).column(0))
else:
converted.append(
cls.convert(
col,
field,
timezone=timezone,
safecheck=safecheck,
arrow_cast=arrow_cast,
prefers_large_types=prefers_large_types,
int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled,
ignore_unexpected_complex_type_values=ignore_unexpected_complex_type_values,
use_legacy_error_handling=use_legacy_error_handling,
)
)

series = col
field_name = field.name
ret_type = field.dataType
# 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)

if isinstance(series.dtype, pd.CategoricalDtype):
series = series.astype(series.dtype.categories.dtype)
@classmethod
def convert(
cls,
series: "pd.Series",
field: StructField,
*,
timezone: Optional[str] = None,
safecheck: bool = True,
arrow_cast: bool = False,
prefers_large_types: bool = False,
int_to_decimal_coercion_enabled: bool = False,
ignore_unexpected_complex_type_values: bool = False,
use_legacy_error_handling: bool = False,
) -> Union["pa.Array", "pa.ChunkedArray"]:
"""Convert a pandas Series to an Arrow Array or ChunkedArray."""
return cls.convert_legacy(
series,
field,
timezone=timezone,
safecheck=safecheck,
arrow_cast=arrow_cast,
prefers_large_types=prefers_large_types,
int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled,
ignore_unexpected_complex_type_values=ignore_unexpected_complex_type_values,
use_legacy_error_handling=use_legacy_error_handling,
)

arrow_type = to_arrow_type(
ret_type, timezone=timezone, prefers_large_types=prefers_large_types
)
series = _create_converter_from_pandas(
ret_type,
timezone=timezone,
error_on_duplicated_field_names=False,
int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled,
ignore_unexpected_complex_type_values=ignore_unexpected_complex_type_values,
)(series)
@classmethod
def convert_legacy(
cls,
series: "pd.Series",
field: StructField,
*,
timezone: Optional[str] = None,
safecheck: bool = True,
arrow_cast: bool = False,
prefers_large_types: bool = False,
int_to_decimal_coercion_enabled: bool = False,
ignore_unexpected_complex_type_values: bool = False,
use_legacy_error_handling: bool = False,
) -> Union["pa.Array", "pa.ChunkedArray"]:
"""Convert a pandas Series to an Arrow Array or ChunkedArray."""
import pandas as pd
import pyarrow as pa

from pyspark.errors import PySparkTypeError, PySparkValueError
from pyspark.sql.pandas.types import _create_converter_from_pandas, to_arrow_type

field_name = field.name
ret_type = field.dataType

if isinstance(series.dtype, pd.CategoricalDtype):
series = series.astype(series.dtype.categories.dtype)

mask = None if hasattr(series.array, "__arrow_array__") else series.isnull()
arrow_type = to_arrow_type(
ret_type, timezone=timezone, prefers_large_types=prefers_large_types
)
series = _create_converter_from_pandas(
ret_type,
timezone=timezone,
error_on_duplicated_field_names=False,
int_to_decimal_coercion_enabled=int_to_decimal_coercion_enabled,
ignore_unexpected_complex_type_values=ignore_unexpected_complex_type_values,
)(series)

if is_legacy:
# Legacy pandas conversion path: broad ArrowException catch so
# that both ArrowInvalid AND ArrowTypeError (e.g. dict->struct)
# trigger the cast fallback.
mask = None if hasattr(series.array, "__arrow_array__") else series.isnull()

if use_legacy_error_handling:
# Legacy error handling: both ArrowInvalid and ArrowTypeError can
# trigger the cast fallback when arrow_cast is enabled.
try:
try:
try:
return pa.Array.from_pandas(
series, mask=mask, type=arrow_type, safe=safecheck
return pa.Array.from_pandas(series, mask=mask, type=arrow_type, safe=safecheck)
except pa.lib.ArrowException: # broad: includes ArrowTypeError
if arrow_cast:
return pa.Array.from_pandas(series, mask=mask).cast(
target_type=arrow_type, safe=safecheck
)
except pa.lib.ArrowException: # broad: includes ArrowTypeError
if arrow_cast:
return pa.Array.from_pandas(series, mask=mask).cast(
target_type=arrow_type, safe=safecheck
)
raise
except pa.lib.ArrowException as e:
error_msg = (
"Exception thrown when converting pandas.Series (%s) "
"with name '%s' to Arrow Array (%s)."
% (series.dtype, field_name, arrow_type)
raise
except pa.lib.ArrowException as e:
error_msg = (
"Exception thrown when converting pandas.Series (%s) "
"with name '%s' to Arrow Array (%s)." % (series.dtype, field_name, arrow_type)
)
if isinstance(e, TypeError):
raise PySparkTypeError(error_msg) from e
if safecheck:
error_msg += (
" It can be caused by overflows or other "
"unsafe conversions warned by Arrow. Arrow safe "
"type check can be disabled by using SQL config "
"`spark.sql.execution.pandas."
"convertToArrowArraySafely`."
)
if isinstance(e, TypeError):
raise PySparkTypeError(error_msg) from e
if safecheck:
error_msg += (
" It can be caused by overflows or other "
"unsafe conversions warned by Arrow. Arrow safe "
"type check can be disabled by using SQL config "
"`spark.sql.execution.pandas."
"convertToArrowArraySafely`."
)
raise PySparkValueError(error_msg) from e
else:
# Non-legacy path: only ArrowInvalid triggers the cast fallback.
# ArrowTypeError (e.g. string->decimal) must NOT be silently cast.
raise PySparkValueError(error_msg) from e
else:
# Non-legacy path: only ArrowInvalid triggers the cast fallback.
# ArrowTypeError must NOT be silently cast.
try:
try:
try:
return pa.Array.from_pandas(
series, mask=mask, type=arrow_type, safe=safecheck
return pa.Array.from_pandas(series, mask=mask, type=arrow_type, safe=safecheck)
except pa.lib.ArrowInvalid: # narrow: skip ArrowTypeError
if arrow_cast:
return pa.Array.from_pandas(series, mask=mask).cast(
target_type=arrow_type, safe=safecheck
)
except pa.lib.ArrowInvalid: # narrow: skip ArrowTypeError
if arrow_cast:
return pa.Array.from_pandas(series, mask=mask).cast(
target_type=arrow_type, safe=safecheck
)
raise
except TypeError as e:
raise PySparkTypeError(
f"Cannot convert the output value of the column "
f"'{field_name}' with type '{series.dtype}' to the "
f"specified return type of the column: '{arrow_type}'."
f" Please check if the data types match and try again."
) from e
except ValueError as e:
error_msg = (
f"Failed to convert the value of the column "
f"'{field_name}' with type '{series.dtype}' to Arrow "
f"type '{arrow_type}'."
raise
except TypeError as e:
raise PySparkTypeError(
f"Cannot convert the output value of the column "
f"'{field_name}' with type '{series.dtype}' to the "
f"specified return type of the column: '{arrow_type}'."
f" Please check if the data types match and try again."
) from e
except ValueError as e:
error_msg = (
f"Failed to convert the value of the column "
f"'{field_name}' with type '{series.dtype}' to Arrow "
f"type '{arrow_type}'."
)
if safecheck:
error_msg += (
" It can be caused by overflows or other unsafe "
"conversions warned by Arrow. Arrow safe type "
"check can be disabled by using SQL config "
"`spark.sql.execution.pandas."
"convertToArrowArraySafely`."
)
if safecheck:
error_msg += (
" It can be caused by overflows or other unsafe "
"conversions warned by Arrow. Arrow safe type "
"check can be disabled by using SQL config "
"`spark.sql.execution.pandas."
"convertToArrowArraySafely`."
)
raise PySparkValueError(error_msg) from e

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)
raise PySparkValueError(error_msg) from e


class LocalDataToArrowConversion:
Expand Down
Loading