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 0910fc925db [FLINK-40187][python] Introduce DataType class in
DataFrame API (#28843)
0910fc925db is described below
commit 0910fc925dbe28965f1e15aff8aba5faa651630a
Author: Liu Liu <[email protected]>
AuthorDate: Mon Aug 3 10:53:14 2026 +0800
[FLINK-40187][python] Introduce DataType class in DataFrame API (#28843)
---
.../docs/reference/pyflink.dataframe/datatype.rst | 20 ++
flink-python/pyflink/dataframe/datatype.py | 387 ++++++++++++++++++++-
.../pyflink/dataframe/tests/test_datatype.py | 372 +++++++++++++++++++-
flink-python/pyflink/table/expression.py | 11 +-
flink-python/pyflink/table/tests/test_types.py | 22 +-
flink-python/pyflink/table/types.py | 28 +-
6 files changed, 807 insertions(+), 33 deletions(-)
diff --git a/flink-python/docs/reference/pyflink.dataframe/datatype.rst
b/flink-python/docs/reference/pyflink.dataframe/datatype.rst
index 9d9b19b5cf5..c439b3a4a47 100644
--- a/flink-python/docs/reference/pyflink.dataframe/datatype.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/datatype.rst
@@ -34,5 +34,25 @@ Example::
:toctree: api/
DataType
+ DataType.int8
+ DataType.int16
+ DataType.int32
DataType.int64
+ DataType.float32
+ DataType.float64
+ DataType.decimal
DataType.string
+ DataType.fixed_size_string
+ DataType.binary
+ DataType.fixed_size_binary
+ DataType.bool
+ DataType.null
+ DataType.date
+ DataType.time
+ DataType.timestamp
+ DataType.timestamp_ltz
+ DataType.list
+ DataType.map
+ DataType.struct
+ DataType.not_null
+ DataType.nullable
diff --git a/flink-python/pyflink/dataframe/datatype.py
b/flink-python/pyflink/dataframe/datatype.py
index c22f4ba6d68..84676885a6f 100644
--- a/flink-python/pyflink/dataframe/datatype.py
+++ b/flink-python/pyflink/dataframe/datatype.py
@@ -16,11 +16,33 @@
# limitations under the License.
################################################################################
-from pyflink.table.types import DataType as TableDataType, DataTypes
+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 pyflink.table.types import DataType as TableDataType, DataTypes, NullType
from pyflink.util.api_stability_decorators import PublicEvolving
__all__ = ["DataType"]
+_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()
class DataType:
@@ -39,8 +61,13 @@ class DataType:
"""
def __init__(self, table_data_type: TableDataType):
+ if isinstance(table_data_type, NullType) and not
table_data_type._nullable:
+ raise ValueError("NULL data type must be nullable")
self._table_data_type = table_data_type
+ def __repr__(self) -> str:
+ return f"DataType({self._table_data_type!r})"
+
@PublicEvolving()
def __eq__(self, other: object) -> bool:
if not isinstance(other, DataType):
@@ -49,7 +76,55 @@ class DataType:
@PublicEvolving()
def __hash__(self) -> int:
- return hash(str(self._table_data_type))
+ return hash(repr(self._table_data_type))
+
+ @PublicEvolving()
+ def not_null(self) -> "DataType":
+ """
+ Return a non-nullable version of this data type.
+
+ .. versionadded:: 2.4.0
+ """
+ return DataType(self._table_data_type.not_null())
+
+ @PublicEvolving()
+ def nullable(self) -> "DataType":
+ """
+ Return a nullable version of this data type.
+
+ .. versionadded:: 2.4.0
+ """
+ return DataType(self._table_data_type.nullable())
+
+ @classmethod
+ @PublicEvolving()
+ def int8(cls) -> "DataType":
+ """
+ Create an 8-bit signed integer type.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.TINYINT())
+
+ @classmethod
+ @PublicEvolving()
+ def int16(cls) -> "DataType":
+ """
+ Create a 16-bit signed integer type.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.SMALLINT())
+
+ @classmethod
+ @PublicEvolving()
+ def int32(cls) -> "DataType":
+ """
+ Create a 32-bit signed integer type.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.INT())
@classmethod
@PublicEvolving()
@@ -66,6 +141,39 @@ class DataType:
"""
return cls(DataTypes.BIGINT())
+ @classmethod
+ @PublicEvolving()
+ def float32(cls) -> "DataType":
+ """
+ Create a 32-bit floating point type.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.FLOAT())
+
+ @classmethod
+ @PublicEvolving()
+ def float64(cls) -> "DataType":
+ """
+ Create a 64-bit floating point type.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.DOUBLE())
+
+ @classmethod
+ @PublicEvolving()
+ def decimal(cls, precision: int, scale: int) -> "DataType":
+ """
+ Create a decimal type with the given precision and scale.
+
+ :param precision: Total number of digits.
+ :param scale: Number of digits to the right of the decimal point.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.DECIMAL(precision, scale))
+
@classmethod
@PublicEvolving()
def string(cls) -> "DataType":
@@ -81,5 +189,280 @@ class DataType:
"""
return cls(DataTypes.STRING())
+ @classmethod
+ @PublicEvolving()
+ def fixed_size_string(cls, length: int) -> "DataType":
+ """
+ Create a fixed-length character string type.
+
+ :param length: Number of characters.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.CHAR(length))
+
+ @classmethod
+ @PublicEvolving()
+ def binary(cls) -> "DataType":
+ """
+ Create a variable-length binary string type.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.BYTES())
+
+ @classmethod
+ @PublicEvolving()
+ def fixed_size_binary(cls, length: int) -> "DataType":
+ """
+ Create a fixed-length binary string type.
+
+ :param length: Number of bytes.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.BINARY(length))
+
+ @classmethod
+ @PublicEvolving()
+ def bool(cls) -> "DataType":
+ """
+ Create a boolean type.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.BOOLEAN())
+
+ @classmethod
+ @PublicEvolving()
+ def null(cls) -> "DataType":
+ """
+ Create a null type.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.NULL())
+
+ @classmethod
+ @PublicEvolving()
+ def date(cls) -> "DataType":
+ """
+ Create a date type.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.DATE())
+
+ @classmethod
+ @PublicEvolving()
+ def time(cls, precision: int = 0) -> "DataType":
+ """
+ Create a time type.
+
+ :param precision: Number of fractional-second digits.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.TIME(precision))
+
+ @classmethod
+ @PublicEvolving()
+ def timestamp(cls, precision: int = 6) -> "DataType":
+ """
+ Create a timestamp type without a time zone.
+
+ :param precision: Number of fractional-second digits.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.TIMESTAMP(precision))
+
+ @classmethod
+ @PublicEvolving()
+ def timestamp_ltz(cls, precision: int = 6) -> "DataType":
+ """
+ Create a timestamp type with a local time zone.
+
+ :param precision: Number of fractional-second digits.
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.TIMESTAMP_LTZ(precision))
+
+ @classmethod
+ @PublicEvolving()
+ def list(cls, dtype: "DataType") -> "DataType":
+ """
+ Create a list type.
+
+ :param dtype: Type of each list element.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> scores_type = pf.DataType.list(pf.DataType.int32())
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(DataTypes.ARRAY(dtype._to_table_data_type()))
+
+ @classmethod
+ @PublicEvolving()
+ def map(cls, key_type: "DataType", value_type: "DataType") -> "DataType":
+ """
+ Create a map type.
+
+ :param key_type: Type of each map key.
+ :param value_type: Type of each map value.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> config_type = pf.DataType.map(
+ ... pf.DataType.string(),
+ ... pf.DataType.int64(),
+ ... )
+
+ .. versionadded:: 2.4.0
+ """
+ return cls(
+ DataTypes.MAP(
+ key_type._to_table_data_type(),
+ value_type._to_table_data_type(),
+ )
+ )
+
+ @classmethod
+ @PublicEvolving()
+ def struct(
+ cls,
+ fields: Union[
+ Dict[str, "DataType"],
+ List[Tuple[str, "DataType"]],
+ ],
+ ) -> "DataType":
+ """
+ Create a struct type with named fields.
+
+ ``fields`` may be an insertion-ordered dictionary or a list of name
and type pairs.
+
+ :param fields: Field names and their data types.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> person_type = pf.DataType.struct({
+ ... "name": pf.DataType.string(),
+ ... "age": pf.DataType.int32(),
+ ... })
+
+ Fields may also be passed as a list of name and type pairs::
+
+ >>> person_type = pf.DataType.struct([
+ ... ("name", pf.DataType.string()),
+ ... ("age", pf.DataType.int32()),
+ ... ])
+
+ .. versionadded:: 2.4.0
+ """
+ field_items = fields.items() if isinstance(fields, dict) else fields
+
+ return cls(
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD(name, data_type._to_table_data_type())
+ for name, data_type in field_items
+ ]
+ )
+ )
+
+ @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)
+
+ @classmethod
+ def _from_sql(cls, sql_type: str) -> "DataType":
+ """
+ Create a data type from its SQL representation.
+
+ :param sql_type: SQL data type string, such as ``INT`` or
``ARRAY<BIGINT>``.
+ :raises ValueError: If the SQL data type cannot be parsed.
+ """
+ from py4j.protocol import Py4JJavaError
+
+ from pyflink.java_gateway import get_gateway
+ from pyflink.table.types import _from_java_data_type
+ from pyflink.util.exceptions import JavaException
+
+ try:
+ gateway = get_gateway()
+ j_logical_type = (
+
gateway.jvm.org.apache.flink.table.types.logical.utils.LogicalTypeParser.parse(
+ sql_type,
+ gateway.jvm.Thread.currentThread().getContextClassLoader(),
+ )
+ )
+ j_data_type = (
+ gateway.jvm.org.apache.flink.table.types.utils.TypeConversions
+ .fromLogicalToDataType(j_logical_type)
+ )
+ return cls(_from_java_data_type(j_data_type))
+ except (JavaException, Py4JJavaError) as exc:
+ raise ValueError(str(exc)) from None
+
def _to_table_data_type(self) -> TableDataType:
return self._table_data_type
diff --git a/flink-python/pyflink/dataframe/tests/test_datatype.py
b/flink-python/pyflink/dataframe/tests/test_datatype.py
index 756bba371ea..f129f9a5eb6 100644
--- a/flink-python/pyflink/dataframe/tests/test_datatype.py
+++ b/flink-python/pyflink/dataframe/tests/test_datatype.py
@@ -16,26 +16,128 @@
# limitations under the License.
################################################################################
+import datetime
+import decimal
+import sys
import unittest
+from typing import Any, List, Optional, Union
import pyflink.dataframe as pf
from pyflink.table import DataTypes
-from pyflink.util.api_stability_decorators import PublicEvolving
+from pyflink.testing.test_case_utils import PyFlinkTestCase
class DataTypeTests(unittest.TestCase):
- def test_public_factory_surface(self):
+ def test_public_api_surface(self):
public_methods = {
name for name in dir(pf.DataType) if not name.startswith("_")
}
- self.assertEqual(public_methods, {"int64", "string"})
+ self.assertEqual(
+ public_methods,
+ {
+ "binary",
+ "bool",
+ "date",
+ "decimal",
+ "fixed_size_binary",
+ "fixed_size_string",
+ "float32",
+ "float64",
+ "int8",
+ "int16",
+ "int32",
+ "int64",
+ "list",
+ "map",
+ "not_null",
+ "null",
+ "nullable",
+ "string",
+ "struct",
+ "time",
+ "timestamp",
+ "timestamp_ltz",
+ },
+ )
- def test_int64_maps_to_table_bigint(self):
- self.assertEqual(pf.DataType.int64()._to_table_data_type(),
DataTypes.BIGINT())
+ def test_scalar_factories_map_to_table_types(self):
+ expected_types = {
+ "binary": DataTypes.BYTES(),
+ "bool": DataTypes.BOOLEAN(),
+ "date": DataTypes.DATE(),
+ "float32": DataTypes.FLOAT(),
+ "float64": DataTypes.DOUBLE(),
+ "int8": DataTypes.TINYINT(),
+ "int16": DataTypes.SMALLINT(),
+ "int32": DataTypes.INT(),
+ "int64": DataTypes.BIGINT(),
+ "null": DataTypes.NULL(),
+ "string": DataTypes.STRING(),
+ }
+
+ for factory_name, expected_type in expected_types.items():
+ with self.subTest(factory_name=factory_name):
+ data_type = getattr(pf.DataType, factory_name)()
+ self.assertEqual(data_type._to_table_data_type(),
expected_type)
+
+ def test_parameterized_factories_map_to_table_types(self):
+ test_cases = [
+ (
+ "decimal",
+ pf.DataType.decimal,
+ {"precision": 10, "scale": 3},
+ DataTypes.DECIMAL(10, 3),
+ ),
+ (
+ "fixed_size_binary",
+ pf.DataType.fixed_size_binary,
+ {"length": 16},
+ DataTypes.BINARY(16),
+ ),
+ (
+ "fixed_size_string",
+ pf.DataType.fixed_size_string,
+ {"length": 12},
+ DataTypes.CHAR(12),
+ ),
+ ("time_default", pf.DataType.time, {}, DataTypes.TIME(0)),
+ (
+ "time_precision",
+ pf.DataType.time,
+ {"precision": 3},
+ DataTypes.TIME(3),
+ ),
+ (
+ "timestamp_default",
+ pf.DataType.timestamp,
+ {},
+ DataTypes.TIMESTAMP(6),
+ ),
+ (
+ "timestamp_precision",
+ pf.DataType.timestamp,
+ {"precision": 3},
+ DataTypes.TIMESTAMP(3),
+ ),
+ (
+ "timestamp_ltz_default",
+ pf.DataType.timestamp_ltz,
+ {},
+ DataTypes.TIMESTAMP_LTZ(6),
+ ),
+ (
+ "timestamp_ltz_precision",
+ pf.DataType.timestamp_ltz,
+ {"precision": 3},
+ DataTypes.TIMESTAMP_LTZ(3),
+ ),
+ ]
- def test_string_maps_to_table_string(self):
- self.assertEqual(pf.DataType.string()._to_table_data_type(),
DataTypes.STRING())
+ for factory_name, factory, arguments, expected_type in test_cases:
+ with self.subTest(factory_name=factory_name):
+ data_type = factory(**arguments)
+ self.assertEqual(data_type._to_table_data_type(),
expected_type)
def test_logically_equal_types_compare_and_hash_equally(self):
first_int = pf.DataType.int64()
@@ -46,19 +148,253 @@ class DataTypeTests(unittest.TestCase):
self.assertNotEqual(first_int, string)
self.assertEqual(len({first_int, second_int, string}), 2)
- def test_equality_and_hash_are_public_evolving(self):
- for method in [pf.DataType.__eq__, pf.DataType.__hash__]:
- with self.subTest(method=method.__name__):
- self.assertIn(
- PublicEvolving,
- getattr(method, "__stability_decorators", set()),
+ def test_nullability_participates_in_equality_and_hashing(self):
+ nullable = pf.DataType.int32()
+ first_non_nullable = nullable.not_null()
+ second_non_nullable = pf.DataType.int32().not_null()
+
+ self.assertNotEqual(nullable, first_non_nullable)
+ self.assertEqual(first_non_nullable, second_non_nullable)
+ self.assertEqual(hash(first_non_nullable), hash(second_non_nullable))
+ self.assertEqual(len({nullable, first_non_nullable}), 2)
+
+ def test_nullability_modifiers_preserve_the_original_type(self):
+ original = pf.DataType.int32()
+
+ non_nullable = original.not_null()
+ nullable_again = non_nullable.nullable()
+
+ self.assertEqual(original._to_table_data_type(), DataTypes.INT())
+ self.assertEqual(
+ non_nullable._to_table_data_type(),
+ DataTypes.INT().not_null(),
+ )
+ self.assertEqual(nullable_again._to_table_data_type(), DataTypes.INT())
+
+ def test_null_type_cannot_be_made_non_nullable(self):
+ with self.assertRaisesRegex(ValueError, "NULL"):
+ pf.DataType.null().not_null()
+
+ def test_list_preserves_its_element_type(self):
+ list_type = pf.DataType.list(dtype=pf.DataType.int32().not_null())
+
+ self.assertEqual(
+ list_type._to_table_data_type(),
+ DataTypes.ARRAY(DataTypes.INT().not_null()),
+ )
+
+ def test_map_preserves_its_key_and_value_types(self):
+ map_type = pf.DataType.map(
+ key_type=pf.DataType.string(),
+ value_type=pf.DataType.int64(),
+ )
+
+ self.assertEqual(
+ map_type._to_table_data_type(),
+ DataTypes.MAP(DataTypes.STRING(), DataTypes.BIGINT()),
+ )
+
+ def test_struct_preserves_dict_insertion_order(self):
+ struct_type = pf.DataType.struct(
+ fields={
+ "name": pf.DataType.string(),
+ "age": pf.DataType.int32().not_null(),
+ }
+ )
+
+ self.assertEqual(
+ struct_type._to_table_data_type(),
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD("name", DataTypes.STRING()),
+ DataTypes.FIELD("age", DataTypes.INT().not_null()),
+ ]
+ ),
+ )
+
+ def test_struct_preserves_list_field_order(self):
+ struct_type = pf.DataType.struct(
+ fields=[
+ ("age", pf.DataType.int32().not_null()),
+ ("name", pf.DataType.string()),
+ ]
+ )
+
+ self.assertEqual(
+ struct_type._to_table_data_type(),
+ DataTypes.ROW(
+ [
+ DataTypes.FIELD("age", DataTypes.INT().not_null()),
+ DataTypes.FIELD("name", DataTypes.STRING()),
+ ]
+ ),
+ )
+
+ def test_from_basic_python_type_hints(self):
+ expected_types = {
+ bool: pf.DataType.bool(),
+ int: pf.DataType.int64(),
+ float: pf.DataType.float64(),
+ str: pf.DataType.string(),
+ bytes: pf.DataType.binary(),
+ bytearray: pf.DataType.binary(),
+ decimal.Decimal: pf.DataType.decimal(38, 18),
+ datetime.date: pf.DataType.date(),
+ datetime.time: pf.DataType.time(),
+ 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,
+ )
+
+ def test_from_optional_type_hint(self):
+ self.assertEqual(
+ pf.DataType._from_type_hint(Optional[int]),
+ pf.DataType.int64(),
+ )
+
+ @unittest.skipIf(
+ sys.version_info < (3, 10),
+ "PEP 604 union types require Python 3.10 or later",
+ )
+ def test_from_pep_604_union_type_hint(self):
+ self.assertEqual(
+ pf.DataType._from_type_hint(int | None),
+ pf.DataType.int64(),
+ )
+ self.assertEqual(
+ pf.DataType._from_type_hint(list[int | None]),
+ pf.DataType.list(dtype=pf.DataType.int64()),
+ )
+ with self.assertRaises(TypeError):
+ pf.DataType._from_type_hint(int | str)
+
+ def test_from_list_type_hint(self):
+ self.assertEqual(
+ pf.DataType._from_type_hint(list[int]),
+ pf.DataType.list(dtype=pf.DataType.int64()),
+ )
+
+ 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(),
+ ),
+ )
+
+ def test_from_type_hint_rejects_ambiguous_or_incomplete_hints(self):
+ invalid_hints = [
+ List,
+ complex,
+ ]
+
+ for type_hint in invalid_hints:
+ with self.subTest(type_hint=type_hint):
+ with self.assertRaises(TypeError):
+ pf.DataType._from_type_hint(type_hint)
+
+ def test_from_type_hint_reports_ambiguous_union_error(self):
+ with self.assertRaises(TypeError) as context:
+ pf.DataType._from_type_hint(Union[int, str])
+
+ self.assertEqual(
+ "Cannot infer DataType from type hint 'typing.Union[int, str]'. "
+ "Please specify the data type explicitly.",
+ str(context.exception),
+ )
+
+ def test_repr_preserves_type_parameters_and_nested_nullability(self):
+ self.assertEqual(
+ repr(
+ pf.DataType.list(
+ pf.DataType.struct(
+ [
+ (
+ "amount",
+ pf.DataType.decimal(10, 2).not_null(),
+ )
+ ]
+ ).not_null()
+ )
+ ),
+ "DataType(ArrayType("
+ "RowType(RowField(amount, DecimalType(10, 2, false), ...), false),
"
+ "true))",
+ )
+
+
+class DataTypeExpressionTests(PyFlinkTestCase):
+ def test_expression_casts_accept_dataframe_data_types(self):
+ test_cases = [
+ ("cast", "cast(value, DOUBLE)"),
+ ("try_cast", "TRY_CAST(value, DOUBLE)"),
+ ]
+
+ for operation, expected_expression in test_cases:
+ with self.subTest(operation=operation):
+ expression = getattr(pf.col("value"), operation)(
+ pf.DataType.float64()
)
- def test_nullability_modifiers_are_not_exposed(self):
- for data_type in [pf.DataType.int64(), pf.DataType.string()]:
- with self.subTest(data_type=data_type):
- self.assertFalse(hasattr(data_type, "not_null"))
- self.assertFalse(hasattr(data_type, "nullable"))
+ self.assertEqual(str(expression), expected_expression)
+
+
+class DataTypeSqlTests(PyFlinkTestCase):
+ def test_from_sql_parses_scalar_and_nested_types(self):
+ test_cases = [
+ ("INT", pf.DataType.int32()),
+ (
+ "DECIMAL(10, 3) NOT NULL",
+ pf.DataType.decimal(10, 3).not_null(),
+ ),
+ (
+ "ROW<name STRING, scores ARRAY<DOUBLE NOT NULL>>",
+ pf.DataType.struct(
+ fields=[
+ ("name", pf.DataType.string()),
+ (
+ "scores",
+ pf.DataType.list(
+ dtype=pf.DataType.float64().not_null()
+ ),
+ ),
+ ]
+ ),
+ ),
+ ]
+
+ for sql_type, expected_type in test_cases:
+ with self.subTest(sql_type=sql_type):
+ self.assertEqual(pf.DataType._from_sql(sql_type),
expected_type)
+
+ def test_from_sql_preserves_timestamp_precision(self):
+ self.assertEqual(
+ pf.DataType._from_sql("TIMESTAMP(9)"),
+ pf.DataType.timestamp(precision=9),
+ )
+
+ def test_from_sql_preserves_timestamp_ltz_precision(self):
+ self.assertEqual(
+ pf.DataType._from_sql("TIMESTAMP_LTZ(3)"),
+ pf.DataType.timestamp_ltz(precision=3),
+ )
+
+ def test_from_sql_supports_null_type(self):
+ self.assertEqual(
+ pf.DataType._from_sql("NULL"),
+ pf.DataType.null(),
+ )
+
+ def test_from_sql_reports_parser_errors_as_value_errors(self):
+ with self.assertRaises(ValueError):
+ pf.DataType._from_sql("VARCHAR(test)")
if __name__ == "__main__":
diff --git a/flink-python/pyflink/table/expression.py
b/flink-python/pyflink/table/expression.py
index 82c49f2edc6..aa3b952f215 100644
--- a/flink-python/pyflink/table/expression.py
+++ b/flink-python/pyflink/table/expression.py
@@ -20,7 +20,12 @@ from typing import Union, TypeVar, Generic, Any
from pyflink import add_version_doc
from pyflink.java_gateway import get_gateway
-from pyflink.table.types import DataType, DataTypes, _to_java_data_type
+from pyflink.table.types import (
+ DataType,
+ DataTypes,
+ _TableDataTypeLike,
+ _to_java_data_type,
+)
from pyflink.util.api_stability_decorators import PublicEvolving
from pyflink.util.java_utils import to_jarray
@@ -886,7 +891,7 @@ class Expression(Generic[T]):
"""
return _binary_op("asArgument")(self, name)
- def cast(self, data_type: DataType) -> 'Expression':
+ def cast(self, data_type: _TableDataTypeLike) -> 'Expression':
"""
Returns a new value being cast to type type.
A cast error throws an exception and fails the job.
@@ -900,7 +905,7 @@ class Expression(Generic[T]):
"""
return _binary_op("cast")(self, _to_java_data_type(data_type))
- def try_cast(self, data_type: DataType) -> 'Expression':
+ def try_cast(self, data_type: _TableDataTypeLike) -> 'Expression':
"""
Like cast, but in case of error, returns NULL rather than failing the
job.
diff --git a/flink-python/pyflink/table/tests/test_types.py
b/flink-python/pyflink/table/tests/test_types.py
index d3bd37c6ffd..f8124ba5e8c 100644
--- a/flink-python/pyflink/table/tests/test_types.py
+++ b/flink-python/pyflink/table/tests/test_types.py
@@ -128,6 +128,18 @@ class UTCOffsetTimezone(datetime.tzinfo):
class TypesTests(PyFlinkTestCase):
+ def test_row_type_repr_includes_nullability(self):
+ row_type = RowType([RowField("id", BigIntType())])
+
+ self.assertEqual(
+ "RowType(RowField(id, BigIntType(true), ...), true)",
+ repr(row_type),
+ )
+ self.assertEqual(
+ "RowType(RowField(id, BigIntType(true), ...), false)",
+ repr(row_type.not_null()),
+ )
+
def test_infer_schema(self):
from decimal import Decimal
@@ -170,14 +182,14 @@ class TypesTests(PyFlinkTestCase):
'DoubleType(true)',
"ArrayType(DoubleType(false), true)",
"ArrayType(BigIntType(true), true)",
- 'RowType(RowField(_1, BigIntType(true), ...))',
- 'RowType(RowField(x, DoubleType(true), ...),RowField(y,
DoubleType(true), ...))',
+ 'RowType(RowField(_1, BigIntType(true), ...), true)',
+ 'RowType(RowField(x, DoubleType(true), ...),RowField(y,
DoubleType(true), ...), true)',
'MapType(VarCharType(2147483647, false), BigIntType(true), true)',
'VarBinaryType(2147483647, true)',
'DecimalType(38, 18, true)',
- 'RowType(RowField(a, BigIntType(true), ...))',
- 'RowType(RowField(a, BigIntType(true), ...))',
- 'RowType(RowField(a, BigIntType(true), ...))',
+ 'RowType(RowField(a, BigIntType(true), ...), true)',
+ 'RowType(RowField(a, BigIntType(true), ...), true)',
+ 'RowType(RowField(a, BigIntType(true), ...), true)',
]
schema = _infer_schema_from_data([data])
diff --git a/flink-python/pyflink/table/types.py
b/flink-python/pyflink/table/types.py
index b62b55af9e1..964097f25e2 100644
--- a/flink-python/pyflink/table/types.py
+++ b/flink-python/pyflink/table/types.py
@@ -29,7 +29,7 @@ from functools import reduce
from threading import RLock
from py4j.java_gateway import get_java_class
-from typing import List, Union
+from typing import List, Protocol, Union
from pyflink.common.types import _create_row
from pyflink.util.api_stability_decorators import PublicEvolving
@@ -123,6 +123,14 @@ class DataType(object):
return obj
+class _SupportsToTableDataType(Protocol):
+ def _to_table_data_type(self) -> DataType:
+ ...
+
+
+_TableDataTypeLike = Union[DataType, _SupportsToTableDataType]
+
+
class AtomicType(DataType):
"""
An internal type used to represent everything that is not
@@ -1216,7 +1224,8 @@ class RowType(DataType):
raise TypeError('RowType keys should be strings, integers or
slices')
def __repr__(self):
- return "RowType(%s)" % ",".join(repr(field) for field in self)
+ fields = ",".join(repr(field) for field in self)
+ return f"RowType({fields}, {str(self._nullable).lower()})"
def field_names(self):
"""
@@ -1709,7 +1718,8 @@ def _from_java_data_type(j_data_type):
elif is_instance_of(logical_type, gateway.jvm.TimeType):
data_type = DataTypes.TIME(logical_type.getPrecision(),
logical_type.isNullable())
elif is_instance_of(logical_type, gateway.jvm.TimestampType):
- data_type = DataTypes.TIMESTAMP(precision=3,
nullable=logical_type.isNullable())
+ data_type = DataTypes.TIMESTAMP(
+ precision=logical_type.getPrecision(),
nullable=logical_type.isNullable())
elif is_instance_of(logical_type, gateway.jvm.BooleanType):
data_type = DataTypes.BOOLEAN(logical_type.isNullable())
elif is_instance_of(logical_type, gateway.jvm.TinyIntType):
@@ -1729,7 +1739,8 @@ def _from_java_data_type(j_data_type):
TypeError("Unsupported type: %s, ZonedTimestampType is not
supported yet."
% j_data_type)
elif is_instance_of(logical_type, gateway.jvm.LocalZonedTimestampType):
- data_type =
DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(nullable=logical_type.isNullable())
+ data_type = DataTypes.TIMESTAMP_WITH_LOCAL_TIME_ZONE(
+ precision=logical_type.getPrecision(),
nullable=logical_type.isNullable())
elif is_instance_of(logical_type, gateway.jvm.DayTimeIntervalType) or \
is_instance_of(logical_type,
gateway.jvm.YearMonthIntervalType):
data_type = _from_java_interval_type(logical_type)
@@ -1759,6 +1770,8 @@ def _from_java_data_type(j_data_type):
% type_info)
elif is_instance_of(logical_type, gateway.jvm.RawType):
data_type = RawType()
+ elif is_instance_of(logical_type, gateway.jvm.NullType):
+ data_type = DataTypes.NULL()
else:
raise TypeError("Unsupported type: %s, it is not supported yet in
current python type"
" system" % j_data_type)
@@ -1821,10 +1834,15 @@ def _from_java_data_type(j_data_type):
TypeError("Unsupported data type: %s" % j_data_type)
-def _to_java_data_type(data_type: DataType):
+def _to_java_data_type(data_type: _TableDataTypeLike):
"""
Converts the specified Python DataType to Java DataType.
"""
+ if not isinstance(data_type, DataType):
+ to_table_data_type = getattr(data_type, "_to_table_data_type", None)
+ if callable(to_table_data_type):
+ data_type = to_table_data_type()
+
gateway = get_gateway()
JDataTypes = gateway.jvm.org.apache.flink.table.api.DataTypes
