This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new 3f4e1a3f43f [FLINK-40436][python] Add UDAF support to the DataFrame
API (#29431)
3f4e1a3f43f is described below
commit 3f4e1a3f43f762ddf9bb7ad7ff2abef850798e74
Author: Victor Wang <[email protected]>
AuthorDate: Fri Oct 16 13:16:31 2026 +0800
[FLINK-40436][python] Add UDAF support to the DataFrame API (#29431)
---
.../docs/reference/pyflink.dataframe/udf.rst | 11 +
flink-python/pyflink/dataframe/__init__.py | 2 +
flink-python/pyflink/dataframe/sql.py | 15 +-
flink-python/pyflink/dataframe/tests/test_udaf.py | 473 +++++++++++++++++++++
flink-python/pyflink/dataframe/udaf.py | 318 ++++++++++++++
flink-python/pyflink/dataframe/udf.py | 14 +-
.../pyflink/fn_execution/tests/test_coders.py | 26 ++
.../pyflink/table/tests/test_pandas_udaf.py | 19 +
flink-python/pyflink/table/utils.py | 7 +-
9 files changed, 870 insertions(+), 15 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/udf.rst
b/flink-python/docs/reference/pyflink.dataframe/udf.rst
index 798297eea33..7618b887cb7 100644
--- a/flink-python/docs/reference/pyflink.dataframe/udf.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/udf.rst
@@ -45,6 +45,16 @@ DataFrame table UDFs support synchronous functions, callable
classes and instanc
declaration forms, type inference, lifecycle, and examples, and
:meth:`~pyflink.dataframe.DataFrame.flat_map` for row input and output
semantics.
+Aggregate Functions
+===================
+
+Use :func:`pyflink.dataframe.udaf` to declare an aggregate for
+:meth:`~pyflink.dataframe.DataFrame.agg`,
+:meth:`~pyflink.dataframe.GroupedDataFrame.agg`, or
:func:`pyflink.dataframe.sql`.
+General UDAFs use the Table API's ``AggregateFunction`` accumulator protocol.
+Pandas UDAFs also accept ordinary functions and callable classes or instances.
+See :func:`pyflink.dataframe.udaf` for type declarations, execution modes, and
examples.
+
API Reference
=============
@@ -55,3 +65,4 @@ API Reference
udf
udtf
+ udaf
diff --git a/flink-python/pyflink/dataframe/__init__.py
b/flink-python/pyflink/dataframe/__init__.py
index 5c0b3af3941..561f2ea318f 100644
--- a/flink-python/pyflink/dataframe/__init__.py
+++ b/flink-python/pyflink/dataframe/__init__.py
@@ -86,6 +86,7 @@ from pyflink.dataframe.over_window import following, preceding
from pyflink.dataframe.sql import sql
from pyflink.dataframe.udf import udf
from pyflink.dataframe.udtf import udtf
+from pyflink.dataframe.udaf import udaf
__all__ = [
"CloseableIterator",
@@ -103,6 +104,7 @@ __all__ = [
"lit",
"udf",
"udtf",
+ "udaf",
"from_arrow",
"from_dict",
"from_pandas",
diff --git a/flink-python/pyflink/dataframe/sql.py
b/flink-python/pyflink/dataframe/sql.py
index 9a9e11e2134..46d66cc60de 100644
--- a/flink-python/pyflink/dataframe/sql.py
+++ b/flink-python/pyflink/dataframe/sql.py
@@ -27,6 +27,7 @@ from pyflink.dataframe.context import
get_or_create_table_environment
from pyflink.dataframe.dataframe import DataFrame
from pyflink.dataframe.udf import _DataFrameUDFWrapper
from pyflink.dataframe.udtf import _DataFrameUDTFWrapper
+from pyflink.dataframe.udaf import _DataFrameUDAFWrapper
from pyflink.java_gateway import get_gateway
from pyflink.table import Table, TableEnvironment
from pyflink.table.expression import Expression
@@ -41,9 +42,9 @@ _LOG = logging.getLogger(__name__)
# Scalar UDFs use the declared return type of :func:`pyflink.dataframe.udf`
# so that its result type-checks as a binding; at runtime a binding must be an
actual
# UDF object, which is what _BINDABLE_TYPES enforces.
-_FunctionBinding = Union[_DataFrameUDFWrapper, _DataFrameUDTFWrapper]
+_FunctionBinding = Union[_DataFrameUDFWrapper, _DataFrameUDTFWrapper,
_DataFrameUDAFWrapper]
_Binding = Union[DataFrame, Callable[..., Expression], _DataFrameUDTFWrapper]
-_FUNCTION_TYPES = (_DataFrameUDFWrapper, _DataFrameUDTFWrapper)
+_FUNCTION_TYPES = (_DataFrameUDFWrapper, _DataFrameUDTFWrapper,
_DataFrameUDAFWrapper)
_BINDABLE_TYPES = (DataFrame, *_FUNCTION_TYPES)
@@ -55,7 +56,8 @@ def sql(query: str, *, auto_bind: bool = True, **bindings:
_Binding) -> DataFram
The query must be a single statement that returns a result, such as SELECT
or
VALUES (no INSERT / DDL; use :meth:`TableEnvironment.execute_sql` for
those).
The referenced DataFrames are registered as temporary views and the
referenced
- UDFs (created with :func:`~pyflink.dataframe.udf` or
:func:`~pyflink.dataframe.udtf`)
+ UDFs (created with :func:`~pyflink.dataframe.udf`,
:func:`~pyflink.dataframe.udtf`,
+ or :func:`~pyflink.dataframe.udaf`)
as temporary system functions for the duration of the call, and both are
dropped
afterwards. The result can be further transformed with the DataFrame API.
@@ -95,8 +97,8 @@ def sql(query: str, *, auto_bind: bool = True, **bindings:
_Binding) -> DataFram
auto-bound candidates belong to different
TableEnvironments
when there are no explicit bindings.
:raises TypeError: If an explicit binding is neither a :class:`DataFrame`
nor a
- UDF created with :func:`~pyflink.dataframe.udf` or
- :func:`~pyflink.dataframe.udtf`.
+ UDF created with :func:`~pyflink.dataframe.udf`,
+ :func:`~pyflink.dataframe.udtf`, or
:func:`~pyflink.dataframe.udaf`.
Example::
@@ -148,7 +150,8 @@ def sql(query: str, *, auto_bind: bool = True, **bindings:
_Binding) -> DataFram
if not isinstance(value, _BINDABLE_TYPES):
raise TypeError(
f"sql() binding '{name}' must be a DataFrame or a UDF created
with "
- f"pyflink.dataframe.udf or pyflink.dataframe.udtf, got
{type(value).__name__}"
+ "pyflink.dataframe.udf, pyflink.dataframe.udtf or
pyflink.dataframe.udaf, "
+ f"got {type(value).__name__}"
)
explicit_frames = _get_dataframes(bindings)
explicit_udfs = _get_udfs(bindings)
diff --git a/flink-python/pyflink/dataframe/tests/test_udaf.py
b/flink-python/pyflink/dataframe/tests/test_udaf.py
new file mode 100644
index 00000000000..86f51b65571
--- /dev/null
+++ b/flink-python/pyflink/dataframe/tests/test_udaf.py
@@ -0,0 +1,473 @@
+################################################################################
+# 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.
+################################################################################
+
+import datetime
+import functools
+import os
+import unittest
+from typing import TypedDict
+
+import cloudpickle
+import pandas as pd
+import pyflink.dataframe as pf
+from pyflink.common import Row, Types
+from pyflink.table import DataTypes, ListView, MapView, Schema
+from pyflink.table.expressions import col, lit
+from pyflink.table.udf import AggregateFunction, ScalarFunction, udaf as
table_udaf
+from pyflink.table.window import Session
+from pyflink.testing.test_case_utils import (
+ PyFlinkBatchTableTestCase, PyFlinkDataFrameUTTestCase,
PyFlinkStreamDataFrameTestCase,
+)
+from pyflink.util.exceptions import TableException
+
+
+class Sum(AggregateFunction):
+ def create_accumulator(self):
+ return [0]
+
+ def accumulate(self, accumulator, value):
+ if value is not None:
+ accumulator[0] += value
+
+ def retract(self, accumulator, value):
+ if value is not None:
+ accumulator[0] -= value
+
+ def merge(self, accumulator, accumulators):
+ for other in accumulators:
+ accumulator[0] += other[0]
+
+ def get_value(self, accumulator) -> int:
+ return accumulator[0]
+
+ def get_accumulator_type(self):
+ return DataTypes.ARRAY(DataTypes.BIGINT())
+
+
+class Statistics(TypedDict):
+ total: int
+ count: int
+
+
+def _aggregate(declaration, *columns):
+ adapter =
cloudpickle.loads(cloudpickle.dumps(declaration._table_udf_wrapper._func))
+ adapter.open(None)
+ try:
+ accumulator = adapter.create_accumulator()
+ adapter.accumulate(accumulator, *columns)
+ return adapter.get_value(accumulator)
+ finally:
+ adapter.close()
+
+
+class DataFrameUDAFDeclarationTests(unittest.TestCase):
+ def test_declaration_forms(self):
+ for declaration in (pf.udaf(Sum), pf.udaf()(Sum),
pf.udaf(return_dtype=int)(Sum),
+ pf.udaf(Sum())):
+ self.assertEqual(declaration.return_dtype,
pf.DataType.int64().not_null())
+ self.assertEqual(declaration.__name__, "Sum")
+
+ def test_general_requires_aggregate_function(self):
+ class CallableSum:
+ def __call__(self, values) -> int:
+ return sum(values)
+
+ for func in (sum, CallableSum, CallableSum()):
+ with self.subTest(func=func), self.assertRaisesRegex(TypeError,
"AggregateFunction"):
+ pf.udaf(func, return_dtype=int, accumulator_type=list[int])
+
+ def test_type_precedence(self):
+ class TypedSum(Sum):
+ def get_result_type(self):
+ return DataTypes.INT()
+
+ declaration = pf.udaf(TypedSum)
+ self.assertEqual(declaration.return_dtype, pf.DataType.int32())
+ self.assertEqual(declaration.accumulator_type,
pf.DataType.list(pf.DataType.int64()))
+ declaration = pf.udaf(TypedSum, return_dtype=float,
accumulator_type=list[float])
+ self.assertEqual(declaration.return_dtype,
pf.DataType.float64().not_null())
+ self.assertEqual(declaration.accumulator_type,
+
pf.DataType.list(pf.DataType.float64().not_null()).not_null())
+
+ class ExplicitSum(Sum):
+ def get_result_type(self):
+ raise AssertionError("Explicit return_dtype must take
precedence")
+
+ def get_accumulator_type(self):
+ raise AssertionError("Explicit accumulator_type must take
precedence")
+
+ self.assertEqual(pf.udaf(ExplicitSum, return_dtype=int,
+ accumulator_type=list[int]).return_dtype,
+ pf.DataType.int64().not_null())
+
+ def test_return_annotation_is_resolved_independently(self):
+ def total(values):
+ return values.sum()
+
+ total.__annotations__ = {"values": "UnknownInput", "return": int}
+ self.assertEqual(pf.udaf(total, func_type="pandas").return_dtype,
+ pf.DataType.int64().not_null())
+ total.__annotations__["return"] = "UnknownOutput"
+ self.assertEqual(pf.udaf(total, return_dtype=int,
func_type="pandas").return_dtype,
+ pf.DataType.int64().not_null())
+
+ def test_missing_and_invalid_types(self):
+ class UntypedSum(Sum):
+ get_accumulator_type = AggregateFunction.get_accumulator_type
+
+ def create_accumulator(self) -> list[int]:
+ raise AssertionError("Type resolution must not execute
create_accumulator")
+
+ with self.assertRaisesRegex(TypeError, "accumulator_type"):
+ pf.udaf(UntypedSum)
+ with self.assertRaisesRegex(TypeError, "return_dtype"):
+ pf.udaf(lambda values: values.sum(), func_type="pandas")
+ with self.assertRaisesRegex(TypeError, "MAP"):
+ pf.udaf(lambda values: {}, return_dtype=dict[str, int],
func_type="pandas")
+ with self.assertRaises(TypeError):
+ pf.udaf(Sum, return_dtype=object())
+
+ class InvalidType(Sum):
+ def get_result_type(self):
+ return int
+
+ with self.assertRaisesRegex(TypeError, "get_result_type"):
+ pf.udaf(InvalidType)
+
+ def test_invalid_functions_and_async_methods(self):
+ class Scalar(ScalarFunction):
+ def eval(self, value):
+ return value
+
+ class NeedsArgument(Sum):
+ def __init__(self, argument):
+ self.argument = argument
+
+ for func in (42, Scalar, Scalar(), NeedsArgument):
+ with self.subTest(func=func), self.assertRaises(TypeError):
+ pf.udaf(func, return_dtype=int, func_type="pandas")
+ with self.assertRaises(ValueError):
+ pf.udaf(Sum, func_type="arrow")
+ with self.assertRaises(TypeError):
+ pf.udaf(Sum, deterministic=1)
+
+ async def coroutine(values):
+ return 1
+
+ async def generator(values):
+ yield 1
+
+ @functools.wraps(generator)
+ def wrapped(values):
+ return generator(values)
+
+ for func in (coroutine, generator, wrapped):
+ with self.subTest(func=func), self.assertRaisesRegex(TypeError,
"async"):
+ pf.udaf(func, return_dtype=int, func_type="pandas")
+ for method in ("accumulate", "get_value", "create_accumulator",
"open", "merge"):
+ cls = type("AsyncAggregate", (Sum,), {method: coroutine})
+ with self.subTest(method=method),
self.assertRaisesRegex(TypeError, "async"):
+ pf.udaf(cls)
+
+ def test_names_and_determinism(self):
+ for name in (None, "", "my_sum"):
+ self.assertEqual(pf.udaf(Sum, name=name).__name__, name or "Sum")
+
+ class RandomSum(Sum):
+ def is_deterministic(self):
+ return False
+
+ with self.assertRaisesRegex(ValueError, "Inconsistent deterministic"):
+ pf.udaf(RandomSum)
+ adapter = pf.udaf(RandomSum,
deterministic=False)._table_udf_wrapper._func
+ self.assertFalse(adapter.is_deterministic())
+
+ def test_client_constructor_and_worker_lifecycle(self):
+ events = []
+
+ class LifecycleSum(Sum):
+ def __init__(self):
+ events.append("init")
+
+ def open(self, context):
+ events.append(("open", context))
+
+ def create_accumulator(self):
+ events.append("accumulator")
+ return super().create_accumulator()
+
+ def close(self):
+ events.append("close")
+ raise RuntimeError("close failure")
+
+ declaration = pf.udaf(LifecycleSum)
+ adapter = declaration._table_udf_wrapper._func
+ self.assertEqual(events, ["init"])
+ adapter.open("worker")
+ accumulator = adapter.create_accumulator()
+ adapter.accumulate(accumulator, 5)
+ adapter.retract(accumulator, 2)
+ adapter.merge(accumulator, [[7], [1]])
+ self.assertEqual(accumulator, [11])
+ self.assertEqual(adapter.get_value(accumulator), 11)
+ with self.assertRaisesRegex(RuntimeError, "close failure"):
+ adapter.close()
+ self.assertEqual(events, ["init", ("open", "worker"), "accumulator",
"close"])
+ with self.assertRaisesRegex(RuntimeError, "before open"):
+ adapter.get_value(accumulator)
+ adapter.close()
+
+ def test_callable_classes_construct_on_worker(self):
+ events = []
+
+ class Total:
+ def __init__(self, offset=0):
+ events.append("init")
+ self.offset = offset
+
+ def __call__(self, values) -> int:
+ return int(values.sum()) + self.offset
+
+ declaration = pf.udaf(Total, func_type="pandas")
+ adapter = declaration._table_udf_wrapper._func
+ self.assertEqual(events, [])
+ adapter.open(None)
+ try:
+ acc = adapter.create_accumulator()
+ adapter.accumulate(acc, pd.Series([1, 2]))
+ self.assertEqual(adapter.get_value(acc), 3)
+ finally:
+ adapter.close()
+ self.assertEqual(events, ["init"])
+ for func in (Total(2), functools.partial(lambda offset, values:
values.sum() + offset, 2)):
+ self.assertEqual(_aggregate(pf.udaf(func, return_dtype=int,
func_type="pandas"),
+ pd.Series([1, 2])), 5)
+
+ def test_struct_normalization_and_user_failures(self):
+ dtype = pf.DataType.struct({
+ "nested": pf.DataType.list(pf.DataType.struct({"value":
pf.DataType.int64()})),
+ "missing": pf.DataType.string(),
+ })
+
+ class Nested(Sum):
+ def get_value(self, accumulator):
+ return {"nested": [{"value": accumulator[0]}]}
+
+ for func_type, expected in (
+ ("general", Row(nested=[Row(value=3)], missing=None)),
+ ("pandas", ([(3,)], None)),
+ ):
+ with self.subTest(func_type=func_type):
+ self.assertEqual(_aggregate(pf.udaf(
+ Nested, return_dtype=dtype, func_type=func_type), 3),
expected)
+
+ def fail(values) -> int:
+ raise RuntimeError("user failure")
+
+ with self.assertRaisesRegex(RuntimeError, "user failure"):
+ _aggregate(pf.udaf(fail, func_type="pandas"), pd.Series([1]))
+
+
+class DataFrameUDAFPlanningTests(PyFlinkDataFrameUTTestCase):
+ def test_sql_type_declarations(self):
+ class SqlSum(Sum):
+ def get_result_type(self):
+ return "BIGINT"
+
+ def get_accumulator_type(self):
+ return "ARRAY<BIGINT>"
+
+ for declaration in (pf.udaf(SqlSum), pf.udaf(
+ Sum, return_dtype="BIGINT", accumulator_type="ARRAY<BIGINT>")):
+ self.assertEqual(declaration.return_dtype, pf.DataType.int64())
+ self.assertEqual(declaration.accumulator_type,
pf.DataType.list(pf.DataType.int64()))
+
+
+class DataFrameUDAFBatchTests(PyFlinkBatchTableTestCase):
+ def setUp(self):
+ previous = pf.get_table_environment()
+ self.addCleanup(pf.set_table_environment, previous)
+ pf.set_table_environment(self.t_env)
+
+ def test_pandas_callables_classes_and_multiple_columns(self):
+ client_pid = os.getpid()
+
+ class Total:
+ def __init__(self):
+ if os.getpid() == client_pid:
+ raise AssertionError("Callable class must be constructed
on the worker")
+
+ def __call__(self, values, weights) -> int:
+ return int((values * weights).sum())
+
+ class WeightedSum(Sum):
+ def __init__(self):
+ self.pid = os.getpid()
+
+ def open(self, context):
+ if self.pid != client_pid:
+ raise AssertionError("AggregateFunction must be
constructed on the client")
+ self.opened = True
+
+ def accumulate(self, accumulator, values, weights):
+ if not self.opened:
+ raise AssertionError("open must be called")
+ accumulator[0] += int((values * weights).sum())
+
+ source = pf.from_records([("a", 1, 2), ("a", 3, 4), ("b", 2, 3)],
+ schema=["category", "value", "weight"])
+ weighted = pf.udaf(Total, func_type="pandas")
+ result = source.group_by("category").agg(
+ total=weighted(pf.col("value"), pf.col("weight")))
+ self.assertEqual(result.columns, ["category", "total"])
+ self.assertEqual(result.schema.get_field_data_types(),
+ [DataTypes.STRING(), DataTypes.BIGINT().not_null()])
+ second = source.group_by("category").agg(
+ total=pf.udaf(WeightedSum, func_type="pandas")(pf.col("value"),
pf.col("weight")))
+ self.assertCountEqual(result.union_all(second).collect(),
+ [Row("a", 14), Row("b", 6)] * 2)
+
+ def test_sql_bindings_and_struct_fields(self):
+ @pf.udaf(func_type="pandas")
+ def stats(values) -> Statistics:
+ return {"count": len(values), "total": int(values.sum())}
+
+ source = pf.from_dict({"value": [1, 2, 3]})
+ direct = source.agg(stats(pf.col("value")).alias("summary"))
+ explicit = pf.sql("SELECT s(`value`) AS summary FROM src",
auto_bind=False,
+ src=source, s=stats)
+ automatic = pf.sql("SELECT stats(`value`) AS summary FROM source")
+ for name in ("stats", "s"):
+ self.assertNotIn(name, self.t_env.list_user_defined_functions())
+ combined = direct.union_all(explicit).union_all(automatic)
+ projected = combined.select(pf.col("summary").get("total"),
+ pf.col("summary").get("count"))
+ self.assertCountEqual(projected.collect(), [Row(6, 3)] * 3)
+
+ def test_pandas_nested_struct_results(self):
+ def summarize(values):
+ return {"nested": {
+ "sum": {"label": "total", "value": int(values.sum())},
+ }}
+
+ class Summary(AggregateFunction):
+ def create_accumulator(self):
+ return [None]
+
+ def accumulate(self, accumulator, values):
+ accumulator[0] = summarize(values)
+
+ def get_value(self, accumulator):
+ return accumulator[0]
+
+ dtype = "ROW<nested MAP<STRING, ROW<value BIGINT, label STRING>>>"
+ function = pf.udaf(summarize, return_dtype=dtype, func_type="pandas")
+ aggregate = pf.udaf(Summary, return_dtype=dtype, func_type="pandas")
+ source = pf.from_dict({"value": [1, 2, 3]})
+ result = source.agg(function(pf.col("value")),
aggregate(pf.col("value")))
+ expected = Row({"sum": Row(6, "total")})
+ self.assertEqual(result.collect(), [Row(expected, expected)])
+
+ def test_empty_nulls_and_table_udaf_compatibility(self):
+ def total(values) -> int:
+ return int(values.sum())
+
+ source = pf.from_dict({"value": [1, None, 3]})
+ declaration = pf.udaf(total, func_type="pandas")
+ native = table_udaf(total, result_type=DataTypes.BIGINT(),
func_type="pandas")
+
self.assertEqual(source.agg(total=declaration(pf.col("value"))).collect(),
[Row(4)])
+ empty = source.filter(pf.col("value") < 0)
+
self.assertEqual(empty.agg(total=declaration(pf.col("value"))).collect(),
+ empty.agg(total=native(pf.col("value"))).collect())
+ with self.assertRaisesRegex(TableException, "non-Pandas UDAFs"):
+
source.agg(total=pf.udaf(Sum)(pf.col("value"))).to_table().explain()
+
+
+class DataFrameUDAFStreamTests(PyFlinkStreamDataFrameTestCase):
+ def test_grouped_global_and_retracting_inputs(self):
+ total = pf.udaf(Sum)
+ source = pf.from_records([("a", 1), ("a", 2), ("b", 3), ("b", None)],
+ schema=["category", "value"])
+ grouped = source.group_by("category").agg(total=total(pf.col("value")))
+ self.assertEqual(self._materialize(grouped, key=["category"]), [("a",
3), ("b", 3)])
+
self.assertEqual(self._materialize(source.agg(total=total(pf.col("value"))),
key=[]),
+ [(6,)])
+ intermediate =
source.group_by("category").agg(subtotal=pf.col("value").sum)
+ result = intermediate.agg(total=total(pf.col("subtotal")))
+ self.assertEqual(self._materialize(result, key=[]), [(6,)])
+ with self.assertRaisesRegex(TableException, "Pandas UDAFs"):
+ source.agg(total=pf.udaf(lambda values: values.sum(),
return_dtype=int,
+
func_type="pandas")(pf.col("value"))).to_table().explain()
+
+ def test_general_sql_and_empty_input(self):
+ total = pf.udaf(Sum())
+ source = pf.from_dict({"value": [1, 2, 3]})
+ result = pf.sql("SELECT total(`value`) AS total_value FROM source",
+ auto_bind=False, source=source, total=total)
+ self.assertEqual(self._materialize(result, key=[]), [(6,)])
+ empty = source.filter(pf.col("value") < 0)
+ native = table_udaf(Sum(), result_type=DataTypes.BIGINT())
+ self.assertEqual(self._materialize(empty.agg(total(pf.col("value"))),
key=[]),
+ self._materialize(empty.agg(native(pf.col("value"))),
key=[]))
+
+ def test_dataview_accumulator_is_preserved(self):
+ class DistinctCount(AggregateFunction):
+ def create_accumulator(self):
+ return Row(MapView(), ListView())
+
+ def accumulate(self, accumulator, value):
+ if value not in accumulator[0]:
+ accumulator[0][value] = True
+ accumulator[1].add(value)
+
+ def get_value(self, accumulator) -> int:
+ return len(list(accumulator[1].get()))
+
+ def get_accumulator_type(self):
+ return DataTypes.ROW([
+ DataTypes.FIELD("seen",
DataTypes.MAP_VIEW(DataTypes.STRING(),
+
DataTypes.BOOLEAN())),
+ DataTypes.FIELD("values",
DataTypes.LIST_VIEW(DataTypes.STRING())),
+ ])
+
+ source = pf.from_records([("a", "x"), ("a", "x"), ("a", "y"), ("b",
"x")],
+ schema=["category", "value"])
+ result =
source.group_by("category").agg(count=pf.udaf(DistinctCount)(pf.col("value")))
+ self.assertEqual(self._materialize(result, key=["category"]), [("a",
2), ("b", 1)])
+
+ def test_session_window_merges_accumulators(self):
+ data = self.env.from_collection([
+ Row("a", 1, datetime.datetime(2020, 1, 1, 0, 0, 1)),
+ Row("a", 3, datetime.datetime(2020, 1, 1, 0, 0, 3)),
+ Row("a", 2, datetime.datetime(2020, 1, 1, 0, 0, 2)),
+ ], type_info=Types.ROW_NAMED(["category", "value", "ts"],
+ [Types.STRING(), Types.INT(),
Types.SQL_TIMESTAMP()]))
+ schema = Schema.new_builder().column("category", DataTypes.STRING()) \
+ .column("value", DataTypes.INT()) \
+ .column("ts",
DataTypes.TIMESTAMP(3).bridged_to("java.sql.Timestamp")) \
+ .watermark("ts", "ts - INTERVAL '10' SECOND").build()
+ table = self.t_env.from_data_stream(data, schema)
+ total = pf.udaf(Sum)
+ result =
table.window(Session.with_gap(lit(2).seconds).on(col("ts")).alias("w")) \
+ .group_by(col("category"), col("w")) \
+ .select(col("category"), total(col("value")).alias("total"))
+ self.assertEqual(pf.from_table(result).collect(), [Row("a", 6)])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/flink-python/pyflink/dataframe/udaf.py
b/flink-python/pyflink/dataframe/udaf.py
new file mode 100644
index 00000000000..dd82b267f6f
--- /dev/null
+++ b/flink-python/pyflink/dataframe/udaf.py
@@ -0,0 +1,318 @@
+################################################################################
+# 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.
+################################################################################
+
+"""User-defined aggregate functions for the DataFrame API."""
+
+import inspect
+from dataclasses import dataclass, field
+from typing import Any, Callable, Optional, Tuple, Type, Union, cast, overload
+
+from pyflink.dataframe.datatype import DataType
+from pyflink.dataframe.udf import (
+ _DataTypeLike,
+ _UDFDeclarationContext,
+ _convert_to_dtype,
+ _create_declaration_context,
+ _create_result_normalizer,
+ _infer_return_dtype,
+ _resolve_udf,
+ _validate_determinism_agreement,
+ _validate_zero_argument_class,
+)
+from pyflink.table.expression import Expression
+from pyflink.table.types import DataType as TableDataType, MapType
+from pyflink.table.udf import (
+ AggregateFunction,
+ DelegatingPandasAggregateFunction,
+ UserDefinedAggregateFunctionWrapper,
+ UserDefinedFunction,
+ udaf as table_udaf,
+)
+from pyflink.util.api_stability_decorators import PublicEvolving
+
+__all__ = ["udaf"]
+
+_UDAFInput = Union[Callable[..., Any], AggregateFunction, Type]
+
+
+@dataclass(frozen=True)
+class _DataFrameUDAFWrapper:
+ _func: _UDAFInput
+ return_dtype: DataType
+ accumulator_type: Optional[DataType]
+ _func_type: str
+ _deterministic: bool
+ _name: Optional[str]
+ _cached_table_udf_wrapper: Optional[UserDefinedAggregateFunctionWrapper] =
field(
+ default=None, init=False, repr=False, compare=False)
+
+ @property
+ def __name__(self) -> str:
+ return self._table_udf_wrapper._name
+
+ def __call__(self, *args: Any) -> Expression:
+ return self._table_udf_wrapper(*args)
+
+ @property
+ def _table_udf_wrapper(self) -> UserDefinedAggregateFunctionWrapper:
+ if self._cached_table_udf_wrapper is None:
+ wrapper = table_udaf(
+ _DataFrameAggregateFunctionAdapter(
+ self._func, self.return_dtype, self._deterministic,
self._func_type),
+ result_type=self.return_dtype._to_table_data_type(),
+ accumulator_type=(self.accumulator_type._to_table_data_type()
+ if self.accumulator_type is not None else
None),
+ func_type=self._func_type,
+ deterministic=self._deterministic,
+ name=self._name,
+ )
+ object.__setattr__(self, "_cached_table_udf_wrapper", wrapper)
+ return cast(UserDefinedAggregateFunctionWrapper,
self._cached_table_udf_wrapper)
+
+
+@overload
+def udaf(
+ func: _UDAFInput, *, return_dtype: Optional[_DataTypeLike] = ...,
+ accumulator_type: Optional[_DataTypeLike] = ..., func_type: str = ...,
+ deterministic: bool = ..., name: Optional[str] = ...,
+) -> _DataFrameUDAFWrapper:
+ ...
+
+
+@overload
+def udaf(
+ func: None = ..., *, return_dtype: Optional[_DataTypeLike] = ...,
+ accumulator_type: Optional[_DataTypeLike] = ..., func_type: str = ...,
+ deterministic: bool = ..., name: Optional[str] = ...,
+) -> Callable[[_UDAFInput], _DataFrameUDAFWrapper]:
+ ...
+
+
+@PublicEvolving()
+def udaf(
+ func: Optional[_UDAFInput] = None,
+ *,
+ return_dtype: Optional[_DataTypeLike] = None,
+ accumulator_type: Optional[_DataTypeLike] = None,
+ func_type: str = "general",
+ deterministic: bool = True,
+ name: Optional[str] = None,
+) -> Union[_DataFrameUDAFWrapper, Callable[[_UDAFInput],
_DataFrameUDAFWrapper]]:
+ """
+ Declare a function that aggregates multiple rows into one result.
+
+ General UDAFs require an :class:`~pyflink.table.udf.AggregateFunction`
instance
+ or class. Pandas UDAFs also accept functions and callable instances or
classes;
+ each argument contains a group or window's column values as pandas data.
+ Calling the declaration with column expressions produces an aggregation
expression
+ for ``agg()``. Declarations can also be bound in
:func:`pyflink.dataframe.sql`.
+
+ ``AggregateFunction`` classes are constructed on the client. Plain
callable classes
+ are constructed on workers. Classes must have a zero-argument constructor;
+ provided instances are serialized with the job. ``AggregateFunction.open``
and
+ ``AggregateFunction.close`` run on workers.
+
+ Explicit types take precedence over ``get_result_type()`` and
+ ``get_accumulator_type()``. Without a result type, the return annotation of
+ ``get_value()`` or the callable is used. Accumulator types are not
inferred.
+ Aggregation results follow the nested value conversion rules of DataFrame
scalar UDFs.
+
+ Ordinary grouped and global aggregations support general UDAFs in
streaming mode
+ and pandas UDAFs in batch mode. Window aggregations follow Table API
support.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> from pyflink.table import EnvironmentSettings, TableEnvironment
+ >>> from pyflink.table.udf import AggregateFunction
+ >>> pf.set_table_environment(TableEnvironment.create(
+ ... EnvironmentSettings.in_streaming_mode()))
+ >>> class Sum(AggregateFunction):
+ ... def create_accumulator(self):
+ ... return [0]
+ ... def accumulate(self, accumulator, value):
+ ... accumulator[0] += value
+ ... def get_value(self, accumulator) -> int:
+ ... return accumulator[0]
+ >>> total = pf.udaf(Sum, accumulator_type=list[int])
+ >>> df = pf.from_dict({"category": ["a", "a", "b"], "value": [1, 2,
3]})
+ >>> result = df.group_by("category").agg(total=total(pf.col("value")))
+ >>> result.columns
+ ['category', 'total']
+
+ A pandas callable can be used in batch aggregations::
+
+ >>> pf.set_table_environment(TableEnvironment.create(
+ ... EnvironmentSettings.in_batch_mode()))
+ >>> df = pf.from_dict({"category": ["a", "a", "b"], "value": [1, 2,
3]})
+ >>> @pf.udaf(return_dtype=float, func_type="pandas")
+ ... def mean(values):
+ ... return values.mean()
+ >>> result = df.group_by("category").agg(mean=mean(pf.col("value")))
+ >>> result = pf.sql(
+ ... "SELECT category, avg_value(`value`) AS mean FROM src GROUP BY
category",
+ ... auto_bind=False, src=df, avg_value=mean)
+
+ :param func: AggregateFunction instance/class, or a pandas callable
instance/class/function.
+ :param return_dtype: Result type as a DataFrame DataType, Python type, or
SQL type string.
+ :param accumulator_type: General UDAF state type, in the same formats as
``return_dtype``.
+ Omit when ``get_accumulator_type()`` supplies it.
Pandas UDAFs
+ use the accumulator type managed by the Table API.
+ :param func_type: ``"general"`` (default) or ``"pandas"``.
+ :param deterministic: Must agree with
``AggregateFunction.is_deterministic()``.
+ :param name: Optional function name. None or an empty string uses the
callable's name.
+ :return: A reusable aggregate declaration, or a decorator when ``func`` is
omitted.
+
+ .. versionadded:: 2.4.0
+ """
+ if func_type not in ("general", "pandas"):
+ raise ValueError("func_type must be 'general' or 'pandas'.")
+ if not isinstance(deterministic, bool):
+ raise TypeError("deterministic must be a bool.")
+
+ def decorator(f: _UDAFInput) -> _DataFrameUDAFWrapper:
+ actual_func, context = _resolve_udaf(f, func_type)
+ result = _resolve_aggregate_type(actual_func, return_dtype,
"get_result_type")
+ if result is None:
+ result = _infer_return_dtype(context, None, name or "UDAF")
+ if func_type == "pandas" and isinstance(result._to_table_data_type(),
MapType):
+ raise TypeError("Pandas UDAFs do not support MAP results.")
+ accumulator = _resolve_aggregate_type(
+ actual_func, accumulator_type, "get_accumulator_type"
+ ) if func_type == "general" or accumulator_type is not None else None
+ if func_type == "general" and accumulator is None:
+ raise TypeError("Specify accumulator_type or implement
get_accumulator_type().")
+ if isinstance(actual_func, AggregateFunction):
+ _validate_determinism_agreement(deterministic,
actual_func.is_deterministic())
+ return _DataFrameUDAFWrapper(
+ actual_func, result, accumulator, func_type, deterministic, name)
+
+ return decorator if func is None else decorator(func)
+
+
+def _resolve_udaf(func: _UDAFInput, func_type: str) -> Tuple[_UDAFInput,
_UDFDeclarationContext]:
+ if inspect.isclass(func) and issubclass(func, AggregateFunction):
+ _validate_zero_argument_class(func)
+ func = func()
+ if not isinstance(func, AggregateFunction):
+ raise TypeError(
+ "An AggregateFunction class must construct an
AggregateFunction instance.")
+ if isinstance(func, AggregateFunction):
+ for method_name in (
+ "create_accumulator", "accumulate", "get_value", "retract",
"merge",
+ "open", "close", "get_result_type", "get_accumulator_type",
"is_deterministic",
+ ):
+ _validate_sync_aggregate_method(getattr(func, method_name))
+ context = _create_declaration_context(func.get_value,
partial_source=func.get_value)
+ else:
+ if func_type == "general":
+ raise TypeError("General UDAFs require an AggregateFunction
instance or class.")
+ if isinstance(func, UserDefinedFunction) or (
+ inspect.isclass(func) and issubclass(func, UserDefinedFunction)
+ ):
+ raise TypeError("Pandas UDAFs require an AggregateFunction or a
Python callable.")
+ context = _resolve_udf(func).declaration_context
+ _validate_sync_aggregate_method(context.annotation_target)
+ return func, context
+
+
+def _validate_sync_aggregate_method(func: Callable[..., Any]) -> None:
+ if not callable(func):
+ raise TypeError("UDAF methods must be callable.")
+ for candidate in (func, inspect.unwrap(func)):
+ if inspect.iscoroutinefunction(candidate) or
inspect.isasyncgenfunction(candidate):
+ raise TypeError("UDAFs must be synchronous; async functions are
not supported.")
+
+
+def _resolve_aggregate_type(
+ func: _UDAFInput, explicit_type: Optional[_DataTypeLike], method_name: str,
+) -> Optional[DataType]:
+ if explicit_type is not None:
+ return _convert_to_dtype(explicit_type)
+ if isinstance(func, AggregateFunction):
+ method = getattr(func, method_name)
+ if getattr(method, "__func__", None) is not getattr(AggregateFunction,
method_name):
+ declared_type = method()
+ if isinstance(declared_type, TableDataType):
+ return DataType(declared_type)
+ if isinstance(declared_type, str):
+ return _convert_to_dtype(declared_type)
+ raise TypeError(f"{method_name}() must return a Table DataType or
SQL type string.")
+ return None
+
+
+class _DataFrameAggregateFunctionAdapter(AggregateFunction):
+ def __init__(
+ self, func: _UDAFInput, return_dtype: DataType, deterministic: bool,
func_type: str,
+ ) -> None:
+ self._func = func
+ self._return_dtype = return_dtype
+ self._deterministic = deterministic
+ self._func_type = func_type
+ self._active_func: Optional[AggregateFunction] = None
+ self._normalizer: Optional[Callable[[Any], Any]] = None
+ self.__name__ = getattr(func, "__name__", type(func).__name__)
+
+ def open(self, function_context: Any) -> None:
+ if isinstance(self._func, AggregateFunction):
+ active_func = self._func
+ else:
+ invoke = self._func() if inspect.isclass(self._func) else
self._func
+ if not callable(invoke):
+ raise TypeError("UDAF class must construct a callable
instance.")
+ active_func = DelegatingPandasAggregateFunction(invoke)
+ # Arrow requires tuples or dictionaries for structs, including nested
ones.
+ normalizer = _create_result_normalizer(
+ self._return_dtype._to_table_data_type(),
row_as_tuple=self._func_type == "pandas")
+ active_func.open(function_context)
+ self._active_func = active_func
+ self._normalizer = normalizer
+
+ def close(self) -> None:
+ try:
+ if self._active_func is not None:
+ self._active_func.close()
+ finally:
+ self._active_func = None
+ self._normalizer = None
+
+ def is_deterministic(self) -> bool:
+ if isinstance(self._func, AggregateFunction):
+ return self._func.is_deterministic()
+ return self._deterministic
+
+ def _active(self) -> AggregateFunction:
+ if self._active_func is None:
+ raise RuntimeError("UDAF was invoked before open().")
+ return self._active_func
+
+ def create_accumulator(self) -> Any:
+ return self._active().create_accumulator()
+
+ def accumulate(self, accumulator: Any, *args: Any) -> None:
+ self._active().accumulate(accumulator, *args)
+
+ def retract(self, accumulator: Any, *args: Any) -> None:
+ self._active().retract(accumulator, *args)
+
+ def merge(self, accumulator: Any, accumulators: Any) -> None:
+ self._active().merge(accumulator, accumulators)
+
+ def get_value(self, accumulator: Any) -> Any:
+ result: Any = self._active().get_value(accumulator)
+ return self._normalizer(result) if self._normalizer is not None else
result
diff --git a/flink-python/pyflink/dataframe/udf.py
b/flink-python/pyflink/dataframe/udf.py
index df0b92da8d3..2d036454edb 100644
--- a/flink-python/pyflink/dataframe/udf.py
+++ b/flink-python/pyflink/dataframe/udf.py
@@ -1555,12 +1555,13 @@ def _wrap_map_arrow(func, input_type: RowType,
output_type: RowType):
def _create_result_normalizer(
- data_type: Any, *, require_row_fields: bool = False,
+ data_type: Any, *, require_row_fields: bool = False, row_as_tuple: bool =
False,
) -> Optional[Callable[[Any], Any]]:
if isinstance(data_type, RowType):
field_names = data_type.field_names()
field_normalizers = tuple(
- _create_result_normalizer(field.data_type,
require_row_fields=require_row_fields)
+ _create_result_normalizer(
+ field.data_type, require_row_fields=require_row_fields,
row_as_tuple=row_as_tuple)
for field in data_type
)
@@ -1577,6 +1578,8 @@ def _create_result_normalizer(
field_values, field_normalizers
)
]
+ if row_as_tuple:
+ return tuple(normalized_fields)
row = Row(*normalized_fields)
row.set_field_names(field_names)
if isinstance(value, Row):
@@ -1586,7 +1589,8 @@ def _create_result_normalizer(
return normalize_row
if isinstance(data_type, ArrayType):
element_normalizer = _create_result_normalizer(
- data_type.element_type, require_row_fields=require_row_fields)
+ data_type.element_type, require_row_fields=require_row_fields,
+ row_as_tuple=row_as_tuple)
if element_normalizer is None:
def normalize_leaf_array(value: Any) -> Any:
@@ -1602,9 +1606,9 @@ def _create_result_normalizer(
return normalize_array
if isinstance(data_type, MapType):
key_normalizer = _create_result_normalizer(
- data_type.key_type, require_row_fields=require_row_fields)
+ data_type.key_type, require_row_fields=require_row_fields,
row_as_tuple=row_as_tuple)
value_normalizer = _create_result_normalizer(
- data_type.value_type, require_row_fields=require_row_fields)
+ data_type.value_type, require_row_fields=require_row_fields,
row_as_tuple=row_as_tuple)
def normalize_map(value: Any) -> Any:
if value is None:
diff --git a/flink-python/pyflink/fn_execution/tests/test_coders.py
b/flink-python/pyflink/fn_execution/tests/test_coders.py
index 3bbc8cf859f..50a8c0690b1 100644
--- a/flink-python/pyflink/fn_execution/tests/test_coders.py
+++ b/flink-python/pyflink/fn_execution/tests/test_coders.py
@@ -36,6 +36,32 @@ from pyflink.testing.test_case_utils import PyFlinkTestCase
class ArrowSchemaTests(unittest.TestCase):
+ def test_pandas_row_preserves_field_nullability(self):
+ import pandas as pd
+ from pyflink.table import DataTypes
+ from pyflink.table.types import create_arrow_schema
+ from pyflink.table.utils import pandas_to_arrow
+
+ row_type = DataTypes.ROW([
+ DataTypes.FIELD("total", DataTypes.BIGINT().not_null()),
+ DataTypes.FIELD("label", DataTypes.STRING()),
+ ]).not_null()
+ schema = create_arrow_schema(["summary"], [row_type])
+ values = [{"total": 6, "label": None}, {"total": 9, "label": "ok"}]
+ dataframe = pd.DataFrame(values)
+ for size in (0, len(values)):
+ with self.subTest(size=size):
+ batch = pandas_to_arrow(schema, pytz.UTC, [row_type],
[dataframe.iloc[:size]])
+ with pa.BufferOutputStream() as output:
+ with pa.ipc.new_stream(output, schema) as writer:
+ writer.write_batch(batch)
+ decoded =
pa.ipc.open_stream(output.getvalue()).read_next_batch()
+ self.assertEqual(decoded.schema, schema)
+
self.assertFalse(decoded.column(0).type.field("total").nullable)
+ self.assertTrue(decoded.column(0).type.field("label").nullable)
+ self.assertEqual(decoded.to_pylist(),
+ [{"summary": value} for value in
values[:size]])
+
def test_pandas_positional_rows_require_matching_arity(self):
import pandas as pd
from pyflink.common import Row
diff --git a/flink-python/pyflink/table/tests/test_pandas_udaf.py
b/flink-python/pyflink/table/tests/test_pandas_udaf.py
index 09e17dd1fca..e9bca487db0 100644
--- a/flink-python/pyflink/table/tests/test_pandas_udaf.py
+++ b/flink-python/pyflink/table/tests/test_pandas_udaf.py
@@ -125,6 +125,25 @@ class BatchPandasUDAFITTests(PyFlinkBatchTableTestCase):
actual = source_sink_utils.results()
self.assert_equals(actual, ["+I[5]"])
+ def test_row_result_with_non_nullable_fields(self):
+ source = self.t_env.from_elements(
+ [(1,), (2,), (3,)],
+ DataTypes.ROW([DataTypes.FIELD("value", DataTypes.BIGINT())]))
+
+ @udaf(result_type=DataTypes.ROW([
+ DataTypes.FIELD("total", DataTypes.BIGINT().not_null()),
+ DataTypes.FIELD("count", DataTypes.BIGINT().not_null()),
+ DataTypes.FIELD("label", DataTypes.STRING()),
+ ]).not_null(), func_type="pandas")
+ def stats(values):
+ return int(values.sum()), len(values), None
+
+ summary = source.select(stats(source.value).alias("summary"))
+ result = summary.select(col("summary").get("total"),
col("summary").get("count"),
+ col("summary").get("label"))
+ with result.execute().collect() as rows:
+ self.assertEqual(list(rows), [Row(6, 3, None)])
+
def test_group_aggregate_with_aux_group(self):
t = self.t_env.from_elements(
[(1, 2, 3), (3, 2, 3), (2, 1, 3), (1, 5, 4), (1, 8, 6), (2, 3, 4)],
diff --git a/flink-python/pyflink/table/utils.py
b/flink-python/pyflink/table/utils.py
index 7001dc9dc9a..f2467e41802 100644
--- a/flink-python/pyflink/table/utils.py
+++ b/flink-python/pyflink/table/utils.py
@@ -217,10 +217,9 @@ def pandas_to_arrow(schema, timezone, field_types, series):
field_type = field_types[i]
schema_type = schema.types[i]
if type(s) == pd.DataFrame:
- array_names = [(create_array(s[s.columns[j]], field.type),
field.name)
- for j, field in enumerate(schema_type)]
- struct_arrays, struct_names = zip(*array_names)
- arrays.append(pa.StructArray.from_arrays(struct_arrays,
struct_names))
+ struct_arrays = [create_array(s[s.columns[j]], field.type)
+ for j, field in enumerate(schema_type)]
+ arrays.append(pa.StructArray.from_arrays(struct_arrays,
fields=list(schema_type)))
else:
arrays.append(create_array(
tz_convert_to_internal(s, field_type, timezone), schema_type))
