You are viewing a plain text version of this content. The canonical link for it is here.
Posted to reviews@spark.apache.org by "Yicong-Huang (via GitHub)" <gi...@apache.org> on 2026/04/23 00:30:36 UTC
[PR] [SPARK-56477][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_UDF to be self-contained in read_udfs [spark]
Yicong-Huang opened a new pull request, #55495:
URL: https://github.com/apache/spark/pull/55495
### What changes were proposed in this pull request?
Refactor `SQL_GROUPED_MAP_PANDAS_UDF` to be self-contained in `read_udfs()`.
### Why are the changes needed?
Part of SPARK-55388 (Refactor PythonEvalType processing logic). Making each eval type self-contained in `read_udfs()` improves readability and makes it easier to reason about the data flow for each eval type independently.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Existing tests. No behavior change.
### 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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56477][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_UDF [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on PR #55495:
URL: https://github.com/apache/spark/pull/55495#issuecomment-4352316590
merged to master
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56477][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_UDF [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng closed pull request #55495: [SPARK-56477][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_UDF
URL: https://github.com/apache/spark/pull/55495
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org
Re: [PR] [SPARK-56477][PYTHON] Refactor SQL_GROUPED_MAP_PANDAS_UDF [spark]
Posted by "zhengruifeng (via GitHub)" <gi...@apache.org>.
zhengruifeng commented on code in PR #55495:
URL: https://github.com/apache/spark/pull/55495#discussion_r3165269271
##########
python/pyspark/worker.py:
##########
@@ -2943,6 +2914,77 @@ def grouped_func(
# profiling is not supported for UDF
return grouped_func, None, ser, ser
+ if eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF:
+ import pyarrow as pa
+ import pandas as pd
+
+ assert num_udfs == 1, "One GROUPED_MAP_PANDAS UDF expected here."
+ grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0]
+ parsed_offsets = extract_key_value_indexes(arg_offsets)
+ assert len(parsed_offsets) == 1, "Expected one pair of offsets for GROUPED_MAP_PANDAS UDF."
+
+ key_offsets = parsed_offsets[0][0]
+ value_offsets = parsed_offsets[0][1]
+ output_schema = StructType([StructField("_0", return_type)])
+
+ def grouped_func(
+ split_index: int,
+ data: Iterator[Iterator[pa.RecordBatch]],
+ ) -> Iterator[pa.RecordBatch]:
+ """Apply groupBy Pandas UDF (non-iterator variant).
+
+ The explicit ``del`` calls below keep peakmem bounded across
+ groups. Without them, generator locals from the previous
+ iteration stay bound on the frame until each statement in
+ the next iteration rebinds its slot, so the input-side
+ DataFrames overlap with the next group's allocations and
+ the working set grows unbounded on wide-column, large-group
+ inputs. ``del result`` runs on resume from yield, before
+ ``data.__next__()`` is asked for the next group.
+ """
+ for group in data:
+ all_batches = list(group)
+ if all_batches:
+ table = pa.Table.from_batches(all_batches).combine_chunks()
+ else:
+ table = pa.table({})
+ all_series = ArrowBatchTransformer.to_pandas(
+ table,
+ timezone=runner_conf.timezone,
+ prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
+ )
+ value_df = pd.concat([all_series[o] for o in value_offsets], axis=1)
+
+ if num_udf_args == 1:
+ result = grouped_udf(value_df)
+ else:
+ key = tuple(all_series[o].iloc[0] for o in key_offsets)
+ result = grouped_udf(key, value_df)
+
+ del all_batches, table, all_series, value_df
+
+ verify_pandas_result(
+ result,
+ return_type,
+ runner_conf.assign_cols_by_name,
+ truncate_return_schema=False,
+ )
+
+ yield PandasToArrowConversion.convert(
+ [result],
+ output_schema,
+ timezone=runner_conf.timezone,
+ safecheck=runner_conf.safecheck,
+ arrow_cast=True,
+ prefers_large_types=runner_conf.use_large_var_types,
Review Comment:
The OLD output path went through `GroupPandasUDFSerializer.__init__`, which omits `prefers_large_types` when calling `super().__init__()` and so defaulted to `False`. The new code respects `spark.sql.execution.arrow.useLargeVarTypes`. Since `ArrowUtils.fromArrowSchema` maps both `Utf8` and `LargeUtf8` to `StringType` (`sql/api/.../ArrowUtils.scala:82,84`), this is wire-format-only and user-invisible at the Spark type level — but the PR description's "No" to user-facing change is no longer strictly correct. Either:
- Note in the PR description that this aligns with the Arrow analogue (SPARK-55608) and is intentional, or
- Pass `prefers_large_types=False` here to preserve the exact pre-PR wire format.
Note also the resulting divergence with `SQL_GROUPED_MAP_PANDAS_ITER_UDF`, which still uses `GroupPandasUDFSerializer` with the hardcoded `False` default (worker.py:2391-2397). Once the iter variant migrates as part of SPARK-55388, this divergence resolves.
##########
python/pyspark/worker.py:
##########
@@ -2413,10 +2387,7 @@ def read_udfs(pickleSer, infile, eval_type, runner_conf, eval_conf):
prefer_int_ext_dtype=runner_conf.prefer_int_ext_dtype,
int_to_decimal_coercion_enabled=runner_conf.int_to_decimal_coercion_enabled,
)
- elif (
- eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF
- or eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF
- ):
+ elif eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_ITER_UDF:
ser = GroupPandasUDFSerializer(
Review Comment:
Heads up — after this PR, `GroupPandasUDFSerializer` is only used by `SQL_GROUPED_MAP_PANDAS_ITER_UDF`. The class comment at `python/pyspark/sql/pandas/serializers.py:657` (`# Serializer for SQL_GROUPED_MAP_PANDAS_UDF, SQL_GROUPED_MAP_PANDAS_ITER_UDF`) is now stale. Worth updating in this PR or noting as a follow-up.
##########
python/pyspark/worker.py:
##########
@@ -2943,6 +2914,77 @@ def grouped_func(
# profiling is not supported for UDF
return grouped_func, None, ser, ser
+ if eval_type == PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF:
+ import pyarrow as pa
+ import pandas as pd
+
+ assert num_udfs == 1, "One GROUPED_MAP_PANDAS UDF expected here."
+ grouped_udf, arg_offsets, return_type, num_udf_args = udfs[0]
+ parsed_offsets = extract_key_value_indexes(arg_offsets)
+ assert len(parsed_offsets) == 1, "Expected one pair of offsets for GROUPED_MAP_PANDAS UDF."
+
+ key_offsets = parsed_offsets[0][0]
+ value_offsets = parsed_offsets[0][1]
+ output_schema = StructType([StructField("_0", return_type)])
+
+ def grouped_func(
+ split_index: int,
+ data: Iterator[Iterator[pa.RecordBatch]],
+ ) -> Iterator[pa.RecordBatch]:
+ """Apply groupBy Pandas UDF (non-iterator variant).
+
+ The explicit ``del`` calls below keep peakmem bounded across
+ groups. Without them, generator locals from the previous
+ iteration stay bound on the frame until each statement in
+ the next iteration rebinds its slot, so the input-side
+ DataFrames overlap with the next group's allocations and
+ the working set grows unbounded on wide-column, large-group
+ inputs. ``del result`` runs on resume from yield, before
+ ``data.__next__()`` is asked for the next group.
+ """
Review Comment:
The peakmem rationale is plausible (input-side `del` before the convert call lowers peak vs. a helper-function approach), but the PR description benchmarks are wall-clock only. Could you add a peakmem comparison (e.g., `tracemalloc.get_traced_memory()` peaks or ASV's `peakmem_*` benchmarks) for the wide-column / large-group scenarios this docstring describes? That would close the loop on why the inline-with-`del` form was preferred over the simpler helper-function variant from the first commit, and protect against future edits inadvertently dropping a `del` and regressing peakmem.
--
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: reviews-unsubscribe@spark.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org