[FLINK-40472][python] Support arrow vectorized UDFs in DataFrame API (#29210)

Author: auroflow

flink-python/docs/reference/pyflink.dataframe/udf.rst

@@ -26,9 +26,9 @@ columns. A scalar UDF produces one logical output column and can be used in
 :meth:`~pyflink.dataframe.DataFrame.with_columns`, and
 :meth:`~pyflink.dataframe.DataFrame.select`.
 
-DataFrame scalar UDFs support synchronous, asynchronous, and pandas-vectorized
-callables. See :func:`pyflink.dataframe.udf` for declaration forms, type
-inference, execution modes, and examples.
+DataFrame scalar UDFs support general synchronous and asynchronous callables,
+and synchronous pandas or Arrow vectorized callables. See :func:`pyflink.dataframe.udf`
+for declaration forms, type inference, execution modes, and examples.
 
 API Reference
 =============

flink-python/pyflink/dataframe/tests/test_udf.py

@@ -61,6 +61,86 @@ def _call_module_alias_function(value):
 
 
 class DataFrameUDFDeclarationTests(unittest.TestCase):
+    def test_arrow_annotation_inference_and_overrides(self):
+        def arrow_identity(values: pa.Array) -> pa.ChunkedArray:
+            return pa.chunked_array([values])
+
+        def integer_array(values: pa.Int64Array) -> pa.Int64Array:
+            return values
+
+        def list_array(values: pa.ListArray):
+            return values.flatten()
+
+        def struct_array(values: "pa.StructArray"):
+            return values.field("value")
+
+        def concrete_return(values) -> pa.Int64Array:
+            return values
+
+        def mixed(values: pd.Series) -> pa.Int64Array:
+            return pa.array(values)
+
+        def captured(context: pd.Series, values: pa.Array) -> pa.Array:
+            return values
+
+        class ArrowCallable:
+            def __call__(self, values: "pa.Array") -> "pa.Array":
+                return values
+
+        class ArrowScalar(ScalarFunction):
+            def eval(self, values: pa.Array) -> pa.Array:
+                return values
+
+        for func in (arrow_identity, integer_array, list_array, struct_array, concrete_return,
+                     ArrowCallable, ArrowCallable(), ArrowScalar, ArrowScalar(),
+                     functools.partial(captured, pd.Series([1]))):
+            with self.subTest(func=func):
+                declaration = pf.udf(func, return_dtype=pf.DataType.int64())
+                self.assertEqual(cast(Any, declaration)._func_type, "arrow")
+
+        for mode in ("general", "pandas", "arrow"):
+            with self.subTest(mode=mode):
+                declaration = pf.udf(mixed, return_dtype=pf.DataType.int64(), func_type=mode)
+                self.assertEqual(cast(Any, declaration)._func_type, mode)
+
+        with self.assertRaisesRegex(ValueError, "pandas.*Arrow.*func_type"):
+            pf.udf(mixed, return_dtype=pf.DataType.int64())
+        with self.assertRaisesRegex(TypeError, "return_dtype is required for arrow"):
+            pf.udf(arrow_identity)
+
+    def test_explicit_arrow_declarations(self):
+        from pyflink.table.udf import udf as table_udf
+
+        def identity(values):
+            return values
+
+        declaration = pf.udf(identity, return_dtype=pf.DataType.string(), func_type="arrow")
+        self.assertEqual(_return_dtype(declaration), pf.DataType.string())
+        table_udf(identity, result_type=TableDataTypes.STRING(), func_type="arrow")
+
+        with self.assertRaisesRegex(TypeError, "return_dtype is required for arrow"):
+            pf.udf(identity, func_type="arrow")
+
+        async def async_identity(values):
+            return values
+
+        class AsyncCallable:
+            async def __call__(self, values):
+                return values
+
+        for declare in (
+            lambda: pf.udf(async_identity, return_dtype=pf.DataType.string(), func_type="arrow"),
+            lambda: table_udf(async_identity, result_type=TableDataTypes.STRING(),
+                              func_type="arrow"),
+            lambda: table_udf(AsyncCallable(), result_type=TableDataTypes.STRING(),
+                              func_type="arrow"),
+            lambda: table_udf(functools.partial(AsyncCallable()),
+                              result_type=TableDataTypes.STRING(), func_type="arrow"),
+        ):
+            with self.subTest(declare=declare):
+                with self.assertRaisesRegex(ValueError, "Async.*arrow"):
+                    declare()
+
     def test_function_declarations_return_types_and_metadata(self):
         class Details(TypedDict):
             label: str
@@ -547,9 +627,9 @@ def __call__(self, context: pd.Series, value: int) -> int:
                 False,
             ),
             (
-                "pyarrow annotations remain general",
+                "inferred arrow",
                 lambda: pf.udf(arrow_add_one, return_dtype=pf.DataType.int64()),
-                "general",
+                "arrow",
                 False,
             ),
             (
@@ -1178,11 +1258,11 @@ def eval(self, value):
                 "name must not be empty",
             ),
             (
-                "arrow func type",
+                "unsupported func type",
                 lambda: pf.udf(
                     missing_return,
                     return_dtype=pf.DataType.int64(),
-                    func_type="arrow",
+                    func_type="unsupported",
                 ),
                 ValueError,
                 "func_type must be one of",
@@ -1480,6 +1560,21 @@ def close(self):
 
 
 class DataFrameUDFPlannerTests(PyFlinkDataFrameUTTestCase):
+    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"])
+        for args in ((), (1,), (pf.lit(1),), (identity(),)):
+            with self.subTest(args=args):
+                with self.assertRaisesRegex(Exception, "at least one column-valued argument"):
+                    result = dataframe.with_columns(
+                        valid=identity(pf.col("id")), invalid=identity(*args))
+                    result.to_table().explain(ExplainDetail.JSON_EXECUTION_PLAN)
+
     def test_with_columns_binds_expressions_and_resolves_output_schema(self):
         @pf.udf(name="render_value")
         def render(value: int, suffix: str) -> str:
@@ -1521,6 +1616,25 @@ def describe(value):
 
 class DataFrameUDFITCase(PyFlinkStreamDataFrameTestCase):
     def test_supported_scalar_udfs_in_one_job(self):
+        import pyarrow.compute as pc
+
+        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")
+
+        @pf.udf(return_dtype=pf.DataType.string())
+        def normalize_name(names: pa.Array) -> pa.Array:
+            return pc.utf8_upper(names)
+
+        @pf.udf(return_dtype=pf.DataType.struct({"value": pf.DataType.int64().not_null()}))
+        def describe(values: pa.Array) -> pa.ChunkedArray:
+            result = pa.StructArray.from_arrays([pc.multiply(values, 2)], names=["value"])
+            return pa.chunked_array([result.slice(0, 1), result.slice(1)])
+
+        @pf.udf(return_dtype=pf.DataType.int64())
+        def struct_value(values: pa.Array) -> pa.Array:
+            return pc.struct_field(values, "value")
+
         @dataclass
         class Details:
             doubled: int
@@ -1561,19 +1675,27 @@ def eval(self, *values: int) -> int:
         opened_scalar_class = pf.udf(OpenedScalarFunction)
 
         result = (
-            pf.from_records([(1,)], schema=["id"])
+            pf.from_records([(1, "alice"), (2, None), (3, "Bob")], schema=["id", "name"])
             .with_columns(async_value=add_two(pf.col("id")))
             .with_columns(
                 pandas_value=add_three(pf.col("id")),
                 details=details(pf.col("id")),
                 deferred_value=deferred(pf.col("id")),
                 scalar_value=opened_scalar_class(pf.col("id")),
+                normalized_name=normalize_name(pf.col("name")),
+                arrow_details=describe(pf.col("id")),
+                arrow_after_pandas=struct_value(describe(add_three(pf.col("id")))),
+                pandas_after_arrow=add_three(struct_value(describe(pf.col("id")))),
             )
         )
 
         self.assertEqual(
-            result.collect(),
-            [Row(1, 3, 4, Row(2, ["1"]), 5, 6)],
+            sorted(result.collect(), key=lambda row: row[0]),
+            [
+                Row(1, "alice", 3, 4, Row(2, ["1"]), 5, 6, "ALICE", Row(2), 8, 5),
+                Row(2, None, 4, 5, Row(4, ["2"]), 6, 7, None, Row(4), 10, 7),
+                Row(3, "Bob", 5, 6, Row(6, ["3"]), 7, 8, "BOB", Row(6), 12, 9),
+            ],
         )
 
 

flink-python/pyflink/dataframe/udf.py

@@ -410,6 +410,27 @@ def udf(
         ... def inferred_pandas_add_one(values: pd.Series) -> pd.Series:
         ...     return values + 1
 
+    Arrow UDFs always require an explicit logical ``return_dtype`` and support
+    synchronous functions. Each column argument is received as a ``pyarrow.Array``;
+    a ``ROW``-typed column is received as a ``pyarrow.StructArray`` with one child
+    array per field. Results should be returned as a ``pyarrow.Array`` or
+    ``pyarrow.ChunkedArray`` of the declared logical type, with the same number of
+    rows as the input batch. A ``ROW``-typed result uses a ``pyarrow.StructArray``
+    or a chunked array of structs. Arrow mode can be selected explicitly, or
+    inferred from an Arrow container annotation on any unbound parameter or the
+    return value::
+
+        >>> import pyarrow as pa
+        >>> import pyarrow.compute as pc
+
+        >>> @pf.udf(return_dtype=pf.DataType.int64(), func_type="arrow")
+        ... def arrow_add_one(values):
+        ...     return pc.add(values, 1)
+
+        >>> @pf.udf(return_dtype=pf.DataType.string())
+        ... def normalize_name(names: pa.Array) -> pa.Array:
+        ...     return pc.utf8_upper(names)
+
     A declared UDF is called with DataFrame expressions or Python literals to
     produce a single-column expression::
 
@@ -424,12 +445,13 @@ def udf(
                  callable/scalar-UDF class.
     :param return_dtype: DataFrame logical type, Python type, or SQL type string.
                          General UDFs may infer it from a return annotation;
-                         pandas UDFs require it.
+                         pandas and Arrow UDFs require it.
     :param deterministic: Whether equal inputs always produce equal results.
                           Must agree with scalar-function metadata.
     :param name: Non-empty function identity used by the Table planner.
-    :param func_type: ``"general"`` or ``"pandas"``. If omitted, any unbound
-                      pandas container annotation selects pandas mode.
+    :param func_type: ``"general"``, ``"pandas"``, or ``"arrow"``. If omitted,
+                      unbound container annotations select pandas or Arrow mode;
+                      otherwise general mode is used.
     :return: A callable that accepts DataFrame expressions or Python literals and
              returns an :class:`~pyflink.table.expression.Expression`, or a decorator
              producing such a callable when ``func`` is omitted.
@@ -487,18 +509,18 @@ def _validate_scalar_udf_options(
     return_dtype: Optional[_DataTypeLike],
     is_async: bool,
 ) -> None:
-    if func_type not in ("general", "pandas"):
+    if func_type not in ("general", "pandas", "arrow"):
         raise ValueError(
-            f"The func_type must be one of 'general, pandas', got {func_type}."
+            f"The func_type must be one of 'general, pandas, arrow', got {func_type}."
         )
-    if return_dtype is None and func_type == "pandas":
+    if return_dtype is None and func_type in ("pandas", "arrow"):
         raise TypeError(
-            "return_dtype is required for pandas UDFs because pandas container "
+            f"return_dtype is required for {func_type} UDFs because {func_type} container "
             "annotations do not describe the logical result type."
         )
-    if is_async and func_type == "pandas":
+    if is_async and func_type in ("pandas", "arrow"):
         raise ValueError(
-            "Async scalar functions do not support pandas func_type. "
+            f"Async scalar functions do not support {func_type} func_type. "
             "Use func_type='general'."
         )
 
@@ -1003,30 +1025,44 @@ def _data_type_from_type_hint(type_hint: Any) -> DataType:
 
 
 def _detect_func_type(declaration_context: _UDFDeclarationContext) -> str:
-    """Detect pandas mode from an unbound pandas container annotation."""
+    """Detect a unique vectorized mode from unbound container annotations."""
     hint_func = declaration_context.annotation_target
+    container_types: Dict[str, Tuple[Type, ...]] = {}
+    container_globalns: Dict[str, Any] = {}
     try:
         import pandas as pd
+        container_types["pandas"] = (pd.Series, pd.DataFrame)
+        container_globalns.update(pandas=pd, pd=pd)
     except ImportError:
-        return "general"
+        pass
+    try:
+        import pyarrow as pa
+        container_types["arrow"] = (pa.Array, pa.ChunkedArray)
+        container_globalns.update(pyarrow=pa, pa=pa)
+    except ImportError:
+        pass
 
-    pandas_types = (pd.Series, pd.DataFrame)
-    pandas_globalns = {
-        "pandas": pd,
-        "pd": pd,
-        **declaration_context.globalns,
-    }
+    modes: set[str] = set()
     for name in getattr(hint_func, "__annotations__", {}):
         if name in declaration_context.ignored_hint_names:
             continue
         hint = _resolve_callable_annotation(
             declaration_context,
             name,
-            globalns=pandas_globalns,
+            globalns={**container_globalns, **declaration_context.globalns},
+        )
+        modes.update(
+            mode for mode, types in container_types.items()
+            if hint in types or (
+                mode == "arrow" and isinstance(hint, type) and issubclass(hint, types)
+            )
+        )
+    if len(modes) > 1:
+        raise ValueError(
+            "UDF annotations contain both pandas and Arrow containers; "
+            "specify func_type explicitly."
         )
-        if hint in pandas_types:
-            return "pandas"
-    return "general"
+    return next(iter(modes), "general")
 
 
 # ======================== Worker Adapters ========================

flink-python/pyflink/fn_execution/coder_impl_fast.pxd

@@ -90,13 +90,14 @@ cdef class RowCoderImpl(FieldCoderImpl):
     cdef MaskUtils _mask_utils
 
 cdef class ArrowCoderImpl(FieldCoderImpl):
+    cdef object _batch_format
     cdef object _schema
     cdef list _field_types
     cdef object _timezone
     cdef object _resettable_io
     cdef object _batch_reader
 
-    cdef list decode_one_batch_from_stream(self, InputStream in_stream, size_t size)
+    cdef decode_one_batch_from_stream(self, InputStream in_stream, size_t size)
 
 cdef class OverWindowArrowCoderImpl(FieldCoderImpl):
     cdef ArrowCoderImpl _arrow_coder

flink-python/pyflink/fn_execution/coder_impl_fast.pyx

@@ -36,7 +36,7 @@ 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
+from pyflink.table.utils import pandas_to_arrow, arrow_to_pandas, validate_arrow_batch
 
 ROW_KIND_BIT_SIZE = 2
 
@@ -431,7 +431,8 @@ cdef class ArrowCoderImpl(FieldCoderImpl):
     A coder for arrow format data.
     """
 
-    def __init__(self, schema, row_type, timezone):
+    def __init__(self, schema, row_type, timezone, batch_format="PANDAS"):
+        self._batch_format = batch_format
         self._schema = schema
         self._field_types = row_type.field_types()
         self._timezone = timezone
@@ -443,16 +444,22 @@ cdef class ArrowCoderImpl(FieldCoderImpl):
 
         self._resettable_io.set_output_stream(out_stream)
         batch_writer = pa.RecordBatchStreamWriter(self._resettable_io, self._schema)
-        batch_writer.write_batch(
-            pandas_to_arrow(self._schema, self._timezone, self._field_types, cols))
+        if self._batch_format == "ARROW":
+            batch = validate_arrow_batch(cols, self._schema, self._field_types)
+        else:
+            batch = pandas_to_arrow(self._schema, self._timezone, self._field_types, cols)
+        batch_writer.write_batch(batch)
 
     cpdef decode_from_stream(self, InputStream in_stream, size_t size):
         return self.decode_one_batch_from_stream(in_stream, size)
 
-    cdef list decode_one_batch_from_stream(self, InputStream in_stream, size_t size):
+    cdef decode_one_batch_from_stream(self, InputStream in_stream, size_t size):
         self._resettable_io.set_input_bytes(in_stream.read(size))
         # there is only one arrow batch in the underlying input stream
-        return arrow_to_pandas(self._timezone, self._field_types, [next(self._batch_reader)])
+        batch = next(self._batch_reader)
+        if self._batch_format == "ARROW":
+            return batch
+        return arrow_to_pandas(self._timezone, self._field_types, [batch])
 
     def _load_from_stream(self, stream):
         import pyarrow as pa