diff --git a/flink-python/docs/reference/pyflink.dataframe/datatype.rst b/flink-python/docs/reference/pyflink.dataframe/datatype.rst index 9d9b19b5cf5505..8eb094857a1797 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.not_null + DataType.nullable + 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 diff --git a/flink-python/pyflink/dataframe/datatype.py b/flink-python/pyflink/dataframe/datatype.py index c22f4ba6d6819b..22bae652af85cb 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 @@ def __eq__(self, other: object) -> bool: @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 @@ def int64(cls) -> "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,260 @@ def string(cls) -> "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. + + .. 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. + + .. 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(), + ... }) + + .. 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``. + :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 756bba371ea873..f129f9a5eb6f8d 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 @@ def test_logically_equal_types_compare_and_hash_equally(self): 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>", + 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 82c49f2edc63dd..aa3b952f215f21 100644 --- a/flink-python/pyflink/table/expression.py +++ b/flink-python/pyflink/table/expression.py @@ -20,7 +20,12 @@ 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 @@ def as_argument(self, name: str) -> 'Expression': """ 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 @@ def cast(self, data_type: DataType) -> 'Expression': """ 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 d3bd37c6ffd8fa..f8124ba5e8cdf7 100644 --- a/flink-python/pyflink/table/tests/test_types.py +++ b/flink-python/pyflink/table/tests/test_types.py @@ -128,6 +128,18 @@ def dst(self, dt): 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 @@ def __init__(self): '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 b62b55af9e1f3c..964097f25e2723 100644 --- a/flink-python/pyflink/table/types.py +++ b/flink-python/pyflink/table/types.py @@ -29,7 +29,7 @@ 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 @@ def from_sql_type(self, obj): 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 @@ def __getitem__(self, key): 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