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 8441f2ae3b [python] Support Variant projections in streaming Torch row
datasets (#10858)
8441f2ae3b is described below
commit 8441f2ae3b428b1e0b066a4bd15a19594c29ce12
Author: QuakeWang <[email protected]>
AuthorDate: Thu Oct 29 11:08:11 2026 +0800
[python] Support Variant projections in streaming Torch row datasets
(#10858)
---
.../pypaimon/read/datasource/torch_dataset.py | 6 +-
paimon-python/pypaimon/read/table_read.py | 11 +-
.../tests/named_projection_iterator_test.py | 13 +-
.../pypaimon/tests/native_torch_variant_test.py | 218 +++++++++++++++++++++
paimon-python/pypaimon/tests/torch_read_test.py | 16 ++
paimon-python/pypaimon/tests/variant_test.py | 24 ++-
6 files changed, 271 insertions(+), 17 deletions(-)
diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py
b/paimon-python/pypaimon/read/datasource/torch_dataset.py
index cfb102a0b6..9bebc92772 100644
--- a/paimon-python/pypaimon/read/datasource/torch_dataset.py
+++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py
@@ -588,7 +588,8 @@ class TorchIterDataset(_BaseTorchIterDataset):
worker_info = torch.utils.data.get_worker_info()
splits_to_process = self._worker_splits(worker_info)
- if self.prefetch_concurrency > 1:
+ # A binding limit needs one reader quota across all selected splits.
+ if self.prefetch_concurrency > 1 and self._limit_covers_all_splits():
yield from self._iter_rows(splits_to_process)
return
@@ -1004,7 +1005,8 @@ class TorchShuffledIterDataset(_TorchShuffleMixin,
_BaseTorchIterDataset):
worker_id = worker_info.id if worker_info is not None else 0
splits_to_process = self._worker_splits(worker_info)
- if self.max_buffer_input_splits == 1:
+ # Select the limited prefix before shuffling, as in the batch dataset.
+ if self.max_buffer_input_splits == 1 or not
self._limit_covers_all_splits():
rows = self._iter_ordered_rows(splits_to_process)
else:
rows = self._iter_interleaved_rows(splits_to_process)
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index 45f50bd5da..90abe65126 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -483,10 +483,12 @@ class TableRead:
eligible data-evolution reads fetch map-style batches lazily
by row ID.
prefetch_concurrency: Reader threads per DataLoader worker in row
- format.
+ format. Binding limits read serially to preserve the selected
+ prefix.
batch_format: ``"row"``, ``"pyarrow"``, or ``"torch"``. Batch
formats require streaming. Streaming row format also supports
- named column and path projections.
+ named columns, paths and float32 VARIANT extractions through
+ native reading.
batch_size: Rows per batch; ``None`` preserves reader batches when
unshuffled, or emits shuffle blocks of at most buffer_size
rows.
to_tensor_fn: Optional RecordBatch converter for Torch batches.
@@ -514,10 +516,7 @@ class TableRead:
):
raise ValueError("batch_size must be a positive int or None")
if batch_format == "row":
- from pypaimon.read.variant_read_type import has_variant_extractions
-
- if (self.output_projection is not None and
self.output_projection.named
- and (not streaming or
has_variant_extractions(self.read_type))):
+ if not streaming and self.output_projection is not None and
self.output_projection.named:
raise RuntimeError(
"Projection expressions are not supported by Torch row
format; "
"use streaming=True with batch_format='pyarrow' or
'torch'")
diff --git a/paimon-python/pypaimon/tests/named_projection_iterator_test.py
b/paimon-python/pypaimon/tests/named_projection_iterator_test.py
index 1f54d3dc16..85528d0555 100644
--- a/paimon-python/pypaimon/tests/named_projection_iterator_test.py
+++ b/paimon-python/pypaimon/tests/named_projection_iterator_test.py
@@ -258,7 +258,7 @@ 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_variant_iterator_requires_native_and_keeps_torch_restriction(tmp_path,
function, nested, mixed):
+def test_variant_iterator_and_torch_rows_require_native(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'))
@@ -275,7 +275,12 @@ def
test_variant_iterator_requires_native_and_keeps_torch_restriction(tmp_path,
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')
+ with pytest.raises(RuntimeError, match='not supported by Torch row
format'):
+ read.to_torch([], streaming=False, batch_format='row')
+ native.assert_not_called()
+
+ pytest.importorskip('torch')
+ dataset = read.to_torch([], streaming=True, batch_format='row')
+ with pytest.raises(RuntimeError, match='Variant extraction requires a
compatible native reader'):
+ list(dataset)
native.assert_not_called()
diff --git a/paimon-python/pypaimon/tests/native_torch_variant_test.py
b/paimon-python/pypaimon/tests/native_torch_variant_test.py
new file mode 100644
index 0000000000..30b09c4f47
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_torch_variant_test.py
@@ -0,0 +1,218 @@
+# 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.
+
+"""Streaming Torch rows retain Native Variant projection semantics."""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+pytest.importorskip('torch')
+
+from torch.utils.data import DataLoader
+
+from pypaimon.execution.native.read.table_read import native_read
+from pypaimon.execution.native.runtime import native_runtime_available
+from pypaimon.tests import native_variant_iterator_test
+
+
+pytestmark = pytest.mark.skipif(not native_runtime_available(),
reason='pypaimon-rust not installed')
+variant_table = native_variant_iterator_test.variant_table
+variant_builder = native_variant_iterator_test.variant_builder
+
+
[email protected](params=[
+ {}, {'prefetch_concurrency': 2},
+ {'shuffle': True, 'seed': 17, 'buffer_size': 1, 'max_buffer_input_splits':
1},
+ {'shuffle': True, 'seed': 17, 'buffer_size': 1, 'max_buffer_input_splits':
3},
+], ids=['serial', 'prefetch', 'shuffle_ordered', 'shuffle_interleaved'])
+def row_options(request):
+ return request.param
+
+
+def _expected_rows():
+ names = ['score', 'identifier', 'score_copy', 'good', 'identifier_copy']
+ values = [(1.25, 1, 1.25, 2.5, 1), (None, 2, None, 3.5, 2)]
+ values += [(None, value, None, None, value) for value in range(3, 7)]
+ return [dict(zip(names, row)) for row in values]
+
+
+def _assert_rows(actual, expected):
+ assert sorted(actual, key=lambda row: row['identifier']) == expected
+ assert all(list(row) == list(_expected_rows()[0]) for row in actual)
+ assert all(row['score'] is None or isinstance(row['score'], float) for row
in actual)
+
+
+def test_streaming_variant_rows(variant_builder, row_options):
+ builder, _ = variant_builder
+ splits = builder.new_scan().plan().splits()
+ assert len(splits) > 1
+ dataset = builder.new_read().to_torch(splits, streaming=True,
**row_options)
+ _assert_rows(list(DataLoader(dataset, batch_size=None, num_workers=0)),
_expected_rows())
+ if row_options.get('shuffle'):
+ assert list(dataset) == list(dataset)
+
+
+def test_streaming_variant_only_projection(variant_builder, row_options):
+ builder, source = variant_builder
+ builder.with_projection({'score': "try_variant_get(%s, '$.score',
'float')" % source})
+ splits = builder.new_scan().plan().splits()
+ actual = list(builder.new_read().to_torch(splits, streaming=True,
**row_options))
+ expected = [{'score': row['score']} for row in _expected_rows()]
+ assert sorted(actual, key=repr) == sorted(expected, key=repr)
+
+
[email protected]('limit', [0, 1, 3, 6])
[email protected]('plan_before_limit', [False, True])
+def test_streaming_variant_limit(variant_builder, row_options, limit,
plan_before_limit):
+ builder, _ = variant_builder
+ splits = builder.new_scan().plan().splits() if plan_before_limit else None
+ builder.with_limit(limit)
+ if splits is None:
+ splits = builder.new_scan().plan().splits()
+ dataset = builder.new_read().to_torch(splits, streaming=True,
**row_options)
+ _assert_rows(list(dataset), _expected_rows()[:limit])
+
+
[email protected]('limit', [None, 2])
+def test_streaming_variant_filter(variant_builder, row_options, limit):
+ builder, _ = variant_builder
+ predicate =
builder.table.new_read_builder().new_predicate_builder().greater_than('id', 2)
+ builder.with_filter(predicate)
+ if limit is not None:
+ builder.with_limit(limit)
+ splits = builder.new_scan().plan().splits()
+ dataset = builder.new_read().to_torch(splits, streaming=True,
**row_options)
+ _assert_rows(list(dataset), _expected_rows()[2:][:limit])
+
+
+def test_streaming_variant_empty_input(variant_builder, row_options):
+ builder, _ = variant_builder
+ assert list(builder.new_read().to_torch([], streaming=True,
**row_options)) == []
+
+
+def test_streaming_variant_strict_cast_errors(variant_builder, row_options):
+ builder, source = variant_builder
+ builder.with_projection({'score': "variant_get(%s, '$.score', 'float')" %
source})
+ splits = builder.new_scan().plan().splits()
+ dataset = builder.new_read().to_torch(splits, streaming=True,
**row_options)
+ with pytest.raises(Exception, match='Cannot cast Variant'):
+ list(dataset)
+
+
[email protected]('failure', ['disabled', 'missing_runtime', 'setup',
'unsupported_split'])
+def test_streaming_variant_requires_native(variant_builder, row_options,
failure):
+ builder, _ = variant_builder
+ splits = builder.new_scan().plan().splits()
+ target = 'pypaimon.execution.native.read.table_read.native_read'
+ if failure == 'disabled':
+ builder.table = builder.table.copy({'read.native.enabled': 'false'})
+ options = {'wraps': native_read}
+ elif failure == 'missing_runtime':
+ target = 'pypaimon.execution.native.runtime.native_runtime_available'
+ options = {'return_value': False}
+ elif failure == 'setup':
+ options = {'side_effect': RuntimeError('Native setup failed')}
+ else:
+ target =
'pypaimon.execution.native.read.table_read.NativeTableRead._split_files_supported'
+ options = {'return_value': False}
+ dataset = builder.new_read().to_torch(splits, streaming=True,
**row_options)
+ with patch(target, **options):
+ with pytest.raises(RuntimeError, match='Variant extraction requires a
compatible native reader'):
+ list(dataset)
+
+
[email protected]('limit', [None, 0, 1, 6])
+def test_streaming_variant_rank_worker_sharding(variant_builder, row_options,
limit):
+ builder, _ = variant_builder
+ if limit is not None:
+ builder.with_limit(limit)
+ read, splits = builder.new_read(), builder.new_scan().plan().splits()
+ actual = []
+ for rank in range(2):
+ dataset = read.to_torch(splits, streaming=True, sharding_rank=rank,
sharding_world_size=2,
+ **row_options)
+ for worker in range(2):
+ with patch('torch.utils.data.get_worker_info',
return_value=SimpleNamespace(id=worker, num_workers=2)):
+ if limit == 1:
+ with pytest.raises(ValueError, match='limit is not
supported with distributed Torch sharding'):
+ list(dataset)
+ else:
+ actual.extend(dataset)
+ if limit != 1:
+ _assert_rows(actual, _expected_rows()[:limit])
+
+
+def test_streaming_variant_dataloader_workers(variant_builder):
+ builder, _ = variant_builder
+ splits = builder.new_scan().plan().splits()
+ dataset = builder.new_read().to_torch(splits, streaming=True)
+ loader = DataLoader(dataset, batch_size=None, num_workers=2,
multiprocessing_context='spawn')
+ _assert_rows(list(loader), _expected_rows())
+
+
[email protected]('stop', ['unstarted', 'close', 'exhaust', 'failure'])
[email protected]('limit', [None, 3])
+def test_streaming_variant_closes_native_streams(variant_builder, row_options,
stop, limit):
+ builder, _ = variant_builder
+ if limit is not None:
+ builder.with_limit(limit)
+ splits = builder.new_scan().plan().splits()
+ opened, closed = [], []
+ failure = OSError('Native stream failed')
+
+ def tracked_native(*args, **kwargs):
+ batches = native_read(*args, **kwargs)
+ token = object()
+ opened.append(token)
+
+ def tracked_batches():
+ try:
+ for batch in batches:
+ yield batch
+ if stop == 'failure':
+ raise failure
+ finally:
+ batches.close()
+ closed.append(token)
+ return tracked_batches()
+
+ with patch('pypaimon.execution.native.read.table_read.native_read',
side_effect=tracked_native):
+ dataset = builder.new_read().to_torch(splits, streaming=True,
**row_options)
+ rows = iter(dataset)
+ try:
+ if stop == 'failure':
+ with pytest.raises(OSError) as error:
+ list(rows)
+ assert error.value is failure
+ elif stop == 'exhaust':
+ _assert_rows(list(rows), _expected_rows()[:limit])
+ elif stop == 'close':
+ assert next(rows) in _expected_rows()[:limit]
+ finally:
+ rows.close()
+ assert set(opened) == set(closed)
+ assert len(opened) == len(closed)
+ assert bool(opened) == (stop != 'unstarted')
+
+
+def test_non_streaming_variant_rows_remain_unsupported(variant_builder):
+ builder, _ = variant_builder
+ splits = builder.new_scan().plan().splits()
+ with pytest.raises(RuntimeError, match='Projection expressions are not
supported by Torch row format'):
+ builder.new_read().to_torch(splits)
diff --git a/paimon-python/pypaimon/tests/torch_read_test.py
b/paimon-python/pypaimon/tests/torch_read_test.py
index 2478d53dbf..192b2b7537 100644
--- a/paimon-python/pypaimon/tests/torch_read_test.py
+++ b/paimon-python/pypaimon/tests/torch_read_test.py
@@ -651,6 +651,22 @@ class TorchReadTest(unittest.TestCase):
with self.assertRaisesRegex(NotImplementedError, 'with query
authorization'):
list(dataset)
+ def test_torch_streaming_row_limit_with_prefetch_and_shuffle(self):
+ table, _ = self._create_named_projection_table('default.row_limit',
'append')
+ builder = table.new_read_builder().with_projection({'identifier':
'id'})
+ splits = builder.new_scan().plan().splits()
+ self.assertGreater(len(splits), 1)
+ for limit in [0, 1, 3, 5]:
+ builder.with_limit(limit)
+ read = builder.new_read()
+ expected = read.to_arrow(splits, parallelism=1).to_pylist()
+ for options in [{'prefetch_concurrency': 2},
+ {'shuffle': True, 'buffer_size': 1,
'max_buffer_input_splits': 1},
+ {'shuffle': True, 'buffer_size': 1,
'max_buffer_input_splits': 3}]:
+ with self.subTest(limit=limit, options=options):
+ dataset = read.to_torch(splits, streaming=True, **options)
+ self.assertCountEqual(list(dataset), expected)
+
@parameterized.expand([True, False])
def test_torch_read(self, is_streaming: bool = False):
schema = Schema.from_pyarrow_schema(self.pa_schema,
partition_keys=['user_id'])
diff --git a/paimon-python/pypaimon/tests/variant_test.py
b/paimon-python/pypaimon/tests/variant_test.py
index b83694fca3..c8e2067ff2 100644
--- a/paimon-python/pypaimon/tests/variant_test.py
+++ b/paimon-python/pypaimon/tests/variant_test.py
@@ -41,6 +41,7 @@ Sections
"""
import asyncio
+import importlib.util
import io
import json
import os
@@ -1262,7 +1263,8 @@ class TestVariantPaimonTable(unittest.TestCase):
pa.field('payload', _variant_arrow_type()),
])
- def test_variant_projection_keeps_torch_row_restriction(self):
+ @unittest.skipUnless(importlib.util.find_spec('torch') is not None, 'torch
is not installed')
+ def test_variant_projection_torch_row_backend_boundary(self):
for mode, options in (
('append', {}),
('lazy', {
@@ -1289,17 +1291,29 @@ class TestVariantPaimonTable(unittest.TestCase):
writer.close()
commit.close()
- builder = table.new_read_builder().with_projection(
- {'ratio': "try_variant_get(payload, '$.ratio', 'float')"})
+ projection = {'ratio': "try_variant_get(payload, '$.ratio',
'float')"}
+ builder = table.copy({'read.native.enabled':
'false'}).new_read_builder().with_projection(projection)
splits = builder.new_scan().plan().splits()
self.assertTrue(splits)
read = builder.new_read()
- with self.assertRaisesRegex(RuntimeError, 'Torch row format'):
- read.to_torch(splits, streaming=True)
with self.assertRaisesRegex(RuntimeError, 'Torch row format'):
read.to_torch(splits, streaming=False)
+ dataset = read.to_torch(splits, streaming=True)
+ with self.assertRaisesRegex(RuntimeError, 'Variant extraction
requires a compatible native reader'):
+ list(dataset)
+
+ builder = table.copy({'read.native.enabled':
'true'}).new_read_builder().with_projection(projection)
+ read = builder.new_read()
+ with
patch('pypaimon.execution.native.runtime.native_runtime_available',
return_value=False):
+ dataset = read.to_torch(splits, streaming=True)
+ with self.assertRaisesRegex(RuntimeError, 'Variant
extraction requires a compatible native reader'):
+ list(dataset)
+
+ if _native_variant_projection_available():
+ self.assertEqual(list(read.to_torch(splits,
streaming=True)), [{'ratio': 1.25}])
+
def _check_variant_projection_literal_top_level_names(
self, suffix, native_read):
pa_schema = pa.schema([
