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 20172254a50 [FLINK-40433][python] Introduce map and map_batches in
DataFrame API
20172254a50 is described below
commit 20172254a5006dcdcc449808778be260bfacc00d
Author: auroflow <[email protected]>
AuthorDate: Mon Oct 12 19:27:53 2026 +0800
[FLINK-40433][python] Introduce map and map_batches in DataFrame API
This closes #29469.
---
.../docs/reference/pyflink.dataframe/dataframe.rst | 2 +
flink-python/pyflink/dataframe/dataframe.py | 148 +++
flink-python/pyflink/dataframe/tests/test_udf.py | 1015 +++++++++++++++++++-
flink-python/pyflink/dataframe/tests/test_udtf.py | 12 +
flink-python/pyflink/dataframe/udf.py | 371 ++++++-
.../pyflink/fn_execution/coder_impl_fast.pxd | 1 +
.../pyflink/fn_execution/coder_impl_fast.pyx | 10 +-
.../pyflink/fn_execution/coder_impl_slow.py | 10 +-
flink-python/pyflink/fn_execution/coders.py | 19 +-
.../pyflink/fn_execution/flink_fn_execution_pb2.py | 22 +-
.../fn_execution/flink_fn_execution_pb2.pyi | 6 +-
.../pyflink/fn_execution/tests/test_coders.py | 113 ++-
.../pyflink/proto/flink-fn-execution.proto | 2 +
flink-python/pyflink/table/table.py | 9 +-
.../pyflink/table/tests/test_pandas_udf.py | 38 +-
.../table/tests/test_row_based_operation.py | 24 +-
flink-python/pyflink/table/udf.py | 3 +
flink-python/pyflink/table/utils.py | 40 +-
.../org/apache/flink/python/util/ProtoUtils.java | 11 +
.../AbstractPythonScalarFunctionOperator.java | 14 +
.../arrow/ArrowPythonScalarFunctionOperator.java | 6 +-
.../operations/utils/OperationTreeBuilder.java | 3 +-
.../python/PythonAsyncScalarFunction.java | 71 ++
.../table/functions/python/PythonFunction.java | 5 +
.../functions/python/PythonScalarFunction.java | 71 ++
.../plan/rules/logical/PythonMapMergeRule.java | 8 +
.../PythonCalcSplitFunctionKindRuleTest.java | 37 +
.../PythonCalcSplitFunctionKindRuleTest.xml | 42 +
28 files changed, 2040 insertions(+), 73 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
index d03d62b27e6..1fc52491ec5 100644
--- a/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/dataframe.rst
@@ -73,6 +73,8 @@ Transformations
DataFrame.rename_columns
DataFrame.rename
DataFrame.filter
+ DataFrame.map
+ DataFrame.map_batches
DataFrame.where
DataFrame.explode
DataFrame.drop_duplicates
diff --git a/flink-python/pyflink/dataframe/dataframe.py
b/flink-python/pyflink/dataframe/dataframe.py
index e81bfee1d95..1d3eedd0973 100644
--- a/flink-python/pyflink/dataframe/dataframe.py
+++ b/flink-python/pyflink/dataframe/dataframe.py
@@ -184,6 +184,154 @@ class DataFrame:
# ======================== Core Operations ========================
+ @PublicEvolving()
+ def map(self, func, *, return_dtype=None) -> "DataFrame":
+ """
+ Apply a function to each row, producing a new DataFrame.
+
+ Pass a Python function, a callable class or instance, a
``ScalarFunction`` or
+ ``AsyncScalarFunction`` class or instance, or a UDF created with
+ :func:`pyflink.dataframe.udf`. The function (or its ``__call__`` or
``eval`` method)
+ receives a dictionary whose keys are input column names and whose
values are the
+ values in that row. It must return one dictionary whose keys name the
output columns
+ and whose values form the output row. The output replaces the input
columns.
+
+ Specify the output schema with ``return_dtype``, or omit it when the
callable has a
+ ``TypedDict`` return annotation. For an already declared UDF, its
schema is reused
+ and ``return_dtype`` must be omitted. Output columns follow the
schema's order;
+ extra dictionary keys are ignored, and missing output fields raise an
error.
+
+ :param func: A callable or a DataFrame UDF declared with
``func_type="general"``.
+ See :func:`pyflink.dataframe.udf` for callable lifecycle
behavior.
+ :param return_dtype: Output schema as a struct DataType, TypedDict
type, or SQL ROW
+ type string. May be inferred from a callable's
return annotation.
+ :return: A new DataFrame with one output row per input row.
+
+ Examples:
+
+ A Python function with an explicit output schema::
+
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([(1, "alice")], schema=["id", "name"])
+ >>> output_type = pf.DataType.struct({"label":
pf.DataType.string()})
+ >>> result = df.map(
+ ... lambda row: {"label": row["name"].upper()},
+ ... return_dtype=output_type,
+ ... )
+
+ Infer the output schema from a TypedDict return annotation::
+
+ >>> from typing import TypedDict
+ >>> class Output(TypedDict):
+ ... label: str
+ >>> def uppercase(row) -> Output:
+ ... return {"label": row["name"].upper()}
+ >>> result = df.map(uppercase)
+
+ Pass a callable class, or configure an instance before passing it::
+
+ >>> class PrefixName:
+ ... def __init__(self, prefix="user_"):
+ ... self.prefix = prefix
+ ... def __call__(self, row) -> Output:
+ ... return {"label": self.prefix + row["name"]}
+ >>> result = df.map(PrefixName)
+ >>> result = df.map(PrefixName("member_"))
+
+ Use a ScalarFunction's eval method::
+
+ >>> from pyflink.table.udf import ScalarFunction
+ >>> class UppercaseName(ScalarFunction):
+ ... def eval(self, row) -> Output:
+ ... return {"label": row["name"].upper()}
+ >>> result = df.map(UppercaseName())
+
+ Reuse a declared DataFrame UDF without repeating its output schema::
+
+ >>> declared_uppercase = pf.udf(uppercase)
+ >>> result = df.map(declared_uppercase)
+
+ .. versionadded:: 2.4.0
+ """
+ from pyflink.dataframe.udf import _resolve_map_udf
+
+ function = _resolve_map_udf(func, return_dtype, "general",
self.schema.to_row_data_type())
+ return DataFrame(self._table.map(function(*[table_col(name) for name
in self.columns])))
+
+ @PublicEvolving()
+ def map_batches(self, func, *, return_dtype=None, batch_format: str) ->
"DataFrame":
+ """
+ Apply a function to batches of rows, producing a new DataFrame.
+
+ Flink automatically groups input rows into batches and calls ``func``
once per batch.
+ Pass a Python function, a callable class or instance, a
``ScalarFunction`` class or
+ instance, or a UDF created with :func:`pyflink.dataframe.udf`. The
callable receives
+ and returns a dictionary keyed by column name: each non-ROW value is a
``pandas.Series``
+ when ``batch_format="pandas"``, or a ``pyarrow.Array`` when
``batch_format="arrow"``.
+ Input values contain the batched column values; returned values form
the output
+ columns, replacing the input columns.
+
+ Each output row must correspond to the input row at the same position,
preserving
+ the number and order of rows. Specify
+ the output struct schema with ``return_dtype``; for a declared UDF,
omit it to reuse
+ that UDF's schema. Output columns follow the schema's order; extra
dictionary keys
+ are ignored, and missing output fields raise an error. Pandas output
columns align
+ by index.
+
+ .. note::
+
+ For ROW input and output columns, use ``pandas.DataFrame`` in
pandas mode
+ and ``pyarrow.StructArray`` in Arrow mode.
+
+ :param func: A callable or a DataFrame UDF declared for the selected
batch format.
+ :param return_dtype: Output struct schema. Required for raw callables
in both formats;
+ must be omitted for an already declared UDF.
+ :param batch_format: ``"pandas"`` or ``"arrow"``. Required even for
declared UDFs.
+ :return: A new DataFrame with the same number of rows as the input.
+
+ Examples:
+
+ Transform pandas batches with a Python function::
+
+ >>> import pandas as pd
+ >>> import pyflink.dataframe as pf
+ >>> df = pf.from_records([(1,), (2,), (3,)], schema=["value"])
+ >>> output_type = pf.DataType.struct({"doubled":
pf.DataType.int64()})
+ >>> def double_pandas(batch: dict[str, pd.Series]) -> dict[str,
pd.Series]:
+ ... return {"doubled": batch["value"] * 2}
+ >>> result = df.map_batches(
+ ... double_pandas, batch_format="pandas",
return_dtype=output_type,
+ ... )
+
+ Use a callable class to transform Arrow batches::
+
+ >>> import pyarrow as pa
+ >>> import pyarrow.compute as pc
+ >>> class DoubleArrow:
+ ... def __call__(self, batch: dict[str, pa.Array]) ->
dict[str, pa.Array]:
+ ... return {"doubled": pc.multiply(batch["value"], 2)}
+ >>> result = df.map_batches(
+ ... DoubleArrow, batch_format="arrow",
return_dtype=output_type,
+ ... )
+
+ Reuse a declared UDF, specifying its batch format at the call site::
+
+ >>> declared_double = pf.udf(
+ ... double_pandas, return_dtype=output_type,
func_type="pandas",
+ ... )
+ >>> result = df.map_batches(declared_double, batch_format="pandas")
+
+ .. versionadded:: 2.4.0
+ """
+ from pyflink.dataframe.udf import _resolve_map_udf
+
+ if batch_format not in ("pandas", "arrow"):
+ raise ValueError("batch_format must be 'pandas' or 'arrow'.")
+ function = _resolve_map_udf(
+ func, return_dtype, batch_format, self.schema.to_row_data_type()
+ )
+ return DataFrame(self._table.map(function(*[table_col(name) for name
in self.columns])))
+
@PublicEvolving()
def filter(
self,
diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py
b/flink-python/pyflink/dataframe/tests/test_udf.py
index 66fe1334bf0..a8b2e0e1a9c 100644
--- a/flink-python/pyflink/dataframe/tests/test_udf.py
+++ b/flink-python/pyflink/dataframe/tests/test_udf.py
@@ -33,7 +33,7 @@ import pyflink.dataframe as pf
from pyflink.common import Row, RowKind
from pyflink.table import DataTypes as TableDataTypes
from pyflink.table.expression import Expression
-from pyflink.table.types import RowType
+from pyflink.table.types import MapType, RowType
from pyflink.table.udf import AsyncScalarFunction, ScalarFunction,
TableFunction
from pyflink.testing.test_case_utils import (
PyFlinkDataFrameUTTestCase,
@@ -934,11 +934,8 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
False,
),
("class default", lambda: pf.udf(DefaultDeterministic), True),
- (
- "class matching metadata",
- lambda: pf.udf(NonDeterministic, deterministic=False),
- False,
- ),
+ ("matching class metadata", lambda: pf.udf(NonDeterministic,
deterministic=False),
+ False),
]
for case_name, declare, expected in declarations:
with self.subTest(case=case_name):
@@ -950,12 +947,40 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
)
with self.assertRaisesRegex(ValueError, "Inconsistent deterministic"):
pf.udf(instance)
- self.assertTrue(pf.udf(NonDeterministic)._deterministic)
-
named = pf.udf(instance, deterministic=False, name="identity")
self.assertEqual(named.__name__, "identity")
self.assertEqual(named._table_udf_wrapper._name, "identity")
+ def test_scalar_function_classes_defer_custom_determinism_validation(self):
+ for base in (ScalarFunction, AsyncScalarFunction):
+ for deterministic in (True, False):
+ constructed = []
+
+ class CustomDeterminism(base):
+ def __init__(self):
+ constructed.append(True)
+
+ def is_deterministic(self):
+ return deterministic
+
+ if base is AsyncScalarFunction:
+ async def eval(self, value):
+ return value
+ else:
+ def eval(self, value):
+ return value
+
+ class InheritedDeterminism(CustomDeterminism):
+ pass
+
+ for source in (CustomDeterminism, InheritedDeterminism):
+ with self.subTest(base=base, deterministic=deterministic,
source=source):
+ pf.udf(source, return_dtype=int,
deterministic=deterministic)
+ self.assertEqual(constructed, [])
+ declaration = pf.udf(
+ CustomDeterminism(), return_dtype=int,
deterministic=deterministic)
+ self.assertEqual(declaration._deterministic, deterministic)
+
def test_general_structured_results_are_normalized_recursively(self):
from pyflink.dataframe.udf import _create_result_normalizer
@@ -1296,6 +1321,634 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
class DataFrameUDFAdapterTests(unittest.TestCase):
+ def pandas_map_output(self, adapter, dtype, *inputs):
+ from pyflink.table.utils import create_pandas_table_row_schema,
pandas_table_row_to_arrow
+ import pytz
+
+ row_type = dtype._to_table_data_type()
+ schema = create_pandas_table_row_schema(TableDataTypes.ROW([
+ TableDataTypes.FIELD("result", row_type)]))
+ return pandas_table_row_to_arrow(
+ schema, pytz.UTC, row_type,
adapter.eval(*inputs)).column(0).to_pylist()
+
+ def test_pandas_map_batches_returns_flattened_dataframe(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+
+ dtype = pf.DataType.struct({
+ "v": pf.DataType.int64(),
+ "payload": pf.DataType.struct({"v": pf.DataType.int64(),
+ "tail": pf.DataType.int64()}),
+ "other": pf.DataType.struct({"v": pf.DataType.string()}),
+ "tail": pf.DataType.int64(),
+ })
+ input_type = TableDataTypes.ROW([
+ TableDataTypes.FIELD("tag", TableDataTypes.STRING()),
+ TableDataTypes.FIELD("id", TableDataTypes.BIGINT())])
+
+ def transform(batch):
+ self.assertEqual(batch["tag"].tolist(), ["x"] * len(batch["id"]))
+ return {"tail": batch["id"] + 10,
+ "other": pd.DataFrame({"v": ["a", None]},
index=batch["id"].index),
+ "payload": pd.DataFrame({
+ "v": pd.Series([9007199254740993, None],
+ index=batch["id"].index, dtype=object),
+ "tail": batch["id"] + 20}),
+ "v": batch["id"].sort_index(ascending=False)}
+
+ adapter = _resolve_map_udf(transform, dtype, "pandas",
input_type)._func
+ adapter.open(object())
+ try:
+ result = adapter.eval("x", pd.Series([1, 2]))
+ self.assertIsInstance(result, pd.DataFrame)
+ pd.testing.assert_frame_equal(result, pd.concat([
+ pd.Series([1, 2], name="v"),
+ pd.Series([9007199254740993, None], name="v", dtype=object),
+ pd.Series([21, 22], name="tail"),
+ pd.Series(["a", None], name="v"),
+ pd.Series([11, 12], name="tail"),
+ ], axis=1))
+ self.assertEqual(self.pandas_map_output(adapter, dtype, "x",
pd.Series([1, 2])), [
+ {"v": 1, "payload": {"v": 9007199254740993, "tail": 21},
+ "other": {"v": "a"}, "tail": 11},
+ {"v": 2, "payload": {"v": None, "tail": 22},
+ "other": {"v": None}, "tail": 12},
+ ])
+ finally:
+ adapter.close()
+
+ def test_map_bindings_leave_other_uses_of_a_declaration_unchanged(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+
+ @pf.udf(return_dtype=pf.DataType.struct({"total":
pf.DataType.int64()}))
+ def total(value):
+ return {"total": sum(value.values()) if isinstance(value, dict)
else value}
+
+ expression = total._table_udf_wrapper._func
+ maps = [_resolve_map_udf(total, None, "general", TableDataTypes.ROW([
+ TableDataTypes.FIELD(name, TableDataTypes.BIGINT()) for name in
names
+ ]))._func for names in (["left", "right"], ["value"])]
+ for adapter in [expression] + maps:
+ adapter.open(object())
+ self.addCleanup(adapter.close)
+ self.assertEqual(maps[0].eval(1, 2), Row(total=3))
+ self.assertEqual(maps[1].eval(4), Row(total=4))
+ self.assertEqual(maps[0].eval(5, 6), Row(total=11))
+ self.assertEqual(expression.eval(7), Row(total=7))
+
+ def test_map_requires_a_dictionary_result(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+
+ dtype = pf.DataType.struct({"value": pf.DataType.int64()})
+ for result in (None, (1,), 1):
+ with self.subTest(result=result):
+ adapter = _resolve_map_udf(
+ lambda row: result, dtype, "general",
dtype._to_table_data_type())._func
+ adapter.open(object())
+ try:
+ with self.assertRaisesRegex(TypeError, "dictionary"):
+ adapter.eval(1)
+ finally:
+ adapter.close()
+
+ def
test_pandas_map_batches_preserves_row_columns_for_transported_and_fused_inputs(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+ from pyflink.table.types import create_arrow_schema
+ from pyflink.table.utils import arrow_to_pandas
+ import pytz
+
+ input_type = TableDataTypes.ROW([
+ TableDataTypes.FIELD("id", TableDataTypes.BIGINT()),
+ TableDataTypes.FIELD("payload", TableDataTypes.ROW([
+ TableDataTypes.FIELD("left", TableDataTypes.BIGINT()),
+ TableDataTypes.FIELD("right", TableDataTypes.BIGINT())])),
+ TableDataTypes.FIELD("tail", TableDataTypes.BIGINT()),
+ ])
+ dtype = pf.DataType.struct({"score": pf.DataType.int64()})
+
+ def transform(batch):
+ self.assertEqual(list(batch), ["id", "payload", "tail"])
+ payload = batch["payload"]
+ self.assertIsInstance(payload, pd.DataFrame)
+ self.assertEqual(list(payload), ["left", "right"])
+ pd.testing.assert_frame_equal(payload, pd.DataFrame(
+ {"left": [2, None, 4], "right": [3, None, None]},
+ index=payload.index))
+ return {"score": batch["tail"] + batch["id"]}
+
+ adapter = _resolve_map_udf(transform, dtype, "pandas",
input_type)._func
+ adapter.open(object())
+ try:
+ grouped = pd.DataFrame({
+ "id": [1, 2, 3], "tail": [10, 20, 30], "payload": [
+ {"left": 2, "right": 3}, None,
+ {"left": 4, "right": None},
+ ],
+ })[["id", "payload", "tail"]]
+ schema = create_arrow_schema(input_type.field_names(),
input_type.field_types())
+ batch = pa.RecordBatch.from_pandas(grouped, schema=schema)
+ transported = arrow_to_pandas(pytz.UTC, input_type.field_types(),
[batch])
+ self.assertEqual(self.pandas_map_output(adapter, dtype,
*transported),
+ [{"score": 11}, {"score": 22}, {"score": 33}])
+ for names in (["unrelated1", "unrelated2"], ["right", "left"]):
+ fused_payload = transported[1].set_axis(names, axis=1)
+ self.assertEqual(self.pandas_map_output(
+ adapter, dtype, transported[0], fused_payload,
transported[2]),
+ [{"score": 11}, {"score": 22}, {"score": 33}])
+ self.assertEqual(list(fused_payload), names)
+ grouped.index = [2, 0, 1]
+ with self.assertRaisesRegex(TypeError, "ROW.*pandas.DataFrame"):
+ adapter.eval(*(grouped[name] for name in
input_type.field_names()))
+ fused = [column.set_axis(grouped.index, axis=0) for column in
transported]
+ self.assertEqual(self.pandas_map_output(adapter, dtype, *fused),
+ [{"score": 11}, {"score": 22}, {"score": 33}])
+ finally:
+ adapter.close()
+
+ def
test_pandas_map_batches_preserves_nested_collections_with_repeated_names(self):
+ import cloudpickle
+ from pyflink.dataframe.udf import _resolve_map_udf
+ from pyflink.fn_execution import flink_fn_execution_pb2 as proto
+ from pyflink.fn_execution.table.operations import
ScalarFunctionOperation
+ from pyflink.table.udf import DelegatingScalarFunction
+ from pyflink.table.utils import arrow_to_pandas
+ import pytz
+
+ row_type = TableDataTypes.ROW([TableDataTypes.FIELD(
+ "payload", TableDataTypes.MAP(TableDataTypes.STRING(),
TableDataTypes.BIGINT()))])
+ input_type = TableDataTypes.ROW([TableDataTypes.FIELD("payload",
row_type)])
+ dtype = pf.DataType.struct({"size": pf.DataType.int64()})
+
+ def transform(batch):
+ self.assertEqual(batch["payload"]["payload"].tolist(), [[("x", 1),
("y", 2)], []])
+ return {"size": batch["payload"]["payload"].map(len)}
+
+ batch = pa.record_batch([pa.array([
+ {"payload": [("x", 1), ("y", 2)]}, {"payload": []},
+ ], type=pa.struct([("payload", pa.map_(pa.string(), pa.int64()))]))],
names=["payload"])
+ columns = arrow_to_pandas(pytz.UTC, [row_type], [batch])
+ adapter = _resolve_map_udf(transform, dtype, "pandas",
input_type)._func
+ adapter.open(object())
+ try:
+ self.assertEqual(self.pandas_map_output(adapter, dtype,
columns[0]),
+ [{"size": 2}, {"size": 0}])
+ nested_dtype = pf.DataType.struct({"payload": pf.DataType.struct({
+ "payload": pf.DataType.map(pf.DataType.string(),
pf.DataType.int64())})})
+ producer = _resolve_map_udf(
+ lambda batch: {"payload": batch["payload"]},
+ nested_dtype, "pandas", input_type)._func
+ producer_proto = proto.UserDefinedFunction(
+ payload=cloudpickle.dumps(producer), is_pandas_udf=True,
+ inputs=[proto.Input(inputOffset=0)])
+ producer.open(object())
+ try:
+ self.assertEqual(self.pandas_map_output(producer,
nested_dtype, columns[0]),
+ [{"payload": {"payload": [("x", 1), ("y",
2)]}},
+ {"payload": {"payload": []}}])
+ finally:
+ producer.close()
+ for values, expected in ((columns, [2, 0]),
+ ([column.head(0) for column in columns],
[])):
+ for shared in (False, True):
+ with self.subTest(empty_batch=not expected, shared=shared):
+ consumer = proto.UserDefinedFunction(
+ payload=cloudpickle.dumps(DelegatingScalarFunction(
+ lambda frame: frame["payload"].map(len))),
+ is_pandas_udf=True, takes_row_as_input=True,
+ inputs=[proto.Input(refIndex=0)] if shared
+ else [proto.Input(udf=producer_proto)])
+ operation =
ScalarFunctionOperation(proto.UserDefinedFunctions(
+ udfs=[producer_proto, consumer] if shared else
[consumer],
+ output_indices=[1] if shared else []))
+ operation.open()
+ try:
+
self.assertEqual(operation.process_element(values)[0].tolist(),
+ expected)
+ finally:
+ operation.close()
+ finally:
+ adapter.close()
+
+ def test_pandas_map_batches_preserves_map_inputs(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+ from pyflink.table.types import create_arrow_schema
+ from pyflink.table.utils import arrow_to_pandas, pandas_to_arrow
+ import pytz
+
+ lookup = TableDataTypes.MAP(TableDataTypes.STRING(),
TableDataTypes.BIGINT())
+ payload = TableDataTypes.ROW([
+ TableDataTypes.FIELD("lookup", lookup),
+ ])
+ input_type = TableDataTypes.ROW([
+ TableDataTypes.FIELD("lookup", lookup),
TableDataTypes.FIELD("payload", payload)])
+ dtype = pf.DataType.struct({"total": pf.DataType.int64()})
+
+ def transform(batch):
+ self.assertIs(batch["lookup"], inputs[0])
+ self.assertIs(batch["payload"], inputs[1])
+ return {"total": batch["lookup"].map(
+ lambda value: 0 if value is None else sum(value.values())
+ if isinstance(value, dict) else sum(item for _, item in
value))}
+
+ columns = [pd.Series([{"top": 1}, {}, None], dtype=object),
pd.DataFrame({
+ "lookup": [{"field": 2}, [("same", 3), ("same", 4)], None],
+ })]
+ schema = create_arrow_schema(input_type.field_names(),
input_type.field_types())
+ transported = arrow_to_pandas(pytz.UTC, input_type.field_types(),
[pandas_to_arrow(
+ schema, pytz.UTC, input_type.field_types(), columns)])
+ adapter = _resolve_map_udf(transform, dtype, "pandas",
input_type)._func
+ adapter.open(object())
+ try:
+ for inputs in (columns, transported):
+ self.assertEqual(self.pandas_map_output(adapter, dtype,
*inputs),
+ [{"total": 1}, {"total": 0}, {"total": 0}])
+ finally:
+ adapter.close()
+
+ def test_arrow_map_batches_preserves_nested_and_chunked_columns(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+ from pyflink.fn_execution.utils.operation_utils import
check_arrow_udf_result
+
+ dtype = pf.DataType.struct({
+ "id": pf.DataType.int64(),
+ "payload": pf.DataType.struct({"x": pf.DataType.int64()}),
+ })
+ batch = pa.StructArray.from_arrays([
+ pa.array([1, 2]), pa.array([{"x": None}, None],
type=pa.struct([("x", pa.int64())])),
+ pa.array([{"inner": {"v": 1}}, {"inner": {"v": 2}}],
+ type=pa.struct([("inner", pa.struct([("v",
pa.int64())]))])),
+ pa.array([[{"v": 1}], []], type=pa.list_(pa.struct([("v",
pa.int64())])))
+ ], names=["id", "payload", "deep", "items"])
+
+ def transform(columns):
+ self.assertEqual(columns["tag"].to_pylist(), ["x"] *
len(columns["id"]))
+ self.assertIsInstance(columns["payload"], pa.StructArray)
+
self.assertEqual(columns["deep"].field("inner").field("v").to_pylist(),
+ [1, 2][:len(columns["id"])])
+ self.assertEqual(columns["items"].to_pylist(),
+ [[{"v": 1}], []][:len(columns["id"])])
+ return {"payload": columns["payload"], "ignored": pa.array([0, 0]),
+ "id": pa.chunked_array([columns["id"].slice(0, 1),
+ columns["id"].slice(1)])}
+
+ input_type = TableDataTypes.ROW([
+ TableDataTypes.FIELD("tag", TableDataTypes.STRING())
+ ] + dtype._to_table_data_type().fields + [
+ TableDataTypes.FIELD("deep", TableDataTypes.ROW([
+ TableDataTypes.FIELD("inner", TableDataTypes.ROW([
+ TableDataTypes.FIELD("v", TableDataTypes.BIGINT())]))])),
+ TableDataTypes.FIELD("items",
TableDataTypes.ARRAY(TableDataTypes.ROW([
+ TableDataTypes.FIELD("v", TableDataTypes.BIGINT())]))),
+ ])
+ adapter = _resolve_map_udf(transform, dtype, "arrow", input_type)._func
+ adapter.open(object())
+ try:
+ result = check_arrow_udf_result(adapter.eval, "x",
*batch.flatten())
+ self.assertEqual(result.to_pylist(),
+ [{"id": 1, "payload": {"x": None}}, {"id": 2,
"payload": None}])
+ self.assertEqual(adapter.eval("x", *batch.slice(0,
0).flatten()).to_pylist(), [])
+ finally:
+ adapter.close()
+
+ def test_pandas_map_batches_aligns_indexes_and_checks_cardinality(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+
+ dtype = pf.DataType.struct({"score": pf.DataType.int64(), "id":
pf.DataType.int64()})
+
+ def align(batch):
+ self.assertIsInstance(batch, dict)
+ self.assertIsInstance(batch["id"], pd.Series)
+ return {"score": pd.Series([100, 200],
index=batch["id"].index[::-1]).rename_axis(None),
+ "id": batch["id"],
+ "ignored": pd.Series([0, 0])}
+
+ def wrong_length(batch):
+ return {"id": batch["id"], "score": pd.Series([100, 200],
index=[1, 2])}
+
+ def too_short(batch):
+ return {"id": batch["id"].head(1), "score": batch["id"].head(1)}
+
+ def same_order(batch):
+ return {"score": pd.Series([200, 100],
index=batch["id"].index).rename_axis(None),
+ "id": batch["id"]}
+
+ def partial(batch):
+ return {"score": pd.Series([9007199254740993], index=[0]), "id":
batch["id"]}
+
+ input_type = TableDataTypes.ROW([TableDataTypes.FIELD("id",
TableDataTypes.BIGINT())])
+ for source, index in ((align, pd.RangeIndex(2)), (align, pd.Index([2,
0], name="rows")),
+ (align, pd.to_datetime(["2024-01-03",
"2024-01-01"])),
+ (same_order, pd.Index([2, 0], name="rows")),
+ (partial, pd.RangeIndex(2)),
+ (wrong_length, pd.RangeIndex(2)), (too_short,
pd.RangeIndex(2)),
+ (functools.partial(too_short),
pd.RangeIndex(2))):
+ adapter = _resolve_map_udf(source, dtype, "pandas",
input_type)._func
+ adapter.open(object())
+ try:
+ batch = pd.DataFrame({"id": [10, 20]}, index=index)
+ if source in (wrong_length, too_short, partial) \
+ or isinstance(source, functools.partial):
+ with self.assertRaisesRegex(ValueError, "length"):
+ adapter.eval(batch["id"])
+ else:
+ self.assertEqual(self.pandas_map_output(adapter, dtype,
batch["id"]),
+ [{"score": 200, "id": 10}, {"score": 100,
"id": 20}])
+ finally:
+ adapter.close()
+
+ def
test_pandas_map_batches_requires_row_fields_including_collections(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+ from pyflink.table.utils import create_pandas_table_row_schema,
pandas_table_row_to_arrow
+ import pytz
+
+ details = pf.DataType.struct({"name": pf.DataType.string()})
+ dtype = pf.DataType.struct({
+ "details": details,
+ "items": pf.DataType.list(
+ pf.DataType.map(pf.DataType.string(), details)),
+ "plain_items": pf.DataType.list(
+ pf.DataType.map(pf.DataType.string(), pf.DataType.int64())),
+ "lookup": pf.DataType.map(pf.DataType.string(), details),
+ })
+ input_type = TableDataTypes.ROW([TableDataTypes.FIELD("id",
TableDataTypes.BIGINT())])
+ valid = {"details": pd.DataFrame({"name": [None]}),
+ "items": [None, {"item key": {"name": "a"}}],
+ "plain_items": [None, {"plain key": 1}, {}],
+ "lookup": {"any key": {"name": None}}, "extra": 2}
+ cases = [(valid, None), ({"items": [], "lookup": {}}, KeyError),
+ (dict(valid, details=pd.DataFrame(index=[0])), ValueError),
+ (dict(valid, details=pd.DataFrame({"name": [None], "extra":
[1]})), ValueError),
+ (dict(valid, items=[{"item key": {}}]), KeyError),
+ (dict(valid, lookup={"any key": {}}), KeyError),
+ (dict(valid, lookup={"any key": types.SimpleNamespace()}),
AttributeError),
+ (dict(valid, lookup={"any key": ()}), ValueError),
+ (dict(valid, lookup={"any key": (None, "extra")}),
ValueError)]
+ for result, error in cases:
+ with self.subTest(result=result):
+ adapter = _resolve_map_udf(
+ lambda batch: {name: value if isinstance(value,
pd.DataFrame)
+ else pd.Series([value]) for name, value in
result.items()},
+ dtype, "pandas", input_type)._func
+ adapter.open(object())
+ try:
+ if result is valid:
+ output = adapter.eval(pd.Series([1]))
+ self.assertEqual(output["items"].iloc[0],
+ [None, [("item key", {"name": "a"})]])
+ self.assertEqual(output["plain_items"].iloc[0],
+ [None, [("plain key", 1)], []])
+ self.assertEqual(valid["items"], [None, {"item key":
{"name": "a"}}])
+ self.assertEqual(valid["plain_items"], [None, {"plain
key": 1}, {}])
+ row_type = dtype._to_table_data_type()
+ schema =
create_pandas_table_row_schema(TableDataTypes.ROW([
+ TableDataTypes.FIELD("result", row_type)]))
+ batch = pandas_table_row_to_arrow(schema, pytz.UTC,
row_type, output)
+ self.assertEqual(batch.to_pylist(), [{"result": {
+ "details": {"name": None},
+ "items": [None, [("item key", {"name": "a"})]],
+ "plain_items": [None, [("plain key", 1)], []],
+ "lookup": [("any key", {"name": None})]}}])
+ else:
+ with self.assertRaises(error):
+ adapter.eval(pd.Series([1]))
+ finally:
+ adapter.close()
+
+ def test_pandas_map_batches_preserves_nullable_structured_columns(self):
+ import tracemalloc
+ from pyflink.dataframe.udf import _create_pandas_result_normalizer,
_resolve_map_udf
+
+ row_type = TableDataTypes.ROW([TableDataTypes.FIELD("v",
TableDataTypes.BIGINT())])
+ for data_type, payload in (
+ (TableDataTypes.ARRAY(row_type), [{"v": i} for i in range(10000)]),
+ (TableDataTypes.MAP(TableDataTypes.STRING(), row_type),
+ [(str(i), {"v": i}) for i in range(10000)]),
+ ):
+ normalize = _create_pandas_result_normalizer(data_type)
+ normalize(payload)
+ tracemalloc.start()
+ try:
+ normalized = normalize(payload)
+ peak = tracemalloc.get_traced_memory()[1]
+ finally:
+ tracemalloc.stop()
+ self.assertIs(normalized, payload)
+ self.assertLess(peak, 16384)
+ original = Row(v=1)
+ if isinstance(data_type, MapType):
+ source = [("same", payload[0][1]), ("changed", original)]
+ converted = normalize(source)
+ self.assertIs(converted[0][1], source[0][1])
+ self.assertEqual(converted[1], ("changed", {"v": 1}))
+ self.assertIs(source[1][1], original)
+ else:
+ source = [payload[0], original]
+ converted = normalize(source)
+ self.assertIs(converted[0], source[0])
+ self.assertEqual(converted[1], {"v": 1})
+ self.assertIs(source[1], original)
+
+ dtype = pf.DataType.struct({
+ "details": pf.DataType.struct({"name": pf.DataType.string()}),
+ "items": pf.DataType.list(pf.DataType.int64()),
+ "lookup": pf.DataType.map(pf.DataType.string(),
pf.DataType.int64()),
+ })
+ input_type = TableDataTypes.ROW([TableDataTypes.FIELD("id",
TableDataTypes.BIGINT())])
+
+ class RowFrame(pd.DataFrame):
+ pass
+
+ for frame_type in (pd.DataFrame, RowFrame):
+ with self.subTest(frame_type=frame_type):
+ adapter = _resolve_map_udf(
+ lambda batch: {"details": frame_type({"name": [None,
float("nan"), pd.NA]}),
+ "items": pd.Series([None, float("nan"),
pd.NA]),
+ "lookup": pd.Series([None, float("nan"),
pd.NA])},
+ dtype, "pandas", input_type)._func
+ adapter.open(object())
+ try:
+ self.assertEqual(self.pandas_map_output(adapter, dtype,
pd.Series([1, 2, 3])),
+ [{"details": {"name": None}, "items":
None, "lookup": None}]
+ * 3)
+ finally:
+ adapter.close()
+
+ def test_pandas_map_batches_requires_matching_output_containers(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+
+ input_type = TableDataTypes.ROW([TableDataTypes.FIELD("id",
TableDataTypes.BIGINT())])
+ for field_type, columns, container in (
+ (pf.DataType.struct({"name": pf.DataType.string()}),
+ (pd.Series([{"name": "a"}]), None, [1]), "DataFrame"),
+ (pf.DataType.int64(), (pd.DataFrame({"inner": [1]}), None, [1]),
"Series"),
+ ):
+ dtype = pf.DataType.struct({"payload": field_type})
+ for column in columns:
+ with self.subTest(column=column, field_type=field_type):
+ adapter = _resolve_map_udf(
+ lambda batch: {"payload": column}, dtype, "pandas",
input_type)._func
+ adapter.open(object())
+ try:
+ with self.assertRaisesRegex(TypeError,
f"pandas.{container}"):
+ adapter.eval(pd.Series([1]))
+ finally:
+ adapter.close()
+
+ def test_pandas_map_batches_rejects_one_shot_structured_collections(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+
+ details = pf.DataType.struct({"name": pf.DataType.string()})
+ input_type = TableDataTypes.ROW([TableDataTypes.FIELD("id",
TableDataTypes.BIGINT())])
+ for field_type, value in (
+ (pf.DataType.list(pf.DataType.map(pf.DataType.string(), details)),
+ iter([{"key": {"name": "a"}}])),
+ (pf.DataType.map(pf.DataType.string(), details), iter([("key",
{"name": "a"})])),
+ ):
+ with self.subTest(field_type=field_type):
+ dtype = pf.DataType.struct({"value": field_type})
+ adapter = _resolve_map_udf(
+ lambda batch: {"value": pd.Series([value])}, dtype,
"pandas", input_type)._func
+ adapter.open(object())
+ try:
+ with self.assertRaisesRegex(TypeError, "iterator"):
+ adapter.eval(pd.Series([1]))
+ finally:
+ adapter.close()
+
+ def test_map_requires_output_fields_even_when_nullable(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+
+ class Output(TypedDict):
+ value: Optional[str]
+
+ def transform(row) -> Output:
+ return {"value": None} if row["id"] else {}
+
+ input_type = TableDataTypes.ROW([TableDataTypes.FIELD("id",
TableDataTypes.BIGINT())])
+ explicit = pf.DataType.struct({"value": pf.DataType.string()})
+ for source, dtype in (
+ (transform, None), (transform, Output),
+ (pf.udf(transform), None), (transform, explicit),
+ ):
+ with self.subTest(source=source, dtype=dtype):
+ adapter = _resolve_map_udf(source, dtype, "general",
input_type)._func
+ adapter.open(object())
+ try:
+ self.assertEqual(adapter.eval(1), Row(value=None))
+ with self.assertRaisesRegex(KeyError, "value"):
+ adapter.eval(0)
+ finally:
+ adapter.close()
+
+ def test_map_requires_nested_row_fields_including_collections(self):
+ from pyflink.dataframe.udf import _resolve_map_udf
+
+ details = pf.DataType.struct({"name": pf.DataType.string()})
+ dtype = pf.DataType.struct({
+ "details": details,
+ "items": pf.DataType.list(details),
+ "lookup": pf.DataType.map(pf.DataType.string(), details),
+ })
+ input_type = TableDataTypes.ROW([TableDataTypes.FIELD("id",
TableDataTypes.BIGINT())])
+ valid = {"details": {"name": None, "extra": 1}, "items": [None,
{"name": "a"}],
+ "lookup": {"any key": {"name": None}}, "extra": 2}
+ for result in (valid, dict(valid, details={}), dict(valid, items=[{}]),
+ dict(valid, lookup={"any key": {}}),
+ dict(valid, details=types.SimpleNamespace())):
+ def transform(row):
+ return result
+
+ async def transform_async(row):
+ return result
+
+ for source in (transform, transform_async):
+ with self.subTest(result=result, source=source):
+ adapter = _resolve_map_udf(source, dtype, "general",
input_type)._func
+ adapter.open(object())
+ try:
+ def invoke():
+ value = adapter.eval(1)
+ return asyncio.run(value) if
inspect.isawaitable(value) else value
+
+ if result is valid:
+ self.assertEqual(invoke(), Row(
+ details=Row(name=None), items=[None,
Row(name="a")],
+ lookup={"any key": Row(name=None)}))
+ else:
+ with self.assertRaisesRegex((KeyError,
AttributeError), "name"):
+ invoke()
+ finally:
+ adapter.close()
+
+ def test_row_maps_receive_named_fields_and_project_output_schema(self):
+ from pyflink.dataframe.udf import (
+ _DataFrameAsyncScalarFunctionAdapter,
+ _DataFrameScalarFunctionAdapter,
+ _UDFUsage,
+ _resolve_udf,
+ )
+
+ def transform(row):
+ self.assertEqual(row, {"value": 3, "label": "a"})
+ return {"label": row["label"].upper(), "extra": 99, "value":
row["value"] + 1}
+
+ async def transform_async(row):
+ return transform(row)
+
+ dtype = pf.DataType.struct({"value": pf.DataType.int64(), "label":
pf.DataType.string()})
+ for source, adapter_type in (
+ (transform, _DataFrameScalarFunctionAdapter),
+ (transform_async, _DataFrameAsyncScalarFunctionAdapter),
+ ):
+ with self.subTest(source=source):
+ adapter = adapter_type(
+ _resolve_udf(source).runtime_source, dtype, True,
_UDFUsage.MAP, "general",
+ input_type=dtype._to_table_data_type(),
+ )
+ adapter.open(object())
+ try:
+ result = adapter.eval(3, "a")
+ if inspect.isawaitable(result):
+ result = asyncio.run(result)
+ self.assertEqual(result, Row(value=4, label="A"))
+ finally:
+ adapter.close()
+
+ def test_row_maps_preserve_resolved_column_inputs(self):
+ import cloudpickle
+ from pyflink.dataframe.udf import _resolve_map_udf
+ from pyflink.fn_execution import flink_fn_execution_pb2 as proto
+ from pyflink.fn_execution.table.operations import
ScalarFunctionOperation
+ from pyflink.serializers import PickleSerializer
+ from pyflink.table.udf import DelegatingScalarFunction
+
+ literal = proto.Input(inputConstant=b"\x00" +
PickleSerializer().dumps("xyz"))
+ increment = proto.UserDefinedFunction(
+ payload=cloudpickle.dumps(DelegatingScalarFunction(lambda value:
value + 1)),
+ inputs=[proto.Input(inputOffset=0)])
+ for fields, inputs, values, expected in (
+ ({"tag": pf.DataType.string()}, [literal], [], Row(tag="xyz")),
+ ({"id": pf.DataType.int64(), "tag": pf.DataType.string()},
+ [proto.Input(inputOffset=0), literal], [1], Row(id=1, tag="xyz")),
+ ({"first": pf.DataType.int64(), "second": pf.DataType.int64()},
+ [proto.Input(udf=increment), proto.Input(udf=increment)], [1],
+ Row(first=2, second=2)),
+ ({"payload": pf.DataType.struct({"tag": pf.DataType.string()})},
+ [proto.Input(inputOffset=0)], [Row(tag="xyz")],
Row(payload=Row(tag="xyz"))),
+ ):
+ with self.subTest(fields=fields):
+ dtype = pf.DataType.struct(fields)
+ function = _resolve_map_udf(
+ lambda row: row, dtype, "general",
dtype._to_table_data_type())
+ operation =
ScalarFunctionOperation(proto.UserDefinedFunctions(udfs=[
+ proto.UserDefinedFunction(
+ payload=cloudpickle.dumps(function._func),
inputs=inputs,
+ takes_row_as_input=function._takes_row_as_input)]))
+ operation.open()
+ try:
+ self.assertEqual(operation.process_element(values),
[expected])
+ finally:
+ operation.close()
+
def test_general_result_normalizers_are_bound_by_return_type(self):
from pyflink.dataframe.udf import (
_DataFrameAsyncScalarFunctionAdapter,
@@ -1564,7 +2217,7 @@ class DataFrameUDFAdapterTests(unittest.TestCase):
"general",
)
for _ in range(2):
- with self.assertRaisesRegex(NotImplementedError, "'map'"):
+ with self.assertRaisesRegex(ValueError, "input schema"):
adapter.open(object())
adapter.close()
@@ -1575,14 +2228,127 @@ class DataFrameUDFAdapterTests(unittest.TestCase):
class DataFrameUDFPlannerTests(PyFlinkDataFrameUTTestCase):
+ def test_map_function_classes_require_explicit_custom_determinism(self):
+ class DefaultDeterminism(ScalarFunction):
+ def eval(self, value):
+ return {"value": value["value"]}
+
+ class CustomDeterminism(DefaultDeterminism):
+ def is_deterministic(self):
+ return False
+
+ class InheritedDeterminism(CustomDeterminism):
+ pass
+
+ class DeferredDeterminism(CustomDeterminism):
+ def __init__(self):
+ raise AssertionError("Function classes must be constructed on
the worker.")
+
+ dataframe = pf.from_records([(1,)], schema=["value"])
+ dtype = pf.DataType.struct({"value": pf.DataType.int64()})
+ for mode in ("general", "pandas", "arrow"):
+ def bind(source, **kwargs):
+ return (dataframe.map(source, **kwargs) if mode == "general"
else
+ dataframe.map_batches(source, batch_format=mode,
**kwargs))
+
+ for source in (CustomDeterminism, InheritedDeterminism,
DeferredDeterminism):
+ with self.subTest(mode=mode, source=source):
+ with self.assertRaisesRegex(TypeError,
"deterministic=.*instance"):
+ bind(source, return_dtype=dtype)
+ declared = pf.udf(source, return_dtype=dtype,
func_type=mode,
+ deterministic=False)
+ self.assert_dataframe_schema(bind(declared), ["value"])
+ for source in (DefaultDeterminism, CustomDeterminism()):
+ self.assert_dataframe_schema(bind(source, return_dtype=dtype),
["value"])
+
+ class CustomAsyncDeterminism(AsyncScalarFunction):
+ async def eval(self, value):
+ return {"value": value["value"]}
+
+ def is_deterministic(self):
+ return False
+
+ class InheritedAsyncDeterminism(CustomAsyncDeterminism):
+ pass
+
+ for source in (CustomAsyncDeterminism, InheritedAsyncDeterminism):
+ with self.assertRaisesRegex(TypeError, "deterministic=.*instance"):
+ dataframe.map(source, return_dtype=dtype)
+ declared = pf.udf(source, return_dtype=dtype, deterministic=False)
+ self.assert_dataframe_schema(dataframe.map(declared), ["value"])
+ self.assert_dataframe_schema(dataframe.map(source(),
return_dtype=dtype), ["value"])
+
+ def test_map_batches_requires_matching_format_and_struct_schema(self):
+ dataframe = pf.from_records([(1,)], schema=["id"])
+ dtype = pf.DataType.struct({"value": pf.DataType.int64()})
+
+ def transform(batch):
+ return {"value": batch["id"]}
+
+ for source in (transform, pf.udf(transform, return_dtype=dtype,
func_type="pandas")):
+ result = dataframe.map_batches(
+ source, batch_format="pandas",
+ **({"return_dtype": dtype} if source is transform else {}),
+ )
+ self.assertEqual(result.columns, ["value"])
+ for call, error, message in (
+ (lambda: dataframe.map_batches(transform), TypeError,
"batch_format"),
+ (lambda: dataframe.map_batches(transform, batch_format="numpy"),
ValueError, "format"),
+ (lambda: dataframe.map_batches(transform, batch_format="pandas"),
TypeError,
+ "return_dtype"),
+ (lambda: dataframe.map_batches(transform, batch_format="pandas",
return_dtype=int),
+ TypeError, "struct"),
+ (lambda: dataframe.map_batches(
+ pf.udf(transform, return_dtype=dtype), batch_format="pandas"),
ValueError, "mode"),
+ ):
+ with self.subTest(message=message):
+ with self.assertRaisesRegex(error, message):
+ call()
+
+ for field_type in (
+ pf.DataType.struct({"inner": pf.DataType.struct({"v":
pf.DataType.int64()})}),
+ pf.DataType.struct({"ts": pf.DataType.timestamp_ltz(3)}),
+ pf.DataType.list(pf.DataType.struct({"v": pf.DataType.int64()})),
+ pf.DataType.list(pf.DataType.timestamp_ltz(3)),
+ ):
+ with self.subTest(field_type=field_type):
+ output_type = pf.DataType.struct({"value": field_type})
+ with self.assertRaises((TypeError, ValueError)):
+ dataframe.map_batches(transform, batch_format="pandas",
+ return_dtype=output_type)
+ nested = dataframe.with_columns(value=pf.udf(
+ lambda value: value,
return_dtype=field_type)(pf.col("id"))).select("value")
+ with self.assertRaises((TypeError, ValueError)):
+ nested.map_batches(transform, batch_format="pandas",
return_dtype=dtype)
+ self.assertEqual(dataframe.map_batches(
+ transform, batch_format="arrow",
return_dtype=output_type).columns, ["value"])
+
+ def test_map_infers_schema_and_rejects_conflicting_declarations(self):
+ class Output(TypedDict):
+ label: str
+ value: int
+
+ def transform(row) -> Output:
+ return {"value": row["id"] + 1, "label": "mapped"}
+
+ dataframe = pf.from_records([(1,)], schema=["id"])
+ declaration = pf.udf(transform)
+ for source in (transform, declaration):
+ result = dataframe.map(source)
+ self.assertEqual(result.columns, ["label", "value"])
+ with self.assertRaisesRegex(ValueError, "return_dtype"):
+ dataframe.map(declaration, return_dtype=Output)
+ with self.assertRaisesRegex(TypeError, "struct"):
+ dataframe.map(lambda row: row["id"], return_dtype=int)
+ with self.assertRaisesRegex(ValueError, "mode"):
+ dataframe.map(pf.udf(transform, return_dtype=Output,
func_type="pandas"))
+
def test_arrow_calls_require_a_column_argument(self):
from pyflink.table import ExplainDetail
- @pf.udf(return_dtype=pf.DataType.int64(), func_type="arrow")
- def identity(*values):
- return values[0]
-
dataframe = pf.from_records([(1,)], schema=["id"])
+ identity = pf.udf(lambda *values: values[0],
return_dtype=pf.DataType.int64(),
+ func_type="arrow")
for args in ((), (1,), (pf.lit(1),), (identity(),)):
with self.subTest(args=args):
with self.assertRaisesRegex(Exception, "at least one
column-valued argument"):
@@ -1630,6 +2396,229 @@ class
DataFrameUDFPlannerTests(PyFlinkDataFrameUTTestCase):
class DataFrameUDFITCase(PyFlinkStreamDataFrameTestCase):
+ def test_row_maps_preserve_literal_and_fused_inputs(self):
+ self.env.set_parallelism(1)
+
+ def identity(row):
+ return row
+
+ async def async_identity(row):
+ return row
+
+ @pf.udf
+ def increment(value: int) -> int:
+ return value + 1
+
+ source = pf.from_records([(1,), (2,)], schema=["id"])
+ tagged = source.with_columns(tag=pf.lit("xyz"))
+ for dataframe, fields, expected in (
+ (tagged.select("tag"), {"tag": pf.DataType.string()},
+ [Row("xyz"), Row("xyz")]),
+ (tagged, {"id": pf.DataType.int64(), "tag": pf.DataType.string()},
+ [Row(1, "xyz"), Row(2, "xyz")]),
+ (source.select(increment(pf.col("id")).alias("first"),
+ increment(pf.col("id")).alias("second")),
+ {"first": pf.DataType.int64(), "second": pf.DataType.int64()},
+ [Row(2, 2), Row(3, 3)]),
+ ):
+ dtype = pf.DataType.struct(fields)
+ for function in (identity, async_identity):
+ with self.subTest(fields=fields, function=function):
+ result = dataframe.map(function, return_dtype=dtype).map(
+ identity, return_dtype=dtype)
+ self.assertEqual(result.collect(), expected)
+
+ def test_batch_maps_broadcast_literal_columns(self):
+ self.env.set_parallelism(1)
+ self.t_env.get_config().set("python.fn-execution.bundle.size", "3")
+ self.t_env.get_config().set("python.fn-execution.arrow.batch.size",
"2")
+ dtype = pf.DataType.struct({
+ "tag": pf.DataType.string(), "empty": pf.DataType.int64(),
+ "enabled": pf.DataType.bool(), "id": pf.DataType.int64()})
+ dataframe = pf.from_records([(1,), (2,), (3,)],
schema=["id"]).with_columns(
+ tag=pf.lit("x"), empty=pf.lit(None, pf.DataType.int64()),
+ enabled=pf.lit(True)).select("tag", "empty", "enabled", "id")
+ for batch_format in ("pandas", "arrow"):
+ with self.subTest(batch_format=batch_format):
+ result = dataframe.map_batches(
+ lambda batch: batch, batch_format=batch_format,
return_dtype=dtype)
+ self.assertEqual(result.collect(),
+ [Row("x", None, True, value) for value in (1,
2, 3)])
+
+ def test_pandas_output_alignment_preserves_input_order_and_ltz(self):
+ import datetime
+
+ self.env.set_parallelism(1)
+ self.t_env.get_config().set("python.fn-execution.bundle.size", "3")
+ self.t_env.get_config().set("python.fn-execution.arrow.batch.size",
"2")
+ timestamp = datetime.datetime(2020, 1, 1, 12)
+
+ def transform(batch):
+ return {"score": (batch["id"] * 10).sort_index(ascending=False),
+ "id": batch["id"], "ts": batch["ts"]}
+
+ timezone = self.t_env.get_config().get_local_timezone()
+ try:
+ for zone, epoch in [("UTC", 1577880000000), ("Asia/Shanghai",
1577851200000)]:
+ with self.subTest(timezone=zone):
+ self.t_env.get_config().set_local_timezone(zone)
+ from pyflink.table.expressions import call_sql, col
+
+ source = self.t_env.from_elements(
+ [(10, epoch), (20, epoch), (30, None)], ["id",
"epoch"])
+ dataframe = pf.from_table(source.select(
+ col("id"), call_sql("TO_TIMESTAMP_LTZ(epoch,
3)").alias("ts")))
+ result = dataframe.map_batches(
+ transform, batch_format="pandas",
return_dtype=pf.DataType.struct({
+ "score": pf.DataType.int64(), "id":
pf.DataType.int64(),
+ "ts": pf.DataType.timestamp_ltz(3)}))
+ result = result.select(pf.col("score"), pf.col("id"),
+
pf.col("ts").cast(pf.DataType.timestamp(3)).alias("ts"))
+ self.assertEqual(result.collect(), [
+ Row(100, 10, timestamp), Row(200, 20, timestamp),
Row(300, 30, None)])
+
+ finally:
+ self.t_env.get_config().set_local_timezone(timezone)
+
+ def
test_pandas_map_preserves_row_columns_across_execution_boundaries(self):
+ from pyflink.table.udf import udf as table_udf
+
+ self.env.set_parallelism(1)
+ self.t_env.get_config().set("python.fn-execution.bundle.size", "3")
+ self.t_env.get_config().set("python.fn-execution.arrow.batch.size",
"2")
+ payload_type = pf.DataType.struct({
+ "lookup": pf.DataType.map(pf.DataType.string().not_null(),
pf.DataType.int64()),
+ "v": pf.DataType.int64(),
+ })
+ input_type = pf.DataType.struct({"id": pf.DataType.int64(), "payload":
payload_type})
+ value_type = pf.DataType.struct({"v": pf.DataType.int64()})
+ output_type = pf.DataType.struct({"total": pf.DataType.int64(),
"value": value_type})
+ flat_type = pf.DataType.struct({"total": pf.DataType.int64(), "v":
pf.DataType.int64()})
+
+ def produce(values):
+ rows = [
+ {"lookup": [("long-key", int(value))], "v": 11} if value > 0
else
+ {"lookup": [], "v": None} if value == 0 else
+ {"lookup": None, "v": None}
+ for value in values
+ ]
+ return pd.DataFrame({
+ label: pd.Series([row[field] for row in rows],
index=values.index[::-1],
+ dtype=object)
+ for label, field in zip(("unrelated1", "unrelated2"),
("lookup", "v"))
+ })
+
+ def consume(batch):
+ assert isinstance(batch["payload"], pd.DataFrame)
+ return {"total": batch["payload"]["lookup"].map(
+ lambda pairs: -1 if pairs is None else sum(value for _, value
in pairs)
+ ) + batch["id"] * 10, "value":
batch["payload"][["v"]].sort_index(ascending=False)}
+
+ def extract(batch):
+ return {"total": batch["total"], "v": batch["value"].field("v")}
+
+ @table_udf(result_type=flat_type._to_table_data_type(),
func_type="pandas")
+ def table_extract(batch):
+ assert isinstance(batch["total"], pd.Series)
+ return pd.DataFrame({name: batch[name] for name in ("total", "v")})
+
+ producer = pf.udf(produce, return_dtype=payload_type,
func_type="pandas")
+ for transported in (False, True):
+ with self.subTest(transported=transported):
+ dataframe = pf.from_records([(1,), (0,), (-1,)],
schema=["id"]).with_columns(
+ payload=producer(pf.col("id")))
+ if transported:
+ dataframe = dataframe.map_batches(
+ lambda batch: batch, return_dtype=input_type,
batch_format="arrow")
+ mapped = dataframe.map_batches(
+ consume, return_dtype=output_type,
batch_format="pandas").map_batches(
+ lambda batch: batch, return_dtype=output_type,
batch_format="pandas")
+ for table_map in (False, True):
+ result =
(pf.from_table(mapped.to_table().map(table_extract)) if table_map
+ else mapped.map_batches(extract,
batch_format="arrow",
+ return_dtype=flat_type))
+ self.assertEqual(sorted(result.collect(), key=lambda row:
row[0]),
+ [Row(-11, None), Row(0, None), Row(11,
11)])
+
+ def test_row_and_batch_maps_in_one_job(self):
+ self.env.set_parallelism(1)
+ self.t_env.get_config().set("python.fn-execution.bundle.size", "3")
+ self.t_env.get_config().set("python.fn-execution.arrow.batch.size",
"2")
+ dtype = pf.DataType.struct({"id": pf.DataType.int64(), "value":
pf.DataType.int64()})
+
+ class Prepare(ScalarFunction):
+ def eval(self, row):
+ assert isinstance(row, dict)
+ return {"value": row["id"] * 2, "id": row["id"], "ignored": 0}
+
+ def is_deterministic(self):
+ return False
+
+ prepare = pf.udf(Prepare, return_dtype=dtype, deterministic=False)
+
+ async def enrich(row):
+ assert isinstance(row, dict)
+ return {"id": row["id"], "value": row["value"] + 1}
+
+ def pandas_add(batch):
+ assert isinstance(batch, dict)
+ assert all(isinstance(column, pd.Series) for column in
batch.values())
+ assert 1 <= len(batch["id"]) <= 2
+ return {"id": batch["id"], "value": batch["value"] + 10}
+
+ def arrow_double(batch):
+ import pyarrow.compute as pc
+
+ assert isinstance(batch, dict)
+ assert all(isinstance(column, pa.Array) for column in
batch.values())
+ return {"id": pa.chunked_array([batch["id"].slice(0, 1),
batch["id"].slice(1)]),
+ "value": pc.multiply(batch["value"], 2)}
+
+ flat_dtype = pf.DataType.struct({
+ "id": pf.DataType.int64(), "value": pf.DataType.int64(),
+ "inner": pf.DataType.int64(),
+ })
+ nested_dtype = pf.DataType.struct({
+ "id": pf.DataType.int64(),
+ "payload": pf.DataType.struct({
+ "value": pf.DataType.int64(),
+ "inner": pf.DataType.struct({"value": pf.DataType.int64()}),
+ }),
+ })
+
+ def pandas_prepare(batch):
+ return {"id": batch["id"], "value": batch["value"],
+ "inner": pd.Series([None if identifier == 3 else 3
+ for identifier in batch["id"]],
dtype=object)}
+
+ def arrow_nest(batch):
+ return {"id": batch["id"], "payload": pa.array([
+ None if identifier == 5 else {"value": value, "inner":
+ None if identifier == 4 else {"value": inner}}
+ for identifier, value, inner in zip(
+ batch["id"].to_pylist(), batch["value"].to_pylist(),
batch["inner"].to_pylist())
+ ], type=pa.struct([("value", pa.int64()),
+ ("inner", pa.struct([("value", pa.int64())]))]))}
+
+ def arrow_unnest(batch):
+ assert isinstance(batch["payload"], pa.StructArray)
+ return {"id": batch["id"], "value": pa.array([
+ -1 if payload is None else -2 if payload["inner"] is None else
-3
+ if payload["inner"]["value"] is None
+ else payload["value"] + payload["inner"]["value"]
+ for payload in batch["payload"].to_pylist()
+ ], type=pa.int64())}
+
+ result = (pf.from_records([(1,), (2,), (3,), (4,), (5,)],
schema=["id"])
+ .map(prepare).map(enrich, return_dtype=dtype)
+ .map_batches(pandas_add, return_dtype=dtype,
batch_format="pandas")
+ .map_batches(arrow_double, return_dtype=dtype,
batch_format="arrow")
+ .map_batches(pandas_prepare, return_dtype=flat_dtype,
batch_format="pandas")
+ .map_batches(arrow_nest, return_dtype=nested_dtype,
batch_format="arrow")
+ .map_batches(arrow_unnest, return_dtype=dtype,
batch_format="arrow"))
+ self.assertEqual(sorted(result.collect(), key=lambda row: row[0]),
+ [Row(1, 29), Row(2, 33), Row(3, -3), Row(4, -2),
Row(5, -1)])
+
def test_supported_scalar_udfs_in_one_job(self):
import pyarrow.compute as pc
diff --git a/flink-python/pyflink/dataframe/tests/test_udtf.py
b/flink-python/pyflink/dataframe/tests/test_udtf.py
index 4ea27607be8..3ce842c54f0 100644
--- a/flink-python/pyflink/dataframe/tests/test_udtf.py
+++ b/flink-python/pyflink/dataframe/tests/test_udtf.py
@@ -58,6 +58,18 @@ def _eval_udtf(declaration, *args):
class DataFrameUDTFDeclarationTests(unittest.TestCase):
+ def test_missing_row_fields_remain_nullable_in_expression_results(self):
+ details = pf.DataType.struct({"label": pf.DataType.string()})
+ dtype = pf.DataType.struct({
+ "missing": pf.DataType.string(),
+ "details": details,
+ "items": pf.DataType.list(details),
+ })
+ declaration = pf.udtf(lambda value: [{"details": {}, "items": [{}]}],
+ return_dtype=dtype)
+ self.assertEqual(_eval_udtf(declaration, 1), [
+ Row(missing=None, details=Row(label=None),
items=[Row(label=None)])])
+
def test_decorator_forms_and_type_inference(self):
def emit(value) -> Iterator[_Output]:
yield {"value": value, "label": "ok"}
diff --git a/flink-python/pyflink/dataframe/udf.py
b/flink-python/pyflink/dataframe/udf.py
index e2f3955f435..df0b92da8d3 100644
--- a/flink-python/pyflink/dataframe/udf.py
+++ b/flink-python/pyflink/dataframe/udf.py
@@ -20,7 +20,7 @@
import functools
import inspect
-from collections.abc import Mapping
+from collections.abc import Iterator, Mapping
from dataclasses import dataclass, field
from enum import Enum
from typing import (
@@ -50,6 +50,8 @@ from pyflink.table.udf import (
ScalarFunction,
UserDefinedFunction,
UserDefinedFunctionWrapper,
+ UserDefinedAsyncScalarFunctionWrapper,
+ UserDefinedScalarFunctionWrapper,
udf as table_udf,
)
from pyflink.util.api_stability_decorators import PublicEvolving
@@ -245,7 +247,7 @@ class _DataFrameUDFWrapper:
return cast(UserDefinedFunctionWrapper, self._cached_table_udf_wrapper)
def _create_table_udf_wrapper(
- self, usage: _UDFUsage
+ self, usage: _UDFUsage, input_type: Optional[RowType] = None
) -> UserDefinedFunctionWrapper:
adapter_type = (
_DataFrameAsyncScalarFunctionAdapter
@@ -260,9 +262,10 @@ class _DataFrameUDFWrapper:
self._deterministic,
usage,
self._func_type,
+ input_type,
),
)
- return cast(
+ wrapper = cast(
UserDefinedFunctionWrapper,
table_udf(
actual_func,
@@ -272,12 +275,60 @@ class _DataFrameUDFWrapper:
func_type=self._func_type,
),
)
+ wrapper._returns_table_row = usage in (_UDFUsage.MAP,
_UDFUsage.MAP_BATCHES)
+ return wrapper
@property
def return_dtype(self) -> DataType:
return self._return_dtype
+def _resolve_map_udf(
+ func: Union[_UDFInput, _DataFrameUDFWrapper],
+ return_dtype: Optional[_DataTypeLike],
+ func_type: str,
+ input_type: RowType,
+) -> Union[UserDefinedScalarFunctionWrapper,
UserDefinedAsyncScalarFunctionWrapper]:
+ if isinstance(func, _DataFrameUDFWrapper):
+ if return_dtype is not None:
+ raise ValueError("return_dtype must not be specified for an
existing DataFrame UDF.")
+ declaration = func
+ else:
+ if (inspect.isclass(func) and issubclass(func, (ScalarFunction,
AsyncScalarFunction))
+ and inspect.getattr_static(func, "is_deterministic")
+ is not inspect.getattr_static(UserDefinedFunction,
"is_deterministic")):
+ raise TypeError(
+ f"Cannot infer determinism from class {func.__name__!r}. "
+ "Declare it with pf.udf(..., deterministic=...), or pass an
instance."
+ )
+ deterministic = (
+ func.is_deterministic()
+ if isinstance(func, (ScalarFunction, AsyncScalarFunction))
+ else True
+ )
+ declaration = cast(_DataFrameUDFWrapper, udf(
+ func, return_dtype=return_dtype, func_type=func_type,
deterministic=deterministic
+ ))
+ if declaration._func_type != func_type:
+ raise ValueError(
+ f"The declared UDF mode {declaration._func_type!r} does not match "
+ f"the required mode {func_type!r}."
+ )
+ output_type = declaration.return_dtype._to_table_data_type()
+ if not isinstance(output_type, RowType):
+ raise TypeError("return_dtype for map and map_batches must be a struct
type.")
+ usage = _UDFUsage.MAP if func_type == "general" else _UDFUsage.MAP_BATCHES
+ if usage is _UDFUsage.MAP_BATCHES:
+ if not input_type.fields:
+ raise ValueError("map_batches requires at least one input column.")
+ if func_type == "pandas":
+ _validate_pandas_map_schema(input_type, output_type)
+ return cast(
+ Union[UserDefinedScalarFunctionWrapper,
UserDefinedAsyncScalarFunctionWrapper],
+ declaration._create_table_udf_wrapper(usage, input_type),
+ )
+
+
@overload
def udf(
func: _UDFInput,
@@ -528,6 +579,16 @@ def _validate_scalar_udf_options(
)
+def _validate_pandas_map_schema(input_type: RowType, output_type: RowType) ->
None:
+ from pyflink.table.types import create_arrow_schema
+
+ if not output_type.fields:
+ raise ValueError("Pandas map_batches requires at least one output
column.")
+ # Validate physical columns even when fusion bypasses Arrow transport.
+ for row_type in (input_type, output_type):
+ create_arrow_schema(row_type.field_names(), row_type.field_types())
+
+
# ======================== Callable Inspection and Resolution
========================
# ---- Invocation target resolution ----
@@ -1080,10 +1141,13 @@ class _DataFrameUDFAdapterBase:
deterministic: bool,
usage: _UDFUsage,
func_type: str,
+ input_type: Optional[RowType] = None,
) -> None:
self._runtime_source = runtime_source
self._worker_udf: Optional[_WorkerUDF] = None
- self._return_dtype = return_dtype if func_type == "general" else None
+ self._return_dtype = return_dtype
+ # MAP and MAP_BATCHES need the input DataFrame schema to build their
input dictionaries.
+ self._input_type = input_type
self._deterministic = deterministic
self._usage = usage
self._func_type = func_type
@@ -1108,13 +1172,34 @@ class _DataFrameUDFAdapterBase:
raise
def _bind_func(self, invoke_func: Callable[..., Any]) -> Callable[...,
Any]:
- if self._usage is not _UDFUsage.EXPRESSION:
- raise NotImplementedError(
- f"DataFrame UDF usage {self._usage.value!r} is not supported
yet."
- )
+ if self._usage in (_UDFUsage.MAP, _UDFUsage.MAP_BATCHES):
+ if self._input_type is None:
+ raise ValueError("A map binding requires an input schema.")
+ output_type = cast(RowType,
self._return_dtype._to_table_data_type())
+ if self._usage is _UDFUsage.MAP:
+ input_names = self._input_type.field_names()
+ normalize_result = cast(Callable[[Any], Any],
_create_result_normalizer(
+ output_type, require_row_fields=True))
+
+ def normalize_map_result(value):
+ if not isinstance(value, Mapping):
+ raise TypeError("A map function must return a
dictionary.")
+ return normalize_result(value)
+
+ normalized_func = _wrap_scalar_general_result(
+ invoke_func, normalize_map_result,
self._runtime_source.is_async,
+ )
+
+ @functools.wraps(invoke_func)
+ def map_row(*values):
+ return normalized_func(dict(zip(input_names, values)))
+
+ return map_row
+ wrap = _wrap_map_pandas if self._func_type == "pandas" else
_wrap_map_arrow
+ return wrap(invoke_func, self._input_type, output_type)
if self._func_type == "general":
result_normalizer = _create_result_normalizer(
- cast(DataType, self._return_dtype)._to_table_data_type()
+ self._return_dtype._to_table_data_type()
)
if result_normalizer is None:
return invoke_func
@@ -1148,7 +1233,7 @@ class
_DataFrameScalarFunctionAdapter(_DataFrameUDFAdapterBase, ScalarFunction):
def eval(self, *args: Any) -> Any:
invoke_func = self._invocation()
- if self._func_type == "pandas":
+ if self._func_type == "pandas" and self._usage is not
_UDFUsage.MAP_BATCHES:
from pyflink.fn_execution.utils.operation_utils import (
check_pandas_udf_result,
)
@@ -1169,8 +1254,12 @@ class _DataFrameAsyncScalarFunctionAdapter(
# ======================== Result Normalization ========================
-def _row_field_values(value: Any, field_names: List[str]) -> List[Any]:
+def _row_field_values(
+ value: Any, field_names: List[str], *, require_row_fields: bool = False,
+) -> List[Any]:
if isinstance(value, Mapping):
+ if require_row_fields:
+ return [value[field_name] for field_name in field_names]
return [value.get(field_name) for field_name in field_names]
if isinstance(value, Row) and hasattr(value, "_fields"):
field_indices: Dict[str, int] = {}
@@ -1200,13 +1289,13 @@ def _row_field_values(value: Any, field_names:
List[str]) -> List[Any]:
)
return list(value)
return [
- _object_row_field_value(value, field_name, field_names)
+ _object_row_field_value(value, field_name, field_names,
require_row_fields)
for field_name in field_names
]
def _object_row_field_value(
- value: Any, field_name: str, field_names: List[str]
+ value: Any, field_name: str, field_names: List[str], require_row_fields:
bool
) -> Any:
try:
return getattr(value, field_name)
@@ -1217,6 +1306,11 @@ def _object_row_field_value(
attributes = getattr(value, "__dict__", None)
has_slots = any("__slots__" in cls.__dict__ for cls in
type(value).__mro__)
if isinstance(attributes, Mapping) or has_slots:
+ if require_row_fields:
+ raise AttributeError(
+ f"Field name {field_name!r} does not exist in "
+ f"{type(value).__name__}."
+ ) from None
return None
else:
raise
@@ -1226,19 +1320,255 @@ def _object_row_field_value(
) from None
+def _is_pandas_null(value):
+ import pandas as pd
+
+ return bool(pd.isna(value)) if pd.api.types.is_scalar(value) else False
+
+
+def _wrap_map_pandas(func, input_type: RowType, output_type: RowType):
+ import pandas as pd
+
+ input_converters = [
+ _create_pandas_row_column_converter(field.data_type, [None] *
len(field.data_type))
+ if isinstance(field.data_type, RowType) else None for field in
input_type]
+ output_converters =
[_create_pandas_output_column_converter(field.data_type)
+ for field in output_type]
+
+ @functools.wraps(func)
+ def wrapper(*input_columns):
+ batch_size = next((len(column) for column in input_columns
+ if isinstance(column, (pd.Series, pd.DataFrame))),
None)
+ if batch_size is None:
+ raise ValueError("Pandas map_batches requires a column-valued
input.")
+ input_index = pd.RangeIndex(batch_size)
+ columns = {}
+ for input_field, column, convert in zip(input_type, input_columns,
input_converters):
+ # The planner can fold an input column into a scalar literal
argument.
+ if not isinstance(column, (pd.Series, pd.DataFrame)):
+ if isinstance(input_field.data_type, RowType):
+ names = input_field.data_type.field_names()
+ values = ([None] * len(names) if column is None
+ else _row_field_values(column, names))
+ column = pd.DataFrame([values] * batch_size, columns=names)
+ else:
+ column = pd.Series([column] * batch_size,
index=input_index,
+ name=input_field.name)
+ if convert is not None:
+ column = convert(column)
+ index = pd.RangeIndex(len(column))
+ if not column.index.identical(index):
+ # Arrow transports values positionally and discards upstream
pandas indexes.
+ column = column.copy(deep=False)
+ column.index = index
+ columns[input_field.name] = column
+ result = func(columns)
+ output = []
+ for output_field, convert in zip(output_type, output_converters):
+ column = result[output_field.name]
+ container = pd.DataFrame if isinstance(output_field.data_type,
RowType) else pd.Series
+ if not isinstance(column, container):
+ raise TypeError(f"Pandas map_batches output column
{output_field.name!r} "
+ f"must be a pandas.{container.__name__}.")
+ if len(column) != len(input_index):
+ raise ValueError(
+ f"The result length '{len(column)}' of Pandas UDF
'{wrapper.__name__}' "
+ f"column {output_field.name!r} is not equal to the input
length "
+ f"'{len(input_index)}'")
+ if container is pd.DataFrame and type(column) is not pd.DataFrame:
+ # The shared Arrow converter expects the base DataFrame
produced by concat.
+ column = pd.DataFrame(column, copy=False)
+ if convert is not None:
+ column = convert(column)
+ if isinstance(column, pd.Series) and column.name !=
output_field.name:
+ column = column.copy(deep=False)
+ column.name = output_field.name
+ output.append(column)
+ output = pd.concat(output, axis=1, copy=False)
+ index = output.index
+ if index.equals(input_index):
+ index = input_index
+ elif len(index) == len(input_index):
+ # Restore input order without discarding extra labels or padding
short batches.
+ extra_labels = index[~index.isin(input_index)]
+ index = input_index if extra_labels.empty else
input_index.append(extra_labels)
+ if len(index) != len(input_index):
+ raise ValueError(
+ f"The result length '{len(index)}' of Pandas UDF
'{wrapper.__name__}' "
+ f"is not equal to the input length '{len(input_index)}'")
+ if not output.index.identical(index):
+ if output.index.equals(index):
+ output.index = index
+ else:
+ output = output.reindex(index)
+ return output
+
+ return wrapper
+
+
+def _normalize_pandas_series(column, normalize):
+ import pandas as pd
+
+ converted = None
+ for index, value in enumerate(column):
+ normalized = normalize(value)
+ if normalized is not value:
+ if converted is None:
+ converted = list(column)
+ converted[index] = normalized
+ return (column if converted is None
+ else pd.Series(converted, name=column.name, index=column.index,
dtype=object))
+
+
+def _create_pandas_row_column_converter(data_type, converters):
+ import pandas as pd
+
+ names = data_type.field_names()
+
+ def convert_row(column):
+ if not isinstance(column, pd.DataFrame):
+ raise TypeError("Pandas ROW columns must be pandas.DataFrame. "
+ "Return a pandas.DataFrame from the upstream
pandas UDF.")
+ if list(column.columns) != names:
+ # ROW fields are positional in Arrow transport, regardless of
producer labels.
+ column = column.copy(deep=False)
+ column.columns = names
+ converted_columns = None
+ for index, convert in enumerate(converters):
+ if convert is not None:
+ original = column.iloc[:, index]
+ normalized = convert(original)
+ if normalized is not original:
+ if converted_columns is None:
+ converted_columns = [column.iloc[:, i] for i in
range(len(converters))]
+ converted_columns[index] = normalized
+ return (column if converted_columns is None
+ else pd.concat(converted_columns, axis=1, copy=False))
+
+ return convert_row
+
+
+def _create_pandas_output_column_converter(data_type):
+ if isinstance(data_type, RowType):
+ return _create_pandas_row_column_converter(
+ data_type, [_create_pandas_output_column_converter(field.data_type)
+ for field in data_type])
+ normalize = _create_pandas_result_normalizer(data_type)
+ return (None if normalize is None
+ else lambda column: _normalize_pandas_series(column, normalize))
+
+
+def _create_pandas_result_normalizer(data_type):
+ """Validate map ROW fields and adapt only values that Arrow cannot consume
directly."""
+ if isinstance(data_type, RowType):
+ names = data_type.field_names()
+ normalizers = [_create_pandas_result_normalizer(field.data_type) for
field in data_type]
+
+ def normalize_row(value):
+ if _is_pandas_null(value):
+ return None
+ values = _row_field_values(value, names, require_row_fields=True)
+ normalized = [item if normalize is None else normalize(item)
+ for item, normalize in zip(values, normalizers)]
+ if isinstance(value, dict) and all(item is original
+ for item, original in
zip(normalized, values)):
+ return value
+ return dict(zip(names, normalized))
+
+ return normalize_row
+ if isinstance(data_type, ArrayType):
+ normalize_element =
_create_pandas_result_normalizer(data_type.element_type)
+ if normalize_element is not None:
+ def normalize_array(value):
+ if isinstance(value, Iterator):
+ raise TypeError("Pandas ARRAY values containing ROW must
not be an iterator; "
+ "return a list instead.")
+ if _is_pandas_null(value):
+ return None
+ converted = None
+ for index, item in enumerate(value):
+ normalized = normalize_element(item)
+ if normalized is not item:
+ if converted is None:
+ converted = list(value)
+ converted[index] = normalized
+ return value if converted is None else converted
+
+ return normalize_array
+ if isinstance(data_type, MapType):
+ normalize_key = _create_pandas_result_normalizer(data_type.key_type)
+ normalize_value =
_create_pandas_result_normalizer(data_type.value_type)
+
+ def normalize_map(value):
+ if isinstance(value, Iterator):
+ raise TypeError("Pandas MAP values must not be an iterator; "
+ "return a dictionary or list of pairs
instead.")
+ if _is_pandas_null(value):
+ return None
+ # Older Arrow versions require pair sequences for MAPs nested in
collections.
+ converted = list(value.items()) if isinstance(value, Mapping) else
None
+ if normalize_key is None and normalize_value is None:
+ return value if converted is None else converted
+ items = value if converted is None else converted
+ for index, (key, item) in enumerate(items):
+ normalized_key = key if normalize_key is None else
normalize_key(key)
+ normalized_item = item if normalize_value is None else
normalize_value(item)
+ if normalized_key is not key or normalized_item is not item:
+ if converted is None:
+ converted = list(value)
+ converted[index] = (normalized_key, normalized_item)
+ return value if converted is None else converted
+
+ return normalize_map
+ return None
+
+
+def _wrap_map_arrow(func, input_type: RowType, output_type: RowType):
+ import pyarrow as pa
+ from pyflink.table.types import to_arrow_type
+
+ input_names = input_type.field_names()
+ input_types = [to_arrow_type(field.data_type, allow_nested=True) for field
in input_type]
+ output_names = output_type.field_names()
+
+ @functools.wraps(func)
+ def wrapper(*input_columns):
+ batch_size = next((len(column) for column in input_columns
+ if isinstance(column, (pa.Array,
pa.ChunkedArray))), None)
+ if batch_size is None:
+ raise ValueError("Arrow map_batches requires a column-valued
input.")
+ inputs = {
+ name: column if isinstance(column, (pa.Array, pa.ChunkedArray))
+ else pa.array([column] * batch_size, type=data_type)
+ for name, data_type, column in zip(input_names, input_types,
input_columns)
+ }
+ result = func(inputs)
+ columns = []
+ for name in output_names:
+ column = result[name]
+ if isinstance(column, pa.ChunkedArray):
+ column = column.chunk(0) if column.num_chunks == 1 else
column.combine_chunks()
+ columns.append(column)
+ return pa.StructArray.from_arrays(columns, names=output_names)
+
+ return wrapper
+
+
def _create_result_normalizer(
- data_type: Any,
+ data_type: Any, *, require_row_fields: 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) for field in data_type
+ _create_result_normalizer(field.data_type,
require_row_fields=require_row_fields)
+ for field in data_type
)
def normalize_row(value: Any) -> Any:
if value is None:
return None
- field_values = _row_field_values(value, field_names)
+ field_values = _row_field_values(
+ value, field_names, require_row_fields=require_row_fields)
normalized_fields = [
field_value
if field_normalizer is None
@@ -1255,7 +1585,8 @@ def _create_result_normalizer(
return normalize_row
if isinstance(data_type, ArrayType):
- element_normalizer = _create_result_normalizer(data_type.element_type)
+ element_normalizer = _create_result_normalizer(
+ data_type.element_type, require_row_fields=require_row_fields)
if element_normalizer is None:
def normalize_leaf_array(value: Any) -> Any:
@@ -1270,8 +1601,10 @@ def _create_result_normalizer(
return normalize_array
if isinstance(data_type, MapType):
- key_normalizer = _create_result_normalizer(data_type.key_type)
- value_normalizer = _create_result_normalizer(data_type.value_type)
+ key_normalizer = _create_result_normalizer(
+ data_type.key_type, require_row_fields=require_row_fields)
+ value_normalizer = _create_result_normalizer(
+ data_type.value_type, require_row_fields=require_row_fields)
def normalize_map(value: Any) -> Any:
if value is None:
diff --git a/flink-python/pyflink/fn_execution/coder_impl_fast.pxd
b/flink-python/pyflink/fn_execution/coder_impl_fast.pxd
index aeacced5dc5..ff018ecd21f 100644
--- a/flink-python/pyflink/fn_execution/coder_impl_fast.pxd
+++ b/flink-python/pyflink/fn_execution/coder_impl_fast.pxd
@@ -91,6 +91,7 @@ cdef class RowCoderImpl(FieldCoderImpl):
cdef class ArrowCoderImpl(FieldCoderImpl):
cdef object _batch_format
+ cdef bint _returns_table_row
cdef object _schema
cdef list _field_types
cdef object _timezone
diff --git a/flink-python/pyflink/fn_execution/coder_impl_fast.pyx
b/flink-python/pyflink/fn_execution/coder_impl_fast.pyx
index c25ac8cfdd7..ff2cff868ad 100644
--- a/flink-python/pyflink/fn_execution/coder_impl_fast.pyx
+++ b/flink-python/pyflink/fn_execution/coder_impl_fast.pyx
@@ -36,7 +36,8 @@ from pyflink.datastream.window import CountWindow,
TimeWindow, GlobalWindow
from pyflink.fn_execution.formats.avro import FlinkAvroDecoder,
FlinkAvroDatumReader, \
FlinkAvroBufferWrapper, FlinkAvroEncoder, FlinkAvroDatumWriter
from pyflink.fn_execution.ResettableIO import ResettableIO
-from pyflink.table.utils import pandas_to_arrow, arrow_to_pandas,
validate_arrow_batch
+from pyflink.table.utils import pandas_to_arrow, arrow_to_pandas,
validate_arrow_batch, \
+ pandas_table_row_to_arrow
ROW_KIND_BIT_SIZE = 2
@@ -431,8 +432,10 @@ cdef class ArrowCoderImpl(FieldCoderImpl):
A coder for arrow format data.
"""
- def __init__(self, schema, row_type, timezone, batch_format="PANDAS"):
+ def __init__(self, schema, row_type, timezone, batch_format="PANDAS",
+ returns_table_row=False):
self._batch_format = batch_format
+ self._returns_table_row = returns_table_row
self._schema = schema
self._field_types = row_type.field_types()
self._timezone = timezone
@@ -446,6 +449,9 @@ cdef class ArrowCoderImpl(FieldCoderImpl):
batch_writer = pa.RecordBatchStreamWriter(self._resettable_io,
self._schema)
if self._batch_format == "ARROW":
batch = validate_arrow_batch(cols, self._schema, self._field_types)
+ elif self._returns_table_row:
+ batch = pandas_table_row_to_arrow(
+ self._schema, self._timezone, self._field_types[0], cols[0])
else:
batch = pandas_to_arrow(self._schema, self._timezone,
self._field_types, cols)
batch_writer.write_batch(batch)
diff --git a/flink-python/pyflink/fn_execution/coder_impl_slow.py
b/flink-python/pyflink/fn_execution/coder_impl_slow.py
index a7df600d8fb..bf66f6e4ad3 100644
--- a/flink-python/pyflink/fn_execution/coder_impl_slow.py
+++ b/flink-python/pyflink/fn_execution/coder_impl_slow.py
@@ -31,7 +31,8 @@ from pyflink.fn_execution.ResettableIO import ResettableIO
from pyflink.fn_execution.formats.avro import FlinkAvroDecoder,
FlinkAvroDatumReader, \
FlinkAvroBufferWrapper, FlinkAvroEncoder, FlinkAvroDatumWriter
from pyflink.fn_execution.stream_slow import InputStream, OutputStream
-from pyflink.table.utils import pandas_to_arrow, arrow_to_pandas,
validate_arrow_batch
+from pyflink.table.utils import pandas_to_arrow, arrow_to_pandas,
validate_arrow_batch, \
+ pandas_table_row_to_arrow
ROW_KIND_BIT_SIZE = 2
@@ -278,8 +279,10 @@ class ArrowCoderImpl(FieldCoderImpl):
A coder for arrow format data.
"""
- def __init__(self, schema, row_type, timezone, batch_format="PANDAS"):
+ def __init__(self, schema, row_type, timezone, batch_format="PANDAS",
+ returns_table_row=False):
self._batch_format = batch_format
+ self._returns_table_row = returns_table_row
self._schema = schema
self._field_types = row_type.field_types()
self._timezone = timezone
@@ -293,6 +296,9 @@ class ArrowCoderImpl(FieldCoderImpl):
batch_writer = pa.RecordBatchStreamWriter(self._resettable_io,
self._schema)
if self._batch_format == "ARROW":
batch = validate_arrow_batch(cols, self._schema, self._field_types)
+ elif self._returns_table_row:
+ batch = pandas_table_row_to_arrow(
+ self._schema, self._timezone, self._field_types[0], cols[0])
else:
batch = pandas_to_arrow(self._schema, self._timezone,
self._field_types, cols)
batch_writer.write_batch(batch)
diff --git a/flink-python/pyflink/fn_execution/coders.py
b/flink-python/pyflink/fn_execution/coders.py
index d75afcc1ff9..7003cfbb484 100644
--- a/flink-python/pyflink/fn_execution/coders.py
+++ b/flink-python/pyflink/fn_execution/coders.py
@@ -92,9 +92,15 @@ class LengthPrefixBaseCoder(ABC):
# are not recognized by pytz, so resolve the timezone only for
pandas batches.
timezone = (None if batch_format == "ARROW"
else
pytz.timezone(os.environ['TABLE_LOCAL_TIME_ZONE']))
- schema = create_arrow_schema(row_type.field_names(),
row_type.field_types(),
- allow_nested=batch_format == "ARROW")
- return ArrowCoder(schema, row_type, timezone, batch_format)
+ returns_table_row =
coder_info_descriptor_proto.arrow_type.returns_table_row
+ if batch_format == "PANDAS" and returns_table_row:
+ from pyflink.table.utils import create_pandas_table_row_schema
+
+ schema = create_pandas_table_row_schema(row_type)
+ else:
+ schema = create_arrow_schema(row_type.field_names(),
row_type.field_types(),
+ allow_nested=batch_format ==
"ARROW")
+ return ArrowCoder(schema, row_type, timezone, batch_format,
returns_table_row)
elif coder_info_descriptor_proto.HasField('over_window_arrow_type'):
timezone = pytz.timezone(os.environ['TABLE_LOCAL_TIME_ZONE'])
schema_proto =
coder_info_descriptor_proto.over_window_arrow_type.schema
@@ -243,15 +249,18 @@ class ArrowCoder(FieldCoder):
Coder for Arrow.
"""
- def __init__(self, schema, row_type, timezone, batch_format="PANDAS"):
+ def __init__(self, schema, row_type, timezone, batch_format="PANDAS",
+ returns_table_row=False):
self._batch_format = batch_format
+ self._returns_table_row = returns_table_row
self._schema = schema
self._row_type = row_type
self._timezone = timezone
def get_impl(self):
return coder_impl.ArrowCoderImpl(
- self._schema, self._row_type, self._timezone, self._batch_format)
+ self._schema, self._row_type, self._timezone, self._batch_format,
+ self._returns_table_row)
def __repr__(self):
return 'ArrowCoder[%s]' % self._schema
diff --git a/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.py
b/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.py
index 54016508148..70183d79f23 100644
--- a/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.py
+++ b/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.py
@@ -41,7 +41,7 @@ _sym_db = _symbol_database.Default()
-DESCRIPTOR =
_descriptor_pool.Default().AddSerializedFile(b'\n\x18\x66link-fn-execution.proto\x12
org.apache.flink.fn_execution.v1\"*\n\x0cJobParameter\x12\x0b\n\x03key\x18\x01
\x01(\t\x12\r\n\x05value\x18\x02
\x01(\t\"\x9a\x01\n\x05Input\x12\x44\n\x03udf\x18\x01
\x01(\x0b\x32\x35.org.apache.flink.fn_execution.v1.UserDefinedFunctionH\x00\x12\x15\n\x0binputOffset\x18\x02
\x01(\x05H\x00\x12\x17\n\rinputConstant\x18\x03
\x01(\x0cH\x00\x12\x12\n\x08refIndex\x18\x04 \x01(\x05H\x00\x42\x07\n\x [...]
+DESCRIPTOR =
_descriptor_pool.Default().AddSerializedFile(b'\n\x18\x66link-fn-execution.proto\x12
org.apache.flink.fn_execution.v1\"*\n\x0cJobParameter\x12\x0b\n\x03key\x18\x01
\x01(\t\x12\r\n\x05value\x18\x02
\x01(\t\"\x9a\x01\n\x05Input\x12\x44\n\x03udf\x18\x01
\x01(\x0b\x32\x35.org.apache.flink.fn_execution.v1.UserDefinedFunctionH\x00\x12\x15\n\x0binputOffset\x18\x02
\x01(\x05H\x00\x12\x17\n\rinputConstant\x18\x03
\x01(\x0cH\x00\x12\x12\n\x08refIndex\x18\x04 \x01(\x05H\x00\x42\x07\n\x [...]
_globals = globals()
_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
@@ -150,19 +150,19 @@ if not _descriptor._USE_C_DESCRIPTORS:
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_TTLTIMECHARACTERISTIC']._serialized_start=9709
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_TTLTIMECHARACTERISTIC']._serialized_end=9752
_globals['_CODERINFODESCRIPTOR']._serialized_start=9755
- _globals['_CODERINFODESCRIPTOR']._serialized_end=10902
+ _globals['_CODERINFODESCRIPTOR']._serialized_end=10929
_globals['_CODERINFODESCRIPTOR_FLATTENROWTYPE']._serialized_start=10348
_globals['_CODERINFODESCRIPTOR_FLATTENROWTYPE']._serialized_end=10422
_globals['_CODERINFODESCRIPTOR_ROWTYPE']._serialized_start=10424
_globals['_CODERINFODESCRIPTOR_ROWTYPE']._serialized_end=10491
_globals['_CODERINFODESCRIPTOR_ARROWTYPE']._serialized_start=10494
- _globals['_CODERINFODESCRIPTOR_ARROWTYPE']._serialized_end=10700
-
_globals['_CODERINFODESCRIPTOR_ARROWTYPE_BATCHFORMAT']._serialized_start=10664
- _globals['_CODERINFODESCRIPTOR_ARROWTYPE_BATCHFORMAT']._serialized_end=10700
- _globals['_CODERINFODESCRIPTOR_OVERWINDOWARROWTYPE']._serialized_start=10702
- _globals['_CODERINFODESCRIPTOR_OVERWINDOWARROWTYPE']._serialized_end=10781
- _globals['_CODERINFODESCRIPTOR_RAWTYPE']._serialized_start=10783
- _globals['_CODERINFODESCRIPTOR_RAWTYPE']._serialized_end=10855
- _globals['_CODERINFODESCRIPTOR_MODE']._serialized_start=10857
- _globals['_CODERINFODESCRIPTOR_MODE']._serialized_end=10889
+ _globals['_CODERINFODESCRIPTOR_ARROWTYPE']._serialized_end=10727
+
_globals['_CODERINFODESCRIPTOR_ARROWTYPE_BATCHFORMAT']._serialized_start=10691
+ _globals['_CODERINFODESCRIPTOR_ARROWTYPE_BATCHFORMAT']._serialized_end=10727
+ _globals['_CODERINFODESCRIPTOR_OVERWINDOWARROWTYPE']._serialized_start=10729
+ _globals['_CODERINFODESCRIPTOR_OVERWINDOWARROWTYPE']._serialized_end=10808
+ _globals['_CODERINFODESCRIPTOR_RAWTYPE']._serialized_start=10810
+ _globals['_CODERINFODESCRIPTOR_RAWTYPE']._serialized_end=10882
+ _globals['_CODERINFODESCRIPTOR_MODE']._serialized_start=10884
+ _globals['_CODERINFODESCRIPTOR_MODE']._serialized_end=10916
# @@protoc_insertion_point(module_scope)
diff --git a/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.pyi
b/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.pyi
index 520077f395b..201882de23d 100644
--- a/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.pyi
+++ b/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.pyi
@@ -647,7 +647,7 @@ class CoderInfoDescriptor(_message.Message):
schema: Schema
def __init__(self, schema: _Optional[_Union[Schema, _Mapping]] = ...)
-> None: ...
class ArrowType(_message.Message):
- __slots__ = ("schema", "batch_format")
+ __slots__ = ("schema", "batch_format", "returns_table_row")
class BatchFormat(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
__slots__ = ()
PANDAS: _ClassVar[CoderInfoDescriptor.ArrowType.BatchFormat]
@@ -656,9 +656,11 @@ class CoderInfoDescriptor(_message.Message):
ARROW: CoderInfoDescriptor.ArrowType.BatchFormat
SCHEMA_FIELD_NUMBER: _ClassVar[int]
BATCH_FORMAT_FIELD_NUMBER: _ClassVar[int]
+ RETURNS_TABLE_ROW_FIELD_NUMBER: _ClassVar[int]
schema: Schema
batch_format: CoderInfoDescriptor.ArrowType.BatchFormat
- def __init__(self, schema: _Optional[_Union[Schema, _Mapping]] = ...,
batch_format: _Optional[_Union[CoderInfoDescriptor.ArrowType.BatchFormat, str]]
= ...) -> None: ...
+ returns_table_row: bool
+ def __init__(self, schema: _Optional[_Union[Schema, _Mapping]] = ...,
batch_format: _Optional[_Union[CoderInfoDescriptor.ArrowType.BatchFormat, str]]
= ..., returns_table_row: bool = ...) -> None: ...
class OverWindowArrowType(_message.Message):
__slots__ = ("schema",)
SCHEMA_FIELD_NUMBER: _ClassVar[int]
diff --git a/flink-python/pyflink/fn_execution/tests/test_coders.py
b/flink-python/pyflink/fn_execution/tests/test_coders.py
index 8cdb473b623..3bbc8cf859f 100644
--- a/flink-python/pyflink/fn_execution/tests/test_coders.py
+++ b/flink-python/pyflink/fn_execution/tests/test_coders.py
@@ -36,25 +36,49 @@ from pyflink.testing.test_case_utils import PyFlinkTestCase
class ArrowSchemaTests(unittest.TestCase):
- def test_nested_types_require_arrow_mode(self):
+ def test_pandas_positional_rows_require_matching_arity(self):
+ import pandas as pd
+ from pyflink.common import Row
+ 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("a", DataTypes.BIGINT()),
+ DataTypes.FIELD("b", DataTypes.BIGINT()),
+ ])
+ for data_type, wrap in (
+ (row_type, lambda value: value),
+ (DataTypes.MAP(DataTypes.STRING(), row_type), lambda value:
{"key": value}),
+ ):
+ schema = create_arrow_schema(["value"], [data_type])
+ for value in ((1,), (1, 2, 3), [1], [1, 2, 3], Row(1), Row(1, 2,
3)):
+ with self.subTest(data_type=data_type, value=value):
+ with self.assertRaises(RuntimeError):
+ pandas_to_arrow(schema, pytz.UTC, [data_type],
+ [pd.Series([wrap(value)],
dtype=object)])
+
+ schema = create_arrow_schema(["value"], [row_type])
+ self.assertEqual(pandas_to_arrow(
+ schema, pytz.UTC, [row_type], [pd.Series([{"a": 1}], dtype=object)]
+ ).to_pylist(), [{"value": {"a": 1, "b": None}}])
+
+ def test_pandas_row_ltz_and_arrow_array_ltz_schemas(self):
from pyflink.table import DataTypes
from pyflink.table.types import create_arrow_schema
- for data_type, error, expected in (
- (DataTypes.ARRAY(DataTypes.ROW([DataTypes.FIELD("v",
DataTypes.INT())])), ValueError,
- pa.list_(pa.field("item", pa.struct([pa.field("v",
pa.int32())])))),
- (DataTypes.ARRAY(DataTypes.TIMESTAMP_LTZ(3)), ValueError,
pa.list_(pa.timestamp('ms'))),
- (DataTypes.ROW([DataTypes.FIELD("r", DataTypes.ROW([
- DataTypes.FIELD("v", DataTypes.INT())]))]), TypeError,
- pa.struct([pa.field("r", pa.struct([pa.field("v",
pa.int32())]))])),
- (DataTypes.ROW([DataTypes.FIELD("t",
DataTypes.TIMESTAMP_LTZ(3))]), TypeError,
- pa.struct([pa.field("t", pa.timestamp('ms'))])),
+ row = DataTypes.ROW([DataTypes.FIELD("v", DataTypes.BIGINT())])
+ for data_type, error in (
+ (DataTypes.ROW([DataTypes.FIELD("t",
DataTypes.TIMESTAMP_LTZ(3))]), TypeError),
+ (DataTypes.ROW([DataTypes.FIELD("inner", row)]), TypeError),
+ (DataTypes.ARRAY(row), ValueError),
+ (DataTypes.ARRAY(DataTypes.TIMESTAMP_LTZ(3)), ValueError),
):
with self.subTest(data_type=data_type):
with self.assertRaises(error):
create_arrow_schema(["value"], [data_type])
- schema = create_arrow_schema(["value"], [data_type],
allow_nested=True)
- self.assertEqual(schema.field("value").type, expected)
+ self.assertEqual(len(create_arrow_schema(["value"],
[data_type],
+ allow_nested=True)),
1)
def test_pandas_collection_schema_and_round_trip(self):
import pandas as pd
@@ -253,6 +277,71 @@ class ArrowCodersTests(unittest.TestCase):
result = coder.decode(coder.encode([values]))
self.assertEqual(result[0].tolist(), [[("value", 10)], [("value",
20)], None, []])
+ def test_pandas_table_row_results_round_trip(self):
+ import datetime
+ import pandas as pd
+ from pyflink.table import DataTypes
+ from pyflink.table.types import create_arrow_schema
+
+ result_type = DataTypes.ROW([
+ DataTypes.FIELD("ts", DataTypes.TIMESTAMP_LTZ(3)),
+ DataTypes.FIELD("payload", DataTypes.ROW([
+ DataTypes.FIELD("v", DataTypes.BIGINT()),
+ DataTypes.FIELD("other", DataTypes.BIGINT())])),
+ DataTypes.FIELD("other", DataTypes.ROW([DataTypes.FIELD("v",
DataTypes.STRING())])),
+ ])
+ timezone = pytz.timezone("Asia/Shanghai")
+ columns = [
+ pd.Series(pd.to_datetime(["2020-01-01T00:00:00Z",
None]).tz_convert(timezone)),
+ pd.DataFrame({"v": pd.Series([9007199254740993, None],
dtype=object),
+ "other": pd.Series([7, None], dtype=object)}),
+ pd.DataFrame({"v": ["other", None]}),
+ ]
+ cases = [
+ (columns, [
+ {"result": {"ts": datetime.datetime(2020, 1, 1, 8),
+ "payload": {"v": 9007199254740993, "other": 7},
+ "other": {"v": "other"}}},
+ {"result": {"ts": None, "payload": {"v": None, "other": None},
+ "other": {"v": None}}},
+ ]),
+ ([column.head(0) for column in columns], []),
+ ]
+ for result, expected in cases:
+ with self.subTest(empty_batch=not expected):
+ row_type = DataTypes.ROW([DataTypes.FIELD("result",
result_type)])
+ schema = create_arrow_schema(row_type.field_names(),
row_type.field_types(),
+ allow_nested=True)
+ coder = self.implementation.ArrowCoderImpl(
+ schema, row_type, timezone, "PANDAS", True)
+ restored = self.arrow_coder(schema, row_type).decode(
+ coder.encode([pd.concat(result, axis=1)])).to_pylist()
+ self.assertEqual(restored, expected)
+
+ def test_pandas_row_collections_round_trip(self):
+ import numpy as np
+ import pandas as pd
+ from pyflink.table import DataTypes
+ from pyflink.table.types import create_arrow_schema
+
+ mapping = DataTypes.MAP(DataTypes.STRING(), DataTypes.BIGINT())
+ row_type = DataTypes.ROW([
+ DataTypes.FIELD("payload",
DataTypes.ROW([DataTypes.FIELD("payload", mapping)])),
+ DataTypes.FIELD("legacy", DataTypes.ARRAY(DataTypes.BIGINT()))])
+ schema = create_arrow_schema(row_type.field_names(),
row_type.field_types())
+ coder = self.implementation.ArrowCoderImpl(schema, row_type, pytz.UTC)
+ columns = [pd.Series([{"payload": [("x", 1), ("y", 2)]}, {"payload":
[]}]),
+ pd.Series([[1, None], []])]
+ encoded = coder.encode(columns)
+ restored = coder.decode(encoded)
+ self.assertEqual(restored[0]["payload"].tolist(), [[("x", 1), ("y",
2)], []])
+ self.assertEqual(restored[1].iloc[0].dtype, np.dtype("float64"))
+ self.assertEqual(restored[1].iloc[0][0], 1.0)
+ self.assertTrue(np.isnan(restored[1].iloc[0][1]))
+ arrow_coder = self.arrow_coder(schema, row_type)
+
self.assertEqual(arrow_coder.decode(coder.encode(restored)).to_pylist(),
+ arrow_coder.decode(encoded).to_pylist())
+
def test_inferred_struct_nullability(self):
from pyflink.table import DataTypes
from pyflink.table.types import create_arrow_schema
diff --git a/flink-python/pyflink/proto/flink-fn-execution.proto
b/flink-python/pyflink/proto/flink-fn-execution.proto
index a9300ba4949..b2a3bcc2e62 100644
--- a/flink-python/pyflink/proto/flink-fn-execution.proto
+++ b/flink-python/pyflink/proto/flink-fn-execution.proto
@@ -534,6 +534,8 @@ message CoderInfoDescriptor {
ARROW = 1;
}
BatchFormat batch_format = 2;
+ // The sole output field is a result ROW containing physical table columns.
+ bool returns_table_row = 3;
}
// only used in batch over window
diff --git a/flink-python/pyflink/table/table.py
b/flink-python/pyflink/table/table.py
index 820600efd63..6195d560817 100644
--- a/flink-python/pyflink/table/table.py
+++ b/flink-python/pyflink/table/table.py
@@ -34,7 +34,8 @@ from pyflink.table.table_schema import TableSchema
from pyflink.table.catalog import ResolvedSchema
from pyflink.table.types import create_arrow_schema
from pyflink.table.udf import UserDefinedScalarFunctionWrapper, \
- UserDefinedAggregateFunctionWrapper, UserDefinedTableFunctionWrapper
+ UserDefinedAggregateFunctionWrapper, UserDefinedTableFunctionWrapper, \
+ UserDefinedAsyncScalarFunctionWrapper
from pyflink.table.utils import tz_convert_from_internal, to_expression_jarray
from pyflink.table.window import OverWindow, GroupWindow
@@ -770,7 +771,11 @@ class Table(object):
return Table(self._j_table.dropColumns(to_expression_jarray(fields)),
self._t_env)
- def map(self, func: Union[Expression, UserDefinedScalarFunctionWrapper])
-> 'Table':
+ def map(
+ self,
+ func: Union[Expression, UserDefinedScalarFunctionWrapper,
+ UserDefinedAsyncScalarFunctionWrapper],
+ ) -> 'Table':
"""
Performs a map operation with a user-defined scalar function.
diff --git a/flink-python/pyflink/table/tests/test_pandas_udf.py
b/flink-python/pyflink/table/tests/test_pandas_udf.py
index 84d25d4307d..bdc1d565f64 100644
--- a/flink-python/pyflink/table/tests/test_pandas_udf.py
+++ b/flink-python/pyflink/table/tests/test_pandas_udf.py
@@ -22,7 +22,7 @@ import pytz
from pyflink.common import Row
from pyflink.table import DataTypes
from pyflink.table.tests.test_udf import SubtractOne
-from pyflink.table.udf import udf
+from pyflink.table.udf import ScalarFunction, udf
from pyflink.testing import source_sink_utils
from pyflink.testing.test_case_utils import PyFlinkBatchTableTestCase, \
PyFlinkStreamTableTestCase, PyFlinkTestCase
@@ -399,6 +399,42 @@ class PandasUDFITTests(object):
[list(r) for r in
self.t_env.sql_query(sql).execute().collect()])
self.assertEqual(actual, expected, sql)
+ import pandas as pd
+
+ class Seven(ScalarFunction):
+ def eval(self):
+ return pd.Series([7])
+
+ zero_argument = udf(Seven(), result_type=DataTypes.BIGINT(),
+ func_type="pandas")
+ batch_size =
self.t_env.get_config().get("python.fn-execution.arrow.batch.size", "1000")
+ try:
+
self.t_env.get_config().set("python.fn-execution.arrow.batch.size", "1")
+ result = t.select(plus_one(t.a),
zero_argument()).execute().collect()
+ self.assertEqual(sorted(result, key=lambda row: row[0]),
+ [Row(2, 7), Row(3, 7), Row(4, 7)])
+ finally:
+
self.t_env.get_config().set("python.fn-execution.arrow.batch.size", batch_size)
+
+ import numpy as np
+
+ square_root = udf(np.sqrt, result_type=DataTypes.DOUBLE(),
func_type="pandas")
+ nullable = self.t_env.from_elements(
+ [(1,), (None,), (4,)], DataTypes.ROW([DataTypes.FIELD("v",
DataTypes.BIGINT())]))
+
self.assertEqual(list(nullable.select(square_root(nullable.v)).execute().collect()),
+ [Row(1.0), Row(None), Row(2.0)])
+
+ result_type = DataTypes.MAP(DataTypes.STRING().not_null(),
DataTypes.ROW([
+ DataTypes.FIELD("v", DataTypes.BIGINT())]))
+ for as_pairs in (False, True):
+ def produce_rows(values):
+ return pd.Series([{"key": [("v", int(value))] if as_pairs
+ else {"v": int(value)}} for value in
values])
+
+ rows = udf(produce_rows, result_type=result_type,
func_type="pandas")
+ self.assertEqual(list(t.select(rows(t.a)).execute().collect()),
+ [Row({"key": Row(1)}), Row({"key": Row(2)}),
Row({"key": Row(3)})])
+
def test_data_types(self):
import pandas as pd
diff --git a/flink-python/pyflink/table/tests/test_row_based_operation.py
b/flink-python/pyflink/table/tests/test_row_based_operation.py
index 06d8008ecfe..dcbb11828c0 100644
--- a/flink-python/pyflink/table/tests/test_row_based_operation.py
+++ b/flink-python/pyflink/table/tests/test_row_based_operation.py
@@ -25,7 +25,29 @@ from pyflink.table.types import DataTypes
from pyflink.table.udf import udf, udtf, udaf, AggregateFunction,
TableAggregateFunction, udtaf
from pyflink.testing import source_sink_utils
from pyflink.testing.test_case_utils import PyFlinkBatchTableTestCase, \
- PyFlinkStreamTableTestCase
+ PyFlinkStreamTableTestCase, PyFlinkUTTestCase
+
+
+class RowBasedOperationPlanningTests(PyFlinkUTTestCase):
+ def test_async_map_resolves_scalar_and_row_outputs(self):
+ table = self.t_env.from_elements([(1,)], ["input"])
+
+ async def identity(value):
+ return value
+
+ for result_type, names, types in (
+ (DataTypes.BIGINT(), ["f0"], [DataTypes.BIGINT()]),
+ (DataTypes.ROW([DataTypes.FIELD("label", DataTypes.STRING()),
+ DataTypes.FIELD("value", DataTypes.BIGINT())]),
+ ["label", "value"], [DataTypes.STRING(), DataTypes.BIGINT()]),
+ ):
+ for as_expression in (True, False):
+ with self.subTest(result_type=result_type,
as_expression=as_expression):
+ function = udf(identity, result_type=result_type)
+ result = table.map(function(table.input) if as_expression
else function)
+ schema = result.get_resolved_schema()
+ self.assertEqual(list(schema.get_column_names()), names)
+ self.assertEqual(schema.get_column_data_types(), types)
class RowBasedOperationTests(object):
diff --git a/flink-python/pyflink/table/udf.py
b/flink-python/pyflink/table/udf.py
index 733b8a79178..222cfb3be17 100644
--- a/flink-python/pyflink/table/udf.py
+++ b/flink-python/pyflink/table/udf.py
@@ -483,6 +483,7 @@ class UserDefinedFunctionWrapper(object):
self._func_type = func_type
self._judf_placeholder = None
self._takes_row_as_input = False
+ self._returns_table_row = False
def __call__(self, *args) -> Expression:
from pyflink.table import expressions as expr
@@ -570,6 +571,7 @@ class
UserDefinedScalarFunctionWrapper(UserDefinedFunctionWrapper):
j_function_kind,
self._deterministic,
self._takes_row_as_input,
+ self._returns_table_row,
_get_python_env())
return j_scalar_function
@@ -609,6 +611,7 @@ class
UserDefinedAsyncScalarFunctionWrapper(UserDefinedFunctionWrapper):
j_function_kind,
self._deterministic,
self._takes_row_as_input,
+ self._returns_table_row,
_get_python_env())
return j_async_scalar_function
diff --git a/flink-python/pyflink/table/utils.py
b/flink-python/pyflink/table/utils.py
index be1e71f64b2..7001dc9dc9a 100644
--- a/flink-python/pyflink/table/utils.py
+++ b/flink-python/pyflink/table/utils.py
@@ -22,7 +22,7 @@ from pyflink.common.types import RowKind
from pyflink.java_gateway import get_gateway
from pyflink.table.types import DataType, LocalZonedTimestampType, Row,
RowType, \
- TimeType, DateType, ArrayType, MapType, TimestampType, FloatType, RawType
+ TimeType, DateType, ArrayType, MapType, TimestampType, FloatType, RawType,
create_arrow_schema
from pyflink.util.api_stability_decorators import Internal
from pyflink.util.java_utils import to_jarray
import datetime
@@ -160,6 +160,44 @@ def _get_offsets(column):
[None, column.buffers()[1]],
offset=column.offset)
+@Internal()
+def create_pandas_table_row_schema(row_type):
+ import pyarrow as pa
+
+ if len(row_type) != 1 or not isinstance(row_type[0].data_type, RowType):
+ raise TypeError("A table-row result must have exactly one ROW field.")
+ field = row_type[0]
+ # Validate its children as physical columns, not as fields of a user ROW
column.
+ child_schema = create_arrow_schema(
+ field.data_type.field_names(), field.data_type.field_types())
+ return pa.schema([pa.field(field.name, pa.struct(list(child_schema)),
+ nullable=field.data_type._nullable)])
+
+
+@Internal()
+def pandas_table_row_to_arrow(schema, timezone, row_type, dataframe):
+ import pyarrow as pa
+
+ field_types = row_type.field_types()
+ widths = [len(field_type) if isinstance(field_type, RowType) else 1
+ for field_type in field_types]
+ field_count = sum(widths)
+ if dataframe.shape[1] != field_count:
+ raise ValueError(f"Pandas table-row result has {dataframe.shape[1]}
columns, "
+ f"expected {field_count}.")
+ columns = []
+ offset = 0
+ for field_type, width in zip(field_types, widths):
+ # Schema positions retain ROW boundaries even when flattened names are
duplicated.
+ columns.append(dataframe.iloc[:, offset:offset + width]
+ if isinstance(field_type, RowType) else
dataframe.iloc[:, offset])
+ offset += width
+ batch = pandas_to_arrow(
+ pa.schema(list(schema[0].type)), timezone, field_types, columns)
+ array = pa.StructArray.from_arrays(batch.columns,
fields=list(schema[0].type))
+ return pa.RecordBatch.from_arrays([array], schema=schema)
+
+
@Internal()
def pandas_to_arrow(schema, timezone, field_types, series):
import pyarrow as pa
diff --git
a/flink-python/src/main/java/org/apache/flink/python/util/ProtoUtils.java
b/flink-python/src/main/java/org/apache/flink/python/util/ProtoUtils.java
index 6da72d75bfc..af45fd9c804 100644
--- a/flink-python/src/main/java/org/apache/flink/python/util/ProtoUtils.java
+++ b/flink-python/src/main/java/org/apache/flink/python/util/ProtoUtils.java
@@ -121,12 +121,23 @@ public enum ProtoUtils {
FlinkFnApi.CoderInfoDescriptor.Mode mode,
boolean separatedWithEndMessage,
FlinkFnApi.CoderInfoDescriptor.ArrowType.BatchFormat batchFormat) {
+ return createArrowTypeCoderInfoDescriptorProto(
+ rowType, mode, separatedWithEndMessage, batchFormat, false);
+ }
+
+ public static FlinkFnApi.CoderInfoDescriptor
createArrowTypeCoderInfoDescriptorProto(
+ RowType rowType,
+ FlinkFnApi.CoderInfoDescriptor.Mode mode,
+ boolean separatedWithEndMessage,
+ FlinkFnApi.CoderInfoDescriptor.ArrowType.BatchFormat batchFormat,
+ boolean returnsTableRow) {
return createCoderInfoDescriptorProto(
null,
null,
FlinkFnApi.CoderInfoDescriptor.ArrowType.newBuilder()
.setSchema(toProtoType(rowType).getRowSchema())
.setBatchFormat(batchFormat)
+ .setReturnsTableRow(returnsTableRow)
.build(),
null,
null,
diff --git
a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/scalar/AbstractPythonScalarFunctionOperator.java
b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/scalar/AbstractPythonScalarFunctionOperator.java
index 571b7148e14..a5d7bcef8f8 100644
---
a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/scalar/AbstractPythonScalarFunctionOperator.java
+++
b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/scalar/AbstractPythonScalarFunctionOperator.java
@@ -159,6 +159,20 @@ public abstract class AbstractPythonScalarFunctionOperator
config.get(PYTHON_PROFILE_ENABLED));
}
+ protected boolean returnsTableRow() {
+ final int outputCount =
+ udfOutputIndices == null ? scalarFunctions.length :
udfOutputIndices.length;
+ for (int i = 0; i < outputCount; i++) {
+ final int functionIndex = udfOutputIndices == null ? i :
udfOutputIndices[i];
+ if
(scalarFunctions[functionIndex].getPythonFunction().returnsTableRow()) {
+ Preconditions.checkArgument(
+ outputCount == 1, "A table-row result must be the sole
UDF output.");
+ return true;
+ }
+ }
+ return false;
+ }
+
@Override
public String getFunctionUrn() {
return SCALAR_FUNCTION_URN;
diff --git
a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/scalar/arrow/ArrowPythonScalarFunctionOperator.java
b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/scalar/arrow/ArrowPythonScalarFunctionOperator.java
index bee19392571..6e4d77eab2f 100644
---
a/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/scalar/arrow/ArrowPythonScalarFunctionOperator.java
+++
b/flink-python/src/main/java/org/apache/flink/table/runtime/operators/python/scalar/arrow/ArrowPythonScalarFunctionOperator.java
@@ -115,7 +115,11 @@ public class ArrowPythonScalarFunctionOperator extends
AbstractPythonScalarFunct
public FlinkFnApi.CoderInfoDescriptor createOutputCoderInfoDescriptor(
RowType runnerOutputType) {
return createArrowTypeCoderInfoDescriptorProto(
- runnerOutputType, FlinkFnApi.CoderInfoDescriptor.Mode.SINGLE,
false, batchFormat);
+ runnerOutputType,
+ FlinkFnApi.CoderInfoDescriptor.Mode.SINGLE,
+ false,
+ batchFormat,
+ returnsTableRow());
}
@Override
diff --git
a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java
b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java
index 28978469b6e..1170cdecd22 100644
---
a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java
+++
b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/operations/utils/OperationTreeBuilder.java
@@ -552,7 +552,8 @@ public final class OperationTreeBuilder {
final ExpressionResolver resolver = getResolverBuilder(child).build();
final ResolvedExpression resolvedCall =
resolveSingleExpression(mapFunction, resolver);
- if (!isFunctionOfKind(resolvedCall, FunctionKind.SCALAR)) {
+ if (!isFunctionOfKind(resolvedCall, FunctionKind.SCALAR)
+ && !isFunctionOfKind(resolvedCall, FunctionKind.ASYNC_SCALAR))
{
throw new ValidationException(
"Only a scalar function can be used in the map operator.");
}
diff --git
a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonAsyncScalarFunction.java
b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonAsyncScalarFunction.java
index 85c166424c0..e2498ab50a1 100644
---
a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonAsyncScalarFunction.java
+++
b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonAsyncScalarFunction.java
@@ -43,6 +43,7 @@ public class PythonAsyncScalarFunction extends
AsyncScalarFunction implements Py
private final boolean deterministic;
private final PythonEnv pythonEnv;
private final boolean takesRowAsInput;
+ private final boolean returnsTableRow;
private DataType[] inputTypes;
private String[] inputTypesString;
@@ -61,9 +62,32 @@ public class PythonAsyncScalarFunction extends
AsyncScalarFunction implements Py
this(
name,
serializedAsyncScalarFunction,
+ inputTypes,
+ resultType,
pythonFunctionKind,
deterministic,
takesRowAsInput,
+ false,
+ pythonEnv);
+ }
+
+ public PythonAsyncScalarFunction(
+ String name,
+ byte[] serializedAsyncScalarFunction,
+ DataType[] inputTypes,
+ DataType resultType,
+ PythonFunctionKind pythonFunctionKind,
+ boolean deterministic,
+ boolean takesRowAsInput,
+ boolean returnsTableRow,
+ PythonEnv pythonEnv) {
+ this(
+ name,
+ serializedAsyncScalarFunction,
+ pythonFunctionKind,
+ deterministic,
+ takesRowAsInput,
+ returnsTableRow,
pythonEnv);
this.inputTypes = inputTypes;
this.resultType = resultType;
@@ -81,9 +105,32 @@ public class PythonAsyncScalarFunction extends
AsyncScalarFunction implements Py
this(
name,
serializedAsyncScalarFunction,
+ inputTypesString,
+ resultTypeString,
pythonFunctionKind,
deterministic,
takesRowAsInput,
+ false,
+ pythonEnv);
+ }
+
+ public PythonAsyncScalarFunction(
+ String name,
+ byte[] serializedAsyncScalarFunction,
+ String[] inputTypesString,
+ String resultTypeString,
+ PythonFunctionKind pythonFunctionKind,
+ boolean deterministic,
+ boolean takesRowAsInput,
+ boolean returnsTableRow,
+ PythonEnv pythonEnv) {
+ this(
+ name,
+ serializedAsyncScalarFunction,
+ pythonFunctionKind,
+ deterministic,
+ takesRowAsInput,
+ returnsTableRow,
pythonEnv);
this.inputTypesString = inputTypesString;
this.resultTypeString = resultTypeString;
@@ -96,12 +143,31 @@ public class PythonAsyncScalarFunction extends
AsyncScalarFunction implements Py
boolean deterministic,
boolean takesRowAsInput,
PythonEnv pythonEnv) {
+ this(
+ name,
+ serializedAsyncScalarFunction,
+ pythonFunctionKind,
+ deterministic,
+ takesRowAsInput,
+ false,
+ pythonEnv);
+ }
+
+ public PythonAsyncScalarFunction(
+ String name,
+ byte[] serializedAsyncScalarFunction,
+ PythonFunctionKind pythonFunctionKind,
+ boolean deterministic,
+ boolean takesRowAsInput,
+ boolean returnsTableRow,
+ PythonEnv pythonEnv) {
this.name = name;
this.serializedAsyncScalarFunction = serializedAsyncScalarFunction;
this.pythonFunctionKind = pythonFunctionKind;
this.deterministic = deterministic;
this.pythonEnv = pythonEnv;
this.takesRowAsInput = takesRowAsInput;
+ this.returnsTableRow = returnsTableRow;
}
public void eval(CompletableFuture<Object> future, Object... args) {
@@ -129,6 +195,11 @@ public class PythonAsyncScalarFunction extends
AsyncScalarFunction implements Py
return takesRowAsInput;
}
+ @Override
+ public boolean returnsTableRow() {
+ return returnsTableRow;
+ }
+
@Override
public boolean isDeterministic() {
return deterministic;
diff --git
a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunction.java
b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunction.java
index a1a77764cdb..482dad404c9 100644
---
a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunction.java
+++
b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunction.java
@@ -44,4 +44,9 @@ public interface PythonFunction extends Serializable {
default boolean takesRowAsInput() {
return false;
}
+
+ /** Whether the result ROW contains the physical columns of an output
table. */
+ default boolean returnsTableRow() {
+ return false;
+ }
}
diff --git
a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonScalarFunction.java
b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonScalarFunction.java
index 9d14db265f9..59c2f3594a9 100644
---
a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonScalarFunction.java
+++
b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonScalarFunction.java
@@ -44,6 +44,7 @@ public class PythonScalarFunction extends ScalarFunction
implements PythonFuncti
private final boolean deterministic;
private final PythonEnv pythonEnv;
private final boolean takesRowAsInput;
+ private final boolean returnsTableRow;
private DataType[] inputTypes;
private String[] inputTypesString;
@@ -62,9 +63,32 @@ public class PythonScalarFunction extends ScalarFunction
implements PythonFuncti
this(
name,
serializedScalarFunction,
+ inputTypes,
+ resultType,
pythonFunctionKind,
deterministic,
takesRowAsInput,
+ false,
+ pythonEnv);
+ }
+
+ public PythonScalarFunction(
+ String name,
+ byte[] serializedScalarFunction,
+ DataType[] inputTypes,
+ DataType resultType,
+ PythonFunctionKind pythonFunctionKind,
+ boolean deterministic,
+ boolean takesRowAsInput,
+ boolean returnsTableRow,
+ PythonEnv pythonEnv) {
+ this(
+ name,
+ serializedScalarFunction,
+ pythonFunctionKind,
+ deterministic,
+ takesRowAsInput,
+ returnsTableRow,
pythonEnv);
this.inputTypes = inputTypes;
this.resultType = resultType;
@@ -82,9 +106,32 @@ public class PythonScalarFunction extends ScalarFunction
implements PythonFuncti
this(
name,
serializedScalarFunction,
+ inputTypesString,
+ resultTypeString,
pythonFunctionKind,
deterministic,
takesRowAsInput,
+ false,
+ pythonEnv);
+ }
+
+ public PythonScalarFunction(
+ String name,
+ byte[] serializedScalarFunction,
+ String[] inputTypesString,
+ String resultTypeString,
+ PythonFunctionKind pythonFunctionKind,
+ boolean deterministic,
+ boolean takesRowAsInput,
+ boolean returnsTableRow,
+ PythonEnv pythonEnv) {
+ this(
+ name,
+ serializedScalarFunction,
+ pythonFunctionKind,
+ deterministic,
+ takesRowAsInput,
+ returnsTableRow,
pythonEnv);
this.inputTypesString = inputTypesString;
this.resultTypeString = resultTypeString;
@@ -97,12 +144,31 @@ public class PythonScalarFunction extends ScalarFunction
implements PythonFuncti
boolean deterministic,
boolean takesRowAsInput,
PythonEnv pythonEnv) {
+ this(
+ name,
+ serializedScalarFunction,
+ pythonFunctionKind,
+ deterministic,
+ takesRowAsInput,
+ false,
+ pythonEnv);
+ }
+
+ public PythonScalarFunction(
+ String name,
+ byte[] serializedScalarFunction,
+ PythonFunctionKind pythonFunctionKind,
+ boolean deterministic,
+ boolean takesRowAsInput,
+ boolean returnsTableRow,
+ PythonEnv pythonEnv) {
this.name = name;
this.serializedScalarFunction = serializedScalarFunction;
this.pythonFunctionKind = pythonFunctionKind;
this.deterministic = deterministic;
this.pythonEnv = pythonEnv;
this.takesRowAsInput = takesRowAsInput;
+ this.returnsTableRow = returnsTableRow;
}
public Object eval(Object... args) {
@@ -130,6 +196,11 @@ public class PythonScalarFunction extends ScalarFunction
implements PythonFuncti
return takesRowAsInput;
}
+ @Override
+ public boolean returnsTableRow() {
+ return returnsTableRow;
+ }
+
@Override
public boolean isDeterministic() {
return deterministic;
diff --git
a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRule.java
b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRule.java
index c93c643301e..27a953415ba 100644
---
a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRule.java
+++
b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonMapMergeRule.java
@@ -20,6 +20,7 @@ package org.apache.flink.table.planner.plan.rules.logical;
import org.apache.flink.table.functions.python.PythonFunctionKind;
import org.apache.flink.table.planner.plan.nodes.logical.FlinkLogicalCalc;
+import org.apache.flink.table.planner.plan.utils.AsyncUtil;
import org.apache.flink.table.planner.plan.utils.PythonUtil;
import org.apache.calcite.plan.RelOptRule;
@@ -87,6 +88,13 @@ public class PythonMapMergeRule extends RelOptRule {
return false;
}
+ // Async calls are split into separate operators. Removing the flatten
calc would
+ // make a subsequent row-input function receive the enclosing row
instead of its fields.
+ if (AsyncUtil.containsAsyncCall(topProjects.get(0))
+ || AsyncUtil.containsAsyncCall(bottomProjects.get(0))) {
+ return false;
+ }
+
// only Python functions with same Python function kind can be merged
together.
if (Arrays.stream(PythonFunctionKind.values())
.noneMatch(
diff --git
a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRuleTest.java
b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRuleTest.java
index 09cc9d10dd8..4f52355d586 100644
---
a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRuleTest.java
+++
b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRuleTest.java
@@ -21,6 +21,7 @@ package org.apache.flink.table.planner.plan.rules.logical;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.api.Schema;
import org.apache.flink.table.api.Table;
+import org.apache.flink.table.functions.python.PythonAsyncScalarFunction;
import org.apache.flink.table.functions.python.PythonEnv;
import org.apache.flink.table.functions.python.PythonFunctionKind;
import org.apache.flink.table.functions.python.PythonScalarFunction;
@@ -59,6 +60,42 @@ class PythonCalcSplitFunctionKindRuleTest extends
TableTestBase {
verifyMapComposition(javaBatchTestUtil());
}
+ @Test
+ void testStreamingAsyncMapComposition() {
+ verifyAsyncMapComposition(javaStreamTestUtil());
+ }
+
+ @Test
+ void testBatchAsyncMapComposition() {
+ verifyAsyncMapComposition(javaBatchTestUtil());
+ }
+
+ private void verifyAsyncMapComposition(JavaTableTestUtil util) {
+ util.addTableSource("T", Schema.newBuilder().column("a",
DataTypes.INT()).build());
+ final DataType resultType = DataTypes.ROW(DataTypes.FIELD("a",
DataTypes.INT()));
+ registerFunctions(util, resultType, true);
+ util.tableEnv()
+ .createTemporarySystemFunction(
+ "async_udf",
+ new PythonAsyncScalarFunction(
+ "async_udf",
+ new byte[0],
+ new DataType[] {DataTypes.INT()},
+ resultType,
+ PythonFunctionKind.GENERAL,
+ true,
+ true,
+ new PythonEnv(PythonEnv.ExecType.PROCESS)));
+ final Table result =
+ util.tableEnv()
+ .from("T")
+ .map(call("general_udf", $("a")))
+ .map(call("async_udf", $("a")))
+ .map(call("async_udf", $("a")))
+ .map(call("general_udf", $("a")));
+ util.verifyRelPlan(result);
+ }
+
private void verifyComposition(JavaTableTestUtil util) {
util.addTableSource("T", Schema.newBuilder().column("a",
DataTypes.INT()).build());
registerFunctions(util, DataTypes.INT(), false);
diff --git
a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRuleTest.xml
b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRuleTest.xml
index 021138e012a..5cb697af821 100644
---
a/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRuleTest.xml
+++
b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRuleTest.xml
@@ -16,6 +16,27 @@ See the License for the specific language governing
permissions and
limitations under the License.
-->
<Root>
+ <TestCase name="testBatchAsyncMapComposition">
+ <Resource name="ast">
+ <![CDATA[
+LogicalProject(a=[AS(general_udf(AS(async_udf(AS(async_udf(AS(general_udf($0).a,
_UTF-16LE'a')).a, _UTF-16LE'a')).a, _UTF-16LE'a')).a, _UTF-16LE'a')])
++- LogicalTableScan(table=[[default_catalog, default_database, T]])
+]]>
+ </Resource>
+ <Resource name="optimized rel plan">
+ <![CDATA[
+Calc(select=[f0.a AS a])
++- PythonCalc(select=[general_udf(a) AS f0])
+ +- Calc(select=[f0.a AS a])
+ +- PythonAsyncCalc(select=[async_udf(a) AS f0])
+ +- Calc(select=[f0.a AS a])
+ +- PythonAsyncCalc(select=[async_udf(a) AS f0])
+ +- Calc(select=[f0.a AS a])
+ +- PythonCalc(select=[general_udf(a) AS f0])
+ +- TableSourceScan(table=[[default_catalog,
default_database, T]], fields=[a])
+]]>
+ </Resource>
+ </TestCase>
<TestCase name="testBatchComposition">
<Resource name="sql">
<![CDATA[SELECT arrow_udf(a), pandas_udf(a), general_udf(a),
arrow_udf(arrow_udf(a)), arrow_udf(pandas_udf(a)), pandas_udf(arrow_udf(a)),
arrow_udf(general_udf(a)), general_udf(arrow_udf(a)) FROM T]]>
@@ -37,6 +58,27 @@ Calc(select=[f0 AS EXPR$0, f1 AS EXPR$1, f00 AS EXPR$2, f2
AS EXPR$3, f3 AS EXPR
+- PythonCalc(select=[a, f0, general_udf(a) AS f1])
+- PythonCalc(select=[a, pandas_udf(a) AS f0])
+- TableSourceScan(table=[[default_catalog,
default_database, T]], fields=[a])
+]]>
+ </Resource>
+ </TestCase>
+ <TestCase name="testStreamingAsyncMapComposition">
+ <Resource name="ast">
+ <![CDATA[
+LogicalProject(a=[AS(general_udf(AS(async_udf(AS(async_udf(AS(general_udf($0).a,
_UTF-16LE'a')).a, _UTF-16LE'a')).a, _UTF-16LE'a')).a, _UTF-16LE'a')])
++- LogicalTableScan(table=[[default_catalog, default_database, T]])
+]]>
+ </Resource>
+ <Resource name="optimized rel plan">
+ <![CDATA[
+Calc(select=[f0.a AS a])
++- PythonCalc(select=[general_udf(a) AS f0])
+ +- Calc(select=[f0.a AS a])
+ +- PythonAsyncCalc(select=[async_udf(a) AS f0])
+ +- Calc(select=[f0.a AS a])
+ +- PythonAsyncCalc(select=[async_udf(a) AS f0])
+ +- Calc(select=[f0.a AS a])
+ +- PythonCalc(select=[general_udf(a) AS f0])
+ +- TableSourceScan(table=[[default_catalog,
default_database, T]], fields=[a])
]]>
</Resource>
</TestCase>
