This is an automated email from the ASF dual-hosted git repository.
HyukjinKwon pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new e068d117232c [SPARK-58890][SQL][PYTHON][CONNECT] Support incremental
Python aggregators over window functions
e068d117232c is described below
commit e068d117232c267365437f1b5fb33b1459355447
Author: Hyukjin Kwon <[email protected]>
AuthorDate: Thu Aug 20 12:58:35 2026 +0900
[SPARK-58890][SQL][PYTHON][CONNECT] Support incremental Python aggregators
over window functions
### What changes were proposed in this pull request?
A follow-up to
[SPARK-58736](https://issues.apache.org/jira/browse/SPARK-58736) (incremental
Python
`Aggregator` / `udaf`). This lets an incremental aggregator be used as a
**window function** --
`df.withColumn("m", udaf(agg)(col).over(window))` -- which previously
failed analysis with
`UNSUPPORTED_EXPR_FOR_WINDOW`.
A window has no shuffle, so it reuses neither the map-side PARTIAL nor the
post-shuffle FINAL
two-stage eval type. A new eval type `SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF`
(257) routes the
aggregator through the **existing Arrow window operator**
(`ArrowWindowPythonExec`), which sends
each frame to the Python worker -- the whole partition for an unbounded
frame, or per-row
`[begin, end)` slices for a bounded one. The worker folds the frame's rows
with `reduce` from a
fresh `zero` and produces the value with `finish` (one output value per
input row); `merge` is not
used on the window path.
Concretely:
- Route `PythonAggregate` to the Python window path:
`PythonUDF.isWindowPandasUDF` and
`WindowFunctionType.pythonEvalType` now recognize it, and the
`WindowResolution` reject guard is
removed.
- Generalize `ArrowWindowPythonEvaluatorFactory` from
`asInstanceOf[PythonUDAF]` to
`PythonFuncExpression` (both `PythonUDAF` and `PythonAggregate` extend
it) and add the new eval
type to `ArrowWindowPythonExec`'s supported set. All the frame-boundary
machinery is reused.
- Add the worker handler for the new eval type in `worker.py`.
No Spark Connect-specific change is needed: a window-over-aggregator
resolves through the same
server-side Catalyst rules.
This PR also **relocates the `udaf` factory to `pyspark.sql.functions`**
(next to `udf` /
`pandas_udf`), matching Scala's `functions.udaf`, and stops exposing it from
`pyspark.sql.aggregator`. `Aggregator` (the base class users subclass)
stays in
`pyspark.sql.aggregator`. `udaf` is unreleased (added on `master` /
`branch-4.x` for SPARK-58736),
so this only changes the import path within unreleased branches:
```python
from pyspark.sql.aggregator import Aggregator # base class to subclass
from pyspark.sql.functions import udaf # was: from
pyspark.sql.aggregator import udaf
```
### Why are the changes needed?
Parity. Grouped-agg pandas/arrow UDFs already work as window functions, and
so does the Scala
`Aggregator` (via the JVM SQL window path). The incremental Python
aggregator is presented as the
analog of the Scala `Aggregator`, so rejecting it over a window was an
inconsistent limitation.
### Does this PR introduce _any_ user-facing change?
Yes. `udaf(...).over(window)` now works; previously it raised
`UNSUPPORTED_EXPR_FOR_WINDOW` at
analysis. Additionally, `udaf` is now imported from `pyspark.sql.functions`
rather than
`pyspark.sql.aggregator`. Both are relative to the unreleased
incremental-aggregator feature
(SPARK-58736); no released behavior changes.
### How was this patch tested?
New tests in
`python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py`, run for both
classic
and Spark Connect (via the parity suite):
- `test_window_unbounded` -- unbounded partition frame.
- `test_window_running_frame` -- ordered growing frame (unbounded preceding
.. current row).
- `test_window_sliding_frame` -- sliding frame (1 preceding .. 1 following)
with a custom
single-field buffer aggregator.
Each cross-checks the aggregator's result against the equivalent SQL window
aggregate. Ran locally:
classic 26/26 and Spark Connect parity 24/24 pass.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 4.8)
Closes #58142 from HyukjinKwon/python-aggregator-window.
Authored-by: Hyukjin Kwon <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
---
.../src/main/resources/error/error-conditions.json | 5 +
.../org/apache/spark/api/python/PythonRunner.scala | 6 +
.../source/reference/pyspark.sql/functions.rst | 1 +
python/pyspark/sql/aggregator.py | 92 +--------------
python/pyspark/sql/connect/functions/builtin.py | 6 +-
python/pyspark/sql/functions/__init__.py | 1 +
python/pyspark/sql/functions/builtin.py | 126 +++++++++++++++++++++
python/pyspark/sql/pandas/_typing/__init__.pyi | 1 +
.../tests/arrow/test_arrow_python_aggregator.py | 93 ++++++++++++++-
python/pyspark/sql/tests/test_functions.py | 1 -
python/pyspark/util.py | 5 +
python/pyspark/worker.py | 98 ++++++++++++++++
.../sql/catalyst/analysis/WindowResolution.scala | 10 --
.../spark/sql/catalyst/expressions/PythonUDF.scala | 9 +-
.../catalyst/expressions/windowExpressions.scala | 3 +
.../spark/sql/errors/QueryCompilationErrors.scala | 6 +
.../python/ArrowWindowPythonEvaluatorFactory.scala | 8 +-
.../execution/python/ArrowWindowPythonExec.scala | 17 ++-
18 files changed, 371 insertions(+), 117 deletions(-)
diff --git a/common/utils/src/main/resources/error/error-conditions.json
b/common/utils/src/main/resources/error/error-conditions.json
index 2b7c146fcf99..c8bdaa4a48f7 100644
--- a/common/utils/src/main/resources/error/error-conditions.json
+++ b/common/utils/src/main/resources/error/error-conditions.json
@@ -8886,6 +8886,11 @@
"Multiple bucket TRANSFORMs."
]
},
+ "MULTIPLE_PYTHON_UDF_TYPES_IN_WINDOW" : {
+ "message" : [
+ "Cannot use Python user-defined functions of different types
together over a single window: <functionList>. Use a separate window
specification for each."
+ ]
+ },
"MULTI_ACTION_ALTER" : {
"message" : [
"The target JDBC server hosting table <tableName> does not support
ALTER TABLE with multiple actions. Split the ALTER TABLE up into individual
actions to avoid this error."
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
index 015067766c9c..e643045134f4 100644
--- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
+++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -95,6 +95,11 @@ private[spark] object PythonEvalType {
val SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF = 255
val SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF = 256
+ // Window aggregation with an incremental Arrow aggregator. A window has no
shuffle, so it needs
+ // neither the PARTIAL nor the FINAL eval type above: the operator sends
each frame's rows to the
+ // worker, which folds them with `reduce` (from `zero`) and produces the
value with `finish`.
+ val SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF = 257
+
val SQL_TABLE_UDF = 300
val SQL_ARROW_TABLE_UDF = 301
val SQL_ARROW_UDTF = 302
@@ -142,6 +147,7 @@ private[spark] object PythonEvalType {
"SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF"
case SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF =>
"SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF"
+ case SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF =>
"SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF"
}
// The eval types produced by ExtractPythonUDFFromLambda: a scalar UDF
lifted out of a
diff --git a/python/docs/source/reference/pyspark.sql/functions.rst
b/python/docs/source/reference/pyspark.sql/functions.rst
index a79922c8051d..229fcaf6c381 100644
--- a/python/docs/source/reference/pyspark.sql/functions.rst
+++ b/python/docs/source/reference/pyspark.sql/functions.rst
@@ -769,6 +769,7 @@ UDF, UDTF and UDT
arrow_udtf
call_udf
pandas_udf
+ udaf
udf
udtf
unwrap_udt
diff --git a/python/pyspark/sql/aggregator.py b/python/pyspark/sql/aggregator.py
index 77110b8f88a3..8ce383b9b4f2 100644
--- a/python/pyspark/sql/aggregator.py
+++ b/python/pyspark/sql/aggregator.py
@@ -22,11 +22,10 @@ Incremental user-defined aggregators for PySpark, the
Python analog of Scala's
from abc import ABC, abstractmethod
from typing import Any, Tuple
-from pyspark.errors import PySparkNotImplementedError, PySparkTypeError,
PySparkValueError
+from pyspark.errors import PySparkNotImplementedError
from pyspark.sql.types import DataType, StructType
-from pyspark.util import PythonEvalType
-__all__ = ["Aggregator", "udaf"]
+__all__ = ["Aggregator"]
class Aggregator(ABC):
@@ -52,7 +51,8 @@ class Aggregator(ABC):
--------
A mean aggregator::
- from pyspark.sql.aggregator import Aggregator, udaf
+ from pyspark.sql.aggregator import Aggregator
+ from pyspark.sql.functions import udaf
from pyspark.sql.types import StructType, StructField, DoubleType,
LongType
class Mean(Aggregator):
@@ -137,87 +137,3 @@ class Aggregator(ABC):
errorClass="NOT_IMPLEMENTED",
messageParameters={"feature": "calling an Aggregator directly;
wrap it with udaf(...)"},
)
-
-
-def udaf(agg: "Aggregator") -> Any:
- """
- Turn an :class:`Aggregator` instance into a callable usable in
``groupBy().agg(...)``, the
- Python counterpart of Scala's ``functions.udaf``.
-
- The aggregator is executed with true incremental (partial) aggregation and
transfers its
- intermediate buffer as Arrow; PyArrow is therefore required.
-
- .. versionadded:: 4.4.0
-
- Parameters
- ----------
- agg : :class:`Aggregator`
- The aggregator instance.
-
- Returns
- -------
- function
- A callable that, applied to input columns, produces an aggregate
:class:`Column`.
-
- Raises
- ------
- :class:`PySparkImportError`
- If a supported version of PyArrow is not installed.
- :class:`PySparkTypeError`
- If ``agg`` is not an :class:`Aggregator`, or its ``bufferSchema`` is
not a
- :class:`StructType`.
- """
- from pyspark.sql.pandas.utils import require_minimum_pyarrow_version
- from pyspark.sql.utils import is_remote
-
- require_minimum_pyarrow_version()
-
- if is_remote():
- from pyspark.sql.connect.udf import UserDefinedFunction
- else:
- # The classic UserDefinedFunction is a distinct class from the Connect
one above;
- # both provide the same interface used below, so silence mypy's
reassignment check.
- from pyspark.sql.udf import UserDefinedFunction # type:
ignore[assignment]
-
- if not isinstance(agg, Aggregator):
- raise PySparkTypeError(
- errorClass="NOT_EXPECTED_TYPE",
- messageParameters={
- "arg_name": "agg",
- "expected_type": "Aggregator",
- "arg_type": type(agg).__name__,
- },
- )
- if not isinstance(agg.bufferSchema, StructType):
- raise PySparkTypeError(
- errorClass="NOT_EXPECTED_TYPE",
- messageParameters={
- "arg_name": "bufferSchema",
- "expected_type": "StructType",
- "arg_type": type(agg.bufferSchema).__name__,
- },
- )
- # The buffer crosses the shuffle as an Arrow struct whose children are
matched by name, and the
- # worker keys the buffer tuple back by field name. Duplicate names would
silently collapse
- # fields on the map side and then fail with an opaque Arrow error
post-shuffle, so reject them
- # up front where the aggregator is created.
- field_names = [field.name for field in agg.bufferSchema.fields]
- if len(field_names) != len(set(field_names)):
- duplicates = sorted({name for name in field_names if
field_names.count(name) > 1})
- raise PySparkValueError(
- errorClass="DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT",
- messageParameters={"field_names": ", ".join(duplicates)},
- )
-
- # ``bufferSchema`` is a first-class ``UserDefinedFunction`` field
(threaded to the JVM in
- # ``_create_judf`` so ``PythonAggregate`` can plan the two-stage
aggregation), so it survives
- # ``_wrapped()`` and ``spark.udf.register`` without being re-attached.
- udf_obj = UserDefinedFunction(
- agg,
- returnType=agg.outputType,
- name=agg.__class__.__name__,
- evalType=PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
- deterministic=True,
- bufferSchema=agg.bufferSchema,
- )
- return udf_obj._wrapped()
diff --git a/python/pyspark/sql/connect/functions/builtin.py
b/python/pyspark/sql/connect/functions/builtin.py
index 8451ae386770..9ad1bf43cb63 100644
--- a/python/pyspark/sql/connect/functions/builtin.py
+++ b/python/pyspark/sql/connect/functions/builtin.py
@@ -68,9 +68,9 @@ from pyspark.sql.types import (
)
from pyspark.sql.utils import enum_to_value as _enum_to_value
-# The implementation of pandas_udf is embedded in
pyspark.sql.function.pandas_udf
-# for code reuse.
-from pyspark.sql.functions import arrow_udf, pandas_udf # noqa: F401
+# The implementations of pandas_udf, arrow_udf and udaf are embedded in
pyspark.sql.functions
+# (they select the classic vs Connect UserDefinedFunction internally), so
reuse them here.
+from pyspark.sql.functions import arrow_udf, pandas_udf, udaf # noqa: F401
if TYPE_CHECKING:
from pyspark.sql.connect._typing import (
diff --git a/python/pyspark/sql/functions/__init__.py
b/python/pyspark/sql/functions/__init__.py
index d20790f20bcb..001a7d4784ca 100644
--- a/python/pyspark/sql/functions/__init__.py
+++ b/python/pyspark/sql/functions/__init__.py
@@ -617,6 +617,7 @@ __all__ = [ # noqa: F405
# Call Functions
"call_udf",
"pandas_udf",
+ "udaf",
"udf",
"udtf",
"arrow_udtf",
diff --git a/python/pyspark/sql/functions/builtin.py
b/python/pyspark/sql/functions/builtin.py
index d5ec94fa6d62..a906b884872b 100644
--- a/python/pyspark/sql/functions/builtin.py
+++ b/python/pyspark/sql/functions/builtin.py
@@ -78,6 +78,7 @@ from pyspark.sql.utils import (
if TYPE_CHECKING:
from pyspark import SparkContext
+ from pyspark.sql.aggregator import Aggregator
from pyspark.sql.dataframe import DataFrame
from pyspark.sql._typing import (
ColumnOrName,
@@ -33059,6 +33060,130 @@ def bitmap_and_agg(col: "ColumnOrName") -> Column:
# ---------------------------- User Defined Function
----------------------------------
+def udaf(agg: "Aggregator") -> "UserDefinedFunctionLike":
+ """Turn an :class:`~pyspark.sql.aggregator.Aggregator` instance into a
callable usable in
+ ``groupBy().agg(...)`` (and as a window function), the Python counterpart
of Scala's
+ ``functions.udaf``.
+
+ The aggregator is executed with true incremental (partial) aggregation and
transfers its
+ intermediate buffer as Arrow; PyArrow is therefore required.
+
+ .. versionadded:: 4.4.0
+
+ Parameters
+ ----------
+ agg : :class:`~pyspark.sql.aggregator.Aggregator`
+ The aggregator instance.
+
+ Returns
+ -------
+ function
+ A callable that, applied to input columns, produces an aggregate
+ :class:`~pyspark.sql.Column`.
+
+ Raises
+ ------
+ :class:`PySparkImportError`
+ If a supported version of PyArrow is not installed.
+ :class:`PySparkTypeError`
+ If ``agg`` is not an :class:`~pyspark.sql.aggregator.Aggregator`, or
its ``bufferSchema`` is
+ not a :class:`StructType`.
+
+ Examples
+ --------
+ >>> from pyspark.sql.aggregator import Aggregator
+ >>> from pyspark.sql.functions import udaf
+ >>> from pyspark.sql.types import StructType, StructField, DoubleType,
LongType
+ >>> class Mean(Aggregator):
+ ... @property
+ ... def bufferSchema(self):
+ ... return StructType(
+ ... [StructField("sum", DoubleType()), StructField("count",
LongType())]
+ ... )
+ ...
+ ... @property
+ ... def outputType(self):
+ ... return DoubleType()
+ ...
+ ... def zero(self):
+ ... return (0.0, 0)
+ ...
+ ... def reduce(self, buffer, value):
+ ... (v,) = value
+ ... return buffer if v is None else (buffer[0] + v, buffer[1] + 1)
+ ...
+ ... def merge(self, b1, b2):
+ ... return (b1[0] + b2[0], b1[1] + b2[1])
+ ...
+ ... def finish(self, buffer):
+ ... return buffer[0] / buffer[1] if buffer[1] else None
+ >>> df = spark.createDataFrame([(1, 1.0), (1, 2.0), (2, 3.0)], ("k", "v"))
+ >>> df.groupBy("k").agg(udaf(Mean())(df.v).alias("m")).orderBy("k").show()
+ +---+---+
+ | k| m|
+ +---+---+
+ | 1|1.5|
+ | 2|3.0|
+ +---+---+
+ """
+ from pyspark.sql.aggregator import Aggregator
+ from pyspark.sql.pandas.utils import require_minimum_pyarrow_version
+ from pyspark.sql.utils import is_remote
+ from pyspark.util import PythonEvalType
+
+ require_minimum_pyarrow_version()
+
+ if is_remote():
+ from pyspark.sql.connect.udf import UserDefinedFunction
+ else:
+ # The classic UserDefinedFunction is a distinct class from the Connect
one above;
+ # both provide the same interface used below, so silence mypy's
reassignment check.
+ from pyspark.sql.udf import UserDefinedFunction # type:
ignore[assignment]
+
+ if not isinstance(agg, Aggregator):
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "arg_name": "agg",
+ "expected_type": "Aggregator",
+ "arg_type": type(agg).__name__,
+ },
+ )
+ if not isinstance(agg.bufferSchema, StructType):
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "arg_name": "bufferSchema",
+ "expected_type": "StructType",
+ "arg_type": type(agg.bufferSchema).__name__,
+ },
+ )
+ # The buffer crosses the shuffle as an Arrow struct whose children are
matched by name, and the
+ # worker keys the buffer tuple back by field name. Duplicate names would
silently collapse
+ # fields on the map side and then fail with an opaque Arrow error
post-shuffle, so reject them
+ # up front where the aggregator is created.
+ field_names = [field.name for field in agg.bufferSchema.fields]
+ if len(field_names) != len(set(field_names)):
+ duplicates = sorted({name for name in field_names if
field_names.count(name) > 1})
+ raise PySparkValueError(
+ errorClass="DUPLICATED_FIELD_NAME_IN_ARROW_STRUCT",
+ messageParameters={"field_names": ", ".join(duplicates)},
+ )
+
+ # ``bufferSchema`` is a first-class ``UserDefinedFunction`` field
(threaded to the JVM in
+ # ``_create_judf`` so ``PythonAggregate`` can plan the two-stage
aggregation), so it survives
+ # ``_wrapped()`` and ``spark.udf.register`` without being re-attached.
+ udf_obj = UserDefinedFunction(
+ agg,
+ returnType=agg.outputType,
+ name=agg.__class__.__name__,
+ evalType=PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
+ deterministic=True,
+ bufferSchema=agg.bufferSchema,
+ )
+ return udf_obj._wrapped()
+
+
@overload
def udf(
f: Callable[..., Any],
@@ -33743,6 +33868,7 @@ def _test() -> None:
if not have_pandas or not have_pyarrow:
del pyspark.sql.functions.builtin.udf.__doc__
del pyspark.sql.functions.builtin.arrow_udtf.__doc__
+ del pyspark.sql.functions.builtin.udaf.__doc__
spark = (
SparkSession.builder.master("local[4]").appName("sql.functions.builtin
tests").getOrCreate()
diff --git a/python/pyspark/sql/pandas/_typing/__init__.pyi
b/python/pyspark/sql/pandas/_typing/__init__.pyi
index 442849bcbf0a..8e8d4f34783c 100644
--- a/python/pyspark/sql/pandas/_typing/__init__.pyi
+++ b/python/pyspark/sql/pandas/_typing/__init__.pyi
@@ -70,6 +70,7 @@ ArrowWindowAggUDFType = Literal[253]
ArrowGroupedAggIterUDFType = Literal[254]
ArrowGroupedAggIncrementalPartialUDFType = Literal[255]
ArrowGroupedAggIncrementalFinalUDFType = Literal[256]
+ArrowWindowAggIncrementalUDFType = Literal[257]
# Arrow stream types
# A single group of Arrow batches (e.g., one key group in groupBy).
diff --git a/python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py
b/python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py
index 3965d804ae84..692c2256edcc 100644
--- a/python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py
+++ b/python/pyspark/sql/tests/arrow/test_arrow_python_aggregator.py
@@ -36,7 +36,8 @@ from pyspark.testing.utils import (
if have_pyarrow:
- from pyspark.sql.aggregator import Aggregator, udaf
+ from pyspark.sql.aggregator import Aggregator
+ from pyspark.sql.functions import udaf
class Mean(Aggregator):
@property
@@ -280,12 +281,92 @@ class ArrowPythonAggregatorTestsMixin:
with self.assertRaises(AnalysisException):
df.groupBy("k").pivot("p").agg(udaf(Mean())(sf.col("v"))).collect()
- def test_window_rejected(self):
- # The incremental aggregator has no window operator; using it over a
window must fail
- # analysis rather than crash with an internal error at execution.
+ def test_window_unbounded(self):
+ # Unbounded partition frame: every row gets its whole group's
aggregate. Cross-checked
+ # against the equivalent SQL window aggregate.
df = self._data()
- with self.assertRaises(AnalysisException):
-
df.select(udaf(Mean())(sf.col("v")).over(Window.partitionBy("k"))).collect()
+ w = Window.partitionBy("k")
+ result = df.withColumn("m",
udaf(Mean())(sf.col("v")).over(w)).orderBy("k", "v").collect()
+ expected = df.withColumn("m", sf.avg("v").over(w)).orderBy("k",
"v").collect()
+ self.assertEqual(len(result), len(expected))
+ for r, e in zip(result, expected):
+ self.assertAlmostEqual(r["m"], e["m"], places=6)
+
+ def test_window_running_frame(self):
+ # Ordered, growing frame (unbounded preceding .. current row): a
running aggregate that
+ # exercises the per-row bounded-frame path in the worker.
+ df = self._data()
+ w = (
+ Window.partitionBy("k")
+ .orderBy("v")
+ .rowsBetween(Window.unboundedPreceding, Window.currentRow)
+ )
+ result = df.withColumn("m",
udaf(Mean())(sf.col("v")).over(w)).orderBy("k", "v").collect()
+ expected = df.withColumn("m", sf.avg("v").over(w)).orderBy("k",
"v").collect()
+ self.assertEqual(len(result), len(expected))
+ for r, e in zip(result, expected):
+ self.assertAlmostEqual(r["m"], e["m"], places=6)
+
+ def test_window_sliding_frame(self):
+ # Sliding frame (1 preceding .. 1 following) with a custom
single-field buffer aggregator.
+ df = self._data()
+ w = Window.partitionBy("k").orderBy("v").rowsBetween(-1, 1)
+ result = (
+ df.withColumn("s",
udaf(SumSquares())(sf.col("v")).over(w)).orderBy("k", "v").collect()
+ )
+ expected = (
+ df.withColumn("s", sf.sum(sf.col("v") * sf.col("v")).over(w))
+ .orderBy("k", "v")
+ .collect()
+ )
+ self.assertEqual(len(result), len(expected))
+ for r, e in zip(result, expected):
+ self.assertAlmostEqual(r["s"], e["s"], places=4)
+
+ def test_window_bounded_preceding_frame(self):
+ # A fixed number of preceding rows exercises both branches of the
running-buffer
+ # optimization: the lower bound is clamped to 0 for the first rows
(running buffer is
+ # extended in place) and then advances (each frame is refolded from
zero).
+ df = self._data()
+ w = Window.partitionBy("k").orderBy("v").rowsBetween(-3,
Window.currentRow)
+ result = df.withColumn("m",
udaf(Mean())(sf.col("v")).over(w)).orderBy("k", "v").collect()
+ expected = df.withColumn("m", sf.avg("v").over(w)).orderBy("k",
"v").collect()
+ self.assertEqual(len(result), len(expected))
+ for r, e in zip(result, expected):
+ self.assertAlmostEqual(r["m"], e["m"], places=6)
+
+ def test_window_decimal_output(self):
+ # A non-trivial (Decimal) output type over a window, exercising the
explicit
+ # ``pa.array(..., type=result_type)`` typing on the window path.
+ df = self._data()
+ w = Window.partitionBy("k")
+ result = (
+ df.withColumn("s",
udaf(DecimalSum())(sf.col("v")).over(w)).orderBy("k", "v").collect()
+ )
+ expected = df.withColumn("s", sf.sum("v").over(w)).orderBy("k",
"v").collect()
+ self.assertEqual(len(result), len(expected))
+ for r, e in zip(result, expected):
+ self.assertIsInstance(r["s"], Decimal)
+ self.assertAlmostEqual(float(r["s"]), e["s"], places=4)
+
+ def test_window_mixed_python_udf_rejected(self):
+ # An incremental aggregator and a grouped-agg pandas UDF over the same
window are both
+ # Python window functions but use different eval types, so they cannot
share one operator.
+ # This must raise a clear analysis error rather than an internal
assertion.
+ from pyspark.sql.functions import pandas_udf, PandasUDFType
+
+ @pandas_udf("double", PandasUDFType.GROUPED_AGG)
+ def pandas_mean(v):
+ return v.mean()
+
+ df = self._data()
+ w = Window.partitionBy("k")
+ with self.assertRaises(AnalysisException) as ctx:
+ df.select(udaf(Mean())(sf.col("v")).over(w),
pandas_mean(sf.col("v")).over(w)).collect()
+ self.assertEqual(
+ ctx.exception.getCondition(),
+ "UNSUPPORTED_FEATURE.MULTIPLE_PYTHON_UDF_TYPES_IN_WINDOW",
+ )
def test_mixed_with_other_aggregate_rejected(self):
# An incremental aggregator mixed with another aggregate in one
Aggregate is unsupported;
diff --git a/python/pyspark/sql/tests/test_functions.py
b/python/pyspark/sql/tests/test_functions.py
index 44c1047cbe3c..da8c13537ab8 100644
--- a/python/pyspark/sql/tests/test_functions.py
+++ b/python/pyspark/sql/tests/test_functions.py
@@ -65,7 +65,6 @@ class FunctionsTestsMixin:
"not", # equivalent to python ~expression
"any", # equivalent to python ~some
"len", # equivalent to python ~length
- "udaf", # used for creating UDAF's which are not supported in
PySpark
"partitioning$", # partitioning expressions for DSv2
]
diff --git a/python/pyspark/util.py b/python/pyspark/util.py
index 81bea6152f98..8b37a2577f93 100644
--- a/python/pyspark/util.py
+++ b/python/pyspark/util.py
@@ -88,6 +88,7 @@ if typing.TYPE_CHECKING:
ArrowGroupedAggIterUDFType,
ArrowGroupedAggIncrementalPartialUDFType,
ArrowGroupedAggIncrementalFinalUDFType,
+ ArrowWindowAggIncrementalUDFType,
ArrowWindowAggUDFType,
)
from pyspark.sql._typing import (
@@ -708,6 +709,10 @@ class PythonEvalType:
SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF:
"ArrowGroupedAggIncrementalPartialUDFType" = 255
SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF:
"ArrowGroupedAggIncrementalFinalUDFType" = 256
+ # Window aggregation with an incremental ``Aggregator``. A window has no
shuffle, so each
+ # frame's rows are folded with ``reduce`` (from ``zero``) and finished
with ``finish``.
+ SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF: "ArrowWindowAggIncrementalUDFType" =
257
+
SQL_TABLE_UDF: "SQLTableUDFType" = 300
SQL_ARROW_TABLE_UDF: "SQLArrowTableUDFType" = 301
SQL_ARROW_UDTF: "SQLArrowUDTFType" = 302
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 565323365c33..d92824f5a053 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -679,6 +679,7 @@ def read_single_udf(pickleSer, udf_info, eval_type,
runner_conf, udf_index):
if eval_type in (
PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
+ PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF,
):
return chained_func, args_offsets, kwargs_offsets, return_type
@@ -2066,6 +2067,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
PythonEvalType.SQL_TRANSFORM_WITH_STATE_PYTHON_ROW_INIT_STATE_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_PARTIAL_UDF,
PythonEvalType.SQL_GROUPED_AGG_ARROW_INCREMENTAL_FINAL_UDF,
+ PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF,
):
# NOTE: if timezone is set here, that implies respectSessionTimeZone
is True
if eval_type in (
@@ -2083,6 +2085,7 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
PythonEvalType.SQL_GROUPED_MAP_PANDAS_UDF,
PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF,
PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
+ PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF,
):
ser = ArrowStreamGroupSerializer(write_start_stream=True)
elif eval_type in (
@@ -2647,6 +2650,101 @@ def read_udfs(pickleSer, udf_info_list, eval_type,
runner_conf, eval_conf):
# profiling is not supported for UDF
return grouped_func, None, ser, ser
+ if eval_type == PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF:
+ import pyarrow as pa
+
+ # Window aggregation with an incremental ``Aggregator``. The operator
sends each frame --
+ # the whole partition for an unbounded frame, or per-row ``[begin,
end)`` slices for a
+ # bounded one -- and the worker folds the frame's rows with ``reduce``
(from a fresh
+ # ``zero``) and produces the value with ``finish``, one output value
per input row. A window
+ # has no shuffle, so the intermediate buffer never leaves the worker
(unlike the two-stage
+ # groupBy path); ``merge`` is not used here.
+ window_bound_types_str = runner_conf.get("window_bound_types")
+ window_bound_types = [t.strip().lower() for t in
window_bound_types_str.split(",")]
+
+ col_names = ["_%d" % i for i in range(len(udfs))]
+ return_schema = to_arrow_schema(
+ StructType([StructField(name, rt) for name, (_, _, _, rt) in
zip(col_names, udfs)]),
+ timezone="UTC",
+ prefers_large_types=runner_conf.use_large_var_types,
+ )
+
+ def fold(agg: Any, buffer: Any, value_cols: list, start: int, end:
int) -> Any:
+ # Fold rows ``[start, end)`` (each a tuple across ``value_cols``,
matching the call-site
+ # argument order) into ``buffer`` via the aggregator's ``reduce``.
+ for r in range(start, end):
+ buffer = agg.reduce(buffer, tuple(c[r] for c in value_cols))
+ return buffer
+
+ def grouped_func(
+ split_index: int, data: Iterator["GroupedBatch"]
+ ) -> Iterator[pa.RecordBatch]:
+ for group in data:
+ batch_list = list(group)
+ if not batch_list:
+ continue
+ if hasattr(pa, "concat_batches"):
+ concatenated = pa.concat_batches(batch_list)
+ else:
+ # pyarrow.concat_batches not supported before 19.0.0
+ # remove this once we drop support for old versions
+ concatenated = pa.RecordBatch.from_struct_array(
+ pa.concat_arrays([b.to_struct_array() for b in
batch_list])
+ )
+ num_rows = concatenated.num_rows
+
+ result_arrays = []
+ for udf_index, (agg, args_offsets, kwargs_offsets, _) in
enumerate(udfs):
+ bound_type = window_bound_types[udf_index]
+ result_type = return_schema.field(udf_index).type
+ if bound_type == "unbounded":
+ # One frame spanning the whole partition: compute
once, repeat per row.
+ value_cols = [concatenated.column(o).to_pylist() for o
in args_offsets] + [
+ concatenated.column(v).to_pylist() for v in
kwargs_offsets.values()
+ ]
+ result = agg.finish(fold(agg, agg.zero(), value_cols,
0, num_rows))
+ result_arrays.append(pa.array([result] * num_rows,
type=result_type))
+ elif bound_type == "bounded":
+ # Per-row frame ``[begin, end)``. Materialize the
aggregator's input columns
+ # once; frames index into them by row.
+ begin_col = concatenated.column(args_offsets[0])
+ end_col = concatenated.column(args_offsets[1])
+ data_offsets = list(args_offsets[2:]) +
list(kwargs_offsets.values())
+ value_cols = [concatenated.column(o).to_pylist() for o
in data_offsets]
+ # When consecutive frames share the same lower bound
and only grow on the
+ # right (e.g. rowsBetween(unboundedPreceding,
currentRow)), extend the
+ # running buffer by the newly-included rows instead of
refolding from
+ # ``zero`` -- O(n) overall rather than O(n^2).
Otherwise -- the lower bound
+ # advanced (a row left the window, which ``reduce``
cannot subtract) or the
+ # frame shrank -- refold the frame from ``zero``.
+ results = []
+ have_running = False
+ running: Any = None
+ prev_begin = -1
+ prev_end = 0
+ for i in range(num_rows):
+ begin = begin_col[i].as_py()
+ end = end_col[i].as_py()
+ if have_running and begin == prev_begin and end >=
prev_end:
+ running = fold(agg, running, value_cols,
prev_end, end)
+ else:
+ running = fold(agg, agg.zero(), value_cols,
begin, end)
+ have_running = True
+ prev_begin, prev_end = begin, end
+ results.append(agg.finish(running))
+ result_arrays.append(pa.array(results,
type=result_type))
+ else:
+ raise PySparkRuntimeError(
+ errorClass="INVALID_WINDOW_BOUND_TYPE",
+ messageParameters={"window_bound_type":
bound_type},
+ )
+
+ batch = pa.RecordBatch.from_arrays(result_arrays, col_names)
+ yield ArrowBatchTransformer.enforce_schema(batch,
return_schema)
+
+ # profiling is not supported for UDF
+ return grouped_func, None, ser, ser
+
if eval_type == PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF:
import pyarrow as pa
import pandas as pd
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/WindowResolution.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/WindowResolution.scala
index 6c81d2ab7f52..a477fe4398f6 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/WindowResolution.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/WindowResolution.scala
@@ -21,7 +21,6 @@ import org.apache.spark.sql.catalyst.expressions.{
AggregateWindowFunction,
CurrentRow,
FrameLessOffsetWindowFunction,
- PythonAggregate,
RangeFrame,
RankLike,
RowFrame,
@@ -164,15 +163,6 @@ object WindowResolution {
errorClass = "INVALID_WINDOW_SPEC_FOR_AGGREGATION_FUNC",
messageParameters = Map("aggFunc" ->
toSQLExpr(agg.aggregateFunction))
)
- case AggregateExpression(_: PythonAggregate, _, _, _, _) =>
- // The incremental Python aggregator has a dedicated two-stage
physical operator and is not
- // supported as a window function. Reject it here with a clear
analysis error, rather than
- // letting it be misclassified as a SQL aggregate
(`WindowFunctionType.functionType` only
- // routes `PythonUDAF` to the Python path) and fail with an internal
error at execution.
- windowExpression.failAnalysis(
- errorClass = "UNSUPPORTED_EXPR_FOR_WINDOW",
- messageParameters = Map("sqlExpr" ->
toSQLExpr(windowExpression.windowFunction))
- )
case _: AggregateExpression | _: FrameLessOffsetWindowFunction | _:
AggregateWindowFunction =>
case other =>
other.failAnalysis(
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
index d62131183c49..b820172bc2b5 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/PythonUDF.scala
@@ -240,10 +240,11 @@ object PythonUDF {
}
def isWindowPandasUDF(e: PythonFuncExpression): Boolean = {
- // This is currently only `PythonUDAF` (which means
SQL_GROUPED_AGG_PANDAS_UDF or
- // SQL_GROUPED_AGG_ARROW_UDF), but we might
- // support new types in the future, e.g, N -> N transform.
- e.isInstanceOf[PythonUDAF]
+ // `PythonUDAF` (SQL_GROUPED_AGG_PANDAS_UDF or SQL_GROUPED_AGG_ARROW_UDF)
and the incremental
+ // `PythonAggregate` are the Python aggregate functions that run over a
window through the
+ // Python window operator, rather than the JVM SQL window path. We might
support new types in
+ // the future, e.g. N -> N transform.
+ e.isInstanceOf[PythonUDAF] || e.isInstanceOf[PythonAggregate]
}
def correctEvalType(udf: PythonUDF, pythonUDFArrowFallbackOnUDT: Boolean):
Int = {
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/windowExpressions.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/windowExpressions.scala
index f26a4dc9a0c0..e6c52c08e3b1 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/windowExpressions.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/windowExpressions.scala
@@ -410,6 +410,9 @@ object WindowFunctionType {
case PythonEvalType.SQL_GROUPED_AGG_PANDAS_UDF =>
PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF
case PythonEvalType.SQL_GROUPED_AGG_ARROW_UDF =>
PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF
}
+ // The incremental aggregator has a single window eval type: the window
operator sends each
+ // frame's rows to the worker, which folds them with `reduce` and
produces `finish`.
+ case _: PythonAggregate =>
PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF
}
}
}
diff --git
a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
index 5db80e86d50c..ada71af43d3b 100644
---
a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
+++
b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryCompilationErrors.scala
@@ -3029,6 +3029,12 @@ private[sql] object QueryCompilationErrors extends
QueryErrorsBase with Compilat
messageParameters = Map("windowExpressions" ->
windowExpressions.toString()))
}
+ def multiplePythonUDFTypesInWindowError(functionNames: Seq[String]):
Throwable = {
+ new AnalysisException(
+ errorClass = "UNSUPPORTED_FEATURE.MULTIPLE_PYTHON_UDF_TYPES_IN_WINDOW",
+ messageParameters = Map("functionList" ->
functionNames.map(toSQLId).mkString(", ")))
+ }
+
def escapeCharacterInTheMiddleError(pattern: String, char: String):
Throwable = {
new AnalysisException(
errorClass = "INVALID_FORMAT.ESC_IN_THE_MIDDLE",
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala
index ab9671c022f9..7a6c90478a80 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonEvaluatorFactory.scala
@@ -26,7 +26,7 @@ import org.apache.spark.{JobArtifactSet, PartitionEvaluator,
PartitionEvaluatorF
import org.apache.spark.api.python.ChainedPythonFunctions
import org.apache.spark.internal.config.Python.PYTHON_UDF_PIPELINED_EXECUTION
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Attribute,
AttributeReference, BoundReference, EmptyRow, Expression, JoinedRow,
NamedArgumentExpression, NamedExpression, PythonFuncExpression, PythonUDAF,
SortOrder, SpecificInternalRow, UnsafeProjection, UnsafeRow, WindowExpression}
+import org.apache.spark.sql.catalyst.expressions.{Attribute,
AttributeReference, BoundReference, EmptyRow, Expression, JoinedRow,
NamedArgumentExpression, NamedExpression, PythonFuncExpression, SortOrder,
SpecificInternalRow, UnsafeProjection, UnsafeRow, WindowExpression}
import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.execution.ExternalAppendOnlyUnsafeRowArray
@@ -156,8 +156,12 @@ class ArrowWindowPythonEvaluatorFactory(
// Extract window expressions and window functions
private val windowExpressions = expressions.flatMap(_.collect { case e:
WindowExpression => e })
+ // The window aggregate function is either a `PythonUDAF` (grouped-agg
pandas/arrow UDF) or the
+ // incremental `PythonAggregate`; both are `PythonFuncExpression`s and
share the per-frame Arrow
+ // window path (only the worker-side per-frame computation and eval type
differ).
private val udfExpressions = windowExpressions.map { e =>
-
e.windowFunction.asInstanceOf[AggregateExpression].aggregateFunction.asInstanceOf[PythonUDAF]
+ e.windowFunction.asInstanceOf[AggregateExpression].aggregateFunction
+ .asInstanceOf[PythonFuncExpression]
}
// We shouldn't be chaining anything here.
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonExec.scala
index dd360bbf230b..92956e4ecb5d 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/python/ArrowWindowPythonExec.scala
@@ -22,6 +22,7 @@ import org.apache.spark.api.python.PythonEvalType
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.errors.QueryCompilationErrors
import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
import org.apache.spark.sql.execution.window._
@@ -33,6 +34,7 @@ import org.apache.spark.sql.execution.window._
* <ul>
* <li> SQL_WINDOW_AGG_ARROW_UDF for Arrow UDF
* <li> SQL_WINDOW_AGG_PANDAS_UDF for Pandas UDF
+ * <li> SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF for the incremental Arrow
aggregator (`udaf`)
* </ul>
*
* This is similar to [[WindowExec]]. The main difference is that this node
does not compute
@@ -127,7 +129,8 @@ case class ArrowWindowPythonExec(
private def supportedPythonEvalTypes: Array[Int] =
Array(
PythonEvalType.SQL_WINDOW_AGG_ARROW_UDF,
- PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF)
+ PythonEvalType.SQL_WINDOW_AGG_PANDAS_UDF,
+ PythonEvalType.SQL_WINDOW_AGG_ARROW_INCREMENTAL_UDF)
}
object ArrowWindowPythonExec {
@@ -139,8 +142,16 @@ object ArrowWindowPythonExec {
val evalTypes = windowExpression.flatMap(w =>
WindowFunctionType.pythonEvalType(w))
assert(evalTypes.nonEmpty,
"Cannot extract eval type from PythonUDAFs in ArrowWindowPythonExec")
- assert(evalTypes.distinct.size == 1,
- "All window functions must have the same eval type in
ArrowWindowPythonExec")
+ // Distinct eval types here means incompatible Python window UDFs were
grouped into one window
+ // (e.g. an incremental aggregator mixed with a grouped-agg pandas/arrow
UDAF). They share the
+ // `WindowFunctionType.Python` type, so they pass the check in
`PhysicalWindow`, but cannot run
+ // in a single operator; surface a clear analysis error rather than an
internal AssertionError.
+ if (evalTypes.distinct.size != 1) {
+ val functionNames = windowExpression.flatMap(_.collect {
+ case e: PythonFuncExpression => e.name
+ }).distinct
+ throw
QueryCompilationErrors.multiplePythonUDFTypesInWindowError(functionNames)
+ }
ArrowWindowPythonExec(windowExpression, partitionSpec, orderSpec, child,
evalTypes.head)
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
