This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new ad978f45d6cb [SPARK-58230][PS] Batch pandas-on-Spark column projections
ad978f45d6cb is described below
commit ad978f45d6cb9179b4d46da69f99a37adba56a8a
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Tue Jul 21 15:58:28 2026 +0800
[SPARK-58230][PS] Batch pandas-on-Spark column projections
### What changes were proposed in this pull request?
Batch independent pandas-on-Spark column projections with `withColumns`:
- Add all missing columns together in outer `concat`, while skipping the
call when no columns are missing.
- Add the three percentile columns together for each aggregate column in
`DataFrameGroupBy.describe`.
- Add group-key aliases together in `GroupBy.idxmax` and `GroupBy.idxmin`.
### Why are the changes needed?
In classic mode, the previous loops made one Py4J call and one projection
per missing, percentile, or group-key column. Batching the independent
expressions reduces driver overhead and logical-plan construction and analysis
work.
In Spark Connect mode, each batch produces one protobuf transformation
instead of one transformation per column, reducing serialized-plan and
transformation/analyze overhead. Results are unchanged.
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
CI
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-5)
Closes #57390 from zhengruifeng/pandas-concat-withcolumns-dev2.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit cb252172034a6fe5832bd0cea257657c0891e1a1)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
python/pyspark/pandas/groupby.py | 26 ++++++++++++++++----------
python/pyspark/pandas/namespace.py | 6 ++++--
2 files changed, 20 insertions(+), 12 deletions(-)
diff --git a/python/pyspark/pandas/groupby.py b/python/pyspark/pandas/groupby.py
index 9e02b73d5c64..9907d9553245 100644
--- a/python/pyspark/pandas/groupby.py
+++ b/python/pyspark/pandas/groupby.py
@@ -2459,8 +2459,10 @@ class GroupBy(Generic[FrameLike], metaclass=ABCMeta):
groupkey_names = ["__groupkey_{}__".format(i) for i in
range(len(self._groupkeys))]
sdf = self._psdf._internal.spark_frame
- for s, name in zip(self._groupkeys, groupkey_names):
- sdf = sdf.withColumn(name, s.spark.column)
+ if groupkey_names:
+ sdf = sdf.withColumns(
+ {name: s.spark.column for s, name in zip(self._groupkeys,
groupkey_names)}
+ )
index = self._psdf._internal.index_spark_column_names[0]
index_spark_type = self._psdf._internal.index_fields[0].spark_type
@@ -2554,8 +2556,10 @@ class GroupBy(Generic[FrameLike], metaclass=ABCMeta):
groupkey_names = ["__groupkey_{}__".format(i) for i in
range(len(self._groupkeys))]
sdf = self._psdf._internal.spark_frame
- for s, name in zip(self._groupkeys, groupkey_names):
- sdf = sdf.withColumn(name, s.spark.column)
+ if groupkey_names:
+ sdf = sdf.withColumns(
+ {name: s.spark.column for s, name in zip(self._groupkeys,
groupkey_names)}
+ )
index = self._psdf._internal.index_spark_column_names[0]
index_spark_type = self._psdf._internal.index_fields[0].spark_type
@@ -4067,12 +4071,14 @@ class DataFrameGroupBy(GroupBy[DataFrame]):
# Split "quartiles" columns into first, second, and third quartiles.
for label in agg_column_labels:
quartiles_col = name_like_string(tuple(list(label) +
["quartiles"]))
- for i, percentile in enumerate(formatted_percentiles):
- sdf = sdf.withColumn(
- name_like_string(tuple(list(label) + [percentile])),
- scol_for(sdf, quartiles_col)[i],
- )
- sdf = sdf.drop(quartiles_col)
+ sdf = sdf.withColumns(
+ {
+ name_like_string(tuple(list(label) + [percentile])):
scol_for(
+ sdf, quartiles_col
+ )[i]
+ for i, percentile in enumerate(formatted_percentiles)
+ }
+ ).drop(quartiles_col)
# Reorder columns lexicographically by agg column followed by stats.
stats = ["count", "mean", "std", "min"] + formatted_percentiles +
["max"]
diff --git a/python/pyspark/pandas/namespace.py
b/python/pyspark/pandas/namespace.py
index 15dd70c0234d..05875347a375 100644
--- a/python/pyspark/pandas/namespace.py
+++ b/python/pyspark/pandas/namespace.py
@@ -2732,8 +2732,10 @@ def concat(
# TODO: NaN and None difference for missing values. pandas
seems to be filling NaN.
sdf = psdf._internal.resolved_copy.spark_frame
- for label in columns_to_add:
- sdf = sdf.withColumn(name_like_string(label), F.lit(None))
+ if columns_to_add:
+ sdf = sdf.withColumns(
+ {name_like_string(label): F.lit(None) for label in
columns_to_add}
+ )
data_columns = psdf._internal.data_spark_column_names + [
name_like_string(label) for label in columns_to_add
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
