This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 957ba37f42 [python] Fix Arrow large_string compatibility in PyPaimon
(#10002)
957ba37f42 is described below
commit 957ba37f42de01541e7e156cf58d3fb7d08f41fe
Author: Zouxxyy <[email protected]>
AuthorDate: Mon Sep 21 11:33:11 2026 +0800
[python] Fix Arrow large_string compatibility in PyPaimon (#10002)
---
paimon-python/dev/lint-python.sh | 2 +-
paimon-python/pypaimon/daft/daft_catalog.py | 24 +---
.../pypaimon/multimodal/lerobot/metadata.py | 5 +-
.../pypaimon/multimodal/lerobot/writer.py | 4 +-
paimon-python/pypaimon/ray/shuffle.py | 22 +---
paimon-python/pypaimon/schema/arrow_schema.py | 92 ++++++++++++++
paimon-python/pypaimon/schema/data_types.py | 2 +-
paimon-python/pypaimon/tests/arrow_schema_test.py | 139 +++++++++++++++++++++
paimon-python/pypaimon/tests/data_types_test.py | 23 ++++
paimon-python/pypaimon/tests/ray_sink_test.py | 15 +++
.../pypaimon/tests/test_ray_shuffle_helper.py | 30 -----
.../pypaimon/tests/write/table_write_test.py | 66 ++++++++++
.../pypaimon/write/postpone_batch_table_write.py | 4 +-
paimon-python/pypaimon/write/ray_datasink.py | 24 ++--
paimon-python/pypaimon/write/table_write.py | 29 ++---
15 files changed, 366 insertions(+), 115 deletions(-)
diff --git a/paimon-python/dev/lint-python.sh b/paimon-python/dev/lint-python.sh
index 79dc14ede9..8ce00549e3 100755
--- a/paimon-python/dev/lint-python.sh
+++ b/paimon-python/dev/lint-python.sh
@@ -198,7 +198,7 @@ function pytest_check() {
# 3.6/3.7 run a curated core subset (their dep ceiling rules out the
# vector/index/multimodal/blob suites); 3.10+ run the full suite.
if [ "$PYTHON_VERSION" = "3.6" ] || [ "$PYTHON_VERSION" = "3.7" ]; then
- TEST_DIR="pypaimon/tests/py36 pypaimon/tests/file_io_test.py"
+ TEST_DIR="pypaimon/tests/py36 pypaimon/tests/file_io_test.py
pypaimon/tests/arrow_schema_test.py"
echo "Running core test subset for Python $PYTHON_VERSION: $TEST_DIR"
else
TEST_DIR="pypaimon/tests pypaimon/acceptance
--ignore=pypaimon/tests/py36 --ignore=pypaimon/tests/e2e
--ignore=pypaimon/tests/torch_read_test.py"
diff --git a/paimon-python/pypaimon/daft/daft_catalog.py
b/paimon-python/pypaimon/daft/daft_catalog.py
index d59df52dac..0f67098d7e 100644
--- a/paimon-python/pypaimon/daft/daft_catalog.py
+++ b/paimon-python/pypaimon/daft/daft_catalog.py
@@ -22,7 +22,6 @@ from __future__ import annotations
from typing import TYPE_CHECKING, Any
-import pyarrow as pa
from pypaimon.catalog.catalog import Catalog as InnerCatalog
from pypaimon.catalog.catalog_exception import (
DatabaseNotExistException,
@@ -85,7 +84,7 @@ class PaimonCatalog(Catalog):
) -> Table:
import pypaimon
- pa_schema = _cast_large_types(schema.to_pyarrow_schema())
+ pa_schema = schema.to_pyarrow_schema()
partition_keys = [pf.field.name for pf in (partition_fields or [])]
primary_keys = list((properties or {}).get("primary_keys", []))
options = {k: str(v) for k, v in (properties or {}).items() if k !=
"primary_keys"} if properties else {}
@@ -303,24 +302,3 @@ def _to_paimon_table_ident(ident: Identifier) -> str |
None:
return f"{parts[0]}.{parts[1]}"
return None
return ident
-
-
-def _cast_large_types(arrow_schema: pa.Schema) -> pa.Schema:
- """Convert PyArrow schema to be compatible with pypaimon.
-
- pypaimon doesn't support large_string, so we convert it to regular string.
- large_binary is kept as-is because pypaimon 1.4+ maps it to the BLOB type.
- """
- new_fields = []
- need_conversion = False
-
- for field in arrow_schema:
- field_type = field.type
- if pa.types.is_large_string(field_type):
- field_type = pa.string()
- need_conversion = True
- new_fields.append(pa.field(field.name, field_type,
nullable=field.nullable, metadata=field.metadata))
-
- if need_conversion:
- return pa.schema(new_fields, metadata=arrow_schema.metadata)
- return arrow_schema
diff --git a/paimon-python/pypaimon/multimodal/lerobot/metadata.py
b/paimon-python/pypaimon/multimodal/lerobot/metadata.py
index df75ce9ab0..08c405191b 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/metadata.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/metadata.py
@@ -25,6 +25,7 @@ import pyarrow as pa
import pyarrow.parquet as pq
from pypaimon import Schema as PaimonSchema
+from pypaimon.schema.arrow_schema import arrow_schemas_compatible
from pypaimon.catalog.catalog_exception import (
TableAlreadyExistException,
TagNotExistException,
@@ -361,7 +362,7 @@ def _append_arrow(table, data):
def _overwrite_arrow(table, data):
target_schema = _target_schema(table)
- if not data.schema.equals(target_schema, check_metadata=False):
+ if not arrow_schemas_compatible(data.schema, target_schema):
raise ValueError(
"LeRobot component schema %s does not match target %s."
% (data.schema, target_schema))
@@ -400,7 +401,7 @@ def _append_arrow_tables(table, tables):
for data in tables:
if data.num_rows == 0:
continue
- if not data.schema.equals(target_schema, check_metadata=False):
+ if not arrow_schemas_compatible(data.schema, target_schema):
raise ValueError(
"LeRobot component schema %s does not match target %s."
% (data.schema, target_schema))
diff --git a/paimon-python/pypaimon/multimodal/lerobot/writer.py
b/paimon-python/pypaimon/multimodal/lerobot/writer.py
index 8fc152e6f1..adc5dd54d3 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/writer.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/writer.py
@@ -29,6 +29,7 @@ from pypaimon.catalog.catalog_exception import (
DatabaseNotExistException,
TableNotExistException,
)
+from pypaimon.schema.arrow_schema import arrow_schemas_compatible
from pypaimon.multimodal.arrow_utils import strict_arrow_table
from pypaimon.multimodal.hdf5 import _SnapshotRecorder
from pypaimon.multimodal.lerobot.metadata import (
@@ -420,8 +421,7 @@ class PaimonLeRobotWriter:
for name, identifier in identifiers.items()
}
for name, table in tables.items():
- if not _target_schema(table).equals(
- expected[name], check_metadata=False):
+ if not arrow_schemas_compatible(_target_schema(table),
expected[name]):
raise ValueError(
"LeRobot %s companion schema does not match "
"PaimonLeRobotWriter." % name)
diff --git a/paimon-python/pypaimon/ray/shuffle.py
b/paimon-python/pypaimon/ray/shuffle.py
index a3d5c2165d..f1a7fd6898 100644
--- a/paimon-python/pypaimon/ray/shuffle.py
+++ b/paimon-python/pypaimon/ray/shuffle.py
@@ -198,27 +198,7 @@ def _sort_by_partition_bucket_primary_key(
def _identity_batch(batch: pa.Table) -> pa.Table:
- # Some Ray versions promote ``string`` to ``large_string`` (and
- # ``binary`` to ``large_binary``) while materialising blocks for
- # ``groupby().map_groups``. Paimon's writer compares schemas with a
- # strict ``!=`` and rejects the large variants, so coerce them back
- # to the regular types here. Other Arrow types pass through.
- return _coerce_large_string_types(batch)
-
-
-def _coerce_large_string_types(batch: pa.Table) -> pa.Table:
- needs_cast = False
- fields = []
- for field in batch.schema:
- if pa.types.is_large_string(field.type):
- fields.append(field.with_type(pa.string()))
- needs_cast = True
- elif pa.types.is_large_binary(field.type):
- fields.append(field.with_type(pa.binary()))
- needs_cast = True
- else:
- fields.append(field)
- return batch.cast(pa.schema(fields)) if needs_cast else batch
+ return batch
def _make_bucket_udf(extractor, bucket_col):
diff --git a/paimon-python/pypaimon/schema/arrow_schema.py
b/paimon-python/pypaimon/schema/arrow_schema.py
new file mode 100644
index 0000000000..98a232033b
--- /dev/null
+++ b/paimon-python/pypaimon/schema/arrow_schema.py
@@ -0,0 +1,92 @@
+# 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.
+
+"""STRING compatibility and canonical Arrow input for core writers."""
+
+import pyarrow as pa
+
+
+def _normalize_string_type(data_type):
+ if pa.types.is_large_string(data_type):
+ return pa.string()
+ if pa.types.is_struct(data_type):
+ return pa.struct([field.with_type(_normalize_string_type(field.type))
for field in data_type])
+ if (pa.types.is_list(data_type) or pa.types.is_large_list(data_type)
+ or pa.types.is_fixed_size_list(data_type)):
+ field =
data_type.value_field.with_type(_normalize_string_type(data_type.value_type))
+ if pa.types.is_large_list(data_type):
+ return pa.large_list(field)
+ if pa.types.is_fixed_size_list(data_type):
+ return pa.list_(field, data_type.list_size)
+ return pa.list_(field)
+ if pa.types.is_map(data_type):
+ # Arrow 6/12 have no MapType.keys_sorted property, but type equality
+ # includes the sorting flag. Preserve it along with the child fields.
+ sorted_keys = data_type == pa.map_(data_type.key_field,
data_type.item_field, keys_sorted=True)
+ return pa.map_(
+
data_type.key_field.with_type(_normalize_string_type(data_type.key_type)),
+
data_type.item_field.with_type(_normalize_string_type(data_type.item_type)),
+ keys_sorted=sorted_keys)
+ return data_type
+
+
+def arrow_schemas_compatible(source, target, check_top_level_nullability=True,
allow_binary_compatibility=False):
+ """Compare ordered fields, ignoring metadata and STRING offset width.
+
+ The writer opts into its existing top-level nullability and binary /
+ fixed_size_binary policies. BLOB and numeric types remain distinct.
+ """
+ if source.names != target.names:
+ return False
+ for actual, expected in zip(source, target):
+ if check_top_level_nullability and actual.nullable !=
expected.nullable:
+ return False
+ if actual.type == expected.type or _normalize_string_type(actual.type)
== _normalize_string_type(expected.type):
+ continue
+ if allow_binary_compatibility and all(
+ pa.types.is_binary(t) or pa.types.is_fixed_size_binary(t)
+ for t in (actual.type, expected.type)):
+ continue
+ return False
+ return True
+
+
+def normalize_arrow_strings(data):
+ """Safely normalize large_string input before routing or buffering rows.
+
+ Only STRING layouts change. Keep all other types, names, nullability and
+ metadata. Conversion limits and nested cast support belong to PyArrow;
+ unsupported inputs fail before reaching a data writer.
+ """
+ schema = pa.schema([
+ field.with_type(_normalize_string_type(field.type)) for field in
data.schema
+ ], metadata=data.schema.metadata)
+ if data.schema.equals(schema, check_metadata=False):
+ return data
+ try:
+ if isinstance(data, pa.RecordBatch):
+ # RecordBatch.cast is unavailable in Arrow 6/12. Casting through
+ # Table produces one chunk per column, or no chunks for empty
input.
+ table = pa.Table.from_batches([data]).cast(schema, safe=True)
+ arrays = [column.chunk(0) if column.num_chunks else pa.array([],
type=field.type)
+ for column, field in zip(table.columns, schema)]
+ return pa.RecordBatch.from_arrays(arrays, schema=schema)
+ return data.cast(schema, safe=True)
+ except (ValueError, TypeError, NotImplementedError) as error:
+ raise ValueError(
+ "Cannot convert large_string input to string with PyArrow %s: %s"
+ % (pa.__version__, error)) from error
diff --git a/paimon-python/pypaimon/schema/data_types.py
b/paimon-python/pypaimon/schema/data_types.py
index 5cf0b3eef6..c73223b883 100755
--- a/paimon-python/pypaimon/schema/data_types.py
+++ b/paimon-python/pypaimon/schema/data_types.py
@@ -783,7 +783,7 @@ class PyarrowFieldParser:
type_name = 'DOUBLE'
elif types.is_boolean(pa_type):
type_name = 'BOOLEAN'
- elif types.is_string(pa_type):
+ elif types.is_string(pa_type) or types.is_large_string(pa_type):
type_name = 'STRING'
elif types.is_fixed_size_binary(pa_type):
type_name = f'BINARY({pa_type.byte_width})'
diff --git a/paimon-python/pypaimon/tests/arrow_schema_test.py
b/paimon-python/pypaimon/tests/arrow_schema_test.py
new file mode 100644
index 0000000000..7a555f8a27
--- /dev/null
+++ b/paimon-python/pypaimon/tests/arrow_schema_test.py
@@ -0,0 +1,139 @@
+# 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 glob
+import os
+import tempfile
+import unittest
+from unittest.mock import Mock
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+from parameterized import parameterized
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.schema.arrow_schema import arrow_schemas_compatible,
normalize_arrow_strings
+from pypaimon.write.table_write import TableWrite
+
+
+class ArrowSchemaTest(unittest.TestCase):
+ def test_compatibility_keeps_non_string_contracts(self):
+ target = pa.schema([
+ pa.field('id', pa.int32(), nullable=False),
+ ('text', pa.struct([pa.field('value', pa.string(),
nullable=False)])),
+ ('blob', pa.large_binary()),
+ ])
+ source = target.set(1, pa.field('text', pa.struct([
+ pa.field('value', pa.large_string(), nullable=False),
+ ])))
+ self.assertTrue(arrow_schemas_compatible(source, target))
+ for invalid in [
+ source.set(0, pa.field('id', pa.int64(), nullable=False)),
+ source.set(1, pa.field('text', pa.struct([('value',
pa.large_string())]))),
+ source.set(2, pa.field('blob', pa.binary())),
+ pa.schema(list(reversed(list(source)))),
+ ]:
+ self.assertFalse(arrow_schemas_compatible(invalid, target))
+
+ @parameterized.expand([(False,), (True,)])
+ def test_map_sorting_contract_is_preserved(self, sorted_keys):
+ small = pa.schema([('v', pa.map_(pa.string(), pa.string(),
keys_sorted=sorted_keys))])
+ large = pa.schema([('v', pa.map_(pa.large_string(), pa.large_string(),
keys_sorted=sorted_keys))])
+ different_sorting = pa.schema([('v', pa.map_(pa.string(), pa.string(),
keys_sorted=not sorted_keys))])
+ self.assertTrue(arrow_schemas_compatible(small, large))
+ self.assertFalse(arrow_schemas_compatible(large, different_sorting))
+
+ @parameterized.expand([(False, False), (False, True), (True, False),
(True, True)])
+ def test_normalization_preserves_schema_and_other_buffers(self,
record_batch, empty):
+ schema = pa.schema([
+ pa.field('id', pa.int64(), nullable=False),
+ pa.field('text', pa.large_string(), metadata={b'description':
b'label'}),
+ ('blob', pa.large_binary()),
+ ('bytes', pa.binary(1)),
+ ], metadata={b'source': b'input'})
+ values = [[], [], [], []] if empty else [[1, 2], ['中文', None],
[b'blob', None], [b'x', b'y']]
+ arrays = [pa.array(items, type=field.type) for items, field in
zip(values, schema)]
+ factory = pa.RecordBatch if record_batch else pa.Table
+ source = factory.from_arrays(arrays, schema=schema)
+ result = normalize_arrow_strings(source)
+ self.assertIsInstance(result, factory)
+ self.assertEqual(result.schema, schema.set(1,
schema[1].with_type(pa.string())))
+ self.assertEqual(result.schema.metadata, schema.metadata)
+ self.assertEqual(result.schema[1].metadata, schema[1].metadata)
+ self.assertEqual(result.to_pydict(), source.to_pydict())
+ self.assertIs(normalize_arrow_strings(result), result)
+ if not empty:
+ blob = result.column(2) if record_batch else
result.column(2).chunk(0)
+ self.assertEqual(blob.buffers()[2].address,
arrays[2].buffers()[2].address)
+
+ @parameterized.expand([('table',), ('batch',), ('bucket',), ('postpone',),
('postpone_batch',)])
+ def test_core_writes_canonical_strings(self, entry):
+ with tempfile.TemporaryDirectory() as directory:
+ source_schema = pa.schema([
+ pa.field('id', pa.int64(), nullable=False), ('text',
pa.large_string()),
+ ])
+ catalog = CatalogFactory.create({'warehouse': directory})
+ catalog.create_database('default', False)
+ catalog.create_table('default.strings', Schema.from_pyarrow_schema(
+ source_schema, primary_keys=['id'] if
entry.startswith('postpone') else [],
+ options={'bucket': '-2' if entry.startswith('postpone') else
'1', 'file.format': 'parquet'},
+ ), False)
+ table = catalog.get_table('default.strings')
+ builder = (table.new_postpone_fixed_bucket_write_builder()
+ if entry.startswith('postpone') else
table.new_batch_write_builder())
+ writer, commit = builder.new_write(), builder.new_commit()
+ expected = {'id': [1, 2], 'text': ['中文', None]}
+ source = pa.Table.from_pydict(expected, schema=source_schema)
+ try:
+ if entry in ('batch', 'postpone_batch'):
+ writer.write_arrow_batch(source.to_batches()[0])
+ elif entry == 'bucket':
+ writer.write_arrow_batch_to_bucket(source.to_batches()[0],
0)
+ else:
+ writer.write_arrow(source)
+ commit.commit(writer.prepare_commit())
+ except Exception:
+ writer.abort()
+ raise
+ finally:
+ writer.close()
+ commit.close()
+ files = glob.glob(os.path.join(directory, '**', '*.parquet'),
recursive=True)
+ self.assertTrue(files)
+ for path in files:
+ self.assertEqual(pq.read_schema(path).field('text').type,
pa.string())
+ reader = table.new_read_builder()
+ result =
reader.new_read().to_arrow(reader.new_scan().plan().splits())
+ actual = result.to_pydict()
+ self.assertEqual(sorted(zip(actual['id'], actual['text'])), [(1,
'中文'), (2, None)])
+
+ @unittest.skipUnless(int(pa.__version__.split('.')[0]) == 6, 'Arrow 6
lacks struct cast kernels')
+ def test_unsupported_nested_cast_fails_before_routing(self):
+ writer = object.__new__(TableWrite)
+ writer.file_store_write = Mock(write_cols=None)
+ writer.row_key_extractor = Mock()
+ writer.table_pyarrow_schema = pa.schema([('nested',
pa.struct([('text', pa.string())]))])
+ source = pa.Table.from_pydict(
+ {'nested': [{'text': '中文'}]},
+ schema=pa.schema([('nested', pa.struct([('text',
pa.large_string())]))]),
+ )
+ with self.assertRaisesRegex(ValueError, 'Cannot convert large_string
input to string'):
+ writer.write_arrow(source)
+ with self.assertRaisesRegex(ValueError, 'Cannot convert large_string
input to string'):
+ writer.write_arrow_batch(source.to_batches()[0])
+
writer.row_key_extractor.extract_partition_bucket_groups.assert_not_called()
+ writer.file_store_write.write.assert_not_called()
diff --git a/paimon-python/pypaimon/tests/data_types_test.py
b/paimon-python/pypaimon/tests/data_types_test.py
index e96c612fb8..8ccef3a165 100755
--- a/paimon-python/pypaimon/tests/data_types_test.py
+++ b/paimon-python/pypaimon/tests/data_types_test.py
@@ -25,6 +25,29 @@ from pypaimon.schema.data_types import (DataField,
AtomicType, ArrayType, Multis
class DataTypesTest(unittest.TestCase):
+ def test_large_string_schema_preserves_field_contract(self):
+ from pypaimon.schema.schema import Schema
+
+ arrow_schema = pa.schema([
+ pa.field('text', pa.large_string(), nullable=False,
+ metadata={b'description': b'task label'}),
+ pa.field('nested', pa.struct([
+ pa.field('labels', pa.list_(pa.large_string())),
+ pa.field('mapping', pa.map_(pa.large_string(),
pa.large_string())),
+ ])),
+ ])
+ schema = Schema.from_pyarrow_schema(arrow_schema)
+ restored = PyarrowFieldParser.from_paimon_schema(schema.fields)
+ self.assertEqual(restored, pa.schema([
+ pa.field('text', pa.string(), nullable=False,
+ metadata={b'description': b'task label'}),
+ pa.field('nested', pa.struct([
+ pa.field('labels', pa.list_(pa.string())),
+ pa.field('mapping', pa.map_(pa.string(), pa.string())),
+ ])),
+ ]))
+ self.assertEqual(schema.fields[0].description, 'task label')
+
def test_atomic_type(self):
self.assertEqual(str(AtomicType("BLOB")), "BLOB")
self.assertEqual(str(AtomicType("TINYINT", nullable=False)), "TINYINT
NOT NULL")
diff --git a/paimon-python/pypaimon/tests/ray_sink_test.py
b/paimon-python/pypaimon/tests/ray_sink_test.py
index 41e7365e5f..d5db6f3074 100644
--- a/paimon-python/pypaimon/tests/ray_sink_test.py
+++ b/paimon-python/pypaimon/tests/ray_sink_test.py
@@ -27,6 +27,7 @@ from pypaimon import CatalogFactory, Schema
from pypaimon.write.ray_datasink import (
PaimonDatasink,
_consume_write_results,
+ _cast_binary_to_table_schema,
)
from pypaimon.write.commit_message import CommitMessage
from pypaimon.write.table_write import TableWrite
@@ -78,6 +79,20 @@ class RaySinkTest(unittest.TestCase):
data_files.append(os.path.join(root, file_name))
return data_files
+ def
test_binary_restoration_uses_table_contract_and_preserves_strings(self):
+ batch = pa.table({
+ 'blob': pa.array([b'x', None], type=pa.binary()),
+ 'bytes': pa.array([b'y', b'z'], type=pa.large_binary()),
+ 'text': pa.array(['中文', ''], type=pa.large_string()),
+ }).replace_schema_metadata({b'source': b'ray'})
+ target = pa.schema([('blob', pa.large_binary()), ('bytes',
pa.binary()), ('text', pa.string())])
+ result = _cast_binary_to_table_schema(batch, target)
+ self.assertEqual(result.to_pydict(), batch.to_pydict())
+ self.assertEqual(result.schema.field('blob').type, pa.large_binary())
+ self.assertEqual(result.schema.field('bytes').type, pa.binary())
+ self.assertEqual(result.schema.field('text').type, pa.large_string())
+ self.assertEqual(result.schema.metadata, batch.schema.metadata)
+
def test_init_and_serialization(self):
"""Test initialization, serialization, and table name."""
datasink = PaimonDatasink(self.table, overwrite=False)
diff --git a/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
b/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
index 6cfa5ea5f6..8e66795f6d 100644
--- a/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
+++ b/paimon-python/pypaimon/tests/test_ray_shuffle_helper.py
@@ -43,7 +43,6 @@ _SHUFFLE = importlib.util.module_from_spec(_SHUFFLE_SPEC)
_SHUFFLE_SPEC.loader.exec_module(_SHUFFLE)
BUCKET_KEY_COL = _SHUFFLE.BUCKET_KEY_COL
-_coerce_large_string_types = _SHUFFLE._coerce_large_string_types
_make_bucket_udf = _SHUFFLE._make_bucket_udf
_pick_bucket_col_name = _SHUFFLE._pick_bucket_col_name
maybe_apply_repartition = _SHUFFLE.maybe_apply_repartition
@@ -118,35 +117,6 @@ class PickBucketColNameTest(unittest.TestCase):
self.assertNotIn(name, {"id", BUCKET_KEY_COL})
-class CoerceLargeStringTypesTest(unittest.TestCase):
- """``_identity_batch`` casts back the large_string / large_binary
- types that some Ray versions introduce when materialising blocks
- during ``groupby().map_groups``. The Paimon writer's strict schema
- check would otherwise reject those rows."""
-
- def test_pass_through_when_no_large_variants(self):
- batch = pa.table({"id": pa.array([1, 2], type=pa.int32()),
- "name": pa.array(["a", "b"], type=pa.string())})
- out = _coerce_large_string_types(batch)
- self.assertEqual(out.schema, batch.schema)
-
- def test_casts_large_string_back_to_string(self):
- batch = pa.table({
- "id": pa.array([1, 2], type=pa.int32()),
- "name": pa.array(["x", "y"], type=pa.large_string()),
- })
- out = _coerce_large_string_types(batch)
- self.assertEqual(out.schema.field("name").type, pa.string())
- self.assertEqual(out.column("name").to_pylist(), ["x", "y"])
-
- def test_casts_large_binary_back_to_binary(self):
- batch = pa.table({
- "blob": pa.array([b"x", b"y"], type=pa.large_binary()),
- })
- out = _coerce_large_string_types(batch)
- self.assertEqual(out.schema.field("blob").type, pa.binary())
-
-
class BucketModeDispatchTest(unittest.TestCase):
"""``maybe_apply_repartition`` clusters only supported HASH_FIXED
writes and rejects unsafe primary-key Ray writes."""
diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py
b/paimon-python/pypaimon/tests/write/table_write_test.py
index 04e555a73c..f2875f4030 100644
--- a/paimon-python/pypaimon/tests/write/table_write_test.py
+++ b/paimon-python/pypaimon/tests/write/table_write_test.py
@@ -37,6 +37,72 @@ from pypaimon.write.writer.append_only_data_writer import
AppendOnlyDataWriter
class TableWriteTest(unittest.TestCase):
+ def test_large_string_compatibility_rejects_other_type_changes(self):
+ writer = object.__new__(TableWrite)
+ writer.file_store_write = Mock(write_cols=None)
+ writer.table_pyarrow_schema = pa.schema([
+ ('id', pa.int32()),
+ ('nested', pa.struct([pa.field('text', pa.string(),
nullable=False)])),
+ ])
+ valid = pa.schema([
+ ('id', pa.int32()),
+ ('nested', pa.struct([pa.field('text', pa.large_string(),
nullable=False)])),
+ ])
+ writer._validate_pyarrow_schema(valid)
+ invalid = [
+ valid.set(0, pa.field('id', pa.int64())),
+ pa.schema(list(reversed(list(valid)))),
+ valid.set(1, pa.field('nested', pa.struct([
+ pa.field('text', pa.large_string(), nullable=True)]))),
+ valid.set(1, pa.field('nested', pa.struct([
+ pa.field('renamed', pa.large_string(), nullable=False)]))),
+ ]
+ for schema in invalid:
+ with self.subTest(schema=schema):
+ with self.assertRaisesRegex(ValueError, 'consistent'):
+ writer._validate_pyarrow_schema(schema)
+
+ @parameterized.expand([('append', False, '-1'), ('primary', True, '1'),
('postpone', True, '-2')])
+ def test_large_string_input_roundtrip(self, mode, primary_key, bucket):
+ name = 'default.large_strings_' + mode
+
+ def arrow_schema(string_type):
+ return pa.schema([
+ pa.field('id', pa.int64(), nullable=False),
+ pa.field('text', string_type),
+ pa.field('nested', pa.struct([
+ pa.field('labels', pa.list_(string_type)),
+ pa.field('mapping', pa.map_(string_type, string_type)),
+ ])),
+ ])
+
+ self.catalog.create_table(name, Schema.from_pyarrow_schema(
+ arrow_schema(pa.large_string()),
+ primary_keys=['id'] if primary_key else [],
+ options={'bucket': bucket},
+ ), False)
+ table = self.catalog.get_table(name)
+ builder = (table.new_postpone_fixed_bucket_write_builder()
+ if mode == 'postpone' else table.new_batch_write_builder())
+ write, commit = builder.new_write(), builder.new_commit()
+ expected = []
+ try:
+ for index, (dtype, text) in enumerate([
+ (pa.string(), '抓笔'), (pa.large_string(), ''),
+ (pa.string(), None), (pa.large_string(), 'x' * 100)]):
+ row = {'id': index, 'text': text, 'nested': {
+ 'labels': [text, None], 'mapping': [('任务', text)],
+ }}
+ expected.append(row)
+ write.write_arrow(pa.Table.from_pylist([row],
schema=arrow_schema(dtype)))
+ commit.commit(write.prepare_commit())
+ finally:
+ write.close()
+ commit.close()
+ actual = self._read_sorted(table, [('id', 'ascending')])
+ self.assertEqual(actual.to_pylist(), expected)
+ self.assertEqual(actual.schema.remove_metadata(),
arrow_schema(pa.string()))
+
@classmethod
def setUpClass(cls):
cls.tempdir = tempfile.mkdtemp()
diff --git a/paimon-python/pypaimon/write/postpone_batch_table_write.py
b/paimon-python/pypaimon/write/postpone_batch_table_write.py
index 7ef2ab9e2b..f4fc24f70b 100644
--- a/paimon-python/pypaimon/write/postpone_batch_table_write.py
+++ b/paimon-python/pypaimon/write/postpone_batch_table_write.py
@@ -118,16 +118,16 @@ class PostponeFixedBucketBatchTableWrite(BatchTableWrite):
)
def write_arrow(self, table: pa.Table):
+ table = self._prepare_arrow_data(table)
if not self._buffer_input(table):
return super().write_arrow(table)
- self._validate_pyarrow_schema(table.schema)
self._pending_inputs.extend(
("batch", batch) for batch in table.to_batches())
def write_arrow_batch(self, data: pa.RecordBatch):
+ data = self._prepare_arrow_data(data)
if not self._buffer_input(data):
return super().write_arrow_batch(data)
- self._validate_pyarrow_schema(data.schema)
self._pending_inputs.append(("batch", data))
def write_row(self, row):
diff --git a/paimon-python/pypaimon/write/ray_datasink.py
b/paimon-python/pypaimon/write/ray_datasink.py
index baf8d49597..1fb56bbcac 100644
--- a/paimon-python/pypaimon/write/ray_datasink.py
+++ b/paimon-python/pypaimon/write/ray_datasink.py
@@ -67,15 +67,17 @@ class _TaskCommitMessages(list):
def _cast_binary_to_table_schema(table: pa.Table, target_schema: pa.Schema) ->
pa.Table:
- """Cast binary to large_binary for BLOB fields.
+ """Restore binary layouts using the table's BYTES/BLOB contract.
- When map_batches returns Python dicts, PyArrow infers bytes as binary,
- losing the original large_binary (BLOB) type. Cast back before writing.
+ Python dict inference loses BLOB's large_binary layout, while Ray block
+ promotion may widen BYTES. Neither transformation changes the table type.
"""
cast_indices = []
for i, field in enumerate(table.schema):
target_field = target_schema.field(field.name) if field.name in
target_schema.names else None
- if target_field and pa.types.is_binary(field.type) and
pa.types.is_large_binary(target_field.type):
+ if (target_field and field.type != target_field.type
+ and all(pa.types.is_binary(t) or pa.types.is_large_binary(t)
+ for t in (field.type, target_field.type))):
cast_indices.append(i)
if not cast_indices:
@@ -83,10 +85,10 @@ def _cast_binary_to_table_schema(table: pa.Table,
target_schema: pa.Schema) -> p
columns = table.columns
for i in cast_indices:
- columns[i] = columns[i].cast(pa.large_binary())
+ columns[i] =
columns[i].cast(target_schema.field(table.schema[i].name).type, safe=True)
fields = [target_schema.field(f.name) if i in cast_indices else f
for i, f in enumerate(table.schema)]
- return pa.table(columns, schema=pa.schema(fields))
+ return pa.table(columns, schema=pa.schema(fields,
metadata=table.schema.metadata))
# Python 3.8 / Ray 2.10: Datasink is not subscriptable at runtime
try:
@@ -395,7 +397,6 @@ def _write_postpone_primary_key_blocks(
import pickle
from pypaimon.ray.shuffle import (
- _coerce_large_string_types,
_sort_by_partition_bucket_primary_key,
)
@@ -415,9 +416,7 @@ def _write_postpone_primary_key_blocks(
error_col: pa.array([], type=pa.string()),
})
- rows = _coerce_large_string_types(
- batch.drop_columns(routing_columns)
- )
+ rows = batch.drop_columns(routing_columns)
worker_sink = PaimonDatasink(
captured_table,
overwrite=overwrite,
@@ -570,7 +569,6 @@ def _write_primary_key_groups(
import pickle
from pypaimon.ray.shuffle import (
- _coerce_large_string_types,
_group_by_partition_bucket,
)
@@ -590,9 +588,7 @@ def _write_primary_key_groups(
error_col: pa.array([], type=pa.string()),
})
- rows = _coerce_large_string_types(
- group.drop_columns([bucket_col])
- )
+ rows = group.drop_columns([bucket_col])
worker_sink = PaimonDatasink(
captured_table,
overwrite=overwrite,
diff --git a/paimon-python/pypaimon/write/table_write.py
b/paimon-python/pypaimon/write/table_write.py
index fd09282c1d..eea1319e12 100644
--- a/paimon-python/pypaimon/write/table_write.py
+++ b/paimon-python/pypaimon/write/table_write.py
@@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional
import pyarrow as pa
+from pypaimon.schema.arrow_schema import arrow_schemas_compatible,
normalize_arrow_strings
from pypaimon.schema.data_types import PyarrowFieldParser
from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.table.row.blob import BlobConsumer
@@ -53,13 +54,13 @@ class TableWrite:
)
def write_arrow(self, table: pa.Table):
- self._validate_pyarrow_schema(table.schema)
+ table = self._prepare_arrow_data(table)
batches_iterator = table.to_batches()
for batch in batches_iterator:
self.write_arrow_batch(batch)
def write_arrow_batch(self, data: pa.RecordBatch):
- self._validate_pyarrow_schema(data.schema)
+ data = self._prepare_arrow_data(data)
for partition, bucket, row_indices in \
self.row_key_extractor.extract_partition_bucket_groups(data):
@@ -141,7 +142,7 @@ class TableWrite:
if not isinstance(self.row_key_extractor,
DynamicBucketRowKeyExtractor):
if bucket_mode == BucketMode.HASH_DYNAMIC:
raise RuntimeError("Dynamic bucket extractor is not
configured")
- self._validate_pyarrow_schema(data.schema)
+ data = self._prepare_arrow_data(data)
if bucket_mode == BucketMode.HASH_DYNAMIC:
if key_hashes is None:
partition =
self.row_key_extractor.notify_precomputed_bucket_batch(
@@ -344,6 +345,10 @@ class TableWrite:
if release is not None:
release()
+ def _prepare_arrow_data(self, data):
+ self._validate_pyarrow_schema(data.schema)
+ return normalize_arrow_strings(data)
+
def _validate_pyarrow_schema(self, data_schema: pa.Schema):
if self._is_compatible_pyarrow_schema(data_schema,
self.table_pyarrow_schema):
return
@@ -358,18 +363,8 @@ class TableWrite:
def _is_compatible_pyarrow_schema(
self, data_schema: pa.Schema, expected_schema: pa.Schema) -> bool:
- # Allow compatible binary types: binary, fixed_size_binary[N] are
interchangeable
- if data_schema.names != expected_schema.names:
- return False
- for i in range(len(data_schema)):
- input_type = data_schema.field(i).type
- expected_type = expected_schema.field(i).type
- if input_type == expected_type:
- continue
- if self._is_binary_family(input_type) and
self._is_binary_family(expected_type):
- continue
- return False
- return True
+ return arrow_schemas_compatible(
+ data_schema, expected_schema, check_top_level_nullability=False,
allow_binary_compatibility=True)
def _write_cols_pyarrow_schema(self, write_cols: List[str]) -> pa.Schema:
table_fields = {
@@ -383,10 +378,6 @@ class TableWrite:
f"Table schema is: {self.table_pyarrow_schema} "
f"Write cols is: {self.file_store_write.write_cols}")
- @staticmethod
- def _is_binary_family(arrow_type) -> bool:
- return pa.types.is_binary(arrow_type) or
pa.types.is_fixed_size_binary(arrow_type)
-
class BatchTableWrite(TableWrite):
def __init__(self, table, commit_user, static_partition: Optional[dict] =
None):
