diff --git a/python/benchmarks/bench_arrow_to_pandas.py b/python/benchmarks/bench_arrow_to_pandas.py index 80e954b8ae0ec..ef12a32224ae5 100644 --- a/python/benchmarks/bench_arrow_to_pandas.py +++ b/python/benchmarks/bench_arrow_to_pandas.py @@ -16,7 +16,7 @@ # """ -Microbenchmarks for ``ArrowBatchTransformer.to_pandas``, the hot path of pandas +Microbenchmarks for ``ArrowToPandasConversion.to_pandas``, the hot path of pandas UDF inputs: every pandas UDF eval type calls it once per batch to build the Series it passes to the user's function. @@ -24,7 +24,7 @@ count, so ``n_cols`` is swept alongside ``n_rows``: wide batches and small batches are the shapes where that fixed cost dominates. -``ArrowArrayToPandasConversion.convert`` routes each column by type. ``long`` and +``ArrowToPandasConversion.convert`` routes each column by type. ``long`` and ``timestamp`` are in the ``_prefer_convert_numpy`` allowlist and take ``convert_numpy``; ``string`` is not, and takes ``convert_legacy``. The two allowlist types differ by an order of magnitude in conversion cost -- a timestamp @@ -37,7 +37,7 @@ class ArrowBatchToPandasBenchmark: - """Benchmark ``ArrowBatchTransformer.to_pandas`` over a whole RecordBatch.""" + """Benchmark ``ArrowToPandasConversion.to_pandas`` over a whole RecordBatch.""" params = [ [128, 10000], @@ -47,7 +47,7 @@ class ArrowBatchToPandasBenchmark: param_names = ["n_rows", "n_cols", "col_type"] def setup(self, n_rows, n_cols, col_type): - from pyspark.sql.conversion import ArrowBatchTransformer + from pyspark.sql.conversion import ArrowToPandasConversion from pyspark.sql.types import ( LongType, StringType, @@ -74,7 +74,7 @@ def setup(self, n_rows, n_cols, col_type): names = [f"c{i}" for i in range(n_cols)] self.batch = pa.RecordBatch.from_arrays([column] * n_cols, names) self.schema = StructType([StructField(name, spark_type) for name in names]) - self.to_pandas = ArrowBatchTransformer.to_pandas + self.to_pandas = ArrowToPandasConversion.to_pandas def time_batch_to_pandas(self, n_rows, n_cols, col_type): self.to_pandas(self.batch, timezone="UTC", schema=self.schema) diff --git a/python/pyspark/sql/conversion.py b/python/pyspark/sql/conversion.py index fc03acf026436..1af7bbd1278a0 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -74,7 +74,7 @@ class ArrowBatchTransformer: """ - Pure functions that transform RecordBatch -> RecordBatch. + Pure functions that transform Arrow RecordBatches and Tables. They should have no side effects (no I/O, no writing to streams). """ @@ -244,64 +244,6 @@ def enforce_schema( return pa.Table.from_arrays(coerced_arrays, names=output_names) return pa.RecordBatch.from_arrays(coerced_arrays, names=output_names) - @classmethod - def to_pandas( - cls, - batch: Union["pa.RecordBatch", "pa.Table"], - timezone: str, - schema: Optional["StructType"] = None, - struct_in_pandas: str = "dict", - ndarray_as_list: bool = False, - prefer_int_ext_dtype: bool = False, - df_for_struct: bool = False, - ) -> List[Union["pd.Series", "pd.DataFrame"]]: - """ - Convert a RecordBatch or Table to a list of pandas Series. - - Parameters - ---------- - batch : pa.RecordBatch or pa.Table - The Arrow RecordBatch or Table to convert. - timezone : str - Timezone for timestamp conversion. - schema : StructType, optional - Spark schema for type conversion. If None, types are inferred from Arrow. - struct_in_pandas : str - How to represent struct in pandas ("dict", "row", etc.) - ndarray_as_list : bool - Whether to convert ndarray as list. - prefer_int_ext_dtype : bool, optional - Whether to convert integers to Pandas ExtensionDType. - df_for_struct : bool - If True, convert struct columns to DataFrame instead of Series. - - Returns - ------- - List[Union[pd.Series, pd.DataFrame]] - List of pandas Series (or DataFrame if df_for_struct=True), one for each column. - """ - import pandas as pd - - if batch.num_columns == 0: - return [pd.Series([pyspark._NoValue] * batch.num_rows)] - - if schema is None: - schema = from_arrow_schema(batch.schema) - - return [ - ArrowArrayToPandasConversion.convert( - batch.column(i), - schema[i].dataType, - ser_name=schema[i].name, - timezone=timezone, - struct_in_pandas=struct_in_pandas, - ndarray_as_list=ndarray_as_list, - prefer_int_ext_dtype=prefer_int_ext_dtype, - df_for_struct=df_for_struct, - ) - for i in range(batch.num_columns) - ] - class PandasToArrowConversion: """ @@ -1750,16 +1692,68 @@ def convert_func(arr: pa.Array) -> pa.Array: ) -class ArrowArrayToPandasConversion: +class ArrowToPandasConversion: + """ + Conversion utilities from Arrow batches and arrays to pandas for UDF execution. """ - Conversion utilities for converting PyArrow Arrays and ChunkedArrays to pandas. - This class provides methods to convert PyArrow columnar data structures to pandas - Series or DataFrames, with support for Spark-specific type handling and conversions. + @classmethod + def to_pandas( + cls, + batch: Union["pa.RecordBatch", "pa.Table"], + timezone: str, + schema: Optional["StructType"] = None, + struct_in_pandas: str = "dict", + ndarray_as_list: bool = False, + prefer_int_ext_dtype: bool = False, + df_for_struct: bool = False, + ) -> List[Union["pd.Series", "pd.DataFrame"]]: + """ + Convert a RecordBatch or Table to a list of pandas Series. - The class is primarily used by PySpark's Arrow-based serializers for UDF execution, - where Arrow data needs to be converted to pandas for Python UDF processing. - """ + Parameters + ---------- + batch : pa.RecordBatch or pa.Table + The Arrow RecordBatch or Table to convert. + timezone : str + Timezone for timestamp conversion. + schema : StructType, optional + Spark schema for type conversion. If None, types are inferred from Arrow. + struct_in_pandas : str + How to represent struct in pandas ("dict", "row", etc.) + ndarray_as_list : bool + Whether to convert ndarray as list. + prefer_int_ext_dtype : bool, optional + Whether to convert integers to Pandas ExtensionDType. + df_for_struct : bool + If True, convert struct columns to DataFrame instead of Series. + + Returns + ------- + List[Union[pd.Series, pd.DataFrame]] + List of pandas Series (or DataFrame if df_for_struct=True), one for each column. + """ + import pandas as pd + + if batch.num_columns == 0: + return [pd.Series([pyspark._NoValue] * batch.num_rows)] + + if schema is None: + schema = from_arrow_schema(batch.schema) + + return [ + cls.convert( + batch.column(i), + schema[i].dataType, + ser_name=schema[i].name, + timezone=timezone, + struct_in_pandas=struct_in_pandas, + ndarray_as_list=ndarray_as_list, + prefer_int_ext_dtype=prefer_int_ext_dtype, + df_for_struct=df_for_struct, + ) + for i in range(batch.num_columns) + ] @classmethod def convert( diff --git a/python/pyspark/sql/tests/test_conversion.py b/python/pyspark/sql/tests/test_conversion.py index f660ffa0e7d32..1eaa75c29d2f2 100644 --- a/python/pyspark/sql/tests/test_conversion.py +++ b/python/pyspark/sql/tests/test_conversion.py @@ -23,9 +23,9 @@ from pyspark.errors import PySparkRuntimeError, PySparkTypeError, PySparkValueError from pyspark.sql.conversion import ( ArrowArrayConversion, - ArrowArrayToPandasConversion, ArrowBatchTransformer, ArrowTableToRowsConversion, + ArrowToPandasConversion, LocalDataToArrowConversion, PandasToArrowConversion, ) @@ -719,7 +719,7 @@ def test_arrow_array_localize_tz(self): @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) -class ArrowArrayToPandasConversionTests(unittest.TestCase): +class ArrowToPandasConversionTests(unittest.TestCase): def test_convert_numpy_ser_name_survives_preprocess_time(self): # convert_numpy reads the Arrow field name before preprocess_time, because the # pa.compute kernels it runs for timestamps return a new array with no field name. @@ -731,7 +731,7 @@ def test_convert_numpy_ser_name_survives_preprocess_time(self): ) col = pa.RecordBatch.from_arrays([ts], ["tscol"]).column(0) spark_type = TimestampType() if pa_type.tz is not None else TimestampNTZType() - result = ArrowArrayToPandasConversion.convert_numpy(col, spark_type, timezone="UTC") + result = ArrowToPandasConversion.convert_numpy(col, spark_type, timezone="UTC") self.assertEqual(result.name, "tscol", f"name lost for {pa_type}") def test_udt_convert_numpy(self): @@ -741,7 +741,7 @@ def test_udt_convert_numpy(self): # basic conversion with nulls arr = pa.array([[1.0, 2.0], None, [3.0, 4.0]], type=pa.list_(pa.float64())) - result = ArrowArrayToPandasConversion.convert_numpy(arr, udt, ser_name="my_point") + result = ArrowToPandasConversion.convert_numpy(arr, udt, ser_name="my_point") self.assertIsInstance(result.iloc[0], ExamplePoint) self.assertEqual(result.iloc[0], ExamplePoint(1.0, 2.0)) self.assertIsNone(result.iloc[1]) @@ -749,13 +749,13 @@ def test_udt_convert_numpy(self): self.assertEqual(result.name, "my_point") # empty - result = ArrowArrayToPandasConversion.convert_numpy( + result = ArrowToPandasConversion.convert_numpy( pa.array([], type=pa.list_(pa.float64())), udt ) self.assertEqual(len(result), 0) # PythonOnlyUDT - result = ArrowArrayToPandasConversion.convert_numpy( + result = ArrowToPandasConversion.convert_numpy( pa.array([[5.0, 6.0]], type=pa.list_(pa.float64())), PythonOnlyUDT() ) self.assertIsInstance(result.iloc[0], PythonOnlyPoint) @@ -767,7 +767,7 @@ def test_udt_chunked_array(self): chunk1 = pa.array([[1.0, 2.0]], type=pa.list_(pa.float64())) chunk2 = pa.array([[3.0, 4.0]], type=pa.list_(pa.float64())) chunked = pa.chunked_array([chunk1, chunk2]) - result = ArrowArrayToPandasConversion.convert_numpy(chunked, ExamplePointUDT()) + result = ArrowToPandasConversion.convert_numpy(chunked, ExamplePointUDT()) self.assertEqual(result.iloc[0], ExamplePoint(1.0, 2.0)) self.assertEqual(result.iloc[1], ExamplePoint(3.0, 4.0)) @@ -790,7 +790,7 @@ def test_variant_convert_numpy(self): ], type=variant_type, ) - result = ArrowArrayToPandasConversion.convert_numpy(arr, VariantType(), ser_name="v") + result = ArrowToPandasConversion.convert_numpy(arr, VariantType(), ser_name="v") self.assertIsInstance(result.iloc[0], VariantVal) self.assertEqual(result.iloc[0].value, b"\x01") self.assertEqual(result.iloc[0].metadata, b"\x02") @@ -800,7 +800,7 @@ def test_variant_convert_numpy(self): self.assertEqual(result.name, "v") # empty - result = ArrowArrayToPandasConversion.convert_numpy( + result = ArrowToPandasConversion.convert_numpy( pa.array([], type=variant_type), VariantType() ) self.assertEqual(len(result), 0) @@ -832,14 +832,14 @@ def test_geography_convert_numpy(self): ], type=geography_type, ) - result = ArrowArrayToPandasConversion.convert_numpy(arr, GeographyType(4326), ser_name="g") + result = ArrowToPandasConversion.convert_numpy(arr, GeographyType(4326), ser_name="g") self.assertEqual(result.iloc[0], Geography(wkb1, 4326)) self.assertIsNone(result.iloc[1]) self.assertEqual(result.iloc[2], Geography(wkb2, 4326)) self.assertEqual(result.name, "g") # empty - result = ArrowArrayToPandasConversion.convert_numpy( + result = ArrowToPandasConversion.convert_numpy( pa.array([], type=geography_type), GeographyType(4326) ) self.assertEqual(len(result), 0) @@ -871,14 +871,14 @@ def test_geometry_convert_numpy(self): ], type=geometry_type, ) - result = ArrowArrayToPandasConversion.convert_numpy(arr, GeometryType(0), ser_name="g") + result = ArrowToPandasConversion.convert_numpy(arr, GeometryType(0), ser_name="g") self.assertEqual(result.iloc[0], Geometry(wkb1, 0)) self.assertIsNone(result.iloc[1]) self.assertEqual(result.iloc[2], Geometry(wkb2, 0)) self.assertEqual(result.name, "g") # empty - result = ArrowArrayToPandasConversion.convert_numpy( + result = ArrowToPandasConversion.convert_numpy( pa.array([], type=geometry_type), GeometryType(0) ) self.assertEqual(len(result), 0) diff --git a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py index ef35868bba08f..64994b3dbdc5e 100644 --- a/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py +++ b/python/pyspark/tests/upstream/pyarrow/test_pyarrow_arrow_to_pandas_non_default.py @@ -239,7 +239,7 @@ class PyArrowArrayToPandasZeroCopyTests(_PyArrowToPandasTestBase): asks pandas to keep pointing at the Arrow buffers instead of materializing them into NumPy -- avoiding the copy is the point of that backend. Its own golden file records the result, so the two can be read side by side. PySpark takes this path - in ``ArrowArrayToPandasConversion.convert_numpy`` + in ``ArrowToPandasConversion.convert_numpy`` (``python/pyspark/sql/conversion.py``). """ @@ -479,7 +479,7 @@ class PyArrowArrayToPandasIntegerObjectNullsTests(_PyArrowToPandasTestBase): large value silently changes. ``integer_object_nulls=True`` keeps ``object`` dtype (Python ``int`` and ``None``) instead, preserving the values. - PySpark passes it in ``ArrowArrayToPandasConversion.convert_legacy`` + PySpark passes it in ``ArrowToPandasConversion.convert_legacy`` (``python/pyspark/sql/conversion.py``), bundled with ``date_as_object`` and ``coerce_temporal_nanoseconds``, then narrows the object Series to a nullable extension dtype (``Int8Dtype`` .. ``Int64Dtype``) -- the only bridge from Arrow to diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py index 878b7bb63a126..465b54efcdd3a 100644 --- a/python/pyspark/worker.py +++ b/python/pyspark/worker.py @@ -73,6 +73,7 @@ from pyspark.sql.conversion import ( ArrowBatchTransformer, ArrowTableToRowsConversion, + ArrowToPandasConversion, LocalDataToArrowConversion, PandasToArrowConversion, ) @@ -1499,7 +1500,7 @@ def func(split_index: int, data: Iterator["pa.RecordBatch"]) -> Iterator["pa.Rec for batch in data: # Deserialize the Arrow batch into a list of pandas Series (one per # input column), then call eval once per input row. - series_list = ArrowBatchTransformer.to_pandas( + series_list = ArrowToPandasConversion.to_pandas( batch, timezone=runner_conf.timezone, schema=eval_conf.input_type, @@ -2000,9 +2001,8 @@ def _elementwise_flatten_column(flat, element_type, is_pandas, runner_conf): """ if not is_pandas: return flat - from pyspark.sql.conversion import ArrowArrayToPandasConversion - return ArrowArrayToPandasConversion.convert( + return ArrowToPandasConversion.convert( flat, element_type, timezone=runner_conf.timezone, @@ -2529,7 +2529,7 @@ def grouped_func( if not batch_list: continue table = pa.Table.from_batches(batch_list).combine_chunks() - all_series = ArrowBatchTransformer.to_pandas( + all_series = ArrowToPandasConversion.to_pandas( table, timezone=runner_conf.timezone, prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, @@ -2570,7 +2570,7 @@ def extract_series( # Convert one RecordBatch to a pandas Series per column, then select args: # - pd.Series for a single column # - tuple[pd.Series, ...] for multiple columns - all_series = ArrowBatchTransformer.to_pandas( + all_series = ArrowToPandasConversion.to_pandas( batch, timezone=runner_conf.timezone, prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, @@ -2781,7 +2781,7 @@ def grouped_func( if not batch_list: continue table = pa.Table.from_batches(batch_list).combine_chunks() - all_series = ArrowBatchTransformer.to_pandas( + all_series = ArrowToPandasConversion.to_pandas( table, timezone=runner_conf.timezone, prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, @@ -3005,7 +3005,7 @@ def grouped_func( table = pa.Table.from_batches(all_batches).combine_chunks() else: table = pa.table({}) - all_series = ArrowBatchTransformer.to_pandas( + all_series = ArrowToPandasConversion.to_pandas( table, timezone=runner_conf.timezone, prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, @@ -3070,7 +3070,7 @@ def grouped_func( for group in data: group_iter = iter(group) # Read the first batch to extract grouping keys. - first_series = ArrowBatchTransformer.to_pandas( + first_series = ArrowToPandasConversion.to_pandas( next(group_iter), timezone=runner_conf.timezone, prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, @@ -3079,7 +3079,7 @@ def grouped_func( def dataframe_iter(): yield pd.concat([first_series[o] for o in value_offsets], axis=1) for batch in group_iter: - series = ArrowBatchTransformer.to_pandas( + series = ArrowToPandasConversion.to_pandas( batch, timezone=runner_conf.timezone, prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, @@ -3191,7 +3191,7 @@ def dataframe_iter(): # MapInBatchEvaluatorFactory); convert lazily so peakmem stays # bounded by one batch. for batch in data: - yield ArrowBatchTransformer.to_pandas( + yield ArrowToPandasConversion.to_pandas( batch, timezone=runner_conf.timezone, prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, @@ -3264,12 +3264,12 @@ def cogrouped_func( for left_batches, right_batches in data: left_table = pa.Table.from_batches(left_batches) right_table = pa.Table.from_batches(right_batches) - left_series = ArrowBatchTransformer.to_pandas( + left_series = ArrowToPandasConversion.to_pandas( left_table, timezone=runner_conf.timezone, prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, ) - right_series = ArrowBatchTransformer.to_pandas( + right_series = ArrowToPandasConversion.to_pandas( right_table, timezone=runner_conf.timezone, prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, @@ -3452,7 +3452,7 @@ def _evaluate_batch_udf_legacy(udf_func, rows): def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.RecordBatch]: for input_batch in data: # --- Input: Arrow -> pandas columns --- - pandas_columns = ArrowBatchTransformer.to_pandas( + pandas_columns = ArrowToPandasConversion.to_pandas( input_batch, timezone=runner_conf.timezone, schema=eval_conf.input_type, @@ -3931,7 +3931,7 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record num_rows = input_batch.num_rows # --- Input: Arrow -> pandas Series (struct columns become DataFrames) --- - pandas_columns = ArrowBatchTransformer.to_pandas( + pandas_columns = ArrowToPandasConversion.to_pandas( input_batch, timezone=runner_conf.timezone, struct_in_pandas="dict", @@ -4003,7 +4003,7 @@ def func(split_index: int, data: Iterator[pa.RecordBatch]) -> Iterator[pa.Record def extract_args(batch: pa.RecordBatch): nonlocal num_input_rows # Input: Arrow -> pandas Series (struct columns become DataFrames) - pandas_columns = ArrowBatchTransformer.to_pandas( + pandas_columns = ArrowToPandasConversion.to_pandas( batch, timezone=runner_conf.timezone, struct_in_pandas="dict", @@ -4120,7 +4120,7 @@ def row_stream(): ) total_rows += batch.num_rows average_arrow_row_size = total_bytes / total_rows - data_pandas = ArrowBatchTransformer.to_pandas( + data_pandas = ArrowToPandasConversion.to_pandas( batch, timezone=runner_conf.timezone, prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, @@ -4265,7 +4265,7 @@ def flatten_columns(cur_batch: "pa.RecordBatch", col_name: str) -> "pa.Table": return pa.Table.from_arrays(field_arrays, names=field_names) def to_pandas(table: "pa.Table") -> list: - return ArrowBatchTransformer.to_pandas( + return ArrowToPandasConversion.to_pandas( table, timezone=runner_conf.timezone, prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype, @@ -4455,7 +4455,7 @@ def convert_results( ) def to_pandas(batch: "pa.RecordBatch") -> list: - return ArrowBatchTransformer.to_pandas( + return ArrowToPandasConversion.to_pandas( batch, timezone=runner_conf.timezone, struct_in_pandas="dict",