This is an automated email from the ASF dual-hosted git repository. dianfu pushed a commit to branch release-2.2 in repository https://gitbox.apache.org/repos/asf/flink.git
commit d8df14c35255c5adc0f6cc723970370cb9b12b27 Author: Dian Fu <[email protected]> AuthorDate: Wed Oct 28 17:22:39 2026 +0800 [FLINK-41245][python] Fix missing allow_nested support (#29641) --- flink-python/pyflink/table/tests/test_types.py | 43 +++++++++++++++++++++++++- flink-python/pyflink/table/types.py | 29 ++++++++++------- 2 files changed, 60 insertions(+), 12 deletions(-) diff --git a/flink-python/pyflink/table/tests/test_types.py b/flink-python/pyflink/table/tests/test_types.py index 2997a3ea63f..6e934f43993 100644 --- a/flink-python/pyflink/table/tests/test_types.py +++ b/flink-python/pyflink/table/tests/test_types.py @@ -35,7 +35,7 @@ from pyflink.table.types import (_infer_schema_from_data, _infer_type, _create_type_verifier, UserDefinedType, DataTypes, Row, RowField, RowType, ArrayType, BigIntType, VarCharType, MapType, DataType, _from_java_data_type, ZonedTimestampType, - LocalZonedTimestampType, _to_java_data_type) + LocalZonedTimestampType, _to_java_data_type, create_arrow_schema) from pyflink.testing.test_case_utils import PyFlinkTestCase @@ -814,6 +814,38 @@ class DataTypeVerificationTests(PyFlinkTestCase): class DataTypeConvertTests(PyFlinkTestCase): + def test_create_arrow_schema_with_nested_types(self): + import pyarrow as pa + + row_type = DataTypes.ROW([DataTypes.FIELD("value", DataTypes.INT().not_null())]) + arrow_row_type = pa.struct([pa.field("value", pa.int32(), nullable=False)]) + nested_row_type = DataTypes.ROW([DataTypes.FIELD("nested", row_type)]) + arrow_nested_row_type = pa.struct([pa.field("nested", arrow_row_type)]) + test_types = [ + (DataTypes.ARRAY(row_type), pa.list_(arrow_row_type), ValueError), + (nested_row_type, arrow_nested_row_type, TypeError), + (DataTypes.ARRAY(DataTypes.TIMESTAMP_LTZ(3)), + pa.list_(pa.timestamp('ms')), ValueError), + (DataTypes.ROW([DataTypes.FIELD("ts", DataTypes.TIMESTAMP_LTZ(3))]), + pa.struct([pa.field("ts", pa.timestamp('ms'))]), TypeError), + (DataTypes.ARRAY(DataTypes.ARRAY(row_type)), + pa.list_(pa.list_(arrow_row_type)), ValueError), + (DataTypes.MAP(DataTypes.STRING(), DataTypes.ARRAY(nested_row_type)), + pa.map_(pa.string(), pa.list_(arrow_nested_row_type)), ValueError), + (DataTypes.MAP(nested_row_type, DataTypes.STRING()), + pa.map_(arrow_nested_row_type, pa.string()), TypeError), + ] + for data_type, arrow_type, error in test_types: + with self.subTest(data_type=data_type): + expected = pa.schema([pa.field("payload", arrow_type, nullable=False)]) + actual = create_arrow_schema( + ["payload"], [data_type.not_null()], allow_nested=True) + self.assertEqual(expected, actual) + for kwargs in [{}, {'allow_nested': False}]: + with self.subTest(kwargs=kwargs): + with self.assertRaises(error): + create_arrow_schema(["payload"], [data_type], **kwargs) + def test_basic_type(self): test_types = [DataTypes.STRING(), DataTypes.BOOLEAN(), @@ -834,6 +866,15 @@ class DataTypeConvertTests(PyFlinkTestCase): self.assertEqual(test_types, converted_python_types) + def test_timestamp_precision_round_trip(self): + for timestamp_type in [DataTypes.TIMESTAMP, DataTypes.TIMESTAMP_LTZ]: + for precision in range(10): + for nullable in [True, False]: + data_type = timestamp_type(precision, nullable) + with self.subTest(data_type=data_type): + converted = _from_java_data_type(_to_java_data_type(data_type)) + self.assertEqual(data_type, converted) + def test_atomic_type_with_data_type_with_parameters(self): gateway = get_gateway() JDataTypes = gateway.jvm.DataTypes diff --git a/flink-python/pyflink/table/types.py b/flink-python/pyflink/table/types.py index 5ef5644d0cd..6c9cf6d2a34 100644 --- a/flink-python/pyflink/table/types.py +++ b/flink-python/pyflink/table/types.py @@ -1709,7 +1709,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 +1730,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) @@ -2195,12 +2197,15 @@ def _create_type_verifier(data_type: DataType, name: str = None): return verify -def create_arrow_schema(field_names: List[str], field_types: List[DataType]): +def create_arrow_schema(field_names: List[str], field_types: List[DataType], *, allow_nested=False): """ - Create an Arrow schema with the specified filed names and types. + Create an Arrow schema with the specified field names and types. + + Set ``allow_nested`` to include nested rows and timestamps when converting empty tables. """ import pyarrow as pa - fields = [pa.field(field_name, to_arrow_type(field_type), field_type._nullable) + fields = [pa.field(field_name, to_arrow_type(field_type, allow_nested=allow_nested), + field_type._nullable) for field_name, field_type in zip(field_names, field_types)] return pa.schema(fields) @@ -2271,7 +2276,7 @@ def from_arrow_type(arrow_type, nullable: bool = True) -> DataType: raise TypeError("Unsupported data type to convert to Arrow type: " + str(dt)) -def to_arrow_type(data_type: DataType): +def to_arrow_type(data_type: DataType, *, allow_nested=False): """ Converts the specified Flink data type to pyarrow data type. """ @@ -2319,18 +2324,20 @@ def to_arrow_type(data_type: DataType): else: return pa.timestamp('ns') elif isinstance(data_type, MapType): - return pa.map_(to_arrow_type(data_type.key_type), to_arrow_type(data_type.value_type)) + return pa.map_(to_arrow_type(data_type.key_type, allow_nested=allow_nested), + to_arrow_type(data_type.value_type, allow_nested=allow_nested)) elif isinstance(data_type, ArrayType): - if type(data_type.element_type) in [LocalZonedTimestampType, RowType]: + if not allow_nested and type(data_type.element_type) in [LocalZonedTimestampType, RowType]: raise ValueError("%s is not supported to be used as the element type of ArrayType." % data_type.element_type) - return pa.list_(to_arrow_type(data_type.element_type)) + return pa.list_(to_arrow_type(data_type.element_type, allow_nested=allow_nested)) elif isinstance(data_type, RowType): for field in data_type: - if type(field.data_type) in [LocalZonedTimestampType, RowType]: + if not allow_nested and type(field.data_type) in [LocalZonedTimestampType, RowType]: raise TypeError("%s is not supported to be used as the field type of RowType" % field.data_type) - fields = [pa.field(field.name, to_arrow_type(field.data_type), field.data_type._nullable) + fields = [pa.field(field.name, to_arrow_type(field.data_type, allow_nested=allow_nested), + field.data_type._nullable) for field in data_type] return pa.struct(fields) elif isinstance(data_type, NullType):
