Skip to content

[python][torch] Add batch-first streaming datasets - #9365

Draft
XiaoHongbo-Hope wants to merge 1 commit into
apache:masterfrom
XiaoHongbo-Hope:codex/pypaimon-torch-batches
Draft

[python][torch] Add batch-first streaming datasets#9365
XiaoHongbo-Hope wants to merge 1 commit into
apache:masterfrom
XiaoHongbo-Hope:codex/pypaimon-torch-batches

Conversation

@XiaoHongbo-Hope

@XiaoHongbo-Hope XiaoHongbo-Hope commented Aug 23, 2026

Copy link
Copy Markdown
Contributor

Purpose

Avoid per-row Python dictionary conversion in PyPaimon Torch streaming by allowing the dataset to yield PyArrow RecordBatch objects or dictionaries of Torch tensors.

  • batch_format="row" remains the compatible default; "pyarrow" and "torch" enable batch streaming.
  • batch_size=None preserves native reader batches; otherwise batches are combined or sliced to the requested size.
  • The default converter supports non-null numeric, boolean, and numeric fixed-size-list columns. Other types can use to_tensor_fn.
  • Use DataLoader(batch_size=None) to avoid batching the data again.

Implementation notes

The queue markers and timeouts (_SENTINEL, _ITEM, _ERR, and the put/get/join timeouts) already existed in row streaming and are moved to the shared base class. _ITEM is the generalized name of the previous _ROW marker. The only new buffering constant is _PREFETCH_BATCH_QUEUE_MAXSIZE = 16, which bounds each DataLoader worker's Python prefetch queue to 16 RecordBatch objects. Native reader and active producer buffers are additional.

This follows the same batch-first pattern as Lance's PyTorch LanceDataset, which provides batch_size, to_tensor_fn, Arrow batch buffering, and a default batch_readahead=16. PyIceberg provides to_arrow_batch_reader() for streaming Arrow batches, but does not currently provide a PyTorch dataset or Tensor conversion layer.

Tests

  • PYTHONPATH=paimon-python pytest paimon-python/pypaimon/tests/torch_read_test.py -q (21 passed)
  • Flake8, Python 3.6 py_compile, and git diff --check

@XiaoHongbo-Hope
XiaoHongbo-Hope marked this pull request as draft August 23, 2026 15:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant