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 4eb8666832f [FLINK-40472][python] Support arrow vectorized UDFs in
DataFrame API (#29210)
4eb8666832f is described below
commit 4eb8666832f7fb6fd0a08c01ac93ed135b1dc11d
Author: Liu Liu <[email protected]>
AuthorDate: Fri Sep 18 21:43:39 2026 +0800
[FLINK-40472][python] Support arrow vectorized UDFs in DataFrame API
(#29210)
---
.../docs/reference/pyflink.dataframe/udf.rst | 6 +-
flink-python/pyflink/dataframe/tests/test_udf.py | 136 +++++++++-
flink-python/pyflink/dataframe/udf.py | 78 ++++--
.../pyflink/fn_execution/coder_impl_fast.pxd | 3 +-
.../pyflink/fn_execution/coder_impl_fast.pyx | 19 +-
.../pyflink/fn_execution/coder_impl_slow.py | 19 +-
flink-python/pyflink/fn_execution/coders.py | 31 ++-
.../pyflink/fn_execution/flink_fn_execution_pb2.py | 222 +++++++--------
.../fn_execution/flink_fn_execution_pb2.pyi | 18 +-
.../pyflink/fn_execution/table/operations.py | 17 +-
.../pyflink/fn_execution/tests/test_coders.py | 300 +++++++++++++++++++++
.../pyflink/fn_execution/utils/operation_utils.py | 38 +++
.../pyflink/proto/flink-fn-execution.proto | 8 +
.../pyflink/table/tests/test_pandas_udf.py | 2 +-
flink-python/pyflink/table/tests/test_udf.py | 169 +++++++++++-
flink-python/pyflink/table/types.py | 30 ++-
flink-python/pyflink/table/udf.py | 22 +-
flink-python/pyflink/table/utils.py | 132 +++++++++
.../org/apache/flink/python/util/ProtoUtils.java | 16 ++
.../arrow/ArrowPythonScalarFunctionOperator.java | 12 +-
.../flink/streaming/api/utils/ProtoUtilsTest.java | 51 ++++
.../table/functions/python/PythonFunctionKind.java | 4 +-
.../nodes/exec/common/CommonExecPythonCalc.java | 32 ++-
.../logical/PythonCalcSplitFunctionKindRule.java | 75 ++++++
.../plan/rules/logical/PythonMapMergeRule.java | 8 +-
.../plan/rules/logical/PythonCalcSplitRule.scala | 51 +---
.../PythonCalcSplitFunctionKindRuleTest.java | 112 ++++++++
.../PythonCalcSplitFunctionKindRuleTest.xml | 67 +++++
28 files changed, 1432 insertions(+), 246 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/udf.rst
b/flink-python/docs/reference/pyflink.dataframe/udf.rst
index 73d1ff2abe7..c0084bee45d 100644
--- a/flink-python/docs/reference/pyflink.dataframe/udf.rst
+++ b/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
=============
diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py
b/flink-python/pyflink/dataframe/tests/test_udf.py
index 12be3306aeb..dfd0a396551 100644
--- a/flink-python/pyflink/dataframe/tests/test_udf.py
+++ b/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 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
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 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
"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 @@ class DataFrameUDFAdapterTests(unittest.TestCase):
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 @@ class
DataFrameUDFPlannerTests(PyFlinkDataFrameUTTestCase):
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 @@ class
DataFrameUDFITCase(PyFlinkStreamDataFrameTestCase):
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),
+ ],
)
diff --git a/flink-python/pyflink/dataframe/udf.py
b/flink-python/pyflink/dataframe/udf.py
index b5e55c3c14c..06ad69cd6dc 100644
--- a/flink-python/pyflink/dataframe/udf.py
+++ b/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 ========================
diff --git a/flink-python/pyflink/fn_execution/coder_impl_fast.pxd
b/flink-python/pyflink/fn_execution/coder_impl_fast.pxd
index 05d3ba5fd6b..aeacced5dc5 100644
--- a/flink-python/pyflink/fn_execution/coder_impl_fast.pxd
+++ b/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
diff --git a/flink-python/pyflink/fn_execution/coder_impl_fast.pyx
b/flink-python/pyflink/fn_execution/coder_impl_fast.pyx
index 92dff893fe9..c25ac8cfdd7 100644
--- a/flink-python/pyflink/fn_execution/coder_impl_fast.pyx
+++ b/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
diff --git a/flink-python/pyflink/fn_execution/coder_impl_slow.py
b/flink-python/pyflink/fn_execution/coder_impl_slow.py
index 769720dc277..a7df600d8fb 100644
--- a/flink-python/pyflink/fn_execution/coder_impl_slow.py
+++ b/flink-python/pyflink/fn_execution/coder_impl_slow.py
@@ -31,7 +31,7 @@ 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
+from pyflink.table.utils import pandas_to_arrow, arrow_to_pandas,
validate_arrow_batch
ROW_KIND_BIT_SIZE = 2
@@ -278,7 +278,8 @@ 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
@@ -290,16 +291,22 @@ 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)
def decode_from_stream(self, in_stream: InputStream, length=0):
return self.decode_one_batch_from_stream(in_stream, length)
- def decode_one_batch_from_stream(self, in_stream: InputStream, size: int)
-> List:
+ def decode_one_batch_from_stream(self, in_stream: InputStream, size: int):
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])
@staticmethod
def _load_from_stream(stream):
diff --git a/flink-python/pyflink/fn_execution/coders.py
b/flink-python/pyflink/fn_execution/coders.py
index 838528f24d1..d75afcc1ff9 100644
--- a/flink-python/pyflink/fn_execution/coders.py
+++ b/flink-python/pyflink/fn_execution/coders.py
@@ -36,8 +36,8 @@ from pyflink.common.typeinfo import TypeInformation,
BasicTypeInfo, BasicType, D
ExternalTypeInfo
from pyflink.table.types import TinyIntType, SmallIntType, IntType,
BigIntType, BooleanType, \
FloatType, DoubleType, VarCharType, VarBinaryType, DecimalType, DateType,
TimeType, \
- LocalZonedTimestampType, RowType, RowField, to_arrow_type, TimestampType,
ArrayType, MapType, \
- BinaryType, NullType, CharType
+ LocalZonedTimestampType, RowType, RowField, create_arrow_schema,
TimestampType, ArrayType, \
+ MapType, BinaryType, NullType, CharType
__all__ = ['FlattenRowCoder', 'RowCoder', 'BigIntCoder', 'TinyIntCoder',
'BooleanCoder',
'SmallIntCoder', 'IntCoder', 'FloatCoder', 'DoubleCoder',
'BinaryCoder', 'CharCoder',
@@ -84,16 +84,24 @@ class LengthPrefixBaseCoder(ABC):
field_names = [f.name for f in schema_proto.fields]
return RowCoder(field_coders, field_names)
elif coder_info_descriptor_proto.HasField('arrow_type'):
- timezone = pytz.timezone(os.environ['TABLE_LOCAL_TIME_ZONE'])
schema_proto = coder_info_descriptor_proto.arrow_type.schema
row_type = cls._to_row_type(schema_proto)
- return ArrowCoder(cls._to_arrow_schema(row_type), row_type,
timezone)
+ batch_format =
coder_info_descriptor_proto.arrow_type.BatchFormat.Name(
+ coder_info_descriptor_proto.arrow_type.batch_format)
+ # Native Arrow does not use pandas timezone conversion. Some valid
JVM zone IDs
+ # 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)
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
row_type = cls._to_row_type(schema_proto)
return OverWindowArrowCoder(
- cls._to_arrow_schema(row_type), row_type, timezone)
+ create_arrow_schema(row_type.field_names(),
row_type.field_types()),
+ row_type, timezone)
elif coder_info_descriptor_proto.HasField('raw_type'):
type_info_proto = coder_info_descriptor_proto.raw_type.type_info
field_coder = from_type_info_proto(type_info_proto)
@@ -101,13 +109,6 @@ class LengthPrefixBaseCoder(ABC):
else:
raise ValueError("Unexpected coder type %s" %
coder_info_descriptor_proto)
- @classmethod
- def _to_arrow_schema(cls, row_type):
- import pyarrow as pa
-
- return pa.schema([pa.field(n, to_arrow_type(t), t._nullable)
- for n, t in zip(row_type.field_names(),
row_type.field_types())])
-
@classmethod
def _to_data_type(cls, field_type):
from pyflink.fn_execution import flink_fn_execution_pb2
@@ -242,13 +243,15 @@ class ArrowCoder(FieldCoder):
Coder for Arrow.
"""
- 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._row_type = row_type
self._timezone = timezone
def get_impl(self):
- return coder_impl.ArrowCoderImpl(self._schema, self._row_type,
self._timezone)
+ return coder_impl.ArrowCoderImpl(
+ self._schema, self._row_type, self._timezone, self._batch_format)
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 6fabbefcc7c..54016508148 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)
@@ -54,113 +54,115 @@ if not _descriptor._USE_C_DESCRIPTORS:
_globals['_INPUT']._serialized_start=107
_globals['_INPUT']._serialized_end=261
_globals['_USERDEFINEDFUNCTION']._serialized_start=264
- _globals['_USERDEFINEDFUNCTION']._serialized_end=432
- _globals['_ASYNCOPTIONS']._serialized_start=435
- _globals['_ASYNCOPTIONS']._serialized_end=579
- _globals['_USERDEFINEDFUNCTIONS']._serialized_start=582
- _globals['_USERDEFINEDFUNCTIONS']._serialized_end=1057
- _globals['_OVERWINDOW']._serialized_start=1060
- _globals['_OVERWINDOW']._serialized_end=1409
- _globals['_OVERWINDOW_WINDOWTYPE']._serialized_start=1201
- _globals['_OVERWINDOW_WINDOWTYPE']._serialized_end=1409
- _globals['_USERDEFINEDAGGREGATEFUNCTION']._serialized_start=1412
- _globals['_USERDEFINEDAGGREGATEFUNCTION']._serialized_end=2191
- _globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC']._serialized_start=1677
- _globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC']._serialized_end=2191
-
_globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC_LISTVIEW']._serialized_start=1940
-
_globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC_LISTVIEW']._serialized_end=2024
-
_globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC_MAPVIEW']._serialized_start=2027
-
_globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC_MAPVIEW']._serialized_end=2178
- _globals['_GROUPWINDOW']._serialized_start=2194
- _globals['_GROUPWINDOW']._serialized_end=2750
- _globals['_GROUPWINDOW_WINDOWTYPE']._serialized_start=2558
- _globals['_GROUPWINDOW_WINDOWTYPE']._serialized_end=2649
- _globals['_GROUPWINDOW_WINDOWPROPERTY']._serialized_start=2651
- _globals['_GROUPWINDOW_WINDOWPROPERTY']._serialized_end=2750
- _globals['_USERDEFINEDAGGREGATEFUNCTIONS']._serialized_start=2753
- _globals['_USERDEFINEDAGGREGATEFUNCTIONS']._serialized_end=3464
- _globals['_SCHEMA']._serialized_start=3467
- _globals['_SCHEMA']._serialized_end=5505
- _globals['_SCHEMA_MAPINFO']._serialized_start=3542
- _globals['_SCHEMA_MAPINFO']._serialized_end=3693
- _globals['_SCHEMA_TIMEINFO']._serialized_start=3695
- _globals['_SCHEMA_TIMEINFO']._serialized_end=3724
- _globals['_SCHEMA_TIMESTAMPINFO']._serialized_start=3726
- _globals['_SCHEMA_TIMESTAMPINFO']._serialized_end=3760
- _globals['_SCHEMA_LOCALZONEDTIMESTAMPINFO']._serialized_start=3762
- _globals['_SCHEMA_LOCALZONEDTIMESTAMPINFO']._serialized_end=3806
- _globals['_SCHEMA_ZONEDTIMESTAMPINFO']._serialized_start=3808
- _globals['_SCHEMA_ZONEDTIMESTAMPINFO']._serialized_end=3847
- _globals['_SCHEMA_DECIMALINFO']._serialized_start=3849
- _globals['_SCHEMA_DECIMALINFO']._serialized_end=3896
- _globals['_SCHEMA_BINARYINFO']._serialized_start=3898
- _globals['_SCHEMA_BINARYINFO']._serialized_end=3926
- _globals['_SCHEMA_VARBINARYINFO']._serialized_start=3928
- _globals['_SCHEMA_VARBINARYINFO']._serialized_end=3959
- _globals['_SCHEMA_CHARINFO']._serialized_start=3961
- _globals['_SCHEMA_CHARINFO']._serialized_end=3987
- _globals['_SCHEMA_VARCHARINFO']._serialized_start=3989
- _globals['_SCHEMA_VARCHARINFO']._serialized_end=4018
- _globals['_SCHEMA_FIELDTYPE']._serialized_start=4021
- _globals['_SCHEMA_FIELDTYPE']._serialized_end=5093
- _globals['_SCHEMA_FIELD']._serialized_start=5095
- _globals['_SCHEMA_FIELD']._serialized_end=5203
- _globals['_SCHEMA_TYPENAME']._serialized_start=5206
- _globals['_SCHEMA_TYPENAME']._serialized_end=5505
- _globals['_TYPEINFO']._serialized_start=5508
- _globals['_TYPEINFO']._serialized_end=6855
- _globals['_TYPEINFO_MAPTYPEINFO']._serialized_start=6002
- _globals['_TYPEINFO_MAPTYPEINFO']._serialized_end=6141
- _globals['_TYPEINFO_ROWTYPEINFO']._serialized_start=6144
- _globals['_TYPEINFO_ROWTYPEINFO']._serialized_end=6328
- _globals['_TYPEINFO_ROWTYPEINFO_FIELD']._serialized_start=6237
- _globals['_TYPEINFO_ROWTYPEINFO_FIELD']._serialized_end=6328
- _globals['_TYPEINFO_TUPLETYPEINFO']._serialized_start=6330
- _globals['_TYPEINFO_TUPLETYPEINFO']._serialized_end=6410
- _globals['_TYPEINFO_AVROTYPEINFO']._serialized_start=6412
- _globals['_TYPEINFO_AVROTYPEINFO']._serialized_end=6442
- _globals['_TYPEINFO_TYPENAME']._serialized_start=6445
- _globals['_TYPEINFO_TYPENAME']._serialized_end=6842
- _globals['_USERDEFINEDDATASTREAMFUNCTION']._serialized_start=6858
- _globals['_USERDEFINEDDATASTREAMFUNCTION']._serialized_end=7835
-
_globals['_USERDEFINEDDATASTREAMFUNCTION_RUNTIMECONTEXT']._serialized_start=7353
-
_globals['_USERDEFINEDDATASTREAMFUNCTION_RUNTIMECONTEXT']._serialized_end=7659
-
_globals['_USERDEFINEDDATASTREAMFUNCTION_FUNCTIONTYPE']._serialized_start=7662
- _globals['_USERDEFINEDDATASTREAMFUNCTION_FUNCTIONTYPE']._serialized_end=7835
- _globals['_STATEDESCRIPTOR']._serialized_start=7838
- _globals['_STATEDESCRIPTOR']._serialized_end=9730
- _globals['_STATEDESCRIPTOR_STATETTLCONFIG']._serialized_start=7970
- _globals['_STATEDESCRIPTOR_STATETTLCONFIG']._serialized_end=9730
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES']._serialized_start=8441
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES']._serialized_end=9539
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_INCREMENTALCLEANUPSTRATEGY']._serialized_start=8619
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_INCREMENTALCLEANUPSTRATEGY']._serialized_end=8707
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_ROCKSDBCOMPACTFILTERCLEANUPSTRATEGY']._serialized_start=8709
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_ROCKSDBCOMPACTFILTERCLEANUPSTRATEGY']._serialized_end=8784
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_MAPSTRATEGIESENTRY']._serialized_start=8787
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_MAPSTRATEGIESENTRY']._serialized_end=9395
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_STRATEGIES']._serialized_start=9397
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_STRATEGIES']._serialized_end=9495
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_EMPTYCLEANUPSTRATEGY']._serialized_start=9497
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_EMPTYCLEANUPSTRATEGY']._serialized_end=9539
- _globals['_STATEDESCRIPTOR_STATETTLCONFIG_UPDATETYPE']._serialized_start=9541
- _globals['_STATEDESCRIPTOR_STATETTLCONFIG_UPDATETYPE']._serialized_end=9609
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_STATEVISIBILITY']._serialized_start=9611
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_STATEVISIBILITY']._serialized_end=9685
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_TTLTIMECHARACTERISTIC']._serialized_start=9687
-
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_TTLTIMECHARACTERISTIC']._serialized_end=9730
- _globals['_CODERINFODESCRIPTOR']._serialized_start=9733
- _globals['_CODERINFODESCRIPTOR']._serialized_end=10742
- _globals['_CODERINFODESCRIPTOR_FLATTENROWTYPE']._serialized_start=10326
- _globals['_CODERINFODESCRIPTOR_FLATTENROWTYPE']._serialized_end=10400
- _globals['_CODERINFODESCRIPTOR_ROWTYPE']._serialized_start=10402
- _globals['_CODERINFODESCRIPTOR_ROWTYPE']._serialized_end=10469
- _globals['_CODERINFODESCRIPTOR_ARROWTYPE']._serialized_start=10471
- _globals['_CODERINFODESCRIPTOR_ARROWTYPE']._serialized_end=10540
- _globals['_CODERINFODESCRIPTOR_OVERWINDOWARROWTYPE']._serialized_start=10542
- _globals['_CODERINFODESCRIPTOR_OVERWINDOWARROWTYPE']._serialized_end=10621
- _globals['_CODERINFODESCRIPTOR_RAWTYPE']._serialized_start=10623
- _globals['_CODERINFODESCRIPTOR_RAWTYPE']._serialized_end=10695
- _globals['_CODERINFODESCRIPTOR_MODE']._serialized_start=10697
- _globals['_CODERINFODESCRIPTOR_MODE']._serialized_end=10729
+ _globals['_USERDEFINEDFUNCTION']._serialized_end=454
+ _globals['_ASYNCOPTIONS']._serialized_start=457
+ _globals['_ASYNCOPTIONS']._serialized_end=601
+ _globals['_USERDEFINEDFUNCTIONS']._serialized_start=604
+ _globals['_USERDEFINEDFUNCTIONS']._serialized_end=1079
+ _globals['_OVERWINDOW']._serialized_start=1082
+ _globals['_OVERWINDOW']._serialized_end=1431
+ _globals['_OVERWINDOW_WINDOWTYPE']._serialized_start=1223
+ _globals['_OVERWINDOW_WINDOWTYPE']._serialized_end=1431
+ _globals['_USERDEFINEDAGGREGATEFUNCTION']._serialized_start=1434
+ _globals['_USERDEFINEDAGGREGATEFUNCTION']._serialized_end=2213
+ _globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC']._serialized_start=1699
+ _globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC']._serialized_end=2213
+
_globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC_LISTVIEW']._serialized_start=1962
+
_globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC_LISTVIEW']._serialized_end=2046
+
_globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC_MAPVIEW']._serialized_start=2049
+
_globals['_USERDEFINEDAGGREGATEFUNCTION_DATAVIEWSPEC_MAPVIEW']._serialized_end=2200
+ _globals['_GROUPWINDOW']._serialized_start=2216
+ _globals['_GROUPWINDOW']._serialized_end=2772
+ _globals['_GROUPWINDOW_WINDOWTYPE']._serialized_start=2580
+ _globals['_GROUPWINDOW_WINDOWTYPE']._serialized_end=2671
+ _globals['_GROUPWINDOW_WINDOWPROPERTY']._serialized_start=2673
+ _globals['_GROUPWINDOW_WINDOWPROPERTY']._serialized_end=2772
+ _globals['_USERDEFINEDAGGREGATEFUNCTIONS']._serialized_start=2775
+ _globals['_USERDEFINEDAGGREGATEFUNCTIONS']._serialized_end=3486
+ _globals['_SCHEMA']._serialized_start=3489
+ _globals['_SCHEMA']._serialized_end=5527
+ _globals['_SCHEMA_MAPINFO']._serialized_start=3564
+ _globals['_SCHEMA_MAPINFO']._serialized_end=3715
+ _globals['_SCHEMA_TIMEINFO']._serialized_start=3717
+ _globals['_SCHEMA_TIMEINFO']._serialized_end=3746
+ _globals['_SCHEMA_TIMESTAMPINFO']._serialized_start=3748
+ _globals['_SCHEMA_TIMESTAMPINFO']._serialized_end=3782
+ _globals['_SCHEMA_LOCALZONEDTIMESTAMPINFO']._serialized_start=3784
+ _globals['_SCHEMA_LOCALZONEDTIMESTAMPINFO']._serialized_end=3828
+ _globals['_SCHEMA_ZONEDTIMESTAMPINFO']._serialized_start=3830
+ _globals['_SCHEMA_ZONEDTIMESTAMPINFO']._serialized_end=3869
+ _globals['_SCHEMA_DECIMALINFO']._serialized_start=3871
+ _globals['_SCHEMA_DECIMALINFO']._serialized_end=3918
+ _globals['_SCHEMA_BINARYINFO']._serialized_start=3920
+ _globals['_SCHEMA_BINARYINFO']._serialized_end=3948
+ _globals['_SCHEMA_VARBINARYINFO']._serialized_start=3950
+ _globals['_SCHEMA_VARBINARYINFO']._serialized_end=3981
+ _globals['_SCHEMA_CHARINFO']._serialized_start=3983
+ _globals['_SCHEMA_CHARINFO']._serialized_end=4009
+ _globals['_SCHEMA_VARCHARINFO']._serialized_start=4011
+ _globals['_SCHEMA_VARCHARINFO']._serialized_end=4040
+ _globals['_SCHEMA_FIELDTYPE']._serialized_start=4043
+ _globals['_SCHEMA_FIELDTYPE']._serialized_end=5115
+ _globals['_SCHEMA_FIELD']._serialized_start=5117
+ _globals['_SCHEMA_FIELD']._serialized_end=5225
+ _globals['_SCHEMA_TYPENAME']._serialized_start=5228
+ _globals['_SCHEMA_TYPENAME']._serialized_end=5527
+ _globals['_TYPEINFO']._serialized_start=5530
+ _globals['_TYPEINFO']._serialized_end=6877
+ _globals['_TYPEINFO_MAPTYPEINFO']._serialized_start=6024
+ _globals['_TYPEINFO_MAPTYPEINFO']._serialized_end=6163
+ _globals['_TYPEINFO_ROWTYPEINFO']._serialized_start=6166
+ _globals['_TYPEINFO_ROWTYPEINFO']._serialized_end=6350
+ _globals['_TYPEINFO_ROWTYPEINFO_FIELD']._serialized_start=6259
+ _globals['_TYPEINFO_ROWTYPEINFO_FIELD']._serialized_end=6350
+ _globals['_TYPEINFO_TUPLETYPEINFO']._serialized_start=6352
+ _globals['_TYPEINFO_TUPLETYPEINFO']._serialized_end=6432
+ _globals['_TYPEINFO_AVROTYPEINFO']._serialized_start=6434
+ _globals['_TYPEINFO_AVROTYPEINFO']._serialized_end=6464
+ _globals['_TYPEINFO_TYPENAME']._serialized_start=6467
+ _globals['_TYPEINFO_TYPENAME']._serialized_end=6864
+ _globals['_USERDEFINEDDATASTREAMFUNCTION']._serialized_start=6880
+ _globals['_USERDEFINEDDATASTREAMFUNCTION']._serialized_end=7857
+
_globals['_USERDEFINEDDATASTREAMFUNCTION_RUNTIMECONTEXT']._serialized_start=7375
+
_globals['_USERDEFINEDDATASTREAMFUNCTION_RUNTIMECONTEXT']._serialized_end=7681
+
_globals['_USERDEFINEDDATASTREAMFUNCTION_FUNCTIONTYPE']._serialized_start=7684
+ _globals['_USERDEFINEDDATASTREAMFUNCTION_FUNCTIONTYPE']._serialized_end=7857
+ _globals['_STATEDESCRIPTOR']._serialized_start=7860
+ _globals['_STATEDESCRIPTOR']._serialized_end=9752
+ _globals['_STATEDESCRIPTOR_STATETTLCONFIG']._serialized_start=7992
+ _globals['_STATEDESCRIPTOR_STATETTLCONFIG']._serialized_end=9752
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES']._serialized_start=8463
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES']._serialized_end=9561
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_INCREMENTALCLEANUPSTRATEGY']._serialized_start=8641
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_INCREMENTALCLEANUPSTRATEGY']._serialized_end=8729
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_ROCKSDBCOMPACTFILTERCLEANUPSTRATEGY']._serialized_start=8731
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_ROCKSDBCOMPACTFILTERCLEANUPSTRATEGY']._serialized_end=8806
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_MAPSTRATEGIESENTRY']._serialized_start=8809
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_MAPSTRATEGIESENTRY']._serialized_end=9417
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_STRATEGIES']._serialized_start=9419
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_STRATEGIES']._serialized_end=9517
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_EMPTYCLEANUPSTRATEGY']._serialized_start=9519
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_CLEANUPSTRATEGIES_EMPTYCLEANUPSTRATEGY']._serialized_end=9561
+ _globals['_STATEDESCRIPTOR_STATETTLCONFIG_UPDATETYPE']._serialized_start=9563
+ _globals['_STATEDESCRIPTOR_STATETTLCONFIG_UPDATETYPE']._serialized_end=9631
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_STATEVISIBILITY']._serialized_start=9633
+
_globals['_STATEDESCRIPTOR_STATETTLCONFIG_STATEVISIBILITY']._serialized_end=9707
+
_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_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
# @@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 554c04492d9..520077f395b 100644
--- a/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.pyi
+++ b/flink-python/pyflink/fn_execution/flink_fn_execution_pb2.pyi
@@ -45,18 +45,20 @@ class Input(_message.Message):
def __init__(self, udf: _Optional[_Union[UserDefinedFunction, _Mapping]] =
..., inputOffset: _Optional[int] = ..., inputConstant: _Optional[bytes] = ...,
refIndex: _Optional[int] = ...) -> None: ...
class UserDefinedFunction(_message.Message):
- __slots__ = ("payload", "inputs", "window_index", "takes_row_as_input",
"is_pandas_udf")
+ __slots__ = ("payload", "inputs", "window_index", "takes_row_as_input",
"is_pandas_udf", "is_arrow_udf")
PAYLOAD_FIELD_NUMBER: _ClassVar[int]
INPUTS_FIELD_NUMBER: _ClassVar[int]
WINDOW_INDEX_FIELD_NUMBER: _ClassVar[int]
TAKES_ROW_AS_INPUT_FIELD_NUMBER: _ClassVar[int]
IS_PANDAS_UDF_FIELD_NUMBER: _ClassVar[int]
+ IS_ARROW_UDF_FIELD_NUMBER: _ClassVar[int]
payload: bytes
inputs: _containers.RepeatedCompositeFieldContainer[Input]
window_index: int
takes_row_as_input: bool
is_pandas_udf: bool
- def __init__(self, payload: _Optional[bytes] = ..., inputs:
_Optional[_Iterable[_Union[Input, _Mapping]]] = ..., window_index:
_Optional[int] = ..., takes_row_as_input: bool = ..., is_pandas_udf: bool =
...) -> None: ...
+ is_arrow_udf: bool
+ def __init__(self, payload: _Optional[bytes] = ..., inputs:
_Optional[_Iterable[_Union[Input, _Mapping]]] = ..., window_index:
_Optional[int] = ..., takes_row_as_input: bool = ..., is_pandas_udf: bool =
..., is_arrow_udf: bool = ...) -> None: ...
class AsyncOptions(_message.Message):
__slots__ = ("max_concurrent_operations", "timeout_ms", "retry_enabled",
"retry_max_attempts", "retry_delay_ms")
@@ -645,10 +647,18 @@ class CoderInfoDescriptor(_message.Message):
schema: Schema
def __init__(self, schema: _Optional[_Union[Schema, _Mapping]] = ...)
-> None: ...
class ArrowType(_message.Message):
- __slots__ = ("schema",)
+ __slots__ = ("schema", "batch_format")
+ class BatchFormat(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+ __slots__ = ()
+ PANDAS: _ClassVar[CoderInfoDescriptor.ArrowType.BatchFormat]
+ ARROW: _ClassVar[CoderInfoDescriptor.ArrowType.BatchFormat]
+ PANDAS: CoderInfoDescriptor.ArrowType.BatchFormat
+ ARROW: CoderInfoDescriptor.ArrowType.BatchFormat
SCHEMA_FIELD_NUMBER: _ClassVar[int]
+ BATCH_FORMAT_FIELD_NUMBER: _ClassVar[int]
schema: Schema
- def __init__(self, schema: _Optional[_Union[Schema, _Mapping]] = ...)
-> None: ...
+ batch_format: CoderInfoDescriptor.ArrowType.BatchFormat
+ def __init__(self, schema: _Optional[_Union[Schema, _Mapping]] = ...,
batch_format: _Optional[_Union[CoderInfoDescriptor.ArrowType.BatchFormat, str]]
= ...) -> None: ...
class OverWindowArrowType(_message.Message):
__slots__ = ("schema",)
SCHEMA_FIELD_NUMBER: _ClassVar[int]
diff --git a/flink-python/pyflink/fn_execution/table/operations.py
b/flink-python/pyflink/fn_execution/table/operations.py
index 35ff6a09d68..6459602aa31 100644
--- a/flink-python/pyflink/fn_execution/table/operations.py
+++ b/flink-python/pyflink/fn_execution/table/operations.py
@@ -141,9 +141,11 @@ class ScalarFunctionOperation(BaseOperation):
Generates a UDF execution function. Uses sequential execution with
result references
when refIndex is present (CSE mode), otherwise uses lambda-based
approach.
"""
+ is_arrow = all(udf.is_arrow_udf for udf in serialized_fn.udfs)
+ one_arg_optimization = self._one_arg_optimization and not is_arrow
udf_infos = [
operation_utils.extract_user_defined_function(
- udf, one_arg_optimization=self._one_arg_optimization)
+ udf, one_arg_optimization=one_arg_optimization)
for udf in serialized_fn.udfs]
variable_dict = {}
@@ -154,6 +156,9 @@ class ScalarFunctionOperation(BaseOperation):
user_defined_funcs.extend(funcs)
func_strs.append(func_str)
+ if is_arrow:
+ variable_dict['create_record_batch'] =
operation_utils.create_record_batch
+
output_indices = list(serialized_fn.output_indices)
# Result references require sequential evaluation. A non-empty
output_indices does too:
# it may repeat or reorder results even when none of the UDFs
references another result.
@@ -163,7 +168,10 @@ class ScalarFunctionOperation(BaseOperation):
if not requires_sequential_execution:
# Keep original lambda-based approach for backward compatibility
scalar_functions = ','.join(func_strs)
- if self._one_result_optimization:
+ if is_arrow:
+ func_str = (f'lambda value: create_record_batch('
+ f'[{scalar_functions}], value.num_rows)')
+ elif self._one_result_optimization:
func_str = 'lambda value: %s' % scalar_functions
else:
func_str = 'lambda value: [%s]' % scalar_functions
@@ -186,7 +194,10 @@ class ScalarFunctionOperation(BaseOperation):
code_lines.append(' results = [None] * %d' % len(func_strs))
for i, fn in enumerate(func_strs):
code_lines.append(' results[%d] = %s' % (i, fn))
- if self._one_result_optimization:
+ if is_arrow:
+ outputs = ','.join('results[%d]' % i for i in output_indices)
+ code_lines.append(f' return create_record_batch([{outputs}],
value.num_rows)')
+ elif self._one_result_optimization:
code_lines.append(' return results[%d]' % output_indices[0])
else:
code_lines.append(
diff --git a/flink-python/pyflink/fn_execution/tests/test_coders.py
b/flink-python/pyflink/fn_execution/tests/test_coders.py
index f04a33f1649..8cdb473b623 100644
--- a/flink-python/pyflink/fn_execution/tests/test_coders.py
+++ b/flink-python/pyflink/fn_execution/tests/test_coders.py
@@ -19,7 +19,12 @@
"""Tests common to all coder implementations."""
import decimal
import logging
+import os
import unittest
+from unittest import mock
+
+import pyarrow as pa
+import pytz
from pyflink.fn_execution.coders import BigIntCoder, TinyIntCoder,
BooleanCoder, \
SmallIntCoder, IntCoder, FloatCoder, DoubleCoder, BinaryCoder, CharCoder,
DateCoder, \
@@ -30,6 +35,301 @@ from pyflink.datastream.window import TimeWindow,
CountWindow
from pyflink.testing.test_case_utils import PyFlinkTestCase
+class ArrowSchemaTests(unittest.TestCase):
+ def test_nested_types_require_arrow_mode(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'))])),
+ ):
+ 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)
+
+ def test_pandas_collection_schema_and_round_trip(self):
+ import pandas as pd
+ from pyflink.table import DataTypes
+ from pyflink.table.types import create_arrow_schema
+ from pyflink.table.utils import arrow_to_pandas, pandas_to_arrow
+
+ names = ["values", "lookup", "record"]
+ types = [
+ DataTypes.ARRAY(DataTypes.INT().not_null()),
+ DataTypes.MAP(DataTypes.STRING(),
DataTypes.ARRAY(DataTypes.INT()).not_null()),
+ DataTypes.ROW([DataTypes.FIELD("values",
DataTypes.ARRAY(DataTypes.INT().not_null()))])]
+ schema = create_arrow_schema(names, types)
+ self.assertEqual(schema.field("values").type.value_field.name, "item")
+ self.assertFalse(schema.field("values").type.value_field.nullable)
+ if hasattr(pa.MapType, 'item_field'):
+ self.assertFalse(schema.field("lookup").type.key_field.nullable)
+ self.assertFalse(schema.field("lookup").type.item_field.nullable)
+
self.assertTrue(schema.field("lookup").type.item_type.value_field.nullable)
+
self.assertFalse(schema.field("record").type[0].type.value_field.nullable)
+
+ # Correcting schema metadata must not add null checks to the pandas
conversion path.
+ columns = [pd.Series([[1, 2], [], None, [None]]),
+ pd.Series([[('a', [1])], [], None, [('b', None)]]),
+ pd.DataFrame({"values": [[3], [], None, [None]]})]
+ batch = pandas_to_arrow(schema, pytz.UTC, types, columns)
+ with pa.BufferOutputStream() as output:
+ with pa.ipc.new_stream(output, schema) as writer:
+ writer.write_batch(batch)
+ decoded = pa.ipc.open_stream(output.getvalue()).read_next_batch()
+ expected = [
+ {"values": [1, 2], "lookup": [('a', [1])], "record": {"values":
[3]}},
+ {"values": [], "lookup": [], "record": {"values": []}},
+ {"values": None, "lookup": None, "record": {"values": None}},
+ {"values": [None], "lookup": [('b', None)], "record": {"values":
[None]}}]
+ expected = {name: [row[name] for row in expected] for name in names}
+ self.assertEqual(decoded.to_pydict(), expected)
+ restored = arrow_to_pandas(pytz.UTC, types, [decoded])
+ self.assertEqual(pandas_to_arrow(schema, pytz.UTC, types,
restored).to_pydict(), expected)
+
+ def test_arrow_descriptor_preserves_pandas_default(self):
+ from pyflink.fn_execution import flink_fn_execution_pb2 as proto
+ from pyflink.fn_execution.coders import LengthPrefixBaseCoder
+ import pandas as pd
+
+ arrow_type =
proto.CoderInfoDescriptor.ArrowType(schema=proto.Schema(fields=[
+ proto.Schema.Field(name="name", type=proto.Schema.FieldType(
+ type_name=proto.Schema.VARCHAR, nullable=True,
+ var_char_info=proto.Schema.VarCharInfo(length=2147483647)))]))
+ descriptor = proto.CoderInfoDescriptor(arrow_type=arrow_type)
+ with mock.patch.dict(os.environ, {"TABLE_LOCAL_TIME_ZONE": "UTC"}):
+ pandas_coder =
LengthPrefixBaseCoder._to_field_coder(descriptor).get_impl()
+ result = pandas_coder.decode(pandas_coder.encode([pd.Series(["a",
None])]))
+ self.assertIsInstance(result[0], pd.Series)
+ self.assertEqual(result[0].tolist(), ["a", None])
+ descriptor.arrow_type.batch_format =
proto.CoderInfoDescriptor.ArrowType.ARROW
+ arrow_coder =
LengthPrefixBaseCoder._to_field_coder(descriptor).get_impl()
+ batch = pa.record_batch([pa.array(["a", None])], names=["name"])
+ self.assertEqual(arrow_coder.decode(arrow_coder.encode(batch)),
batch)
+
+
+class ArrowCodersTests(unittest.TestCase):
+ from pyflink.fn_execution import coder_impl_slow as implementation
+
+ def arrow_coder(self, schema, row_type):
+ return self.implementation.ArrowCoderImpl(schema, row_type, pytz.UTC,
"ARROW")
+
+ @staticmethod
+ def with_parent_nulls(column, nulls):
+ validity = pa.array([not value for value in nulls]).buffers()[1]
+ if pa.types.is_struct(column.type):
+ buffers = [validity]
+ children = [column.field(index) for index in
range(column.type.num_fields)]
+ else:
+ buffers = [validity, column.buffers()[1]]
+ if pa.types.is_map(column.type):
+ children = [pa.StructArray.from_arrays(
+ [column.keys, column.items],
+ fields=[pa.field("key", column.type.key_type,
nullable=False),
+ pa.field("value", column.type.item_type)])]
+ else:
+ children = [column.values]
+ return pa.Array.from_buffers(column.type, len(column), buffers,
children=children)
+
+ def test_native_arrow_timezone(self):
+ from pyflink.fn_execution import flink_fn_execution_pb2 as proto
+ from pyflink.fn_execution.coders import LengthPrefixBaseCoder
+
+ schema_proto = proto.Schema(fields=[
+ proto.Schema.Field(name="number", type=proto.Schema.FieldType(
+ type_name=proto.Schema.BIGINT, nullable=True)),
+ proto.Schema.Field(name="timestamp", type=proto.Schema.FieldType(
+ type_name=proto.Schema.TIMESTAMP, nullable=True,
+ timestamp_info=proto.Schema.TimestampInfo(precision=3))),
+ proto.Schema.Field(name="local_timestamp",
type=proto.Schema.FieldType(
+ type_name=proto.Schema.LOCAL_ZONED_TIMESTAMP, nullable=True,
+
local_zoned_timestamp_info=proto.Schema.LocalZonedTimestampInfo(precision=3)))])
+ descriptor =
proto.CoderInfoDescriptor(arrow_type=proto.CoderInfoDescriptor.ArrowType(
+ schema=schema_proto,
batch_format=proto.CoderInfoDescriptor.ArrowType.ARROW))
+ batch = pa.record_batch([pa.array([1, None]), pa.array([0, None],
type=pa.timestamp('ms')),
+ pa.array([0, None], type=pa.timestamp('ms'))],
+ names=["number", "timestamp",
"local_timestamp"])
+ for timezone in ("UTC", "GMT+08:00", "SystemV/PST8PDT"):
+ with self.subTest(timezone=timezone):
+ with mock.patch.dict(os.environ, {"TABLE_LOCAL_TIME_ZONE":
timezone}), \
+ mock.patch('pyflink.fn_execution.coders.coder_impl',
self.implementation):
+ coder =
LengthPrefixBaseCoder._to_field_coder(descriptor).get_impl()
+ self.assertEqual(coder.decode(coder.encode(batch)), batch)
+
+ def test_arrow_nested_nullability(self):
+ from pyflink.table import DataTypes
+ from pyflink.table.types import to_arrow_type
+
+ row_type = DataTypes.ROW([
+ DataTypes.FIELD("values",
DataTypes.ARRAY(DataTypes.INT().not_null()))])
+ schema = pa.schema([pa.field("values",
to_arrow_type(row_type.field_types()[0]))])
+ coder = self.arrow_coder(schema, row_type)
+ batch = pa.record_batch([pa.array([[1, None]],
type=pa.list_(pa.int32()))], schema=schema)
+ with self.assertRaisesRegex(ValueError, "values.*not nullable"):
+ coder.encode(batch)
+
+ valid = pa.record_batch([pa.array([[1, 2], None],
type=pa.list_(pa.int32()))],
+ schema=schema)
+ self.assertEqual(coder.decode(coder.encode(valid)), valid)
+
+ def test_struct_map_and_temporal_results(self):
+ import datetime
+ from pyflink.table import DataTypes
+ from pyflink.table.types import create_arrow_schema
+
+ row_type = DataTypes.ROW([
+ DataTypes.FIELD("record", DataTypes.ROW([
+ DataTypes.FIELD("inner", DataTypes.ROW([
+ DataTypes.FIELD("value", DataTypes.INT().not_null())]))])),
+ DataTypes.FIELD("lookup", DataTypes.MAP(
+ DataTypes.STRING().not_null(), DataTypes.INT().not_null())),
+ DataTypes.FIELD("amount", DataTypes.DECIMAL(6, 2)),
+ DataTypes.FIELD("time", DataTypes.TIMESTAMP(3))])
+ schema = create_arrow_schema(row_type.field_names(),
row_type.field_types(),
+ allow_nested=True)
+ coder = self.arrow_coder(schema, row_type)
+ rows = [
+ {"record": {"inner": {"value": 7}}, "lookup": [("a", 1)],
+ "amount": decimal.Decimal("12.34"), "time":
datetime.datetime(2020, 1, 2)},
+ {"record": None, "lookup": None, "amount": None, "time": None},
+ {"record": {"inner": None}, "lookup": [],
+ "amount": decimal.Decimal("-0.50"), "time":
datetime.datetime(2021, 3, 4)}]
+ expected = {name: [row[name] for row in rows] for name in schema.names}
+ batch = pa.record_batch([pa.array(expected[field.name],
type=field.type)
+ for field in schema], schema=schema)
+ self.assertEqual(coder.decode(coder.encode(batch)).to_pydict(),
expected)
+
self.assertEqual(coder.decode(coder.encode(batch.slice(1))).to_pydict(),
+ {name: values[1:] for name, values in
expected.items()})
+
+ for field, value, message in (
+ ("record", {"inner": {"value": None}}, "record.inner.value.*not
nullable"),
+ ("lookup", [("a", None)], "lookup.value.*not nullable"),
+ ):
+ with self.subTest(field=field):
+ invalid_values = {**rows[0], field: value}
+ invalid = pa.record_batch([
+ pa.array([invalid_values[child.name]], type=child.type)
for child in schema],
+ schema=schema)
+ with self.assertRaisesRegex(ValueError, message):
+ coder.encode(invalid)
+
+ wrong = pa.StructArray.from_arrays([pa.array([1], type=pa.int64())],
names=["value"])
+ invalid = pa.record_batch([pa.StructArray.from_arrays([wrong],
names=["inner"]),
+ batch.column(1).slice(0, 1),
batch.column(2).slice(0, 1),
+ batch.column(3).slice(0, 1)],
names=schema.names)
+ with self.assertRaisesRegex(TypeError,
"record.inner.value.*int64.*int32"):
+ coder.encode(invalid)
+
+ def test_native_arrow_round_trip(self):
+ from pyflink.table import DataTypes
+
+ row_type = DataTypes.ROW([DataTypes.FIELD("name", DataTypes.STRING())])
+ schema = pa.schema([pa.field("name", pa.string())])
+ coder = self.arrow_coder(schema, row_type)
+ batch = pa.record_batch([pa.array(["ALICE", None, "BOB"])],
schema=schema)
+ self.assertEqual(coder.decode(coder.encode(batch)), batch)
+
+ with self.assertRaisesRegex(TypeError, "name.*string"):
+ coder.encode(pa.record_batch([pa.array([1, 2])], names=["name"]))
+
+ def test_pandas_map_round_trip(self):
+ import pandas as pd
+ from pyflink.table import DataTypes
+ from pyflink.table.types import create_arrow_schema
+
+ row_type = DataTypes.ROW([
+ DataTypes.FIELD("lookup", DataTypes.MAP(DataTypes.STRING(),
DataTypes.BIGINT()))])
+ schema = create_arrow_schema(row_type.field_names(),
row_type.field_types())
+ coder = self.implementation.ArrowCoderImpl(schema, row_type, pytz.UTC)
+ values = pd.Series([[("value", 10)], [("value", 20)], None, []])
+ result = coder.decode(coder.encode([values]))
+ self.assertEqual(result[0].tolist(), [[("value", 10)], [("value",
20)], None, []])
+
+ def test_inferred_struct_nullability(self):
+ from pyflink.table import DataTypes
+ from pyflink.table.types import create_arrow_schema
+
+ row_type = DataTypes.ROW([DataTypes.FIELD("wrapped", DataTypes.ROW([
+ DataTypes.FIELD("v", DataTypes.BIGINT().not_null())]))])
+ schema = create_arrow_schema(row_type.field_names(),
row_type.field_types(),
+ allow_nested=True)
+ coder = self.arrow_coder(schema, row_type)
+ column = pa.StructArray.from_arrays([pa.array([0, 1, 2])], names=["v"])
+ with_nulls = pa.StructArray.from_arrays([pa.array([None, 1, 2])],
names=["v"])
+ for values in (column, column.slice(1), column.slice(1, 0),
+ with_nulls.slice(1), with_nulls.slice(1, 0)):
+ with self.subTest(offset=values.offset, length=len(values)):
+ result = coder.decode(coder.encode(pa.record_batch([values],
names=["wrapped"])))
+ self.assertEqual(result.schema, schema)
+ self.assertEqual(result.column(0).to_pylist(),
values.to_pylist())
+ invalid = pa.StructArray.from_arrays([pa.array([None],
type=pa.int64())], names=["v"])
+ with self.assertRaisesRegex(ValueError, "wrapped.v.*not nullable"):
+ coder.encode(pa.record_batch([invalid], names=["wrapped"]))
+
+ def test_sliced_container_nullability(self):
+ from pyflink.table import DataTypes
+ from pyflink.table.types import create_arrow_schema
+
+ item_type = DataTypes.ROW([DataTypes.FIELD("required",
DataTypes.INT().not_null())])
+ items = pa.array([{"required": None}, {"required": 1}, {"required":
None},
+ {"required": None}, None, {"required": None}],
+ type=pa.struct([pa.field("required", pa.int32())]))
+ booleans = pa.array([None, True, None, None, False, None])
+ nulls = [False, False, True, False, False, False]
+ offsets = pa.array([0, 1, 2, 2, 4, 5, 6], type=pa.int32())
+ for data_type, column, first, last in (
+ (DataTypes.ROW([DataTypes.FIELD("value", item_type)]),
+ pa.StructArray.from_arrays([items], names=["value"]),
+ {"value": {"required": 1}}, {"value": None}),
+ (DataTypes.ARRAY(item_type), pa.ListArray.from_arrays(offsets,
items),
+ [{"required": 1}], [None]),
+ (DataTypes.MAP(DataTypes.STRING().not_null(), item_type),
+ pa.MapArray.from_arrays(offsets, pa.array(['k'] * 6), items),
+ [('k', {"required": 1})], [('k', None)]),
+ (DataTypes.ARRAY(DataTypes.BOOLEAN().not_null()),
+ pa.ListArray.from_arrays(offsets, booleans), [True], [False]),
+ (DataTypes.MAP(DataTypes.STRING().not_null(),
DataTypes.BOOLEAN().not_null()),
+ pa.MapArray.from_arrays(offsets, pa.array(['k'] * 6), booleans),
+ [('k', True)], [('k', False)]),
+ ):
+ with self.subTest(data_type=data_type):
+ column = self.with_parent_nulls(column, nulls)
+ row_type = DataTypes.ROW([DataTypes.FIELD("record",
DataTypes.ROW([
+ DataTypes.FIELD("container", data_type)]))])
+ schema = create_arrow_schema(row_type.field_names(),
row_type.field_types(),
+ allow_nested=True)
+ outer = self.with_parent_nulls(
+ pa.StructArray.from_arrays([column], names=["container"]),
+ [False, False, False, True, False, False])
+ batch = pa.record_batch([outer], names=["record"])
+ coder = self.arrow_coder(schema, row_type)
+ self.assertEqual(coder.decode(coder.encode(batch.slice(1,
4))).to_pydict(), {
+ "record": [{"container": first}, {"container": None},
+ None, {"container": last}]})
+ self.assertEqual(coder.decode(coder.encode(batch.slice(0,
0))).num_rows, 0)
+ with self.assertRaisesRegex(ValueError, "not nullable"):
+ coder.encode(batch.slice(0, 4))
+
+
+try:
+ from pyflink.fn_execution import coder_impl_fast
+except ImportError:
+ coder_impl_fast = None
+
+
[email protected](coder_impl_fast is None, "Compiled coders are not installed")
+class FastArrowCodersTests(ArrowCodersTests):
+ implementation = coder_impl_fast
+
+
class CodersTest(PyFlinkTestCase):
def check_coder(self, coder, *values):
diff --git a/flink-python/pyflink/fn_execution/utils/operation_utils.py
b/flink-python/pyflink/fn_execution/utils/operation_utils.py
index ff96585e4c2..5bb4ff18226 100644
--- a/flink-python/pyflink/fn_execution/utils/operation_utils.py
+++ b/flink-python/pyflink/fn_execution/utils/operation_utils.py
@@ -93,6 +93,35 @@ def check_pandas_udf_result(f, *input_args):
return output
+def check_arrow_udf_result(func, *args):
+ import pyarrow as pa
+
+ result = func(*args)
+ name = getattr(func, "__qualname__", type(func).__name__)
+ if not isinstance(result, (pa.Array, pa.ChunkedArray)):
+ raise TypeError(
+ f"Arrow UDF '{name}' must return a pyarrow.Array or
pyarrow.ChunkedArray, "
+ f"got {type(result).__name__}.")
+ for arg in args:
+ if isinstance(arg, (pa.Array, pa.ChunkedArray)) and len(result) !=
len(arg):
+ raise ValueError(
+ f"Arrow UDF '{name}' returned {len(result)} rows, expected
{len(arg)}.")
+ if isinstance(result, pa.ChunkedArray):
+ result = result.chunk(0) if result.num_chunks == 1 else
result.combine_chunks()
+ return result
+
+
+def create_record_batch(results, row_count):
+ import pyarrow as pa
+
+ columns = []
+ for result in results:
+ if len(result) != row_count:
+ raise ValueError(f"Arrow UDF returned {len(result)} rows, expected
{row_count}.")
+ columns.append(result)
+ return pa.RecordBatch.from_arrays(columns, names=[f"f{i}" for i in
range(len(columns))])
+
+
def extract_over_window_user_defined_function(user_defined_function_proto):
window_index = user_defined_function_proto.window_index
return (*extract_user_defined_function(user_defined_function_proto, True),
window_index)
@@ -160,6 +189,8 @@ def
extract_user_defined_function(user_defined_function_proto, pandas_udaf=False
variable_dict[func_name] = user_defined_func.func
else:
variable_dict[func_name] = user_defined_func.eval
+ if user_defined_function_proto.is_arrow_udf:
+ variable_dict[func_name] = partial(check_arrow_udf_result,
variable_dict[func_name])
user_defined_funcs.append(user_defined_func)
func_args, input_variable_dict, input_funcs =
_extract_input(user_defined_function_proto.inputs)
@@ -179,6 +210,13 @@ def
extract_user_defined_function(user_defined_function_proto, pandas_udaf=False
# receives a previously computed intermediate result rather than
# the original input row. We must use func_args instead of `value`.
func_str = "%s(%s)" % (func_name, func_args)
+ elif user_defined_function_proto.is_arrow_udf:
+ import pyarrow as pa
+
+ variable_dict['create_struct_array'] = pa.StructArray.from_arrays
+ offsets = [arg.inputOffset for arg in
user_defined_function_proto.inputs]
+ func_str = (f"{func_name}(create_struct_array([{func_args}], "
+ f"fields=[value.schema[i] for i in {offsets}]))")
else:
# directly use `value` as input argument
# e.g.
diff --git a/flink-python/pyflink/proto/flink-fn-execution.proto
b/flink-python/pyflink/proto/flink-fn-execution.proto
index 4644909bad9..a9300ba4949 100644
--- a/flink-python/pyflink/proto/flink-fn-execution.proto
+++ b/flink-python/pyflink/proto/flink-fn-execution.proto
@@ -64,6 +64,9 @@ message UserDefinedFunction {
// Whether it's pandas UDF
bool is_pandas_udf = 5;
+
+ // Whether the scalar UDF consumes and returns Arrow arrays directly
+ bool is_arrow_udf = 6;
}
// Async execution configuration for async functions
@@ -526,6 +529,11 @@ message CoderInfoDescriptor {
message ArrowType {
Schema schema = 1;
+ enum BatchFormat {
+ PANDAS = 0;
+ ARROW = 1;
+ }
+ BatchFormat batch_format = 2;
}
// only used in batch over window
diff --git a/flink-python/pyflink/table/tests/test_pandas_udf.py
b/flink-python/pyflink/table/tests/test_pandas_udf.py
index 192ce9c65b4..84d25d4307d 100644
--- a/flink-python/pyflink/table/tests/test_pandas_udf.py
+++ b/flink-python/pyflink/table/tests/test_pandas_udf.py
@@ -32,7 +32,7 @@ class PandasUDFTests(PyFlinkTestCase):
def test_non_exist_func_type(self):
with self.assertRaisesRegex(ValueError,
- 'The func_type must be one of \'general,
pandas\''):
+ "The func_type must be one of 'general,
pandas, arrow'"):
udf(lambda i: i + 1, result_type=DataTypes.BIGINT(),
func_type="non-exist")
diff --git a/flink-python/pyflink/table/tests/test_udf.py
b/flink-python/pyflink/table/tests/test_udf.py
index eb3e1447c94..3a5d42107e1 100644
--- a/flink-python/pyflink/table/tests/test_udf.py
+++ b/flink-python/pyflink/table/tests/test_udf.py
@@ -17,15 +17,21 @@
################################################################################
import datetime
import os
+import pickle
import unittest
import uuid
+import cloudpickle
+import pyarrow as pa
+import pyarrow.compute as pc
import pytz
from pyflink.common import Row
+from pyflink.fn_execution import flink_fn_execution_pb2 as proto
+from pyflink.fn_execution.table.operations import ScalarFunctionOperation
from pyflink.table import DataTypes, expressions as expr
from pyflink.table.expressions import call
-from pyflink.table.udf import ScalarFunction, udf, FunctionContext
+from pyflink.table.udf import DelegatingScalarFunction, ScalarFunction, udf,
FunctionContext
from pyflink.testing import source_sink_utils
from pyflink.testing.test_case_utils import PyFlinkStreamTableTestCase, \
PyFlinkBatchTableTestCase
@@ -35,6 +41,167 @@ def generate_random_table_name():
return "Table{0}".format(str(uuid.uuid1()).replace("-", "_"))
+class ArrowScalarOperationTests(unittest.TestCase):
+ def operation(self, func, inputs, takes_row_as_input=False, preceding=(),
output_indices=()):
+ function = proto.UserDefinedFunction(
+ payload=cloudpickle.dumps(DelegatingScalarFunction(func)),
+ is_arrow_udf=True, inputs=inputs,
takes_row_as_input=takes_row_as_input)
+ operation = ScalarFunctionOperation(proto.UserDefinedFunctions(
+ udfs=[*preceding, function], output_indices=output_indices))
+ operation.open()
+ self.addCleanup(operation.close)
+ return operation
+
+ def test_literals_columns_and_chunked_results(self):
+ def add(offset, left, right):
+ if not isinstance(offset, int) or not isinstance(left, pa.Array) \
+ or not isinstance(right, pa.Array):
+ raise TypeError("Expected a scalar literal followed by two
Arrow arrays.")
+ values = pc.add(pc.add(left, right), offset)
+ return pa.chunked_array([values.slice(0, 1), values.slice(1)])
+
+ operation = self.operation(add, [
+ proto.Input(inputConstant=b"\x00" + pickle.dumps(10)),
+ proto.Input(inputOffset=0), proto.Input(inputOffset=1)])
+ result = operation.process_element(pa.record_batch(
+ [pa.array([1, None, 3]), pa.array([4, 5, 6])], names=["a", "b"]))
+ self.assertEqual(result.column(0).to_pylist(), [15, None, 19])
+
+ def test_single_chunk_result_preserves_buffers(self):
+ values = pa.array([0, 1, None, 3]).slice(1, 2)
+ operation = self.operation(lambda column: pa.chunked_array([column]),
+ [proto.Input(inputOffset=0)])
+ result = operation.process_element(pa.record_batch([values],
names=["value"]))
+ column = result.column(0)
+ self.assertEqual(column.to_pylist(), [1, None])
+ self.assertEqual(column.offset, values.offset)
+ self.assertEqual([buffer.address for buffer in column.buffers()],
+ [buffer.address for buffer in values.buffers()])
+
+ def test_empty_chunked_results(self):
+ values = pa.array([], type=pa.int64())
+ for chunks in ([], [values]):
+ with self.subTest(chunks=len(chunks)):
+ operation = self.operation(
+ lambda column: pa.chunked_array(chunks, type=pa.int64()),
+ [proto.Input(inputOffset=0)])
+ result = operation.process_element(pa.record_batch([values],
names=["value"]))
+ self.assertEqual(result.column(0), values)
+
+ def test_nested_chunked_results(self):
+ values = pa.StructArray.from_arrays([pa.array([0, 1, None, 3])],
names=["value"])
+ for chunks in (0, 1, 2):
+ with self.subTest(chunks=chunks):
+ column = values.slice(1, 0 if chunks == 0 else 3)
+
+ def chunk_result(array):
+ parts = [] if chunks == 0 else [array] if chunks == 1 else
[
+ array.slice(0, 1), array.slice(1)]
+ return pa.chunked_array(parts, type=array.type)
+
+ inner = proto.UserDefinedFunction(
+
payload=cloudpickle.dumps(DelegatingScalarFunction(chunk_result)),
+ is_arrow_udf=True, inputs=[proto.Input(inputOffset=0)])
+ for shared in (False, True):
+ with self.subTest(shared=shared):
+ operation = self.operation(
+ lambda row: row.field("value"),
+ [proto.Input(refIndex=0)] if shared else
[proto.Input(udf=inner)],
+ preceding=[inner] if shared else [],
+ output_indices=[1, 0, 1] if shared else [])
+ result = operation.process_element(
+ pa.record_batch([column], names=["record"]))
+ self.assertEqual(result.num_columns, 3 if shared else
1)
+ self.assertEqual(result.column(0).to_pylist(),
+ [] if chunks == 0 else [1, None, 3])
+ if shared:
+ self.assertEqual(result.column(1), column)
+ self.assertEqual(result.column(2),
result.column(0))
+ if chunks == 1:
+ self.assertEqual(result.column(0).offset,
column.field(0).offset)
+ self.assertEqual(
+ [buffer.address for buffer in
result.column(0).buffers()],
+ [buffer.address for buffer in
column.field(0).buffers()])
+
+ def test_whole_row_input(self):
+ batch = pa.record_batch([pa.array(["alice", None]), pa.array([1, 2])],
+ names=["name", "count"])
+
+ def increment(row):
+ return pa.StructArray.from_arrays(
+ [pc.struct_field(row, "name"), pc.add(pc.struct_field(row,
"count"), 1)],
+ names=["name", "count"])
+
+ inputs = [proto.Input(inputOffset=0), proto.Input(inputOffset=1)]
+ identity = proto.UserDefinedFunction(
+ payload=cloudpickle.dumps(DelegatingScalarFunction(lambda row:
row)),
+ is_arrow_udf=True, inputs=inputs, takes_row_as_input=True)
+ for arguments in (inputs, [proto.Input(udf=identity)],
[proto.Input(refIndex=0)]):
+ with self.subTest(input_kind=arguments[0].WhichOneof("input")):
+ shared = arguments[0].HasField("refIndex")
+ operation = self.operation(
+ increment, arguments, takes_row_as_input=True,
+ preceding=[identity] if shared else [], output_indices=[1]
if shared else [])
+ result = operation.process_element(batch)
+ self.assertEqual(result.num_columns, 1)
+ self.assertEqual(result.column(0).to_pylist(),
+ [{"name": "alice", "count": 2}, {"name":
None, "count": 3}])
+
+ def test_whole_row_argument_offsets(self):
+ batch = pa.record_batch([pa.array([0, 1, None, 3]), pa.array([0, 100,
200, 300])],
+ names=["a", "b"]).slice(1)
+ for offsets in ([0, 0, 1], [1, 0]):
+ with self.subTest(offsets=offsets):
+ operation = self.operation(
+ lambda row: row.field(1),
+ [proto.Input(inputOffset=offset) for offset in offsets],
+ takes_row_as_input=True)
+ result = operation.process_element(batch)
+ self.assertEqual(result.column(0).to_pylist(), [1, None, 3])
+
+ def test_invalid_scalar_results(self):
+ batch = pa.record_batch([pa.array([1, 2, 3])], names=["value"])
+ for result, error, message in (
+ ([1, 2, 3], TypeError, "Array or pyarrow.ChunkedArray"),
+ (None, TypeError, "NoneType"),
+ (pa.scalar(1), TypeError, "Scalar"),
+ (batch, TypeError, "RecordBatch"),
+ (pa.Table.from_batches([batch]), TypeError, "Table"),
+ (pa.array([1]), ValueError, "returned 1 rows, expected 3"),
+ (pa.chunked_array([[1], [2]]), ValueError, "returned 2 rows,
expected 3"),
+ ):
+ with self.subTest(result=result):
+ operation = self.operation(lambda values: result,
[proto.Input(inputOffset=0)])
+ with self.assertRaisesRegex(error, message):
+ operation.process_element(batch)
+
+ def test_invalid_intermediate_result_is_not_consumed(self):
+ inner = proto.UserDefinedFunction(
+ payload=cloudpickle.dumps(DelegatingScalarFunction(lambda values:
values.slice(0, 1))),
+ is_arrow_udf=True, inputs=[proto.Input(inputOffset=0)])
+ # The outer result has the correct batch length, but must not hide the
invalid inner result.
+ operation = self.operation(lambda values: pa.array([1, 2, 3]),
[proto.Input(udf=inner)])
+ with self.assertRaisesRegex(ValueError, "returned 1 rows, expected 3"):
+ operation.process_element(pa.record_batch([pa.array([1, 2, 3])],
names=["value"]))
+
+ def test_arrow_scalar_operation(self):
+ class Uppercase(ScalarFunction):
+ def eval(self, values):
+ return pc.utf8_upper(values)
+
+ function = proto.UserDefinedFunction(
+ payload=cloudpickle.dumps(Uppercase()), is_arrow_udf=True,
+ inputs=[proto.Input(inputOffset=0)])
+ operation = ScalarFunctionOperation(
+ proto.UserDefinedFunctions(udfs=[function]),
one_arg_optimization=True)
+ operation.open()
+ self.addCleanup(operation.close)
+ result = operation.process_element(pa.record_batch(
+ [pa.array(["alice", None, "Bob"])], names=["name"]))
+ self.assertIsInstance(result, pa.RecordBatch)
+ self.assertEqual(result.column(0).to_pylist(), ["ALICE", None, "BOB"])
+
+
class UserDefinedFunctionTests(object):
def test_scalar_function(self):
diff --git a/flink-python/pyflink/table/types.py
b/flink-python/pyflink/table/types.py
index 16d92c3c1ff..ac8397e007d 100644
--- a/flink-python/pyflink/table/types.py
+++ b/flink-python/pyflink/table/types.py
@@ -2213,12 +2213,16 @@ def _create_type_verifier(data_type: DataType, name:
str = None):
return verify
-def create_arrow_schema(field_names: List[str], field_types: List[DataType]):
+def create_arrow_schema(field_names: List[str], field_types: List[DataType],
*, allow_nested=False):
"""
- Create an Arrow schema with the specified filed names and types.
+ Create an Arrow schema with the specified field names and types.
+
+ By default, retain the nested-type restrictions of the pandas conversion
path.
+ Arrow-native transport can set ``allow_nested`` to include nested rows and
timestamps.
"""
import pyarrow as pa
- fields = [pa.field(field_name, to_arrow_type(field_type),
field_type._nullable)
+ fields = [pa.field(field_name, to_arrow_type(field_type,
allow_nested=allow_nested),
+ field_type._nullable)
for field_name, field_type in zip(field_names, field_types)]
return pa.schema(fields)
@@ -2296,7 +2300,7 @@ def from_arrow_type(arrow_type, nullable: bool = True) ->
DataType:
raise TypeError("Unsupported data type to convert from Arrow type: " +
str(arrow_type))
-def to_arrow_type(data_type: DataType):
+def to_arrow_type(data_type: DataType, *, allow_nested=False):
"""
Converts the specified Flink data type to pyarrow data type.
"""
@@ -2344,18 +2348,26 @@ def to_arrow_type(data_type: DataType):
else:
return pa.timestamp('ns')
elif isinstance(data_type, MapType):
- return pa.map_(to_arrow_type(data_type.key_type),
to_arrow_type(data_type.value_type))
+ key_type = to_arrow_type(data_type.key_type, allow_nested=allow_nested)
+ value_type = to_arrow_type(data_type.value_type,
allow_nested=allow_nested)
+ # PyArrow 5 only accepts data types and always makes map values
nullable.
+ if hasattr(pa.MapType, 'item_field'):
+ value_type = pa.field("value", value_type,
nullable=data_type.value_type._nullable)
+ return pa.map_(key_type, value_type)
elif isinstance(data_type, ArrayType):
- if type(data_type.element_type) in [LocalZonedTimestampType, RowType]:
+ if not allow_nested and type(data_type.element_type) in
[LocalZonedTimestampType, RowType]:
raise ValueError("%s is not supported to be used as the element
type of ArrayType." %
data_type.element_type)
- return pa.list_(to_arrow_type(data_type.element_type))
+ return pa.list_(pa.field(
+ "item", to_arrow_type(data_type.element_type,
allow_nested=allow_nested),
+ nullable=data_type.element_type._nullable))
elif isinstance(data_type, RowType):
for field in data_type:
- if type(field.data_type) in [LocalZonedTimestampType, RowType]:
+ if not allow_nested and type(field.data_type) in
[LocalZonedTimestampType, RowType]:
raise TypeError("%s is not supported to be used as the field
type of RowType" %
field.data_type)
- fields = [pa.field(field.name, to_arrow_type(field.data_type),
field.data_type._nullable)
+ fields = [pa.field(field.name, to_arrow_type(field.data_type,
allow_nested=allow_nested),
+ field.data_type._nullable)
for field in data_type]
return pa.struct(fields)
elif isinstance(data_type, NullType):
diff --git a/flink-python/pyflink/table/udf.py
b/flink-python/pyflink/table/udf.py
index 92bd180d189..733b8a79178 100644
--- a/flink-python/pyflink/table/udf.py
+++ b/flink-python/pyflink/table/udf.py
@@ -507,6 +507,8 @@ class UserDefinedFunctionWrapper(object):
return JPythonFunctionKind.GENERAL
elif self._func_type == "pandas":
return JPythonFunctionKind.PANDAS
+ elif self._func_type == "arrow":
+ return JPythonFunctionKind.ARROW
else:
raise TypeError("Unsupported func_type: %s." %
self._func_type)
@@ -760,10 +762,19 @@ def _get_python_env():
def _create_udf(f, input_types, result_type, func_type, deterministic, name):
+ if func_type in ('pandas', 'arrow'):
+ target = f
+ while isinstance(target, functools.partial):
+ target = target.func
+ if isinstance(target, ScalarFunction):
+ target = target.eval
+ if inspect.iscoroutinefunction(target) or inspect.iscoroutinefunction(
+ getattr(target, '__call__', None)):
+ raise ValueError(f"Async scalar functions do not support
{func_type} func_type.")
if isinstance(f, AsyncScalarFunction) or inspect.iscoroutinefunction(f):
- if func_type == 'pandas':
+ if 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.
"
"Please use func_type='general' (default) for async
functions.")
return UserDefinedAsyncScalarFunctionWrapper(
f, input_types, result_type, func_type, deterministic, name)
@@ -838,16 +849,15 @@ def udf(f: Union[Callable, ScalarFunction,
AsyncScalarFunction, Type] = None,
this function is guaranteed to always return the
same result given the
same parameters. (default True)
:param name: the function name.
- :param func_type: the type of the python function, available value:
general, pandas,
+ :param func_type: the type of the python function, available value:
general, pandas, arrow,
(default: general)
:return: UserDefinedScalarFunctionWrapper,
UserDefinedAsyncScalarFunctionWrapper, or function.
.. versionadded:: 1.10.0
"""
- if func_type not in ('general', 'pandas'):
- raise ValueError("The func_type must be one of 'general, pandas', got
%s."
- % func_type)
+ if func_type not in ('general', 'pandas', 'arrow'):
+ raise ValueError(f"The func_type must be one of 'general, pandas,
arrow', got {func_type}.")
# decorator
if f is None:
diff --git a/flink-python/pyflink/table/utils.py
b/flink-python/pyflink/table/utils.py
index bc4444e66eb..be1e71f64b2 100644
--- a/flink-python/pyflink/table/utils.py
+++ b/flink-python/pyflink/table/utils.py
@@ -16,6 +16,7 @@
# limitations under the License.
################################################################################
import ast
+from functools import cache
from pyflink.common.types import RowKind
@@ -28,6 +29,137 @@ import datetime
import pickle
+@Internal()
+def validate_arrow_batch(batch, schema, field_types):
+ import pyarrow as pa
+
+ if not isinstance(batch, pa.RecordBatch):
+ raise TypeError("Arrow transport expects a pyarrow.RecordBatch.")
+ if batch.num_columns != len(schema):
+ raise ValueError(f"Arrow result has {batch.num_columns} columns,
expected {len(schema)}.")
+ # PyArrow 5 cannot express map-value nullability, so retain the logical
constraint there.
+ legacy_map_fields = not hasattr(pa.MapType, 'item_field')
+ columns = []
+ for index, (column, field) in enumerate(zip(batch.columns, schema)):
+ _validate_array(column, field, field.name,
+ data_type=field_types[index] if legacy_map_fields else
None)
+ columns.append(_apply_arrow_type(column, field.type))
+ return pa.RecordBatch.from_arrays(columns, schema=schema)
+
+
+def _apply_arrow_type(column, expected_type):
+ import pyarrow as pa
+
+ if column.type == expected_type:
+ return column
+ # Old Arrow casts reject valid nullability changes, and views inspect
unsliced children.
+ # After validation, rebuild only container metadata while keeping the
payload buffers.
+ if pa.types.is_struct(expected_type):
+ return pa.StructArray.from_arrays(
+ [_apply_arrow_type(column.field(index), field.type)
+ for index, field in enumerate(expected_type)],
+ fields=list(expected_type), mask=column.is_null() if
column.null_count else None)
+ if pa.types.is_list(expected_type):
+ children = [_apply_arrow_type(column.values, expected_type.value_type)]
+ else:
+ value_field = (expected_type.item_field if hasattr(expected_type,
'item_field')
+ else pa.field("value", expected_type.item_type))
+ children = [pa.StructArray.from_arrays(
+ [_apply_arrow_type(column.keys, expected_type.key_type),
+ _apply_arrow_type(column.items, expected_type.item_type)],
+ fields=[pa.field("key", expected_type.key_type, nullable=False),
value_field])]
+ return pa.Array.from_buffers(
+ expected_type, len(column),
column.buffers()[:expected_type.num_buffers],
+ null_count=column.null_count, offset=column.offset, children=children)
+
+
+def _validate_array(column, field, path, parent_validity=None, data_type=None):
+ import pyarrow as pa
+ import pyarrow.compute as pc
+ expected_type = field.type
+
+ def wrong_type():
+ raise TypeError(
+ f"Arrow result field '{path}' has type {column.type}, expected
{expected_type}.")
+
+ if not field.nullable and column.null_count:
+ validity = parent_validity() if parent_validity is not None else None
+ if validity is None or pc.any(pc.and_(validity,
column.is_null())).as_py():
+ raise ValueError(f"Arrow result field '{path}' is not nullable.")
+ if pa.types.is_struct(expected_type):
+ if (not pa.types.is_struct(column.type)
+ or column.type.num_fields != expected_type.num_fields):
+ wrong_type()
+ validity = cache(lambda: _get_validity(column, parent_validity))
+ for index, child_field in enumerate(expected_type):
+ if column.type[index].name != child_field.name:
+ wrong_type()
+ _validate_array(column.field(index), child_field,
+ f"{path}.{child_field.name}", validity,
+ data_type.fields[index].data_type if data_type is
not None else None)
+ elif pa.types.is_list(expected_type):
+ if not pa.types.is_list(column.type):
+ wrong_type()
+ start, end = column.offsets[0].as_py(), column.offsets[-1].as_py()
+ _validate_array(column.values.slice(start, end - start),
+ expected_type.value_field, f"{path}[]",
+ cache(lambda: _get_child_validity(column,
parent_validity)),
+ data_type.element_type if data_type is not None else
None)
+ elif pa.types.is_map(expected_type):
+ if not pa.types.is_map(column.type):
+ wrong_type()
+ validity = cache(lambda: _get_child_validity(column, parent_validity))
+ offsets = _get_offsets(column)
+ start, end = offsets[0].as_py(), offsets[-1].as_py()
+ key_field = (expected_type.key_field if data_type is None
+ else pa.field("key", expected_type.key_type,
nullable=False))
+ value_field = (expected_type.item_field if data_type is None
+ else pa.field("value", expected_type.item_type,
+ nullable=data_type.value_type._nullable))
+ _validate_array(column.keys.slice(start, end - start), key_field,
+ f"{path}.key", validity,
+ data_type.key_type if data_type is not None else None)
+ _validate_array(column.items.slice(start, end - start), value_field,
+ f"{path}.value", validity,
+ data_type.value_type if data_type is not None else
None)
+ elif column.type != expected_type:
+ wrong_type()
+
+
+def _get_validity(column, parent_validity):
+ import pyarrow.compute as pc
+
+ # Resolve ancestor visibility only when a NOT NULL descendant contains
physical nulls.
+ parent_validity = parent_validity() if parent_validity is not None else
None
+ if column.null_count:
+ validity = column.is_valid()
+ return validity if parent_validity is None else
pc.and_(parent_validity, validity)
+ return parent_validity
+
+
+def _get_child_validity(column, parent_validity):
+ import numpy as np
+ import pyarrow as pa
+
+ validity = _get_validity(column, parent_validity)
+ if validity is None:
+ return None
+ # Expand visibility directly instead of allocating an integer parent index
for every child.
+ offsets = _get_offsets(column).to_numpy(zero_copy_only=True)
+ visible = validity.to_numpy(zero_copy_only=False)
+ return pa.array(np.repeat(visible, np.diff(offsets)), type=pa.bool_())
+
+
+def _get_offsets(column):
+ import pyarrow as pa
+
+ if hasattr(column, 'offsets'):
+ return column.offsets
+ # PyArrow 5 does not expose MapArray.offsets, but uses the same int32
offsets as lists.
+ return pa.Array.from_buffers(pa.int32(), len(column) + 1,
+ [None, column.buffers()[1]],
offset=column.offset)
+
+
@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 a4cee4f9ebc..6da72d75bfc 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
@@ -109,11 +109,24 @@ public enum ProtoUtils {
RowType rowType,
FlinkFnApi.CoderInfoDescriptor.Mode mode,
boolean separatedWithEndMessage) {
+ return createArrowTypeCoderInfoDescriptorProto(
+ rowType,
+ mode,
+ separatedWithEndMessage,
+ FlinkFnApi.CoderInfoDescriptor.ArrowType.BatchFormat.PANDAS);
+ }
+
+ public static FlinkFnApi.CoderInfoDescriptor
createArrowTypeCoderInfoDescriptorProto(
+ RowType rowType,
+ FlinkFnApi.CoderInfoDescriptor.Mode mode,
+ boolean separatedWithEndMessage,
+ FlinkFnApi.CoderInfoDescriptor.ArrowType.BatchFormat batchFormat) {
return createCoderInfoDescriptorProto(
null,
null,
FlinkFnApi.CoderInfoDescriptor.ArrowType.newBuilder()
.setSchema(toProtoType(rowType).getRowSchema())
+ .setBatchFormat(batchFormat)
.build(),
null,
null,
@@ -229,6 +242,9 @@ public enum ProtoUtils {
builder.setIsPandasUdf(
pythonFunctionInfo.getPythonFunction().getPythonFunctionKind()
== PythonFunctionKind.PANDAS);
+ builder.setIsArrowUdf(
+ pythonFunctionInfo.getPythonFunction().getPythonFunctionKind()
+ == PythonFunctionKind.ARROW);
return builder.build();
}
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 ef5bd046390..bee19392571 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
@@ -25,6 +25,7 @@ import org.apache.flink.fnexecution.v1.FlinkFnApi;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.functions.ScalarFunction;
import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.functions.python.PythonFunctionKind;
import org.apache.flink.table.runtime.arrow.serializers.ArrowSerializer;
import org.apache.flink.table.runtime.generated.GeneratedProjection;
import
org.apache.flink.table.runtime.operators.python.scalar.AbstractPythonScalarFunctionOperator;
@@ -49,6 +50,8 @@ public class ArrowPythonScalarFunctionOperator extends
AbstractPythonScalarFunct
private transient ArrowSerializer arrowSerializer;
+ private final FlinkFnApi.CoderInfoDescriptor.ArrowType.BatchFormat
batchFormat;
+
public ArrowPythonScalarFunctionOperator(
Configuration config,
PythonFunctionInfo[] scalarFunctions,
@@ -86,6 +89,11 @@ public class ArrowPythonScalarFunctionOperator extends
AbstractPythonScalarFunct
udfOutputType,
udfInputGeneratedProjection,
forwardedFieldGeneratedProjection);
+ batchFormat =
+ scalarFunctions[0].getPythonFunction().getPythonFunctionKind()
+ == PythonFunctionKind.ARROW
+ ?
FlinkFnApi.CoderInfoDescriptor.ArrowType.BatchFormat.ARROW
+ :
FlinkFnApi.CoderInfoDescriptor.ArrowType.BatchFormat.PANDAS;
}
@Override
@@ -100,14 +108,14 @@ public class ArrowPythonScalarFunctionOperator extends
AbstractPythonScalarFunct
@Override
public FlinkFnApi.CoderInfoDescriptor
createInputCoderInfoDescriptor(RowType runnerInputType) {
return createArrowTypeCoderInfoDescriptorProto(
- runnerInputType, FlinkFnApi.CoderInfoDescriptor.Mode.MULTIPLE,
false);
+ runnerInputType, FlinkFnApi.CoderInfoDescriptor.Mode.MULTIPLE,
false, batchFormat);
}
@Override
public FlinkFnApi.CoderInfoDescriptor createOutputCoderInfoDescriptor(
RowType runnerOutputType) {
return createArrowTypeCoderInfoDescriptorProto(
- runnerOutputType, FlinkFnApi.CoderInfoDescriptor.Mode.SINGLE,
false);
+ runnerOutputType, FlinkFnApi.CoderInfoDescriptor.Mode.SINGLE,
false, batchFormat);
}
@Override
diff --git
a/flink-python/src/test/java/org/apache/flink/streaming/api/utils/ProtoUtilsTest.java
b/flink-python/src/test/java/org/apache/flink/streaming/api/utils/ProtoUtilsTest.java
index 78a10489b31..1331226b200 100644
---
a/flink-python/src/test/java/org/apache/flink/streaming/api/utils/ProtoUtilsTest.java
+++
b/flink-python/src/test/java/org/apache/flink/streaming/api/utils/ProtoUtilsTest.java
@@ -21,8 +21,18 @@ package org.apache.flink.streaming.api.utils;
import org.apache.flink.api.common.state.StateTtlConfig;
import org.apache.flink.fnexecution.v1.FlinkFnApi;
import org.apache.flink.python.util.ProtoUtils;
+import org.apache.flink.table.api.DataTypes;
+import org.apache.flink.table.functions.python.InputRef;
+import org.apache.flink.table.functions.python.PythonEnv;
+import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.functions.python.PythonFunctionInput;
+import org.apache.flink.table.functions.python.PythonFunctionKind;
+import org.apache.flink.table.functions.python.PythonScalarFunction;
+import org.apache.flink.table.types.logical.RowType;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
import java.time.Duration;
@@ -33,6 +43,47 @@ import static org.assertj.core.api.Assertions.assertThat;
* protobuf objects.
*/
class ProtoUtilsTest {
+ @Test
+ void testArrowCoderBatchFormat() {
+ final RowType rowType = RowType.of(DataTypes.INT().getLogicalType());
+
assertThat(FlinkFnApi.CoderInfoDescriptor.ArrowType.getDefaultInstance().getBatchFormat())
+
.isEqualTo(FlinkFnApi.CoderInfoDescriptor.ArrowType.BatchFormat.PANDAS);
+ assertThat(
+ ProtoUtils.createArrowTypeCoderInfoDescriptorProto(
+ rowType,
FlinkFnApi.CoderInfoDescriptor.Mode.SINGLE, false)
+ .getArrowType()
+ .getBatchFormat())
+
.isEqualTo(FlinkFnApi.CoderInfoDescriptor.ArrowType.BatchFormat.PANDAS);
+ assertThat(
+ ProtoUtils.createArrowTypeCoderInfoDescriptorProto(
+ rowType,
+
FlinkFnApi.CoderInfoDescriptor.Mode.SINGLE,
+ false,
+
FlinkFnApi.CoderInfoDescriptor.ArrowType.BatchFormat.ARROW)
+ .getArrowType()
+ .getBatchFormat())
+
.isEqualTo(FlinkFnApi.CoderInfoDescriptor.ArrowType.BatchFormat.ARROW);
+ }
+
+ @ParameterizedTest
+ @EnumSource(PythonFunctionKind.class)
+ void testScalarFunctionKind(PythonFunctionKind kind) {
+ final PythonScalarFunction function =
+ new PythonScalarFunction(
+ "identity",
+ new byte[0],
+ kind,
+ true,
+ false,
+ new PythonEnv(PythonEnv.ExecType.PROCESS));
+ final FlinkFnApi.UserDefinedFunction proto =
+ ProtoUtils.createUserDefinedFunctionProto(
+ new PythonFunctionInfo(
+ function, new PythonFunctionInput[] {new
InputRef(0)}));
+ assertThat(proto.getIsArrowUdf()).isEqualTo(kind ==
PythonFunctionKind.ARROW);
+ assertThat(proto.getIsPandasUdf()).isEqualTo(kind ==
PythonFunctionKind.PANDAS);
+ }
+
@Test
void testParseStateTtlConfigFromProto() {
FlinkFnApi.StateDescriptor.StateTTLConfig.CleanupStrategies
cleanupStrategiesProto =
diff --git
a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunctionKind.java
b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunctionKind.java
index cd7a27ef8a8..e8afcdfaaa1 100644
---
a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunctionKind.java
+++
b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/python/PythonFunctionKind.java
@@ -25,5 +25,7 @@ import org.apache.flink.annotation.Internal;
public enum PythonFunctionKind {
GENERAL,
- PANDAS
+ PANDAS,
+
+ ARROW
}
diff --git
a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonCalc.java
b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonCalc.java
index b0fcd6e463a..8a15da8d3c5 100644
---
a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonCalc.java
+++
b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/nodes/exec/common/CommonExecPythonCalc.java
@@ -28,8 +28,11 @@ import
org.apache.flink.streaming.api.transformations.OneInputTransformation;
import org.apache.flink.table.api.TableException;
import org.apache.flink.table.connector.Projection;
import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.functions.python.InputRef;
import org.apache.flink.table.functions.python.PythonFunctionInfo;
+import org.apache.flink.table.functions.python.PythonFunctionInput;
import org.apache.flink.table.functions.python.PythonFunctionKind;
+import org.apache.flink.table.functions.python.ResultRef;
import org.apache.flink.table.planner.calcite.FlinkTypeFactory;
import org.apache.flink.table.planner.codegen.CodeGeneratorContext;
import org.apache.flink.table.planner.codegen.ProjectionCodeGenerator;
@@ -156,7 +159,10 @@ public abstract class CommonExecPythonCalc extends
ExecNodeBase<RowData>
extractPythonScalarFunctionInfos(cseResult, classLoader);
int[] pythonUdfInputOffsets = extractResult.f0;
PythonFunctionInfo[] pythonFunctionInfos = extractResult.f1;
-
+ boolean[] hasColumn = new boolean[pythonFunctionInfos.length];
+ for (int i = 0; i < pythonFunctionInfos.length; i++) {
+ hasColumn[i] = validateArrowInputs(pythonFunctionInfos[i],
hasColumn);
+ }
LogicalType[] inputLogicalTypes =
((InternalTypeInfo<RowData>)
inputTransform.getOutputType()).toRowFieldTypes();
InternalTypeInfo<RowData> pythonOperatorInputTypeInfo =
@@ -192,7 +198,9 @@ public abstract class CommonExecPythonCalc extends
ExecNodeBase<RowData>
.anyMatch(
x ->
PythonUtil.containsPythonCall(
- x,
PythonFunctionKind.PANDAS)));
+ x,
PythonFunctionKind.PANDAS)
+ ||
PythonUtil.containsPythonCall(
+ x,
PythonFunctionKind.ARROW)));
return ExecNodeUtil.createOneInputTransformation(
inputTransform,
@@ -203,6 +211,26 @@ public abstract class CommonExecPythonCalc extends
ExecNodeBase<RowData>
false);
}
+ private static boolean validateArrowInputs(
+ PythonFunctionInfo function, boolean[] resultHasColumn) {
+ boolean hasColumn = false;
+ for (PythonFunctionInput input : function.getInputs()) {
+ if (input instanceof PythonFunctionInfo) {
+ hasColumn |= validateArrowInputs((PythonFunctionInfo) input,
resultHasColumn);
+ } else if (input instanceof InputRef) {
+ hasColumn = true;
+ } else if (input instanceof ResultRef) {
+ hasColumn |= resultHasColumn[((ResultRef) input).getIndex()];
+ }
+ }
+ if (function.getPythonFunction().getPythonFunctionKind() ==
PythonFunctionKind.ARROW
+ && !hasColumn) {
+ throw new TableException(
+ "Arrow scalar UDFs require at least one column-valued
argument.");
+ }
+ return hasColumn;
+ }
+
private Tuple2<int[], PythonFunctionInfo[]>
extractPythonScalarFunctionInfos(
PythonCallCseResult cseResult, ClassLoader classLoader) {
List<RexCall> rexCalls = cseResult.getDeduplicatedCalls();
diff --git
a/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRule.java
b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRule.java
new file mode 100644
index 00000000000..0e42290ed70
--- /dev/null
+++
b/flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRule.java
@@ -0,0 +1,75 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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.PythonUtil;
+
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexProgram;
+
+import java.util.Arrays;
+
+import scala.Option;
+
+/** Separates scalar Python functions whose user-facing batch representations
differ. */
+public class PythonCalcSplitFunctionKindRule extends
RemoteCalcSplitProjectionRuleBase<Void> {
+
+ public PythonCalcSplitFunctionKindRule(RemoteCallFinder callFinder) {
+ super("PythonCalcSplitFunctionKindRule", callFinder);
+ }
+
+ @Override
+ public boolean matches(RelOptRuleCall call) {
+ final FlinkLogicalCalc calc = call.rel(0);
+ final RexProgram program = calc.getProgram();
+ return Arrays.stream(PythonFunctionKind.values())
+ .filter(
+ kind ->
+ program.getProjectList().stream()
+ .map(program::expandLocalRef)
+ .anyMatch(
+ node ->
+
PythonUtil.containsPythonCall(
+ node,
kind)))
+ .count()
+ > 1;
+ }
+
+ @Override
+ public boolean needConvert(RexProgram program, RexNode node, Option<Void>
matchState) {
+ // Keep one top-level kind above the split and extract calls of every
other kind.
+ final PythonFunctionKind topLevelKind =
+ Arrays.stream(PythonFunctionKind.values())
+ .filter(
+ kind ->
+ program.getProjectList().stream()
+ .map(program::expandLocalRef)
+ .anyMatch(
+ project ->
+
PythonUtil.isPythonCall(
+
project, kind)))
+ .findFirst()
+ .orElseThrow(
+ () -> new IllegalStateException("Missing
top-level Python call."));
+ return PythonUtil.isPythonCall(node) && !PythonUtil.isPythonCall(node,
topLevelKind);
+ }
+}
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 57d2f2cea68..c93c643301e 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
@@ -32,6 +32,7 @@ import org.apache.calcite.rex.RexNode;
import org.apache.calcite.rex.RexProgram;
import org.apache.calcite.rex.RexProgramBuilder;
+import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;
@@ -87,8 +88,11 @@ public class PythonMapMergeRule extends RelOptRule {
}
// only Python functions with same Python function kind can be merged
together.
- if (PythonUtil.isPythonCall(topProjects.get(0),
PythonFunctionKind.GENERAL)
- ^ PythonUtil.isPythonCall(bottomProjects.get(0),
PythonFunctionKind.GENERAL)) {
+ if (Arrays.stream(PythonFunctionKind.values())
+ .noneMatch(
+ kind ->
+ PythonUtil.isPythonCall(topProjects.get(0),
kind)
+ &&
PythonUtil.isPythonCall(bottomProjects.get(0), kind))) {
return false;
}
diff --git
a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRule.scala
b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRule.scala
index 642ca5f5691..7983ea96a1a 100644
---
a/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRule.scala
+++
b/flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitRule.scala
@@ -17,53 +17,10 @@
*/
package org.apache.flink.table.planner.plan.rules.logical
-import org.apache.flink.table.functions.ScalarFunction
-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.{InputRefVisitor, PythonUtil,
RexDefaultVisitor}
-import
org.apache.flink.table.planner.plan.utils.PythonUtil.{containsNonPythonCall,
containsPythonCall, isNonPythonCall, isPythonCall}
+import org.apache.flink.table.planner.plan.utils.PythonUtil
-import org.apache.calcite.plan.{RelOptRule, RelOptRuleCall}
-import org.apache.calcite.plan.RelOptRule.{any, operand}
-import org.apache.calcite.rex.{RexBuilder, RexCall, RexCorrelVariable,
RexFieldAccess, RexInputRef, RexLocalRef, RexNode, RexProgram}
-import org.apache.calcite.sql.validate.SqlValidatorUtil
-
-import java.util.function.Function
-
-import scala.collection.JavaConversions._
-import scala.collection.JavaConverters._
-import scala.collection.mutable
-
-/**
- * Rule that splits [[FlinkLogicalCalc]]s which contain both general Python
functions and pandas
- * Python functions in the projection into multiple [[FlinkLogicalCalc]]s.
After this rule is
- * applied, it will only contain general Python functions or pandas Python
functions in the
- * projection of each [[FlinkLogicalCalc]].
- */
-class PythonCalcSplitPandasInProjectionRule(callFinder: RemoteCallFinder)
- extends
RemoteCalcSplitProjectionRuleBase("PythonCalcSplitPandasInProjectionRule",
callFinder) {
-
- override def matches(call: RelOptRuleCall): Boolean = {
- val calc: FlinkLogicalCalc = call.rel(0).asInstanceOf[FlinkLogicalCalc]
- val projects =
calc.getProgram.getProjectList.map(calc.getProgram.expandLocalRef)
-
- // matches if it contains both general Python functions and
- // pandas Python functions in the projection
- projects.exists(containsPythonCall(_, PythonFunctionKind.GENERAL)) &&
- projects.exists(containsPythonCall(_, PythonFunctionKind.PANDAS))
- }
-
- override def needConvert(
- program: RexProgram,
- node: RexNode,
- matchState: Option[Nothing]): Boolean = {
- program.getProjectList
- .map(program.expandLocalRef)
- .exists(isPythonCall(_, PythonFunctionKind.GENERAL)) == isPythonCall(
- node,
- PythonFunctionKind.PANDAS)
- }
-}
+import org.apache.calcite.plan.RelOptRule
+import org.apache.calcite.rex.RexNode
class PythonRemoteCallFinder extends RemoteCallFinder {
override def containsRemoteCall(node: RexNode): Boolean = {
@@ -105,7 +62,7 @@ object PythonCalcSplitRule {
val CONDITION_PROJECTION_CSE: RelOptRule =
RemoteCalcConditionProjectionCseRule.Config.DEFAULT.withRemoteCallFinder(callFinder).toRule()
val SPLIT_PROJECT: RelOptRule = new RemoteCalcSplitProjectionRule(callFinder)
- val SPLIT_PANDAS_IN_PROJECT: RelOptRule = new
PythonCalcSplitPandasInProjectionRule(callFinder)
+ val SPLIT_PANDAS_IN_PROJECT: RelOptRule = new
PythonCalcSplitFunctionKindRule(callFinder)
val SPLIT_PROJECTION_REX_FIELD: RelOptRule = new
RemoteCalcSplitProjectionRexFieldRule(callFinder)
val SPLIT_CONDITION_REX_FIELD: RelOptRule = new
RemoteCalcSplitConditionRexFieldRule(callFinder)
val EXPAND_PROJECT: RelOptRule = new RemoteCalcExpandProjectRule(callFinder)
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
new file mode 100644
index 00000000000..09cc9d10dd8
--- /dev/null
+++
b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRuleTest.java
@@ -0,0 +1,112 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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.PythonEnv;
+import org.apache.flink.table.functions.python.PythonFunctionKind;
+import org.apache.flink.table.functions.python.PythonScalarFunction;
+import org.apache.flink.table.planner.utils.JavaScalaConversionUtil;
+import org.apache.flink.table.planner.utils.JavaTableTestUtil;
+import org.apache.flink.table.planner.utils.TableTestBase;
+import org.apache.flink.table.types.DataType;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+
+import static org.apache.flink.table.api.Expressions.$;
+import static org.apache.flink.table.api.Expressions.call;
+
+/** Plans compositions of scalar UDFs with different Python argument
representations. */
+class PythonCalcSplitFunctionKindRuleTest extends TableTestBase {
+
+ @Test
+ void testStreamingComposition() {
+ verifyComposition(javaStreamTestUtil());
+ }
+
+ @Test
+ void testBatchComposition() {
+ verifyComposition(javaBatchTestUtil());
+ }
+
+ @Test
+ void testStreamingMapComposition() {
+ verifyMapComposition(javaStreamTestUtil());
+ }
+
+ @Test
+ void testBatchMapComposition() {
+ verifyMapComposition(javaBatchTestUtil());
+ }
+
+ private void verifyComposition(JavaTableTestUtil util) {
+ util.addTableSource("T", Schema.newBuilder().column("a",
DataTypes.INT()).build());
+ registerFunctions(util, DataTypes.INT(), false);
+ util.verifyExecPlan(
+ "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");
+ }
+
+ private void verifyMapComposition(JavaTableTestUtil util) {
+ util.addTableSource("T", Schema.newBuilder().column("a",
DataTypes.INT()).build());
+ registerFunctions(util, DataTypes.ROW(DataTypes.FIELD("a",
DataTypes.INT())), true);
+ final Table result =
+ util.tableEnv()
+ .from("T")
+ .map(call("general_udf", $("a")))
+ .map(call("arrow_udf", $("a")))
+ .map(call("arrow_udf", $("a")))
+ .map(call("pandas_udf", $("a")))
+ .map(call("pandas_udf", $("a")))
+ .map(call("arrow_udf", $("a")))
+ .map(call("general_udf", $("a")));
+ util.verifyRelPlanExpected(
+ result,
+ JavaScalaConversionUtil.toScala(
+ Arrays.asList(
+ "PythonCalc(select=[arrow_udf(arrow_udf(a)) AS
f0])",
+ "PythonCalc(select=[pandas_udf(pandas_udf(a))
AS f0])",
+ "PythonCalc(select=[arrow_udf(a) AS f0])",
+ "PythonCalc(select=[general_udf(a) AS f0])")));
+ }
+
+ private void registerFunctions(
+ JavaTableTestUtil util, DataType resultType, boolean
takesRowAsInput) {
+ for (PythonFunctionKind kind : PythonFunctionKind.values()) {
+ util.tableEnv()
+ .createTemporarySystemFunction(
+ kind.name().toLowerCase() + "_udf",
+ new PythonScalarFunction(
+ kind.name(),
+ new byte[0],
+ new DataType[] {DataTypes.INT()},
+ resultType,
+ kind,
+ true,
+ takesRowAsInput,
+ new
PythonEnv(PythonEnv.ExecType.PROCESS)));
+ }
+ }
+}
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
new file mode 100644
index 00000000000..021138e012a
--- /dev/null
+++
b/flink-table/flink-table-planner/src/test/resources/org/apache/flink/table/planner/plan/rules/logical/PythonCalcSplitFunctionKindRuleTest.xml
@@ -0,0 +1,67 @@
+<?xml version="1.0" ?>
+<!--
+Licensed to the Apache Software Foundation (ASF) under one or more
+contributor license agreements. See the NOTICE file distributed with
+this work for additional information regarding copyright ownership.
+The ASF licenses this file to you under the Apache License, Version 2.0
+(the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+-->
+<Root>
+ <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]]>
+ </Resource>
+ <Resource name="ast">
+ <![CDATA[
+LogicalProject(EXPR$0=[arrow_udf($0)], EXPR$1=[pandas_udf($0)],
EXPR$2=[general_udf($0)], EXPR$3=[arrow_udf(arrow_udf($0))],
EXPR$4=[arrow_udf(pandas_udf($0))], EXPR$5=[pandas_udf(arrow_udf($0))],
EXPR$6=[arrow_udf(general_udf($0))], EXPR$7=[general_udf(arrow_udf($0))])
++- LogicalTableScan(table=[[default_catalog, default_database, T]])
+]]>
+ </Resource>
+ <Resource name="optimized exec plan">
+ <![CDATA[
+Calc(select=[f0 AS EXPR$0, f1 AS EXPR$1, f00 AS EXPR$2, f2 AS EXPR$3, f3 AS
EXPR$4, f4 AS EXPR$5, f5 AS EXPR$6, f10 AS EXPR$7])
++- PythonCalc(select=[f0, f1, f2, f3, f4, f5, general_udf(a) AS f00,
general_udf(f6) AS f10])
+ +- Calc(select=[a, f0, f00 AS f1, f1 AS f2, f2 AS f3, f10 AS f4, f4 AS f5,
f5 AS f6])
+ +- PythonCalc(select=[a, f0, f1, f2, f4, f5, pandas_udf(a) AS f00,
pandas_udf(f3) AS f10])
+ +- Calc(select=[f0 AS a, f1 AS f0, f2 AS f1, f3 AS f2, f1 AS f3, f4,
f1 AS f5])
+ +- PythonCalc(select=[a AS f0, arrow_udf(a) AS f1,
arrow_udf(arrow_udf(a)) AS f2, arrow_udf(f0) AS f3, arrow_udf(f1) AS f4])
+ +- 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="testStreamingComposition">
+ <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]]>
+ </Resource>
+ <Resource name="ast">
+ <![CDATA[
+LogicalProject(EXPR$0=[arrow_udf($0)], EXPR$1=[pandas_udf($0)],
EXPR$2=[general_udf($0)], EXPR$3=[arrow_udf(arrow_udf($0))],
EXPR$4=[arrow_udf(pandas_udf($0))], EXPR$5=[pandas_udf(arrow_udf($0))],
EXPR$6=[arrow_udf(general_udf($0))], EXPR$7=[general_udf(arrow_udf($0))])
++- LogicalTableScan(table=[[default_catalog, default_database, T]])
+]]>
+ </Resource>
+ <Resource name="optimized exec plan">
+ <![CDATA[
+Calc(select=[f0 AS EXPR$0, f1 AS EXPR$1, f00 AS EXPR$2, f2 AS EXPR$3, f3 AS
EXPR$4, f4 AS EXPR$5, f5 AS EXPR$6, f10 AS EXPR$7])
++- PythonCalc(select=[f0, f1, f2, f3, f4, f5, general_udf(a) AS f00,
general_udf(f6) AS f10])
+ +- Calc(select=[a, f0, f00 AS f1, f1 AS f2, f2 AS f3, f10 AS f4, f4 AS f5,
f5 AS f6])
+ +- PythonCalc(select=[a, f0, f1, f2, f4, f5, pandas_udf(a) AS f00,
pandas_udf(f3) AS f10])
+ +- Calc(select=[f0 AS a, f1 AS f0, f2 AS f1, f3 AS f2, f1 AS f3, f4,
f1 AS f5])
+ +- PythonCalc(select=[a AS f0, arrow_udf(a) AS f1,
arrow_udf(arrow_udf(a)) AS f2, arrow_udf(f0) AS f3, arrow_udf(f1) AS f4])
+ +- 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>
+</Root>
