[Preview][Data] Avoid concatenating multi-chunk tensor columns during row takes

[OneSizeFitsQuorum]

Fork-only preview: This draft PR is for reviewing the code and proposed
contribution in OneSizeFitsQuorum/ray. It has not been submitted to
ray-project/ray.

Description

This change avoids concatenating eligible multi-chunk Ray tensor extension
columns before row takes. Local shuffle is the primary use case, but the
one-shot path is shared by other PyArrow-backed Ray Data operations.

For a shuffle buffer containing N tensor rows, each with width D, producing
a K-row batch currently pays for a full O(N × D) column concatenation before
the required O(K × D) output copy. The proposed path gathers selected rows
directly from the source chunks, bringing the primary copy cost closer to the
unavoidable O(K × D).

There are two integration levels:

• take_table() prepares and executes a one-shot direct gather for general
  Arrow row-take operations.
• ShufflingBatcher prepares validated, generation-scoped chunk metadata once
  per compaction and reuses it for each output batch. The plan is discarded and
  rebuilt after the next compaction.

Only local shuffle currently reuses a prepared plan across multiple takes.

Implementation

• Add PreparedChunkedTensorTakePlan for validated Ray tensor V1/V2 chunk
  metadata.
• Retain zero-copy NumPy views of source chunks and their global row
  boundaries.
• Map global row indices to chunk-local rows and gather directly into the final
  contiguous tensor output.
• Divide wide-tensor gathers into bounded subbatches.
• Cap additional gather and index scratch memory at 8 MiB per subbatch.
• Integrate direct gathering with take_table().
• Reuse prepared plans from ShufflingBatcher.next_batch().
• Add RAY_DATA_ENABLE_CHUNKED_TENSOR_TAKE as a default-on operational
  fallback switch for both integration levels.
• Continue combining ordinary Arrow columns through the existing path.
• Preserve list-index parsing, exception behavior, and conservative fallback.

Ray Data paths

Because ArrowBlockAccessor.take() delegates to take_table(), the one-shot
fast path is already reachable from:

• general Arrow block row takes;
• hash partitioning and hash shuffle;
• PyArrow sort;
• random shuffle and shuffle-based repartition;
• concat_and_sort();
• projection of Arrow-join-unsupported tensor columns after a join.

Local shuffle has the largest expected gain because the same prepared plan is
reused to emit multiple batches and typically selects K < N rows at a time.
General takes and join projections can have the same favorable K < N
relationship.

Sort, random shuffle, and hash partition usually perform a full permutation
with K = N. They still avoid the extra full-source concatenation and its
temporary allocation, but must copy the full output tensor once. The performance
results below measure the local-shuffle use case; they do not claim measured
speedups for every reachable operator.

Pandas blocks, Polars sort paths, and direct pyarrow.Table.take() calls do not
use this fast path. Contiguous Arrow slicing remains on its existing zero-copy
path.

Safety and fallback behavior

The fast path is limited to non-null numeric fixed-shape ArrowTensorType and
ArrowTensorTypeV2 columns with multiple chunks and validated contiguous child
buffers.

The existing implementation remains in use for:

• null tensor rows or child values;
• variable-shaped and native PyArrow tensors;
• single-chunk tensor columns;
• all tensor rows narrower than 256 bytes;
• unsupported scalar dtypes and shapes;
• zero-sized tensor shapes;
• non-zero child offsets;
• unrepresentable output offsets;
• masked, nullable, non-native, non-integral, negative, out-of-range, or
  multidimensional indices.

Recognized unsupported layouts return to the existing concatenate/take path.
Setting RAY_DATA_ENABLE_CHUNKED_TENSOR_TAKE=0 disables the fast path
entirely. Unexpected allocation and internal errors are propagated.

Tests

The focused and existing regression suites cover:

• Ray tensor V1/V2 with 2, 17, and 257 source chunks;
• empty chunks and multiple 8 MiB scratch subbatches;
• shuffled, duplicated, and cross-boundary indices;
• prepared-plan and output buffer lifetime;
• local-shuffle compaction and plan rebuilding;
• fixed-seed row-order parity;
• hash partitioning with a multi-chunk tensor data column;
• nullable, child-offset, variable-shaped, native tensor, narrow, single-chunk,
  zero-sized-shape, and zero-row fallbacks;
• non-native-index and truncated-buffer-layout fallbacks;
• environment-controlled fallback for general takes and local shuffle;
• invalid-index exception parity;
• prepared-take fallback and unexpected-error propagation.

Commands:

python -m pytest \
  python/ray/data/tests/unit/test_chunked_tensor_take.py \
  python/ray/data/tests/unit/test_transform_pyarrow.py \
  python/ray/data/tests/test_batcher.py \
  -q

