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 f2aad3f8a8 [python][torch] Support shuffle for Arrow and tensor batch
streaming (#10012)
f2aad3f8a8 is described below
commit f2aad3f8a82b0404afcf5972ecbaa8cf7030c730
Author: chaoyang <[email protected]>
AuthorDate: Mon Sep 21 10:41:58 2026 +0800
[python][torch] Support shuffle for Arrow and tensor batch streaming
(#10012)
---
docs/docs/pypaimon/pytorch.md | 52 +++++-
.../pypaimon/read/datasource/torch_dataset.py | 192 +++++++++++++++-----
paimon-python/pypaimon/read/table_read.py | 17 +-
paimon-python/pypaimon/tests/torch_read_test.py | 197 ++++++++++++++++++++-
4 files changed, 397 insertions(+), 61 deletions(-)
diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md
index 26770bdbba..320b824c85 100644
--- a/docs/docs/pypaimon/pytorch.md
+++ b/docs/docs/pypaimon/pytorch.md
@@ -133,12 +133,58 @@ for batch in dataloader:
converter supports non-null numeric, boolean, and numeric fixed-size-list
columns. Use `to_tensor_fn` for other types or custom conversion.
-Omit `batch_size` to preserve native reader batches. Otherwise, batches are
-combined or sliced to the requested size. Use `DataLoader(batch_size=None)` to
-disable a second batching step. Batch streaming does not support
`shuffle=True`.
+Without shuffle, omit `batch_size` to preserve native reader batches.
Otherwise,
+batches are combined or sliced to the requested size. Use
+`DataLoader(batch_size=None)` to disable a second batching step.
Numeric tensors may share read-only Arrow buffers; clone them before in-place
mutation. Batch formats currently require `prefetch_concurrency=1`.
+### Shuffled Batch Streaming
+
+Set `shuffle=True` to mix rows across input batches while keeping values in
+Arrow until the final Tensor conversion. The same options work with
+`batch_format="pyarrow"` and custom `to_tensor_fn` converters:
+
+```python
+dataset = table_read.to_torch(
+ splits,
+ streaming=True,
+ batch_format="torch",
+ batch_size=256,
+ shuffle=True,
+ seed=42,
+ buffer_size=4096,
+ max_buffer_input_splits=4,
+)
+loader = DataLoader(dataset, batch_size=None, num_workers=2)
+
+for epoch in range(10):
+ dataset.set_epoch(epoch)
+ for batch in loader:
+ train(batch["features"], batch["label"])
+```
+
+Each worker retains at most `buffer_size` rows in a rolling shuffle buffer and
+replaces randomly selected slots with rows from incoming Arrow blocks. Input
+blocks contain at most `buffer_size` rows; gathering replacements, output
+batching, and each open format reader use additional memory. This is a row
+bound, not a byte bound. The buffer drains early if combining Arrow blocks
+would overflow a 32-bit offset. Without `batch_size`, output blocks contain at
+most `buffer_size` rows.
+
+`max_buffer_input_splits` bounds the number of interleaved split readers per
+worker; `1` reads splits in order. A binding limit keeps the existing ordered
+read path so it selects the same rows before shuffling. Filters and distributed
+worker/rank sharding also precede shuffle. Each selected row is emitted once.
+
+The shuffle is local to each worker's buffer, not a uniform permutation of the
+whole dataset. The same seed, epoch, input batches and worker/rank
configuration
+reproduce the same order; `set_epoch()` also reaches persistent workers.
Changing
+reader batching or worker configuration can change the order. Arrow and Tensor
+formats share the same shuffle path, while row-format split interleaving can
+produce a different order. The existing `prefetch_concurrency=1` requirement
+also applies to shuffled batches.
+
## Video frame descriptors
For a multimodal frame table, use the higher-level scan API:
diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py
b/paimon-python/pypaimon/read/datasource/torch_dataset.py
index 541dd1d987..c83d6da630 100644
--- a/paimon-python/pypaimon/read/datasource/torch_dataset.py
+++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py
@@ -25,6 +25,7 @@ import queue
import random
import threading
import warnings
+from collections import deque
from typing import Any, Callable, Iterator, List, Optional
import pyarrow as pa
@@ -763,6 +764,52 @@ def _default_to_tensor(batch: pa.RecordBatch) -> dict:
return tensors
+class _TorchShuffleMixin:
+
+ def _init_shuffle(self, seed, buffer_size, max_buffer_input_splits):
+ self.seed = self._require_int(seed, "seed")
+ self.buffer_size = self._require_positive_int(buffer_size,
"buffer_size")
+ self.max_buffer_input_splits = self._require_positive_int(
+ max_buffer_input_splits, "max_buffer_input_splits")
+ self._epoch = _share_epoch_with_torch_workers(0)
+
+ def __setstate__(self, state):
+ self.__dict__ = state
+ self._epoch = _share_epoch_with_torch_workers(self._epoch)
+
+ @property
+ def epoch(self) -> int:
+ return int(self._epoch)
+
+ @epoch.setter
+ def epoch(self, epoch: int) -> None:
+ epoch = self._require_int(epoch, "epoch")
+ self._epoch += epoch - self._epoch
+
+ @staticmethod
+ def _require_int(value: int, name: str) -> int:
+ if not isinstance(value, int):
+ raise ValueError("%s must be an int" % name)
+ return value
+
+ @staticmethod
+ def _require_positive_int(value: int, name: str) -> int:
+ if not isinstance(value, int) or value <= 0:
+ raise ValueError("%s must be a positive int" % name)
+ return value
+
+ def set_epoch(self, epoch: int):
+ self.epoch = epoch
+ return self
+
+ def _shuffle_rng(self, worker_id):
+ rank, world_size = self._distributed_context()
+ rng_seed = self.seed + self.epoch * 1000003 + worker_id
+ if world_size > 1:
+ rng_seed = "%d:%d" % (rng_seed, rank)
+ return random.Random(rng_seed)
+
+
class TorchBatchIterDataset(_BaseTorchIterDataset):
"""Streaming IterableDataset which yields Arrow or Tensor batches."""
@@ -792,16 +839,20 @@ class TorchBatchIterDataset(_BaseTorchIterDataset):
worker_info = torch.utils.data.get_worker_info()
splits_to_process = self._worker_splits(worker_info)
raw_batches = self._arrow_batches_for_splits(splits_to_process)
+ worker_id = worker_info.id if worker_info is not None else 0
+ try:
+ for batch in self._prepare_batches(raw_batches, worker_id):
+ if self.batch_format == "torch":
+ converter = self.to_tensor_fn or _default_to_tensor
+ yield converter(batch)
+ else:
+ yield batch
+ finally:
+ raw_batches.close()
- batches = _sized_record_batches(
- self._limit_batches(raw_batches), self.batch_size
- )
- for batch in batches:
- if self.batch_format == "torch":
- converter = self.to_tensor_fn or _default_to_tensor
- yield converter(batch)
- else:
- yield batch
+ def _prepare_batches(self, raw_batches, worker_id):
+ return _sized_record_batches(
+ self._limit_batches(raw_batches), self.batch_size)
def _arrow_batches_for_splits(
self, splits: List[Split]
@@ -830,7 +881,87 @@ class TorchBatchIterDataset(_BaseTorchIterDataset):
yield batch
-class TorchShuffledIterDataset(_BaseTorchIterDataset):
+class TorchShuffledBatchIterDataset(_TorchShuffleMixin, TorchBatchIterDataset):
+ """Shuffle Arrow rows before sizing batches and converting to tensors."""
+
+ def __init__(self, table_read, splits, seed=0, buffer_size=1000,
+ max_buffer_input_splits=10, **kwargs):
+ super().__init__(table_read, splits, **kwargs)
+ self._init_shuffle(seed, buffer_size, max_buffer_input_splits)
+
+ def _prepare_batches(self, raw_batches, worker_id):
+ shuffled = _shuffle_record_batches(
+ self._limit_batches(raw_batches), self.buffer_size,
+ self._shuffle_rng(worker_id))
+ return _sized_record_batches(shuffled, self.batch_size)
+
+ def _arrow_batches_for_splits(self, splits):
+ # Keep the selected prefix of a binding limit unchanged before shuffle.
+ if self.max_buffer_input_splits == 1 or not
self._limit_covers_all_splits():
+ yield from super()._arrow_batches_for_splits(splits)
+ return
+
+ split_iter = iter(splits)
+ active = deque()
+ read_batches = super()._arrow_batches_for_splits
+
+ def add_next():
+ split = next(split_iter, None)
+ if split is not None:
+ active.append(read_batches([split]))
+
+ try:
+ for _ in range(min(self.max_buffer_input_splits, len(splits))):
+ add_next()
+ while active:
+ try:
+ yield next(active[0])
+ except StopIteration:
+ active.popleft().close()
+ add_next()
+ else:
+ active.rotate(-1)
+ finally:
+ for batches in active:
+ batches.close()
+
+
+def _shuffle_record_batches(batches, buffer_size, rng):
+ # Only row positions enter Python. Payloads, including nested values and
+ # nulls, stay in Arrow. Each input block replaces random reservoir slots.
+ buffer = None
+ for batch in _sized_record_batches(batches, buffer_size):
+ if buffer is None:
+ buffer = batch
+ continue
+ usage = _batch_offset_usage(buffer)
+ if any(usage.get(path, 0) + value > _MAX_ARROW_OFFSET
+ for path, value in _batch_offset_usage(batch).items()):
+ # Drain before concatenation would overflow a 32-bit Arrow offset.
+ yield _shuffled_record_batch(buffer, rng)
+ buffer = batch
+ continue
+ combined = _concat_record_batches([buffer, batch])
+ retained = list(range(buffer.num_rows))
+ selected = []
+ for row in range(batch.num_rows):
+ slot = rng.randrange(buffer.num_rows)
+ selected.append(retained[slot])
+ retained[slot] = buffer.num_rows + row
+ output = combined.take(pa.array(selected, type=pa.int64()))
+ buffer = combined.take(pa.array(retained, type=pa.int64()))
+ yield output
+ if buffer is not None:
+ yield _shuffled_record_batch(buffer, rng)
+
+
+def _shuffled_record_batch(batch, rng):
+ order = list(range(batch.num_rows))
+ rng.shuffle(order)
+ return batch.take(pa.array(order, type=pa.int64()))
+
+
+class TorchShuffledIterDataset(_TorchShuffleMixin, _BaseTorchIterDataset):
"""
PyTorch IterableDataset with Paimon-controlled streaming shuffle.
@@ -857,40 +988,7 @@ class TorchShuffledIterDataset(_BaseTorchIterDataset):
sharding_rank,
sharding_world_size,
)
- self.seed = self._require_int(seed, "seed")
- self.buffer_size = self._require_positive_int(buffer_size,
"buffer_size")
- self.max_buffer_input_splits = self._require_positive_int(
- max_buffer_input_splits, "max_buffer_input_splits")
- self._epoch = _share_epoch_with_torch_workers(0)
-
- def __setstate__(self, state):
- self.__dict__ = state
- self._epoch = _share_epoch_with_torch_workers(self._epoch)
-
- @property
- def epoch(self) -> int:
- return int(self._epoch)
-
- @epoch.setter
- def epoch(self, epoch: int) -> None:
- epoch = self._require_int(epoch, "epoch")
- self._epoch += epoch - self._epoch
-
- @staticmethod
- def _require_int(value: int, name: str) -> int:
- if not isinstance(value, int):
- raise ValueError("%s must be an int" % name)
- return value
-
- @staticmethod
- def _require_positive_int(value: int, name: str) -> int:
- if not isinstance(value, int) or value <= 0:
- raise ValueError("%s must be a positive int" % name)
- return value
-
- def set_epoch(self, epoch: int) -> "TorchShuffledIterDataset":
- self.epoch = epoch
- return self
+ self._init_shuffle(seed, buffer_size, max_buffer_input_splits)
def __iter__(self):
worker_info = torch.utils.data.get_worker_info()
@@ -957,11 +1055,7 @@ class TorchShuffledIterDataset(_BaseTorchIterDataset):
rows: Iterator[dict],
worker_id: int,
) -> Iterator[dict]:
- rank, world_size = self._distributed_context()
- rng_seed = self.seed + self.epoch * 1000003 + worker_id
- if world_size > 1:
- rng_seed = "%d:%d" % (rng_seed, rank)
- rng = random.Random(rng_seed)
+ rng = self._shuffle_rng(worker_id)
buffer = []
for row in rows:
if len(buffer) < self.buffer_size:
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index c003a828e9..48b8c1127c 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -1264,9 +1264,10 @@ class TableRead:
format.
batch_format: ``"row"``, ``"pyarrow"``, or ``"torch"``. Batch
formats require streaming.
- batch_size: Rows per batch; ``None`` preserves reader batches.
+ 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.
- shuffle: Whether to shuffle rows; supported only in row format.
+ shuffle: Whether to shuffle rows in a bounded buffer before output.
auto_detect_rank: Whether streaming reads shard by DDP rank.
sharding_rank: Explicit rank in the intended DDP process group.
sharding_world_size: Explicit size of that process group.
@@ -1301,10 +1302,6 @@ class TableRead:
raise ValueError(
"batch_format=%r requires streaming=True" % batch_format
)
- if shuffle:
- raise ValueError(
- "shuffle=True only supports batch_format='row'"
- )
if batch_format == "pyarrow" and to_tensor_fn is not None:
raise ValueError("to_tensor_fn requires batch_format='torch'")
if to_tensor_fn is not None and not callable(to_tensor_fn):
@@ -1320,8 +1317,13 @@ class TableRead:
from pypaimon.read.datasource.torch_dataset import (
TorchBatchIterDataset,
+ TorchShuffledBatchIterDataset,
)
- return TorchBatchIterDataset(
+ dataset_type = TorchShuffledBatchIterDataset if shuffle else
TorchBatchIterDataset
+ shuffle_options = (dict(seed=seed, buffer_size=buffer_size,
+
max_buffer_input_splits=max_buffer_input_splits)
+ if shuffle else {})
+ return dataset_type(
self,
splits,
batch_format=batch_format,
@@ -1330,6 +1332,7 @@ class TableRead:
auto_detect_rank=auto_detect_rank,
sharding_rank=sharding_rank,
sharding_world_size=sharding_world_size,
+ **shuffle_options,
)
if shuffle:
diff --git a/paimon-python/pypaimon/tests/torch_read_test.py
b/paimon-python/pypaimon/tests/torch_read_test.py
index bb14bd0320..620c90fff5 100644
--- a/paimon-python/pypaimon/tests/torch_read_test.py
+++ b/paimon-python/pypaimon/tests/torch_read_test.py
@@ -19,6 +19,7 @@ import json
import multiprocessing
import os
import pickle
+import random
import shutil
import subprocess
import sys
@@ -39,8 +40,11 @@ from pypaimon.multimodal.table import MultimodalTable
from pypaimon.read.datasource.torch_dataset import (
SplitRangeIndex,
TorchIterDataset,
+ TorchBatchIterDataset,
+ TorchShuffledBatchIterDataset,
TorchShuffledIterDataset,
_resolve_distributed_context,
+ _shuffle_record_batches,
)
from pypaimon.read.table_read import TableRead
from pypaimon.table.file_store_table import FileStoreTable
@@ -428,6 +432,92 @@ class TorchDistributedShardingTest(unittest.TestCase):
)
+class TorchBatchShuffleTest(unittest.TestCase):
+ @parameterized.expand([1, 3, 8, 50])
+ def test_columnar_shuffle_matches_existing_row_shuffle(self, buffer_size):
+ rows = [{"id": i, "nested": None if i % 4 == 0 else {"values": [i,
None]},
+ "text": None if i % 3 == 0 else str(i)} for i in range(25)]
+ table = pa.Table.from_pylist(rows)
+ source = table.to_batches(max_chunksize=4)
+ source.insert(1, source[0].slice(0, 0))
+ actual = pa.Table.from_batches(list(_shuffle_record_batches(
+ iter(source), buffer_size, random.Random(17))))
+
+ legacy = TorchShuffledIterDataset(
+ SimpleNamespace(read_type=[], limit=None), [], seed=17,
buffer_size=buffer_size)
+ expected = list(legacy._iter_buffer_shuffled_rows(iter(rows), 0))
+ self.assertEqual(actual.schema, table.schema)
+ self.assertEqual(actual.to_pylist(), expected)
+
+ def test_shuffle_does_not_collect_whole_input(self):
+ consumed = []
+
+ def source():
+ for start in range(0, 300, 3):
+ consumed.extend(range(start, start + 3))
+ yield pa.record_batch([pa.array(range(start, start + 3))],
names=["id"])
+
+ batches = _shuffle_record_batches(source(), 5, random.Random(0))
+ self.assertEqual(next(batches).num_rows, 5)
+ self.assertLessEqual(len(consumed), 12)
+ batches.close()
+
+ def test_shuffle_drains_before_arrow_offset_overflow(self):
+ table = pa.table({"text": [str(i).zfill(4) for i in range(12)]})
+ with patch("pypaimon.read.datasource.torch_dataset._MAX_ARROW_OFFSET",
12):
+ batches = list(_shuffle_record_batches(
+ iter(table.to_batches(max_chunksize=3)), 3, random.Random(3)))
+ self.assertEqual(sorted(row["text"] for b in batches for row in
b.to_pylist()),
+ table["text"].to_pylist())
+ self.assertTrue(all(sum(len(v) for v in b["text"].to_pylist()) <= 12
for b in batches))
+
+ @parameterized.expand(["complete", "close", "read_error", "convert_error"])
+ def test_interleaved_readers_are_bounded_and_closed(self, mode):
+ opened, closed, active = [], [], set()
+ maximum = [0]
+
+ def read(dataset, splits):
+ split = splits[0]
+ opened.append(split)
+ active.add(split)
+ maximum[0] = max(maximum[0], len(active))
+ try:
+ for i in range(10):
+ if mode == "read_error" and split == 1 and i == 1:
+ raise RuntimeError("read failed")
+ yield pa.record_batch([pa.array([split * 10 + i])],
names=["id"])
+ finally:
+ active.remove(split)
+ closed.append(split)
+
+ def convert(batch):
+ if mode == "convert_error":
+ raise RuntimeError("conversion failed")
+ return batch
+
+ dataset = TorchShuffledBatchIterDataset(
+ SimpleNamespace(read_type=[], limit=None), list(range(8)),
+ buffer_size=3, max_buffer_input_splits=2, batch_size=2,
+ batch_format="torch", to_tensor_fn=convert)
+ with patch.object(TorchBatchIterDataset, "_arrow_batches_for_splits",
read):
+ batches = iter(dataset)
+ if mode == "close":
+ next(batches)
+ batches.close()
+ elif mode.endswith("error"):
+ with self.assertRaisesRegex(RuntimeError, "failed"):
+ list(batches)
+ else:
+ rows = [row["id"] for batch in batches for row in
batch.to_pylist()]
+ self.assertEqual(sorted(rows), list(range(80)))
+ self.assertFalse(active)
+ self.assertCountEqual(opened, closed)
+ self.assertEqual(maximum[0], 2)
+
+ def test_empty_input(self):
+ self.assertEqual(list(_shuffle_record_batches(iter([]), 5,
random.Random(0))), [])
+
+
class TorchReadTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
@@ -1084,12 +1174,13 @@ class TorchReadTest(unittest.TestCase):
)
with self.assertRaisesRegex(ValueError, 'batch_size requires'):
table_read.to_torch(splits, streaming=True, batch_size=2)
- with self.assertRaisesRegex(ValueError, 'only supports batch_format'):
+ with self.assertRaisesRegex(ValueError, 'buffer_size'):
table_read.to_torch(
splits,
streaming=True,
batch_format='torch',
shuffle=True,
+ buffer_size=0,
)
for invalid in [0, -1, 1.9, True, 2]:
with self.subTest(prefetch_concurrency=invalid):
@@ -1128,13 +1219,14 @@ class TorchReadTest(unittest.TestCase):
splits,
streaming=True,
batch_format=batch_format,
- shuffle=batch_format == 'row' and shuffle,
+ shuffle=shuffle,
auto_detect_rank=True,
)
for batch_format, shuffle in [
('row', False),
('row', True),
('pyarrow', False),
+ ('pyarrow', True),
]
]
expected = splits[(len(splits) + 1) // 2:]
@@ -1718,6 +1810,107 @@ class TorchReadTest(unittest.TestCase):
print("✓ All predicate test cases passed!")
print(f"{'=' * 60}\n")
+ @parameterized.expand([(1, None), (7, None), (7, 3), (11, 17)])
+ def test_torch_batch_shuffle_formats_and_epochs(self, buffer_size,
batch_size):
+ table = self._create_shuffle_append_table(
+ 'default.batch_shuffle_%s_%s' % (buffer_size, batch_size))
+ builder = table.new_read_builder().with_projection(['user_id',
'item_id'])
+ splits = builder.new_scan().plan().splits()
+ read = builder.new_read()
+ outputs = []
+ for batch_format in ['pyarrow', 'torch']:
+ dataset = read.to_torch(
+ splits, streaming=True, shuffle=True,
batch_format=batch_format,
+ batch_size=batch_size, seed=23, buffer_size=buffer_size,
+ max_buffer_input_splits=3)
+
+ def collect():
+ batches = list(dataset)
+ if batch_format == 'pyarrow':
+ return pa.Table.from_batches(batches).to_pydict(),
[b.num_rows for b in batches]
+ return {name: torch.cat([b[name] for b in batches]).tolist()
+ for name in ['user_id', 'item_id']},
[len(b['user_id']) for b in batches]
+
+ with patch.object(TableRead, 'to_arrow',
side_effect=AssertionError('whole table')), \
+ patch.object(TableRead, 'to_iterator',
side_effect=AssertionError('row conversion')):
+ data, sizes = collect()
+ self.assertEqual(data, collect()[0])
+ self.assertEqual(sorted(data['user_id']), list(range(80)))
+ self.assertEqual(data['item_id'], [1000 + i for i in
data['user_id']])
+ self.assertTrue(all(size <= (batch_size or buffer_size) for
size in sizes))
+ if batch_size:
+ self.assertTrue(all(size == batch_size for size in
sizes[:-1]))
+ outputs.append(data)
+ dataset.set_epoch(1)
+ if buffer_size > 1:
+ self.assertNotEqual(data, collect()[0])
+ dataset.set_epoch(0)
+ self.assertEqual(data, collect()[0])
+ self.assertEqual(outputs[0], outputs[1])
+
+ def test_torch_batch_shuffle_filter_limit_and_empty(self):
+ table =
self._create_shuffle_append_table('default.batch_shuffle_limit')
+ builder = table.new_read_builder().with_projection(['user_id'])
+ predicate =
builder.new_predicate_builder().greater_or_equal('user_id', 10)
+ builder.with_filter(predicate).with_limit(9)
+ splits = builder.new_scan().plan().splits()
+ read = builder.new_read()
+ expected = read.to_arrow(splits)['user_id'].to_pylist()
+ dataset = read.to_torch(splits, streaming=True, shuffle=True,
batch_format='pyarrow',
+ batch_size=4, buffer_size=3,
max_buffer_input_splits=3)
+ actual = [row['user_id'] for batch in dataset for row in
batch.to_pylist()]
+ self.assertCountEqual(actual, expected)
+ for limit in [0, 9]:
+ read.limit = limit
+ empty = read.to_torch([], streaming=True, shuffle=True,
batch_format='pyarrow', buffer_size=3)
+ self.assertEqual(list(empty), [])
+
+ def test_torch_batch_shuffle_persistent_spawn_workers(self):
+ table =
self._create_shuffle_append_table('default.batch_shuffle_spawn',
partition_keys=['dt'])
+ builder = table.new_read_builder().with_projection(['user_id'])
+ splits = builder.new_scan().plan().splits()
+ self.assertGreaterEqual(len(splits), 4)
+ all_ids = []
+ for rank in range(2):
+ dataset = builder.new_read().to_torch(
+ splits, streaming=True, shuffle=True, batch_format='torch',
+ batch_size=4, buffer_size=7, seed=23,
+ sharding_rank=rank, sharding_world_size=2)
+ loader = DataLoader(dataset, batch_size=None, num_workers=2,
+ persistent_workers=True,
multiprocessing_context='spawn')
+ try:
+ epoch0 = self._collect_torch_user_ids_from_dataloader(loader)
+ self.assertEqual(epoch0,
self._collect_torch_user_ids_from_dataloader(loader))
+ dataset.set_epoch(1)
+ epoch1 = self._collect_torch_user_ids_from_dataloader(loader)
+ self.assertCountEqual(epoch0, epoch1)
+ self.assertNotEqual(epoch0, epoch1)
+ all_ids.extend(epoch0)
+ finally:
+ if loader._iterator is not None:
+ loader._iterator._shutdown_workers()
+ self.assertEqual(sorted(all_ids), list(range(80)))
+
+ def test_torch_batch_shuffle_vectors_deletions_and_snapshot(self):
+ import pypaimon.multimodal as pm
+
+ schema = pa.schema([('id', pa.int64()), ('vector',
pa.list_(pa.float32(), 2))])
+ table = pm.connect(options={'warehouse': self.warehouse}).create_table(
+ 'batch_shuffle_vectors', schema=schema, options={'file.format':
'parquet'})
+ table.add(pa.table({'id': list(range(8)), 'vector': [[i, i + 1] for i
in range(8)]}, schema=schema))
+ snapshot = table.raw_table.snapshot_manager().get_latest_snapshot().id
+ table.delete('id = 2')
+ for snapshot_id, expected in [(snapshot, list(range(1, 8))), (None,
[1, 3, 4, 5, 6, 7])]:
+ dataset = table.scan(snapshot_id=snapshot_id).where('id >=
1').to_torch(
+ streaming=True, batch_format='torch', batch_size=3,
shuffle=True,
+ seed=7, buffer_size=4, max_buffer_input_splits=1)
+ batches = list(dataset)
+ ids = torch.cat([batch['id'] for batch in batches]).tolist()
+ vectors = torch.cat([batch['vector'] for batch in batches])
+ self.assertEqual(vectors.dtype, torch.float32)
+ self.assertEqual(vectors.tolist(), [[i, i + 1] for i in ids])
+ self.assertEqual(sorted(ids), expected)
+
def test_torch_streaming_shuffle_single_worker(self):
table =
self._create_shuffle_append_table('default.test_torch_shuffle_single')
read_builder = table.new_read_builder().with_projection(['user_id'])
