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 b7a2e6d38bf [FLINK-40608][python] Extract type inference mechanism to
shared component (#29267)
b7a2e6d38bf is described below
commit b7a2e6d38bf1778d4d30836f505a8621175ded5f
Author: Mika Naylor <[email protected]>
AuthorDate: Fri Oct 9 03:59:58 2026 +0200
[FLINK-40608][python] Extract type inference mechanism to shared component
(#29267)
---
flink-python/pyflink/dataframe/datatype.py | 78 +-----
.../pyflink/dataframe/tests/test_datatype.py | 25 +-
flink-python/pyflink/dataframe/tests/test_sql.py | 14 +-
flink-python/pyflink/dataframe/tests/test_udf.py | 53 ++--
flink-python/pyflink/dataframe/tests/test_udtf.py | 17 +-
flink-python/pyflink/dataframe/udf.py | 35 +--
flink-python/pyflink/dataframe/udtf.py | 5 +-
flink-python/pyflink/table/tests/test_typehints.py | 276 +++++++++++++++++++++
flink-python/pyflink/table/typehints.py | 194 +++++++++++++++
9 files changed, 551 insertions(+), 146 deletions(-)
diff --git a/flink-python/pyflink/dataframe/datatype.py
b/flink-python/pyflink/dataframe/datatype.py
index d604f2c7fda..9e7802b0803 100644
--- a/flink-python/pyflink/dataframe/datatype.py
+++ b/flink-python/pyflink/dataframe/datatype.py
@@ -16,13 +16,10 @@
# limitations under the License.
################################################################################
-import datetime
-import decimal
-import types
-from functools import partial
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union,
get_args, get_origin
+from typing import Any, Dict, List, Tuple, Union
from pyflink.table.types import DataType as TableDataType, DataTypes, NullType
+from pyflink.table.typehints import _from_python_type
from pyflink.util.api_stability_decorators import PublicEvolving
__all__ = ["DataType"]
@@ -31,21 +28,6 @@ _INT_MIN = -(1 << 31)
_INT_MAX = (1 << 31) - 1
_BIGINT_MIN = -(1 << 63)
_BIGINT_MAX = (1 << 63) - 1
-_PEP_604_UNION_TYPE = getattr(types, "UnionType", None)
-
-_BASIC_TYPE_HINT_FACTORIES: Dict[Any, Callable[[], TableDataType]] = {
- bool: DataTypes.BOOLEAN,
- int: DataTypes.BIGINT,
- float: DataTypes.DOUBLE,
- str: DataTypes.STRING,
- bytes: DataTypes.BYTES,
- bytearray: DataTypes.BYTES,
- decimal.Decimal: partial(DataTypes.DECIMAL, 38, 18),
- datetime.date: DataTypes.DATE,
- datetime.time: DataTypes.TIME,
- datetime.datetime: DataTypes.TIMESTAMP,
- Any: DataTypes.STRING,
-}
@PublicEvolving()
@@ -382,61 +364,7 @@ class DataType:
@classmethod
def _from_type_hint(cls, type_hint: Any) -> "DataType":
- def infer_union_type(hint: Any, arguments: Tuple[Any, ...]) ->
"DataType":
- non_none_types = [
- argument for argument in arguments if argument is not
type(None)
- ]
- if len(non_none_types) == 1:
- return infer(non_none_types[0]).nullable()
-
- raise TypeError(
- f"Cannot infer DataType from type hint '{hint}'. "
- "Please specify the data type explicitly."
- )
-
- def infer_basic_type(hint: Any) -> Optional["DataType"]:
- factory = _BASIC_TYPE_HINT_FACTORIES.get(hint)
- return cls(factory()) if factory is not None else None
-
- def infer(hint: Any) -> "DataType":
- origin = get_origin(hint)
- arguments = get_args(hint)
-
- if origin is Union or (
- _PEP_604_UNION_TYPE is not None
- and origin is _PEP_604_UNION_TYPE
- ):
- return infer_union_type(hint, arguments)
-
- if origin is list:
- if not arguments:
- raise TypeError(
- "Cannot infer DataType from list without type
argument. "
- "Use list[T], for example list[int]."
- )
- return cls.list(infer(arguments[0]))
-
- if origin is dict:
- if len(arguments) != 2:
- raise TypeError(
- "Cannot infer DataType from dict without key and value
type arguments. "
- "Use dict[K, V], for example dict[str, int]."
- )
- return cls.map(
- infer(arguments[0]),
- infer(arguments[1]),
- )
-
- data_type = infer_basic_type(hint)
- if data_type is not None:
- return data_type
-
- raise TypeError(
- f"Cannot infer DataType from type hint '{hint}'. "
- "Please specify the data type explicitly."
- )
-
- return infer(type_hint)
+ return cls(_from_python_type(type_hint))
@classmethod
def _from_sql(cls, sql_type: str) -> "DataType":
diff --git a/flink-python/pyflink/dataframe/tests/test_datatype.py
b/flink-python/pyflink/dataframe/tests/test_datatype.py
index f129f9a5eb6..9d0eed98119 100644
--- a/flink-python/pyflink/dataframe/tests/test_datatype.py
+++ b/flink-python/pyflink/dataframe/tests/test_datatype.py
@@ -240,19 +240,28 @@ class DataTypeTests(unittest.TestCase):
bytearray: pf.DataType.binary(),
decimal.Decimal: pf.DataType.decimal(38, 18),
datetime.date: pf.DataType.date(),
- datetime.time: pf.DataType.time(),
+ datetime.time: pf.DataType.time(3),
datetime.datetime: pf.DataType.timestamp(),
- Any: pf.DataType.string(),
}
for python_type, expected_type in expected_types.items():
with self.subTest(python_type=python_type):
self.assertEqual(
pf.DataType._from_type_hint(python_type),
- expected_type,
+ expected_type.not_null(),
)
+ # `Any` opts out of the type system, so it stays nullable STRING.
+ self.assertEqual(
+ pf.DataType._from_type_hint(Any),
+ pf.DataType.string(),
+ )
+
def test_from_optional_type_hint(self):
+ self.assertEqual(
+ pf.DataType._from_type_hint(int),
+ pf.DataType.int64().not_null(),
+ )
self.assertEqual(
pf.DataType._from_type_hint(Optional[int]),
pf.DataType.int64(),
@@ -269,7 +278,7 @@ class DataTypeTests(unittest.TestCase):
)
self.assertEqual(
pf.DataType._from_type_hint(list[int | None]),
- pf.DataType.list(dtype=pf.DataType.int64()),
+ pf.DataType.list(dtype=pf.DataType.int64()).not_null(),
)
with self.assertRaises(TypeError):
pf.DataType._from_type_hint(int | str)
@@ -277,16 +286,16 @@ class DataTypeTests(unittest.TestCase):
def test_from_list_type_hint(self):
self.assertEqual(
pf.DataType._from_type_hint(list[int]),
- pf.DataType.list(dtype=pf.DataType.int64()),
+ pf.DataType.list(dtype=pf.DataType.int64().not_null()).not_null(),
)
def test_from_dict_type_hint(self):
self.assertEqual(
pf.DataType._from_type_hint(dict[str, float]),
pf.DataType.map(
- key_type=pf.DataType.string(),
- value_type=pf.DataType.float64(),
- ),
+ key_type=pf.DataType.string().not_null(),
+ value_type=pf.DataType.float64().not_null(),
+ ).not_null(),
)
def test_from_type_hint_rejects_ambiguous_or_incomplete_hints(self):
diff --git a/flink-python/pyflink/dataframe/tests/test_sql.py
b/flink-python/pyflink/dataframe/tests/test_sql.py
index be8d9033c94..c0b8cc07847 100644
--- a/flink-python/pyflink/dataframe/tests/test_sql.py
+++ b/flink-python/pyflink/dataframe/tests/test_sql.py
@@ -388,7 +388,7 @@ class SqlTests(PyFlinkDataFrameUTTestCase):
explicit = pf.sql(query, auto_bind=False, source=source, chars=chars)
automatic = pf.sql(query)
self.assertEqual(explicit.columns, ["ch"])
- self.assertEqual(explicit.schema.get_field_data_types(),
[DataTypes.STRING()])
+ self.assertEqual(explicit.schema.get_field_data_types(),
[DataTypes.STRING().not_null()])
self.assertCountEqual(explicit.union_all(automatic).collect(),
[Row("a"), Row("b")] * 2)
self.assertNotIn("chars", self.t_env.list_user_defined_functions())
self.assertNotIn("source", self.t_env.list_temporary_views())
@@ -402,9 +402,12 @@ class SqlTests(PyFlinkDataFrameUTTestCase):
yield {"word": text, "size": len(text)}
source = pf.from_dict({"text": ["ab"]})
- for return_dtype in (None, "ROW<word STRING, size BIGINT>",
- pf.DataType.struct({"word": pf.DataType.string(),
- "size":
pf.DataType.int64()})):
+ nullable = [DataTypes.STRING(), DataTypes.BIGINT()]
+ for return_dtype, expected_types in (
+ (None, [DataTypes.STRING().not_null(),
DataTypes.BIGINT().not_null()]),
+ ("ROW<word STRING, size BIGINT>", nullable),
+ (pf.DataType.struct({"word": pf.DataType.string(), "size":
pf.DataType.int64()}),
+ nullable)):
with self.subTest(return_dtype=return_dtype):
declaration = pf.udtf(expand, return_dtype=return_dtype)
named = pf.sql(
@@ -415,8 +418,7 @@ class SqlTests(PyFlinkDataFrameUTTestCase):
auto_bind=False, src=source, expand=declaration)
self.assertEqual(named.columns, ["word", "size"])
self.assertEqual(aliased.columns, ["w", "n"])
- self.assertEqual(named.schema.get_field_data_types(),
- [DataTypes.STRING(), DataTypes.BIGINT()])
+ self.assertEqual(named.schema.get_field_data_types(),
expected_types)
self.assertEqual(named.union_all(aliased).collect(),
[Row("ab", 2)] * 2)
self.assertNotIn("expand",
self.t_env.list_user_defined_functions())
diff --git a/flink-python/pyflink/dataframe/tests/test_udf.py
b/flink-python/pyflink/dataframe/tests/test_udf.py
index dfd0a396551..66fe1334bf0 100644
--- a/flink-python/pyflink/dataframe/tests/test_udf.py
+++ b/flink-python/pyflink/dataframe/tests/test_udf.py
@@ -24,7 +24,7 @@ import operator
import types
import unittest
from dataclasses import dataclass
-from typing import Any, Callable, TypedDict, cast
+from typing import Any, Callable, Optional, TypedDict, cast
from unittest import mock
import pandas as pd
@@ -154,6 +154,9 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
"""Add one to a value."""
return value + 1
+ def optional_return(value: int) -> Optional[int]:
+ return value if value % 2 == 0 else None
+
def identity(value):
return value
@@ -184,7 +187,7 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
self.assertFalse(hasattr(pf, "DataFrameUDFWrapper"))
udf_module = importlib.import_module("pyflink.dataframe.udf")
self.assertFalse(hasattr(udf_module, "DataFrameUDFWrapper"))
- self.assertEqual(_return_dtype(decorated), pf.DataType.int64())
+ self.assertEqual(_return_dtype(decorated),
pf.DataType.int64().not_null())
self.assertEqual(decorated.__name__, "add_one")
self.assertEqual(decorated.__qualname__, add_one.__qualname__)
self.assertEqual(decorated.__module__, add_one.__module__)
@@ -208,7 +211,7 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
direct = pf.udf(functools.partial(add_one), name="partial_add_one")
self.assertEqual(_return_dtype(configured), pf.DataType.string())
- self.assertEqual(_return_dtype(direct), pf.DataType.int64())
+ self.assertEqual(_return_dtype(direct), pf.DataType.int64().not_null())
self.assertEqual(direct.__name__, "partial_add_one")
with mock.patch.object(
pf.DataType,
@@ -221,19 +224,26 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
expected_result_dtype = pf.DataType.struct(
{
- "id": pf.DataType.int64(),
+ "id": pf.DataType.int64().not_null(),
"details": pf.DataType.struct(
{
- "label": pf.DataType.string(),
- "scores": pf.DataType.list(pf.DataType.int64()),
+ "label": pf.DataType.string().not_null(),
+ "scores": pf.DataType.list(
+ pf.DataType.int64().not_null()
+ ).not_null(),
}
- ),
+ ).not_null(),
}
- )
+ ).not_null()
declarations = [
(
"Python type",
lambda: pf.udf(identity, return_dtype=int),
+ pf.DataType.int64().not_null(),
+ ),
+ (
+ "Optional return annotation widens to nullable",
+ lambda: pf.udf(optional_return),
pf.DataType.int64(),
),
(
@@ -249,12 +259,12 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
(
"concrete return with unresolved input",
lambda: pf.udf(concrete_return_with_unresolved_input),
- pf.DataType.int64(),
+ pf.DataType.int64().not_null(),
),
(
"postponed return with unresolved input",
lambda: pf.udf(postponed_return_with_unresolved_input),
- pf.DataType.int64(),
+ pf.DataType.int64().not_null(),
),
]
for case_name, declare, expected in declarations:
@@ -326,7 +336,9 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
for source in callables:
with self.subTest(source=source):
decorated = pf.udf(source)
- self.assertEqual(_return_dtype(decorated), pf.DataType.int64())
+ self.assertEqual(
+ _return_dtype(decorated), pf.DataType.int64().not_null()
+ )
self.assertEqual(plain_constructor_calls, [])
self.assertEqual(scalar_constructor_calls, [])
@@ -343,7 +355,7 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
def __call__(self, value: int) -> "Output":
return {"value": value}
- expected = pf.DataType.struct({"value": pf.DataType.int64()})
+ expected = pf.DataType.struct({"value":
pf.DataType.int64().not_null()}).not_null()
for source in (Describe, Describe()):
with self.subTest(source=source):
self.assertEqual(_return_dtype(pf.udf(source)), expected)
@@ -373,7 +385,7 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
def __call__(self, value: int) -> "SelfQualified.Output":
return {"value": value}
- expected = pf.DataType.struct({"value": pf.DataType.int64()})
+ expected = pf.DataType.struct({"value":
pf.DataType.int64().not_null()}).not_null()
bound_method = BoundMethodOwner().describe
for source in (
bound_method,
@@ -398,7 +410,8 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
__call__ = _module_alias_method
self.assertEqual(
- _return_dtype(pf.udf(ReceivingCallable)), pf.DataType.int64()
+ _return_dtype(pf.udf(ReceivingCallable)),
+ pf.DataType.int64().not_null(),
)
class OverriddenScalarFunction(ScalarFunction):
@@ -411,7 +424,7 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
overridden = OverriddenScalarFunction()
overridden.eval = types.MethodType(_module_alias_method, overridden)
self.assertEqual(
- _return_dtype(pf.udf(overridden)), pf.DataType.int64()
+ _return_dtype(pf.udf(overridden)), pf.DataType.int64().not_null()
)
def test_wrapped_callable_annotations_and_partial_validation(self):
@@ -443,7 +456,9 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
pf.udf(functools.partial(add, missing=1))
uninspectable = pf.udf(operator.itemgetter(0), return_dtype=int)
- self.assertEqual(_return_dtype(uninspectable), pf.DataType.int64())
+ self.assertEqual(
+ _return_dtype(uninspectable), pf.DataType.int64().not_null()
+ )
wrapped_callable_instance = WrappedCallableClass()
for source in (
@@ -471,7 +486,7 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
functools.update_wrapper(cross_namespace_wrapper,
_module_alias_function)
self.assertEqual(
_return_dtype(pf.udf(cross_namespace_wrapper)),
- pf.DataType.int64(),
+ pf.DataType.int64().not_null(),
)
def test_func_type_resolution_and_async_detection(self):
@@ -817,7 +832,7 @@ class DataFrameUDFDeclarationTests(unittest.TestCase):
source(), return_dtype=int, func_type="general"
)
self.assertEqual(
- _return_dtype(declaration), pf.DataType.int64()
+ _return_dtype(declaration), pf.DataType.int64().not_null()
)
def test_unresolved_typed_dict_fields_have_actionable_errors(self):
@@ -1601,7 +1616,7 @@ class
DataFrameUDFPlannerTests(PyFlinkDataFrameUTTestCase):
["id", "rendered", "description"],
[
TableDataTypes.BIGINT(),
- TableDataTypes.STRING(),
+ TableDataTypes.STRING().not_null(),
TableDataTypes.ROW(
[
TableDataTypes.FIELD("value", TableDataTypes.BIGINT()),
diff --git a/flink-python/pyflink/dataframe/tests/test_udtf.py
b/flink-python/pyflink/dataframe/tests/test_udtf.py
index 80623aaaae1..711050be7c7 100644
--- a/flink-python/pyflink/dataframe/tests/test_udtf.py
+++ b/flink-python/pyflink/dataframe/tests/test_udtf.py
@@ -58,7 +58,10 @@ class DataFrameUDTFDeclarationTests(unittest.TestCase):
def emit(value) -> Iterator[_Output]:
yield {"value": value, "label": "ok"}
- expected = pf.DataType.struct({"value": pf.DataType.int64(), "label":
pf.DataType.string()})
+ expected = pf.DataType.struct({
+ "value": pf.DataType.int64().not_null(),
+ "label": pf.DataType.string().not_null(),
+ }).not_null()
for declaration in (pf.udtf(emit), pf.udtf()(emit),
pf.udtf(return_dtype=expected)(emit),
pf.udtf(emit, return_dtype=_Output)):
self.assertEqual(declaration.return_dtype, expected)
@@ -71,15 +74,16 @@ class DataFrameUDTFDeclarationTests(unittest.TestCase):
return [value]
emit.__annotations__ = {"return": hint}
with self.subTest(hint=hint):
- self.assertEqual(pf.udtf(emit).return_dtype,
pf.DataType.int64())
+ self.assertEqual(pf.udtf(emit).return_dtype,
pf.DataType.int64().not_null())
def test_explicit_type_does_not_resolve_annotations(self):
def emit(value):
return [value]
emit.__annotations__ = {"value": "UnavailableInput", "return":
"UnavailableOutput"}
- self.assertEqual(pf.udtf(emit, return_dtype=int).return_dtype,
pf.DataType.int64())
+ self.assertEqual(
+ pf.udtf(emit, return_dtype=int).return_dtype,
pf.DataType.int64().not_null())
emit.__annotations__["return"] = Iterator[int]
- self.assertEqual(pf.udtf(emit).return_dtype, pf.DataType.int64())
+ self.assertEqual(pf.udtf(emit).return_dtype,
pf.DataType.int64().not_null())
def test_callable_instance_and_partial(self):
class Repeat:
@@ -543,7 +547,7 @@ class _DataFrameFlatMapTests:
result = source.flat_map(expand)
self.assertEqual(result.columns, ["value", "label"])
self.assertEqual(result.schema.get_field_data_types(),
- [DataTypes.BIGINT(), DataTypes.STRING()])
+ [DataTypes.BIGINT().not_null(),
DataTypes.STRING().not_null()])
rows = result.filter(pf.col("value") >= 0).select("label",
"value").collect()
self.assertCountEqual(rows, [Row("a", 0), Row("b", 0), Row("b", 1)])
@@ -686,7 +690,8 @@ class _DataFrameFlatMapTests:
result = pf.from_dict({"x": [1]}).flat_map(expand)
self.assertEqual(result.columns, ["f0"])
self.assertEqual(result.schema.get_field_data_types(),
- [DataTypes.MAP(DataTypes.STRING(),
DataTypes.BIGINT())])
+ [DataTypes.MAP(DataTypes.STRING().not_null(),
+
DataTypes.BIGINT().not_null()).not_null()])
self.assertCountEqual(result.collect(), [Row({"value": 1}),
Row({"value": 2})])
diff --git a/flink-python/pyflink/dataframe/udf.py
b/flink-python/pyflink/dataframe/udf.py
index 06ad69cd6dc..e2f3955f435 100644
--- a/flink-python/pyflink/dataframe/udf.py
+++ b/flink-python/pyflink/dataframe/udf.py
@@ -43,6 +43,7 @@ from pyflink.common import Row
from pyflink.dataframe.datatype import DataType
from pyflink.table.expression import Expression
from pyflink.table.expressions import call as table_call
+from pyflink.table.typehints import _is_typed_dict
from pyflink.table.types import ArrayType, MapType, RowType
from pyflink.table.udf import (
AsyncScalarFunction,
@@ -319,7 +320,9 @@ def udf(
A UDF can be declared with a bare decorator, a configured decorator, or a
direct call. General UDFs may infer ``return_dtype`` from the return
annotation of the function, ``__call__``, or ``eval``. A ``TypedDict``
- return annotation becomes a struct column::
+ return annotation becomes a struct column. Inferred return types are
+ ``NOT NULL``. Annotate ``Optional[T]`` (or ``T | None``) to make the result
+ nullable::
>>> import pyflink.dataframe as pf
@@ -971,7 +974,7 @@ def _infer_return_dtype(
"Specify return_dtype explicitly."
)
try:
- return _data_type_from_type_hint(return_hint)
+ return DataType._from_type_hint(return_hint)
except (NameError, AttributeError, SyntaxError, TypeError) as exc:
raise TypeError(
f"Cannot infer return_dtype for '{udf_name}' from its return
annotation.\n"
@@ -985,7 +988,7 @@ def _convert_to_dtype(dtype_like: _DataTypeLike) ->
DataType:
if isinstance(dtype_like, str):
return DataType._from_sql(dtype_like)
try:
- return _data_type_from_type_hint(dtype_like)
+ return DataType._from_type_hint(dtype_like)
except (NameError, AttributeError, SyntaxError, TypeError) as exc:
if _is_typed_dict(dtype_like):
raise TypeError(
@@ -998,32 +1001,6 @@ def _convert_to_dtype(dtype_like: _DataTypeLike) ->
DataType:
) from exc
-def _is_typed_dict(type_hint: Any) -> bool:
- try:
- from typing import is_typeddict
-
- if is_typeddict(type_hint):
- return True
- except ImportError:
- pass
- return (
- isinstance(type_hint, type)
- and issubclass(type_hint, dict)
- and hasattr(type_hint, "__required_keys__")
- )
-
-
-def _data_type_from_type_hint(type_hint: Any) -> DataType:
- if _is_typed_dict(type_hint):
- return DataType.struct(
- {
- name: _data_type_from_type_hint(field_hint)
- for name, field_hint in get_type_hints(type_hint).items()
- }
- )
- return DataType._from_type_hint(type_hint)
-
-
def _detect_func_type(declaration_context: _UDFDeclarationContext) -> str:
"""Detect a unique vectorized mode from unbound container annotations."""
hint_func = declaration_context.annotation_target
diff --git a/flink-python/pyflink/dataframe/udtf.py
b/flink-python/pyflink/dataframe/udtf.py
index f1b92a1839e..60c1ed4bf71 100644
--- a/flink-python/pyflink/dataframe/udtf.py
+++ b/flink-python/pyflink/dataframe/udtf.py
@@ -37,7 +37,6 @@ from pyflink.dataframe.udf import (
_convert_to_dtype,
_create_declaration_context,
_create_result_normalizer,
- _data_type_from_type_hint,
_get_callable_inspection_target,
_get_callable_return_type_hint,
_is_typed_dict,
@@ -315,12 +314,12 @@ def _infer_udtf_return_dtype(
if not fields or Ellipsis in fields:
raise TypeError("UDTF tuple outputs must have a fixed number
of fields.")
dtype = DataType.struct([
- (f"f{index}", _data_type_from_type_hint(field))
+ (f"f{index}", DataType._from_type_hint(field))
for index, field in enumerate(fields)
])
has_named_fields = False
else:
- dtype = _data_type_from_type_hint(hint)
+ dtype = DataType._from_type_hint(hint)
has_named_fields = isinstance(dtype._to_table_data_type(), RowType)
table_type = dtype._to_table_data_type()
if isinstance(table_type, RowType) and not table_type.fields:
diff --git a/flink-python/pyflink/table/tests/test_typehints.py
b/flink-python/pyflink/table/tests/test_typehints.py
new file mode 100644
index 00000000000..deddb03055b
--- /dev/null
+++ b/flink-python/pyflink/table/tests/test_typehints.py
@@ -0,0 +1,276 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+################################################################################
+
+import datetime
+import decimal
+import sys
+import unittest
+from typing import Any, Awaitable, Coroutine, List, Optional, TypedDict, Union
+
+import typing_extensions
+
+from pyflink.table.types import DataTypes
+from pyflink.table.typehints import _from_python_type, _unwrap_awaitable
+
+
+class _Point(TypedDict):
+ x: int
+ y: int
+
+
+class _Nested(TypedDict):
+ name: str
+ point: _Point
+
+
+class _WithOptional(TypedDict):
+ id: int
+ label: Optional[str]
+
+
+class _Partial(TypedDict, total=False):
+ x: int
+ y: int
+
+
+# typing.TypedDict ignores theese markers when computing optional keys prior
to python 3.11
+class _WithNotRequired(TypedDict):
+ id: int
+ label: typing_extensions.NotRequired[str]
+ note: typing_extensions.NotRequired[Optional[str]]
+
+
+class _ExtensionsWithNotRequired(typing_extensions.TypedDict):
+ id: int
+ label: typing_extensions.NotRequired[str]
+ note: typing_extensions.NotRequired[Optional[str]]
+
+
+class _PartialWithRequired(TypedDict, total=False):
+ id: typing_extensions.Required[int]
+ label: str
+
+
+class _ExtensionsPartialWithRequired(typing_extensions.TypedDict, total=False):
+ id: typing_extensions.Required[int]
+ label: str
+
+
+class _LinkedNode(TypedDict):
+ value: int
+ next: Optional["_LinkedNode"]
+
+
+class _TreeNode(TypedDict):
+ # get_type_hints leaves forward references inside builtin generics
unresolved before 3.10.
+ children: List["_TreeNode"]
+
+
+class _Employee(TypedDict):
+ department: Optional["_Department"]
+
+
+class _Department(TypedDict):
+ manager: Optional[_Employee]
+
+
+class FromPythonTypeTests(unittest.TestCase):
+
+ def test_basic_types_infer_as_not_null(self):
+ expected = {
+ bool: DataTypes.BOOLEAN(),
+ int: DataTypes.BIGINT(),
+ float: DataTypes.DOUBLE(),
+ str: DataTypes.STRING(),
+ bytes: DataTypes.BYTES(),
+ bytearray: DataTypes.BYTES(),
+ decimal.Decimal: DataTypes.DECIMAL(38, 18),
+ datetime.date: DataTypes.DATE(),
+ datetime.time: DataTypes.TIME(3),
+ datetime.datetime: DataTypes.TIMESTAMP(6),
+ }
+ for hint, data_type in expected.items():
+ with self.subTest(hint=hint):
+ self.assertEqual(_from_python_type(hint), data_type.not_null())
+
+ def test_any_infers_as_nullable_string(self):
+ self.assertEqual(_from_python_type(Any), DataTypes.STRING())
+
+ def test_optional_widens_to_nullable(self):
+ self.assertEqual(_from_python_type(Optional[int]), DataTypes.BIGINT())
+ self.assertEqual(_from_python_type(Optional[str]), DataTypes.STRING())
+
+ @unittest.skipIf(
+ sys.version_info < (3, 10), "PEP 604 union types require Python 3.10
or later"
+ )
+ def test_pep_604_optional_widens_to_nullable(self):
+ self.assertEqual(_from_python_type(int | None), DataTypes.BIGINT())
+
+ def test_list_is_not_null_with_not_null_element(self):
+ self.assertEqual(
+ _from_python_type(list[int]),
+ DataTypes.ARRAY(DataTypes.BIGINT().not_null()).not_null(),
+ )
+
+ def test_list_of_optional_has_nullable_element(self):
+ self.assertEqual(
+ _from_python_type(list[Optional[int]]),
+ DataTypes.ARRAY(DataTypes.BIGINT()).not_null(),
+ )
+
+ def test_optional_list_is_nullable_with_not_null_element(self):
+ self.assertEqual(
+ _from_python_type(Optional[list[int]]),
+ DataTypes.ARRAY(DataTypes.BIGINT().not_null()),
+ )
+
+ def test_dict_is_not_null_with_not_null_key_and_value(self):
+ self.assertEqual(
+ _from_python_type(dict[str, float]),
+ DataTypes.MAP(
+ DataTypes.STRING().not_null(), DataTypes.DOUBLE().not_null()
+ ).not_null(),
+ )
+
+ def test_typed_dict_maps_to_not_null_row(self):
+ self.assertEqual(
+ _from_python_type(_Point),
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD("x", DataTypes.BIGINT().not_null()),
+ DataTypes.FIELD("y", DataTypes.BIGINT().not_null()),
+ ]
+ ).not_null(),
+ )
+
+ def test_nested_typed_dict_field_is_not_null_row(self):
+ self.assertEqual(
+ _from_python_type(_Nested),
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD("name", DataTypes.STRING().not_null()),
+ DataTypes.FIELD(
+ "point",
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD("x",
DataTypes.BIGINT().not_null()),
+ DataTypes.FIELD("y",
DataTypes.BIGINT().not_null()),
+ ]
+ ).not_null(),
+ ),
+ ]
+ ).not_null(),
+ )
+
+ def test_typed_dict_optional_field_is_nullable(self):
+ self.assertEqual(
+ _from_python_type(_WithOptional),
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD("id", DataTypes.BIGINT().not_null()),
+ DataTypes.FIELD("label", DataTypes.STRING()),
+ ]
+ ).not_null(),
+ )
+
+ def test_typed_dict_not_required_field_is_nullable(self):
+ for hint in (_WithNotRequired, _ExtensionsWithNotRequired):
+ with self.subTest(hint=hint):
+ self.assertEqual(
+ _from_python_type(hint),
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD("id",
DataTypes.BIGINT().not_null()),
+ DataTypes.FIELD("label", DataTypes.STRING()),
+ DataTypes.FIELD("note", DataTypes.STRING()),
+ ]
+ ).not_null(),
+ )
+
+ def test_non_total_typed_dict_fields_are_nullable(self):
+ self.assertEqual(
+ _from_python_type(_Partial),
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD("x", DataTypes.BIGINT()),
+ DataTypes.FIELD("y", DataTypes.BIGINT()),
+ ]
+ ).not_null(),
+ )
+
+ def test_required_field_in_non_total_typed_dict_is_not_null(self):
+ for hint in (_PartialWithRequired, _ExtensionsPartialWithRequired):
+ with self.subTest(hint=hint):
+ self.assertEqual(
+ _from_python_type(hint),
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD("id",
DataTypes.BIGINT().not_null()),
+ DataTypes.FIELD("label", DataTypes.STRING()),
+ ]
+ ).not_null(),
+ )
+
+ def test_typed_dict_resolves_in_container_position(self):
+ self.assertEqual(
+ _from_python_type(list[_Point]),
+ DataTypes.ARRAY(
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD("x", DataTypes.BIGINT().not_null()),
+ DataTypes.FIELD("y", DataTypes.BIGINT().not_null()),
+ ]
+ ).not_null()
+ ).not_null(),
+ )
+
+ def test_self_referencing_typed_dict_raises(self):
+ for hint, name in (
+ (_LinkedNode, "_LinkedNode"),
+ (_TreeNode, "_TreeNode"),
+ (_Employee, "_Employee"),
+ (list[_Department], "_Department"),
+ ):
+ with self.subTest(hint=hint):
+ with self.assertRaisesRegex(TypeError, f"'{name}'.*references
itself"):
+ _from_python_type(hint)
+
+ def test_ambiguous_union_raises(self):
+ with self.assertRaises(TypeError):
+ _from_python_type(Union[int, str])
+
+ def test_unsupported_hints_raise(self):
+ for hint in (complex, list, dict[str]):
+ with self.subTest(hint=hint):
+ with self.assertRaises(TypeError):
+ _from_python_type(hint)
+
+
+class UnwrapAwaitableTests(unittest.TestCase):
+
+ def test_unwraps_coroutine_and_awaitable(self):
+ self.assertIs(_unwrap_awaitable(Coroutine[Any, Any, int]), int)
+ self.assertIs(_unwrap_awaitable(Awaitable[str]), str)
+
+ def test_passes_through_non_awaitable(self):
+ self.assertIs(_unwrap_awaitable(int), int)
+ self.assertEqual(_unwrap_awaitable(Optional[int]), Optional[int])
+
+
+if __name__ == "__main__":
+ unittest.main()
diff --git a/flink-python/pyflink/table/typehints.py
b/flink-python/pyflink/table/typehints.py
new file mode 100644
index 00000000000..f1e08783ab5
--- /dev/null
+++ b/flink-python/pyflink/table/typehints.py
@@ -0,0 +1,194 @@
+################################################################################
+# 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.
+################################################################################
+"""Shared inference from Python type hints to table :class:`DataType`s.
+
+This is the neutral core used to resolve a standard Python annotation into a
+:mod:`pyflink.table.types` ``DataType``. It is consumed by the DataFrame API
and
+by UDF type-hint inference so both derive types from a single mapping.
+"""
+
+import collections.abc
+import datetime
+import decimal
+import types
+from functools import partial
+from typing import Any, Callable, Dict, Set, Union, get_args, get_origin,
get_type_hints
+
+from pyflink.table.types import DataType, DataTypes
+
+_PEP_604_UNION_TYPE = getattr(types, "UnionType", None)
+_AWAITABLE_ORIGINS = (collections.abc.Coroutine, collections.abc.Awaitable)
+
+try:
+ # Prior to py 3.11 get_type_hints leaves these markers on the field type
itself, and
+ # typing.TypedDict does not account for them in __optional_keys__
+ from typing_extensions import NotRequired, Required
+
+ _TYPED_DICT_KEY_MARKER_OPTIONALITY: Dict[Any, bool] = {NotRequired: True,
Required: False}
+except ImportError:
+ _TYPED_DICT_KEY_MARKER_OPTIONALITY = {}
+
+_BASIC_TYPE_HINT_FACTORIES: Dict[Any, Callable[[], DataType]] = {
+ bool: DataTypes.BOOLEAN,
+ int: DataTypes.BIGINT,
+ float: DataTypes.DOUBLE,
+ str: DataTypes.STRING,
+ bytes: DataTypes.BYTES,
+ bytearray: DataTypes.BYTES,
+ decimal.Decimal: partial(DataTypes.DECIMAL, 38, 18),
+ datetime.date: DataTypes.DATE,
+ # TIME is stored at runtime as an int number of milliseconds of the day,
so 3 is
+ # the highest fractional-second precision that survives a round trip.
+ datetime.time: partial(DataTypes.TIME, 3),
+ datetime.datetime: DataTypes.TIMESTAMP,
+}
+
+
+def _is_typed_dict(type_hint: Any) -> bool:
+ try:
+ from typing import is_typeddict
+
+ if is_typeddict(type_hint):
+ return True
+ except ImportError:
+ pass
+ return (
+ isinstance(type_hint, type)
+ and issubclass(type_hint, dict)
+ and hasattr(type_hint, "__required_keys__")
+ )
+
+
+def _from_python_type(type_hint: Any) -> DataType:
+ """Resolve a Python type hint into a table :class:`DataType`.
+
+ Supports the basic scalar types, ``list[T]``/``dict[K, V]`` containers, and
+ ``TypedDict`` (mapped to ``ROW``). Inferred types are ``NOT NULL``;
+ ``Optional[T]``/``T | None`` is the marker that widens a type to nullable,
+ and ``Any`` (an opt-out of the type system) stays nullable ``STRING``.
+ Raises :class:`TypeError` for hints that cannot be resolved unambiguously.
+ """
+
+ typed_dicts_in_progress: Set[Any] = set()
+
+ def infer_typed_dict(hint: Any) -> DataType:
+ if hint in typed_dicts_in_progress:
+ raise TypeError(
+ f"Cannot infer DataType from TypedDict '{hint.__name__}': it
references "
+ "itself, which ROW types cannot represent. "
+ "Please specify the data type explicitly."
+ )
+ typed_dicts_in_progress.add(hint)
+ try:
+ return DataTypes.ROW(
+ [
+ DataTypes.FIELD(
+ name,
+ infer_typed_dict_field(
+ field_hint, name in hint.__optional_keys__
+ ),
+ )
+ for name, field_hint in get_type_hints(hint).items()
+ ]
+ )
+ finally:
+ typed_dicts_in_progress.discard(hint)
+
+ def infer_typed_dict_field(hint: Any, is_optional_key: bool) -> DataType:
+ origin = get_origin(hint)
+ if origin in _TYPED_DICT_KEY_MARKER_OPTIONALITY:
+ is_optional_key = _TYPED_DICT_KEY_MARKER_OPTIONALITY[origin]
+ hint = get_args(hint)[0]
+ data_type = infer(hint)
+ return data_type.nullable() if is_optional_key else data_type
+
+ def infer_union(hint: Any, arguments) -> DataType:
+ non_none_types = [
+ argument for argument in arguments if argument is not type(None)
+ ]
+ if len(non_none_types) == 1:
+ return infer(non_none_types[0]).nullable()
+
+ raise TypeError(
+ f"Cannot infer DataType from type hint '{hint}'. "
+ "Please specify the data type explicitly."
+ )
+
+ def infer_basic(hint: Any):
+ factory = _BASIC_TYPE_HINT_FACTORIES.get(hint)
+ return factory() if factory is not None else None
+
+ def infer_not_null(hint: Any, origin: Any, arguments) -> DataType:
+ if _is_typed_dict(hint):
+ return infer_typed_dict(hint)
+
+ if origin is list:
+ if not arguments:
+ raise TypeError(
+ "Cannot infer DataType from list without type argument. "
+ "Use list[T], for example list[int]."
+ )
+ return DataTypes.ARRAY(infer(arguments[0]))
+
+ if origin is dict:
+ if len(arguments) != 2:
+ raise TypeError(
+ "Cannot infer DataType from dict without key and value
type arguments. "
+ "Use dict[K, V], for example dict[str, int]."
+ )
+ return DataTypes.MAP(infer(arguments[0]), infer(arguments[1]))
+
+ data_type = infer_basic(hint)
+ if data_type is not None:
+ return data_type
+
+ raise TypeError(
+ f"Cannot infer DataType from type hint '{hint}'. "
+ "Please specify the data type explicitly."
+ )
+
+ def infer(hint: Any) -> DataType:
+ origin = get_origin(hint)
+ arguments = get_args(hint)
+
+ if origin is Union or (
+ _PEP_604_UNION_TYPE is not None and origin is _PEP_604_UNION_TYPE
+ ):
+ return infer_union(hint, arguments)
+
+ # `Any` opts out of the type system, so it stays nullable rather than
+ # acquiring the NOT NULL default applied to the concrete hints.
+ if hint is Any:
+ return DataTypes.STRING()
+
+ return infer_not_null(hint, origin, arguments).not_null()
+
+ return infer(type_hint)
+
+
+def _unwrap_awaitable(type_hint: Any) -> Any:
+ """Unwrap ``Coroutine[Any, Any, T]``/``Awaitable[T]`` to ``T``.
+
+ Returns the hint unchanged when it is not an awaitable wrapper, so it can
be
+ applied unconditionally before resolving an async function's result type.
+ """
+ if get_origin(type_hint) in _AWAITABLE_ORIGINS:
+ arguments = get_args(type_hint)
+ if arguments:
+ return arguments[-1]
+ return type_hint
