From 43ad3675e4e9473430312b8014420e0cbff30e83 Mon Sep 17 00:00:00 2001 From: Spenser Sun Date: Thu, 3 Sep 2026 19:27:23 +0000 Subject: [PATCH] [SPARK-59219][PYTHON] Avoid copying the Series to name it in convert_numpy `ArrowArrayToPandasConversion.convert_numpy` ended with `series.rename(ser_name)`. `Series.rename` is not an in-place rename: pandas implements it as `self.copy(deep=False)` followed by setting the name, so every column paid a shallow Series copy - 4 new objects and ~69 Python calls - just to attach a name. The Series is created inside `convert_numpy`, so it can be named in place instead. Removes ~20-38us of fixed cost per column, which does not scale with row count. --- python/benchmarks/bench_arrow_to_pandas.py | 80 +++++++++++++++++++++ python/pyspark/sql/conversion.py | 4 +- python/pyspark/sql/tests/test_conversion.py | 15 ++++ 3 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 python/benchmarks/bench_arrow_to_pandas.py diff --git a/python/benchmarks/bench_arrow_to_pandas.py b/python/benchmarks/bench_arrow_to_pandas.py new file mode 100644 index 0000000000000..80e954b8ae0ec --- /dev/null +++ b/python/benchmarks/bench_arrow_to_pandas.py @@ -0,0 +1,80 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one or more +# contributor license agreements. See the NOTICE file distributed with +# this work for additional information regarding copyright ownership. +# The ASF licenses this file to You under the Apache License, Version 2.0 +# (the "License"); you may not use this file except in compliance with +# the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +""" +Microbenchmarks for ``ArrowBatchTransformer.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. + +Part of the per-batch cost is fixed per COLUMN and does not scale with row +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 +``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 +column is localized by pyarrow compute kernels -- so sweeping both shows how much +of a change is fixed per-column cost rather than per-row work. +""" + +import numpy as np +import pyarrow as pa + + +class ArrowBatchToPandasBenchmark: + """Benchmark ``ArrowBatchTransformer.to_pandas`` over a whole RecordBatch.""" + + params = [ + [128, 10000], + [1, 50], + ["long", "timestamp", "string"], + ] + 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.types import ( + LongType, + StringType, + StructField, + StructType, + TimestampType, + ) + + if col_type == "long": + column = pa.array(np.arange(n_rows, dtype=np.int64)) + spark_type = LongType() + elif col_type == "timestamp": + base = np.datetime64("2020-01-01T00:00:00", "us") + column = pa.array(base + np.arange(n_rows) * np.timedelta64(1, "s")).cast( + pa.timestamp("us", tz="UTC") + ) + spark_type = TimestampType() + elif col_type == "string": + column = pa.array([f"s{i:07d}" for i in range(n_rows)], type=pa.string()) + spark_type = StringType() + else: + raise ValueError(f"unknown col_type: {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 + + 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 0585a730fb087..f2fd2c6b0c3bf 100644 --- a/python/pyspark/sql/conversion.py +++ b/python/pyspark/sql/conversion.py @@ -2045,4 +2045,6 @@ def convert_numpy( else: # pragma: no cover assert False, f"Need converter for {spark_type} but failed to find one." - return series.rename(ser_name) + # `series` is created in this method, so naming it in place is safe; rename() copies it. + series.name = ser_name + return series diff --git a/python/pyspark/sql/tests/test_conversion.py b/python/pyspark/sql/tests/test_conversion.py index c5e9fb55aba20..f660ffa0e7d32 100644 --- a/python/pyspark/sql/tests/test_conversion.py +++ b/python/pyspark/sql/tests/test_conversion.py @@ -46,6 +46,7 @@ StringType, StructField, StructType, + TimestampNTZType, TimestampType, UserDefinedType, VariantType, @@ -719,6 +720,20 @@ def test_arrow_array_localize_tz(self): @unittest.skipIf(not have_pyarrow, pyarrow_requirement_message) class ArrowArrayToPandasConversionTests(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. + import pyarrow as pa + + for pa_type in [pa.timestamp("us", tz="UTC"), pa.timestamp("s"), pa.timestamp("ns")]: + ts = pa.array([datetime.datetime(2020, 6, 15, 12, 30)], type=pa.timestamp("us")).cast( + pa_type + ) + 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") + self.assertEqual(result.name, "tscol", f"name lost for {pa_type}") + def test_udt_convert_numpy(self): import pyarrow as pa