This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch release-1.20
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/release-1.20 by this push:
new 1b9b42b2208 [FLINK-41032][python] Preserve column dtypes when
converting empty tables to pandas (#29526)
1b9b42b2208 is described below
commit 1b9b42b220885fffdd7644bf5590fe2d65df17e4
Author: Biao Geng <[email protected]>
AuthorDate: Tue Oct 27 15:44:57 2026 +0800
[FLINK-41032][python] Preserve column dtypes when converting empty tables
to pandas (#29526)
---
flink-python/pyflink/table/table.py | 45 ++++++++++-------
.../pyflink/table/tests/test_pandas_conversion.py | 59 +++++++++++++++++++++-
flink-python/pyflink/table/tests/test_types.py | 9 ++++
flink-python/pyflink/table/types.py | 28 ++++++----
4 files changed, 110 insertions(+), 31 deletions(-)
diff --git a/flink-python/pyflink/table/table.py
b/flink-python/pyflink/table/table.py
index 34cf15357f2..5a46aa0c932 100644
--- a/flink-python/pyflink/table/table.py
+++ b/flink-python/pyflink/table/table.py
@@ -16,8 +16,10 @@
# limitations under the License.
################################################################################
+from datetime import tzinfo
+
from py4j.java_gateway import get_method
-from typing import Union
+from typing import Optional, Union
from pyflink.java_gateway import get_gateway
from pyflink.table import ExplainDetail
@@ -27,7 +29,7 @@ from pyflink.table.serializers import ArrowSerializer
from pyflink.table.table_descriptor import TableDescriptor
from pyflink.table.table_result import TableResult
from pyflink.table.table_schema import TableSchema
-from pyflink.table.types import create_arrow_schema
+from pyflink.table.types import LocalZonedTimestampType, create_arrow_schema
from pyflink.table.udf import UserDefinedScalarFunctionWrapper, \
UserDefinedAggregateFunctionWrapper, UserDefinedTableFunctionWrapper
from pyflink.table.utils import tz_convert_from_internal, to_expression_jarray
@@ -935,27 +937,32 @@ class Table(object):
.get(gateway.jvm.org.apache.flink.python.PythonOptions.MAX_ARROW_BATCH_SIZE)
batches_iterator =
gateway.jvm.org.apache.flink.table.runtime.arrow.ArrowUtils\
.collectAsPandasDataFrame(self._j_table, max_arrow_batch_size)
- if batches_iterator.hasNext():
+ schema = self.get_schema()
+ field_types = schema.get_field_data_types()
+ has_rows = batches_iterator.hasNext()
+ arrow_schema = create_arrow_schema(
+ schema.get_field_names(), field_types, allow_nested=not has_rows)
+ timezone: Optional[tzinfo] = None
+ if any(isinstance(field_type, LocalZonedTimestampType) for field_type
in field_types):
import pytz
- timezone = pytz.timezone(
-
self._j_table.getTableEnvironment().getConfig().getLocalTimeZone().getId())
- serializer = ArrowSerializer(
- create_arrow_schema(self.get_schema().get_field_names(),
- self.get_schema().get_field_data_types()),
- self.get_schema().to_row_data_type(),
- timezone)
+ j_timezone =
self._j_table.getTableEnvironment().getConfig().getLocalTimeZone()
+ timezone_id = j_timezone.getId()
+ if timezone_id.startswith(('GMT+', 'GMT-')):
+ timezone =
pytz.FixedOffset(j_timezone.normalized().getTotalSeconds() // 60)
+ else:
+ timezone = pytz.timezone(timezone_id)
+ if has_rows:
+ serializer = ArrowSerializer(arrow_schema,
schema.to_row_data_type(), timezone)
import pyarrow as pa
table =
pa.Table.from_batches(serializer.load_from_iterator(batches_iterator))
- pdf = table.to_pandas()
-
- schema = self.get_schema()
- for field_name in schema.get_field_names():
- pdf[field_name] = tz_convert_from_internal(
- pdf[field_name], schema.get_field_data_type(field_name),
timezone)
- return pdf
else:
- import pandas as pd
- return pd.DataFrame.from_records([],
columns=self.get_schema().get_field_names())
+ table = arrow_schema.empty_table()
+ pdf = table.to_pandas()
+
+ for field_name in schema.get_field_names():
+ pdf[field_name] = tz_convert_from_internal(
+ pdf[field_name], schema.get_field_data_type(field_name),
timezone)
+ return pdf
def get_schema(self) -> TableSchema:
"""
diff --git a/flink-python/pyflink/table/tests/test_pandas_conversion.py
b/flink-python/pyflink/table/tests/test_pandas_conversion.py
index 77924dbf8e7..2e436e30622 100644
--- a/flink-python/pyflink/table/tests/test_pandas_conversion.py
+++ b/flink-python/pyflink/table/tests/test_pandas_conversion.py
@@ -164,7 +164,64 @@ class PandasConversionITTests(PandasConversionTestBase):
def test_empty_to_pandas(self):
table = self.t_env.from_pandas(self.pdf, self.data_type)
pdf = table.filter(table.f1 < 0).to_pandas()
- self.assertTrue(pdf.empty)
+ assert_frame_equal(pdf, table.to_pandas().iloc[:0])
+
+ def test_to_pandas_with_nested_types(self):
+ import pandas as pd
+
+ for expression, error in [('ARRAY[ROW(1)]', ValueError),
('ROW(ROW(1))', TypeError)]:
+ with self.subTest(expression=expression):
+ query = 'SELECT 1 AS a, %s AS payload FROM (VALUES (1)) AS
t(x)' % expression
+ result = self.t_env.sql_query(query + ' WHERE
FALSE').to_pandas()
+ expected = pd.DataFrame({'a': pd.Series(dtype='int32'),
+ 'payload': pd.Series(dtype='object')})
+ assert_frame_equal(result, expected)
+ with self.assertRaisesRegex(error, 'ROW is not supported'):
+ self.t_env.sql_query(query).to_pandas()
+
+ def test_empty_to_pandas_with_timestamps(self):
+ original_timezone = self.t_env.get_config().get_local_timezone()
+ self.addCleanup(self.t_env.get_config().set_local_timezone,
original_timezone)
+ for timezone in ['UTC', 'Asia/Shanghai', 'Etc/GMT-8']:
+ with self.subTest(timezone=timezone):
+ self.t_env.get_config().set_local_timezone(timezone)
+ table = self.t_env.sql_query("""
+ SELECT
+ 1 AS a,
+ CAST(TIMESTAMP '2024-01-01 00:00:00.123456789'
+ AS TIMESTAMP(0)) AS ts_seconds,
+ CAST(TIMESTAMP '2024-01-01 00:00:00.123456789'
+ AS TIMESTAMP(3)) AS ts_millis,
+ CAST(TIMESTAMP '2024-01-01 00:00:00.123456789'
+ AS TIMESTAMP(6)) AS ts_micros,
+ TIMESTAMP '2024-01-01 00:00:00.123456789' AS ts_nanos,
+ TO_TIMESTAMP_LTZ(1704067200123, 3) AS ltz
+ """)
+ expected = table.to_pandas().iloc[:0]
+ result = table.filter(table.a < 0).to_pandas()
+ assert_frame_equal(result, expected)
+ self.assertEqual(timezone, str(result['ltz'].dt.tz))
+
+ def test_to_pandas_with_fixed_offset_timezones(self):
+ original_timezone = self.t_env.get_config().get_local_timezone()
+ self.addCleanup(self.t_env.get_config().set_local_timezone,
original_timezone)
+ for timezone, named_timezone, offset_seconds in [
+ ('GMT+08:00', 'Asia/Shanghai', 28800),
+ ('GMT+05:30', 'Asia/Kolkata', 19800),
+ ('GMT-03:30', 'America/St_Johns', -12600)]:
+ with self.subTest(timezone=timezone):
+ self.t_env.get_config().set_local_timezone(named_timezone)
+ table = self.t_env.sql_query("""
+ SELECT 1 AS a, TO_TIMESTAMP_LTZ(1704067200123, 3) AS ltz
+ """)
+ expected_timestamp = table.to_pandas()['ltz'].iloc[0]
+ self.t_env.get_config().set_local_timezone(timezone)
+ result = table.to_pandas()
+ timestamp = result['ltz'].iloc[0]
+ self.assertEqual(expected_timestamp, timestamp)
+ self.assertEqual(datetime.timedelta(seconds=offset_seconds),
timestamp.utcoffset())
+ empty_result = table.filter(table.a < 0).to_pandas()
+ assert_frame_equal(empty_result, result.iloc[:0])
def test_to_pandas_preserves_local_zoned_timestamp_instants(self):
import pandas as pd
diff --git a/flink-python/pyflink/table/tests/test_types.py
b/flink-python/pyflink/table/tests/test_types.py
index 24685e50af1..19c25a69764 100644
--- a/flink-python/pyflink/table/tests/test_types.py
+++ b/flink-python/pyflink/table/tests/test_types.py
@@ -819,6 +819,15 @@ class DataTypeVerificationTests(PyFlinkTestCase):
class DataTypeConvertTests(PyFlinkTestCase):
+ def test_timestamp_precision(self):
+ for type_factory in (DataTypes.TIMESTAMP, DataTypes.TIMESTAMP_LTZ):
+ for precision in (0, 3, 6, 9):
+ for nullable in (True, False):
+ data_type = type_factory(precision, nullable)
+ with self.subTest(data_type=data_type,
precision=precision, nullable=nullable):
+ self.assertEqual(
+ data_type,
_from_java_data_type(_to_java_data_type(data_type)))
+
def test_basic_type(self):
test_types = [DataTypes.STRING(),
DataTypes.BOOLEAN(),
diff --git a/flink-python/pyflink/table/types.py
b/flink-python/pyflink/table/types.py
index 2e5354ba67b..f39d2fe5f1b 100644
--- a/flink-python/pyflink/table/types.py
+++ b/flink-python/pyflink/table/types.py
@@ -1707,7 +1707,7 @@ 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(logical_type.getPrecision(),
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):
@@ -1727,7 +1727,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(
+ logical_type.getPrecision(), 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)
@@ -2186,12 +2187,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.
+
+ By default, retain the nested-type restrictions of the pandas conversion
path.
"""
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)
@@ -2262,7 +2266,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.
"""
@@ -2310,18 +2314,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):
