diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 6ab0af5173ca..1a33680f767f 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -59,6 +59,80 @@ when it is false, it will read the full amount of data into memory. **`prefetch_concurrency`** (default: 1): When streaming is true, number of threads used for parallel prefetch within each DataLoader worker. Set to a value greater than 1 to partition splits across threads and increase read throughput. Has no effect when streaming is false. +### Batch-first Streaming + +For training pipelines where Python row conversion is the bottleneck, set +`batch_format` to make the `IterableDataset` yield whole batches. Disable +DataLoader auto-batching with `batch_size=None`. + +To consume PyArrow `RecordBatch` objects directly: + +```python +dataset = table_read.to_torch( + splits, + streaming=True, + batch_format="pyarrow", + batch_size=1024, + prefetch_concurrency=4, +) +dataloader = DataLoader(dataset, batch_size=None, num_workers=2) + +for record_batch in dataloader: + train_on_arrow(record_batch) +``` + +To convert numeric columns to Tensor batches inside each DataLoader worker: + +```python +dataset = table_read.to_torch( + splits, + streaming=True, + batch_format="torch", + batch_size=1024, +) +dataloader = DataLoader(dataset, batch_size=None, num_workers=2) + +for batch in dataloader: + # batch is a dict[str, torch.Tensor] + train(batch["features"], batch["label"]) +``` + +The default Tensor converter supports non-null numeric, boolean, and numeric +fixed-size-list columns. Use `to_tensor_fn` for strings, variable-length +features, BLOB decoding, null handling, or application-specific device and +dtype conversion: + +```python +import torch + +def convert(batch): + return { + "label": torch.from_numpy( + batch.column("label").to_numpy(zero_copy_only=False) + ), + "text": batch.column("text").to_pylist(), + } + +dataset = table_read.to_torch( + splits, + streaming=True, + batch_format="torch", + batch_size=1024, + to_tensor_fn=convert, +) +``` + +The default converter can share numeric Arrow buffers with CPU tensors. Treat +those tensors as read-only, or clone them before an in-place mutation. + +If `batch_size` is omitted, PyPaimon preserves native reader batch sizes and +avoids batch resizing. When it is set, batches are combined or sliced to that +size; only the final batch produced by each DataLoader worker may be smaller. +Batch-first streaming partitions splits across DataLoader workers in the same +way as row streaming. Row-level `shuffle=True` is not supported in batch-first +mode; shuffle split order during planning or shuffle batches in the training +pipeline instead. + ## File Format Metadata Cache Reusable PyArrow Dataset metadata is cached across reads. Configure its estimated diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py b/paimon-python/pypaimon/read/datasource/torch_dataset.py index 5eb3485dddd1..721e8a3300f1 100644 --- a/paimon-python/pypaimon/read/datasource/torch_dataset.py +++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py @@ -21,8 +21,10 @@ import queue import random import threading -from typing import Iterator, List +import warnings +from typing import Any, Callable, Iterator, List, Optional +import pyarrow as pa import torch from torch.utils.data import Dataset, IterableDataset @@ -88,6 +90,13 @@ class _BaseTorchIterDataset(IterableDataset): Shared helpers for streaming PyTorch datasets backed by Paimon splits. """ + _SENTINEL = 0 + _ITEM = 1 + _ERR = 2 + _PREFETCH_PUT_TIMEOUT_SEC = 30.0 + _PREFETCH_GET_TIMEOUT_SEC = 300.0 + _PREFETCH_JOIN_TIMEOUT_SEC = 5.0 + def __init__(self, table_read: TableRead, splits: List[Split]): self.table_read = table_read self.splits = splits @@ -119,6 +128,66 @@ def _worker_splits(self, worker_info) -> List[Split]: return self.splits[start_idx:end_idx] + def _iter_concurrently( + self, + splits: List[Split], + concurrency: int, + item_iterator: Callable[[List[Split]], Iterator], + queue_maxsize: int, + ) -> Iterator: + n = min(concurrency, len(splits)) + if n == 0: + return + split_groups = [splits[i::n] for i in range(n)] + + q = queue.Queue(maxsize=queue_maxsize) + stop = threading.Event() + + def put_item(tag: int, payload): + while not stop.is_set(): + try: + q.put((tag, payload), timeout=self._PREFETCH_PUT_TIMEOUT_SEC) + return True + except queue.Full: + continue + return False + + def producer(split_group: List[Split]): + try: + for item in item_iterator(split_group): + if stop.is_set() or not put_item(self._ITEM, item): + break + put_item(self._SENTINEL, None) + except Exception as e: + put_item(self._ERR, e) + + threads = [ + threading.Thread(target=producer, args=(group,), daemon=True) + for group in split_groups + ] + for thread in threads: + thread.start() + + try: + done = 0 + while done < n: + try: + tag, payload = q.get(timeout=self._PREFETCH_GET_TIMEOUT_SEC) + except queue.Empty: + if stop.is_set(): + break + continue + if tag == self._SENTINEL: + done += 1 + elif tag == self._ERR: + raise payload + else: + yield payload + finally: + stop.set() + for thread in threads: + thread.join(timeout=self._PREFETCH_JOIN_TIMEOUT_SEC) + class TorchIterDataset(_BaseTorchIterDataset): """ @@ -129,13 +198,7 @@ class TorchIterDataset(_BaseTorchIterDataset): rather than loading everything into memory upfront. """ - _SENTINEL = 0 - _ROW = 1 - _ERR = 2 _PREFETCH_QUEUE_MAXSIZE = 512 - _PREFETCH_PUT_TIMEOUT_SEC = 30.0 - _PREFETCH_GET_TIMEOUT_SEC = 300.0 - _PREFETCH_JOIN_TIMEOUT_SEC = 5.0 def __init__(self, table_read: TableRead, splits: List[Split], prefetch_concurrency: int = 1): """ @@ -165,8 +228,12 @@ def __iter__(self): splits_to_process = self._worker_splits(worker_info) if self.prefetch_concurrency > 1: - for row in self._iter_rows(splits_to_process): - yield row + yield from self._iter_concurrently( + splits_to_process, + self.prefetch_concurrency, + self._rows_for_splits, + self._PREFETCH_QUEUE_MAXSIZE, + ) return worker_iterator = self.table_read.to_iterator(splits_to_process) @@ -174,60 +241,172 @@ def __iter__(self): for offset_row in worker_iterator: yield self._row_to_dict(offset_row) - def _iter_rows(self, splits: List[Split]): - n = min(self.prefetch_concurrency, len(splits)) - if n == 0: - return - split_groups = [splits[i::n] for i in range(n)] - - q = queue.Queue(maxsize=self._PREFETCH_QUEUE_MAXSIZE) - stop = threading.Event() + def _rows_for_splits(self, splits: List[Split]) -> Iterator[dict]: + for offset_row in self.table_read.to_iterator(splits): + yield self._row_to_dict(offset_row) - def put_item(tag: int, payload): - while not stop.is_set(): - try: - q.put((tag, payload), timeout=self._PREFETCH_PUT_TIMEOUT_SEC) - return True - except queue.Full: - continue - return False - def producer(split_group: List): - try: - for offset_row in self.table_read.to_iterator(split_group): - if stop.is_set(): - break - row_dict = self._row_to_dict(offset_row) - if not put_item(self._ROW, row_dict): - break - put_item(self._SENTINEL, None) - except Exception as e: - put_item(self._ERR, e) +def _concat_record_batches(batches: List[pa.RecordBatch]) -> pa.RecordBatch: + if len(batches) == 1: + return batches[0] + return pa.RecordBatch.from_arrays( + [ + pa.concat_arrays([batch.column(i) for batch in batches]) + for i in range(batches[0].num_columns) + ], + schema=batches[0].schema, + ) + + +def _sized_record_batches( + batches: Iterator[pa.RecordBatch], + batch_size: Optional[int], +) -> Iterator[pa.RecordBatch]: + if batch_size is None: + yield from batches + return + + pending: List[pa.RecordBatch] = [] + pending_rows = 0 + for batch in batches: + offset = 0 + while offset < batch.num_rows: + take = min(batch_size - pending_rows, batch.num_rows - offset) + pending.append(batch.slice(offset, take)) + pending_rows += take + offset += take + if pending_rows == batch_size: + yield _concat_record_batches(pending) + pending = [] + pending_rows = 0 + + if pending: + yield _concat_record_batches(pending) + + +def _default_to_tensor(batch: pa.RecordBatch) -> dict: + tensors = {} + for name, array in zip(batch.schema.names, batch.columns): + if array.null_count: + raise ValueError( + "Torch tensor conversion does not support null values in " + "column %r; provide to_tensor_fn to handle them." % name + ) + + if pa.types.is_fixed_size_list(array.type): + value_type = array.type.value_type + if not ( + pa.types.is_integer(value_type) + or pa.types.is_floating(value_type) + or pa.types.is_boolean(value_type) + ): + raise ValueError( + "Torch tensor conversion does not support column %r with " + "type %s; provide to_tensor_fn." % (name, array.type) + ) + values = array.values.slice( + array.offset * array.type.list_size, + len(array) * array.type.list_size, + ) + if values.null_count: + raise ValueError( + "Torch tensor conversion does not support null list values " + "in column %r; provide to_tensor_fn to handle them." % name + ) + numpy_array = values.to_numpy(zero_copy_only=False).reshape( + len(array), array.type.list_size + ) + elif ( + pa.types.is_integer(array.type) + or pa.types.is_floating(array.type) + or pa.types.is_boolean(array.type) + ): + numpy_array = array.to_numpy(zero_copy_only=False) + else: + raise ValueError( + "Torch tensor conversion only supports numeric, boolean, and " + "fixed-size-list columns; column %r has type %s. Select " + "batch_format='pyarrow' or provide to_tensor_fn." + % (name, array.type) + ) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="The given NumPy array is not writable", + category=UserWarning, + ) + tensors[name] = torch.from_numpy(numpy_array) + return tensors + + +class TorchBatchIterDataset(_BaseTorchIterDataset): + """Streaming IterableDataset which yields Arrow or Tensor batches.""" + + _PREFETCH_BATCH_QUEUE_MAXSIZE = 16 - threads = [threading.Thread(target=producer, args=(split_groups[i],), daemon=True) - for i in range(n)] - for t in threads: - t.start() + def __init__( + self, + table_read: TableRead, + splits: List[Split], + batch_format: str, + batch_size: Optional[int], + prefetch_concurrency: int = 1, + to_tensor_fn: Optional[Callable[[pa.RecordBatch], Any]] = None, + ): + super().__init__(table_read, splits) + self.batch_format = batch_format + self.batch_size = batch_size + self.prefetch_concurrency = max(1, int(prefetch_concurrency)) + self.to_tensor_fn = to_tensor_fn + def __iter__(self): + worker_info = torch.utils.data.get_worker_info() + splits_to_process = self._worker_splits(worker_info) + if self.prefetch_concurrency > 1: + raw_batches = self._iter_concurrently( + splits_to_process, + self.prefetch_concurrency, + self._arrow_batches_for_splits, + self._PREFETCH_BATCH_QUEUE_MAXSIZE, + ) + else: + raw_batches = self._arrow_batches_for_splits(splits_to_process) + + 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 _arrow_batches_for_splits( + self, splits: List[Split] + ) -> Iterator[pa.RecordBatch]: + reader = self.table_read.to_arrow_batch_reader(splits) try: - done = 0 - while done < n: - try: - tag, payload = q.get(timeout=self._PREFETCH_GET_TIMEOUT_SEC) - except queue.Empty: - if stop.is_set(): - break - continue - if tag == self._SENTINEL: - done += 1 - elif tag == self._ERR: - raise payload - else: - yield payload + for batch in iter(reader.read_next_batch, None): + if batch.num_rows: + yield batch finally: - stop.set() - for t in threads: - t.join(timeout=self._PREFETCH_JOIN_TIMEOUT_SEC) + close = getattr(reader, "close", None) + if close is not None: + close() + + def _limit_batches( + self, batches: Iterator[pa.RecordBatch] + ) -> Iterator[pa.RecordBatch]: + remaining = self.table_read.limit + for batch in batches: + if remaining is not None: + if remaining <= 0: + return + if batch.num_rows > remaining: + batch = batch.slice(0, remaining) + remaining -= batch.num_rows + yield batch class TorchShuffledIterDataset(_BaseTorchIterDataset): diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index c2ba44545a40..94d41fae608b 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -18,7 +18,7 @@ import os import threading from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, Iterator, List, Optional +from typing import Any, Callable, Dict, Iterator, List, Optional import pandas import pyarrow @@ -654,12 +654,84 @@ def to_torch( streaming: bool = False, prefetch_concurrency: int = 1, *, + batch_format: str = "row", + batch_size: Optional[int] = None, + to_tensor_fn: Optional[Callable] = None, shuffle: bool = False, seed: int = 0, buffer_size: int = 1000, max_buffer_input_splits: int = 10, ) -> "torch.utils.data.Dataset": - """Wrap Paimon table data to PyTorch Dataset.""" + """Wrap Paimon table data in a PyTorch Dataset. + + ``batch_format="row"`` preserves the original behavior and yields one + Python dictionary per row. Streaming reads can instead yield native + PyArrow ``RecordBatch`` objects or dictionaries of Torch tensors, + avoiding per-row Python conversion. Use ``DataLoader(batch_size=None)`` + for either batch-first format so PyTorch does not batch the data again. + + Args: + splits: Scan-plan splits to read. + streaming: Whether to stream through an ``IterableDataset``. + prefetch_concurrency: Split-read threads inside each DataLoader + worker. + batch_format: ``"row"`` (default), ``"pyarrow"``, or ``"torch"``. + Batch formats require ``streaming=True``. + batch_size: Optional number of rows per yielded batch. ``None`` + preserves native Paimon reader batch sizes. A positive value + combines or slices RecordBatches, with only the last batch in + each DataLoader worker allowed to be smaller. + to_tensor_fn: Optional callable receiving a PyArrow RecordBatch in + ``batch_format="torch"`` mode. By default, numeric, boolean, + and fixed-size-list columns are converted to tensors. + shuffle: Whether to apply Paimon's row-level streaming shuffle. + This currently supports only ``batch_format="row"``. + """ + valid_batch_formats = {"row", "pyarrow", "torch"} + if batch_format not in valid_batch_formats: + raise ValueError( + "batch_format must be one of %s, got %r" + % (sorted(valid_batch_formats), batch_format) + ) + if batch_size is not None and ( + isinstance(batch_size, bool) + or not isinstance(batch_size, int) + or batch_size <= 0 + ): + raise ValueError("batch_size must be a positive int or None") + if batch_format == "row": + if batch_size is not None: + raise ValueError( + "batch_size requires batch_format='pyarrow' or 'torch'" + ) + if to_tensor_fn is not None: + raise ValueError("to_tensor_fn requires batch_format='torch'") + else: + if not streaming: + 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): + raise ValueError("to_tensor_fn must be callable") + + from pypaimon.read.datasource.torch_dataset import ( + TorchBatchIterDataset, + ) + return TorchBatchIterDataset( + self, + splits, + batch_format=batch_format, + batch_size=batch_size, + prefetch_concurrency=prefetch_concurrency, + to_tensor_fn=to_tensor_fn, + ) + if shuffle: if not streaming: raise ValueError("shuffle=True only supports streaming=True") diff --git a/paimon-python/pypaimon/tests/torch_read_test.py b/paimon-python/pypaimon/tests/torch_read_test.py index 5f55cb2bc892..a292b3b52530 100644 --- a/paimon-python/pypaimon/tests/torch_read_test.py +++ b/paimon-python/pypaimon/tests/torch_read_test.py @@ -22,6 +22,7 @@ import pyarrow as pa from parameterized import parameterized +import torch from torch.utils.data import DataLoader from pypaimon import CatalogFactory, Schema @@ -143,6 +144,199 @@ def test_torch_streaming_prefetch_concurrency(self): self.assertEqual(sorted_user_ids, expected_user_ids) self.assertEqual(sorted_behaviors, expected_behaviors) + def test_torch_streaming_pyarrow_batches(self): + schema = Schema.from_pyarrow_schema( + self.pa_schema, partition_keys=['user_id'] + ) + self.catalog.create_table( + 'default.test_torch_pyarrow_batches', schema, False + ) + table = self.catalog.get_table( + 'default.test_torch_pyarrow_batches' + ) + self._write_test_table(table) + + read_builder = table.new_read_builder().with_projection( + ['user_id', 'behavior'] + ) + splits = read_builder.new_scan().plan().splits() + dataset = read_builder.new_read().to_torch( + splits, + streaming=True, + batch_format='pyarrow', + batch_size=3, + ) + dataloader = DataLoader( + dataset, + batch_size=None, + num_workers=2, + shuffle=False, + ) + + batches = list(dataloader) + self.assertTrue(batches) + self.assertTrue( + all(isinstance(batch, pa.RecordBatch) for batch in batches) + ) + self.assertTrue(all(0 < batch.num_rows <= 3 for batch in batches)) + result = pa.Table.from_batches(batches).sort_by('user_id').to_pydict() + self.assertEqual(result['user_id'], list(range(1, 9))) + self.assertEqual(result['behavior'], list('abcdefgh')) + + def test_torch_streaming_tensor_batches(self): + schema = Schema.from_pyarrow_schema( + self.pa_schema, partition_keys=['user_id'] + ) + self.catalog.create_table( + 'default.test_torch_tensor_batches', schema, False + ) + table = self.catalog.get_table( + 'default.test_torch_tensor_batches' + ) + self._write_test_table(table) + + read_builder = table.new_read_builder().with_projection( + ['user_id', 'item_id'] + ) + splits = read_builder.new_scan().plan().splits() + dataset = read_builder.new_read().to_torch( + splits, + streaming=True, + prefetch_concurrency=4, + batch_format='torch', + batch_size=3, + ) + + batches = list(dataset) + self.assertEqual([len(batch['user_id']) for batch in batches], [3, 3, 2]) + self.assertTrue( + all(batch['user_id'].dtype == torch.int32 for batch in batches) + ) + self.assertTrue( + all(batch['item_id'].dtype == torch.int64 for batch in batches) + ) + user_ids = torch.cat( + [batch['user_id'] for batch in batches] + ).sort().values.tolist() + self.assertEqual(user_ids, list(range(1, 9))) + + def test_torch_streaming_batches_respect_limit_with_prefetch(self): + schema = Schema.from_pyarrow_schema( + self.pa_schema, partition_keys=['user_id'] + ) + self.catalog.create_table( + 'default.test_torch_batch_limit', schema, False + ) + table = self.catalog.get_table('default.test_torch_batch_limit') + self._write_test_table(table) + + read_builder = table.new_read_builder().with_projection( + ['user_id'] + ).with_limit(5) + splits = read_builder.new_scan().plan().splits() + dataset = read_builder.new_read().to_torch( + splits, + streaming=True, + prefetch_concurrency=4, + batch_format='pyarrow', + batch_size=3, + ) + batches = list(dataset) + self.assertEqual([batch.num_rows for batch in batches], [3, 2]) + + def test_default_tensor_converter_supports_fixed_size_list(self): + from pypaimon.read.datasource.torch_dataset import _default_to_tensor + + values = pa.array([1, 2, 3, 4, 5, 6], type=pa.int32()) + features = pa.FixedSizeListArray.from_arrays(values, 3) + batch = pa.RecordBatch.from_arrays([features], ['features']) + + result = _default_to_tensor(batch) + + self.assertEqual(result['features'].dtype, torch.int32) + self.assertEqual(result['features'].tolist(), [[1, 2, 3], [4, 5, 6]]) + + def test_torch_streaming_custom_tensor_conversion(self): + schema = Schema.from_pyarrow_schema(self.pa_schema) + self.catalog.create_table( + 'default.test_torch_custom_tensor_batch', schema, False + ) + table = self.catalog.get_table( + 'default.test_torch_custom_tensor_batch' + ) + self._write_test_table(table) + + read_builder = table.new_read_builder().with_projection( + ['user_id', 'behavior'] + ) + splits = read_builder.new_scan().plan().splits() + + def to_tensor(batch): + return { + 'user_id': torch.from_numpy( + batch.column('user_id').to_numpy(zero_copy_only=False) + ), + 'behavior': batch.column('behavior').to_pylist(), + } + + dataset = read_builder.new_read().to_torch( + splits, + streaming=True, + batch_format='torch', + batch_size=5, + to_tensor_fn=to_tensor, + ) + batches = list(dataset) + self.assertEqual([len(batch['user_id']) for batch in batches], [5, 3]) + self.assertEqual( + sorted(value for batch in batches for value in batch['behavior']), + list('abcdefgh'), + ) + + default_dataset = read_builder.new_read().to_torch( + splits, + streaming=True, + batch_format='torch', + ) + with self.assertRaisesRegex(ValueError, "batch_format='pyarrow'"): + next(iter(default_dataset)) + + def test_torch_batch_options_validation(self): + schema = Schema.from_pyarrow_schema(self.pa_schema) + self.catalog.create_table( + 'default.test_torch_batch_validation', schema, False + ) + table = self.catalog.get_table( + 'default.test_torch_batch_validation' + ) + self._write_test_table(table) + read_builder = table.new_read_builder().with_projection(['user_id']) + splits = read_builder.new_scan().plan().splits() + table_read = read_builder.new_read() + + with self.assertRaisesRegex(ValueError, 'batch_format must be one of'): + table_read.to_torch( + splits, streaming=True, batch_format='numpy' + ) + with self.assertRaisesRegex(ValueError, 'requires streaming=True'): + table_read.to_torch(splits, batch_format='pyarrow') + with self.assertRaisesRegex(ValueError, 'batch_size must be'): + table_read.to_torch( + splits, + streaming=True, + batch_format='torch', + batch_size=0, + ) + 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'): + table_read.to_torch( + splits, + streaming=True, + batch_format='torch', + shuffle=True, + ) + def test_blob_torch_read(self): """Test end-to-end blob functionality using blob descriptors.""" import random