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 1d23d3d23a [python] Fix lossy VARIANT shredding in Python writer
(#10633)
1d23d3d23a is described below
commit 1d23d3d23a2e379e8722e301a7cd2414c867330a
Author: XiaoHongbo <[email protected]>
AuthorDate: Sun Oct 18 08:59:52 2026 +0800
[python] Fix lossy VARIANT shredding in Python writer (#10633)
---
.../pypaimon/benchmark/variant_shredding.py | 109 ++++++++++++++++
paimon-python/pypaimon/data/variant_shredding.py | 51 +++++++-
paimon-python/pypaimon/tests/variant_test.py | 144 +++++++++++++++++++++
3 files changed, 299 insertions(+), 5 deletions(-)
diff --git a/paimon-python/pypaimon/benchmark/variant_shredding.py
b/paimon-python/pypaimon/benchmark/variant_shredding.py
new file mode 100644
index 0000000000..9302f007be
--- /dev/null
+++ b/paimon-python/pypaimon/benchmark/variant_shredding.py
@@ -0,0 +1,109 @@
+# 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.
+
+"""Benchmark matching DOUBLE shredding and Parquet serialization in memory.
+
+Example: python -m pypaimon.benchmark.variant_shredding --rows 20000 --fields
10
+Use --reference label=/path/to/variant_shredding.py to compare revisions.
+"""
+
+import argparse
+import hashlib
+import importlib.util
+import json
+import platform
+import statistics
+import sys
+import time
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+from pypaimon.data.generic_variant import GenericVariant
+from pypaimon.data import variant_shredding
+
+
+def main():
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument('--rows', type=int, default=20000)
+ parser.add_argument('--fields', type=int, default=10)
+ parser.add_argument('--repeats', type=int, default=5)
+ parser.add_argument('--warmup', type=int, default=1)
+ parser.add_argument('--reference', action='append', default=[],
metavar='LABEL=FILE')
+ args = parser.parse_args()
+ if min(args.rows, args.fields, args.repeats) < 1 or args.warmup < 0:
+ parser.error('rows, fields and repeats must be positive; warmup must
be non-negative')
+
+ modules = {'current': variant_shredding}
+ sources = {}
+ for reference in args.reference:
+ label, path = reference.split('=', 1)
+ if label in modules:
+ parser.error('duplicate label: ' + label)
+ spec = importlib.util.spec_from_file_location('_benchmark_' + label,
path)
+ module = importlib.util.module_from_spec(spec)
+ sys.modules[spec.name] = module
+ spec.loader.exec_module(module)
+ modules[label] = module
+ for label, module in modules.items():
+ with open(module.__file__, 'rb') as source:
+ sources[label] = hashlib.sha256(source.read()).hexdigest()
+
+ names = ['f{}'.format(i) for i in range(args.fields)]
+ column = GenericVariant.to_arrow_array([
+ GenericVariant.from_python({name: float(row + i) + 0.25
+ for i, name in enumerate(names)})
+ for row in range(args.rows)
+ ])
+ option = json.dumps({'type': 'ROW', 'fields': [{'name': 'payload', 'type':
{
+ 'type': 'ROW', 'fields': [{'name': name, 'type': 'DOUBLE'} for name in
names]
+ }}]})
+ schemas = {label: module.parse_shredding_schema_option(option)['payload']
+ for label, module in modules.items()}
+ types = {label: module.shredding_schema_to_arrow_type(schemas[label])
+ for label, module in modules.items()}
+ times = {label: [] for label in modules}
+ expected = None
+ labels = list(modules)
+ for repeat in range(args.warmup + args.repeats):
+ # Rotate execution order to reduce order bias.
+ order = labels[repeat % len(labels):] + labels[:repeat % len(labels)]
+ for label in order:
+ start = time.perf_counter()
+ shredded = modules[label].shred_variant_column(column,
schemas[label], types[label])
+ sink = pa.BufferOutputStream()
+ pq.write_table(pa.table({'payload': shredded}), sink,
compression='zstd')
+ output = sink.getvalue()
+ elapsed = time.perf_counter() - start
+ fingerprint = (output.size, hashlib.sha256(output).hexdigest())
+ if expected is None:
+ expected = fingerprint
+ if fingerprint != expected:
+ raise AssertionError('Parquet output differs: ' + label)
+ if repeat >= args.warmup:
+ times[label].append(elapsed)
+ print(json.dumps({
+ 'python': platform.python_version(), 'pyarrow': pa.__version__,
+ 'platform': platform.platform(), 'parameters': vars(args),
+ 'source_sha256': sources, 'parquet_bytes': expected[0],
+ 'parquet_sha256': expected[1], 'seconds': times,
+ 'median_seconds': {label: statistics.median(values) for label, values
in times.items()},
+ }, indent=2))
+
+
+if __name__ == '__main__':
+ main()
diff --git a/paimon-python/pypaimon/data/variant_shredding.py
b/paimon-python/pypaimon/data/variant_shredding.py
index ae95d5f800..8a033a6e71 100644
--- a/paimon-python/pypaimon/data/variant_shredding.py
+++ b/paimon-python/pypaimon/data/variant_shredding.py
@@ -52,6 +52,8 @@ from typing import Dict, List, Optional, Tuple
import pyarrow as pa
+from pypaimon.data.generic_variant import GenericVariant, _Type,
_variant_get_type
+
from pypaimon.data._variant_binary import (
_OBJECT,
_U8_MAX, _U32_SIZE,
@@ -856,6 +858,47 @@ def shredding_schema_to_arrow_type(obj_fields:
List[ObjectField]) -> pa.StructTy
])
+def _try_shred_scalar(field_bytes: bytes, metadata: bytes, arrow_type:
pa.DataType):
+ kind = _variant_get_type(field_bytes, 0)
+ if kind == _Type.DOUBLE:
+ compatible = pa.types.is_float64(arrow_type)
+ elif kind == _Type.FLOAT:
+ compatible = pa.types.is_float32(arrow_type)
+ elif kind == _Type.LONG:
+ compatible = pa.types.is_integer(arrow_type)
+ elif kind == _Type.BOOLEAN:
+ compatible = pa.types.is_boolean(arrow_type)
+ elif kind == _Type.STRING:
+ compatible = pa.types.is_string(arrow_type) or
pa.types.is_large_string(arrow_type)
+ elif kind == _Type.BINARY:
+ compatible = pa.types.is_binary(arrow_type) or
pa.types.is_large_binary(arrow_type)
+ elif kind == _Type.DATE:
+ compatible = pa.types.is_date32(arrow_type)
+ elif kind == _Type.DECIMAL:
+ compatible = pa.types.is_decimal(arrow_type)
+ elif kind in (_Type.TIMESTAMP, _Type.TIMESTAMP_NTZ):
+ compatible = (pa.types.is_timestamp(arrow_type)
+ and bool(arrow_type.tz) == (kind == _Type.TIMESTAMP))
+ else:
+ # JSON null stays in residual; two null columns mean absent.
+ return None
+ if not compatible:
+ return None
+
+ value = GenericVariant(field_bytes, metadata).to_python()
+ if kind == _Type.LONG:
+ bits = arrow_type.bit_width
+ lower = 0 if pa.types.is_unsigned_integer(arrow_type) else -(1 <<
(bits - 1))
+ if not lower <= value < lower + (1 << bits):
+ return None
+ elif kind in (_Type.DECIMAL, _Type.TIMESTAMP, _Type.TIMESTAMP_NTZ):
+ # Keep Arrow validation for scale, timestamp precision and overflow.
+ value = pa.scalar(value, type=arrow_type).as_py()
+ if _encode_scalar_to_value_bytes(value, arrow_type) != field_bytes:
+ return None
+ return value
+
+
def _decompose_field_bytes(
field_bytes: bytes,
schema: VariantSchema,
@@ -875,14 +918,12 @@ def _decompose_field_bytes(
A Python dict ``{'value': bytes_or_none, 'typed_value':
val_or_dict_or_none}``
suitable for building a PyArrow struct array row.
"""
- from pypaimon.data.generic_variant import GenericVariant
-
if schema.scalar_arrow_type is not None:
try:
- py_val = GenericVariant(field_bytes, metadata).to_python()
+ typed_val = _try_shred_scalar(field_bytes, metadata,
schema.scalar_arrow_type)
except Exception:
- return {'value': field_bytes, 'typed_value': None}
- return {'value': None, 'typed_value': py_val}
+ typed_val = None
+ return {'value': field_bytes if typed_val is None else None,
'typed_value': typed_val}
if schema.object_fields is not None:
if not field_bytes or (field_bytes[0] & 0x3) != _OBJECT:
diff --git a/paimon-python/pypaimon/tests/variant_test.py
b/paimon-python/pypaimon/tests/variant_test.py
index 4a27543fe0..658d766dd7 100644
--- a/paimon-python/pypaimon/tests/variant_test.py
+++ b/paimon-python/pypaimon/tests/variant_test.py
@@ -51,6 +51,7 @@ import unittest
import datetime
import uuid
from decimal import Decimal
+from unittest.mock import patch
import pyarrow as pa
import pyarrow.parquet as pq
@@ -60,6 +61,7 @@ from pypaimon.data.generic_variant import GenericVariant
from pypaimon.data.variant_shredding import (
VARIANT_ARROW_TYPE,
VariantSchema,
+ _decompose_field_bytes,
_NULL_VALUE_BYTES,
_build_array_value,
_build_object_value,
@@ -1161,6 +1163,57 @@ class TestShreddingWrite(unittest.TestCase):
overflow_py = GenericVariant(overflow,
result['metadata']).to_python()
self.assertIn('extra', overflow_py)
+ def test_scalar_shredding_preserves_types_and_values(self):
+ residual_cases = [
+ (0, pa.float64()), (2.5, pa.int64()), (2.0, pa.int64()),
+ (True, pa.int64()), (False, pa.float64()), (1, pa.bool_()),
+ (None, pa.int64()), ('12', pa.int64()), ([], pa.string()),
+ ({'x': 1}, pa.string()), (128, pa.int8()), (-129, pa.int8()),
+ (2 ** 31, pa.int32()), (2 ** 63 - 1, pa.float64()),
+ (1.25, pa.float32()), (Decimal('1.234'), pa.decimal128(4, 2)),
+ (Decimal('123.45'), pa.decimal128(4, 2)),
+ (datetime.datetime(2024, 1, 1, microsecond=1), pa.timestamp('ms')),
+ (datetime.datetime(2500, 1, 1), pa.timestamp('ns')),
+ (datetime.datetime(2024, 1, 1), pa.timestamp('us', tz='UTC')),
+ (datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc),
pa.timestamp('us')),
+ (True, pa.decimal128(4, 0)), (1, pa.decimal128(4, 0)),
+ (Decimal('1'), pa.int64()),
+ ]
+ typed_cases = [
+ (True, pa.bool_()), (-128, pa.int8()), (127, pa.int8()),
+ (2 ** 63 - 1, pa.int64()), (1.25, pa.float64()),
+ (-0.0, pa.float64()), ('hello', pa.string()), (b'abc',
pa.binary()),
+ (Decimal('12.34'), pa.decimal128(4, 2)),
+ (datetime.date(2024, 1, 1), pa.date32()),
+ (datetime.datetime(2024, 1, 1, microsecond=1), pa.timestamp('us')),
+ ]
+ for residual, cases in ((True, residual_cases), (False, typed_cases)):
+ for value, arrow_type in cases:
+ with self.subTest(value=value, arrow_type=arrow_type):
+ gv = GenericVariant.from_python(value)
+ schema = VariantSchema(value_idx=0, typed_idx=1,
+ scalar_arrow_type=arrow_type)
+ with patch('pypaimon.data.variant_shredding.pa.scalar',
wraps=pa.scalar) as scalar:
+ row = _decompose_field_bytes(gv.value(), schema,
gv.metadata())
+ if not (pa.types.is_decimal(arrow_type) or
pa.types.is_timestamp(arrow_type)):
+ scalar.assert_not_called()
+ self.assertEqual(row['value'], gv.value() if residual else
None)
+ if residual:
+ self.assertIsNone(row['typed_value'])
+ self.assertEqual(rebuild_value(row, schema, {}),
gv.value())
+
+ def test_float_wire_type_is_preserved(self):
+ metadata = GenericVariant.from_python(0).metadata()
+ for source in (pa.float32(), pa.float64()):
+ for target in (pa.float32(), pa.float64()):
+ with self.subTest(source=source, target=target):
+ raw = _encode_scalar_to_value_bytes(1.25, source)
+ schema = VariantSchema(value_idx=0, typed_idx=1,
+ scalar_arrow_type=target)
+ row = _decompose_field_bytes(raw, schema, metadata)
+ self.assertEqual(row['value'], None if source == target
else raw)
+ self.assertEqual(rebuild_value(row, schema, {}), raw)
+
# ===========================================================================
# 8. Full-chain via Paimon API
@@ -1424,6 +1477,97 @@ class TestVariantPaimonTable(unittest.TestCase):
with self.assertRaisesRegex(Exception, 'Cannot cast Variant'):
strict.new_read().to_arrow(strict.new_scan().plan().splits())
+ def _check_shredding_type_preservation(self, native):
+ def typed(value):
+ if isinstance(value, dict):
+ return ('object', sorted((k, typed(v)) for k, v in
value.items()))
+ if isinstance(value, list):
+ return ('array', [typed(v) for v in value])
+ return (type(value).__name__, value)
+
+ values = [
+ {'score': 1.5, 'count': 2, 'small': 127, 'nested': {'x': 3.25}},
+ {'score': 0, 'count': 2.5, 'small': 128, 'nested': {'x': 0}},
+ {'score': True, 'count': False, 'small': -129, 'nested': {'x':
False}},
+ {'score': None, 'count': None, 'nested': {'x': None}},
+ {}, {'nested': {}}, {'nested': None}, {'nested': [1, None]},
+ [1, None], 42, None, {'count': 2 ** 63 - 1},
+ ]
+ gvs = [GenericVariant.from_python(value) for value in values] + [None]
+ expected = {i: ('sql-null',) if gv is None else typed(gv.to_python())
+ for i, gv in enumerate(gvs)}
+ shredding = _schema_json('payload', [
+ ('score', 'DOUBLE'), ('count', 'BIGINT'), ('small', 'TINYINT'),
+ ('nested', {'type': 'ROW', 'fields': [{'name': 'x', 'type':
'DOUBLE'}]}),
+ ])
+ for shredded in (False, True):
+ name = 'default.type_preservation_{}_{}'.format(native, shredded)
+ options = {'bucket': '-1', 'file.format': 'parquet',
+ 'write.native.enabled': str(native).lower(),
+ 'read.native.enabled': 'false'}
+ if shredded:
+ options['parquet.variant.shreddingSchema'] = shredding
+ schema = Schema.from_pyarrow_schema(self._pa_schema())
+ schema.options.update(options)
+ self.catalog.create_table(name, schema, False)
+ table = self.catalog.get_table(name)
+ builder = table.new_batch_write_builder()
+ writer = builder.new_write()
+ self.assertEqual(type(writer).__name__,
+ 'NativeTableWrite' if native else
'BatchTableWrite')
+ commit = builder.new_commit()
+ try:
+ for start in range(0, len(gvs), 3):
+ batch = pa.RecordBatch.from_arrays([
+ pa.array(range(start, min(start + 3, len(gvs))),
type=pa.int64()),
+ GenericVariant.to_arrow_array(gvs[start:start + 3]),
+ ], schema=self._pa_schema())
+ writer.write_arrow_batch(batch)
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+
+ reader = table.new_read_builder()
+ rows =
reader.new_read().to_arrow(reader.new_scan().plan().splits()).to_pylist()
+ actual = {row['id']: ('sql-null',) if row['payload'] is None else
+
typed(GenericVariant.from_arrow_struct(row['payload']).to_python())
+ for row in rows}
+ self.assertEqual(actual, expected)
+ if not shredded:
+ continue
+ physical = {}
+ for root, _, files in os.walk(table.table_path):
+ for filename in files:
+ if filename.endswith('.parquet'):
+ for row in pq.read_table(os.path.join(root,
filename)).to_pylist():
+ physical[row['id']] = row['payload']
+ self.assertEqual(len(physical), len(gvs))
+ for index, field in ((1, 'score'), (1, 'count'), (1, 'small'),
+ (2, 'score'), (2, 'count'), (2, 'small'), (3,
'score')):
+ leaf = physical[index]['typed_value'][field]
+ self.assertIsNone(leaf['typed_value'])
+ self.assertEqual(leaf['value'],
GenericVariant.from_python(values[index][field]).value())
+ missing = physical[4]['typed_value']['score']
+ self.assertEqual(missing, {'value': None, 'typed_value': None})
+
self.assertEqual(physical[3]['typed_value']['nested']['typed_value']['x']['value'],
+ _NULL_VALUE_BYTES)
+ self.assertEqual(physical[10]['value'], _NULL_VALUE_BYTES)
+ self.assertIsNone(physical[12])
+
self.assertEqual(physical[0]['typed_value']['score']['typed_value'], 1.5)
+ self.assertIsNone(physical[0]['typed_value']['score']['value'])
+
+ def test_python_writer_preserves_shredded_variant_types(self):
+ self._check_shredding_type_preservation(native=False)
+
+ def test_native_writer_preserves_shredded_variant_types(self):
+ try:
+ import pypaimon_rust # noqa: F401
+ from pypaimon.write.native_write import NativeTableWrite # noqa:
F401
+ except ImportError:
+ self.skipTest('native writer is not available')
+ self._check_shredding_type_preservation(native=True)
+
def test_plain_variant_write_and_read(self):
"""Plain VARIANT: GenericVariant → write_arrow → to_arrow →
GenericVariant."""
schema = Schema.from_pyarrow_schema(self._pa_schema())