Result:

137 passed in 77.66s

The complete repository hooks were run against all four changed files:

pre-commit run

All applicable hooks passed, including Ruff, Black, pydoclint, Semgrep,
docstyle, AST, and Ray import-order checks.

The correctness probe produced byte-identical baseline, patched, and
packaged-runtime outputs.

Synthetic performance

The primary local-shuffle scenario uses 100,000 rows of float32[3000],
10,000-row batches, a 100,000-row shuffle buffer, and four independently paired
baseline/patch processes.

Case	Result
Primary geometric-mean speedup	9.80x
Bootstrap 95% CI	[9.68x, 9.89x]
Patch / baseline peak USS	0.139
Width-64 streaming	4.92x speedup
100 source chunks	8.97x speedup
Narrow fallback patch / baseline	0.964
Plain Arrow patch / baseline	0.972
Single-chunk control	Fast path disabled; no regression observed

The summary was regenerated from 22 raw result files and matched the recorded
output byte-for-byte. All predefined performance, variance, and fallback
regression gates passed.

End-to-end Ray Train + Ray Data validation

In the original Ray Train + Ray Data use case, workers consume batches
containing wide multi-chunk tensors from a locally shuffled dataset. The
workload configuration was:

Setting	Value
Training rows	60,074,872
Validation rows	6,444,594
Row layout	float32[30, 100] tensor plus one float32 label
Bytes per row	approximately 12,004
Training blocks	5,373 (approximately 11,181 rows per block)
Validation blocks	approximately 577
Training payload	approximately 721 GB (672 GiB)
Validation payload	approximately 77 GB (72 GiB)
Trainer workers	8, using 8 GPUs
Batch size	10,000 rows
Local shuffle buffer	100,000 rows
Epochs completed	5
Shuffle settings	shuffle=True, shuffle_ratio=1.0, block_random_shuffle=True
Software	Ray 2.56.1, Python 3.11.9, PyArrow 17.0.0, NumPy 1.26.4

The workload did not include an application-level filter or sort. The baseline
and patch runs used hardware-equivalent nodes. The end-to-end comparison
produced:

Metric	Patch / baseline	Equivalent change
Epoch duration geometric mean	0.221	approximately 4.52x faster
Final epoch duration	0.208	approximately 4.80x faster
Trial duration	0.233	approximately 4.30x faster
Parent duration	0.259	approximately 3.85x faster
Worker peak USS	0.918	approximately 8.2% lower
Object-store peak usage	0.870	approximately 13.0% lower

The validation completed without object spilling, object-store fallback, data
task failures, retries, worker startup failures, node failures, or container
restarts.

Duplicate-work check

On 2026-07-23, I searched open Ray Issues and PRs for chunked tensor,
multi-chunk tensor, tensor take, local shuffle tensor, and
ShufflingBatcher. I found no open work addressing this optimization.

PR ray-project#64576 touches transform_pyarrow.py, but only changes an assertion error
message and unrelated operator bounds checks. It does not overlap this design or
implementation.

AI assistance was used in preparing this change.

[Copilot]

Pull request overview

This PR introduces a faster row-take path for eligible multi-chunk Ray tensor extension columns by gathering selected rows directly from the original chunks instead of first concatenating the entire column. This targets Ray Data’s Arrow-backed take paths (notably local shuffle) to avoid an avoidable full-source O(N × D) concatenation cost when only K rows are needed.

Changes:

• Add a prepared/validated “chunked tensor take” implementation that can gather rows directly from multi-chunk Ray tensor columns with bounded scratch memory.
• Integrate the one-shot fast path into take_table() for general Arrow row-take operations when indices can be normalized to native int64.
• Teach ShufflingBatcher to build and reuse prepared tensor take plans per shuffle-buffer generation, while continuing to combine/take non-tensor columns via existing mechanisms.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
python/ray/data/tests/unit/test_chunked_tensor_take.py	Adds focused unit/regression coverage for prepared/unprepared chunked tensor takes, fallbacks, error propagation, and local-shuffle reuse behavior.
python/ray/data/_internal/tensor_extensions/chunked_tensor_take.py	Implements prepared-plan validation, zero-copy chunk view checks, and bounded-scratch gathering for multi-chunk Ray tensor extension columns.
python/ray/data/_internal/batcher.py	Adds local-shuffle preparation/take helpers to reuse prepared tensor plans across batches within a shuffle-buffer generation.
python/ray/data/_internal/arrow_ops/transform_pyarrow.py	Adds conservative index normalization for enabling the tensor fast path inside take_table() while preserving fallback behavior.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.