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 3ffb4e4ce8 [python] Support Variant extraction in native row iterators
(#10842)
3ffb4e4ce8 is described below
commit 3ffb4e4ce8524d91ccc7ad6d27544ab324c61cac
Author: QuakeWang <[email protected]>
AuthorDate: Wed Oct 28 15:27:50 2026 +0800
[python] Support Variant extraction in native row iterators (#10842)
---
paimon-python/pypaimon/read/read_dispatch.py | 7 -
paimon-python/pypaimon/read/table_read.py | 2 +-
.../tests/named_projection_iterator_test.py | 12 +-
.../pypaimon/tests/native_variant_iterator_test.py | 210 +++++++++++++++++++++
paimon-python/pypaimon/tests/variant_test.py | 4 +-
5 files changed, 219 insertions(+), 16 deletions(-)
diff --git a/paimon-python/pypaimon/read/read_dispatch.py
b/paimon-python/pypaimon/read/read_dispatch.py
index 9c3fc4bcbc..6a78ae8fce 100644
--- a/paimon-python/pypaimon/read/read_dispatch.py
+++ b/paimon-python/pypaimon/read/read_dispatch.py
@@ -92,13 +92,6 @@ def to_iterator(self, splits: List[Split]) -> Iterator:
Python split reader, the iterator reuses its row view and visits splits
in order. Close the iterator when stopping before exhaustion.
"""
- from pypaimon.read.variant_read_type import has_variant_extractions
-
- if (self.output_projection is not None and self.output_projection.named
- and has_variant_extractions(self.read_type)):
- raise RuntimeError(
- "Variant extraction expressions are not supported by
to_iterator(); "
- "use to_arrow() or to_arrow_batch_reader()")
self._begin_auth_read(splits)
limit = self.limit
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index 7761aea4d4..45f50bd5da 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -107,7 +107,7 @@ class TableRead:
in order. Close the iterator when stopping before exhaustion.
Named columns and nested paths retain their projection order. Named
- VARIANT extraction expressions require the Arrow APIs.
+ VARIANT extraction expressions require native reading.
"""
from pypaimon.read.read_dispatch import to_iterator
return to_iterator(self, splits)
diff --git a/paimon-python/pypaimon/tests/named_projection_iterator_test.py
b/paimon-python/pypaimon/tests/named_projection_iterator_test.py
index 9f0936af40..1f54d3dc16 100644
--- a/paimon-python/pypaimon/tests/named_projection_iterator_test.py
+++ b/paimon-python/pypaimon/tests/named_projection_iterator_test.py
@@ -258,21 +258,23 @@ def
test_named_projection_keeps_non_streaming_torch_row_boundary(projected_table
@pytest.mark.parametrize('function', ['variant_get', 'try_variant_get'])
@pytest.mark.parametrize('nested', [False, True])
@pytest.mark.parametrize('mixed', [False, True])
-def test_named_iterator_rejects_variant_extractions(tmp_path, function,
nested, mixed):
+def
test_variant_iterator_requires_native_and_keeps_torch_restriction(tmp_path,
function, nested, mixed):
catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
catalog.create_database('db', True)
payload = DataField(2 if nested else 1, 'payload', AtomicType('VARIANT'))
fields = [DataField(0, 'id', AtomicType('BIGINT')),
DataField(1, 'profile', RowType(True, [payload])) if nested else
payload]
- catalog.create_table('db.t', Schema(fields, options={'bucket': '-1'}),
False)
+ catalog.create_table('db.t', Schema(fields, options={
+ 'bucket': '-1', 'read.native.enabled': 'false',
'scan.native-plan.enabled': 'false'}), False)
source = 'profile.payload' if nested else 'payload'
projection = {'x': "%s(%s, '$.x', 'float')" % (function, source)}
if mixed:
projection['user_id'] = 'id'
read =
catalog.get_table('db.t').new_read_builder().with_projection(projection).new_read()
- with patch('pypaimon.read.read_dispatch._native_read_execution') as native:
- with pytest.raises(RuntimeError, match='Variant extraction expressions
are not supported by to_iterator'):
- read.to_iterator([])
+ with patch.object(NativeTableRead, 'read_rows') as native, \
+ patch.object(PythonTableRead, 'read_rows',
side_effect=AssertionError('Python Variant execution')):
+ with pytest.raises(RuntimeError, match='Variant extraction requires a
compatible native reader'):
+ list(read.to_iterator([]))
for streaming in (False, True):
with pytest.raises(RuntimeError, match='not supported by Torch row
format'):
read.to_torch([], streaming=streaming, batch_format='row')
diff --git a/paimon-python/pypaimon/tests/native_variant_iterator_test.py
b/paimon-python/pypaimon/tests/native_variant_iterator_test.py
new file mode 100644
index 0000000000..7cd902fcbd
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_variant_iterator_test.py
@@ -0,0 +1,210 @@
+# 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.
+
+"""Native Variant projections retain the public InternalRow iterator
contract."""
+
+from unittest.mock import patch
+
+import pyarrow as pa
+import pytest
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.data.generic_variant import GenericVariant
+from pypaimon.execution.native.read.table_read import native_read
+from pypaimon.execution.native.runtime import native_runtime_available
+from pypaimon.execution.python.read.table_read import PythonTableRead
+from pypaimon.schema.arrow_schema import read_arrow_schema
+from pypaimon.table.row.row_kind import RowKind
+
+
+pytestmark = pytest.mark.skipif(not native_runtime_available(),
reason='pypaimon-rust not installed')
+
+
[email protected](params=['append', 'pk', 'de'])
+def variant_table(tmp_path, request):
+ mode = request.param
+ payload = GenericVariant.to_arrow_array([
+ GenericVariant.from_python({'score': 1.25, 'good': 2.5}),
+ GenericVariant.from_python({'score': 'bad', 'good': 3.5}),
+ GenericVariant.from_python({}), None,
+ GenericVariant.from_python({'score': None}), None,
+ ])
+ options = {
+ 'bucket': '1' if mode == 'pk' else '-1',
+ 'file.format': 'parquet', 'read.batch-size': '2',
+ 'source.split.target-size': '1b',
+ 'read.native.enabled': 'true', 'scan.native-plan.enabled': 'false',
+ 'write.native.enabled': 'false', 'commit.native.enabled': 'false',
+ }
+ if mode == 'de':
+ options.update({'data-evolution.enabled': 'true',
'row-tracking.enabled': 'true'})
+ data = pa.Table.from_arrays([
+ pa.array(range(1, 7), type=pa.int64()), payload,
+ pa.StructArray.from_arrays([payload], names=['payload'],
mask=pa.array([False] * 5 + [True])),
+ ], names=['id', 'payload', 'profile'])
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('db', True)
+ catalog.create_table('db.t', Schema.from_pyarrow_schema(
+ data.schema, primary_keys=['id'] if mode == 'pk' else None,
options=options), False)
+ table = catalog.get_table('db.t')
+ for offset in range(0, data.num_rows, 2):
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow(data.slice(offset, 2))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ with patch.object(PythonTableRead, 'read_rows',
side_effect=AssertionError('Python fallback')), \
+ patch.object(PythonTableRead, 'read_batches',
side_effect=AssertionError('Python fallback')):
+ yield table
+
+
[email protected](params=['payload', 'profile.payload'])
+def variant_builder(variant_table, request):
+ source = request.param
+ builder = variant_table.new_read_builder().with_projection({
+ 'score': "try_variant_get(%s, '$.score', 'float')" % source,
+ 'identifier': 'id',
+ 'score_copy': "try_variant_get(%s, '$.score', 'float')" % source,
+ 'good': "variant_get(%s, '$.good', 'float')" % source,
+ 'identifier_copy': 'id',
+ })
+ return builder, source
+
+
+def _rows(iterator):
+ try:
+ result = []
+ for row in iterator:
+ assert row.get_row_kind() == RowKind.INSERT
+ result.append(tuple(row.get_field(index) for index in
range(len(row))))
+ return result
+ finally:
+ iterator.close()
+
+
[email protected]('limit', [None, 0, 1, 3])
+def test_variant_iterator_matches_arrow(variant_builder, limit):
+ builder, _ = variant_builder
+ splits = builder.new_scan().plan().splits()
+ assert len(splits) >= 2
+ if limit is not None:
+ builder.with_limit(limit)
+ read = builder.new_read()
+ expected = [(1.25, 1, 1.25, 2.5, 1), (None, 2, None, 3.5, 2)]
+ expected += [(None, value, None, None, value) for value in range(3, 7)]
+ if limit is not None:
+ expected = expected[:limit]
+ assert _rows(read.to_iterator(splits)) == expected
+ arrow = read.to_arrow(splits, parallelism=1)
+ assert arrow.column_names == [alias for alias, _ in
read.output_projection.columns]
+ assert arrow.schema.field('score').type == pa.float32()
+ assert list(zip(*(column.to_pylist() for column in arrow.columns))) ==
expected
+ with read.to_arrow_batch_reader(splits, parallelism=1) as batches:
+ assert [tuple(row.values()) for batch in batches for row in
batch.to_pylist()] == expected
+ assert _rows(read.to_iterator([])) == []
+ assert read.include_row_kind is False
+
+
+def test_variant_iterator_propagates_strict_cast_errors(variant_builder):
+ builder, source = variant_builder
+ builder.with_projection({'score': "variant_get(%s, '$.score', 'float')" %
source})
+ splits = builder.new_scan().plan().splits()
+ read = builder.new_read()
+ with pytest.raises(Exception, match='Cannot cast Variant'):
+ read.to_arrow(splits, parallelism=1)
+ with pytest.raises(Exception, match='Cannot cast Variant'):
+ _rows(read.to_iterator(splits))
+
+
[email protected]('failure', ['setup', 'unsupported_split'])
+def test_variant_iterator_never_falls_back_to_python(variant_builder, failure):
+ builder, _ = variant_builder
+ splits = builder.new_scan().plan().splits()
+ target, options = (
+ ('native_read', {'side_effect': RuntimeError('Native setup failed')})
if failure == 'setup' else
+ ('NativeTableRead._split_files_supported', {'return_value': False}))
+ with patch('pypaimon.execution.native.read.table_read.' + target,
**options):
+ with pytest.raises(RuntimeError, match='Variant extraction requires a
compatible native reader'):
+ _rows(builder.new_read().to_iterator(splits))
+
+
[email protected]('stop', ['unstarted', 'close', 'failure', 'exhaust'])
+def test_variant_iterator_batch_lifetime(variant_builder, stop):
+ fetched, closed = [], []
+ failure = OSError('Native stream failed')
+
+ def tracked_native(*args, **kwargs):
+ batches = native_read(*args, **kwargs)
+
+ def tracked_batches():
+ try:
+ for batch in batches:
+ fetched.append(batch.num_rows)
+ yield batch
+ if stop == 'failure':
+ raise failure
+ finally:
+ batches.close()
+ closed.append(True)
+
+ return tracked_batches()
+
+ builder, _ = variant_builder
+ splits = builder.new_scan().plan().splits()
+ with patch('pypaimon.execution.native.read.table_read.native_read',
side_effect=tracked_native) as native:
+ iterator = builder.new_read().to_iterator(splits)
+ native.assert_not_called()
+ try:
+ if stop != 'unstarted':
+ assert next(iterator).get_field(0) == 1.25
+ assert len(fetched) == 1
+ assert closed == []
+ if stop == 'failure':
+ with pytest.raises(OSError) as raised:
+ list(iterator)
+ assert raised.value is failure
+ elif stop == 'exhaust':
+ assert len(list(iterator)) == 5
+ finally:
+ iterator.close()
+ assert native.call_count == int(stop != 'unstarted')
+ assert closed == ([] if stop == 'unstarted' else [True])
+ if stop in ('close', 'failure'):
+ assert len(fetched) == 1
+ if stop == 'exhaust':
+ assert sum(fetched) == 6
+
+
+def test_variant_iterator_preserves_row_kinds(variant_builder):
+ builder, path = variant_builder
+ read = builder.new_read()
+ source = {'0': 1.25, '1': 1.25, '2': 2.5}
+ field = path.split('.')[0]
+ payload = {'payload': source} if field == 'profile' else source
+ kinds = ['+I', '-U', '+U', '-D']
+ batch = pa.RecordBatch.from_pylist([
+ {'rowkind': kind, field: payload, 'id': index + 1} for index, kind in
enumerate(kinds)
+ ], schema=pa.schema([pa.field('rowkind', pa.string())] +
list(read_arrow_schema(read.read_type))))
+ splits = builder.new_scan().plan().splits()
+ with patch('pypaimon.execution.native.read.table_read.native_read',
return_value=iter([batch])):
+ actual = [(tuple(row.get_field(index) for index in range(len(row))),
row.get_row_kind())
+ for row in read.to_iterator(splits)]
+ assert actual == [((1.25, index + 1, 1.25, 2.5, index + 1),
RowKind.from_string(kind))
+ for index, kind in enumerate(kinds)]
diff --git a/paimon-python/pypaimon/tests/variant_test.py
b/paimon-python/pypaimon/tests/variant_test.py
index dccaf07c85..b83694fca3 100644
--- a/paimon-python/pypaimon/tests/variant_test.py
+++ b/paimon-python/pypaimon/tests/variant_test.py
@@ -1262,7 +1262,7 @@ class TestVariantPaimonTable(unittest.TestCase):
pa.field('payload', _variant_arrow_type()),
])
- def test_native_variant_projection_rejects_row_readers(self):
+ def test_variant_projection_keeps_torch_row_restriction(self):
for mode, options in (
('append', {}),
('lazy', {
@@ -1295,8 +1295,6 @@ class TestVariantPaimonTable(unittest.TestCase):
self.assertTrue(splits)
read = builder.new_read()
- with self.assertRaisesRegex(RuntimeError, 'to_iterator'):
- read.to_iterator(splits)
with self.assertRaisesRegex(RuntimeError, 'Torch row format'):
read.to_torch(splits, streaming=True)
with self.assertRaisesRegex(RuntimeError, 'Torch row format'):
