diff --git a/ci/lint/pyrefly-excluded-files.txt b/ci/lint/pyrefly-excluded-files.txt index 807c740ed38e..62d75513a352 100644 --- a/ci/lint/pyrefly-excluded-files.txt +++ b/ci/lint/pyrefly-excluded-files.txt @@ -200,6 +200,7 @@ python/ray/data/tests/block_batching/test_iter_batches.py python/ray/data/tests/block_batching/test_util.py python/ray/data/tests/conftest.py python/ray/data/tests/datasource/test_arrow.py +python/ray/data/tests/datasource/test_arrow_rs_parquet_reader.py python/ray/data/tests/datasource/test_bigquery.py python/ray/data/tests/datasource/test_clickhouse.py python/ray/data/tests/datasource/test_csv.py diff --git a/docker/base-extra/Dockerfile b/docker/base-extra/Dockerfile index af81e865d884..9f033f519d28 100644 --- a/docker/base-extra/Dockerfile +++ b/docker/base-extra/Dockerfile @@ -48,7 +48,7 @@ wget -O - https://packages.cloud.google.com/apt/doc/apt-key.gpg \ # Add gdb since ray dashboard uses `memray attach`, which requires gdb. APT_PKGS=( - google-cloud-sdk + google-cloud-cli supervisor vim zsh diff --git a/pyrefly.toml b/pyrefly.toml index 873516f90193..5d835b13add0 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -45,6 +45,7 @@ ignore-missing-imports = [ "pytest.*", "pytest_lazy_fixtures.*", "ray.core.generated.*", + "ray_data_arrow_rs.*", "raydp.*", "rich.*", "snappy.*", diff --git a/python/ray/air/BUILD.bazel b/python/ray/air/BUILD.bazel index 0c119c16cf2b..57dbaaecab92 100644 --- a/python/ray/air/BUILD.bazel +++ b/python/ray/air/BUILD.bazel @@ -1,6 +1,29 @@ load("@rules_python//python:defs.bzl", "py_library", "py_test") load("//bazel:python.bzl", "doctest") +# Shared env for the python/ray/air/tests targets. ``RAY_DATA_PARQUET_FOOTER_ +# NUM_ACTORS=1`` keeps the Parquet footer-reader pool tiny: the production +# default of 32 actors times out and trips Ray's "too many worker processes" +# warning when many test targets run in parallel under CI. Mirrored by a +# setdefault in python/ray/air/tests/conftest.py for non-bazel pytest runs. +_AIR_TEST_ENV = { + "RAY_DATA_PARQUET_FOOTER_NUM_ACTORS": "1", + "RAY_TRAIN_V2_ENABLED": "1", +} + +# Same, for the targets that pin Train V1 telemetry/behaviour. +_AIR_TEST_ENV_TRAIN_V1 = { + "RAY_DATA_PARQUET_FOOTER_NUM_ACTORS": "1", + "RAY_TRAIN_V2_ENABLED": "0", +} + +# Same, plus legacy Keras for the Keras callback test. +_AIR_TEST_ENV_TF = { + "RAY_DATA_PARQUET_FOOTER_NUM_ACTORS": "1", + "RAY_TRAIN_V2_ENABLED": "1", + "TF_USE_LEGACY_KERAS": "1", +} + doctest( name = "py_doctest[air]", env = {"RAY_TRAIN_V2_ENABLED": "1"}, @@ -31,7 +54,7 @@ py_test( size = "small", srcs = ["tests/test_air_usage.py"], # NOTE: This tests Train V1 telemetry. - env = {"RAY_TRAIN_V2_ENABLED": "0"}, + env = _AIR_TEST_ENV_TRAIN_V1, tags = [ "exclusive", "team:ml", @@ -44,7 +67,7 @@ py_test( size = "large", srcs = ["tests/test_new_dataset_config.py"], # NOTE: Relevant tests moved to train/v2/tests/test_data_integration.py - env = {"RAY_TRAIN_V2_ENABLED": "0"}, + env = _AIR_TEST_ENV_TRAIN_V1, tags = [ "exclusive", "team:ml", @@ -60,7 +83,7 @@ py_test( "tests/test_experiment_restore.py", ], # NOTE: This tests Tune and Train V1 restoration. - env = {"RAY_TRAIN_V2_ENABLED": "0"}, + env = _AIR_TEST_ENV_TRAIN_V1, tags = [ "exclusive", "team:ml", @@ -73,7 +96,7 @@ py_test( size = "medium", srcs = ["tests/test_errors.py"], # NOTE: This tests Tune (Train V1) error propagation logic. - env = {"RAY_TRAIN_V2_ENABLED": "0"}, + env = _AIR_TEST_ENV_TRAIN_V1, tags = [ "exclusive", "team:ml", @@ -86,7 +109,7 @@ py_test( size = "small", srcs = ["tests/test_integration_comet.py"], # NOTE: This tests the Tune Comet callback. - env = {"RAY_TRAIN_V2_ENABLED": "0"}, + env = _AIR_TEST_ENV_TRAIN_V1, tags = [ "exclusive", "team:ml", @@ -99,7 +122,7 @@ py_test( size = "small", srcs = ["tests/test_integration_wandb.py"], # NOTE: This tests the Tune wandb callback. - env = {"RAY_TRAIN_V2_ENABLED": "0"}, + env = _AIR_TEST_ENV_TRAIN_V1, tags = [ "exclusive", "team:ml", @@ -112,7 +135,7 @@ py_test( size = "medium", srcs = ["tests/test_integration_mlflow.py"], # NOTE: This tests the Tune mlflow callback. - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -124,10 +147,7 @@ py_test( name = "test_keras_callback", size = "medium", srcs = ["tests/test_keras_callback.py"], - env = { - "RAY_TRAIN_V2_ENABLED": "1", - "TF_USE_LEGACY_KERAS": "1", - }, + env = _AIR_TEST_ENV_TF, tags = [ "exclusive", "team:ml", @@ -139,7 +159,7 @@ py_test( name = "test_remote_storage_hdfs", size = "small", srcs = ["tests/test_remote_storage_hdfs.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "hdfs", @@ -152,7 +172,7 @@ py_test( name = "test_tracebacks", size = "small", srcs = ["tests/test_tracebacks.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -164,7 +184,7 @@ py_test( name = "test_utils", size = "small", srcs = ["tests/test_utils.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -183,7 +203,7 @@ py_test( name = "test_barrier", size = "small", srcs = ["tests/execution/test_barrier.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -195,7 +215,7 @@ py_test( name = "test_e2e_train_flow", size = "medium", srcs = ["tests/execution/test_e2e_train_flow.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -207,7 +227,7 @@ py_test( name = "test_e2e_tune_flow", size = "medium", srcs = ["tests/execution/test_e2e_tune_flow.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -219,7 +239,7 @@ py_test( name = "test_event_manager", size = "medium", srcs = ["tests/execution/test_event_manager.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -231,7 +251,7 @@ py_test( name = "test_resource_manager_fixed", size = "small", srcs = ["tests/execution/test_resource_manager_fixed.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -243,7 +263,7 @@ py_test( name = "test_resource_manager_placement_group", size = "medium", srcs = ["tests/execution/test_resource_manager_placement_group.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -255,7 +275,7 @@ py_test( name = "test_resource_request", size = "small", srcs = ["tests/execution/test_resource_request.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -267,7 +287,7 @@ py_test( name = "test_tracked_actor", size = "small", srcs = ["tests/execution/test_tracked_actor.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -279,7 +299,7 @@ py_test( name = "test_tracked_actor_task", size = "small", srcs = ["tests/execution/test_tracked_actor_task.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _AIR_TEST_ENV, tags = [ "exclusive", "team:ml", diff --git a/python/ray/air/tests/conftest.py b/python/ray/air/tests/conftest.py index 6595beb4b40b..e5aaa5062e90 100644 --- a/python/ray/air/tests/conftest.py +++ b/python/ray/air/tests/conftest.py @@ -1,11 +1,18 @@ # Trigger pytest hook to automatically zip test cluster logs to archive dir on failure import copy +import os import pytest import ray from ray.tests.conftest import pytest_runtest_makereport # noqa +# Keep the Parquet footer-reader pool tiny for these tests. The production +# default of 32 actors times out under CI parallelism; tests that need a larger +# pool can override with monkeypatch.setenv. Mirrored in +# python/ray/air/BUILD.bazel for bazel test targets. +os.environ.setdefault("RAY_DATA_PARQUET_FOOTER_NUM_ACTORS", "1") + @pytest.fixture def restore_data_context(request): diff --git a/python/ray/data/.claude/CLAUDE.md b/python/ray/data/.claude/CLAUDE.md index 000f82642435..c264402e845a 100644 --- a/python/ray/data/.claude/CLAUDE.md +++ b/python/ray/data/.claude/CLAUDE.md @@ -8,3 +8,39 @@ ## Gotchas + +## Project docs (local-only, gitignored) + +Long-running investigations keep their notes in a directory under `python/ray/data/` +excluded via `.git/info/exclude` — currently `arrow_rs_docs/` (the arrow-rs Parquet +reader migration, draft PR #65117). **Start at that directory's `README.md`**: it holds +the read order, the doc index with a staleness column, the ratio conventions, and the +rules below. Do not read the other docs front to back; arrive at them from a link. + +## Doc rules for those directories + +Each rule exists because its absence already cost something real. + +1. **No markdown file over 1000 lines.** The moment one crosses it, move a whole section + out into an existing doc if one owns the topic, else a new file leaving a stub that + says where it went and keeping section numbers unchanged so old cross-references still + resolve. Do it in the same edit that crosses the line. Two docs reached 1910 lines and + became unreadable and un-editable at once; splitting after the fact cost far more. +2. **A measured number is a row in `findings.md`** — one tabular registry with permanent + IDs, a status column (LIVE / RE-MEASURE / DEAD / RETRACTED / OPEN) and a caveat column, + so a new confound is recorded by editing one row. Never renumber, never delete a row: a + wrong finding becomes RETRACTED with the reason, because the wrong version is what + older docs still say. The prose that qualifies a number goes in the topic doc, never in + the plan — a plan that quotes numbers goes stale silently. +3. **Closing a work item splits it three ways, in the same edit:** numbers → `findings.md`, + prose → the topic doc, and the *decision* (one row: disposition, why, and for parked + items what would revive it) → `todo_archive.md` under its original number, which is + never reused. It leaves `TODO.md` entirely, which holds open work only. **Before + archiving, check the closed item for open work hiding inside it** and promote that — + this has now bitten four times, including a P0 that had already shipped and a stale + docstring nobody could see. +4. **Verify against the tree before writing a doc claim.** Env vars, defaults, file paths + and which branch a script lives on all drift; each audit of these docs has found + documented knobs that do not exist and shipped defaults documented three values stale. +5. After every message, check if there is anything to be updated/added/removed in any doc + and make the apt change. diff --git a/python/ray/data/BUILD.bazel b/python/ray/data/BUILD.bazel index 4255d9eddde5..0467d707ad35 100644 --- a/python/ray/data/BUILD.bazel +++ b/python/ray/data/BUILD.bazel @@ -1,5 +1,5 @@ -load("@rules_python//python:defs.bzl", "py_library", "py_test") -load("//bazel:python.bzl", "doctest", "py_test_module_list") +load("@rules_python//python:defs.bzl", "py_library") +load(":test.bzl", "doctest", "py_test", "py_test_module_list") # Export pytest plugin so it can be used in the documentation tests. exports_files( @@ -1030,7 +1030,7 @@ py_test( py_test( name = "test_predicate_pushdown", - size = "small", + size = "medium", srcs = ["tests/test_predicate_pushdown.py"], tags = [ "exclusive", diff --git a/python/ray/data/_internal/datasource_v2/chunkers/file_chunker.py b/python/ray/data/_internal/datasource_v2/chunkers/file_chunker.py index 0f8e504965c0..a592f4693d6c 100644 --- a/python/ray/data/_internal/datasource_v2/chunkers/file_chunker.py +++ b/python/ray/data/_internal/datasource_v2/chunkers/file_chunker.py @@ -19,7 +19,7 @@ get_type_hints, ) -from ray.data._internal.util import GiB, MiB, infer_compression +from ray.data._internal.util import MiB, infer_compression from ray.util.annotations import DeveloperAPI @@ -54,16 +54,25 @@ class LineDelimitedFileChunkMetadata(ChunkMetadata): chunk_byte_end_idx: int -class ParquetFileChunkMetadata(ChunkMetadata): - """Metadata for Parquet file chunks. +class ParquetRowGroupChunkMetadata(ChunkMetadata): + """Metadata for a Parquet chunk described by explicit row-group indices. - For a parquet file, the chunks are based on the total size of the file, not on the - underlying row groups. We will split a file into potentially many chunks of the - target chunk size. This may correspond to 0, 1, or more row groups per chunk. + Produced by the footer-based chunking path (``ListFiles`` reads each file's + footer, prunes/bin-packs its row groups, and emits one manifest row per file + per bin), so it carries the exact physical row groups the reader should scan + for the file in that bin -- no size-based reconciliation needed. + + ``row_group_ids`` are physical row-group indices into the file; any + coalescing/splitting the bin packer applied is already expanded away here. + ``num_rows`` is the summed footer row count of those groups (for sizing / + limit accounting). ``uncompressed_size`` is their summed, projection-scoped + uncompressed byte size, carried so the reader can size batches without + re-reading the footer ``ListFiles`` already read. """ - chunk_idx: int - total_num_chunks: int + row_group_ids: Tuple[int, ...] + num_rows: int + uncompressed_size: int @DeveloperAPI @@ -146,64 +155,3 @@ def generate_chunk_metadatas( ), chunk_size, ) - - -@DeveloperAPI -class ParquetFileChunker(FileChunker): - """File chunker for Parquet files. - - This chunker splits Parquet files into an estimated number of chunks. We do not - fetch the metadata for the file, so we may overestimate the number of chunks - compared to the actual number of underlying row groups. The partitioner creates - groupings based on these estimates, and the reader fetches the metadata and - ensures that all row groups are read / any overestimated row groups are ignored. - """ - - # Chosen so that we can effectively chunk files but will not result in OOMs if - # the compression ratio is high. - # - # If the compression ratio is high and this chunk size is large, we end up with - # larger chunks than we need and reading can OOM. Reducing the chunk size gives - # better memory performance by reading a smaller fraction of row groups at a time. - # - # We also want to keep this large enough such that we do not end up reading too - # much data if we underestimate the number of chunks. If row groups are larger - # than the chunk size and we place many of them in the same read task, the total - # amount of data read might be larger than expected. By increasing the chunk - # size we are less likely to put many such row groups in the same task. - _DEFAULT_TARGET_CHUNK_SIZE = 1 * GiB - - def __init__(self, target_chunk_size: Optional[int] = None): - from ray.data.context import DataContext - - ctx = DataContext.get_current() - if target_chunk_size is not None: - self._target_chunk_size = target_chunk_size - elif ctx.parquet_chunker_target_chunk_size is not None: - self._target_chunk_size = ctx.parquet_chunker_target_chunk_size - else: - self._target_chunk_size = self._DEFAULT_TARGET_CHUNK_SIZE - - def generate_chunk_metadatas( - self, path: str, file_size: int - ) -> Iterable[Tuple[Optional[ChunkMetadata], int]]: - if file_size <= self._target_chunk_size: - # Do not chunk if the file is smaller than the target chunk size; when - # we read the file, this prevents additional metadata fetching since we - # want to read the entire file. - yield None, file_size - return - - num_chunks = math.ceil(file_size / self._target_chunk_size) - for i in range(num_chunks): - chunk_start = self._target_chunk_size * i - chunk_end = min(self._target_chunk_size * (i + 1), file_size) - chunk_size = chunk_end - chunk_start - yield ( - create_chunk_metadata( - ParquetFileChunkMetadata, - chunk_idx=i, - total_num_chunks=num_chunks, - ), - chunk_size, - ) diff --git a/python/ray/data/_internal/datasource_v2/chunkers/parquet_file_chunking_utils.py b/python/ray/data/_internal/datasource_v2/chunkers/parquet_file_chunking_utils.py index 2ae89d926c3a..4d84c864c087 100644 --- a/python/ray/data/_internal/datasource_v2/chunkers/parquet_file_chunking_utils.py +++ b/python/ray/data/_internal/datasource_v2/chunkers/parquet_file_chunking_utils.py @@ -1,113 +1,71 @@ -"""Parquet file-level chunking helpers for DataSourceV2. +"""Parquet chunk helpers for DataSourceV2. -Maps planner chunk metadata (``ParquetFileChunkMetadata``) to row-group -ranges and PyArrow ``ParquetFileFragment`` subsets for parallel reads. +Maps ``ParquetRowGroupChunkMetadata`` (the explicit surviving row groups a bin +assigns to a file) to PyArrow ``ParquetFileFragment`` subsets for reading. """ -from typing import List, Optional, Tuple +from typing import Callable, Iterable, List, Tuple, TypeVar import pyarrow.dataset as pds -from ray.data._internal.datasource_v2.chunkers.file_chunker import ( - ParquetFileChunkMetadata, -) +from ray._common.retry import call_with_retry +R = TypeVar("R") -def _calculate_row_group_range( - chunk_idx: int, total_num_chunks: int, total_row_groups: int -) -> Optional[Tuple[int, int]]: - """Compute the half-open row-group range for a given chunk. - Distributes row groups as evenly as possible across chunks. If row groups - don't divide evenly, earlier chunks get the extra row groups. +def _with_io_retry(f: Callable[[], R], description: str) -> R: + """Run ``f``, retrying the transient IO errors configured on the context. - Example: - - 10 row groups, 3 chunks -> [0:4), [4:7), [7:10) - - 11 row groups, 3 chunks -> [0:4), [4:8), [8:11) - - Args: - chunk_idx: Index of the current chunk (0-based). - total_num_chunks: Total number of chunks. - total_row_groups: Total number of row groups to distribute. - - Returns: - Tuple ``(start_row_group, end_row_group)`` where ``end`` is exclusive, - or ``None`` if ``chunk_idx`` falls beyond the actual number of row - groups (i.e. the planner over-estimated the chunk count). + ``ParquetFileFragment.subset`` and ``.metadata`` both open the file to read + its footer, so on remote storage they fail with the same transient errors + (S3 timeouts, throttling) the rest of the read path already retries. """ - assert ( - total_row_groups >= 0 - ), f"total_row_groups must be non-negative, got {total_row_groups}" - assert ( - total_num_chunks > 0 - ), f"total_num_chunks must be positive, got {total_num_chunks}" - assert ( - chunk_idx < total_num_chunks - ), f"chunk_idx must be less than total_num_chunks, got {chunk_idx} and {total_num_chunks}" - assert chunk_idx >= 0, f"chunk_idx must be non-negative, got {chunk_idx}" - - # Handle the case where ``chunk_idx`` exceeds the actual number of chunks - # needed. This happens when the planner overestimated the number of chunks - # (the chunker doesn't fetch metadata). - if chunk_idx >= total_row_groups: - return None + from ray.data.context import DataContext - base_row_groups_per_chunk = total_row_groups // total_num_chunks - remainder = total_row_groups % total_num_chunks - - # Chunks 0 through (remainder-1) get one extra row group. - if chunk_idx < remainder: - row_groups_in_this_chunk = base_row_groups_per_chunk + 1 - start = chunk_idx * row_groups_in_this_chunk - else: - row_groups_in_this_chunk = base_row_groups_per_chunk - start = ( - remainder * (base_row_groups_per_chunk + 1) - + (chunk_idx - remainder) * base_row_groups_per_chunk - ) - - end = start + row_groups_in_this_chunk - - assert ( - 0 <= start <= end <= total_row_groups - ), f"Invalid range [{start}, {end}) for {total_row_groups} row groups" - - return start, end + return call_with_retry( + f, + description=description, + match=DataContext.get_current().retried_io_errors, + ) -def _fragments_from_chunk_metadata( +def _fragments_from_row_group_ids( fragment: pds.ParquetFileFragment, - chunk_metadata: ParquetFileChunkMetadata, + row_group_ids: Iterable[int], + *, + per_row_group_offsets: bool, ) -> List[Tuple[pds.ParquetFileFragment, int]]: - """Slice ``fragment`` into per-row-group sub-fragments per chunk metadata. - - Returns one ``(ParquetFileFragment, file_row_offset)`` pair per row group - covered by the chunk, where ``file_row_offset`` is the sum of ``num_rows`` - across all row groups that precede the sub-fragment in the underlying - file. Callers seed per-fragment hashing offsets with this value so - sub-fragments of the same file don't collide on ``(path, 0, n)``. - - Returns an empty list when the chunk index falls beyond the file's actual - row-group count (the planner over-estimated; we silently drop the slice). + """Slice ``fragment`` to the explicit physical ``row_group_ids`` of one bin. + + Used by the footer-based chunking path, where ``ParquetRowGroupChunkMetadata`` + names the exact surviving row groups for a file (predicate pruning + bin + packing already happened upstream), so no size-based reconciliation is needed. + + When ``per_row_group_offsets`` is False (the common case) the file's groups + are scanned together as a single sub-fragment with a row offset of 0 -- this + lets PyArrow coalesce reads across the groups. When True (``include_row_hash`` + is on), one sub-fragment per row group is returned, each paired with its + cumulative pre-filter row offset within the file, so row hashes stay unique + and match the physical row positions even when pruned groups make the + surviving set non-contiguous. """ - chunk_idx = chunk_metadata["chunk_idx"] - total_num_chunks = chunk_metadata["total_num_chunks"] - metadata = fragment.metadata - total_row_groups = metadata.num_row_groups - - row_group_range = _calculate_row_group_range( - chunk_idx, total_num_chunks, total_row_groups - ) - - if row_group_range is None: + ids = sorted(row_group_ids) + if not ids: return [] - start, end = row_group_range - - file_row_offset = sum(metadata.row_group(i).num_rows for i in range(start)) - sub_fragments: List[Tuple[pds.ParquetFileFragment, int]] = [] - for row_group_index in range(start, end): - sub_fragments.append( - (fragment.subset(row_group_ids=[row_group_index]), file_row_offset) + def _subset(rg_ids: List[int]) -> pds.ParquetFileFragment: + return _with_io_retry( + lambda: fragment.subset(row_group_ids=rg_ids), + f"subset row groups {rg_ids} of {fragment.path}", ) - file_row_offset += metadata.row_group(row_group_index).num_rows - return sub_fragments + + if not per_row_group_offsets: + return [(_subset(ids), 0)] + + metadata = _with_io_retry( + lambda: fragment.metadata, f"read Parquet footer for {fragment.path}" + ) + # Cumulative pre-filter row offset at the start of each physical row group. + prefix = [0] * (metadata.num_row_groups + 1) + for i in range(metadata.num_row_groups): + prefix[i + 1] = prefix[i] + metadata.row_group(i).num_rows + return [(_subset([rg_id]), prefix[rg_id]) for rg_id in ids] diff --git a/python/ray/data/_internal/datasource_v2/chunkers/parquet_footer_types.py b/python/ray/data/_internal/datasource_v2/chunkers/parquet_footer_types.py new file mode 100644 index 000000000000..13c00dc4b999 --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/chunkers/parquet_footer_types.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Tuple + + +@dataclass(frozen=True) +class RowGroupInfo: + """A chunk of one file: a contiguous run of ``rg_count`` physical row groups. + + ``uncompressed_size`` / ``num_rows`` are summed over the run. For a single + physical row group (``rg_count == 1``) the run is its own atom and needs no + breakdown, so ``rg_sizes`` / ``rg_rows`` stay empty; they're populated only + for coalesced runs (``rg_count > 1``) so the bin packer can split them back + at exact byte/row boundaries. + """ + + # Start row-group index (== the row group's index when rg_count == 1). + rg_idx: int + uncompressed_size: int # summed over the run + num_rows: int # summed over the run + # True when every row in the run is guaranteed to satisfy the filter (or there + # is no filter), so ``num_rows`` is an exact survivor count and the limit can + # be pushed down on it. False for partially-matching groups, whose ``num_rows`` + # overestimates survivors. Coalescing never merges across this flag. + fully_matched: bool = True + # Number of consecutive physical row groups this chunk covers. + rg_count: int = 1 + # Per-physical-row-group uncompressed sizes / row counts, in ``rg_idx`` order. + # Populated only for coalesced runs (``rg_count > 1``). + rg_sizes: Tuple[int, ...] = () + rg_rows: Tuple[int, ...] = () + + +@dataclass(frozen=True) +class FileChunks: + """The footer-derived chunks for a single file.""" + + path: str + size: int # on-disk file size, from the file listing + row_groups: Tuple[RowGroupInfo, ...] + + +@dataclass(frozen=True) +class BinItem: + """One file's row-group chunk placed into a bin. + + Same contiguous-run shape as :class:`RowGroupInfo`; ``path`` is the packer's + "colour". A (possibly-split) item covers physical row groups + ``range(rg_idx, rg_idx + rg_count)``. + """ + + path: str # colour + rg_idx: int # start row-group index (see RowGroupInfo.rg_idx) + uncompressed_size: int + num_rows: int + fully_matched: bool = True # see RowGroupInfo.fully_matched + rg_count: int = 1 # number of consecutive physical row groups this item covers + # Per-physical-row-group breakdown (rg_idx order), populated only for + # coalesced runs (rg_count > 1); lets the packer split at row-group boundaries. + rg_sizes: Tuple[int, ...] = () + rg_rows: Tuple[int, ...] = () + + +@dataclass(frozen=True) +class Bin: + """A sealed bin: a set of row-group chunks (across one or more files) whose + combined uncompressed size targets one bin budget. Becomes one + ``FileManifest`` block == one downstream read task.""" + + items: Tuple[BinItem, ...] + total_uncompressed_size: int diff --git a/python/ray/data/_internal/datasource_v2/listing/file_indexer.py b/python/ray/data/_internal/datasource_v2/listing/file_indexer.py index 66cea801ce6c..c63fdd23c5e4 100644 --- a/python/ray/data/_internal/datasource_v2/listing/file_indexer.py +++ b/python/ray/data/_internal/datasource_v2/listing/file_indexer.py @@ -1,7 +1,7 @@ import logging from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Callable, Iterable, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Callable, Iterable, List, Optional, Tuple, Union from pyarrow.fs import FileSystem @@ -21,6 +21,9 @@ from ray.data.block import BlockColumn from ray.data.datasource.path_util import _resolve_paths_and_filesystem +if TYPE_CHECKING: + from ray.data.expressions import Expr + logger = logging.getLogger(__name__) @@ -31,6 +34,19 @@ def file_chunker(self) -> FileChunker: """The file chunker that this indexer uses.""" ... + @property + def yields_read_units(self) -> bool: + """Whether ``list_files`` already emits complete read units. + + ``False`` (default) means the indexer yields per-file / per-chunk + manifests that still need size-balanced partitioning downstream. An + indexer that bin-packs internally (e.g. the footer-based Parquet indexer, + which reads footers and packs row groups into ~one-block manifests) + returns ``True``; ``ListFiles`` then skips the partitioner and runs + listing as a single task so packing sees the whole file stream. + """ + return False + @abstractmethod def list_files( self, @@ -39,6 +55,9 @@ def list_files( filesystem: "FileSystem", pruners: Optional[List[FilePruner]] = None, preserve_order: bool = False, + predicate: Optional["Expr"] = None, + limit: Optional[int] = None, + projected_columns: Optional[List[str]] = None, ) -> Iterable[FileManifest]: """List files and their on-disk sizes for the given path. @@ -47,6 +66,13 @@ def list_files( filesystem: A PyArrow filesystem object. pruners: A list of file pruners to apply. preserve_order: Whether to preserve order in file listing. + predicate: Pushed-down row filter. Indexers that read file + metadata (e.g. the footer-based Parquet indexer) use it to skip + row groups; others ignore it. + limit: Pushed-down row limit, for indexers that can stop listing + early. Others ignore it. + projected_columns: Pushed-down column projection, for metadata-aware + sizing. Others ignore it. Returns: An iterator of `FileManifest` objects, each of which contains a file path @@ -54,6 +80,24 @@ def list_files( """ ... + @abstractmethod + def list_file_infos( + self, + paths: "BlockColumn", + *, + filesystem: "FileSystem", + pruners: Optional[List[FilePruner]] = None, + preserve_order: bool = False, + ) -> Iterable["FileInfo"]: + """List files as raw ``FileInfo``\\ s (path + on-disk size). + + Unlike :meth:`list_files`, this yields the pre-chunk file stream. The + footer-based Parquet path consumes it directly -- it reads each file's + footer and bin-packs row groups itself, so it needs paths + sizes rather + than pre-chunked manifest rows. + """ + ... + @dataclass(frozen=True) class FileInfo: @@ -139,6 +183,26 @@ def file_chunker(self) -> FileChunker: """ return self._file_chunker + def as_whole_file_indexer(self) -> "NonSamplingFileIndexer": + """A plain per-file indexer sharing this one's traversal config. + + Metadata-only consumers (the ``PushdownCountFiles`` rule) need a listing + that emits each file exactly once and does no per-file IO while listing. + Subclasses that override :meth:`list_files` with a metadata-aware + strategy -- e.g. ``FooterFileIndexer``, which footer-reads every file on + an actor pool and bin-packs row groups, emitting one manifest row per + file *per bin* -- would both duplicate that IO and risk emitting a path + more than once. So this deliberately returns a base + ``NonSamplingFileIndexer`` rather than ``type(self)``, carrying over only + the traversal settings. + """ + return NonSamplingFileIndexer( + ignore_missing_paths=self._ignore_missing_paths, + num_workers=self._num_workers, + max_paths_per_output=self._max_paths_per_output, + file_chunker=WholeFileChunker(), + ) + def list_files( self, paths: "BlockColumn", @@ -146,16 +210,35 @@ def list_files( filesystem: "FileSystem", pruners: Optional[List[FilePruner]] = None, preserve_order: bool = False, + predicate: Optional["Expr"] = None, + limit: Optional[int] = None, + projected_columns: Optional[List[str]] = None, ) -> Iterable[FileManifest]: - file_info_iterator = ( - self._get_file_info_iterator_threaded(paths, filesystem, preserve_order) - if self._num_workers > 1 - else self._get_file_info_iterator_sequential(paths, filesystem) + # This per-file listing path ignores predicate/limit/projected_columns; + # they're consumed by metadata-aware indexers (e.g. the footer indexer). + # ``list_file_infos`` already skips zero-size files and applies pruners, + # so the manifest builder only has to chunk. + file_infos = self.list_file_infos( + paths, + filesystem=filesystem, + pruners=pruners, + preserve_order=preserve_order, ) + yield from self._process_file_infos_to_manifests(file_infos) - yield from self._process_file_infos_to_manifests( - file_info_iterator, pruners or [] - ) + def _get_file_info_iterator( + self, + paths: "BlockColumn", + filesystem: "FileSystem", + preserve_order: bool, + ) -> Iterable[FileInfo]: + """Threaded (work-stealing) traversal when ``num_workers > 1``, else + sequential. Shared by :meth:`list_files` and :meth:`list_file_infos`.""" + if self._num_workers > 1: + return self._get_file_info_iterator_threaded( + paths, filesystem, preserve_order + ) + return self._get_file_info_iterator_sequential(paths, filesystem) def _get_file_info_iterator_sequential( self, @@ -262,11 +345,41 @@ def _ordered_result_key(result: OrderedFileResult) -> Tuple[int, str]: num_workers=num_workers, ) + def list_file_infos( + self, + paths: "BlockColumn", + *, + filesystem: "FileSystem", + pruners: Optional[List[FilePruner]] = None, + preserve_order: bool = False, + ) -> Iterable[FileInfo]: + """Yield pruned, non-empty ``FileInfo``\\ s (path + on-disk size). + + The raw file-info stream that :meth:`list_files` chunks into manifests. + The footer-based Parquet path consumes this directly -- it reads each + file's footer and bin-packs row groups itself, so it needs paths + sizes + rather than pre-chunked manifest rows. Zero-size files are skipped and + ``pruners`` (file-extension / partition filters) are applied here, so + both listing paths share one filtering point. + """ + pruners = pruners or [] + file_info_iterator = self._get_file_info_iterator( + paths, filesystem, preserve_order + ) + for file_info in file_info_iterator: + if file_info.size is None or file_info.size == 0: + logger.warning(f"Skipping zero-size file: {file_info.path!r}") + continue + if not all(pruner.should_include(file_info.path) for pruner in pruners): + continue + yield file_info + def _process_file_infos_to_manifests( self, file_infos: Iterable[FileInfo], - pruners: List[FilePruner], ) -> Iterable[FileManifest]: + # ``file_infos`` are already filtered (zero-size skipped, pruners applied) + # by ``list_file_infos``; this method only chunks them into manifests. running_paths: List[str] = [] running_file_sizes: List[int] = [] running_chunk_metadatas: List[Optional[ChunkMetadata]] = [] @@ -274,15 +387,10 @@ def _process_file_infos_to_manifests( chunks_count = 0 for file_info in file_infos: + # ``list_file_infos`` already dropped zero/None-size files. + assert file_info.size is not None path, file_size = file_info.path, file_info.size - if file_size is None or file_size == 0: - logger.warning(f"Skipping zero-size file: {path!r}") - continue - - if not all(pruner.should_include(path) for pruner in pruners): - continue - # Drive the chunker once per file; emit one manifest row per chunk. # ``chunk_metadata`` is ``None`` for whole-file chunks (default # ``WholeFileChunker`` behavior and ``ParquetFileChunker`` for files diff --git a/python/ray/data/_internal/datasource_v2/listing/footer_file_indexer.py b/python/ray/data/_internal/datasource_v2/listing/footer_file_indexer.py new file mode 100644 index 000000000000..bb6109bc4dd0 --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/listing/footer_file_indexer.py @@ -0,0 +1,254 @@ +from __future__ import annotations + +import logging +from collections import deque +from typing import TYPE_CHECKING, Deque, Iterable, Iterator, List, Optional, Tuple + +import ray +from ray._common.utils import env_integer +from ray.data._internal.datasource_v2.listing.file_indexer import ( + NonSamplingFileIndexer, +) +from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest +from ray.data._internal.datasource_v2.listing.footer_reader import FooterReaderActor +from ray.data._internal.datasource_v2.partitioners.online_bin_packer import ( + OnlineBinPacker, +) +from ray.data._internal.util import MiB + +if TYPE_CHECKING: + from pyarrow.fs import FileSystem + + from ray.actor import ActorProxy + from ray.data._internal.datasource_v2.listing.file_indexer import FileInfo + from ray.data._internal.datasource_v2.listing.file_pruners import FilePruner + from ray.data._internal.datasource_v2.listing.footer_reader import FooterReader + from ray.data.block import BlockColumn + from ray.data.expressions import Expr + +logger = logging.getLogger(__name__) + +# A pool of footer-reading actors spread across the cluster, provisioned once +# per ``list_files`` call. Footer reads are network-bound, so several actors each +# driving many concurrent reads keeps IO from bottlenecking on a single node. +_DEFAULT_NUM_ACTORS = env_integer("RAY_DATA_PARQUET_FOOTER_NUM_ACTORS", 32) +_DEFAULT_IO_CONCURRENCY = env_integer("RAY_DATA_PARQUET_FOOTER_IO_CONCURRENCY", 128) +# Files per ``read_footers`` call. Small footers -> batch several per task to +# amortize the per-task and per-result object-store overhead. +_DEFAULT_BATCH_SIZE = env_integer("RAY_DATA_PARQUET_FOOTER_BATCH_SIZE", 10) +# ``FileChunks`` per streamed result. The driver pays one object-store fetch per +# yielded list, so a large directory costs one fetch per file at ``1``. Raising +# it trades a little latency-to-first-chunk for far fewer driver-side fetches. +_DEFAULT_RESULT_BATCH_SIZE = env_integer("RAY_DATA_PARQUET_FOOTER_RESULT_BATCH_SIZE", 1) +# Max in-flight footer batches. ``0`` -> auto (``num_actors * 2``). A smaller +# window reads fewer footers before an early ``limit`` stop cancels the pool, at +# the cost of less pipelining on full reads. +_DEFAULT_MAX_INFLIGHT_BATCHES = env_integer( + "RAY_DATA_PARQUET_FOOTER_MAX_INFLIGHT_BATCHES", 0 +) +# Fallback bin budget (uncompressed bytes per read task) when +# ``target_max_block_size`` is unset. +_DEFAULT_BIN_PACKING_BYTES = env_integer( + "RAY_DATA_PARQUET_BIN_PACKING_BYTES", 128 * MiB +) +# Shared (mixed-colour) bins the packer keeps open at once. A wider pool packs +# tighter; a narrower one seals bins sooner, and since each sealed bin becomes a +# read task, that's what lets reads start before every footer has landed. +_DEFAULT_MAX_SHARED_OPEN_BINS = env_integer( + "RAY_DATA_PARQUET_BIN_PACKING_MAX_SHARED_OPEN_BINS", 16 +) + + +class FooterFileIndexer(NonSamplingFileIndexer): + """Lists files, then footer-reads + bin-packs their row groups into manifests. + + Inherits directory traversal and ``list_file_infos`` from + :class:`NonSamplingFileIndexer`; overrides :meth:`list_files` to emit + bin-packed read units instead of per-file chunks. + """ + + def __init__( + self, + *, + ignore_missing_paths: bool, + num_workers: Optional[int] = None, + max_paths_per_output: Optional[int] = None, + coalesce_bytes: int = 0, + split_coalesced: bool = False, + io_concurrency: Optional[int] = None, + footer_batch_size: Optional[int] = None, + result_batch_size: Optional[int] = None, + max_inflight_batches: Optional[int] = None, + max_shared_open_bins: Optional[int] = None, + ): + super().__init__( + ignore_missing_paths=ignore_missing_paths, + num_workers=num_workers, + max_paths_per_output=max_paths_per_output, + ) + self._coalesce_bytes = coalesce_bytes + self._split_coalesced = split_coalesced + # Re-read at construction so ``monkeypatch.setenv`` / release-test env + # overrides work after this module has already been imported. + self._num_actors = env_integer( + "RAY_DATA_PARQUET_FOOTER_NUM_ACTORS", _DEFAULT_NUM_ACTORS + ) + self._io_concurrency = ( + io_concurrency if io_concurrency is not None else _DEFAULT_IO_CONCURRENCY + ) + self._footer_batch_size = ( + footer_batch_size if footer_batch_size is not None else _DEFAULT_BATCH_SIZE + ) + self._result_batch_size = ( + result_batch_size + if result_batch_size is not None + else _DEFAULT_RESULT_BATCH_SIZE + ) + # In-flight footer-batch window; 0/None -> auto (``num_actors * 2``). + _inflight = ( + max_inflight_batches + if max_inflight_batches is not None + else _DEFAULT_MAX_INFLIGHT_BATCHES + ) + self._max_inflight_batches = _inflight if _inflight else self._num_actors * 2 + self._max_shared_open_bins = ( + max_shared_open_bins + if max_shared_open_bins is not None + else _DEFAULT_MAX_SHARED_OPEN_BINS + ) + self._max_bin_bytes = env_integer( + "RAY_DATA_PARQUET_BIN_PACKING_BYTES", _DEFAULT_BIN_PACKING_BYTES + ) + + @property + def yields_read_units(self) -> bool: + # list_files already emits bin-packed read units, so ListFiles skips the + # partitioner and lists in a single task (global packing + one pool). + return True + + def list_files( + self, + paths: "BlockColumn", + *, + filesystem: "FileSystem", + pruners: Optional[List["FilePruner"]] = None, + preserve_order: bool = False, + predicate: Optional["Expr"] = None, + limit: Optional[int] = None, + projected_columns: Optional[List[str]] = None, + ) -> Iterable[FileManifest]: + max_bin_bytes = self._max_bin_bytes + file_infos = self.list_file_infos( + paths, + filesystem=filesystem, + pruners=pruners, + preserve_order=preserve_order, + ) + actors: List[ActorProxy[FooterReader]] = [ + FooterReaderActor.options(scheduling_strategy="SPREAD").remote( + filesystem, + self._io_concurrency, + predicate, + projected_columns, + self._coalesce_bytes, + ) + for _ in range(self._num_actors) + ] + logger.debug( + "Provisioned %d FooterReader actors (io_concurrency=%d)", + self._num_actors, + self._io_concurrency, + ) + try: + yield from self._read_and_pack(actors, file_infos, max_bin_bytes, limit) + finally: + for actor in actors: + # ``ActorProxy`` is ``ActorHandle | type[T]``; kill wants a handle. + # pyrefly: ignore[bad-argument-type] + ray.kill(actor) + + def _read_and_pack( + self, + actors: List["ActorProxy[FooterReader]"], + file_infos: "Iterable[FileInfo]", + max_bin_bytes: int, + limit: Optional[int], + ) -> Iterator[FileManifest]: + packer = OnlineBinPacker( + max_bin_bytes, + max_shared_open_bins=self._max_shared_open_bins, + split_coalesced=self._split_coalesced, + ) + # Bound the number of in-flight footer batches so listing stays roughly + # demand-driven (matters under a limit) and memory stays flat. + window = max(1, self._max_inflight_batches) + batches = self._batches(file_infos) + # FIFO of in-flight streaming generators, one per dispatched footer batch. + pending: Deque[ray.ObjectRefGenerator] = deque() + batch_no = 0 + delivered_fully_matched_rows = 0 + + def dispatch_next() -> bool: + nonlocal batch_no + batch = next(batches, None) + if batch is None: + return False + actor: ActorProxy[FooterReader] = actors[batch_no % len(actors)] + # Streaming ``@ray.method``; stub types ``.remote`` as ``ObjectRef`` + # and don't preserve the method's keyword args. + # pyrefly: ignore[bad-assignment] + gen: ray.ObjectRefGenerator = actor.read_footers.remote( + batch, + result_batch_size=self._result_batch_size, # pyrefly: ignore[unexpected-keyword] + ) + pending.append(gen) + batch_no += 1 + return True + + # Prime the window. + for _ in range(window): + if not dispatch_next(): + break + + while pending: + gen = pending.popleft() + for ref in gen: # blocks until this generator's next result lands + for file_chunks in ray.get(ref): + packer.add_file_chunks(file_chunks) + if limit is not None: + # Count only fully-matched (exact-survivor) rows so + # stopping can never under-deliver under a filter. + delivered_fully_matched_rows += sum( + rg.num_rows + for rg in file_chunks.row_groups + if rg.fully_matched + ) + while packer.has_partition(): + yield packer.next_partition() + if limit is not None and delivered_fully_matched_rows >= limit: + # Flush open bins so a small limit yields promptly; abandon + # in-flight generators (the actor teardown cancels them). + packer.finalize() + while packer.has_partition(): + yield packer.next_partition() + return + # This generator drained; keep the window full. + dispatch_next() + + packer.finalize() + while packer.has_partition(): + yield packer.next_partition() + + def _batches( + self, file_infos: "Iterable[FileInfo]" + ) -> Iterator[List[Tuple[str, int]]]: + batch: List[Tuple[str, int]] = [] + for file_info in file_infos: + if file_info.size is None: + continue + batch.append((file_info.path, file_info.size)) + if len(batch) >= self._footer_batch_size: + yield batch + batch = [] + if batch: + yield batch diff --git a/python/ray/data/_internal/datasource_v2/listing/footer_reader.py b/python/ray/data/_internal/datasource_v2/listing/footer_reader.py new file mode 100644 index 000000000000..54e8fcc60486 --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/listing/footer_reader.py @@ -0,0 +1,250 @@ +from __future__ import annotations + +import logging +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import replace +from typing import TYPE_CHECKING, Iterable, Iterator, List, Optional, Set, Tuple + +import pyarrow as pa +import pyarrow.dataset as pds +import pyarrow.fs as pafs +from pyarrow.parquet import RowGroupMetaData + +import ray +from ray.data._internal.datasource_v2.chunkers.parquet_footer_types import ( + FileChunks, + RowGroupInfo, +) + +if TYPE_CHECKING: + import pyarrow.compute as pc + + from ray.data.expressions import Expr + +logger = logging.getLogger(__name__) + + +def coalesce_row_groups( + per_rg: List[RowGroupInfo], target: int +) -> Tuple[RowGroupInfo, ...]: + """Merge runs of consecutive row groups into ~``target``-byte chunks. + + A run breaks on: a change in ``fully_matched`` (never merge across the + match-class boundary, or limit push-down would miscount), a gap in the + ``rg_idx`` sequence (e.g. filter-pruned groups), or once the accumulator + reaches ``target``. A single row group larger than ``target`` forms its own + chunk. ``target == 0`` disables coalescing entirely (one chunk per physical + row group). Coalesced chunks (``rg_count > 1``) carry per-row-group + ``rg_sizes`` / ``rg_rows`` so the packer can split them back at exact + boundaries. Pure function -- unit-tested directly. + """ + if not target: + return tuple(per_rg) + out: List[RowGroupInfo] = [] + cur: Optional[RowGroupInfo] = None + cur_sizes: List[int] = [] + cur_rows: List[int] = [] + + def _flush() -> None: + if cur is None: + return + if len(cur_sizes) > 1: + out.append(replace(cur, rg_sizes=tuple(cur_sizes), rg_rows=tuple(cur_rows))) + else: + out.append(cur) + + for rg in per_rg: # ascending rg_idx; each is a single physical row group + if ( + cur is not None + and cur.fully_matched == rg.fully_matched + and cur.rg_idx + cur.rg_count == rg.rg_idx # contiguous + and cur.uncompressed_size < target # not full yet + ): + cur = replace( + cur, + rg_count=cur.rg_count + rg.rg_count, + uncompressed_size=cur.uncompressed_size + rg.uncompressed_size, + num_rows=cur.num_rows + rg.num_rows, + ) + cur_sizes.append(rg.uncompressed_size) + cur_rows.append(rg.num_rows) + else: + _flush() + cur = rg + cur_sizes = [rg.uncompressed_size] + cur_rows = [rg.num_rows] + _flush() + return tuple(out) + + +class FooterReader: + """Reads Parquet footers and chunks files into row-group runs. + + Run as a pool of Ray actors (see :data:`FooterReaderActor`) to spread footer + IO across the cluster instead of bottlenecking on the driver. ``read_footers`` + is a streaming generator that yields ``FileChunks`` in small batches as their + footers land, so the driver does far fewer object-store fetches than one per + file. + + Kept as a plain class (the actor is created via the functional + ``ray.remote(FooterReader)`` form below) so callers can type actor handles as + ``ActorProxy[FooterReader]`` -- see Ray's type-hint docs + (https://docs.ray.io/en/latest/ray-core/type-hint.html). + """ + + def __init__( + self, + filesystem: pafs.FileSystem, + io_concurrency: int = 128, + filter_expr: Optional["Expr"] = None, + projected_cols: Optional[List[str]] = None, + coalesce_bytes: int = 0, + ): + # Coalescing target: merge contiguous row groups into chunks of + # ~coalesce_bytes uncompressed before returning them, so the driver packs + # fewer items. 0 disables coalescing (one chunk per physical row group). + self.coalesce_bytes = coalesce_bytes + self.filesystem = filesystem + self.pool = ThreadPoolExecutor(max_workers=io_concurrency) + # Match Arrow's process-wide pools to the actor's IO concurrency so + # nested S3/footer work isn't bottlenecked on the default 8 threads. + pa.set_io_thread_count(io_concurrency) + pa.set_cpu_count(io_concurrency) + # Lower the Ray Data predicate to a native PyArrow compute expression once + # (per actor) so it can be pushed into split_by_row_group for row-group + # pruning. ``None`` means "no predicate" -> keep every row group. + self.filter: Optional["pc.Expression"] = ( + filter_expr.to_pyarrow() if filter_expr is not None else None + ) + # Projection pushdown: when set, row-group byte sizes are accounted over + # only these top-level columns, since the reader will only fetch those. + self.projected_cols: Optional[Set[str]] = ( + set(projected_cols) if projected_cols is not None else None + ) + # Reused across files; make_fragment is what lets us apply + # split_by_row_group. + self.file_format = pds.ParquetFileFormat() + + def _projected_leaf_indices(self, row_group) -> Optional[List[int]]: + # Map the requested top-level column names to Parquet leaf-column indices + # (a nested field expands to several leaves, e.g. "a.b.list.element"). + # Returns None to signal "all columns" so callers can take the cheap path. + if self.projected_cols is None: + return None + return [ + j + for j in range(row_group.num_columns) + if row_group.column(j).path_in_schema.split(".", 1)[0] + in self.projected_cols + ] + + def _row_group_info( + self, + row_group: RowGroupMetaData, + rg_idx: int, + leaf_indices: Optional[List[int]], + fully_matched: bool = True, + ) -> RowGroupInfo: + if leaf_indices is None: + # total_byte_size is a single cheap accessor for the whole row group, + # so we avoid walking columns entirely on the no-projection path. + uncompressed = row_group.total_byte_size + else: + # Sum only the projected leaves so bin packing reflects the bytes the + # downstream reader will actually pull for this row group. + uncompressed = sum( + row_group.column(j).total_uncompressed_size for j in leaf_indices + ) + return RowGroupInfo( + rg_idx=rg_idx, + uncompressed_size=uncompressed, + num_rows=row_group.num_rows, + fully_matched=fully_matched, + ) + + def _read_and_chunk(self, path: str, size: int) -> FileChunks: + fragment = self.file_format.make_fragment(path, filesystem=self.filesystem) + metadata = fragment.metadata # reads + caches the footer on the fragment + if self.filter is not None: + # Predicate pushdown: drop row groups whose Parquet statistics + # contradict the filter, so we never emit chunks the reader would + # skip anyway. Each returned fragment views a single surviving group. + surviving = fragment.split_by_row_group(self.filter) + rg_indices: Iterable[int] = [sub.row_groups[0].id for sub in surviving] + # Classify surviving groups as fully- vs partially-matching via + # predicate negation (a la DataFusion): a group is fully matched iff + # ~filter cannot match any of its rows, i.e. the negation prunes it + # out. For those, num_rows is an exact survivor count and can drive + # the limit push-down. Any failure defaults to "not fully matched", + # which is always safe (falls back to the post-filter stop). + try: + not_fully = { + sub.row_groups[0].id + for sub in fragment.split_by_row_group(~self.filter) + } + except Exception: + not_fully = set(rg_indices) + fully_by_idx: Optional[dict] = {i: i not in not_fully for i in rg_indices} + else: + rg_indices = range(metadata.num_row_groups) + fully_by_idx = None # no predicate -> every group is fully matched + + leaf_indices = ( + self._projected_leaf_indices(metadata.row_group(0)) + if metadata.num_row_groups + else None + ) + per_rg = [ + self._row_group_info( + metadata.row_group(i), + i, + leaf_indices, + fully_matched=(True if fully_by_idx is None else fully_by_idx[i]), + ) + for i in rg_indices + ] + # Coalesce contiguous row groups into ~coalesce_bytes chunks (no-op when + # coalesce_bytes == 0) so only a handful of descriptors per file reach the + # driver's bin-packer instead of one per physical row group. + row_groups = coalesce_row_groups(per_rg, self.coalesce_bytes) + return FileChunks(path=path, size=size, row_groups=row_groups) + + @ray.method(num_returns="streaming") + def read_footers( + self, files: List[Tuple[str, int]], *, result_batch_size: int = 1 + ) -> Iterator[List[FileChunks]]: + """Read the footers of ``files`` concurrently, yielding ``FileChunks``. + + Yields lists of ``FileChunks`` (not single results): each list the driver + receives costs one object-store fetch, so batching cuts driver-side + deserialization overhead ~``result_batch_size``-fold. At the default of + ``1`` a directory of N files costs N fetches on the single listing task; + callers set it via ``RAY_DATA_PARQUET_FOOTER_RESULT_BATCH_SIZE``. + ``num_returns`` is fixed to ``"streaming"`` via ``@ray.method`` so + results stream out as footers land. + """ + futures = [ + self.pool.submit(self._read_and_chunk, path, size) for path, size in files + ] + buffer: List[FileChunks] = [] + total_row_groups = 0 + for finished in as_completed(futures): + chunk = finished.result() + total_row_groups += len(chunk.row_groups) + buffer.append(chunk) + if len(buffer) >= result_batch_size: + yield buffer + buffer = [] + if buffer: + yield buffer + logger.debug( + "FooterReader batch: %d files, %d row groups", + len(files), + total_row_groups, + ) + + +# The Ray actor class. Built via the functional ``ray.remote(...)`` form (rather +# than the ``@ray.remote`` decorator) so ``FooterReader`` stays a plain class and +# actor handles can be typed ``ActorProxy[FooterReader]``. +FooterReaderActor = ray.remote(FooterReader) diff --git a/python/ray/data/_internal/datasource_v2/listing/listing_utils.py b/python/ray/data/_internal/datasource_v2/listing/listing_utils.py index 7287e59bed8b..01c464b14c4d 100644 --- a/python/ray/data/_internal/datasource_v2/listing/listing_utils.py +++ b/python/ray/data/_internal/datasource_v2/listing/listing_utils.py @@ -16,7 +16,7 @@ ) from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder from ray.data._internal.execution.interfaces.task_context import TaskContext -from ray.data.block import Block, BlockAccessor +from ray.data.block import Block if TYPE_CHECKING: from pyarrow.fs import FileSystem @@ -24,6 +24,7 @@ from ray.data._internal.datasource_v2.listing.file_indexer import FileIndexer from ray.data.datasource.file_based_datasource import FileShuffleConfig from ray.data.datasource.partitioning import PathPartitionFilter + from ray.data.expressions import Expr def partition_files( @@ -62,6 +63,9 @@ def list_files_for_each_block( file_extensions: Optional[List[str]] = None, partition_filter: Optional["PathPartitionFilter"] = None, preserve_order: bool = False, + predicate: Optional["Expr"] = None, + limit: Optional[int] = None, + projected_columns: Optional[List[str]] = None, ) -> Iterable[Block]: """Expand path blocks into ``FileManifest`` blocks. @@ -71,6 +75,11 @@ def list_files_for_each_block( Pruners are constructed once per task from ``file_extensions`` / ``partition_filter`` — keeps pruner construction out of the ``_read_datasource_v2`` entry point. + + Pushed-down ``predicate`` / ``limit`` / ``projected_columns`` are forwarded + to the indexer; metadata-aware indexers (e.g. the footer-based Parquet + indexer) use them to skip row groups, stop early, and size projected + columns, while the plain indexer ignores them. """ pruners = _build_pruners(file_extensions, partition_filter) for block in blocks: @@ -79,6 +88,9 @@ def list_files_for_each_block( filesystem=filesystem, pruners=pruners, preserve_order=preserve_order, + predicate=predicate, + limit=limit, + projected_columns=projected_columns, ): if len(manifest) > 0: yield manifest.as_block() @@ -120,40 +132,37 @@ def sample_files( pruners: Optional[List[FilePruner]] = None, max_files: int = 16, ) -> FileManifest: - """Drive the indexer until up to ``max_files`` files arrive; return them. - - Used for driver-side schema inference in ``_read_datasource_v2``. - Sampling more than one file lets callers unify schemas (e.g., if the - first file has an all-null column, later files' non-null types can - promote it). No caching — the returned manifest is discarded after - schema inference, and the ``ListFiles`` op lists the same paths - again on workers at execution time. + """List up to ``max_files`` files and return them as a whole-file manifest. + + Used for driver-side schema inference in ``_read_datasource_v2``. Sampling + more than one file lets callers unify schemas (e.g., if the first file has an + all-null column, later files' non-null types can promote it). No caching -- + the returned manifest is discarded after schema inference, and the + ``ListFiles`` op lists the same paths again on workers at execution time. + + Uses ``list_file_infos`` (raw path + size), not ``list_files``, so that + metadata-heavy indexers (e.g. the footer-based Parquet indexer) don't do + their footer-read + bin-pack work on the driver just to sample a schema -- + schema inference only needs the paths. """ assert max_files >= 1 paths_column = pa.array(paths, type=pa.string()) - collected: List[FileManifest] = [] - collected_rows = 0 - for manifest in indexer.list_files( + sampled_paths: List[str] = [] + sampled_sizes: List[int] = [] + for file_info in indexer.list_file_infos( paths_column, filesystem=filesystem, pruners=pruners or [], preserve_order=True, ): - if len(manifest) == 0: + if file_info.size is None: continue - remaining = max_files - collected_rows - if len(manifest) <= remaining: - collected.append(manifest) - collected_rows += len(manifest) - else: - collected.append( - FileManifest( - BlockAccessor.for_block(manifest.as_block()).slice(0, remaining) - ) - ) - collected_rows = max_files - if collected_rows >= max_files: + sampled_paths.append(file_info.path) + sampled_sizes.append(file_info.size) + if len(sampled_paths) >= max_files: break - if not collected: - return FileManifest.construct_manifest(paths=[], sizes=[], chunk_metadatas=[]) - return FileManifest.concat(collected) + return FileManifest.construct_manifest( + paths=sampled_paths, + sizes=sampled_sizes, + chunk_metadatas=[None] * len(sampled_paths), + ) diff --git a/python/ray/data/_internal/datasource_v2/logical_optimizers.py b/python/ray/data/_internal/datasource_v2/logical_optimizers.py index d6852424509f..ef347a73310a 100644 --- a/python/ray/data/_internal/datasource_v2/logical_optimizers.py +++ b/python/ray/data/_internal/datasource_v2/logical_optimizers.py @@ -31,6 +31,18 @@ def push_filters(self, predicate: "Expr") -> Tuple["Scanner", Optional["Expr"]]: """ ... + @abstractmethod + def pushed_predicate(self) -> Optional["Expr"]: + """The predicate this scanner will apply at read time, if any. + + This is the accepted result of :meth:`push_filters`, not the predicate + that was offered to it. Planning derives upstream listing-time pruning + from this value, so it must never be stronger than what the scanner + actually evaluates -- returning ``None`` is always safe, returning a + predicate the reader does not apply drops rows. + """ + ... + @DeveloperAPI class SupportsColumnPruning(ABC): @@ -84,6 +96,16 @@ def push_limit(self, limit: int) -> "Scanner": """ ... + @abstractmethod + def pushed_limit(self) -> Optional[int]: + """The row limit this scanner will stop at, if any. + + This is the accepted result of :meth:`push_limit`. Planning derives + early-stop listing from it, so it must never be smaller than the limit + the scanner actually honors. + """ + ... + @DeveloperAPI class SupportsPartitionPruning(ABC): @@ -120,3 +142,35 @@ def prune_partitions(self, predicate: "Expr") -> "Scanner": New Scanner instance with partition pruning applied. """ ... + + +def derive_list_files_pushdown( + scanner: Optional["Scanner"], +) -> Tuple[Optional["Expr"], Optional[List[str]], Optional[int]]: + """Read the pushed-down state a scanner accepted, for upstream listing. + + Returns ``(predicate, projected_columns, limit)`` -- the constraints a + ``ListFiles`` feeding this scanner's ``ReadFiles`` may safely apply while + listing (see :class:`~ray.data._internal.logical.rules. + derive_list_files_pushdown.DeriveListFilesPushdown`). Each element is + ``None`` unless the scanner both implements the corresponding ``Supports*`` + mixin and reports state it actually accepted, so a datasource that ignores + a pushdown can never cause listing-time pruning. + + ``scanner`` may be ``None`` (no downstream reader), which yields all-``None``: + nothing downstream applies these constraints, so listing must not either. + """ + predicate = ( + scanner.pushed_predicate() + if isinstance(scanner, SupportsFilterPushdown) + else None + ) + if isinstance(scanner, SupportsColumnPruning): + pruned = scanner.pruned_column_names() + projected_columns = list(pruned) if pruned is not None else None + else: + projected_columns = None + limit = ( + scanner.pushed_limit() if isinstance(scanner, SupportsLimitPushdown) else None + ) + return predicate, projected_columns, limit diff --git a/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/.gitignore b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/.gitignore new file mode 100644 index 000000000000..f2a64375d09c --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/.gitignore @@ -0,0 +1,12 @@ +/target +# maturin build output. Platform-specific wheels must never be committed — one +# slipped into a rebase on 2026-08-10 via a broad `git add -A`. +/dist +# Cargo.lock IS committed: this crate builds a wheel for reproducible CI, and +# ci/docker/dataarrowrs.build.wanda.yaml lists Cargo.lock in its build srcs. +# patch_crate_parquet.sh (arrow_rs_probe, TODO 1o A/B) build-local state: the +# vendored+patched parquet copy, the [patch.crates-io] override, and the stock +# Cargo.lock backup. None of these may ever be committed. +/vendor +/.cargo +/Cargo.lock.stock diff --git a/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/Cargo.lock b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/Cargo.lock new file mode 100644 index 000000000000..a8eeee651726 --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/Cargo.lock @@ -0,0 +1,2637 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "const-random", + "getrandom 0.3.4", + "once_cell", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e76a019e91224d279006ff972f1e984179a6e9feb050adba6ce8274aef23195" +dependencies = [ + "alloc-no-stdlib", +] + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "arrow" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b952ca5a8046ad741b60f142d6eca4aeebcad615694202bc64c5341f23e32c5b" +dependencies = [ + "arrow-arith", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-csv", + "arrow-data", + "arrow-ipc", + "arrow-json", + "arrow-ord", + "arrow-row", + "arrow-schema", + "arrow-select", + "arrow-string", +] + +[[package]] +name = "arrow-arith" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64a13b8d3008c4e9063c597a08f46446fe3fd5789277127672d6c0bdbb43b1ff" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "num-traits", +] + +[[package]] +name = "arrow-array" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9486151b2f0785bafc6fa04fc5c99fcb4495455662e58787ea32eaaed33c4192" +dependencies = [ + "ahash", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "hashbrown", + "num-complex", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-buffer" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4776577a87794bfdf0b4e90e2ea12454fa7738ea2823c4be5b9d1851da7b434" +dependencies = [ + "bytes", + "half", + "num-bigint", + "num-traits", +] + +[[package]] +name = "arrow-cast" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9ad451ce4f98710828a455b96991b8f031deb2e67f5fcad6773f017e4a69c3a" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-ord", + "arrow-schema", + "arrow-select", + "atoi", + "base64", + "chrono", + "half", + "lexical-core", + "num-traits", + "ryu", +] + +[[package]] +name = "arrow-csv" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8aa7bf96d6141a7bcca2eed57c7c9767d2a2175281857b8a7b68308992864784" +dependencies = [ + "arrow-array", + "arrow-cast", + "arrow-schema", + "chrono", + "csv", + "csv-core", + "regex", +] + +[[package]] +name = "arrow-data" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b38fe43e2e8704360f1464e6e8cc4fc381ef02cc4fb0192afa8df1aaa0115c66" +dependencies = [ + "arrow-buffer", + "arrow-schema", + "half", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-ipc" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29dac499fcbc6ba74ee0324057821d381929a48526a3966bd9dffb44aa06d98c" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "flatbuffers", +] + +[[package]] +name = "arrow-json" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fe05e916ddc50f4c7a363cd69c0ef5894fcee063517e9a0b8582f0c56746af6" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-ord", + "arrow-schema", + "arrow-select", + "chrono", + "half", + "indexmap", + "itoa", + "lexical-core", + "memchr", + "num-traits", + "ryu", + "serde_core", + "serde_json", + "simdutf8", +] + +[[package]] +name = "arrow-ord" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e13dbdc2a9c053c10c7baa6e30faee04a180aa7ce88e471835850ce37abd20b" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", +] + +[[package]] +name = "arrow-row" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d5a1f8c733d15260b305683472ee8ad89c62cbd706703ca873b90d051b41592" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "half", +] + +[[package]] +name = "arrow-schema" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9e4969dc350d571766247143ab36a5187d095d3d3690970408bc630d47c69e5" +dependencies = [ + "bitflags", +] + +[[package]] +name = "arrow-select" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "402770dba90865359d98d1ef92ef16e23d75c0cca9c2c880c8a05468b7743bf9" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "num-traits", +] + +[[package]] +name = "arrow-string" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2b0afbb8b9016700938291123df30838b89decc3213dba00852021988b170d3" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "memchr", + "num-traits", + "regex", + "regex-syntax", +] + +[[package]] +name = "async-trait" +version = "0.1.91" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.2", +] + +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "brotli" +version = "8.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc91aac060a7a1e25823bdccbfb6af1875b88f17c6daac97894eed8207166b3" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a32acac15fe1967bc3986b2a6347dffc965602354ea6f450ad07e8bfd253583" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "cc" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89588d05638b5b4594a3348a2d6c20277e43a7f5c5202b05cc56888475a47b8" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" + +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures", + "rand_core", +] + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "num-traits", + "serde", + "windows-link", +] + +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom 0.2.17", + "once_cell", + "tiny-keccak", +] + +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "csv" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", +] + +[[package]] +name = "csv-core" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +dependencies = [ + "memchr", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" +dependencies = [ + "bitflags", + "rustc_version", +] + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "miniz_oxide", + "zlib-rs", +] + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a88cf1f829d945f548cf8fec32c61b1f202b6d93b45848602fc02af4b12ad218" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" + +[[package]] +name = "futures-executor" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6754879cc9f2c66f88c6e5c35344bb0bdb0708b0352b1201815667c7eabc7458" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a" + +[[package]] +name = "futures-macro" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "futures-sink" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307" + +[[package]] +name = "futures-task" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" + +[[package]] +name = "futures-util" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "wasi", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi 6.0.0", + "rand_core", + "wasm-bindgen", +] + +[[package]] +name = "h2" +version = "0.4.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap", + "slab", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "humantime" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15cdd26707701c53297e2fa6afb323d55fbc1d0810c3aec078ae3ef0424c3c15" + +[[package]] +name = "hyper" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d22053281f852e11534f5198498373cbb59295120a20771d90f7ed1897490a72" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "h2", + "http", + "http-body", + "httparse", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" +dependencies = [ + "http", + "hyper", + "hyper-util", + "rustls", + "rustls-native-certs", + "tokio", + "tokio-rustls", + "tower-service", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2", + "tokio", + "tower-service", + "tracing", +] + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown", +] + +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "jobserver" +version = "0.1.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c00acbd29eabad4a2392fa0e921c874934dbbf4194312ad20f04a0ed67a3cb3" +dependencies = [ + "getrandom 0.4.3", + "libc", +] + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "lexical-core" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d8d125a277f807e55a77304455eb7b1cb52f2b18c143b60e766c120bd64a594" +dependencies = [ + "lexical-parse-float", + "lexical-parse-integer", + "lexical-util", + "lexical-write-float", + "lexical-write-integer", +] + +[[package]] +name = "lexical-parse-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52a9f232fbd6f550bc0137dcb5f99ab674071ac2d690ac69704593cb4abbea56" +dependencies = [ + "lexical-parse-integer", + "lexical-util", +] + +[[package]] +name = "lexical-parse-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a7a039f8fb9c19c996cd7b2fcce303c1b2874fe1aca544edc85c4a5f8489b34" +dependencies = [ + "lexical-util", +] + +[[package]] +name = "lexical-util" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2604dd126bb14f13fb5d1bd6a66155079cb9fa655b37f875b3a742c705dbed17" + +[[package]] +name = "lexical-write-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50c438c87c013188d415fbabbb1dceb44249ab81664efbd31b14ae55dabb6361" +dependencies = [ + "lexical-util", + "lexical-write-integer", +] + +[[package]] +name = "lexical-write-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "409851a618475d2d5796377cad353802345cba92c867d9fbcde9cf4eac4e14df" +dependencies = [ + "lexical-util", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "lz4_flex" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" +dependencies = [ + "twox-hash", +] + +[[package]] +name = "md-5" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" +dependencies = [ + "cfg-if", + "digest", +] + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "mio" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" +dependencies = [ + "libc", + "wasi", + "windows-sys 0.61.2", +] + +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", + "libm", +] + +[[package]] +name = "object_store" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "622acbc9100d3c10e2ee15804b0caa40e55c933d5aa53814cd520805b7958a49" +dependencies = [ + "async-trait", + "base64", + "bytes", + "chrono", + "form_urlencoded", + "futures-channel", + "futures-core", + "futures-util", + "http", + "http-body-util", + "humantime", + "hyper", + "itertools", + "md-5", + "parking_lot", + "percent-encoding", + "quick-xml", + "rand", + "reqwest", + "ring", + "serde", + "serde_json", + "serde_urlencoded", + "thiserror", + "tokio", + "tracing", + "url", + "walkdir", + "wasm-bindgen-futures", + "web-time", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "parquet" +version = "59.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5302d4da74d6596a1f11f9928767995b53bca657cbeea1e4e8c5074f8a1157dd" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-ipc", + "arrow-schema", + "arrow-select", + "base64", + "brotli", + "bytes", + "chrono", + "crc32fast", + "flate2", + "futures", + "half", + "hashbrown", + "lz4_flex", + "num-bigint", + "num-integer", + "num-traits", + "object_store", + "paste", + "seq-macro", + "simdutf8", + "snap", + "tokio", + "twox-hash", + "zstd", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "portable-atomic" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f402062616ab18202ae8319da13fa4279883a2b8a9d9f83f20dbade813ce1884" +dependencies = [ + "cfg-if", + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b14b5775b5ff446dd1056212d778012cbe8a0fbffd368029fd9e25b514479c38" +dependencies = [ + "once_cell", + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ab5bcf04a2cdcbb50c7d6105de943f543f9ed92af55818fd17b660390fc8636" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fd24d897903a9e6d80b968368a34e1525aeb719d568dba8b3d4bfa5dc67d453" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.22.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36c011a03ba1e50152b4b394b479826cad97e7a21eb52df179cd91ac411cbfbe" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "quick-xml" +version = "0.39.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdcc8dd4e2f670d309a5f0e83fe36dfdc05af317008fea29144da1a2ac858e5e" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "quinn" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash", + "rustls", + "socket2", + "thiserror", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4bfc015262b9df63c8845072ce59068853ff5872180c2ce2f13038b970e560" +dependencies = [ + "bytes", + "getrandom 0.4.3", + "lru-slab", + "rand", + "rand_pcg", + "ring", + "rustc-hash", + "rustls", + "rustls-pki-types", + "slab", + "thiserror", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35a133f956daabe89a61a685c2649f13d82d5aa4bd5d12d1277e1072a21c0694" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2", + "tracing", + "windows-sys 0.61.2", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom 0.4.3", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "rand_pcg" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caa0f4137e1c0a72f4c651489402276c8e8e1cf081f3b0ba156d2cbeef09e86a" +dependencies = [ + "rand_core", +] + +[[package]] +name = "ray_data_arrow_rs" +version = "0.1.0" +dependencies = [ + "arrow", + "bytes", + "futures", + "object_store", + "parquet", + "pyo3", + "serde", + "serde_json", + "tokio", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64", + "bytes", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-native-certs", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rustc-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" + +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + +[[package]] +name = "rustls" +version = "0.23.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" +dependencies = [ + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "764899a24af3980067ee14bc143654f297b22eaebfe3c7b6b211920a5a59b046" +dependencies = [ + "web-time", + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c429a8649f110dddef65e2a5ad240f747e85f7758a6bccc7e5777bd33f756e" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "semver" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" + +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.2", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "simd-adler32" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a219298ac11a56ea9a6d2120044824d6f01aeb034955e7af7bc16858527deea" + +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "snap" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "199905e6153d6405f9728fe44daace35f8f837bbf830bb6e85fbd5828709a886" + +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a207d6d6a2b7fc470b80443726053f18a2481b7e1eee970597051596567987a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "target-lexicon" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61c41af27dd6d1e27b1b16b489db798443478cef1f06a660c96db617ba5de3b1" + +[[package]] +name = "thiserror" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.2", +] + +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinyvec" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokio" +version = "1.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "socket2", + "tokio-macros", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6328af13490e73a9b4694030fafd93f8c8c6a9dede33e821c3fc63eddf8042ba" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", +] + +[[package]] +name = "tower-http" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" +dependencies = [ + "bitflags", + "bytes", + "futures-util", + "http", + "http-body", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", + "url", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "twox-hash" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9" + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c62df1340f32221cb9c54d6a27b030e3dba64361d4a95bed55f9aacb44da291d" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 2.0.119", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "web-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8622dcb61c0bcc9fffa6938bed81210af2da9a7e4a1a834b2e37a59b6dfb6141" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5a105cd7b140f6eeec8acff2ea38135d3cab283ada58540f629fe51e46696eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fe976fb70c78cd64cccfe3a6fc142244e8a77b70959b30faf9d0ac37ee228eb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zlib-rs" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b142a20ec14a91d5bc708c1dc21b080c550113d8aa77afa29635673a65dd02c5" + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" + +[[package]] +name = "zstd" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e91ee311a569c327171651566e07972200e76fcfe2242a4fa446149a3881c08a" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "7.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f49c4d5f0abb602a93fb8736af2a4f4dd9512e36f7f570d66e65ff867ed3b9d" +dependencies = [ + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "2.0.16+zstd.1.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e19ebc2adc8f83e43039e79776e3fda8ca919132d68a1fed6a5faca2683748" +dependencies = [ + "cc", + "pkg-config", +] diff --git a/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/Cargo.toml b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/Cargo.toml new file mode 100644 index 000000000000..2e66faabccc5 --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/Cargo.toml @@ -0,0 +1,60 @@ +[package] +name = "ray_data_arrow_rs" +version = "0.1.0" +edition = "2021" +description = "Experimental arrow-rs Parquet reader for Ray Data (prototype)" + +[lib] +name = "ray_data_arrow_rs" +# cdylib for the Python extension module; rlib so the logic can be unit-tested +# from Rust too. +crate-type = ["cdylib", "rlib"] + +[dependencies] +# Pin to the same major the surviving standalone benchmark used, so behavior +# matches the numbers in Agents.md. object_store MUST match the version parquet +# pins (0.13.x) or ParquetObjectReader's API won't line up. +# `crc`: verify each data page's stored CRC32 during decode (decode_page in +# parquet's serialized_reader, used by BOTH the sync/local and async/S3 paths). +# Compile-time only — there is no runtime toggle — and a no-op for pages that +# don't store a CRC (pyarrow doesn't write them by default). Enables native +# support for `page_checksum_verification=True`; an explicit `False` (read +# *despite* bad checksums) still falls back to PyArrow, the only reader that +# can skip the check. +parquet = { version = "59", features = ["arrow", "async", "object_store", "crc"] } +# Only the `ffi` feature — we hand batches to Python via the Arrow C-stream +# (FFI_ArrowArrayStream + PyCapsule), NOT arrow's `pyarrow` conversions. Pulling +# `pyarrow` would drag in arrow-pyarrow → pyo3 0.28, conflicting with our pyo3 0.22. +arrow = { version = "59", features = ["ffi"] } +object_store = { version = "0.13", features = ["aws"] } +# `extension-module` is behind a default feature so `cargo test` can run the +# pure-Rust unit tests (predicate pruning) by linking libpython: +# cargo test --no-default-features +# maturin builds with default features on, so the shipped extension is unchanged. +pyo3 = { version = "0.22", features = ["abi3-py39"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync"] } +futures = "0.3" +# Already in the tree via parquet/object_store; direct dep only so the +# column-prefetch path can hold/slice fetched buffers (`Bytes` is refcounted, +# so slices share the underlying allocation — no copies). +bytes = "1" +# Predicate IR is passed from Python as a small JSON string (built from the Ray +# `Expr` AST) and parsed here into a `Pred` tree for row-group stats pruning. +# JSON keeps the boundary language-agnostic and lets the pruning logic be +# unit-tested in pure Rust with string literals (no parquet fixtures needed). +serde = { version = "1", features = ["derive"] } +serde_json = "1" + +# No optional allocator features: an allocator-retention theory was tested and +# disproven (Agents.md §3.5.1/§7.8) — the system allocator is correct here, and +# mimalloc-as-global-allocator segfaulted Ray workers across the FFI boundary. +# Allocator A/B experiments use LD_PRELOAD, which needs no crate support. + +[features] +default = ["extension-module"] +extension-module = ["pyo3/extension-module"] + +[profile.release] +opt-level = 3 +lto = true +codegen-units = 1 diff --git a/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/README.md b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/README.md new file mode 100644 index 000000000000..9aee8fe913db --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/README.md @@ -0,0 +1,54 @@ +# `ray_data_arrow_rs` — experimental arrow-rs Parquet reader + +Native PyO3 extension backing `ArrowRsParquetFileReader` +(`_internal/datasource_v2/readers/arrow_rs_parquet_file_reader.py`). Selected at +runtime by `DataContext.use_arrow_rs_parquet_reader` (only under +`use_datasource_v2=True`). + +## Status + +This source is a **reconstruction** from the surviving standalone benchmark +(`main.rs`). A prebuilt macOS/cpython-3.12 `.so` (v0.1.0) is already installed in +the dev venv and is the behavioral source of truth for local (macOS) work — the +integration and benchmark run against it today. This crate exists so the reader +can be: + +1. **rebuilt for Linux/x86-64** (the "deciding experiment" in `Agents.md` §7 must + run on Ray's Linux runtime with USS metrics + real S3), and +2. **evolved** to expose the tuning knobs the standalone benchmark had but the + v0.1.0 API doesn't (K intra-row-group split, fetch window, byte budget). + +It has **not** been compiled in-session. Before relying on it, build and validate: + +```bash +cd python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs +maturin develop --release # installs ray_data_arrow_rs into the venv +pytest ../../../../tests/datasource/test_arrow_rs_parquet_reader.py -v +``` + +If the parity tests pass, the reconstruction matches PyArrow. Then run the +benchmark (`release/nightly_tests/dataset/arrow_rs_read_benchmark.py`). + +## API + +``` +read_row_groups(path, row_groups=None, columns=None, batch_size=131072) +read_row_groups_s3(bucket, key, region, anonymous, + row_groups=None, columns=None, batch_size=131072) +``` + +Both return an object implementing `__arrow_c_stream__`, consumed on the Python +side with `pa.RecordBatchReader.from_stream(...)`. + +## Open items for the S3/Linux phase + +- **Expose the knobs** (`decode_budget_bytes`, `k`, `fetch_window_mb`) and port + `build_units` / `read_unit_windowed` / `read_all_async` from `main.rs` so a + single big row group splits into K parallel range reads — the mechanism + `Agents.md` credits for the 4–5× S3 speedup. The current `read_row_groups` + does a single streaming pass (memory win only, no intra-fragment K). +- **Un-gate S3** in `_arrow_rs_supported` once `read_row_groups_s3` is validated. +- **Reconcile FFI/dep versions** on first build (arrow/parquet 59, object_store + 0.13, pyo3 0.22 — adjust to whatever resolves). +- **Avoid double parallelism**: with intra-fragment K in the crate, set + `RAY_DATA_READ_FILES_NUM_THREADS=1` for the arrow-rs path (see the reader). diff --git a/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/pyproject.toml b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/pyproject.toml new file mode 100644 index 000000000000..3afcb043472d --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/pyproject.toml @@ -0,0 +1,15 @@ +[build-system] +requires = ["maturin>=1.5,<2.0"] +build-backend = "maturin" + +[project] +name = "ray_data_arrow_rs" +version = "0.1.0" +description = "Experimental arrow-rs Parquet reader for Ray Data (prototype)" +requires-python = ">=3.9" +classifiers = ["Programming Language :: Rust"] + +[tool.maturin] +# Build a plain extension module (no separate Python package dir). +module-name = "ray_data_arrow_rs.ray_data_arrow_rs" +features = ["pyo3/extension-module"] diff --git a/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/src/lib.rs b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/src/lib.rs new file mode 100644 index 000000000000..8cf9f41be49d --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/src/lib.rs @@ -0,0 +1,3276 @@ +//! Experimental arrow-rs Parquet reader for Ray Data (PyO3 extension). +//! +//! RECONSTRUCTION STATUS +//! --------------------- +//! Reconstructed from the surviving standalone benchmark (`main.rs`). The local +//! path (`read_row_groups`) ports two of `main.rs`'s modes: +//! +//! 1. **Byte-budgeted per-group streaming** (`row_group_loop_bb`, threads==1): +//! read one row group at a time with a batch size computed *by bytes* from the +//! footer (`byte_budget_rows`), so the decoded working set stays flat across +//! schemas (wide-string groups get few rows/batch, numeric groups many). A +//! single reader streams each group in file order and drops each batch, so peak +//! memory ~= one budget while row order is preserved. +//! +//! 2. **Intra-fragment K-split** (`build_units` / `read_range_fixed`): when a call +//! covers a *single* row group larger than `split_threshold_bytes`, split its +//! rows into K contiguous ranges decoded by K threads and merge them back in +//! range order (`ParallelRangeReader`). This is the case Ray can't parallelize +//! (a big row group is a lone fragment → thread pool of 1), so PyArrow decodes +//! it ~single-threaded; K gives us parallel decode without regressing speed. +//! Multi-row-group / small-row-group calls use path 1 (K=1) because Ray's +//! fragment thread pool already parallelizes those — so the two parallelism +//! layers never multiply. +//! +//! Unlike `main.rs` (which only sums a commutative checksum, so range order is +//! irrelevant), we return real data, so the K-split merge is strictly order +//! preserving: one bounded channel per range, drained in range order. +//! +//! The S3 path (`read_row_groups_s3`) ports `main.rs`'s windowed-async reader +//! (`read_all_async` / `read_unit_windowed`) but tuned **memory-first**: peak RSS +//! is `≈ (fetch window compressed) + (decode budget)`, both knobs, flat regardless +//! of row-group size — not `main.rs`'s always-K-way fan-out (which multiplies +//! in-flight memory by K for speed). We fan out to K concurrent GET streams ONLY +//! for a lone row group above `split_threshold_bytes` (the case Ray's fragment +//! pool can't parallelize) — exactly mirroring the local K-split rule — so crate-K +//! and Ray's 4-thread pool never multiply. Every other layout is a single windowed +//! stream (K=1). Output is order-preserving (per-unit channels drained in order), +//! and the decode batch is byte-budgeted just like the local path. +//! +//! Within a single stream, every read shape decomposes into prefetchable *units* +//! — row windows (`~fetch_window_mb` compressed of all projected columns) for +//! ordinary groups, column groups (`~column_fetch_mb` compressed of some columns, +//! all rows) for wide ones — and all units flow through ONE mechanism (`drive_s3`): +//! a byte-denominated semaphore ("the bucket", `prefetch_budget_mb`) admits ranged +//! GETs concurrently until the bucket is full, while a single decoder consumes the +//! prefetched bytes strictly in order; a decoded unit dropping its bytes releases +//! its permits and wakes the next fetch. Fetch concurrency thus self-adjusts to +//! the fetch:decode speed ratio and peak prefetch memory is the bucket size by +//! construction. This is orthogonal to K: K adds parallel streams (spatial +//! split), each with its own bucket (`≈ k * prefetch_budget` compressed in +//! flight). +//! +//! Public API (consumed by `ArrowRsParquetFileReader` on the Python side via +//! `pa.RecordBatchReader.from_stream(...)`): +//! +//! read_row_groups(path, row_groups=None, columns=None, batch_size=131072, +//! decode_budget_bytes=2*1024*1024, k=1, +//! split_threshold_bytes=128*1024*1024) +//! read_row_groups_s3(bucket, key, region, anonymous, endpoint=None, ...creds..., +//! row_groups=None, columns=None, batch_size=131072, +//! decode_budget_bytes=2*1024*1024, fetch_window_mb=16, k=1, +//! split_threshold_bytes=128*1024*1024, predicate_json=None, +//! column_fetch_mb=16, prefetch_budget_mb=64) +//! +//! Per-file handles (TODO 1r — open once, decode many; see the "Per-file native +//! handles" section for why): +//! +//! connect_s3(bucket, region, anonymous, ...creds...) -> NativeS3Store +//! NativeS3Store.open_file(key, page_index=True) -> NativeParquetFile +//! open_parquet_file(path, page_index=False) -> NativeParquetFile +//! NativeParquetFile.metadata() -> ParquetFileMetadata (no I/O) +//! NativeParquetFile.read_row_groups(...) -> Arrow C-stream + +mod predicate; + +use std::collections::{HashMap, HashSet}; +use std::fs::File; +use std::sync::mpsc::{sync_channel, Receiver}; +use std::thread; + +use crate::predicate::{can_match, ColStats, Pred, Value}; +use parquet::basic::{ConvertedType, Encoding, LogicalType, Type as PhysicalType}; +use parquet::file::metadata::RowGroupMetaData; +use parquet::file::statistics::Statistics; +use parquet::schema::types::ColumnDescriptor; + +use arrow::array::{ArrayRef, RecordBatch}; +use arrow::datatypes::{Schema, SchemaRef}; +use arrow::error::ArrowError; +use arrow::ffi::FFI_ArrowSchema; +use arrow::ffi_stream::FFI_ArrowArrayStream; +use arrow::record_batch::RecordBatchReader; +use parquet::arrow::arrow_reader::{ + ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader, + ParquetRecordBatchReaderBuilder, RowSelection, RowSelector, +}; +use parquet::arrow::async_reader::{ + AsyncFileReader, ParquetObjectReader, ParquetRecordBatchStreamBuilder, +}; +use parquet::arrow::ProjectionMask; +use parquet::errors::ParquetError; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; +use std::ffi::CString; + +use bytes::Bytes; +use futures::future::BoxFuture; +use futures::StreamExt; +use object_store::aws::AmazonS3Builder; +use object_store::path::Path as ObjPath; +use object_store::ObjectStore; +use std::ops::Range; +use std::sync::{Arc, OnceLock}; +use tokio::sync::{mpsc, Semaphore}; + +// NOTE on allocators: earlier prototypes carried optional mimalloc/jemalloc +// global-allocator features to chase a suspected allocator-retention gap vs +// PyArrow. Measurement killed the theory (jemalloc LD_PRELOAD inert, per-worker +// high-water lower than PyArrow's on the same fixtures), and mimalloc as a +// cdylib global allocator segfaulted Ray workers across the Arrow C-stream FFI +// boundary. Both features were removed to keep the dependency tree minimal; +// the system allocator is correct here. A/B experiments can still use +// LD_PRELOAD without recompiling. + +// --------------------------------------------------------------------------- // +// Shared tokio runtime +// --------------------------------------------------------------------------- // +/// One process-wide multi-thread runtime, lazily built, shared by every +/// `read_row_groups_s3` call. Previously each fragment read built and tore down +/// its own 2-thread runtime — churn that scales with the file count. The async +/// work is IO-bound (awaiting range GETs), so a small fixed worker pool drives +/// many concurrent fetches. `worker_threads(4)` matches Ray's per-worker fragment +/// pool so decode never oversubscribes cores: either 4 fragments × K=1 unit, or +/// 1 lone fragment × K units — never both at once. +fn shared_runtime() -> &'static tokio::runtime::Runtime { + static RT: OnceLock = OnceLock::new(); + RT.get_or_init(|| { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .enable_all() + .build() + .expect("build shared tokio runtime") + }) +} + +// --------------------------------------------------------------------------- // +// Byte-budget batch sizing (ported from main.rs `byte_budget_rows`) +// --------------------------------------------------------------------------- // +/// Absolute floor on batch rows: protects against degenerate few-row batches +/// (per-batch allocation overhead) without overriding the byte budget. The old +/// floor of 2048 silently voided `decode_budget_bytes` for any schema over +/// ~16 KiB/row (2048 × 16 KiB = the 32 MiB default budget — findings K8): a +/// 1 MiB/row fat-string group decoded 2 GiB batches no matter what the knob +/// said. At 32 the floor only overrides the budget above `budget/32` per row +/// (1 MiB/row at the default), where a 32-row batch is ~one budget anyway. +const MIN_BATCH_ROWS: usize = 32; + +/// Choose a batch row count so `rows * bytes_per_row ~= budget_bytes`, using the +/// row group's uncompressed size / row count from the footer. `requested` is the +/// upper clamp (a narrow schema never grows past the caller's ask) and +/// [`MIN_BATCH_ROWS`] the lower clamp. This is what keeps the decoded working +/// set flat across schemas. +fn byte_budget_rows( + uncompressed_bytes: i64, + num_rows: i64, + requested: usize, + budget_bytes: u64, +) -> usize { + if num_rows <= 0 { + return requested; + } + let bpr = (uncompressed_bytes as f64 / num_rows as f64).max(1.0); + let budget_rows = (budget_bytes as f64 / bpr) as usize; + budget_rows.clamp(MIN_BATCH_ROWS, requested.max(MIN_BATCH_ROWS)) +} + +/// Estimate the row group's DECODED (in-memory Arrow) size from the footer +/// alone. The footer's `total_byte_size` is the *encoded*-uncompressed size: +/// for dictionary/RLE-encoded columns that is dict values + indices, and the +/// decoded batch is larger by the expansion ratio — sizing batches by it made +/// decoded batches ≈ budget × expansion (findings M43/M45/M49: 7.3× overshoot +/// on 9.4×-dictionary tensor data, 4× on RLE-heavy numerics, on BOTH readers' +/// estimators). +/// +/// For fixed-width physical types the decoded size is exactly +/// `num_values × width`, no decoding needed, so take +/// `max(encoded, num_values × width)` per column chunk: +/// - dict/RLE fixed-width (the entire measured loss family — float tensors, +/// repeated ints): the fixed-width term wins → exact; +/// - PLAIN fixed-width: encoded ≥ decoded (rep/def levels) → unchanged; +/// - BYTE_ARRAY (strings/binary): no footer-exact decoded size (a dict +/// chunk's value lengths aren't recorded), width 0 → encoded fallback, +/// i.e. exactly today's behavior; +/// - BOOLEAN: bit-packed on both sides → encoded fallback. +/// +/// The estimate can only grow, so batches can only shrink vs. the old sizing — +/// strictly safer against the decode budget. +fn decoded_estimate_bytes(rgm: &RowGroupMetaData) -> i64 { + let mut total: i64 = 0; + for col in rgm.columns() { + let descr = col.column_descr(); + let width: i64 = match descr.physical_type() { + PhysicalType::INT32 | PhysicalType::FLOAT => 4, + PhysicalType::INT64 | PhysicalType::DOUBLE => 8, + PhysicalType::INT96 => 12, + PhysicalType::FIXED_LEN_BYTE_ARRAY => descr.type_length() as i64, + PhysicalType::BOOLEAN | PhysicalType::BYTE_ARRAY => 0, + }; + let mut decoded = col.num_values().saturating_mul(width); + // Nested leaves also materialize one i32 offset buffer per repetition + // level (~4 B/row/level at the outermost lists; inner fan-out is + // ignored, keeping this a floor). Not noise: on a 5000-list-column + // schema the offsets alone are ~20 KB/row, and omitting them left + // measured batches at 1.5x the budget — exactly on the G1 gate line. + if descr.max_rep_level() > 0 { + decoded = + decoded.saturating_add(4 * rgm.num_rows().max(0) * descr.max_rep_level() as i64); + } + total = total.saturating_add(col.uncompressed_size().max(decoded)); + } + total +} + +/// [`byte_budget_rows`] with the decoded-aware estimate for a whole row group: +/// the one batch-sizing entry point for every read path (local sequential, +/// K-split ranges, S3 windowed units). +fn group_batch_rows(rgm: &RowGroupMetaData, requested: usize, budget_bytes: u64) -> usize { + byte_budget_rows( + decoded_estimate_bytes(rgm), + rgm.num_rows(), + requested, + budget_bytes, + ) +} + +/// A row group is *estimator-blind* when its decoded size cannot be derived +/// from the footer: dictionary-encoded BYTE_ARRAY (strings/binary) chunks, +/// where [`decoded_estimate_bytes`] falls back to the encoded size and batches +/// sized from it overshoot the decode budget by the dictionary expansion +/// ratio (M50 residual: 4x on dict-string lone_big_rg — the one measured +/// shape the footer-exact estimator can't cover). Checked against ALL columns, +/// not the projection, mirroring `decoded_estimate_bytes`. +fn group_is_estimator_blind(rgm: &RowGroupMetaData) -> bool { + rgm.columns().iter().any(|c| { + c.column_descr().physical_type() == PhysicalType::BYTE_ARRAY + && c + .encodings() + .any(|e| matches!(e, Encoding::PLAIN_DICTIONARY | Encoding::RLE_DICTIONARY)) + }) +} + +/// Running decoded-size measurement for mid-stream batch adaptation: +/// cumulative (bytes, rows) over every batch yielded from blind groups, so +/// `bytes_per_row` amortizes per-batch buffer-capacity rounding (a lone +/// [`MIN_BATCH_ROWS`]-row probe can overcount from allocator doubling; the +/// cumulative average self-corrects as full-size batches arrive). +#[derive(Default, Clone, Copy)] +struct BprTracker { + bytes: u64, + rows: u64, +} + +impl BprTracker { + fn record(&mut self, batch: &RecordBatch) { + self.bytes = self + .bytes + .saturating_add(batch.get_array_memory_size() as u64); + self.rows = self.rows.saturating_add(batch.num_rows() as u64); + } + fn bytes_per_row(&self) -> Option { + (self.rows > 0).then(|| self.bytes as f64 / self.rows as f64) + } +} + +/// Re-size a batch row count from the MEASURED decoded bytes/row. Shrink-only: +/// the static estimate stays the upper clamp, so a non-expanding shape (or a +/// noisy low measurement) can never produce batches bigger than today's, and +/// [`MIN_BATCH_ROWS`] stays the floor. `None` (nothing measured yet) keeps the +/// static size. +fn adapted_rows(static_rows: usize, budget_bytes: u64, measured_bpr: Option) -> usize { + match measured_bpr { + Some(bpr) if bpr > 0.0 => ((budget_bytes as f64 / bpr) as usize) + .clamp(MIN_BATCH_ROWS, static_rows.max(MIN_BATCH_ROWS)), + _ => static_rows, + } +} + +// --------------------------------------------------------------------------- // +// Row-group statistics pruning (predicate pushdown, part 1) +// --------------------------------------------------------------------------- // +/// Whether an INT32/INT64 physical column is a *plain signed integer*, i.e. the +/// raw statistic value can be compared directly as a signed `i64`. Everything +/// else stored in an INT32/INT64 physical slot is rejected: +/// - UINT_8/16/32/64: Parquet orders these by *unsigned* comparison, so a +/// value ≥ 2^(bits-1) is stored with its high bit set and reads back as a +/// negative `i64` — min/max can invert (e.g. u32 max reads as -1). Comparing +/// as signed would prune row groups that actually match. +/// - DECIMAL: the stat is the *unscaled* integer, which doesn't match the +/// decimal literal a predicate carries. +/// - DATE / TIME / TIMESTAMP: the encoded units may not match the predicate +/// literal's encoding. +/// For any of these we return no bound → `can_match` keeps the row group. +fn is_plain_signed_int(descr: &ColumnDescriptor) -> bool { + match descr.logical_type_ref() { + Some(LogicalType::Integer(int)) => int.is_signed, + // Decimal / Date / Time / Timestamp / etc. on an int physical type. + Some(_) => false, + // Legacy files predate LogicalType; fall back to the converted type. + None => matches!( + descr.converted_type(), + ConvertedType::NONE + | ConvertedType::INT_8 + | ConvertedType::INT_16 + | ConvertedType::INT_32 + | ConvertedType::INT_64 + ), + } +} + +/// Map one column chunk's Parquet statistics to `(min, max)` in `Value` terms. +/// Types we can't soundly order for pruning — Int96, fixed-len byte arrays, +/// non-UTF8 byte arrays, and INT32/INT64 columns that aren't plain signed +/// integers (unsigned / decimal / date-time; see [`is_plain_signed_int`]) — +/// return `None`, which `can_match` treats as "keep". `descr` supplies the +/// logical type that `Statistics` (keyed by physical type) can't. +fn stat_min_max(stats: &Statistics, descr: &ColumnDescriptor) -> (Option, Option) { + match stats { + Statistics::Boolean(v) => ( + v.min_opt().map(|b| Value::Bool(*b)), + v.max_opt().map(|b| Value::Bool(*b)), + ), + Statistics::Int32(v) if is_plain_signed_int(descr) => ( + v.min_opt().map(|x| Value::Int(*x as i64)), + v.max_opt().map(|x| Value::Int(*x as i64)), + ), + Statistics::Int64(v) if is_plain_signed_int(descr) => ( + v.min_opt().map(|x| Value::Int(*x)), + v.max_opt().map(|x| Value::Int(*x)), + ), + Statistics::Float(v) => ( + v.min_opt().map(|x| Value::Float(*x as f64)), + v.max_opt().map(|x| Value::Float(*x as f64)), + ), + Statistics::Double(v) => ( + v.min_opt().map(|x| Value::Float(*x)), + v.max_opt().map(|x| Value::Float(*x)), + ), + Statistics::ByteArray(v) => ( + v.min_opt() + .and_then(|b| b.as_utf8().ok().map(|s| Value::Str(s.to_string()))), + v.max_opt() + .and_then(|b| b.as_utf8().ok().map(|s| Value::Str(s.to_string()))), + ), + // Int96, FixedLenByteArray, and non-signed-int INT32/INT64 columns + // (guards above fell through): not ordered here → keep. + _ => (None, None), + } +} + +/// Build the per-column statistics map for one row group. Keyed by the leaf +/// column path string, which equals the field name for the flat columns +/// predicates push on; nested paths (dotted) simply won't match a top-level +/// predicate column and are left un-pruned (conservative). +fn row_group_col_stats(rg: &RowGroupMetaData) -> HashMap { + let num_rows = rg.num_rows(); + let mut map = HashMap::with_capacity(rg.num_columns()); + for i in 0..rg.num_columns() { + let col = rg.column(i); + if let Some(stats) = col.statistics() { + let (min, max) = stat_min_max(stats, col.column_descr()); + let null_count = stats.null_count_opt().map(|n| n as i64); + map.insert( + col.column_path().string(), + ColStats { + min, + max, + null_count, + num_rows, + }, + ); + } + } + map +} + +/// Drop the row groups in `selected` that `pred` proves cannot contain a match. +/// Conservative by construction (see `predicate::can_match`): a row group is +/// removed only when provably empty for the predicate, so this never drops a +/// group that could have contributed rows. +fn prune_row_groups(meta: &ArrowReaderMetadata, selected: Vec, pred: &Pred) -> Vec { + let md = meta.metadata(); + selected + .into_iter() + .filter(|&rg| can_match(pred, &row_group_col_stats(md.row_group(rg)))) + .collect() +} + +/// Parse the optional predicate IR and prune `selected` in one step; `None` +/// (no pushdown) returns `selected` unchanged. +fn apply_predicate( + meta: &ArrowReaderMetadata, + selected: Vec, + predicate_json: &Option, +) -> Vec { + match predicate_json { + None => selected, + Some(j) => prune_row_groups(meta, selected, &Pred::from_json(j)), + } +} + +// --------------------------------------------------------------------------- // +// Column projection helper +// --------------------------------------------------------------------------- // +/// Build a leaf-column ProjectionMask from column names using the parquet schema +/// descriptor. Names not present are ignored (Python already resolved the read +/// set). Flat schemas only — the Python `_arrow_rs_supported` gate rejects nested +/// columns before we get here. +fn projection_mask( + parquet_schema: &parquet::schema::types::SchemaDescriptor, + columns: &Option>, +) -> ProjectionMask { + match columns { + None => ProjectionMask::all(), + Some(names) => { + let root = parquet_schema.root_schema(); + let mut indices = Vec::new(); + for (i, f) in root.get_fields().iter().enumerate() { + if names.iter().any(|n| n == f.name()) { + indices.push(i); + } + } + ProjectionMask::roots(parquet_schema, indices) + } + } +} + +/// Ordered ROOT (top-level field) indices for a projection. Mirrors +/// `projection_mask`'s name matching so the two always agree on which columns +/// are read. Note these are root indices, not leaf/column-chunk indices — use +/// [`leaves_under_roots`] to expand to the chunks a root projection touches +/// (identity for flat schemas, several leaves per root for structs/lists). +fn projected_root_indices( + parquet_schema: &parquet::schema::types::SchemaDescriptor, + columns: &Option>, +) -> Vec { + let root = parquet_schema.root_schema(); + match columns { + None => (0..root.get_fields().len()).collect(), + Some(names) => root + .get_fields() + .iter() + .enumerate() + .filter(|(_, f)| names.iter().any(|n| n == f.name())) + .map(|(i, _)| i) + .collect(), + } +} + +/// All leaf-column (column-chunk) indices under the given root fields, in +/// ascending leaf order — exactly the chunks the decoder requests for a root +/// projection. Flat schemas: identity. +fn leaves_under_roots( + schema: &parquet::schema::types::SchemaDescriptor, + roots: &[usize], +) -> Vec { + (0..schema.num_columns()) + .filter(|&l| roots.contains(&schema.get_column_root_idx(l))) + .collect() +} + +/// `(root index, compressed size)` for the projected top-level columns of a row +/// group, in ascending root order — the input to `partition_columns_by_budget`. +/// A root's size is the sum of its leaf chunks, so a struct column is weighed +/// (and later fetched/hstacked) as one indivisible unit. +fn projected_root_sizes( + schema: &parquet::schema::types::SchemaDescriptor, + rgm: &RowGroupMetaData, + roots: &[usize], +) -> Vec<(usize, u64)> { + roots + .iter() + .map(|&r| { + let sz = leaves_under_roots(schema, &[r]) + .iter() + .map(|&l| rgm.column(l).compressed_size().max(0) as u64) + .sum(); + (r, sz) + }) + .collect() +} + +/// Partition projected columns into contiguous groups whose per-group compressed +/// size stays under `budget_bytes`, so the S3 reader can fetch+decode ONE group at +/// a time and hold only that group's compressed chunks resident — the wide-schema +/// memory fix (the async reader's default `InMemoryRowGroup` otherwise fetches every +/// projected column chunk for the row group up front). `cols` is `(leaf, size)` in +/// ascending leaf order and groups preserve that order, so hstacking the groups +/// reproduces file/schema column order. `budget_bytes == 0` (or ≤1 column) => a +/// single group (disabled). A lone oversized column still gets its own group — a +/// column can't be split below itself (that would be row-windowing, handled +/// elsewhere). +fn partition_columns_by_budget(cols: &[(usize, u64)], budget_bytes: u64) -> Vec> { + if budget_bytes == 0 || cols.len() <= 1 { + return vec![cols.iter().map(|(i, _)| *i).collect()]; + } + let mut groups: Vec> = Vec::new(); + let mut cur: Vec = Vec::new(); + let mut acc: u64 = 0; + for &(idx, sz) in cols { + // Start a new group when the current one is non-empty and adding this column + // would exceed the budget. A single column always fits (never split below 1). + if !cur.is_empty() && acc.saturating_add(sz) > budget_bytes { + groups.push(std::mem::take(&mut cur)); + acc = 0; + } + cur.push(idx); + acc = acc.saturating_add(sz); + } + if !cur.is_empty() { + groups.push(cur); + } + groups +} + +/// Probe the projected output schema with an empty (zero row group) reader, so +/// `schema()` is available to the FFI stream before any batch is pulled. +fn probe_schema( + path: &str, + meta: &ArrowReaderMetadata, + mask: &ProjectionMask, +) -> Result { + Ok( + ParquetRecordBatchReaderBuilder::new_with_metadata(File::open(path)?, meta.clone()) + .with_projection(mask.clone()) + .with_row_groups(vec![]) + .build()? + .schema(), + ) +} + +// --------------------------------------------------------------------------- // +// Arrow C-stream wrapper returned to Python +// --------------------------------------------------------------------------- // +/// Holds an FFI stream until Python pulls it out via `__arrow_c_stream__`. +#[pyclass] +struct ArrowStream { + inner: Option, +} + +#[pymethods] +impl ArrowStream { + /// PyCapsule protocol: PyArrow's `RecordBatchReader.from_stream` calls this. + #[pyo3(signature = (_requested_schema=None))] + fn __arrow_c_stream__<'py>( + &mut self, + py: Python<'py>, + _requested_schema: Option, + ) -> PyResult> { + let stream = self + .inner + .take() + .ok_or_else(|| PyRuntimeError::new_err("stream already consumed"))?; + let name = CString::new("arrow_array_stream").unwrap(); + PyCapsule::new_bound(py, stream, Some(name)) + } +} + +fn into_py_stream(reader: Box) -> ArrowStream { + ArrowStream { + inner: Some(FFI_ArrowArrayStream::new(reader)), + } +} + +// --------------------------------------------------------------------------- // +// Footer metadata returned to Python (Track 1: arrow-rs owns the footer read) +// --------------------------------------------------------------------------- // +/// The parts of the Parquet footer the Python reader needs so PyArrow no longer +/// has to open the file for supported fragments: the full Arrow schema (exposed +/// zero-copy via the Arrow C-schema PyCapsule, so it round-trips extension/field +/// metadata for the UDT path) plus per-row-group row counts and byte sizes (for +/// chunking, row-offset bookkeeping, `count()`, and the split threshold). Column +/// statistics for row-group pruning are a follow-up on this same struct. +#[pyclass] +struct ParquetFileMetadata { + schema: SchemaRef, + #[pyo3(get)] + num_rows: i64, + #[pyo3(get)] + num_row_groups: usize, + #[pyo3(get)] + row_group_num_rows: Vec, + #[pyo3(get)] + row_group_byte_sizes: Vec, + // Per-row-group *compressed* (on-disk) byte size — the sum of each column + // chunk's compressed size. This is what the Python chunker bundles by, so it + // must match PyArrow's `sum(col.total_compressed_size)`. Distinct from + // `row_group_byte_sizes`, which is the *uncompressed* `total_byte_size()`. + #[pyo3(get)] + row_group_compressed_sizes: Vec, + // Root (top-level) column names that contain an INT96-physical leaf. The + // support gate needs this because parquet-rs honors an embedded Arrow-schema + // unit hint for INT96 (→ us/ms/s) whereas PyArrow always forces ns: a column + // in this list whose decoded unit isn't ns would diverge from PyArrow, so the + // gate must fall it back. INT96 columns that come out as ns already match. + #[pyo3(get)] + int96_columns: Vec, + // True when the embedded Arrow schema (`ARROW:schema` footer metadata) could + // not be parsed and was skipped (see `load_meta_local`), so `schema` here is + // the parquet-inferred *storage* schema rather than the Arrow logical schema. + // The Python reader reconstructs any lost extension types (e.g. Ray's + // cloudpickle-serialized tensor type) from the file's own footer schema. + #[pyo3(get)] + arrow_schema_skipped: bool, +} + +#[pymethods] +impl ParquetFileMetadata { + /// Arrow PyCapsule protocol: `pa.schema(obj)` / `Schema._import_from_c_capsule` + /// pull the schema through this. Rebuilt each call (cheap, non-consuming). + fn __arrow_c_schema__<'py>(&self, py: Python<'py>) -> PyResult> { + let ffi = FFI_ArrowSchema::try_from(self.schema.as_ref()).map_err(to_py)?; + let name = CString::new("arrow_schema").unwrap(); + PyCapsule::new_bound(py, ffi, Some(name)) + } +} + +/// Reader options shared by *every* entry point (metadata reads, local decode, +/// S3 decode), so the schema a metadata read reports can never diverge from what +/// a decode produces. `page_index` varies by caller: pruning paths ask for the +/// page index (`Optional`); a bare footer/metadata read skips it (`Skip`). +/// +/// INT96 note: parquet-rs decodes the legacy INT96 timestamp physical type to +/// `Timestamp(Nanosecond, None)` by default (arrow/schema/primitive.rs), which is +/// exactly what PyArrow produces for INT96 by default — so a Spark/Hive/Impala +/// INT96 file (the common producers, which embed no Arrow schema) decodes +/// byte-identically on both paths and takes the native path with no coercion. +/// The one divergence is a file that embeds an Arrow schema pinning a *non-ns* +/// unit (e.g. a PyArrow writer with `use_deprecated_int96_timestamps=True` over a +/// `timestamp[us]` column): parquet-rs honors that embedded hint (→ us) while +/// PyArrow forces ns. That mismatch is caught by the support gate's +/// per-file-vs-unified type check, which falls the file back to PyArrow — correct, +/// if not yet native. parquet 59 has no `with_coerce_int96`, so forcing ns there +/// would need a per-column `with_schema` override; deferred (narrow case, and the +/// fallback is already correct). +fn reader_options(page_index: PageIndexPolicy) -> ArrowReaderOptions { + ArrowReaderOptions::new().with_page_index_policy(page_index) +} + +/// arrow-rs's IPC verifier rejects an embedded `ARROW:schema` footer whose field +/// `custom_metadata` values aren't valid UTF-8. Ray files written by 2.49-2.54 +/// store the tensor extension type's metadata as a cloudpickle blob (binary), so +/// `ArrowReaderMetadata::load` fails parsing that embedded schema +/// (`Unable to get root as message stored in ARROW:schema: Utf8Error`). Detect +/// exactly that failure so the loaders below can retry with the embedded arrow +/// schema skipped — decoding the parquet-inferred storage types instead — after +/// which the Python reader re-applies the extension type from the pinned dataset +/// schema (a `list<..>`→`extension<..>` cast). Any other load error propagates. +fn is_embedded_arrow_schema_error(e: &E) -> bool { + e.to_string().contains("ARROW:schema") +} + +/// Load footer metadata for a local file, retrying with the embedded arrow schema +/// skipped when (and only when) it fails to parse (see +/// [`is_embedded_arrow_schema_error`]). Files whose footer parses normally are +/// untouched, so INT96 hints and valid embedded schemas behave exactly as before. +/// The returned bool is `true` when the retry fired (embedded arrow schema +/// skipped → the reported schema is the parquet-inferred storage type), which the +/// Python reader uses to reconstruct the extension type from the pinned schema. +fn load_meta_local( + file: &File, + page_index: PageIndexPolicy, +) -> Result<(ArrowReaderMetadata, bool), ParquetError> { + match ArrowReaderMetadata::load(file, reader_options(page_index)) { + Err(e) if is_embedded_arrow_schema_error(&e) => ArrowReaderMetadata::load( + file, + reader_options(page_index).with_skip_arrow_metadata(true), + ) + .map(|m| (m, true)), + other => other.map(|m| (m, false)), + } +} + +/// S3 counterpart of [`load_meta_local`]: same targeted retry (and same +/// skipped-bool contract), building a fresh `ParquetObjectReader` for the second +/// attempt so no half-consumed reader state carries over. +async fn load_meta_s3( + store: Arc, + path: ObjPath, + page_index: PageIndexPolicy, +) -> Result<(ArrowReaderMetadata, bool), ParquetError> { + let mut probe = ParquetObjectReader::new(store.clone(), path.clone()); + match ArrowReaderMetadata::load_async(&mut probe, reader_options(page_index)).await { + Err(e) if is_embedded_arrow_schema_error(&e) => { + let mut retry = ParquetObjectReader::new(store, path); + ArrowReaderMetadata::load_async( + &mut retry, + reader_options(page_index).with_skip_arrow_metadata(true), + ) + .await + .map(|m| (m, true)) + } + other => other.map(|m| (m, false)), + } +} + +/// Pull the fields Python needs out of an already-loaded `ArrowReaderMetadata`. +/// Local and S3 both funnel through here so the shape is identical. +fn build_file_metadata( + meta: &ArrowReaderMetadata, + arrow_schema_skipped: bool, +) -> ParquetFileMetadata { + let md = meta.metadata(); + let n = md.num_row_groups(); + let mut row_group_num_rows = Vec::with_capacity(n); + let mut row_group_byte_sizes = Vec::with_capacity(n); + let mut row_group_compressed_sizes = Vec::with_capacity(n); + let mut num_rows = 0i64; + for i in 0..n { + let rg = md.row_group(i); + row_group_num_rows.push(rg.num_rows()); + row_group_byte_sizes.push(rg.total_byte_size()); + row_group_compressed_sizes.push(rg.compressed_size()); + num_rows += rg.num_rows(); + } + + // Collect the root column names backing an INT96 leaf. Walk the flat leaf + // descriptors and key by the first path component so a nested INT96 (e.g. a + // struct field) still surfaces its top-level column to the gate. + let mut int96_roots: HashSet = HashSet::new(); + for col in md.file_metadata().schema_descr().columns() { + if col.physical_type() == PhysicalType::INT96 { + if let Some(root) = col.path().parts().first() { + int96_roots.insert(root.clone()); + } + } + } + + ParquetFileMetadata { + schema: meta.schema().clone(), + num_rows, + num_row_groups: n, + row_group_num_rows, + row_group_byte_sizes, + row_group_compressed_sizes, + int96_columns: int96_roots.into_iter().collect(), + arrow_schema_skipped, + } +} + +// --------------------------------------------------------------------------- // +// Local read (sync): per-group byte-budgeted sequential reader (K=1 path) +// --------------------------------------------------------------------------- // +/// Streams the selected row groups in order, building one `ParquetRecordBatchReader` +/// per group with a byte-budgeted batch size. Row order is preserved (single +/// reader, groups in ascending order) and peak memory stays ~one decode budget +/// because each batch is dropped as Python pulls the next. +struct RowGroupSeqReader { + path: String, + meta: ArrowReaderMetadata, + mask: ProjectionMask, + budget_bytes: u64, + batch_clamp: usize, + row_groups: Vec, + pos: usize, + current: Option, + schema: SchemaRef, + /// Decoded bytes/row observed so far (recorded for blind groups only — + /// see [`group_is_estimator_blind`]); carries across row groups. + bpr: BprTracker, + /// Whether `current` reads a blind group (= record its batches). + cur_blind: bool, + /// Rest-of-group continuation after a probe reader: (rg, rows to skip). + pending: Option<(usize, usize)>, +} + +impl RowGroupSeqReader { + fn new( + path: String, + meta: ArrowReaderMetadata, + mask: ProjectionMask, + schema: SchemaRef, + row_groups: Vec, + batch_clamp: usize, + budget_bytes: u64, + ) -> Self { + Self { + path, + meta, + mask, + budget_bytes, + batch_clamp, + row_groups, + pos: 0, + current: None, + schema, + bpr: BprTracker::default(), + cur_blind: false, + pending: None, + } + } + + /// Build a reader for rows `[skip, skip+take)` of `rg` at `batch_rows`. + /// A partial range uses a `RowSelection`; skipping needs no page index — + /// the parquet reader decode-skips leading rows, and every skip here is at + /// most one probe's worth. + fn build_group_reader( + &self, + rg: usize, + skip: usize, + take: usize, + total: usize, + batch_rows: usize, + ) -> Result { + let mut builder = ParquetRecordBatchReaderBuilder::new_with_metadata( + File::open(&self.path)?, + self.meta.clone(), + ) + .with_batch_size(batch_rows) + .with_row_groups(vec![rg]) + .with_projection(self.mask.clone()); + if skip > 0 || skip + take < total { + builder = builder.with_row_selection(RowSelection::from(vec![ + RowSelector::skip(skip), + RowSelector::select(take), + ])); + } + builder.build() + } +} + +impl Iterator for RowGroupSeqReader { + type Item = Result; + fn next(&mut self) -> Option { + loop { + if let Some(reader) = self.current.as_mut() { + match reader.next() { + Some(Ok(batch)) => { + if self.cur_blind { + self.bpr.record(&batch); + } + return Some(Ok(batch)); + } + Some(Err(e)) => return Some(Err(e)), + None => self.current = None, + } + } + // Next reader: the rest of a probed group, else the next group. + let (rg, skip) = match self.pending.take() { + Some(cont) => cont, + None => { + if self.pos >= self.row_groups.len() { + return None; + } + let rg = self.row_groups[self.pos]; + self.pos += 1; + (rg, 0) + } + }; + let (blind, static_eff, total) = { + let rgm = self.meta.metadata().row_group(rg); + ( + group_is_estimator_blind(rgm), + group_batch_rows(rgm, self.batch_clamp, self.budget_bytes), + rgm.num_rows().max(0) as usize, + ) + }; + self.cur_blind = blind; + let remaining = total.saturating_sub(skip); + if remaining == 0 { + continue; + } + let (batch_rows, take) = if !blind { + (static_eff, remaining) + } else if self.bpr.bytes_per_row().is_some() { + ( + adapted_rows(static_eff, self.budget_bytes, self.bpr.bytes_per_row()), + remaining, + ) + } else if remaining > MIN_BATCH_ROWS { + // First blind group, nothing measured yet: open with a tiny + // probe reader (bounded by the same argument that sets + // MIN_BATCH_ROWS), then continue the group adapted. + self.pending = Some((rg, skip + MIN_BATCH_ROWS)); + (MIN_BATCH_ROWS, MIN_BATCH_ROWS) + } else { + // Group no bigger than a probe: just read it; its batches + // still feed the tracker for the groups after it. + (static_eff, remaining) + }; + match self.build_group_reader(rg, skip, take, total, batch_rows) { + Ok(reader) => self.current = Some(reader), + Err(e) => return Some(Err(ArrowError::ExternalError(Box::new(e)))), + } + } + } +} + +impl RecordBatchReader for RowGroupSeqReader { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +// --------------------------------------------------------------------------- // +// Local read (sync): intra-fragment K-split for one big row group +// --------------------------------------------------------------------------- // +/// Splits one row group's rows into K contiguous ranges decoded by K threads, +/// merging them back in range order so output row order matches a sequential read. +/// Each range has its own bounded channel (backpressure), and the consumer drains +/// channels in ascending range order — so at most `k * channel_depth` batches are +/// resident and rows come out in file order. Requires the offset/page index so a +/// `RowSelection` fetches only its range's pages (else each worker would decode the +/// whole column chunk); the caller checks this before choosing this path. +struct ParallelRangeReader { + schema: SchemaRef, + receivers: Vec>>, + cur: usize, +} + +fn build_range_reader( + path: &str, + meta: &ArrowReaderMetadata, + mask: &ProjectionMask, + rg: usize, + start: usize, + len: usize, + batch: usize, +) -> Result { + let sel = RowSelection::from(vec![RowSelector::skip(start), RowSelector::select(len)]); + ParquetRecordBatchReaderBuilder::new_with_metadata(File::open(path)?, meta.clone()) + .with_row_groups(vec![rg]) + .with_row_selection(sel) + .with_batch_size(batch) + .with_projection(mask.clone()) + .build() +} + +impl ParallelRangeReader { + fn spawn( + path: String, + meta: ArrowReaderMetadata, + mask: ProjectionMask, + schema: SchemaRef, + rg: usize, + total_rows: usize, + k: usize, + batch: usize, + ) -> Self { + let chunk = total_rows.div_ceil(k.max(1)).max(1); + let mut receivers = Vec::new(); + let mut start = 0usize; + while start < total_rows { + let len = chunk.min(total_rows - start); + // Depth 2: a worker may run one batch ahead of the consumer, no more. + let (tx, rx) = sync_channel::>(2); + receivers.push(rx); + let (path, meta, mask) = (path.clone(), meta.clone(), mask.clone()); + thread::spawn(move || { + match build_range_reader(&path, &meta, &mask, rg, start, len, batch) { + Ok(reader) => { + for batch in reader { + if tx.send(batch).is_err() { + break; // consumer dropped + } + } + } + Err(e) => { + let _ = tx.send(Err(ArrowError::ExternalError(Box::new(e)))); + } + } + }); + start += len; + } + Self { + schema, + receivers, + cur: 0, + } + } +} + +impl Iterator for ParallelRangeReader { + type Item = Result; + fn next(&mut self) -> Option { + while self.cur < self.receivers.len() { + match self.receivers[self.cur].recv() { + Ok(item) => return Some(item), + Err(_) => self.cur += 1, // this range's channel closed → next range + } + } + None + } +} + +impl RecordBatchReader for ParallelRangeReader { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +// --------------------------------------------------------------------------- // +// Local entry point: choose sequential vs K-split +// --------------------------------------------------------------------------- // +#[allow(clippy::too_many_arguments)] +fn open_local_reader( + path: String, + row_groups: Option>, + columns: Option>, + batch_size: usize, + budget_bytes: u64, + k: usize, + split_threshold_bytes: u64, + predicate_json: Option, +) -> Result, ParquetError> { + // Lean footer parse (#6): the page index is only needed for the K-split + // RowSelection (to skip pages by byte range). K-split can only fire when k > 1, + // so for the common k == 1 local path we Skip the page index entirely — a + // cheaper footer parse that matters on many-row-group files. When k > 1 we load + // it Optional so the lone-big-row-group split can use it if present. + let policy = if k > 1 { + PageIndexPolicy::Optional + } else { + PageIndexPolicy::Skip + }; + let (meta, _skipped) = load_meta_local(&File::open(&path)?, policy)?; + build_local_reader( + path, + meta, + row_groups, + columns, + batch_size, + budget_bytes, + k, + split_threshold_bytes, + predicate_json, + ) +} + +/// The metadata-independent half of the local read: everything after the footer +/// load. Shared by [`open_local_reader`] (loads the footer per call — the +/// original API) and [`NativeParquetFile::read_row_groups`] (footer loaded once +/// at open, reused across calls — TODO 1r). Whether the K-split can fire depends +/// on the page-index policy the *caller* loaded `meta` under, exactly as before. +#[allow(clippy::too_many_arguments)] +fn build_local_reader( + path: String, + meta: ArrowReaderMetadata, + row_groups: Option>, + columns: Option>, + batch_size: usize, + budget_bytes: u64, + k: usize, + split_threshold_bytes: u64, + predicate_json: Option, +) -> Result, ParquetError> { + let mask = projection_mask(meta.metadata().file_metadata().schema_descr(), &columns); + let selected: Vec = match row_groups { + Some(v) => v, + None => (0..meta.metadata().num_row_groups()).collect(), + }; + // Statistics-based row-group pruning (conservative — see predicate.rs). This + // is the mechanism that replaces PyArrow's `fragment.subset(filter=...)` so + // pruned groups are never fetched or decoded. + let selected = apply_predicate(&meta, selected, &predicate_json); + let schema = probe_schema(&path, &meta, &mask)?; + + // K-split only for a *single* row group above the threshold, and only when the + // page index is present (else each range would decode the whole column chunk). + // This is exactly the lone-big-fragment case Ray's pool can't parallelize; every + // other layout uses the sequential path so crate-K and Ray's pool never multiply. + let split = k > 1 + && selected.len() == 1 + && meta.metadata().row_group(selected[0]).total_byte_size() as u64 >= split_threshold_bytes + && meta.metadata().offset_index().is_some(); + + if split { + let rg = selected[0]; + let (total_rows, static_eff, blind) = { + let rgm = meta.metadata().row_group(rg); + ( + rgm.num_rows().max(0) as usize, + group_batch_rows(rgm, batch_size, budget_bytes), + group_is_estimator_blind(rgm), + ) + }; + let mut eff = static_eff; + if blind && total_rows > MIN_BATCH_ROWS { + // K range readers can't re-size mid-range, so measure BEFORE the + // fan-out: a measure-only decode of the group's first + // MIN_BATCH_ROWS rows (range 0 decodes them again — 32 rows, + // negligible) gives the real decoded bytes/row. + let mut tracker = BprTracker::default(); + let probe = build_range_reader( + &path, + &meta, + &mask, + rg, + 0, + MIN_BATCH_ROWS, + MIN_BATCH_ROWS, + )?; + for batch in probe { + let batch = batch.map_err(|e| ParquetError::External(Box::new(e)))?; + tracker.record(&batch); + } + eff = adapted_rows(static_eff, budget_bytes, tracker.bytes_per_row()); + } + Ok(Box::new(ParallelRangeReader::spawn( + path, meta, mask, schema, rg, total_rows, k, eff, + ))) + } else { + Ok(Box::new(RowGroupSeqReader::new( + path, + meta, + mask, + schema, + selected, + batch_size, + budget_bytes, + ))) + } +} + +#[pyfunction] +#[pyo3(signature = (path, row_groups=None, columns=None, batch_size=131072, decode_budget_bytes=2*1024*1024, k=1, split_threshold_bytes=134217728, predicate_json=None))] +#[allow(clippy::too_many_arguments)] +fn read_row_groups( + py: Python<'_>, + path: String, + row_groups: Option>, + columns: Option>, + batch_size: usize, + decode_budget_bytes: u64, + k: usize, + split_threshold_bytes: u64, + // Optional predicate IR (JSON, built from the Ray `Expr`) for statistics + // row-group pruning. None = no pushdown. Row-level filtering still happens + // in Python post-decode, so this only avoids IO/decode, never changes rows. + predicate_json: Option, +) -> PyResult { + // Footer read + reader construction is blocking file I/O; release the GIL so + // sibling Python read threads (Ray's fragment pool) run in parallel. + let reader = py + .allow_threads(|| { + open_local_reader( + path, + row_groups, + columns, + batch_size, + decode_budget_bytes, + k, + split_threshold_bytes, + predicate_json, + ) + }) + .map_err(to_py)?; + Ok(into_py_stream(reader)) +} + +// --------------------------------------------------------------------------- // +// S3 read (async, windowed, byte-budgeted, order-preserving) +// --------------------------------------------------------------------------- // +/// Number of decoded batches a unit task may run ahead of the consumer. Depth 2 +/// bounds resident memory while still letting a task fetch/decode one batch ahead. +const S3_CHANNEL_DEPTH: usize = 2; + +/// A sync `RecordBatchReader` fed by K background tokio tasks (one per row-range +/// unit), each draining its unit into a bounded async channel. The consumer drains +/// channels in ascending unit order — so at most `k * S3_CHANNEL_DEPTH` batches are +/// resident and rows come out in file order (K units are contiguous ascending +/// ranges). `blocking_recv` is called from the Python thread (outside the runtime), +/// which is exactly what tokio's mpsc supports. +struct S3ChannelReader { + schema: SchemaRef, + receivers: Vec>>, + cur: usize, +} + +impl Iterator for S3ChannelReader { + type Item = Result; + fn next(&mut self) -> Option { + while self.cur < self.receivers.len() { + match self.receivers[self.cur].blocking_recv() { + Some(item) => return Some(item), + None => self.cur += 1, // this unit's channel closed → next unit + } + } + None + } +} + +impl RecordBatchReader for S3ChannelReader { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +/// Rows per fetch window from a byte budget over the row group's *compressed* +/// bytes/row — this bounds IN-FLIGHT NETWORK bytes (what we fetch before decode). +/// 0 means "whole range in one shot" (no window cap). +fn window_rows_for(rgm: &parquet::file::metadata::RowGroupMetaData, fetch_window_mb: u64) -> usize { + if fetch_window_mb == 0 { + return 0; + } + let comp = rgm.compressed_size().max(1) as f64; + let rows = rgm.num_rows().max(1) as f64; + let comp_bpr = (comp / rows).max(1.0); + (((fetch_window_mb as f64) * 1024.0 * 1024.0) / comp_bpr).max(1.0) as usize +} + +/// Largest single data page, in rows, across the projected columns of row group +/// `rg` (from the offset index). This is the floor a fetch window must not go +/// below: a `RowSelection` can only skip whole *pages*, never rows within a page, +/// so a window narrower than a column's page forces every window overlapping that +/// page to decode the whole page — re-decoding the row group once per window (the +/// wide/short-row-group pathology: few rows → one page/column → N-way re-decode). +/// +/// Returns the row group's total row count when there is no offset index (windowing +/// can't skip pages without it anyway) or the index is empty — either way the +/// caller then collapses to a single window (no split), which is correct. +fn max_page_rows(md: &parquet::file::metadata::ParquetMetaData, rg: usize) -> usize { + let num_rows = md.row_group(rg).num_rows().max(0); + let rg_oi = match md.offset_index().and_then(|oi| oi.get(rg)) { + Some(rg_oi) if !rg_oi.is_empty() => rg_oi, + _ => return num_rows as usize, + }; + let mut max_rows: i64 = 0; + for col in rg_oi { + let locs = col.page_locations(); + for (i, loc) in locs.iter().enumerate() { + let end = locs + .get(i + 1) + .map(|next| next.first_row_index) + .unwrap_or(num_rows); + max_rows = max_rows.max(end - loc.first_row_index); + } + } + if max_rows <= 0 { + num_rows as usize + } else { + max_rows as usize + } +} + +/// The row step a fetch window should advance by, given the byte-budget +/// `window_rows` (0 = no cap) and the coarsest column's `max_page_rows`. Clamping +/// the window up to at least one full page means each page is decoded by ~one +/// window instead of every overlapping window; a group whose largest page spans +/// the whole range (wide/short) collapses to `len` (a single window == parity with +/// no windowing), while a tall multi-page group still splits and keeps the +/// bounded-working-set win. Pure so it is unit-tested without a Parquet fixture. +fn effective_window_step(window_rows: usize, max_page_rows: usize, len: usize) -> usize { + if window_rows == 0 { + return len.max(1); + } + window_rows.max(max_page_rows).max(1) +} + +/// One prefetched unit's compressed bytes, plus the byte-budget permits they +/// hold. A unit is either a ROW WINDOW (`sel = Some((skip, take))`, all projected +/// columns) or a COLUMN GROUP (`sel = None`, a slice of the projection over all +/// rows) — the decode side reassembles accordingly; the fetch side treats both +/// identically. The permits travel with the bytes (into the decode stream's +/// `PrefetchedReader`), so dropping the decoded stream frees the bytes AND +/// releases the permits — which is what wakes the admission loop to launch the +/// next unit's fetch. That drop-to-wake handoff is the whole backpressure +/// mechanism: memory pressure stays ~constant at `prefetch_budget` compressed +/// bytes without any explicit signalling code. The permit is `Arc`-shared: +/// a column-group episode's n units are admitted under ONE summed permit +/// (their streams are all held open for the lockstep hstack, so per-unit +/// permits would deadlock the admission loop), released when the last of the +/// group's streams drops; a row window is an episode of one, so its Arc is +/// sole-owner and drops exactly as before. +struct PrefetchedUnit { + mask: ProjectionMask, + sel: Option<(usize, usize)>, + ranges: Vec>, + data: Vec, + permit: Arc, +} + +/// Permits (KiB) one admission episode (a row window, or a column-group set +/// admitted together) may hold from the prefetch bucket: its compressed +/// size, capped at HALF the budget so a single oversized episode can never +/// drain the whole semaphore. When one unit held the full budget, no other fetch +/// could be admitted while it decoded — fetch stopped overlapping decode and +/// the read went strictly serial (findings T6 measured this at up to 5× wall). +/// Capping at half guarantees at least two units can be in flight, restoring +/// the overlap; the cost is that the bucket under-accounts a unit whose real +/// size exceeds half the budget (its full bytes are fetched regardless — a +/// unit can't be split below a page/column), so peak in-flight compressed +/// bytes is bounded by `budget + 2 * max_oversized_unit_excess` rather than +/// `budget` exactly. Oversized units are rare after the windows-first planning +/// rule (a lone column bigger than `colwindow_budget` in a genuinely wide +/// group, or a no-offset-index chunk fallback). A 0/1-KiB budget still +/// degrades to strict fetch→decode→fetch as before. +fn unit_permit_kib(kib: u64, budget_kib: u64) -> u32 { + kib.clamp(1, (budget_kib / 2).max(1)) as u32 +} + +/// Find `want` inside one of the prefetched `ranges` and return the matching +/// slice of its `Bytes` (refcounted view, no copy). `None` if no prefetched +/// range fully contains it. Pure so it unit-tests without an object store. +fn slice_prefetched(ranges: &[Range], data: &[Bytes], want: &Range) -> Option { + for (r, b) in ranges.iter().zip(data) { + if want.start >= r.start && want.end <= r.end { + let s = (want.start - r.start) as usize; + let e = (want.end - r.start) as usize; + return Some(b.slice(s..e)); + } + } + None +} + +/// `AsyncFileReader` that serves a unit's page reads from its prefetched +/// buffers instead of S3. The decode stream built on top of this never touches +/// the network — the fetch already happened, budget-gated, in the admission +/// loop. Requests are always sub-ranges of the prefetched ranges (page reads +/// within a chunk, or within a window's page span), so containment lookup +/// suffices. Owning the permit ties the budget release to the stream's drop. +struct PrefetchedReader { + ranges: Vec>, + data: Vec, + meta: Arc, + _permit: Arc, +} + +impl AsyncFileReader for PrefetchedReader { + fn get_bytes(&mut self, range: Range) -> BoxFuture<'_, parquet::errors::Result> { + let res = slice_prefetched(&self.ranges, &self.data, &range).ok_or_else(|| { + ParquetError::General(format!( + "column-prefetch: byte range {range:?} was not prefetched" + )) + }); + Box::pin(futures::future::ready(res)) + } + + fn get_metadata<'a>( + &'a mut self, + _options: Option<&'a ArrowReaderOptions>, + ) -> BoxFuture<'a, parquet::errors::Result>> { + // Never hit in practice (streams are built with `new_with_metadata`), + // but the trait requires it. + let meta = self.meta.clone(); + Box::pin(futures::future::ready(Ok(meta))) + } +} + +/// Byte ranges + compressed KiB of one column group's chunks, straight from +/// the footer metadata — the "conservative estimate from parquet metadata" +/// that sizes each fetch exactly (chunk offsets/lengths are exact, not +/// estimates, so the budget accounting is exact too). +fn group_fetch_plan( + rgm: &parquet::file::metadata::RowGroupMetaData, + group: &[usize], +) -> (Vec>, u64) { + let mut ranges = Vec::with_capacity(group.len()); + let mut bytes = 0u64; + for &leaf in group { + let cc = rgm.column(leaf); + let (start, len) = cc.byte_range(); + ranges.push(start..start + len); + bytes += len; + } + (ranges, bytes.div_ceil(1024)) +} + +/// Exact byte ranges a ROW-WINDOW decode will request from the store: for each +/// projected leaf column of row group `rg`, the dictionary page (when present) +/// plus the contiguous span of data pages overlapping rows `[w, w+wlen)` — read +/// off the same offset index the decoder's `RowSelection` uses to skip pages, +/// so prefetch and decode always agree on which bytes are needed. A column +/// without an offset index, or a window covering the whole group, falls back to +/// the whole column chunk — also exactly what the decoder requests in that +/// case. Returns the ranges plus their compressed KiB (exact, from the footer) +/// for budget accounting, like [`group_fetch_plan`]. +fn window_fetch_plan( + md: &parquet::file::metadata::ParquetMetaData, + rg: usize, + leaves: &[usize], + w: usize, + wlen: usize, +) -> (Vec>, u64) { + let rgm = md.row_group(rg); + let num_rows = rgm.num_rows().max(0); + let rg_oi = md.offset_index().and_then(|oi| oi.get(rg)); + let whole = w == 0 && wlen as i64 >= num_rows; + let (w0, w1) = (w as i64, (w + wlen) as i64); + let mut ranges = Vec::with_capacity(leaves.len()); + let mut bytes = 0u64; + for &leaf in leaves { + let cc = rgm.column(leaf); + let (chunk_start, chunk_len) = cc.byte_range(); + let locs = rg_oi + .and_then(|oi| oi.get(leaf)) + .map(|c| c.page_locations()) + .filter(|l| !l.is_empty()); + let locs = match (whole, locs) { + (false, Some(l)) => l, + _ => { + bytes += chunk_len; + ranges.push(chunk_start..chunk_start + chunk_len); + continue; + } + }; + // Dictionary page: sits between the chunk start and the first data + // page, and every window of a dictionary-encoded chunk needs it. + let first_data = locs[0].offset.max(0) as u64; + if cc.dictionary_page_offset().is_some() && first_data > chunk_start { + bytes += first_data - chunk_start; + ranges.push(chunk_start..first_data); + } + // Contiguous span of the data pages overlapping [w, w+wlen). Pages are + // laid out back-to-back within a chunk, so one range covers the run. + let mut span: Option<(u64, u64)> = None; + for (i, loc) in locs.iter().enumerate() { + let rows_end = locs + .get(i + 1) + .map(|next| next.first_row_index) + .unwrap_or(num_rows); + if loc.first_row_index < w1 && rows_end > w0 { + let s = loc.offset.max(0) as u64; + let e = s + loc.compressed_page_size.max(0) as u64; + span = Some(match span { + None => (s, e), + Some((a, _)) => (a, e), + }); + } + } + if let Some((s, e)) = span { + bytes += e - s; + ranges.push(s..e); + } + } + (ranges, bytes.div_ceil(1024)) +} + +/// How one row group's prefetched units reassemble into output batches. +enum RgDecode { + /// N row-window units: decode each in order, stream every batch straight out. + Windows(usize), + /// N column-group units: decode all N in LOCKSTEP (their batch boundaries + /// are row-aligned by construction), hstacking one batch-slice at a time to + /// full width and emitting it before the next slice is decoded. + Hstack(usize), +} + +/// Decode structure of one (row group, sub-range), in output order. +struct RgPlan { + rg: usize, + batch_rows: usize, + decode: RgDecode, +} + +/// Fetch plan of one unit, in decode order: what to project, which rows +/// (`None` = all), which bytes, and their compressed KiB for budget accounting. +struct UnitFetch { + mask: ProjectionMask, + sel: Option<(usize, usize)>, + ranges: Vec>, + kib: u64, +} + +/// Plan every `(row group, start, len)` sub-range into prefetchable units. +/// The split axis is chosen per row group by its shape: +/// * ROW WINDOWS whenever they can actually split the range (the effective +/// window step is smaller than the range): each unit is ~`fetch_window_mb` +/// compressed bytes of ALL projected columns, page-aligned via the offset +/// index (see `effective_window_step`). Windows stream batches straight +/// out, so decoded retention is ~one decode budget regardless of group +/// size — which is why they are preferred whenever possible. +/// * COLUMN GROUPS only when row-windowing is inert (the step covers the +/// whole range — a wide/short group whose every column is a single page, a +/// missing offset index, or a group smaller than one fetch window) AND the +/// projected roots' compressed bytes partition into >1 group under +/// `colwindow_budget`: each unit is a slice of the projection over all the +/// group's rows. The async reader's `InMemoryRowGroup` otherwise stages +/// every projected chunk of the group at once, which for a 5000-column +/// schema was the entire S3 memory regression; fetching group-sized units +/// bounds each unit to ~`colwindow_budget`. The decode side +/// (`RgDecode::Hstack`) holds all the groups' COMPRESSED bytes and decodes +/// them in lockstep, one row-aligned batch-slice at a time, so decoded +/// retention is ~one full-width batch — it no longer parks whole decoded +/// groups (the old TODO-1u behavior, which degenerated to PyArrow's +/// whole-group retention). Tall-fat-column groups (few projected columns, +/// each over the budget) used to mis-select this axis and retain the +/// entire decoded row group (findings M20: 3.47 GB pinned, +/// `fetch_window_mb`-inert); the windows-first rule above fixed the +/// selection — they window instead. +/// Units are returned as admission EPISODES: each inner `Vec` is admitted +/// under one summed byte-budget permit in [`drive_s3`]. A row window is an +/// episode of one; an Hstack group's n units form one episode because the +/// decoder holds all n streams open for the lockstep hstack — admitting them +/// under per-unit permits would deadlock the admission loop once the bucket +/// ran dry mid-group. Planning globally (not per row group) lets the +/// prefetcher run ahead across row-group boundaries. +#[allow(clippy::too_many_arguments)] +fn plan_s3_units( + meta: &ArrowReaderMetadata, + full_mask: &ProjectionMask, + roots: &[usize], + subranges: &[(usize, usize, usize)], + batch_clamp: usize, + decode_budget: u64, + fetch_window_mb: u64, + colwindow_budget: u64, +) -> (Vec, Vec>) { + let md = meta.metadata(); + let schema_descr = md.file_metadata().schema_descr(); + let leaves = leaves_under_roots(schema_descr, roots); + let mut rg_plans: Vec = Vec::with_capacity(subranges.len()); + let mut units: Vec> = Vec::new(); + for &(rg, start, len) in subranges { + let rgm = md.row_group(rg); + let batch_rows = group_batch_rows(rgm, batch_clamp, decode_budget); + // Never window below the coarsest column's largest page (see + // max_page_rows): a sub-page window re-decodes that page in every + // window it overlaps. Wide/short groups (one page per column) collapse + // to a single window; tall multi-page groups still split. + let step = effective_window_step( + window_rows_for(rgm, fetch_window_mb), + max_page_rows(md, rg), + len, + ); + // Column groups only apply to whole-group reads (a K-split sub-range is + // by definition a tall group being split by rows, not columns) and only + // when row windows can't split the range (step covers it whole) — see + // the doc comment above for why windows always win when they can fire. + let whole = start == 0 && len == rgm.num_rows().max(0) as usize; + let groups = if whole && step >= len && colwindow_budget > 0 { + partition_columns_by_budget( + &projected_root_sizes(schema_descr, rgm, roots), + colwindow_budget, + ) + } else { + Vec::new() + }; + if groups.len() > 1 { + let n = groups.len(); + let mut episode = Vec::with_capacity(n); + for g in groups { + let (ranges, kib) = group_fetch_plan(rgm, &leaves_under_roots(schema_descr, &g)); + episode.push(UnitFetch { + mask: ProjectionMask::roots(schema_descr, g), + sel: None, + ranges, + kib, + }); + } + units.push(episode); + rg_plans.push(RgPlan { + rg, + batch_rows, + decode: RgDecode::Hstack(n), + }); + } else { + let (mut w, end) = (start, start + len); + let mut n = 0usize; + while w < end { + let wlen = step.min(end - w); + let (ranges, kib) = window_fetch_plan(md, rg, &leaves, w, wlen); + units.push(vec![UnitFetch { + mask: full_mask.clone(), + sel: Some((w, wlen)), + ranges, + kib, + }]); + w += wlen; + n += 1; + } + rg_plans.push(RgPlan { + rg, + batch_rows, + decode: RgDecode::Windows(n), + }); + } + } + (rg_plans, units) +} + +/// Receive the next prefetched unit (in plan order) and build its decode +/// stream, which serves every page read from the prefetched bytes — never the +/// network. The unit's budget permits ride inside the stream's +/// `PrefetchedReader`, so dropping the stream is what releases them. +async fn next_unit_stream( + hrx: &mut mpsc::Receiver>>, + meta: &ArrowReaderMetadata, + rg: usize, + batch_rows: usize, +) -> Result, ArrowError> { + let handle = hrx.recv().await.ok_or_else(|| { + ArrowError::ExternalError(Box::new(ParquetError::General( + "prefetch: admission loop ended early".to_string(), + ))) + })?; + let unit = match handle.await { + Ok(Ok(u)) => u, + Ok(Err(e)) => return Err(e), + Err(e) => return Err(ArrowError::ExternalError(Box::new(e))), // task panicked + }; + let PrefetchedUnit { + mask, + sel, + ranges, + data, + permit, + } = unit; + let reader = PrefetchedReader { + ranges, + data, + meta: Arc::clone(meta.metadata()), + _permit: permit, + }; + let mut builder = ParquetRecordBatchStreamBuilder::new_with_metadata(reader, meta.clone()) + .with_row_groups(vec![rg]) + .with_batch_size(batch_rows) + .with_projection(mask); + if let Some((skip, take)) = sel { + builder = builder.with_row_selection(RowSelection::from(vec![ + RowSelector::skip(skip), + RowSelector::select(take), + ])); + } + builder + .build() + .map_err(|e| ArrowError::ExternalError(Box::new(e))) +} + +/// THE single S3 decode driver: every read shape — row-windowed streaming, +/// wide-schema column groups, and each K-split row partition — flows through +/// here. Two halves, connected by a byte-denominated semaphore ("the bucket"): +/// +/// * admission loop (spawned): for each planned unit IN ORDER, acquire +/// permits equal to its compressed size (capped at half the budget so one +/// oversized unit never drains the bucket and serializes fetch behind +/// decode — see [`unit_permit_kib`]), then spawn its ranged GET. Fetches +/// whose permits fit run CONCURRENTLY — that's what overlaps S3 latency +/// with decode — while `acquire` blocks the loop the moment the bucket is +/// spent. +/// * decoder (this task): strictly one unit at a time (bounds decode +/// scratch), served entirely from the prefetched bytes. Dropping a decoded +/// unit's stream frees its bytes AND releases its permits, waking the +/// admission loop: constant memory pressure, fetch concurrency +/// self-adjusting to the fetch:decode speed ratio, no rate estimation +/// anywhere. +/// +/// `prefetch_budget_mb == 0` degrades to strict fetch→decode→fetch (every +/// acquire is for the full 1-permit budget). Output batches are sent to `tx` +/// in row order. +#[allow(clippy::too_many_arguments)] +async fn drive_s3( + store: Arc, + path: ObjPath, + meta: ArrowReaderMetadata, + out_schema: SchemaRef, + full_mask: ProjectionMask, + roots: Vec, + subranges: Vec<(usize, usize, usize)>, + batch_clamp: usize, + decode_budget: u64, + fetch_window_mb: u64, + colwindow_budget: u64, + prefetch_budget_mb: u64, + tx: mpsc::Sender>, +) { + // Send an error downstream and stop. + macro_rules! send_err { + ($e:expr) => {{ + let _ = tx.send(Err($e)).await; + return; + }}; + } + + let (rg_plans, units) = plan_s3_units( + &meta, + &full_mask, + &roots, + &subranges, + batch_clamp, + decode_budget, + fetch_window_mb, + colwindow_budget, + ); + + // --- admission loop: budget-gated concurrent prefetch --- + let budget_kib = prefetch_budget_mb.saturating_mul(1024).max(1); + let budget_kib = budget_kib.min(u32::MAX as u64 / 2); + let sem = Arc::new(Semaphore::new(budget_kib as usize)); + // Handles are tiny; the byte budget is what actually bounds prefetch. The + // channel only keeps the admission loop from racing unboundedly far ahead + // in *task count* when units are small. + let (htx, mut hrx) = + mpsc::channel::>>(64); + { + let store = store.clone(); + let path = path.clone(); + tokio::spawn(async move { + for episode in units { + // One summed permit per episode: a window is an episode of + // one (identical accounting to per-unit admission), while an + // Hstack group's units MUST be co-admitted — the decoder holds + // all their streams open for the lockstep hstack, so per-unit + // permits would deadlock once the bucket ran dry mid-group. + let total_kib: u64 = episode.iter().map(|u| u.kib).sum(); + let want = unit_permit_kib(total_kib, budget_kib); + let permit = match sem.clone().acquire_many_owned(want).await { + Ok(p) => Arc::new(p), + Err(_) => return, // semaphore closed = consumer gone + }; + for UnitFetch { + mask, + sel, + ranges, + kib: _, + } in episode + { + let store = store.clone(); + let path = path.clone(); + let permit = Arc::clone(&permit); + let handle = tokio::spawn(async move { + match store.get_ranges(&path, &ranges).await { + Ok(data) => Ok(PrefetchedUnit { + mask, + sel, + ranges, + data, + permit, + }), + Err(e) => Err(ArrowError::ExternalError(Box::new(e))), + } + }); + if htx.send(handle).await.is_err() { + return; // decoder dropped (error path) — stop admitting + } + } + } + }); + } + + // --- decoder: strictly one unit at a time (bounds decode scratch; + // concurrency lives ONLY on the fetch side above) --- + // Mid-stream batch adaptation (see `group_is_estimator_blind`): measured + // decoded bytes/row carries across units and row groups, re-sizing each + // window unit's batch rows at stream build time. Residual: the FIRST unit + // of a blind file still decodes at the static (encoded-fallback) size — + // adapting inside an already-built stream would need re-buildable + // prefetched bytes. The Hstack path is excluded: its units are column + // groups (partial rows), so their bytes/row is not comparable across + // units or with window units. + let mut bpr = BprTracker::default(); + for plan in rg_plans { + let blind = group_is_estimator_blind(meta.metadata().row_group(plan.rg)); + match plan.decode { + RgDecode::Windows(n) => { + for _ in 0..n { + let batch_rows = if blind { + adapted_rows(plan.batch_rows, decode_budget, bpr.bytes_per_row()) + } else { + plan.batch_rows + }; + let mut stream = + match next_unit_stream(&mut hrx, &meta, plan.rg, batch_rows).await { + Ok(s) => s, + Err(e) => send_err!(e), + }; + while let Some(item) = stream.next().await { + if blind { + if let Ok(b) = &item { + bpr.record(b); + } + } + let is_err = item.is_err(); + let msg = item.map_err(|e| ArrowError::ExternalError(Box::new(e))); + if tx.send(msg).await.is_err() { + return; // consumer dropped + } + if is_err { + return; + } + } + // `stream` drops here -> window bytes freed, permits + // released, next fetch admitted. + } + } + RgDecode::Hstack(n) => { + // Incremental hstack (the TODO-1u fix): open ALL n column-group + // streams at once — their COMPRESSED bytes stay resident, + // co-admitted under one shared permit — and decode in lockstep: + // batch-slice i of every group is glued to full width and + // emitted before slice i+1 is decoded. Every stream is built + // with the same batch_rows over the same whole-group rows with + // no row selection, so batch boundaries align by construction. + // Peak decoded retention is ~one full-width batch instead of + // the whole decoded row group the old parked-groups hstack + // held (which degenerated to PyArrow's whole-group retention). + let mut streams = Vec::with_capacity(n); + for _ in 0..n { + match next_unit_stream(&mut hrx, &meta, plan.rg, plan.batch_rows).await { + Ok(s) => streams.push(s), + Err(e) => send_err!(e), + } + } + loop { + let mut slices: Vec = Vec::with_capacity(n); + let mut ended = 0usize; + for stream in &mut streams { + match stream.next().await { + Some(Ok(b)) => slices.push(b), + Some(Err(e)) => send_err!(ArrowError::ExternalError(Box::new(e))), + None => ended += 1, + } + } + if ended == n { + break; // all groups exhausted together + } + if ended != 0 { + send_err!(ArrowError::ComputeError(format!( + "column-window batch-count mismatch in row group {}: \ + {ended} of {n} groups ended early", + plan.rg + ))); + } + let mut cols: Vec = Vec::with_capacity(out_schema.fields().len()); + for s in &slices { + cols.extend(s.columns().iter().cloned()); + } + // try_new re-checks row alignment: unequal column lengths + // across groups fail here rather than emitting a torn batch. + match RecordBatch::try_new(out_schema.clone(), cols) { + Ok(b) => { + if tx.send(Ok(b)).await.is_err() { + return; // consumer dropped + } + } + Err(e) => send_err!(e), + } + } + // Streams drop here -> the group's bytes are freed and the + // shared permit is released, waking the admission loop. + } + } + } +} + +/// Build an S3 `ObjectStore` from the config recovered from the pyarrow +/// `S3FileSystem` on the Python side (`fs.__reduce__()[1][0]`) so credentialed / +/// custom-endpoint (MinIO, moto) / anonymous buckets all connect identically to +/// PyArrow. Empty/None fields are treated as unset. Shared by `read_row_groups_s3` +/// and `read_metadata_s3`. +#[allow(clippy::too_many_arguments)] +fn build_s3_store( + bucket: &str, + region: &str, + anonymous: bool, + endpoint: Option, + access_key_id: Option, + secret_access_key: Option, + session_token: Option, + allow_http: bool, + virtual_hosted_style: bool, +) -> Result, object_store::Error> { + let mut sb = AmazonS3Builder::new() + .with_bucket_name(bucket) + .with_region(region) + .with_virtual_hosted_style_request(virtual_hosted_style); + if let Some(ep) = endpoint.filter(|s| !s.is_empty()) { + sb = sb.with_endpoint(ep); + } + if allow_http { + sb = sb.with_allow_http(true); + } + if anonymous { + // No signing — public buckets. Any creds are irrelevant. + sb = sb.with_skip_signature(true); + } else { + // Explicit static creds if the S3FileSystem carried them; otherwise the + // builder falls back to the AWS credential chain (env / IMDS role). + if let Some(kid) = access_key_id.filter(|s| !s.is_empty()) { + sb = sb.with_access_key_id(kid); + } + if let Some(s) = secret_access_key.filter(|s| !s.is_empty()) { + sb = sb.with_secret_access_key(s); + } + if let Some(t) = session_token.filter(|s| !s.is_empty()) { + sb = sb.with_token(t); + } + } + Ok(Arc::new(sb.build()?)) +} + +#[pyfunction] +#[pyo3(signature = (bucket, key, region, anonymous, endpoint=None, access_key_id=None, + secret_access_key=None, session_token=None, allow_http=false, + virtual_hosted_style=false, row_groups=None, columns=None, + batch_size=131072, decode_budget_bytes=2*1024*1024, + fetch_window_mb=16, k=1, split_threshold_bytes=134217728, + predicate_json=None, column_fetch_mb=16, + prefetch_budget_mb=64))] +#[allow(clippy::too_many_arguments)] +fn read_row_groups_s3( + py: Python<'_>, + bucket: String, + key: String, + region: String, + anonymous: bool, + // Full S3 config, recovered from the pyarrow S3FileSystem on the Python side + // (fs.__reduce__()[1][0]) so credentialed / custom-endpoint (MinIO, moto) / + // anonymous buckets all decode identically to PyArrow. Empty/None → unset. + endpoint: Option, + access_key_id: Option, + secret_access_key: Option, + session_token: Option, + allow_http: bool, + virtual_hosted_style: bool, + row_groups: Option>, + columns: Option>, + batch_size: usize, + decode_budget_bytes: u64, + fetch_window_mb: u64, + k: usize, + split_threshold_bytes: u64, + // See `read_row_groups`: statistics row-group pruning only. + predicate_json: Option, + // Compressed-byte budget per column group for the wide-schema column-windowing + // split axis (0 disables it). See `plan_s3_units`. + column_fetch_mb: u64, + // Compressed bytes the prefetcher may hold in flight/buffered ahead of the + // (single) decoder — "the bucket", shared by every unit kind. Bounds memory + // by construction while letting enough GETs run concurrently to keep the + // decoder fed regardless of the fetch:decode speed ratio. 0 = strictly + // sequential (no overlap). See `drive_s3`. + prefetch_budget_mb: u64, +) -> PyResult { + let store = build_s3_store( + &bucket, + ®ion, + anonymous, + endpoint, + access_key_id, + secret_access_key, + session_token, + allow_http, + virtual_hosted_style, + ) + .map_err(to_py)?; + let obj_path = ObjPath::from(key); + + // Load footer + page index ONCE (Optional so a window's RowSelection can skip + // unselected pages by byte range). Blocking async footer fetch; release the + // GIL so sibling Python read threads (Ray's fragment pool) issue their own + // S3 requests in parallel. + let reader = py + .allow_threads(|| { + let (meta, _skipped) = shared_runtime().block_on(load_meta_s3( + store.clone(), + obj_path.clone(), + PageIndexPolicy::Optional, + ))?; + build_s3_reader( + store, + obj_path, + meta, + row_groups, + columns, + batch_size, + decode_budget_bytes, + fetch_window_mb, + k, + split_threshold_bytes, + predicate_json, + column_fetch_mb, + prefetch_budget_mb, + ) + }) + .map_err(to_py)?; + Ok(into_py_stream(Box::new(reader))) +} + +/// The metadata-independent half of the S3 read: everything after the store +/// construction and footer load. Shared by [`read_row_groups_s3`] (fresh client +/// + footer per call — the original API) and +/// [`NativeParquetFile::read_row_groups`] (client and parsed footer opened once +/// and reused across a read task's calls — TODO 1r, the fix for the +/// per-file S3 setup cost on multi-file bins). Builds the projected output +/// schema up front from an empty stream (no network); reporting the projected +/// schema is what keeps it matching the projected batches at the FFI boundary. +#[allow(clippy::too_many_arguments)] +fn build_s3_reader( + store: Arc, + obj_path: ObjPath, + meta: ArrowReaderMetadata, + row_groups: Option>, + columns: Option>, + batch_size: usize, + decode_budget_bytes: u64, + fetch_window_mb: u64, + k: usize, + split_threshold_bytes: u64, + predicate_json: Option, + column_fetch_mb: u64, + prefetch_budget_mb: u64, +) -> Result { + let rt = shared_runtime(); + let mask = projection_mask(meta.metadata().file_metadata().schema_descr(), &columns); + let schema = ParquetRecordBatchStreamBuilder::new_with_metadata( + ParquetObjectReader::new(store.clone(), obj_path.clone()), + meta.clone(), + ) + .with_projection(mask.clone()) + .with_row_groups(vec![]) + .build()? + .schema() + .clone(); + + let selected: Vec = match row_groups { + Some(v) => v, + None => (0..meta.metadata().num_row_groups()).collect(), + }; + // Statistics-based pruning (conservative) before any range GET is issued, so + // pruned groups cost no S3 traffic. Same mechanism as the local path. + let selected = apply_predicate(&meta, selected, &predicate_json); + + // K-split ONLY for a lone row group above the threshold with a page index — + // the case Ray's fragment pool can't parallelize. Mirrors the local rule so + // crate-K and Ray's pool never multiply. Otherwise a single driver (K=1) + // over all selected groups in order; Ray's pool parallelizes files. Each of + // the K streams gets its own prefetch bucket (decode parallelism is the + // point of the split), so ~`k * prefetch_budget` compressed may be in + // flight for this one deliberately-parallel shape. + let split = k > 1 + && selected.len() == 1 + && meta.metadata().row_group(selected[0]).total_byte_size() as u64 >= split_threshold_bytes + && meta.metadata().offset_index().is_some(); + + // Build the per-stream sub-range lists (each becomes one drive_s3 task + + // one channel, drained in order). How a sub-range further splits into + // prefetchable units (row windows vs column groups) is decided per row + // group inside the driver — see `plan_s3_units`. + let stream_ranges: Vec> = if split { + let rg = selected[0]; + let total_rows = meta.metadata().row_group(rg).num_rows().max(0) as usize; + let chunk = total_rows.div_ceil(k.max(1)).max(1); + let mut ranges = Vec::new(); + let mut start = 0usize; + while start < total_rows { + let len = chunk.min(total_rows - start); + ranges.push(vec![(rg, start, len)]); + start += len; + } + ranges + } else { + // One stream: every selected group, whole, in order. + let subranges = selected + .iter() + .map(|&rg| { + ( + rg, + 0usize, + meta.metadata().row_group(rg).num_rows().max(0) as usize, + ) + }) + .collect(); + vec![subranges] + }; + + let roots = projected_root_indices(meta.metadata().file_metadata().schema_descr(), &columns); + let colwindow_budget = column_fetch_mb.saturating_mul(1024 * 1024); + + // Spawn one driver per stream on the shared runtime; collect receivers in order. + let mut receivers = Vec::with_capacity(stream_ranges.len()); + for subranges in stream_ranges { + let (tx, rx) = mpsc::channel::>(S3_CHANNEL_DEPTH); + receivers.push(rx); + rt.spawn(drive_s3( + store.clone(), + obj_path.clone(), + meta.clone(), + schema.clone(), + mask.clone(), + roots.clone(), + subranges, + batch_size, + decode_budget_bytes, + fetch_window_mb, + colwindow_budget, + prefetch_budget_mb, + tx, + )); + } + + Ok(S3ChannelReader { + schema, + receivers, + cur: 0, + }) +} + +// --------------------------------------------------------------------------- // +// Per-file native handles (TODO 1r): open once, decode many +// --------------------------------------------------------------------------- // +// The original entry points above pay a fixed setup cost on EVERY call: a fresh +// `AmazonS3Builder` client (new connection pool, no TLS session reuse) plus a +// footer — and, for decode calls, page-index — fetch. PyArrow reads one footer +// per file and shares one HTTP client across the whole read. On the #64985 +// planner a read task's fragment is a multi-file *bin*, so the reader makes +// 2 calls per file (metadata at plan time, decode at read time): a 16-file S3 +// bin paid 32 client builds and 32 footer round trips, measured as a 3.5× +// read-op loss vs PyArrow (findings T10). The handles below restore parity of +// mechanism: `connect_s3` builds ONE client per (bucket, config) for the whole +// task, `open_file` fetches the footer+page index ONCE per file, and +// `read_row_groups` / `metadata` reuse both. +// +// Deliberately NO global/process-level cache behind these: the handle's +// lifetime is owned by the Python caller (one read task), so there is no +// staleness (rotated credentials, replaced objects) and no unbounded growth in +// a long-lived reused worker. + +/// Where a [`NativeParquetFile`]'s bytes live. Local files re-open the path per +/// reader (cheap, and `File` isn't shareable across the K-split threads anyway); +/// S3 files hold the shared client. +enum FileSource { + Local(String), + S3 { + store: Arc, + path: ObjPath, + }, +} + +/// One S3 client (connection pool + credentials) for one bucket, shared across +/// every file opened through it. Construct via [`connect_s3`]. +#[pyclass] +struct NativeS3Store { + store: Arc, +} + +#[pymethods] +impl NativeS3Store { + /// Open one object as a [`NativeParquetFile`]: fetches and parses the + /// footer (and, when `page_index` — the page index) exactly once, on this + /// store's shared client. `page_index=true` is what the S3 decode path + /// needs (row windows skip pages via the offset index); pass `false` for + /// metadata-only handles. + #[pyo3(signature = (key, page_index=true))] + fn open_file( + &self, + py: Python<'_>, + key: String, + page_index: bool, + ) -> PyResult { + let policy = if page_index { + PageIndexPolicy::Optional + } else { + PageIndexPolicy::Skip + }; + let store = self.store.clone(); + let path = ObjPath::from(key); + // Blocking async footer fetch; release the GIL for sibling read threads. + let (meta, skipped) = py + .allow_threads(|| { + shared_runtime().block_on(load_meta_s3(store.clone(), path.clone(), policy)) + }) + .map_err(to_py)?; + Ok(NativeParquetFile { + source: FileSource::S3 { store, path }, + meta, + arrow_schema_skipped: skipped, + }) + } +} + +/// Build the per-bucket S3 client once. Same config contract as +/// [`read_row_groups_s3`] (recovered from the pyarrow `S3FileSystem` on the +/// Python side); the returned store is what every `open_file` shares. +#[pyfunction] +#[pyo3(signature = (bucket, region, anonymous, endpoint=None, access_key_id=None, + secret_access_key=None, session_token=None, allow_http=false, + virtual_hosted_style=false))] +#[allow(clippy::too_many_arguments)] +fn connect_s3( + bucket: String, + region: String, + anonymous: bool, + endpoint: Option, + access_key_id: Option, + secret_access_key: Option, + session_token: Option, + allow_http: bool, + virtual_hosted_style: bool, +) -> PyResult { + let store = build_s3_store( + &bucket, + ®ion, + anonymous, + endpoint, + access_key_id, + secret_access_key, + session_token, + allow_http, + virtual_hosted_style, + ) + .map_err(to_py)?; + Ok(NativeS3Store { store }) +} + +/// Local counterpart of [`NativeS3Store::open_file`]: parse the footer once and +/// reuse it across `metadata()` and every `read_row_groups` call. `page_index` +/// mirrors the lean-footer-parse rule of [`open_local_reader`]: only the K-split +/// needs it, so callers pass `k > 1`. +#[pyfunction] +#[pyo3(signature = (path, page_index=false))] +fn open_parquet_file( + py: Python<'_>, + path: String, + page_index: bool, +) -> PyResult { + let policy = if page_index { + PageIndexPolicy::Optional + } else { + PageIndexPolicy::Skip + }; + // Blocking file I/O; release the GIL for sibling read threads. + let (meta, skipped) = py + .allow_threads(|| load_meta_local(&File::open(&path)?, policy)) + .map_err(to_py)?; + Ok(NativeParquetFile { + source: FileSource::Local(path), + meta, + arrow_schema_skipped: skipped, + }) +} + +/// One opened Parquet file: the parsed footer plus (for S3) the shared client. +/// `metadata()` is free (no I/O); `read_row_groups` skips the footer fetch the +/// original entry points pay per call. +#[pyclass] +struct NativeParquetFile { + source: FileSource, + meta: ArrowReaderMetadata, + arrow_schema_skipped: bool, +} + +#[pymethods] +impl NativeParquetFile { + /// The same footer summary `read_metadata` / `read_metadata_s3` return, + /// built from the already-parsed footer — zero I/O. + fn metadata(&self) -> ParquetFileMetadata { + build_file_metadata(&self.meta, self.arrow_schema_skipped) + } + + /// Replace this handle's Arrow output schema with a caller-supplied one + /// (an Arrow C schema capsule, i.e. `pa.Schema.__arrow_c_schema__()`), + /// rebuilding the reader metadata against the already-parsed footer — + /// zero additional I/O on either transport. + /// + /// Purpose (findings M52): when the embedded arrow schema was skipped + /// (non-UTF8 extension metadata, `arrow_schema_skipped`), the inferred + /// storage types can differ from the extension types' storage + /// (`list` vs `large_list`), forcing the Python reader + /// into a per-batch per-column `Table.cast` (~7 µs/col/batch). Supplying + /// the exact storage schema here makes the crate decode directly into + /// those types, so Python re-attaches the extension labels with one + /// zero-copy C-interface import per batch instead. + /// + /// The supplied schema must be plain storage types with no binary field + /// metadata (the cloudpickle label bytes stay in Python — Rust never + /// holds them). parquet-rs validates the schema against the parquet + /// footer and errors on any mismatch, in which case the caller keeps + /// today's cast path; `self.meta` is only replaced on success. + fn with_schema_override(&mut self, schema_capsule: Bound<'_, PyCapsule>) -> PyResult<()> { + let valid_name = schema_capsule + .name() + .map_err(to_py)? + .map(|n| n.to_bytes() == b"arrow_schema") + .unwrap_or(false); + if !valid_name { + return Err(PyRuntimeError::new_err( + "with_schema_override expects an 'arrow_schema' PyCapsule \ + (pa.Schema.__arrow_c_schema__())", + )); + } + let ptr = schema_capsule.pointer() as *const FFI_ArrowSchema; + if ptr.is_null() { + return Err(PyRuntimeError::new_err("null arrow_schema capsule")); + } + // Borrowed read of the C struct: the capsule keeps ownership and will + // run its own release callback; try_from copies into Rust types. + let schema = Schema::try_from(unsafe { &*ptr }).map_err(to_py)?; + let options = ArrowReaderOptions::new().with_schema(Arc::new(schema)); + self.meta = ArrowReaderMetadata::try_new(Arc::clone(self.meta.metadata()), options) + .map_err(to_py)?; + Ok(()) + } + + /// Decode row groups through the held footer + client. Argument semantics + /// are identical to `read_row_groups` / `read_row_groups_s3`; the S3-only + /// knobs (`fetch_window_mb`, `column_fetch_mb`, `prefetch_budget_mb`) are + /// inert on a local handle, so one uniform signature serves both. + #[pyo3(signature = (row_groups=None, columns=None, batch_size=131072, + decode_budget_bytes=2*1024*1024, k=1, + split_threshold_bytes=134217728, predicate_json=None, + fetch_window_mb=16, column_fetch_mb=16, + prefetch_budget_mb=64))] + #[allow(clippy::too_many_arguments)] + fn read_row_groups( + &self, + py: Python<'_>, + row_groups: Option>, + columns: Option>, + batch_size: usize, + decode_budget_bytes: u64, + k: usize, + split_threshold_bytes: u64, + predicate_json: Option, + fetch_window_mb: u64, + column_fetch_mb: u64, + prefetch_budget_mb: u64, + ) -> PyResult { + match &self.source { + FileSource::Local(path) => { + let (path, meta) = (path.clone(), self.meta.clone()); + let reader = py + .allow_threads(|| { + build_local_reader( + path, + meta, + row_groups, + columns, + batch_size, + decode_budget_bytes, + k, + split_threshold_bytes, + predicate_json, + ) + }) + .map_err(to_py)?; + Ok(into_py_stream(reader)) + } + FileSource::S3 { store, path } => { + let (store, path, meta) = (store.clone(), path.clone(), self.meta.clone()); + let reader = py + .allow_threads(|| { + build_s3_reader( + store, + path, + meta, + row_groups, + columns, + batch_size, + decode_budget_bytes, + fetch_window_mb, + k, + split_threshold_bytes, + predicate_json, + column_fetch_mb, + prefetch_budget_mb, + ) + }) + .map_err(to_py)?; + Ok(into_py_stream(Box::new(reader))) + } + } + } +} + +// --------------------------------------------------------------------------- // +// Footer-only metadata reads (Track 1): one footer parse, no data decode. +// --------------------------------------------------------------------------- // +/// Read just the Parquet footer of a local file and return schema + per-row-group +/// counts. Page index is skipped (not needed for metadata). Lets the Python reader +/// stop building a PyArrow dataset to learn the schema / row-group layout. +#[pyfunction] +fn read_metadata(py: Python<'_>, path: String) -> PyResult { + // Blocking footer read; release the GIL for sibling Python threads. + let (meta, skipped) = py + .allow_threads(|| load_meta_local(&File::open(&path)?, PageIndexPolicy::Skip)) + .map_err(to_py)?; + Ok(build_file_metadata(&meta, skipped)) +} + +/// S3 counterpart of [`read_metadata`]: one async footer fetch via `object_store`, +/// same connection config recovery as `read_row_groups_s3`. +#[pyfunction] +#[pyo3(signature = (bucket, key, region, anonymous, endpoint=None, access_key_id=None, + secret_access_key=None, session_token=None, allow_http=false, + virtual_hosted_style=false))] +#[allow(clippy::too_many_arguments)] +fn read_metadata_s3( + py: Python<'_>, + bucket: String, + key: String, + region: String, + anonymous: bool, + endpoint: Option, + access_key_id: Option, + secret_access_key: Option, + session_token: Option, + allow_http: bool, + virtual_hosted_style: bool, +) -> PyResult { + let store = build_s3_store( + &bucket, + ®ion, + anonymous, + endpoint, + access_key_id, + secret_access_key, + session_token, + allow_http, + virtual_hosted_style, + ) + .map_err(to_py)?; + let obj_path = ObjPath::from(key); + let rt = shared_runtime(); + // Blocking async footer fetch; release the GIL for sibling Python threads. + let (meta, skipped) = py + .allow_threads(|| { + rt.block_on(load_meta_s3( + store.clone(), + obj_path.clone(), + PageIndexPolicy::Skip, + )) + }) + .map_err(to_py)?; + Ok(build_file_metadata(&meta, skipped)) +} + +/// Return the row-group ids of `path` that survive `predicate_json`'s statistics +/// pruning (all of them when it's None). This is the exact selection +/// `read_row_groups` would decode; exposed so callers (and tests) can observe +/// pruning without decoding, and so the pyarrow-free reader can learn the read +/// set up front. Page index is skipped (stats live in the footer). +#[pyfunction] +#[pyo3(signature = (path, predicate_json=None))] +fn select_row_groups( + py: Python<'_>, + path: String, + predicate_json: Option, +) -> PyResult> { + // Blocking footer read; release the GIL for sibling Python threads. + let (meta, _skipped) = py + .allow_threads(|| load_meta_local(&File::open(&path)?, PageIndexPolicy::Skip)) + .map_err(to_py)?; + let all: Vec = (0..meta.metadata().num_row_groups()).collect(); + Ok(apply_predicate(&meta, all, &predicate_json)) +} + +fn to_py(e: E) -> PyErr { + PyRuntimeError::new_err(e.to_string()) +} + +#[pymodule] +fn ray_data_arrow_rs(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(read_row_groups, m)?)?; + m.add_function(wrap_pyfunction!(read_row_groups_s3, m)?)?; + m.add_function(wrap_pyfunction!(read_metadata, m)?)?; + m.add_function(wrap_pyfunction!(read_metadata_s3, m)?)?; + m.add_function(wrap_pyfunction!(select_row_groups, m)?)?; + m.add_function(wrap_pyfunction!(connect_s3, m)?)?; + m.add_function(wrap_pyfunction!(open_parquet_file, m)?)?; + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use parquet::schema::types::{ColumnPath, Type}; + use std::sync::Arc; + + /// Build a leaf `ColumnDescriptor` for one physical/logical/converted type + /// combo (the only inputs `is_plain_signed_int` reads). + fn descr( + physical: PhysicalType, + logical: Option, + converted: ConvertedType, + ) -> ColumnDescriptor { + let ty = Type::primitive_type_builder("c", physical) + .with_logical_type(logical) + .with_converted_type(converted) + .build() + .unwrap(); + ColumnDescriptor::new(Arc::new(ty), 0, 0, ColumnPath::new(vec!["c".to_string()])) + } + + #[test] + fn plain_signed_ints_are_comparable() { + // Plain INT32/INT64 (no logical/converted type). + assert!(is_plain_signed_int(&descr( + PhysicalType::INT64, + None, + ConvertedType::NONE + ))); + assert!(is_plain_signed_int(&descr( + PhysicalType::INT32, + None, + ConvertedType::NONE + ))); + // Explicit signed-integer logical type. + assert!(is_plain_signed_int(&descr( + PhysicalType::INT32, + Some(LogicalType::integer(32, true)), + ConvertedType::NONE + ))); + // Legacy signed converted type. + assert!(is_plain_signed_int(&descr( + PhysicalType::INT32, + None, + ConvertedType::INT_16 + ))); + } + + #[test] + fn window_step_clamps_up_to_a_full_page() { + // The wide/short pathology: byte budget wants a 40-row window but the + // (single) page spans all 200 rows -> clamp up to 200 so the whole range + // is one window (no per-window re-decode of that page). + assert_eq!(effective_window_step(40, 200, 200), 200); + // step >= len => the caller emits exactly one window (parity, no split). + assert!(effective_window_step(40, 200, 200) >= 200); + } + + #[test] + fn window_step_keeps_splitting_a_tall_group() { + // Tall multi-page group: window (4096 rows) already spans several 512-row + // pages, so it is left as-is and the range still splits into windows. + assert_eq!(effective_window_step(4096, 512, 1_000_000), 4096); + // A window smaller than the page still gets clamped up to the page, so a + // page is never split across windows (bounds boundary re-decode to O(1)). + assert_eq!(effective_window_step(256, 512, 1_000_000), 512); + } + + #[test] + fn column_partition_splits_wide_group_under_budget() { + // 5 columns of 100 bytes each, budget 250 -> groups of [0,1],[2,3],[4]. + let cols: Vec<(usize, u64)> = (0..5).map(|i| (i, 100)).collect(); + assert_eq!( + partition_columns_by_budget(&cols, 250), + vec![vec![0, 1], vec![2, 3], vec![4]] + ); + // Ascending leaf order is preserved within and across groups (so hstack + // reproduces schema order). + let flat: Vec = partition_columns_by_budget(&cols, 250) + .into_iter() + .flatten() + .collect(); + assert_eq!(flat, vec![0, 1, 2, 3, 4]); + } + + #[test] + fn column_partition_disabled_or_narrow_is_one_group() { + let cols: Vec<(usize, u64)> = (0..5).map(|i| (i, 100)).collect(); + // budget 0 disables -> single group. + assert_eq!( + partition_columns_by_budget(&cols, 0), + vec![vec![0, 1, 2, 3, 4]] + ); + // budget larger than the total -> single group (narrow/small reads). + assert_eq!( + partition_columns_by_budget(&cols, 10_000), + vec![vec![0, 1, 2, 3, 4]] + ); + // <=1 column -> single group regardless of budget. + assert_eq!(partition_columns_by_budget(&[(7, 999)], 1), vec![vec![7]]); + } + + #[test] + fn column_partition_never_splits_below_one_column() { + // A single oversized column exceeds the budget but still gets its own group + // (can't split a column below itself). Neighbours don't merge into it. + let cols = vec![(0, 10), (1, 500), (2, 10)]; + assert_eq!( + partition_columns_by_budget(&cols, 100), + vec![vec![0], vec![1], vec![2]] + ); + } + + #[test] + fn slice_prefetched_serves_contained_subranges() { + // Two prefetched chunk ranges; page reads are sub-ranges of a chunk. + let ranges = vec![100u64..200, 300u64..350]; + let data = vec![ + Bytes::from((0..100u8).collect::>()), + Bytes::from((0..50u8).collect::>()), + ]; + // Exact chunk. + assert_eq!( + slice_prefetched(&ranges, &data, &(100..200)).unwrap(), + data[0] + ); + // Interior page of the first chunk: bytes at offsets 10..15 within it. + assert_eq!( + slice_prefetched(&ranges, &data, &(110..115)) + .unwrap() + .as_ref(), + &[10, 11, 12, 13, 14] + ); + // Sub-range of the second chunk. + assert_eq!( + slice_prefetched(&ranges, &data, &(340..350)) + .unwrap() + .as_ref(), + &(40..50u8).collect::>()[..] + ); + } + + #[test] + fn slice_prefetched_rejects_uncached_or_straddling_ranges() { + let ranges = vec![100u64..200, 300u64..350]; + let data = vec![Bytes::from(vec![0u8; 100]), Bytes::from(vec![0u8; 50])]; + // Not prefetched at all. + assert!(slice_prefetched(&ranges, &data, &(0..10)).is_none()); + // Straddles the gap between the two chunks -> contained in neither. + assert!(slice_prefetched(&ranges, &data, &(150..320)).is_none()); + // Runs past the end of a chunk. + assert!(slice_prefetched(&ranges, &data, &(190..201)).is_none()); + } + + #[test] + fn window_step_zero_means_whole_range() { + // fetch_window_mb == 0 -> window_rows == 0 -> one window over the range, + // regardless of page size. + assert_eq!(effective_window_step(0, 512, 8192), 8192); + assert_eq!(effective_window_step(0, 0, 1), 1); + } + + /// Write a 2-column, 1000-row parquet (data pages capped at ~100 rows, page + /// index on) into memory and load its metadata WITH the page index — the + /// fixture for the window/unit planning tests below. + fn windowed_fixture() -> (ArrowReaderMetadata, Bytes) { + use arrow::array::{Int64Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use parquet::arrow::arrow_writer::ArrowWriter; + use parquet::file::properties::WriterProperties; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Utf8, false), + ])); + let a = Int64Array::from((0..1000i64).collect::>()); + let b = StringArray::from((0..1000).map(|i| format!("row-{i}")).collect::>()); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap(); + let props = WriterProperties::builder() + .set_data_page_row_count_limit(100) + .set_write_batch_size(100) + .build(); + let mut buf = Vec::new(); + let mut w = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + let buf = Bytes::from(buf); + let meta = + ArrowReaderMetadata::load(&buf, reader_options(PageIndexPolicy::Required)).unwrap(); + (meta, buf) + } + + #[test] + fn window_plan_whole_group_equals_whole_chunks() { + let (meta, _buf) = windowed_fixture(); + let md = meta.metadata(); + let leaves = vec![0usize, 1]; + let (ranges, kib) = window_fetch_plan(md, 0, &leaves, 0, 1000); + let (want_ranges, want_kib) = group_fetch_plan(md.row_group(0), &leaves); + assert_eq!(ranges, want_ranges); + assert_eq!(kib, want_kib); + } + + #[test] + fn window_plan_subwindow_fetches_less_than_the_chunk() { + let (meta, _buf) = windowed_fixture(); + let md = meta.metadata(); + let leaves = vec![0usize, 1]; + let (_, whole_kib) = window_fetch_plan(md, 0, &leaves, 0, 1000); + // A 100-row window out of 1000 (10 pages/column) must fetch strictly + // less than the whole chunks... + let (ranges, kib) = window_fetch_plan(md, 0, &leaves, 400, 100); + assert!(kib < whole_kib, "window kib {kib} >= whole {whole_kib}"); + // ...and every planned range must sit inside one of the column chunks. + for r in &ranges { + let contained = leaves.iter().any(|&l| { + let (s, len) = md.row_group(0).column(l).byte_range(); + r.start >= s && r.end <= s + len + }); + assert!(contained, "range {r:?} escapes the column chunks"); + } + } + + /// The load-bearing guarantee: decode served ONLY from a window's planned + /// ranges must succeed (a request outside the plan is a hard + /// "not prefetched" error) and reproduce exactly the window's rows — for + /// windows that tile the group at non-page-aligned offsets. + #[test] + fn window_plans_serve_every_decoder_request() { + use arrow::array::Int64Array; + + let (meta, buf) = windowed_fixture(); + let md = meta.metadata(); + let leaves = vec![0usize, 1]; + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let sem = Arc::new(Semaphore::new(usize::MAX >> 3)); + let mut next_val = 0i64; + for (w, wlen) in [(0usize, 250usize), (250, 250), (500, 400), (900, 100)] { + let (ranges, _kib) = window_fetch_plan(md, 0, &leaves, w, wlen); + let data: Vec = ranges + .iter() + .map(|r| buf.slice(r.start as usize..r.end as usize)) + .collect(); + let permit = rt.block_on(sem.clone().acquire_many_owned(1)).unwrap(); + let reader = PrefetchedReader { + ranges, + data, + meta: Arc::clone(meta.metadata()), + _permit: Arc::new(permit), + }; + let batches: Vec = rt.block_on(async { + let mut stream = + ParquetRecordBatchStreamBuilder::new_with_metadata(reader, meta.clone()) + .with_row_groups(vec![0]) + .with_batch_size(97) + .with_projection(ProjectionMask::all()) + .with_row_selection(RowSelection::from(vec![ + RowSelector::skip(w), + RowSelector::select(wlen), + ])) + .build() + .unwrap(); + let mut out = Vec::new(); + while let Some(b) = stream.next().await { + out.push(b.unwrap()); + } + out + }); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, wlen, "window ({w},{wlen}) yielded {rows} rows"); + for b in &batches { + let col = b.column(0).as_any().downcast_ref::().unwrap(); + for i in 0..col.len() { + assert_eq!(col.value(i), next_val); + next_val += 1; + } + } + } + assert_eq!(next_val, 1000); + } + + /// Write a single-column Utf8 parquet of `n` rows cycling through 4 + /// distinct 4 KiB values — dictionary-encoded by default, so the footer's + /// encoded size (~dict + indices) understates the decoded size ~500x: the + /// estimator-blind expansion shape for the adaptation tests below. + fn dict_string_file(n: usize, rows_per_group: usize, dictionary: bool) -> Vec { + use arrow::array::StringArray; + use arrow::datatypes::{DataType, Field, Schema}; + use parquet::arrow::arrow_writer::ArrowWriter; + use parquet::file::properties::WriterProperties; + + let vals: Vec = (0..4u8).map(|i| i.to_string().repeat(4096)).collect(); + let col: Vec<&str> = (0..n).map(|i| vals[i % 4].as_str()).collect(); + let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)])); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(col))]).unwrap(); + let props = WriterProperties::builder() + .set_dictionary_enabled(dictionary) + .set_max_row_group_size(rows_per_group) + .build(); + let mut buf = Vec::new(); + let mut w = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + buf + } + + fn write_temp(name: &str, bytes: &[u8]) -> std::path::PathBuf { + let path = std::env::temp_dir().join(format!("rrs_{}_{}.parquet", name, std::process::id())); + std::fs::write(&path, bytes).unwrap(); + path + } + + /// Every value is 4096 repeats of one ASCII digit; check each row against + /// its expected cycle position so parity failures point at a row index. + fn assert_dict_string_content(batches: &[RecordBatch], n: usize) { + use arrow::array::{Array, StringArray}; + let mut row = 0usize; + for b in batches { + let col = b.column(0).as_any().downcast_ref::().unwrap(); + for i in 0..col.len() { + let v = col.value(i); + assert_eq!(v.len(), 4096, "row {row}: bad length"); + assert_eq!( + v.as_bytes()[0], + b'0' + (row % 4) as u8, + "row {row}: wrong value" + ); + row += 1; + } + } + assert_eq!(row, n, "row count mismatch"); + } + + #[test] + fn estimator_blind_detects_dict_byte_array_only() { + // Dict strings -> blind. + let buf = Bytes::from(dict_string_file(1000, 1_000_000, true)); + let meta = ArrowReaderMetadata::load(&buf, reader_options(PageIndexPolicy::Skip)).unwrap(); + assert!(group_is_estimator_blind(meta.metadata().row_group(0))); + // Same data with dictionary encoding off -> not blind. + let buf = Bytes::from(dict_string_file(1000, 1_000_000, false)); + let meta = ArrowReaderMetadata::load(&buf, reader_options(PageIndexPolicy::Skip)).unwrap(); + assert!(!group_is_estimator_blind(meta.metadata().row_group(0))); + } + + #[test] + fn adapted_rows_clamps_both_ways() { + let mib = 1024 * 1024u64; + // 4 KiB/row measured, 1 MiB budget -> 256 rows. + assert_eq!(adapted_rows(131_072, mib, Some(4096.0)), 256); + // Measurement smaller than the static estimate implies -> upper clamp. + assert_eq!(adapted_rows(500, mib, Some(1.0)), 500); + // Fat rows -> floor. + assert_eq!(adapted_rows(131_072, mib, Some(mib as f64)), MIN_BATCH_ROWS); + // Nothing measured -> static. + assert_eq!(adapted_rows(777, mib, None), 777); + } + + /// M50 residual (a), the fix under test: on a dict-string group the static + /// estimator falls back to encoded bytes (~8 B/row here), so it would + /// decode ALL rows as one ~80 MiB batch against a 1 MiB budget. Adaptation + /// must instead yield one MIN_BATCH_ROWS probe, then budget-sized batches. + #[test] + fn seq_reader_adapts_blind_dict_string_batches() { + let n = 20_000usize; + let path = write_temp("adapt1", &dict_string_file(n, 1_000_000, true)); + let budget = 1024 * 1024u64; + let reader = open_local_reader( + path.to_str().unwrap().to_string(), + None, + None, + 131_072, + budget, + 1, + u64::MAX, + None, + ) + .unwrap(); + let batches: Vec = reader.map(|b| b.unwrap()).collect(); + std::fs::remove_file(&path).ok(); + + assert_eq!(batches[0].num_rows(), MIN_BATCH_ROWS, "probe first"); + for (i, b) in batches[1..].iter().enumerate() { + let sz = b.get_array_memory_size() as u64; + assert!( + sz <= budget + budget / 2, + "batch {}: {} bytes > 1.5x budget", + i + 1, + sz + ); + } + // ...and not degenerate: ~budget/4KiB = 256 rows, not another probe. + assert!(batches[1].num_rows() >= 128, "over-shrunk: {}", batches[1].num_rows()); + assert_dict_string_content(&batches, n); + } + + /// The measurement carries across row groups: only the very first blind + /// group pays a probe; later groups open already adapted. + #[test] + fn seq_reader_probe_carries_across_groups() { + let n = 20_000usize; + let path = write_temp("adapt2", &dict_string_file(n, 5_000, true)); // 4 groups + let budget = 1024 * 1024u64; + let reader = open_local_reader( + path.to_str().unwrap().to_string(), + None, + None, + 131_072, + budget, + 1, + u64::MAX, + None, + ) + .unwrap(); + let batches: Vec = reader.map(|b| b.unwrap()).collect(); + std::fs::remove_file(&path).ok(); + + let probes = batches + .iter() + .filter(|b| b.num_rows() == MIN_BATCH_ROWS) + .count(); + assert_eq!(probes, 1, "exactly one probe across 4 groups"); + for (i, b) in batches[1..].iter().enumerate() { + let sz = b.get_array_memory_size() as u64; + assert!( + sz <= budget + budget / 2, + "batch {}: {} bytes > 1.5x budget", + i + 1, + sz + ); + } + assert_dict_string_content(&batches, n); + } + + /// Non-blind groups keep today's static sizing exactly: with dictionary + /// encoding off the encoded size ~= decoded size, one reader per group, + /// no probe batch. + #[test] + fn seq_reader_leaves_plain_groups_alone() { + let n = 2_000usize; + let path = write_temp("adapt3", &dict_string_file(n, 1_000_000, false)); + let budget = 1024 * 1024u64; + let reader = open_local_reader( + path.to_str().unwrap().to_string(), + None, + None, + 131_072, + budget, + 1, + u64::MAX, + None, + ) + .unwrap(); + let batches: Vec = reader.map(|b| b.unwrap()).collect(); + std::fs::remove_file(&path).ok(); + // Static sizing: ~4 KiB/row encoded -> 256-row batches, and the first + // batch is NOT a 32-row probe. + assert!(batches[0].num_rows() > MIN_BATCH_ROWS, "no probe on plain groups"); + assert_dict_string_content(&batches, n); + } + + #[test] + fn byte_budget_rows_floor_yields_to_the_budget() { + // The K8 defect: a 1 MiB/row fat-string group under the 32 MiB default + // budget must decode ~32-row batches — the old 2048-row floor made this + // 2048 rows (a 2 GiB batch), silently voiding the knob. + let mib = 1024 * 1024; + assert_eq!( + byte_budget_rows(1000 * mib, 1000, 131072, 32 * mib as u64), + 32 + ); + // 64 KiB/row, 1 MiB budget -> 16 rows... clamped up to the 32-row floor + // (per-batch overhead guard), the only place the floor still binds. + assert_eq!( + byte_budget_rows(1000 * 64 * 1024, 1000, 131072, mib as u64), + 32 + ); + // Narrow schema: budget_rows huge -> clamped to the caller's ask. + assert_eq!( + byte_budget_rows(8000, 1000, 131072, 32 * mib as u64), + 131072 + ); + // Unknown rows -> requested (unchanged behavior). + assert_eq!(byte_budget_rows(0, 0, 4096, 1), 4096); + } + + /// M43/M49: the footer's encoded size understates decoded size by the + /// dictionary expansion ratio, so batches sized from it overshot the + /// decode budget by that ratio. For fixed-width types the decoded size is + /// exactly `num_values * width` from the footer — verify the estimator + /// uses it, and that batch sizing shrinks accordingly. + #[test] + fn decoded_estimate_expands_dict_encoded_fixed_width() { + use arrow::array::Float64Array; + use arrow::datatypes::{DataType, Field, Schema}; + use parquet::arrow::arrow_reader::ArrowReaderMetadata; + use parquet::arrow::ArrowWriter; + use parquet::file::properties::WriterProperties; + + // 100k float64 rows drawn from a 4-value pool: dictionary-encodes to + // ~indices (2 bits/row RLE) + a 32-byte dict, so encoded-uncompressed + // is tiny while decoded is exactly 800 KB. + let rows = 100_000usize; + let schema = Arc::new(Schema::new(vec![Field::new("f", DataType::Float64, false)])); + let vals = Float64Array::from((0..rows).map(|i| (i % 4) as f64).collect::>()); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(vals)]).unwrap(); + let mut buf = Vec::new(); + let mut w = + ArrowWriter::try_new(&mut buf, schema, Some(WriterProperties::builder().build())) + .unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + let meta = + ArrowReaderMetadata::load(&Bytes::from(buf), reader_options(PageIndexPolicy::Skip)) + .unwrap(); + let rgm = meta.metadata().row_group(0); + + let decoded = rows as i64 * 8; + assert!( + rgm.total_byte_size() < decoded / 4, + "fixture not dict-compressed: encoded {} vs decoded {decoded}", + rgm.total_byte_size() + ); + assert!(decoded_estimate_bytes(rgm) >= decoded); + + // Budget = 1/10 of decoded -> ~rows/10 per batch. The old encoded-based + // sizing would have clamped to `requested` (expansion-sized batches). + let eff = group_batch_rows(rgm, 1 << 20, (decoded / 10) as u64); + assert!( + (rows / 20..=rows / 5).contains(&eff), + "eff {eff} not within 2x of rows/10" + ); + assert!( + byte_budget_rows( + rgm.total_byte_size(), + rgm.num_rows(), + 1 << 20, + (decoded / 10) as u64 + ) > 4 * eff + ); + } + + /// M53: the schema-override path (`with_schema_override`) relies on two + /// parquet-rs behaviors — (1) a supplied schema may widen a list column to + /// `LargeList` (i64 offsets) and the decode honors it, and (2) the supplied + /// schema's nested field names must match the parquet-inferred ones + /// (`element`), which is why the override keeps the crate's names and the + /// Python side relabels afterwards (names live in the schema, not the + /// buffers). If a parquet upgrade changes either, this fails loudly. + #[test] + fn schema_override_widens_list_to_large_list() { + use arrow::array::{Array, Float32Array, LargeListArray, ListArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::{DataType, Field, Schema}; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::arrow::ArrowWriter; + use parquet::file::properties::WriterProperties; + + // A list column — the parquet-inferred storage of a + // Ray tensor column whose embedded arrow schema was skipped. + let child = Arc::new(Field::new("element", DataType::Float32, true)); + let values = Float32Array::from((0..40).map(|i| i as f32).collect::>()); + let offsets = OffsetBuffer::from_lengths(std::iter::repeat(4).take(10)); + let list = ListArray::new(child.clone(), offsets, Arc::new(values), None); + let schema = Arc::new(Schema::new(vec![Field::new( + "t", + DataType::List(child.clone()), + false, + )])); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(list)]).unwrap(); + let mut buf = Vec::new(); + let mut w = + ArrowWriter::try_new(&mut buf, schema, Some(WriterProperties::builder().build())) + .unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + let bytes = Bytes::from(buf); + + // Load like the tensor path does: embedded arrow schema skipped. + let meta = ArrowReaderMetadata::load( + &bytes, + reader_options(PageIndexPolicy::Skip).with_skip_arrow_metadata(true), + ) + .unwrap(); + assert_eq!( + meta.schema().field(0).data_type(), + &DataType::List(child.clone()) + ); + + // (1) Supplied storage schema: same layout, i64 offsets, crate's child + // name. Decode must emit LargeList with the values intact. + let big = Schema::new(vec![Field::new( + "t", + DataType::LargeList(child.clone()), + false, + )]); + let meta2 = ArrowReaderMetadata::try_new( + Arc::clone(meta.metadata()), + ArrowReaderOptions::new().with_schema(Arc::new(big)), + ) + .unwrap(); + let reader = ParquetRecordBatchReaderBuilder::new_with_metadata(bytes.clone(), meta2) + .build() + .unwrap(); + let out: Vec = reader.map(|b| b.unwrap()).collect(); + assert_eq!( + out[0].schema().field(0).data_type(), + &DataType::LargeList(child.clone()) + ); + let ll = out[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(ll.len(), 10); + let vals = ll.values().as_any().downcast_ref::().unwrap(); + assert_eq!(vals.len(), 40); + assert_eq!(vals.value(7), 7.0); + + // (2) A child field name differing from the inferred one ("item", the + // pyarrow extension-storage name) is rejected at some stage of the + // decode — the exact stage moved across parquet versions, so accept + // any of try_new / build / first-batch failing. + let renamed = Arc::new(Field::new("item", DataType::Float32, true)); + let bad = Schema::new(vec![Field::new("t", DataType::LargeList(renamed), false)]); + let failed = ArrowReaderMetadata::try_new( + Arc::clone(meta.metadata()), + ArrowReaderOptions::new().with_schema(Arc::new(bad)), + ) + .map(|m| { + ParquetRecordBatchReaderBuilder::new_with_metadata(bytes.clone(), m) + .build() + .map(|r| { + r.collect::>() + .into_iter() + .collect::, _>>() + }) + }); + let ok = match failed { + Err(_) => true, + Ok(Err(_)) => true, + Ok(Ok(Err(_))) => true, + Ok(Ok(Ok(_))) => false, + }; + assert!(ok, "child-name mismatch unexpectedly accepted"); + } + + #[test] + fn oversized_unit_never_drains_the_bucket() { + // T6: a unit >= the budget used to take the whole semaphore, serializing + // fetch behind decode. It must now cap at half so a second unit fits. + assert_eq!(unit_permit_kib(500_000, 64 * 1024), 32 * 1024); + assert_eq!(unit_permit_kib(64 * 1024, 64 * 1024), 32 * 1024); + // A normal-sized unit takes exactly its size. + assert_eq!(unit_permit_kib(16 * 1024, 64 * 1024), 16 * 1024); + // Zero-size unit still needs one permit to be ordered by the bucket. + assert_eq!(unit_permit_kib(0, 64 * 1024), 1); + // Degenerate budget (prefetch disabled) -> strict one-at-a-time, as before. + assert_eq!(unit_permit_kib(500, 1), 1); + } + + /// Fixture for the M20 mis-selection: a TALL row group with one fat string + /// column (1 KiB/row, uncompressed, multi-page) and one small int column. + /// Every projected root exceeds a small column budget, so the old planner + /// column-grouped it (Hstack = whole decoded group retained); the fixed + /// planner must row-window it because the windows actually split. + fn tall_fat_col_fixture() -> ArrowReaderMetadata { + use arrow::array::{Int64Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use parquet::arrow::arrow_writer::ArrowWriter; + use parquet::file::properties::WriterProperties; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("fat", DataType::Utf8, false), + ])); + let n = 2000usize; + let a = Int64Array::from((0..n as i64).collect::>()); + let payload = "x".repeat(1024); + let fat = StringArray::from((0..n).map(|_| payload.clone()).collect::>()); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(fat)]).unwrap(); + let props = WriterProperties::builder() + .set_data_page_row_count_limit(100) + .set_write_batch_size(100) + // Dictionary encoding would collapse the repeated payload to a + // few KiB on disk; plain encoding keeps the column fat COMPRESSED + // too, which is what the byte-denominated window math sees. + .set_dictionary_enabled(false) + .build(); + let mut buf = Vec::new(); + let mut w = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + let buf = Bytes::from(buf); + ArrowReaderMetadata::load(&buf, reader_options(PageIndexPolicy::Required)).unwrap() + } + + #[test] + fn tall_fat_columns_row_window_instead_of_hstack() { + // The M20 regression test: ~2 MiB of 1 KiB rows, column budget 1 byte + // (every root its own group — the strongest Hstack trigger), fetch + // window 1 MiB. Windows split the range, so they must win. + let meta = tall_fat_col_fixture(); + let mask = ProjectionMask::all(); + let roots = vec![0usize, 1]; + let (plans, units) = + plan_s3_units(&meta, &mask, &roots, &[(0, 0, 2000)], 131072, 2 << 20, 1, 1); + assert_eq!(plans.len(), 1); + match plans[0].decode { + RgDecode::Windows(n) => assert!(n > 1, "expected a real split, got {n} window(s)"), + RgDecode::Hstack(_) => panic!("tall fat-column group mis-selected Hstack again"), + } + // Every unit is a row window over ALL projected columns, each its own + // single-unit admission episode. + assert!(units.iter().all(|e| e.len() == 1)); + assert!(units.iter().flatten().all(|u| u.sel.is_some())); + + // Escape hatch: windowing explicitly disabled (fetch_window_mb == 0) + // makes windows inert -> the column-group axis may fire again. + let (plans, _units) = + plan_s3_units(&meta, &mask, &roots, &[(0, 0, 2000)], 131072, 2 << 20, 0, 1); + assert!(matches!(plans[0].decode, RgDecode::Hstack(2))); + } + + #[test] + fn plan_chooses_windows_for_narrow_and_hstack_for_wide() { + let (meta, _buf) = windowed_fixture(); + let roots = vec![0usize, 1]; + let mask = ProjectionMask::all(); + // Generous column budget -> not wide -> row windows (single window here: + // the fixture is tiny, so the byte-budget window covers all rows). + let (plans, units) = plan_s3_units( + &meta, + &mask, + &roots, + &[(0, 0, 1000)], + 131072, + 2 << 20, + 16, + 1 << 30, + ); + assert_eq!(plans.len(), 1); + assert!(matches!(plans[0].decode, RgDecode::Windows(1))); + assert_eq!(units[0][0].sel, Some((0, 1000))); + // 1-byte column budget -> every root its own group -> hstack of 2. + let (plans, units) = plan_s3_units( + &meta, + &mask, + &roots, + &[(0, 0, 1000)], + 131072, + 2 << 20, + 16, + 1, + ); + assert!(matches!(plans[0].decode, RgDecode::Hstack(2))); + // The 2 column-group units form ONE admission episode (co-admitted + // under a single summed permit for the lockstep hstack). + assert_eq!(units.len(), 1); + assert_eq!(units[0].len(), 2); + assert!(units.iter().flatten().all(|u| u.sel.is_none() && u.kib > 0)); + // A K-split style partial sub-range must NEVER column-window (it is a + // tall group split by rows), even under a tiny budget. + let (plans, units) = plan_s3_units( + &meta, + &mask, + &roots, + &[(0, 500, 500)], + 131072, + 2 << 20, + 16, + 1, + ); + assert!(matches!(plans[0].decode, RgDecode::Windows(_))); + assert_eq!(units[0][0].sel, Some((500, 500))); + } + + #[test] + fn unsigned_and_nonint_logical_types_are_not_comparable() { + // Unsigned: a u32 max is stored as the i32 bit pattern -1, so reading + // the stat as signed inverts min/max and would wrongly prune. Reject. + assert!(!is_plain_signed_int(&descr( + PhysicalType::INT32, + Some(LogicalType::integer(32, false)), + ConvertedType::NONE + ))); + assert!(!is_plain_signed_int(&descr( + PhysicalType::INT64, + Some(LogicalType::integer(64, false)), + ConvertedType::NONE + ))); + // Legacy unsigned converted type (logical absent). + assert!(!is_plain_signed_int(&descr( + PhysicalType::INT32, + None, + ConvertedType::UINT_32 + ))); + // Date is INT32-backed but not a plain integer value — reject. + assert!(!is_plain_signed_int(&descr( + PhysicalType::INT32, + Some(LogicalType::Date), + ConvertedType::NONE + ))); + } +} diff --git a/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/src/predicate.rs b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/src/predicate.rs new file mode 100644 index 000000000000..7321d3309a7d --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/src/predicate.rs @@ -0,0 +1,615 @@ +//! Predicate IR for row-group statistics pruning (Track 4, part 1). +//! +//! The Python reader lowers the pushed-down Ray `Expr` predicate into the small +//! JSON IR parsed here (see `_predicate_to_ir` on the Python side), and this +//! module decides which row groups *cannot possibly* contain a matching row and +//! so may be skipped before any data is fetched or decoded. This replaces the +//! PyArrow `fragment.subset(filter=...)` pruning the reader used to depend on, +//! which is what let PyArrow stop opening supported files. +//! +//! SOUNDNESS CONTRACT +//! ------------------ +//! [`can_match`] is *conservative*: it returns `true` (keep the row group) +//! unless the predicate is provably false for **every** row in the group given +//! its column statistics. Every source of uncertainty — a missing column, absent +//! statistics, a cross-type or NaN comparison it can't order, a `NOT`, or any op +//! it doesn't model — resolves to `true`. Over-pruning (dropping a group that +//! *could* have matched) is the only way stats pruning can silently lose rows, so +//! it is made impossible by construction: the worst a bug in here can do is keep +//! a group we could have skipped (a performance miss), never drop a live one. +//! +//! The Python reader additionally re-applies the full predicate post-decode, so +//! row-level correctness never rests on this module at all — this is purely an +//! IO/decode-avoidance optimization on top of an already-correct result. + +use std::cmp::Ordering; +use std::collections::HashMap; + +use serde::Deserialize; + +/// A scalar literal or a statistic value, in the few types Parquet statistics +/// and pushed predicates actually use. Comparisons across numeric types promote +/// to f64; every other cross-type comparison is defined as incomparable +/// (`partial_cmp` -> `None`), which the pruning logic treats as "keep". +#[derive(Debug, Clone, Deserialize)] +#[serde(tag = "vt", content = "v", rename_all = "snake_case")] +pub enum Value { + Int(i64), + Float(f64), + Str(String), + Bool(bool), + Null, +} + +impl Value { + /// Partial order used for pruning. Returns `None` (incomparable) for NaN, + /// null operands, and any cross-type pair that isn't int/float — every such + /// case makes the caller keep the row group. + pub fn partial_cmp(&self, other: &Value) -> Option { + // Integers with |v| above this lose precision when promoted to f64 + // (f64 has a 53-bit mantissa), so a mixed int/float comparison could + // silently flip. Treat those as incomparable (None) instead — the + // caller keeps the row group, which is always sound. + const MAX_SAFE_INT: i64 = 9_007_199_254_740_991; // 2^53 - 1 + const MIN_SAFE_INT: i64 = -9_007_199_254_740_991; + let in_safe_range = |v: i64| (MIN_SAFE_INT..=MAX_SAFE_INT).contains(&v); + match (self, other) { + (Value::Int(a), Value::Int(b)) => Some(a.cmp(b)), + (Value::Float(a), Value::Float(b)) => a.partial_cmp(b), + // Mixed int/float: promote to f64 (e.g. `float_col > 5`), but only + // when the integer is exactly representable in f64. + (Value::Int(a), Value::Float(b)) if in_safe_range(*a) => { + (*a as f64).partial_cmp(b) + } + (Value::Float(a), Value::Int(b)) if in_safe_range(*b) => { + a.partial_cmp(&(*b as f64)) + } + (Value::Int(_), Value::Float(_)) | (Value::Float(_), Value::Int(_)) => None, + (Value::Str(a), Value::Str(b)) => Some(a.cmp(b)), + (Value::Bool(a), Value::Bool(b)) => Some(a.cmp(b)), + _ => None, + } + } + + fn is_null(&self) -> bool { + matches!(self, Value::Null) + } +} + +/// Comparison operator in a `cmp` predicate atom. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CmpOp { + Gt, + Lt, + Ge, + Le, + Eq, + Ne, +} + +/// The predicate IR. `unknown` is the explicit catch-all the Python translator +/// emits for anything it can't lower; it always keeps the row group. +#[derive(Debug, Clone, Deserialize)] +#[serde(tag = "t", rename_all = "snake_case")] +pub enum Pred { + And { + preds: Vec, + }, + Or { + preds: Vec, + }, + // `can_match` treats any negation conservatively (keep) without inspecting + // the child, so the field is unread today; retained for a future sound + // NOT-pushdown and to keep the wire format stable. + Not { + #[allow(dead_code)] + pred: Box, + }, + Cmp { + col: String, + op: CmpOp, + value: Value, + }, + IsNull { + col: String, + }, + IsNotNull { + col: String, + }, + In { + col: String, + values: Vec, + negated: bool, + }, + Unknown, +} + +impl Pred { + /// Parse the IR from the JSON string passed across the FFI boundary. On any + /// parse error we fall back to the always-keep predicate so a malformed IR + /// degrades to "no pruning", never to an error or a wrong result. + pub fn from_json(s: &str) -> Pred { + serde_json::from_str(s).unwrap_or(Pred::Unknown) + } +} + +/// Per-column statistics for one row group, in `Value` terms. `min`/`max` are +/// `None` when the column has no statistics (Parquet may omit them). `null_count` +/// is `None` when unknown. `num_rows` is the row group's row count. +#[derive(Debug, Clone)] +pub struct ColStats { + pub min: Option, + pub max: Option, + pub null_count: Option, + pub num_rows: i64, +} + +impl ColStats { + fn all_null(&self) -> bool { + self.num_rows > 0 && self.null_count == Some(self.num_rows) + } +} + +/// True if the row group described by `stats` *could* contain a row satisfying +/// `pred`. Conservative: any uncertainty returns true. See the module contract. +pub fn can_match(pred: &Pred, stats: &HashMap) -> bool { + match pred { + // A AND B can match only if every conjunct can match: if no row can + // satisfy one conjunct, none can satisfy the whole. + Pred::And { preds } => preds.iter().all(|p| can_match(p, stats)), + // A OR B can match if any disjunct can match. + Pred::Or { preds } => preds.iter().any(|p| can_match(p, stats)), + // Negation over ranges isn't soundly prunable from min/max alone; keep. + Pred::Not { .. } => true, + Pred::Unknown => true, + Pred::Cmp { col, op, value } => match stats.get(col) { + None => true, // column absent from this group's stats -> keep + Some(cs) => cmp_can_match(cs, *op, value), + }, + Pred::IsNull { col } => match stats.get(col) { + None => true, + // Keep unless we know there are zero nulls. + Some(cs) => cs.null_count != Some(0), + }, + Pred::IsNotNull { col } => match stats.get(col) { + None => true, + // Keep unless every row is null. + Some(cs) => !cs.all_null(), + }, + Pred::In { + col, + values, + negated, + } => { + if *negated { + return true; // NOT IN can't be pruned from min/max; keep. + } + match stats.get(col) { + None => true, + Some(cs) => { + if cs.all_null() { + return false; // all rows null -> none in the set + } + // Keep if any listed value could fall within [min, max]. + values.iter().any(|v| value_in_range(cs, v)) + } + } + } + } +} + +/// Whether `v` is possibly within `[min, max]` (used for IN and EQ). Unknown +/// bounds or incomparable values -> possible (keep). +fn value_in_range(cs: &ColStats, v: &Value) -> bool { + if v.is_null() { + return true; // null-in-set semantics are subtle; stay conservative. + } + let below_min = match &cs.min { + Some(mn) => matches!(v.partial_cmp(mn), Some(Ordering::Less)), + None => false, + }; + let above_max = match &cs.max { + Some(mx) => matches!(v.partial_cmp(mx), Some(Ordering::Greater)), + None => false, + }; + !(below_min || above_max) +} + +/// Row-group verdict for a single `col OP value` atom. Returns true to keep. +fn cmp_can_match(cs: &ColStats, op: CmpOp, v: &Value) -> bool { + // A comparison against a null literal is null (=> false) for every row; + // but this shape is unusual, so keep rather than reason about it. + if v.is_null() { + return true; + } + // Every row is null => every comparison is null => false for all => prune. + if cs.all_null() { + return false; + } + match op { + // Keep iff some value can exceed / reach the bound. + CmpOp::Gt => keep_if(&cs.max, v, |o| o == Ordering::Greater), + CmpOp::Ge => keep_if(&cs.max, v, |o| o != Ordering::Less), + CmpOp::Lt => keep_if(&cs.min, v, |o| o == Ordering::Less), + CmpOp::Le => keep_if(&cs.min, v, |o| o != Ordering::Greater), + // v in [min, max] is a necessary condition for equality to be possible. + CmpOp::Eq => value_in_range(cs, v), + // col != v is false for all only if the column is the constant v with no + // nulls; that needs min==max==v && null_count==0. Otherwise keep. + CmpOp::Ne => { + let is_constant_v = matches!( + (&cs.min, &cs.max), + (Some(mn), Some(mx)) + if matches!(mn.partial_cmp(v), Some(Ordering::Equal)) + && matches!(mx.partial_cmp(v), Some(Ordering::Equal)) + ); + !(is_constant_v && cs.null_count == Some(0)) + } + } +} + +/// Keep the group iff `bound` is known and `pred(bound.cmp(v))` holds. An unknown +/// bound or an incomparable pair (`partial_cmp` -> None, e.g. NaN) keeps it. +fn keep_if(bound: &Option, v: &Value, pred: impl Fn(Ordering) -> bool) -> bool { + match bound { + None => true, + Some(b) => match b.partial_cmp(v) { + Some(ord) => pred(ord), + None => true, + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn cs(min: i64, max: i64, nulls: i64, rows: i64) -> ColStats { + ColStats { + min: Some(Value::Int(min)), + max: Some(Value::Int(max)), + null_count: Some(nulls), + num_rows: rows, + } + } + + fn stats(pairs: &[(&str, ColStats)]) -> HashMap { + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.clone())) + .collect() + } + + fn cmp(col: &str, op: CmpOp, v: i64) -> Pred { + Pred::Cmp { + col: col.into(), + op, + value: Value::Int(v), + } + } + + // ---- the sorted-id / row-group example from the Python pruning test ---- + #[test] + fn ge_prunes_groups_below_threshold() { + // Four row groups of a sorted id column: [0,999] [1000,1999] ... + let groups = [(0, 999), (1000, 1999), (2000, 2999), (3000, 3999)]; + let pred = cmp("id", CmpOp::Ge, 3000); + let kept: Vec = groups + .iter() + .enumerate() + .filter(|(_, (lo, hi))| can_match(&pred, &stats(&[("id", cs(*lo, *hi, 0, 1000))]))) + .map(|(i, _)| i) + .collect(); + assert_eq!( + kept, + vec![3], + "only the [3000,3999] group can satisfy id>=3000" + ); + } + + #[test] + fn ge_prunes_everything_when_out_of_range() { + let pred = cmp("id", CmpOp::Ge, 1_000_000_000); + assert!(!can_match(&pred, &stats(&[("id", cs(0, 3999, 0, 4000))]))); + } + + // ---- per-op boundary behavior ---- + #[test] + fn gt_boundary() { + // max == v: no row is strictly greater -> prune. + assert!(!can_match( + &cmp("x", CmpOp::Gt, 10), + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + assert!(can_match( + &cmp("x", CmpOp::Gt, 9), + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + } + + #[test] + fn lt_le_use_min() { + assert!(!can_match( + &cmp("x", CmpOp::Lt, 0), + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + assert!(can_match( + &cmp("x", CmpOp::Lt, 1), + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + assert!(!can_match( + &cmp("x", CmpOp::Le, -1), + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + assert!(can_match( + &cmp("x", CmpOp::Le, 0), + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + } + + #[test] + fn eq_outside_range_prunes() { + assert!(!can_match( + &cmp("x", CmpOp::Eq, 11), + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + assert!(!can_match( + &cmp("x", CmpOp::Eq, -1), + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + assert!(can_match( + &cmp("x", CmpOp::Eq, 5), + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + } + + #[test] + fn ne_only_prunes_constant_column() { + // min==max==v, no nulls -> col != v false for all -> prune. + assert!(!can_match( + &cmp("x", CmpOp::Ne, 7), + &stats(&[("x", cs(7, 7, 0, 3))]) + )); + // a null present -> keep (null != v is null, but rows differ anyway keep) + assert!(can_match( + &cmp("x", CmpOp::Ne, 7), + &stats(&[("x", cs(7, 7, 1, 3))]) + )); + // range wider than {v} -> keep + assert!(can_match( + &cmp("x", CmpOp::Ne, 7), + &stats(&[("x", cs(0, 10, 0, 3))]) + )); + } + + // ---- conservative fallbacks: every uncertainty keeps ---- + #[test] + fn missing_column_keeps() { + assert!(can_match( + &cmp("absent", CmpOp::Gt, 100), + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + } + + #[test] + fn unknown_bounds_keep() { + let s = stats(&[( + "x", + ColStats { + min: None, + max: None, + null_count: None, + num_rows: 5, + }, + )]); + assert!(can_match(&cmp("x", CmpOp::Gt, 10_000), &s)); + assert!(can_match(&cmp("x", CmpOp::Eq, 10_000), &s)); + } + + #[test] + fn cross_type_compare_keeps() { + // string literal against an int column -> incomparable -> keep. + let pred = Pred::Cmp { + col: "x".into(), + op: CmpOp::Gt, + value: Value::Str("abc".into()), + }; + assert!(can_match(&pred, &stats(&[("x", cs(0, 10, 0, 5))]))); + } + + #[test] + fn nan_bounds_keep() { + let s = stats(&[( + "f", + ColStats { + min: Some(Value::Float(f64::NAN)), + max: Some(Value::Float(f64::NAN)), + null_count: Some(0), + num_rows: 5, + }, + )]); + let pred = Pred::Cmp { + col: "f".into(), + op: CmpOp::Gt, + value: Value::Float(1.0), + }; + assert!(can_match(&pred, &s)); + } + + #[test] + fn int_float_promotion() { + // float column, int literal. + let s = stats(&[( + "f", + ColStats { + min: Some(Value::Float(0.0)), + max: Some(Value::Float(2.5)), + null_count: Some(0), + num_rows: 5, + }, + )]); + assert!(!can_match( + &Pred::Cmp { + col: "f".into(), + op: CmpOp::Gt, + value: Value::Int(3) + }, + &s + )); + assert!(can_match( + &Pred::Cmp { + col: "f".into(), + op: CmpOp::Gt, + value: Value::Int(2) + }, + &s + )); + } + + // ---- null-count driven ---- + #[test] + fn all_null_group_prunes_comparisons_and_in() { + let s = stats(&[("x", cs(0, 0, 5, 5))]); // all 5 rows null + assert!(!can_match(&cmp("x", CmpOp::Gt, -100), &s)); + assert!(!can_match(&cmp("x", CmpOp::Eq, 0), &s)); + assert!(!can_match( + &Pred::In { + col: "x".into(), + values: vec![Value::Int(0)], + negated: false + }, + &s + )); + } + + #[test] + fn is_null_and_is_not_null() { + // no nulls + assert!(!can_match( + &Pred::IsNull { col: "x".into() }, + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + assert!(can_match( + &Pred::IsNotNull { col: "x".into() }, + &stats(&[("x", cs(0, 10, 0, 5))]) + )); + // all null + assert!(can_match( + &Pred::IsNull { col: "x".into() }, + &stats(&[("x", cs(0, 0, 5, 5))]) + )); + assert!(!can_match( + &Pred::IsNotNull { col: "x".into() }, + &stats(&[("x", cs(0, 0, 5, 5))]) + )); + // some null + assert!(can_match( + &Pred::IsNull { col: "x".into() }, + &stats(&[("x", cs(0, 10, 2, 5))]) + )); + assert!(can_match( + &Pred::IsNotNull { col: "x".into() }, + &stats(&[("x", cs(0, 10, 2, 5))]) + )); + } + + // ---- IN ---- + #[test] + fn in_prunes_when_no_value_in_range() { + let s = stats(&[("x", cs(0, 10, 0, 5))]); + assert!(!can_match( + &Pred::In { + col: "x".into(), + values: vec![Value::Int(20), Value::Int(30)], + negated: false + }, + &s + )); + assert!(can_match( + &Pred::In { + col: "x".into(), + values: vec![Value::Int(20), Value::Int(5)], + negated: false + }, + &s + )); + // NOT IN never prunes. + assert!(can_match( + &Pred::In { + col: "x".into(), + values: vec![Value::Int(20)], + negated: true + }, + &s + )); + } + + // ---- boolean composition ---- + #[test] + fn and_prunes_if_any_conjunct_prunes() { + let s = stats(&[("x", cs(0, 10, 0, 5)), ("y", cs(0, 10, 0, 5))]); + // x > 5 (possible) AND y > 100 (impossible) -> prune + let p = Pred::And { + preds: vec![cmp("x", CmpOp::Gt, 5), cmp("y", CmpOp::Gt, 100)], + }; + assert!(!can_match(&p, &s)); + // x > 5 AND y > 5 -> both possible -> keep + let p2 = Pred::And { + preds: vec![cmp("x", CmpOp::Gt, 5), cmp("y", CmpOp::Gt, 5)], + }; + assert!(can_match(&p2, &s)); + } + + #[test] + fn or_keeps_if_any_disjunct_possible() { + let s = stats(&[("x", cs(0, 10, 0, 5))]); + // x > 100 (impossible) OR x < 5 (possible) -> keep + let p = Pred::Or { + preds: vec![cmp("x", CmpOp::Gt, 100), cmp("x", CmpOp::Lt, 5)], + }; + assert!(can_match(&p, &s)); + // x > 100 OR x > 200 -> both impossible -> prune + let p2 = Pred::Or { + preds: vec![cmp("x", CmpOp::Gt, 100), cmp("x", CmpOp::Gt, 200)], + }; + assert!(!can_match(&p2, &s)); + } + + #[test] + fn not_and_unknown_keep() { + let s = stats(&[("x", cs(0, 10, 0, 5))]); + assert!(can_match( + &Pred::Not { + pred: Box::new(cmp("x", CmpOp::Gt, 100)) + }, + &s + )); + assert!(can_match(&Pred::Unknown, &s)); + } + + // ---- JSON round-trip of the wire format ---- + #[test] + fn parses_wire_json() { + let j = r#"{"t":"and","preds":[ + {"t":"cmp","col":"id","op":"ge","value":{"vt":"int","v":3000}}, + {"t":"cmp","col":"x","op":"lt","value":{"vt":"float","v":1.5}}, + {"t":"is_not_null","col":"id"}, + {"t":"in","col":"g","values":[{"vt":"str","v":"a"}],"negated":false} + ]}"#; + let p = Pred::from_json(j); + // sanity: with id in [0,999] the id>=3000 conjunct prunes. + let s = stats(&[ + ("id", cs(0, 999, 0, 1000)), + ("x", cs(0, 0, 0, 1000)), + ("g", cs(0, 0, 0, 1000)), + ]); + assert!(!can_match(&p, &s)); + } + + #[test] + fn malformed_json_becomes_keep_all() { + assert!(matches!(Pred::from_json("not json"), Pred::Unknown)); + assert!(matches!(Pred::from_json(r#"{"t":"bogus"}"#), Pred::Unknown)); + } +} diff --git a/python/ray/data/_internal/datasource_v2/native_metadata.py b/python/ray/data/_internal/datasource_v2/native_metadata.py new file mode 100644 index 000000000000..b7ee6cfc4d17 --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/native_metadata.py @@ -0,0 +1,200 @@ +"""Shared helpers for reading Parquet footers through the arrow-rs native crate. + +Used by the read path (``ArrowRsParquetFileReader``), which reads each file's +footer via the crate so a supported file is opened by arrow-rs *end to end* — no +PyArrow footer read for Local/S3 files when +``DataContext.use_arrow_rs_parquet_reader`` is on. The logic is factored out here +(filesystem eligibility, S3 config bridging, and the actual ``read_metadata`` +call) rather than inlined in the reader because it is also what a listing-stage +caller would need. + +**There is currently only one caller.** This module's docstring previously named +``ParquetFileChunker`` as a second one; that was never true on any branch. The +listing stage does now read footers — the footer-chunking path reads every file's +footer on a pool of ``FooterReader`` actors to prune and bin-pack row groups — but +it does so through PyArrow, not through this module. Wiring it to the crate is a +separate, unstarted piece of work; until then, a native read still pays one footer +read here on top of the one ``ListFiles`` already did. +""" + +import os +import threading +from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple + +if TYPE_CHECKING: + from pyarrow.fs import FileSystem, S3FileSystem + + +# --------------------------------------------------------------------------- # +# Process-level native S3 client cache (findings M96/M97) +# --------------------------------------------------------------------------- # +# Client construction is the expensive part of a native S3 read's setup: a +# fresh `object_store` client is a fresh connection pool, so the first requests +# pay DNS + full TLS handshakes with no session reuse. When the planner emits +# single-row-group tasks (64 MiB bins, the release `read_large_parquet` +# regime), a per-TASK client meant ~5.8k cold client builds per job for 104 +# distinct files — measured as read workers spending ~2/3 of their wall time +# blocked in `open_file`'s metadata fetch while the box sat idle (M97). +# Caching per process gives the client the same lifetime pyarrow's serialized +# `S3FileSystem` client already has in a reused Ray worker. Staleness is +# handled by the KEY, not a TTL: the key includes the full connection config, +# credentials included, so rotated credentials miss the cache and build a +# fresh client while the stale entry ages out of the size-capped table. +_S3_STORE_CACHE: Dict[Tuple, Any] = {} +_S3_STORE_CACHE_LOCK = threading.Lock() +# Real jobs hold one or two entries (one per bucket x config); the cap only +# bounds a long-lived worker against credential-rotation churn. +_S3_STORE_CACHE_MAX_ENTRIES = 8 + + +def native_metadata_supported_filesystem( + filesystem: Optional["FileSystem"], +) -> bool: + """Whether the crate's ``read_metadata`` can read footers on this filesystem. + + The crate's ``object_store`` backend is compiled with the local and AWS + features only, so it can footer-read local files and S3 objects. ``None`` + means the caller will use the default local filesystem. Any other + filesystem (GCS, Azure, HDFS, an fsspec wrapper) must stay on PyArrow. + """ + from pyarrow.fs import LocalFileSystem, S3FileSystem + + return filesystem is None or isinstance(filesystem, (LocalFileSystem, S3FileSystem)) + + +def s3_config(fs: "S3FileSystem") -> dict: + """Recover the full S3 connection config from a pyarrow ``S3FileSystem`` so + the native crate connects *identically* — same endpoint, credentials, region, + addressing style — instead of rebuilding a default client from the ambient env + (which would silently ignore an explicit endpoint override or static creds and + break MinIO / moto / custom-endpoint / credentialed buckets). + + pyarrow round-trips the whole config through ``__reduce__`` (verified to include + ``secret_key``/``session_token``/``endpoint_override``/``scheme`` across the + pyarrow versions Ray supports), so that is the source of truth. Empty strings + (pyarrow's "unset" sentinel) are normalized to ``None``. + """ + try: + opts = fs.__reduce__()[1][0] + except Exception: + opts = {} + + def _val(key): + v = opts.get(key) + return v if v else None + + endpoint = _val("endpoint_override") + # object_store refuses plain-HTTP endpoints unless explicitly allowed + # (moto / MinIO are http). pyarrow's `scheme` defaults to "https" even when + # the endpoint override is an http:// URL, so trust the endpoint URL first. + allow_http = (str(endpoint).startswith("http://")) or opts.get("scheme") == "http" + + return { + "region": _val("region") or "us-east-1", + "anonymous": bool(opts.get("anonymous", False)), + "endpoint": endpoint, + "access_key_id": _val("access_key"), + "secret_access_key": _val("secret_key"), + "session_token": _val("session_token"), + "allow_http": allow_http, + "virtual_hosted_style": bool(opts.get("force_virtual_addressing", False)), + } + + +def split_s3_path(path: str) -> tuple: + """Split a pyarrow-style S3 path into ``(bucket, key)``. + + pyarrow filesystem paths are normally scheme-less (``bucket/key``), but a + leading ``s3://`` is stripped defensively so it can never be split into a + bogus ``s3:`` bucket. + """ + if path.startswith("s3://"): + path = path[len("s3://") :] + bucket, _, key = path.partition("/") + return bucket, key + + +def connect_native_s3(bucket: str, filesystem: "S3FileSystem"): + """Build or reuse the crate's per-bucket S3 client (``NativeS3Store``), + configured identically to the pyarrow ``S3FileSystem`` (see + :func:`s3_config`). + + Fresh-client construction is the expensive per-call setup the original + ``read_metadata_s3`` / ``read_row_groups_s3`` entry points paid on *every* + call (new HTTP client, no TLS session reuse — findings T10). Clients are + cached per process, keyed by (bucket, full connection config including + credentials) — see the cache comment above for why per-task scoping was + not enough (M97) and why the key handles credential rotation. The store is + an immutable ``Arc`` handle, safe to share across tasks. + ``RAY_DATA_ARROW_RS_S3_CLIENT_CACHE=0`` restores the uncached behavior. + """ + import ray_data_arrow_rs + + cfg = s3_config(filesystem) + + def _build(): + return ray_data_arrow_rs.connect_s3( + bucket, + cfg["region"], + cfg["anonymous"], + endpoint=cfg["endpoint"], + access_key_id=cfg["access_key_id"], + secret_access_key=cfg["secret_access_key"], + session_token=cfg["session_token"], + allow_http=cfg["allow_http"], + virtual_hosted_style=cfg["virtual_hosted_style"], + ) + + if os.environ.get("RAY_DATA_ARROW_RS_S3_CLIENT_CACHE", "1") == "0": + return _build() + + key = (bucket, tuple(sorted(cfg.items(), key=lambda kv: kv[0]))) + with _S3_STORE_CACHE_LOCK: + store = _S3_STORE_CACHE.get(key) + if store is not None: + return store + store = _build() # outside the lock: connect does network setup + with _S3_STORE_CACHE_LOCK: + existing = _S3_STORE_CACHE.get(key) + if existing is not None: + return existing # benign build race: reuse the winner + while len(_S3_STORE_CACHE) >= _S3_STORE_CACHE_MAX_ENTRIES: + _S3_STORE_CACHE.pop(next(iter(_S3_STORE_CACHE))) + _S3_STORE_CACHE[key] = store + return store + + +def read_native_metadata(path: str, filesystem: Optional["FileSystem"]): + """Read one Parquet file's footer via the native crate. + + Returns the crate's ``ParquetFileMetadata`` pyclass (exposing ``num_rows``, + ``num_row_groups``, ``row_group_num_rows``, ``row_group_byte_sizes``, + ``row_group_compressed_sizes``, and ``__arrow_c_schema__``). Raises on a + missing extension or any footer-read failure — the caller decides whether to + fall back. ``filesystem`` must already be native-eligible (see + :func:`native_metadata_supported_filesystem`). + """ + import ray_data_arrow_rs + from pyarrow.fs import S3FileSystem + + if isinstance(filesystem, S3FileSystem): + # pyarrow filesystem paths are normally scheme-less ("bucket/key"), but + # strip a leading "s3://" defensively so we never split it into a bogus + # "s3:" bucket. + if path.startswith("s3://"): + path = path[len("s3://") :] + bucket, _, key = path.partition("/") + cfg = s3_config(filesystem) + return ray_data_arrow_rs.read_metadata_s3( + bucket, + key, + cfg["region"], + cfg["anonymous"], + endpoint=cfg["endpoint"], + access_key_id=cfg["access_key_id"], + secret_access_key=cfg["secret_access_key"], + session_token=cfg["session_token"], + allow_http=cfg["allow_http"], + virtual_hosted_style=cfg["virtual_hosted_style"], + ) + return ray_data_arrow_rs.read_metadata(path) diff --git a/python/ray/data/_internal/datasource_v2/parquet_datasource_v2.py b/python/ray/data/_internal/datasource_v2/parquet_datasource_v2.py index d064e8401ea2..2abc8c5851ba 100644 --- a/python/ray/data/_internal/datasource_v2/parquet_datasource_v2.py +++ b/python/ray/data/_internal/datasource_v2/parquet_datasource_v2.py @@ -14,22 +14,16 @@ import pyarrow as pa from typing_extensions import override +from ray._common.utils import env_bool, env_integer from ray.data._internal.datasource.parquet_datasource import ( ParquetDatasource, check_for_legacy_tensor_type, ) -from ray.data._internal.datasource_v2.chunkers.file_chunker import ( - FileChunker, - ParquetFileChunker, -) from ray.data._internal.datasource_v2.datasource_v2 import ( DatasourceCategory, DataSourceV2, ) -from ray.data._internal.datasource_v2.listing.file_indexer import ( - FileIndexer, - NonSamplingFileIndexer, -) +from ray.data._internal.datasource_v2.listing.file_indexer import FileIndexer from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest from ray.data._internal.datasource_v2.readers.file_reader import ( INCLUDE_PATHS_COLUMN_NAME, @@ -80,7 +74,6 @@ def __init__( arrow_parquet_args: Optional[dict] = None, schema: Optional[pa.Schema] = None, parquet_format_kwargs: Optional[dict] = None, - file_chunker: Optional[FileChunker] = None, ): super().__init__(name="ParquetV2", category=DatasourceCategory.FILE_BASED) # Capture the ``local://`` check against the *original* paths; @@ -111,14 +104,6 @@ def __init__( # footers, and the scanner pins it on the pyarrow dataset so files # are cast to these types at scan time. self._user_schema = schema - # Chunker that splits each listed Parquet file into one or more - # row-group-aligned read units. Defaults to ``ParquetFileChunker`` - # (1 GiB target chunk size, or whatever ``DataContext`` configures). - # Callers can inject an alternative for tests or shuffle-aware - # planning code that wants whole-file reads. - self._file_chunker: FileChunker = ( - file_chunker if file_chunker is not None else ParquetFileChunker() - ) @property def paths(self) -> List[str]: @@ -149,9 +134,18 @@ def shuffle(self) -> Optional[Union[Literal["files"], "FileShuffleConfig"]]: return self._shuffle def _get_file_indexer(self) -> FileIndexer: - return NonSamplingFileIndexer( + # Parquet V2 reads always use the footer-based indexer: it reads each + # file's footer on a ``FooterReader`` actor pool and packs row groups + # with the online bin packer (predicate/limit push-down + row-group + # skipping), producing accurate row-group-aligned read units. + from ray.data._internal.datasource_v2.listing.footer_file_indexer import ( + FooterFileIndexer, + ) + + return FooterFileIndexer( ignore_missing_paths=self._ignore_missing_paths, - file_chunker=self._file_chunker, + coalesce_bytes=env_integer("RAY_DATA_PARQUET_FOOTER_COALESCE_BYTES", 0), + split_coalesced=env_bool("RAY_DATA_PARQUET_FOOTER_SPLIT_COALESCED", False), ) def get_size_estimator(self) -> ParquetInMemorySizeEstimator: diff --git a/python/ray/data/_internal/datasource_v2/partitioners/online_bin_packer.py b/python/ray/data/_internal/datasource_v2/partitioners/online_bin_packer.py new file mode 100644 index 000000000000..49b08bc5c762 --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/partitioners/online_bin_packer.py @@ -0,0 +1,299 @@ +from __future__ import annotations + +import bisect +from collections import defaultdict, deque +from dataclasses import dataclass, field +from typing import Deque, List, Optional, Tuple, cast + +from ray.data._internal.datasource_v2.chunkers.file_chunker import ( + ChunkMetadata, + ParquetRowGroupChunkMetadata, + create_chunk_metadata, +) +from ray.data._internal.datasource_v2.chunkers.parquet_footer_types import ( + Bin, + BinItem, + FileChunks, +) +from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest + + +@dataclass +class _OpenBin: + items: List[BinItem] = field(default_factory=list) + used_bytes: int = 0 + + def add(self, item: BinItem) -> None: + self.items.append(item) + self.used_bytes += item.uncompressed_size + + def seal(self) -> Bin: + return Bin(tuple(self.items), self.used_bytes) + + +def _prefix_sums(unit_sizes: List[int]) -> List[int]: + # prefix[i] == sum of the first i unit sizes (prefix[0] == 0). + prefix = [0] + for size in unit_sizes: + prefix.append(prefix[-1] + size) + return prefix + + +def _largest_prefix_fit(prefix: List[int], start: int, cap_left: int) -> int: + # Largest end (exclusive), end >= start, such that the row groups [start, end) + # sum to <= cap_left. prefix[i] is the cumulative size of the first i row + # groups, so sum(sizes[start:end]) == prefix[end] - prefix[start]. Binary + # search for the largest end with prefix[end] <= cap_left + prefix[start]. + # Returns ``start`` when not even one row group fits (caller treats that as + # "nothing fits here"). + end = bisect.bisect_right(prefix, cap_left + prefix[start]) - 1 + return max(end, start) + + +def _slice_bin_item(item: BinItem, a: int, b: int) -> BinItem: + # Row groups [a, b) of a coalesced item as a new contiguous BinItem. Uses the + # exact per-RG sizes/rows so num_rows stays an exact survivor count for the + # limit push-down, and rg_idx shifts by ``a`` because the run is contiguous. + sizes = item.rg_sizes[a:b] + rows = item.rg_rows[a:b] + count = b - a + return BinItem( + path=item.path, + rg_idx=item.rg_idx + a, + uncompressed_size=sum(sizes), + num_rows=sum(rows), + fully_matched=item.fully_matched, + rg_count=count, + rg_sizes=sizes if count > 1 else (), + rg_rows=rows if count > 1 else (), + ) + + +def _subitem(item: BinItem, num_units: int, start: int, end: int) -> BinItem: + # The item covering units [start, end). When that is the whole item, return it + # unchanged (so a non-split item keeps its original rg_count/rg_sizes); + # otherwise carve out the row-group range -- only reached for splittable runs. + if start == 0 and end == num_units: + return item + return _slice_bin_item(item, start, end) + + +def _best_open_bin( + bins: List[_OpenBin], prefix: List[int], start: int, cap: int +) -> Tuple[Optional[_OpenBin], int]: + # Among open bins, the one that swallows the largest prefix of units[start:] + # with the least leftover space (best fit). Returns (bin, end); (None, start) + # if no open bin can take even one unit. + best: Optional[_OpenBin] = None + best_end, best_gap = start, 0 + for b in bins: + room = cap - b.used_bytes + end = _largest_prefix_fit(prefix, start, room) + if end > start: + gap = room - (prefix[end] - prefix[start]) + if best is None or gap < best_gap: + best, best_end, best_gap = b, end, gap + return best, best_end + + +class OnlineBinPacker: + """Streaming coloured bin packer over row-group chunks. + + Feed ``FileChunks`` via :meth:`add_file_chunks`; drain sealed bins (as + :class:`FileManifest` blocks) via :meth:`has_partition` / :meth:`next_partition` + as they become available; call :meth:`finalize` once all chunks are added to + flush the still-open bins. + """ + + def __init__( + self, + max_bin_bytes: int, + *, + max_shared_open_bins: int = 16, + split_coalesced: bool = False, + ): + # ``max_bin_bytes`` doubles as the "colour turns heavy" isolate threshold. + self._cap = max_bin_bytes + self._max_shared_open_bins = max_shared_open_bins + # When True, a coalesced item (rg_count > 1) that does not fit whole is + # split at physical-row-group boundaries to fill residual bin space + # instead of opening a fresh bin. Single row groups stay atomic, so with + # coalescing off (every rg_count == 1) this is a no-op and the packer + # behaves exactly as when the flag is False. + self._split_coalesced = split_coalesced + + self._seen_bytes_by_path: dict = {} # running w(c) per colour + self._shared_bins: List[_OpenBin] = [] # non-isolated bins (mixed colours) + self._heavy_path: Optional[str] = None # current heavy colour + self._heavy_bin: Optional[_OpenBin] = None # its open monochromatic bin + self._output: Deque[Bin] = deque() # sealed bins awaiting drain + + # === Feeding === + + def add_file_chunks(self, file_chunks: FileChunks) -> None: + path = file_chunks.path + for row_group in file_chunks.row_groups: + self._place( + BinItem( + path=path, + rg_idx=row_group.rg_idx, + uncompressed_size=row_group.uncompressed_size, + num_rows=row_group.num_rows, + fully_matched=row_group.fully_matched, + rg_count=row_group.rg_count, + rg_sizes=row_group.rg_sizes, + rg_rows=row_group.rg_rows, + ) + ) + + def _units(self, item: BinItem) -> List[int]: + # The row-group boundaries an item may be cut between, as unit sizes. A + # splittable coalesced run (split_coalesced and rg_count > 1) yields one + # unit per physical row group; anything else yields a single indivisible + # unit (the whole item). Placement only ever cuts at unit boundaries. + if self._split_coalesced and item.rg_count > 1: + return list(item.rg_sizes) + return [item.uncompressed_size] + + def _place(self, item: BinItem) -> None: + item_bytes = item.uncompressed_size + seen_bytes = self._seen_bytes_by_path.get(item.path, 0) + self._seen_bytes_by_path[item.path] = seen_bytes + item_bytes + + if item_bytes > self._cap and len(self._units(item)) == 1: + # Relaxation: an indivisible chunk bigger than a whole bin gets its own + # bin. A splittable oversized run instead falls through and is cut into + # bin-sized pieces by the placers. + self._output.append(Bin((item,), item_bytes)) + elif seen_bytes < self._cap: + self._place_light(item) + else: + self._place_heavy(item) + + def _seal_if_full(self, bin_: _OpenBin) -> None: + # A shared bin at (or over) cap can never take another positive-size item: + # ``_best_open_bin`` gives it end == start and the whole-item fast path + # fails its ``used_bytes + total <= cap`` test. Leaving it in the pool just + # burns one of the ``_max_shared_open_bins`` slots, so seal and evict it. + if bin_.used_bytes >= self._cap: + self._shared_bins.remove(bin_) + self._output.append(bin_.seal()) + + def _place_light(self, item: BinItem) -> None: + # LIGHT colour -> shared First-Fit bins. First try to place the WHOLE item + # in the first bin it fits. If it fits nowhere, cut it at unit boundaries + # and best-fit each piece into the tightest open bin, opening a fresh bin + # only when no open bin can take even one unit. (With splitting off the + # item is a single unit, so this reduces to the original First Fit.) + cap = self._cap + units = self._units(item) + prefix = _prefix_sums(units) + total = prefix[-1] + target = next( + (b for b in self._shared_bins if b.used_bytes + total <= cap), None + ) + if target is not None: + target.add(item) + self._seal_if_full(target) + return + start = 0 + while start < len(units): + target, end = _best_open_bin(self._shared_bins, prefix, start, cap) + if target is None: + if len(self._shared_bins) >= self._max_shared_open_bins: + fullest = max(self._shared_bins, key=lambda b: b.used_bytes) + self._shared_bins.remove(fullest) + self._output.append(fullest.seal()) + target = _OpenBin() + self._shared_bins.append(target) + # A lone unit larger than a whole bin gets its own (over-sized) bin. + end = max(_largest_prefix_fit(prefix, start, cap), start + 1) + target.add(_subitem(item, len(units), start, end)) + self._seal_if_full(target) + start = end + + def _place_heavy(self, item: BinItem) -> None: + # HEAVY colour -> dedicated monochromatic bins. Fill the open bin at a unit + # boundary, seal it once full, and carry any remnant into the next bin. + # (With splitting off the item is a single unit, so this reduces to the + # original Next Fit.) + cap = self._cap + if self._heavy_path != item.path or self._heavy_bin is None: + if self._heavy_bin is not None: + self._output.append(self._heavy_bin.seal()) + self._heavy_path = item.path + self._heavy_bin = _OpenBin() + heavy_bin = self._heavy_bin + units = self._units(item) + prefix = _prefix_sums(units) + start = 0 + while start < len(units): + end = _largest_prefix_fit(prefix, start, cap - heavy_bin.used_bytes) + if end == start: # nothing fits the open bin + if heavy_bin.items: # seal it and retry on a fresh bin + self._output.append(heavy_bin.seal()) + heavy_bin = _OpenBin() + self._heavy_bin = heavy_bin + continue + end = start + 1 # empty bin, lone unit > cap -> relaxation + heavy_bin.add(_subitem(item, len(units), start, end)) + start = end + if start < len(units): # remnant remains -> bin is full, seal + self._output.append(heavy_bin.seal()) + heavy_bin = _OpenBin() + self._heavy_bin = heavy_bin + + # === Draining === + + def has_partition(self) -> bool: + return len(self._output) > 0 + + def next_partition(self) -> FileManifest: + return self._bin_to_manifest(self._output.popleft()) + + def finalize(self) -> None: + # Flush everything still open. + if self._heavy_bin is not None and self._heavy_bin.items: + self._output.append(self._heavy_bin.seal()) + self._heavy_bin = None + self._heavy_path = None + for open_bin in self._shared_bins: + if open_bin.items: + self._output.append(open_bin.seal()) + self._shared_bins = [] + + @staticmethod + def _bin_to_manifest(bin_: Bin) -> FileManifest: + # One manifest row per distinct file in the bin. A file's (possibly-split) + # items cover disjoint contiguous runs, so union their physical row-group + # ids into the read unit for that file. + ids_by_path: defaultdict = defaultdict(list) + rows_by_path: defaultdict = defaultdict(int) + size_by_path: defaultdict = defaultdict(int) + for item in bin_.items: + ids_by_path[item.path].extend( + range(item.rg_idx, item.rg_idx + item.rg_count) + ) + rows_by_path[item.path] += item.num_rows + size_by_path[item.path] += item.uncompressed_size + + paths: List[str] = [] + sizes: List[int] = [] + chunk_metadatas: List[ParquetRowGroupChunkMetadata] = [] + for path, ids in ids_by_path.items(): + paths.append(path) + sizes.append(size_by_path[path]) + chunk_metadatas.append( + create_chunk_metadata( + ParquetRowGroupChunkMetadata, + row_group_ids=tuple(sorted(ids)), + num_rows=rows_by_path[path], + uncompressed_size=size_by_path[path], + ) + ) + # TypedDict invariance: ``ParquetRowGroupChunkMetadata`` has extra keys + # beyond the empty ``ChunkMetadata`` base, so the concrete list is not + # assignable to ``List[Optional[ChunkMetadata]]`` without a cast. + return FileManifest.construct_manifest( + paths, sizes, cast(List[Optional[ChunkMetadata]], chunk_metadatas) + ) diff --git a/python/ray/data/_internal/datasource_v2/readers/arrow_rs_parquet_file_reader.py b/python/ray/data/_internal/datasource_v2/readers/arrow_rs_parquet_file_reader.py new file mode 100644 index 000000000000..703a681ff7a0 --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/readers/arrow_rs_parquet_file_reader.py @@ -0,0 +1,2404 @@ +"""Experimental arrow-rs Parquet reader (prototype). + +Subclasses :class:`ParquetFileReader` and swaps *only* the per-fragment decode +step (:meth:`_iter_fragment_tables`) for the ``ray_data_arrow_rs`` PyO3 +extension (a thin wrapper over the Rust ``parquet``/``arrow`` crates). +Everything above the seam — chunking / row-group fan-out, column projection +resolution, ``path`` / ``row_hash`` synthesis, ``limit`` slicing, block sizing, +per-fragment retry — is inherited unchanged from :class:`FileReader` / +:class:`ParquetFileReader`. + +Selected via ``DataContext.use_arrow_rs_parquet_reader`` (only takes effect when +``use_datasource_v2`` is also set). Switched in +:meth:`ParquetScanner.create_reader`. + +How it reads +------------ +The native extension exposes two entry points, both returning an Arrow +C-stream (consumed zero-copy via ``pa.RecordBatchReader.from_stream``): + +- ``read_row_groups(path, row_groups, columns, batch_size, ...)`` — local files. +- ``read_row_groups_s3(bucket, key, region, anonymous, ...creds..., row_groups, + columns, batch_size, decode_budget_bytes, fetch_window_mb, k, + split_threshold_bytes, predicate_json, column_fetch_mb, prefetch_budget_mb)`` + — S3 via the Rust ``object_store`` crate. Every read decomposes into + prefetchable *units* — row windows of ``≈fetch_window_mb`` compressed bytes + whenever windows can split the group, column groups of ``≈column_fetch_mb`` + only where they can't (wide/short groups: one page per column, or no page + index) — and all units flow through one byte-budget prefetcher: concurrent + ranged GETs + are admitted until ``prefetch_budget_mb`` compressed bytes are in flight ahead + of the (single, in-order) decoder, so S3 peak RSS is a knob, not a property of + the file layout, and S3 latency hides behind decode without staging whole row + groups. + +All of the crate's performance knobs (decode budget, K-split, fetch window, +column window, prefetch budget) are settable per read through +``dataset_kwargs`` under an ``arrow_rs_`` prefix — see the "Tuning knobs" +section below the imports. + +Byte-budgeted decode (no reader-side accumulation) +-------------------------------------------------- +The native reader sizes each decode batch *by bytes, not rows*: it reads each +row group's uncompressed size / row count from the footer and picks a row count +so ``rows × bytes_per_row ≈ decode_budget_bytes`` (default = the current +``DataContext.target_max_block_size``; see :func:`_default_decode_budget_bytes`). +A wide-string group gets few rows/batch, +a numeric group many — both land near the budget, so the decoded working set is +flat across schemas (this is *why* arrow-rs memory doesn't scale with the data +the way PyArrow's whole-row-group materialization does). The ``batch_size`` we +pass is only the upper *clamp*. + +We yield each budget-sized batch straight through — exactly like the base +PyArrow path yields one table per scanner batch (:meth:`FileReader. +_iter_fragment_tables`). Coalescing to ``target_max_block_size`` is done once, +downstream, by the read op's :class:`BlockOutputBuffer`. Accumulating a full +block *here* as well (an earlier prototype did) just stacks a second +block-sized buffer on top of the output buffer's, roughly doubling per-worker +peak RSS relative to PyArrow — so we don't. The decode transient stays bounded +by the byte budget; the single ~128 MiB coalesce buffer lives downstream, shared +with the PyArrow path. + +Prototype limitations (documented, not hidden) +---------------------------------------------- +- Predicate handling prunes at row-group granularity *natively*: the pushed + Ray ``Expr`` is lowered to a small JSON IR (:func:`_predicate_to_ir`) and + handed to the crate, which drops row groups whose footer statistics prove no + row can match (``predicate.rs``) before fetching or decoding them. This + replaces PyArrow's ``fragment.subset(filter=...)``. Pruning is conservative + (a missing column / absent stats / uncomparable type keeps the group), so it + can only avoid IO/decode, never change results. *Row-level* filtering is then + applied post-decode in Python via PyArrow (the final authority) — the crate + has no in-decode ``RowFilter`` yet, so rows inside a surviving row group are + decoded before being dropped. +- The native path covers local **and S3** files whose columns the crate + decodes byte-identically to PyArrow: flat types, ``dictionary``, ``map``, + and ``extension`` types (registered like Ray's tensor types or not — the + crate passes the embedded arrow-schema field metadata straight through FFI, + so pyarrow reconstructs them exactly as it would on its own read path), plus + struct / list / map nesting of all those to any depth. Where the crate's + decode *differs* from what the pyarrow scanner would output but the + difference is mechanical, the planned ``read()`` stays native and realigns + post-decode via a per-file :class:`_ColumnAlignment`: schema-evolution + columns are null-filled, per-file type drift is cast to the unified schema, + INT96 hint units are upcast to pyarrow's default ns (but a file decoding + INT96 under ``coerce_int96_timestamp_unit`` falls back — decode-time + coercion floors where a cast truncates, splitting on pre-1970 values), and + forced ``dictionary_columns`` reads are dictionary-cast. An empty projection + with no predicate (count-style scan) decodes nothing at all — footer row + counts answer it (:class:`_NativeCountFragment`). +- Still gated to PyArrow: non-local/S3 filesystems, Parquet-format kwargs + outside the native allowlist (anything not I/O-only, not reproduced by the + alignment, and not footer-verified — e.g. decryption, an explicit + ``page_checksum_verification=False``, or ``binary_type`` / ``list_type`` + without a pinned dataset schema; see + :meth:`ArrowRsParquetFileReader._blocking_format_kwargs`; the thrift footer + limits stay native via a metadata-only pyarrow probe, + :meth:`_verify_footer_limits`), extension-typed + schema drift, and tz-carrying / nested INT96 oddities. There is no per-type + gate: every type Parquet can encode decodes byte-identically through the + crate, and Arrow's in-memory-only types (``union``, ``list_view``, …) cannot + appear in a Parquet footer at all. (Nested-column + *projection* via dotted names is NOT a gate: V2 discards dotted names in + ``FileReader._split_columns`` before any reader sees them — both paths + silently drop the column, and a flat column literally named ``"a.b"`` + decodes natively; see ``test_dotted_nested_projection_native_parity``.) + Everything gated transparently falls back to the + PyArrow reader, so correctness is never at risk — but benchmarks must + confirm the arrow-rs path actually ran (see the + ``RAY_DATA_USE_ARROW_RS_PARQUET_READER`` verification). +""" + +import json +import logging +import os +import time +from functools import cached_property +from typing import TYPE_CHECKING, Any, Dict, Iterator, List, NamedTuple, Optional, Tuple + +import pyarrow as pa +import pyarrow.dataset as pds +from typing_extensions import override + +from ray._common.utils import env_integer +from ray.data._internal.datasource_v2.native_metadata import ( + connect_native_s3 as _connect_native_s3, + s3_config as _s3_config, + split_s3_path as _split_s3_path, +) +from ray.data._internal.datasource_v2.readers.file_reader import ( + _ARROW_DEFAULT_BATCH_SIZE, + _ARROW_SCANNER_BATCH_READAHEAD, +) +from ray.data._internal.datasource_v2.readers.parquet_file_reader import ( + ParquetFileReader, + _estimate_batch_size_from_chunk_stats, + _estimate_batch_size_from_metadata, +) +from ray.data._internal.object_extensions.arrow import ( + raise_on_pickle_object_columns, +) +from ray.data._internal.util import MiB +from ray.util.annotations import DeveloperAPI + +if TYPE_CHECKING: + import pyarrow.compute as pc # noqa: F401 + + from ray.data._internal.datasource_v2.listing.file_manifest import ( # noqa: F401 + FileManifest, + ) + from ray.data.expressions import Expr # noqa: F401 + +logger = logging.getLogger(__name__) + +# One mallopt attempt per worker process (see _maybe_enable_malloc_trim). +_MALLOC_TRIM_ATTEMPTED = False + + +def _maybe_enable_malloc_trim() -> None: + """Ask glibc to return freed pages eagerly (``M_TRIM_THRESHOLD = 0``). + + Why (findings M48): under task churn on fused read→write shapes, glibc + retains the reader's freed decode heap — worker idle USS climbed +492 MiB + over ~100 tasks in the soak, converting a per-task memory win into the + release suite's measured loss as workers age. Capping arenas barely helped + (the long-time suspect), but ``MALLOC_TRIM_THRESHOLD_=0`` removed the climb + entirely (idle floor 969 → 154 MiB, 0.08× PyArrow's). ``mallopt`` is the + runtime equivalent of that env var and reaches already-started Ray workers, + where an env only reaches processes started after it is set. + + Gated by ``DataContext.arrow_rs_malloc_trim`` + (env ``RAY_DATA_ARROW_RS_MALLOC_TRIM``), default off until the trim arm's + wall cost is certified (trim-on-every-free is a syscall-per-large-free + trade). Linux/glibc only; everywhere else this is a silent no-op. + """ + global _MALLOC_TRIM_ATTEMPTED + if _MALLOC_TRIM_ATTEMPTED: + return + from ray.data.context import DataContext + + if not DataContext.get_current().arrow_rs_malloc_trim: + return + # Only mark attempted once the knob is on: a driver-side reader construction + # before the context propagates must not burn the one attempt. + _MALLOC_TRIM_ATTEMPTED = True + import sys + + if sys.platform != "linux": + return + try: + import ctypes + + libc = ctypes.CDLL("libc.so.6", use_errno=True) + M_TRIM_THRESHOLD = -1 # glibc malloc.h + if libc.mallopt(M_TRIM_THRESHOLD, 0) != 1: + logger.warning( + "mallopt(M_TRIM_THRESHOLD, 0) rejected; retention " + "lever inactive in this worker" + ) + except Exception: + logger.warning( + "arrow_rs_malloc_trim requested but mallopt unavailable " + "(non-glibc libc?); continuing without it", + exc_info=True, + ) + + +# glibc ``malloc_trim`` resolved once per worker process: None = not yet +# looked up, False = unavailable (non-Linux / non-glibc), else the callable. +_LIBC_MALLOC_TRIM = None + + +def _malloc_trim_now() -> None: + """glibc ``malloc_trim(0)``: return freed heap to the OS, once, right now.""" + global _LIBC_MALLOC_TRIM + if _LIBC_MALLOC_TRIM is None: + import sys + + _LIBC_MALLOC_TRIM = False + if sys.platform == "linux": + try: + import ctypes + + _LIBC_MALLOC_TRIM = ctypes.CDLL("libc.so.6", use_errno=True).malloc_trim + except Exception: + logger.warning( + "arrow_rs_malloc_trim_eos requested but malloc_trim unavailable " + "(non-glibc libc?); continuing without it", + exc_info=True, + ) + if _LIBC_MALLOC_TRIM: + _LIBC_MALLOC_TRIM(0) + + +def _maybe_trim_at_stream_end() -> float: + """End-of-stream allocator lever (``DataContext.arrow_rs_malloc_trim_eos``). + + Why: under task churn glibc keeps the reader's freed decode heap resident + (findings M48: idle-worker USS +492 MiB over ~100 tasks). The mallopt lever + (:func:`_maybe_enable_malloc_trim`) removes that floor but trims on every + free AND disables the dynamic mmap threshold, costing 24-36% wall (M61/M64). + Trimming ONCE per task stream releases the same retained pages while the + allocator behaves normally during decode. On by default since 2026-09-08 + (``RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS=0`` ablates it); no-op off glibc. + + Returns the wall seconds the trim took (0.0 when off). The time is surfaced + per task as ``ReadFilesTaskStats.trim_wall_s`` so a "trim cost" reading can + be told apart from placement in result.json: the one 2.37x wall cell seen + under this lever (findings M107) did not replicate x3 (M124). + """ + from ray.data.context import DataContext + + if not DataContext.get_current().arrow_rs_malloc_trim_eos: + return 0.0 + start_s = time.perf_counter() + _malloc_trim_now() + return time.perf_counter() - start_s + + +# Set the first time this worker actually decodes a fragment natively, so the +# "native decode ACTIVE" confirmation is emitted once per worker process rather +# than once per fragment (which would flood the logs). +_LOGGED_NATIVE_ACTIVE = False + +# --------------------------------------------------------------------------- +# Tuning knobs +# --------------------------------------------------------------------------- +# Every knob below is settable *per read* via ``dataset_kwargs`` — the same +# channel PyArrow's own I/O-tuning kwargs (``pre_buffer``, ``buffer_size``, +# ...) travel — under an ``arrow_rs_`` prefix: +# +# ray.data.read_parquet( +# path, dataset_kwargs={"arrow_rs_fetch_window_mb": 64, "arrow_rs_k": 4} +# ) +# +# Precedence: ``dataset_kwargs`` value > ``RAY_DATA_ARROW_RS_*`` env var > +# built-in default. The env vars remain the cluster-wide lever (benchmark +# sweeps set them per worker); the kwargs are the per-read override. +# +# The PyArrow reader ignores the ``arrow_rs_*`` keys — they're popped out of +# the format kwargs in ``ParquetFileReader.__init__`` before PyArrow can see +# them — exactly as this reader ignores PyArrow's I/O-only kwargs +# (:data:`_FORMAT_KWARGS_PERF_ONLY`). So a call carrying either family of +# perf knobs stays valid whichever reader ``use_arrow_rs_parquet_reader`` +# selects. Resolution + validation: :meth:`ArrowRsParquetFileReader._tuning`. + +# Knob ``arrow_rs_decode_budget_bytes`` — crate arg ``decode_budget_bytes``. +# Byte budget for a single arrow-rs decode batch. Sizing decode batches by +# bytes (not a fixed row count) keeps the transient working set flat across +# schemas: a wide file gets few rows/batch, a narrow file gets many. Kept far +# below ``target_block_size`` so the decode transient is bounded while output +# blocks are still coalesced to the normal Ray block size. +# +# History: the default was 32 MiB, raised from 2 MiB after the 2026-08-07 +# Linux + real-S3 sweep, then replaced by the block-target default described at +# the end of this comment (findings M59/M63). The 2026-08-07 evidence: +# (arrow_rs_docs/regression_testing.md §8.2). The old default came from the +# standalone benchmark, where the budget looked like a pure floor knob — it moved +# peak RSS only ~12 MB across its range, so the smallest value that still held +# throughput was the obvious pick. **Inside Ray that reasoning inverts.** Sweeping +# the budget on ``write_parquet`` over S3, per-task USS and wall time as a ratio of +# PyArrow's: +# +# budget avg USS max USS wall +# 2 MiB 1.00x 1.47x 0.99x <- the old default, worst on 2 of 3 +# 32 MiB 0.83x 1.03x 0.87x <- better on all three +# 128 MiB 0.90x 1.29x 1.28x +# +# 32 / 64 / 128 MiB agree within 4% on memory, so 32 is a knee and not a sharp +# optimum; it also wins on the local-disk arm (0.81 / 0.81 / 0.93). +# +# **The memory half of that no longer stands.** Every arm above was fused (a write +# attached to the read) at four fragment threads. Re-run UNFUSED at one thread — +# the configuration this reader now ships — on local disk, 128 MiB blocks: +# +# budget avg USS max USS wall +# 2 MiB 443 MiB 449 MiB 6.7 s +# 8 MiB 489 MiB 496 MiB 6.2 s +# 32 MiB 462 MiB 466 MiB 5.8 s +# 128 MiB 463 MiB 466 MiB 5.9 s +# +# Memory is FLAT: 1.10x across a 64x sweep, non-monotone, and the minimum sits on +# the *old* default. The 1.47x max/avg that condemned 2 MiB is gone too — every arm +# is within 1.5% of its own average, so that spike was the writer's, not the +# decoder's. Wall time keeps the shape it had fused (32 MiB fastest, 0.87x against +# 2 MiB, from two independent configurations) — so the budget is justified on +# throughput and on the row-floor argument below, not on memory. +# +# It also settles why small batches seemed to cost more: they do not. The suspect +# was that handing Ray's block builder many sub-block batches forces an +# accumulate-then-concatenate needing inputs and output alive together. If that +# were it, budget=128 MiB — one batch per 128 MiB block, nothing to concatenate — +# would collapse the cost. It reads 463 MiB against 2 MiB's 443. Refuted; the +# block-size cost measured in exp6 phase G is Ray's block layer (PyArrow scales +# with block size at the same marginal rate) rather than anything this knob +# reaches. +# +# The crate's own batch-row floor used to void this knob for any schema over +# ~16 KiB/row (a 2048-row floor — findings K8); it is now 32 rows +# (``MIN_BATCH_ROWS`` in the crate), so the knob holds up to ~1 MiB/row at this +# default. The 2048-row floor *below* is different: it floors only the +# *requested* batch size handed to the crate, which is the crate's upper clamp, +# never the decoded batch itself. +# Tuning: 32-128 MiB are within noise of each other on memory; go below 8 MiB only +# to reproduce the pre-2026-08 behaviour. +# +# The DEFAULT (env unset, no kwarg) follows ``DataContext.target_max_block_size`` +# rather than a constant. Rationale (findings M59/M63): read tasks coalesce +# decode batches through ``BlockOutputBuffer`` to ~one block anyway, so batches +# below the block size buy NO resident memory (the block-in-progress dominates) +# while paying the per-batch × per-column dispatch cost more often — on 5,000-col +# schemas that cost was the whole in-Ray wall loss (M59: wall R 1.40 → 0.99 at +# 128 MiB, tUSS unchanged). The 10-shape gate at 128 passed the memory gate on +# every cell, including decoded-rg ≫ budget shapes (M63). Tying the default to +# the block target hands the granularity decision to the optimizer/DataContext +# instead of this reader. +_ARROW_RS_DECODE_BUDGET_BYTES: Optional[int] = ( + env_integer("RAY_DATA_ARROW_RS_DECODE_BUDGET_BYTES", 32 * MiB) + if "RAY_DATA_ARROW_RS_DECODE_BUDGET_BYTES" in os.environ + else None +) + + +def _default_decode_budget_bytes() -> int: + """Decode budget when neither the env var nor the kwarg overrides it: + the current ``DataContext.target_max_block_size``, falling back to 128 MiB + when the block target is unset (blocks unbounded ≠ decode unbounded).""" + if _ARROW_RS_DECODE_BUDGET_BYTES is not None: + return _ARROW_RS_DECODE_BUDGET_BYTES + from ray.data.context import DataContext + + target = DataContext.get_current().target_max_block_size + return target if target else 128 * MiB + + +# Floor on the estimated *requested* batch size (rows) handed to the crate. +# The crate treats it as an upper clamp and re-derives the byte-budgeted row +# count per row group (with its own small 32-row floor), so this cannot force +# large decoded batches — it only stops a coarse Python-side estimate from +# capping the crate below what the budget would allow. +_ARROW_RS_MIN_DECODE_BATCH_ROWS = 2048 + +# Was ``RAY_DATA_READ_FILES_NUM_THREADS`` set explicitly? This reader's fragment +# pool defaults to one-worker-per-fragment, parity with the base (see +# ``_num_fragment_read_threads``), but an explicit value must still win — and +# ``env_integer`` cannot distinguish "4 because the user asked for 4" from "4 +# because that is the fallback". Read once at import so that mutating the +# environment after import has no effect, matching how the base reader's own +# module-level knobs behave. +# +# The base no longer reads this variable at all: the footer-chunking path deleted +# ``_DEFAULT_NUM_THREADS`` and goes one-worker-per-fragment, unbounded. So deferring +# to ``super()`` would silently *ignore* an explicit setting rather than honour it, +# which would also break the benchmark harness, whose thread sweep sets exactly this +# variable. We therefore resolve the value here instead of delegating. +_READ_FILES_NUM_THREADS_IS_EXPLICIT = "RAY_DATA_READ_FILES_NUM_THREADS" in os.environ +_READ_FILES_NUM_THREADS_EXPLICIT_VALUE = env_integer( + "RAY_DATA_READ_FILES_NUM_THREADS", 1 +) + +# Knob ``arrow_rs_k`` — crate arg ``k``. +# Intra-fragment parallelism: when a fragment is a *single* row group larger than +# the block-size target (the lone-big-fragment case Ray's thread pool can't split), +# the native reader decodes it in ``K`` parallel row-range workers and merges them +# back in order. Every other layout (multiple / small row groups) uses K=1 because +# Ray's fragment thread pool already parallelizes those — so crate-K and Ray's pool +# never multiply. +# +# Default K=1: locally, K-split costs memory (each range holds its own decode +# transient) for ~no speed, since there is no network latency to hide (benchmarks: +# Agents.md §5.1, §6.3). K>1 is opt-in and is reserved for the S3 phase, where +# concurrent range GETs hide request latency. +# Tuning: try 2–8 for big single-row-group files on S3 when read throughput is +# latency-bound; expect per-task peak memory to scale ~linearly with K (each +# range worker holds its own fetch window + decode transient). Keep 1 for +# local reads and many-small-row-group layouts. +_ARROW_RS_K = env_integer("RAY_DATA_ARROW_RS_K", 1) + +# Knob ``arrow_rs_split_threshold_bytes`` — crate arg ``split_threshold_bytes``. +# A lone row group is only K-split when its uncompressed size exceeds this +# threshold (smaller groups decode sequentially — splitting them buys nothing +# and costs merge overhead). When the knob is unset the reader uses its +# ``target_block_size`` (the Ray block-size target), falling back to the +# default below when it has neither. +# Tuning: rarely needed — lower it (e.g. to 0) only to force the K-split on +# for testing, or raise it to keep K-splitting away from mid-size groups. +_ARROW_RS_DEFAULT_SPLIT_THRESHOLD_BYTES = 128 * MiB + +# Knob ``arrow_rs_fetch_window_mb`` — crate arg ``fetch_window_mb``. +# S3 fetch window (MiB of *compressed* bytes in flight per stream). This is the +# memory knob for the S3 path: the native reader slices each row group's rows into +# windows sized so only ~this many compressed bytes are fetched+buffered before +# decode, so peak RSS is `≈ fetch_window + decode_budget` — flat regardless of +# row-group size, instead of PyArrow's whole-row-group pre-buffer. 0 = no window +# cap (fetch the whole range at once). Swept on the Linux/S3 run (Agents.md §7.1). +# Tuning: this is the primary memory<->throughput trade on S3. Raise (64+) to +# amortize request latency over fewer, larger GETs when memory is plentiful; +# lower toward 4 to cap per-task RSS on memory-tight clusters. No effect on +# local reads. +_ARROW_RS_FETCH_WINDOW_MB = env_integer("RAY_DATA_ARROW_RS_FETCH_WINDOW_MB", 16) + +# Knob ``arrow_rs_column_fetch_mb`` — crate arg ``column_fetch_mb``. +# S3 column-fetch budget (MiB of *compressed* bytes per column group). This is the +# memory knob for WIDE schemas. arrow-rs's async reader fetches every projected +# column chunk of a row group into memory up front (``InMemoryRowGroup``) and holds +# them all while decoding — so a 5000-column group's whole compressed footprint is +# resident, which PyArrow avoids (it releases each column chunk as it decodes). When +# a single row group's projected columns exceed this budget, the native reader reads +# them in sequential column groups (each ≲ this many compressed bytes), holding only +# one group's compressed chunks at a time, so peak ≈ ``column_fetch_mb`` + the fully +# decoded row group (the output, which PyArrow holds too) — asymptotes to PyArrow +# parity as the budget shrinks. 0 disables (fetch the whole row group at once, the +# pre-fix behavior). Only affects the S3 path, and only engages where row windows +# can't split the group (wide/short groups — one page per column — or no page +# index): a tall group with fat columns row-windows instead, because the column-group +# decode retains the whole decoded group (the M20 retention, findings) while windows +# stream it. Narrow/small reads partition to a single group and are untouched. Measured on the +# Linux/S3 run (Agents.md §7.1): on the 5000-column fixture, cf=16 cut per-task USS +# below PyArrow -- fanned out over 4 files arrow_rs peaked at 4.30 GB vs PyArrow's +# 6.78 GB (~37% less) and finished faster; the sweep was monotone in the budget +# (256->3.16, 64->2.30, 16->1.85, 4->1.76 at concurrency=1, byte-identical output), +# so 16 is the knee -- minimum RSS at wall-time parity. fetch_window_mb is its +# row-axis dual. Tuning: the default is already low; RAISE it (64-256) if a +# high-latency S3 layout shows a wall regression from the many small sequential GETs, +# or set 0 to disable (fetch the whole row group at once -- the pre-fix behavior). +_ARROW_RS_COLUMN_FETCH_MB = env_integer("RAY_DATA_ARROW_RS_COLUMN_FETCH_MB", 16) + +# Knob ``arrow_rs_prefetch_budget_mb`` — crate arg ``prefetch_budget_mb``. +# Compressed bytes the S3 path may prefetch AHEAD of its (single) decoder — the +# "bucket" it tries to keep full. ONE mechanism for every unit kind (row +# windows and column groups alike). DERIVED by default: 4 × the larger unit-size +# knob (``max(fetch_window_mb, column_fetch_mb)``), i.e. a bucket of about four +# units, so tuning a unit-size knob scales the bucket with it and there is +# nothing separate to tune. Units are fetched concurrently, gated by a +# byte-denominated semaphore: each in-order fetch acquires permits equal to its +# compressed size (exact, from the footer) and releases them when the decoder +# finishes (drops) that unit — so fetch concurrency self-adjusts to the +# fetch:decode speed ratio (a slow network gets ~4 parallel GETs; a slow decoder +# backpressures fetching to a halt). This is a THROUGHPUT control, not a memory +# knob: the semaphore does cap in-flight *compressed prefetch* bytes, but that +# term is small against decode scratch + retained output, so sweeping the +# budget does not move per-task USS (non-monotone, inside the noise floor — +# findings K2; ``fetch_window_mb``/``column_fetch_mb`` are the levers that do +# move memory, by ablation). Decode itself stays one-unit-at-a-time +# (bounds decode scratch; see ``column_fetch_mb`` for why that matters on wide +# schemas). Explicitly setting this overrides the 4× derivation (benchmark +# escape hatch); 0 = strictly sequential fetch->decode->fetch. ``-1`` sentinel = +# unset -> derive. +_ARROW_RS_PREFETCH_BUDGET_MB = env_integer("RAY_DATA_ARROW_RS_PREFETCH_BUDGET_MB", -1) + +# Parquet-format kwargs (``pds.ParquetFileFormat``) that tune PyArrow's I/O +# strategy only — they cannot change decoded bytes, so the native path (which +# has its own I/O strategy: byte-budgeted streaming + bounded fetch window) may +# safely ignore them. Every other format kwarg either has a native equivalent +# planned per file (see ``_FORMAT_KWARGS_ALIGNED``) or forces a PyArrow +# fallback (see ``ArrowRsParquetFileReader._blocking_format_kwargs``). +_FORMAT_KWARGS_PERF_ONLY = frozenset( + {"pre_buffer", "buffer_size", "use_buffered_stream", "cache_options"} +) +# Parquet-format kwargs whose *semantic* effect the planned native read +# reproduces post-decode via ``_ColumnAlignment`` casts. Membership here means +# the kwarg never blocks the reader; the per-file plan decides. For +# ``coerce_int96_timestamp_unit`` the plan is a *fallback* whenever the file +# actually decodes an INT96 column: pyarrow's decode-time coercion floors +# (parquet types.h divides the unsigned nanos-of-day) while a post-decode cast +# truncates toward zero — off by one unit on every pre-1970 value, so the cast +# cannot reproduce it. Files without INT96 stay native (the kwarg is inert). +_FORMAT_KWARGS_ALIGNED = frozenset( + {"coerce_int96_timestamp_unit", "dictionary_columns"} +) +# Parquet-format kwargs the planned native read enforces via a metadata-only +# pyarrow footer probe (:meth:`ArrowRsParquetFileReader._verify_footer_limits`): +# the thrift limits only decide whether a file's *footer* is accepted or +# rejected — they can never change decoded bytes — so running pyarrow's own +# footer parse with the limits applied reproduces the accept/reject behavior +# (and the raised ``OSError``) exactly, after which the decode stays native. +_FORMAT_KWARGS_FOOTER_VERIFIED = frozenset( + {"thrift_string_size_limit", "thrift_container_size_limit"} +) +# Schema-shaping kwargs (pyarrow 21+): on a file *without* an embedded arrow +# schema they change decoded types (``binary_type=large_binary`` flips +# binary→large_binary AND string→large_string; ``list_type=LargeListType`` +# flips list→large_list); on embedded-schema files (all Ray-written files) +# they are inert. On the V2 pipeline the pinned unified schema — computed by +# the listing via ``pq.read_schema``, which is blind to these kwargs — is the +# final authority: the base reader's pinned-schema cast silently *undoes* +# them (verified empirically, pyarrow 24). So with a pinned schema, parity is +# simply "output the pinned schema", which the native path's per-file +# :class:`_ColumnAlignment` drift casts already guarantee — admit natively. +# WITHOUT a pinned schema the kwargs do change the output and the crate +# doesn't reproduce them — fall back. See +# :meth:`ArrowRsParquetFileReader._blocking_format_kwargs`. +_FORMAT_KWARGS_SCHEMA_SHAPED = frozenset({"binary_type", "list_type"}) + + +class _ArrowRsTuning(NamedTuple): + """Resolved values of the tuning knobs above for one reader instance + (kwarg > env var > default; see :meth:`ArrowRsParquetFileReader._tuning`). + ``split_threshold_bytes=None`` means "derive from ``target_block_size``" + at the call site.""" + + decode_budget_bytes: int + k: int + split_threshold_bytes: Optional[int] + fetch_window_mb: int + column_fetch_mb: int + # None means "derive at the call site": 4 x max(fetch_window_mb, column_fetch_mb). + prefetch_budget_mb: Optional[int] + + +# There is deliberately NO per-type support gate: every type Parquet can store +# (flat types, ``dictionary``, ``map``, ``extension`` — registered like Ray's +# tensor types or not — plus struct / list / map nesting of all those to any +# depth) decodes byte-identically to PyArrow through the crate's C-data +# interface (verified in ``test_extension_types_native_parity`` and the type +# probes behind it). The Arrow types the crate has NOT been verified against +# (``union``, ``list_view`` / ``large_list_view``, ``run_end_encoded``, …) are +# in-memory-only types with no Parquet encoding — PyArrow itself refuses to +# write them ("Unhandled type for Arrow to Parquet schema conversion") — so a +# schema read from a Parquet footer can never contain one, and a gate on them +# was unreachable dead code (removed 2026-07-28). + + +def _pyarrow_fragment_int96_roots(fragment: "pds.ParquetFileFragment") -> set: + """Root (top-level) column names backing an INT96 leaf in a PyArrow Parquet + fragment, read from the *parquet* schema (``fragment.metadata.schema``). + + The fragment's ``physical_schema`` is PyArrow's post-coercion Arrow schema, in + which INT96 already shows up as ``timestamp[ns]`` — so it can't reveal which + columns were INT96 on disk. The parquet schema descriptor can, via each leaf + column's ``physical_type``. Used by the conservative re-gate so an INT96 file + is never handed to the crate through the per-fragment path.""" + roots: set = set() + try: + schema = fragment.metadata.schema + for i in range(len(schema)): + col = schema.column(i) + if col.physical_type == "INT96": + roots.add(col.path.split(".", 1)[0]) + except Exception: # noqa: BLE001 - missing/odd metadata => treat as none + pass + return roots + + +def _raise_if_strict_no_fallback(reason: str) -> None: + """Correctness-harness guard: when ``RAY_DATA_ARROW_RS_STRICT`` is set (to + anything but ``0``/``false``), any decision to serve part of a read through + the PyArrow fallback raises instead of proceeding. A large-scale validation + run flips this on to *guarantee* every byte it checked came off the native + arrow-rs path — a silent fallback would make the run validate PyArrow. + Read per call (not at import) so a harness can toggle it within a session. + Inert by default: production reads never set the variable. + + Regardless of strict mode, emit a visible warning naming the reason so a + benchmark run can confirm from the logs exactly where (and why) a read + dropped off the native path onto PyArrow. + """ + logger.warning("Ray Data ARROW-RS: falling back to PyArrow — %s", reason) + if os.environ.get("RAY_DATA_ARROW_RS_STRICT", "").lower() in ("", "0", "false"): + return + raise RuntimeError( + "RAY_DATA_ARROW_RS_STRICT is set, but this read requires the PyArrow " + f"fallback: {reason}. Strict mode exists for validation harnesses that " + "must prove the native arrow-rs path ran; unset the env var to allow " + "the fallback." + ) + + +def _trace_reader_path(supported: bool) -> None: + """Benchmark instrumentation (inert unless ``RAY_DATA_ARROW_RS_PATH_TRACE`` + names a directory): append ``native``/``fallback`` for each fragment to a + per-pid file so a harness can assert which path the support gate chose. Never + raises into the read path. + """ + trace_dir = os.environ.get("RAY_DATA_ARROW_RS_PATH_TRACE") + if not trace_dir: + return + try: + import socket + + # Namespace by hostname so nodes writing to a shared trace dir (multi-node + # verification) don't collide on pid; the harness's ``path_*.log`` glob + # still matches. Single-node is unaffected. + line = "native\n" if supported else "fallback\n" + fname = f"path_{socket.gethostname()}_{os.getpid()}.log" + with open(os.path.join(trace_dir, fname), "a") as fh: + fh.write(line) + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Predicate lowering: Ray Expr -> native pruning IR (predicate pushdown, part 1) +# --------------------------------------------------------------------------- +# The native crate does statistics-based row-group pruning from a small JSON IR +# (parsed by ``predicate.rs``). We lower the *pushed* Ray ``Expr`` predicate into +# that IR here rather than translating the PyArrow expression, because the Ray +# AST (ColumnExpr / LiteralExpr / BinaryExpr / UnaryExpr) is directly +# introspectable. +# +# The lowering is **total**: any node it can't represent becomes ``{"t": +# "unknown"}``, which the crate treats as "keep this row group". So a partially +# understood predicate like ``a > 5 AND some_udf(b)`` still lowers to +# ``And[cmp(a>5), unknown]`` and prunes soundly on the ``a > 5`` conjunct instead +# of giving up. Pruning is conservative on the Rust side (a group is dropped only +# when provably empty), and the reader re-applies the full predicate post-decode, +# so this only ever avoids IO/decode — it can never change which rows are +# returned. + +# Ray comparison Operation -> IR op string. +_CMP_OP_TO_IR = { + "gt": "gt", + "lt": "lt", + "ge": "ge", + "le": "le", + "eq": "eq", + "ne": "ne", +} +# When the column is on the *right* of a comparison (``5 < col``), flip the op so +# the IR always reads ``col OP literal``. +_CMP_OP_FLIP = {"gt": "lt", "lt": "gt", "ge": "le", "le": "ge", "eq": "eq", "ne": "ne"} + +_IR_UNKNOWN: Dict[str, Any] = {"t": "unknown"} + + +def _literal_to_ir_value(value: Any) -> Optional[Dict[str, Any]]: + """Lower a Python literal to a tagged IR value, or None if the crate can't + order it for pruning (bytes, datetimes, decimals, ...), which makes the + enclosing atom ``unknown``. ``bool`` is checked before ``int`` because + ``bool`` is an ``int`` subclass.""" + if isinstance(value, bool): + return {"vt": "bool", "v": value} + if isinstance(value, int): + return {"vt": "int", "v": value} + if isinstance(value, float): + return {"vt": "float", "v": value} + if isinstance(value, str): + return {"vt": "str", "v": value} + if value is None: + return {"vt": "null"} + return None + + +def _predicate_to_ir(expr: "Expr") -> Dict[str, Any]: + """Lower a Ray Data predicate ``Expr`` to the native pruning IR (see above). + + Total by construction: unrepresentable subtrees become ``_IR_UNKNOWN``. + """ + from ray.data.expressions import ( + AliasExpr, + BinaryExpr, + ColumnExpr, + LiteralExpr, + Operation, + UnaryExpr, + ) + + def unwrap(e: "Expr") -> "Expr": + # Aliasing doesn't change the value being compared. + while isinstance(e, AliasExpr): + e = e.expr + return e + + def lower(e: "Expr") -> Dict[str, Any]: + e = unwrap(e) + + if isinstance(e, UnaryExpr): + if e.op == Operation.NOT: + return {"t": "not", "pred": lower(e.operand)} + if e.op in (Operation.IS_NULL, Operation.IS_NOT_NULL): + operand = unwrap(e.operand) + if isinstance(operand, ColumnExpr): + tag = "is_null" if e.op == Operation.IS_NULL else "is_not_null" + return {"t": tag, "col": operand.name} + return _IR_UNKNOWN + + if isinstance(e, BinaryExpr): + if e.op in (Operation.AND, Operation.OR): + tag = "and" if e.op == Operation.AND else "or" + return {"t": tag, "preds": [lower(e.left), lower(e.right)]} + + if e.op in (Operation.IN, Operation.NOT_IN): + col = unwrap(e.left) + rhs = unwrap(e.right) + if isinstance(col, ColumnExpr) and isinstance(rhs, LiteralExpr): + raw = rhs.value + raw = raw if isinstance(raw, list) else [raw] + values = [_literal_to_ir_value(v) for v in raw] + if all(v is not None for v in values): + return { + "t": "in", + "col": col.name, + "values": values, + "negated": e.op == Operation.NOT_IN, + } + return _IR_UNKNOWN + + ir_op = _CMP_OP_TO_IR.get(e.op.value) + if ir_op is not None: + left = unwrap(e.left) + right = unwrap(e.right) + if isinstance(left, ColumnExpr) and isinstance(right, LiteralExpr): + val = _literal_to_ir_value(right.value) + if val is not None: + return {"t": "cmp", "col": left.name, "op": ir_op, "value": val} + elif isinstance(left, LiteralExpr) and isinstance(right, ColumnExpr): + val = _literal_to_ir_value(left.value) + if val is not None: + return { + "t": "cmp", + "col": right.name, + "op": _CMP_OP_FLIP[ir_op], + "value": val, + } + return _IR_UNKNOWN + + return _IR_UNKNOWN + + return lower(expr) + + +def _predicate_json(predicate: "Optional[Expr]") -> Optional[str]: + """Serialize the pushed predicate's pruning IR for the crate, or None when + there's nothing prunable (no predicate, or it lowered entirely to + ``unknown``) so we skip the pushdown argument altogether.""" + if predicate is None: + return None + ir = _predicate_to_ir(predicate) + if ir == _IR_UNKNOWN: + return None + return json.dumps(ir) + + +def _is_extension_type(t: pa.DataType) -> bool: + """Two-way extension detection: ``isinstance`` for registered extensions, + ``extension_name`` for canonical pyarrow extensions (e.g. + ``fixed_shape_tensor``) that aren't ``pa.ExtensionType`` instances on every + pyarrow version.""" + return isinstance(t, pa.ExtensionType) or ( + getattr(t, "extension_name", None) is not None + ) + + +class _ColumnAlignment(NamedTuple): + """Per-file post-decode fixups that make a native decode byte-match what the + pyarrow scanner would have produced for the same file against the unified + dataset schema. Built once per file at plan time + (:meth:`ArrowRsParquetFileReader._plan_column_alignment`), applied to every + decoded batch (:func:`_apply_column_alignment`) *before* the post-decode + filter, so the predicate evaluates against the same types pyarrow's scanner + filters on. + + - ``null_fill``: columns absent from this file (schema evolution) appended + as typed all-null columns — exactly pyarrow's null-fill under a pinned + dataset schema. + - ``casts``: columns whose crate-decoded type differs from the expected + output type (per-file type drift vs the unified schema, INT96 unit + realignment, forced ``dictionary_columns`` decode). The bool is + ``allow_time_truncate``, set only for INT96 unit coercion where pyarrow + itself truncates (``coerce_int96_timestamp_unit``); every other cast is + safe, so lossy data errors loudly — the same outcome pyarrow's own + scanner cast produces. + - ``order``: final column order, or ``None`` to keep the crate's order. + Set only when ``null_fill`` is non-empty (appended columns must land in + read order, matching the scanner's projected-column order). + """ + + null_fill: Tuple[Tuple[str, pa.DataType], ...] + casts: Tuple[Tuple[str, pa.DataType, bool], ...] + order: Optional[Tuple[str, ...]] + + @property + def is_noop(self) -> bool: + return not self.null_fill and not self.casts + + +_NOOP_ALIGNMENT = _ColumnAlignment(null_fill=(), casts=(), order=None) + + +def _cast_table_to(table: pa.Table, target_fields: List[pa.Field]) -> pa.Table: + """One ``Table.cast`` against a prebuilt positional schema. A per-column + ``ChunkedArray.cast`` + ``set_column`` loop costs ~50 µs of Python dispatch + per column *and* rebuilds the schema each time (O(columns²) field copies); + on the 5000-column tensor shape that is ~1 s per decoded batch where this + call is ~40 ms for the identical result (T23).""" + return table.cast(pa.schema(target_fields, metadata=table.schema.metadata)) + + +def _apply_column_alignment( + table: pa.Table, alignment: Optional[_ColumnAlignment] +) -> pa.Table: + """Apply a plan-time :class:`_ColumnAlignment` to one decoded batch.""" + if alignment is None or alignment.is_noop: + return table + import pyarrow.compute as pc + + # allow_time_truncate needs per-column CastOptions, which Table.cast can't + # carry — apply those first (INT96 unit coercion only, so at most a few). + for name, target, allow_time_truncate in alignment.casts: + if not allow_time_truncate: + continue + idx = table.schema.get_field_index(name) + if idx == -1: + continue + column = table.column(idx).cast( + options=pc.CastOptions(target, allow_time_truncate=True) + ) + table = table.set_column(idx, pa.field(name, target), column) + + # Everything else in one Table.cast (safe: lossy values raise, like + # pyarrow). Cast only each name's first occurrence — get_field_index + # semantics of the old per-column loop. + targets = { + name: target for name, target, truncate in alignment.casts if not truncate + } + target_fields, seen, changed = [], set(), False + for field in table.schema: + target = targets.get(field.name) if field.name not in seen else None + seen.add(field.name) + if target is not None and field.type != target: + target_fields.append(pa.field(field.name, target)) + changed = True + else: + target_fields.append(field) + if changed: + table = _cast_table_to(table, target_fields) + + for name, fill_type in alignment.null_fill: + table = table.append_column( + pa.field(name, fill_type), pa.nulls(table.num_rows, type=fill_type) + ) + if alignment.order is not None: + table = table.select([c for c in alignment.order if c in table.column_names]) + return table + + +def _reconcile_to_expected(table: pa.Table, expected_schema: pa.Schema) -> pa.Table: + """Cast any decoded column whose type differs from ``expected_schema`` to the + expected type. Used by the per-fragment path (:meth:`_iter_native_tables` + with ``expected_schema``), where the crate may hand back a parquet *storage* + type for an extension column whose non-UTF8 embedded metadata it had to skip + — this restores the extension type the base pyarrow scanner would produce. + The per-fragment gate already withholds every other kind of drift, so the + only divergence reaching here is a safe storage → extension wrap (a lossy + cast would raise, exactly as pyarrow's own scanner cast does).""" + target_fields, seen, changed = [], set(), False + for field in table.schema: + exp_idx = ( + expected_schema.get_field_index(field.name) + if field.name not in seen + else -1 + ) + seen.add(field.name) + exp_type = expected_schema.field(exp_idx).type if exp_idx != -1 else None + if exp_type is not None and field.type != exp_type: + target_fields.append(pa.field(field.name, exp_type)) + changed = True + else: + target_fields.append(field) + if not changed: + return table + return _cast_table_to(table, target_fields) + + +def _widen_storage_type( + inferred: pa.DataType, storage: pa.DataType +) -> Optional[pa.DataType]: + """Adopt ``storage``'s container *kinds* while keeping ``inferred``'s nested + field names. parquet-rs validates a supplied schema's nested field names + against its own inference (list child name ``element``), while a pyarrow + extension's storage type uses pyarrow's (``item``) — so the schema handed + to ``with_schema_override`` must carry the crate's names with the storage's + layout (e.g. ``list`` → ``large_list``, i.e. i32 → i64 offsets). Names live + in the schema, not the buffers, so the later relabel supplies the + extension's own names over the same buffers. Any mismatch beyond container + kind returns ``None`` (no override, today's cast path).""" + if inferred.equals(storage): + return storage + inferred_listish = pa.types.is_list(inferred) or pa.types.is_large_list(inferred) + storage_listish = pa.types.is_list(storage) or pa.types.is_large_list(storage) + if inferred_listish and storage_listish: + child = _widen_storage_type(inferred.value_type, storage.value_type) + if child is None: + return None + child_field = inferred.value_field + make = pa.large_list if pa.types.is_large_list(storage) else pa.list_ + return make(pa.field(child_field.name, child, nullable=child_field.nullable)) + if pa.types.is_fixed_size_list(inferred) and pa.types.is_fixed_size_list(storage): + if inferred.list_size != storage.list_size: + return None + child = _widen_storage_type(inferred.value_type, storage.value_type) + if child is None: + return None + child_field = inferred.value_field + return pa.list_( + pa.field(child_field.name, child, nullable=child_field.nullable), + storage.list_size, + ) + if pa.types.is_struct(inferred) and pa.types.is_struct(storage): + if inferred.num_fields != storage.num_fields: + return None + children = [] + for i in range(inferred.num_fields): + child_field = inferred.field(i) + child = _widen_storage_type(child_field.type, storage.field(i).type) + if child is None: + return None + children.append( + pa.field(child_field.name, child, nullable=child_field.nullable) + ) + return pa.struct(children) + return None + + +def _storage_override_schema( + inferred: pa.Schema, target: pa.Schema +) -> Optional[pa.Schema]: + """Build the storage-typed schema handed to the crate's + ``with_schema_override`` when the embedded arrow schema was skipped: for + every column the pyarrow footer types as an extension, the extension's + *storage* layout (via :func:`_widen_storage_type`); everything else the + crate's own inference, metadata-stripped (``Field`` metadata is UTF-8-only + on the Rust side — the whole reason the embedded schema was skipped). + Returns ``None`` when there is nothing to change or a storage type can't + be reconciled (caller keeps today's cast path).""" + fields, changed = [], False + for i in range(len(inferred)): + field = inferred.field(i) + target_idx = target.get_field_index(field.name) + target_type = target.field(target_idx).type if target_idx != -1 else None + if target_type is not None and _is_extension_type(target_type): + storage = getattr(target_type, "storage_type", None) + if storage is None: + return None + widened = _widen_storage_type(field.type, storage) + if widened is None: + return None + if not widened.equals(field.type): + changed = True + fields.append(pa.field(field.name, widened, nullable=field.nullable)) + else: + fields.append(pa.field(field.name, field.type, nullable=field.nullable)) + if not changed: + return None + return pa.schema(fields) + + +def _storage_layout_matches(actual: pa.DataType, storage: pa.DataType) -> bool: + """True when ``actual`` (the crate's decoded type) and ``storage`` (an + extension type's storage) share the same physical buffer layout, ignoring + nested field *names* — names live in the schema, not the buffers, and the + C-import relabel supplies the target's names.""" + if actual.equals(storage): + return True + if pa.types.is_large_list(actual) and pa.types.is_large_list(storage): + return _storage_layout_matches(actual.value_type, storage.value_type) + if pa.types.is_list(actual) and pa.types.is_list(storage): + return _storage_layout_matches(actual.value_type, storage.value_type) + if pa.types.is_fixed_size_list(actual) and pa.types.is_fixed_size_list(storage): + return actual.list_size == storage.list_size and _storage_layout_matches( + actual.value_type, storage.value_type + ) + if pa.types.is_struct(actual) and pa.types.is_struct(storage): + if actual.num_fields != storage.num_fields: + return False + return all( + _storage_layout_matches(actual.field(i).type, storage.field(i).type) + for i in range(actual.num_fields) + ) + return False + + +def _ffi_relabel_enabled() -> bool: + """Kill switch for the zero-copy relabel fast path (M53).""" + return os.environ.get("RAY_DATA_ARROW_RS_FFI_RELABEL", "1") != "0" + + +_PYCAPSULE_GET_POINTER = None + + +def _ffi_relabel_batch(batch: pa.RecordBatch, schema: pa.Schema) -> pa.RecordBatch: + """Zero-copy re-type of a decoded batch: export via the Arrow C Data + Interface, re-import against a prebuilt target ``schema`` object. Same + buffers, new labels — this is how the storage → extension realign avoids + both the per-batch ``Table.cast`` (~40 ms at 5000 columns) *and* pyarrow's + per-batch re-deserialization of pickled extension type params (the capsule + schema route pays that; a schema *object* import does not — measured + 2-4 ms/batch, bit-identical, M53). Only valid when the target schema's + layout matches the batch's buffers exactly (:func:`_storage_layout_matches` + per column) — the import validates structure, not values.""" + global _PYCAPSULE_GET_POINTER + if _PYCAPSULE_GET_POINTER is None: + import ctypes + + fn = ctypes.pythonapi.PyCapsule_GetPointer + fn.restype = ctypes.c_void_p + fn.argtypes = [ctypes.py_object, ctypes.c_char_p] + _PYCAPSULE_GET_POINTER = fn + _, array_capsule = batch.__arrow_c_array__() + addr = _PYCAPSULE_GET_POINTER(array_capsule, b"arrow_array") + if not addr: + raise RuntimeError("null arrow_array capsule") + # _import_from_c MOVES the C struct (nulls its release callback), so the + # capsule's own destructor no-ops — but the capsule must stay alive until + # the import returns or it releases the arrays first (measured: it does). + imported = pa.RecordBatch._import_from_c(addr, schema) + del array_capsule + return imported + + +def _plan_batch_relabel( + stream_schema: pa.Schema, + alignment: Optional[_ColumnAlignment], + expected_schema: Optional[pa.Schema], +) -> Optional[Tuple[pa.Schema, Optional[_ColumnAlignment], bool]]: + """Decide once per stream whether the storage → extension realign can run + as a single zero-copy C-import relabel per batch instead of per-batch + casts. Returns ``(relabel_schema, residual_alignment, skip_reconcile)`` or + ``None`` to keep the cast path unchanged. + + Applies when every extension-typed realign target's storage layout + structurally matches what the crate decoded (true after + ``with_schema_override``; false without it, where list offsets still + differ and the cast path must widen them). Non-extension casts (INT96 + units, dictionary_columns, type drift), null-fill and reordering stay in + ``residual_alignment`` and run on the relabeled table as before. + ``skip_reconcile`` is ``True`` when the relabel alone already lands every + column on ``expected_schema``'s type, making the per-fragment + :func:`_reconcile_to_expected` pass redundant.""" + if not hasattr(pa.RecordBatch, "_import_from_c"): + return None + # Final extension-typed targets per stream column (first occurrence), + # from the alignment casts (planned path) and/or expected_schema (the + # per-fragment path) — mirroring _apply_column_alignment / + # _reconcile_to_expected's get_field_index semantics. + ext_targets: Dict[str, pa.DataType] = {} + if alignment is not None: + for name, target, truncate in alignment.casts: + if not truncate and _is_extension_type(target): + ext_targets.setdefault(name, target) + if expected_schema is not None: + seen = set() + for field in stream_schema: + if field.name in seen: + continue + seen.add(field.name) + exp_idx = expected_schema.get_field_index(field.name) + if exp_idx == -1: + continue + exp_type = expected_schema.field(exp_idx).type + if exp_type != field.type and _is_extension_type(exp_type): + if ext_targets.setdefault(field.name, exp_type) != exp_type: + return None # alignment and expected disagree — stay safe + if not ext_targets: + return None + + fields, covered, seen = [], set(), set() + for field in stream_schema: + target = ext_targets.get(field.name) if field.name not in seen else None + seen.add(field.name) + if target is None: + fields.append(field) + continue + storage = getattr(target, "storage_type", None) + if storage is None or not _storage_layout_matches(field.type, storage): + return None + fields.append(pa.field(field.name, target, nullable=field.nullable)) + covered.add(field.name) + relabel_schema = pa.schema(fields) + + residual = alignment + if alignment is not None: + residual_casts = tuple(c for c in alignment.casts if c[0] not in covered) + residual = _ColumnAlignment( + alignment.null_fill, residual_casts, alignment.order + ) + if residual.is_noop: + residual = None + + skip_reconcile = False + if expected_schema is not None and (residual is None or not residual.casts): + skip_reconcile = True + seen = set() + for field in relabel_schema: + if field.name in seen: + continue + seen.add(field.name) + exp_idx = expected_schema.get_field_index(field.name) + if exp_idx != -1 and expected_schema.field(exp_idx).type != field.type: + skip_reconcile = False + break + return relabel_schema, residual, skip_reconcile + + +class _NativeParquetFragment(NamedTuple): + """A native (pyarrow-free) unit of work for one file's row-group slice. + + The arrow-rs ``read()`` builds these instead of pyarrow ``ParquetFileFragment`` + objects for files the native reader handles, so pyarrow never opens a + supported file. ``row_groups is None`` means "all row groups in the file" + (whole-file read). Exposes ``.path`` so it flows through the same + :meth:`FileReader._dispatch_fragment_reads` threading/retry machinery as a + pyarrow fragment; :meth:`ArrowRsParquetFileReader._iter_fragment_tables` + dispatches on the type. ``alignment`` carries this file's post-decode + fixups (:class:`_ColumnAlignment`), ``None`` when the decode already + matches the expected output. + """ + + path: str + row_groups: Optional[List[int]] + alignment: Optional[_ColumnAlignment] = None + # The crate's per-file ``NativeParquetFile`` handle (TODO 1r): the parsed + # footer + (for S3) the task's shared client, opened once at plan time. + # Decode goes through ``handle.read_row_groups`` so it never re-fetches the + # footer or rebuilds an HTTP client. ``None`` only in tests that build + # fragments by hand; the decode then falls back to the per-call entry + # points. Never serialized — fragments live and die inside one read task. + handle: Optional[Any] = None + + +class _NativeCountFragment(NamedTuple): + """A zero-decode work unit for an empty projection (count-style scan) with + no predicate: the footer row counts are exact, so the read yields a + zero-column table with the right ``num_rows`` and never touches a data + page. The base pyarrow path instead scans a stub column; + :meth:`FileReader._postprocess`'s stub guard re-adds the row-preserving + stub downstream, identically for both paths.""" + + path: str + num_rows: int + + +@DeveloperAPI +class ArrowRsParquetFileReader(ParquetFileReader): + """Parquet reader that decodes each fragment via the arrow-rs extension. + + See the module docstring for the design. Only :meth:`_iter_fragment_tables`, + :meth:`_resolve_batch_size`, :meth:`_num_fragment_read_threads`, and + :meth:`_on_batch_read` are overridden; the rest of the read pipeline is + inherited from :class:`ParquetFileReader`. + """ + + @override + def _num_fragment_read_threads(self, num_fragments: int) -> int: + """``num_fragments`` — pool-width **parity** with the base path's + one-worker-per-fragment pool (decided 2026-08-12, revising the earlier + ``min(4, num_fragments)`` cap). + + Why parity rather than a cap: + + - A sub-fragment is one *file's* bin-assigned row groups, so + ``num_fragments`` = files spanned by the bin. At realistic bin + budgets (64 MiB-1.25 GiB in the release suite) that is a handful — + the "unbounded" pool is bounded by bin geometry in practice, which + is the same reason the base path tolerates it. + - A narrower pool than the base turns every multi-fragment A/B cell + into a pool-width comparison instead of a decode comparison. With + parity, an arrow-rs wall loss is decode; before, it was ambiguous. + - The memory cost of an extra in-flight fragment on this path is a + decode-budget transient (plus the S3 fetch window), not a whole + decoded row group — the multiplier the old cap guarded against is + the *base* path's failure mode, not ours. + + History (findings K6, K10 in ``arrow_rs_docs/findings.md``): K6 + (old row-group-fragment base) found serial free → default 1; K10 + (GE1, multi-file-bin base) found threads=4 vs 1 cuts read-op time + 1.6-3.3x at flat-to-+22% memory → default 4; 2026-08-12 → parity. + 4-vs-unbounded was never the measured comparison; the bin sweep's + arrow-rs-flat prediction (TODO item 10) now doubles as this default's + regression check — if arrow-rs USS *grows* with bin size, suspect + pool width first and re-cap via the env var. + + A 1-fragment task still takes the sequential branch + (``_dispatch_fragment_reads`` on ``num_workers <= 1`` never constructs + ``make_async_gen``), where the crate alone owns parallelism — the + lone-big-row-group case. An explicit + ``RAY_DATA_READ_FILES_NUM_THREADS`` still wins: a user who set it + meant it, and the benchmark harness sweeps it. Shapes where a bin + genuinely spans very many tiny files can use it to re-cap. + """ + if _READ_FILES_NUM_THREADS_IS_EXPLICIT: + return _READ_FILES_NUM_THREADS_EXPLICIT_VALUE + return max(1, num_fragments) + + @override + def _resolve_batch_size( + self, dataset: pds.Dataset, manifest: "FileManifest" + ) -> int: + """Size the decode batch to the arrow-rs byte budget, not the block size. + + Priority: explicit ``batch_size`` > footer-stat estimate from the manifest + (no I/O) > byte-budget estimate from a row-group footer read > default. + Unlike the base reader this targets the resolved decode budget (kwarg > + env > :func:`_default_decode_budget_bytes`, which follows + ``DataContext.target_max_block_size``), because each decode batch is + yielded straight through in :meth:`_iter_fragment_tables` (the + downstream ``BlockOutputBuffer`` does the coalescing to the block size). + + Preferring the manifest is not just a tidier source: ``ListFiles`` already + read every footer to prune and pack the row groups, and it recorded the + projection-scoped uncompressed size and row count of the groups it assigned + to this chunk. Reading a footer again here would be a second round trip per + read task purely to recompute a number we were handed — and on the + many-small-files shapes behind the release regressions that round trip is + the dominant per-task cost. + """ + if self._explicit_batch_size is not None: + return self._explicit_batch_size + + if self._target_block_size is None: + return _ARROW_DEFAULT_BATCH_SIZE + + budget = self._tuning.decode_budget_bytes + + # Footer stats off the manifest, when the chunker supplied them. Sum across + # the split's chunks so the average row size reflects everything this task + # will decode rather than whichever file happens to be first. + total_size = total_rows = 0 + for chunk_metadata in manifest.file_chunk_metadatas: + if not chunk_metadata: + continue + size = chunk_metadata.get("uncompressed_size") + rows = chunk_metadata.get("num_rows") + if size and rows: + total_size += size + total_rows += rows + if total_size and total_rows: + estimated = _estimate_batch_size_from_chunk_stats( + total_size, total_rows, budget + ) + if estimated is not None: + return max(estimated, _ARROW_RS_MIN_DECODE_BATCH_ROWS) + + # No chunk stats (e.g. a whole-file manifest from ``WholeFileChunker``): + # fall back to reading the first fragment's footer, as before. + first_fragment = next(dataset.get_fragments(), None) + if first_fragment is None: + return _ARROW_DEFAULT_BATCH_SIZE + + estimated = _estimate_batch_size_from_metadata( + first_fragment, self._columns, budget + ) + if estimated is None: + return _ARROW_DEFAULT_BATCH_SIZE + return max(estimated, _ARROW_RS_MIN_DECODE_BATCH_ROWS) + + @override + def _on_batch_read(self, table: pa.Table) -> None: + """No-op: the decode batch size is fixed by the byte budget, so there is + nothing to refine from actual data (unlike the base reader).""" + return None + + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + # Wall seconds of end-of-stream trims since the last pop_task_stats(). + self._trim_wall_s = 0.0 + + @override + def pop_task_stats(self) -> Dict[str, float]: + trim_wall_s, self._trim_wall_s = self._trim_wall_s, 0.0 + return {"trim_wall_s": trim_wall_s} + + @override + def read(self, input_split: "FileManifest") -> "Iterator[pa.Table]": + """Pyarrow-free Parquet read for supported files. + + For files the native reader handles (local/S3, flat + struct/list types, + no schema evolution / dictionary / extension), the footer, row-group + layout, and statistics come from the crate's ``read_metadata`` and decode + from ``read_row_groups`` — pyarrow never opens the file. Files (or whole + splits) the native path can't handle transparently fall back to the base + pyarrow ``read()`` / scanner, so correctness is never at risk. The + format-agnostic finishing (limit, partition/``path``/``row_hash`` + synthesis, projection) is shared with the base reader via + :meth:`_postprocess`. + """ + if len(input_split) == 0: + return + + # Worker-process allocator lever (no-op unless the knob is on). + _maybe_enable_malloc_trim() + try: + yield from self._read_split(input_split) + finally: + # Once per task stream, whichever path served it (no-op unless on); + # timed, drained per task by :meth:`pop_task_stats`. + self._trim_wall_s += _maybe_trim_at_stream_end() + + def _read_split(self, input_split: "FileManifest") -> "Iterator[pa.Table]": + # Reader-wide ineligibility: unsupported filesystem, or a Parquet-format + # kwarg outside the native allowlist (anything not perf-only, not + # reproduced by the per-file :class:`_ColumnAlignment`, and not + # footer-verified — e.g. decryption, ``binary_type``) — use the base + # pyarrow read() unchanged, which honors every format kwarg. + blocked = self._blocking_format_kwargs(aligned_ok=True) + if not self._filesystem_supported() or blocked: + if blocked: + _raise_if_strict_no_fallback( + f"unsupported parquet format kwargs {sorted(blocked)}" + ) + else: + _raise_if_strict_no_fallback( + f"unsupported filesystem {type(self._filesystem).__name__}" + ) + yield from super().read(input_split) + return + + plan = self._plan_native_read(input_split) + if plan is None: + # A file's footer couldn't be read natively (corrupt / unsupported + # footer); fall the whole split back to pyarrow rather than reason + # about a partially-known layout. + _raise_if_strict_no_fallback( + "a file's footer could not be read via the native crate" + ) + yield from super().read(input_split) + return + + fragments_with_offsets, columns_to_synthesize, scanner_kwargs = plan + global _LOGGED_NATIVE_ACTIVE + if not _LOGGED_NATIVE_ACTIVE: + _LOGGED_NATIVE_ACTIVE = True + logger.warning( + "Ray Data ARROW-RS: native decode ACTIVE on this worker — " + "Parquet fragments are being read via the Rust ray_data_arrow_rs " + "crate, not PyArrow." + ) + triples = self._dispatch_fragment_reads(fragments_with_offsets, scanner_kwargs) + yield from self._postprocess(triples, columns_to_synthesize) + + def _verify_footer_limits(self, paths: List[str]) -> None: + """Enforce ``thrift_string_size_limit`` / ``thrift_container_size_limit`` + on the planned native read with a metadata-only pyarrow footer parse. + + The limits guard *footer deserialization* (accept vs reject a file's + metadata) and can never change decoded bytes, so running pyarrow's own + thrift parser with the limits applied reproduces the base path's + accept/reject behavior — and the exact ``OSError`` it raises — while + the data decode stays native. This is the one deliberate exception to + "pyarrow never opens a supported file": a footer-only read (a few KB), + and only when the user actually set a limit. A raised error is the + *correct* outcome, not a fallback trigger — the base path would raise + the same error, so we let it propagate.""" + limits = { + key: self._parquet_format_kwargs[key] + for key in _FORMAT_KWARGS_FOOTER_VERIFIED + if self._parquet_format_kwargs.get(key) is not None + } + if not limits: + return + import pyarrow.parquet as pq + from pyarrow.fs import LocalFileSystem + + fs = self._filesystem or LocalFileSystem() + for path in paths: + with fs.open_input_file(path) as source: + # Constructing ParquetFile parses the footer under the limits; + # a violation raises pyarrow's usual thrift OSError. + pq.ParquetFile(source, **limits) + + def _read_pyarrow_footer_schema(self, path: str) -> Optional[pa.Schema]: + """Read a file's footer schema via pyarrow, which parses binary field + metadata that arrow-rs's IPC verifier rejects. Used only to recover the + Arrow *logical* schema — including reconstructed extension types like + Ray's cloudpickle-serialized tensor type — when the crate had to skip the + embedded arrow schema (:meth:`_open_native_file`). Footer-only (a few + KB); ``None`` on failure so the caller falls the file back to pyarrow.""" + import pyarrow.parquet as pq + + try: + return pq.read_schema(path, filesystem=self._filesystem) + except Exception as e: # noqa: BLE001 - any footer failure => fallback + logger.debug("pyarrow footer schema read failed for %s: %s", path, e) + return None + + def _open_native_file( + self, path: str, s3_stores: Dict[str, Any] + ) -> Optional[Tuple[Any, pa.Schema, List[int], List[str], Optional[pa.Schema]]]: + """Open one file through the crate's per-file handle (TODO 1r): + ``(handle, arrow schema, per-row-group row counts, int96 root columns, + extension-target schema)``, or ``None`` if the native footer read fails + (caller then falls the whole split back to pyarrow). Does *not* swallow + a missing extension — :meth:`_import_extension` raises that loudly. + + The handle holds the parsed footer and (for S3) a shared client, so + the footer is fetched exactly once per file and the decode call never + rebuilds an HTTP client — the fix for the per-file S3 setup cost on + multi-file bins (findings T10). ``s3_stores`` is the caller's + per-bucket lookup for this one planned read; behind it, + :func:`connect_native_s3` caches clients per process (keyed by bucket + + full connection config) so single-file tasks don't pay a cold + DNS+TLS client build each (findings M97). For S3 the page index is fetched at open (the decode's + row windows need it — a file that later falls back to pyarrow wastes + one range GET, which is cheaper than the footer re-fetch every native + file used to pay); locally it follows the same rule as the per-call + path: only a possible K-split (``k > 1``) needs it. + + The int96 list lets :meth:`_plan_column_alignment` realign the crate's + decoded unit for those columns to what PyArrow produces. The last + element is a per-file *target* schema, non-``None`` only when the crate + reports it had to skip the embedded arrow schema (non-UTF8 field + metadata, e.g. Ray's cloudpickle tensor type): then the crate decodes + the parquet *storage* types, and this pyarrow-read footer schema + carries the reconstructed extension types the base path would produce, + so :meth:`_plan_column_alignment` can cast storage → extension per + file.""" + # Surfaces a missing extension loudly (import inside the crate call); + # any *footer-read* failure below becomes a whole-split pyarrow fallback. + ray_data_arrow_rs = self._import_extension() + + from pyarrow.fs import S3FileSystem + + try: + if isinstance(self._filesystem, S3FileSystem): + bucket, key = _split_s3_path(path) + store = s3_stores.get(bucket) + if store is None: + store = _connect_native_s3(bucket, self._filesystem) + s3_stores[bucket] = store + handle = store.open_file(key, page_index=True) + else: + handle = ray_data_arrow_rs.open_parquet_file( + path, page_index=self._tuning.k > 1 + ) + md = handle.metadata() + except Exception as e: # noqa: BLE001 - any footer failure => fallback + logger.debug("arrow-rs native file open failed for %s: %s", path, e) + return None + + target_override: Optional[pa.Schema] = None + if getattr(md, "arrow_schema_skipped", False): + # The crate skipped a non-UTF8 embedded arrow schema; recover the + # logical (extension-typed) schema via pyarrow so the decode can be + # realigned to it. If pyarrow can't read the footer either, fall the + # whole split back rather than emit storage-typed columns. + target_override = self._read_pyarrow_footer_schema(path) + if target_override is None: + return None + # M53: hand the crate the exact storage schema (extension storage + # layout, crate's own nested field names) so it decodes e.g. + # large_list offsets directly — turning the per-batch storage → + # extension realign from a Table.cast into a zero-copy relabel + # (_plan_batch_relabel). Zero IO: the footer is already parsed. + # Any rejection (older crate module, unsupported shape) keeps + # today's cast path unchanged. + override = _storage_override_schema(pa.schema(md), target_override) + if override is not None: + try: + handle.with_schema_override(override.__arrow_c_schema__()) + md = handle.metadata() + except Exception as e: # noqa: BLE001 - override is optional + logger.debug( + "arrow-rs schema override rejected for %s: %s", path, e + ) + return ( + handle, + pa.schema(md), + list(md.row_group_num_rows), + list(md.int96_columns), + target_override, + ) + + def _plan_native_read( + self, manifest: "FileManifest" + ) -> Optional[Tuple[List[Tuple[Any, int]], Optional[set], dict]]: + """Plan a native read: footer-read every file, decide native vs pyarrow + per file, and build the ordered ``[(fragment, file_row_offset)]`` list + plus the shared column split and scanner kwargs. Returns ``None`` to + signal a whole-split pyarrow fallback (some file's footer read failed).""" + import pyarrow.dataset as pds + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.chunkers.parquet_file_chunking_utils import ( # noqa: E501 + _fragments_from_row_group_ids, + ) + + unique_paths = list(dict.fromkeys(list(manifest.paths))) + + # Thrift footer limits, when set, decide whether each file is accepted + # or REJECTED — enforce them first with pyarrow's own parser so a + # too-large footer raises the identical OSError the base path would + # raise (parity-of-error), before any native work happens. + self._verify_footer_limits(unique_paths) + + # One footer read per file, through a per-file handle that the decode + # step reuses; S3 files additionally share one client per bucket, + # cached per process behind ``connect_native_s3`` (``s3_stores`` is + # just this plan's lookup). See :meth:`_open_native_file` / + # findings T10+M97. + s3_stores: Dict[str, Any] = {} + handle_by_path: Dict[str, Any] = {} + native_md: Dict[ + str, Tuple[pa.Schema, List[int], List[str], Optional[pa.Schema]] + ] = {} + for path in unique_paths: + opened = self._open_native_file(path, s3_stores) + if opened is None: + return None # whole-split pyarrow fallback + handle_by_path[path] = opened[0] + native_md[path] = opened[1:] + + # Column split, mirroring the base reader's ``dataset.schema.names``: + # with a pinned unified schema, ``pds.dataset(schema=...)`` reports + # exactly that schema — so a unified column absent from every file in + # this split still counts as on-disk (and gets null-filled per file by + # the alignment), instead of being silently dropped. Without a unified + # schema, fall back to the union of the files' footer schemas. + # Partition / path / row_hash columns aren't on disk anywhere and so + # land in the synthesize set either way. + if self._file_dataset_schema is not None: + on_disk_names = set(self._file_dataset_schema.names) + else: + on_disk_names = set() + for schema, _, _, _ in native_md.values(): + on_disk_names.update(schema.names) + columns_to_read_from_file, columns_to_synthesize = self._split_columns( + on_disk_names + ) + + scanner_kwargs = { + "columns": columns_to_read_from_file, + "filter": ( + self._predicate.to_pyarrow() if self._predicate is not None else None + ), + # The native decode re-derives its per-batch size by byte budget from + # the footer, so this is only an upper clamp; pyarrow-fallback + # fragments (nested/dictionary/extension) further clamp it themselves. + "batch_size": self._explicit_batch_size or _ARROW_DEFAULT_BATCH_SIZE, + "batch_readahead": _ARROW_SCANNER_BATCH_READAHEAD, + } + scanner_kwargs.update(self._arrow_scanner_kwargs()) + + read_columns = self._resolve_read_columns_for(scanner_kwargs) + + # Empty projection with no predicate (count-style scan): zero decode — + # the footer row counts already read above are exact, so emit + # count fragments for every file and never touch a data page. (With a + # predicate the count depends on the data; that case falls through to + # the per-file verdict below, which rejects empty projections.) + if ( + read_columns is not None + and len(read_columns) == 0 + and scanner_kwargs["filter"] is None + ): + count_fragments: List[Tuple[Any, int]] = [] + for path, chunk_metadata in zip( + manifest.paths, manifest.file_chunk_metadatas + ): + count_fragments.extend( + self._native_count_fragments( + path, + chunk_metadata, + native_md[path][1], + per_row_group_offsets=self._include_row_hash, + ) + ) + return count_fragments, columns_to_synthesize, scanner_kwargs + + # Per-file verdict: native decode (with an optional post-decode + # alignment plan) vs pyarrow fallback. + alignment_by_path: Dict[str, Optional[_ColumnAlignment]] = {} + for path, (schema, _, int96_cols, target_override) in native_md.items(): + alignment = self._plan_column_alignment( + schema, read_columns, int96_cols, target_schema=target_override + ) + if alignment is not None: + alignment_by_path[path] = None if alignment.is_noop else alignment + native_paths = set(alignment_by_path) + fallback_paths = [p for p in unique_paths if p not in native_paths] + if fallback_paths: + _raise_if_strict_no_fallback( + "no native column-alignment plan for file(s) " + f"{fallback_paths} (unplannable schema drift or unsupported " + "read-time coercion)" + ) + + # Build pyarrow fragments for the fallback files only (pyarrow never opens + # native files). One dataset over the fallback paths; the per-file fan-out + # reuses the base chunker helper so offsets / row-group slicing match the + # base path exactly. + fallback_fragment_by_path: dict = {} + if fallback_paths: + fb_dataset = pds.dataset( + source=fallback_paths, + format=self._make_format(), + filesystem=self._filesystem or LocalFileSystem(), + schema=self._file_dataset_schema, + ignore_prefixes=self._ignore_prefixes, + ) + fallback_fragment_by_path = { + frag.path: frag for frag in fb_dataset.get_fragments() + } + + fragments_with_offsets: List[Tuple[Any, int]] = [] + for path, chunk_metadata in zip(manifest.paths, manifest.file_chunk_metadatas): + if path in native_paths: + fragments_with_offsets.extend( + self._native_fragments_for_file( + path, + chunk_metadata, + native_md[path][1], + alignment_by_path[path], + handle_by_path[path], + per_row_group_offsets=self._include_row_hash, + ) + ) + else: + fragment = fallback_fragment_by_path[path] + if chunk_metadata is None: + fragments_with_offsets.append((fragment, 0)) + else: + fragments_with_offsets.extend( + _fragments_from_row_group_ids( + fragment, + chunk_metadata["row_group_ids"], + per_row_group_offsets=self._include_row_hash, + ) + ) + + return fragments_with_offsets, columns_to_synthesize, scanner_kwargs + + @staticmethod + def _native_fragments_for_file( + path: str, + chunk_metadata: Optional[dict], + row_group_num_rows: List[int], + alignment: Optional[_ColumnAlignment] = None, + handle: Optional[Any] = None, + *, + per_row_group_offsets: bool = False, + ) -> List[Tuple[_NativeParquetFragment, int]]: + """Build native fragments for one file, matching the base reader's + granularity so ``row_hash`` offsets are identical: + + - whole file (``chunk_metadata is None``) → one fragment over *all* row + groups at offset 0 (the base emits one whole-file fragment); + - a bin, ``per_row_group_offsets=False`` (**the common case**) → **one** + fragment naming all of the bin's row groups, at offset 0. The base + coalesces here so PyArrow can merge reads across the groups; we + coalesce so the crate makes one call instead of N, which is the whole + of old TODO 1l, obtained by following the base rather than inventing + our own coalescing; + - a bin, ``per_row_group_offsets=True`` (``include_row_hash``) → one + fragment per row group, each seeded with that group's **absolute** + pre-filter file row offset. + + The offsets in the fan-out case are absolute prefix sums indexed by + physical row-group id, *not* an accumulation across the bin's groups. + Upstream statistics pruning can leave the surviving set non-contiguous + (e.g. groups 0, 3, 7), and only the absolute position makes a row hash + match the row's true physical location — accumulating would silently + renumber every group after a pruned one. This mirrors ``prefix[rg_id]`` + in :func:`_fragments_from_row_group_ids`. + + ``alignment`` is the file's post-decode fixup plan, embedded in every + fragment so it survives the threaded fragment dispatch. + """ + if chunk_metadata is None: + return [(_NativeParquetFragment(path, None, alignment, handle), 0)] + + # The bin names the exact physical row groups for this file: predicate + # pruning and packing already happened upstream in ``ListFiles``, so there + # is no relative chunk descriptor left to reconcile against a row-group + # count, and no over-estimate that could silently drop a slice. + row_group_ids = sorted(chunk_metadata["row_group_ids"]) + if not row_group_ids: + return [] + + if not per_row_group_offsets: + return [(_NativeParquetFragment(path, row_group_ids, alignment, handle), 0)] + + # Absolute pre-filter offset at the start of each physical row group. + # The per-row-group fragments share ONE handle (one parsed footer), + # which is exactly the case the handle exists for. + prefix = [0] * (len(row_group_num_rows) + 1) + for i, num_rows in enumerate(row_group_num_rows): + prefix[i + 1] = prefix[i] + num_rows + return [ + (_NativeParquetFragment(path, [rg], alignment, handle), prefix[rg]) + for rg in row_group_ids + ] + + @staticmethod + def _native_count_fragments( + path: str, + chunk_metadata: Optional[dict], + row_group_num_rows: List[int], + *, + per_row_group_offsets: bool = False, + ) -> List[Tuple[_NativeCountFragment, int]]: + """Build zero-decode count fragments for one file (empty projection, no + predicate), at the same granularity/offsets as + :meth:`_native_fragments_for_file` so ``limit`` slicing and any + synthesized columns (``path``, partitions) behave identically. + + Row counts come from the footer we already read, summed over the bin's + named groups — so a count still touches no data page. + """ + if chunk_metadata is None: + return [(_NativeCountFragment(path, sum(row_group_num_rows)), 0)] + + # See ``_native_fragments_for_file``: the bin names the exact physical row + # groups, and offsets must be absolute so pruning gaps don't renumber. + row_group_ids = sorted(chunk_metadata["row_group_ids"]) + if not row_group_ids: + return [] + + if not per_row_group_offsets: + total = sum(row_group_num_rows[rg] for rg in row_group_ids) + return [(_NativeCountFragment(path, total), 0)] + + prefix = [0] * (len(row_group_num_rows) + 1) + for i, num_rows in enumerate(row_group_num_rows): + prefix[i + 1] = prefix[i] + num_rows + return [ + (_NativeCountFragment(path, row_group_num_rows[rg]), prefix[rg]) + for rg in row_group_ids + ] + + @cached_property + def _tuning(self) -> _ArrowRsTuning: + """Resolve the arrow-rs tuning knobs for this reader. + + Each knob comes from the ``arrow_rs_*`` key in ``dataset_kwargs`` when + present (popped into ``self._arrow_rs_tuning`` by + ``ParquetFileReader.__init__``), else from its ``RAY_DATA_ARROW_RS_*`` + env var, else the built-in default — see the "Tuning knobs" section at + the top of this module for what each knob does and how to tune it. A + ``None`` value means "use the default", consistent with the + format-kwarg convention. Invalid values raise loudly (a mis-set perf + knob must not silently degrade a benchmark or production read). + """ + + def resolve_optional( + key: str, default: Optional[int], minimum: int + ) -> Optional[int]: + value = self._arrow_rs_tuning.get(key) + if value is None: + return default + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError( + f"'{key}' in 'dataset_kwargs' must be an int, got {value!r}" + ) + if value < minimum: + raise ValueError( + f"'{key}' in 'dataset_kwargs' must be >= {minimum}, got {value}" + ) + return value + + def resolve(key: str, default: int, minimum: int) -> int: + # A non-None default guarantees a concrete int (the value is either + # that default or a validated int), so this never returns None. + resolved = resolve_optional(key, default, minimum) + assert resolved is not None + return resolved + + return _ArrowRsTuning( + decode_budget_bytes=resolve( + "arrow_rs_decode_budget_bytes", _default_decode_budget_bytes(), 1 + ), + k=resolve("arrow_rs_k", _ARROW_RS_K, 1), + split_threshold_bytes=resolve_optional( + "arrow_rs_split_threshold_bytes", None, 0 + ), + fetch_window_mb=resolve( + "arrow_rs_fetch_window_mb", _ARROW_RS_FETCH_WINDOW_MB, 0 + ), + column_fetch_mb=resolve( + "arrow_rs_column_fetch_mb", _ARROW_RS_COLUMN_FETCH_MB, 0 + ), + prefetch_budget_mb=resolve_optional( + "arrow_rs_prefetch_budget_mb", + None + if _ARROW_RS_PREFETCH_BUDGET_MB < 0 + else _ARROW_RS_PREFETCH_BUDGET_MB, + 0, + ), + ) + + @cached_property + def _pushdown_predicate_json(self) -> Optional[str]: + """The pushed predicate lowered to the native pruning IR (JSON), or + ``None`` when there's nothing prunable. Depends only on + ``self._predicate``, so it's computed once and reused for every + fragment. See :func:`_predicate_to_ir` for the (total, conservative) + lowering and the soundness argument.""" + return _predicate_json(self._predicate) + + def _filesystem_supported(self) -> bool: + """Whether the native crate can read from this filesystem at all. + Local and S3 are wired in `_iter_fragment_tables` / the native `read()` + (S3 uses the windowed, byte-budgeted native path). Any other + filesystem (GCS, ABFS, HTTP, …) falls back to PyArrow.""" + from pyarrow.fs import LocalFileSystem, S3FileSystem + + # ``None`` means the default local filesystem (matching + # ``native_metadata_supported_filesystem`` and the non-S3 native read + # path); treat it as supported so eligible local reads don't silently + # fall back to PyArrow. + return self._filesystem is None or isinstance( + self._filesystem, (LocalFileSystem, S3FileSystem) + ) + + def _blocking_format_kwargs(self, aligned_ok: bool) -> Dict[str, Any]: + """Parquet-format kwargs (the ``dataset_kwargs`` payload spread into + ``pds.ParquetFileFormat``) that the native path cannot honor — a + non-empty result forces a PyArrow fallback, which honors them all. + + The audit rule is an explicit ALLOWLIST, so a format kwarg added by a + future pyarrow version is *unsupported until proven supported* — never + silently ignored (e.g. pyarrow 21+'s ``binary_type`` / ``list_type`` / + ``arrow_extensions_enabled`` change the decoded schema, and + ``thrift_string_size_limit`` changes which files are *rejected*, so + ignoring any of them would diverge from the PyArrow paths): + + - :data:`_FORMAT_KWARGS_PERF_ONLY` (``pre_buffer``, ``buffer_size``, + ``use_buffered_stream``, ``cache_options``) tune PyArrow's I/O + strategy only and cannot change decoded bytes; the crate has its own + I/O strategy (byte-budgeted streaming + fetch window), so they are + safely ignorable natively. + - :data:`_FORMAT_KWARGS_ALIGNED` (``coerce_int96_timestamp_unit``, + ``dictionary_columns``) are reproduced by the *planned* path via + :class:`_ColumnAlignment`, and + :data:`_FORMAT_KWARGS_FOOTER_VERIFIED` (the thrift limits) are + enforced by the planned path's pyarrow footer probe + (:meth:`_verify_footer_limits`) — both admitted only with + ``aligned_ok=True``. The per-fragment re-gate can plan neither an + alignment nor a probe (see :meth:`_reader_level_supported`), so + there they block (``aligned_ok=False``). + - :data:`_FORMAT_KWARGS_SCHEMA_SHAPED` (``binary_type``, + ``list_type``) are admitted on the planned path only when a unified + dataset schema is pinned: the pin is the output-type authority (it + silently *undoes* these kwargs on the base path too), and the + alignment's drift casts already produce the pinned types. Without a + pinned schema the kwargs genuinely change output types — fall back. + - ``page_checksum_verification=True`` is admitted everywhere: the + crate is built with parquet's ``crc`` feature and always verifies + stored page CRCs, so ``True`` *is* the native behavior. An explicit + ``False`` — the opt-out for reading a file despite corrupt + checksums — is something the crate build cannot honor, so it falls + back to PyArrow (the only reader that can skip the check). + - A ``None`` value means "pyarrow default" for every format kwarg, so + ``None``-valued keys never block. + """ + allowed = _FORMAT_KWARGS_PERF_ONLY | ( + (_FORMAT_KWARGS_ALIGNED | _FORMAT_KWARGS_FOOTER_VERIFIED) + if aligned_ok + else frozenset() + ) + if aligned_ok and self._file_dataset_schema is not None: + allowed |= _FORMAT_KWARGS_SCHEMA_SHAPED + blocked: Dict[str, Any] = {} + for key, value in self._parquet_format_kwargs.items(): + if value is None or key in allowed: + continue + if key == "page_checksum_verification" and value is True: + continue + blocked[key] = value + return blocked + + def _reader_level_supported(self) -> bool: + """Reader-wide half of the *per-fragment* re-gate (the pyarrow-fragment + path in :meth:`_iter_fragment_tables`): filesystem + Parquet-format + kwargs. The aligned-kwarg checks stay here — not in the planned native + ``read()`` — because the per-fragment path has no crate footer metadata + to plan a :class:`_ColumnAlignment` from (a pyarrow + ``physical_schema`` already reflects ``coerce_int96_timestamp_unit`` / + ``dictionary_columns``, so an alignment computed from it would be a + false no-op). The planned path handles both kwargs natively via + :meth:`_plan_column_alignment`.""" + if not self._filesystem_supported(): + return False + if self._blocking_format_kwargs(aligned_ok=False): + return False + return True + + def _columns_supported( + self, + physical_schema: pa.Schema, + read_columns: Optional[List[str]], + int96_columns: Optional[List[str]] = None, + ) -> bool: + """Per-fragment re-gate verdict: native only when the decode needs *no* + post-decode fixups. Used by the pyarrow-fragment path, where the + alignment can't be trusted (see :meth:`_reader_level_supported`); the + planned ``read()`` instead admits any file with a plannable + :class:`_ColumnAlignment`.""" + alignment = self._plan_column_alignment( + physical_schema, read_columns, int96_columns + ) + return alignment is not None and alignment.is_noop + + def _plan_column_alignment( + self, + physical_schema: pa.Schema, + read_columns: Optional[List[str]], + int96_columns: Optional[List[str]] = None, + target_schema: Optional[pa.Schema] = None, + ) -> Optional[_ColumnAlignment]: + """Per-file half of the support gate, upgraded from a yes/no verdict to + a *plan*: how to make this file's native decode match what the pyarrow + scanner would produce. Returns ``None`` for a pyarrow fallback, a no-op + alignment for a byte-identical native decode, or a fixup plan + (null-fill / cast / reorder) the decode path applies per batch. + + Takes a ``pa.Schema`` (the crate's ``read_metadata`` schema — i.e. what + the crate will actually decode) plus, optionally, the root column names + the crate reports as INT96-physical. Still conservative — anything not + covered falls back to PyArrow, so correctness is never at risk: + + - empty projection (count scan) → handled upstream by the zero-decode + count path (:class:`_NativeCountFragment`) when there's no predicate; + ``None`` here so the per-fragment re-gate keeps PyArrow's stub dance; + - a column absent from this file (schema evolution) → **null-fill** + with the unified type (``None`` when there's no unified schema to + take the type from, or the fill type is an extension); + - an INT96 column → **cast** to timestamp[ns] (PyArrow's default; an + exact upcast from any embedded hint unit). When + ``coerce_int96_timestamp_unit`` is set the file **falls back** + instead: decode-time coercion floors, a post-decode cast truncates + toward zero — irreconcilable on pre-1970 values. Non-timestamp / + tz-carrying INT96 oddities stay on PyArrow; + - a forced ``dictionary_columns`` read → **cast** to + ``dictionary`` (what PyArrow's forced-dict decode + yields); non-string/binary targets stay on PyArrow; + - a per-file type that differs from the unified schema → **cast** to + the unified type (the scanner's implicit cast under a pinned + schema); extension-typed drift stays on PyArrow. + """ + # ``target_schema`` overrides the reader-wide pin for this one file. It's + # supplied when the crate skipped the embedded arrow schema (non-UTF8 + # extension metadata): the file's own pyarrow-read footer schema carries + # the reconstructed extension types the decode must be realigned to, + # even though the reader-wide pin (``_file_dataset_schema``) is ``None`` + # for extension-bearing reads. + unified_schema = ( + target_schema if target_schema is not None else self._file_dataset_schema + ) + int96 = set(int96_columns or ()) + # The expected output columns: the explicit read set when projected; + # otherwise the *unified* schema's columns (the scanner outputs the + # pinned dataset schema, null-filling what a file lacks — using the + # file's own names here would silently drop evolved columns); the + # file's names only when there is no unified schema to pin. + if read_columns is not None: + names = read_columns + elif unified_schema is not None: + names = list(unified_schema.names) + else: + names = list(physical_schema.names) + + if read_columns is not None and len(read_columns) == 0: + return None + + coerce_unit = self._parquet_format_kwargs.get("coerce_int96_timestamp_unit") + dictionary_columns = set( + self._parquet_format_kwargs.get("dictionary_columns") or () + ) + + null_fill: List[Tuple[str, pa.DataType]] = [] + casts: List[Tuple[str, pa.DataType, bool]] = [] + for name in names: + idx = physical_schema.get_field_index(name) + if idx == -1: + # Column absent from this file (schema evolution): null-fill + # with the unified type — exactly pyarrow's behavior under a + # pinned dataset schema. Without a unified schema the fill type + # is unknowable — defer to PyArrow. + if unified_schema is None: + return None + unified_idx = unified_schema.get_field_index(name) + if unified_idx == -1: + return None + fill_type = unified_schema.field(unified_idx).type + if _is_extension_type(fill_type): + return None + null_fill.append((name, fill_type)) + continue + + field_type = physical_schema.field(idx).type + target = field_type + allow_time_truncate = False + if name in int96: + if coerce_unit is not None: + # A cast cannot reproduce pyarrow's decode-time coercion: + # pyarrow FLOORS (parquet types.h Int96GetXxx divide the + # unsigned nanos-of-day before adding the signed day + # offset) while a post-decode cast truncates the signed + # total toward zero — one unit apart on every pre-1970 + # value with a sub-unit remainder (measured: all 1715 + # negative values in the 1964 corpus fixture). The kwarg + # is honored by decoding this file via pyarrow instead. + return None + # No kwarg: pyarrow decodes INT96 to timestamp[ns, no tz]; the + # crate instead honors an embedded non-ns arrow-schema hint. + # Realign by casting to ns — an exact upcast (multiplication), + # never a truncation. + if not (pa.types.is_timestamp(target) and target.tz is None): + return None # nested/tz-carrying INT96 oddity — stay safe + target = pa.timestamp("ns") + if name in dictionary_columns: + # PyArrow's forced dictionary decode yields + # dictionary. Only string/binary + # columns are dictionary-read by pyarrow's parquet layer. + if not (pa.types.is_string(target) or pa.types.is_binary(target)): + return None + target = pa.dictionary(pa.int32(), target) + if unified_schema is not None: + unified_idx = unified_schema.get_field_index(name) + if unified_idx != -1: + unified_type = unified_schema.field(unified_idx).type + if unified_type != target: + # Per-file drift vs the pinned unified schema: the + # scanner casts implicitly; mirror it. + if _is_extension_type(unified_type) and not _is_extension_type( + target + ): + # The crate decoded the parquet *storage* type because + # the embedded arrow schema was skipped (non-UTF8 + # extension metadata, e.g. Ray's cloudpickle tensor + # type). Reconstruct the extension by casting storage + # -> extension: pyarrow does the offset-width change + # and the extension wrap in one cast, matching what + # the base scanner reconstructs from the same footer + # (verified byte-identical). + target = unified_type + elif _is_extension_type(unified_type) or _is_extension_type( + target + ): + # extension<->extension drift (e.g. per-file tensor + # shapes) isn't a safe cast — fall back. + return None + else: + target = unified_type + + if target != field_type: + casts.append((name, target, allow_time_truncate)) + + if not null_fill and not casts: + return _NOOP_ALIGNMENT + # Appended null-fill columns must land in read order (the scanner's + # projected-column order); reordering is only needed when a column was + # appended. + order = tuple(names) if null_fill else None + return _ColumnAlignment( + null_fill=tuple(null_fill), casts=tuple(casts), order=order + ) + + def _arrow_rs_supported( + self, + fragment: pds.ParquetFileFragment, + read_columns: Optional[List[str]], + ) -> bool: + """Whole-gate verdict for a pyarrow fragment: reader-level checks plus + the per-file column/type checks against the fragment's physical schema. + Used by the per-fragment ``_iter_fragment_tables`` path. + + A fragment's ``physical_schema`` is PyArrow's *post-coercion* Arrow schema, + so an INT96 column already reads as ``timestamp[ns]`` and can't reveal + whether the crate would decode it differently (it honors an embedded + non-ns hint). This re-gate can't see the crate's output, so it is + conservative: any INT96-physical read column falls the fragment back to + PyArrow. The authoritative plan-time gate (:meth:`_columns_supported` with + the crate's ``int96_columns``) is what admits INT96→ns files to the native + path; this path only ever *withholds*, never wrongly admits. + """ + if not self._reader_level_supported(): + return False + int96_roots = _pyarrow_fragment_int96_roots(fragment) + if int96_roots: + names = ( + read_columns + if read_columns is not None + else list(fragment.physical_schema.names) + ) + if any(name in int96_roots for name in names): + return False + return self._columns_supported(fragment.physical_schema, read_columns) + + @staticmethod + def _import_extension(): + """Import the native extension, raising a clear, actionable error if it + isn't built. Called on every native entry point so a missing module + surfaces loudly (never a silent fall back to PyArrow, which would + corrupt benchmark attribution).""" + try: + import ray_data_arrow_rs + + return ray_data_arrow_rs + except ImportError as e: + raise ImportError( + "use_arrow_rs_parquet_reader=True requires the " + "'ray_data_arrow_rs' extension. Build it with " + "`maturin develop --release` from " + "python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs/." + ) from e + + def _resolve_read_columns_for(self, scanner_kwargs: dict) -> Optional[List[str]]: + """The set of columns the native decode must read from the file: the + projected columns plus any columns referenced only by the pushed filter + (which we still filter on post-decode). ``None`` means all columns.""" + from ray.data._internal.datasource.parquet_datasource import ( + _resolve_read_columns, + ) + from ray.data._internal.planner.plan_expression.expression_visitors import ( + get_column_references, + ) + + columns = scanner_kwargs.get("columns") + filter_expr = scanner_kwargs.get("filter") + filter_columns = ( + get_column_references(self._predicate) + if self._predicate is not None + else None + ) + return _resolve_read_columns(columns, filter_expr, filter_columns) + + @override + def _iter_fragment_tables( + self, + fragment: pds.Fragment, + scanner_kwargs: dict, + ) -> "Iterator[pa.Table]": + # Native front-end (arrow-rs ``read()``) hands us pyarrow-free work units. + if isinstance(fragment, _NativeCountFragment): + # Empty projection, no predicate: the footer count is exact — yield + # a zero-column table with the right num_rows and decode nothing. + # (``Table.select([])`` preserves ``num_rows``; the base path's + # zero-column tables flow through ``_postprocess`` identically.) + _trace_reader_path(True) + if fragment.num_rows > 0: + yield pa.table({"__num_rows": pa.nulls(fragment.num_rows)}).select([]) + return + if isinstance(fragment, _NativeParquetFragment): + _trace_reader_path(True) + yield from self._iter_native_tables( + fragment.path, + fragment.row_groups, + scanner_kwargs, + alignment=fragment.alignment, + handle=fragment.handle, + ) + return + + # Pyarrow fragment (used when ``read()`` is not overridden, e.g. the + # reader-level-unsupported delegate, and by unit tests that drive this + # method directly). Re-check the per-fragment gate and either decode + # natively or fall back to the PyArrow scanner. + read_columns = self._resolve_read_columns_for(scanner_kwargs) + supported = self._arrow_rs_supported(fragment, read_columns) + _trace_reader_path(supported) + if not supported: + _raise_if_strict_no_fallback( + f"fragment {fragment.path!r} rejected by the per-fragment " + "support gate" + ) + yield from super()._iter_fragment_tables(fragment, scanner_kwargs) + return + + row_groups = ( + [rg.id for rg in fragment.row_groups] + if fragment.row_groups is not None + else None + ) + # The gate admits this fragment as a byte-identical native decode, but the + # crate may still return a storage type for an extension column whose + # binary metadata it had to skip. Reconcile to the fragment's own arrow + # (post-coercion) schema so that legacy path can't silently emit storage. + yield from self._iter_native_tables( + fragment.path, + row_groups, + scanner_kwargs, + expected_schema=fragment.physical_schema, + ) + + def _iter_native_tables( + self, + path: str, + row_groups: Optional[List[int]], + scanner_kwargs: dict, + alignment: Optional[_ColumnAlignment] = None, + expected_schema: Optional[pa.Schema] = None, + handle: Optional[Any] = None, + ) -> "Iterator[pa.Table]": + """Decode ``row_groups`` of ``path`` via the native crate and yield + ``pa.Table`` batches, applying the file's :class:`_ColumnAlignment` + (null-fill / cast / reorder) to each batch *before* the post-decode + filter so the predicate sees the same types pyarrow's scanner filters + on. + + Row-group pruning is native (``predicate.rs``), replacing PyArrow's + ``fragment.subset(filter=...)``: we hand the crate the row-group ids plus + the pushed predicate lowered to a JSON IR, and it drops the groups whose + footer statistics prove no row can match before fetching or decoding + them. Pruning is conservative by construction — a missing column, absent + stats, or an uncomparable type all *keep* the group — so it can only ever + avoid IO/decode, never change which rows surface. Row-level filtering + then runs post-decode here (the final authority), and a fully-pruned file + simply yields nothing. + """ + ray_data_arrow_rs = self._import_extension() + + from pyarrow.fs import S3FileSystem + + batch_size = scanner_kwargs.get("batch_size") or _ARROW_DEFAULT_BATCH_SIZE + read_columns = self._resolve_read_columns_for(scanner_kwargs) + predicate_json = self._pushdown_predicate_json + + tuning = self._tuning + split_threshold = tuning.split_threshold_bytes + if split_threshold is None: + split_threshold = ( + self._target_block_size + if self._target_block_size is not None + else _ARROW_RS_DEFAULT_SPLIT_THRESHOLD_BYTES + ) + + # The prefetch bucket defaults to ~4 units: one decoding + ~3 in flight + # keeps the (single) decoder fed across fetch:decode ratios without a + # second knob to tune. Units are row windows (fetch_window_mb) or column + # groups (column_fetch_mb), so the bucket scales with the larger + # unit-size knob. + prefetch_budget_mb = ( + tuning.prefetch_budget_mb + if tuning.prefetch_budget_mb is not None + else 4 * max(tuning.fetch_window_mb, tuning.column_fetch_mb) + ) + + if handle is not None: + # Planned-path decode (TODO 1r): the footer was parsed — and for S3 + # the client built — once at plan time; this call reuses both. The + # handle knows its own transport, so there is no local-vs-S3 branch. + reader = handle.read_row_groups( + row_groups=row_groups, + columns=read_columns, + batch_size=batch_size, + decode_budget_bytes=tuning.decode_budget_bytes, + k=tuning.k, + split_threshold_bytes=split_threshold, + predicate_json=predicate_json, + fetch_window_mb=tuning.fetch_window_mb, + column_fetch_mb=tuning.column_fetch_mb, + prefetch_budget_mb=prefetch_budget_mb, + ) + yield from self._yield_native_batches( + reader, scanner_kwargs, alignment, expected_schema + ) + return + + fs = self._filesystem + if isinstance(fs, S3FileSystem): + bucket, key = _split_s3_path(path) + cfg = _s3_config(fs) + reader = ray_data_arrow_rs.read_row_groups_s3( + bucket, + key, + cfg["region"], + cfg["anonymous"], + endpoint=cfg["endpoint"], + access_key_id=cfg["access_key_id"], + secret_access_key=cfg["secret_access_key"], + session_token=cfg["session_token"], + allow_http=cfg["allow_http"], + virtual_hosted_style=cfg["virtual_hosted_style"], + row_groups=row_groups, + columns=read_columns, + batch_size=batch_size, + decode_budget_bytes=tuning.decode_budget_bytes, + fetch_window_mb=tuning.fetch_window_mb, + k=tuning.k, + split_threshold_bytes=split_threshold, + predicate_json=predicate_json, + column_fetch_mb=tuning.column_fetch_mb, + prefetch_budget_mb=prefetch_budget_mb, + ) + else: + reader = ray_data_arrow_rs.read_row_groups( + path, + row_groups, + read_columns, + batch_size, + tuning.decode_budget_bytes, + tuning.k, + split_threshold, + predicate_json, + ) + + yield from self._yield_native_batches( + reader, scanner_kwargs, alignment, expected_schema + ) + + def _yield_native_batches( + self, + reader: Any, + scanner_kwargs: dict, + alignment: Optional[_ColumnAlignment], + expected_schema: Optional[pa.Schema], + ) -> "Iterator[pa.Table]": + """Consume a crate stream and yield aligned, filtered ``pa.Table`` + batches. Shared by the handle path and the per-call entry points.""" + columns = scanner_kwargs.get("columns") + filter_expr = scanner_kwargs.get("filter") + + record_batch_reader = pa.RecordBatchReader.from_stream(reader) + + # M53 fast path: when the only type realign is storage → extension + # labels over layout-identical buffers (true after the crate schema + # override in _open_native_file), relabel each batch through the C + # Data Interface (~2-4 ms at 5000 columns) instead of casting (~40 ms). + # Any per-batch relabel failure downgrades the rest of the stream to + # the cast path — both produce bit-identical tables. + relabel_schema: Optional[pa.Schema] = None + residual_alignment = alignment + skip_reconcile = False + if _ffi_relabel_enabled(): + relabel_plan = _plan_batch_relabel( + record_batch_reader.schema, alignment, expected_schema + ) + if relabel_plan is not None: + relabel_schema, residual_alignment, skip_reconcile = relabel_plan + + # Yield each budget-sized batch straight through. The read op's + # BlockOutputBuffer coalesces to target_max_block_size downstream (same + # as the PyArrow path) — accumulating a full block here too would just + # stack a second block-sized buffer on top of it. See module docstring. + pickle_checked = False + for batch in record_batch_reader: + if relabel_schema is not None: + try: + batch = _ffi_relabel_batch(batch, relabel_schema) + except Exception as e: # noqa: BLE001 - cast path still correct + logger.debug("arrow-rs FFI relabel failed, cast fallback: %s", e) + relabel_schema = None + residual_alignment = alignment + skip_reconcile = False + if relabel_schema is not None: + table = pa.Table.from_batches([batch]) + table = _apply_column_alignment(table, residual_alignment) + else: + table = pa.Table.from_batches( + [batch], schema=record_batch_reader.schema + ) + table = _apply_column_alignment(table, alignment) + if expected_schema is not None and not skip_reconcile: + table = _reconcile_to_expected(table, expected_schema) + # Same opt-in gate as the pyarrow path: unpickling an + # ArrowPythonObjectType column executes arbitrary code, so serving + # one requires the explicit env opt-in. raise_on_pickle_object_columns + # itself no-ops when RAY_DATA_AUTOLOAD_PICKLE_OBJECT_SCALAR=1. Before + # the row filter — the check is schema-based, and pyarrow's scanner + # raises even for batches the filter would empty out. Once per + # stream, not per batch: a crate stream has one schema by + # construction and the alignment plan is fixed per file, so every + # batch here carries the schema the first one did (walking 5000 + # extension fields per batch was 18% of the tensor-shape read, T23). + if not pickle_checked: + raise_on_pickle_object_columns(table) + pickle_checked = True + if filter_expr is not None: + table = table.filter(filter_expr) + if table.num_rows == 0: + continue + if columns is not None: + table = table.select([c for c in columns if c in table.column_names]) + yield table diff --git a/python/ray/data/_internal/datasource_v2/readers/base_reader.py b/python/ray/data/_internal/datasource_v2/readers/base_reader.py index a117c5524a06..3c936943426d 100644 --- a/python/ray/data/_internal/datasource_v2/readers/base_reader.py +++ b/python/ray/data/_internal/datasource_v2/readers/base_reader.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import Generic, Iterator +from typing import Dict, Generic, Iterator import pyarrow as pa @@ -18,6 +18,15 @@ class Reader(ABC, Generic[InputSplit]): pushdown optimizations (columns, predicates, limits) that were applied. """ + def pop_task_stats(self) -> Dict[str, float]: + """Reader-specific per-task counters, drained (reset to zero) by the call. + + Keys are ``ReadFilesTaskStats`` field names (currently ``trim_wall_s``); + the ``ReadFiles`` transform calls this once each ``read()`` stream has + ended and adds the values into the task's stats. Default: nothing. + """ + return {} + @abstractmethod def read(self, input_split: InputSplit) -> Iterator[pa.Table]: """Read data from the input bucket and yield Arrow tables. diff --git a/python/ray/data/_internal/datasource_v2/readers/file_reader.py b/python/ray/data/_internal/datasource_v2/readers/file_reader.py index b78fc5a8eb4b..2ae9666abbdd 100644 --- a/python/ray/data/_internal/datasource_v2/readers/file_reader.py +++ b/python/ray/data/_internal/datasource_v2/readers/file_reader.py @@ -35,14 +35,6 @@ "RAY_DATA_ARROW_SCANNER_BATCH_READAHEAD", 8 ) -# Number of worker threads used to read fragments concurrently per task. -# Defaults to 4 to overlap remote-filesystem I/O latency across multiple -# fragments. ``_read_fragment_batches`` caps this to ``len(fragments)`` -# at runtime so single-fragment tasks don't spin up extra workers, and -# falls back to the sequential path entirely when -# ``DataContext.execution_options.preserve_order`` is set. -_DEFAULT_NUM_THREADS = env_integer("RAY_DATA_READ_FILES_NUM_THREADS", 4) - ROW_HASH_COLUMN_NAME = "row_hash" @@ -229,33 +221,62 @@ def read(self, input_split: FileManifest) -> Iterator[pa.Table]: # Split the requested columns into ones the on-disk file has # (pyarrow reads these) and ones we need to synthesize post-read - # (hive partition keys, "path"). ``self._columns is None`` means - # "no projection" — read every file column and synthesize every - # available partition/path column. + # (hive partition keys, "path"). on_disk_column_names = set(dataset.schema.names) - if self._columns is None: - columns_to_read_from_file: Optional[List[str]] = None - columns_to_synthesize: Optional[Set[str]] = None - else: - columns_to_read_from_file = [ - c for c in self._columns if c in on_disk_column_names - ] - columns_to_synthesize = set(self._columns) - on_disk_column_names + columns_to_read_from_file, columns_to_synthesize = self._split_columns( + on_disk_column_names + ) scanner_kwargs = { "columns": columns_to_read_from_file, "filter": ( self._predicate.to_pyarrow() if self._predicate is not None else None ), - "batch_size": self._resolve_batch_size(dataset), + "batch_size": self._resolve_batch_size(dataset, input_split), "batch_readahead": _ARROW_SCANNER_BATCH_READAHEAD, } scanner_kwargs.update(self._arrow_scanner_kwargs()) + triples = self._read_fragment_batches(dataset, scanner_kwargs, input_split) + yield from self._postprocess(triples, columns_to_synthesize) + + def _split_columns( + self, on_disk_column_names: Set[str] + ) -> Tuple[Optional[List[str]], Optional[Set[str]]]: + """Split the requested ``columns`` into (read-from-file, synthesize). + + ``self._columns is None`` means "no projection" — read every file + column and synthesize every available partition/path column, so both + halves are ``None``. Otherwise the requested columns are partitioned by + whether they exist on disk; the rest (hive partition keys, ``path``, + ``row_hash``) are synthesized post-read. Shared by the base pyarrow + ``read()`` and subclass front-ends (e.g. the arrow-rs reader). + """ + if self._columns is None: + return None, None + columns_to_read_from_file = [ + c for c in self._columns if c in on_disk_column_names + ] + columns_to_synthesize = set(self._columns) - on_disk_column_names + return columns_to_read_from_file, columns_to_synthesize + + def _postprocess( + self, + triples: Iterator[Tuple[pa.Table, str, int]], + columns_to_synthesize: Optional[Set[str]], + ) -> Iterator[pa.Table]: + """Apply the format-agnostic per-table finishing steps to a stream of + ``(table, fragment_path, fragment_row_offset)`` triples: ``limit`` + slicing, partition/``path`` synthesis, ``row_hash`` synthesis, column + projection/reordering, and the zero-column stub guard. Yields the + finished tables. + + Factored out of :meth:`read` so any front-end that can produce the same + triples — pyarrow scanners *or* the native arrow-rs decode path — reuses + identical post-read semantics (row hashes, partition columns, limit). + """ rows_read = 0 - for table, fragment_path, fragment_row_offset in self._read_fragment_batches( - dataset, scanner_kwargs, input_split - ): + for table, fragment_path, fragment_row_offset in triples: if self._limit is not None: if rows_read >= self._limit: break @@ -326,10 +347,11 @@ def read(self, input_split: FileManifest) -> Iterator[pa.Table]: rows_read += len(table) yield table - def _resolve_batch_size(self, dataset: pds.Dataset) -> int: + def _resolve_batch_size(self, dataset: pds.Dataset, manifest: FileManifest) -> int: """Return the batch size to use for scanning. - Subclasses can override this to implement adaptive batch sizing. + Subclasses can override this to implement adaptive batch sizing, using + the ``manifest`` (e.g. footer-derived per-chunk stats) to avoid re-I/O. """ return self._batch_size @@ -420,8 +442,6 @@ def _read_fragment_batches( which we materialize below anyway. File data is still read lazily by the worker threads. """ - ctx = DataContext.get_current() - # ``preserve_ordering=True`` would drain the input iterator # eagerly anyway, so materialize once here to (a) cap # ``num_workers`` at the actual fragment count and (b) avoid @@ -430,10 +450,53 @@ def _read_fragment_batches( # ``_get_fragments_to_read`` to fan out chunk-level # sub-fragments from the manifest's chunk metadata. fragments_with_offsets = self._get_fragments_to_read(dataset, manifest) + yield from self._dispatch_fragment_reads(fragments_with_offsets, scanner_kwargs) + + def _num_fragment_read_threads(self, num_fragments: int) -> int: + """How many fragments this reader decodes concurrently within one read task. + + Overridable because the right answer is reader-specific rather than a + property of the file format. PyArrow's per-fragment decode gives a read + task no intra-task parallelism of its own, so this pool is the only source + of it; a reader that already parallelises *inside* a fragment gets nothing + from the pool but still pays for every fragment it keeps in flight. See + ``ArrowRsParquetFileReader._num_fragment_read_threads``. + + Returning 1 is not merely "less concurrency": ``_dispatch_fragment_reads`` + takes the sequential branch and ``make_async_gen`` is never constructed at + all. + + The base returns ``num_fragments`` — one worker per fragment, unbounded — + which is exactly what the footer-chunking base path does inline. It takes + the count as an argument for that reason: "unbounded" is not expressible as + a constant, and the previous ``_DEFAULT_NUM_THREADS`` (4) no longer exists + upstream. Hard-coding any cap here would quietly throttle the PyArrow path + below its own default and bias every A/B against it. + """ + return num_fragments + + def _dispatch_fragment_reads( + self, + fragments_with_offsets: List[Tuple[pds.Fragment, int]], + scanner_kwargs: dict, + ) -> Iterator[Tuple[pa.Table, str, int]]: + """Read a pre-materialized ``[(fragment, file_row_offset)]`` list, either + sequentially or across a bounded thread pool, preserving fragment order. + + Split out of :meth:`_read_fragment_batches` so front-ends that build the + fragment list some other way — e.g. the arrow-rs reader assembling native + and pyarrow-fallback fragments per file — reuse the same threading, retry, + and ordering semantics. ``fragment`` need only expose ``.path`` and be + accepted by :meth:`_iter_fragment_tables`. + """ if not fragments_with_offsets: return - num_workers = min(_DEFAULT_NUM_THREADS, len(fragments_with_offsets)) + ctx = DataContext.get_current() + num_workers = min( + self._num_fragment_read_threads(len(fragments_with_offsets)), + len(fragments_with_offsets), + ) if num_workers <= 1 or ctx.execution_options.preserve_order: yield from self._read_fragments_sequential( iter(fragments_with_offsets), scanner_kwargs diff --git a/python/ray/data/_internal/datasource_v2/readers/parquet_file_reader.py b/python/ray/data/_internal/datasource_v2/readers/parquet_file_reader.py index 38f4c7b7ea86..d62791c49ae8 100644 --- a/python/ray/data/_internal/datasource_v2/readers/parquet_file_reader.py +++ b/python/ray/data/_internal/datasource_v2/readers/parquet_file_reader.py @@ -1,5 +1,6 @@ import logging import math +import os from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Set, Tuple @@ -14,7 +15,7 @@ from ray._common.utils import env_integer from ray.data._internal.datasource_v2.chunkers.parquet_file_chunking_utils import ( - _fragments_from_chunk_metadata, + _fragments_from_row_group_ids, ) from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest from ray.data._internal.datasource_v2.readers.file_reader import ( @@ -40,6 +41,25 @@ _UNSET = object() +# arrow-rs tuning knobs. These travel on the V2 ``dataset_kwargs`` payload (see +# ``arrow_rs_parquet_file_reader.py`` for semantics/defaults) but must never +# reach ``pds.ParquetFileFormat``, so ``__init__`` pops them out of +# ``parquet_format_kwargs`` before PyArrow can see them. The PyArrow reader then +# ignores them entirely — the mirror image of the native reader ignoring +# PyArrow's I/O-only kwargs — so a ``read_parquet`` call carrying these keys +# stays valid whichever reader the ``use_arrow_rs_parquet_reader`` flag selects. +ARROW_RS_TUNING_KWARGS = frozenset( + { + "arrow_rs_decode_budget_bytes", + "arrow_rs_k", + "arrow_rs_split_threshold_bytes", + "arrow_rs_fetch_window_mb", + "arrow_rs_column_fetch_mb", + "arrow_rs_prefetch_budget_mb", + } +) +_ARROW_RS_KWARG_PREFIX = "arrow_rs_" + # Per-stream read-ahead buffer for ``use_buffered_stream=True``. PyArrow's # default (~8 KiB) produces many tiny range requests on S3; 8 MiB # amortizes per-request latency across meaningful payload sizes. @@ -47,6 +67,26 @@ "RAY_DATA_PARQUET_FRAGMENT_BUFFER_SIZE", 8 * MiB ) +# Experimental override for pyarrow's ``pre_buffer`` scan option. Unset (the +# default) leaves pyarrow's own default (``True``) in place — identical +# behavior to before this knob existed. ``0`` disables the coalesced +# whole-fragment prefetch; with the footer-based planner a fragment is all of +# a file's bin-assigned row groups, so ``pre_buffer=True`` stages that entire +# compressed span per fragment. Exists for the bin-size sweep's memory +# attribution; disabling it trades memory for (potentially many) uncoalesced +# range requests, so it is not a recommended production setting. +_PARQUET_PRE_BUFFER_OVERRIDE: Optional[int] = ( + env_integer("RAY_DATA_PARQUET_PRE_BUFFER", -1) + if "RAY_DATA_PARQUET_PRE_BUFFER" in os.environ + else None +) + +# Arrow process-wide IO / CPU thread pools for the read task. Arrow's default +# (~num cores, 8 IO threads) throttles the concurrent column/range fetches a +# Parquet scan issues against S3, especially for row-group-scoped fragments. +_READER_IO_THREAD_COUNT = env_integer("RAY_DATA_PARQUET_READER_IO_THREAD_COUNT", 128) +_READER_CPU_COUNT = env_integer("RAY_DATA_PARQUET_READER_CPU_COUNT", 128) + def _estimate_batch_size_from_metadata( fragment: pds.ParquetFileFragment, @@ -137,6 +177,32 @@ def _estimate_batch_size_from_metadata( return target_batch_size +def _estimate_batch_size_from_chunk_stats( + uncompressed_size: int, + num_rows: int, + target_block_size: int, +) -> Optional[int]: + """Estimate batch size from footer-derived chunk stats, without any I/O. + + ``ListFiles`` already read each file's footer and recorded the + projection-scoped uncompressed byte size and row count of the row groups it + assigned to this chunk (:class:`ParquetRowGroupChunkMetadata`). Sizing from + those avoids the extra footer read that + :func:`_estimate_batch_size_from_metadata` incurs. Mirrors that function's + math but over the whole chunk (its row-group average) rather than the first + row group; the estimate is refined from real data after the first batch. + """ + if num_rows <= 0 or uncompressed_size <= 0: + return None + estimated_in_mem_row_size = ( + uncompressed_size * PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT / num_rows + ) + if estimated_in_mem_row_size == 0: + return None + # Never request more rows than the chunk actually contains. + return min(math.ceil(target_block_size / estimated_in_mem_row_size), num_rows) + + @DeveloperAPI class ParquetFileReader(FileReader, SupportsMetadata): """Parquet-specific file reader with adaptive batch sizing. @@ -215,17 +281,46 @@ def __init__( ) self._explicit_batch_size = batch_size self._target_block_size = target_block_size - self._parquet_format_kwargs: Dict[str, Any] = parquet_format_kwargs or {} + self._parquet_format_kwargs: Dict[str, Any] = dict(parquet_format_kwargs or {}) + # Split out the arrow-rs tuning knobs (see ARROW_RS_TUNING_KWARGS): they + # must never be spread into ``pds.ParquetFileFormat``. This base reader + # ignores them; ``ArrowRsParquetFileReader`` resolves them in + # ``_tuning``. An unrecognized ``arrow_rs_*`` key can only be a typo'd + # tuning knob — fail loudly here rather than let PyArrow raise a baffling + # ``ParquetFileFormat`` TypeError (or let the native gate treat it as an + # unsupported format kwarg and silently fall back). + self._arrow_rs_tuning: Dict[str, Any] = { + key: self._parquet_format_kwargs.pop(key) + for key in ARROW_RS_TUNING_KWARGS + if key in self._parquet_format_kwargs + } + unknown_arrow_rs_keys = sorted( + key + for key in self._parquet_format_kwargs + if key.startswith(_ARROW_RS_KWARG_PREFIX) + ) + if unknown_arrow_rs_keys: + raise ValueError( + f"Unknown arrow-rs tuning kwargs {unknown_arrow_rs_keys} in " + f"'dataset_kwargs'. Valid keys: {sorted(ARROW_RS_TUNING_KWARGS)}." + ) self._sampled_batch_size: int | object = ( _UNSET # pyrefly: ignore[bad-assignment] ) + # Size Arrow's process-wide IO/CPU pools for the read task so a Parquet + # scan can issue many concurrent column/range fetches against S3 instead + # of being capped at Arrow's small default. + if _READER_IO_THREAD_COUNT > 0: + pa.set_io_thread_count(_READER_IO_THREAD_COUNT) + if _READER_CPU_COUNT > 0: + pa.set_cpu_count(_READER_CPU_COUNT) @override def _make_format(self) -> pds.ParquetFileFormat: return pds.ParquetFileFormat(**self._parquet_format_kwargs) @override - def _resolve_batch_size(self, dataset: pds.Dataset) -> int: + def _resolve_batch_size(self, dataset: pds.Dataset, manifest: FileManifest) -> int: """Determine batch size from explicit setting, metadata, or default. Priority: explicit batch_size > sampled estimate > metadata estimate > default. @@ -234,6 +329,12 @@ def _resolve_batch_size(self, dataset: pds.Dataset) -> int: through to the metadata estimate and seed ``_sampled_batch_size`` with the result. ``_on_batch_read`` later refines it from actual data, and subsequent ``read()`` calls on the same instance use the refined value. + + The metadata estimate prefers the footer-derived stats ``ListFiles`` + already recorded on the manifest (:class:`ParquetRowGroupChunkMetadata`), + so the common footer-chunking path sizes batches without re-reading the + footer. It falls back to reading the first fragment's metadata only when + the manifest carries no such stats (e.g. the whole-file path). """ if self._explicit_batch_size is not None: return self._explicit_batch_size @@ -243,22 +344,41 @@ def _resolve_batch_size(self, dataset: pds.Dataset) -> int: batch_size = _ARROW_DEFAULT_BATCH_SIZE if self._target_block_size is not None: - first_fragment = next(dataset.get_fragments(), None) - if first_fragment is not None: - estimated = _estimate_batch_size_from_metadata( - first_fragment, self._columns, self._target_block_size + estimated = self._estimate_batch_size(dataset, manifest) + if estimated is not None: + logger.debug( + "Estimated Parquet batch size: %d rows (target_block_size=%d)", + estimated, + self._target_block_size, ) - if estimated is not None: - logger.debug( - "Estimated Parquet batch size: %d rows (target_block_size=%d)", - estimated, - self._target_block_size, - ) - batch_size = estimated + batch_size = estimated self._sampled_batch_size = batch_size return batch_size + def _estimate_batch_size( + self, dataset: pds.Dataset, manifest: FileManifest + ) -> Optional[int]: + assert self._target_block_size is not None + # Prefer footer stats already on the manifest; fall back to a footer read. + chunk = next( + (md for md in manifest.file_chunk_metadatas if md is not None), None + ) + if chunk is not None and "uncompressed_size" in chunk: + estimated = _estimate_batch_size_from_chunk_stats( + chunk["uncompressed_size"], + chunk["num_rows"], + self._target_block_size, + ) + if estimated is not None: + return estimated + first_fragment = next(dataset.get_fragments(), None) + if first_fragment is None: + return None + return _estimate_batch_size_from_metadata( + first_fragment, self._columns, self._target_block_size + ) + @override def _on_batch_read(self, table: pa.Table) -> None: """Refine batch size estimate from actual in-memory data.""" @@ -273,41 +393,44 @@ def _get_fragments_to_read( dataset: pds.Dataset, manifest: FileManifest, ) -> List[Tuple[pds.Fragment, int]]: - """Fan file fragments into chunk-level sub-fragments per manifest row. + """Fan file fragments into read-level sub-fragments per manifest row. For each manifest row, looks up the file's fragment by path and: - If ``chunk_metadata`` is ``None`` (whole-file case), the file - fragment is yielded as-is with a row offset of 0. This matches - ``ParquetFileChunker``'s behavior for files at or below - ``target_chunk_size`` and the default ``WholeFileChunker`` for - non-chunking callers. - - Otherwise the row carries a :class:`ParquetFileChunkMetadata`; + fragment is yielded as-is with a row offset of 0 (the default + ``WholeFileChunker`` for non-chunking callers). + - Otherwise the row carries a :class:`ParquetRowGroupChunkMetadata` + naming the exact physical row groups the bin assigned to this file + (predicate pruning + bin packing already happened in ``ListFiles``); we slice the fragment via - :func:`~ray.data._internal.datasource_v2.chunkers.parquet_file_chunking_utils._fragments_from_chunk_metadata` - which returns one sub-fragment per row group in the chunk's - row-group range, paired with the cumulative pre-filter row - offset of that row group within the file. The downstream - ``_compute_row_hashes`` call uses this offset so row hashes - remain unique across sub-fragments that share ``fragment.path``. + :func:`~ray.data._internal.datasource_v2.chunkers.parquet_file_chunking_utils._fragments_from_row_group_ids`. + When ``include_row_hash`` is on it fans out one sub-fragment per row + group, each paired with its cumulative pre-filter row offset, so the + downstream ``_compute_row_hashes`` call keeps hashes unique across + sub-fragments that share ``fragment.path``. Paths are deduped by :meth:`FileReader.read` before the dataset is built, so the dataset has exactly one fragment per file. The per-row chunk metadata drives the fan-out here, not the dataset itself — multiple manifest rows can share a single path with - different chunk indices. + different row-group sets. """ path_to_fragment = { fragment.path: fragment for fragment in dataset.get_fragments() } fragments: List[Tuple[pds.Fragment, int]] = [] for path, chunk_metadata in zip(manifest.paths, manifest.file_chunk_metadatas): - fragment = path_to_fragment[path] + fragment: pds.ParquetFileFragment = path_to_fragment[path] if chunk_metadata is None: fragments.append((fragment, 0)) else: fragments.extend( - _fragments_from_chunk_metadata(fragment, chunk_metadata) + _fragments_from_row_group_ids( + fragment, + chunk_metadata["row_group_ids"], + per_row_group_offsets=self._include_row_hash, + ) ) return fragments @@ -513,11 +636,29 @@ def _arrow_scanner_kwargs(self) -> dict: # meaningful bytes per round-trip. Tunable via env var for # workloads that need a different point on the latency/memory- # peak curve. + scan_opts: Dict[str, Any] = { + "use_buffered_stream": True, + "buffer_size": _PARQUET_FRAGMENT_BUFFER_SIZE, + } + # Experimental: honor RAY_DATA_PARQUET_PRE_BUFFER when set (see the + # module-level knob). Unset leaves pyarrow's default in place, so the + # explicit ``fragment_scan_options`` below stays byte-identical to the + # pre-knob behavior. + if _PARQUET_PRE_BUFFER_OVERRIDE is not None: + scan_opts["pre_buffer"] = bool(_PARQUET_PRE_BUFFER_OVERRIDE) + # ``page_checksum_verification`` is a *scan* option, not a format option. + # ``_make_format`` sets it on the format's ``default_fragment_scan_options``, + # but the explicit ``fragment_scan_options`` we hand the scanner below + # overrides that default — so an explicit request would be silently + # dropped (pyarrow's default is ``False``, i.e. no verification). Thread + # it through here so ``page_checksum_verification=True`` is actually + # honored on the PyArrow read path (matching the arrow-rs crate, which + # always verifies). + pcv = self._parquet_format_kwargs.get("page_checksum_verification") + if pcv is not None: + scan_opts["page_checksum_verification"] = pcv kwargs: dict = { - "fragment_scan_options": pds.ParquetFragmentScanOptions( - use_buffered_stream=True, - buffer_size=_PARQUET_FRAGMENT_BUFFER_SIZE, - ), + "fragment_scan_options": pds.ParquetFragmentScanOptions(**scan_opts), "fragment_readahead": 1, } return kwargs diff --git a/python/ray/data/_internal/datasource_v2/scanners/arrow_file_scanner.py b/python/ray/data/_internal/datasource_v2/scanners/arrow_file_scanner.py index 2641e0439a3f..907f3b089162 100644 --- a/python/ray/data/_internal/datasource_v2/scanners/arrow_file_scanner.py +++ b/python/ray/data/_internal/datasource_v2/scanners/arrow_file_scanner.py @@ -110,6 +110,10 @@ def push_filters( return replace(self, predicate=combined), None + @override + def pushed_predicate(self) -> Optional["Expr"]: + return self.predicate + @override def prune_columns(self, columns: List[str]) -> "ArrowFileScanner": """Prune to only the specified columns. @@ -144,6 +148,10 @@ def push_limit(self, limit: int) -> "ArrowFileScanner": new_limit = min(current, limit) if current is not None else limit return replace(self, limit=new_limit) + @override + def pushed_limit(self) -> Optional[int]: + return self.limit + @override def prune_partitions(self, predicate: "Expr") -> "ArrowFileScanner": """Store a partition predicate for file-level pruning during plan(). diff --git a/python/ray/data/_internal/datasource_v2/scanners/parquet_scanner.py b/python/ray/data/_internal/datasource_v2/scanners/parquet_scanner.py index 5ff7744e9210..f12c7a252d4b 100644 --- a/python/ray/data/_internal/datasource_v2/scanners/parquet_scanner.py +++ b/python/ray/data/_internal/datasource_v2/scanners/parquet_scanner.py @@ -1,3 +1,4 @@ +import logging from dataclasses import dataclass, field from typing import Any, Dict, Optional @@ -18,6 +19,8 @@ ) from ray.util.annotations import DeveloperAPI +logger = logging.getLogger(__name__) + @DeveloperAPI @dataclass(frozen=True) @@ -67,12 +70,37 @@ def read_schema(self) -> pa.Schema: return schema def create_reader(self) -> ParquetFileReader: - """Create a ParquetFileReader configured for this scanner. + """Create the Parquet reader configured for this scanner. + + Returns the experimental arrow-rs reader when + ``DataContext.use_arrow_rs_parquet_reader`` is set, otherwise the + PyArrow ``ParquetFileReader``. The two share an identical constructor, + so the reader class is the only thing that changes. Returns: - ParquetFileReader with all pushdowns and adaptive batch sizing. + A reader with all pushdowns and adaptive batch sizing. """ - return ParquetFileReader( + from ray.data.context import DataContext + + reader_cls = ParquetFileReader + if DataContext.get_current().use_arrow_rs_parquet_reader: + # Import lazily so the default PyArrow path never imports the + # optional native extension. A missing extension raises inside + # the reader's ``_iter_fragment_tables`` with a build hint rather + # than silently falling back — silent fallback would corrupt + # benchmark attribution. + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( # noqa: E501 + ArrowRsParquetFileReader, + ) + + reader_cls = ArrowRsParquetFileReader + logger.warning( + "Ray Data ARROW-RS: using the native arrow-rs Parquet reader " + "(ray_data_arrow_rs) for this read_parquet. Per-file native vs " + "PyArrow-fallback decisions are logged from the read tasks." + ) + + return reader_cls( batch_size=self.batch_size, columns=list(self.columns) if self.columns is not None else None, predicate=self.predicate, diff --git a/python/ray/data/_internal/datasource_v2/tests/test_footer_chunking.py b/python/ray/data/_internal/datasource_v2/tests/test_footer_chunking.py new file mode 100644 index 000000000000..e8dcfeb5b067 --- /dev/null +++ b/python/ray/data/_internal/datasource_v2/tests/test_footer_chunking.py @@ -0,0 +1,313 @@ +"""Tests for the footer-based Parquet chunking path. + +Covers the pure-logic pieces (row-group coalescing and the online bin packer, +including the split-coalesced no-op invariant) without a Ray cluster, plus a few +end-to-end reads that exercise ``FooterFileIndexer`` through +``ray.data.read_parquet`` with predicate / limit / projection push-down. +""" + +import pytest + +from ray.data._internal.datasource_v2.chunkers.parquet_footer_types import ( + FileChunks, + RowGroupInfo, +) +from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest +from ray.data._internal.datasource_v2.listing.footer_reader import coalesce_row_groups +from ray.data._internal.datasource_v2.partitioners.online_bin_packer import ( + OnlineBinPacker, +) + + +def _rg(idx, size, rows=10, fully_matched=True): + return RowGroupInfo( + rg_idx=idx, uncompressed_size=size, num_rows=rows, fully_matched=fully_matched + ) + + +# --------------------------------------------------------------------------- +# coalesce_row_groups (pure) +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "per_rg, target, expected", + [ + pytest.param( + [_rg(0, 10), _rg(1, 20), _rg(2, 30)], + 0, + [(0, 1, 10, 10), (1, 1, 20, 10), (2, 1, 30, 10)], + id="disabled-is-identity", + ), + pytest.param( + [_rg(0, 10, 1), _rg(1, 10, 2), _rg(2, 10, 3)], + 25, + [(0, 3, 30, 6)], + id="merge-contiguous-until-target", + ), + pytest.param( + [_rg(0, 10, fully_matched=True), _rg(1, 10, fully_matched=False)], + 1000, + [(0, 1, 10, 10), (1, 1, 10, 10)], + id="break-on-fully-matched-change", + ), + pytest.param( + [_rg(0, 10), _rg(2, 10)], + 1000, + [(0, 1, 10, 10), (2, 1, 10, 10)], + id="break-on-index-gap", + ), + ], +) +def test_coalesce(per_rg, target, expected): + out = coalesce_row_groups(per_rg, target) + + assert [ + (c.rg_idx, c.rg_count, c.uncompressed_size, c.num_rows) for c in out + ] == expected + # Every physical row group is covered exactly once. + covered = [i for c in out for i in range(c.rg_idx, c.rg_idx + c.rg_count)] + assert sorted(covered) == sorted(r.rg_idx for r in per_rg) + # Per-group breakdown is attached iff the chunk is coalesced (rg_count > 1). + for c in out: + if c.rg_count > 1: + assert len(c.rg_sizes) == c.rg_count and len(c.rg_rows) == c.rg_count + else: + assert c.rg_sizes == () and c.rg_rows == () + + +# --------------------------------------------------------------------------- +# OnlineBinPacker (pure) +# --------------------------------------------------------------------------- + + +def _manifest_map(manifest: FileManifest): + """A sealed bin's manifest as ``{path: sorted physical row-group ids}``.""" + return { + str(path): sorted(meta["row_group_ids"]) + for path, meta in zip(manifest.paths, manifest.file_chunk_metadatas) + } + + +def _pack(file_chunks_list, max_bin_bytes, **kwargs): + packer = OnlineBinPacker(max_bin_bytes, **kwargs) + bins = [] + for file_chunks in file_chunks_list: + packer.add_file_chunks(file_chunks) + while packer.has_partition(): + bins.append(_manifest_map(packer.next_partition())) + packer.finalize() + while packer.has_partition(): + bins.append(_manifest_map(packer.next_partition())) + return bins + + +def _pairs(bins): + """All ``(path, rg_id)`` pairs across bins, sorted.""" + return sorted((p, i) for b in bins for p, ids in b.items() for i in ids) + + +@pytest.mark.parametrize( + "files, max_bin, expected_bins", + [ + pytest.param( + [FileChunks("a", 10, (_rg(0, 10),)), FileChunks("b", 10, (_rg(0, 10),))], + 1000, + [{"a": [0], "b": [0]}], + id="light-colours-share-a-bin", + ), + pytest.param( + [FileChunks("a", 500, (_rg(0, 500),))], + 100, + [{"a": [0]}], + id="oversize-group-gets-own-bin", + ), + ], +) +def test_packer_placement(files, max_bin, expected_bins): + assert _pack(files, max_bin) == expected_bins + + +def test_packer_heavy_colour_spans_multiple_bins(): + # A file far heavier than one bin spills into several bins (exact split point + # depends on the light->heavy threshold, so assert coverage, not layout). + files = [FileChunks("a", 400, tuple(_rg(i, 100) for i in range(4)))] + bins = _pack(files, max_bin_bytes=100) + assert len(bins) == 4 + assert _pairs(bins) == [("a", i) for i in range(4)] + + +@pytest.mark.parametrize("split_coalesced", [False, True]) +def test_packer_covers_every_row_group_exactly_once(split_coalesced): + files = [ + FileChunks("a", 120, tuple(_rg(i, 30) for i in range(4))), + FileChunks("b", 90, tuple(_rg(i, 30) for i in range(3))), + ] + pairs = _pairs(_pack(files, max_bin_bytes=100, split_coalesced=split_coalesced)) + expected = sorted([("a", i) for i in range(4)] + [("b", i) for i in range(3)]) + assert pairs == expected + assert len(pairs) == len(set(pairs)) # no duplicates + + +def test_split_coalesced_is_noop_without_coalescing(): + # With every rg_count == 1, the split flag must not change the packing. + files = [ + FileChunks("a", 120, tuple(_rg(i, 40) for i in range(3))), + FileChunks("b", 80, tuple(_rg(i, 40) for i in range(2))), + ] + assert _pack(files, 100, split_coalesced=False) == _pack( + files, 100, split_coalesced=True + ) + + +def test_split_coalesced_splits_oversize_run_at_boundaries(): + # A coalesced chunk (rg 0..2) that can't fit whole in a 50-byte bin is cut at + # physical row-group boundaries; every group still appears exactly once. + coalesced = RowGroupInfo( + rg_idx=0, + uncompressed_size=90, + num_rows=30, + rg_count=3, + rg_sizes=(30, 30, 30), + rg_rows=(10, 10, 10), + ) + bins = _pack([FileChunks("a", 90, (coalesced,))], 50, split_coalesced=True) + assert _pairs(bins) == [("a", 0), ("a", 1), ("a", 2)] + assert len(bins) >= 2 # 90 bytes across 50-byte bins + + +# --------------------------------------------------------------------------- +# End-to-end through ray.data.read_parquet (footer path) +# --------------------------------------------------------------------------- + +_N_PER_FILE = 400 +_N_FILES = 3 + + +@pytest.fixture +def footer_parquet(tmp_path, monkeypatch): + """Write a small multi-row-group Parquet dataset; enable the footer path.""" + import pyarrow as pa + import pyarrow.parquet as pq + + from ray.data.context import DataContext + + # Footer chunking is the default for Parquet V2; just keep the pool small. + monkeypatch.setenv("RAY_DATA_PARQUET_FOOTER_NUM_ACTORS", "2") + monkeypatch.setenv("RAY_DATA_PARQUET_FOOTER_BATCH_SIZE", "2") + + ctx = DataContext.get_current() + prev_v2 = ctx.use_datasource_v2 + ctx.use_datasource_v2 = True + + for f in range(_N_FILES): + start = f * _N_PER_FILE + table = pa.table( + { + "id": list(range(start, start + _N_PER_FILE)), + "val": [f"v{i}" for i in range(_N_PER_FILE)], + } + ) + pq.write_table(table, str(tmp_path / f"part_{f}.parquet"), row_group_size=100) + + try: + yield str(tmp_path) + finally: + ctx.use_datasource_v2 = prev_v2 + + +def test_e2e_footer_read_matches_expected(footer_parquet): + import ray + + total = _N_PER_FILE * _N_FILES + ds = ray.data.read_parquet(footer_parquet) + assert ds.count() == total + assert sorted(r["id"] for r in ds.take_all()) == list(range(total)) + + +@pytest.mark.parametrize( + "op, expected", + [ + pytest.param(lambda ds: ds.filter(expr="id < 50").count(), 50, id="filter"), + pytest.param(lambda ds: ds.limit(10).count(), 10, id="limit"), + pytest.param( + lambda ds: ds.select_columns(["id"]).schema().names, + ["id"], + id="projection", + ), + ], +) +def test_e2e_footer_pushdowns(footer_parquet, op, expected): + import ray + + assert op(ray.data.read_parquet(footer_parquet)) == expected + + +# Filter and limit push down together: the limit stops listing early once the +# ``num_rows`` of *fully matched* row groups reaches it, so that classification +# has to be exact. Nulls are the interesting case -- Parquet min/max statistics +# are computed over non-null values only, so a group whose non-null values all +# satisfy the filter looks fully matched by bounds alone while its null rows do +# not survive. Deliberately lopsided at 10 survivors per 100 rows: the stop is +# evaluated per file, so a fixture with a small shortfall can pass by luck when +# the last file's overshoot covers the deficit. +_NULL_FILES = 20 +_NULL_ROWS_PER_FILE = 100 +_NULL_TOTAL_SURVIVORS = _NULL_FILES * 10 + + +@pytest.fixture +def nullable_parquet(tmp_path, monkeypatch): + """Multi-file, multi-row-group data whose filtered column holds nulls.""" + import pyarrow as pa + import pyarrow.parquet as pq + + from ray.data.context import DataContext + + monkeypatch.setenv("RAY_DATA_PARQUET_FOOTER_NUM_ACTORS", "2") + monkeypatch.setenv("RAY_DATA_PARQUET_FOOTER_BATCH_SIZE", "2") + + ctx = DataContext.get_current() + prev_v2 = ctx.use_datasource_v2 + ctx.use_datasource_v2 = True + + for f in range(_NULL_FILES): + ids = [ + 3 + f * 1000 + i if i % 10 == 0 else None + for i in range(_NULL_ROWS_PER_FILE) + ] + pq.write_table( + pa.table({"id": pa.array(ids, pa.int64())}), + str(tmp_path / f"part_{f}.parquet"), + row_group_size=25, + ) + + try: + yield str(tmp_path) + finally: + ctx.use_datasource_v2 = prev_v2 + + +@pytest.mark.parametrize( + "limit", [1, 10, 100, _NULL_TOTAL_SURVIVORS, 10 * _NULL_TOTAL_SURVIVORS] +) +def test_e2e_filter_then_limit_with_nulls(nullable_parquet, limit): + """``filter(...).limit(n)`` delivers ``n`` rows whenever ``n`` survivors exist. + + If nulls were ever counted as survivors, listing would stop short and + ``Limit`` would return fewer rows than asked for, with no error. + """ + import ray + from ray.data.expressions import col + + ds = ray.data.read_parquet(nullable_parquet) + rows = ds.filter(expr=col("id") > 2).limit(limit).take_all() + + assert len(rows) == min(limit, _NULL_TOTAL_SURVIVORS) + assert all(r["id"] is not None and r["id"] > 2 for r in rows) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main(["-v", __file__])) diff --git a/python/ray/data/_internal/datasource_v2/tests/test_parquet_datasource_v2.py b/python/ray/data/_internal/datasource_v2/tests/test_parquet_datasource_v2.py index 447a405a8573..ccc8fe0edbbf 100644 --- a/python/ray/data/_internal/datasource_v2/tests/test_parquet_datasource_v2.py +++ b/python/ray/data/_internal/datasource_v2/tests/test_parquet_datasource_v2.py @@ -9,19 +9,15 @@ import pyarrow as pa import pyarrow.parquet as pq -import pytest from ray.data._internal.datasource_v2.chunkers.file_chunker import ( - ParquetFileChunker, - ParquetFileChunkMetadata, - WholeFileChunker, + ParquetRowGroupChunkMetadata, create_chunk_metadata, ) -from ray.data._internal.datasource_v2.chunkers.parquet_file_chunking_utils import ( - _calculate_row_group_range, - _fragments_from_chunk_metadata, -) from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest +from ray.data._internal.datasource_v2.listing.footer_file_indexer import ( + FooterFileIndexer, +) from ray.data._internal.datasource_v2.parquet_datasource_v2 import ( ParquetDatasourceV2, ) @@ -209,72 +205,13 @@ def test_nested_fallback_handles_schema_evolution(tmp_path, monkeypatch): assert rows_by_fragment == {"with_b.parquet": 2, "without_b.parquet": 0} -def test_datasource_defaults_to_parquet_file_chunker(tmp_path): - """``ParquetDatasourceV2`` plugs ``ParquetFileChunker`` into its indexer.""" +def test_datasource_uses_footer_indexer(tmp_path): + """``ParquetDatasourceV2`` uses the footer-based indexer for row-group reads.""" file_path = tmp_path / "data.parquet" _write_parquet(str(file_path), pa.table({"a": [1, 2, 3]})) datasource = ParquetDatasourceV2([str(file_path)]) - indexer = datasource._get_file_indexer() - assert isinstance(indexer.file_chunker, ParquetFileChunker) - - -def test_datasource_accepts_custom_chunker(tmp_path): - """An explicit ``file_chunker`` override propagates to the indexer.""" - file_path = tmp_path / "data.parquet" - _write_parquet(str(file_path), pa.table({"a": [1, 2, 3]})) - - custom = WholeFileChunker() - datasource = ParquetDatasourceV2([str(file_path)], file_chunker=custom) - indexer = datasource._get_file_indexer() - assert indexer.file_chunker is custom - - -@pytest.mark.parametrize( - "total_row_groups,total_num_chunks,expected_ranges", - [ - # Even distribution. - (10, 2, [(0, 5), (5, 10)]), - (12, 3, [(0, 4), (4, 8), (8, 12)]), - (20, 4, [(0, 5), (5, 10), (10, 15), (15, 20)]), - # Uneven distribution: earlier chunks get extra row groups. - (10, 3, [(0, 4), (4, 7), (7, 10)]), - (11, 3, [(0, 4), (4, 8), (8, 11)]), - (13, 4, [(0, 4), (4, 7), (7, 10), (10, 13)]), - # Edge cases — over-estimated chunk counts must produce ``None``. - (1, 1, [(0, 1)]), - (1, 2, [(0, 1), None]), - (0, 1, [None]), - (5, 10, [(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)] + [None] * 5), - ], -) -def test_calculate_row_group_range_distribution( - total_row_groups, total_num_chunks, expected_ranges -): - """Row-group distribution across chunks is even and covers everything.""" - for chunk_idx in range(total_num_chunks): - result = _calculate_row_group_range( - chunk_idx, total_num_chunks, total_row_groups - ) - expected = ( - expected_ranges[chunk_idx] if chunk_idx < len(expected_ranges) else None - ) - assert ( - result == expected - ), f"Chunk {chunk_idx}: expected {expected}, got {result}" - - # No gaps, no overlaps, every row group covered exactly once. - covered = set() - for chunk_idx in range(total_num_chunks): - result = _calculate_row_group_range( - chunk_idx, total_num_chunks, total_row_groups - ) - if result is not None: - start, end = result - chunk_rows = set(range(start, end)) - assert not (covered & chunk_rows) - covered.update(chunk_rows) - assert covered == set(range(total_row_groups)) + assert isinstance(datasource._get_file_indexer(), FooterFileIndexer) def _write_multi_row_group_parquet(path, num_rows: int, row_group_size: int): @@ -283,139 +220,53 @@ def _write_multi_row_group_parquet(path, num_rows: int, row_group_size: int): return table -def test_fragments_from_chunk_metadata_subsets_by_row_group(tmp_path): - """``_fragments_from_chunk_metadata`` slices a file fragment per chunk.""" - import pyarrow.dataset as pds - - file_path = str(tmp_path / "multi.parquet") - # 1000 rows, 100 row groups (row_group_size=10). - _write_multi_row_group_parquet(file_path, num_rows=1000, row_group_size=10) - - dataset = pds.dataset(file_path, format="parquet") - (fragment,) = dataset.get_fragments() - assert fragment.metadata.num_row_groups == 100 - - # 100 row groups split into 4 chunks -> 25 row groups each. - chunk_md = create_chunk_metadata( - ParquetFileChunkMetadata, chunk_idx=1, total_num_chunks=4 +def _row_group_manifest(path, row_group_ids, num_rows): + """A one-row footer-path manifest selecting explicit row groups of a file.""" + return FileManifest.construct_manifest( + [path], + [0], + [ + create_chunk_metadata( + ParquetRowGroupChunkMetadata, + row_group_ids=tuple(row_group_ids), + num_rows=num_rows, + # Nominal projected uncompressed size (8-byte int64 ids); only + # used for footer-free batch sizing, not row selection. + uncompressed_size=num_rows * 8, + ) + ], ) - sub_fragments = _fragments_from_chunk_metadata(fragment, chunk_md) - assert len(sub_fragments) == 25 - # chunk_idx=1, 25 rows/chunk * 10 rows/row_group = starting offset 250. - expected_offset = 250 - for sub, offset in sub_fragments: - assert len(sub.row_groups) == 1 - assert offset == expected_offset - expected_offset += sub.metadata.row_group(sub.row_groups[0].id).num_rows -def test_fragments_from_chunk_metadata_returns_empty_for_out_of_range_chunk( - tmp_path, -): - """Over-estimated chunk indices fall off the end → no sub-fragments.""" - import pyarrow.dataset as pds - - file_path = str(tmp_path / "single.parquet") - # 5 rows, single row group. - _write_multi_row_group_parquet(file_path, num_rows=5, row_group_size=5) - - dataset = pds.dataset(file_path, format="parquet") - (fragment,) = dataset.get_fragments() - - # chunk_idx=4 with 5 chunks but only 1 row group → no sub-fragments. - chunk_md = create_chunk_metadata( - ParquetFileChunkMetadata, chunk_idx=4, total_num_chunks=5 - ) - assert _fragments_from_chunk_metadata(fragment, chunk_md) == [] - - -def _read_via_reader(reader, manifest): - return list(reader.read(manifest)) - - -def test_parquet_file_reader_reads_chunked_manifest(tmp_path): - """End-to-end: a manifest with per-chunk rows is read into the same rows - as a single whole-file manifest.""" +def test_parquet_file_reader_reads_selected_row_groups(tmp_path): + """The reader reads exactly the row groups named by the footer metadata.""" file_path = str(tmp_path / "data.parquet") - expected_rows = 200 - _write_multi_row_group_parquet(file_path, num_rows=expected_rows, row_group_size=20) - file_size = os.path.getsize(file_path) - - reader_whole = ParquetFileReader() - whole_manifest = FileManifest.construct_manifest([file_path], [file_size], [None]) - whole_tables = _read_via_reader(reader_whole, whole_manifest) - whole_rows = pa.concat_tables(whole_tables).column("id").to_pylist() - - chunker = ParquetFileChunker(target_chunk_size=1024) - chunks = list(chunker.generate_chunk_metadatas(file_path, file_size)) - assert len(chunks) > 1, "test setup expects ParquetFileChunker to chunk" - - paths = [file_path] * len(chunks) - chunk_metadatas = [md for md, _ in chunks] - chunk_sizes = [sz for _, sz in chunks] - chunked_manifest = FileManifest.construct_manifest( - paths, chunk_sizes, chunk_metadatas - ) - - reader_chunked = ParquetFileReader() - chunked_tables = _read_via_reader(reader_chunked, chunked_manifest) - chunked_rows = pa.concat_tables(chunked_tables).column("id").to_pylist() + # 200 rows, 10 row groups of 20 rows each. + _write_multi_row_group_parquet(file_path, num_rows=200, row_group_size=20) - assert sorted(chunked_rows) == sorted(whole_rows) == list(range(expected_rows)) + # Select row groups 0, 2, 5 -> rows [0,20) + [40,60) + [100,120). + manifest = _row_group_manifest(file_path, row_group_ids=[0, 2, 5], num_rows=60) + tables = list(ParquetFileReader().read(manifest)) + rows = sorted(pa.concat_tables(tables).column("id").to_pylist()) + expected = list(range(0, 20)) + list(range(40, 60)) + list(range(100, 120)) + assert rows == sorted(expected) -def test_parquet_file_reader_chunked_row_hashes_are_unique(tmp_path): - """Row hashes must remain unique across chunked sub-fragments of the - same file. +def test_parquet_file_reader_row_group_row_hashes_are_unique(tmp_path): + """Row hashes stay unique across per-row-group sub-fragments of one file. - Regression: ``_read_fragments_sequential`` previously reseeded - ``offset=0`` for every fragment. Since chunked sub-fragments share - ``fragment.path``, ``_compute_row_hashes(path, 0, n)`` collided across - row groups of the same file. + With ``include_row_hash`` the footer path fans one sub-fragment per row + group, each seeded with its cumulative file row offset, so hashes can't + collide across row groups that share ``fragment.path``. """ file_path = str(tmp_path / "data.parquet") expected_rows = 200 _write_multi_row_group_parquet(file_path, num_rows=expected_rows, row_group_size=20) - file_size = os.path.getsize(file_path) - chunker = ParquetFileChunker(target_chunk_size=1024) - chunks = list(chunker.generate_chunk_metadatas(file_path, file_size)) - assert len(chunks) > 1, "test setup expects ParquetFileChunker to chunk" - - paths = [file_path] * len(chunks) - chunk_metadatas = [md for md, _ in chunks] - chunk_sizes = [sz for _, sz in chunks] - chunked_manifest = FileManifest.construct_manifest( - paths, chunk_sizes, chunk_metadatas + manifest = _row_group_manifest( + file_path, row_group_ids=range(10), num_rows=expected_rows ) - reader = ParquetFileReader(include_row_hash=True) - chunked_tables = list(reader.read(chunked_manifest)) - hashes = pa.concat_tables(chunked_tables).column("row_hash").to_pylist() + hashes = pa.concat_tables(reader.read(manifest)).column("row_hash").to_pylist() assert len(hashes) == expected_rows - assert ( - len(set(hashes)) == expected_rows - ), "row_hash must be unique across chunked sub-fragments of one file" - - -def test_parquet_file_reader_handles_out_of_range_chunks(tmp_path): - """Out-of-range chunk metadata is silently dropped — no exception, no rows.""" - file_path = str(tmp_path / "tiny.parquet") - _write_multi_row_group_parquet(file_path, num_rows=5, row_group_size=5) - - file_size = os.path.getsize(file_path) - # ``ParquetFileChunker`` over-estimates here; keep only the over-estimated - # tail (chunk_idx >= 1) to assert the reader yields no tables. - chunker = ParquetFileChunker(target_chunk_size=128) - chunks = list(chunker.generate_chunk_metadatas(file_path, file_size)) - assert len(chunks) > 1 - - paths = [file_path] * (len(chunks) - 1) - out_of_range_metadatas = [md for md, _ in chunks[1:]] - sizes = [sz for _, sz in chunks[1:]] - manifest = FileManifest.construct_manifest(paths, sizes, out_of_range_metadatas) - - reader = ParquetFileReader() - tables = list(reader.read(manifest)) - # All sub-fragments are out-of-range -> 0 tables emitted. - assert sum(t.num_rows for t in tables) == 0 + assert len(set(hashes)) == expected_rows diff --git a/python/ray/data/_internal/execution/interfaces/op_runtime_metrics.py b/python/ray/data/_internal/execution/interfaces/op_runtime_metrics.py index 54609f4be62d..27cd03151b74 100644 --- a/python/ray/data/_internal/execution/interfaces/op_runtime_metrics.py +++ b/python/ray/data/_internal/execution/interfaces/op_runtime_metrics.py @@ -17,7 +17,7 @@ ) from ray.data._internal.execution.interfaces.ref_bundle import RefBundle from ray.data._internal.memory_tracing import trace_allocation -from ray.data.block import BlockMetadata, TaskExecWorkerStats +from ray.data.block import BlockMetadata, ReadFilesTaskStats, TaskExecWorkerStats if TYPE_CHECKING: from ray.data._internal.execution.interfaces.physical_operator import ( @@ -576,6 +576,16 @@ def __init__(self, op: "PhysicalOperator"): self.block_size_rows = RuntimeMetricsHistogram(histogram_bucket_rows) self._op_task_duration_stats = DistributionTracker() self._max_uss_bytes = DistributionTracker() + self._max_rss_bytes = DistributionTracker() + # Per-task reader-level aggregates, reported by ReadFiles tasks via + # ``TaskExecWorkerStats.custom_op_stats`` (see ``ReadFilesTaskStats``). + # Empty for non-read operators. + self._read_task_decoded_bytes = DistributionTracker() + self._read_task_decode_wall_s = DistributionTracker() + self._read_task_peak_batch_bytes = DistributionTracker() + self._read_task_trim_wall_s = DistributionTracker() + self._read_task_yield_wall_s = DistributionTracker() + self._read_task_first_table_wall_s = DistributionTracker() @property def extra_metrics(self) -> Dict[str, Any]: @@ -904,6 +914,81 @@ def op_task_duration_stats(self) -> DistributionTracker: def max_uss_bytes(self) -> DistributionTracker: return self._max_uss_bytes + @metric_property( + description="Distribution across read tasks of bytes the reader decoded.", + metrics_group=MetricsGroup.TASKS, + metrics_type=MetricsType.Unsupported, + ) + def read_task_decoded_bytes(self) -> DistributionTracker: + return self._read_task_decoded_bytes + + @metric_property( + description=( + "Distribution across read tasks of wall seconds spent inside the " + "reader's decode iterator." + ), + metrics_group=MetricsGroup.TASKS, + metrics_type=MetricsType.Unsupported, + ) + def read_task_decode_wall_s(self) -> DistributionTracker: + return self._read_task_decode_wall_s + + @metric_property( + description=( + "Distribution across read tasks of the largest single table the " + "reader yielded (decode working-set proxy)." + ), + metrics_group=MetricsGroup.TASKS, + metrics_type=MetricsType.Unsupported, + ) + def read_task_peak_batch_bytes(self) -> DistributionTracker: + return self._read_task_peak_batch_bytes + + @metric_property( + description=( + "Distribution across read tasks of wall seconds spent in the " + "reader's end-of-stream finalizer (arrow-rs malloc_trim); part of " + "read_task_decode_wall_s." + ), + metrics_group=MetricsGroup.TASKS, + metrics_type=MetricsType.Unsupported, + ) + def read_task_trim_wall_s(self) -> DistributionTracker: + return self._read_task_trim_wall_s + + @metric_property( + description=( + "Distribution across read tasks of wall seconds spent inside the " + "read task's yield: output-buffer shaping, block build, object-store " + "put and streaming-generator backpressure. Disjoint from " + "read_task_decode_wall_s." + ), + metrics_group=MetricsGroup.TASKS, + metrics_type=MetricsType.Unsupported, + ) + def read_task_yield_wall_s(self) -> DistributionTracker: + return self._read_task_yield_wall_s + + @metric_property( + description=( + "Distribution across read tasks of wall seconds from task start to " + "the first decoded table (reader construction + first next())." + ), + metrics_group=MetricsGroup.TASKS, + metrics_type=MetricsType.Unsupported, + ) + def read_task_first_table_wall_s(self) -> DistributionTracker: + return self._read_task_first_table_wall_s + + @metric_property( + description="Average bytes decoded by the reader per read task.", + metrics_group=MetricsGroup.TASKS, + ) + def average_decoded_bytes_per_read_task(self) -> Optional[float]: + if self.read_task_decoded_bytes.num_samples == 0: + return None + return self.read_task_decoded_bytes.mean + @metric_property( description="Average USS usage of tasks.", metrics_group=MetricsGroup.TASKS, @@ -914,6 +999,44 @@ def average_max_uss_per_task(self) -> Optional[float]: return None return self.max_uss_bytes.mean + @metric_property( + description="Max USS usage across tasks (worst task).", + metrics_group=MetricsGroup.TASKS, + ) + def max_uss_per_task(self) -> Optional[float]: + """Peak USS of the single worst task — the number a per-worker memory + budget must survive (the average can hide one task decoding an outlier + file).""" + return self.max_uss_bytes.max + + @metric_property( + description="Distribution of max RSS bytes across tasks.", + metrics_group=MetricsGroup.TASKS, + metrics_type=MetricsType.Unsupported, + ) + def max_rss_bytes(self) -> DistributionTracker: + return self._max_rss_bytes + + @metric_property( + description="Average RSS usage of tasks.", + metrics_group=MetricsGroup.TASKS, + ) + def average_max_rss_per_task(self) -> Optional[float]: + """Average max RSS usage of tasks. RSS counts the shared pages + (e.g. mapped object-store blocks) that USS excludes, so USS vs RSS + separates a task's private working set from its OS-visible footprint.""" + if self.max_rss_bytes.num_samples == 0: + return None + return self.max_rss_bytes.mean + + @metric_property( + description="Max RSS usage across tasks (worst task).", + metrics_group=MetricsGroup.TASKS, + ) + def max_rss_per_task(self) -> Optional[float]: + """Peak RSS of the single worst task.""" + return self.max_rss_bytes.max + @metric_property( description="Indicates if the operator is hanging.", metrics_group=MetricsGroup.MISC, @@ -1165,6 +1288,41 @@ def on_task_finished( if task_exec_stats is not None and task_exec_stats.max_uss_bytes is not None: self._max_uss_bytes.add_sample(task_exec_stats.max_uss_bytes) + if task_exec_stats is not None and task_exec_stats.max_rss_bytes is not None: + self._max_rss_bytes.add_sample(task_exec_stats.max_rss_bytes) + + # Fold reader-level per-task aggregates (reported by ReadFiles tasks) + # into their distributions. Fused Read->X tasks report exactly one + # entry (the read transform); the sum handles hypothetical multiples. + if task_exec_stats is not None and task_exec_stats.custom_op_stats: + read_stats = [ + s + for s in task_exec_stats.custom_op_stats + # manifests == 0 means the task read nothing (e.g. every + # manifest pruned away) — sampling its zeros would skew the + # per-task decode distributions. + if isinstance(s, ReadFilesTaskStats) and s.manifests > 0 + ] + if read_stats: + self._read_task_decoded_bytes.add_sample( + sum(s.decoded_bytes for s in read_stats) + ) + self._read_task_decode_wall_s.add_sample( + sum(s.decode_wall_s for s in read_stats) + ) + self._read_task_peak_batch_bytes.add_sample( + max(s.peak_batch_bytes for s in read_stats) + ) + self._read_task_yield_wall_s.add_sample( + sum(s.yield_wall_s for s in read_stats) + ) + self._read_task_first_table_wall_s.add_sample( + sum(s.first_table_wall_s for s in read_stats) + ) + self._read_task_trim_wall_s.add_sample( + sum(s.trim_wall_s for s in read_stats) + ) + task_output_backpressure_s = ( task_exec_driver_stats.task_output_backpressure_s if task_exec_driver_stats diff --git a/python/ray/data/_internal/execution/operators/map_operator.py b/python/ray/data/_internal/execution/operators/map_operator.py index 5a72596d3af2..ddaddf3dbc41 100644 --- a/python/ray/data/_internal/execution/operators/map_operator.py +++ b/python/ray/data/_internal/execution/operators/map_operator.py @@ -886,6 +886,7 @@ def build_metadata(block_ser_time_s): task_exec_stats=TaskExecWorkerStats( task_wall_time_s=task_dur_s, max_uss_bytes=profiler.estimate_max_uss(), + max_rss_bytes=profiler.max_rss(), # Reported by producing transforms through the # per-task reporter; empty if the op reports nothing. custom_op_stats=op_stats_reporter.get_stats(), diff --git a/python/ray/data/_internal/logical/operators/read_operator.py b/python/ray/data/_internal/logical/operators/read_operator.py index 03d373c1f615..03aa88632e0a 100644 --- a/python/ray/data/_internal/logical/operators/read_operator.py +++ b/python/ray/data/_internal/logical/operators/read_operator.py @@ -448,6 +448,16 @@ class ListFiles(LogicalOperator, SourceOperator): shuffle_config_factory: Callable[[], Optional["FileShuffleConfig"]] = field( default=lambda: None ) + # Pushed-down read constraints, populated by the optimizer rules + # (``predicate_pushdown`` / ``projection_pushdown`` / ``limit_pushdown``). + # A ``StreamingFileChunker`` (e.g. the Parquet footer chunker) uses them to + # prune row groups, size only projected columns, and stop listing early; + # the per-file listing path ignores them. Whether footer-based chunking runs + # is decided by the indexer's chunker type, not a flag here -- this op stays + # format-agnostic. + predicate: Optional[Expr] = None + projected_columns: Optional[List[str]] = None + limit: Optional[int] = None _name: str = field(init=False, repr=False) _input_dependencies: List[LogicalOperator] = field( init=False, repr=False, default_factory=list diff --git a/python/ray/data/_internal/logical/optimizers.py b/python/ray/data/_internal/logical/optimizers.py index a8b1f7b14a0d..59c0952d15f0 100644 --- a/python/ray/data/_internal/logical/optimizers.py +++ b/python/ray/data/_internal/logical/optimizers.py @@ -16,6 +16,7 @@ CombineShuffles, CommonSubExprElimination, ConfigureMapTaskMemoryUsingOutputSize, + DeriveListFilesPushdown, FuseOperators, InheritTargetMaxBlockSizeRule, LimitPushdownRule, @@ -67,7 +68,13 @@ def _post_optimize(self, plan: LogicalPlan) -> LogicalPlan: # CommonSubExprElimination is only supposed to run once # isolated from the optimizer rule loop as it applies to # a single Projection operator not a chain of operators. - return CommonSubExprElimination().apply(plan) + plan = CommonSubExprElimination().apply(plan) + # Must run last, over the final plan: it derives each ``ListFiles``' + # listing-time constraints from the scanner of the ``ReadFiles`` that + # consumes it. Running it earlier (or as a ruleset entry, which a + # caller-added rule could follow) would let a later rewrite leave + # ``ListFiles`` pruning by a predicate the reader no longer applies. + return DeriveListFilesPushdown().apply(plan) class PhysicalOptimizer(Optimizer): diff --git a/python/ray/data/_internal/logical/rules/__init__.py b/python/ray/data/_internal/logical/rules/__init__.py index ae93a545fe4e..405105849abb 100644 --- a/python/ray/data/_internal/logical/rules/__init__.py +++ b/python/ray/data/_internal/logical/rules/__init__.py @@ -6,6 +6,7 @@ ConfigureMapTaskMemoryRule, ConfigureMapTaskMemoryUsingOutputSize, ) +from .derive_list_files_pushdown import DeriveListFilesPushdown from .inherit_target_max_block_size import InheritTargetMaxBlockSizeRule from .limit_pushdown import LimitPushdownRule from .operator_fusion import FuseOperators, are_remote_args_compatible @@ -22,6 +23,7 @@ "CommonSubExprElimination", "ConfigureMapTaskMemoryRule", "ConfigureMapTaskMemoryUsingOutputSize", + "DeriveListFilesPushdown", "FuseOperators", "InheritTargetMaxBlockSizeRule", "LimitPushdownRule", diff --git a/python/ray/data/_internal/logical/rules/derive_list_files_pushdown.py b/python/ray/data/_internal/logical/rules/derive_list_files_pushdown.py new file mode 100644 index 000000000000..c19b482f1a8d --- /dev/null +++ b/python/ray/data/_internal/logical/rules/derive_list_files_pushdown.py @@ -0,0 +1,89 @@ +"""Derive ``ListFiles`` pushdown state from its consuming ``ReadFiles`` scanner. + +This is the single source of truth for the read constraints a ``ListFiles`` +applies while listing. It runs once, after every other logical rule, so no rule +has to remember to keep the two operators in sync. +""" +from dataclasses import replace +from typing import List + +from ray.data._internal.datasource_v2.logical_optimizers import ( + derive_list_files_pushdown, +) +from ray.data._internal.logical.interfaces import LogicalOperator, LogicalPlan, Rule +from ray.data._internal.logical.operators.read_operator import ListFiles, ReadFiles + +__all__ = [ + "DeriveListFilesPushdown", +] + + +class DeriveListFilesPushdown(Rule): + """Recompute every ``ListFiles``' pushed-down read constraints from scratch. + + A metadata-aware indexer (e.g. the footer-based Parquet indexer) reads + ``predicate`` / ``projected_columns`` / ``limit`` off ``ListFiles`` at + planning time to prune row groups by their statistics, size only projected + columns, and stop listing early. Those constraints are only sound if they + are no stronger than what the downstream ``ReadFiles`` actually applies -- + a predicate on ``ListFiles`` that the reader does not evaluate prunes row + groups nobody re-checks, silently dropping rows. + + Rather than have each pushdown rule mirror its own state onto ``ListFiles`` + (an invariant every future rule would have to re-establish by hand), this + rule derives the state from the consuming ``ReadFiles`` scanner and + overwrites whatever was there. Deriving is unconditional in both + directions: a ``ListFiles`` whose consumer is not a ``ReadFiles`` -- e.g. + after ``PushdownCountFiles`` rewrites the plan -- is reset to no + constraints. So a rule that weakens or drops a scanner's predicate + automatically weakens listing too, and the worst a future rule can cause is + listing more than it needs to. + + Runs in ``LogicalOptimizer._post_optimize``, i.e. after the rule loop has + reached a fixed point, so it observes each scanner's final state. + """ + + def apply(self, plan: LogicalPlan) -> LogicalPlan: # pyrefly: ignore[bad-override] + def transform(node: LogicalOperator) -> LogicalOperator: + inputs = node.input_dependencies + if not any(isinstance(input_op, ListFiles) for input_op in inputs): + return node + + # ``ReadFiles`` is the only consumer that applies these constraints + # downstream; for anything else they must be dropped. + scanner = node.scanner if isinstance(node, ReadFiles) else None + predicate, projected_columns, limit = derive_list_files_pushdown(scanner) + + new_inputs: List[LogicalOperator] = [] + changed = False + for input_op in inputs: + if isinstance(input_op, ListFiles) and ( + input_op.predicate is not predicate + or input_op.projected_columns != projected_columns + or input_op.limit != limit + ): + input_op = replace( + input_op, + predicate=predicate, + projected_columns=projected_columns, + limit=limit, + ) + changed = True + new_inputs.append(input_op) + + if not changed: + return node + return node._with_new_input_dependencies(new_inputs) + + dag = plan.dag._apply_transform(transform) + + # A bare ``ListFiles`` root has no consumer at all, so nothing above + # applies its constraints either. + if isinstance(dag, ListFiles) and ( + dag.predicate is not None + or dag.projected_columns is not None + or dag.limit is not None + ): + dag = replace(dag, predicate=None, projected_columns=None, limit=None) + + return LogicalPlan(dag=dag, context=plan.context) diff --git a/python/ray/data/_internal/logical/rules/limit_pushdown.py b/python/ray/data/_internal/logical/rules/limit_pushdown.py index 70f410679047..7f6780599b95 100644 --- a/python/ray/data/_internal/logical/rules/limit_pushdown.py +++ b/python/ray/data/_internal/logical/rules/limit_pushdown.py @@ -196,9 +196,15 @@ def _push_limit_down(self, limit_op: Limit) -> LogicalOperator: num_rows_preserving_ops.append(current_op) current_op = current_op.input_dependencies[0] - # If we couldn't push through any operators, return original - if not num_rows_preserving_ops: + # If we couldn't push through any operators, the Limit sits directly on its + # input, and only a V2 ``ReadFiles`` source is worth rewriting: pushing the + # per-block limit into its scanner also reaches the upstream ``ListFiles`` + # (via ``DeriveListFilesPushdown``) so a footer-based indexer can stop + # listing early. The ``Limit`` stays on top for exact enforcement either + # way. Other ops are left untouched. + if not num_rows_preserving_ops and not isinstance(current_op, ReadFiles): return limit_op + # Apply per-block limit to the deepest operator if it supports it limit_input = self._apply_per_block_limit_if_supported( current_op, limit_op.limit @@ -235,10 +241,12 @@ def _apply_per_block_limit_if_supported( ) if isinstance(op.scanner, SupportsLimitPushdown): - return replace( - op, - scanner=op.scanner.push_limit(limit), - ) + # The pushed limit reaches the upstream ``ListFiles`` -- + # letting a footer-based indexer stop listing early once + # enough exact-survivor rows are found -- via + # ``DeriveListFilesPushdown``, which reads it off the + # scanner once the plan is final. + return replace(op, scanner=op.scanner.push_limit(limit)) return op assert len(op.input_dependencies) == 1, len(op.input_dependencies) return replace( diff --git a/python/ray/data/_internal/logical/rules/predicate_pushdown.py b/python/ray/data/_internal/logical/rules/predicate_pushdown.py index 44ac6e62eb8a..edc56eca95d2 100644 --- a/python/ray/data/_internal/logical/rules/predicate_pushdown.py +++ b/python/ray/data/_internal/logical/rules/predicate_pushdown.py @@ -286,6 +286,11 @@ def _try_push_down_predicate(cls, op: LogicalOperator) -> LogicalOperator: if result_op is input_op: return filter_op + # The pushed (data-column) predicate reaches an upstream + # ``ListFiles`` -- letting a footer-based indexer skip row groups by + # their statistics -- via ``DeriveListFilesPushdown``, which reads it + # off the scanner once the plan is final. Nothing to mirror here. + # Convertible conjuncts were pushed into the read. Re-apply any # residual (non-convertible) conjuncts as a Filter above it. if split.residual is None: diff --git a/python/ray/data/_internal/logical/rules/projection_pushdown.py b/python/ray/data/_internal/logical/rules/projection_pushdown.py index 40c70dd87416..c03b88b617c4 100644 --- a/python/ray/data/_internal/logical/rules/projection_pushdown.py +++ b/python/ray/data/_internal/logical/rules/projection_pushdown.py @@ -471,6 +471,10 @@ def _push_projection_into_read_op(cls, op: LogicalOperator) -> LogicalOperator: if required_columns is None else {name: name for name in required_columns} ) + # The pruned columns reach an upstream ``ListFiles`` -- letting a + # footer-based indexer size only projected columns -- via + # ``DeriveListFilesPushdown``, which reads them off the scanner once + # the plan is final. Nothing to mirror here. projected_input_op = input_op.apply_projection(projection_map) # If the ``Project`` is a pure-prune (only ``col()`` refs, diff --git a/python/ray/data/_internal/logical/rules/pushdown_count_files.py b/python/ray/data/_internal/logical/rules/pushdown_count_files.py index b12c4d0db03a..a2e551c71d25 100644 --- a/python/ray/data/_internal/logical/rules/pushdown_count_files.py +++ b/python/ray/data/_internal/logical/rules/pushdown_count_files.py @@ -1,10 +1,7 @@ -import copy import dataclasses +import logging from typing import TYPE_CHECKING, Optional -from ray.data._internal.datasource_v2.chunkers.file_chunker import ( - WholeFileChunker, -) from ray.data._internal.datasource_v2.listing.file_indexer import ( NonSamplingFileIndexer, ) @@ -27,6 +24,8 @@ if TYPE_CHECKING: import pyarrow as pa +logger = logging.getLogger(__name__) + class PushdownCountFiles(Rule): """Answer ``Dataset.count()`` from file metadata instead of reading data. @@ -39,9 +38,15 @@ class PushdownCountFiles(Rule): Count(ReadFiles(ListFiles)) -> MapBatches(count_rows, ListFiles) ``count_rows`` sums ``read_metadata()`` (e.g. Parquet-footer row counts), - so no data columns are read. The upstream ``ListFiles`` is rebuilt to list - each file exactly once (``WholeFileChunker``, no partitioner) so footers are - read once, in the parallel count pass rather than during listing. + so no data columns are read. The upstream ``ListFiles`` is rebuilt with a + plain whole-file indexer and no partitioner, so (a) each file appears in + exactly one manifest row -- no over-counting -- and (b) listing does no + metadata IO of its own: footers are read once, in the parallel count pass. + + Note this *replaces* metadata-aware indexers such as the footer-based + Parquet one, rather than reconfiguring them: those override ``list_files`` + outright, so reconfiguring is a silent no-op and they would footer-sweep + every file during listing and pack them into a single read unit. """ # Default CPU allocation per task is 1; lower it so at least 2 footer-read @@ -91,16 +96,27 @@ def apply(self, plan: LogicalPlan) -> LogicalPlan: # pyrefly: ignore[bad-overri assert isinstance(list_files, ListFiles), list_files # Rebuild ``ListFiles`` to list each file exactly once: disable - # partitioning and use the ``WholeFileChunker`` (otherwise a file could - # appear once per chunk, in different batches, and be over-counted). - # ``ListFiles`` is frozen, so ``replace`` a copy with a fresh indexer. - count_indexer = copy.deepcopy(list_files.file_indexer) - assert isinstance(count_indexer, NonSamplingFileIndexer), type(count_indexer) - count_indexer._file_chunker = WholeFileChunker() + # partitioning and swap in a plain whole-file indexer. Mutating the + # existing indexer's chunker isn't enough -- a metadata-aware indexer + # (e.g. the footer-based Parquet one) overrides ``list_files`` outright + # and ignores its chunker, so it would keep footer-sweeping during + # listing and could emit a path once per bin, over-counting it. + base_indexer = list_files.file_indexer + if not isinstance(base_indexer, NonSamplingFileIndexer): + # A third-party indexer may chunk files or do its own metadata IO; + # we can't prove one-row-per-file, so leave the plan alone and let + # ``count()`` fall back to the regular read path. + logger.debug( + "Skipping count pushdown: %s is not a NonSamplingFileIndexer", + type(base_indexer).__name__, + ) + return plan + + # ``ListFiles`` is frozen, so ``replace`` it with a fresh indexer. list_files = dataclasses.replace( list_files, file_partitioner=None, - file_indexer=count_indexer, + file_indexer=base_indexer.as_whole_file_indexer(), ) # ``reader`` is narrowed to ``SupportsMetadata`` by the guard above, but diff --git a/python/ray/data/_internal/planner/plan_list_files_op.py b/python/ray/data/_internal/planner/plan_list_files_op.py index 7a01125acb12..360bb1a79a41 100644 --- a/python/ray/data/_internal/planner/plan_list_files_op.py +++ b/python/ray/data/_internal/planner/plan_list_files_op.py @@ -70,6 +70,11 @@ def plan_list_files_op( shuffle_config = op.shuffle_config_factory() + # Some indexers (e.g. the footer-based Parquet indexer) already emit + # bin-packed read units from ``list_files`` -- they need the whole file + # stream on one task to pack globally, and there's nothing left to partition. + yields_read_units = indexer.yields_read_units + transform_fns: List[MapTransformFn] = [ BlockMapTransformFn( partial( @@ -79,6 +84,11 @@ def plan_list_files_op( file_extensions=file_extensions, partition_filter=partition_filter, preserve_order=data_context.execution_options.preserve_order, + # Pushed-down read constraints; metadata-aware indexers use them + # to prune row groups, stop early, and size projected columns. + predicate=op.predicate, + limit=op.limit, + projected_columns=op.projected_columns, ), # Disable block-shaping: produce manifest blocks as-is. disable_block_shaping=True, @@ -97,7 +107,7 @@ def plan_list_files_op( ) ) - if partitioner is not None: + if partitioner is not None and not yields_read_units: transform_fns.append( BlockMapTransformFn( partial(partition_files, partitioner=partitioner), @@ -112,9 +122,10 @@ def plan_list_files_op( _create_input_data_buffer( op, data_context, - # Shuffle needs every manifest on a single task to compute one - # global RNG over the full listing. - should_parallelize=shuffle_config is None, + # A single task is required when shuffle needs one global RNG over + # the full listing, or when the indexer bin-packs read units itself + # (it must see the whole file stream to pack globally). + should_parallelize=shuffle_config is None and not yields_read_units, ), data_context, name="ListFiles", diff --git a/python/ray/data/_internal/planner/plan_read_files_op.py b/python/ray/data/_internal/planner/plan_read_files_op.py index 4b3ff408e394..e89fe187ace3 100644 --- a/python/ray/data/_internal/planner/plan_read_files_op.py +++ b/python/ray/data/_internal/planner/plan_read_files_op.py @@ -21,6 +21,7 @@ from __future__ import annotations import logging +import time from typing import Iterable, List from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest @@ -30,11 +31,12 @@ from ray.data._internal.execution.operators.map_operator import MapOperator from ray.data._internal.execution.operators.map_transformer import ( BlockMapTransformFn, + CustomOpStatsReportFn, MapTransformer, ) from ray.data._internal.logical.operators import ReadFiles from ray.data._internal.output_buffer import OutputBlockSizeOption -from ray.data.block import Block +from ray.data.block import Block, ReadFilesTaskStats from ray.data.context import DataContext logger = logging.getLogger(__name__) @@ -57,8 +59,33 @@ def plan_read_files_op( scanner = op.scanner block_udf = op.block_udf - def do_read(blocks: Iterable[Block], _: TaskContext) -> Iterable[Block]: + def do_read( + blocks: Iterable[Block], + _: TaskContext, + report_custom_op_stats: CustomOpStatsReportFn, + ) -> Iterable[Block]: + task_start_s = time.perf_counter() reader = scanner.create_reader() + # Reader-level per-task aggregates (bytes/batches the decoder actually + # produced, time spent inside its iterator, largest single table), + # folded into per-task distributions by + # ``OpRuntimeMetrics.on_task_finished``. The driver reads them off the + # FINAL output block's ``TaskExecWorkerStats``, and whether a block is + # emitted after this generator's last resume depends on how the + # shaping buffer's flush happens to align with the reader's batch + # sizes — so a single report at end-of-task is silently dropped on + # some (reader-dependent!) shapes. Instead, report ONE stats object up + # front and update it in place as batches flow: every block's + # snapshot carries the totals so far, and the final block's carries + # (at least) everything up to the last yielded table. + task_stats = ReadFilesTaskStats() + report_custom_op_stats(task_stats) + decode_wall_s = 0.0 + decoded_bytes = decoded_batches = decoded_rows = peak_batch_bytes = 0 + manifests = 0 + trim_wall_s = 0.0 + yield_wall_s = 0.0 + first_table_wall_s = 0.0 # File-level predicate pruning (partition predicates pushed down # onto the scanner) runs per incoming manifest block. Only # ``FileScanner`` subclasses expose ``prune_manifest``; the base @@ -70,10 +97,73 @@ def do_read(blocks: Iterable[Block], _: TaskContext) -> Iterable[Block]: manifest = scanner.prune_manifest(manifest) if len(manifest) == 0: continue - for table in reader.read(manifest): - if block_udf is not None: - table = block_udf(table) - yield table + manifests += 1 + table_iter = iter(reader.read(manifest)) + # One-table lookahead: a block's stats snapshot is pickled when + # the block leaves the task, and a table that completes a block + # (buffer >= target, below the 1.5x slice limit -> emitted whole, + # no remainder) is followed by NO flush block. Anything learned + # only when the stream ends -- the reader's end-of-stream drain + # (arrow-rs malloc_trim wall) -- therefore has to be folded in + # BEFORE the last table is yielded, which means pulling the next + # table first. Memory-neutral: the previous table stayed bound in + # this frame during the next decode anyway. + pending = None + while True: + start_s = time.perf_counter() + try: + table = next(table_iter) + except StopIteration: + # The reader's finalizer (incl. the eos trim) ran inside + # this next(); its wall is inside decode_wall_s. + decode_wall_s += time.perf_counter() - start_s + trim_wall_s += reader.pop_task_stats().get("trim_wall_s", 0.0) + task_stats._update( + decode_wall_s=decode_wall_s, + trim_wall_s=trim_wall_s, + yield_wall_s=yield_wall_s, + ) + break + except BaseException: + # A table already decoded when a LATER next() fails must + # still reach the output (same as before the lookahead, + # when it had been yielded before that next() ran) so a + # task tolerated via ``max_errored_blocks`` keeps it. + if pending is not None: + y0 = time.perf_counter() + yield pending + yield_wall_s += time.perf_counter() - y0 + pending = None + raise + decode_wall_s += time.perf_counter() - start_s + if decoded_batches == 0: + first_table_wall_s = time.perf_counter() - task_start_s + nbytes = table.nbytes + decoded_bytes += nbytes + decoded_batches += 1 + decoded_rows += table.num_rows + if nbytes > peak_batch_bytes: + peak_batch_bytes = nbytes + task_stats._update( + decode_wall_s=decode_wall_s, + decoded_bytes=decoded_bytes, + decoded_batches=decoded_batches, + decoded_rows=decoded_rows, + peak_batch_bytes=peak_batch_bytes, + manifests=manifests, + yield_wall_s=yield_wall_s, + first_table_wall_s=first_table_wall_s, + ) + if pending is not None: + y0 = time.perf_counter() + yield pending + yield_wall_s += time.perf_counter() - y0 + pending = block_udf(table) if block_udf is not None else table + if pending is not None: + y0 = time.perf_counter() + yield pending + yield_wall_s += time.perf_counter() - y0 + task_stats._update(yield_wall_s=yield_wall_s) return MapOperator.create( MapTransformer( @@ -84,6 +174,7 @@ def do_read(blocks: Iterable[Block], _: TaskContext) -> Iterable[Block]: output_block_size_option=OutputBlockSizeOption.of( target_max_block_size=data_context.target_max_block_size, ), + should_report_custom_op_stats=True, ), ] ), diff --git a/python/ray/data/_internal/util.py b/python/ray/data/_internal/util.py index a916427eca62..9b7945277848 100644 --- a/python/ray/data/_internal/util.py +++ b/python/ray/data/_internal/util.py @@ -1719,6 +1719,10 @@ def __init__(self, poll_interval_s: Optional[float]): self._process = psutil.Process(os.getpid()) self._max_uss = None + # Peak RSS, sampled by the same poll (memory_info() returns both). RSS + # counts shared pages (e.g. mapped plasma blocks) that USS excludes, so + # the pair distinguishes private decode memory from OS-visible footprint. + self._max_rss = None self._max_uss_lock = threading.Lock() self._uss_poll_thread = None @@ -1752,17 +1756,31 @@ def estimate_max_uss(self) -> Optional[int]: return None with self._max_uss_lock: - if self._max_uss is None: - self._max_uss = self._estimate_uss() - else: - self._max_uss = max(self._max_uss, self._estimate_uss()) + self._sample() assert self._max_uss is not None return self._max_uss + def max_rss(self) -> Optional[int]: + """Get the max RSS of the current process observed by the poll. + + Sampled by the same poll as :meth:`estimate_max_uss` (Linux-only, like + USS, so the two metrics always appear together). Returns ``None`` when + unavailable. + """ + if not self._can_estimate_uss(): + assert self._max_rss is None + return None + + with self._max_uss_lock: + self._sample() + + return self._max_rss + def reset(self): with self._max_uss_lock: self._max_uss = None + self._max_rss = None def _start_uss_poll_thread(self) -> Tuple[threading.Thread, threading.Event]: assert self._poll_interval_s is not None @@ -1773,10 +1791,7 @@ def _start_uss_poll_thread(self) -> Tuple[threading.Thread, threading.Event]: def poll_uss(): while not stop_event.is_set(): with self._max_uss_lock: - if self._max_uss is None: - self._max_uss = self._estimate_uss() - else: - self._max_uss = max(self._max_uss, self._estimate_uss()) + self._sample() stop_event.wait(self._poll_interval_s) thread = threading.Thread(target=poll_uss, daemon=True) @@ -1789,13 +1804,25 @@ def _stop_uss_poll_thread(self): self._stop_uss_poll_event.set() self._uss_poll_thread.join() - def _estimate_uss(self) -> int: + def _sample(self): + """Take one memory sample and fold it into the running maxima. + + Caller must hold ``self._max_uss_lock``. + """ assert self._can_estimate_uss() memory_info = self._process.memory_info() # Estimate the USS (the amount of memory that'd be free if we killed the # process right now) as the difference between the RSS (total physical memory) # and amount of shared physical memory. - return memory_info.rss - memory_info.shared + uss = memory_info.rss - memory_info.shared + if self._max_uss is None: + self._max_uss = uss + else: + self._max_uss = max(self._max_uss, uss) + if self._max_rss is None: + self._max_rss = memory_info.rss + else: + self._max_rss = max(self._max_rss, memory_info.rss) @staticmethod @functools.cache diff --git a/python/ray/data/block.py b/python/ray/data/block.py index 02f8d42d3fa2..06e396f8106c 100644 --- a/python/ray/data/block.py +++ b/python/ray/data/block.py @@ -204,6 +204,64 @@ def __post_init__(self): raise TypeError("CustomOpStats cannot be instantiated directly") +@DeveloperAPI +@dataclass(frozen=True) +class ReadFilesTaskStats(CustomOpStats): + """Per-task aggregate of a ``ReadFiles`` task's decode loop. + + Reported once per read task by the ``ReadFiles`` transform (see + ``plan_read_files_op.py``) and folded into per-task distributions by + ``OpRuntimeMetrics.on_task_finished``. These are the reader-level facts the + node- and task-memory metrics cannot see: how many bytes the decoder + actually produced, how long the task spent inside the reader's iterator + (vs. downstream block shaping / fused compute), and the largest single + table the reader handed over (a proxy for the decode working set). + Reader-implementation agnostic: identical for the PyArrow and arrow-rs + Parquet readers, and for every other file-based V2 datasource. + + Worker-side mutability: the driver reads these off the FINAL output + block's ``TaskExecWorkerStats``, but whether any block is emitted after + the read generator's last resume depends on buffer/batch alignment — a + report made only at end-of-task is dropped on some shapes. The transform + therefore reports one instance up front and keeps it current via + ``_update`` as batches flow; every block snapshot then carries the totals + so far. Frozen like its base class — ``_update`` is the single sanctioned + writer (worker-side only; the driver must treat instances as immutable). + """ + + def _update(self, **fields: Any) -> None: + for key, value in fields.items(): + object.__setattr__(self, key, value) + + # Wall-clock seconds spent inside the reader's table iterator (pure + # decode + IO wait), excluding downstream consumption of the tables. + decode_wall_s: float = 0.0 + # Sum of ``table.nbytes`` over every table the reader yielded. + decoded_bytes: int = 0 + # Number of tables the reader yielded. + decoded_batches: int = 0 + # Sum of rows over those tables. + decoded_rows: int = 0 + # ``nbytes`` of the single largest yielded table. + peak_batch_bytes: int = 0 + # Number of (non-empty, post-pruning) file manifests the task read. + manifests: int = 0 + # Wall seconds the reader spent in its end-of-stream finalizer — for the + # arrow-rs reader, the ``arrow_rs_malloc_trim_eos`` glibc ``malloc_trim(0)`` + # (0 for readers without one). CONTAINED in ``decode_wall_s``: the + # finalizer runs inside the iterator's final ``next()``. Readers hand it + # over through ``Reader.pop_task_stats``. + trim_wall_s: float = 0.0 + # Wall seconds the read task spent handing tables DOWNSTREAM — inside the + # planner's ``yield``: output-buffer shaping, block build, the object-store + # put and any streaming-generator backpressure wait. Disjoint from + # ``decode_wall_s``; task duration minus the two is start-up/teardown. + yield_wall_s: float = 0.0 + # Wall seconds from task start to the first decoded table (reader + # construction + the first ``next()``): the per-task fixed cost. + first_table_wall_s: float = 0.0 + + @DeveloperAPI @dataclass(frozen=True) class TaskExecWorkerStats: @@ -216,6 +274,12 @@ class TaskExecWorkerStats: # or None if USS measurement is unavailable (e.g., non-Linux platforms). max_uss_bytes: Optional[int] = None + # Peak RSS (Resident Set Size) memory in bytes observed during the task, + # sampled by the same poll as ``max_uss_bytes`` (None off-Linux likewise). + # RSS counts shared pages (e.g. mapped object-store blocks) that USS + # excludes, so the pair separates private working set from OS footprint. + max_rss_bytes: Optional[int] = None + # Operator-specific worker-reported stats: one CustomOpStats entry per # reporting transform (fused transforms each contribute one). Empty for # operators that do not report any extra stats. diff --git a/python/ray/data/context.py b/python/ray/data/context.py index 1ec774a3ac23..c2e5d1657c2b 100644 --- a/python/ray/data/context.py +++ b/python/ray/data/context.py @@ -84,6 +84,39 @@ class ShuffleStrategy(str, enum.Enum): DEFAULT_USE_DATASOURCE_V2 = env_bool("RAY_DATA_USE_DATASOURCE_V2", True) +# Prototype flag: route the V2 Parquet read path through the experimental +# arrow-rs (Rust) reader instead of PyArrow. Only takes effect when +# ``use_datasource_v2`` is also set. Requires the native ``ray_data_arrow_rs`` +# module to be installed. Defaults to False. +DEFAULT_USE_ARROW_RS_PARQUET_READER = env_bool( + # TESTING COMMIT ONLY — default flipped to True so every read_parquet in the + # release suite exercises the Rust reader. Revert this commit before opening + # the PR: the shipping default MUST stay False (opt-in). Paired with + # release/ray_release/byod/byod_arrow_rs_parquet.sh (installs the crate). + "RAY_DATA_USE_ARROW_RS_PARQUET_READER", + True, +) + +# With the arrow-rs reader: set glibc's M_TRIM_THRESHOLD to 0 (via mallopt) in +# each worker process on first native read, so freed decode heap returns to +# the OS immediately instead of accumulating as idle-worker USS under task +# churn (measured +492 MiB over ~100 tasks on fused read->write; the trim +# lever removed it entirely — arrow_rs_docs findings M48). Linux/glibc only; +# a no-op elsewhere. Default False until the lever's wall cost is certified. +DEFAULT_ARROW_RS_MALLOC_TRIM = env_bool("RAY_DATA_ARROW_RS_MALLOC_TRIM", False) + +# With the arrow-rs reader: call glibc ``malloc_trim(0)`` ONCE at the end of +# each read task's stream, handing that task's freed decode heap back to the OS +# without touching the allocator's thresholds while it decodes. The mallopt +# variant above collapses the same idle-worker floor but costs 24-36% wall, +# because any explicit M_TRIM_THRESHOLD also disables glibc's dynamic mmap +# threshold so every large decode buffer goes mmap/munmap (arrow_rs_docs +# findings M61/M64). Linux/glibc only; a no-op elsewhere. Default True since +# 2026-09-08: on the release fleet it closed every arrow-rs retention row at +# ~rs wall (findings M101, build 106096) and replicated x3 on the gate cells +# with no wall price (M124, build 106284); set the env var to 0 to ablate. +DEFAULT_ARROW_RS_MALLOC_TRIM_EOS = env_bool("RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS", True) + # Default target chunk size for ``ParquetFileChunker``. ``None`` means the chunker # uses its built-in default (currently 1 GiB). DEFAULT_PARQUET_CHUNKER_TARGET_CHUNK_SIZE: Optional[int] = None @@ -591,6 +624,12 @@ class DataContext: override with ``RAY_DATA_USE_DATASOURCE_V2`` (``0`` for V1, ``1`` for V2). Parquet is the only reader migrated to V2 so far; the others read through V1 for now regardless of this flag. + use_arrow_rs_parquet_reader: Prototype flag. When True (and + ``use_datasource_v2`` is also True), ``ParquetScanner.create_reader()`` + returns the experimental arrow-rs (Rust) reader + (``ArrowRsParquetFileReader``) instead of the PyArrow + ``ParquetFileReader``. Requires the native ``ray_data_arrow_rs`` + module. Defaults to False. parquet_chunker_target_chunk_size: Target chunk size in bytes used by ``ParquetFileChunker`` when splitting large Parquet files into multiple read tasks. When ``None``, the chunker's built-in default @@ -899,6 +938,9 @@ class DataContext: min_parallelism: int = DEFAULT_MIN_PARALLELISM read_op_min_num_blocks: int = DEFAULT_READ_OP_MIN_NUM_BLOCKS use_datasource_v2: bool = DEFAULT_USE_DATASOURCE_V2 + use_arrow_rs_parquet_reader: bool = DEFAULT_USE_ARROW_RS_PARQUET_READER + arrow_rs_malloc_trim: bool = DEFAULT_ARROW_RS_MALLOC_TRIM + arrow_rs_malloc_trim_eos: bool = DEFAULT_ARROW_RS_MALLOC_TRIM_EOS # Target chunk size in bytes for ``ParquetFileChunker``. When ``None``, the # chunker uses its built-in default (currently 1 GiB). parquet_chunker_target_chunk_size: Optional[ diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py index 330d803d1293..970526d65bbc 100644 --- a/python/ray/data/dataset.py +++ b/python/ray/data/dataset.py @@ -6070,10 +6070,11 @@ def write_lance( .. testcode:: import ray import pandas as pd + from ray.data import SaveMode docs = [{"title": "Lance data sink test"} for key in range(4)] ds = ray.data.from_pandas(pd.DataFrame(docs)) - ds.write_lance("/tmp/data/") + ds.write_lance("/tmp/lance_data/", mode=SaveMode.OVERWRITE) Args: path: The path to the destination Lance dataset. Ignored when namespace diff --git a/python/ray/data/read_api.py b/python/ray/data/read_api.py index a47ffa74b8c5..7e0c9f606567 100644 --- a/python/ray/data/read_api.py +++ b/python/ray/data/read_api.py @@ -542,11 +542,17 @@ def _read_datasource_v2( # (``-1`` when unset). Honoring it here per-read avoids mutating the # process-global ``DataContext.read_op_min_num_blocks``. num_buckets = parallelism if parallelism != -1 else ctx.read_op_min_num_blocks - partitioner = RoundRobinPartitioner( - in_memory_size_estimator=datasource.get_size_estimator(), - min_bucket_size=min_bucket_size, - max_bucket_size=max_bucket_size, - num_buckets=num_buckets, + # An indexer that already emits bin-packed read units (e.g. the footer-based + # Parquet indexer) doesn't use the size-estimate ``RoundRobinPartitioner``. + partitioner = ( + None + if getattr(indexer, "yields_read_units", False) + else RoundRobinPartitioner( + in_memory_size_estimator=datasource.get_size_estimator(), + min_bucket_size=min_bucket_size, + max_bucket_size=max_bucket_size, + num_buckets=num_buckets, + ) ) # NOTE: We're using shuffle config factory to fix the seed at the planning diff --git a/python/ray/data/test.bzl b/python/ray/data/test.bzl new file mode 100644 index 000000000000..1d123050cd73 --- /dev/null +++ b/python/ray/data/test.bzl @@ -0,0 +1,32 @@ +"""Wrappers that inject default Ray Data test environment variables. + +Keeps the Parquet footer-reader actor pool tiny for unit/integration tests +that exercise small fixtures. Release/perf jobs that need the production +default should not use these wrappers (or should override ``env``). +""" + +load("@rules_python//python:defs.bzl", _py_test = "py_test") +load("//bazel:python.bzl", _doctest = "doctest", _py_test_module_list = "py_test_module_list") + +_DATA_TEST_ENV = { + # Default 32-actor footer pool times out / warns under CI parallelism. + "RAY_DATA_PARQUET_FOOTER_NUM_ACTORS": "1", +} + +def _merge_env(env): + merged = dict(_DATA_TEST_ENV) + if env: + merged.update(env) + return merged + +def py_test(**kwargs): + kwargs["env"] = _merge_env(kwargs.pop("env", None)) + _py_test(**kwargs) + +def py_test_module_list(**kwargs): + kwargs["env"] = _merge_env(kwargs.pop("env", None)) + _py_test_module_list(**kwargs) + +def doctest(**kwargs): + kwargs["env"] = _merge_env(kwargs.pop("env", None)) + _doctest(**kwargs) diff --git a/python/ray/data/tests/conftest.py b/python/ray/data/tests/conftest.py index f3905bd854a1..0479b737cfcc 100644 --- a/python/ray/data/tests/conftest.py +++ b/python/ray/data/tests/conftest.py @@ -30,6 +30,12 @@ from ray.util.debug import reset_log_once from ray.util.state import list_actors +# Keep the footer-reader pool tiny for unit/integration tests. The default +# 32-actor pool times out under CI parallelism; tests that need a larger pool +# can override with monkeypatch.setenv. Mirrored in python/ray/data/test.bzl +# for bazel test targets. +os.environ.setdefault("RAY_DATA_PARQUET_FOOTER_NUM_ACTORS", "1") + def mock_all_to_all_op(input_op, name="MockAllToAll"): """Create a mock AllToAllOperator for testing. diff --git a/python/ray/data/tests/datasource/test_arrow_rs_parquet_reader.py b/python/ray/data/tests/datasource/test_arrow_rs_parquet_reader.py new file mode 100644 index 000000000000..29305e19f54a --- /dev/null +++ b/python/ray/data/tests/datasource/test_arrow_rs_parquet_reader.py @@ -0,0 +1,3431 @@ +"""Correctness + integration tests for the experimental arrow-rs Parquet reader. + +These run only when the native ``ray_data_arrow_rs`` extension is importable +(built via ``maturin`` from the crate under +``_internal/datasource_v2/native/ray_data_arrow_rs/``); otherwise the whole +module is skipped. They confirm that: + +- reading through the arrow-rs path yields byte-identical columns to PyArrow, +- the native decode path actually runs (not the PyArrow fallback), and +- unsupported schemas transparently fall back to PyArrow and stay correct. +""" +import os + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +import pyarrow.parquet as pq +import pytest + +import ray +from ray.data.context import DataContext +from ray.data.datasource.path_util import _unwrap_protocol + +ray_data_arrow_rs = pytest.importorskip("ray_data_arrow_rs") + + +@pytest.fixture +def restore_ctx(): + ctx = DataContext.get_current() + v2, arrow_rs = ctx.use_datasource_v2, ctx.use_arrow_rs_parquet_reader + try: + ctx.use_datasource_v2 = True + yield ctx + finally: + ctx.use_datasource_v2 = v2 + ctx.use_arrow_rs_parquet_reader = arrow_rs + + +def _whole_file_manifest(): + """FileManifest stand-in for tests that call ``_resolve_batch_size`` directly. + + Carries no footer chunk stats (like a ``WholeFileChunker`` manifest), so the + reader takes its footer-probe fallback — the pre-#64985 behaviour these tests + were written against. Tests that go through ``read()`` never need this; the + real manifest arrives there. + """ + from types import SimpleNamespace + + return SimpleNamespace(file_chunk_metadatas=[None]) + + +def _flat_table(num_rows=20_000): + rng = np.random.default_rng(0) + return pa.table( + { + "id": pa.array(np.arange(num_rows, dtype=np.int64)), + "x": pa.array(rng.random(num_rows)), + "label": pa.array((np.arange(num_rows) % 5).astype(np.int32)), + "name": pa.array([f"row-{i}" for i in range(num_rows)]), + } + ) + + +def _read_sorted(path, use_arrow_rs, restore_ctx, **read_kwargs): + restore_ctx.use_arrow_rs_parquet_reader = use_arrow_rs + ds = ray.data.read_parquet(str(path), **read_kwargs) + return pa.Table.from_pandas(ds.to_pandas()).sort_by("id") + + +def _read_arrow_sorted(path, use_arrow_rs, restore_ctx, **read_kwargs): + """Like :func:`_read_sorted` but materializes straight to Arrow (no pandas + round-trip). Required for columns whose Arrow type can't round-trip through + pandas — e.g. multi-dimensional tensor extensions, which ``from_pandas`` + rejects ("Can only convert 1-dimensional array values"). Both readers go + through this identically, so the ``.equals`` comparison stays a fair parity + check on the reader output itself.""" + restore_ctx.use_arrow_rs_parquet_reader = use_arrow_rs + ds = ray.data.read_parquet(str(path), **read_kwargs) + return pa.concat_tables(ray.get(ds.to_arrow_refs())).sort_by("id") + + +@pytest.mark.parametrize("row_group_size", [20_000, 5_000]) +def test_arrow_rs_parity_full_scan(tmp_path, restore_ctx, row_group_size): + """arrow-rs and PyArrow produce identical tables (full scan).""" + path = tmp_path / "data.parquet" + table = _flat_table() + pq.write_table( + table, str(path), write_page_index=True, row_group_size=row_group_size + ) + + pa_tbl = _read_sorted(path, False, restore_ctx) + rs_tbl = _read_sorted(path, True, restore_ctx) + + assert pa_tbl.num_rows == rs_tbl.num_rows == table.num_rows + assert pa_tbl.equals(rs_tbl) + + +def test_arrow_rs_parity_with_projection(tmp_path, restore_ctx): + path = tmp_path / "data.parquet" + table = _flat_table() + pq.write_table(table, str(path), write_page_index=True) + + pa_tbl = _read_sorted(path, False, restore_ctx, columns=["id", "x"]) + rs_tbl = _read_sorted(path, True, restore_ctx, columns=["id", "x"]) + assert rs_tbl.column_names == ["id", "x"] + assert pa_tbl.equals(rs_tbl) + + +@pytest.mark.parametrize("row_group_size", [20_000, 5_000]) +def test_arrow_rs_parity_sum(tmp_path, restore_ctx, row_group_size): + """The aggregation workload benchmarked in Agents.md §3.3 (``ds.sum()``) must + return identical results via the arrow-rs decode path and PyArrow. This is a + decode-heavy / output-light consumer: the read decodes every value and the + aggregation collapses it to a scalar, so it exercises full-column decode + correctness end-to-end through Ray's aggregation, not just a raw table read. + """ + path = tmp_path / "data.parquet" + table = _flat_table() + pq.write_table( + table, str(path), write_page_index=True, row_group_size=row_group_size + ) + + # Ground truth from the source table, independent of either reader. + expected_id = pc.sum(table["id"]).as_py() + expected_label = pc.sum(table["label"]).as_py() + + restore_ctx.use_arrow_rs_parquet_reader = False + pa_sum = ray.data.read_parquet(str(path)).sum(["id", "label"]) + restore_ctx.use_arrow_rs_parquet_reader = True + rs_sum = ray.data.read_parquet(str(path)).sum(["id", "label"]) + + assert rs_sum == pa_sum + assert rs_sum["sum(id)"] == expected_id + assert rs_sum["sum(label)"] == expected_label + + +def _read_crate_stream(path, **kwargs): + """Read a file straight through the crate (bypassing the reader) into a + single table, so we can force the K-split path via ``split_threshold_bytes`` + / ``k`` explicitly.""" + stream = ray_data_arrow_rs.read_row_groups(str(path), **kwargs) + return pa.RecordBatchReader.from_stream(stream).read_all() + + +@pytest.mark.parametrize("k", [2, 4, 8]) +def test_kspilt_parity_and_order(tmp_path, k): + """The intra-fragment K-split path (single big row group, forced via + ``split_threshold_bytes=0``) must be byte-identical to both the sequential + (k=1) crate path and PyArrow, and preserve row order across the K parallel + range workers. + + Row order is the load-bearing property here: the split decodes K disjoint + row ranges on separate threads and merges them back. A merge bug would + surface as a shuffled ``id`` column even when the row *set* is correct, so + we assert the ``id`` column is exactly ``0..n-1`` in order. + """ + num_rows = 50_000 + path = tmp_path / "big_single_rg.parquet" + table = _flat_table(num_rows) + # One row group covering all rows → a lone fragment Ray's pool can't split. + pq.write_table(table, str(path), write_page_index=True, row_group_size=num_rows) + assert pq.ParquetFile(str(path)).num_row_groups == 1 + + # k=1 sequential (never splits) vs forced K-split (threshold=0). + seq = _read_crate_stream(path, k=1) + split = _read_crate_stream(path, k=k, split_threshold_bytes=0) + + # Byte-identical to the sequential path and to the source table. + assert split.equals(seq) + assert split.equals(table) + # Order preserved across ranges: id is exactly 0..n-1, not just the right set. + assert split.column("id").to_pylist() == list(range(num_rows)) + + +def test_native_path_actually_runs(tmp_path): + """Directly exercise the reader and confirm it calls the native extension + rather than silently falling back to PyArrow.""" + import pyarrow.dataset as pds + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + path = tmp_path / "data.parquet" + table = _flat_table() + pq.write_table(table, str(path), write_page_index=True) + + calls = {"n": 0} + orig = ray_data_arrow_rs.read_row_groups + + def wrapped(*a, **k): + calls["n"] += 1 + return orig(*a, **k) + + ray_data_arrow_rs.read_row_groups = wrapped + try: + reader = ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), target_block_size=128 * 1024 * 1024 + ) + dataset = pds.dataset(str(path), format="parquet", filesystem=LocalFileSystem()) + fragment = next(dataset.get_fragments()) + scanner_kwargs = { + "columns": None, + "filter": None, + "batch_size": reader._resolve_batch_size(dataset, _whole_file_manifest()), + } + got = pa.concat_tables( + list(reader._iter_fragment_tables(fragment, scanner_kwargs)) + ) + finally: + ray_data_arrow_rs.read_row_groups = orig + + assert calls["n"] > 0, "native read_row_groups was not called (fell back)" + assert got.sort_by("id").equals(table.sort_by("id")) + + +def _make_manifest(paths, sizes, chunk_metadatas): + from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest + + return FileManifest.construct_manifest(paths, sizes, chunk_metadatas) + + +class _HandleProxy: + """Wraps a crate ``NativeParquetFile`` handle so tests can count decode + calls — pyo3 methods can't be monkeypatched, but the reader only sees the + object ``open_parquet_file`` / ``open_file`` returned, so a forwarding + proxy is observationally identical.""" + + def __init__(self, handle, counters, on_decode=None): + self._handle = handle + self._counters = counters + self._on_decode = on_decode + + def read_row_groups(self, *a, **k): + self._counters["decode"] += 1 + self._counters["decode_calls"].append((a, k)) + if self._on_decode is not None: + self._on_decode(self._counters) + return self._handle.read_row_groups(*a, **k) + + def __getattr__(self, name): + return getattr(self._handle, name) + + +def _spy_native_decode(monkeypatch, on_decode=None): + """Count planned-path native activity on local files: ``open`` = per-file + handle opens (footer parses), ``decode`` = ``read_row_groups`` calls on + those handles. Since the per-file handle API (TODO 1r) this — not the + module-level ``read_row_groups`` — is how a planned read decodes. + ``on_decode(counters)`` runs before each decode, so tests can inject + failures (raise) at exactly the point the crate would start reading. + ``decode_calls`` holds each decode's ``(args, kwargs)`` so tests can + assert which tuning knobs reached the crate.""" + counters = {"open": 0, "decode": 0, "decode_calls": []} + orig_open = ray_data_arrow_rs.open_parquet_file + + def spy_open(*a, **k): + counters["open"] += 1 + return _HandleProxy(orig_open(*a, **k), counters, on_decode) + + monkeypatch.setattr(ray_data_arrow_rs, "open_parquet_file", spy_open) + return counters + + +def test_native_read_is_pyarrow_free(tmp_path, monkeypatch): + """The whole point of the ``read()`` rewrite: for a file the native reader + supports, PyArrow must *never open it*. The footer, row-group layout, and + decode all come from the crate — since the per-file handle API (TODO 1r), + via ``open_parquet_file`` (one footer parse per file) whose handle then + serves both ``metadata()`` and ``read_row_groups()``; the old per-call + entry points (``read_metadata`` / ``read_row_groups``) must NOT run on a + planned read, or the footer is being parsed twice. ``pyarrow.dataset.dataset`` + — the only way the base reader opens a Parquet file — must not be called + at all. + + We drive ``reader.read(manifest)`` directly (not through + ``ray.data.read_parquet``) so the assertion is scoped to the *read* stage: + the listing/indexing stage does call pyarrow to enumerate row groups, and + counting over a full pipeline execution would conflate the two. Here the + manifest is handed in pre-built, so any ``pds.dataset`` call can only come + from the reader itself. + """ + import pyarrow.dataset as pds + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + path = tmp_path / "data.parquet" + table = _flat_table() + pq.write_table(table, str(path), write_page_index=True) + + # Spy: pyarrow.dataset.dataset must NOT be called for a supported native read. + ds_calls = {"n": 0} + orig_dataset = pds.dataset + + def spy_dataset(*a, **k): + ds_calls["n"] += 1 + return orig_dataset(*a, **k) + + monkeypatch.setattr(pds, "dataset", spy_dataset) + + # Spy: the native crate must actually run (one handle open per file), so a + # "0 pyarrow calls" result can't be a silent no-op — and the per-call entry + # points must stay cold (each would re-parse the footer the handle holds). + native_calls = {"open_parquet_file": 0, "read_metadata": 0, "read_row_groups": 0} + for name in native_calls: + orig = getattr(ray_data_arrow_rs, name) + + def make_spy(orig, name): + def spy(*a, **k): + native_calls[name] += 1 + return orig(*a, **k) + + return spy + + monkeypatch.setattr(ray_data_arrow_rs, name, make_spy(orig, name)) + + reader = ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), target_block_size=128 * 1024 * 1024 + ) + manifest = _make_manifest([str(path)], [os.path.getsize(path)], [None]) + got = pa.concat_tables(list(reader.read(manifest))) + + assert ds_calls["n"] == 0, ( + "pyarrow.dataset.dataset was called during a supported native read " + "(pyarrow opened the file — the read is not pyarrow-free)" + ) + assert native_calls["open_parquet_file"] == 1, ( + "expected exactly one native handle open for the one-file read, got " + f"{native_calls['open_parquet_file']}" + ) + assert native_calls["read_metadata"] == 0, ( + "per-call footer entry point ran on a planned read — the footer was " + "parsed twice instead of reused from the handle" + ) + assert ( + native_calls["read_row_groups"] == 0 + ), "per-call decode entry point ran on a planned read instead of the handle" + assert got.sort_by("id").equals(table.sort_by("id")) + + +def _write_pickle_object_file(path, objs): + import pickle + + from ray.data._internal.object_extensions.arrow import ArrowPythonObjectType + + ext_type = ArrowPythonObjectType() + storage = pa.array([pickle.dumps(o) for o in objs], type=ext_type.storage_type) + table = pa.table( + { + "id": pa.array(range(len(objs)), type=pa.int64()), + "obj": pa.ExtensionArray.from_storage(ext_type, storage), + } + ) + pq.write_table(table, str(path), write_page_index=True) + + +def test_native_read_rejects_pickle_object_columns(tmp_path, monkeypatch): + """The pyarrow path refuses to serve pickled-object columns without the + explicit env opt-in — unpickling executes arbitrary code, so the guard is + a security boundary, not a convenience. The native path must enforce the + same gate. Regression test for the corpus `pickle_default` finding, where + the native decode served the column with no error.""" + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + marker = tmp_path / "exploit_marker" + + class Exploit: + def __reduce__(self): + return (os.system, (f"touch {marker}",)) + + path = tmp_path / "data.parquet" + _write_pickle_object_file(path, [Exploit()]) + + monkeypatch.delenv("RAY_DATA_AUTOLOAD_PICKLE_OBJECT_SCALAR", raising=False) + + # Spy: the raise must come from the NATIVE path — if the file fell back, + # the pyarrow reader's own guard would fire and this test would prove + # nothing about the native one. + native_calls = _spy_native_decode(monkeypatch) + + reader = ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), target_block_size=128 * 1024 * 1024 + ) + manifest = _make_manifest([str(path)], [os.path.getsize(path)], [None]) + with pytest.raises(ValueError, match="arrow_pickled_object"): + pa.concat_tables(list(reader.read(manifest))) + + assert native_calls["decode"] > 0, "native decode never ran (pyarrow fallback?)" + assert not marker.exists(), "pickle.load executed attacker code" + + +def test_native_read_allows_pickle_object_columns_with_env_var(tmp_path, monkeypatch): + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + path = tmp_path / "data.parquet" + _write_pickle_object_file(path, [{"key": "value"}, {"key": "other"}]) + + monkeypatch.setenv("RAY_DATA_AUTOLOAD_PICKLE_OBJECT_SCALAR", "1") + + reader = ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), target_block_size=128 * 1024 * 1024 + ) + manifest = _make_manifest([str(path)], [os.path.getsize(path)], [None]) + got = pa.concat_tables(list(reader.read(manifest))).sort_by("id") + assert got.column("obj").to_pylist() == [{"key": "value"}, {"key": "other"}] + + +def _write_int96_file(path, num_rows=1_000): + """INT96-physical timestamps, all pre-1970 and 1ns past a microsecond + boundary — the values where decode-time unit coercion (floors) and a + post-decode cast (truncates toward zero) differ by exactly one unit.""" + us_vals = [(i - num_rows) * 86_400_000_000 + i for i in range(num_rows)] + table = pa.table( + { + "id": pa.array(range(num_rows), type=pa.int64()), + "ts": pa.array([v * 1000 + 1 for v in us_vals], type=pa.timestamp("ns")), + } + ) + pq.write_table( + table, + str(path), + use_deprecated_int96_timestamps=True, + store_schema=False, + write_page_index=True, + ) + + +def test_coerce_int96_kwarg_parity(tmp_path, restore_ctx): + """`coerce_int96_timestamp_unit` must yield exactly what the base V2 + pyarrow reader yields. On the full V2 pipeline that is subtle: the pinned + unified schema (kwarg-blind ``pq.read_schema``) casts the coerced values + BACK to the inferred ns — so the kwarg's observable effect is + ms-quantized-and-FLOORED *values* in an ns-typed column. A native ns + decode plus cast can't reproduce the floor on pre-1970 values, so the + reader falls back per file; this test pins the observable contract, + however it's met. Regression test for the corpus `int96_coerce_ms` + finding (native path returned raw ns, ignoring the kwarg).""" + path = tmp_path / "int96.parquet" + _write_int96_file(path) + + kw = {"dataset_kwargs": {"coerce_int96_timestamp_unit": "ms"}} + expected = _read_arrow_sorted( + path, use_arrow_rs=False, restore_ctx=restore_ctx, **kw + ) + got = _read_arrow_sorted(path, use_arrow_rs=True, restore_ctx=restore_ctx, **kw) + assert got.equals(expected) + + # The kwarg must have had its observable effect (guards against a "parity" + # where both readers ignored it): values are ms-quantized and floored — + # the raw values sit 1ns past a µs boundary, so flooring to ms lands on + # the boundary and truncation toward zero would not. + ts = expected.column("ts").cast(pa.int64()).to_pylist() + raw = _read_arrow_sorted(path, use_arrow_rs=False, restore_ctx=restore_ctx) + raw_ts = raw.column("ts").cast(pa.int64()).to_pylist() + assert all(v % 1_000_000 == 0 for v in ts), "values not ms-quantized" + assert all(v <= r for v, r in zip(ts, raw_ts)), "not floored (truncated?)" + + +def test_coerce_int96_kwarg_routes_int96_file_to_fallback(tmp_path, monkeypatch): + """The kwarg-honoring mechanism: a file that decodes an INT96 column under + `coerce_int96_timestamp_unit` must NOT go native; the same file without + the kwarg must.""" + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + path = tmp_path / "int96.parquet" + _write_int96_file(path) + + native_calls = _spy_native_decode(monkeypatch) + + def run(parquet_format_kwargs): + reader = ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), + target_block_size=128 * 1024 * 1024, + parquet_format_kwargs=parquet_format_kwargs, + ) + manifest = _make_manifest([str(path)], [os.path.getsize(path)], [None]) + return pa.concat_tables(list(reader.read(manifest))) + + native_calls["decode"] = 0 + with_kwarg = run({"coerce_int96_timestamp_unit": "ms"}) + assert native_calls["decode"] == 0, "int96 file went native despite the kwarg" + assert with_kwarg.schema.field("ts").type == pa.timestamp("ms") + + native_calls["decode"] = 0 + without_kwarg = run(None) + assert native_calls["decode"] > 0, "int96 file without the kwarg should stay native" + assert without_kwarg.schema.field("ts").type == pa.timestamp("ns") + + +def test_native_chunked_read_row_hash_parity(tmp_path): + """A binned file (one manifest row per bin, each naming explicit physical + ``row_group_ids``) must produce byte-identical ``row_hash`` values via the + native path and PyArrow. + + ``row_hash`` is seeded by ``(fragment_path, file_row_offset)`` per sub- + fragment (:func:`_compute_row_hashes`), so this is the load-bearing test for + :meth:`ArrowRsParquetFileReader._native_fragments_for_file`: under + ``include_row_hash`` it must emit one native fragment *per row group* seeded + with that group's **absolute** pre-filter row offset, mirroring ``prefix[rg_id]`` + in the base :func:`_fragments_from_row_group_ids`. + + The bins here are deliberately **non-contiguous** — ``(0, 2)`` and ``(1, 3)`` — + because that is what upstream statistics pruning produces, and it is the case + that distinguishes a correct implementation from one that accumulates offsets + across the bin's own groups. Accumulating would place group 2 at offset 5_000 + instead of its true 10_000, shifting every hash in that group. A contiguous bin + cannot tell the two apart. + """ + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.chunkers.file_chunker import ( + ParquetRowGroupChunkMetadata, + create_chunk_metadata, + ) + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + from ray.data._internal.datasource_v2.readers.parquet_file_reader import ( + ParquetFileReader, + ) + + path = tmp_path / "data.parquet" + table = _flat_table(20_000) + # 4 row groups of 5k rows each. + pq.write_table(table, str(path), write_page_index=True, row_group_size=5_000) + assert pq.ParquetFile(str(path)).num_row_groups == 4 + + # Two bins over the same file, each holding an interleaved pair of groups. + # Their union is all 4 groups, so the read is still lossless and comparable. + chunks = [ + create_chunk_metadata( + ParquetRowGroupChunkMetadata, + row_group_ids=(0, 2), + num_rows=10_000, + uncompressed_size=1, + ), + create_chunk_metadata( + ParquetRowGroupChunkMetadata, + row_group_ids=(1, 3), + num_rows=10_000, + uncompressed_size=1, + ), + ] + size = os.path.getsize(path) + manifest = _make_manifest([str(path), str(path)], [size, size], chunks) + + def read_all(reader_cls): + reader = reader_cls( + filesystem=LocalFileSystem(), + target_block_size=128 * 1024 * 1024, + include_row_hash=True, + ) + return pa.concat_tables(list(reader.read(manifest))).sort_by("id") + + rs_tbl = read_all(ArrowRsParquetFileReader) + pa_tbl = read_all(ParquetFileReader) + + assert "row_hash" in rs_tbl.column_names + assert rs_tbl.num_rows == table.num_rows + assert rs_tbl.equals(pa_tbl) + + +def test_native_bin_coalesces_into_one_call_without_row_hash(): + """Without ``include_row_hash``, a bin's row groups become **one** native + fragment, not one per group. + + This is the whole of old TODO 1l ("coalesce a chunk's contiguous row groups + into one native call"), obtained by following the base path rather than + inventing our own coalescing: the footer-chunking base collapses a file's + bin-assigned groups into a single sub-fragment so PyArrow can merge the reads, + and we collapse so the crate makes one call instead of N. Each extra call means + a fresh S3 client and its own footer/page-index fetch, so the fan-out is the + expensive shape on the transport every release regression came from. + + Asserted at the fragment level rather than end-to-end because the row *data* is + identical either way — only the call count differs, and that is invisible to a + table comparison. The ``include_row_hash`` arm is covered by + :func:`test_native_chunked_read_row_hash_parity`; here it is the control showing + the fan-out still happens when offsets are actually needed. + """ + from ray.data._internal.datasource_v2.chunkers.file_chunker import ( + ParquetRowGroupChunkMetadata, + create_chunk_metadata, + ) + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + chunk = create_chunk_metadata( + ParquetRowGroupChunkMetadata, + row_group_ids=(0, 2, 3), + num_rows=15_000, + uncompressed_size=1, + ) + row_group_num_rows = [5_000, 5_000, 5_000, 5_000] + + coalesced = ArrowRsParquetFileReader._native_fragments_for_file( + "f.parquet", chunk, row_group_num_rows, None, per_row_group_offsets=False + ) + assert len(coalesced) == 1, "bin fanned out into per-row-group native calls" + fragment, offset = coalesced[0] + assert offset == 0 + assert fragment.row_groups == [0, 2, 3], "coalesced fragment lost a row group" + + fanned = ArrowRsParquetFileReader._native_fragments_for_file( + "f.parquet", chunk, row_group_num_rows, None, per_row_group_offsets=True + ) + # Absolute offsets, so the pruning gap at group 1 is preserved rather than + # closed up: 0, 10_000, 15_000 — not 0, 5_000, 10_000. + assert [(f.row_groups, off) for f, off in fanned] == [ + ([0], 0), + ([2], 10_000), + ([3], 15_000), + ] + + # Counts follow the same granularity, and the coalesced count is the sum over + # the named groups only — never the whole file. + counts = ArrowRsParquetFileReader._native_count_fragments( + "f.parquet", chunk, row_group_num_rows, per_row_group_offsets=False + ) + assert len(counts) == 1 + assert counts[0][0].num_rows == 15_000 + + +def test_native_read_include_paths_parity(tmp_path): + """``include_paths`` synthesis (the ``path`` column) must match PyArrow when + the file is read natively. Driven on the same whole-file manifest through + both readers.""" + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + from ray.data._internal.datasource_v2.readers.parquet_file_reader import ( + ParquetFileReader, + ) + + path = tmp_path / "data.parquet" + table = _flat_table() + pq.write_table(table, str(path), write_page_index=True) + manifest = _make_manifest([str(path)], [os.path.getsize(path)], [None]) + + def read_all(reader_cls): + reader = reader_cls( + filesystem=LocalFileSystem(), + target_block_size=128 * 1024 * 1024, + include_paths=True, + ) + return pa.concat_tables(list(reader.read(manifest))).sort_by("id") + + rs_tbl = read_all(ArrowRsParquetFileReader) + pa_tbl = read_all(ParquetFileReader) + + assert "path" in rs_tbl.column_names + assert set(rs_tbl.column("path").to_pylist()) == {str(path)} + assert rs_tbl.equals(pa_tbl) + + +def test_native_read_partitioning_parity(tmp_path, restore_ctx): + """A partition column (encoded in the directory path, absent from the file's + on-disk schema) must be synthesized identically on the native path. + + This exercises native-path-specific planning: ``_plan_native_read`` derives + ``on_disk_names`` from the *crate's* footer schema, which won't contain the + partition column, so it must land in the synthesize set (not be read from the + file). End-to-end through ``read_parquet`` so Ray's hive-partition detection + drives the layout. + """ + base = tmp_path / "parts" + table = _flat_table(6_000) + for g in range(3): + sub = base / f"grp={g}" + sub.mkdir(parents=True) + part = table.slice(g * 2_000, 2_000) + pq.write_table(part, str(sub / "data.parquet"), write_page_index=True) + + pa_tbl = _read_sorted(base, False, restore_ctx) + rs_tbl = _read_sorted(base, True, restore_ctx) + + assert "grp" in rs_tbl.column_names + assert pa_tbl.equals(rs_tbl) + + +def test_count_is_metadata_only_under_arrow_rs(tmp_path, restore_ctx): + """``ds.count()`` is answered from listing metadata and never invokes the + reader — so it is correct under the arrow-rs flag by construction, with zero + native decode (a count scan reads no data columns, so there is no working set + to shrink; handling it natively would buy no memory). This guards that the + metadata short-circuit stays intact when the flag is on.""" + from ray.data._internal.datasource_v2.readers import ( + arrow_rs_parquet_file_reader as mod, + ) + + path = tmp_path / "data.parquet" + table = _flat_table() + pq.write_table(table, str(path), write_page_index=True) + + read_calls = {"n": 0} + orig_read = mod.ArrowRsParquetFileReader.read + + def spy_read(self, *a, **k): + read_calls["n"] += 1 + return orig_read(self, *a, **k) + + mod.ArrowRsParquetFileReader.read = spy_read + try: + restore_ctx.use_arrow_rs_parquet_reader = True + count = ray.data.read_parquet(str(path)).count() + finally: + mod.ArrowRsParquetFileReader.read = orig_read + + assert count == table.num_rows + assert read_calls["n"] == 0, "count unexpectedly invoked the reader" + + +def test_empty_projection_counts_natively_with_zero_decode(tmp_path, monkeypatch): + """A column-less read (empty projection, no predicate) is answered from the + footer row counts alone: no crate decode, no ``pds.dataset`` — strictly + less work than PyArrow's stub-column scan. The yielded tables are + zero-column with the right ``num_rows``; ``_postprocess``'s stub guard + re-adds the row-preserving stub, exactly as on the base path.""" + import pyarrow.dataset as pds + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + path = tmp_path / "data.parquet" + table = _flat_table() + pq.write_table(table, str(path), write_page_index=True) + manifest = _make_manifest([str(path)], [os.path.getsize(path)], [None]) + + calls = _spy_native_decode(monkeypatch) + pds_calls = {"pds": 0} + orig_dataset = pds.dataset + + def dataset_spy(*a, **k): + pds_calls["pds"] += 1 + return orig_dataset(*a, **k) + + monkeypatch.setattr(pds, "dataset", dataset_spy) + + reader = ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), + target_block_size=128 * 1024 * 1024, + columns=[], + ) + tables = list(reader.read(manifest)) + assert sum(t.num_rows for t in tables) == table.num_rows + # The stub-column guard in ``_postprocess`` preserves row counts. + assert all(len(t.column_names) == 1 for t in tables) + # The whole answer came from the footer: nothing was decoded, and pyarrow + # never opened the file. + assert calls["decode"] == 0 + assert pds_calls["pds"] == 0 + + +def test_native_fragment_read_retries_transient_error(tmp_path, monkeypatch): + """A transient I/O failure during native decode must be retried and + recovered exactly like the PyArrow path — the native ``_NativeParquetFragment`` + flows through the same ``iterate_with_retry`` wrapper in + ``_read_fragments_sequential``. We inject a one-shot retryable error (matching + a default ``retried_io_errors`` pattern) into the handle's decode call and + assert the read still returns byte-correct data and that the crate was + re-invoked (the retry fired) — reusing the same handle, so the retry does + not pay a second footer parse.""" + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + path = tmp_path / "data.parquet" + table = _flat_table() + pq.write_table(table, str(path), write_page_index=True) + manifest = _make_manifest([str(path)], [os.path.getsize(path)], [None]) + + def fail_first_decode(counters): + if counters["decode"] == 1: + # A default-retryable message (context.DEFAULT_RETRIED_IO_ERRORS), so + # no context mutation is needed. Raised before any batch is yielded. + raise OSError("AWS Error SLOW_DOWN: injected transient failure") + + calls = _spy_native_decode(monkeypatch, on_decode=fail_first_decode) + + reader = ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), target_block_size=128 * 1024 * 1024 + ) + got = pa.concat_tables(list(reader.read(manifest))) + + assert calls["decode"] >= 2, "native read was not retried after a transient error" + assert got.sort_by("id").equals(table.sort_by("id")) + + +def test_filter_pushdown_prunes_row_groups(tmp_path, restore_ctx): + """A pushed-down predicate is lowered to the native pruning IR and handed to + the crate, which drops the row groups whose footer statistics prove no row + can match — replacing PyArrow's ``fragment.subset``. On a sorted ``id`` over + 4 row groups, ``id >= 3000`` reaches the crate as a non-None ``predicate_json`` + that prunes to row group ``[3]`` (verified via ``select_row_groups``); a + predicate no row group can satisfy yields an empty result. Results stay + byte-correct because the Python post-filter is the final authority.""" + import pyarrow.dataset as pds + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + from ray.data.expressions import col + + path = tmp_path / "sorted.parquet" + n = 4000 + table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "x": pa.array(np.arange(n) * 0.5), + } + ) + pq.write_table(table, str(path), write_page_index=True, row_group_size=1000) + + # `read_row_groups(path, row_groups, columns, batch_size, budget, k, + # split_threshold, predicate_json)` — capture the row_groups handed in and + # the predicate_json (8th positional) the reader lowered. + seen = [] + orig = ray_data_arrow_rs.read_row_groups + + def wrapped(path_, row_groups, *a, **k): + predicate_json = a[5] if len(a) > 5 else k.get("predicate_json") + seen.append((row_groups, predicate_json)) + return orig(path_, row_groups, *a, **k) + + def _run(predicate): + # Mirror `read()`: the reader's Ray `Expr` predicate drives native + # pruning, and its pyarrow form is the scanner filter used for the + # post-decode row filter. + reader = ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), + target_block_size=128 * 1024 * 1024, + predicate=predicate, + ) + dataset = pds.dataset(str(path), format="parquet", filesystem=LocalFileSystem()) + fragment = next(dataset.get_fragments()) + batch_size = reader._resolve_batch_size(dataset, _whole_file_manifest()) + return list( + reader._iter_fragment_tables( + fragment, + { + "columns": None, + "filter": predicate.to_pyarrow(), + "batch_size": batch_size, + }, + ) + ) + + ray_data_arrow_rs.read_row_groups = wrapped + try: + got = pa.concat_tables(_run(col("id") >= 3000)) + pruned_all = _run(col("id") >= 10**9) + finally: + ray_data_arrow_rs.read_row_groups = orig + + # The crate receives the fragment's full row-group list plus the lowered + # predicate; pruning happens *inside* the crate now, not before the call. + (rg0, pj0), (rg1, pj1) = seen + assert rg0 == [0, 1, 2, 3], f"expected all row groups handed to crate: {rg0}" + assert pj0 is not None, "predicate did not lower to a pruning IR" + # Prove the IR actually prunes to the single satisfying row group. + assert ray_data_arrow_rs.select_row_groups(str(path), pj0) == [3] + assert got.sort_by("id").equals(table.slice(3000)) + # A fully-unsatisfiable predicate prunes every group inside the crate, so + # the stream is empty (the crate is still invoked — one cheap footer read). + assert pruned_all == [] + assert ray_data_arrow_rs.select_row_groups(str(path), pj1) == [] + + +def test_filter_pushdown_e2e_parity(tmp_path, restore_ctx): + """``ds.filter(expr=...)`` goes through the PredicatePushdown rule into the + read; both readers must agree (and match ground truth) on a sorted + multi-row-group file where pruning actually kicks in.""" + path = tmp_path / "sorted_e2e.parquet" + n = 4000 + table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "x": pa.array(np.arange(n) * 0.5), + } + ) + pq.write_table(table, str(path), write_page_index=True, row_group_size=1000) + + def _read(use_arrow_rs): + restore_ctx.use_arrow_rs_parquet_reader = use_arrow_rs + ds = ray.data.read_parquet(str(path)).filter(expr="id >= 3500") + return pa.Table.from_pandas(ds.to_pandas()).sort_by("id") + + pa_tbl = _read(False) + rs_tbl = _read(True) + assert pa_tbl.num_rows == rs_tbl.num_rows == 500 + assert pa_tbl.equals(rs_tbl) + + +@pytest.mark.parametrize( + "expr", + [ + # Partial pruning across the middle two row groups (0 and 3 prune). + "id >= 1500 and id < 2500", + # Compound over multiple types, with a string equality conjunct. + 'id >= 1500 and id < 2500 and g == "g2"', + # OR of two disjoint ranges — neither end row group prunes. + "id < 500 or id >= 3500", + # Float comparison drives the pruning column. + "x >= 900.0 and x < 1100.0", + # A predicate no row group can satisfy (fully pruned → empty). + "id >= 100000000", + ], +) +def test_filter_pushdown_compound_parity(tmp_path, restore_ctx, expr): + """Native row-group pruning + Python post-filter must be byte-identical to + the PyArrow v2 reader across compound predicates spanning int/float/string + columns and multiple row groups. This is the backstop for native pruning + being the sole mechanism: any over-pruning bug would drop rows PyArrow keeps + and fail here. + """ + path = tmp_path / "compound.parquet" + n = 4000 + table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "x": pa.array(np.arange(n) * 0.5), + "g": pa.array([f"g{i % 5}" for i in range(n)]), + } + ) + pq.write_table(table, str(path), write_page_index=True, row_group_size=1000) + + def _read(use_arrow_rs): + restore_ctx.use_arrow_rs_parquet_reader = use_arrow_rs + # take_all() over sort() preserves the schema even for empty results + # (to_pandas() on an empty dataset yields a 0-column frame). + ds = ray.data.read_parquet(str(path)).filter(expr=expr).sort("id") + return ds.take_all() + + pa_rows = _read(False) + rs_rows = _read(True) + # Row-for-row parity across every column (dicts compare all keys/values). + assert rs_rows == pa_rows, ( + f"arrow-rs diverged from pyarrow for `{expr}`: " + f"{len(rs_rows)} vs {len(pa_rows)} rows" + ) + + +def _gate_verdict(path, read_columns=None): + """The support gate's verdict for a file, via a real fragment.""" + import pyarrow.dataset as pds + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + reader = ArrowRsParquetFileReader(filesystem=LocalFileSystem()) + frag = next( + pds.dataset( + str(path), format="parquet", filesystem=LocalFileSystem() + ).get_fragments() + ) + return reader._arrow_rs_supported(frag, read_columns) + + +@pytest.mark.parametrize( + "colname,builder", + [ + ("vals", lambda n: pa.array([[i, i + 1, None] for i in range(n)])), + ( + "st", + lambda n: pa.StructArray.from_arrays( + [pa.array(np.arange(n)), pa.array(np.arange(n) * 0.5)], + names=["a", "b"], + ), + ), + ( + "st_nested", + lambda n: pa.array( + [{"a": [i, i + 1], "b": {"c": f"row-{i}"}} for i in range(n)] + ), + ), + ], +) +def test_nested_column_native_parity(tmp_path, restore_ctx, colname, builder): + """List, struct, and deeper struct/list nesting decode NATIVELY (the gate + admits them) and stay byte-identical to PyArrow.""" + path = tmp_path / f"{colname}.parquet" + n = 2000 + table = pa.table( + {"id": pa.array(np.arange(n, dtype=np.int64)), colname: builder(n)} + ) + pq.write_table(table, str(path), write_page_index=True) + + assert _gate_verdict(path) is True, f"{colname} should be native now" + pa_tbl = _read_sorted(path, False, restore_ctx) + rs_tbl = _read_sorted(path, True, restore_ctx) + assert pa_tbl.equals(rs_tbl) + + +def _list_col_from_lengths(lengths, *, value_type=pa.int64(), null_rows=None, seed=0): + """Build a ``list`` array where row ``i`` holds ``lengths[i]`` + elements (deterministic but arbitrary values). ``null_rows`` is an optional + boolean mask marking rows that are NULL lists — distinct from empty lists; + those rows are forced to zero length because Parquet cannot store a null list + that spans elements.""" + lengths = np.asarray(lengths, dtype=np.int64).copy() + if null_rows is not None: + null_rows = np.asarray(null_rows, dtype=bool) + lengths[null_rows] = 0 + offsets = np.zeros(len(lengths) + 1, dtype=np.int32) + np.cumsum(lengths, out=offsets[1:]) + total = int(offsets[-1]) + rng = np.random.default_rng(seed) + if pa.types.is_string(value_type): + values = pa.array([f"e{v}" for v in rng.integers(0, 10**6, total)]) + else: + values = pa.array(rng.integers(0, 10**9, total), type=value_type) + mask = None if null_rows is None else pa.array(null_rows) + return pa.ListArray.from_arrays(pa.array(offsets), values, mask=mask) + + +@pytest.mark.parametrize( + "shape,lengths_fn", + [ + ("empty", lambda n: np.zeros(n, dtype=np.int64)), + ("singleton", lambda n: np.ones(n, dtype=np.int64)), + ("small_fixed", lambda n: np.full(n, 4, dtype=np.int64)), + ("big_fixed", lambda n: np.full(n, 512, dtype=np.int64)), + ("ascending", lambda n: np.arange(n, dtype=np.int64)), + ("descending", lambda n: np.arange(n, dtype=np.int64)[::-1]), + ("random", lambda n: np.random.default_rng(0).integers(0, 128, n)), + ], +) +def test_list_length_shapes_native_parity(tmp_path, restore_ctx, shape, lengths_fn): + """``list`` columns across every length distribution — all-empty, + singleton, small/big fixed width, monotonically ascending and descending + ramps, and random sizes — decode NATIVELY and stay byte-identical to PyArrow. + + The 2k-row fixed-length-3 parity test alone leaves the offset handling barely + exercised: it never sees wide lists, empty rows, or non-uniform offsets, which + is exactly where the crate's offset buffer and its footer bytes-per-row + estimate are most likely to slip.""" + n = 1000 + table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "vals": _list_col_from_lengths(lengths_fn(n)), + } + ) + path = tmp_path / f"list_{shape}.parquet" + pq.write_table(table, str(path), write_page_index=True) + + assert _gate_verdict(path) is True, f"list<{shape}> should decode natively" + pa_tbl = _read_sorted(path, False, restore_ctx) + rs_tbl = _read_sorted(path, True, restore_ctx) + assert pa_tbl.equals(rs_tbl) + + +def test_list_null_rows_and_string_values_parity(tmp_path, restore_ctx): + """Null lists (list-level nulls, distinct from empty lists) and a + ``list`` column with variable sizes both decode natively and match + PyArrow — covering non-int element types and the validity buffer alongside + the offsets.""" + n = 500 + lengths = np.random.default_rng(1).integers(0, 20, n) + null_rows = np.arange(n) % 7 == 0 # every 7th row is a NULL list + table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "ints": _list_col_from_lengths(lengths, null_rows=null_rows, seed=2), + "strs": _list_col_from_lengths( + lengths, value_type=pa.string(), null_rows=null_rows, seed=3 + ), + } + ) + path = tmp_path / "list_nulls.parquet" + pq.write_table(table, str(path), write_page_index=True) + + assert _gate_verdict(path) is True + pa_tbl = _read_sorted(path, False, restore_ctx) + rs_tbl = _read_sorted(path, True, restore_ctx) + assert pa_tbl.equals(rs_tbl) + + +def test_large_nested_byte_budget_batching(tmp_path, monkeypatch): + """A LARGE nested read (variable-length lists + struct-of-string, decoded + size many times the decode budget, one lone row group) must stream through + the byte-budget path as many small batches — not one giant table — while + staying byte-identical and in row order. + + The small parity tests (2k rows) fit in a single budget batch, so they never + prove the batching math holds for nested types, where the footer's + bytes-per-row estimate is coarser than for flat columns. The budget is + shrunk to 1 MiB via monkeypatch (it's a module global read at call time), so + this must run the reader in-process rather than through Ray workers. + """ + import pyarrow.dataset as pds + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers import ( + arrow_rs_parquet_file_reader as reader_mod, + ) + + n = 200_000 + rng = np.random.default_rng(0) + # Variable-length list (avg ~6 elems, some empty) + struct{int, str}. + lens = rng.integers(0, 12, n) + offsets = np.zeros(n + 1, dtype=np.int32) + np.cumsum(lens, out=offsets[1:]) + values = pa.array(rng.integers(0, 10**9, offsets[-1]), type=pa.int64()) + table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "vals": pa.ListArray.from_arrays(pa.array(offsets), values), + "st": pa.StructArray.from_arrays( + [ + pa.array(np.arange(n, dtype=np.int64)), + pa.array([f"payload-string-{i:012d}" for i in range(n)]), + ], + names=["a", "b"], + ), + } + ) + path = tmp_path / "large_nested.parquet" + # One row group covering all rows: the whole decode must be paced by the + # byte budget, with no row-group boundaries helping out. + pq.write_table(table, str(path), write_page_index=True, row_group_size=n) + assert pq.ParquetFile(str(path)).num_row_groups == 1 + + budget = 1024 * 1024 + monkeypatch.setattr(reader_mod, "_ARROW_RS_DECODE_BUDGET_BYTES", budget) + assert table.nbytes > 10 * budget # "large": decoded size >> budget + + calls = {"n": 0} + orig = ray_data_arrow_rs.read_row_groups + + def wrapped(*a, **k): + calls["n"] += 1 + return orig(*a, **k) + + monkeypatch.setattr(ray_data_arrow_rs, "read_row_groups", wrapped) + reader = reader_mod.ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), target_block_size=128 * 1024 * 1024 + ) + dataset = pds.dataset(str(path), format="parquet", filesystem=LocalFileSystem()) + fragment = next(dataset.get_fragments()) + scanner_kwargs = { + "columns": None, + "filter": None, + "batch_size": reader._resolve_batch_size(dataset, _whole_file_manifest()), + } + batches = list(reader._iter_fragment_tables(fragment, scanner_kwargs)) + + assert calls["n"] > 0, "native read_row_groups was not called (fell back)" + # Streaming, not slurping: many batches, each near the budget (generous 8x + # slack because the crate sizes rows from the footer's bytes-per-row, which + # is approximate for variable-width nested data). + assert len(batches) >= 5, f"expected many budget batches, got {len(batches)}" + assert max(b.nbytes for b in batches) <= 8 * budget + got = pa.concat_tables(batches) + # In-order and byte-identical (no sort: order is part of the contract). + assert got.equals(table) + + +def test_fat_row_decode_budget_not_voided_by_batch_floor(tmp_path): + """A fat-row group (~64 KiB/row) under a small decode budget must stream + many small batches. The crate's old 2048-row batch floor overrode + ``decode_budget_bytes`` for any schema above ~16 KiB/row (findings K8) — + here it would have decoded the whole 32 MiB group as a single batch. + Direct crate call: the floor lives in the crate's batch sizing, below the + Ray layer.""" + n = 512 + table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "fat": pa.array(["x" * 65536] * n), + } + ) + path = tmp_path / "fat_rows.parquet" + # Dictionary encoding would shrink the footer's uncompressed size (what the + # budget math reads) to nothing; plain encoding keeps ~64 KiB/row. + pq.write_table( + table, str(path), write_page_index=True, row_group_size=n, use_dictionary=False + ) + + reader = pa.RecordBatchReader.from_stream( + ray_data_arrow_rs.read_row_groups(str(path), decode_budget_bytes=1024 * 1024) + ) + batches = list(reader) + # 1 MiB budget / 64 KiB rows -> 16 rows, clamped up to the 32-row floor: + # 16 batches. The old floor produced ONE 512-row (32 MiB) batch. + assert len(batches) >= 8, f"expected many budget-sized batches, got {len(batches)}" + assert max(b.num_rows for b in batches) <= 64 + got = pa.Table.from_batches(batches, schema=table.schema) + assert got.equals(table) + + +def test_scanner_leak_signature_and_arrow_rs_avoids_it(tmp_path, monkeypatch): + """Reproduce the PyArrow accumulation behind ray#49158 / apache/arrow#39808 and + show the arrow-rs reader sidesteps it. + + The reported "leak" is ``pyarrow.dataset`` ``to_batches`` (the Scanner) + accumulating ~the whole file in the Arrow allocator (the issue also reported + this as independent of ``batch_size``; that half now varies by pyarrow + version/platform and is recorded rather than asserted — see (1)), whereas + ``pq.ParquetFile.iter_batches`` (what the V2 reader uses) holds only ~a row + group. We measure exactly as the issue did: ``pa.total_allocated_bytes()`` + tracked as a running max across the batch iteration. + + The arrow-rs reader decodes in Rust and hands batches across a zero-copy FFI + boundary, so decoded data never enters the Arrow allocator at all — its + Arrow-side peak is ~0, categorically below the row-group floor iter_batches + pays. (This asserts Arrow-allocator accumulation only, which is deterministic; + the Rust working set is invisible here and is validated by RSS/USS in the + bench suite. It is a regression guard: the pyarrow tiers must stay ordered and + arrow-rs must not start buffering decoded data on the Arrow side.) + """ + import pyarrow.dataset as pds + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers import ( + arrow_rs_parquet_file_reader as reader_mod, + ) + + # 800k rows of fat strings split into many small row groups, so "whole file" + # (to_batches) and "one row group" (iter_batches) are clearly different scales. + n = 800_000 + rng = np.random.default_rng(0) + cols = {"id": pa.array(np.arange(n, dtype=np.int64))} + for i in range(4): + cols[f"s{i}"] = pa.array( + [f"val-{rng.integers(0, 10**6)}-{'x' * 8}" for _ in range(n)] + ) + table = pa.table(cols) + file_mb = table.nbytes / 1024 / 1024 + path = tmp_path / "leaky.parquet" + pq.write_table( + table, + str(path), + row_group_size=50_000, + write_page_index=True, + compression="snappy", + ) + n_rg = pq.ParquetFile(str(path)).num_row_groups + assert n_rg >= 8 # need many groups for file-scale vs row-group-scale to differ + del table, cols + import gc + + gc.collect() + + def _peak_mb(make_iter): + """Max Arrow-allocator bytes held at once across the iteration (MB) — the + issue's own metric. Baseline-subtracted so unrelated allocations don't + count; each batch is dropped so only what the reader *retains* is seen.""" + base = pa.total_allocated_bytes() + mx = 0 + for batch in make_iter(): + mx = max(mx, pa.total_allocated_bytes() - base) + del batch + return mx / 1024 / 1024 + + def _fragment(): + return next(pds.dataset(str(path), format="parquet").get_fragments()) + + # (1) Scanner/to_batches accumulates ~the whole file — the leak signature, and + # the load-bearing half of this test. + to_small = _peak_mb(lambda: _fragment().to_batches(batch_size=256)) + to_big = _peak_mb(lambda: _fragment().to_batches(batch_size=2048)) + assert to_small > 0.4 * file_mb, (to_small, file_mb) + # The issue's other half — "and it does NOT shrink with batch_size" — is + # upstream behaviour that has since diverged by platform/version, so it is + # RECORDED, not asserted (same treatment as (2) below, for the same reason). + # macOS pyarrow 21: 44.5 vs 44.4 MB, insensitive as the issue described. + # Linux, 2026-08-13: 44.5 vs 4.75 MB — batch_size=2048 no longer accumulates. + # This bears on finding C1, so it must stay visible rather than be tuned away: + # if PyArrow's scanner really has become batch_size-bounded, the honest + # version of C1 is narrower than "batch_size does not bound it". Note the + # metric here is the Arrow *allocator* on one fragment in-process, which is + # not what M28 measures (per-read-task USS, where PyArrow still retained + # 1.5x the decoded bin on the same Linux box the same week). + print( + f"[C1 datum] pyarrow {pa.__version__}: to_batches peak MB " + f"batch_size=256 -> {to_small:.1f}, batch_size=2048 -> {to_big:.1f}" + ) + + # (2) ParquetFile.iter_batches (the ARROW-5030 fallback path) HISTORICALLY held + # only ~a row group — measured ~5x below to_batches on the pyarrow this was + # written against (macOS venv, 2026-07). Newer pyarrow buffers more + # aggressively and can hold ~the whole file here too (first seen on Linux, + # 2026-08-11: it_small 43.8 vs to_small 44.5), so this is upstream behaviour + # we record but no longer assert. What we DO still require is sanity: the + # fallback path is not materially WORSE than the scanner. + it_small = _peak_mb(lambda: pq.ParquetFile(str(path)).iter_batches(batch_size=256)) + assert it_small < 1.2 * to_small, (it_small, to_small) + + # (3) arrow-rs: decoded data is Rust-owned and crosses via zero-copy FFI, so ~0 + # lands in the Arrow allocator — below whichever floor THIS pyarrow + # version exhibits (row-group or whole-file). + monkeypatch.setattr(reader_mod, "_ARROW_RS_DECODE_BUDGET_BYTES", 1024 * 1024) + reader = reader_mod.ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), target_block_size=128 * 1024 * 1024 + ) + dataset = pds.dataset(str(path), format="parquet", filesystem=LocalFileSystem()) + frag = next(dataset.get_fragments()) + scanner_kwargs = { + "columns": None, + "filter": None, + "batch_size": reader._resolve_batch_size(dataset, _whole_file_manifest()), + } + rs_peak = _peak_mb(lambda: reader._iter_fragment_tables(frag, scanner_kwargs)) + assert rs_peak < 0.5 * min(it_small, to_small), (rs_peak, it_small, to_small) + + +def test_map_column_native_parity(tmp_path, restore_ctx): + """Map columns decode NATIVELY (the gate admits them) and byte-identically to + PyArrow — including null and empty map entries spread across multiple row + groups. Verified empirically: the crate emits an identical ``MapType`` (same + key/value field names) and the same values, so ``pa.Table.equals`` holds.""" + n = 500 + map_path = tmp_path / "map.parquet" + map_table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "m": pa.array( + [ + None + if i % 50 == 0 + else ([] if i % 9 == 0 else [(f"k{i}", i), ("b", i * 2)]) + for i in range(n) + ], + type=pa.map_(pa.string(), pa.int64()), + ), + } + ) + pq.write_table(map_table, str(map_path), write_page_index=True, row_group_size=100) + assert _gate_verdict(map_path) is True + pa_tbl = _read_sorted(map_path, False, restore_ctx) + rs_tbl = _read_sorted(map_path, True, restore_ctx) + assert pa_tbl.equals(rs_tbl) + + +def test_dictionary_column_native_parity(tmp_path, restore_ctx): + """A *naturally* dictionary-typed column (the file embeds an arrow dictionary + type) decodes natively and byte-identically to PyArrow, even with per-row-group + dictionaries and nulls — Parquet dictionaries are per-row-group, the classic + index-divergence trap, so this spans several row groups on purpose. + + Distinct from the ``dictionary_columns`` *forced* read (columns coerced to + dictionary output at read time), which the planned path handles via an + alignment cast (see ``test_forced_dictionary_columns_native_parity``).""" + n = 500 + path = tmp_path / "dict.parquet" + vals = [None if i % 37 == 0 else f"cat{(i * 7) % 13}" for i in range(n)] + table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "d": pa.array(vals, type=pa.dictionary(pa.int32(), pa.string())), + } + ) + pq.write_table(table, str(path), write_page_index=True, row_group_size=100) + assert _gate_verdict(path) is True + pa_tbl = _read_sorted(path, False, restore_ctx) + rs_tbl = _read_sorted(path, True, restore_ctx) + assert pa_tbl.equals(rs_tbl) + + +def test_extension_types_native_parity(tmp_path, restore_ctx): + """Extension types decode NATIVELY and byte-identically to PyArrow: Ray's + tensor extension, its variable-shaped (ragged) tensor, and pyarrow's canonical + ``fixed_shape_tensor``. + + The crate carries the embedded arrow-schema field metadata + (``ARROW:extension:name`` / ``:metadata``) straight through the C data + interface, so pyarrow reconstructs the *registered* extension identically on + the native and PyArrow paths. (This is the case a blanket ``extension_name`` + rejection used to fall back — empirically it round-trips, so the gate now + recurses into the storage type instead.)""" + from ray.data.extensions import ArrowTensorArray, ArrowVariableShapedTensorArray + + n = 400 + + # Ray fixed-shape tensor, spanning several row groups. + tpath = tmp_path / "tensor.parquet" + tens = ArrowTensorArray.from_numpy( + np.arange(4 * n, dtype=np.float32).reshape(n, 2, 2) + ) + pq.write_table( + pa.table({"id": pa.array(np.arange(n, dtype=np.int64)), "t": tens}), + str(tpath), + write_page_index=True, + row_group_size=100, + ) + assert _gate_verdict(tpath) is True + assert _read_arrow_sorted(tpath, False, restore_ctx).equals( + _read_arrow_sorted(tpath, True, restore_ctx) + ) + + # Ray variable-shaped (ragged) tensor — storage is a struct(data, shape). + vpath = tmp_path / "vtensor.parquet" + ragged = np.array( + [np.arange(i % 5 + 1, dtype=np.int64) for i in range(n)], dtype=object + ) + vst = ArrowVariableShapedTensorArray.from_numpy(ragged) + pq.write_table( + pa.table({"id": pa.array(np.arange(n, dtype=np.int64)), "t": vst}), + str(vpath), + write_page_index=True, + row_group_size=100, + ) + assert _gate_verdict(vpath) is True + assert _read_arrow_sorted(vpath, False, restore_ctx).equals( + _read_arrow_sorted(vpath, True, restore_ctx) + ) + + # pyarrow canonical fixed_shape_tensor (not isinstance(ExtensionType) on some + # versions — caught via extension_name). + if hasattr(pa, "fixed_shape_tensor"): + cpath = tmp_path / "canonical.parquet" + flat = pa.array(np.arange(n * 4, dtype=np.float32), type=pa.float32()) + storage = pa.FixedSizeListArray.from_arrays(flat, 4) + tarr = pa.ExtensionArray.from_storage( + pa.fixed_shape_tensor(pa.float32(), [4]), storage + ) + pq.write_table( + pa.table({"id": pa.array(np.arange(n, dtype=np.int64)), "t": tarr}), + str(cpath), + write_page_index=True, + row_group_size=100, + ) + assert _gate_verdict(cpath) is True + assert _read_arrow_sorted(cpath, False, restore_ctx).equals( + _read_arrow_sorted(cpath, True, restore_ctx) + ) + + # There is no per-type support gate to unit-check anymore: every + # Parquet-encodable type is admitted (proven by the parity reads above), and + # Arrow's in-memory-only types (union, list_view, run_end_encoded, ...) have + # no Parquet encoding — PyArrow refuses to write them ("Unhandled type for + # Arrow to Parquet schema conversion") — so they can never appear in a + # footer-derived schema. The old ``_arrow_rs_type_supported`` gate on them + # was unreachable dead code, removed 2026-07-28. + union_array = pa.UnionArray.from_dense( + pa.array([0], pa.int8()), + pa.array([0], pa.int32()), + [pa.array([1], pa.int64())], + ) + with pytest.raises(pa.lib.ArrowNotImplementedError, match="Unhandled type"): + pq.write_table(pa.table({"u": union_array}), str(tmp_path / "union.parquet")) + + +def test_cloudpickle_tensor_metadata_native_parity(tmp_path, monkeypatch): + """A Ray tensor column serialized with the legacy *cloudpickle* format (files + written by Ray 2.49-2.54) stores non-UTF8 bytes in the parquet ``ARROW:schema`` + field metadata, which arrow-rs's IPC verifier rejects — the crate used to + crash (``Unable to get root as message stored in ARROW:schema: Utf8Error``). + + The crate now retries the footer load with the embedded arrow schema skipped + (decoding the parquet *storage* type), and the reader reconstructs the + extension from the file's own pyarrow-read footer schema. So the file decodes + NATIVELY and byte-identically to PyArrow instead of crashing or silently + falling back. Regression test for the release ``wide_schema_pipeline_tensors`` + failure (Bug 2).""" + import ray.data._internal.tensor_extensions.arrow as tx + from ray.data.extensions import ArrowTensorArray + + # Write with the legacy cloudpickle serialization → binary metadata value. + monkeypatch.setattr( + tx, + "ARROW_EXTENSION_SERIALIZATION_FORMAT", + tx._SerializationFormat.CLOUDPICKLE, + ) + n = 200 + tens = ArrowTensorArray.from_numpy( + np.arange(4 * n, dtype=np.float32).reshape(n, 2, 2) + ) + path = tmp_path / "cp_tensor.parquet" + pq.write_table( + pa.table({"id": pa.array(np.arange(n, dtype=np.int64)), "t": tens}), + str(path), + write_page_index=True, + row_group_size=50, + ) + # Reading the cloudpickle-serialized extension back requires the opt-in + # autoload (set before any pyarrow read below, which reconstructs the type). + monkeypatch.setattr(tx, "_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA", True) + + # Confirm the footer really carries the ``ARROW:schema`` metadata whose + # (cloudpickle) value isn't valid UTF-8 — the case the crate used to crash on. + kv = pq.read_metadata(str(path)).metadata or {} + assert b"ARROW:schema" in kv + + rs, pa_tbl, native_decodes = _read_both_in_process([path], monkeypatch) + # The crate actually decoded it (didn't fall back to PyArrow). + assert native_decodes > 0 + assert rs.equals(pa_tbl) + assert isinstance(rs.schema.field("t").type, pa.ExtensionType) + + +def test_cloudpickle_tensor_ffi_relabel_fast_path(tmp_path, monkeypatch): + """M53: for a skipped-embedded-schema tensor file the reader hands the crate + the extension's *storage* schema (``with_schema_override``, so the decode + emits large_list offsets directly) and re-types each batch with a zero-copy + C-Data-Interface relabel instead of a per-batch ``Table.cast`` (~2-4 ms vs + ~40 ms at 5000 columns). Asserts the fast path actually engages, that the + ``RAY_DATA_ARROW_RS_FFI_RELABEL=0`` kill switch restores the cast path, and + that both are byte-identical to PyArrow.""" + import ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader as m + import ray.data._internal.tensor_extensions.arrow as tx + from ray.data.extensions import ArrowTensorArray + + monkeypatch.setattr( + tx, + "ARROW_EXTENSION_SERIALIZATION_FORMAT", + tx._SerializationFormat.CLOUDPICKLE, + ) + n = 200 + tens = ArrowTensorArray.from_numpy( + np.arange(4 * n, dtype=np.float32).reshape(n, 2, 2) + ) + path = tmp_path / "cp_tensor.parquet" + pq.write_table( + pa.table({"id": pa.array(np.arange(n, dtype=np.int64)), "t": tens}), + str(path), + write_page_index=True, + row_group_size=50, + ) + monkeypatch.setattr(tx, "_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA", True) + + relabels = {"n": 0} + orig_relabel = m._ffi_relabel_batch + + def counting_relabel(batch, schema): + relabels["n"] += 1 + return orig_relabel(batch, schema) + + monkeypatch.setattr(m, "_ffi_relabel_batch", counting_relabel) + + rs, pa_tbl, native_decodes = _read_both_in_process([path], monkeypatch) + assert native_decodes > 0 + assert relabels["n"] > 0, "FFI relabel fast path did not engage" + assert rs.equals(pa_tbl) + assert isinstance(rs.schema.field("t").type, pa.ExtensionType) + + # Kill switch: the cast path still runs and produces the same bytes. + relabels["n"] = 0 + monkeypatch.setenv("RAY_DATA_ARROW_RS_FFI_RELABEL", "0") + rs2, pa_tbl2, _ = _read_both_in_process([path], monkeypatch) + assert relabels["n"] == 0 + assert rs2.equals(pa_tbl2) + assert rs2.equals(rs) + + +def test_arrow_rs_supported_gate(tmp_path): + """Unit-check the fallback gate: local flat AND struct/list = supported; + empty projection / unknown-not-on-disk column (no unified schema to + null-fill from) / non-local filesystem = unsupported.""" + import pyarrow.dataset as pds + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + flat = tmp_path / "flat.parquet" + pq.write_table(_flat_table(1000), str(flat), write_page_index=True) + nested = tmp_path / "nested.parquet" + pq.write_table( + pa.table({"id": pa.array([1, 2]), "v": pa.array([[1], [2]])}), str(nested) + ) + + reader = ArrowRsParquetFileReader(filesystem=LocalFileSystem()) + + flat_frag = next( + pds.dataset( + str(flat), format="parquet", filesystem=LocalFileSystem() + ).get_fragments() + ) + nested_frag = next( + pds.dataset( + str(nested), format="parquet", filesystem=LocalFileSystem() + ).get_fragments() + ) + + assert reader._arrow_rs_supported(flat_frag, None) is True + assert reader._arrow_rs_supported(flat_frag, ["id", "x"]) is True + # Empty projection → fall back. + assert reader._arrow_rs_supported(flat_frag, []) is False + # List column → native (ungated 2026-07-21). + assert reader._arrow_rs_supported(nested_frag, None) is True + # A requested column that isn't on disk (here a dotted name — in real reads + # ``_split_columns`` filters these out before the gate) is treated as a + # missing column; with no unified schema to null-fill from → fall back. + assert reader._arrow_rs_supported(nested_frag, ["v.item"]) is False + + # ``filesystem=None`` is the default local filesystem → supported + # (see ``_filesystem_supported``). A genuinely foreign filesystem (neither + # local nor S3, e.g. a ``SubTreeFileSystem``) → fall back. + from pyarrow.fs import SubTreeFileSystem + + reader_default_fs = ArrowRsParquetFileReader(filesystem=None) + assert reader_default_fs._arrow_rs_supported(flat_frag, None) is True + + reader_foreign_fs = ArrowRsParquetFileReader( + filesystem=SubTreeFileSystem(str(tmp_path), LocalFileSystem()) + ) + assert reader_foreign_fs._arrow_rs_supported(flat_frag, None) is False + + +# --------------------------------------------------------------------------- +# read_metadata FFI (Track 1): arrow-rs owns the footer read. The crate returns +# the Arrow schema (via the C-schema PyCapsule) plus per-row-group counts, so the +# Python reader no longer has to build a PyArrow dataset to learn the layout. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("row_group_size", [20_000, 5_000]) +def test_read_metadata_matches_pyarrow(tmp_path, row_group_size): + import ray_data_arrow_rs + + table = _flat_table(20_000) + path = tmp_path / "meta.parquet" + pq.write_table( + table, str(path), row_group_size=row_group_size, write_page_index=True + ) + + pf = pq.ParquetFile(str(path)) + md = ray_data_arrow_rs.read_metadata(str(path)) + + # Schema round-trips exactly through __arrow_c_schema__ (types + names). + assert pa.schema(md).equals(pf.schema_arrow) + assert md.num_rows == pf.metadata.num_rows + assert md.num_row_groups == pf.metadata.num_row_groups + assert md.row_group_num_rows == [ + pf.metadata.row_group(i).num_rows for i in range(pf.metadata.num_row_groups) + ] + assert len(md.row_group_byte_sizes) == pf.metadata.num_row_groups + assert all(b > 0 for b in md.row_group_byte_sizes) + + +# --------------------------------------------------------------------------- +# Native predicate pushdown, part 1: statistics-based row-group pruning. +# +# These cover the seam the Rust unit tests (which use synthetic ColStats) can't: +# lowering a Ray Expr to the IR (`_predicate_to_ir`) and reading *real* Parquet +# statistics into that pruning (`select_row_groups` / `read_row_groups`'s +# `predicate_json`). Pruning is row-group granular — a surviving group is decoded +# whole; row-level filtering is the reader's post-decode job, exercised +# separately by the end-to-end filter parity tests above. +# --------------------------------------------------------------------------- + + +def _sorted_rg_table(n=4000): + """A sorted int ``id`` (+ float ``x``, string ``g``) so that with + ``row_group_size=1000`` each of the 4 row groups holds a disjoint id range, + making stats pruning observable.""" + return pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "x": pa.array(np.arange(n) * 0.5), + "g": pa.array([f"k{i // 1000}" for i in range(n)]), # k0..k3 per group + } + ) + + +def _ir(expr): + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + _predicate_to_ir, + ) + + return _predicate_to_ir(expr) + + +def test_predicate_to_ir_lowering(): + """The Ray Expr -> IR lowering produces the expected shapes, flips the op + when the column is on the right, and degrades unrepresentable nodes to + ``unknown`` (keeping the rest of a conjunction prunable).""" + from ray.data.expressions import col, lit + + assert _ir(col("id") >= 3000) == { + "t": "cmp", + "col": "id", + "op": "ge", + "value": {"vt": "int", "v": 3000}, + } + # literal-on-left flips ge -> le + assert _ir(lit(3000) <= col("id")) == { + "t": "cmp", + "col": "id", + "op": "ge", + "value": {"vt": "int", "v": 3000}, + } + # conjunction; float + string literals tagged by type + assert _ir((col("x") < 1.5) & (col("g") == "k2")) == { + "t": "and", + "preds": [ + {"t": "cmp", "col": "x", "op": "lt", "value": {"vt": "float", "v": 1.5}}, + {"t": "cmp", "col": "g", "op": "eq", "value": {"vt": "str", "v": "k2"}}, + ], + } + assert _ir(col("id").is_null()) == {"t": "is_null", "col": "id"} + assert _ir(col("id").is_in([1, 2, 3])) == { + "t": "in", + "col": "id", + "values": [ + {"vt": "int", "v": 1}, + {"vt": "int", "v": 2}, + {"vt": "int", "v": 3}, + ], + "negated": False, + } + # a UDF conjunct is unrepresentable -> unknown, but the other conjunct stays. + part = _ir((col("id") >= 3000) & (col("x").abs() > 1.0)) + assert part["t"] == "and" + assert part["preds"][0] == { + "t": "cmp", + "col": "id", + "op": "ge", + "value": {"vt": "int", "v": 3000}, + } + assert part["preds"][1] == {"t": "unknown"} + + +def test_predicate_json_skips_when_nothing_prunable(): + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + _predicate_json, + ) + from ray.data.expressions import col + + assert _predicate_json(None) is None + # A bare UDF predicate lowers entirely to unknown -> no pushdown arg. + assert _predicate_json(col("x").abs() > 1.0) is None + assert _predicate_json(col("id") >= 3000) is not None + + +def test_select_row_groups_pruning_on_real_stats(tmp_path): + """``select_row_groups`` prunes using the file's actual footer statistics, + across int / float / string columns, and never over-prunes.""" + import json + + path = tmp_path / "sorted.parquet" + table = _sorted_rg_table(4000) + pq.write_table(table, str(path), write_page_index=True, row_group_size=1000) + + def sel(expr): + return ray_data_arrow_rs.select_row_groups(str(path), json.dumps(_ir(expr))) + + from ray.data.expressions import col + + # No predicate -> all four groups. + assert ray_data_arrow_rs.select_row_groups(str(path), None) == [0, 1, 2, 3] + # id in [0,999],[1000,1999],[2000,2999],[3000,3999] per group. + assert sel(col("id") >= 3000) == [3] + assert sel(col("id") >= 3500) == [3] # group-granular: whole group 3 survives + assert sel(col("id") < 1000) == [0] + assert sel(col("id") >= 10**9) == [] # nothing can match + assert sel((col("id") >= 1500) & (col("id") < 2500)) == [1, 2] + # float column (x = id*0.5, so group 3 is [1500, 1999.5]) + assert sel(col("x") >= 1800.0) == [3] + # string column pruning (g = "k0".."k3") + assert sel(col("g") == "k2") == [2] + assert sel(col("g") == "zzz") == [] + # is_in over the string groups + assert sel(col("g").is_in(["k0", "k3"])) == [0, 3] + + +def test_read_row_groups_predicate_json_decodes_only_surviving_groups(tmp_path): + """End to end through the crate: ``predicate_json`` prunes row groups before + decode, so the stream contains exactly the surviving groups' rows (whole + groups — row-level filtering is applied by the reader, not here) and stays + byte-correct and in order.""" + import json + + path = tmp_path / "sorted.parquet" + table = _sorted_rg_table(4000) + pq.write_table(table, str(path), write_page_index=True, row_group_size=1000) + + from ray.data.expressions import col + + # id >= 3000 keeps only group 3 (rows 3000..3999). + got = _read_crate_stream(path, predicate_json=json.dumps(_ir(col("id") >= 3000))) + assert got.equals(table.slice(3000, 1000)) + + # id >= 3500 keeps group 3 whole (pruning is row-group granular). + got = _read_crate_stream(path, predicate_json=json.dumps(_ir(col("id") >= 3500))) + assert got.equals(table.slice(3000, 1000)) + + # A fully-pruning predicate yields an empty (schema-correct) stream. + got = _read_crate_stream(path, predicate_json=json.dumps(_ir(col("id") >= 10**9))) + assert got.num_rows == 0 + assert got.column_names == ["id", "x", "g"] + + +# --------------------------------------------------------------------------- +# S3 (moto server). The crate reads S3 through the Rust `object_store` client, +# so it needs a real HTTP endpoint — Ray Data's `s3_server` fixture (a moto +# server) provides one. These prove the native path (a) connects with the same +# endpoint/credentials/region PyArrow used (recovered from the S3FileSystem via +# `_s3_config`), and (b) returns byte-identical data. +# --------------------------------------------------------------------------- + + +def _s3_write(table, s3_fs, s3_path, name="data.parquet"): + """Write ``table`` as one file into the moto S3 dir, return the s3:// URI.""" + base = _unwrap_protocol(s3_path) # strip s3:// → bucket/key + key = os.path.join(base, name) + pq.write_table(table, key, filesystem=s3_fs, write_page_index=True) + return f"s3://{key}" + + +def _read_s3_sorted(uri, s3_fs, use_arrow_rs, restore_ctx, **read_kwargs): + restore_ctx.use_arrow_rs_parquet_reader = use_arrow_rs + ds = ray.data.read_parquet(uri, filesystem=s3_fs, **read_kwargs) + return pa.Table.from_pandas(ds.to_pandas()).sort_by("id") + + +def test_s3_config_recovers_endpoint_and_creds(s3_fs): + """`_s3_config` must recover the moto endpoint, static creds, region, and + (critically) allow_http from an http:// endpoint whose `scheme` field still + reads 'https' — otherwise the crate can't reach moto/MinIO.""" + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + _s3_config, + ) + + cfg = _s3_config(s3_fs) + assert cfg["endpoint"] and cfg["endpoint"].startswith("http://") + assert cfg["allow_http"] is True + assert cfg["region"] == "us-west-2" + assert cfg["access_key_id"] == "testing" + assert cfg["secret_access_key"] == "testing" + assert cfg["anonymous"] is False + + +def test_read_metadata_s3_matches_pyarrow(s3_fs, s3_path): + """`read_metadata_s3` fetches the footer over the moto endpoint (same config + recovery as the data path) and returns the same schema + row-group counts as + PyArrow reading the same object.""" + import ray_data_arrow_rs + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + _s3_config, + ) + + table = _flat_table(20_000) + uri = _s3_write(table, s3_fs, s3_path, name="meta.parquet") + bucket, _, key = _unwrap_protocol(uri).partition("/") + + pf = pq.ParquetFile(_unwrap_protocol(uri), filesystem=s3_fs) + + cfg = _s3_config(s3_fs) + md = ray_data_arrow_rs.read_metadata_s3( + bucket, + key, + cfg["region"], + cfg["anonymous"], + endpoint=cfg["endpoint"], + access_key_id=cfg["access_key_id"], + secret_access_key=cfg["secret_access_key"], + session_token=cfg["session_token"], + allow_http=cfg["allow_http"], + virtual_hosted_style=cfg["virtual_hosted_style"], + ) + + assert pa.schema(md).equals(pf.schema_arrow) + assert md.num_rows == pf.metadata.num_rows + assert md.num_row_groups == pf.metadata.num_row_groups + assert md.row_group_num_rows == [ + pf.metadata.row_group(i).num_rows for i in range(pf.metadata.num_row_groups) + ] + + +def test_arrow_rs_s3_parity(s3_fs, s3_path, restore_ctx): + """Full-scan parity reading from (moto) S3 via the native path vs PyArrow.""" + table = _flat_table() + uri = _s3_write(table, s3_fs, s3_path) + + pa_tbl = _read_s3_sorted(uri, s3_fs, False, restore_ctx) + rs_tbl = _read_s3_sorted(uri, s3_fs, True, restore_ctx) + + assert rs_tbl.num_rows == table.num_rows + assert pa_tbl.equals(rs_tbl) + + +def test_arrow_rs_s3_parity_with_projection(s3_fs, s3_path, restore_ctx): + table = _flat_table() + uri = _s3_write(table, s3_fs, s3_path) + + pa_tbl = _read_s3_sorted(uri, s3_fs, False, restore_ctx, columns=["id", "x"]) + rs_tbl = _read_s3_sorted(uri, s3_fs, True, restore_ctx, columns=["id", "x"]) + assert rs_tbl.column_names == ["id", "x"] + assert pa_tbl.equals(rs_tbl) + + +def test_arrow_rs_s3_fat_columns_window_parity(s3_fs, s3_path, restore_ctx): + """The M20 shape: a TALL row group whose every projected column exceeds the + column-group budget. The old planner mis-selected the column-group (Hstack) + path here and retained the whole decoded row group (per-task USS pinned, + ``fetch_window_mb``-inert); it now row-windows the group instead — the + selection itself is asserted by the crate's + ``tall_fat_columns_row_window_instead_of_hstack`` unit test. This proves the + multi-window S3 decode is byte-identical on that shape under the same knobs + (small window + small column budget, so the read really splits).""" + n = 2000 + table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "fat": pa.array([f"{i:04d}" + "x" * 4092 for i in range(n)]), + } + ) + base = _unwrap_protocol(s3_path) + key = os.path.join(base, "fat_cols.parquet") + # Plain encoding keeps the column fat on disk (the window math is + # compressed-byte-denominated); small pages let windows actually split. + pq.write_table( + table, + key, + filesystem=s3_fs, + write_page_index=True, + use_dictionary=False, + data_page_size=64 * 1024, + row_group_size=n, + ) + uri = f"s3://{key}" + + knobs = {"arrow_rs_fetch_window_mb": 1, "arrow_rs_column_fetch_mb": 1} + pa_tbl = _read_s3_sorted(uri, s3_fs, False, restore_ctx, dataset_kwargs=knobs) + rs_tbl = _read_s3_sorted(uri, s3_fs, True, restore_ctx, dataset_kwargs=knobs) + assert rs_tbl.num_rows == n + assert pa_tbl.equals(rs_tbl) + + # Escape hatch: windowing explicitly disabled -> windows are inert, so the + # column-group (Hstack) axis fires instead. Its decode side must stay + # byte-identical too (it still serves genuinely wide/short groups). + hstack_knobs = {"arrow_rs_fetch_window_mb": 0, "arrow_rs_column_fetch_mb": 1} + rs_hstack = _read_s3_sorted( + uri, s3_fs, True, restore_ctx, dataset_kwargs=hstack_knobs + ) + assert pa_tbl.equals(rs_hstack) + + +def test_arrow_rs_s3_native_path_runs(s3_fs, s3_path): + """Confirm an S3 fragment actually goes through the native crate's S3 entry + point (``read_row_groups_s3``), not a silent PyArrow fallback. Exercised via + the reader directly (in-process) so the monkeypatch can observe the call — + a ``ray.data.read_parquet`` read would run in a worker the driver can't patch. + """ + import pyarrow.dataset as pds + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + table = _flat_table() + _s3_write(table, s3_fs, s3_path) + base = _unwrap_protocol(s3_path) + + calls = {"n": 0} + orig = ray_data_arrow_rs.read_row_groups_s3 + + def wrapped(*a, **k): + calls["n"] += 1 + return orig(*a, **k) + + ray_data_arrow_rs.read_row_groups_s3 = wrapped + try: + reader = ArrowRsParquetFileReader( + filesystem=s3_fs, target_block_size=128 * 1024 * 1024 + ) + dataset = pds.dataset(base, format="parquet", filesystem=s3_fs) + fragment = next(dataset.get_fragments()) + scanner_kwargs = { + "columns": None, + "filter": None, + "batch_size": reader._resolve_batch_size(dataset, _whole_file_manifest()), + } + got = pa.concat_tables( + list(reader._iter_fragment_tables(fragment, scanner_kwargs)) + ) + finally: + ray_data_arrow_rs.read_row_groups_s3 = orig + + assert calls["n"] > 0, "native read_row_groups_s3 was not called (fell back)" + assert got.sort_by("id").equals(table.sort_by("id")) + + +def test_arrow_rs_s3_planned_read_shares_one_client(s3_fs, s3_path, monkeypatch): + """A planned S3 read (TODO 1r) builds ONE ``object_store`` client for the + bucket via ``connect_s3`` and opens one handle per file (one footer + + page-index fetch); the per-call entry points — which rebuilt the HTTP + client and re-fetched the footer on *every* call (findings T10) — must + stay cold. The process-level client cache (findings M97) is cleared first + so the single ``connect_s3`` build is observable, and cleared after so no + spy-wrapped store leaks to later tests. In-process so the spies can + observe the calls.""" + from ray.data._internal.datasource_v2 import native_metadata + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + native_metadata._S3_STORE_CACHE.clear() + + table = _flat_table() + half = table.num_rows // 2 + uri_a = _s3_write(table.slice(0, half), s3_fs, s3_path, name="a.parquet") + uri_b = _s3_write(table.slice(half), s3_fs, s3_path, name="b.parquet") + + calls = { + "connect": 0, + "open_file": 0, + "read_row_groups_s3": 0, + "read_metadata_s3": 0, + } + + class _StoreProxy: + """Counts per-file handle opens; forwards everything else — pyo3 + methods can't be monkeypatched directly.""" + + def __init__(self, store): + self._store = store + + def open_file(self, *a, **k): + calls["open_file"] += 1 + return self._store.open_file(*a, **k) + + def __getattr__(self, name): + return getattr(self._store, name) + + orig_connect = ray_data_arrow_rs.connect_s3 + + def spy_connect(*a, **k): + calls["connect"] += 1 + return _StoreProxy(orig_connect(*a, **k)) + + def spy_entry(name): + orig = getattr(ray_data_arrow_rs, name) + + def spy(*a, **k): + calls[name] += 1 + return orig(*a, **k) + + return spy + + monkeypatch.setattr(ray_data_arrow_rs, "connect_s3", spy_connect) + monkeypatch.setattr( + ray_data_arrow_rs, "read_row_groups_s3", spy_entry("read_row_groups_s3") + ) + monkeypatch.setattr( + ray_data_arrow_rs, "read_metadata_s3", spy_entry("read_metadata_s3") + ) + + paths = [_unwrap_protocol(uri_a), _unwrap_protocol(uri_b)] + sizes = [s3_fs.get_file_info(p).size for p in paths] + manifest = _make_manifest(paths, sizes, [None, None]) + + reader = ArrowRsParquetFileReader( + filesystem=s3_fs, target_block_size=128 * 1024 * 1024 + ) + try: + got = pa.concat_tables(list(reader.read(manifest))) + + assert calls["connect"] == 1, ( + f"expected ONE S3 client for a 2-file same-bucket read, got " + f"{calls['connect']}" + ) + assert ( + calls["open_file"] == 2 + ), f"expected one handle (footer fetch) per file, got {calls['open_file']}" + assert ( + calls["read_row_groups_s3"] == 0 + ), "per-call S3 decode entry point ran — client/footer were rebuilt" + assert ( + calls["read_metadata_s3"] == 0 + ), "per-call S3 footer entry point ran — footer fetched twice" + assert got.sort_by("id").equals(table.sort_by("id")) + + # M97: a SECOND planned read of the same bucket + config reuses the + # process-cached client — connect_s3 must not run again (this is the + # single-row-group-task cold-start fix; the footer is still re-fetched + # per read, hence open_file grows to 4). + got2 = pa.concat_tables(list(reader.read(manifest))) + assert calls["connect"] == 1, ( + f"expected the second read to reuse the cached S3 client, but " + f"connect_s3 ran {calls['connect']} times" + ) + assert calls["open_file"] == 4 + assert got2.sort_by("id").equals(table.sort_by("id")) + finally: + native_metadata._S3_STORE_CACHE.clear() + + +def test_arrow_rs_s3_client_cache_key_and_kill_switch(monkeypatch): + """The process-level client cache (findings M97) is keyed by (bucket, full + connection config): same config reuses, changed credentials/region/bucket + rebuild (that key-miss is how credential rotation stays safe), and + ``RAY_DATA_ARROW_RS_S3_CLIENT_CACHE=0`` bypasses the cache entirely. + Pure construction — ``connect_s3`` does no network IO at build time.""" + from pyarrow.fs import S3FileSystem + + from ray.data._internal.datasource_v2 import native_metadata + + builds = {"n": 0} + orig_connect = ray_data_arrow_rs.connect_s3 + + def spy_connect(*a, **k): + builds["n"] += 1 + return orig_connect(*a, **k) + + monkeypatch.setattr(ray_data_arrow_rs, "connect_s3", spy_connect) + + def make_fs(**overrides): + kwargs = dict( + access_key="key-a", + secret_key="secret-a", + region="us-east-1", + endpoint_override="http://127.0.0.1:9", + ) + kwargs.update(overrides) + return S3FileSystem(**kwargs) + + native_metadata._S3_STORE_CACHE.clear() + try: + fs = make_fs() + s1 = native_metadata.connect_native_s3("bucket-a", fs) + assert builds["n"] == 1 + # Same bucket + config (even via a distinct but identical fs object): + # cache hit, same store instance. + s2 = native_metadata.connect_native_s3("bucket-a", make_fs()) + assert builds["n"] == 1 and s2 is s1 + # Rotated credentials -> new key -> fresh client. + native_metadata.connect_native_s3("bucket-a", make_fs(secret_key="rotated")) + assert builds["n"] == 2 + # Different region / different bucket -> fresh clients too. + native_metadata.connect_native_s3("bucket-a", make_fs(region="eu-west-1")) + native_metadata.connect_native_s3("bucket-b", fs) + assert builds["n"] == 4 + # Kill switch: every call builds, and the cache is left untouched. + monkeypatch.setenv("RAY_DATA_ARROW_RS_S3_CLIENT_CACHE", "0") + before = dict(native_metadata._S3_STORE_CACHE) + native_metadata.connect_native_s3("bucket-a", fs) + native_metadata.connect_native_s3("bucket-a", fs) + assert builds["n"] == 6 + assert native_metadata._S3_STORE_CACHE == before + finally: + native_metadata._S3_STORE_CACHE.clear() + + +def test_arrow_rs_s3_sum(s3_fs, s3_path, restore_ctx): + """The `ds.sum()` aggregation workload (§3.3) over S3 must match PyArrow and + ground truth — decode-heavy / output-light through the native S3 path.""" + table = _flat_table() + uri = _s3_write(table, s3_fs, s3_path) + expected = pc.sum(table["id"]).as_py() + + restore_ctx.use_arrow_rs_parquet_reader = False + pa_sum = ray.data.read_parquet(uri, filesystem=s3_fs).sum("id") + restore_ctx.use_arrow_rs_parquet_reader = True + rs_sum = ray.data.read_parquet(uri, filesystem=s3_fs).sum("id") + + assert rs_sum == pa_sum == expected + + +@pytest.mark.parametrize("k", [2, 4, 8]) +def test_arrow_rs_s3_kspilt_windowed_order(s3_fs, s3_path, k): + """The windowed K-split S3 path (lone big row group → K concurrent GET streams, + plus a small fetch window slicing each range into sub-windows) must return rows + in EXACT file order — not just the right set. Forces the path by calling the + native entry point directly with a tiny split threshold and window; a monotone + `id` column then makes any range/window mis-ordering a hard failure. + """ + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + _s3_config, + ) + + n = 50_000 + table = pa.table({"id": pa.array(np.arange(n, dtype=np.int64))}) + _s3_write(table, s3_fs, s3_path, name="mono.parquet") + + base = _unwrap_protocol(s3_path) + bucket, _, key = os.path.join(base, "mono.parquet").partition("/") + cfg = _s3_config(s3_fs) + + stream = ray_data_arrow_rs.read_row_groups_s3( + bucket, + key, + cfg["region"], + cfg["anonymous"], + endpoint=cfg["endpoint"], + access_key_id=cfg["access_key_id"], + secret_access_key=cfg["secret_access_key"], + session_token=cfg["session_token"], + allow_http=cfg["allow_http"], + virtual_hosted_style=cfg["virtual_hosted_style"], + row_groups=[0], + columns=["id"], + batch_size=4096, + fetch_window_mb=1, # force sub-window slicing within each range + k=k, + split_threshold_bytes=1, # force the K-split path on a small file + ) + got = pa.RecordBatchReader.from_stream(stream).read_all() + assert got.num_rows == n + # Exact order: id must be 0,1,2,...,n-1 with no reordering across K ranges. + assert got["id"].to_pylist() == list(range(n)) + + +@pytest.mark.parametrize("prefetch_budget_mb", [0, 1, 64]) +def test_arrow_rs_s3_window_prefetch_order(s3_fs, s3_path, prefetch_budget_mb): + """The budget-gated prefetch pipeline (K=1 single stream, many row-window + units) must return rows in EXACT file order at every bucket size. This is + the common S3 path: one row group sliced into many small windows, fetched + concurrently under the byte-budget semaphore while a single decoder drains + them in order. budget=0 is strictly-serial fetch→decode→fetch; a budget + smaller than one window still admits that window alone (clamped); a large + budget lets many window fetches run concurrently. A monotone ``id`` makes + any window mis-ordering a hard failure, so this guards that concurrent + fetches never let a later window's rows overtake an earlier one's. + """ + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + _s3_config, + ) + + n = 50_000 + table = pa.table({"id": pa.array(np.arange(n, dtype=np.int64))}) + _s3_write(table, s3_fs, s3_path, name="mono_pf.parquet") + + base = _unwrap_protocol(s3_path) + bucket, _, key = os.path.join(base, "mono_pf.parquet").partition("/") + cfg = _s3_config(s3_fs) + + stream = ray_data_arrow_rs.read_row_groups_s3( + bucket, + key, + cfg["region"], + cfg["anonymous"], + endpoint=cfg["endpoint"], + access_key_id=cfg["access_key_id"], + secret_access_key=cfg["secret_access_key"], + session_token=cfg["session_token"], + allow_http=cfg["allow_http"], + virtual_hosted_style=cfg["virtual_hosted_style"], + row_groups=[0], + columns=["id"], + batch_size=4096, + fetch_window_mb=1, # force many sub-windows within the single stream + k=1, # single stream: isolate the prefetch pipeline from K-split + prefetch_budget_mb=prefetch_budget_mb, + ) + got = pa.RecordBatchReader.from_stream(stream).read_all() + assert got.num_rows == n + assert got["id"].to_pylist() == list(range(n)) + + +def test_arrow_rs_s3_struct_parity(s3_fs, s3_path, restore_ctx): + """A struct column over S3 decodes natively (the gate admits struct now, + on S3 exactly as it does locally) and stays byte-identical to PyArrow.""" + table = pa.table( + { + "id": pa.array(np.arange(2000, dtype=np.int64)), + "st": pa.StructArray.from_arrays( + [pa.array(np.arange(2000)), pa.array(np.arange(2000) * 0.5)], + names=["a", "b"], + ), + } + ) + uri = _s3_write(table, s3_fs, s3_path, name="struct.parquet") + + restore_ctx.use_arrow_rs_parquet_reader = False + pa_tbl = pa.Table.from_pandas( + ray.data.read_parquet(uri, filesystem=s3_fs).to_pandas() + ).sort_by("id") + restore_ctx.use_arrow_rs_parquet_reader = True + rs_tbl = pa.Table.from_pandas( + ray.data.read_parquet(uri, filesystem=s3_fs).to_pandas() + ).sort_by("id") + assert pa_tbl.equals(rs_tbl) + + +def _write_int96(path, embed_arrow_schema): + """Write a Parquet file storing the timestamp column as the legacy INT96 + physical type. ``embed_arrow_schema=False`` mimics Spark/Hive/Impala (no + embedded Arrow schema); ``True`` mimics a PyArrow writer, which pins the + source ``timestamp[us]`` unit in the file's key-value metadata.""" + import datetime + + ts = pa.array( + [ + datetime.datetime(2021, 6, 1) + datetime.timedelta(minutes=i) + for i in range(2000) + ], + type=pa.timestamp("us"), + ) + table = pa.table({"id": pa.array(np.arange(2000, dtype=np.int64)), "t": ts}) + pq.write_table( + table, + str(path), + use_deprecated_int96_timestamps=True, + store_schema=embed_arrow_schema, + write_page_index=True, + ) + return table + + +def _read_manifest_with_path_verdict(path, monkeypatch): + """Drive ``ArrowRsParquetFileReader.read`` in-process over a whole-file + manifest and report ``(table, took_native_decode)``. + + In-process (not via ``ray.data.read_parquet``) so the crate/pyarrow spies + actually observe the decode — Ray would run it in a worker where a + driver-side monkeypatch is invisible. ``took_native_decode`` is True iff a + native handle decode ran (the file took the native path); a fallback + instead opens the file via ``pyarrow.dataset.dataset``.""" + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + calls = _spy_native_decode(monkeypatch) + + reader = ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), target_block_size=128 * 1024 * 1024 + ) + manifest = _make_manifest([str(path)], [os.path.getsize(path)], [None]) + table = pa.concat_tables(list(reader.read(manifest))).sort_by("id") + return table, calls["decode"] > 0 + + +def test_int96_no_arrow_hint_reads_native_as_ns(tmp_path, restore_ctx, monkeypatch): + """A Spark/Hive/Impala-style INT96 file (no embedded Arrow schema) decodes to + ``timestamp[ns]`` on both paths, so the crate handles it *natively* and stays + byte-identical to PyArrow — bringing the common INT96 case onto the + memory-efficient native path.""" + path = tmp_path / "spark_int96.parquet" + _write_int96(path, embed_arrow_schema=False) + + # The crate's footer read reports the column as INT96 and decodes it to ns. + md = ray_data_arrow_rs.read_metadata(str(path)) + assert "t" in list(md.int96_columns) + assert pa.schema(md).field("t").type == pa.timestamp("ns") + + rs_tbl, took_native = _read_manifest_with_path_verdict(path, monkeypatch) + pa_tbl = pa.Table.from_pandas(ray.data.read_parquet(str(path)).to_pandas()).sort_by( + "id" + ) + + assert took_native, "INT96/ns file should take the native decode path" + assert rs_tbl.schema.field("t").type == pa.timestamp("ns") + assert pa_tbl.equals(rs_tbl) + + +def test_int96_with_non_ns_hint_realigns_natively(tmp_path, restore_ctx, monkeypatch): + """A PyArrow-written INT96 file embeds a ``timestamp[us]`` hint that the crate + honors (→ us) but PyArrow ignores (always ns). The planned native path stays + native and *realigns*: the plan-time ``_ColumnAlignment`` casts the decoded + column to ``ns`` so the result is byte-identical to PyArrow — no fallback.""" + path = tmp_path / "pyarrow_int96.parquet" + _write_int96(path, embed_arrow_schema=True) + + # Crate decodes to us (honoring the embedded hint) — the divergence the + # alignment cast repairs. + assert pa.schema(ray_data_arrow_rs.read_metadata(str(path))).field( + "t" + ).type == pa.timestamp("us") + + rs_tbl, took_native = _read_manifest_with_path_verdict(path, monkeypatch) + pa_tbl = pa.Table.from_pandas(ray.data.read_parquet(str(path)).to_pandas()).sort_by( + "id" + ) + + assert took_native, "INT96/non-ns-hint file should decode natively + realign" + assert rs_tbl.schema.field("t").type == pa.timestamp("ns") + assert pa_tbl.equals(rs_tbl) + + +def test_int96_gate_rejects_non_ns_in_columns_supported(tmp_path): + """Unit-level: ``_columns_supported`` (the per-fragment re-gate, which + requires a *no-op* alignment) stays True for an INT96 column already at + ns/no-tz and False for a non-ns unit — the latter now means "needs an + alignment cast", which the planned ``read()`` handles natively while the + pyarrow-fragment path conservatively falls back.""" + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + reader = ArrowRsParquetFileReader(filesystem=LocalFileSystem()) + ns_schema = pa.schema([("id", pa.int64()), ("t", pa.timestamp("ns"))]) + us_schema = pa.schema([("id", pa.int64()), ("t", pa.timestamp("us"))]) + + # 't' flagged as INT96: ns is admitted, us (non-ns) is rejected. + assert reader._columns_supported(ns_schema, ["id", "t"], ["t"]) is True + assert reader._columns_supported(us_schema, ["id", "t"], ["t"]) is False + # A genuine (non-INT96) timestamp[us] column stays supported — only INT96 + # columns are unit-gated. + assert reader._columns_supported(us_schema, ["id", "t"], []) is True + # Projecting away the INT96 column sidesteps the gate. + assert reader._columns_supported(us_schema, ["id"], ["t"]) is True + + +# --------------------------------------------------------------------------- +# Column alignment: gates closed by the per-file post-decode fixup plan +# (schema-evolution null-fill, unified-schema cast, INT96 coercion, forced +# dictionary_columns). Each test compares the native reader against the base +# PyArrow reader IN-PROCESS over the same manifest, with a crate spy proving +# the native decode actually ran (no silent fallback). +# --------------------------------------------------------------------------- + + +def _read_both_in_process(paths, monkeypatch, **reader_kwargs): + """Read ``paths`` with the arrow-rs reader and the base PyArrow reader over + an identical whole-file manifest; return ``(rs_table, pa_table, + native_decodes)`` where ``native_decodes`` counts crate decode calls.""" + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + from ray.data._internal.datasource_v2.readers.parquet_file_reader import ( + ParquetFileReader, + ) + + calls = _spy_native_decode(monkeypatch) + + paths = [str(p) for p in paths] + manifest = _make_manifest( + paths, [os.path.getsize(p) for p in paths], [None] * len(paths) + ) + rs = pa.concat_tables( + list( + ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), **reader_kwargs + ).read(manifest) + ) + ) + pa_tbl = pa.concat_tables( + list( + ParquetFileReader(filesystem=LocalFileSystem(), **reader_kwargs).read( + manifest + ) + ) + ) + return rs, pa_tbl, calls["decode"] + + +def _evolved_fixture(tmp_path): + """Two-file dataset with schema evolution: file A has (id, x, s[large]); + file B lacks ``x`` and stores ``s`` as plain ``string`` (type drift). The + unified schema is A's.""" + a = tmp_path / "evo_a.parquet" + b = tmp_path / "evo_b.parquet" + pq.write_table( + pa.table( + { + "id": pa.array([1, 2], pa.int64()), + "x": pa.array([1.5, 2.5], pa.float64()), + "s": pa.array(["a", "b"], pa.large_string()), + } + ), + str(a), + write_page_index=True, + ) + pq.write_table( + pa.table( + { + "id": pa.array([3, 4], pa.int64()), + "s": pa.array(["c", "d"], pa.string()), + } + ), + str(b), + write_page_index=True, + ) + unified = pa.schema( + [("id", pa.int64()), ("x", pa.float64()), ("s", pa.large_string())] + ) + return a, b, unified + + +def test_schema_evolution_null_fill_native_parity(tmp_path, monkeypatch): + """A file missing a unified-schema column (schema evolution) plus per-file + type drift decodes NATIVELY: the alignment null-fills the missing column + with the unified type and casts the drifted one, byte-matching PyArrow's + pinned-schema scan. Both files must take the crate path.""" + a, b, unified = _evolved_fixture(tmp_path) + + rs, pa_tbl, native_decodes = _read_both_in_process( + [a, b], monkeypatch, schema=unified + ) + assert native_decodes >= 2, "both files should decode natively" + assert rs.sort_by("id").equals(pa_tbl.sort_by("id")) + assert rs.schema == pa_tbl.schema + + # Same with an explicit projection that includes the evolved column. + rs, pa_tbl, native_decodes = _read_both_in_process( + [a, b], monkeypatch, schema=unified, columns=["x", "id"] + ) + assert native_decodes >= 2 + assert rs.sort_by("id").equals(pa_tbl.sort_by("id")) + + +def test_filter_on_evolved_column_native_parity(tmp_path, monkeypatch): + """A pushed predicate over the null-filled (evolved) column evaluates on + the ALIGNED batch, so rows from the file lacking the column drop exactly + like PyArrow's null-comparison semantics.""" + from ray.data.expressions import col + + a, b, unified = _evolved_fixture(tmp_path) + rs, pa_tbl, native_decodes = _read_both_in_process( + [a, b], monkeypatch, schema=unified, predicate=(col("x") > 2.0) + ) + assert native_decodes >= 1 + assert rs.sort_by("id").equals(pa_tbl.sort_by("id")) + assert rs.num_rows == 1 # only id=2 (x=2.5) survives; file B is all-null x + + +def test_coerce_int96_timestamp_unit_falls_back(tmp_path, monkeypatch): + """A file decoding INT96 under ``coerce_int96_timestamp_unit`` FALLS BACK + (with or without an embedded arrow-schema hint): pyarrow's decode-time + coercion floors (parquet types.h divides the unsigned nanos-of-day) while + a post-decode cast truncates toward zero — one unit apart on every + pre-1970 value, so no cast reproduces the kwarg. Parity still holds + because the fallback IS pyarrow for that file.""" + for embed in (False, True): + path = tmp_path / f"i96_{embed}.parquet" + _write_int96(path, embed_arrow_schema=embed) + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], + monkeypatch, + parquet_format_kwargs={"coerce_int96_timestamp_unit": "ms"}, + ) + assert native_decodes == 0, f"embed={embed} must fall back under the kwarg" + assert rs.schema.field("t").type == pa.timestamp("ms") + assert rs.sort_by("id").equals(pa_tbl.sort_by("id")) + + +def test_forced_dictionary_columns_native_parity(tmp_path, monkeypatch): + """A forced ``dictionary_columns`` read no longer falls back: the crate + decodes the plain column and the alignment dictionary-casts it to exactly + PyArrow's forced-dict output (``dictionary``).""" + path = tmp_path / "forced_dict.parquet" + pq.write_table( + pa.table( + { + "id": pa.array([1, 2, 3], pa.int64()), + "s": pa.array(["x", "y", "x"]), + } + ), + str(path), + write_page_index=True, + ) + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], monkeypatch, parquet_format_kwargs={"dictionary_columns": ["s"]} + ) + assert native_decodes >= 1 + assert rs.schema.field("s").type == pa.dictionary(pa.int32(), pa.string()) + assert rs.sort_by("id").equals(pa_tbl.sort_by("id")) + + +def test_dotted_nested_projection_native_parity(tmp_path, monkeypatch): + """Dotted (nested-field) projection like ``user.name`` stays NATIVE and + matches the base reader exactly — which today means the dotted column is + silently DROPPED by both paths. + + V2 discards dotted names *before* any reader sees them: + ``FileReader._split_columns`` classifies ``user.name`` as not-on-disk (the + footer schema has only the root ``user``), routing it to the synthesize + bucket, where nothing synthesizes it and ``_postprocess`` drops it. The raw + PyArrow scanner *could* resolve it (yielding a leaf column named ``name``), + but Ray never passes it through. So true nested projection is a platform + feature V2 lacks on the PyArrow path too — implementing a Rust + ``ProjectionMask`` for it would be new functionality, not migration parity. + This test pins the parity: no fallback, identical (dropped-column) output. + """ + path = tmp_path / "nested_proj.parquet" + pq.write_table( + pa.table( + { + "id": pa.array([1, 2], pa.int64()), + "user": pa.array( + [{"name": "a", "age": 30}, {"name": "b", "age": 40}], + type=pa.struct([("name", pa.string()), ("age", pa.int64())]), + ), + } + ), + str(path), + write_page_index=True, + ) + + # Dotted projection: native decode, and both readers drop the dotted name. + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], monkeypatch, columns=["user.name", "id"] + ) + assert native_decodes >= 1, "dotted projection must not force a fallback" + assert rs.column_names == ["id"] + assert rs.equals(pa_tbl) + + # Whole-struct projection (the supported way to read nested data): native + # decode with full byte parity. + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], monkeypatch, columns=["user", "id"] + ) + assert native_decodes >= 1 + assert rs.column_names == ["user", "id"] + assert rs.equals(pa_tbl) + + +def test_flat_column_named_with_dot_native_parity(tmp_path, monkeypatch): + """A FLAT top-level column literally named ``"user.name"`` (legal in + Parquet) decodes natively with parity — even when the same file also has a + struct ``user`` with a ``name`` child. Exact flat-name match wins in both + the pyarrow scanner and the crate, so the gate must not fall back on a + dot in a column name (dots only ever mean nested *projection* upstream of + the reader, where V2 discards them).""" + path = tmp_path / "flatdot.parquet" + pq.write_table( + pa.table( + { + "id": pa.array([1, 2], pa.int64()), + "user.name": pa.array(["flat1", "flat2"]), + "user": pa.array( + [{"name": "nested1"}, {"name": "nested2"}], + type=pa.struct([("name", pa.string())]), + ), + } + ), + str(path), + write_page_index=True, + ) + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], monkeypatch, columns=["user.name", "id"] + ) + assert native_decodes >= 1, "flat dotted-named column must decode natively" + assert rs.column_names == ["user.name", "id"] + assert rs.column("user.name").to_pylist() == ["flat1", "flat2"] + assert rs.equals(pa_tbl) + + +def test_perf_only_format_kwargs_stay_native(tmp_path, monkeypatch): + """I/O-tuning format kwargs (``pre_buffer`` / ``buffer_size`` / + ``use_buffered_stream``) cannot change decoded bytes, so the native path + ignores them and stays native; the PyArrow reader honors them and must + produce identical output.""" + path = tmp_path / "perf_kwargs.parquet" + pq.write_table( + pa.table( + { + "id": pa.array([1, 2, 3], pa.int64()), + "s": pa.array(["a", "b", "c"]), + } + ), + str(path), + write_page_index=True, + ) + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], + monkeypatch, + parquet_format_kwargs={ + "pre_buffer": False, + "buffer_size": 64 * 1024, + "use_buffered_stream": True, + }, + ) + assert native_decodes >= 1, "perf-only format kwargs must not force fallback" + assert rs.sort_by("id").equals(pa_tbl.sort_by("id")) + + +def test_unsupported_format_kwarg_falls_back(tmp_path, monkeypatch): + """A format kwarg outside the native allowlist forces a PyArrow fallback + (which honors it) instead of being silently ignored. Exercised with + ``arrow_extensions_enabled`` — a pyarrow 21+ schema-shaping toggle the + native path doesn't reproduce (see TODO.md "arrow_extensions_enabled").""" + path = tmp_path / "audit_kwargs.parquet" + pq.write_table( + pa.table({"id": pa.array([1, 2], pa.int64())}), str(path), write_page_index=True + ) + + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], + monkeypatch, + parquet_format_kwargs={"arrow_extensions_enabled": True}, + ) + assert native_decodes == 0, "unsupported format kwarg must force fallback" + assert rs.equals(pa_tbl) + + +def test_strict_mode_raises_on_fallback_and_passes_native(tmp_path, monkeypatch): + """``RAY_DATA_ARROW_RS_STRICT`` turns every decision to serve a read via + the PyArrow fallback into a hard error — so a large-scale validation run + can *guarantee* it exercised the native path — while leaving natively + supported reads untouched.""" + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + path = tmp_path / "strict.parquet" + pq.write_table( + pa.table({"id": pa.array([1, 2], pa.int64())}), str(path), write_page_index=True + ) + monkeypatch.setenv("RAY_DATA_ARROW_RS_STRICT", "1") + manifest = _make_manifest([str(path)], [os.path.getsize(path)], [None]) + + # Native-supported read: strict mode is a no-op. + table = pa.concat_tables( + list(ArrowRsParquetFileReader(filesystem=LocalFileSystem()).read(manifest)) + ) + assert table["id"].to_pylist() == [1, 2] + + # Fallback-forcing read (format kwarg outside the allowlist): raises + # instead of silently serving PyArrow-decoded bytes. + reader = ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), + parquet_format_kwargs={"arrow_extensions_enabled": True}, + ) + with pytest.raises(RuntimeError, match="RAY_DATA_ARROW_RS_STRICT"): + list(reader.read(manifest)) + + # Per-file gate (no plannable alignment): a file whose read requires an + # unsupported read-time coercion also raises under strict mode. + reader = ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), + parquet_format_kwargs={"dictionary_columns": ["id"]}, + schema=pa.schema([("id", pa.dictionary(pa.int32(), pa.int64()))]), + ) + try: + got = list(reader.read(manifest)) + except RuntimeError as e: + assert "RAY_DATA_ARROW_RS_STRICT" in str(e) + else: + # If the alignment can plan this coercion it stays native — equally + # acceptable; the point is "never a silent fallback under strict". + assert got + + +def test_thrift_limits_native_parity(tmp_path, monkeypatch): + """The thrift footer limits stay NATIVE via the metadata-only pyarrow + footer probe: a generous limit decodes natively with identical output, + and a tiny limit raises the same ``OSError`` the base path raises + (parity-of-error) — from both readers, not just the fallback.""" + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.parquet_file_reader import ( + ParquetFileReader, + ) + + path = tmp_path / "thrift_limits.parquet" + pq.write_table( + pa.table({"id": pa.array([1, 2], pa.int64())}), str(path), write_page_index=True + ) + + # Generous limit: the probe passes and the decode is native. + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], + monkeypatch, + parquet_format_kwargs={"thrift_string_size_limit": 1 << 20}, + ) + assert native_decodes >= 1, "thrift limits must not force fallback anymore" + assert rs.equals(pa_tbl) + + # Tiny limit: both readers must reject the footer with the same error. + tiny = {"thrift_string_size_limit": 10} + manifest = _make_manifest([str(path)], [os.path.getsize(path)], [None]) + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + + with pytest.raises(OSError): + list( + ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), parquet_format_kwargs=dict(tiny) + ).read(manifest) + ) + with pytest.raises(OSError): + list( + ParquetFileReader( + filesystem=LocalFileSystem(), parquet_format_kwargs=dict(tiny) + ).read(manifest) + ) + + +def _write_crc_file(path, corrupt=False): + """Write an uncompressed file with page checksums; optionally flip one + byte inside a data page (detectable only via CRC verification).""" + table = pa.table( + { + "id": pa.array(range(100), pa.int64()), + "s": pa.array([f"crc_sentinel_{i:04d}" for i in range(100)]), + } + ) + pq.write_table( + table, + str(path), + write_page_index=True, + compression="NONE", + write_page_checksum=True, + ) + if corrupt: + data = bytearray(path.read_bytes()) + idx = data.find(b"crc_sentinel_0050") + assert idx > 0 + data[idx] = ord("X") + path.write_bytes(bytes(data)) + return table + + +def test_page_checksum_verification_true_native(tmp_path, monkeypatch): + """``page_checksum_verification=True`` decodes natively: the crate is + built with parquet's ``crc`` feature and always verifies stored page + CRCs, so True *is* the native behavior — clean files match byte-for-byte, + and a corrupt page raises from BOTH readers (parity-of-error).""" + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + from ray.data._internal.datasource_v2.readers.parquet_file_reader import ( + ParquetFileReader, + ) + + clean = tmp_path / "crc_clean.parquet" + _write_crc_file(clean) + rs, pa_tbl, native_decodes = _read_both_in_process( + [clean], + monkeypatch, + parquet_format_kwargs={"page_checksum_verification": True}, + ) + assert native_decodes >= 1, "page_checksum_verification=True must stay native" + assert rs.equals(pa_tbl) + + corrupt = tmp_path / "crc_corrupt.parquet" + _write_crc_file(corrupt, corrupt=True) + manifest = _make_manifest([str(corrupt)], [os.path.getsize(corrupt)], [None]) + kwargs = {"page_checksum_verification": True} + for reader_cls in (ArrowRsParquetFileReader, ParquetFileReader): + with pytest.raises((OSError, pa.lib.ArrowInvalid), match="CRC"): + list( + reader_cls( + filesystem=LocalFileSystem(), parquet_format_kwargs=dict(kwargs) + ).read(manifest) + ) + + +def test_page_checksum_verification_false_falls_back(tmp_path, monkeypatch): + """An explicit ``page_checksum_verification=False`` is the opt-out for + reading a file *despite* corrupt checksums. The crate build always + verifies (compile-time ``crc`` feature, no off-switch), so only PyArrow + can honor the opt-out — the read must fall back and succeed, returning + the same (corrupted) bytes from both paths.""" + path = tmp_path / "crc_optout.parquet" + _write_crc_file(path, corrupt=True) + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], + monkeypatch, + parquet_format_kwargs={"page_checksum_verification": False}, + ) + assert native_decodes == 0, "explicit False must force the PyArrow fallback" + assert rs.equals(pa_tbl) + # The corruption is really there — the flipped byte comes through. + assert rs["s"][50].as_py() == "Xrc_sentinel_0050" + + +def _schema_shaped_fixture(tmp_path): + """A file WITHOUT an embedded arrow schema (``store_schema=False``), so + ``binary_type`` / ``list_type`` genuinely change pyarrow's decoded types + (with an embedded schema they are inert).""" + path = tmp_path / "schema_shaped.parquet" + pq.write_table( + pa.table( + { + "id": pa.array([1, 2], pa.int64()), + "b": pa.array([b"x", b"y"], pa.binary()), + "l": pa.array([[1, 2], [3]], pa.list_(pa.int64())), + "s": pa.array(["a", "bb"], pa.string()), + } + ), + str(path), + write_page_index=True, + store_schema=False, + ) + return path + + +def test_schema_shaped_kwargs_native_with_pinned_schema(tmp_path, monkeypatch): + """``binary_type`` / ``list_type`` with a pinned dataset schema decode + natively: the pin is the output-type authority — on the base path the + pinned-schema cast silently *undoes* these kwargs (the V2 listing infers + the schema via ``pq.read_schema``, which is blind to them), and the + native path's alignment produces the pinned types identically.""" + path = _schema_shaped_fixture(tmp_path) + unified = pq.read_schema(str(path)) + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], + monkeypatch, + schema=unified, + parquet_format_kwargs={ + "binary_type": pa.large_binary(), + "list_type": pa.LargeListType, + }, + ) + assert native_decodes >= 1, "pinned-schema read must stay native" + assert rs.equals(pa_tbl) + # The pin wins: output types are the plain (non-large) footer types. + assert rs.schema.field("b").type == pa.binary() + assert rs.schema.field("s").type == pa.string() + + +def test_schema_shaped_kwargs_fall_back_without_schema(tmp_path, monkeypatch): + """Without a pinned schema, ``binary_type`` / ``list_type`` genuinely + change the decoded types (large_binary / large_string / large_list on a + no-embedded-schema file) — the crate doesn't reproduce that, so the read + falls back to PyArrow and both paths agree on the large types.""" + path = _schema_shaped_fixture(tmp_path) + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], + monkeypatch, + parquet_format_kwargs={ + "binary_type": pa.large_binary(), + "list_type": pa.LargeListType, + }, + ) + assert native_decodes == 0, "schema-shaping kwargs without a pin must fall back" + assert rs.equals(pa_tbl) + assert rs.schema.field("b").type == pa.large_binary() + assert rs.schema.field("s").type == pa.large_string() + assert rs.schema.field("l").type == pa.large_list(pa.int64()) + + +def test_arrow_rs_tuning_kwargs_reach_crate(tmp_path, monkeypatch): + """``arrow_rs_*`` tuning knobs in ``dataset_kwargs`` must (a) reach the + crate's decode call with the requested values, (b) keep the read native, + and (c) be ignored by the base PyArrow reader (popped before + ``pds.ParquetFileFormat`` sees them) — the mirror image of the native + reader ignoring ``pre_buffer``. The base reader reading the same kwargs + without a TypeError is assertion (c).""" + path = tmp_path / "tuning_kwargs.parquet" + table = _flat_table(10_000) + pq.write_table(table, str(path), write_page_index=True, row_group_size=10_000) + + # Capture the crate call kwargs underneath the helper's own decode spy: + # stack a second _spy_native_decode whose proxy records each handle + # decode's (args, kwargs) — the planned path passes all knobs by keyword. + captured = _spy_native_decode(monkeypatch) + + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], + monkeypatch, + parquet_format_kwargs={ + "arrow_rs_decode_budget_bytes": 4 * 1024 * 1024, + "arrow_rs_k": 2, + "arrow_rs_split_threshold_bytes": 0, # force the K-split path + }, + ) + assert native_decodes >= 1, "tuning kwargs must not force fallback" + assert rs.sort_by("id").equals(pa_tbl.sort_by("id")) + assert captured["decode_calls"], "crate decode call was not captured" + for _, kwargs in captured["decode_calls"]: + assert kwargs["decode_budget_bytes"] == 4 * 1024 * 1024 + assert kwargs["k"] == 2 + assert kwargs["split_threshold_bytes"] == 0 + + +def test_arrow_rs_tuning_kwarg_typo_raises(tmp_path): + """A misspelled ``arrow_rs_*`` key must fail loudly at reader construction + (both readers share the check in ``ParquetFileReader.__init__``), not + surface as a baffling pyarrow TypeError or a silent native fallback.""" + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + from ray.data._internal.datasource_v2.readers.parquet_file_reader import ( + ParquetFileReader, + ) + + for reader_cls in (ArrowRsParquetFileReader, ParquetFileReader): + with pytest.raises(ValueError, match="arrow_rs_decode_budget"): + reader_cls( + filesystem=LocalFileSystem(), + parquet_format_kwargs={"arrow_rs_decode_budget": 1}, # typo'd key + ) + + +@pytest.mark.parametrize("bad_value", [0, "four", True]) +def test_arrow_rs_tuning_kwarg_invalid_value_raises(tmp_path, monkeypatch, bad_value): + """An invalid tuning value (wrong type, below minimum, or a bool sneaking + in as int) must raise a ValueError naming the knob — a mis-set perf knob + silently clamped or ignored would corrupt benchmarks.""" + path = tmp_path / "bad_tuning.parquet" + pq.write_table( + pa.table({"id": pa.array([1, 2], pa.int64())}), str(path), write_page_index=True + ) + with pytest.raises(ValueError, match="arrow_rs_k"): + _read_both_in_process( + [path], + monkeypatch, + parquet_format_kwargs={"arrow_rs_k": bad_value}, + ) + + +def test_arrow_rs_tuning_kwargs_end_to_end(tmp_path, restore_ctx): + """The knobs survive the full ``read_parquet(dataset_kwargs=...)`` plumbing + (read_api -> datasource -> scanner -> reader) under BOTH flag settings: + natively they tune the crate; on the PyArrow reader they're inert — the + same call is valid either way.""" + path = tmp_path / "tuning_e2e.parquet" + table = _flat_table(5_000) + pq.write_table(table, str(path), write_page_index=True) + + kwargs = {"dataset_kwargs": {"arrow_rs_decode_budget_bytes": 4 * 1024 * 1024}} + pa_tbl = _read_sorted(path, False, restore_ctx, **kwargs) + rs_tbl = _read_sorted(path, True, restore_ctx, **kwargs) + assert pa_tbl.equals(rs_tbl) + assert rs_tbl.num_rows == table.num_rows + + +def test_unified_only_column_not_dropped_natively(tmp_path, monkeypatch): + """A unified-schema column absent from EVERY file in the split must still + surface (all-null), matching the base reader under a pinned schema — the + column split takes names from the unified schema, not the footers.""" + path = tmp_path / "no_extra.parquet" + pq.write_table( + pa.table({"id": pa.array([1, 2], pa.int64())}), str(path), write_page_index=True + ) + unified = pa.schema([("id", pa.int64()), ("later_col", pa.string())]) + rs, pa_tbl, native_decodes = _read_both_in_process( + [path], monkeypatch, schema=unified + ) + assert native_decodes >= 1 + assert "later_col" in rs.column_names + assert rs.sort_by("id").equals(pa_tbl.sort_by("id")) + + +def _dispatch_fragments(reader, monkeypatch, n=2): + """Run ``_dispatch_fragment_reads`` over ``n`` stub fragments, reporting + whether the concurrent path was taken. + + Asserting on ``_num_fragment_read_threads()`` alone would be too weak: the + value it returns changes the code *path*, because ``num_workers <= 1`` returns + early into ``_read_fragments_sequential`` and ``make_async_gen`` is never + constructed. A future refactor could keep the number and lose the branch. So + spy on ``make_async_gen`` at the module where it is looked up. + """ + from ray.data._internal.datasource_v2.readers import file_reader as fr_mod + + used = {"async": False} + orig = fr_mod.make_async_gen + + def spy(*a, **k): + used["async"] = True + return orig(*a, **k) + + monkeypatch.setattr(fr_mod, "make_async_gen", spy) + monkeypatch.setattr( + type(reader), + "_iter_fragment_tables", + lambda self, frag, kwargs: iter([pa.table({"id": [1]})]), + raising=True, + ) + + class _Frag: + path = "stub.parquet" + + tables = list(reader._dispatch_fragment_reads([(_Frag(), i) for i in range(n)], {})) + return used["async"], tables + + +def test_arrow_rs_defaults_bounded_fragment_pool(monkeypatch): + """Both readers decode fragments on a one-worker-per-fragment pool + (pool-width PARITY, decided 2026-08-12 — a narrower arrow-rs pool turned + every multi-fragment A/B into a pool-width comparison instead of a decode + comparison, and at realistic bin budgets a bin spans few files anyway). + A single-fragment task stays on the sequential branch, where the crate + alone owns parallelism. + + History (findings K6, K10 in arrow_rs_docs/findings.md): K6 default 1 → + K10 default 4 (threads=4 vs 1 cuts read-op time 1.6-3.3x at + flat-to-+22% memory) → 2026-08-12 parity. If the bin sweep shows + arrow-rs per-task USS growing with bin size, suspect this default first + and re-cap via RAY_DATA_READ_FILES_NUM_THREADS. + + The assertions are deliberately ``== num_fragments`` at two sizes rather + than a literal: both paths are unbounded, and a cap leaking onto either + is the regression this test exists to catch. + """ + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + ArrowRsParquetFileReader, + ) + from ray.data._internal.datasource_v2.readers.parquet_file_reader import ( + ParquetFileReader, + ) + + kwargs = dict(filesystem=LocalFileSystem(), target_block_size=128 * 1024 * 1024) + + rs_reader = ArrowRsParquetFileReader(**kwargs) + # Parity with the base: one worker per fragment, no cap. + assert rs_reader._num_fragment_read_threads(1) == 1 + assert rs_reader._num_fragment_read_threads(2) == 2 + assert rs_reader._num_fragment_read_threads(4) == 4 + assert rs_reader._num_fragment_read_threads(64) == 64 + used_async, tables = _dispatch_fragments(rs_reader, monkeypatch, n=2) + assert used_async, "arrow-rs multi-fragment dispatch lost its fragment pool" + assert len(tables) == 2, "concurrent path dropped fragments" + used_async, tables = _dispatch_fragments(rs_reader, monkeypatch, n=1) + assert not used_async, ( + "a single-fragment task took the concurrent path — the sequential " + "branch (crate-owned parallelism, no make_async_gen) is gone" + ) + assert len(tables) == 1 + + pa_reader = ParquetFileReader(**kwargs) + # Unbounded on the base path: one worker per fragment. + assert pa_reader._num_fragment_read_threads(2) == 2 + assert pa_reader._num_fragment_read_threads(17) == 17 + used_async, tables = _dispatch_fragments(pa_reader, monkeypatch, n=2) + assert used_async, "the PyArrow reader lost its fragment pool" + assert len(tables) == 2 + + +def test_explicit_num_threads_env_overrides_arrow_rs_default(monkeypatch): + """An explicitly set ``RAY_DATA_READ_FILES_NUM_THREADS`` beats the per-reader + default — a user who set it meant it, and the benchmark harness sweeps it. + + Both the env read and the "was it explicit?" flag happen at import time, so this + patches the two module attributes rather than ``os.environ``. + + This reader resolves the value itself rather than delegating to ``super()``. It + has to: the footer-chunking base path deleted ``_DEFAULT_NUM_THREADS`` and no + longer reads ``RAY_DATA_READ_FILES_NUM_THREADS`` at all, so delegating would + silently ignore an explicit setting — and the benchmark harness's thread sweep + sets exactly this variable, so that failure would be invisible and would corrupt + a whole sweep into flat lines. + """ + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers import ( + arrow_rs_parquet_file_reader as rs_mod, + ) + + monkeypatch.setattr(rs_mod, "_READ_FILES_NUM_THREADS_IS_EXPLICIT", True) + monkeypatch.setattr(rs_mod, "_READ_FILES_NUM_THREADS_EXPLICIT_VALUE", 4) + + reader = rs_mod.ArrowRsParquetFileReader( + filesystem=LocalFileSystem(), target_block_size=128 * 1024 * 1024 + ) + # 4 regardless of the fragment count — an explicit setting is a cap, and it must + # not be widened by the base's unbounded default. + assert reader._num_fragment_read_threads(2) == 4 + assert reader._num_fragment_read_threads(64) == 4 + + +def test_arrow_rs_decode_budget_default_follows_block_target(monkeypatch): + """Pin the default's semantics: budget = ``DataContext.target_max_block_size``. + + History: 2 MiB -> 32 MiB on the 2026-08-07 sweep (regression_testing.md + §8.2), then -> the block target on findings M59/M63: read tasks coalesce + decode batches through ``BlockOutputBuffer`` to ~one block anyway, so + sub-block batches bought no memory while the per-batch × per-column + dispatch cost was the whole in-Ray wall loss on 5,000-col schemas (M59 + wall R 1.40 -> 0.99 at 128 MiB); the 10-shape gate at 128 MiB passed the + memory gate on every cell (M63). Env var must still win when set, and an + unset block target must NOT mean an unbounded decode budget. + """ + from ray.data._internal.datasource_v2.readers import ( + arrow_rs_parquet_file_reader as reader_mod, + ) + from ray.data.context import DEFAULT_TARGET_MAX_BLOCK_SIZE, DataContext + + ctx = DataContext.get_current() + + # No env override (None unless the var was set at import): follow the + # current block target. + monkeypatch.setattr(reader_mod, "_ARROW_RS_DECODE_BUDGET_BYTES", None) + monkeypatch.setattr(ctx, "target_max_block_size", 64 * 1024 * 1024) + assert reader_mod._default_decode_budget_bytes() == 64 * 1024 * 1024 + + # Unset block target: bounded 128 MiB fallback, never unbounded decode. + monkeypatch.setattr(ctx, "target_max_block_size", None) + assert reader_mod._default_decode_budget_bytes() == 128 * 1024 * 1024 + + # Env var (captured at import into _ARROW_RS_DECODE_BUDGET_BYTES) wins + # over the block target. + monkeypatch.setattr(ctx, "target_max_block_size", 64 * 1024 * 1024) + monkeypatch.setattr(reader_mod, "_ARROW_RS_DECODE_BUDGET_BYTES", 2 * 1024 * 1024) + assert reader_mod._default_decode_budget_bytes() == 2 * 1024 * 1024 + + # The budget only binds below budget/floor bytes per row; at the default + # 128 MiB block target that is ~64 KiB/row against the 2048-row request + # floor (the crate's own floor is 32 rows, so decoded batches stay + # budget-sized regardless). + assert ( + DEFAULT_TARGET_MAX_BLOCK_SIZE / reader_mod._ARROW_RS_MIN_DECODE_BATCH_ROWS + > 8 * 1024 + ) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main(["-v", __file__])) + + +def test_arrow_rs_malloc_trim_eos_default_on(monkeypatch): + """The end-of-stream ``malloc_trim(0)`` ships ON: ``DataContext`` defaults + ``arrow_rs_malloc_trim_eos`` to True and the env var still wins when set. + + History (findings M101 / M124 in arrow_rs_docs/findings.md): on the release + fleet one trim per read-task stream closed every arrow-rs retention row + (rlp per-task USS 1.85 -> 0.99, write_parquet sustained 1.40 -> 0.88) at + ~rs wall, and its single 2.37x wall cell (M107) did not replicate x3, so + the default flipped 2026-09-08. A regression of this default would bring + the retention rows back without touching any decode path, i.e. invisibly + to the decoded-bytes gates -- hence the literal ``is True`` here. + """ + from ray._common.utils import env_bool + from ray.data import context as ctx_mod + + if "RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS" in os.environ: + pytest.skip("env override set in this session; default not observable") + assert ctx_mod.DEFAULT_ARROW_RS_MALLOC_TRIM_EOS is True + # A fresh context (not the process-wide one tests mutate) carries it. + assert ctx_mod.DataContext().arrow_rs_malloc_trim_eos is True + # The mallopt lever stays off: only one allocator lever ships. + assert ctx_mod.DEFAULT_ARROW_RS_MALLOC_TRIM is False + # Ablation path: the env var overrides the default at import time. + monkeypatch.setenv("RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS", "0") + assert env_bool("RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS", True) is False + + +def test_arrow_rs_malloc_trim_eos_fires_once_per_read(tmp_path, monkeypatch): + """``DataContext.arrow_rs_malloc_trim_eos`` (env + ``RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS``) calls glibc ``malloc_trim(0)`` exactly + once per ``read()`` stream — not per file, not per batch — and never when + the knob is off. The libc call itself is Linux-only, so the test counts the + module-level trampoline the ``finally`` invokes.""" + import time + + from pyarrow.fs import LocalFileSystem + + from ray.data._internal.datasource_v2.readers import ( + arrow_rs_parquet_file_reader as mod, + ) + from ray.data.context import DataContext + + paths = [] + for i in range(2): + path = tmp_path / f"eos{i}.parquet" + pq.write_table( + pa.table({"id": pa.array([i, i + 10], pa.int64())}), + str(path), + write_page_index=True, + ) + paths.append(str(path)) + manifest = _make_manifest(paths, [os.path.getsize(p) for p in paths], [None, None]) + + calls = {"n": 0} + + def _fake_trim(): + calls["n"] += 1 + time.sleep(0.002) # measurable, so the timed value is provably > 0 + + monkeypatch.setattr(mod, "_malloc_trim_now", _fake_trim) + ctx = DataContext.get_current() + old = ctx.arrow_rs_malloc_trim_eos + try: + ctx.arrow_rs_malloc_trim_eos = False + reader = mod.ArrowRsParquetFileReader(filesystem=LocalFileSystem()) + assert pa.concat_tables(list(reader.read(manifest))).num_rows == 4 + assert calls["n"] == 0, "knob off must never trim" + assert reader.pop_task_stats() == {"trim_wall_s": 0.0} + + ctx.arrow_rs_malloc_trim_eos = True + assert pa.concat_tables(list(reader.read(manifest))).num_rows == 4 + assert calls["n"] == 1, "one trim per read() stream (2 files, many batches)" + # The trim is timed per stream and drained by pop_task_stats() — what + # the ReadFiles transform folds into ReadFilesTaskStats.trim_wall_s. + assert reader.pop_task_stats()["trim_wall_s"] >= 0.002 + assert reader.pop_task_stats() == {"trim_wall_s": 0.0}, "drained" + + # A consumer that stops early still ends the stream exactly once. + gen = reader.read(manifest) + next(gen) + gen.close() + assert calls["n"] == 2 + finally: + ctx.arrow_rs_malloc_trim_eos = old + + +def test_arrow_rs_eos_trim_wall_reaches_driver_without_flush_block( + tmp_path, restore_ctx +): + """``ReadFilesTaskStats.trim_wall_s`` is only known when the reader's stream + ENDS (the eos trim runs in its finalizer, after the last table came out), + and a block's stats snapshot is pickled when the block leaves the task. So + when the last table itself completes a block (buffer >= target but below the + 1.5x slice limit -> emitted whole, no remainder, no flush block) the planner + must fold the drain in BEFORE yielding that table, or the driver-side + ``read_task_trim_wall_s`` reads 0 while the trim actually ran. + + Shape: one numeric file below the 2048-row batch floor -> one table -> one + task -> exactly one block (asserted), target = 0.8x the table's bytes. + """ + from ray.data.block import BlockAccessor + + path = tmp_path / "one.parquet" + n = 2000 + table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "x": pa.array(np.random.default_rng(0).random(n)), + "label": pa.array((np.arange(n) % 5).astype(np.int32)), + } + ) + pq.write_table(table, str(path)) + nbytes = BlockAccessor.for_block(table).size_bytes() + + ctx = restore_ctx + ctx.use_arrow_rs_parquet_reader = True + old_target = ctx.target_max_block_size + old_knob = ctx.arrow_rs_malloc_trim_eos + try: + ctx.target_max_block_size = int(nbytes * 0.8) + ctx.arrow_rs_malloc_trim_eos = True + mds = ray.data.read_parquet(str(path)).materialize() + assert mds.count() == n + # The regression only shows on this shape: no flush block after the + # table that completed the block. + assert mds.num_blocks() == 1 + trim = mds._raw_stats().extra_metrics["read_task_trim_wall_s"] + assert trim["num_samples"] == 1 + # Linux: real malloc_trim(0); elsewhere the timed no-op trampoline. + # Either way strictly positive iff the drain reached the block. + assert trim["max"] > 0.0 + em = mds._raw_stats().extra_metrics + # The same last-block snapshot carries the downstream and first-table + # walls; the single yield of the single table is timed. + assert em["read_task_yield_wall_s"]["num_samples"] == 1 + assert em["read_task_yield_wall_s"]["max"] > 0.0 + assert em["read_task_first_table_wall_s"]["max"] > 0.0 + finally: + ctx.target_max_block_size = old_target + ctx.arrow_rs_malloc_trim_eos = old_knob + + +def test_lookahead_yields_decoded_table_before_later_decode_error( + tmp_path, restore_ctx +): + """The planner's one-table lookahead holds a decoded table until the NEXT + ``next()`` resolves. When that next call raises (IO error, corrupt page) + instead of returning or ending the stream, the held table must still be + yielded before the error propagates -- exactly what happened before the + lookahead, when it had already left the task. Observable under + ``max_errored_blocks``: the tolerated task keeps its good blocks. + + Shape, chosen so EXACTLY ONE table precedes the error and that table + completes a block on its own (otherwise the old planner still emits the + earlier tables, or the shaping buffer eats the table on either planner): + one file, two row groups; row group 0 is 32 rows (arrow-rs's batch floor) + of a 200-byte constant string (dictionary-encoded, so pyarrow's + encoded-bytes batch sizing also covers the whole group in one batch); + row group 1 has its first data page overwritten with zeros behind an + intact footer, so planning (footer reads) succeeds and the failure + happens mid-stream. Block target = 0.9x row group 0's decoded bytes. + (Two files would not do either: a bin lists files in footer-arrival + order, so the corrupt one may come first.) + + arrow-rs only: the planner loop is reader-agnostic, but pyarrow's scanner + reads a fragment's row groups concurrently and surfaces the row-group-1 + error before row group 0's batch, so on that arm no table ever precedes + the error and the shape cannot be built. + """ + from ray.data.block import BlockAccessor + + good_rows = 32 + n = 1000 + path = tmp_path / "two_groups.parquet" + table = pa.table( + { + "id": pa.array(np.arange(n, dtype=np.int64)), + "s": pa.array(["v" * 200] * n), + } + ) + # Row groups: [0, 32) and [32, 1000) -- write_table splits at the given + # size, so a 32-row cap yields many groups; write the two explicitly. + with pq.ParquetWriter(str(path), table.schema, compression="snappy") as w: + w.write_table(table.slice(0, good_rows)) + w.write_table(table.slice(good_rows)) + md = pq.ParquetFile(str(path)).metadata + assert md.num_row_groups == 2 and md.row_group(0).num_rows == good_rows + col = md.row_group(1).column(0) + first_page = col.dictionary_page_offset or col.data_page_offset + raw = bytearray(path.read_bytes()) + raw[first_page : first_page + 256] = bytes(256) + path.write_bytes(bytes(raw)) + + ctx = restore_ctx + ctx.use_arrow_rs_parquet_reader = True + old_errored = ctx.max_errored_blocks + old_target = ctx.target_max_block_size + try: + ctx.max_errored_blocks = 1 + rg0_bytes = BlockAccessor.for_block(table.slice(0, good_rows)).size_bytes() + ctx.target_max_block_size = int(rg0_bytes * 0.9) + mds = ray.data.read_parquet(str(path)).materialize() + assert mds._raw_stats().extra_metrics["num_tasks_submitted"] == 1 + # Old planner: the held table dies with the exception -> 0 rows. + assert mds.to_pandas()["id"].tolist() == list(range(good_rows)) + finally: + ctx.max_errored_blocks = old_errored + ctx.target_max_block_size = old_target diff --git a/python/ray/data/tests/datasource/test_parquet.py b/python/ray/data/tests/datasource/test_parquet.py index 4798d0bd85e7..97d3e4d1fd8d 100644 --- a/python/ray/data/tests/datasource/test_parquet.py +++ b/python/ray/data/tests/datasource/test_parquet.py @@ -1482,7 +1482,9 @@ def map_batches(batch): # tests should only be carefully reordered to retain this invariant! -def test_parquet_read_spread(ray_start_cluster, tmp_path, restore_data_context): +def test_parquet_read_spread( + ray_start_cluster, tmp_path, restore_data_context, monkeypatch +): ray.shutdown() cluster = ray_start_cluster cluster.add_node( @@ -1516,8 +1518,11 @@ def get_node_id(): df2.to_parquet(path2) # Minimize the block size to prevent Ray Data from reading multiple fragments in a - # single task. + # single task. On the V2 footer path the packer uses + # RAY_DATA_PARQUET_BIN_PACKING_BYTES (not target_max_block_size), so pin that + # too or both files collapse into one read task on one node. ray.data.DataContext.get_current().target_max_block_size = 1 + monkeypatch.setenv("RAY_DATA_PARQUET_BIN_PACKING_BYTES", "1") ds = ray.data.read_parquet(data_path) # Force reads. diff --git a/python/ray/data/tests/datasource/test_read_parquet_v2.py b/python/ray/data/tests/datasource/test_read_parquet_v2.py index 42547e88f5a3..0c2a9e7bffbf 100644 --- a/python/ray/data/tests/datasource/test_read_parquet_v2.py +++ b/python/ray/data/tests/datasource/test_read_parquet_v2.py @@ -11,8 +11,8 @@ import pytest import ray -from ray.data._internal.datasource_v2.partitioners.round_robin_partitioner import ( - RoundRobinPartitioner, +from ray.data._internal.datasource_v2.listing.footer_file_indexer import ( + FooterFileIndexer, ) from ray.data._internal.datasource_v2.scanners.parquet_scanner import ParquetScanner from ray.data._internal.logical.operators import ListFiles, ReadFiles @@ -125,22 +125,156 @@ def test_read_parquet_v2_columns_with_include_paths_preserves_path( assert [expr.name for expr in dag.exprs] == ["a", "path"] -def test_read_parquet_v2_override_num_blocks_drives_partitioner(tmp_path, restore_ctx): +def test_read_parquet_v2_uses_footer_indexer_without_partitioner(tmp_path, restore_ctx): _write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3]})) restore_ctx.use_datasource_v2 = True original = restore_ctx.read_op_min_num_blocks ds = ray.data.read_parquet(str(tmp_path), override_num_blocks=7) - # The override should drive the ListFiles partitioner's bucket count - # for this read only — the global DataContext must not be mutated. + # Parquet V2 uses the footer indexer, which bin-packs read units itself, so + # ``ListFiles`` carries no size-balancing partitioner (``override_num_blocks`` + # no longer drives a partitioner bucket count for Parquet). The global + # DataContext must not be mutated. list_files_op = ds._logical_plan.dag.input_dependencies[0] assert isinstance(list_files_op, ListFiles) - assert isinstance(list_files_op.file_partitioner, RoundRobinPartitioner) - assert list_files_op.file_partitioner.num_buckets == 7 + assert isinstance(list_files_op.file_indexer, FooterFileIndexer) + assert list_files_op.file_partitioner is None assert restore_ctx.read_op_min_num_blocks == original +def _write_row_groups(path, *, num_files, rows_per_file, row_group_size): + """Write ``num_files`` parquet files, each split into several row groups.""" + for i in range(num_files): + pq.write_table( + pa.table({"a": list(range(rows_per_file))}), + str(path / f"f{i}.parquet"), + row_group_size=row_group_size, + ) + + +def _optimized_count_plan(ds): + """The plan ``Dataset.count()`` executes, without executing it. + + ``Dataset.count()`` builds this internally and never exposes it, so tests + reconstruct it to assert on the optimizer's output. + """ + from ray.data._internal.logical.interfaces import LogicalPlan + from ray.data._internal.logical.operators.count_operator import Count + from ray.data._internal.logical.operators.map_operator import Project + from ray.data._internal.logical.optimizers import LogicalOptimizer + + count_op = Count( + input_dependencies=[ + Project(exprs=[], input_dependencies=[ds._logical_plan.dag]) + ] + ) + return LogicalOptimizer().optimize(LogicalPlan(count_op, ds.context)) + + +def _walk(op): + yield op + for child in op.input_dependencies: + yield from _walk(child) + + +def test_count_pushdown_replaces_footer_indexer(tmp_path, restore_ctx): + """``count()`` must not run the footer indexer. + + ``FooterFileIndexer`` subclasses ``NonSamplingFileIndexer`` but overrides + ``list_files``, so it footer-sweeps every file during listing and bin-packs + them into read units -- for a zero-column count projection every row group + measures 0 bytes, collapsing the whole dataset into a single manifest and + therefore a single count task. + """ + from ray.data._internal.datasource_v2.chunkers.file_chunker import WholeFileChunker + from ray.data._internal.datasource_v2.listing.file_indexer import ( + NonSamplingFileIndexer, + ) + from ray.data._internal.logical.operators.map_operator import MapBatches + + _write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3]})) + + restore_ctx.use_datasource_v2 = True + ds = ray.data.read_parquet(str(tmp_path)) + list_files_before = ds._logical_plan.dag.input_dependencies[0] + assert isinstance(list_files_before, ListFiles) + assert isinstance(list_files_before.file_indexer, FooterFileIndexer) + + dag = _optimized_count_plan(ds).dag + + assert isinstance(dag, MapBatches) + assert not any(isinstance(op, ReadFiles) for op in _walk(dag)) + + (list_files_op,) = [op for op in _walk(dag) if isinstance(op, ListFiles)] + # Exact type: an ``isinstance`` check is precisely what let the footer + # indexer through before. + assert type(list_files_op.file_indexer) is NonSamplingFileIndexer + assert isinstance(list_files_op.file_indexer.file_chunker, WholeFileChunker) + assert list_files_op.file_partitioner is None + + +@pytest.mark.parametrize( + "read_kwargs", [{}, {"include_paths": True}], ids=["plain", "include_paths"] +) +def test_count_pushdown_preserves_list_files_fields(tmp_path, restore_ctx, read_kwargs): + _write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3]})) + + restore_ctx.use_datasource_v2 = True + ds = ray.data.read_parquet(str(tmp_path), **read_kwargs) + original = ds._logical_plan.dag.input_dependencies[0] + assert isinstance(original, ListFiles) + + (rebuilt,) = [ + op for op in _walk(_optimized_count_plan(ds).dag) if isinstance(op, ListFiles) + ] + + assert rebuilt.paths == original.paths + assert rebuilt.source_paths == original.source_paths + assert rebuilt.file_extensions == original.file_extensions + assert rebuilt.partition_filter is original.partition_filter + + +@pytest.mark.parametrize("case", ["predicate", "limit"]) +def test_count_pushdown_declines_row_reducing_reads(tmp_path, restore_ctx, case): + """A row-reducing pushdown makes footer ``num_rows`` an overcount.""" + from ray.data.expressions import col + + _write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3, 4]})) + + restore_ctx.use_datasource_v2 = True + ds = ray.data.read_parquet(str(tmp_path)) + ds = ds.filter(expr=col("a") > 2) if case == "predicate" else ds.limit(2) + + assert any(isinstance(op, ReadFiles) for op in _walk(_optimized_count_plan(ds).dag)) + + +@pytest.mark.parametrize( + "num_files,rows_per_file,row_group_size", + [(1, 10, None), (3, 100, 10), (4, 1000, 250)], + ids=["single_file", "multi_file_many_row_groups", "multi_file_large"], +) +def test_count_matches_rows( + tmp_path, restore_ctx, num_files, rows_per_file, row_group_size +): + """End-to-end count correctness, including multi-row-group files. + + Regression guard for over-counting: the bin packer emits one manifest row + per path *per bin*, so a file whose row groups span bins would otherwise be + counted once per bin. + """ + _write_row_groups( + tmp_path, + num_files=num_files, + rows_per_file=rows_per_file, + row_group_size=row_group_size, + ) + + restore_ctx.use_datasource_v2 = True + + assert ray.data.read_parquet(str(tmp_path)).count() == num_files * rows_per_file + + def test_read_parquet_v2_filter_raises(tmp_path, restore_ctx): import pyarrow.dataset as pds diff --git a/python/ray/data/tests/doctest_pytest_plugin.py b/python/ray/data/tests/doctest_pytest_plugin.py index 123daf523379..ef07f7524941 100644 --- a/python/ray/data/tests/doctest_pytest_plugin.py +++ b/python/ray/data/tests/doctest_pytest_plugin.py @@ -1,8 +1,16 @@ """This file is injected for Ray Data doctest targets.""" +import os + import pytest import ray +# Keep the footer-reader pool tiny: doctests read small Parquet fixtures, and +# the default 32-actor pool can trip Ray's "too many worker processes" warning, +# which pollutes Sphinx ``testoutput`` expectations. Mirrored in +# python/ray/data/test.bzl for bazel doctest targets. +os.environ.setdefault("RAY_DATA_PARQUET_FOOTER_NUM_ACTORS", "1") + @pytest.fixture(autouse=True, scope="module") def shutdown_ray(): diff --git a/python/ray/data/tests/test_execution_optimizer_limit_pushdown.py b/python/ray/data/tests/test_execution_optimizer_limit_pushdown.py index 603580089a86..6514ca0880f6 100644 --- a/python/ray/data/tests/test_execution_optimizer_limit_pushdown.py +++ b/python/ray/data/tests/test_execution_optimizer_limit_pushdown.py @@ -7,7 +7,7 @@ import ray from ray.data import Dataset from ray.data._internal.logical.interfaces import LogicalOperator, Plan -from ray.data._internal.logical.operators import Download, Limit +from ray.data._internal.logical.operators import Download, Limit, ListFiles, ReadFiles from ray.data._internal.logical.rules.limit_pushdown import LimitPushdownRule from ray.data._internal.util import rows_same from ray.data.block import BlockMetadata @@ -249,8 +249,8 @@ def test_limit_pushdown_correctness(ray_start_regular_shared_2_cpus): # Test 6: Complex chain with both safe operations (should all get limit pushed) ds = ( ray.data.range(100) - .select_columns(["id"]) # Project - could be safe if it was the immediate input - .map(lambda x: {"id": x["id"] + 1}) # MapRows - NOT safe, stops pushdown + .select_columns(["id"]) # Project - row-preserving, limit pushes through + .map(lambda x: {"id": x["id"] + 1}) # MapRows - row-preserving, pushes through .limit(3) ) result = ds.take_all() @@ -265,6 +265,37 @@ def test_limit_pushdown_correctness(ray_start_regular_shared_2_cpus): ) +def test_limit_pushdown_into_read_files_scanner( + tmp_path, ray_start_regular_shared_2_cpus, restore_data_context +): + """A Limit sitting directly on a V2 ``ReadFiles`` still pushes a limit in. + + The pushed limit reaches the scanner and, through + ``DeriveListFilesPushdown``, the upstream ``ListFiles``, which is what lets + a footer-based indexer stop listing early. None of that shows up in + ``dag_str``, so assert on the operators themselves. + """ + ray.data.DataContext.get_current().use_datasource_v2 = True + + for i in range(3): + pd.DataFrame({"id": range(i * 10, i * 10 + 10)}).to_parquet( + tmp_path / f"part{i}.parquet" + ) + + ds = ray.data.read_parquet(str(tmp_path)).limit(5) + assert len(ds.take_all()) == 5 + + # The Limit stays on top for exact enforcement. + limit_op = ds._logical_plan.dag + assert isinstance(limit_op, Limit), limit_op.dag_str + (read_files,) = limit_op.input_dependencies + assert isinstance(read_files, ReadFiles), read_files + assert read_files.scanner.pushed_limit() == 5 + (list_files,) = read_files.input_dependencies + assert isinstance(list_files, ListFiles), list_files + assert list_files.limit == 5 + + def test_limit_pushdown_scan_efficiency(ray_start_regular_shared_2_cpus): """Test that limit pushdown scans fewer rows from the data source.""" diff --git a/python/ray/data/tests/test_op_runtime_metrics.py b/python/ray/data/tests/test_op_runtime_metrics.py index f73c3ef30766..0f56a8938d01 100644 --- a/python/ray/data/tests/test_op_runtime_metrics.py +++ b/python/ray/data/tests/test_op_runtime_metrics.py @@ -46,6 +46,96 @@ def test_average_max_uss_per_task(): assert metrics.average_max_uss_per_task == 200 # (100 + 300) / 2 +def test_read_files_task_stats_distributions(): + """ReadFilesTaskStats entries on TaskExecWorkerStats.custom_op_stats are + folded into the read_task_* per-task distributions (bytes/wall summed + across entries, peak_batch maxed).""" + from ray.data.block import ReadFilesTaskStats + + op = MagicMock() + op.data_context.enable_get_object_locations_for_metrics = False + metrics = OpRuntimeMetrics(op) + assert metrics.average_decoded_bytes_per_read_task is None + + input_bundle = RefBundle([], owns_blocks=False, schema=None) + + metrics.on_task_submitted(0, input_bundle) + metrics.on_task_finished( + 0, + None, + TaskExecWorkerStats( + task_wall_time_s=1.0, + custom_op_stats=[ + ReadFilesTaskStats( + decode_wall_s=0.5, + decoded_bytes=100, + decoded_batches=2, + decoded_rows=10, + peak_batch_bytes=60, + manifests=1, + trim_wall_s=0.10, + yield_wall_s=1.5, + first_table_wall_s=0.7, + ), + # A hypothetical second reporting transform in a fused task: + # bytes/wall sum, peak_batch maxes. + ReadFilesTaskStats( + decode_wall_s=0.25, + decoded_bytes=50, + decoded_batches=1, + decoded_rows=5, + peak_batch_bytes=50, + manifests=1, + trim_wall_s=0.05, + yield_wall_s=0.5, + ), + ], + ), + TaskExecDriverStats(task_output_backpressure_s=0), + ) + metrics.on_task_submitted(1, input_bundle) + metrics.on_task_finished( + 1, + None, + TaskExecWorkerStats( + task_wall_time_s=1.0, + custom_op_stats=[ + ReadFilesTaskStats( + decode_wall_s=0.25, + decoded_bytes=50, + decoded_batches=1, + decoded_rows=5, + peak_batch_bytes=50, + manifests=1, + ) + ], + ), + TaskExecDriverStats(task_output_backpressure_s=0), + ) + # Tasks without read stats leave the distributions untouched. + metrics.on_task_submitted(2, input_bundle) + metrics.on_task_finished( + 2, + None, + TaskExecWorkerStats(task_wall_time_s=1.0), + TaskExecDriverStats(task_output_backpressure_s=0), + ) + + assert metrics.read_task_decoded_bytes.num_samples == 2 + assert metrics.average_decoded_bytes_per_read_task == 100 # (150 + 50) / 2 + assert metrics.read_task_decoded_bytes.max == 150 + assert metrics.read_task_decode_wall_s.max == 0.75 + assert metrics.read_task_peak_batch_bytes.max == 60 + # trim_wall_s sums across entries like the other wall counter; the second + # task left it at the 0.0 default (a reader without a finalizer). + assert metrics.read_task_trim_wall_s.num_samples == 2 + assert abs(metrics.read_task_trim_wall_s.max - 0.15) < 1e-9 + assert metrics.read_task_trim_wall_s.min == 0.0 + assert metrics.read_task_yield_wall_s.num_samples == 2 + assert abs(metrics.read_task_yield_wall_s.max - 2.0) < 1e-9 + assert abs(metrics.read_task_first_table_wall_s.max - 0.7) < 1e-9 + + def test_task_completion_time_histogram(): """Test task completion time histogram bucket assignment and counting.""" op = MagicMock() diff --git a/python/ray/data/tests/test_stats.py b/python/ray/data/tests/test_stats.py index ccc738d47e85..0964dad3d946 100644 --- a/python/ray/data/tests/test_stats.py +++ b/python/ray/data/tests/test_stats.py @@ -354,6 +354,10 @@ def gen_expected_metrics( "'op_task_duration_stats': {'num_samples': N, 'mean': N, 'variance': N, 'min': N, 'max': N, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P}", "'max_uss_bytes': H", "'average_max_uss_per_task': H", + "'max_uss_per_task': H", + "'max_rss_bytes': H", + "'average_max_rss_per_task': H", + "'max_rss_per_task': H", "'num_inputs_received': N", "'num_row_inputs_received': N", "'bytes_inputs_received': N", @@ -444,6 +448,10 @@ def gen_expected_metrics( "'op_task_duration_stats': {'num_samples': Z, 'mean': Z, 'variance': Z, 'min': None, 'max': None, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P}", "'max_uss_bytes': H", "'average_max_uss_per_task': H", + "'max_uss_per_task': H", + "'max_rss_bytes': H", + "'average_max_rss_per_task': H", + "'max_rss_per_task': H", "'num_inputs_received': N", "'num_row_inputs_received': N", "'bytes_inputs_received': N", @@ -634,7 +642,8 @@ def canonicalize( canonicalized_stats = re.sub("\t", " ", canonicalized_stats) canonicalized_stats = re.sub( - r"(average_max_uss_per_task:|'average_max_uss_per_task':) (?:N|Z|None)\b", + r"((?:average_)?max_[ur]ss_per_task:|'(?:average_)?max_[ur]ss_per_task':)" + r" (?:N|Z|None)\b", r"\g<1> H", canonicalized_stats, ) @@ -645,10 +654,11 @@ def canonicalize( r"\g<1>P", canonicalized_stats, ) - # max_uss_bytes DistributionTracker may have 0 or N samples depending on - # platform (USS measurement only available on Linux). Normalize entire dict. + # max_uss_bytes/max_rss_bytes DistributionTrackers may have 0 or N samples + # depending on platform (USS/RSS measurement only available on Linux). + # Normalize entire dict. canonicalized_stats = re.sub( - r"(max_uss_bytes['\s:]+)\{[^}]+\}", + r"(max_[ur]ss_bytes['\s:]+)\{[^}]+\}", r"\g<1>H", canonicalized_stats, ) @@ -931,6 +941,10 @@ def test_dataset__repr__(ray_start_regular_shared, restore_data_context): " op_task_duration_stats: {'num_samples': N, 'mean': N, 'variance': N, 'min': N, 'max': N, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P},\n" " max_uss_bytes: H,\n" " average_max_uss_per_task: H,\n" + " max_uss_per_task: H,\n" + " max_rss_bytes: H,\n" + " average_max_rss_per_task: H,\n" + " max_rss_per_task: H,\n" " num_inputs_received: N,\n" " num_row_inputs_received: N,\n" " bytes_inputs_received: N,\n" @@ -1096,6 +1110,10 @@ def check_stats(): " op_task_duration_stats: {'num_samples': N, 'mean': N, 'variance': N, 'min': N, 'max': N, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P},\n" " max_uss_bytes: H,\n" " average_max_uss_per_task: H,\n" + " max_uss_per_task: H,\n" + " max_rss_bytes: H,\n" + " average_max_rss_per_task: H,\n" + " max_rss_per_task: H,\n" " num_inputs_received: N,\n" " num_row_inputs_received: N,\n" " bytes_inputs_received: N,\n" @@ -1214,6 +1232,10 @@ def check_stats(): " op_task_duration_stats: {'num_samples': N, 'mean': N, 'variance': N, 'min': N, 'max': N, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P, 'pN': P},\n" " max_uss_bytes: H,\n" " average_max_uss_per_task: H,\n" + " max_uss_per_task: H,\n" + " max_rss_bytes: H,\n" + " average_max_rss_per_task: H,\n" + " max_rss_per_task: H,\n" " num_inputs_received: N,\n" " num_row_inputs_received: N,\n" " bytes_inputs_received: N,\n" diff --git a/python/ray/data/tests/unit/datasource_v2/test_derive_list_files_pushdown.py b/python/ray/data/tests/unit/datasource_v2/test_derive_list_files_pushdown.py new file mode 100644 index 000000000000..939bed4c3f0b --- /dev/null +++ b/python/ray/data/tests/unit/datasource_v2/test_derive_list_files_pushdown.py @@ -0,0 +1,255 @@ +"""Unit tests for :class:`DeriveListFilesPushdown`. + +``ListFiles`` prunes row groups, sizes columns, and stops listing early using +the constraints it carries. Those constraints are only sound while they are no +stronger than what the downstream ``ReadFiles`` scanner actually applies -- a +predicate ``ListFiles`` prunes by but the reader never evaluates drops rows +with no error. These tests pin that invariant, including for plan shapes no +current rule produces but a future one could. +""" +from dataclasses import replace +from pathlib import Path +from typing import List + +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from ray.data._internal.datasource_v2.listing.file_indexer import ( + NonSamplingFileIndexer, +) +from ray.data._internal.datasource_v2.listing.listing_utils import sample_files +from ray.data._internal.datasource_v2.parquet_datasource_v2 import ( + ParquetDatasourceV2, +) +from ray.data._internal.datasource_v2.scanners.arrow_file_scanner import ( + ArrowFileScanner, +) +from ray.data._internal.logical.interfaces import ( + LogicalOperator, + LogicalPlan, + Plan, + Rule, +) +from ray.data._internal.logical.operators import ListFiles, ReadFiles +from ray.data._internal.logical.operators.map_operator import MapBatches +from ray.data._internal.logical.optimizers import LogicalOptimizer, get_logical_ruleset +from ray.data._internal.logical.rules.derive_list_files_pushdown import ( + DeriveListFilesPushdown, +) +from ray.data.context import DataContext +from ray.data.expressions import col + + +def _mk_read_files(tmp_path: Path) -> ReadFiles: + """A minimal ``ListFiles -> ReadFiles`` chain over one Parquet file.""" + f = tmp_path / "data.parquet" + pq.write_table(pa.table({"a": [1, 2, 3], "b": ["x", "y", "z"]}), str(f)) + + datasource = ParquetDatasourceV2([str(f)]) + indexer = NonSamplingFileIndexer(ignore_missing_paths=False) + sample = sample_files(indexer, datasource.paths, datasource.filesystem) + schema = datasource.infer_schema(sample) + + list_files_op = ListFiles( + paths=list(datasource.paths), + file_indexer=indexer, + filesystem=datasource.filesystem, + source_paths=list(datasource.paths), + file_extensions=datasource.file_extensions, + ) + return ReadFiles( + datasource_name=datasource.name, + scanner=datasource.create_scanner(schema=schema), + schema=schema, + parallelism=-1, + input_dependencies=[list_files_op], + ) + + +def _apply(dag: LogicalOperator) -> LogicalPlan: + plan = LogicalPlan(dag=dag, context=DataContext.get_current()) + return DeriveListFilesPushdown().apply(plan) + + +def _list_files_of(plan: Plan) -> ListFiles: + (list_files,) = [ + op for op in plan.dag.post_order_iter() if isinstance(op, ListFiles) + ] + return list_files + + +def _source_list_files(read_files: ReadFiles) -> ListFiles: + """The ``ListFiles`` feeding ``read_files``, typed as such. + + ``input_dependencies`` is declared as plain ``LogicalOperator``. + """ + (list_files,) = read_files.input_dependencies + assert isinstance(list_files, ListFiles), list_files + return list_files + + +def _scanner_of(read_files: ReadFiles) -> ArrowFileScanner: + """The scanner of ``read_files``, typed as the pushdown-capable subclass. + + ``ReadFiles.scanner`` is declared as the base ``Scanner``, which carries + none of the ``Supports*`` pushdown methods these tests drive. + """ + scanner = read_files.scanner + assert isinstance(scanner, ArrowFileScanner), scanner + return scanner + + +def test_derives_state_the_scanner_accepted(tmp_path): + read_files = _mk_read_files(tmp_path) + predicate = col("a") > 2 + scanner, _residual = _scanner_of(read_files).push_filters(predicate) + scanner = scanner.prune_columns(["a"]).push_limit(5) + read_files = replace(read_files, scanner=scanner) + + list_files = _list_files_of(_apply(read_files)) + + assert list_files.predicate is predicate + assert list_files.projected_columns == ["a"] + assert list_files.limit == 5 + + +def test_no_pushdown_leaves_list_files_unconstrained(tmp_path): + list_files = _list_files_of(_apply(_mk_read_files(tmp_path))) + + assert list_files.predicate is None + assert list_files.projected_columns is None + assert list_files.limit is None + + +@pytest.mark.parametrize( + "stale", + [ + {"predicate": col("a") > 2}, + {"projected_columns": ["a"]}, + {"limit": 1}, + {"predicate": col("a") > 2, "projected_columns": ["a"], "limit": 1}, + ], + ids=["predicate", "columns", "limit", "all"], +) +def test_state_the_scanner_does_not_carry_is_cleared(tmp_path, stale): + """The failure the rule exists to prevent. + + A ``ListFiles`` carrying constraints its ``ReadFiles`` does not apply -- + e.g. a rewrite dropped or weakened the scanner's predicate -- would prune + row groups nothing downstream re-checks. + """ + read_files = _mk_read_files(tmp_path) + read_files = replace( + read_files, + input_dependencies=[replace(_source_list_files(read_files), **stale)], + ) + + list_files = _list_files_of(_apply(read_files)) + + assert list_files.predicate is None + assert list_files.projected_columns is None + assert list_files.limit is None + + +def test_state_is_cleared_when_consumer_is_not_read_files(tmp_path): + """``PushdownCountFiles`` rewrites ``ReadFiles`` out of the plan entirely.""" + read_files = _mk_read_files(tmp_path) + list_files = replace( + _source_list_files(read_files), predicate=col("a") > 2, limit=1 + ) + count_rows = MapBatches( + fn=lambda batch: batch, + input_dependencies=[list_files], + batch_format="pyarrow", + can_modify_num_rows=True, + ) + + derived = _list_files_of(_apply(count_rows)) + + assert derived.predicate is None + assert derived.limit is None + + +def test_bare_list_files_root_is_cleared(tmp_path): + read_files = _mk_read_files(tmp_path) + list_files = replace(_source_list_files(read_files), predicate=col("a") > 2) + + assert _list_files_of(_apply(list_files)).predicate is None + + +class _WeakenScannerPredicate(Rule): + """Stand-in for a future rule that rewrites the scanner's predicate.""" + + def apply(self, plan: LogicalPlan) -> LogicalPlan: # pyrefly: ignore[bad-override] + def transform(node: LogicalOperator) -> LogicalOperator: + if isinstance(node, ReadFiles): + scanner = _scanner_of(node) + if scanner.pushed_predicate() is not None: + return replace(node, scanner=replace(scanner, predicate=None)) + return node + + return LogicalPlan( + dag=plan.dag._apply_transform(transform), context=plan.context + ) + + +@pytest.fixture +def weakening_rule(): + ruleset = get_logical_ruleset() + ruleset.add(_WeakenScannerPredicate) + try: + yield + finally: + ruleset.remove(_WeakenScannerPredicate) + + +def test_optimizer_does_not_strand_a_predicate_a_later_rule_dropped( + tmp_path, weakening_rule +): + """A rule that drops the scanner's predicate must weaken listing too. + + The rule runs after the pushdown rules and knows nothing about + ``ListFiles``; the invariant has to hold anyway. + """ + from ray.data._internal.logical.operators import Filter + + read_files = _mk_read_files(tmp_path) + dag = Filter(predicate_expr=col("a") > 2, input_dependencies=[read_files]) + + optimized = LogicalOptimizer().optimize( + LogicalPlan(dag=dag, context=DataContext.get_current()) + ) + + read_files_ops: List[ReadFiles] = [ + op for op in optimized.dag.post_order_iter() if isinstance(op, ReadFiles) + ] + (scanner_predicate,) = [_scanner_of(op).pushed_predicate() for op in read_files_ops] + assert scanner_predicate is None + assert _list_files_of(optimized).predicate is None + + +def test_optimizer_keeps_list_files_in_sync_with_the_scanner(tmp_path): + """Without the weakening rule, the pushed predicate does reach listing.""" + from ray.data._internal.logical.operators import Filter + + read_files = _mk_read_files(tmp_path) + dag = Filter(predicate_expr=col("a") > 2, input_dependencies=[read_files]) + + optimized = LogicalOptimizer().optimize( + LogicalPlan(dag=dag, context=DataContext.get_current()) + ) + + (scanner_predicate,) = [ + _scanner_of(op).pushed_predicate() + for op in optimized.dag.post_order_iter() + if isinstance(op, ReadFiles) + ] + assert scanner_predicate is not None + assert _list_files_of(optimized).predicate is scanner_predicate + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main(["-v", __file__])) diff --git a/python/ray/data/tests/unit/datasource_v2/test_file_chunker.py b/python/ray/data/tests/unit/datasource_v2/test_file_chunker.py index 0d83f1b8a102..8ae1b288ab65 100644 --- a/python/ray/data/tests/unit/datasource_v2/test_file_chunker.py +++ b/python/ray/data/tests/unit/datasource_v2/test_file_chunker.py @@ -7,8 +7,7 @@ ChunkMetadata, LineDelimitedFileChunker, LineDelimitedFileChunkMetadata, - ParquetFileChunker, - ParquetFileChunkMetadata, + ParquetRowGroupChunkMetadata, WholeFileChunker, create_chunk_metadata, ) @@ -17,22 +16,30 @@ class TestCreateChunkMetadata: def test_validates_missing_keys(self): with pytest.raises(ValueError, match="Missing required keys"): - create_chunk_metadata(ParquetFileChunkMetadata, chunk_idx=0) + create_chunk_metadata(ParquetRowGroupChunkMetadata, row_group_ids=(0,)) def test_validates_unexpected_keys(self): with pytest.raises(ValueError, match="Unexpected keys"): create_chunk_metadata( - ParquetFileChunkMetadata, - chunk_idx=0, - total_num_chunks=1, + ParquetRowGroupChunkMetadata, + row_group_ids=(0,), + num_rows=1, + uncompressed_size=10, extra_field="boom", ) def test_returns_dict_with_keys(self): md = create_chunk_metadata( - ParquetFileChunkMetadata, chunk_idx=2, total_num_chunks=5 + ParquetRowGroupChunkMetadata, + row_group_ids=(0, 1), + num_rows=5, + uncompressed_size=10, ) - assert md == {"chunk_idx": 2, "total_num_chunks": 5} + assert md == { + "row_group_ids": (0, 1), + "num_rows": 5, + "uncompressed_size": 10, + } class TestWholeFileChunker: @@ -63,74 +70,20 @@ def test_compressed_file_yields_whole(self): assert chunks == [(None, 1024)] -class TestParquetFileChunker: - def test_small_file_yields_whole(self): - chunker = ParquetFileChunker(target_chunk_size=256 * 1024 * 1024) - chunks = list( - chunker.generate_chunk_metadatas("data.parquet", 100 * 1024 * 1024) - ) - assert chunks == [(None, 100 * 1024 * 1024)] - - def test_at_target_yields_whole(self): - chunker = ParquetFileChunker(target_chunk_size=256 * 1024 * 1024) - chunks = list( - chunker.generate_chunk_metadatas("data.parquet", 256 * 1024 * 1024) - ) - assert chunks == [(None, 256 * 1024 * 1024)] - - @pytest.mark.parametrize( - "file_size,expected_num_chunks", - [ - (257 * 1024 * 1024, 2), - (300 * 1024 * 1024, 2), - (512 * 1024 * 1024, 2), - (600 * 1024 * 1024, 3), - (1024 * 1024 * 1024, 4), - ], - ) - def test_large_files_produce_chunks(self, file_size, expected_num_chunks): - target_chunk_size = 256 * 1024 * 1024 - chunker = ParquetFileChunker(target_chunk_size=target_chunk_size) - chunks = list(chunker.generate_chunk_metadatas("data.parquet", file_size)) - assert len(chunks) == expected_num_chunks - total_size = 0 - for i, (md, chunk_size) in enumerate(chunks): - assert isinstance(md, dict) - assert md["chunk_idx"] == i - assert md["total_num_chunks"] == expected_num_chunks - if i < expected_num_chunks - 1: - assert chunk_size == target_chunk_size - else: - assert chunk_size == file_size - target_chunk_size * i - total_size += chunk_size - assert total_size == file_size - - def test_default_target_chunk_size_from_context(self, restore_data_context): - from ray.data.context import DataContext - - DataContext.get_current().parquet_chunker_target_chunk_size = 1024 - chunker = ParquetFileChunker() - assert chunker._target_chunk_size == 1024 - - def test_ctor_arg_takes_precedence_over_context(self, restore_data_context): - from ray.data.context import DataContext - - DataContext.get_current().parquet_chunker_target_chunk_size = 1024 - chunker = ParquetFileChunker(target_chunk_size=2048) - assert chunker._target_chunk_size == 2048 - - def test_chunk_metadata_subclasses_are_typeddicts(): # Ensures the subclasses don't accidentally inherit unrelated keys. pmd: ChunkMetadata = create_chunk_metadata( - ParquetFileChunkMetadata, chunk_idx=0, total_num_chunks=1 + ParquetRowGroupChunkMetadata, + row_group_ids=(0,), + num_rows=1, + uncompressed_size=10, ) lmd: ChunkMetadata = create_chunk_metadata( LineDelimitedFileChunkMetadata, chunk_byte_start_idx=0, chunk_byte_end_idx=10, ) - assert set(pmd.keys()) == {"chunk_idx", "total_num_chunks"} + assert set(pmd.keys()) == {"row_group_ids", "num_rows", "uncompressed_size"} assert set(lmd.keys()) == {"chunk_byte_start_idx", "chunk_byte_end_idx"} diff --git a/python/ray/data/tests/unit/datasource_v2/test_file_indexer.py b/python/ray/data/tests/unit/datasource_v2/test_file_indexer.py index e8e0dc328f2d..ead8c41475b2 100644 --- a/python/ray/data/tests/unit/datasource_v2/test_file_indexer.py +++ b/python/ray/data/tests/unit/datasource_v2/test_file_indexer.py @@ -6,7 +6,6 @@ from ray.data._internal.datasource_v2.chunkers.file_chunker import ( LineDelimitedFileChunker, - ParquetFileChunker, WholeFileChunker, ) from ray.data._internal.datasource_v2.listing.file_indexer import ( @@ -211,7 +210,7 @@ def test_default_uses_whole_file_chunker(self): assert isinstance(indexer.file_chunker, WholeFileChunker) def test_explicit_chunker_is_exposed(self): - chunker = ParquetFileChunker(target_chunk_size=1024) + chunker = LineDelimitedFileChunker() indexer = NonSamplingFileIndexer( ignore_missing_paths=False, file_chunker=chunker ) @@ -229,29 +228,6 @@ def test_whole_file_chunker_yields_none_chunk_metadata(self, tmp_path): assert list(manifest.file_chunk_metadatas) == [None] assert list(manifest.file_sizes) == [100] - def test_parquet_chunker_splits_large_file_into_many_chunks(self, tmp_path): - # Write a "Parquet" file by name only — the chunker doesn't open it. - (tmp_path / "big.parquet").write_bytes(b"x" * 10_000) - chunker = ParquetFileChunker(target_chunk_size=1024) - indexer = NonSamplingFileIndexer( - ignore_missing_paths=False, - num_workers=1, - file_chunker=chunker, - ) - fs = LocalFileSystem() - manifests = list(indexer.list_files(pa.array([str(tmp_path)]), filesystem=fs)) - rows = [] - for m in manifests: - for path, size, md in zip(m.paths, m.file_sizes, m.file_chunk_metadatas): - rows.append((str(path), int(size), md)) - - # 10000 bytes / 1024 target chunk size -> 10 chunks (ceil). - assert len(rows) == 10 - for i, (_, _, md) in enumerate(rows): - assert md is not None - assert md["chunk_idx"] == i - assert md["total_num_chunks"] == 10 - def test_line_delimited_chunker_byte_ranges(self, tmp_path): (tmp_path / "a.jsonl").write_bytes(b"x" * 10_000) chunker = LineDelimitedFileChunker() @@ -275,6 +251,67 @@ def test_line_delimited_chunker_byte_ranges(self, tmp_path): assert rows[-1][2]["chunk_byte_end_idx"] == 10_000 +class TestAsWholeFileIndexer: + """``as_whole_file_indexer`` must downgrade to a plain per-file lister. + + ``PushdownCountFiles`` relies on this to count each file exactly once + without triggering a metadata-aware subclass's listing strategy. + """ + + def test_returns_base_type_from_metadata_aware_subclass(self): + from ray.data._internal.datasource_v2.listing.footer_file_indexer import ( + FooterFileIndexer, + ) + + downgraded = FooterFileIndexer( + ignore_missing_paths=False + ).as_whole_file_indexer() + + # Exact type, not isinstance: FooterFileIndexer subclasses + # NonSamplingFileIndexer but overrides list_files, so an isinstance + # check here would not catch a regression. + assert type(downgraded) is NonSamplingFileIndexer + + @pytest.mark.parametrize( + "ignore_missing_paths,num_workers,max_paths_per_output", + [(False, 1, 10), (True, 4, 1000)], + ) + def test_carries_over_traversal_config( + self, ignore_missing_paths, num_workers, max_paths_per_output + ): + source = NonSamplingFileIndexer( + ignore_missing_paths=ignore_missing_paths, + num_workers=num_workers, + max_paths_per_output=max_paths_per_output, + ) + + downgraded = source.as_whole_file_indexer() + + assert downgraded._ignore_missing_paths == ignore_missing_paths + assert downgraded._num_workers == num_workers + assert downgraded._max_paths_per_output == max_paths_per_output + # Derived in __init__, so rebuilding must recompute it rather than + # copying a stale value. + assert downgraded._queue_size_per_thread == max_paths_per_output * 4 + + def test_always_uses_whole_file_chunker(self): + source = NonSamplingFileIndexer( + ignore_missing_paths=False, file_chunker=LineDelimitedFileChunker() + ) + + assert isinstance(source.as_whole_file_indexer().file_chunker, WholeFileChunker) + + def test_source_indexer_is_not_mutated(self): + source = NonSamplingFileIndexer( + ignore_missing_paths=False, file_chunker=LineDelimitedFileChunker() + ) + + source.as_whole_file_indexer() + + # Guards against regressing to in-place mutation of the caller's indexer. + assert isinstance(source.file_chunker, LineDelimitedFileChunker) + + if __name__ == "__main__": import sys diff --git a/python/ray/train/BUILD.bazel b/python/ray/train/BUILD.bazel index 3d71751d4233..5bd0faa8ec02 100644 --- a/python/ray/train/BUILD.bazel +++ b/python/ray/train/BUILD.bazel @@ -919,7 +919,9 @@ py_test( py_test( name = "test_e2e_wandb_integration", - size = "small", + # Two full ``TorchTrainer.fit()`` runs (2 workers, 3 epochs) plus a + # ``ray.init`` per parameterization do not fit in small's 60s budget. + size = "medium", srcs = ["tests/test_e2e_wandb_integration.py"], env = {"RAY_TRAIN_V2_ENABLED": "0"}, tags = [ diff --git a/python/ray/train/v2/BUILD.bazel b/python/ray/train/v2/BUILD.bazel index 464a0720c042..fa79fdecfac4 100644 --- a/python/ray/train/v2/BUILD.bazel +++ b/python/ray/train/v2/BUILD.bazel @@ -1,13 +1,23 @@ load("@rules_python//python:defs.bzl", "py_library", "py_test") load("//bazel:python.bzl", "doctest") +# Shared env for train v2 tests. Footer actor pool kept tiny so Data-integrated +# tests do not time out under CI parallelism. +_TRAIN_V2_TEST_ENV = { + "RAY_TRAIN_V2_ENABLED": "1", + "RAY_DATA_PARQUET_FOOTER_NUM_ACTORS": "1", +} + +_TRAIN_V2_TEST_ENV_TF = { + "RAY_TRAIN_V2_ENABLED": "1", + "RAY_DATA_PARQUET_FOOTER_NUM_ACTORS": "1", + "TF_USE_LEGACY_KERAS": "1", +} + doctest( name = "py_doctest[train_v2]", size = "large", - env = { - "RAY_TRAIN_V2_ENABLED": "1", - "TF_USE_LEGACY_KERAS": "1", - }, + env = _TRAIN_V2_TEST_ENV_TF, files = glob( ["**/*.py"], exclude = [ @@ -29,7 +39,7 @@ py_test( name = "test_accelerator_utils", size = "small", srcs = ["tests/test_accelerator_utils.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -45,7 +55,7 @@ py_test( name = "test_autoscaling_coordinator_client", size = "small", srcs = ["tests/test_autoscaling_coordinator_client.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -61,7 +71,7 @@ py_test( name = "test_async_checkpointing_validation", size = "large", srcs = ["tests/test_async_checkpointing_validation.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -77,7 +87,7 @@ py_test( name = "test_checkpoint_manager", size = "small", srcs = ["tests/test_checkpoint_manager.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -93,7 +103,7 @@ py_test( name = "test_config", size = "small", srcs = ["tests/test_config.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -109,7 +119,7 @@ py_test( name = "test_circular_imports", size = "small", srcs = ["tests/test_circular_imports.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -125,7 +135,7 @@ py_test( name = "test_circular_import_linter", size = "small", srcs = ["tests/test_circular_import_linter.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -141,7 +151,7 @@ py_test( name = "test_validation_manager", size = "small", srcs = ["tests/test_validation_manager.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -157,7 +167,7 @@ py_test( name = "test_collective", size = "medium", srcs = ["tests/test_collective.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -173,7 +183,7 @@ py_test( name = "test_callback_manager", size = "small", srcs = ["tests/test_callback_manager.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -189,7 +199,7 @@ py_test( name = "test_controller", size = "medium", srcs = ["tests/test_controller.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -205,7 +215,7 @@ py_test( name = "test_elastic_scaling_policy", size = "small", srcs = ["tests/test_elastic_scaling_policy.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -221,7 +231,7 @@ py_test( name = "test_elastic_e2e", size = "medium", srcs = ["tests/test_elastic_e2e.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -237,7 +247,7 @@ py_test( name = "test_controller_callback_behaviour", size = "medium", srcs = ["tests/test_controller_callback_behaviour.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -253,7 +263,7 @@ py_test( name = "test_data_integration", size = "large", srcs = ["tests/test_data_integration.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "data_integration", "exclusive", @@ -270,7 +280,7 @@ py_test( name = "test_data_parallel_trainer", size = "large", srcs = ["tests/test_data_parallel_trainer.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -286,7 +296,7 @@ py_test( name = "test_report_fault_tolerance", size = "medium", srcs = ["tests/test_report_fault_tolerance.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -302,7 +312,7 @@ py_test( name = "test_data_resource_cleanup", size = "medium", srcs = ["tests/test_data_resource_cleanup.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "data_integration", "exclusive", @@ -319,7 +329,7 @@ py_test( name = "test_dataset_manager", size = "medium", srcs = ["tests/test_dataset_manager.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "data_integration", "exclusive", @@ -336,7 +346,7 @@ py_test( name = "test_env_callbacks", size = "small", srcs = ["tests/test_env_callbacks.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -352,7 +362,7 @@ py_test( name = "test_failure_policy", size = "small", srcs = ["tests/test_failure_policy.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -368,7 +378,7 @@ py_test( name = "test_jax_elastic_e2e", size = "medium", srcs = ["tests/test_jax_elastic_e2e.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -384,7 +394,7 @@ py_test( name = "test_jax_trainer", size = "medium", srcs = ["tests/test_jax_trainer.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -400,7 +410,7 @@ py_test( name = "test_jax_gpu", size = "medium", srcs = ["tests/test_jax_gpu.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ # Temporarily disabled in CI: Ray GPU CI is on CUDA 12.1, but JAX wheels # do not support CUDA 12.1 (and older JAX versions were removed from PyPI). @@ -420,7 +430,7 @@ py_test( name = "test_lightgbm_trainer", size = "small", srcs = ["tests/test_lightgbm_trainer.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -436,7 +446,7 @@ py_test( name = "test_lightning_integration", size = "medium", srcs = ["tests/test_lightning_integration.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -452,7 +462,7 @@ py_test( name = "test_logging", size = "medium", srcs = ["tests/test_logging.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -468,7 +478,7 @@ py_test( name = "test_metrics", size = "small", srcs = ["tests/test_metrics.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -484,7 +494,7 @@ py_test( name = "test_persistence", size = "medium", srcs = ["tests/test_persistence.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -500,7 +510,7 @@ py_test( name = "test_report_handler", size = "small", srcs = ["tests/test_report_handler.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -516,7 +526,7 @@ py_test( name = "test_placement_group_cleaner", size = "small", srcs = ["tests/test_placement_group_cleaner.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -532,7 +542,7 @@ py_test( name = "test_placement_group_handle", size = "medium", srcs = ["tests/test_placement_group_handle.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -548,7 +558,7 @@ py_test( name = "test_preemption_watcher", size = "small", srcs = ["tests/test_preemption_watcher.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -564,7 +574,7 @@ py_test( name = "test_preemption_fault_tolerance", size = "medium", srcs = ["tests/test_preemption_fault_tolerance.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -580,7 +590,7 @@ py_test( name = "test_result", size = "medium", srcs = ["tests/test_result.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -596,7 +606,7 @@ py_test( name = "test_scheduling", size = "medium", srcs = ["tests/test_scheduling.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -612,7 +622,7 @@ py_test( name = "test_serialization", size = "small", srcs = ["tests/test_serialization.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -628,7 +638,7 @@ py_test( name = "test_state", size = "medium", srcs = ["tests/test_state.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -644,7 +654,7 @@ py_test( name = "test_state_export", size = "medium", srcs = ["tests/test_state_export.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -660,7 +670,7 @@ py_test( name = "test_storage", size = "small", srcs = ["tests/test_storage.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -676,7 +686,7 @@ py_test( name = "test_sync_actor", size = "small", srcs = ["tests/test_sync_actor.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -692,7 +702,7 @@ py_test( name = "test_telemetry", size = "medium", srcs = ["tests/test_telemetry.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -708,10 +718,7 @@ py_test( name = "test_tensorflow_trainer", size = "medium", srcs = ["tests/test_tensorflow_trainer.py"], - env = { - "RAY_TRAIN_V2_ENABLED": "1", - "TF_USE_LEGACY_KERAS": "1", - }, + env = _TRAIN_V2_TEST_ENV_TF, tags = [ "exclusive", "team:ml", @@ -727,7 +734,7 @@ py_test( name = "test_thread_runner", size = "small", srcs = ["tests/test_thread_runner.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -743,7 +750,7 @@ py_test( name = "test_torch_gpu", size = "large", srcs = ["tests/test_torch_gpu.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -759,7 +766,7 @@ py_test( name = "test_torch_trainer", size = "medium", srcs = ["tests/test_torch_trainer.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -776,7 +783,7 @@ py_test( name = "test_torch_transformers_train", size = "medium", srcs = ["tests/test_torch_transformers_train.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -792,7 +799,7 @@ py_test( name = "test_util", size = "medium", srcs = ["tests/test_util.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -808,7 +815,7 @@ py_test( name = "test_v2_api", size = "small", srcs = ["tests/test_v2_api.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -824,7 +831,7 @@ py_test( name = "test_worker", size = "small", srcs = ["tests/test_worker.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -840,7 +847,7 @@ py_test( name = "test_worker_group", size = "medium", srcs = ["tests/test_worker_group.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -856,7 +863,7 @@ py_test( name = "test_worker_group_poll_status", size = "small", srcs = ["tests/test_worker_group_poll_status.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -872,7 +879,7 @@ py_test( name = "test_xgboost_trainer", size = "small", srcs = ["tests/test_xgboost_trainer.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "exclusive", "team:ml", @@ -888,10 +895,7 @@ py_test( name = "test_local_mode", size = "medium", srcs = ["tests/test_local_mode.py"], - env = { - "RAY_TRAIN_V2_ENABLED": "1", - "TF_USE_LEGACY_KERAS": "1", - }, + env = _TRAIN_V2_TEST_ENV_TF, tags = [ "exclusive", "team:ml", @@ -907,7 +911,7 @@ py_test( name = "test_data_config", size = "medium", srcs = ["tests/test_data_config.py"], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = _TRAIN_V2_TEST_ENV, tags = [ "data_integration", "exclusive", diff --git a/python/ray/train/v2/tests/conftest.py b/python/ray/train/v2/tests/conftest.py index 0d0802df2d5e..68bdd7528d19 100644 --- a/python/ray/train/v2/tests/conftest.py +++ b/python/ray/train/v2/tests/conftest.py @@ -1,4 +1,5 @@ import logging +import os import pytest @@ -10,6 +11,10 @@ ENABLE_STATE_ACTOR_RECONCILIATION_ENV_VAR, ) +# Keep the footer-reader pool tiny for Data-integrated train tests. Mirrored +# in python/ray/train/v2/BUILD.bazel for bazel test targets. +os.environ.setdefault("RAY_DATA_PARQUET_FOOTER_NUM_ACTORS", "1") + @pytest.fixture() def ray_start_4_cpus(): diff --git a/release/BUILD.bazel b/release/BUILD.bazel index b03048235b52..c7fdd6669b81 100644 --- a/release/BUILD.bazel +++ b/release/BUILD.bazel @@ -111,7 +111,13 @@ py_test( "xgboost", "--smoke-test", ], - env = {"RAY_TRAIN_V2_ENABLED": "1"}, + env = { + # Keep the Parquet footer-reader pool tiny: the production default of 32 + # actors times out and trips Ray's "too many worker processes" warning + # under CI parallelism. + "RAY_DATA_PARQUET_FOOTER_NUM_ACTORS": "1", + "RAY_TRAIN_V2_ENABLED": "1", + }, main = "train_tests/xgboost_lightgbm/train_batch_inference_benchmark.py", tags = [ "exclusive", diff --git a/release/nightly_tests/dataset/arrow_rs_probe/.gitignore b/release/nightly_tests/dataset/arrow_rs_probe/.gitignore new file mode 100644 index 000000000000..9b9a19287ae0 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/.gitignore @@ -0,0 +1,7 @@ +# Probe/experiment output: per-cell logs + summaries, written by run_matrix.py, +# grand_experiment.py and replication_matrix.py. Regenerated by every run and +# tens of MB per run — the analysis lands in arrow_rs_docs/, not here. +matrix_runs/ +grand_runs/ +replication_runs/ +env.sh diff --git a/release/nightly_tests/dataset/arrow_rs_probe/README.md b/release/nightly_tests/dataset/arrow_rs_probe/README.md new file mode 100644 index 000000000000..4368acaa2bab --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/README.md @@ -0,0 +1,203 @@ +# arrow-rs read probe — Linux + S3 regression reproduction + +## TL;DR — the grand experiment (fresh Linux box, one command) + +Maps the whole tuning surface of arrow-rs vs PyArrow on the NEW footer-based +planner (#64985 series): environment setup, the two pytest suites as a +correctness gate, 5 local fixture shapes, then sweeps over the new bin-packing +knob, the decode budget, fragment threads, and (with `ARROW_RS_S3_BUCKET` set) +S3 + fetch window. See the docstrings of `grand_experiment.py` / +`gen_local_fixtures.py` for the stage/shape rationale. + +```bash +git clone https://github.com/AarryaSaraf/ray.git ~/ray && cd ~/ray +git checkout arrow-rs-on-64985 +bash release/nightly_tests/dataset/arrow_rs_probe/run_grand_experiment.sh +# optional S3 stage: export AWS creds + ARROW_RS_S3_BUCKET=s3://... first +# quick smoke run: FIXTURE_SCALE=0.25 STAGES=A bash .../run_grand_experiment.sh +``` + +Results: `grand_runs//summary.md` (ratio tables; R/P > 1.00 = arrow-rs worse), +`summary.json`, one log per cell. + +## Replication matrix — the 2026-08-12 release-A/B trusted signals (TODO 1ab phase 1) + +Replicates the multi-node A/B's *trusted* (wall / decode task-s) good/bad list on one +Linux box: `tensors` (item 1y's negative control — the fsl lookalike that decodes +*faster* native), **`tensorscp`** (the 1y **reproducer**: cloudpickle tensor metadata → +the crate's skip+realign path; macOS read wall R was 5.4 pre-fix, 1.25 after the +`Table.cast` fix), `binsweep` (item 10 — bins from 1 row group up to 10× a file, plus a +PyArrow `pre_buffer=off` arm), **`binbound`** (R2b — is a read task's USS bounded by the +bin budget, or is something retaining?), `write` (item 1aa), `fatcol` (item 1o), +**`oom`** (R5 — the failure-mode demo: same memory ceiling via Ray's own memory monitor, +sweep the bin; PyArrow's arm is *expected* to die with `OutOfMemoryError` at the big +bins — an OOM there is the stage's result, not a broken run). Stage rationale in +`replication_matrix.py`'s docstring; the predictions each stage falsifies are in +`arrow_rs_docs/TODO.md` items 1ab/10. + +## Loss triage — the 2026-08-15 release-A/B losses in 3 parts (M31/M32/M33) + +One command; answers "is each loss the native decoder, Ray integration, or the S3 +path?" by running each loss shape **standalone (no Ray/no S3) → Ray on local files → +Ray on S3**, both readers per part, plus a `MALLOC_ARENA_MAX=2` arm on the arrow-rs +Ray cells (the glibc-arena-retention discriminator). Shapes: `auto` (M31 +read_large_parquet_autoscaling: one ~69 MiB rg/task, re-measured at 20 Hz because the +release 1 Hz poll can't see sub-second tasks), `write` (M32 write_parquet: fused +read→write, ~1.25 GiB churn/task), `tensorscp` (M33 wide_schema tensors: the 1y +skip+realign path — the shape with a history of losing outside Ray). Mapping and +rationale in `loss_triage.py`'s docstring. + +```bash +bash release/nightly_tests/dataset/arrow_rs_probe/run_loss_triage.sh +# with the S3 part (scratch bucket you own; AWS creds exported): +ARROW_RS_S3_BUCKET=s3://arrowrs-bench-xxxx bash .../run_loss_triage.sh +# subsets / smoke: +SHAPES=write PARTS=ray_local FIXTURE_SCALE=0.25 REPEAT=1 bash .../run_loss_triage.sh +``` + +Results: `loss_triage_runs//summary.json` + per-cell logs; the summary table is +R = arrow_rs/pyarrow per (shape, part) — a loss only in `ray_local` with the arena2 +column collapsing ⇒ allocator retention; only in `ray_s3` ⇒ crate S3 path. + +For item 1o's crate-level A/B, `patch_crate_parquet.sh` rebuilds `ray_data_arrow_rs` +against a vendored parquet 59.1.0 carrying the dictionary values-reserve fix +(`patches/parquet-59.1.0-dict-reserve.diff`; REVERT=1 restores stock — see the script +header for the stock→patched fatcol procedure). + +`binbound` is the bound/leak check: one bin is one read task, so it runs at +`--task-concurrency 1` (one bin resident per process) with `--mem-poll-s 0.05` (the 1 Hz +default samples a short task once or not at all) and fits per-task USS against **decoded** +bytes per task — slope ≲0.3 flat, ≲1.1 bounded by the bin, >1.1 unbounded ⇒ leak. It +prints the bin→decoded expansion per cell because the knob is spent in +`row_group.total_byte_size` (encoded bytes), measured 1.64× below decoded Arrow on the +fixture — finding C10. **Needs Linux**: per-task USS is `None` on macOS and the fit is +skipped (the task-count and expansion columns still work). + +Fresh Linux box, one command (same setup skeleton as the grand experiment): + +```bash +git clone https://github.com/AarryaSaraf/ray.git ~/ray && cd ~/ray +git checkout arrow-rs-on-64985 +bash release/nightly_tests/dataset/arrow_rs_probe/run_replication.sh +# quick smoke: FIXTURE_SCALE=0.25 REPEAT=1 bash .../run_replication.sh +# one stage: ONLY=binsweep bash .../run_replication.sh +# after a git pull that touches the crate: FORCE_SETUP=1 (or you benchmark a stale .so) +``` + +Piecemeal (env already set up — `source arrow_rs_probe/env.sh` first): + +```bash +python gen_local_fixtures.py --root ~/arrow_rs_repl_fixtures \ + --shapes bin_sweep,tensors_wide,tensors_cp,fat_col +python replication_matrix.py --fixture-root ~/arrow_rs_repl_fixtures --repeat 3 +``` + +Results: `replication_runs//summary.json` + per-cell logs; the summary block +prints R = arrow_rs/pyarrow per cell pair and the pre_buffer on/off deltas. + +**How a cell is aggregated** (both matrices; `run_matrix.median_cell`): each cell runs +`--warmup` times (default **1**, logged as `.w`, discarded) and then `--repeat` +times (logged as `.r`), and **every metric is medianed independently** over the +measured runs. `summary.json` carries `_n` and `_samples` per cell so spread is visible +without opening a log. Both halves matter: before 2026-08-13 the function returned the +whole dict of whichever run had the median `wall_s`, which made every other metric a +single sample chosen by an unrelated one — that alone reported a 1.98× read-wall ratio +as 1.04. And the first repeat is reproducibly the cold one (page cache, `.so` load): +1.44 s against a 1.03–1.22 s steady state on `fat_col`, repeatable to 0.35%, enough to +drag a 3-sample median. Use `--warmup 0` only when you want the cold number. + +--- + +Single-node harness to measure the two cases where the arrow-rs Parquet reader was +worse than PyArrow in the release run (build 102757), so we can optimize them: + +| release test | axis | gap | why it needs Linux + S3 | +|---|---|---|---| +| `mix.8ds_equal_random_mix` (imagenet, many tiny row groups) | **time** | 1.67× | I/O-bound on S3 (many small serial range GETs). Faster than PyArrow on local disk at every scale — no network to expose it. | +| `wide_schema_pipeline_primitives` (5000 cols) | **memory** | 1.50× | The crate's page-sized working set only engages on the **S3** decode path; and per-worker USS is Linux-only. | + +Both are single-worker read properties, so one node + S3 reproduces them — no cluster. + +## Setup (fresh Linux workspace on branch `arrow-rs-parquet-reader-pr`) + +```bash +# 1. Repo + a commit-matched Ray nightly wheel (a "latest" wheel drifts from this +# branch's compiled protobufs and asserts "out of sync" at import). +git clone https://github.com/AarryaSaraf/ray.git ~/ray && cd ~/ray +git checkout arrow-rs-parquet-reader-pr +git remote add upstream https://github.com/ray-project/ray.git && git fetch upstream master --quiet + +uv venv --python 3.12 ~/ray/.venv && source ~/ray/.venv/bin/activate +# Pick the x86-64 manylinux nightly matching `git merge-base HEAD upstream/master`, +# install --no-deps, then re-link this branch's source: +uv pip install --no-deps +uv pip install "ray[data]" psutil +python python/ray/setup-dev.py -y # symlinks python/ray/ over the wheel + +# 2. Build the native crate for Linux (first time off macOS). +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y && source "$HOME/.cargo/env" +uv pip install maturin +cd python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs +maturin build --release && uv pip install --force-reinstall --no-deps target/wheels/*.whl +cd ~/ray +``` + +### Anyscale-workspace gotchas (none are arrow-rs-related) +- **Run a private cluster:** `export RAY_ADDRESS=local` — never attach to the managed + cluster (different Ray version → version-check failure). +- **Re-activate the venv in every shell** (`source ~/ray/.venv/bin/activate`) — otherwise + `python` is the image's anaconda Ray (the Anyscale runtime), where this branch's + reader/crate/flags don't exist. Confirm with + `python -c "import ray.data, os; print(os.path.realpath(ray.data.__file__))"` → must + resolve into this checkout. +- If `ray.init()` hangs: `unset RAY_RUNTIME_ENV_HOOK RAY_RUNTIME_ENV_PLUGINS` and + `export RAY_task_events_report_interval_ms=0` (dodges a 2026-07 master task-event + SIGSEGV); check `/tmp/ray/session_latest/logs/runtime_env_agent.err` for import errors. + +## Run + +Confirm the S3 prefixes first (they drift): `aws s3 ls s3://ray-benchmark-data-internal-us-west-2/wide_schema/` +and `.../imagenet/`, and export AWS creds/region. Put the box in the bucket's region. + +```bash +cd release/nightly_tests/dataset/arrow_rs_probe +export RAY_ADDRESS=local + +# (A) CPU-bound-vs-IO diagnostic — force one read task, compare cpu_over_wall. +# ~1 => CPU-bound decode; <<1 => I/O-waiting on S3. +python read_probe.py --preset imagenet --reader pyarrow --concurrency 1 +python read_probe.py --preset imagenet --reader arrow_rs --concurrency 1 +python read_probe.py --preset wide_schema --reader pyarrow --concurrency 1 +python read_probe.py --preset wide_schema --reader arrow_rs --concurrency 1 + +# (B) Realistic memory — let it fan out; compare peak_uss_gb (the metric of record). +python read_probe.py --preset wide_schema --reader pyarrow +python read_probe.py --preset wide_schema --reader arrow_rs + +# (C) Allocator A/B on the arrow_rs run (no rebuild) — is any residual mem gap glibc +# arena retention rather than the decoder? +MALLOC_ARENA_MAX=2 python read_probe.py --preset wide_schema --reader arrow_rs +LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2 python read_probe.py --preset wide_schema --reader arrow_rs +``` + +Local scale sweeps (no S3 — separates a scaling effect from the S3/Linux effect; on +Linux they also fill in `peak_uss_gb`): + +```bash +python scale_sweep.py imagenet # bytes sweep: does arrow_rs stay faster as data grows? +python scale_sweep.py wide # row-group-size sweep: USS ratio vs rg size +``` + +## What to look for +- **imagenet time:** if `cpu_over_wall` ≪ 1 on the S3 read, the gap is prefetch/fetch — + the fix is crate-level concurrent column-chunk prefetch within a row group (PyArrow's + `pre_buffer` issues parallel column GETs; the crate's windowed stream fetches serially), + or reading N row groups per task concurrently. +- **wide_schema memory:** compare `peak_uss_gb` (and Ray's `read_avg_max_uss_gb`). The + architectural expectation is arrow-rs **at or below** PyArrow once the S3 windowed path + engages, growing flat in row-group size while PyArrow grows with it. If arrow-rs is still + worse on Linux+S3, it's a genuine integration regression to chase there (not the macOS + allocator/RSS artifact the local sweep shows). + +For richer log output on a full release run, `collect_operator_metrics` (in +`../benchmark.py`) now emits per-operator wall/output-bytes/decode-USS. diff --git a/release/nightly_tests/dataset/arrow_rs_probe/batch_ablation.py b/release/nightly_tests/dataset/arrow_rs_probe/batch_ablation.py new file mode 100644 index 000000000000..4f3c0b87afa3 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/batch_ablation.py @@ -0,0 +1,381 @@ +"""Batch-sizing ablation across the release loss shapes (M43 / item 1y). + +WHAT THIS ANSWERS +----------------- +A/B #4 measured arrow-rs's largest yielded table at 288 MiB on a 32 MiB decode +budget for `wide_schema tensors` (M39). The mechanism (M43, verified standalone +2026-08-18): both batch-sizing layers divide by ENCODED footer bytes/row — + + * Python request: ceil(budget / (enc_bpr * 5)) then max(_, 2048) + (`_estimate_batch_size_from_*`, PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT=5, + `_ARROW_RS_MIN_DECODE_BATCH_ROWS=2048`) + * crate clamp: byte_budget_rows = budget / enc_bpr, upper-clamped by the + request (`byte_budget_rows`, src/lib.rs) + +so on data that dictionary-encodes r×, the decoded batch lands at ~budget×r +(release tensors: r≈9 → 288 MiB). The 2048 floor is an AMPLIFIER on top (the +un-floored request would be ~410 rows ≈ 57 MiB), not the root. + +This probe ablates the request policy — the one lever that upper-clamps the +crate — across multiple shapes, per the user's ask ("ablational test on those +values on multiple test cases like the ones we saw"): + + policies (rs): + floor32 / floor128 / floor512 / floor2048 request = max(est5, N) + (floor2048 = shipping behavior) + decoded request = budget / MEASURED + decoded bytes-per-row (the fix + candidate: decoded-aware sizing) + reference: + pa PyArrow fragment scanner at the + shipping estimate (its peak is + whole-decoded-row-group-bound + regardless — C9) + + shapes (all local fixtures, gen_local_fixtures.py): + tensors_dict release-faithful ~9.4x dictionary expansion — the M39 shape + tensors_cp same schema, ~1.1x expansion (control: budget should hold) + fat_col one fat utf8 column, ~256 KiB/row (T13's shape) + auto_rg thin rows (control: floor is irrelevant, est5 >> 2048) + +Per cell (fresh subprocess): wall, peak RSS (ru_maxrss), yielded-batch dist +(max/p50 MiB, max rows), request used. Verdict logic: + + * tensors_dict floor2048 max-batch >> budget AND decoded max-batch ≈ budget + ⇒ M43 confirmed end-to-end; decoded-aware sizing is the fix. + * peak RSS should track max-batch on rs; pa's stays row-group-bound. + * auto_rg rows/wall across floors shows what the floor actually buys on the + thin shapes it was added for (T22: fewer batches was −34% wall there). + +V2 (2026-08-18, "bigger and newer" per review): the matrix is now +shapes x policies x BUDGETS, the shapes cover an expansion SWEEP +(tensors_cp ~1.1x -> tensors_lo ~2-3x -> tensors_dict ~9.4x -> tensors_hi ~15x+) +plus the structural shapes (fat_col, auto_rg, wide, tiny_rgs, single_rg_files, +lone_big_rg), and the DECISION VARIABLES are outcome gates rather than raw +numbers: + + G1 overshoot max yielded batch bytes / decode budget. The mechanism gate: + shipping (floor2048) overshoot ~= expansion ratio on dict + shapes (M43); a fix must hold overshoot <= 1.5 on EVERY shape. + G2 memory peak RSS <= 1.10x the pa reference cell. + G3 wall <= 1.25x pa, with pa decoding SINGLE-THREADED (use_threads= + False): the crate side is K=1 and in-Ray both arms take their + parallelism from tasks, so the gate compares CPU cost. A + threaded pa baseline fails every cell by ~cores (M35's + standalone artifact — bit the first Linux run). NB the static `decoded` policy is EXPECTED to + fail G3 on the 5000-col tensor shapes (per-batch realign cost, + T22/T23) — that failure is the argument that the real fix must + be crate-side mid-stream adaptation (size from the first + yielded batch, like pa's parquet_file_reader.py:385 refinement), + not a uniformly smaller static request. + G4 rows row-count parity with the pa cell (cheap correctness gate). + +A candidate code change passes this suite when its policy row passes all four +gates on all shapes at every budget. Standalone-only by design: batch sizing is +transport-independent (same math local and S3); the in-Ray and S3 legs of the +same shapes live in loss_triage.py and the retention leg in soak_probe.py — +run_all.sh chains all of them. + +Usage: + python batch_ablation.py --fixtures-root DIR [--scale 0.25] + [--shapes ...] [--policies pa,floor32,floor128,floor512,floor2048,decoded] + [--budgets-mib 32] [--out .] +Missing fixtures are generated (gen_local_fixtures.py) automatically. +Results: table + gate verdict on stdout, ablation.json under --out. +""" + +import argparse +import glob +import json +import math +import os +import resource +import subprocess +import sys +import time + +MiB = 1024 * 1024 +# ru_maxrss is KiB on Linux, bytes on macOS. +_RU_UNIT = 1024 if sys.platform.startswith("linux") else 1 + +_HERE = os.path.dirname(os.path.abspath(__file__)) + +DEFAULT_SHAPES = [ + # expansion sweep (the M43 axis), low -> high + "tensors_cp", + "tensors_lo", + "tensors_dict", + "tensors_hi", + # structural shapes + "fat_col", + "auto_rg", + "wide", + "tiny_rgs", + "single_rg_files", + "lone_big_rg", +] +DEFAULT_POLICIES = ["pa", "floor32", "floor128", "floor512", "floor2048", "decoded"] +# Shapes whose files carry cloudpickle tensor metadata (need the autoload flag +# and the reader's skip+realign path). +_TENSOR_SHAPES = {"tensors_dict", "tensors_cp", "tensors_lo", "tensors_hi"} + + +def _enc_bpr(files): + """Encoded (footer) bytes per row across the fixture — what BOTH shipping + sizing layers divide by.""" + import pyarrow.parquet as pq + + enc = rows = 0 + for f in files: + md = pq.read_metadata(f) + enc += sum(md.row_group(i).total_byte_size for i in range(md.num_row_groups)) + rows += md.num_rows + return max(1, enc // max(1, rows)), rows + + +def _est5(files, budget): + """Mirror the shipping request estimate BEFORE the floor: + ceil(budget / (enc_bpr * PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT)), + capped at the row count like `_estimate_batch_size_from_*` caps at the + row-group/chunk rows.""" + bpr, rows = _enc_bpr(files) + return min(math.ceil(budget / (bpr * 5)), rows) + + +def _decoded_request(f, knobs, realign_fields): + """The fix candidate: measure DECODED bytes/row from one small probe batch, + then request budget/dec_bpr. (In-reader this would be the adaptive + refinement the PyArrow path already has.)""" + from loss_triage import _rs_batches + + for t in _rs_batches(f, 64, knobs, realign_fields=realign_fields): + dec_bpr = max(1, t.nbytes // max(1, t.num_rows)) + return max(32, int(knobs["budget"] // dec_bpr)) + return 32 + + +def run_cell(a): + from loss_triage import _consume, _pa_batches, _reader_knobs, _rs_batches + + files = sorted(glob.glob(os.path.join(a.path, "*.parquet"))) + assert files, f"no parquet under {a.path}" + knobs = _reader_knobs() + if a.budget_mib: + knobs["budget"] = a.budget_mib * MiB + + realign_fields = None + if a.shape in _TENSOR_SHAPES: + import pyarrow.parquet as pq + + realign_fields = list(pq.read_schema(files[0])) + + est5 = _est5(files, knobs["budget"]) + if a.policy == "pa": + request = est5 + elif a.policy == "decoded": + request = _decoded_request(files[0], knobs, realign_fields) + else: + request = max(est5, int(a.policy.removeprefix("floor"))) + + batch_rows, batch_mib = [], [] + + def observed(it): + for t in it: + batch_rows.append(t.num_rows) + batch_mib.append(t.nbytes / MiB) + yield t + + t0 = time.monotonic() + rows = 0 + for f in files: + if a.policy == "pa": + # Single-threaded pa: G3 compares CPU cost, not thread-pool + # fan-out — the crate side is K=1 and in-Ray both arms get their + # parallelism from tasks (M35). Threaded pa here made every wall + # gate fail by ~cores on Linux. + it = _pa_batches(f, request, use_threads=False) + else: + it = _rs_batches(f, request, knobs, realign_fields=realign_fields) + rows += _consume(observed(it), "decode", None) + wall = time.monotonic() - t0 + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * _RU_UNIT + + srt = sorted(batch_mib) + print( + "CELL_JSON " + + json.dumps( + dict( + shape=a.shape, + policy=a.policy, + budget_mib=knobs["budget"] // MiB, + overshoot=round(max(batch_mib) / (knobs["budget"] / MiB), 2) + if batch_mib + else 0, + request=request, + est5=est5, + wall_s=round(wall, 2), + peak_rss_mib=round(peak / MiB, 1), + rows=rows, + n_batches=len(batch_mib), + batch_mib_max=round(max(batch_mib), 1) if batch_mib else 0, + batch_mib_p50=round(srt[len(srt) // 2], 1) if srt else 0, + batch_rows_max=max(batch_rows) if batch_rows else 0, + ) + ) + ) + + +def orchestrate(a): + shapes = a.shapes.split(",") + policies = a.policies.split(",") + + # Fixtures: generate any missing shape (version/scale-gated skip inside). + subprocess.run( + [ + sys.executable, + os.path.join(_HERE, "gen_local_fixtures.py"), + "--root", + a.fixtures_root, + "--shapes", + ",".join(shapes), + "--scale", + str(a.scale), + ], + check=True, + ) + with open(os.path.join(a.fixtures_root, "manifest.json")) as fh: + manifest = json.load(fh) + + budgets = [int(b) for b in str(a.budgets_mib).split(",")] + results = [] + for shape in shapes: + for budget in budgets: + for policy in policies: + env = dict(os.environ) + if shape in _TENSOR_SHAPES: + env["RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA"] = "1" + cmd = [ + sys.executable, + os.path.abspath(__file__), + "cell", + "--shape", + shape, + "--policy", + policy, + "--path", + manifest[shape]["path"], + "--budget-mib", + str(budget), + ] + print(f"== {shape} / {policy} / {budget}MiB", flush=True) + out = subprocess.run( + cmd, env=env, cwd=_HERE, capture_output=True, text=True + ) + line = next( + ( + ln + for ln in out.stdout.splitlines() + if ln.startswith("CELL_JSON ") + ), + None, + ) + if line is None: + print(out.stdout[-2000:]) + print(out.stderr[-2000:]) + raise SystemExit(f"cell failed: {shape}/{policy}/{budget}") + rec = json.loads(line[len("CELL_JSON ") :]) + rec["expansion"] = manifest[shape].get("enc_to_dec_ratio") + results.append(rec) + print(f" {rec}", flush=True) + + out_path = os.path.join(a.out, "ablation.json") + + # Gate thresholds (the suite's decision variables — see docstring). + G_OVERSHOOT, G_RSS, G_WALL = 1.5, 1.10, 1.25 + verdict = {} + print("\n== ablation summary (R vs the pa row at the same budget) ==") + for shape in shapes: + for budget in budgets: + rows = [ + r for r in results if r["shape"] == shape and r["budget_mib"] == budget + ] + if not rows: + continue + pa_row = next((r for r in rows if r["policy"] == "pa"), None) + exp = rows[0].get("expansion") + print(f"\n[{shape} @ {budget} MiB] enc->dec expansion: {exp}") + hdr = ( + f"{'policy':>10} {'request':>8} {'maxbatch':>9} {'over':>6} " + f"{'p50batch':>9} {'peakRSS':>8} {'wall':>6} {'rows':>9}" + " R_rss R_wall gates" + ) + print(hdr) + for r in rows: + rr = rw = gates = "" + if pa_row and r is not pa_row: + r_rss = r["peak_rss_mib"] / pa_row["peak_rss_mib"] + r_wall = r["wall_s"] / max(0.01, pa_row["wall_s"]) + rr, rw = f"{r_rss:.2f}", f"{r_wall:.2f}" + g = [ + "G1" if r["overshoot"] <= G_OVERSHOOT else "g1!", + "G2" if r_rss <= G_RSS else "g2!", + "G3" if r_wall <= G_WALL else "g3!", + "G4" if r["rows"] == pa_row["rows"] else "g4!", + ] + gates = " ".join(g) + verdict[f"{shape}@{budget}/{r['policy']}"] = dict( + overshoot=r["overshoot"], + r_rss=round(r_rss, 2), + r_wall=round(r_wall, 2), + rows_match=r["rows"] == pa_row["rows"], + passed=not any(x.endswith("!") for x in g), + ) + print( + f"{r['policy']:>10} {r['request']:>8} {r['batch_mib_max']:>8.1f}M " + f"{r['overshoot']:>6.2f} {r['batch_mib_p50']:>8.1f}M " + f"{r['peak_rss_mib']:>7.0f}M {r['wall_s']:>5.1f}s " + f"{r['rows']:>9} {rr:>5} {rw:>6} {gates}" + ) + + os.makedirs(os.path.dirname(os.path.abspath(out_path)), exist_ok=True) + with open(out_path, "w") as fh: + json.dump({"cells": results, "verdict": verdict}, fh, indent=2) + fails = {k: v for k, v in verdict.items() if not v["passed"]} + print( + f"\n== gate verdict: {len(verdict) - len(fails)}/{len(verdict)} " + f"policy-cells pass (G1 overshoot<={G_OVERSHOOT} G2 R_rss<={G_RSS} " + f"G3 R_wall<={G_WALL} G4 rows==pa) ==" + ) + for k, v in sorted(fails.items()): + print(f" FAIL {k}: {v}") + print(f"\nresults -> {out_path}") + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + sub = p.add_subparsers(dest="cmd") + + c = sub.add_parser("cell", help="internal: one cell in a fresh process") + c.add_argument("--shape", required=True) + c.add_argument("--policy", required=True) + c.add_argument("--path", required=True) + c.add_argument("--budget-mib", type=int, default=0) + + a_ = p.add_argument + a_("--fixtures-root", default=os.environ.get("FIXTURES_ROOT", "")) + a_("--scale", type=float, default=0.25) + a_("--shapes", default=",".join(DEFAULT_SHAPES)) + a_("--policies", default=",".join(DEFAULT_POLICIES)) + a_("--budgets-mib", default="32", help="comma list, e.g. 16,32,128") + a_("--out", default=".") + + a = p.parse_args() + if a.cmd == "cell": + run_cell(a) + else: + assert a.fixtures_root, "--fixtures-root (or FIXTURES_ROOT) required" + orchestrate(a) + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/concurrency_cliff_probe.py b/release/nightly_tests/dataset/arrow_rs_probe/concurrency_cliff_probe.py new file mode 100644 index 000000000000..cbb961d9486a --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/concurrency_cliff_probe.py @@ -0,0 +1,311 @@ +"""Concurrency-cliff probe (M92): does the native S3 read path fall off a cliff +when concurrent read tasks on one node exceed ~physical cores, and is the node +actually out of runnable-thread capacity when it happens? + +Background (2026-09-01.md §9): in the 2x2 release run, arrow-rs read tasks ran at +pyarrow speed (or faster) up to ~47 concurrent tasks on an m5.24xlarge, then +jumped 5-8x at >=48 (= 96 hyperthreads / 2). Hypothesis: an rs task keeps ~2 +threads busy under real S3 (tokio decode task + fetch/TLS work + the Python +consumer) vs pyarrow's ~1, so Ray's num_cpus=1 accounting oversubscribes the +node at ~cores/2 tasks. This script measures BOTH the wall-vs-concurrency curve +and the runnable-thread count, so it discriminates: + + * cliff in wall AND runnable threads ~= hardware threads at the cliff + -> CPU/thread oversubscription confirmed (fix: honest num_cpus, or make a + task ~1 busy thread: shrink tokio pool / decode on calling thread). + * cliff in wall while runnable threads stay LOW + -> tasks are stalling, not competing for CPU (fetch starvation, scheduler + or channel pathology) -> different fix, profile with + RAY_DATA_ARROW_RS_PROFILE=1. + * no cliff at all on this box + -> the release loss needs something this box lacks; escalate to the + release confirmation cell instead. + +Run on a LINUX box with real S3 credentials (the effect needs real network + +TLS; moto/localhost shows nothing — measured: single-read avg busy threads +0.4-0.6 rs vs 0.9 pa on moto). One command, both arms, sweeps concurrency: + + python concurrency_cliff_probe.py --path s3:///cliff_probe --gen 600 + python concurrency_cliff_probe.py --path s3:///cliff_probe + +Per (arm, N): fresh ray.init, ray.data.read_parquet(path, concurrency=N), +consume via .sum(), while a sampler thread counts runnable (R-state) threads +across all ray:: worker processes at 5 Hz plus 1-min load average. Expected +signature if M92 holds (T = hardware threads of the box): + pyarrow : wall keeps improving (or flat) as N grows through T + arrow-rs: wall improves until N ~= T/2, then DEGRADES; runnable ~= T there. +""" + +import argparse +import glob +import json +import os +import re +import threading +import time + +FLAG = "RAY_DATA_USE_ARROW_RS_PARQUET_READER" + + +# ---------------------------------------------------------------- fixtures +def gen_fixtures(path: str, n_files: int, rows_per_file: int) -> None: + """~64 MiB/file at the default 4M rows: int64 + float64 + short string, + snappy, 4 row groups — the boring many-medium-files shape (like rlp).""" + import numpy as np + import pyarrow as pa + import pyarrow.parquet as pq + from pyarrow import fs as pafs + + fs, root = pafs.FileSystem.from_uri(path) + rng = np.random.default_rng(0) + t = pa.table( + { + "a": rng.integers(0, 1 << 40, rows_per_file), + "b": rng.random(rows_per_file), + "c": pa.array(rng.integers(0, 99999, rows_per_file).astype("U8")), + } + ) + for i in range(n_files): + with fs.open_output_stream(f"{root}/part-{i:05d}.parquet") as f: + pq.write_table( + t, f, row_group_size=rows_per_file // 4, compression="snappy" + ) + if i % 50 == 0: + print(f" wrote {i}/{n_files}", flush=True) + print(f"fixtures: {n_files} files at {path}", flush=True) + + +# ------------------------------------------------- runnable-thread sampler +def _ray_worker_pids(): + pids = [] + for cmdline in glob.glob("/proc/[0-9]*/cmdline"): + try: + with open(cmdline, "rb") as f: + if f.read().split(b"\x00", 1)[0].startswith(b"ray::"): + pids.append(int(cmdline.split("/")[2])) + except OSError: + continue + return pids + + +def _runnable_threads(pids): + n = 0 + for pid in pids: + for stat in glob.glob(f"/proc/{pid}/task/[0-9]*/stat"): + try: + with open(stat) as f: + # state is the first field after the last ')' (comm may + # contain spaces/parens) + if f.read().rsplit(")", 1)[1].split()[0] == "R": + n += 1 + except OSError: + continue + return n + + +class Sampler(threading.Thread): + def __init__(self): + super().__init__(daemon=True) + self.samples = [] + self.stop_evt = threading.Event() + + def run(self): + if not os.path.isdir("/proc"): # macOS smoke runs: wall-only + return + last_scan, pids = 0.0, [] + while not self.stop_evt.is_set(): + if time.time() - last_scan > 2: + pids, last_scan = _ray_worker_pids(), time.time() + self.samples.append((_runnable_threads(pids), os.getloadavg()[0])) + time.sleep(0.2) + + def summary(self): + if not self.samples: + return {"runnable_p50": None, "runnable_p90": None, "load1_max": None} + runnable = sorted(s[0] for s in self.samples) + return { + "runnable_p50": runnable[len(runnable) // 2], + "runnable_p90": runnable[int(len(runnable) * 0.9)], + "load1_max": max(s[1] for s in self.samples), + } + + +# ------------------------------------------------------------------ cells +# +# Each cell runs in a FRESH subprocess. The reader is chosen from the driver's +# DataContext singleton (parquet_scanner.py:86), which is created at first +# `ray.data` import from the env var and then serialized to every task — it +# survives ray.shutdown() and OVERRIDES worker runtime_env env vars. Flipping +# arms inside one interpreter would therefore run every cell with the first +# arm's reader. A fresh interpreter per cell (env set before import, plus an +# explicit DataContext set as belt-and-braces) makes the arm switch real. +def run_cell_subprocess(arm: str, path: str, n: int, col: str, consume: str) -> dict: + import subprocess + import sys + + env = dict(os.environ) + env[FLAG] = "1" if arm == "rs" else "0" + # An Anyscale workspace exports RAY_ADDRESS; connecting to that cluster + # would run tasks from the platform's Ray install, not this checkout. + env.pop("RAY_ADDRESS", None) + proc = subprocess.run( + [ + sys.executable, + os.path.abspath(__file__), + "--path", + path, + "--cell", + arm, + "--n", + str(n), + "--sum-col", + col, + "--consume", + consume, + ], + env=env, + capture_output=True, + text=True, + ) + for line in proc.stdout.splitlines(): + if line.startswith("CELLRESULT "): + return json.loads(line[len("CELLRESULT ") :]) + raise RuntimeError( + f"cell {arm} N={n} produced no result\n--- stdout ---\n{proc.stdout[-2000:]}" + f"\n--- stderr ---\n{proc.stderr[-2000:]}" + ) + + +def run_cell_body(arm: str, path: str, n: int, col: str, consume: str) -> None: + want_rs = arm == "rs" + os.environ[FLAG] = "1" if want_rs else "0" + import ray + + # address="local" forces a NEW local cluster from THIS venv. Plain init() + # auto-discovers a running cluster via /tmp/ray/ray_current_cluster — on an + # Anyscale workspace that is the platform's raylet (different Ray + Python), + # and joining it fails with a version mismatch (or worse, runs there). + ray.init(address="local", include_dashboard=False, logging_level="ERROR") + import ray.data + + ctx = ray.data.DataContext.get_current() + ctx.use_arrow_rs_parquet_reader = want_rs # explicit: no reliance on import order + if want_rs: + import ray_data_arrow_rs # noqa: F401 fail loudly if the crate is absent + + from ray.data.aggregate import Sum + + sampler = Sampler() + sampler.start() + t0 = time.time() + ds = ray.data.read_parquet(path, concurrency=n) + if consume == "iter-bundles": + # Release-faithful consume (read_and_consume_benchmark --iter-bundles): + # full decode of EVERY column. The sum consume is NOT equivalent — the + # optimizer pushes the projection into the read, fetching one column. + for _ in ds.iter_internal_ref_bundles(): + pass + stats_ds = ds + else: + # ds.sum() would discard the executed plan's stats (it builds and + # consumes an internal dataset); hold the aggregate so .stats() works. + stats_ds = ds.groupby(None).aggregate(Sum(col)) + stats_ds.take(1) + wall = time.time() - t0 + sampler.stop_evt.set() + sampler.join() + + stats = stats_ds.stats() + m = re.search(r"ReadFilesParquetV2: (\d+) tasks executed", stats) + tasks = int(m.group(1)) if m else None + ray.shutdown() + result = { + "arm": arm, + "concurrency": n, + "wall_s": round(wall, 2), + "read_tasks": tasks, + **sampler.summary(), + } + print("CELLRESULT " + json.dumps(result), flush=True) + + +def default_concurrencies() -> list: + t = os.cpu_count() or 8 + cand = [t // 8, t // 4, 3 * t // 8, t // 2 - 4, t // 2, t // 2 + 8, 3 * t // 4, t] + return sorted({c for c in cand if c >= 1}) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--path", required=True, help="s3://bucket/prefix with fixture files" + ) + ap.add_argument( + "--gen", type=int, default=0, help="generate N fixture files, then exit" + ) + ap.add_argument("--rows-per-file", type=int, default=4_000_000) + ap.add_argument("--arms", default="pa,rs") + ap.add_argument( + "--concurrencies", default="", help="comma list; default derived from cores" + ) + ap.add_argument("--out", default="cliff_probe_results.jsonl") + ap.add_argument( + "--sum-col", + default="a", + help="numeric column to consume via sum ('a' for --gen fixtures; " + "e.g. column05 for the release large-parquet dataset)", + ) + ap.add_argument( + "--consume", + default="sum", + choices=["sum", "iter-bundles"], + help="'sum' projects to one column (pushdown!); 'iter-bundles' decodes " + "every column, like the release read_large_parquet tests", + ) + ap.add_argument("--cell", default="", help="internal: run one (arm) cell and exit") + ap.add_argument("--n", type=int, default=0, help="internal: concurrency for --cell") + args = ap.parse_args() + + if args.gen: + gen_fixtures(args.path, args.gen, args.rows_per_file) + return + + if args.cell: + run_cell_body(args.cell, args.path, args.n, args.sum_col, args.consume) + return + + ns = [int(x) for x in args.concurrencies.split(",") if x] or default_concurrencies() + print( + f"box: {os.cpu_count()} hw threads | sweep N={ns} | arms={args.arms}", + flush=True, + ) + rows = [] + for n in ns: # interleave arms per N so cluster/S3 weather cancels + for arm in args.arms.split(","): + r = run_cell_subprocess(arm, args.path, n, args.sum_col, args.consume) + rows.append(r) + print( + f"{arm:>3} N={n:<3} wall {r['wall_s']:>7.2f}s tasks {r['read_tasks']} " + f"runnable p50/p90 {r['runnable_p50']}/{r['runnable_p90']} load1max {r['load1_max']}", + flush=True, + ) + with open(args.out, "a") as f: + f.write(json.dumps(r) + "\n") + + print("\n=== wall_s by concurrency (rows=arm) ===") + arms = sorted({r["arm"] for r in rows}) + print("arm | " + " | ".join(f"N={n}" for n in ns)) + for arm in arms: + vals = {r["concurrency"]: r["wall_s"] for r in rows if r["arm"] == arm} + print( + f"{arm:>3} | " + " | ".join(f"{vals.get(n, float('nan')):.1f}" for n in ns) + ) + print( + "\nM92 predicts: pa monotone-improving through N=cores; rs improves to ~cores/2 " + "then degrades, with runnable_p90 ~= hw threads at the cliff. Runnable staying low " + "while rs wall explodes = stall, not CPU -> profile instead." + ) + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/fat_col_bench.py b/release/nightly_tests/dataset/arrow_rs_probe/fat_col_bench.py new file mode 100644 index 000000000000..2158935f504c --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/fat_col_bench.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +"""Tiny standalone A/B: PyArrow vs the arrow-rs crate on the fat_col shape. + +fat_col = 1 file, 1 row group, N rows (default 1024), one ~256 KiB/row column +of RANDOM bytes plus one int64 (~256 MiB file at the default). Random bytes +are incompressible, so encoded == decoded and there is nothing to "decode": +the read is pure allocate-and-copy. That makes this the shape where the +crate's allocator behavior (glibc malloc, fresh pages each batch) loses on +wall clock to PyArrow's jemalloc (warm page reuse) while still winning ~2x on +peak memory — findings M56/M69/M70: rs/pa1 wall 1.7-2.7x at peak RSS 0.4-0.5x +(the wall gap varies with box session; the memory win is stable). + +Usage (needs pyarrow + numpy + the branch-built ray_data_arrow_rs crate): + python fat_col_bench.py # 1024 rows, 5 reps per arm + python fat_col_bench.py --rows 2048 --reps 10 + +Arms (each rep runs in a fresh subprocess so ru_maxrss is a clean per-arm +high-water mark; reps are interleaved pa,pa1,rs to spread machine drift): + pa PyArrow dataset scanner, default threads + pa1 same with use_threads=False — the fair single-thread baseline; + pa1 == pa on this shape proves threads are NOT why PyArrow wins + rs arrow-rs crate at shipping defaults (budget 128 MiB, K=1) +""" +import argparse +import json +import os +import resource +import statistics +import subprocess +import sys +import time + +MiB = 1024 * 1024 +BUDGET = 128 * MiB # shipping default (= DataContext.target_max_block_size) + + +def make_fixture(path, rows): + import numpy as np + import pyarrow as pa + import pyarrow.parquet as pq + + rng = np.random.default_rng(4) # same seed as gen_local_fixtures.gen_fat_col + fat = [rng.bytes(256 * 1024) for _ in range(rows)] + t = pa.table( + { + "fat": pa.array(fat, type=pa.binary()), + "small": pa.array(np.arange(rows, dtype=np.int64)), + } + ) + pq.write_table(t, path, write_page_index=True, row_group_size=rows) + + +def batch_size_for(path): + """Ray's request: budget // encoded-bytes-per-row, floored at 2048 rows.""" + import pyarrow.parquet as pq + + md = pq.read_metadata(path) + total = sum(md.row_group(i).total_byte_size for i in range(md.num_row_groups)) + bpr = max(1, total // max(1, md.num_rows)) + return max(2048, int(BUDGET // bpr)) + + +def leg_pa(path, use_threads): + import pyarrow.dataset as pds + from pyarrow.fs import LocalFileSystem + + fmt = pds.ParquetFileFormat( + default_fragment_scan_options=pds.ParquetFragmentScanOptions(pre_buffer=True) + ) + frag = fmt.make_fragment(path, filesystem=LocalFileSystem()) + scanner = frag.scanner(batch_size=batch_size_for(path), use_threads=use_threads) + rows = nbytes = 0 + for b in scanner.to_batches(): + rows += b.num_rows + nbytes += b.nbytes + return rows, nbytes + + +def leg_rs(path): + import pyarrow as pa + import ray_data_arrow_rs as rs + + handle = rs.open_parquet_file(path, page_index=False) + stream = pa.RecordBatchReader.from_stream( + handle.read_row_groups( + row_groups=None, + columns=None, + batch_size=batch_size_for(path), + decode_budget_bytes=BUDGET, + k=1, + split_threshold_bytes=128 * MiB, + predicate_json=None, + fetch_window_mb=16, + column_fetch_mb=16, + prefetch_budget_mb=64, + ) + ) + rows = nbytes = 0 + for b in stream: + rows += b.num_rows + nbytes += b.nbytes + return rows, nbytes + + +def child(leg, path, rows): + if leg == "create": + if not os.path.exists(path): + make_fixture(path, rows) + print( + json.dumps( + { + "size_mib": os.path.getsize(path) / MiB, + "batch_rows": batch_size_for(path), + } + ) + ) + return + t0 = time.perf_counter() + if leg == "pa": + rows, nb = leg_pa(path, use_threads=True) + elif leg == "pa1": + rows, nb = leg_pa(path, use_threads=False) + else: + rows, nb = leg_rs(path) + wall = time.perf_counter() - t0 + ru = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + peak_mib = ru / MiB if sys.platform == "darwin" else ru / 1024 + print(json.dumps({"wall_s": wall, "peak_rss_mib": peak_mib, "rows": rows})) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--rows", type=int, default=1024) + ap.add_argument("--reps", type=int, default=5) + ap.add_argument("--dir", default=os.path.expanduser("~/fat_col_bench_data")) + ap.add_argument("--leg", choices=["create", "pa", "pa1", "rs"], help="internal") + args = ap.parse_args() + + os.makedirs(args.dir, exist_ok=True) + path = os.path.join(args.dir, f"fat_col_{args.rows}.parquet") + + if args.leg: + child(args.leg, path, args.rows) + return + + # The parent must stay small: on Linux subprocess forks, so each child's + # ru_maxrss high-water starts at the parent's RSS at fork time — building + # the fixture (or importing pyarrow) here would floor every child at the + # parent's footprint and erase the per-arm memory signal. + def spawn(leg): + out = subprocess.run( + [ + sys.executable, + os.path.abspath(__file__), + "--leg", + leg, + "--rows", + str(args.rows), + "--dir", + args.dir, + ], + capture_output=True, + text=True, + check=True, + ) + return json.loads(out.stdout.strip().splitlines()[-1]) + + info = spawn("create") + print(f"file: {path} ({info['size_mib']:.0f} MiB on disk)") + print(f"batch request: {info['batch_rows']} rows, budget {BUDGET // MiB} MiB") + results = {"pa": [], "pa1": [], "rs": []} + for rep in range(args.reps): + for leg in ("pa", "pa1", "rs"): + out = subprocess.run( + [ + sys.executable, + os.path.abspath(__file__), + "--leg", + leg, + "--rows", + str(args.rows), + "--dir", + args.dir, + ], + capture_output=True, + text=True, + check=True, + ) + r = json.loads(out.stdout.strip().splitlines()[-1]) + results[leg].append(r) + print( + f" rep {rep + 1} {leg:>3}: {r['wall_s']:6.3f} s " + f"{r['peak_rss_mib']:7.1f} MiB rows={r['rows']}" + ) + + rows_seen = {r["rows"] for legs in results.values() for r in legs} + assert len(rows_seen) == 1, f"arms read different row counts: {rows_seen}" + + def p50(leg, key): + return statistics.median(r[key] for r in results[leg]) + + print( + f"\n{'arm':>4} {'wall p50 (s)':>12} {'wall min':>9} {'peak RSS p50 (MiB)':>18}" + ) + for leg in ("pa", "pa1", "rs"): + wmin = min(r["wall_s"] for r in results[leg]) + print( + f"{leg:>4} {p50(leg, 'wall_s'):12.3f} {wmin:9.3f} " + f"{p50(leg, 'peak_rss_mib'):18.1f}" + ) + print( + f"\nrs / pa1 (single-thread kernel gap): " + f"wall {p50('rs', 'wall_s') / p50('pa1', 'wall_s'):.2f}x " + f"peak RSS {p50('rs', 'peak_rss_mib') / p50('pa1', 'peak_rss_mib'):.2f}x" + ) + print("expected on Linux (M69/M70): wall ~1.7-2.7x slower, RSS ~0.4-0.5x") + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/gen_2x2_release_tests.py b/release/nightly_tests/dataset/arrow_rs_probe/gen_2x2_release_tests.py new file mode 100644 index 000000000000..ae7a5d6a746c --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/gen_2x2_release_tests.py @@ -0,0 +1,462 @@ +"""Generate the arrow-rs 2x2 release-test entries (item 29). + +Every A/B so far compared multi-node fleets only, and none of the ~20 A/B #5 +P0 regressions reproduces single-node (M78) -- so the surviving suspects are +things only the release regime has: long-lived-worker allocator retention and +autoscaler pool dynamics. This generator builds the discriminating experiment: +each regressed test runs as FOUR release entries on ONE branch/image, + + {arrow-rs, pyarrow} x {its original fleet, one fat node} + +so topology is isolated with zero branch or image skew. Reader is toggled +per-entry by prefixing the run script with RAY_DATA_USE_ARROW_RS_PARQUET_READER +(env_bool in python/ray/data/context.py honors the env var over the branch's +flipped-True default; the flag branches inside the read task, so tasks -- not +images -- decide). Topology is toggled per-entry via cluster_compute: +single-node cells swap the fleet yaml for single_node_{cpu,all_to_all}_compute +.yaml (one m5.24xlarge, 96 vCPU / 384 GiB -- at or above the aggregate of the +m5.2xlarge CPU fleets; the 512-vCPU all-to-all and 800-vCPU joins fleets have +no single-node equal, which under-provisions BOTH arms equally). + +What the generated entries change vs their parents, and nothing else: + * name _2x2_{rs|pa}_{multi|single} + * frequency manual (never scheduled; trigger explicitly) + * script reader env-var prefix + * single cells only: single-node compute yaml, RAYTEST_FAIL_ON_SPILLING=0 + (one node spills where a fleet spreads), wait_for_nodes dropped, + timeout doubled (capped at 4h). + +Reading results: the pa_multi cell must reproduce the A/B baseline (5-run +history) -- if it doesn't, the branch or image contaminated the experiment and +the whole run is void. Then: regression present in multi but absent in single, +with the decoder eliminators (decoded bytes/task, peak batch bytes) at parity +=> release-regime mechanism (retention / pool dynamics), not the decoder. + +Two matrices share this machinery (--matrix); the name scheme stays +_2x2__ for both so one results-DB name history and one +Buildkite filter (name:.*_2x2_.*) cover them: + + 2x2 {rs, pa} x {multi, single} over every target (item 29; build 105711). + alloc DEFAULT. {pa, rs} on the original fleet over the memory, sustained + and control targets plus three wall targets kept as the release-scale + confirmation of the M97/M98 wall fix; single-node cells only where + the single-node reading is itself the question (ALLOC_SINGLE_TOO); + the two cells that must replicate inside one build run x3 + (ALLOC_REPEATED). History: builds 106096 / 106284 ran this matrix + with two extra arms, rstrim (RAY_DATA_ARROW_RS_MALLOC_TRIM=1, the + mallopt mechanism probe, retired M108) and rseos + (RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS=1, one malloc_trim(0) per read + task stream). rseos closed every retention row at ~rs wall (M101) and + its one wall outlier did not replicate (M124), so on 2026-09-08 the + eos trim became the arrow-rs default (context.py) -- the rs arm now + IS the former rseos arm and the extra arms are gone. Reading: pa + cells must reproduce the build history or the window is + contaminated; decoded-bytes/peak-batch dists must match across arms + (the trim touches no decode path); any rs/pa ratio >= 1.15 is a fix + item (arrow_rs_docs/findings.md M104 gates). + +Usage (from anywhere; rewrites its marker block in release_data_tests.yaml): + python gen_2x2_release_tests.py # regenerate alloc + validate + python gen_2x2_release_tests.py --matrix 2x2 # the original 2x2 instead + python gen_2x2_release_tests.py --check # validate the file as it is + +The resolver below (deep_update / matrix / variations / {{var}} substitution) +mirrors release/ray_release/config.py:parse_test_definition, which is not +importable outside Bazel (jsonschema/runfiles deps). +""" + +import argparse +import copy +import itertools +import os +import re + +import yaml + +HERE = os.path.dirname(os.path.abspath(__file__)) +RELEASE_DIR = os.path.abspath(os.path.join(HERE, "..", "..", "..")) +TESTS_YAML = os.path.join(RELEASE_DIR, "release_data_tests.yaml") +DATASET_DIR = os.path.join(RELEASE_DIR, "nightly_tests", "dataset") + +BEGIN = "# === BEGIN arrow-rs 2x2 (generated by arrow_rs_probe/gen_2x2_release_tests.py; do not edit by hand) ===" +END = "# === END arrow-rs 2x2 ===" + +READER_ENV = "RAY_DATA_USE_ARROW_RS_PARQUET_READER" +TIMEOUT_CAP_S = 14400 + +# The A/B #5 P0 ledger (arrow_rs_docs/2026-08-27.md section 11; names as the +# results DB spells them, '+' where the yaml matrix value has spaces), plus +# three controls. Category is documentation only -- all cells run identically. +TARGETS = { + # wall-time regressions + "tpch_q10_autoscaling_hash_shuffle_v2": "wall", + "tpch_q3_autoscaling_hash_shuffle_v2": "wall", + "tpch_q22_autoscaling_hash_shuffle": "wall", + "tpch_q14_fixed_size_hash_shuffle": "wall", + "tpch_q18_fixed_size_hash_shuffle_v2": "wall", + "joins_sf100_right_outer": "wall", + "read_parquet_autoscaling": "wall", + "map_groups_autoscaling_sort_shuffle_pull_based_column02+column14": "wall", + "map_groups_autoscaling_hash_shuffle_column02+column14": "wall", + # peak-memory regressions (tUSS / wUSS peak) + "read_large_parquet_autoscaling": "memory", + "read_large_parquet_fixed_size": "memory", + "write_parquet": "memory", + "tpch_q6_fixed_size_hash_shuffle": "memory", + "tpch_q6_fixed_size_hash_shuffle_v2": "memory", + "tpch_q17_fixed_size_hash_shuffle": "memory", + "tpch_q17_fixed_size_hash_shuffle_v2": "memory", + "map_groups_autoscaling_hash_shuffle_column08+column13+column14": "memory", + "wide_schema_pipeline_objects": "memory", + # sustained-wUSS (retention signature, M77) + "map_groups_fixed_size_hash_shuffle_v2_column08+column13+column14": "sustained", + "map_groups_fixed_size_hash_shuffle_v2_column02+column14": "sustained", + "map_groups_fixed_size_sort_shuffle_pull_based_column02+column14": "sustained", + "joins_sf100_inner": "sustained", + "joins_sf100_left_outer": "sustained", + "joins_sf100_full_outer": "sustained", + "wide_schema_pipeline_primitives": "sustained", + "mix.8ds_equal": "sustained", + # controls: parity or arrow-rs-win in A/B #5 -- if these move, suspect the + # experiment, not the reader + "aggregate_groups_autoscaling_hash_shuffle_column08+column13+column14": "control", + "tpch_q12_autoscaling_hash_shuffle": "control", + "iter_batches_pyarrow": "control", +} + +# Fleet -> single-fat-node analog. all_to_all keeps the big-disk variant (one +# node absorbs the spill a fleet spreads); joins' 100-node fleet maps there +# too for the same reason. +SINGLE_COMPUTE = { + "fixed_size_all_to_all_compute.yaml": "single_node_all_to_all_compute.yaml", + "autoscaling_all_to_all_compute.yaml": "single_node_all_to_all_compute.yaml", + "fixed_size_100_cpu_compute.yaml": "single_node_all_to_all_compute.yaml", + "fixed_size_cpu_compute.yaml": "single_node_cpu_compute.yaml", + "autoscaling_cpu_compute.yaml": "single_node_cpu_compute.yaml", + "dataset_mixing/compute_8_cpu.yaml": "single_node_cpu_compute.yaml", +} + +# Arm -> the env prefix prepended to the parent's run script. The rs arm is +# the reader as shipped: since 2026-09-08 that includes the end-of-stream +# malloc_trim (python/ray/data/context.py DEFAULT_ARROW_RS_MALLOC_TRIM_EOS = +# True), so no allocator knob is spelled out here. Add an explicit +# RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS=0 arm only to ablate it. +ARMS = { + "pa": f"{READER_ENV}=0", + "rs": f"{READER_ENV}=1", +} + +# alloc matrix: wall targets dropped except these three -- the release-scale +# confirmation of the M97/M98 per-task-S3-client fix (box: 2.88->0.88, +# 2.15->0.77, 1.22->0.90) on the three wall shapes with the cleanest history. +ALLOC_WALL_KEEP = ( + "read_parquet_autoscaling", + "tpch_q18_fixed_size_hash_shuffle_v2", + "map_groups_autoscaling_hash_shuffle_column02+column14", +) +# alloc matrix: single-node cells only where single-node IS the question -- +# the rlp per-task USS self-regression (M75/M86) and the col02 OOM cliff (M90). +ALLOC_SINGLE_TOO = ( + "read_large_parquet_autoscaling", + "map_groups_autoscaling_hash_shuffle_column02+column14", +) + +# alloc matrix: gate cells that must replicate inside ONE build (TODO 34f: the +# col02 rseos fleet cell ran 2.37 once; one run cannot separate a slowed task +# from an autoscaler that packed 44 workers/node) -- emitted with +# ``repeated_run`` so the release runner schedules N steps per cell (the form +# has no repeat field; release/ray_release/buildkite/step.py reads this key). +ALLOC_REPEATED = { + "map_groups_autoscaling_hash_shuffle_column02+column14": 3, + "read_large_parquet_fixed_size": 3, +} + +MATRICES = ("2x2", "alloc") +DEFAULT_MATRIX = "alloc" + + +def matrix_cells(matrix): + """(target, arm, topology) triples for one matrix, in emitted order.""" + if matrix == "2x2": + return [ + (t, a, s) + for t in TARGETS + for a in ("rs", "pa") + for s in ("multi", "single") + ] + if matrix == "alloc": + cells = [] + for target, category in TARGETS.items(): + if category == "wall" and target not in ALLOC_WALL_KEEP: + continue + topologies = ( + ("multi", "single") if target in ALLOC_SINGLE_TOO else ("multi",) + ) + cells += [(target, a, s) for a in ARMS for s in topologies] + return cells + raise ValueError(f"unknown matrix {matrix!r}; one of {MATRICES}") + + +def parse_2x2_name(name): + """'foo_2x2_rstrim_single' -> ('foo', 'rstrim', 'single'); None if not ours.""" + m = re.fullmatch(rf"(.+)_2x2_({'|'.join(ARMS)})_(multi|single)", name) + return (m.group(1), m.group(2), m.group(3)) if m else None + + +def block_matrix(text): + """Which matrix the file's marker block was generated with (None if absent).""" + m = re.search(rf"^{re.escape(BEGIN)}\n# matrix: (\w+)$", text, re.M) + return m.group(1) if m else None + + +# -------------------------------------------------------------------------- +# Vendored from release/ray_release/config.py + util.py (see module docstring) +# -------------------------------------------------------------------------- + + +def _deep_update(d, u): + for k, v in u.items(): + if isinstance(v, dict): + d[k] = _deep_update(d.get(k, {}), v) + else: + d[k] = v + return d + + +def _substitute_variable(data, variable, replacement): + data = copy.deepcopy(data) + pattern = r"\{\{\s*" + re.escape(variable) + r"\s*\}\}" + for key, value in data.items(): + if isinstance(value, dict): + data[key] = _substitute_variable(value, variable, replacement) + elif isinstance(value, list): + data[key] = [re.sub(pattern, replacement, s) for s in value] + elif isinstance(value, str): + data[key] = re.sub(pattern, replacement, value) + return data + + +def resolve_tests(test_definitions): + """parse_test_definition, minus jsonschema validation and the Test class.""" + defaults, tests = {}, [] + for test_definition in test_definitions: + if test_definition["name"] == "DEFAULTS": + defaults = copy.deepcopy(test_definition) + continue + test_definition = _deep_update(copy.deepcopy(defaults), test_definition) + if "variations" in test_definition: + variations = test_definition.pop("variations") + for variation in variations: + variation = copy.deepcopy(variation) + test = copy.deepcopy(test_definition) + test["name"] = f'{test["name"]}.{variation.pop("__suffix__")}' + tests.append(_deep_update(test, variation)) + elif "matrix" in test_definition: + matrix = test_definition.pop("matrix") + variables = tuple(matrix["setup"].keys()) + for combination in itertools.product(*matrix["setup"].values()): + test = test_definition + for variable, value in zip(variables, combination): + test = _substitute_variable(test, variable, str(value)) + tests.append(test) + for adjustment in matrix.pop("adjustments", []): + test = test_definition + for variable, value in adjustment["with"].items(): + test = _substitute_variable(test, variable, str(value)) + tests.append(test) + else: + tests.append(test_definition) + return defaults, tests + + +# -------------------------------------------------------------------------- +# Entry generation +# -------------------------------------------------------------------------- + + +def _strip_defaults(entry, defaults): + """Drop what DEFAULTS will deep_update back in identically at parse time.""" + entry = copy.deepcopy(entry) + for key in ("group", "working_dir", "team"): + if entry.get(key) == defaults.get(key): + entry.pop(key, None) + byod = entry.get("cluster", {}).get("byod", {}) + default_byod = defaults.get("cluster", {}).get("byod", {}) + for key in ("type", "post_build_script"): + if byod.get(key) == default_byod.get(key): + byod.pop(key, None) + return entry + + +def make_cell(resolved, defaults, arm, topology, matrix=DEFAULT_MATRIX): + """One generated entry: parent test x arm (env prefix) x topology.""" + base = _strip_defaults(resolved, defaults) + base_name = base["name"].replace(" ", "+") + compute = base["cluster"]["cluster_compute"] + if compute not in SINGLE_COMPUTE: + raise ValueError(f"{base_name}: no single-node analog for {compute}") + + cell = copy.deepcopy(base) + cell["name"] = f"{base_name}_2x2_{arm}_{topology}" + cell["frequency"] = "manual" + # Keys use the '+' form (parent names carry spaces), like TARGETS. + if matrix == "alloc" and base_name in ALLOC_REPEATED: + cell["repeated_run"] = ALLOC_REPEATED[base_name] + # Parent scripts come from ">" folded scalars and carry a trailing + # newline; strip so the emitted string is single-line. + script = cell["run"]["script"].strip() + cell["run"]["script"] = f"{ARMS[arm]} {script}" + if topology == "single": + cell["cluster"]["cluster_compute"] = SINGLE_COMPUTE[compute] + env = cell["cluster"].get("byod", {}).get("runtime_env") + if env is not None: + cell["cluster"]["byod"]["runtime_env"] = [ + "RAYTEST_FAIL_ON_SPILLING=0" + if e.startswith("RAYTEST_FAIL_ON_SPILLING=") + else e + for e in env + ] + cell["run"].pop("wait_for_nodes", None) + cell["run"]["timeout"] = min(cell["run"]["timeout"] * 2, TIMEOUT_CAP_S) + # Reorder for readable yaml; dicts keep insertion order. + return { + key: cell[key] + for key in ("name", "frequency", "repeated_run", "python", "cluster", "run") + if key in cell + } + + +_BLOCK_HEADERS = { + "2x2": [ + "# 4 cells per A/B #5 P0 regression: {rs,pa} x {original fleet, one fat", + "# node}.", + ], + "alloc": [ + "# {pa, rs} on the original fleet over the memory / sustained / control", + "# targets + 3 wall-fix confirmations; single cells only for the two shapes", + "# whose single-node reading is the question; two gate cells x3. rs = the", + "# arrow-rs reader as shipped, which since 2026-09-08 includes the", + "# end-of-stream malloc_trim (former rseos arm; rstrim retired, M108).", + ], +} + + +def render_block(defaults, resolved_tests, matrix): + by_name = {t["name"].replace(" ", "+"): t for t in resolved_tests} + cells = matrix_cells(matrix) + missing = sorted({t for t, _, _ in cells} - set(by_name)) + if missing: + raise SystemExit(f"targets not found in release_data_tests.yaml: {missing}") + + lines = [BEGIN, f"# matrix: {matrix}"] + lines += _BLOCK_HEADERS[matrix] + lines += [ + "# frequency:manual -- trigger explicitly, all arms in one window (M75:", + "# readings drift across windows). Regenerate with:", + "# python nightly_tests/dataset/arrow_rs_probe/gen_2x2_release_tests.py" + + ("" if matrix == DEFAULT_MATRIX else f" --matrix {matrix}"), + "", + ] + current = None + for target, arm, topology in cells: + if target != current: + current = target + lines.append(f"# --- {target} [{TARGETS[target]}] ---") + cell = make_cell(by_name[target], defaults, arm, topology, matrix) + lines.append( + yaml.safe_dump( + [cell], sort_keys=False, default_flow_style=False, width=88 + ).rstrip() + ) + lines.append("") + lines.append(END) + return "\n".join(lines) + "\n" + + +def strip_block(text): + if BEGIN not in text: + return text + head, rest = text.split(BEGIN, 1) + if END not in rest: + raise SystemExit("found BEGIN marker without END marker; fix by hand") + _, tail = rest.split(END, 1) + return head.rstrip("\n") + "\n" + tail.lstrip("\n") + + +def validate(defaults, resolved_tests, matrix): + names = [t["name"] for t in resolved_tests] + dupes = sorted({n for n in names if names.count(n) > 1}) + assert not dupes, f"duplicate resolved test names: {dupes}" + + generated = [t for t in resolved_tests if parse_2x2_name(t["name"])] + expected = {f"{t}_2x2_{a}_{s}" for t, a, s in matrix_cells(matrix)} + got = {t["name"] for t in generated} + assert got == expected, ( + f"generated set mismatch: missing={sorted(expected - got)} " + f"extra={sorted(got - expected)}" + ) + for test in generated: + name = test["name"] + _, arm, topology = parse_2x2_name(name) + assert test["frequency"] == "manual", name + assert test["run"]["script"].startswith(f"{ARMS[arm]} "), name + # The trim is the reader's default now; an allocator knob spelled into + # a cell would silently turn the matrix back into an ablation. + assert "MALLOC_TRIM" not in test["run"]["script"], name + compute = test["cluster"]["cluster_compute"] + assert os.path.exists( + os.path.join(DATASET_DIR, compute) + ), f"{name}: compute file missing: {compute}" + if topology == "single": + assert compute.startswith("single_node_"), name + assert "wait_for_nodes" not in test["run"], name + env = test["cluster"].get("byod", {}).get("runtime_env", []) + assert "RAYTEST_FAIL_ON_SPILLING=1" not in env, name + # DEFAULTS must still supply the crate build + node-mem monitor. + assert test["cluster"]["byod"]["post_build_script"] == ( + defaults["cluster"]["byod"]["post_build_script"] + ), name + env = test["cluster"]["byod"]["runtime_env"] + assert "RAY_DATA_BENCH_NODE_MEM_MONITOR=1" in env, name + return len(generated) + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--check", action="store_true", help="validate the current file, do not rewrite" + ) + parser.add_argument( + "--matrix", + choices=MATRICES, + help=f"which matrix to emit (default {DEFAULT_MATRIX}); with --check, " + "defaults to the one recorded in the file's block header", + ) + args = parser.parse_args() + + text = open(TESTS_YAML).read() + matrix = ( + args.matrix or (block_matrix(text) if args.check else None) or DEFAULT_MATRIX + ) + if not args.check: + stripped = strip_block(text) + defaults, resolved = resolve_tests(yaml.safe_load(stripped)) + block = render_block(defaults, resolved, matrix) + text = stripped.rstrip("\n") + "\n\n" + block + with open(TESTS_YAML, "w") as fh: + fh.write(text) + + defaults, resolved = resolve_tests(yaml.safe_load(open(TESTS_YAML))) + count = validate(defaults, resolved, matrix) + cells = matrix_cells(matrix) + n_targets = len({t for t, _, _ in cells}) + n_single = sum(1 for _, _, s in cells if s == "single") + print( + f"OK [{matrix}]: {count} generated entries ({n_targets} targets, " + f"{len(ARMS)} arms, {count - n_single} multi + " + f"{n_single} single), {len(resolved)} resolved tests total, all names unique" + ) + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/gen_local_fixtures.py b/release/nightly_tests/dataset/arrow_rs_probe/gen_local_fixtures.py new file mode 100644 index 000000000000..98c51e418348 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/gen_local_fixtures.py @@ -0,0 +1,549 @@ +#!/usr/bin/env python3 +"""Generate the local Parquet fixture shapes for grand_experiment.py. + +Each shape isolates one axis of the NEW footer-based planner (#64985 series: +footers read up front, row groups pruned then bin-packed by uncompressed size into +read tasks) and/or one known arrow-rs behaviour (the replication shapes bin_sweep / +tensors_wide / tensors_cp are documented on their generators below): + + lone_big_rg 1 file, ONE ~800 MiB uncompressed row group. Unsplittable by any + planner (a row group is the atom). arrow-rs K-splits it internally; + PyArrow must materialize the whole decoded group. The headline + memory shape. + single_rg_files N files x one ~128 MiB row group each (the parquet_split / + release-regression shape). Under the new planner these PACK — + several files' groups can share one bin/read task. Exercises the + per-native-call fixed cost and the bin-size knob. + tiny_rgs blob column, MANY ~2 MiB row groups per file (imagenet/rg_50k + shape). Bin packing coalesces dozens of groups per task; without + row_hash our reader issues ONE native call per bin (the coalescing + win). Also the shape where allocator retention showed up. + wide 2000 float64 columns, one row group per file (wide_schema shape; + mem was 1.50x worse in the old release run). + fat_col one ~256 KiB/row binary column + one tiny int column, single row + group. The shape that mis-selected the crate's S3 column-group path + (Hstack whole-row-group retention). Local decode control here; + the S3 stage is where it bites. + +All files are written with write_page_index=True (crate requirement) and mostly +incompressible data so on-disk bytes ~ uncompressed bytes (S3 fetch cost stays +honest when synced up). + + python gen_local_fixtures.py --root ~/arrow_rs_grand_fixtures + python gen_local_fixtures.py --root ... --shapes lone_big_rg,wide # subset + python gen_local_fixtures.py --root ... --scale 0.25 # smaller/faster + +Writes /manifest.json mapping shape -> path + stats; grand_experiment.py +reads that. Idempotent per shape: skips a shape whose directory already exists. +""" +import argparse +import json +import os +import shutil + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +# ~34 B/row of int64 + ~256 B of random-ish string payload -> ~260 B/row uncompressed. +_STR_COLS = 4 +_STR_LEN = 64 +_POOL_SIZE = 4096 + + +def _str_pool(rng): + alphabet = np.frombuffer(b"abcdefghijklmnopqrstuvwxyz0123456789", dtype="S1") + idx = rng.integers(0, len(alphabet), size=(_POOL_SIZE, _STR_LEN)) + return np.array([alphabet[row].tobytes().decode() for row in idx], dtype=object) + + +def _string_table(rng, pool, n_rows, id_start=0): + cols = {"id": pa.array(np.arange(id_start, id_start + n_rows, dtype=np.int64))} + for c in range(_STR_COLS): + cols[f"s{c}"] = pa.array( + pool[rng.integers(0, _POOL_SIZE, size=n_rows)], type=pa.string() + ) + return pa.table(cols) + + +def gen_lone_big_rg(d, scale): + """One file, one big row group (~800 MiB uncompressed at scale=1).""" + rng = np.random.default_rng(0) + pool = _str_pool(rng) + n_rows = int(3_200_000 * scale) + t = _string_table(rng, pool, n_rows) + pq.write_table( + t, + os.path.join(d, "part0.parquet"), + write_page_index=True, + row_group_size=n_rows, # >= n_rows -> single row group + ) + return {"files": 1, "rows": n_rows, "uncompressed_bytes": t.nbytes} + + +def gen_single_rg_files(d, scale): + """16 files x one ~128 MiB row group each (parquet_split shape).""" + rng = np.random.default_rng(1) + pool = _str_pool(rng) + n_files = max(2, int(16 * scale)) + rows_per = 500_000 # ~128 MiB uncompressed per file + total = 0 + for f in range(n_files): + t = _string_table(rng, pool, rows_per, id_start=f * rows_per) + pq.write_table( + t, + os.path.join(d, f"part{f}.parquet"), + write_page_index=True, + row_group_size=rows_per, + use_dictionary=False, + ) + total += t.nbytes + return { + "files": n_files, + "rows": n_files * rows_per, + "uncompressed_bytes": total, + "v": 2, + } + + +def gen_auto_rg(d, scale): + """48 files x one ~69 MiB row group each — the M31 shape. + + `read_large_parquet_autoscaling` (release A/B #3 loss M31) reads + s3://.../large-parquet/: ~103 files whose row groups are ~69 MiB uncompressed, + under RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 (64 MiB) — a row group is + the packing atom and already exceeds the bin, so every read task is exactly + ONE ~69 MiB row group and finishes in under a second. That sub-second task is + why the release per-task "max USS" is an end-of-task sample, not a peak + (2026-08-15.md §4b) — the loss_triage stage re-measures it at a 20 Hz poll. + `single_rg_files` (128 MiB) is the parquet_split shape; this one exists so the + per-task decode volume matches M31's. + + v2 (2026-08-18): written with use_dictionary=False. The bin packer budgets + `rg.total_byte_size` = uncompressed-but-still-ENCODED page bytes, and the + 4096-string pool dictionary-encoded ~14.6x (4.9 MiB encoded vs 72 MiB + decoded) — so v1 packed ~13 files per 64 MiB bin and the promised + one-row-group-per-task shape NEVER held (M41). Plain encoding makes + total_byte_size ~= decoded bytes and restores the release shape. + """ + rng = np.random.default_rng(8) + pool = _str_pool(rng) + n_files = max(4, int(48 * scale)) + rows_per = 270_000 # ~69 MiB uncompressed at ~260 B/row + total = 0 + for f in range(n_files): + t = _string_table(rng, pool, rows_per, id_start=f * rows_per) + pq.write_table( + t, + os.path.join(d, f"part{f}.parquet"), + write_page_index=True, + row_group_size=rows_per, + use_dictionary=False, + ) + total += t.nbytes + return { + "files": n_files, + "rows": n_files * rows_per, + "uncompressed_bytes": total, + "v": 2, + } + + +gen_auto_rg._fixture_version = 2 + + +def gen_tiny_rgs(d, scale): + """4 files, blob column, ~2 MiB row groups (32 rows x 64 KiB blobs).""" + rng = np.random.default_rng(2) + n_files = 4 + rows_per = max(64, int(1024 * scale)) + blob_bytes = 64 * 1024 + total = 0 + for f in range(n_files): + blobs = [rng.bytes(blob_bytes) for _ in range(rows_per)] + t = pa.table( + { + "id": pa.array( + np.arange(f * rows_per, (f + 1) * rows_per, dtype=np.int64) + ), + "image": pa.array(blobs, type=pa.binary()), + "label": pa.array(rng.integers(0, 1000, rows_per, dtype=np.int64)), + } + ) + pq.write_table( + t, + os.path.join(d, f"part{f}.parquet"), + write_page_index=True, + row_group_size=32, # ~2 MiB per row group + ) + total += t.nbytes + return {"files": n_files, "rows": n_files * rows_per, "uncompressed_bytes": total} + + +def gen_wide(d, scale): + """4 files, 2000 float64 columns, one row group per file (~64 MiB each).""" + rng = np.random.default_rng(3) + n_files = 4 + n_cols = 2000 + rows_per = max(512, int(4096 * scale)) + total = 0 + for f in range(n_files): + cols = {f"c{c}": rng.random(rows_per) for c in range(n_cols)} + t = pa.table(cols) + pq.write_table( + t, + os.path.join(d, f"part{f}.parquet"), + write_page_index=True, + row_group_size=rows_per, + use_dictionary=False, + ) + total += t.nbytes + return { + "files": n_files, + "rows": n_files * rows_per, + "uncompressed_bytes": total, + "v": 2, + } + + +def gen_fat_col(d, scale): + """1 file: one ~256 KiB/row binary column + one int column, single row group.""" + rng = np.random.default_rng(4) + n_rows = max(128, int(1024 * scale)) + fat = [rng.bytes(256 * 1024) for _ in range(n_rows)] + t = pa.table( + { + "fat": pa.array(fat, type=pa.binary()), + "small": pa.array(np.arange(n_rows, dtype=np.int64)), + } + ) + pq.write_table( + t, + os.path.join(d, "part0.parquet"), + write_page_index=True, + row_group_size=n_rows, + ) + return {"files": 1, "rows": n_rows, "uncompressed_bytes": t.nbytes} + + +def gen_bin_sweep(d, scale): + """8 files x 8 row groups x ~64 MiB each (~512 MiB/file, ~4 GiB total at scale=1). + + The bin-sweep fixture (TODO 1ab/R2, item 10): sized so the sweep grid can span + all three C9 regimes — sub-file bins (1 RG, 4 RGs), exactly one file, and + MULTI-FILE bins (5x / 10x the file size — the release yaml's 1 GiB bin never + crossed a file boundary, so mechanism (i), N sub-fragments on the base's + unbounded thread pool, has never been measured). replication_matrix.py derives + the actual byte grid from this manifest entry's rg/file stats, so --scale + changes sizes without breaking the regimes. + + v2 (2026-08-18): use_dictionary=False, same reason as gen_auto_rg (M41) — + v1's encoded rgs were ~4.9 MiB, so the whole fixture fit ONE release-yaml + write bin (1.29 GiB) and the M32/M38 task shape never ran on a box. + """ + rng = np.random.default_rng(5) + pool = _str_pool(rng) + n_files = 8 + rgs_per_file = 8 + rows_per_rg = max(1024, int(258_000 * scale)) # ~64 MiB uncompressed at scale=1 + total = 0 + for f in range(n_files): + n_rows = rgs_per_file * rows_per_rg + t = _string_table(rng, pool, n_rows, id_start=f * n_rows) + pq.write_table( + t, + os.path.join(d, f"part{f}.parquet"), + write_page_index=True, + row_group_size=rows_per_rg, + use_dictionary=False, + ) + total += t.nbytes + return { + "files": n_files, + "rows": n_files * rgs_per_file * rows_per_rg, + "uncompressed_bytes": total, + "rgs_per_file": rgs_per_file, + "rows_per_rg": rows_per_rg, + "v": 2, + } + + +gen_bin_sweep._fixture_version = 2 + + +def gen_tensors_wide(d, scale): + """4 files, 5000 fixed_size_list columns (~40 KiB/row, ~40 MiB RGs). + + Local lookalike for the wide_schema_pipeline_tensors regression (T15, item 1y: + native decode task-seconds 5.59x). The release dataset lives in + ray-benchmark-data-internal-* (unreadable to us), so parity is mechanism-level: + many small fixed-size-list columns decoded natively. Plain storage type, no + extension metadata — the release run confirmed the native (non-fallback) path, + and extension-tagged columns would fall back and measure nothing. + """ + rng = np.random.default_rng(6) + n_files = 4 + n_cols = 5000 + list_size = 2 + rows_per_file = max(256, int(10_000 * scale)) + row_group_size = max(64, int(1_000 * scale)) # ~40 MiB uncompressed at scale=1 + total = 0 + for f in range(n_files): + cols = {} + for c in range(n_cols): + flat = rng.random(rows_per_file * list_size, dtype=np.float32) + cols[f"t{c}"] = pa.FixedSizeListArray.from_arrays( + pa.array(flat, type=pa.float32()), list_size + ) + t = pa.table(cols) + pq.write_table( + t, + os.path.join(d, f"part{f}.parquet"), + write_page_index=True, + row_group_size=row_group_size, + ) + total += t.nbytes + return { + "files": n_files, + "rows": n_files * rows_per_file, + "uncompressed_bytes": total, + "columns": n_cols, + } + + +def gen_tensors_cp(d, scale, _pool=None, _rows=None, _rg=None, _files=None): + """4 files, 5000 tensor columns with **cloudpickle** extension metadata. + + The fixture that actually reproduces T15/1y (T22): `tensors_wide` above is + the 2026-08-12 lookalike that decodes *faster* native (T19), because the + release dataset was written by Ray 2.49-2.54 with cloudpickle-serialized + tensor metadata — non-UTF8 bytes inside ``ARROW:schema`` that the crate's + IPC verifier rejects. That flips the reader onto its skip+realign path + (decode parquet storage types, re-read the footer via pyarrow, cast every + column storage→extension per batch), which a plain-storage fixture never + touches. Storage is variable-size ``list`` (Ray's ArrowTensorArray), + also unlike the lookalike's ``fixed_size_list``. + + Any run against this fixture needs + ``RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA=1`` (the release yaml sets + it for the same reason); the [tensorscp] stage passes it per cell. + """ + # Writing the legacy metadata format requires the in-process switch; keep + # the import and the flip inside the generator so the other shapes never + # depend on ray.data internals. + import ray.data._internal.tensor_extensions.arrow as tx + from ray.data.extensions import ArrowTensorArray + + prev = tx.ARROW_EXTENSION_SERIALIZATION_FORMAT + tx.ARROW_EXTENSION_SERIALIZATION_FORMAT = tx._SerializationFormat.CLOUDPICKLE + try: + rng = np.random.default_rng(7) + n_files = 4 if _files is None else _files + n_cols = 5000 + rows_per_file = max(200, int(400 * scale)) if _rows is None else _rows + row_group_size = max(100, int(200 * scale)) if _rg is None else _rg + total = 0 + for f in range(n_files): + cols = {} + for c in range(n_cols): + if _pool is None: + vals = rng.random((rows_per_file, 2, 2), dtype=np.float32) + else: + # Low-entropy leaves: draw from a small pool so the f32 + # leaf columns RLE_DICTIONARY-encode (see gen_tensors_dict). + vals = _pool[rng.integers(0, len(_pool), (rows_per_file, 2, 2))] + cols[f"t{c}"] = ArrowTensorArray.from_numpy(vals) + t = pa.table(cols) + pq.write_table( + t, + os.path.join(d, f"part{f}.parquet"), + write_page_index=True, + row_group_size=row_group_size, + ) + total += t.nbytes + finally: + tx.ARROW_EXTENSION_SERIALIZATION_FORMAT = prev + + # The whole point is the crate's skip path — fail loudly if the crate can + # parse this file's embedded schema after all (the fixture would then + # silently measure the ordinary native path, T19's mistake in reverse). + # NB: don't try to check the ARROW:schema *bytes* instead — pyarrow + # base64-encodes that value, so it always decodes as ASCII; only the + # crate's own verdict distinguishes the two paths. + try: + import ray_data_arrow_rs as rs + except ImportError: + print(" tensors_cp: crate not importable here — skip-path check deferred") + else: + handle = rs.open_parquet_file( + os.path.join(d, "part0.parquet"), page_index=False + ) + if not getattr(handle.metadata(), "arrow_schema_skipped", False): + raise SystemExit( + "tensors_cp: crate parsed the embedded arrow schema (no skip) — " + "this fixture no longer reproduces 1y's realign path" + ) + return { + "files": n_files, + "rows": n_files * rows_per_file, + "uncompressed_bytes": total, + "columns": n_cols, + } + + +def gen_tensors_dict(d, scale): + """tensors_cp with DICTIONARY-COMPRESSIBLE leaves — the release wide_schema + tensors shape's missing property (M43). + + The release tensors data must dictionary-encode ~9x — implied by A/B #4's + own numbers: arrow-rs's peak yielded batch was 288 MiB on a 32 MiB decode + budget, and the crate's byte budget divides by ENCODED footer bytes/row + (`byte_budget_rows`, src/lib.rs), so batch bytes = budget x (decoded/encoded). + `tensors_cp`'s random floats encode ~1.1x, which is why no box run ever + reproduced M39's fat batches. Here leaves come from a 16-value pool + (4-bit dictionary indices -> roughly the release's expansion; the achieved + ratio is measured and recorded in the manifest), and rows-per-file is + raised so a 2048-row batch can actually materialize (release row groups + are ~1.6k rows; tensors_cp's 200-row files cap any batch at 200 rows). + Runs need RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA=1, same as + tensors_cp. + """ + return _gen_tensors_pooled(d, scale, pool_size=16) + + +def _gen_tensors_pooled(d, scale, pool_size): + """Shared body for the expansion-sweep shapes: tensors_cp with leaves drawn + from a ``pool_size``-value pool. Smaller pool -> narrower dictionary + indices -> larger decoded/encoded expansion. The achieved ratio is + measured from the footers and recorded in the manifest + (``enc_to_dec_ratio``) — the ablation reads the measured value, so the + pool size only needs to land in the right regime, not hit a target.""" + rng = np.random.default_rng(11) + pool = rng.random(pool_size, dtype=np.float32) + rows = max(2560, int(2560 * scale)) + stats = gen_tensors_cp(d, scale, _pool=pool, _rows=rows, _rg=rows, _files=2) + # Footer reads hydrate the (in-process-registered) extension schema, which + # needs the cloudpickle opt-in; flip the module flag around them, mirroring + # the write-side format flip above. Decoded bytes are already in + # stats["uncompressed_bytes"] (sum of Table.nbytes at write time). + import ray.data._internal.tensor_extensions.arrow as tx + + prev = tx._AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA + tx._AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA = True + try: + enc = 0 + for f in sorted(os.listdir(d)): + if f.endswith(".parquet"): + md = pq.read_metadata(os.path.join(d, f)) + enc += sum( + md.row_group(i).total_byte_size for i in range(md.num_row_groups) + ) + finally: + tx._AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA = prev + stats["enc_to_dec_ratio"] = round(stats["uncompressed_bytes"] / max(1, enc), 2) + print( + f" tensors(pool={pool_size}): enc->dec expansion {stats['enc_to_dec_ratio']}x" + ) + return stats + + +gen_tensors_dict._fixture_version = 1 + + +def gen_tensors_lo(d, scale): + """Expansion-sweep LOW point: ~2-3x decoded/encoded (4096-value pool). + With tensors_cp (~1.1x), tensors_dict (~9.4x) and tensors_hi (~15x+) this + gives four points to check that batch overshoot scales linearly with the + expansion ratio (M43) and that a fix holds across the whole range.""" + return _gen_tensors_pooled(d, scale, pool_size=4096) + + +gen_tensors_lo._fixture_version = 1 + + +def gen_tensors_hi(d, scale): + """Expansion-sweep HIGH point: 2-value pool (1-bit dictionary indices), + the worst plausible expansion — the stress case for any sizing policy.""" + return _gen_tensors_pooled(d, scale, pool_size=2) + + +gen_tensors_hi._fixture_version = 1 + + +SHAPES = { + "lone_big_rg": gen_lone_big_rg, + "single_rg_files": gen_single_rg_files, + "auto_rg": gen_auto_rg, + "tiny_rgs": gen_tiny_rgs, + "wide": gen_wide, + "fat_col": gen_fat_col, + "bin_sweep": gen_bin_sweep, + "tensors_wide": gen_tensors_wide, + "tensors_cp": gen_tensors_cp, + "tensors_dict": gen_tensors_dict, + "tensors_lo": gen_tensors_lo, + "tensors_hi": gen_tensors_hi, +} + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--root", required=True) + p.add_argument("--shapes", default=",".join(SHAPES)) + p.add_argument( + "--scale", + type=float, + default=1.0, + help="shrink every shape by this factor (0.25 for a quick smoke run)", + ) + args = p.parse_args() + + root = os.path.expanduser(args.root) + os.makedirs(root, exist_ok=True) + manifest_path = os.path.join(root, "manifest.json") + manifest = {} + if os.path.exists(manifest_path): + with open(manifest_path) as fh: + manifest = json.load(fh) + + for shape in [s.strip() for s in args.shapes.split(",") if s.strip()]: + gen = SHAPES[shape] + d = os.path.join(root, shape) + if os.path.isdir(d) and shape in manifest: + # Skip ONLY at the same scale. A 0.25-scale smoke run must not leave + # quarter-size fixtures behind for the full run to silently benchmark. + want_v = getattr(gen, "_fixture_version", None) + have_v = manifest[shape].get("v") + if manifest[shape].get("scale") == args.scale and have_v == want_v: + print(f" {shape}: exists, skipping ({manifest[shape]})", flush=True) + continue + print( + f" {shape}: exists at scale={manifest[shape].get('scale')} " + f"v={have_v}, want scale={args.scale} v={want_v} — regenerating", + flush=True, + ) + shutil.rmtree(d) + os.makedirs(d, exist_ok=True) + print(f" {shape}: generating (scale={args.scale}) ...", flush=True) + stats = gen(d, args.scale) + on_disk = sum( + os.path.getsize(os.path.join(d, f)) + for f in os.listdir(d) + if f.endswith(".parquet") + ) + stats.update(path=d, on_disk_bytes=on_disk, scale=args.scale) + manifest[shape] = stats + print( + f" {shape}: {stats['files']} files, {stats['rows']} rows, " + f"{stats['uncompressed_bytes'] / 2**20:.0f} MiB uncompressed, " + f"{on_disk / 2**20:.0f} MiB on disk", + flush=True, + ) + with open(manifest_path, "w") as fh: + json.dump(manifest, fh, indent=2) + + print(f"\nmanifest -> {manifest_path}") + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/gen_s3_fixtures.py b/release/nightly_tests/dataset/arrow_rs_probe/gen_s3_fixtures.py new file mode 100644 index 000000000000..7d413490f6f8 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/gen_s3_fixtures.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python3 +"""Generate the two release-test Parquet layouts on OUR S3 bucket. + +The release regressions live on two layouts we can't tune in the shared bucket, so +we recreate them on a scratch bucket we own (fill/empty at will) and point the probe +at them: + + wide -> wide_schema_pipeline_primitives : many columns, ONE big row group per file. + This is the K=1 lone-big-row-group case where the crate's byte-budgeted + windowed-async S3 decode is supposed to keep the working set to a page + while PyArrow materializes the whole decoded row group. mem was 1.50x worse. + img -> imagenet (mix.8ds_equal_random_mix) : blob column, MANY tiny row groups. + This is the I/O-bound-on-S3 case (many small serial range GETs). time 1.67x. + +Data is written locally then `aws s3 sync`d up (more robust than pyarrow's S3 +region/creds handling). Requires: aws CLI on PATH + creds/region exported, and +write_page_index=True (the crate's page-index requirement). + + python gen_s3_fixtures.py --bucket s3://arrowrs-bench-21f6c795 + python gen_s3_fixtures.py --bucket s3://... --cases wide # just one layout + python gen_s3_fixtures.py --bucket s3://... --wide-files 8 --wide-rows 1500 + +Prints the exact --path values to hand to run_matrix.py at the end. +""" +import argparse +import os +import shutil +import subprocess +import sys + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +# A pool of distinct, incompressible-ish strings so Snappy can't crush the fixture to +# nothing -- we want on-disk (== S3 fetch) bytes to be non-trivial, not just decoded size. +_POOL_SIZE = 4096 + + +def _str_pool(rng, str_len): + alphabet = np.frombuffer(b"abcdefghijklmnopqrstuvwxyz0123456789", dtype="S1") + idx = rng.integers(0, len(alphabet), size=(_POOL_SIZE, str_len)) + return np.array([alphabet[row].tobytes().decode() for row in idx], dtype=object) + + +def gen_wide(d, n_rows, n_cols, str_len, rg, n_files): + """Wide schema, ONE big row group per file (rg == n_rows) -> the K=1 S3 case.""" + os.makedirs(d, exist_ok=True) + rng = np.random.default_rng(0) + pool = _str_pool(rng, str_len) + for f in range(n_files): + cols = {"id": pa.array(np.arange(f * n_rows, (f + 1) * n_rows, dtype=np.int64))} + for c in range(n_cols): + picks = pool[rng.integers(0, _POOL_SIZE, size=n_rows)] + cols[f"c{c}"] = pa.array(picks, type=pa.string()) + pq.write_table( + pa.table(cols), + os.path.join(d, f"part{f}.parquet"), + write_page_index=True, + row_group_size=rg, # >= n_rows -> a single row group + ) + print( + f" wide part{f}: {n_rows}x{n_cols} str{str_len} " f"rg={rg} (1 rg/file)", + flush=True, + ) + + +def gen_imagenet(d, n_rows, blob_kb, rg, n_files): + """Blob column, MANY tiny row groups -> the S3-IO-bound case.""" + os.makedirs(d, exist_ok=True) + rng = np.random.default_rng(0) + per = n_rows // n_files + for f in range(n_files): + images = [rng.bytes(blob_kb * 1024) for _ in range(per)] + labels = rng.integers(0, 1000, per, dtype=np.int64) + ids = np.arange(f * per, (f + 1) * per, dtype=np.int64) + pq.write_table( + pa.table( + { + "id": ids, + "image": pa.array(images, type=pa.binary()), + "label": pa.array(labels), + } + ), + os.path.join(d, f"part{f}.parquet"), + write_page_index=True, + row_group_size=rg, + ) + print( + f" img part{f}: {per} rows blob={blob_kb}KB rg={rg} " + f"-> ~{per // rg} rgs/file", + flush=True, + ) + + +def sync_up(local_dir, s3_prefix): + print(f" aws s3 sync {local_dir} -> {s3_prefix}", flush=True) + subprocess.run( + ["aws", "s3", "sync", "--only-show-errors", local_dir, s3_prefix], + check=True, + ) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--bucket", required=True, help="e.g. s3://arrowrs-bench-21f6c795") + p.add_argument("--cases", default="wide,img", help="comma list of {wide,img}") + p.add_argument("--local-tmp", default="/tmp/arrow_rs_s3_gen") + p.add_argument("--keep-local", action="store_true", help="don't delete local tmp") + # wide knobs + p.add_argument("--wide-rows", type=int, default=1000) + p.add_argument("--wide-cols", type=int, default=5000) + p.add_argument("--wide-str-len", type=int, default=100) + p.add_argument("--wide-files", type=int, default=4) + # imagenet knobs + p.add_argument("--img-rows", type=int, default=24000) + p.add_argument("--img-blob-kb", type=int, default=48) + p.add_argument("--img-rg", type=int, default=32) + p.add_argument("--img-files", type=int, default=8) + args = p.parse_args() + + bucket = args.bucket.rstrip("/") + cases = [c.strip() for c in args.cases.split(",") if c.strip()] + paths = {} + + if "wide" in cases: + print("=== gen wide (1 big row group/file) ===", flush=True) + d = os.path.join(args.local_tmp, "wide_schema", "primitives") + gen_wide( + d, + args.wide_rows, + args.wide_cols, + args.wide_str_len, + rg=args.wide_rows, # one row group per file + n_files=args.wide_files, + ) + s3 = f"{bucket}/wide_schema/primitives" + sync_up(d, s3) + paths["wide"] = s3 + + if "img" in cases: + print("=== gen imagenet (many tiny row groups) ===", flush=True) + d = os.path.join(args.local_tmp, "imagenet", "parquet") + gen_imagenet(d, args.img_rows, args.img_blob_kb, args.img_rg, args.img_files) + s3 = f"{bucket}/imagenet/parquet" + sync_up(d, s3) + paths["img"] = s3 + + if not args.keep_local and os.path.isdir(args.local_tmp): + shutil.rmtree(args.local_tmp, ignore_errors=True) + + print("\n=== S3 fixtures ready ===") + for k, v in paths.items(): + print(f" {k}: {v}") + print("\nNext:") + wide = paths.get("wide", "") + img = paths.get("img", "") + print(f" python run_matrix.py --wide-path {wide} --imagenet-path {img}") + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/release/nightly_tests/dataset/arrow_rs_probe/grand_experiment.py b/release/nightly_tests/dataset/arrow_rs_probe/grand_experiment.py new file mode 100644 index 000000000000..7b1e2425850a --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/grand_experiment.py @@ -0,0 +1,373 @@ +#!/usr/bin/env python3 +"""The grand tuning experiment: arrow-rs vs PyArrow on the NEW footer-based planner. + +The #64985 series changed what there is to tune: footers are read up front, row +groups are statistics-pruned and bin-packed by UNCOMPRESSED size into read tasks +(env ``RAY_DATA_PARQUET_BIN_PACKING_BYTES``, default 128 MiB), and the PyArrow path +runs one fragment thread per fragment, unbounded. This experiment maps that new +tuning surface for both readers in one run, over the shapes from +gen_local_fixtures.py. + +Stages (each cell = one read_probe.py subprocess: fresh Ray + fresh crate, env-injected +knobs, full log kept): + + A headline every shape x {pyarrow, arrow_rs}, all-default knobs. + The number that matters: R/P wall and USS per shape on the new planner. + B bin RAY_DATA_PARQUET_BIN_PACKING_BYTES in {32Mi, 512Mi} (128Mi = stage A) + x both readers x {single_rg_files, tiny_rgs}. The NEW knob: how does + read-task sizing interact with each reader? + C budget RAY_DATA_ARROW_RS_DECODE_BUDGET_BYTES in {8Mi, 128Mi} (32Mi = stage A) + x arrow_rs x {lone_big_rg, fat_col}. Does the decode budget still + bound memory now that bins set task size? + D threads RAY_DATA_READ_FILES_NUM_THREADS in {2, 4} x arrow_rs x + {single_rg_files, tiny_rgs}. Re-decides the "1 fragment thread" + default against the new unbounded-PyArrow baseline (old finding K6 + measured against a 4-thread baseline that no longer exists). + E s3 (only with --s3-bucket) sync fixtures up, rerun stage A on S3, plus + RAY_DATA_ARROW_RS_FETCH_WINDOW_MB in {4, 64} (16 = default) x + arrow_rs x {lone_big_rg, tiny_rgs}. The crate's byte-budgeted + windowed decode only runs on the S3 path, so the memory-knob story + is only visible here. + +arrow_rs cells run with RAY_DATA_ARROW_RS_STRICT=1: if the reader silently falls +back to PyArrow the cell FAILS instead of quietly measuring the wrong engine. + +Usage (venv active, after gen_local_fixtures.py): + + python grand_experiment.py --fixtures-root ~/arrow_rs_grand_fixtures + python grand_experiment.py --fixtures-root ... --repeat 3 # median of 3 + python grand_experiment.py --fixtures-root ... --stages A,B # subset + python grand_experiment.py --fixtures-root ... --s3-bucket s3://arrowrs-bench-21f6c795 + +Outputs /summary.json (all cells), summary.md (ratio tables), and one +.log per cell. +""" +import argparse +import json +import os +import subprocess +import sys +import time + +PROBE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "read_probe.py") +PY = sys.executable +MiB = 1024 * 1024 + +# Metrics pulled from read_probe's RESULT block into the tables. +_METRICS = [ + "wall_s", + "cpu_over_wall", + "read_avg_max_uss_gb", + "peak_uss_gb", + "peak_rss_gb", +] +# Ratio direction: for these, ratio = arrow_rs / pyarrow and >1 means arrow_rs WORSE. +_RATIO_METRICS = ["wall_s", "read_avg_max_uss_gb", "peak_uss_gb"] + + +def run_cell(logdir, name, path, reader, extra_env, columns=None, repeat=1): + """Run read_probe.py `repeat` times; return the median-wall run's RESULT dict.""" + runs = [] + for i in range(repeat): + tag = name if repeat == 1 else f"{name}.r{i}" + cmd = [PY, PROBE, "--path", path, "--reader", reader] + if columns: + cmd += ["--columns", *columns] + env = dict(os.environ) + env.update(extra_env) + if reader == "arrow_rs": + env["RAY_DATA_ARROW_RS_STRICT"] = "1" + + t0 = time.perf_counter() + proc = subprocess.run(cmd, capture_output=True, text=True, env=env) + dur = time.perf_counter() - t0 + + with open(os.path.join(logdir, f"{tag}.log"), "w") as fh: + fh.write(f"# cmd: {' '.join(cmd)}\n# extra_env: {extra_env}\n") + fh.write(f"# wall_including_startup_s: {dur:.1f}\n# ---- STDOUT ----\n") + fh.write(proc.stdout) + fh.write("\n# ---- STDERR ----\n") + fh.write(proc.stderr) + + res = {} + in_result = False + for line in proc.stdout.splitlines(): + if "=== RESULT ===" in line: + in_result = True + continue + if in_result and ":" in line: + k, v = line.strip().split(":", 1) + res[k.strip()] = v.strip() + if not res: + print( + f" !! {tag} FAIL rc={proc.returncode} " + f"(see {tag}.log) {proc.stderr.strip()[-300:]}", + flush=True, + ) + else: + print( + f" {tag:<44} wall={res.get('wall_s')} " + f"cpu/wall={res.get('cpu_over_wall')} " + f"uss={res.get('read_avg_max_uss_gb') or res.get('peak_uss_gb')}", + flush=True, + ) + runs.append(res) + + good = [r for r in runs if r] + if not good: + return {} + good.sort(key=lambda r: _num(r, "wall_s") or float("inf")) + return good[len(good) // 2] + + +def _num(res, key): + try: + return float(res.get(key)) + except (TypeError, ValueError): + return None + + +def _ratio(a, b): + if a is None or b is None or b == 0: + return None + return round(a / b, 3) + + +def _s3_sync(fixtures_root, bucket, shapes): + """aws s3 sync each shape dir up; return shape -> s3 path.""" + out = {} + for shape, info in shapes.items(): + dst = f"{bucket.rstrip('/')}/grand/{shape}" + print(f" syncing {shape} -> {dst}", flush=True) + subprocess.run( + ["aws", "s3", "sync", "--only-show-errors", info["path"], dst], + check=True, + ) + out[shape] = dst + return out + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--fixtures-root", required=True) + p.add_argument("--outdir", default=None) + p.add_argument("--repeat", type=int, default=1) + p.add_argument("--stages", default="A,B,C,D,E", help="comma subset of A,B,C,D,E") + p.add_argument( + "--s3-bucket", + default=os.environ.get("ARROW_RS_S3_BUCKET"), + help="s3://... scratch bucket; enables stage E (also via ARROW_RS_S3_BUCKET)", + ) + args = p.parse_args() + + stages = {s.strip().upper() for s in args.stages.split(",") if s.strip()} + root = os.path.expanduser(args.fixtures_root) + with open(os.path.join(root, "manifest.json")) as fh: + shapes = json.load(fh) + + ts = time.strftime("%Y%m%d_%H%M%S") + outdir = args.outdir or os.path.join( + os.path.dirname(os.path.abspath(__file__)), "grand_runs", ts + ) + os.makedirs(outdir, exist_ok=True) + print(f"shapes: {list(shapes)} logs -> {outdir}\n", flush=True) + + rows = {} # cell name -> RESULT dict + + def cell(name, path, reader, extra_env, columns=None): + rows[name] = run_cell( + outdir, name, path, reader, extra_env, columns=columns, repeat=args.repeat + ) + + # -------- A: headline, all defaults -------- + if "A" in stages: + print( + "=== [A] headline: every shape x both readers, default knobs ===", + flush=True, + ) + for shape, info in shapes.items(): + for reader in ("pyarrow", "arrow_rs"): + cell(f"A.{shape}.{reader}", info["path"], reader, {}) + + # -------- B: bin-packing size (the NEW planner knob) -------- + if "B" in stages: + print( + "=== [B] bin size sweep (RAY_DATA_PARQUET_BIN_PACKING_BYTES) ===", + flush=True, + ) + for shape in ("single_rg_files", "tiny_rgs"): + if shape not in shapes: + continue + for bin_mib in (32, 512): + for reader in ("pyarrow", "arrow_rs"): + cell( + f"B.{shape}.bin{bin_mib}Mi.{reader}", + shapes[shape]["path"], + reader, + {"RAY_DATA_PARQUET_BIN_PACKING_BYTES": str(bin_mib * MiB)}, + ) + + # -------- C: decode budget (arrow-rs) -------- + if "C" in stages: + print( + "=== [C] decode budget sweep (RAY_DATA_ARROW_RS_DECODE_BUDGET_BYTES) ===", + flush=True, + ) + for shape in ("lone_big_rg", "fat_col"): + if shape not in shapes: + continue + for budget_mib in (8, 128): + cell( + f"C.{shape}.budget{budget_mib}Mi.arrow_rs", + shapes[shape]["path"], + "arrow_rs", + {"RAY_DATA_ARROW_RS_DECODE_BUDGET_BYTES": str(budget_mib * MiB)}, + ) + + # -------- D: fragment threads (arrow-rs; PyArrow is unbounded by design) -------- + if "D" in stages: + print( + "=== [D] fragment-thread sweep (RAY_DATA_READ_FILES_NUM_THREADS) ===", + flush=True, + ) + for shape in ("single_rg_files", "tiny_rgs"): + if shape not in shapes: + continue + for threads in (2, 4): + cell( + f"D.{shape}.threads{threads}.arrow_rs", + shapes[shape]["path"], + "arrow_rs", + {"RAY_DATA_READ_FILES_NUM_THREADS": str(threads)}, + ) + + # -------- E: S3 (headline + fetch window) -------- + s3_paths = {} + if "E" in stages and args.s3_bucket: + print(f"=== [E] S3 stage (bucket {args.s3_bucket}) ===", flush=True) + try: + s3_paths = _s3_sync(root, args.s3_bucket, shapes) + except (subprocess.CalledProcessError, FileNotFoundError) as e: + # No creds / no aws CLI must not crash the run AFTER stages A-D + # produced results but BEFORE the summary was written. + print(f" !! S3 sync failed ({e}) — skipping stage E", flush=True) + s3_paths = {} + for shape, s3_path in s3_paths.items(): + for reader in ("pyarrow", "arrow_rs"): + cell(f"E.{shape}.{reader}", s3_path, reader, {}) + for shape in ("lone_big_rg", "tiny_rgs"): + if shape not in s3_paths: + continue + for window in (4, 64): + cell( + f"E.{shape}.window{window}Mi.arrow_rs", + s3_paths[shape], + "arrow_rs", + {"RAY_DATA_ARROW_RS_FETCH_WINDOW_MB": str(window)}, + ) + elif "E" in stages: + print("=== [E] skipped: no --s3-bucket / ARROW_RS_S3_BUCKET ===", flush=True) + + # -------- summary -------- + lines = ["# grand_experiment summary", ""] + lines.append(f"run: {ts} repeat={args.repeat} fixtures={root}") + lines.append("") + lines.append("Ratios are arrow_rs / pyarrow; **> 1.00 means arrow-rs is worse**.") + lines.append("") + + def emit(text=""): + print(text, flush=True) + lines.append(text) + + emit("\n=========== HEADLINE (stage A: default knobs, per shape) ===========") + emit("| shape | wall R/P | USS R/P | wall (P -> R) | USS GB (P -> R) |") + emit("|---|---|---|---|---|") + for shape in shapes: + pa_r = rows.get(f"A.{shape}.pyarrow", {}) + ar_r = rows.get(f"A.{shape}.arrow_rs", {}) + wall_p, wall_a = _num(pa_r, "wall_s"), _num(ar_r, "wall_s") + # Ray's own per-task USS is the metric of record; sampler peak is backup. + uss_key = ( + "read_avg_max_uss_gb" + if _num(pa_r, "read_avg_max_uss_gb") is not None + else "peak_uss_gb" + ) + uss_p, uss_a = _num(pa_r, uss_key), _num(ar_r, uss_key) + emit( + f"| {shape} | {_ratio(wall_a, wall_p)} | {_ratio(uss_a, uss_p)} " + f"| {wall_p} -> {wall_a} | {uss_p} -> {uss_a} |" + ) + + def sweep_table(title, prefix_fmt, axis_values, shapes_subset, readers): + emit(f"\n=========== {title} ===========") + header = "| shape | reader | " + " | ".join(str(v) for v in axis_values) + " |" + emit(header) + emit("|---" * (2 + len(axis_values)) + "|") + for shape in shapes_subset: + for reader in readers: + cells = [] + for v in axis_values: + r = rows.get(prefix_fmt.format(shape=shape, v=v, reader=reader), {}) + cells.append( + f"w={_num(r, 'wall_s')} u={_num(r, 'read_avg_max_uss_gb') or _num(r, 'peak_uss_gb')}" + ) + emit(f"| {shape} | {reader} | " + " | ".join(cells) + " |") + + if "B" in stages: + sweep_table( + "stage B: bin size (32Mi / 512Mi; 128Mi default = stage A)", + "B.{shape}.bin{v}Mi.{reader}", + [32, 512], + [s for s in ("single_rg_files", "tiny_rgs") if s in shapes], + ["pyarrow", "arrow_rs"], + ) + if "C" in stages: + sweep_table( + "stage C: arrow-rs decode budget (8Mi / 128Mi; 32Mi default = stage A)", + "C.{shape}.budget{v}Mi.{reader}", + [8, 128], + [s for s in ("lone_big_rg", "fat_col") if s in shapes], + ["arrow_rs"], + ) + if "D" in stages: + sweep_table( + "stage D: arrow-rs fragment threads (2 / 4; stage A = the " + "min(4, num_fragments) default since b73ee04ce9 — no threads=1 arm)", + "D.{shape}.threads{v}.{reader}", + [2, 4], + [s for s in ("single_rg_files", "tiny_rgs") if s in shapes], + ["arrow_rs"], + ) + if s3_paths: + emit("\n=========== stage E: S3 headline ===========") + emit("| shape | wall R/P | USS R/P |") + emit("|---|---|---|") + for shape in shapes: + pa_r = rows.get(f"E.{shape}.pyarrow", {}) + ar_r = rows.get(f"E.{shape}.arrow_rs", {}) + uss_key = ( + "read_avg_max_uss_gb" + if _num(pa_r, "read_avg_max_uss_gb") is not None + else "peak_uss_gb" + ) + emit( + f"| {shape} | {_ratio(_num(ar_r, 'wall_s'), _num(pa_r, 'wall_s'))} " + f"| {_ratio(_num(ar_r, uss_key), _num(pa_r, uss_key))} |" + ) + sweep_table( + "stage E: S3 fetch window (4Mi / 64Mi; 16Mi default = E headline)", + "E.{shape}.window{v}Mi.{reader}", + [4, 64], + [s for s in ("lone_big_rg", "tiny_rgs") if s in shapes], + ["arrow_rs"], + ) + + with open(os.path.join(outdir, "summary.json"), "w") as fh: + json.dump(rows, fh, indent=2) + with open(os.path.join(outdir, "summary.md"), "w") as fh: + fh.write("\n".join(lines) + "\n") + print(f"\nsummary.json + summary.md + per-cell logs in {outdir}") + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/loss_triage.py b/release/nightly_tests/dataset/arrow_rs_probe/loss_triage.py new file mode 100644 index 000000000000..811fe0c21138 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/loss_triage.py @@ -0,0 +1,675 @@ +#!/usr/bin/env python3 +"""3-part triage of the 2026-08-15 release-A/B losses (findings M31/M32/M33). + +Each loss shape runs in THREE parts so the layer that owns the regression is +read off one table instead of argued about: + + standalone no Ray, no S3 — the decoder (or decode+realign / decode+write + pipeline) alone in a fresh process on local files. If a loss + shows here it is a native-decoder problem. + ray_local the same shape through `ray.data.read_parquet` on local files + (read_probe.py, private local cluster). A loss that appears + only here is Ray-integration (worker reuse / allocator + retention / fusion), not the decoder. + ray_s3 the same Ray read against the scratch S3 bucket. A loss that + appears only here is the crate's S3 path or fetch behaviour. + +The shapes (release test -> fixture, both readers each part): + + auto M31 read_large_parquet_autoscaling: one ~69 MiB row group per task + (bin 64 MiB), sub-second tasks. Release instrument could NOT + measure this (1 Hz poll => end-of-task sample); every Ray cell + here polls at 20 Hz (--mem-poll-s 0.05) so per-task USS is a real + peak. The standalone part runs N sequential one-row-group "tasks" + in ONE process and prints the RSS after each — the end-of-task + retention curve itself. + write M32 write_parquet: fused read->write, ~1.2 GiB decode churn per + task (bin 1342177280 on the bin_sweep fixture). Won on Linux+S3 + pre-#64985 (avg USS 0.83x, reader comment 186-196); regressed + 1.24x/1.64x with whole-file bins. + tensorscp M33 wide_schema tensors: 5000 cloudpickle-tagged tensor columns + (bin 40894464 like the release yaml) — the crate's skip+realign + path. This shape has lost OUTSIDE Ray before (T22: wall 5.4x + pre-fix, 1.25x post), so its standalone part is the decisive leg + the 2026-08-17 macOS matrix (M34) did not cover. + +Every arrow-rs Ray cell also runs a MALLOC_ARENA_MAX=2 variant (drop with +--no-arena-sweep): if the ray_local/ray_s3 losses collapse under it, the +mechanism is glibc arena retention in long-lived workers and the fix is +allocator config / crate-side malloc_trim, not reader logic. + +Usage (Linux box; fixtures + venv via run_loss_triage.sh, or piecemeal): + + python gen_local_fixtures.py --root ~/arrow_rs_repl_fixtures \ + --shapes auto_rg,bin_sweep,tensors_cp + python loss_triage.py --fixture-root ~/arrow_rs_repl_fixtures + # with the S3 part (scratch bucket you own; fixtures are synced up first): + ARROW_RS_S3_BUCKET=s3://arrowrs-bench-xxxx python loss_triage.py \ + --fixture-root ~/arrow_rs_repl_fixtures + python loss_triage.py --fixture-root ... --shapes write --parts ray_local + +Results: /summary.json + per-cell logs; the printed table is +R = arrow_rs / pyarrow per (shape, part), >1.00 = arrow-rs worse. Standalone +memory is peak RSS of the case subprocess; Ray memory is the 50 Hz worker +sampler's peak USS plus Ray's own per-task USS (20 Hz in-task poll). +""" +import argparse +import gc +import glob +import json +import os +import resource +import subprocess +import sys +import time + +from run_matrix import _median, _num, median_cell, ratio + +PY = sys.executable +HERE = os.path.dirname(os.path.abspath(__file__)) +MiB = 1024 * 1024 +# ru_maxrss is KiB on Linux, bytes on macOS. +_RU_UNIT = 1024 if sys.platform.startswith("linux") else 1 + +# Release-yaml bin sizes per shape (release/release_data_tests.yaml). +SHAPE_BINS = { + "auto": 67_108_864, + "write": 1_342_177_280, + "tensorscp": 40_894_464, + # M43's shape in-Ray: tensors with ~9.4x dictionary expansion (the release + # wide_schema tensors regime). Same release bin as tensorscp. tensorscp + # (~1.1x expansion) stays as the control. + "tensorsdict": 40_894_464, + # POSITIVE control (A/B #4 review): the aggregate_groups family is arrow-rs's + # biggest per-task USS win (R 0.56-0.81, scaling with decoded bytes — pa's + # scanner keeps the whole decoded task working set resident, we keep ~the + # budget). Projected read -> groupby -> count, Ray parts only. Any batch- + # sizing change must NOT regress this cell: the win comes from exactly the + # small-resident-batch behaviour the fix candidate touches. + "agg": 67_108_864, + # 1o: fat_col (1 file, 1 rg, 1024 rows, binary ~256 KiB/row + int64) — the + # incompressible-BYTE_ARRAY kernel-throughput shape. Bin size is moot (one + # file = one task either way); default 64 MiB. + "fatcol": 67_108_864, +} +SHAPE_FIXTURE = { + "auto": "auto_rg", + "write": "bin_sweep", + "tensorscp": "tensors_cp", + "tensorsdict": "tensors_dict", + "agg": "auto_rg", + "fatcol": "fat_col", +} +# Ray-only shapes: no standalone part (agg's point is the read-op USS inside a +# read->shuffle->aggregate pipeline, which has no standalone analogue). +RAY_ONLY_SHAPES = {"agg"} +# Shapes that decode through the crate's skip+realign path (cloudpickle +# ARROW:schema) — they share the env flag and the per-batch realign cast. +TENSOR_SHAPES = {"tensorscp", "tensorsdict"} +# Env every cell of a shape needs, standalone included (the cloudpickle flag +# must be set before ray.data is imported). +SHAPE_ENV = { + s: {"RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA": "1"} for s in TENSOR_SHAPES +} + + +# -------------------------------------------------------------------------- +# Part 1: the standalone (no Ray, no S3) case, run as its own subprocess. +# -------------------------------------------------------------------------- + + +def _reader_knobs(): + """Ray's shipped arrow-rs defaults, imported from the reader so this stays + a single source of truth; falls back to the documented values when the + reader isn't importable (e.g. running against a stock wheel).""" + try: + from ray.data._internal.datasource_v2.readers import ( + arrow_rs_parquet_file_reader as r, + ) + + budget = r._ARROW_RS_DECODE_BUDGET_BYTES + if budget is None: + # Post-a1b095ec03 the module attr is None unless the env var is set + # (the default now follows DataContext.target_max_block_size); + # resolve it the way the reader would. + budget = r._default_decode_budget_bytes() + return dict( + budget=budget, + k=r._ARROW_RS_K, + split=r._ARROW_RS_DEFAULT_SPLIT_THRESHOLD_BYTES, + window=r._ARROW_RS_FETCH_WINDOW_MB, + column=r._ARROW_RS_COLUMN_FETCH_MB, + min_rows=r._ARROW_RS_MIN_DECODE_BATCH_ROWS, + ) + except Exception: + return dict( + budget=32 * MiB, k=1, split=128 * MiB, window=16, column=16, min_rows=2048 + ) + + +def _batch_size(md, knobs): + total = sum(md.row_group(i).total_byte_size for i in range(md.num_row_groups)) + bpr = max(1, total // max(1, md.num_rows)) + return max(knobs["min_rows"], int(knobs["budget"] // bpr)) + + +def _pa_batches(path, batch_size, use_threads=True): + # use_threads=False gives the single-thread CPU-cost baseline (the pa1 + # standalone cell) — the fair + # comparison against the crate's K=1 decode and the in-Ray regime, where + # parallelism comes from tasks and M35 measured wall R 0.81-1.00 vs the + # 1.16-1.54 threaded-standalone artifact. Default True is kept for this + # file's own standalone leg (comparable with prior runs, artifact + # documented in M35). + import pyarrow as pa + import pyarrow.dataset as pds + from pyarrow.fs import LocalFileSystem + + fmt = pds.ParquetFileFormat( + default_fragment_scan_options=pds.ParquetFragmentScanOptions(pre_buffer=True) + ) + frag = fmt.make_fragment(path, filesystem=LocalFileSystem()) + for b in frag.scanner(batch_size=batch_size, use_threads=use_threads).to_batches(): + yield pa.Table.from_batches([b]) + + +def _rs_batches(path, batch_size, knobs, realign_fields=None): + import pyarrow as pa + import ray_data_arrow_rs as rs + + handle = rs.open_parquet_file(path, page_index=False) + stream = pa.RecordBatchReader.from_stream( + handle.read_row_groups( + row_groups=None, + columns=None, + batch_size=batch_size, + decode_budget_bytes=knobs["budget"], + k=knobs["k"], + split_threshold_bytes=knobs["split"], + predicate_json=None, + fetch_window_mb=knobs["window"], + column_fetch_mb=knobs["column"], + prefetch_budget_mb=4 * max(knobs["window"], knobs["column"]), + ) + ) + if realign_fields is None: + for b in stream: + yield pa.Table.from_batches([b]) + else: + # The 1y skip+realign path: the crate decoded parquet STORAGE types + # (it can't parse the cloudpickle ARROW:schema); cast each batch back + # to the extension schema exactly like the reader does. + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( # noqa: E501 + _cast_table_to, + ) + + for b in stream: + yield _cast_table_to(pa.Table.from_batches([b]), realign_fields) + + +def _consume(tables, mode, out_path): + """BlockOutputBuffer emulation: coalesce to ~128 MiB blocks; write mode + feeds the blocks to a ParquetWriter like the fused Read->Write task.""" + import pyarrow as pa + import pyarrow.parquet as pq + + writer = None + buf, buf_bytes, n_rows = [], 0, 0 + + def flush(): + nonlocal buf, buf_bytes, n_rows, writer + if not buf: + return + block = pa.concat_tables(buf) + buf, buf_bytes = [], 0 + n_rows += block.num_rows + if mode == "write": + if writer is None: + writer = pq.ParquetWriter(out_path, block.schema) + writer.write_table(block) + + for t in tables: + buf.append(t) + buf_bytes += t.nbytes + if buf_bytes >= 128 * MiB: + flush() + flush() + if writer is not None: + writer.close() + return n_rows + + +def _cur_rss_mib(): + try: + import psutil + + return psutil.Process().memory_info().rss / MiB + except ImportError: + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * _RU_UNIT / MiB + + +def run_case(a): + import pyarrow.parquet as pq + + knobs = _reader_knobs() + files = sorted(glob.glob(os.path.join(os.path.expanduser(a.path), "*.parquet"))) + if not files: + raise SystemExit(f"no parquet files under {a.path}") + + realign_fields = None + if a.shape in TENSOR_SHAPES: + # Both readers need Ray's tensor extension types registered (the + # release run has them via ray.data); the rs leg additionally realigns + # to the footer's extension schema, exactly like the reader. + import ray.data # noqa: F401 + + if a.reader == "rs": + realign_fields = list(pq.read_schema(files[0])) + + mode = "write" if a.shape == "write" else "decode" + out_path = os.path.join(a.workdir, f"triage_out_{a.reader}.parquet") + + # `auto` = N sequential one-file (one row group) tasks in one process, + # recording RSS after each: the end-of-task retention curve. `write` = one + # release-task-sized unit: files until the bin budget is met (~1.25 GiB of + # footer bytes, like one whole-file-bin fused task). Else one file. + if a.shape == "auto": + task_files = files[: a.tasks] + elif a.shape == "write": + task_files, cum = [], 0 + for path in files: + task_files.append(path) + md = pq.read_metadata(path) + cum += sum( + md.row_group(i).total_byte_size for i in range(md.num_row_groups) + ) + if cum >= SHAPE_BINS["write"]: + break + else: + task_files = files[:1] + + rss_after_task = [] + rows = 0 + t0 = time.perf_counter() + for path in task_files: + md = pq.read_metadata(path) + bs = _batch_size(md, knobs) + if a.reader in ("pa", "pa1"): + # pa1 = use_threads=False: the single-core kernel baseline, so the + # standalone triple reads as pa (8t) / pa1 (1t) / rs (1 core) — + # rs-vs-pa1 is kernel-vs-kernel, pa-vs-pa1 is pure thread speedup. + it = _pa_batches(path, bs, use_threads=(a.reader == "pa")) + else: + it = _rs_batches(path, bs, knobs, realign_fields) + rows += _consume(it, mode, out_path) + del it + gc.collect() + rss_after_task.append(round(_cur_rss_mib(), 1)) + wall = time.perf_counter() - t0 + + if os.path.exists(out_path): + os.unlink(out_path) + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * _RU_UNIT + print("=== CASE RESULT ===") + print( + json.dumps( + dict( + shape=a.shape, + reader=a.reader, + tasks=len(task_files), + rows=rows, + wall_s=round(wall, 3), + peak_rss_mib=round(peak / MiB, 1), + end_rss_mib=rss_after_task[-1], + first_task_rss_mib=rss_after_task[0], + rss_after_task=rss_after_task if a.shape == "auto" else None, + ) + ) + ) + + +# -------------------------------------------------------------------------- +# The matrix: cells across parts x shapes x readers (x allocator arm). +# -------------------------------------------------------------------------- + + +def run_standalone_cell( + logdir, name, shape, reader, path, env_extra, tasks, repeat, warmup +): + """Median-of-N standalone case, one fresh subprocess per run (peak RSS must + not accumulate across repeats).""" + workdir = os.path.join(logdir, "standalone_tmp") + os.makedirs(workdir, exist_ok=True) + cmd = [ + PY, + os.path.abspath(__file__), + "case", + "--shape", + shape, + "--reader", + reader, + "--path", + path, + "--tasks", + str(tasks), + "--workdir", + workdir, + ] + env = dict(os.environ) + env.update(env_extra) + + def one(tag): + t0 = time.perf_counter() + proc = subprocess.run(cmd, capture_output=True, text=True, env=env) + with open(os.path.join(logdir, f"{tag}.log"), "w") as fh: + fh.write(f"# cmd: {' '.join(cmd)}\n# env_extra: {env_extra}\n") + fh.write(f"# wall_including_startup_s: {time.perf_counter() - t0:.1f}\n") + fh.write("# ---- STDOUT ----\n" + proc.stdout) + fh.write("\n# ---- STDERR ----\n" + proc.stderr) + for i, line in enumerate(proc.stdout.splitlines()): + if "=== CASE RESULT ===" in line: + try: + return json.loads(proc.stdout.splitlines()[i + 1]) + except (IndexError, json.JSONDecodeError): + break + print( + f" !! {tag} CASE FAIL rc={proc.returncode} " + f"(see {logdir}/{tag}.log)\n {proc.stderr.strip()[-400:]}", + flush=True, + ) + return {} + + for i in range(warmup): + one(f"{name}.w{i}") + runs = [one(name if repeat == 1 else f"{name}.r{i}") for i in range(repeat)] + good = [r for r in runs if r] + if not good: + return {} + out = dict(good[0]) + samples = {} + for key in list(out): + vals = [_num({k: str(v) for k, v in r.items()}, key) for r in good] + if any(v is None for v in vals): + continue + out[key] = _median(vals) + samples[key] = vals + out["_n"] = len(good) + out["_samples"] = samples + print( + f" {name:<40} wall={out.get('wall_s')} peak_rss_mib={out.get('peak_rss_mib')} " + f"end_rss_mib={out.get('end_rss_mib')}", + flush=True, + ) + return out + + +def s3_sync(local_dir, s3_prefix): + print(f" aws s3 sync {local_dir} -> {s3_prefix}", flush=True) + subprocess.run( + ["aws", "s3", "sync", "--only-show-errors", local_dir, s3_prefix], + check=True, + ) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + sub = p.add_subparsers(dest="cmd") + + c = sub.add_parser("case", help="internal: one standalone case (fresh process)") + c.add_argument("--shape", choices=list(SHAPE_BINS), required=True) + c.add_argument("--reader", choices=["pa", "pa1", "rs"], required=True) + c.add_argument("--path", required=True) + c.add_argument("--tasks", type=int, default=24) + c.add_argument("--workdir", default="/tmp") + + p.add_argument("--fixture-root", default=None) + p.add_argument("--outdir", default=None) + p.add_argument("--shapes", default="auto,write,tensorscp") + p.add_argument( + "--parts", + default=None, + help=( + "comma list of standalone,ray_local,ray_s3 (default: all three when " + "ARROW_RS_S3_BUCKET / --s3-bucket is set, else the first two)" + ), + ) + p.add_argument("--s3-bucket", default=os.environ.get("ARROW_RS_S3_BUCKET")) + p.add_argument("--repeat", type=int, default=3) + p.add_argument("--warmup", type=int, default=1) + p.add_argument("--tasks", type=int, default=24, help="auto standalone task count") + p.add_argument( + "--budget-sweep", + default=None, + help=( + "comma list of MiB values; adds an rs_b Ray arm per value with " + "RAY_DATA_ARROW_RS_DECODE_BUDGET_BYTES= (pareto sweep, TODO 1y)" + ), + ) + p.add_argument( + "--k-sweep", + default=None, + help=( + "comma list of K values; adds an rs_k Ray arm per value with " + "RAY_DATA_ARROW_RS_K= (intra-task K-split sweep, 1o)" + ), + ) + p.add_argument( + "--no-arena-sweep", + action="store_true", + help="skip the MALLOC_ARENA_MAX=2 variant of each arrow_rs Ray cell", + ) + args = p.parse_args() + + if args.cmd == "case": + run_case(args) + return + + if not args.fixture_root: + p.error("--fixture-root is required (see gen_local_fixtures.py)") + fixture_root = os.path.expanduser(args.fixture_root) + with open(os.path.join(fixture_root, "manifest.json")) as fh: + manifest = json.load(fh) + + shapes = [s.strip() for s in args.shapes.split(",") if s.strip()] + if args.parts: + parts = [x.strip() for x in args.parts.split(",") if x.strip()] + else: + parts = ["standalone", "ray_local"] + (["ray_s3"] if args.s3_bucket else []) + if "ray_s3" in parts and not args.s3_bucket: + p.error("ray_s3 needs --s3-bucket or ARROW_RS_S3_BUCKET") + + outdir = args.outdir or os.path.join( + HERE, "loss_triage_runs", time.strftime("%Y%m%d_%H%M%S") + ) + os.makedirs(outdir, exist_ok=True) + print(f"parts={parts} shapes={shapes} outdir={outdir}", flush=True) + + summary = {"parts": parts, "shapes": shapes, "cells": {}} + + def fixture_path(shape): + entry = manifest[SHAPE_FIXTURE[shape]] + return entry["path"] if isinstance(entry, dict) else entry + + # Sync fixtures for the S3 part once, up front. + s3_paths = {} + if "ray_s3" in parts: + bucket = args.s3_bucket.rstrip("/") + for shape in shapes: + s3_paths[shape] = f"{bucket}/loss_triage/{SHAPE_FIXTURE[shape]}" + s3_sync(fixture_path(shape), s3_paths[shape]) + + for shape in shapes: + shape_env = dict(SHAPE_ENV.get(shape, {})) + local_path = fixture_path(shape) + + for part in parts: + if part == "standalone" and shape in RAY_ONLY_SHAPES: + print(f"\n=== [{shape}] standalone — skipped (Ray-only shape)") + continue + print(f"\n=== [{shape}] {part} ===", flush=True) + cells = {} + + if part == "standalone": + for reader in ("pa", "pa1", "rs"): + cells[reader] = run_standalone_cell( + outdir, + f"{shape}.standalone.{reader}", + shape, + reader, + local_path, + shape_env, + args.tasks, + args.repeat, + args.warmup, + ) + else: + path = local_path if part == "ray_local" else s3_paths[shape] + extra_args = ["--mem-poll-s", "0.05"] + columns = None + if shape == "write": + extra_args += ["--consume", "write_parquet"] + elif shape == "agg": + # Projection pushdown (the release aggregates read 2-3 + # columns of a wide table) + the shuffle/aggregate consume. + extra_args += ["--consume", "groupby", "--groupby-key", "s0"] + columns = ["id", "s0"] + env = dict(shape_env) + env["RAY_DATA_PARQUET_BIN_PACKING_BYTES"] = str(SHAPE_BINS[shape]) + arms = [("pyarrow", "pa", {}), ("arrow_rs", "rs", {})] + if not args.no_arena_sweep: + arms.append(("arrow_rs", "rs_arena2", {"MALLOC_ARENA_MAX": "2"})) + # M56: env-only glibc page retention (never munmap, never + # trim) removed 50-65% of standalone decode wall on fat + # shapes by taking fault amplification to ~0. In-Ray arm + # tests whether the win survives short-lived tasks, at the + # known cost of up to 1.5x per-worker RSS. + arms.append( + ( + "arrow_rs", + "rs_retain", + { + "MALLOC_MMAP_MAX_": "0", + "MALLOC_TRIM_THRESHOLD_": "-1", + }, + ) + ) + if args.budget_sweep: + for mib in [ + int(x) for x in args.budget_sweep.split(",") if x.strip() + ]: + arms.append( + ( + "arrow_rs", + f"rs_b{mib}", + { + "RAY_DATA_ARROW_RS_DECODE_BUDGET_BYTES": str( + mib * MiB + ) + }, + ) + ) + if args.k_sweep: + for k in [int(x) for x in args.k_sweep.split(",") if x.strip()]: + arms.append( + ("arrow_rs", f"rs_k{k}", {"RAY_DATA_ARROW_RS_K": str(k)}) + ) + for reader, tag, arm_env in arms: + cells[tag] = median_cell( + outdir, + f"{shape}.{part}.{tag}", + args.repeat, + warmup=args.warmup, + path=path, + reader=reader, + concurrency=None, + columns=columns, + extra_env={**env, **arm_env}, + extra_args=extra_args, + ) + + summary["cells"][f"{shape}.{part}"] = cells + + # ---------------- ratio table ---------------- + print( + "\n\n================ LOSS TRIAGE SUMMARY (R = arrow_rs/pyarrow) ================" + ) + if sys.platform == "darwin": + print("(macOS: peak_uss is None and wall is not certifiable — smoke run only)") + header = ( + f"{'cell':<26} {'wall R':>8} {'peak mem R':>11} {'task USS R':>11} " + f"{'arena2 mem R':>13} {'spill T/B GB':>13}" + ) + print(header) + print("-" * len(header)) + for key, cells in summary["cells"].items(): + pa_res = cells.get("pa") or cells.get("pyarrow") or {} + rs_res = cells.get("rs") or {} + ar_res = cells.get("rs_arena2") or {} + if "standalone" in key: + wall_r = ratio(_num(rs_res, "wall_s"), _num(pa_res, "wall_s")) + mem_r = ratio(_num(rs_res, "peak_rss_mib"), _num(pa_res, "peak_rss_mib")) + task_r = ratio(_num(rs_res, "end_rss_mib"), _num(pa_res, "end_rss_mib")) + ar_r = None + kern_r = ratio( + _num(rs_res, "wall_s"), _num(cells.get("pa1") or {}, "wall_s") + ) + else: + wall_r = ratio(_num(rs_res, "wall_s"), _num(pa_res, "wall_s")) + mem_r = ratio( + _num(rs_res, "peak_uss_gb"), _num(pa_res, "peak_uss_gb") + ) or ratio(_num(rs_res, "peak_rss_gb"), _num(pa_res, "peak_rss_gb")) + task_r = ratio( + _num(rs_res, "read_avg_max_uss_gb"), _num(pa_res, "read_avg_max_uss_gb") + ) + ar_r = ratio(_num(ar_res, "peak_uss_gb"), _num(pa_res, "peak_uss_gb")) + kern_r = None + rt_res = cells.get("rs_retain") or {} + rt_wall = ratio(_num(rt_res, "wall_s"), _num(pa_res, "wall_s")) + rt_uss = ratio(_num(rt_res, "peak_uss_gb"), _num(pa_res, "peak_uss_gb")) + if rt_wall is not None: + suffix_retain = ( + f" retain: wall {rt_wall:.2f} uss {rt_uss:.2f}" + if rt_uss is not None + else f" retain: wall {rt_wall:.2f}" + ) + else: + suffix_retain = "" + fmt = lambda v: f"{v:.2f}" if v is not None else "—" # noqa: E731 + sp_t, sp_b = _num(rs_res, "spilled_gb"), _num(pa_res, "spilled_gb") + spill = ( + f"{sp_t:.1f}/{sp_b:.1f}" if sp_t is not None and sp_b is not None else "—" + ) + suffix = "" + if "standalone" in key and kern_r is not None: + suffix = f" wall R vs pa-1t: {kern_r:.2f}" + elif "standalone" not in key: + suffix = suffix_retain + print( + f"{key:<26} {fmt(wall_r):>8} {fmt(mem_r):>11} {fmt(task_r):>11} " + f"{fmt(ar_r):>13} {spill:>13}" + suffix + ) + if "standalone" not in key: + for t in sorted(t for t in cells if t.startswith(("rs_b", "rs_k"))): + r = cells.get(t) or {} + w = ratio(_num(r, "wall_s"), _num(pa_res, "wall_s")) + m = ratio(_num(r, "peak_uss_gb"), _num(pa_res, "peak_uss_gb")) + tu = ratio( + _num(r, "read_avg_max_uss_gb"), + _num(pa_res, "read_avg_max_uss_gb"), + ) + print( + f" {t:<22} wall {fmt(w)} peak mem {fmt(m)} " + f"task USS {fmt(tu)}" + ) + print( + "\nRead it as: loss in standalone => decoder; only in ray_local => Ray worker/\n" + "allocator (arena2 column collapsing confirms glibc arenas); only in ray_s3 =>\n" + "crate S3 path. Full metrics: " + os.path.join(outdir, "summary.json") + ) + + with open(os.path.join(outdir, "summary.json"), "w") as fh: + json.dump(summary, fh, indent=2) + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/patch_crate_parquet.sh b/release/nightly_tests/dataset/arrow_rs_probe/patch_crate_parquet.sh new file mode 100755 index 000000000000..35c80099bf0b --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/patch_crate_parquet.sh @@ -0,0 +1,107 @@ +#!/usr/bin/env bash +# --------------------------------------------------------------------------- +# A/B tool for TODO item 1o (fat_col wall loss): rebuild ray_data_arrow_rs +# against a locally patched parquet 59.1.0, or revert to the stock crates.io +# build. The patch (patches/parquet-59.1.0-dict-reserve.diff) pre-sizes the +# values buffer in OffsetBuffer::extend_from_dictionary from the dictionary's +# average value length — the omission that makes fat dictionary-encoded binary +# columns pay ~one extra full copy (findings T20; prior art apache/arrow-rs +# #5250, which used an exact per-key sum and regressed small strings — this is +# the O(1) variant that doesn't). +# +# bash patch_crate_parquet.sh # vendor + patch + rebuild (PATCHED arm) +# REVERT=1 bash patch_crate_parquet.sh # restore stock parquet + rebuild +# +# The A/B for the fatcol stage is then: +# ONLY=fatcol bash run_replication.sh # stock arm +# bash patch_crate_parquet.sh +# ONLY=fatcol bash run_replication.sh # patched arm +# REVERT=1 bash patch_crate_parquet.sh # leave box stock +# +# Mechanics: copies the pristine parquet-59.1.0 source out of the cargo +# registry cache into /vendor/, applies the diff, and points +# [patch.crates-io] at it via /.cargo/config.toml (never Cargo.toml, so +# the tree stays clean). Cargo.lock is backed up before the first patch and +# restored on REVERT — do not commit a lock file that lost its parquet +# checksum line. vendor/ and .cargo/ are build-local; never commit them. +# --------------------------------------------------------------------------- +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "$SCRIPT_DIR/../../../.." && pwd)" +CRATE="$REPO/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs" +DIFF="$SCRIPT_DIR/patches/parquet-59.1.0-dict-reserve.diff" +VENDOR="$CRATE/vendor/parquet-59.1.0-dict-reserve" +CARGO_CFG="$CRATE/.cargo/config.toml" +LOCK_BAK="$CRATE/Cargo.lock.stock" + +say() { printf '\n\033[1;35m### %s\033[0m\n' "$*"; } + +# Environment: prefer the probe env.sh (Linux boxes), else whatever venv can +# already see maturin, else the repo-local .venv (macOS dev). +if [ -f "$SCRIPT_DIR/env.sh" ]; then + source "$SCRIPT_DIR/env.sh" +elif ! command -v maturin >/dev/null 2>&1 && [ -f "$REPO/.venv/bin/activate" ]; then + source "$REPO/.venv/bin/activate" +fi +export PATH="$HOME/.cargo/bin:$PATH" +command -v cargo >/dev/null || { echo "cargo not on PATH — run setup.sh first"; exit 1; } +command -v maturin >/dev/null || { echo "maturin not on PATH — 'uv pip install maturin' into the venv"; exit 1; } + +rebuild() { + say "rebuilding ray_data_arrow_rs (maturin develop --release)" + ( cd "$CRATE" && unset CONDA_PREFIX && VIRTUAL_ENV="${VIRTUAL_ENV}" maturin develop --release ) + python - <<'PYEOF' +import ray_data_arrow_rs as rs +print("crate imports OK:", rs.__name__) +PYEOF +} + +if [ "${REVERT:-0}" = "1" ]; then + say "REVERT: removing the parquet patch" + rm -f "$CARGO_CFG" + rmdir "$CRATE/.cargo" 2>/dev/null || true + rm -rf "$CRATE/vendor" + if [ -f "$LOCK_BAK" ]; then + mv "$LOCK_BAK" "$CRATE/Cargo.lock" + fi + rebuild + say "stock crate restored" + exit 0 +fi + +[ -f "$DIFF" ] || { echo "missing $DIFF"; exit 1; } + +say "vendoring pristine parquet 59.1.0 from the cargo registry cache" +( cd "$CRATE" && cargo fetch ) +SRC="$(find "$HOME"/.cargo/registry/src -maxdepth 2 -type d -name parquet-59.1.0 2>/dev/null | head -1)" +[ -n "$SRC" ] || { echo "parquet-59.1.0 not in the registry cache even after cargo fetch"; exit 1; } +rm -rf "$VENDOR" +mkdir -p "$(dirname "$VENDOR")" +cp -R "$SRC" "$VENDOR" + +say "applying $DIFF" +# Diff paths are relative to the arrow-rs repo root (a/parquet/src/...); +# the vendored dir is the parquet crate root, so strip two components. +patch -d "$VENDOR" -p2 < "$DIFF" + +say "pointing [patch.crates-io] at the vendored copy (via .cargo/config.toml)" +[ -f "$LOCK_BAK" ] || cp "$CRATE/Cargo.lock" "$LOCK_BAK" +mkdir -p "$CRATE/.cargo" +cat > "$CARGO_CFG" <<'CFGEOF' +# Written by patch_crate_parquet.sh (TODO 1o A/B) — NEVER COMMIT. +# REVERT=1 bash patch_crate_parquet.sh removes this and restores Cargo.lock. +[patch.crates-io] +parquet = { path = "vendor/parquet-59.1.0-dict-reserve" } +CFGEOF + +rebuild + +say "verifying the patched source is what got built" +if ( cd "$CRATE" && cargo metadata --format-version 1 --offline 2>/dev/null || cd "$CRATE" && cargo metadata --format-version 1 ) | grep -q "parquet-59.1.0-dict-reserve"; then + echo "OK: cargo resolves parquet to vendor/parquet-59.1.0-dict-reserve" +else + echo "WARNING: vendored parquet not visible in cargo metadata — check [[patch.unused]] in Cargo.lock" + exit 1 +fi +say "patched crate installed — run the fatcol stage now" diff --git a/release/nightly_tests/dataset/arrow_rs_probe/patches/parquet-59.1.0-dict-reserve.diff b/release/nightly_tests/dataset/arrow_rs_probe/patches/parquet-59.1.0-dict-reserve.diff new file mode 100644 index 000000000000..c1a295e20745 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/patches/parquet-59.1.0-dict-reserve.diff @@ -0,0 +1,25 @@ +diff --git a/parquet/src/arrow/buffer/offset_buffer.rs b/parquet/src/arrow/buffer/offset_buffer.rs +index ab67694..401a56c 100644 +--- a/parquet/src/arrow/buffer/offset_buffer.rs ++++ b/parquet/src/arrow/buffer/offset_buffer.rs +@@ -97,6 +97,20 @@ impl OffsetBuffer { + ) -> Result<()> { + self.offsets.reserve(keys.len()); + ++ // Pre-size `values` from the dictionary's average value length (O(1), ++ // no pass over the keys). Without this, each key's value lands in an ++ // unreserved Vec and amortized doubling re-copies roughly all gathered ++ // data one extra time — for large dictionary values (e.g. binary image ++ // columns) that copy dominates the decode. The estimate is only a ++ // reservation hint: skewed key distributions still grow or waste at ++ // most one doubling's worth, and the exact-sum alternative costs a ++ // second bounds-checked pass that measurably regresses small values. ++ let dict_len = dict_offsets.len().saturating_sub(1); ++ if dict_len > 0 { ++ let avg_len = dict_values.len() / dict_len; ++ self.values.reserve(avg_len.saturating_mul(keys.len())); ++ } ++ + for key in keys { + let index = key.as_usize(); + if index + 1 >= dict_offsets.len() { diff --git a/release/nightly_tests/dataset/arrow_rs_probe/plasma_put_micro.py b/release/nightly_tests/dataset/arrow_rs_probe/plasma_put_micro.py new file mode 100644 index 000000000000..5496ac695048 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/plasma_put_micro.py @@ -0,0 +1,72 @@ +"""In-worker A/B: N concurrent Ray tasks each decode 5 row groups of a local +lineitem file (crate or PyArrow) into a ~5-chunk table and time ray.put of it. +Usage: python put_micro.py [tasks_per_arm]""" +import sys +import time +import os +import statistics as st +import resource +import pyarrow.parquet as pq +import ray + +p, ncpu = sys.argv[1], int(sys.argv[2]) +ntasks = int(sys.argv[3]) if len(sys.argv) > 3 else ncpu + + +@ray.remote(num_cpus=1) +def task(arm, path, rgs, bs): + import pyarrow as pa + import pyarrow.parquet as pq + import time + + t0 = time.perf_counter() + if arm == "pa": + pf = pq.ParquetFile(path) + t = pa.Table.from_batches(list(pf.iter_batches(batch_size=bs, row_groups=rgs))) + else: + import ray_data_arrow_rs as rs + + h = rs.open_parquet_file(path, page_index=False) + r = h.read_row_groups(row_groups=rgs, batch_size=bs, k=1) + t = pa.Table.from_batches(list(pa.RecordBatchReader.from_stream(r))) + if arm == "rs_combine": + t = t.combine_chunks() + t1 = time.perf_counter() + ru0 = resource.getrusage(resource.RUSAGE_SELF) + c0 = time.process_time() + ray.put(t) + t2 = time.perf_counter() + ru1 = resource.getrusage(resource.RUSAGE_SELF) + c1 = time.process_time() + return { + "decode": t1 - t0, + "put": t2 - t1, + "nbytes": t.nbytes, + "chunks": t.column(0).num_chunks, + "pid": os.getpid(), + "cpu": c1 - c0, + "minflt": ru1.ru_minflt - ru0.ru_minflt, + "majflt": ru1.ru_majflt - ru0.ru_majflt, + "nivcsw": ru1.ru_nivcsw - ru0.ru_nivcsw, + "nvcsw": ru1.ru_nvcsw - ru0.ru_nvcsw, + } + + +ray.init(address="local", num_cpus=ncpu, include_dashboard=False, log_to_driver=False) +nrg = pq.ParquetFile(p).metadata.num_row_groups +rows = pq.ParquetFile(p).metadata.row_group(0).num_rows +windows = [list(range(i, min(i + 5, nrg))) for i in range(0, nrg - 4, 5)] +arms = sys.argv[4].split(",") if len(sys.argv) > 4 else ["pa", "rs", "rs_combine"] +for arm in arms: + t0 = time.perf_counter() + refs = [task.remote(arm, p, windows[i % len(windows)], rows) for i in range(ntasks)] + res = ray.get(refs) + wall = time.perf_counter() - t0 + puts = sorted(r["put"] for r in res) + decs = [r["decode"] for r in res] + print( + f"[{arm:10s} cpus={ncpu} tasks={ntasks}] wall {wall:5.1f}s put p50 {st.median(puts)*1000:7.0f}ms mean {st.mean(puts)*1000:7.0f}ms max {puts[-1]*1000:7.0f}ms | decode p50 {st.median(decs):.2f}s | put cpu p50 {st.median(r['cpu'] for r in res)*1000:6.0f}ms minflt p50 {st.median(r['minflt'] for r in res):7.0f} majflt {st.median(r['majflt'] for r in res):4.0f} nvcsw {st.median(r['nvcsw'] for r in res):5.0f} nivcsw {st.median(r['nivcsw'] for r in res):5.0f} | {res[0]['nbytes']/1e6:.0f}MB chunks {res[0]['chunks']} workers {len({r['pid'] for r in res})}", + flush=True, + ) + del refs, res +ray.shutdown() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/plasma_ser_micro.py b/release/nightly_tests/dataset/arrow_rs_probe/plasma_ser_micro.py new file mode 100644 index 000000000000..6bb391b1211a --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/plasma_ser_micro.py @@ -0,0 +1,101 @@ +"""Single-process A/B: same Parquet file decoded by the arrow-rs crate and by +PyArrow (both as ~5-chunk tables like the release read op), then Ray +serialization (pickle only) and ray.put (pickle + copy into plasma) timed. +Usage: python ser_micro.py [batch_rows]""" +import sys +import time +import gc +import os +import pyarrow as pa +import pyarrow.parquet as pq +import ray +import ray_data_arrow_rs as rs + +p = sys.argv[1] +pf = pq.ParquetFile(p) +RGS = ( + [int(x) for x in os.environ["RG"].split(",")] + if os.environ.get("RG") + else list(range(pf.metadata.num_row_groups)) +) +n = sum(pf.metadata.row_group(i).num_rows for i in RGS) +print( + "file rows", + pf.metadata.num_rows, + "row_groups", + pf.metadata.num_row_groups, + "using", + RGS, + "rows", + n, +) +bs = int(sys.argv[2]) if len(sys.argv) > 2 else max(1, n // 5) + + +def t_pa(): + return pa.Table.from_batches(list(pf.iter_batches(batch_size=bs, row_groups=RGS))) + + +def t_rs(): + h = rs.open_parquet_file(p, page_index=False) + r = h.read_row_groups(row_groups=RGS, batch_size=bs, k=1) + rbr = pa.RecordBatchReader.from_stream(r) + return pa.Table.from_batches(list(rbr)) + + +def describe(name, t): + print( + f"[{name}] rows={t.num_rows} nbytes={t.nbytes/1e6:.1f}MB chunks/col={t.column(0).num_chunks}" + ) + oddities = [] + for col in t.columns: + for ch in col.chunks: + if ch.offset != 0: + oddities.append( + (col._name if hasattr(col, "_name") else "?", "offset", ch.offset) + ) + for b in ch.buffers(): + if b is not None and b.address % 64: + oddities.append(("align", b.address % 64)) + print(f" non-zero offsets / misaligned buffers: {len(oddities)} {oddities[:3]}") + return t + + +def timeit(label, fn, reps=4): + ts = [] + for _ in range(reps): + gc.collect() + t0 = time.perf_counter() + r = fn() + ts.append(time.perf_counter() - t0) + del r + print(f" {label:<28} " + " ".join(f"{x*1000:7.1f}ms" for x in ts)) + + +ray.init(address="auto", log_to_driver=False) if os.environ.get( + "USE_EXISTING" +) else ray.init( + address="local", num_cpus=2, include_dashboard=False, log_to_driver=False +) +ctx = ray._private.worker.global_worker.get_serialization_context() +for name, mk in (("pyarrow", t_pa), ("arrow-rs", t_rs)): + t0 = time.perf_counter() + t = mk() + print(f"\n[{name}] decode {time.perf_counter()-t0:.2f}s") + describe(name, t) + timeit("serialize (pickle only)", lambda: ctx.serialize(t)) + timeit("ray.put (pickle+plasma)", lambda: ray.put(t)) + tc = t.combine_chunks() # copies into pyarrow's pool, 1 chunk + timeit("ray.put combine_chunks", lambda: ray.put(tc)) + tcopy = pa.Table.from_batches( + [ + pa.RecordBatch.from_arrays( + [c.take(pa.array(range(len(c)))) for c in b.columns], + names=b.schema.names, + ) + for b in t.to_batches() + ] + ) + timeit("ray.put pa-realloc copy", lambda: ray.put(tcopy)) + t = tc = tcopy = None # noqa: F841 - release before the next arm +ray.shutdown() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/read_probe.py b/release/nightly_tests/dataset/arrow_rs_probe/read_probe.py new file mode 100644 index 000000000000..05dae1cf2af4 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/read_probe.py @@ -0,0 +1,474 @@ +#!/usr/bin/env python3 +"""Single-node read probe: arrow-rs vs PyArrow, reporting BOTH time and memory. + +Reproduces the *read* portion of a release test on ONE node so we can measure and +optimize the two cases where arrow-rs was worse in the release run without a +multi-node cluster: + - mix.8ds_equal_random_mix : time 1.67x worse (imagenet, many tiny row groups, S3) + - wide_schema_pipeline_primitives : mem 1.50x worse (5000 columns, S3) +Both effects (read wall time and per-worker decode memory) are single-worker +properties, so one node + S3 is enough. + +IMPORTANT — why this must run on Linux + real S3: + * The crate's byte-budgeted windowed-async decode (the "working set is a page, + not the whole row group" property) only runs on the S3 path. On local disk a + lone big row group is read whole (K=1), so the memory win does not appear + locally — you MUST point --path at S3 to exercise it. + * `peak_uss_gb` (the real per-worker private cost) is populated on Linux only; + it stays None on macOS (shared pages make RSS misleading there). + +For each reader it reports: + - wall_s : wall time of read + consume + - worker_cpu_s : summed CPU seconds across Ray worker processes during the read + - cpu_over_wall : worker_cpu_s / wall_s. The diagnostic (use --concurrency 1): + ~1 => CPU-bound decode -> optimize decode + <<1 => I/O-waiting on S3 -> optimize prefetch + - peak_rss_gb : peak summed RSS across Ray workers + - peak_uss_gb : peak summed USS (Linux only) -- the metric of record + - read_wall_s / read_output_gb / read_avg_max_uss_gb : from Ray's own op stats + +Run each reader in its OWN process (Ray + the crate load once per process): + + # CPU-bound-vs-IO diagnostic: force a single read task + python read_probe.py --preset wide_schema --reader pyarrow --concurrency 1 + python read_probe.py --preset wide_schema --reader arrow_rs --concurrency 1 + + # Realistic memory: let it fan out + python read_probe.py --preset imagenet --reader arrow_rs + +Presets encode the release-test read (path + columns); override with --path/--columns +(the exact S3 prefixes drift — confirm with `aws s3 ls` on the box). Set +MALLOC_ARENA_MAX=2, or LD_PRELOAD a jemalloc .so, to A/B the allocator on the +arrow_rs run. +""" +import argparse +import os +import threading +import time +from typing import Any, Dict, List, Optional + +import psutil + +PRESETS = { + # wide_schema_pipeline_primitives: ~550MB, 5000 columns. mem was 1.50x worse. + # data_type variants exist under .../wide_schema/{primitives,tensors,objects,nested_structs}. + "wide_schema": { + "path": "s3://ray-benchmark-data-internal-us-west-2/wide_schema/primitives", + "columns": None, + }, + # mix.8ds_equal_random_mix reads imagenet per dataset; the read is the imagenet + # decode over many tiny row groups (the K=1 case). time was 1.67x worse. + "imagenet": { + "path": "s3://ray-benchmark-data-internal-us-west-2/imagenet/parquet", + "columns": ["image", "label"], + }, +} + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--preset", choices=list(PRESETS), default=None) + p.add_argument("--path", type=str, default=None) + p.add_argument("--columns", nargs="+", default=None) + p.add_argument( + "--reader", + choices=["pyarrow", "arrow_rs"], + required=True, + help="Which V2 reader to use (sets use_arrow_rs_parquet_reader).", + ) + p.add_argument( + "--concurrency", + type=int, + default=None, + help=( + "Cap read tasks AND pin output blocks (sets override_num_blocks too). " + "Use 1 for the CPU-bound-vs-IO diagnostic." + ), + ) + p.add_argument( + "--task-concurrency", + type=int, + default=None, + help=( + "Cap concurrent read tasks WITHOUT touching override_num_blocks — use " + "this (not --concurrency) whenever the point of the run is the bin " + "geometry, so a process holds exactly one bin at a time and per-task " + "USS is attributable to one bin. (On the Parquet V2 path " + "override_num_blocks is inert anyway — the footer indexer sets " + "yields_read_units=True, so read_api.py:546-556 skips the partitioner " + "— but keeping it out of the command line keeps that an observation " + "rather than a dependency.)" + ), + ) + p.add_argument( + "--mem-poll-s", + type=float, + default=None, + help=( + "DataContext.memory_usage_poll_interval_s for the in-task MemoryProfiler " + "(default 1.0s, context.py:984). Short read tasks get one sample or none " + "at 1 Hz, which silently flattens per-task USS — set ~0.05 for any run " + "whose verdict is a USS number." + ), + ) + p.add_argument( + "--consume", + choices=["iter_bundles", "count", "write_parquet", "groupby"], + default="iter_bundles", + help=( + "write_parquet replicates the release write_parquet test (1aa): the " + "read fuses into the write task, stats come from ds._write_ds (which " + "materializes, so no capture_executor race). Output is deleted after " + "the run. groupby replicates the release aggregate_groups pipeline " + "shape (read -> shuffle -> aggregate): ds.groupby(--groupby-key)" + ".count(), consumed with capture_executor=True — the read op's " + "per-task metrics still resolve through the grouped dataset's stats." + ), + ) + p.add_argument( + "--groupby-key", + default="s0", + help="Grouping column for --consume groupby (fixture string cols are s0..).", + ) + p.add_argument( + "--write-path", + default=None, + help="Output dir for --consume write_parquet (default: /probe_write_out).", + ) + # 50 Hz (20 ms): a 5000-column decode's builder-flush spike is short; at the + # old 10 Hz it was caught on some runs and missed on others, giving ~1.5 GB of + # run-to-run variance in peak_uss. The deterministic metric of record is Ray's + # own per-task max USS (``read_avg_max_uss_gb``); this sampler is the backup. + p.add_argument("--sample-hz", type=float, default=50.0) + args = p.parse_args() + if args.preset: + args.path = args.path or PRESETS[args.preset]["path"] + if args.columns is None: + args.columns = PRESETS[args.preset]["columns"] + if not args.path: + p.error("need --path or --preset") + return args + + +class WorkerMemSampler: + """Samples summed RSS/USS/CPU across Ray worker processes at a fixed rate. + + Ray runs each read task in a separate ``ray::`` worker process, so decode memory + lives there, not in this driver. We match those processes by cmdline and track the + peak of their *summed* RSS (and USS on Linux) plus total CPU consumed while sampling. + """ + + def __init__(self, interval_s: float, root_pid: Optional[int] = None): + self._interval_s = interval_s + self._stop = threading.Event() + self._thread: Optional[threading.Thread] = None + self.peak_rss = 0 + self.peak_uss = 0 # stays 0 if USS is unavailable (macOS) + self._uss_ok = True + # PID of OUR node's raylet. When set, we sum only its descendant workers, + # not every ``ray::`` process on the box — otherwise the workspace's managed + # Ray (:6379) workers get summed in, polluting the peak. None => fall back to + # the cmdline heuristic (macOS / when the raylet pid can't be found). + self._root_pid = root_pid + self.matched_workers = 0 + # Track CPU as (last_seen_total) per pid, summed into cpu_seconds on exit. + self._cpu_last: Dict[int, float] = {} + self.cpu_seconds = 0.0 + + def _ray_workers(self) -> List[psutil.Process]: + # Preferred: OUR raylet + its descendant workers only. A private local Ray + # instance owns exactly its read workers here, so this isolates us from the + # workspace's :6379 node cleanly (ray:: proctitles can't be matched by + # session dir — setproctitle overwrites the cmdline). + if self._root_pid is not None: + try: + root = psutil.Process(self._root_pid) + procs = [root] + root.children(recursive=True) + self.matched_workers = len(procs) + return procs + except psutil.NoSuchProcess: + return [] + out = [] + for proc in psutil.process_iter(["name", "cmdline"]): + try: + cmd = " ".join(proc.info.get("cmdline") or []) + if "ray::" in cmd or "raylet" in (proc.info.get("name") or ""): + out.append(proc) + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + self.matched_workers = len(out) + return out + + def _sample(self): + rss = uss = 0 + for proc in self._ray_workers(): + try: + if self._uss_ok: + try: + mi = proc.memory_full_info() + uss += getattr(mi, "uss", 0) + rss += mi.rss + except (psutil.AccessDenied, NotImplementedError): + self._uss_ok = False + rss += proc.memory_info().rss + else: + rss += proc.memory_info().rss + ct = proc.cpu_times() + total = ct.user + ct.system + pid = proc.pid + prev = self._cpu_last.get(pid) + # Only accumulate forward deltas (new pid or grew); ignore pid reuse noise. + if prev is not None and total >= prev: + self.cpu_seconds += total - prev + self._cpu_last[pid] = total + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + self.peak_rss = max(self.peak_rss, rss) + if self._uss_ok: + self.peak_uss = max(self.peak_uss, uss) + + def _run(self): + while not self._stop.wait(self._interval_s): + self._sample() + + def __enter__(self): + self._sample() + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + return self + + def __exit__(self, *exc): + self._stop.set() + if self._thread: + self._thread.join() + self._sample() + + +def _gb(b: Optional[float]) -> Optional[float]: + return round(b / (1024**3), 4) if b else b + + +def collect_read_op_metrics(ds) -> Dict[str, Any]: + """Pull the read operator's wall time / output bytes / per-task USS from + Ray stats. + + GE1 (2026-08-11) returned ``read_avg_max_uss_gb=None`` in most cells. + Root-caused 2026-08-11 (reproduced locally with a faked-USS + MemoryProfiler): the workers always reported USS; the loss was a + driver-side snapshot race. ``Dataset._execute_to_iterator`` caches + ``executor.get_stats()`` right after the FIRST bundle, and + ``iter_internal_ref_bundles()`` passes ``capture_executor=False`` — so + ``get_stats_summary()`` fell back to that mid-execution snapshot, taken + before the last read task's ``on_task_finished`` populated + ``average_max_uss_per_task``. ListFiles always had values because it + finishes long before consumption ends. Fixed in ``main()`` by consuming + via ``_execute_to_iterator(capture_executor=True)`` so this reads the + post-shutdown final stats. The instrumentation stays: (a) per-node + ``uss_debug`` dump, (b) ``max_uss_per_task`` (the worst task, often the + OOM-relevant number), (c) labeled ``uss_fallback_*`` when the read node + has no USS — never silently substituted into ``read_avg_max_uss_gb``, + which stays the metric of record.""" + from ray.data._internal.stats import DatasetStatsSummary + + out: Dict[str, Any] = {} + debug = [] + try: + summary = ds.get_stats_summary(detail=True) + nodes = DatasetStatsSummary._collect_dataset_stats_summaries(summary) + read_hit = None # (op_summary, extra_metrics) of the first Read node + for node in nodes: + extra = getattr(node, "extra_metrics", {}) or {} + op_names = [op.operator_name or "" for op in (node.operators_stats or [])] + debug.append( + { + "operators": op_names, + "average_max_uss_bytes": extra.get("average_max_uss_per_task"), + "max_uss_bytes": extra.get("max_uss_per_task"), + "num_extra_metrics": len(extra), + } + ) + if read_hit is None: + for op in node.operators_stats or []: + if "Read" in (op.operator_name or ""): + read_hit = (op, extra) + break + if read_hit is not None: + op, extra = read_hit + out["read_operator_name"] = op.operator_name + out["read_wall_s"] = op.wall_time.sum if op.wall_time else None + out["read_output_gb"] = _gb( + op.output_size_bytes.sum if op.output_size_bytes else None + ) + out["read_avg_max_uss_gb"] = _gb(extra.get("average_max_uss_per_task")) + out["read_max_uss_gb"] = _gb(extra.get("max_uss_per_task")) + # --- the bin-bound denominators ------------------------------------- + # A read task == one bin, so decoded bytes/task is the *decoded* size of + # a bin (RAY_DATA_PARQUET_BIN_PACKING_BYTES budgets Parquet + # total_uncompressed_size, i.e. pages after decompression but still + # ENCODED — dictionary/RLE columns decode larger, so the knob is a proxy, + # not an identity). Reporting both lets the bound be stated against the + # decoded number and the expansion factor be seen rather than assumed. + ntasks = extra.get("num_tasks_finished") + out["read_num_tasks"] = ntasks + if ntasks: + out["read_bytes_per_task_gb"] = ( + round(out["read_output_gb"] / ntasks, 4) + if out.get("read_output_gb") + else None + ) + # max/avg over tasks: ~1.0 => every task costs the same (bounded); + # rising with task count => the worker is retaining across tasks + # (allocator retention or a real leak), which no bin cap can bound. + a, m = ( + extra.get("average_max_uss_per_task"), + extra.get("max_uss_per_task"), + ) + if a and m: + out["uss_max_over_avg"] = round(m / a, 3) + if out.get("read_avg_max_uss_gb") is None: + for i, row in enumerate(debug): + if row["average_max_uss_bytes"]: + out["uss_fallback_avg_gb"] = _gb(row["average_max_uss_bytes"]) + out["uss_fallback_max_gb"] = _gb(row["max_uss_bytes"]) + out["uss_fallback_source"] = ( + ",".join(row["operators"]) or f"node_{i}" + ) + break + out["uss_debug"] = debug + except Exception as e: # best-effort + out["read_op_metrics_error"] = repr(e) + return out + + +def main(): + args = parse_args() + + import ray + from ray.data.context import DataContext + + ctx = DataContext.get_current() + ctx.use_datasource_v2 = True + ctx.use_arrow_rs_parquet_reader = args.reader == "arrow_rs" + if args.mem_poll_s is not None: + ctx.memory_usage_poll_interval_s = args.mem_poll_s + + ray.init(ignore_reinit_error=True) + + # Pin subsequent Ray/GCS resolution to THIS instance's address. With + # RAY_ADDRESS=local next to the workspace's managed Ray (:6379), the stats + # collection re-resolves "local", finds two running instances, and errors — + # which is why read_avg_max_uss_gb (the deterministic per-task USS) has been + # empty every run. Pinning the explicit gcs address disambiguates it. + try: + gcs = ray.get_runtime_context().gcs_address + if gcs: + os.environ["RAY_ADDRESS"] = gcs + except Exception: + pass + + # Find OUR raylet's pid so the sampler sums only our node's workers. + root_pid: Optional[int] = None + try: + node = ray._private.worker._global_node # noqa: SLF001 + for name, procs in (node.all_processes or {}).items(): + if "raylet" in name.lower() and procs: + root_pid = procs[0].process.pid + break + except Exception: + pass + + read_kwargs: Dict[str, Any] = {} + if args.columns: + read_kwargs["columns"] = args.columns + if args.concurrency is not None: + read_kwargs["concurrency"] = args.concurrency + read_kwargs["override_num_blocks"] = args.concurrency + elif args.task_concurrency is not None: + read_kwargs["concurrency"] = args.task_concurrency + + print( + f"reader={args.reader} path={args.path} columns={args.columns} " + f"concurrency={args.concurrency} arena_max={os.environ.get('MALLOC_ARENA_MAX')} " + f"ld_preload={os.environ.get('LD_PRELOAD')} " + f"raylet_pid={root_pid} sample_hz={args.sample_hz}" + ) + + try: + with WorkerMemSampler(1.0 / args.sample_hz, root_pid=root_pid) as sampler: + t0 = time.perf_counter() + ds = ray.data.read_parquet(args.path, **read_kwargs) + if args.consume == "count": + ds.count() + elif args.consume == "write_parquet": + import shutil + import tempfile + + write_out = args.write_path or os.path.join( + tempfile.gettempdir(), "probe_write_out" + ) + shutil.rmtree(write_out, ignore_errors=True) + try: + ds.write_parquet(write_out) + finally: + # The fused read->write op's stats live on ds._write_ds + # (dataset.py: get_stats_summary falls through to it), which + # is materialized — read them via collect_read_op_metrics + # below as usual; only the bytes on disk need cleanup. + shutil.rmtree(write_out, ignore_errors=True) + elif args.consume == "groupby": + gds = ds.groupby(args.groupby_key).count() + bundle_iter, _, _ = gds._execute_to_iterator(capture_executor=True) + for _ in bundle_iter: + pass + ds = gds # read-op metrics resolve through the grouped stats + else: + # Same zero-copy consumption as ds.iter_internal_ref_bundles(), + # but with capture_executor=True so ds.get_stats_summary() can + # read the executor's post-shutdown final stats. + # iter_internal_ref_bundles() drops the executor + # (capture_executor=False), which pins get_stats_summary() to a + # stats snapshot cached after the FIRST bundle — taken before + # the last read task finishes, so average_max_uss_per_task is + # still empty in it (the GE1 "read_avg_max_uss_gb=None" bug). + bundle_iter, _, _ = ds._execute_to_iterator(capture_executor=True) + for _ in bundle_iter: + pass + wall = time.perf_counter() - t0 + + # Object-store spill for this session (T27: treatment-only spilling on + # autoscaling tpch shuffles) — parse the raylet memory summary; 0.0 when + # the session never spilled. + spilled_gb = 0.0 + try: + import re as _re + + import ray._private.internal_api as _api + + m = _re.search(r"Spilled (\d+) MiB", _api.memory_summary(stats_only=True)) + if m: + spilled_gb = round(int(m.group(1)) / 1024, 3) + except Exception: + spilled_gb = None + + result = { + "spilled_gb": spilled_gb, + "reader": args.reader, + "wall_s": round(wall, 3), + "worker_cpu_s": round(sampler.cpu_seconds, 3), + "cpu_over_wall": round(sampler.cpu_seconds / wall, 3) if wall else None, + "peak_rss_gb": _gb(sampler.peak_rss), + "peak_uss_gb": _gb(sampler.peak_uss) if sampler.peak_uss else None, + "sampled_workers": sampler.matched_workers, + **collect_read_op_metrics(ds), + } + print("\n=== RESULT ===") + for k, v in result.items(): + print(f" {k}: {v}") + finally: + # Shut Ray down cleanly so the next reader's process does not find (and the + # sampler does not sum) a stale second Ray instance's workers. + ray.shutdown() + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/realign_decomp.py b/release/nightly_tests/dataset/arrow_rs_probe/realign_decomp.py new file mode 100644 index 000000000000..bf8c827f6014 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/realign_decomp.py @@ -0,0 +1,224 @@ +"""Decompose the M51 wall loss (wide x high-expansion shapes): WHERE does the +per-batch cost live? (M52 on macOS; this script = the Linux confirmation.) + +Four experiments on the tensors_dict fixture, all standalone (no Ray, no S3 -- +transport is already acquitted: the loss reproduces in zero-IO local cells and +the S3 ratio is BETTER than local): + + 1. layers rs read timed in 4 layers x decode budgets: + A = crate decode + FFI import only (discard batches) + B = A + pa.Table.from_batches wrap + C = B + _cast_table_to extension realign (= the OLD path) + D = crate schema override + zero-copy FFI relabel + wrap + (= the M53 FIX: with_schema_override makes the crate decode + the extension's storage layout, then each batch is re-typed + through the C Data Interface against a prebuilt extension + schema object -- no cast, no per-batch pickle deserialize) + If A ~beats pa and C carries the loss, the decoder is innocent; + D should sit on top of B (realign cost ~gone). D asserts value + equality against C's result before timing. + 2. pa-ctrl pyarrow forced to the SAME batch counts (batch_size in rows). + pa emits extension-typed batches straight from C++ (it parses + ARROW:schema once per file), so its per-batch cost is the + C++ floor for constructing a 5000-col batch. + 3. variants one decoded batch: rebuild-schema+cast vs cached-schema+cast vs + zero-copy Table.from_arrays re-wrap. If none differ, the cost is + per-column construction churn, not the cast kernels. + 4. cols cast cost vs column count (5/50/500/all) -> us/col linearity. + +Usage: + python realign_decomp.py [--fixtures-root ~/arrow_rs_repl_fixtures] [--shape tensors_dict] +""" +import argparse +import json +import os +import sys +import time + +os.environ.setdefault("RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA", "1") +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +MiB = 1024 * 1024 + + +def _rs_stream(path, budget, knobs, override_schema=None): + import pyarrow as pa + import ray_data_arrow_rs as rs + + h = rs.open_parquet_file(path, page_index=False) + if override_schema is not None: + # M53: the crate decodes the extension's storage layout directly + # (large_list offsets), so the per-batch realign is a pure relabel. + h.with_schema_override(override_schema.__arrow_c_schema__()) + return pa.RecordBatchReader.from_stream( + h.read_row_groups( + row_groups=None, + columns=None, + batch_size=131072, + decode_budget_bytes=budget, + k=knobs["k"], + split_threshold_bytes=knobs["split"], + predicate_json=None, + fetch_window_mb=knobs["window"], + column_fetch_mb=knobs["column"], + prefetch_budget_mb=4 * max(knobs["window"], knobs["column"]), + ) + ) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument( + "--fixtures-root", default=os.path.expanduser("~/arrow_rs_repl_fixtures") + ) + ap.add_argument("--shape", default="tensors_dict") + ap.add_argument("--budgets-mib", default="32,128,512") + a = ap.parse_args() + + import pyarrow as pa + import pyarrow.parquet as pq + + from loss_triage import _pa_batches, _reader_knobs + from ray.data._internal.datasource_v2.readers.arrow_rs_parquet_file_reader import ( + _cast_table_to, + _ffi_relabel_batch, + _storage_override_schema, + ) + + man = json.load(open(os.path.join(a.fixtures_root, "manifest.json"))) + fix = man[a.shape] + files = sorted( + os.path.join(fix["path"], f) + for f in os.listdir(fix["path"]) + if f.endswith(".parquet") + ) + realign_fields = list(pq.read_schema(files[0])) + knobs = _reader_knobs() + budgets = [int(x) * MiB for x in a.budgets_mib.split(",")] + print( + f"shape={a.shape}: {fix['rows']} rows x {fix['columns']} cols, " + f"expansion {fix.get('enc_to_dec_ratio', '?')}x, {len(files)} files" + ) + + # ---- 1. layers ------------------------------------------------------- # + target_schema = pa.schema(realign_fields) + import ray_data_arrow_rs as _rs_mod + + inferred = pa.schema( + _rs_mod.open_parquet_file(files[0], page_index=False).metadata() + ) + override = _storage_override_schema(inferred, target_schema) + + def run_layer(layer, budget, collect=False): + t0 = time.perf_counter() + nb = 0 + out = [] + for f in files: + if layer == 3: + # D: override'd decode + zero-copy relabel (the M53 fix). + for b in _rs_stream(f, budget, knobs, override_schema=override): + nb += 1 + t = pa.Table.from_batches([_ffi_relabel_batch(b, target_schema)]) + if collect: + out.append(t) + continue + for b in _rs_stream(f, budget, knobs): + nb += 1 + if layer >= 1: + t = pa.Table.from_batches([b]) + if layer >= 2: + t = _cast_table_to(t, realign_fields) + if collect and layer >= 2: + out.append(t) + return time.perf_counter() - t0, nb, out + + # Correctness first: D's tables must equal C's exactly. + if override is not None: + _, _, c_tabs = run_layer(2, budgets[0], collect=True) + _, _, d_tabs = run_layer(3, budgets[0], collect=True) + eq = pa.concat_tables(d_tabs).equals(pa.concat_tables(c_tabs)) + print(f"\n[1] D-vs-C equality @ {budgets[0] // MiB}Mi: {eq}") + assert eq, "M53 fixed path diverges from the cast path" + del c_tabs, d_tabs + else: + print("\n[1] WARNING: no storage override derivable -- D falls back to C") + + print("[1] rs layers (A crate+FFI / B +wrap / C +cast = old / D = M53 fix)") + print( + f"{'budget':>8} {'batches':>7} {'A':>7} {'B':>7} {'C':>7} {'D':>7}" + " C-B per batch D-B per batch" + ) + counts = {} + for budget in budgets: + la = min((run_layer(0, budget) for _ in range(3)), key=lambda x: x[0]) + lb = min((run_layer(1, budget) for _ in range(3)), key=lambda x: x[0]) + lc = min((run_layer(2, budget) for _ in range(3)), key=lambda x: x[0]) + ld = min((run_layer(3, budget) for _ in range(3)), key=lambda x: x[0]) + nb = lc[1] + counts[budget] = nb + print( + f"{budget // MiB:>6}Mi {nb:>7} {la[0]:>6.2f}s {lb[0]:>6.2f}s " + f"{lc[0]:>6.2f}s {ld[0]:>6.2f}s " + f"{(lc[0] - lb[0]) / nb * 1000:7.1f} ms {(ld[0] - lb[0]) / nb * 1000:8.1f} ms" + ) + + # ---- 2. pa control at matched batch counts --------------------------- # + print( + "\n[2] pa control (single-thread, batch_size chosen to match rs batch counts)" + ) + total_rows = fix["rows"] + for budget in budgets: + nb = counts[budget] + bs = max(1, total_rows // nb) + + def run_pa(): + t0 = time.perf_counter() + n = 0 + typ = None + for f in files: + for t in _pa_batches(f, bs, use_threads=False): + n += 1 + if typ is None: + typ = str(t.schema.field(0).type) + return time.perf_counter() - t0, n, typ + + w, n, typ = min((run_pa() for _ in range(3)), key=lambda x: x[0]) + print(f" bs={bs:>6} rows: {w:5.2f}s batches={n:>3} col0={typ[:36]}") + + # ---- 3. cast variants on one batch ----------------------------------- # + print( + "\n[3] cast variants, one decoded batch @32Mi (per-batch ms; equal => churn, not kernels)" + ) + b = next(iter(_rs_stream(files[0], 32 * MiB, knobs))) + t = pa.Table.from_batches([b]) + cached = pa.schema(realign_fields, metadata=t.schema.metadata) + + def timeit(fn, n=10): + fn() + t0 = time.perf_counter() + for _ in range(n): + fn() + return (time.perf_counter() - t0) / n * 1000 + + print( + f" rebuild schema + cast : {timeit(lambda: t.cast(pa.schema(realign_fields, metadata=t.schema.metadata))):6.1f} ms" + ) + print(f" cached schema + cast : {timeit(lambda: t.cast(cached)):6.1f} ms") + print( + f" zero-copy from_arrays : {timeit(lambda: pa.Table.from_arrays(t.columns, schema=cached)):6.1f} ms" + ) + + # ---- 4. column linearity ---------------------------------------------- # + print("\n[4] cast cost vs column count (expect ~linear us/col)") + ncols_all = len(realign_fields) + for ncols in [5, 50, 500, ncols_all]: + sub = t.select(range(ncols)) + fields = realign_fields[:ncols] + ms = timeit(lambda: sub.cast(pa.schema(fields, metadata=sub.schema.metadata))) + print( + f" {ncols:>5} cols: {ms:8.2f} ms/batch ({ms / ncols * 1000:5.1f} us/col)" + ) + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/release_regression_probe.py b/release/nightly_tests/dataset/arrow_rs_probe/release_regression_probe.py new file mode 100644 index 000000000000..2aaa501a312a --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/release_regression_probe.py @@ -0,0 +1,498 @@ +#!/usr/bin/env python3 +"""Single-box replication of the A/B #5 >1.15 regressions (non-tpch legs). + +User cut 2026-08-28 (arrow_rs_docs/2026-08-27.md §11): every test where wall, +wUSS, or per-task USS R > 1.15 is a P0. This probe runs the single-node analog +of each NON-tpch P0 — the tpch queries go through tpch_probe.py, which already +takes --queries/--strategies (see run_release_regressions.sh for the pairing). + +Each cell runs the RELEASE benchmark script itself (read_and_consume_benchmark +/ groupby_benchmark / join_benchmark — same code, same public bucket) in a +fresh process per (cell, reader), with the same env pins the release yaml sets +(bin-packing bytes, node-mem monitor), toggling only +RAY_DATA_USE_ARROW_RS_PARQUET_READER. TEST_OUTPUT_JSON lands per cell, so every +run carries the full per-task dists (read_max_uss_per_task_dist etc.) that the +release A/B reports — the final table is computed from those, with the same +ratio convention (R = arrow_rs / pyarrow, <1 = arrow-rs better). + +What is deliberately NOT here, and why (doc §12 has the full map): + read_large_parquet_* s3://ray-benchmark-data-internal-* is ACCESS_DENIED + from our account. Analog = run_loss_triage.sh's + auto_rg S3 shape (same ~69 MiB-row-group geometry). + read_parquet_* same internal bucket (imagenet/parquet). The + read_parquet_binned cell below keeps the test's 64 MiB + bin pin but on public tpch lineitem — an approximation, + labeled as such. + wide_schema_objects internal bucket AND no local fixture for the "objects" + data-type (tensors fixtures don't cover it). Cannot be + replicated; release-only signal. + autoscaling variants one box has no autoscaler; each runs as its fixed-size + analog. A loss that lives in pool dynamics (T27) will + NOT show here — that absence is itself the signal. + +Downsizing: release tpch data is sf1000 (write_parquet) / sf100 (joins, +map_groups); one box gets sf100 / sf10 defaults, overridable per family. The +regime caveat from M35 stands: short fresh-session runs sit below the +allocator-churn floor, so a clean table here does NOT clear the retention +cluster — it separates "reproduces anywhere" from "needs the release regime". + +Usage: + python release_regression_probe.py --outdir DIR [--repeat 1] + [--only write_parquet,mapg_hash] [--write-sf 100] [--groupby-sf 10] + [--joins-sf 10] [--join-types right_outer] [--dry-run] + [--arms pa,rs,rseos] [--cpus 24,48,96] [--monitor-interval 1.0] +--arms adds the allocator arms of build 106096 (rstrim / rseos, see ARMS); +--cpus runs every cell once per local num_cpus value (cell name gets an +"@Ncpu" suffix) — the M107 col02 concurrency sweep; --monitor-interval +(seconds) is the node-memory sampler period (0.1 = the 10 Hz per-worker view +for the short q6/wide_schema sustained rows). +Needs AWS creds; ARROW_RS_S3_BUCKET redirects write_parquet's output (the +release write bucket s3://ray-data-write-benchmark may be unwritable to us). +""" +import argparse +import json +import os +import signal +import subprocess +import sys +import time + +HERE = os.path.dirname(os.path.abspath(__file__)) +DATASET_DIR = os.path.abspath(os.path.join(HERE, "..")) +TPCH = "s3://ray-benchmark-data/tpch/parquet" + +# Reader arms, by env, mirroring the release builds (arrow_rs_docs/ +# 2026-09-04.md): pa = PyArrow; rs = arrow-rs as shipped (since 2026-09-08 that +# INCLUDES one glibc malloc_trim(0) per read-task stream, so rs == rseos); +# rsnoeos = arrow-rs with that trim ablated; rstrim = arrow-rs + mallopt( +# M_TRIM_THRESHOLD, 0) once per worker (retired mechanism probe, M108). Both +# knobs are DataContext fields read only inside the arrow-rs read task, so +# they are inert for pa. +_RS = "RAY_DATA_USE_ARROW_RS_PARQUET_READER" +_KNOBS = ("RAY_DATA_ARROW_RS_MALLOC_TRIM", "RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS") +ARMS = { + "pa": {_RS: "0"}, + "rs": {_RS: "1"}, + "rstrim": {_RS: "1", "RAY_DATA_ARROW_RS_MALLOC_TRIM": "1"}, + "rseos": {_RS: "1", "RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS": "1"}, + "rsnoeos": {_RS: "1", "RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS": "0"}, +} + + +def arm_env(env, arm): + """Pin ``env`` to one arm: clear both allocator knobs, then set the arm's.""" + if arm not in ARMS: + raise SystemExit(f"unknown arm {arm!r}; choose from {sorted(ARMS)}") + for knob in _KNOBS: + env.pop(knob, None) + env.update(ARMS[arm]) + return env + + +# One cell, fresh process: import the release script as a module, patch its +# write root if asked, run its own parse_args()+main() under our sys.argv. +SNIPPET = r""" +import importlib, json, os, re, sys, time + +mod_name, argv_json, write_root, dry = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4] == "1" +sched_mem_gb = int(sys.argv[5]) +mod = importlib.import_module(mod_name) +if dry: + print("CELL_JSON " + json.dumps({"dry_run": True})) + raise SystemExit(0) +if write_root and hasattr(mod, "WRITE_PATH"): + mod.WRITE_PATH = write_root +sys.argv = [mod_name] + json.loads(argv_json) +import ray + +# _memory inflates only the SCHEDULING memory resource (nothing is allocated); +# stock, multi-JoinOperator/aggregator cells deadlock on one box: the operators' +# reservations consume the whole ~14.4GiB budget and upstream tasks starve +# (seen on tpch q2 sf10, PyArrow arm too). +# PROBE_NUM_CPUS caps the local instance's CPUs = the read/shuffle task +# concurrency on this box (the --cpus sweep); empty = every core. +num_cpus = int(os.environ.get("PROBE_NUM_CPUS", "0") or 0) +ray.init( + address="local", + **({"num_cpus": num_cpus} if num_cpus else {}), + **({"_memory": sched_mem_gb << 30} if sched_mem_gb else {}), +) +# The workspace's own Ray (:6379) coexists with this cell's local instance, and +# the state API's address autodetection dies on "multiple active Ray instances" +# — which silently emptied every per-task stats dist. Pin it to this cell. +import os + +os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address +t0 = time.monotonic() +mod.main(mod.parse_args()) +wall = time.monotonic() - t0 +spilled_gb = None +try: + import ray._private.internal_api as api + + m = re.search(r"Spilled (\d+) MiB", api.memory_summary(stats_only=True)) + spilled_gb = round(int(m.group(1)) / 1024, 3) if m else 0.0 +except Exception: + pass +print("CELL_JSON " + json.dumps({"wall_s": round(wall, 1), "spilled_gb": spilled_gb})) +""" + + +def cells(a): + """(name, module, argv, extra_env) per non-tpch P0. Env pins mirror + release_data_tests.yaml for the corresponding test.""" + out = [ + # Exact release replica — same script, same sf10 public data. + # P0 trip: tUSS max 1.18 (p50 0.93 win). + ( + "iter_batches_pyarrow", + "read_and_consume_benchmark", + [ + f"{TPCH}/sf10/lineitem", + "--format", + "parquet", + "--iter-batches", + "pyarrow", + ], + {}, + ), + # Release reads sf1000 lineitem; sf100 here. Same 1.25 GiB bin pin + # (~one file per bin). P0 trip: wUSS 1.38, tUSS 1.26/max 1.63 (M38/M74). + ( + "write_parquet", + "read_and_consume_benchmark", + [f"{TPCH}/sf{a.write_sf}/lineitem", "--format", "parquet", "--write"], + {"RAY_DATA_PARQUET_BIN_PACKING_BYTES": "1342177280"}, + ), + # APPROXIMATION of read_parquet_autoscaling (wall 1.28): release data is + # the internal imagenet bucket (denied); keeps the 64 MiB bin pin on + # public lineitem so the many-small-tasks geometry survives. + ( + "read_parquet_binned", + "read_and_consume_benchmark", + [ + f"{TPCH}/sf{a.write_sf}/lineitem", + "--format", + "parquet", + "--iter-bundles", + ], + {"RAY_DATA_PARQUET_BIN_PACKING_BYTES": "67108864"}, + ), + ] + # map_groups: the P0 pair (hash/sort col02+14, wall 1.16/1.20, T spilling + # 53.6 vs B 40.8 GB) plus the spill-asymmetry shapes flagged 2026-08-28: + # col08+13+14 (T 31.6 vs B 19.0 GB autoscaling; the M76 OOM-cliff family) + # and the hash_shuffle_v2 variants (M77 sustained 1.63/1.17). The aggregate + # cells are the M47 positive control — the cleanest decoder WIN in release + # (tUSS R 0.56-0.81): a box run where the losses reproduce but this win + # doesn't (or vice versa) says the regime is off, not the reader. + for name, consume, cols, strat in [ + ( + "mapg_hash_col02+col14", + "--map-groups", + ["column02", "column14"], + "hash_shuffle", + ), + ( + "mapg_hashv2_col02+col14", + "--map-groups", + ["column02", "column14"], + "hash_shuffle_v2", + ), + ( + "mapg_sort_col02+col14", + "--map-groups", + ["column02", "column14"], + "sort_shuffle_pull_based", + ), + ( + "mapg_hash_col08+13+14", + "--map-groups", + ["column08", "column13", "column14"], + "hash_shuffle", + ), + ( + "mapg_hashv2_col08+13+14", + "--map-groups", + ["column08", "column13", "column14"], + "hash_shuffle_v2", + ), + ( + "agg_hash_col02+col14", + "--aggregate", + ["column02", "column14"], + "hash_shuffle", + ), + ( + "agg_hash_col08+13+14", + "--aggregate", + ["column08", "column13", "column14"], + "hash_shuffle", + ), + ]: + out.append( + ( + name, + "groupby_benchmark", + ["--sf", a.groupby_sf, consume, "--group-by"] + + cols + + ["--shuffle-strategy", strat], + {}, + ) + ) + # joins_sf100_right_outer wall 1.58 (first occurrence); inner/left/full ride + # along on request for the sustained-wUSS addendum (all three sat 1.27-1.35). + for jt in a.join_types.split(","): + out.append( + ( + f"joins_{jt}", + "join_benchmark", + [ + "--left_dataset", + f"{TPCH}/sf{a.joins_sf}/lineitem", + "--right_dataset", + f"{TPCH}/sf{a.joins_sf}/orders", + "--left_join_keys", + "column00", + "--right_join_keys", + "column0", + "--join_type", + jt, + "--num_partitions", + "50", + ], + {}, + ) + ) + return out + + +def run_cell(name, module, argv, extra_env, reader, a): + tag = f"{name}.{reader}" + env = dict(os.environ) + env["PYTHONPATH"] = DATASET_DIR + os.pathsep + env.get("PYTHONPATH", "") + arm_env(env, reader) + env["RAY_DATA_BENCH_NODE_MEM_MONITOR"] = "1" + env["RAY_DATA_BENCH_NODE_MEM_INTERVAL"] = str(a.monitor_interval) + # Anyscale pins RAY_OVERRIDE_RESOURCES (memory=14.4GiB on this box) and it + # beats ray.init(_memory=...); rewrite the memory field per cell or the + # multi-JoinOperator/aggregator cells deadlock on reservation starvation. + if a.sched_mem_gb and env.get("RAY_OVERRIDE_RESOURCES"): + ovr = json.loads(env["RAY_OVERRIDE_RESOURCES"]) + ovr["memory"] = a.sched_mem_gb << 30 + env["RAY_OVERRIDE_RESOURCES"] = json.dumps(ovr) + env["TEST_OUTPUT_JSON"] = os.path.join(a.outdir, f"{tag}.benchmark.json") + env.update(extra_env) + write_root = "" + if module == "read_and_consume_benchmark" and "--write" in argv: + bucket = os.environ.get("ARROW_RS_S3_BUCKET", "") + if bucket: + write_root = f"{bucket.rstrip('/')}/regression_probe/{tag}" + cmd = [ + sys.executable, + "-c", + SNIPPET, + module, + json.dumps(argv), + write_root, + "1" if a.dry_run else "0", + str(a.sched_mem_gb), + ] + log_path = os.path.join(a.outdir, f"{tag}.log") + if not a.dry_run: + print(f" -> {tag} running (tail -f {log_path})", flush=True) + t0 = time.perf_counter() + timed_out = False + with open(log_path, "w") as fh: + fh.write(f"# {module} {' '.join(argv)} reader={reader}\n") + fh.flush() + # start_new_session: one process group per cell, so a timeout reaps the + # local Ray daemons too instead of leaking RSS into the next cells. + proc = subprocess.Popen( + cmd, env=env, stdout=fh, stderr=subprocess.STDOUT, start_new_session=True + ) + try: + proc.wait(timeout=None if a.dry_run else a.cell_timeout) + except subprocess.TimeoutExpired: + timed_out = True + os.killpg(proc.pid, signal.SIGTERM) + try: + proc.wait(timeout=30) + except subprocess.TimeoutExpired: + pass + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + proc.wait() + with open(log_path) as fh: + out = fh.read() + if timed_out: + print( + f" !! {tag} TIMEOUT after {a.cell_timeout}s (see {tag}.log)", flush=True + ) + return {"timeout_s": a.cell_timeout} + line = next((ln for ln in out.splitlines() if ln.startswith("CELL_JSON ")), None) + if line is None: + print(f" !! {tag} FAILED rc={proc.returncode} (see {tag}.log)", flush=True) + print(" " + out.strip()[-400:], flush=True) + return None + rec = json.loads(line[len("CELL_JSON ") :]) + # Fold in the benchmark's own metrics (per-task dists, node-mem monitor). + try: + with open(env["TEST_OUTPUT_JSON"]) as fh: + bench = json.load(fh) + rec["bench"] = next(iter(bench.values())) + except Exception: + rec["bench"] = {} + rec["wall_incl_startup_s"] = round(time.perf_counter() - t0, 1) + print( + f" {tag:<34} wall={rec.get('wall_s')}s spill={rec.get('spilled_gb')}GB", + flush=True, + ) + return rec + + +def _g(rec, *path): + cur = rec or {} + for p in path: + cur = cur.get(p) if isinstance(cur, dict) else None + if cur is None: + return None + return cur + + +def _r(rs, pa): + return f"{rs / pa:.2f}" if rs and pa else "—" + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--outdir", required=True) + p.add_argument("--repeat", type=int, default=1) + p.add_argument("--only", default="", help="comma-separated cell-name filter") + p.add_argument( + "--write-sf", default="100", help="sf for write/read cells (release: 1000)" + ) + p.add_argument( + "--groupby-sf", default="10", help="sf for map_groups (release: 100)" + ) + p.add_argument("--joins-sf", default="10", help="sf for joins (release: 100)") + p.add_argument( + "--join-types", + default="right_outer", + help="right_outer, or all four for the sustained-wUSS addendum", + ) + p.add_argument( + "--cell-timeout", + type=int, + default=int(os.environ.get("PROBE_CELL_TIMEOUT", "1800")), + help="kill a cell's whole process group after this many seconds", + ) + p.add_argument( + "--sched-mem-gb", + type=int, + default=int(os.environ.get("PROBE_SCHED_MEM_GB", "64")), + help="scheduling-only memory resource for the local Ray instance " + "(0 = stock; stock deadlocks multi-join/aggregator cells on one box)", + ) + p.add_argument( + "--dry-run", + action="store_true", + help="import-and-exit per cell: validates plumbing offline", + ) + p.add_argument( + "--arms", + default="pa,rs,rseos", + help=f"comma list from {sorted(ARMS)}; pa is the denominator", + ) + p.add_argument( + "--cpus", + default="", + help="comma list of local num_cpus values; each cell runs once per " + "value as @cpu (empty = the box's core count)", + ) + p.add_argument( + "--monitor-interval", + type=float, + default=float(os.environ.get("RAY_DATA_BENCH_NODE_MEM_INTERVAL", "1.0")), + help="node-memory sampler period in seconds (release: 1.0; 0.1 = 10 Hz)", + ) + a = p.parse_args() + arms = [s.strip() for s in a.arms.split(",") if s.strip()] + for arm in arms: + arm_env({}, arm) # validate names up front + os.makedirs(a.outdir, exist_ok=True) + + todo = cells(a) + if a.only: + keep = {s.strip() for s in a.only.split(",")} + todo = [c for c in todo if any(k in c[0] for k in keep)] + if a.cpus: + todo = [ + ( + f"{name}@{c.strip()}cpu", + module, + argv, + {**env, "PROBE_NUM_CPUS": c.strip()}, + ) + for c in a.cpus.split(",") + for name, module, argv, env in todo + ] + + results = {} + for name, module, argv, extra_env in todo: + for reader in arms: + runs = [ + run_cell(name, module, argv, extra_env, reader, a) + for _ in range(a.repeat) + ] + good = sorted( + (r for r in runs if r and r.get("wall_s")), key=lambda r: r["wall_s"] + ) + results[f"{name}.{reader}"] = ( + good[len(good) // 2] if good else (runs[0] if runs else None) + ) + + with open(os.path.join(a.outdir, "summary.json"), "w") as fh: + json.dump(results, fh, indent=2) + if a.dry_run: + print("\ndry run OK — all release script modules import") + return + + # Same columns as arrow_rs_docs/2026-08-27.md §12: the numbers the P0 cut + # was made on, computed from each cell's own TEST_OUTPUT_JSON. + print("\n========== RELEASE-REGRESSION PROBE (R = arrow_rs/pyarrow) ==========") + hdr = ( + f"{'cell [arm]':<34} {'wall R':>7} {'tUSS p50 R':>10} {'tUSS max R':>10} " + f"{'wUSS pk R':>9} {'wUSS sust R':>11} {'pkbatch R':>9} {'n tasks':>8} " + f"{'trim p50 s':>10} {'spill pa/arm':>12}" + ) + print(hdr) + for (name, _, _, _), arm in ((c, m) for c in todo for m in arms if m != "pa"): + pa = results.get(f"{name}.pa") or {} + rs = results.get(f"{name}.{arm}") or {} + bp, br = pa.get("bench", {}), rs.get("bench", {}) + trim = _g(br, "read_trim_wall_s_per_task_dist", "p50") + row = ( + f"{name + ' [' + arm + ']':<34} " + f"{_r(rs.get('wall_s'), pa.get('wall_s')):>7} " + f"{_r(_g(br, 'read_max_uss_per_task_dist', 'p50'), _g(bp, 'read_max_uss_per_task_dist', 'p50')):>10} " + f"{_r(_g(br, 'read_max_uss_per_task_dist', 'max'), _g(bp, 'read_max_uss_per_task_dist', 'max')):>10} " + f"{_r(_g(br, 'node_mem_peak_worker_uss_gb'), _g(bp, 'node_mem_peak_worker_uss_gb')):>9} " + f"{_r(_g(br, 'node_mem_p50_worker_uss_gb'), _g(bp, 'node_mem_p50_worker_uss_gb')):>11} " + f"{_r(_g(br, 'read_peak_batch_bytes_per_task_dist', 'p50'), _g(bp, 'read_peak_batch_bytes_per_task_dist', 'p50')):>9} " + f"{_g(br, 'read_max_uss_per_task_dist', 'num_samples') or '—':>8} " + f"{(f'{trim:.3f}' if isinstance(trim, (int, float)) else '—'):>10} " + f"{str(pa.get('spilled_gb')) + '/' + str(rs.get('spilled_gb')):>12}" + ) + print(row) + print( + "\nRead it as: a cell that reproduces its release ratio here is debuggable" + "\non this box; a clean cell pushes that P0 into the release regime" + "\n(autoscaling pool dynamics / allocator churn floor — M35, items 18/19)." + ) + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/replication_matrix.py b/release/nightly_tests/dataset/arrow_rs_probe/replication_matrix.py new file mode 100644 index 000000000000..a3cb9b430740 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/replication_matrix.py @@ -0,0 +1,708 @@ +#!/usr/bin/env python3 +"""Replicate the 2026-08-12 release A/B's *trusted* signals on one Linux box. + +TODO item 1ab phase 1 (arrow_rs_docs/TODO.md): the multi-node run's wall / decode +task-seconds are trusted, its memory data mostly is not — so before touching the +release harness, each trusted loss (and the headline win) is replicated locally with +the probe, whose per-task USS (`read_avg_max_uss_gb`) and 50 Hz worker sampler ARE +trustworthy on Linux. Fix what reproduces, keep what's right. + +Stages (pick with --skip / --only): + + tensors R1/item 1y — wide_schema tensors decoded natively 5.59x slower (wall + 1.77x). Fixture: 5000 fixed_size_list columns (lookalike; + the S3 original is unreadable to us). Runs both readers at + concurrency=1 (pure decode-speed diagnostic) AND fanned out (adds the + thread-pool asymmetry: base = unbounded fragment threads, arrow-rs = + min(4, fragments) — the "PyArrow scales better?" hypothesis). + DID NOT reproduce (T19) — kept as the negative control. The stage + that does reproduce 1y is tensorscp. + tensorscp The 1y reproducer (T22/T23): same 5000 columns but written with + cloudpickle tensor extension metadata (Ray 2.49-2.54 format, what + the release dataset actually contains). The crate can't parse the + non-UTF8 embedded schema, so the reader decodes storage types and + realigns storage->extension per batch — the path where the loss + lives (macOS pre-fix: read wall 5.4x; post Table.cast fix: 1.25x). + Cells set RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA=1, exactly + like the release yaml. Needs the tensors_cp fixture. + binsweep R2/item 10 — sweep RAY_DATA_PARQUET_BIN_PACKING_BYTES across + {1 RG, 4 RGs, 1 file, 5 files, 10 files} x both readers. The 5x/10x + multi-file bins are the first cells to exercise C9 mechanism (i) + (N sub-fragments on the base's unbounded pool). Plus a PyArrow + pre_buffer=off arm at {1 file, 10 files} for mechanism (ii) + (RAY_DATA_PARQUET_PRE_BUFFER=0 — the knob exists only for this). + Predictions to falsify are in TODO item 10. + binbound R2b — THE BOUND CHECK (user ask 2026-08-12): one bin is one read + task, its whole decoded output lives in that worker process, so the + arrow-rs guarantee is "per-task USS is bounded by the bin budget" — + the property PyArrow lacks (it buffers a whole decoded row group per + fragment, plus a pre_buffer'd compressed span). Same bin grid as + binsweep but with --task-concurrency 1 (one bin resident per process, + so per-task USS is attributable to ONE bin) and --mem-poll-s 0.05 + (at the 1 Hz default a short read task gets one sample or none, which + silently flattens the very number the verdict rests on). Verdict is a + least-squares slope of per-task USS vs DECODED bytes per task: + slope <~0.3 = flat (bounded far below the bin), <~1 = bounded by the + bin, >1 = unbounded => retention or leak. Denominator is decoded + bytes, not the knob: the knob budgets Parquet total_uncompressed_size + (pages decompressed but still dictionary/RLE-ENCODED), so + decoded/knob is an expansion factor >= 1 that the table prints + instead of assuming. + write R3/item 1aa — write_parquet showed per-task USS 1.23x (trusted + instrument) at a wall WIN 0.83x. read bin_sweep fixture -> + write_parquet, both readers; stats come from the materialized + write plan (no teardown race). + fatcol R4/item 1o — known wall ~1.2x on the fat-column shape; rides along + for a fresh number on this base. + oom R5/item 10's oom axis — the failure-mode demonstration this project + exists for. Same box, same memory ceiling, sweep the bin size: the + per-task fits from binbound (M28: pyarrow ~0.58GB + 1.5x decoded, + arrow-rs ~0.27GB + 0.18x decoded) predict PyArrow's read worker + crosses the ceiling once the bin is big enough while arrow-rs never + does. The ceiling is Ray's OWN memory monitor, not a cgroup: each + cell computes RAY_memory_usage_threshold = (used-at-launch + budget) + / total, sets RAY_task_oom_retries=0 (default -1 retries the killed + task forever, ray_config_def.h:145), and a PyArrow kill shows up as + the exact OutOfMemoryError a user would see. Budget defaults to + 0.5 x the whole fixture's DECODED bytes (footer expansion x packer + bytes) — sized so arrow-rs fits at every bin and PyArrow cannot fit + the biggest ones; --oom-budget-gb overrides. A cell that dies OOM + is the stage's DATA, not a harness failure. + +Usage (Linux box, venv active; fixtures first): + + python gen_local_fixtures.py --root ~/arrow_rs_repl_fixtures \ + --shapes bin_sweep,tensors_wide,tensors_cp,fat_col + python replication_matrix.py --fixture-root ~/arrow_rs_repl_fixtures + python replication_matrix.py --fixture-root ... --only binsweep --repeat 3 + +Each cell runs read_probe.py in its own process; logs + summary.json land in +--outdir (default ./replication_runs/). Ratios printed are arrow_rs/pyarrow, +>1 = arrow-rs worse. +""" +import argparse +import json +import os +import re +import subprocess +import sys +import time + +from run_matrix import PROBE, _num, median_cell, ratio + +MiB = 1024 * 1024 + + +def load_manifest(fixture_root): + path = os.path.join(os.path.expanduser(fixture_root), "manifest.json") + with open(path) as fh: + return json.load(fh) + + +def footer_geometry(path): + """Measure the bin packer's OWN accounting unit from the fixture's footers. + + Three things this pins down that no nominal fixture size can (all verified on + the tree, 2026-08-12): + + 1. On a no-projection read the packer prices a row group at + ``row_group.total_byte_size`` (``listing/footer_reader.py:148-151``), NOT at + the summed per-column ``total_uncompressed_size``. Those are equal for a + well-behaved writer but ``total_byte_size`` can carry the *compressed* size + (apache/arrow#48138 — which is why the reader's own batch-size estimator + refuses that accessor, ``readers/parquet_file_reader.py:148-155``). We report + the ratio so a fixture that trips the bug is visible instead of silently + shrinking every bin. + 2. Decoded Arrow bytes are ``expansion`` x the packer's number, because the + footer counts pages that are decompressed but still dictionary/RLE-*encoded*. + Ray itself assumes 5x here (``PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT``, + ``datasource/parquet_datasource.py:109``). So "USS bounded by the bin" can + only ever mean "bounded by ``expansion`` x bin"; this makes the factor a + measured column rather than an unstated assumption. + 3. The grid labels ("1file") then mean a real file's worth of the packer's + bytes, so bin sizes and task counts agree. + """ + import glob as _glob + + import pyarrow.parquet as pq + + files = sorted(_glob.glob(os.path.join(path, "**", "*.parquet"), recursive=True)) + if not files: + raise SystemExit(f"no parquet files under {path}") + tbs_total = unc_total = rg_count = 0 + for f in files: + md = pq.ParquetFile(f).metadata + for i in range(md.num_row_groups): + rg = md.row_group(i) + tbs_total += rg.total_byte_size + unc_total += sum( + rg.column(j).total_uncompressed_size for j in range(rg.num_columns) + ) + rg_count += 1 + # Expansion from ONE file (cheap) — decoded Arrow bytes / the packer's bytes. + md0 = pq.ParquetFile(files[0]).metadata + f0_tbs = sum(md0.row_group(i).total_byte_size for i in range(md0.num_row_groups)) + decoded0 = pq.read_table(files[0]).nbytes + return { + "files": len(files), + "row_groups": rg_count, + "packer_bytes_total": tbs_total, + "rg_bytes": tbs_total // rg_count, + "file_bytes": tbs_total // len(files), + "tbs_over_uncompressed": round(tbs_total / unc_total, 3) if unc_total else None, + "expansion": round(decoded0 / f0_tbs, 3) if f0_tbs else None, + } + + +def bin_grid(geom): + """Bin budgets in the packer's own units: 1 and 4 row groups, then 1/2/5/10 files. + + The 5x/10x-a-file cells are the "much bigger bin" ask. Any cell that already + swallows the whole fixture is collapsed to a single ``all`` cell — two cells + that both pack everything into one bin measure the same thing twice. + """ + rg, fl = geom["rg_bytes"], geom["file_bytes"] + total = geom["packer_bytes_total"] + out, saturated = [], False + for name, size in [ + ("1rg", rg), + ("4rg", 4 * rg), + ("1file", fl), + ("2file", 2 * fl), + ("5file", 5 * fl), + ("10file", 10 * fl), + ]: + if size >= total: + if saturated: + continue + saturated = True + name = f"all({name})" + out.append((name, size)) + return out + + +def binsweep_grid(entry): + """Fallback grid from the fixture manifest's *nominal* sizes, used only if the + footers can't be read. Nominal bytes are the generator's row-width arithmetic, + which ran 1.6x above the footers on the smoke fixture — so labels drift and task + counts won't match the names. Prefer ``bin_grid(footer_geometry(path))``.""" + rg = entry["uncompressed_bytes"] // (entry["files"] * entry["rgs_per_file"]) + fl = entry["uncompressed_bytes"] // entry["files"] + return [ + ("1rg", rg), + ("4rg", 4 * rg), + ("1file", fl), + ("5file", 5 * fl), + ("10file", 10 * fl), + ] + + +# Ray's memory monitor prints this when it kills a task; the driver then raises +# ray.exceptions.OutOfMemoryError. Either string in the output = the kill we asked for. +_OOM_PAT = re.compile( + r"OutOfMemoryError|killed due to the node running low on memory", re.IGNORECASE +) + + +def run_oom_cell(logdir, name, path, reader, bin_bytes, threshold): + """One oom-axis cell: run read_probe under a memory-monitor ceiling; classify. + + Unlike run_cell, a non-zero exit here can be the expected result — the stage + exists to watch PyArrow's arm die — so the outcome (ok / oom / error) is data + and gets returned instead of being reported as a probe failure. + """ + cmd = [ + sys.executable, + PROBE, + "--path", + path, + "--reader", + reader, + # One bin resident per process (binbound's attribution logic), and the + # 20 Hz poll so the surviving arm's USS number is trustworthy. + "--task-concurrency", + "1", + "--mem-poll-s", + "0.05", + ] + env = dict(os.environ) + env.update( + { + "RAY_DATA_PARQUET_BIN_PACKING_BYTES": str(bin_bytes), + "RAY_memory_usage_threshold": f"{threshold:.4f}", + # Default -1 retries a memory-killed task forever (ray_config_def.h:145) + # — the pyarrow arm would loop kill/retry instead of failing. + "RAY_task_oom_retries": "0", + # Default 250 ms; a fast decode spike can blow past the threshold + # between ticks and take the box with it. 100 ms narrows that window. + "RAY_memory_monitor_refresh_ms": "100", + } + ) + t0 = time.perf_counter() + proc = subprocess.run(cmd, capture_output=True, text=True, env=env) + dur = time.perf_counter() - t0 + + logpath = os.path.join(logdir, f"{name}.log") + with open(logpath, "w") as fh: + fh.write(f"# cmd: {' '.join(cmd)}\n") + fh.write( + f"# threshold={threshold:.4f} bin={bin_bytes} rc={proc.returncode} " + f"wall_including_startup_s={dur:.1f}\n" + ) + fh.write("# ---- STDOUT ----\n") + fh.write(proc.stdout) + fh.write("\n# ---- STDERR ----\n") + fh.write(proc.stderr) + + res = {} + in_result = False + for line in proc.stdout.splitlines(): + if "=== RESULT ===" in line: + in_result = True + continue + if in_result and ":" in line: + k, v = line.strip().split(":", 1) + res[k.strip()] = v.strip() + + if proc.returncode == 0 and res: + outcome = "ok" + elif _OOM_PAT.search(proc.stdout + proc.stderr): + outcome = "oom" + else: + outcome = "error" # crashed for some other reason — read the log + res["outcome"] = outcome + res["rc"] = proc.returncode + print( + f" {name:<34} {outcome.upper():<5} wall={res.get('wall_s', round(dur, 1))} " + f"uss={res.get('read_max_uss_gb') or res.get('read_avg_max_uss_gb')}" + + ("" if outcome != "error" else f" !! see {logpath}"), + flush=True, + ) + return res + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--fixture-root", required=True) + p.add_argument("--outdir", default=None) + p.add_argument( + "--repeat", + type=int, + default=1, + help="measured runs per cell; each metric is reported as its own median", + ) + p.add_argument( + "--warmup", + type=int, + default=1, + help=( + "discarded runs per cell before the measured ones. The first run of a " + "cell is reproducibly slow (cold page cache / .so); 0 to disable." + ), + ) + p.add_argument( + "--skip", + default="", + help="comma list: tensors,tensorscp,binsweep,binbound,write,fatcol,oom", + ) + p.add_argument("--only", default="", help="comma list: run only these stages") + p.add_argument( + "--oom-budget-gb", + type=float, + default=None, + help=( + "oom stage: memory the read job may use above the box's at-launch " + "usage before Ray's monitor kills its biggest task. Default: " + "0.5 x the fixture's decoded bytes (footer expansion x packer bytes)." + ), + ) + args = p.parse_args() + + manifest = load_manifest(args.fixture_root) + skip = {s.strip() for s in args.skip.split(",") if s.strip()} + only = {s.strip() for s in args.only.split(",") if s.strip()} + + def enabled(stage): + if only: + return stage in only + return stage not in skip + + def fixture_path(shape): + if shape not in manifest: + raise SystemExit( + f"fixture '{shape}' missing from {args.fixture_root}/manifest.json — " + f"run gen_local_fixtures.py --shapes {shape} first" + ) + return manifest[shape]["path"] + + ts = time.strftime("%Y%m%d_%H%M%S") + outdir = args.outdir or os.path.join( + os.path.dirname(os.path.abspath(__file__)), "replication_runs", ts + ) + os.makedirs(outdir, exist_ok=True) + print(f"logs -> {outdir}\n", flush=True) + + rows = {} + + def cell(name, **kw): + rows[name] = median_cell(outdir, name, args.repeat, args.warmup, **kw) + + # -------- [tensors] R1 / item 1y -------- + if enabled("tensors"): + print( + "=== [tensors] 5000 fixed_size_list cols — decode 5.59x repro ===", + flush=True, + ) + path = fixture_path("tensors_wide") + for reader in ("pyarrow", "arrow_rs"): + cell( + f"tensors.c1.{reader}", + path=path, + reader=reader, + concurrency=1, + columns=None, + extra_env={}, + ) + cell( + f"tensors.fan.{reader}", + path=path, + reader=reader, + concurrency=None, + columns=None, + extra_env={}, + ) + + # -------- [tensorscp] the 1y reproducer (T22/T23) -------- + if enabled("tensorscp"): + print( + "=== [tensorscp] 5000 cloudpickle-metadata tensor cols — the shape " + "that actually reproduces 1y ===", + flush=True, + ) + path = fixture_path("tensors_cp") + # Same opt-in the release yaml carries; without it BOTH arms refuse the + # file at plan time (pyarrow's dataset factory raises deserializing the + # tensor metadata), so it is a fixture prerequisite, not a treatment. + cp_env = {"RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA": "1"} + for reader in ("pyarrow", "arrow_rs"): + cell( + f"tensorscp.c1.{reader}", + path=path, + reader=reader, + concurrency=1, + columns=None, + extra_env=dict(cp_env), + ) + cell( + f"tensorscp.fan.{reader}", + path=path, + reader=reader, + concurrency=None, + columns=None, + extra_env=dict(cp_env), + ) + + # -------- [binsweep] R2 / item 10 -------- + grid = [] + geom = {} + if enabled("binsweep") or enabled("binbound") or enabled("oom"): + try: + geom = footer_geometry(fixture_path("bin_sweep")) + grid = bin_grid(geom) + print( + f"bin_sweep footers: {geom['files']} files / {geom['row_groups']} row " + f"groups, packer prices a row group at {geom['rg_bytes'] // MiB}MiB " + f"(total_byte_size/uncompressed={geom['tbs_over_uncompressed']}" + + ( + " <-- NOT 1.0: apache/arrow#48138, bins are priced in COMPRESSED " + "bytes on this fixture" + if geom["tbs_over_uncompressed"] + and abs(geom["tbs_over_uncompressed"] - 1.0) > 0.02 + else "" + ) + + f"), decoded/packer expansion={geom['expansion']}x " + f"(Ray's own default assumption is 5x)\n", + flush=True, + ) + except Exception as e: # noqa: BLE001 - fall back, don't lose the run + print(f"!! footer_geometry failed ({e!r}); using nominal manifest grid") + grid = binsweep_grid(manifest["bin_sweep"]) + if enabled("binsweep"): + path = fixture_path("bin_sweep") + print( + "=== [binsweep] bins " + + ", ".join(f"{n}={b // MiB}MiB" for n, b in grid) + + " ===", + flush=True, + ) + for bin_name, bin_bytes in grid: + for reader in ("pyarrow", "arrow_rs"): + cell( + f"binsweep.{bin_name}.{reader}", + path=path, + reader=reader, + concurrency=None, + columns=None, + extra_env={"RAY_DATA_PARQUET_BIN_PACKING_BYTES": str(bin_bytes)}, + ) + # pre_buffer=off arm, PyArrow only (C9 mechanism (ii) attribution). + for bin_name, bin_bytes in grid: + if bin_name not in ("1file", "10file"): + continue + cell( + f"binsweep.{bin_name}.pyarrow.nopb", + path=path, + reader="pyarrow", + concurrency=None, + columns=None, + extra_env={ + "RAY_DATA_PARQUET_BIN_PACKING_BYTES": str(bin_bytes), + "RAY_DATA_PARQUET_PRE_BUFFER": "0", + }, + ) + + # -------- [binbound] R2b — is per-task USS bounded by the bin budget? -------- + if enabled("binbound"): + path = fixture_path("bin_sweep") + print( + "=== [binbound] one bin per process (task-concurrency 1, 20 Hz USS) ===", + flush=True, + ) + for bin_name, bin_bytes in grid: + for reader in ("pyarrow", "arrow_rs"): + cell( + f"binbound.{bin_name}.{reader}", + path=path, + reader=reader, + concurrency=None, + columns=None, + extra_env={"RAY_DATA_PARQUET_BIN_PACKING_BYTES": str(bin_bytes)}, + extra_args=[ + "--task-concurrency", + "1", + "--mem-poll-s", + "0.05", + ], + ) + + # -------- [write] R3 / item 1aa -------- + if enabled("write"): + print("=== [write] read -> write_parquet (fused) ===", flush=True) + path = fixture_path("bin_sweep") + for reader in ("pyarrow", "arrow_rs"): + cell( + f"write.{reader}", + path=path, + reader=reader, + concurrency=None, + columns=None, + extra_env={}, + extra_args=["--consume", "write_parquet"], + ) + + # -------- [fatcol] R4 / item 1o -------- + if enabled("fatcol"): + print("=== [fatcol] fat binary column — wall ~1.2x recheck ===", flush=True) + path = fixture_path("fat_col") + for reader in ("pyarrow", "arrow_rs"): + cell( + f"fatcol.{reader}", + path=path, + reader=reader, + concurrency=1, + columns=None, + extra_env={}, + ) + + # -------- [oom] R5 / item 10's oom axis -------- + oom_cfg = {} + if enabled("oom"): + import psutil + + path = fixture_path("bin_sweep") + decoded_gb = None + if geom.get("expansion") and geom.get("packer_bytes_total"): + decoded_gb = geom["expansion"] * geom["packer_bytes_total"] / (1024**3) + budget_gb = args.oom_budget_gb or ( + max(2.0, 0.5 * decoded_gb) if decoded_gb else 4.0 + ) + vm = psutil.virtual_memory() + total_gb = vm.total / (1024**3) + # Same "used" the monitor computes: total minus reclaimable-available. + baseline_gb = (vm.total - vm.available) / (1024**3) + threshold = (baseline_gb + budget_gb) / total_gb + oom_cfg = { + "budget_gb": round(budget_gb, 2), + "baseline_used_gb": round(baseline_gb, 2), + "box_total_gb": round(total_gb, 2), + "threshold": round(threshold, 4), + "fixture_decoded_gb": round(decoded_gb, 2) if decoded_gb else None, + } + rows["oom.config"] = oom_cfg + print( + f"=== [oom] ceiling = {baseline_gb:.1f}GB used-at-launch + " + f"{budget_gb:.1f}GB budget => RAY_memory_usage_threshold=" + f"{threshold:.3f} of {total_gb:.0f}GB ===", + flush=True, + ) + if threshold >= 0.95: + print( + " !! computed threshold >= the 0.95 default — the box is too " + "small (or the fixture too big) for the budget to be the binding " + "constraint; expect kills at the default instead.", + flush=True, + ) + for bin_name, bin_bytes in grid: + for reader in ("pyarrow", "arrow_rs"): + rows[f"oom.{bin_name}.{reader}"] = run_oom_cell( + outdir, + f"oom.{bin_name}.{reader}", + path=path, + reader=reader, + bin_bytes=bin_bytes, + threshold=threshold, + ) + + # -------- summary -------- + print("\n============ SUMMARY (R = arrow_rs / pyarrow, >1 worse) ============") + + def pair_line( + prefix, metrics=("wall_s", "read_wall_s", "read_avg_max_uss_gb", "peak_uss_gb") + ): + pa_r = rows.get(f"{prefix}.pyarrow", {}) + ar_r = rows.get(f"{prefix}.arrow_rs", {}) + parts = [] + for m in metrics: + a, b = _num(ar_r, m), _num(pa_r, m) + r = ratio(a, b) + if r is not None: + parts.append(f"{m}: {b} -> {a} R={r}") + print(f" {prefix:<22} " + (" | ".join(parts) if parts else "(no data)")) + + if enabled("tensors"): + print("[tensors] c1 wall R is the decode-speed verdict; fan adds pool-width") + pair_line("tensors.c1") + pair_line("tensors.fan") + if enabled("tensorscp"): + print( + "[tensorscp] the 1y reproducer (cloudpickle metadata -> skip+realign " + "path). Pre-fix macOS read wall R was 5.4; post Table.cast fix 1.25" + ) + pair_line("tensorscp.c1") + pair_line("tensorscp.fan") + if enabled("binsweep"): + print("[binsweep] prediction: pyarrow USS rises with bin, arrow_rs flat") + for bin_name, bin_bytes in grid: + pair_line(f"binsweep.{bin_name}") + for bin_name in ("1file", "10file"): + nopb = rows.get(f"binsweep.{bin_name}.pyarrow.nopb", {}) + base = rows.get(f"binsweep.{bin_name}.pyarrow", {}) + v, b = _num(nopb, "read_avg_max_uss_gb"), _num(base, "read_avg_max_uss_gb") + if v is not None or b is not None: + print( + f" binsweep.{bin_name}.pyarrow pre_buffer off/on " + f"uss={v}/{b} ratio={ratio(v, b)} " + f"wall={_num(nopb, 'wall_s')}/{_num(base, 'wall_s')}" + ) + if enabled("binbound"): + print( + "\n[binbound] THE BOUND CHECK — per-task USS vs the bin it decoded.\n" + " bin = RAY_DATA_PARQUET_BIN_PACKING_BYTES, in the packer's units\n" + " (row_group.total_byte_size: decompressed but still ENCODED)\n" + " dec/task = decoded Arrow bytes per read task (the real bin size);\n" + " dec/bin is the encoding expansion — measured " + f"{geom.get('expansion')}x on this\n" + " fixture, and Ray's own planner assumes 5x. The knob is a\n" + " proxy for decoded bytes, never an upper bound on them.\n" + " uss/dec = per-task peak USS over decoded bytes. This is the number:\n" + " it includes the fixed ~0.2-0.4 GB python+ray+pyarrow floor,\n" + " so it is large at tiny bins and must FALL toward a constant.\n" + " mx/av = worst task / average task USS. ~1 = every task costs the\n" + " same; rising = the worker retains across tasks (allocator\n" + " retention or a leak), which no bin cap can bound." + ) + for reader in ("pyarrow", "arrow_rs"): + print(f" --- {reader} ---") + pts = [] + for bin_name, bin_bytes in grid: + r = rows.get(f"binbound.{bin_name}.{reader}", {}) + dec = _num(r, "read_bytes_per_task_gb") + uss = _num(r, "read_max_uss_gb") or _num(r, "read_avg_max_uss_gb") + bin_gb = bin_bytes / (1024**3) + if dec is not None and uss is not None: + pts.append((dec, uss)) + print( + f" {bin_name:<11} bin={bin_bytes // MiB:>6}MiB " + f"tasks={r.get('read_num_tasks')} " + f"dec/task={dec}GB dec/bin={ratio(dec, bin_gb)} " + f"uss_max={uss} uss/dec={ratio(uss, dec)} " + f"mx/av={r.get('uss_max_over_avg')} wall={_num(r, 'wall_s')}" + ) + # Least squares on USS = a + b*decoded_bytes. b is the verdict: how many + # bytes of private memory each extra decoded byte in the bin costs. + if len(pts) >= 3: + n = len(pts) + mx = sum(p[0] for p in pts) / n + my = sum(p[1] for p in pts) / n + den = sum((p[0] - mx) ** 2 for p in pts) + if den > 0: + b = sum((p[0] - mx) * (p[1] - my) for p in pts) / den + a = my - b * mx + if b <= 0.3: + verdict = "FLAT — bounded well below the bin" + elif b <= 1.1: + verdict = "BOUNDED by the bin (slope ~1)" + else: + verdict = ( + "UNBOUNDED — grows FASTER than the bin; " + "suspect retention/leak, not just buffering" + ) + print( + f" fit: uss ~= {round(a, 3)}GB + {round(b, 3)} x decoded " + f"-> {verdict}" + ) + else: + print(" fit: not enough cells with USS (Linux only) to fit a slope") + print( + " arrow-rs must be FLAT or BOUNDED; anything else is the leak this stage\n" + " exists to catch. PyArrow is expected to slope up (whole decoded row\n" + " group per fragment x fragments in the bin, + the pre_buffer span)." + ) + if enabled("write"): + print("[write] release said USS R=1.23 at wall R=0.83") + pair_line("write") + if enabled("fatcol"): + print("[fatcol] release-adjacent wall ~1.2x") + pair_line("fatcol") + if enabled("oom"): + print( + "[oom] one ceiling, sweep the bin — where is each reader's cliff?\n" + f" ceiling: {oom_cfg.get('baseline_used_gb')}GB used-at-launch + " + f"{oom_cfg.get('budget_gb')}GB budget (threshold " + f"{oom_cfg.get('threshold')} of {oom_cfg.get('box_total_gb')}GB total)" + ) + for reader in ("pyarrow", "arrow_rs"): + cells, survived = [], [] + for bin_name, _ in grid: + r = rows.get(f"oom.{bin_name}.{reader}", {}) + oc = r.get("outcome", "?") + cells.append(f"{bin_name}:{oc}") + if oc == "ok": + survived.append(bin_name) + print( + f" {reader:<9} " + + " ".join(cells) + + f" biggest surviving bin: {survived[-1] if survived else 'NONE'}" + ) + print( + " prediction (M28 fits): pyarrow dies once 0.58 + 1.5 x decoded(bin) GB\n" + " exceeds the budget; arrow-rs (0.27 + 0.18 x decoded) survives every bin." + ) + + with open(os.path.join(outdir, "summary.json"), "w") as fh: + json.dump(rows, fh, indent=2) + print(f"\nfull JSON + per-cell logs in {outdir}") + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/run_all.sh b/release/nightly_tests/dataset/arrow_rs_probe/run_all.sh new file mode 100755 index 000000000000..69a94c635817 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/run_all.sh @@ -0,0 +1,151 @@ +#!/usr/bin/env bash +# --------------------------------------------------------------------------- +# THE one command for the Linux box: every arrow-rs benchmark stage, in the +# order the evidence chain needs them, into ONE timestamped run directory. +# +# bash release/nightly_tests/dataset/arrow_rs_probe/run_all.sh +# +# Stages (subset via STAGES=mechanism,release — default all): +# setup setup.sh: venv + commit-matched wheel + source symlink + crate +# (skipped when env.sh already imports; FORCE_SETUP=1 after any +# git pull that touches the crate) +# fixtures gen_local_fixtures.py, ALL shapes (incl. the new expansion +# sweep tensors_lo/tensors_dict/tensors_hi) +# mechanism batch_ablation.py — shapes x request-policies x budgets, +# standalone. Decision variables = the four gates (G1 overshoot, +# G2 R_rss, G3 R_wall, G4 rows parity). This is the gate any +# batch-sizing code change must pass first. +# release loss_triage.py — the same loss shapes IN RAY, local + S3 +# (per-task USS at 20 Hz, decoder dists), both readers + arena2 +# arm. S3 legs run iff ARROW_RS_S3_BUCKET is set (fixtures are +# synced up automatically). This is the release-metric gate. +# tpch tpch_probe.py — the two suspect RELEASE TPC-H queries (q9 = +# the T-only spiller T27, q20 = the hash_shuffle_v2-only wall +# loss M46) run via the release scripts themselves at --sf 10, +# matrix strategy x reader, wall + spill per cell. Needs AWS +# creds for s3://ray-benchmark-data (public bucket). +# soak soak_probe.py — long-lived session, O(100) tasks/worker, +# idle-USS floor per round, arms pa/rs/rs_arena2/rs_trim/ +# rs_jemalloc. This is the retention gate (M37/M38/M44). +# tensors tensors_nbytes_probe.py — cheap M39 representation check. +# +# Knobs (env): +# STAGES=... subset of setup,fixtures,mechanism,release,tpch,soak,tensors +# FIXTURES_ROOT= default ~/arrow_rs_repl_fixtures +# FIXTURE_SCALE=1.0 0.25 for a smoke run +# ARROW_RS_S3_BUCKET=s3://... enables the release stage's S3 legs +# BUDGETS=32 mechanism budget sweep, e.g. 16,32,128 +# ABLATION_SHAPES=... mechanism shapes (default: batch_ablation.py's 10) +# TRIAGE_SHAPES=auto,write,tensorscp,tensorsdict,agg (agg = the aggregate-win positive control) +# TPCH_SF=10 TPCH_QUERIES=tpch_q9,tpch_q20 TPCH_STRATEGIES=... +# SOAK_SHAPES=auto,write ARMS=pa,rs,rs_arena2,rs_trim,rs_jemalloc +# REPEAT=3 WARMUP=1 WORKERS=4 +# FORCE_SETUP=1 re-run setup.sh even if the env imports +# +# Results: everything under arrow_rs_probe/suite_runs// — +# ablation.json (+ gate verdict on stdout), loss_triage/summary.json, +# soak/summary.json, tensors_nbytes.log, and stage logs. Each stage also +# prints its own R-table; the final index lists every artifact. +# --------------------------------------------------------------------------- +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH" +STAGES="${STAGES:-setup,fixtures,mechanism,release,tpch,soak,tensors}" +FIXTURES_ROOT="${FIXTURES_ROOT:-$HOME/arrow_rs_repl_fixtures}" +FIXTURE_SCALE="${FIXTURE_SCALE:-1.0}" +BUDGETS="${BUDGETS:-32}" +TRIAGE_SHAPES="${TRIAGE_SHAPES:-auto,write,tensorscp,tensorsdict,agg}" +SOAK_SHAPES="${SOAK_SHAPES:-auto,write}" +ARMS="${ARMS:-pa,rs,rs_arena2,rs_trim,rs_jemalloc}" +TPCH_SF="${TPCH_SF:-10}" +REPEAT="${REPEAT:-3}" +WARMUP="${WARMUP:-1}" +WORKERS="${WORKERS:-4}" +RUN_DIR="$SCRIPT_DIR/suite_runs/$(date +%Y%m%d_%H%M%S)" +mkdir -p "$RUN_DIR" + +say() { printf '\n\033[1;35m### %s\033[0m\n' "$*"; } +has_stage() { case ",$STAGES," in *",$1,"*) return 0;; *) return 1;; esac; } + +say "run dir: $RUN_DIR (stages: $STAGES)" + +if has_stage setup; then + say "stage: setup" + if [ "${FORCE_SETUP:-0}" != "1" ] && [ -f "$SCRIPT_DIR/env.sh" ] && \ + ( source "$SCRIPT_DIR/env.sh" >/dev/null 2>&1 && \ + python -c "import ray, ray_data_arrow_rs" >/dev/null 2>&1 ); then + say "environment already set up (env.sh + imports OK) - skipping setup.sh" + else + bash "$SCRIPT_DIR/setup.sh" + fi +fi +source "$SCRIPT_DIR/env.sh" +export RAY_ADDRESS=local + +# grpcio: needed by benchmark.py (tpch stage) AND read_probe's spilled_gb +# capture (release stage) — ray[data] doesn't ship it and an already-set-up +# box skips setup.sh. Missing it is SILENT in the release stage (spill column +# reads None/'—'), which is exactly what happened on the first box run, so +# top it up before any stage runs. +python -c "import grpc" >/dev/null 2>&1 || { + say "installing missing grpcio (spill capture + benchmark.py dependency)" + uv pip install --python "$(command -v python)" grpcio +} + +if has_stage fixtures; then + say "stage: fixtures (root=$FIXTURES_ROOT scale=$FIXTURE_SCALE, all shapes)" + python "$SCRIPT_DIR/gen_local_fixtures.py" --root "$FIXTURES_ROOT" \ + --scale "$FIXTURE_SCALE" 2>&1 | tee "$RUN_DIR/fixtures.log" +fi + +if has_stage mechanism; then + say "stage: mechanism (batch_ablation: budgets=$BUDGETS)" + python "$SCRIPT_DIR/batch_ablation.py" \ + --fixtures-root "$FIXTURES_ROOT" --scale "$FIXTURE_SCALE" \ + --budgets-mib "$BUDGETS" --out "$RUN_DIR" \ + ${ABLATION_SHAPES:+--shapes "$ABLATION_SHAPES"} \ + 2>&1 | tee "$RUN_DIR/mechanism.log" +fi + +if has_stage release; then + say "stage: release-metric (loss_triage: shapes=$TRIAGE_SHAPES s3=${ARROW_RS_S3_BUCKET:-off})" + python "$SCRIPT_DIR/loss_triage.py" \ + --fixture-root "$FIXTURES_ROOT" --outdir "$RUN_DIR/loss_triage" \ + --shapes "$TRIAGE_SHAPES" --repeat "$REPEAT" --warmup "$WARMUP" \ + 2>&1 | tee "$RUN_DIR/release.log" +fi + +if has_stage tpch; then + say "stage: tpch suspects (sf=$TPCH_SF — q9 spill T27, q20 shuffle_v2 M46)" + python "$SCRIPT_DIR/tpch_probe.py" \ + --outdir "$RUN_DIR/tpch" --sf "$TPCH_SF" \ + ${TPCH_QUERIES:+--queries "$TPCH_QUERIES"} \ + ${TPCH_STRATEGIES:+--strategies "$TPCH_STRATEGIES"} \ + 2>&1 | tee "$RUN_DIR/tpch.log" +fi + +if has_stage soak; then + say "stage: soak/retention (shapes=$SOAK_SHAPES arms=$ARMS workers=$WORKERS)" + python "$SCRIPT_DIR/soak_probe.py" \ + --fixture-root "$FIXTURES_ROOT" --outdir "$RUN_DIR/soak" \ + --shapes "$SOAK_SHAPES" --arms "$ARMS" --workers "$WORKERS" \ + ${ROUNDS:+--rounds "$ROUNDS"} ${PATH_REPEAT:+--path-repeat "$PATH_REPEAT"} \ + 2>&1 | tee "$RUN_DIR/soak.log" +fi + +if has_stage tensors; then + say "stage: tensors nbytes probe (M39 representation check)" + python "$SCRIPT_DIR/tensors_nbytes_probe.py" --fixture-root "$FIXTURES_ROOT" \ + 2>&1 | tee "$RUN_DIR/tensors_nbytes.log" +fi + +say "DONE — artifact index" +for f in "$RUN_DIR/ablation.json" "$RUN_DIR/loss_triage/summary.json" \ + "$RUN_DIR/tpch/summary.json" "$RUN_DIR/soak/summary.json" "$RUN_DIR"/*.log; do + [ -e "$f" ] && echo " $f" +done +echo +echo "verdict lines (grep of the stage tables):" +grep -h "gate verdict" "$RUN_DIR/mechanism.log" 2>/dev/null || true +echo " release + soak R-tables are at the end of release.log / soak.log" diff --git a/release/nightly_tests/dataset/arrow_rs_probe/run_grand_experiment.sh b/release/nightly_tests/dataset/arrow_rs_probe/run_grand_experiment.sh new file mode 100644 index 000000000000..72704751b165 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/run_grand_experiment.sh @@ -0,0 +1,81 @@ +#!/usr/bin/env bash +# --------------------------------------------------------------------------- +# ONE command for a fresh Linux box: environment -> correctness gate -> +# fixtures -> the grand tuning experiment (arrow-rs vs PyArrow on the new +# footer-based planner). +# +# bash release/nightly_tests/dataset/arrow_rs_probe/run_grand_experiment.sh +# +# What it does, in order: +# 1. setup.sh venv + commit-matched Ray wheel + local-source symlink +# + Rust toolchain + native crate build + end-to-end check +# 2. correctness gate the two pytest suites. The port to #64985 has only been +# verified statically — benchmarking an incorrect reader +# is worthless, so red tests ABORT the run (SKIP_TESTS=1 +# to override while debugging). +# 3. fixtures gen_local_fixtures.py (5 shapes, ~3.5 GiB, idempotent) +# 4. experiment grand_experiment.py stages A-D locally; stage E (S3) +# engages iff ARROW_RS_S3_BUCKET is exported. +# +# Knobs (env): +# SKIP_TESTS=1 skip the pytest gate +# FIXTURE_SCALE=0.25 smaller fixtures for a quick smoke run +# FIXTURES_ROOT= default ~/arrow_rs_grand_fixtures +# REPEAT=3 median-of-N per cell (default 1) +# STAGES=A,B subset of A,B,C,D,E (default all) +# ARROW_RS_S3_BUCKET=s3://... scratch bucket -> enables stage E (env.sh exports +# the default one; unset it to disable stage E) +# FORCE_SETUP=1 re-run setup.sh even if the env already imports +# (re-runs otherwise skip it automatically) +# --------------------------------------------------------------------------- +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "$SCRIPT_DIR/../../../.." && pwd)" +export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH" # uv + cargo (setup.sh installs) +FIXTURES_ROOT="${FIXTURES_ROOT:-$HOME/arrow_rs_grand_fixtures}" +FIXTURE_SCALE="${FIXTURE_SCALE:-1.0}" +REPEAT="${REPEAT:-1}" +STAGES="${STAGES:-A,B,C,D,E}" + +say() { printf '\n\033[1;35m### %s\033[0m\n' "$*"; } + +say "1/4 environment (setup.sh)" +# Skip setup on re-runs when the environment already works — setup.sh is +# idempotent but re-runs the Ray wheel install (~minutes) every time. +# FORCE_SETUP=1 forces a full setup.sh run. +if [ "${FORCE_SETUP:-0}" != "1" ] && [ -f "$SCRIPT_DIR/env.sh" ] && \ + ( source "$SCRIPT_DIR/env.sh" >/dev/null 2>&1 && \ + python -c "import ray, ray_data_arrow_rs" >/dev/null 2>&1 ); then + say "environment already set up (env.sh + imports OK) — skipping setup.sh (FORCE_SETUP=1 to redo)" +else + bash "$SCRIPT_DIR/setup.sh" +fi +# env.sh (written by setup.sh) holds venv activation + RAY_ADDRESS=local + +# the memory guard that keeps an OOM from killing the whole node. +source "$SCRIPT_DIR/env.sh" + +say "installing test deps (pytest, pandas)" +uv pip install --python "$(command -v python)" -q pytest pandas + +if [ "${SKIP_TESTS:-0}" != "1" ]; then + say "2/4 correctness gate: arrow-rs suite" + python -m pytest "$REPO/python/ray/data/tests/datasource/test_arrow_rs_parquet_reader.py" \ + -q --tb=short -p no:cacheprovider + say "2/4 correctness gate: parquet V2 suite" + python -m pytest "$REPO/python/ray/data/tests/datasource/test_read_parquet_v2.py" \ + -q --tb=short -p no:cacheprovider +else + say "2/4 SKIP_TESTS=1 — correctness gate skipped" +fi + +say "3/4 fixtures (root=$FIXTURES_ROOT scale=$FIXTURE_SCALE)" +python "$SCRIPT_DIR/gen_local_fixtures.py" --root "$FIXTURES_ROOT" --scale "$FIXTURE_SCALE" + +say "4/4 grand experiment (stages=$STAGES repeat=$REPEAT s3=${ARROW_RS_S3_BUCKET:-off})" +python "$SCRIPT_DIR/grand_experiment.py" \ + --fixtures-root "$FIXTURES_ROOT" \ + --repeat "$REPEAT" \ + --stages "$STAGES" + +say "DONE — summary.md printed above; full logs under $SCRIPT_DIR/grand_runs/" diff --git a/release/nightly_tests/dataset/arrow_rs_probe/run_loss_triage.sh b/release/nightly_tests/dataset/arrow_rs_probe/run_loss_triage.sh new file mode 100644 index 000000000000..24617d790dfc --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/run_loss_triage.sh @@ -0,0 +1,74 @@ +#!/usr/bin/env bash +# --------------------------------------------------------------------------- +# ONE command for the Linux box: the 3-part triage of the 2026-08-15 release +# A/B losses (findings M31 read_large_parquet_autoscaling, M32 write_parquet, +# M33 wide_schema tensors). +# +# bash release/nightly_tests/dataset/arrow_rs_probe/run_loss_triage.sh +# # with the S3 part (scratch bucket you own — fixtures are synced up): +# ARROW_RS_S3_BUCKET=s3://arrowrs-bench-xxxx bash .../run_loss_triage.sh +# +# Each loss shape runs standalone (no Ray/no S3), through Ray on local files, +# and through Ray on S3 — both readers each, plus a MALLOC_ARENA_MAX=2 arm on +# the arrow-rs Ray cells — so one summary table says whether a loss is the +# native decoder, Ray integration (worker/allocator), or the crate's S3 path. +# See loss_triage.py's docstring for the shape -> release-test mapping. +# +# Same setup skeleton as run_replication.sh: +# 1. setup.sh venv + commit-matched wheel + source symlink + crate +# (skipped when env.sh already imports; FORCE_SETUP=1 to +# redo — MANDATORY after a git pull that touches the crate) +# 2. fixtures gen_local_fixtures.py: auto_rg, bin_sweep, tensors_cp, tensors_dict +# 3. matrix loss_triage.py (S3 part auto-enabled when +# ARROW_RS_S3_BUCKET is set; AWS creds must be exported) +# +# Knobs (env): +# FIXTURE_SCALE=0.25 smaller fixtures for a smoke run +# FIXTURES_ROOT= default ~/arrow_rs_repl_fixtures +# REPEAT=3 WARMUP=1 per-cell medians (as in run_replication.sh) +# SHAPES=write subset: auto,write,tensorscp,tensorsdict,agg +# PARTS=ray_local subset: standalone,ray_local,ray_s3 +# FORCE_SETUP=1 re-run setup.sh even if the env already imports +# --------------------------------------------------------------------------- +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH" +FIXTURES_ROOT="${FIXTURES_ROOT:-$HOME/arrow_rs_repl_fixtures}" +FIXTURE_SCALE="${FIXTURE_SCALE:-1.0}" +REPEAT="${REPEAT:-3}" +WARMUP="${WARMUP:-1}" +SHAPES="${SHAPES:-auto,write,tensorscp,tensorsdict,agg}" +PARTS="${PARTS:-}" + +say() { printf '\n\033[1;35m### %s\033[0m\n' "$*"; } + +say "1/3 environment (setup.sh)" +# env.sh (written by setup.sh) exports its own ARROW_RS_S3_BUCKET; preserve a +# value the caller set explicitly so `ARROW_RS_S3_BUCKET=… bash run_loss_triage.sh` +# is honored rather than silently clobbered. +CALLER_S3_BUCKET="${ARROW_RS_S3_BUCKET:-}" +if [ "${FORCE_SETUP:-0}" != "1" ] && [ -f "$SCRIPT_DIR/env.sh" ] && \ + ( source "$SCRIPT_DIR/env.sh" >/dev/null 2>&1 && \ + python -c "import ray, ray_data_arrow_rs" >/dev/null 2>&1 ); then + say "environment already set up (env.sh + imports OK) — skipping setup.sh (FORCE_SETUP=1 to redo)" +else + bash "$SCRIPT_DIR/setup.sh" +fi +source "$SCRIPT_DIR/env.sh" +[ -n "$CALLER_S3_BUCKET" ] && export ARROW_RS_S3_BUCKET="$CALLER_S3_BUCKET" +export RAY_ADDRESS=local + +say "2/3 fixtures (root=$FIXTURES_ROOT scale=$FIXTURE_SCALE: auto_rg, bin_sweep, tensors_cp, tensors_dict)" +python "$SCRIPT_DIR/gen_local_fixtures.py" --root "$FIXTURES_ROOT" \ + --scale "$FIXTURE_SCALE" --shapes auto_rg,bin_sweep,tensors_cp,tensors_dict + +say "3/3 loss triage matrix (repeat=$REPEAT warmup=$WARMUP shapes=$SHAPES s3=${ARROW_RS_S3_BUCKET:-off})" +python "$SCRIPT_DIR/loss_triage.py" \ + --fixture-root "$FIXTURES_ROOT" \ + --repeat "$REPEAT" \ + --warmup "$WARMUP" \ + --shapes "$SHAPES" \ + ${PARTS:+--parts "$PARTS"} + +say "DONE — summary printed above; full logs under $SCRIPT_DIR/loss_triage_runs/" diff --git a/release/nightly_tests/dataset/arrow_rs_probe/run_matrix.py b/release/nightly_tests/dataset/arrow_rs_probe/run_matrix.py new file mode 100644 index 000000000000..e2295fb7a3dc --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/run_matrix.py @@ -0,0 +1,342 @@ +#!/usr/bin/env python3 +"""Run the full arrow_rs-vs-PyArrow A/B matrix against S3, log everything, tabulate. + +One place to reproduce both release regressions instead of typing the README's A/B/C +blocks by hand. Each cell runs read_probe.py in its OWN process (Ray + the crate load +once per process, and the mem sampler must not see a stale Ray instance), captures full +stdout+stderr to a per-cell log file, parses the RESULT block, and prints a comparison +table with arrow_rs/pyarrow ratios so you can see at a glance what is slow / heavy. + +The matrix: + [diag] concurrency=1, both layouts, both readers -> cpu_over_wall + imagenet ~1 => CPU-bound decode ; <<1 => I/O-waiting on S3 (prefetch is the fix) + [mem] wide fanned out, both readers -> peak_uss_gb / read_avg_max_uss_gb + the memory metric of record (Linux USS). expect arrow_rs at-or-below PyArrow. + [alloc] wide arrow_rs, allocator variants -> is any residual gap glibc arenas? + baseline vs MALLOC_ARENA_MAX=2 vs LD_PRELOAD jemalloc (skipped if .so absent). + + python run_matrix.py --wide-path s3://.../wide_schema/primitives \\ + --imagenet-path s3://.../imagenet/parquet + python run_matrix.py --wide-path ... --imagenet-path ... --skip alloc # subset + python run_matrix.py --wide-path ... --imagenet-path ... --repeat 3 # median of N + +Assumes: venv active, RAY_ADDRESS=local, AWS creds/region exported, box in bucket region. +""" +import argparse +import glob +import json +import os +import subprocess +import sys +import time + +PROBE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "read_probe.py") +PY = sys.executable + +# Common jemalloc locations on Debian/Ubuntu (Anyscale base image). +_JEMALLOC_GLOBS = [ + "/usr/lib/x86_64-linux-gnu/libjemalloc.so.2", + "/usr/lib/x86_64-linux-gnu/libjemalloc.so", + "/usr/lib/libjemalloc.so.2", +] + + +def find_jemalloc(): + for g in _JEMALLOC_GLOBS: + hits = glob.glob(g) + if hits: + return hits[0] + return None + + +def run_cell( + logdir, name, path, reader, concurrency, columns, extra_env, extra_args=None +): + """Run one read_probe.py invocation; tee output to a log; return parsed RESULT dict.""" + cmd = [PY, PROBE, "--path", path, "--reader", reader] + if concurrency is not None: + cmd += ["--concurrency", str(concurrency)] + if columns: + cmd += ["--columns", *columns] + if extra_args: + cmd += list(extra_args) + + env = dict(os.environ) + env.update(extra_env) + + logpath = os.path.join(logdir, f"{name}.log") + t0 = time.perf_counter() + proc = subprocess.run(cmd, capture_output=True, text=True, env=env) + dur = time.perf_counter() - t0 + + with open(logpath, "w") as fh: + fh.write(f"# cmd: {' '.join(cmd)}\n") + fh.write(f"# extra_env: {extra_env}\n") + fh.write(f"# wall_including_startup_s: {dur:.1f}\n") + fh.write("# ---- STDOUT ----\n") + fh.write(proc.stdout) + fh.write("\n# ---- STDERR ----\n") + fh.write(proc.stderr) + + res = {} + in_result = False + for line in proc.stdout.splitlines(): + if "=== RESULT ===" in line: + in_result = True + continue + if in_result and ":" in line: + k, v = line.strip().split(":", 1) + res[k.strip()] = v.strip() + if not res: + print( + f" !! {name} PROBE FAIL rc={proc.returncode} (see {logpath})\n" + f" {proc.stderr.strip()[-400:]}", + flush=True, + ) + else: + print( + f" {name:<34} wall={res.get('wall_s')} " + f"cpu/wall={res.get('cpu_over_wall')} " + f"uss={res.get('peak_uss_gb')} rss={res.get('peak_rss_gb')}", + flush=True, + ) + return res + + +def _num(res, key): + try: + return float(res.get(key)) + except (TypeError, ValueError): + return None + + +def _median(vals): + """Median of a numeric list; the mean of the two middles for an even count.""" + s = sorted(vals) + n = len(s) + if not n: + return None + mid = n // 2 + return s[mid] if n % 2 else (s[mid - 1] + s[mid]) / 2.0 + + +def median_cell(logdir, name, repeat, warmup=1, **kw): + """Run a cell `warmup`+`repeat` times; return a per-metric median of the measured runs. + + Two things this has to get right, both of which it got wrong until 2026-08-13 + (findings T25/T26): + + * **Median each metric independently.** Returning the whole dict of the run whose + `wall_s` was the median leaves every *other* metric a single sample chosen by an + unrelated metric. In the item-1o A/B that paired the fastest arrow_rs read with + the slowest PyArrow one and reported 1.04x where the honest ratio was ~1.7x. + * **Discard the warm-up run.** A cell's first repeat is reproducibly slow — cold page + cache for the fixture, cold `.so`/imports — 1.44 s against a 1.03-1.22 s steady + state on the fat_col shape, and reproducible to 0.35% across two independent runs. + One such sample drags a 3-sample median. Warm-ups still run and still log (as + `.w`); they just don't reach the stats. + + Non-numeric fields (reader, path, ...) carry over from the first measured run. + `_n` and `_samples` record the sample count and each metric's values, so spread is + visible in summary.json without opening the per-cell logs. + """ + for i in range(warmup): + run_cell(logdir, f"{name}.w{i}", **kw) + runs = [] + for i in range(repeat): + tag = name if repeat == 1 else f"{name}.r{i}" + runs.append(run_cell(logdir, tag, **kw)) + good = [r for r in runs if r] + if not good: + return {} + + out = dict(good[0]) + samples = {} + for key in list(out): + vals = [_num(r, key) for r in good] + if any(v is None for v in vals): + continue # not numeric in every run — keep the first run's value verbatim + out[key] = _median(vals) + samples[key] = vals + out["_n"] = len(good) + out["_samples"] = samples + return out + + +def ratio(a, b): + if a is None or b is None or b == 0: + return None + return round(a / b, 3) + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--wide-path", required=True) + p.add_argument("--imagenet-path", required=True) + p.add_argument("--imagenet-columns", nargs="+", default=["image", "label"]) + p.add_argument( + "--outdir", default=None, help="log dir (default ./matrix_runs/)" + ) + p.add_argument( + "--repeat", + type=int, + default=1, + help="measured runs per cell; each metric is reported as its own median", + ) + p.add_argument( + "--warmup", + type=int, + default=1, + help=( + "discarded runs per cell before the measured ones. The first run of a " + "cell is reproducibly slow (cold page cache / .so); 0 to disable." + ), + ) + p.add_argument( + "--skip", + default="", + help="comma list of stages to skip: diag,mem,alloc", + ) + args = p.parse_args() + + skip = {s.strip() for s in args.skip.split(",") if s.strip()} + ts = time.strftime("%Y%m%d_%H%M%S") + outdir = args.outdir or os.path.join( + os.path.dirname(os.path.abspath(__file__)), "matrix_runs", ts + ) + os.makedirs(outdir, exist_ok=True) + print(f"logs -> {outdir}\n", flush=True) + + rows = {} # name -> result dict + + def cell(name, **kw): + rows[name] = median_cell(outdir, name, args.repeat, args.warmup, **kw) + + # -------- [diag] cpu_over_wall @ concurrency=1 -------- + if "diag" not in skip: + print("=== [diag] cpu_over_wall @ concurrency=1 ===", flush=True) + for reader in ("pyarrow", "arrow_rs"): + cell( + f"diag.imagenet.{reader}", + path=args.imagenet_path, + reader=reader, + concurrency=1, + columns=args.imagenet_columns, + extra_env={}, + ) + cell( + f"diag.wide.{reader}", + path=args.wide_path, + reader=reader, + concurrency=1, + columns=None, + extra_env={}, + ) + + # -------- [mem] wide fanned out -------- + if "mem" not in skip: + print("=== [mem] wide fanned out (peak_uss_gb) ===", flush=True) + for reader in ("pyarrow", "arrow_rs"): + cell( + f"mem.wide.{reader}", + path=args.wide_path, + reader=reader, + concurrency=None, + columns=None, + extra_env={}, + ) + + # -------- [alloc] allocator A/B on wide arrow_rs -------- + if "alloc" not in skip: + print("=== [alloc] wide arrow_rs allocator A/B ===", flush=True) + cell( + "alloc.baseline", + path=args.wide_path, + reader="arrow_rs", + concurrency=None, + columns=None, + extra_env={}, + ) + cell( + "alloc.arena2", + path=args.wide_path, + reader="arrow_rs", + concurrency=None, + columns=None, + extra_env={"MALLOC_ARENA_MAX": "2"}, + ) + jem = find_jemalloc() + if jem: + cell( + "alloc.jemalloc", + path=args.wide_path, + reader="arrow_rs", + concurrency=None, + columns=None, + extra_env={"LD_PRELOAD": jem}, + ) + else: + print( + " (skip alloc.jemalloc: no libjemalloc.so found; " + "`apt-get install -y libjemalloc2` to enable)", + flush=True, + ) + + # -------- summary -------- + print("\n=================== SUMMARY (arrow_rs / pyarrow) ===================") + + def pair(prefix): + pa_r, ar_r = rows.get(f"{prefix}.pyarrow", {}), rows.get( + f"{prefix}.arrow_rs", {} + ) + for metric, label in [ + ("wall_s", "wall_s"), + ("cpu_over_wall", "cpu/wall"), + ("peak_uss_gb", "peak_uss_gb"), + ("read_avg_max_uss_gb", "read_uss_gb"), + ("peak_rss_gb", "peak_rss_gb"), + ]: + a, b = _num(ar_r, metric), _num(pa_r, metric) + if a is None and b is None: + continue + r = ratio(a, b) + flag = "" + if r is not None and metric in ( + "wall_s", + "peak_uss_gb", + "read_avg_max_uss_gb", + ): + flag = " <-- WORSE" if r > 1.05 else (" <-- win" if r < 0.95 else "") + print( + f" {prefix:<16} {label:<14} " + f"pyarrow={b} arrow_rs={a} ratio={r}{flag}" + ) + + if "diag" not in skip: + print( + "[diag imagenet] cpu/wall<<1 on the S3 read => I/O-bound (prefetch is the fix)" + ) + pair("diag.imagenet") + print("[diag wide]") + pair("diag.wide") + if "mem" not in skip: + print("[mem wide] the metric of record: peak_uss_gb / read_uss_gb") + pair("mem.wide") + if "alloc" not in skip: + print("[alloc wide arrow_rs] is residual mem gap glibc arena retention?") + base = _num(rows.get("alloc.baseline", {}), "peak_uss_gb") + for variant in ("alloc.arena2", "alloc.jemalloc"): + v = _num(rows.get(variant, {}), "peak_uss_gb") + if v is not None: + print( + f" {variant:<16} peak_uss_gb={v} " + f"vs baseline={base} ratio={ratio(v, base)}" + ) + + with open(os.path.join(outdir, "summary.json"), "w") as fh: + json.dump(rows, fh, indent=2) + print(f"\nfull JSON + per-cell logs in {outdir}") + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/run_release_regressions.sh b/release/nightly_tests/dataset/arrow_rs_probe/run_release_regressions.sh new file mode 100644 index 000000000000..3b78c0ea15ac --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/run_release_regressions.sh @@ -0,0 +1,91 @@ +#!/usr/bin/env bash +# --------------------------------------------------------------------------- +# ONE command for a fresh Linux box: replicate every A/B #5 >1.15 regression +# that public data allows, single-node, both readers (user ask 2026-08-28; +# the P0 ledger is arrow_rs_docs/2026-08-27.md §11, the replication map §12). +# +# ARROW_RS_S3_BUCKET=s3://arrowrs-bench-xxxx \ +# bash release/nightly_tests/dataset/arrow_rs_probe/run_release_regressions.sh +# +# Same skeleton as run_replication.sh (known-good on these boxes): +# 1. setup.sh venv + commit-matched wheel + crate build (skipped when +# env.sh already imports; FORCE_SETUP=1 after a git pull +# that touches the crate, or you benchmark a stale .so) +# 2. tpch leg tpch_probe.py over the P0/addendum queries +# (q2,q3,q4,q6,q9,q10,q11,q13,q14,q17,q18,q22 — q6 rides +# for the sustained-wUSS 5.8x row, q9 for the T27 spill +# history) x both shuffle strategies x both readers, +# sf ${TPCH_SF:-10} (release: 1000) +# 3. non-tpch leg release_regression_probe.py: iter_batches_pyarrow +# (exact replica), write_parquet (sf100 for sf1000), +# read_parquet_binned (public-data approximation), +# map_groups hash/hashv2/sort over col02+14 AND +# col08+13+14 (the T-spills-more shapes), aggregate +# positive controls (M47), joins +# 4. rlp analog NOT run here — run_loss_triage.sh's auto_rg S3 shape +# is the read_large_parquet stand-in (internal bucket is +# ACCESS_DENIED); wide_schema_objects has NO analog. +# +# Knobs (env): TPCH_SF, WRITE_SF, GROUPBY_SF, JOINS_SF (downsizes), +# JOIN_TYPES=right_outer[,inner,...], REPEAT, ONLY (non-tpch cell filter), +# SKIP_TPCH=1 / SKIP_PROBE=1, FORCE_SETUP=1, DRY_RUN=1, +# ARMS=pa,rs,rseos (add rstrim; 2026-09-04 allocator arms), CPUS=24,48,96 +# (non-tpch cells once per local num_cpus — the M107 col02 sweep), +# MONITOR_INTERVAL=1.0 (node-mem sampler seconds; 0.1 for the q6/wide rows) +# --------------------------------------------------------------------------- +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH" +OUT_ROOT="${OUT_ROOT:-$HOME/arrow_rs_regression_runs/$(date +%Y%m%d_%H%M%S)}" +mkdir -p "$OUT_ROOT" + +say() { printf '\n\033[1m== %s ==\033[0m\n' "$*"; } + +say "1/3 environment (setup.sh; FORCE_SETUP=${FORCE_SETUP:-0})" +if [ "${FORCE_SETUP:-0}" = "1" ] || ! bash -c "source '$SCRIPT_DIR/env.sh' 2>/dev/null \ + && python -c 'import ray_data_arrow_rs'" 2>/dev/null; then + bash "$SCRIPT_DIR/setup.sh" +fi +# shellcheck disable=SC1091 +source "$SCRIPT_DIR/env.sh" + +DRY_FLAG="" +[ "${DRY_RUN:-0}" = "1" ] && DRY_FLAG="--dry-run" + +if [ "${SKIP_TPCH:-0}" != "1" ]; then + # P0-scoped matrix, per strategy: most P0 rows are hash_shuffle_v2; v1 runs + # only where a P0 row names v1 (q4,q14,q22,q17) + q6 (in both: the 5.8x + # sustained-wUSS row). The full 12x2 matrix is infeasible on one box: v1 + # RPC-shuffles 8-60M rows at ~13k rows/s (q2 v1 pa: >20 min, timed out at + # sf10). q9 (T27) rides in the additions pass. Override via TPCH_QUERIES_*. + say "2/3 tpch leg, hash_shuffle_v2 (sf ${TPCH_SF:-10}; $OUT_ROOT/tpch_v2)" + python "$SCRIPT_DIR/tpch_probe.py" --outdir "$OUT_ROOT/tpch_v2" \ + --sf "${TPCH_SF:-10}" --repeat "${REPEAT:-1}" $DRY_FLAG \ + --cell-timeout "${TPCH_CELL_TIMEOUT:-3600}" \ + --arms "${ARMS:-pa,rs,rseos}" --monitor-interval "${MONITOR_INTERVAL:-1.0}" \ + --queries "${TPCH_QUERIES_V2:-tpch_q2,tpch_q3,tpch_q6,tpch_q10,tpch_q11,tpch_q13,tpch_q17,tpch_q18}" \ + --strategies hash_shuffle_v2 2>&1 | tee "$OUT_ROOT/tpch_v2.out" + say "2/3 tpch leg, hash_shuffle v1 (sf ${TPCH_SF:-10}; $OUT_ROOT/tpch_v1)" + python "$SCRIPT_DIR/tpch_probe.py" --outdir "$OUT_ROOT/tpch_v1" \ + --sf "${TPCH_SF:-10}" --repeat "${REPEAT:-1}" $DRY_FLAG \ + --cell-timeout "${TPCH_CELL_TIMEOUT:-3600}" \ + --arms "${ARMS:-pa,rs,rseos}" --monitor-interval "${MONITOR_INTERVAL:-1.0}" \ + --queries "${TPCH_QUERIES_V1:-tpch_q4,tpch_q6,tpch_q14,tpch_q17,tpch_q22}" \ + --strategies hash_shuffle 2>&1 | tee "$OUT_ROOT/tpch_v1.out" +fi + +if [ "${SKIP_PROBE:-0}" != "1" ]; then + say "3/3 non-tpch leg (logs under $OUT_ROOT/probe)" + python "$SCRIPT_DIR/release_regression_probe.py" --outdir "$OUT_ROOT/probe" \ + --repeat "${REPEAT:-1}" $DRY_FLAG \ + --cell-timeout "${PROBE_CELL_TIMEOUT:-3600}" \ + --write-sf "${WRITE_SF:-100}" --groupby-sf "${GROUPBY_SF:-10}" \ + --joins-sf "${JOINS_SF:-10}" --join-types "${JOIN_TYPES:-right_outer}" \ + --arms "${ARMS:-pa,rs,rseos}" --monitor-interval "${MONITOR_INTERVAL:-1.0}" \ + ${CPUS:+--cpus "$CPUS"} ${ONLY:+--only "$ONLY"} 2>&1 | tee "$OUT_ROOT/probe.out" +fi + +say "done — tables in $OUT_ROOT/{tpch_v2,tpch_v1,probe}.out; per-cell benchmark JSONs beside them" +echo "reminder: the read_large_parquet analog is run_loss_triage.sh (auto_rg, S3);" +echo "wide_schema_objects cannot be replicated (internal bucket, no fixture)." diff --git a/release/nightly_tests/dataset/arrow_rs_probe/run_replication.sh b/release/nightly_tests/dataset/arrow_rs_probe/run_replication.sh new file mode 100755 index 000000000000..b3acd87ada1f --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/run_replication.sh @@ -0,0 +1,90 @@ +#!/usr/bin/env bash +# --------------------------------------------------------------------------- +# ONE command for a fresh Linux box: environment -> correctness gate -> +# fixtures -> the 2026-08-12 release-A/B replication matrix (TODO 1ab phase 1). +# +# bash release/nightly_tests/dataset/arrow_rs_probe/run_replication.sh +# +# Same skeleton as run_grand_experiment.sh (which is the setup that already +# works on these boxes), pointed at replication_matrix.py: +# 1. setup.sh venv + commit-matched Ray wheel + local-source symlink +# + Rust toolchain + native crate build + end-to-end check +# (skipped automatically when env.sh already imports; +# FORCE_SETUP=1 to redo — MANDATORY after a git pull that +# touches the crate, or you benchmark a stale .so) +# 2. correctness gate the two pytest suites; red tests ABORT (SKIP_TESTS=1 +# to override while debugging) +# 3. fixtures gen_local_fixtures.py, replication shapes only +# (bin_sweep ~4 GiB, tensors_wide ~1.6 GiB, fat_col) +# 4. matrix replication_matrix.py: tensors / tensorscp / binsweep / +# binbound / write / fatcol / oom (see its docstring for the +# rationale; binbound is the "is per-task USS bounded by +# the bin budget?" check and needs Linux — USS is None on +# macOS; oom deliberately gets PyArrow's arm OOM-killed +# by Ray's memory monitor, so FAILED pyarrow cells there +# are the result, not a broken run) +# +# Knobs (env): +# SKIP_TESTS=1 skip the pytest gate +# FIXTURE_SCALE=0.25 smaller fixtures for a quick smoke run +# FIXTURES_ROOT= default ~/arrow_rs_repl_fixtures +# REPEAT=3 measured runs per cell (default 3 — these are the +# numbers we act on, so default to medians). Every +# metric is medianed on its own; per-repeat values +# land in summary.json as _samples +# WARMUP=1 discarded runs per cell before the measured ones +# (default 1). The first run of a cell is cold — +# page cache, .so load — and drags a 3-run median. +# WARMUP=0 to measure the cold path deliberately +# ONLY=binsweep,tensors run a subset of stages +# FORCE_SETUP=1 re-run setup.sh even if the env already imports +# --------------------------------------------------------------------------- +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "$SCRIPT_DIR/../../../.." && pwd)" +export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH" # uv + cargo (setup.sh installs) +FIXTURES_ROOT="${FIXTURES_ROOT:-$HOME/arrow_rs_repl_fixtures}" +FIXTURE_SCALE="${FIXTURE_SCALE:-1.0}" +REPEAT="${REPEAT:-3}" +WARMUP="${WARMUP:-1}" +ONLY="${ONLY:-}" + +say() { printf '\n\033[1;35m### %s\033[0m\n' "$*"; } + +say "1/4 environment (setup.sh)" +if [ "${FORCE_SETUP:-0}" != "1" ] && [ -f "$SCRIPT_DIR/env.sh" ] && \ + ( source "$SCRIPT_DIR/env.sh" >/dev/null 2>&1 && \ + python -c "import ray, ray_data_arrow_rs" >/dev/null 2>&1 ); then + say "environment already set up (env.sh + imports OK) — skipping setup.sh (FORCE_SETUP=1 to redo)" +else + bash "$SCRIPT_DIR/setup.sh" +fi +source "$SCRIPT_DIR/env.sh" + +say "installing test deps (pytest, pandas)" +uv pip install --python "$(command -v python)" -q pytest pandas + +if [ "${SKIP_TESTS:-0}" != "1" ]; then + say "2/4 correctness gate: arrow-rs suite" + python -m pytest "$REPO/python/ray/data/tests/datasource/test_arrow_rs_parquet_reader.py" \ + -q --tb=short -p no:cacheprovider + say "2/4 correctness gate: parquet V2 suite" + python -m pytest "$REPO/python/ray/data/tests/datasource/test_read_parquet_v2.py" \ + -q --tb=short -p no:cacheprovider +else + say "2/4 SKIP_TESTS=1 — correctness gate skipped" +fi + +say "3/4 fixtures (root=$FIXTURES_ROOT scale=$FIXTURE_SCALE, replication shapes)" +python "$SCRIPT_DIR/gen_local_fixtures.py" --root "$FIXTURES_ROOT" \ + --scale "$FIXTURE_SCALE" --shapes bin_sweep,tensors_wide,tensors_cp,fat_col + +say "4/4 replication matrix (repeat=$REPEAT warmup=$WARMUP only=${ONLY:-all})" +python "$SCRIPT_DIR/replication_matrix.py" \ + --fixture-root "$FIXTURES_ROOT" \ + --repeat "$REPEAT" \ + --warmup "$WARMUP" \ + ${ONLY:+--only "$ONLY"} + +say "DONE — summary printed above; full logs under $SCRIPT_DIR/replication_runs/" diff --git a/release/nightly_tests/dataset/arrow_rs_probe/run_soak.sh b/release/nightly_tests/dataset/arrow_rs_probe/run_soak.sh new file mode 100755 index 000000000000..963173233e04 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/run_soak.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# --------------------------------------------------------------------------- +# ONE command for the Linux box: the soak/churn discriminator for the A/B #4 +# retention losses (findings M37 read_large_parquet_autoscaling, M38 +# write_parquet) plus the M39 tensors representation check. +# +# bash release/nightly_tests/dataset/arrow_rs_probe/run_soak.sh +# +# Why this exists: every prior box run (replication_matrix, loss_triage) used +# a FRESH Ray session per cell and ~tens of tasks per worker - and none of the +# release losses reproduced (M35). A/B #4 then proved the losses are real at a +# 20 Hz poll with byte-identical decoder work, so the untested variable is +# WORKER LIFETIME x TASK COUNT. soak_probe.py holds one long-lived Ray session +# per arm (pa / rs / rs+MALLOC_ARENA_MAX=2), pushes rounds of the loss shapes +# through a pinned worker pool until each worker has executed O(100+) tasks, +# and reads the idle-USS floor after every round. See soak_probe.py's +# docstring for the verdict table. +# +# Same setup skeleton as run_loss_triage.sh: +# 1. setup.sh venv + commit-matched wheel + source symlink + crate +# (skipped when env.sh already imports; FORCE_SETUP=1 to +# redo - MANDATORY after a git pull that touches the crate) +# 2. fixtures gen_local_fixtures.py: auto_rg, bin_sweep, tensors_cp +# 3. soak matrix soak_probe.py (local files only - the retention +# question is transport-independent per M35/M38) +# 4. tensors nbytes tensors_nbytes_probe.py (M39, cheap, standalone) +# +# Knobs (env): +# FIXTURE_SCALE=0.25 smaller fixtures for a smoke run +# FIXTURES_ROOT= default ~/arrow_rs_repl_fixtures +# SHAPES=auto subset: auto,write +# ARMS=pa,rs subset of: pa,rs,rs_arena2,rs_trim,rs_jemalloc +# WORKERS=4 pinned worker-pool size (num_cpus) +# ROUNDS= / PATH_REPEAT= override soak_probe per-shape defaults +# SKIP_TENSORS=1 skip step 4 +# FORCE_SETUP=1 re-run setup.sh even if the env already imports +# --------------------------------------------------------------------------- +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH" +FIXTURES_ROOT="${FIXTURES_ROOT:-$HOME/arrow_rs_repl_fixtures}" +FIXTURE_SCALE="${FIXTURE_SCALE:-1.0}" +SHAPES="${SHAPES:-auto,write}" +ARMS="${ARMS:-pa,rs,rs_arena2,rs_trim,rs_jemalloc}" +WORKERS="${WORKERS:-4}" + +say() { printf '\n\033[1;35m### %s\033[0m\n' "$*"; } + +say "1/4 environment (setup.sh)" +if [ "${FORCE_SETUP:-0}" != "1" ] && [ -f "$SCRIPT_DIR/env.sh" ] && \ + ( source "$SCRIPT_DIR/env.sh" >/dev/null 2>&1 && \ + python -c "import ray, ray_data_arrow_rs" >/dev/null 2>&1 ); then + say "environment already set up (env.sh + imports OK) - skipping setup.sh (FORCE_SETUP=1 to redo)" +else + bash "$SCRIPT_DIR/setup.sh" +fi +source "$SCRIPT_DIR/env.sh" +export RAY_ADDRESS=local + +say "2/4 fixtures (root=$FIXTURES_ROOT scale=$FIXTURE_SCALE: auto_rg, bin_sweep, tensors_cp)" +python "$SCRIPT_DIR/gen_local_fixtures.py" --root "$FIXTURES_ROOT" \ + --scale "$FIXTURE_SCALE" --shapes auto_rg,bin_sweep,tensors_cp + +say "3/4 soak matrix (shapes=$SHAPES arms=$ARMS workers=$WORKERS)" +python "$SCRIPT_DIR/soak_probe.py" \ + --fixture-root "$FIXTURES_ROOT" \ + --shapes "$SHAPES" \ + --arms "$ARMS" \ + --workers "$WORKERS" \ + ${ROUNDS:+--rounds "$ROUNDS"} \ + ${PATH_REPEAT:+--path-repeat "$PATH_REPEAT"} + +if [ "${SKIP_TENSORS:-0}" != "1" ]; then + say "4/4 tensors nbytes probe (M39: representation vs batch sizing)" + python "$SCRIPT_DIR/tensors_nbytes_probe.py" --fixture-root "$FIXTURES_ROOT" +fi + +say "DONE - soak summary printed above; series + logs under $SCRIPT_DIR/soak_runs/" diff --git a/release/nightly_tests/dataset/arrow_rs_probe/scale_sweep.py b/release/nightly_tests/dataset/arrow_rs_probe/scale_sweep.py new file mode 100644 index 000000000000..40bd5f0bbb73 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/scale_sweep.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 +"""Scale sweep: is the arrow-rs regression an issue of SCALE? + +Generates local parquet fixtures at increasing scale for the two worse-in-release +layouts and runs BOTH readers via read_probe.py at --concurrency 1. Use this to +separate a scaling effect from an environment (S3/Linux) effect. + + imagenet : scale total bytes, tiny row groups fixed. If arrow_rs stays faster as + data grows, the release time gap is pure S3 network (not scale). On the + Mac it stayed 0.65-0.82x of pyarrow at 0.4-2.6 GB -> confirmed network. + wide : scale ROW-GROUP SIZE (the per-rg decode transient). On the Mac the + arrow_rs/pyarrow RSS ratio climbed 0.67x -> 1.04x as rg grew from 200 to + 4000 rows -> the mem regression's lever is row-group size, not row count. + On Linux read `peak_uss_gb` (populated there) for the real magnitude. + +Run: python scale_sweep.py where case in {imagenet, wide} +Fixtures land under $SWEEP_DIR (default /tmp/arrow_rs_sweep); each run prints a table. +""" +import json +import os +import subprocess +import sys + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +SWEEP = os.environ.get("SWEEP_DIR", "/tmp/arrow_rs_sweep") +PROBE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "read_probe.py") +PY = sys.executable +os.makedirs(SWEEP, exist_ok=True) + + +def gen_imagenet(tag, n_rows, blob_kb=48, rg=32, n_files=4): + d = os.path.join(SWEEP, f"imagenet_{tag}") + if os.path.exists(os.path.join(d, "_done")): + return d + os.makedirs(d, exist_ok=True) + rng = np.random.default_rng(0) + per = n_rows // n_files + for f in range(n_files): + images = [rng.bytes(blob_kb * 1024) for _ in range(per)] + labels = rng.integers(0, 1000, per, dtype=np.int64) + ids = np.arange(f * per, (f + 1) * per, dtype=np.int64) + t = pa.table( + { + "id": ids, + "image": pa.array(images, type=pa.binary()), + "label": pa.array(labels), + } + ) + pq.write_table( + t, + os.path.join(d, f"part{f}.parquet"), + write_page_index=True, + row_group_size=rg, + ) + open(os.path.join(d, "_done"), "w").close() + print( + f" gen imagenet_{tag}: {n_rows} rows ~{n_rows*blob_kb/1024:.0f}MB " + f"rg={rg} -> ~{per//rg} rgs/file x{n_files}", + flush=True, + ) + return d + + +def gen_wide(tag, n_rows, n_cols=5000, str_len=100, rg=200): + d = os.path.join(SWEEP, f"wide_{tag}") + if os.path.exists(os.path.join(d, "_done")): + return d + os.makedirs(d, exist_ok=True) + base = np.array(["x" * str_len] * n_rows, dtype=object) + cols = {"id": pa.array(np.arange(n_rows, dtype=np.int64))} + for c in range(n_cols): + cols[f"c{c}"] = pa.array(base, type=pa.string()) + t = pa.table(cols) + pq.write_table( + t, os.path.join(d, "wide.parquet"), write_page_index=True, row_group_size=rg + ) + open(os.path.join(d, "_done"), "w").close() + print( + f" gen wide_{tag}: {n_rows}x{n_cols} ~{n_rows*n_cols*str_len/1e6:.0f}MB rg={rg}", + flush=True, + ) + return d + + +def run_probe(path, reader, columns=None): + cmd = [PY, PROBE, "--path", path, "--reader", reader, "--concurrency", "1"] + if columns: + cmd += ["--columns", *columns] + out = subprocess.run(cmd, capture_output=True, text=True, env=dict(os.environ)) + res = {} + in_result = False + for line in out.stdout.splitlines(): + if "=== RESULT ===" in line: + in_result = True + continue + if in_result and ":" in line: + k, v = line.strip().split(":", 1) + res[k.strip()] = v.strip() + if not res: + print( + f" PROBE FAIL ({reader}) rc={out.returncode}\n" + f"{out.stdout[-800:]}\n{out.stderr[-800:]}" + ) + return res + + +def sweep_imagenet(): + print("=== imagenet scale sweep (cols=[image,label], tiny rgs) ===", flush=True) + scales = [("s", 8000), ("m", 24000), ("l", 56000)] # ~375MB, ~1.1GB, ~2.6GB + rows = [] + for tag, n in scales: + d = gen_imagenet(tag, n) + for reader in ("pyarrow", "arrow_rs"): + r = run_probe(d, reader, columns=["image", "label"]) + rows.append((tag, n, reader, r)) + print( + f" {tag:>2} n={n:<6} {reader:<8} wall={r.get('wall_s')} " + f"rss={r.get('peak_rss_gb')} uss={r.get('peak_uss_gb')} " + f"cpu/wall={r.get('cpu_over_wall')}", + flush=True, + ) + return rows + + +def sweep_wide(): + print("=== wide_schema row-group-size sweep (5000 cols) ===", flush=True) + variants = [("rg200", 4000, 200), ("rg1000", 4000, 1000), ("rg4000", 4000, 4000)] + rows = [] + for tag, n, rg in variants: + d = gen_wide(tag, n, rg=rg) + for reader in ("pyarrow", "arrow_rs"): + r = run_probe(d, reader) + rows.append((tag, rg, reader, r)) + print( + f" {tag:<7} rg={rg:<5} {reader:<8} wall={r.get('wall_s')} " + f"rss={r.get('peak_rss_gb')} uss={r.get('peak_uss_gb')} " + f"read_uss={r.get('read_avg_max_uss_gb')}", + flush=True, + ) + return rows + + +if __name__ == "__main__": + case = sys.argv[1] if len(sys.argv) > 1 else "imagenet" + rows = sweep_imagenet() if case == "imagenet" else sweep_wide() + print("\n=== JSON ===") + print(json.dumps([{"tag": t, "k": k, "reader": rd, **r} for t, k, rd, r in rows])) diff --git a/release/nightly_tests/dataset/arrow_rs_probe/setup.sh b/release/nightly_tests/dataset/arrow_rs_probe/setup.sh new file mode 100755 index 000000000000..1c4c25b4aa06 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/setup.sh @@ -0,0 +1,222 @@ +#!/usr/bin/env bash +# --------------------------------------------------------------------------- +# One-shot environment setup for the arrow-rs Linux + S3 read probe. +# +# Brings a fresh box (Linux x86-64, or macOS/arm64 for dev) to the point where you +# can run gen_s3_fixtures.py and run_matrix.py in this directory, by: +# 1. ensuring a Python 3.12 venv (uv-managed), +# 2. installing a Ray nightly wheel matching this branch's base commit + symlinking +# THIS repo's python/ray over it (so the arrow-rs reader source is live) — a +# "latest" wheel drifts from the branch's compiled protobufs and asserts +# "out of sync" at import. Skip with SKIP_RAY=1. +# 3. installing the Rust toolchain (rustup) + maturin, +# 4. building the native crate `ray_data_arrow_rs` into the venv, +# 5. installing probe deps (psutil, numpy, pyarrow, aiohttp, awscli), +# 6. verifying the arrow-rs read path actually engages end to end. +# +# Idempotent: re-running skips work already done. Everything goes into the venv / +# ~/.cargo — nothing touches the system Python. This is the whole fresh-workspace +# recovery: clone the branch, then run this. +# +# Usage (from anywhere in the checkout): +# bash release/nightly_tests/dataset/arrow_rs_probe/setup.sh +# +# Knobs (env vars): +# RAY_VENV= venv to use/create (default: /.venv) +# RAY_WHEEL_URL= Ray nightly wheel to install (default: cp312 linux/mac) +# SKIP_RAY=1 don't touch Ray (already installed + symlinked) +# SKIP_APT=1 don't apt-get build deps (build-essential, python3-dev) +# SKIP_CRATE=1 don't (re)build the Rust crate +# --------------------------------------------------------------------------- +set -euo pipefail + +# --- locate the repo (this script lives at /release/nightly_tests/dataset/arrow_rs_probe) --- +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO="$(cd "$SCRIPT_DIR/../../../.." && pwd)" +CRATE="$REPO/python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs" +RAY_VENV="${RAY_VENV:-$REPO/.venv}" +OS="$(uname -s)"; ARCH="$(uname -m)" + +say() { printf '\n\033[1;36m==> %s\033[0m\n' "$*"; } + +say "repo=$REPO venv=$RAY_VENV os=$OS/$ARCH" + +# --- -1. git remotes: ensure `upstream` exists + is fetched. The wheel pick below +# uses merge-base with upstream/master; on a fresh clone of the fork the remote +# is missing and the script silently falls back to a hardcoded SHA. --- +if ! git -C "$REPO" remote get-url upstream >/dev/null 2>&1; then + say "adding upstream remote (ray-project/ray)" + git -C "$REPO" remote add upstream https://github.com/ray-project/ray.git +fi +say "fetching upstream master (for wheel merge-base)" +git -C "$REPO" fetch --quiet upstream master || say "WARN: upstream fetch failed; using fallback SHA" + +# --- 0. system build deps (Linux only; the crate links libpython + needs a C toolchain) --- +if [ "$OS" = "Linux" ] && [ "${SKIP_APT:-0}" != "1" ] && command -v apt-get >/dev/null 2>&1; then + say "apt: build-essential + python3-dev + curl (sudo)" + sudo apt-get update -qq + sudo apt-get install -y -qq build-essential python3-dev curl pkg-config +fi + +# --- 1. uv + venv --- +if ! command -v uv >/dev/null 2>&1; then + say "installing uv (official installer)" + curl -LsSf https://astral.sh/uv/install.sh | sh + export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH" +fi +if [ ! -x "$RAY_VENV/bin/python" ]; then + say "creating venv at $RAY_VENV (python 3.12)" + uv venv --python 3.12 "$RAY_VENV" +fi +PY="$RAY_VENV/bin/python" +PIP() { uv pip install --python "$PY" "$@"; } +say "python: $($PY --version)" + +# --- 2. Ray nightly + local-source symlink --- +# CRITICAL: the wheel's compiled protobuf must match the branch's Python source, or +# `custom_types.py` asserts "out of sync" at import. setup-dev.py symlinks THIS repo's +# python/ray over the wheel, so we install the per-commit nightly built from the +# branch's base commit (merge-base with upstream/master), NOT "latest" — latest drifts. +BASE_SHA="$(git -C "$REPO" merge-base HEAD upstream/master 2>/dev/null \ + || git -C "$REPO" merge-base HEAD origin/master 2>/dev/null \ + || echo 7dc67bed3ba2f3504325b206a70adcc470422860)" +if [ "${SKIP_RAY:-0}" != "1" ]; then + if [ -z "${RAY_WHEEL_URL:-}" ]; then + if [ "$OS" = "Linux" ]; then + PYTAG=cp312; PLAT=manylinux2014_x86_64 + elif [ "$ARCH" = "arm64" ]; then + PYTAG=cp312; PLAT=macosx_11_0_arm64 + else + PYTAG=cp312; PLAT=macosx_10_15_x86_64 + fi + RAY_WHEEL_URL="https://s3-us-west-2.amazonaws.com/ray-wheels/master/${BASE_SHA}/ray-3.0.0.dev0-${PYTAG}-${PYTAG}-${PLAT}.whl" + fi + say "installing Ray nightly (base commit ${BASE_SHA:0:12}): $RAY_WHEEL_URL" + # Wipe any prior ray install FIRST. A re-run over a setup-dev'd tree has symlinked + # subpackages (ray/workflow -> local source); pip/uv --force-reinstall dies trying + # to rmdir a symlink ("Not a directory"). Removing the dir just unlinks those + # symlinks (never touches the repo source they point to), so a clean install lands + # the commit-matched wheel — and swaps a same-version "latest" wheel too. + SITE="$("$PY" -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')" + rm -rf "$SITE/ray" "$SITE"/ray-*.dist-info "$SITE"/ray_*.dist-info 2>/dev/null || true + PIP "ray[data] @ $RAY_WHEEL_URL" + # Symlink THIS repo's python/ray over the installed wheel so the local + # arrow-rs reader source is what actually runs (mirrors the mac dev setup). + say "symlinking local python/ray via setup-dev.py" + "$PY" "$REPO/python/ray/setup-dev.py" -y +else + say "SKIP_RAY=1 — assuming Ray is installed and python/ray is symlinked" +fi + +# --- 3. Rust toolchain + maturin --- +if ! command -v cargo >/dev/null 2>&1; then + say "installing Rust via rustup (official installer)" + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y +fi +export PATH="$HOME/.cargo/bin:$PATH" +say "rustc: $(rustc --version)" +PIP maturin + +# --- 4. build the native crate into the venv --- +if [ "${SKIP_CRATE:-0}" != "1" ]; then + say "building ray_data_arrow_rs (maturin develop --release) — compiles arrow/parquet, ~2-5 min" + # maturin refuses if BOTH VIRTUAL_ENV and CONDA_PREFIX are set (common when a + # base conda env is active); unset CONDA_PREFIX for this build only. + ( cd "$CRATE" && unset CONDA_PREFIX && VIRTUAL_ENV="$RAY_VENV" "$RAY_VENV/bin/maturin" develop --release ) +else + say "SKIP_CRATE=1 — assuming ray_data_arrow_rs is already built" +fi + +# --- 5. probe Python deps --- +# aiohttp: NOT pulled in by the ray[data] extra, but the runtime-env agent imports +# it; without it the agent crashes and the raylet fate-shares (`ray.init()` hangs +# indefinitely) — cost a day on the 2026-07-27 workspace run. +# awscli: gen_s3_fixtures.py uploads fixtures via `aws s3 sync`. +# grpcio: NOT pulled by ray[data] (only the [default] extra) but imported +# unconditionally by release/nightly_tests/dataset/benchmark.py:12 +# (get_memory_info_reply -> grpc channel to GCS) — the tpch stage runs the +# release scripts through benchmark.py and died on ModuleNotFoundError +# without it. +say "installing probe deps (psutil, numpy, pyarrow, aiohttp, awscli, grpcio)" +PIP psutil numpy pyarrow aiohttp awscli grpcio +# aiohttp failure mode is silent (runtime-env agent crashes -> ray.init hangs +# forever), so verify the import loudly here instead of debugging a hang later. +"$PY" -c "import aiohttp, psutil, numpy, pyarrow" \ + || { say "FATAL: probe dep import failed (aiohttp/psutil/numpy/pyarrow)"; exit 1; } + +# --- 6. verify the arrow-rs path actually engages --- +say "verifying arrow-rs read path end to end" +RAY_ADDRESS=local RAY_DATA_USE_DATASOURCE_V2=1 RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 \ + RAY_task_events_report_interval_ms=0 \ + "$PY" - <<'PYEOF' +import os, tempfile +os.environ.pop("RAY_RUNTIME_ENV_HOOK", None) # Anyscale platform hook not in this venv +os.environ.pop("RAY_RUNTIME_ENV_PLUGINS", None) # platform cgroup plugin crashes the agent +import numpy as np, pyarrow as pa, pyarrow.parquet as pq +import ray_data_arrow_rs # noqa: F401 -> import must succeed (crate built) +import ray + +d = tempfile.mkdtemp() +p = os.path.join(d, "t.parquet") +pq.write_table(pa.table({"a": np.arange(1000), "b": np.arange(1000) * 1.5}), + p, write_page_index=True) +ray.init(address="local", include_dashboard=False, + ignore_reinit_error=True, log_to_driver=False) +ds = ray.data.read_parquet(p) +assert ds.count() == 1000, ds.count() +assert ds.sum("a") == sum(range(1000)) +print("OK ray", ray.__version__, " arrow-rs read path verified (count + sum match)") +ray.shutdown() +PYEOF + +# --- 7. write env.sh: the complete probe environment in one sourceable file, so a +# node restart never means retyping exports. Also hook it into ~/.bashrc +# (idempotent) so fresh shells come up ready. --- +say "writing $SCRIPT_DIR/env.sh + ~/.bashrc hook" +cat > "$SCRIPT_DIR/env.sh" </dev/null; then + echo "[ -f \"$SCRIPT_DIR/env.sh\" ] && source \"$SCRIPT_DIR/env.sh\"" >> "$HOME/.bashrc" + say "added env.sh source line to ~/.bashrc" +fi + +say "DONE." +cat <write) +through it, so each worker executes hundreds of tasks, then reads the floor: + + - a per-PID USS time series (2 Hz, raylet-descendant workers) written to + series.jsonl - the climb curve itself; + - an idle-floor snapshot after each round (settle + gc) - floor(round); + - Ray's own per-task USS (20 Hz in-task poll) per round - the avg/max trend. + +Arms (the allocator ablation): pa / rs / rs_arena2 (MALLOC_ARENA_MAX=2) / +rs_trim (MALLOC_TRIM_THRESHOLD_=0, eager top-of-heap trim on free) / +rs_jemalloc (LD_PRELOAD system libjemalloc - routes the crate's glibc +allocations through the same allocator family PyArrow bundles; skipped with a +warning if no libjemalloc.so is found - `apt install libjemalloc2`). Every env +lever is a deployable fix candidate on its own. Read the verdict as: + + rs floor climbs round-over-round while pa stays flat => the retention is + real and single-node-reproducible (the release "multi-node" difference + was worker lifetime all along); + ...and rs_arena2 collapses to pa => glibc ARENA retention; fixes: arena + cap via runtime_env, or crate-side malloc_trim(0) at end-of-stream + (glibc-only, cfg(target_os="linux")); + ...and only rs_trim collapses => retention at top-of-heap, few arenas + involved; fix = trim (env or crate call); + ...and only rs_jemalloc collapses => glibc-vs-jemalloc policy generally + (fragmentation across arenas trim can't reach); fix = LD_PRELOAD + jemalloc in the workers' runtime_env (cheap, no crate change); + ...and NONE of the rs_* arms collapse => not the C allocator - Rust-side + caching or genuine fragmentation; back to crate profiling; + rs floor flat like pa even at 100s of tasks/worker => the losses need + something only the release cluster has (autoscaling node churn, plasma + pressure, genuine multi-node) - escalate to TODO item 18's both-arms + release trigger. + +Shapes (fixtures + release-yaml bins exactly as loss_triage.py): + + auto M37 read_large_parquet_autoscaling: one ~69 MiB row group per task + (bin 64 MiB), sub-second tasks - the many-small-tasks regime. + write M38 write_parquet: fused read->write_parquet, ~1.2 GiB decode churn + per task (bin 1342177280 on the bin_sweep fixture). + +Usage (Linux box; venv + fixtures via run_soak.sh, or piecemeal): + + python gen_local_fixtures.py --root ~/arrow_rs_repl_fixtures \ + --shapes auto_rg,bin_sweep + python soak_probe.py --fixture-root ~/arrow_rs_repl_fixtures + python soak_probe.py --fixture-root ... --shapes auto --rounds 4 \ + --path-repeat 4 --workers 4 + +Results: /summary.json + per-cell logs + per-cell series.jsonl +(rows of {t, pid, uss_mib} - plot to SEE the climb). The printed table is +R = arrow_rs / pyarrow on the end-of-run idle floor; >1.00 = arrow-rs worse. +""" +import argparse +import gc +import glob +import json +import os +import re +import shutil +import subprocess +import sys +import threading +import time + +PY = sys.executable +HERE = os.path.dirname(os.path.abspath(__file__)) +MiB = 1024 * 1024 + +# Release-yaml bin sizes per shape (release/release_data_tests.yaml), same as +# loss_triage.py. +SHAPE_BINS = { + "auto": 67_108_864, + "write": 1_342_177_280, +} +SHAPE_FIXTURE = { + "auto": "auto_rg", + "write": "bin_sweep", +} +# rounds x path_repeat sized so tasks/worker lands in the release regime +# (O(100) and up for auto; write churns ~1.2 GiB per task so fewer tasks carry +# the same churn volume per worker). +SHAPE_DEFAULTS = { + "auto": dict(rounds=6, path_repeat=8), + # v2 bin_sweep is ~4.6 GiB decoded = ~3.5 write bins per pass; x12 gives + # ~42 tasks/round -> ~60+ tasks/worker over 6 rounds at 1.29 GiB churn each. + "write": dict(rounds=6, path_repeat=12), +} + + +def _median(vals): + vals = sorted(v for v in vals if v is not None) + if not vals: + return None + n = len(vals) + return vals[n // 2] if n % 2 else (vals[n // 2 - 1] + vals[n // 2]) / 2 + + +class SeriesSampler: + """Per-PID USS/RSS time series over OUR raylet's descendant workers. + + read_probe.WorkerMemSampler keeps only the summed peak; here the ORDER is + the signal (the floor climb across a worker's task sequence), so every + sample row is kept and written to series.jsonl on exit. USS needs Linux + (memory_full_info); on macOS rows fall back to RSS and say so. + """ + + def __init__(self, interval_s, root_pid, out_path): + self._interval_s = interval_s + self._root_pid = root_pid + self._out_path = out_path + self._stop = threading.Event() + self._thread = None + self._lock = threading.Lock() + self._rows = [] + self._t0 = time.monotonic() + self._uss_ok = True + self.peak_uss = 0 + self.peak_rss = 0 + + def _procs(self): + import psutil + + if self._root_pid is None: + return [] + try: + children = psutil.Process(self._root_pid).children(recursive=True) + except psutil.NoSuchProcess: + return [] + # The raylet also parents agents (dashboard_agent, runtime_env_agent); + # keep only task workers: proctitle "ray::" / "ray::IDLE", or the + # unretitled default_worker.py. The smoke run showed an agent pid + # polluting the floor median at >1 GiB. + out = [] + for proc in children: + try: + cmd = " ".join(proc.cmdline() or []) + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + if "ray::" in cmd or "default_worker.py" in cmd: + out.append(proc) + return out + + def _sample(self, label=None): + import psutil + + t = round(time.monotonic() - self._t0, 2) + snap = {} + sum_uss = sum_rss = 0 + for proc in self._procs(): + try: + if self._uss_ok: + try: + mi = proc.memory_full_info() + uss = getattr(mi, "uss", 0) + rss = mi.rss + except (psutil.AccessDenied, NotImplementedError): + self._uss_ok = False + uss, rss = 0, proc.memory_info().rss + else: + uss, rss = 0, proc.memory_info().rss + except (psutil.NoSuchProcess, psutil.AccessDenied): + continue + row = {"t": t, "pid": proc.pid, "rss_mib": round(rss / MiB, 1)} + if self._uss_ok: + row["uss_mib"] = round(uss / MiB, 1) + if label: + row["label"] = label + with self._lock: + self._rows.append(row) + snap[proc.pid] = round((uss if self._uss_ok else rss) / MiB, 1) + sum_uss += uss + sum_rss += rss + self.peak_uss = max(self.peak_uss, sum_uss) + self.peak_rss = max(self.peak_rss, sum_rss) + return snap + + def snapshot(self, label): + """One labeled sample right now; returns {pid: uss_mib} (rss on mac).""" + return self._sample(label=label) + + def _run(self): + while not self._stop.wait(self._interval_s): + self._sample() + + def __enter__(self): + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + return self + + def __exit__(self, *exc): + self._stop.set() + if self._thread: + self._thread.join() + with self._lock, open(self._out_path, "w") as fh: + for row in self._rows: + fh.write(json.dumps(row) + "\n") + + +# -------------------------------------------------------------------------- +# The case: one arm = one fresh process = ONE long-lived Ray session. +# -------------------------------------------------------------------------- + + +def run_case(a): + import ray + from ray.data.context import DataContext + + from read_probe import collect_read_op_metrics + + ctx = DataContext.get_current() + ctx.use_datasource_v2 = True + ctx.use_arrow_rs_parquet_reader = a.reader == "rs" + # The release instrument: 20 Hz per-task USS poll (A/B #4 parity). + ctx.memory_usage_poll_interval_s = 0.05 + + # num_cpus pins the worker pool: W workers live for the WHOLE session, so + # tasks/worker = rounds * tasks_per_round / W - the release regime the + # fresh-session cells could never reach. + ray.init(num_cpus=a.workers, ignore_reinit_error=True) + try: + gcs = ray.get_runtime_context().gcs_address + if gcs: + os.environ["RAY_ADDRESS"] = gcs # read_probe's stats-resolution fix + except Exception: + pass + root_pid = None + try: + node = ray._private.worker._global_node # noqa: SLF001 + for name, procs in (node.all_processes or {}).items(): + if "raylet" in name.lower() and procs: + root_pid = procs[0].process.pid + break + except Exception: + pass + + if a.path.startswith("s3://"): + # The S3 leg (user, 2026-08-21: does the M38 release write loss need + # S3 transport?). aws-cli is the lister so the fixture needs no + # local mirror on the driver. + ls = subprocess.run( + ["aws", "s3", "ls", a.path.rstrip("/") + "/"], + capture_output=True, + text=True, + ).stdout + files = sorted( + a.path.rstrip("/") + "/" + ln.split()[-1] + for ln in ls.splitlines() + if ln.strip().endswith(".parquet") + ) + else: + files = sorted(glob.glob(os.path.join(os.path.expanduser(a.path), "*.parquet"))) + if not files: + raise SystemExit(f"no parquet files under {a.path}") + paths = files * a.path_repeat + write_out = a.write_out or os.path.join(a.workdir, f"soak_write_out_{a.tag}") + + def clean_write_out(): + if write_out.startswith("s3://"): + subprocess.run( + ["aws", "s3", "rm", "--recursive", "--quiet", write_out], + capture_output=True, + ) + else: + shutil.rmtree(write_out, ignore_errors=True) + + series_path = os.path.join(a.workdir, f"{a.tag}.series.jsonl") + + rounds = [] + tasks_total = 0 + sampler = SeriesSampler(a.sample_s, root_pid, series_path) + try: + with sampler: + start_floor = sampler.snapshot("start") + for rnd in range(a.rounds): + t0 = time.perf_counter() + ds = ray.data.read_parquet(paths) + if a.shape == "write": + clean_write_out() + ds.write_parquet(write_out) + clean_write_out() + else: + # capture_executor=True so per-task USS survives (read_probe's + # GE1 snapshot-race fix); bundles are dropped as they stream. + bundle_iter, _, _ = ds._execute_to_iterator(capture_executor=True) + for _ in bundle_iter: + pass + wall = time.perf_counter() - t0 + m = collect_read_op_metrics(ds) + m.pop("uss_debug", None) + del ds + gc.collect() + time.sleep(a.settle_s) # let workers go idle before reading floors + floor = sampler.snapshot(f"after_round_{rnd}") + ntasks = m.get("read_num_tasks") or 0 + tasks_total += ntasks + if rnd == 0 and ntasks < 2 * a.workers: + print( + f" !! only {ntasks} read tasks/round for {a.workers} " + "workers - this is NOT soaking. Check the fixture is v2 " + "(gen_local_fixtures M41 fix) and raise --path-repeat.", + flush=True, + ) + to_mib = lambda gb: round(gb * 1024, 1) if gb else None # noqa: E731 + rounds.append( + dict( + round=rnd, + wall_s=round(wall, 2), + num_tasks=ntasks, + task_uss_avg_mib=to_mib(m.get("read_avg_max_uss_gb")), + task_uss_max_mib=to_mib(m.get("read_max_uss_gb")), + idle_floor_mib=_median(list(floor.values())), + workers_live=len(floor), + idle_floor_by_pid={str(k): v for k, v in sorted(floor.items())}, + ) + ) + print(f" round {rnd}: {rounds[-1]}", flush=True) + finally: + clean_write_out() + import ray as _ray + + _ray.shutdown() + + floors = [r["idle_floor_mib"] for r in rounds if r["idle_floor_mib"] is not None] + stable_pids = None + if len(rounds) > 1: + first = set(rounds[0]["idle_floor_by_pid"]) + last = set(rounds[-1]["idle_floor_by_pid"]) + stable_pids = len(first & last) + result = dict( + shape=a.shape, + reader=a.reader, + workers=a.workers, + rounds=a.rounds, + tasks_total=tasks_total, + tasks_per_worker=round(tasks_total / a.workers, 1) if a.workers else None, + stable_worker_pids=stable_pids, + start_floor_mib=_median(list(start_floor.values())), + first_round_floor_mib=floors[0] if floors else None, + end_floor_mib=floors[-1] if floors else None, + floor_climb_mib=(round(floors[-1] - floors[0], 1) if len(floors) > 1 else None), + peak_uss_gb=round(sampler.peak_uss / 1024**3, 3) + if sampler.peak_uss + else None, + peak_rss_gb=round(sampler.peak_rss / 1024**3, 3), + task_uss_avg_first_mib=rounds[0]["task_uss_avg_mib"] if rounds else None, + task_uss_avg_last_mib=rounds[-1]["task_uss_avg_mib"] if rounds else None, + task_uss_max_last_mib=rounds[-1]["task_uss_max_mib"] if rounds else None, + rounds_detail=rounds, + series_jsonl=series_path, + ) + print("=== CASE RESULT ===") + print(json.dumps(result)) + + +# -------------------------------------------------------------------------- +# The matrix: shapes x arms, one subprocess (= one long session) per cell. +# -------------------------------------------------------------------------- + + +def run_cell(logdir, name, shape, reader, path, env_extra, a, write_out=None): + defaults = SHAPE_DEFAULTS[shape] + cmd = [ + PY, + os.path.abspath(__file__), + "case", + "--shape", + shape, + "--reader", + reader, + "--path", + path, + "--workers", + str(a.workers), + "--rounds", + str(a.rounds or defaults["rounds"]), + "--path-repeat", + str(a.path_repeat or defaults["path_repeat"]), + "--settle-s", + str(a.settle_s), + "--sample-s", + str(a.sample_s), + "--workdir", + logdir, + "--tag", + name, + ] + if write_out: + cmd += ["--write-out", write_out] + env = dict(os.environ) + env["RAY_DATA_PARQUET_BIN_PACKING_BYTES"] = str(SHAPE_BINS[shape]) + env.update(env_extra) + print(f" -> {name} (env_extra={env_extra})", flush=True) + t0 = time.perf_counter() + proc = subprocess.run(cmd, capture_output=True, text=True, env=env) + with open(os.path.join(logdir, f"{name}.log"), "w") as fh: + fh.write(f"# cmd: {' '.join(cmd)}\n# env_extra: {env_extra}\n") + fh.write(f"# wall_including_startup_s: {time.perf_counter() - t0:.1f}\n") + fh.write("# ---- STDOUT ----\n" + proc.stdout) + fh.write("\n# ---- STDERR ----\n" + proc.stderr) + for i, line in enumerate(proc.stdout.splitlines()): + if "=== CASE RESULT ===" in line: + try: + res = json.loads(proc.stdout.splitlines()[i + 1]) + print( + f" floor {res.get('first_round_floor_mib')} -> " + f"{res.get('end_floor_mib')} MiB (climb " + f"{res.get('floor_climb_mib')}), " + f"{res.get('tasks_per_worker')} tasks/worker", + flush=True, + ) + return res + except (IndexError, json.JSONDecodeError): + break + print( + f" !! {name} CASE FAIL rc={proc.returncode} (see {logdir}/{name}.log)\n" + f" {proc.stderr.strip()[-400:]}", + flush=True, + ) + return {} + + +def _find_jemalloc(): + """Locate a system libjemalloc for the rs_jemalloc LD_PRELOAD arm. + + JEMALLOC_PATH env overrides; otherwise probe the usual Linux locations. + Returns None (arm skipped) when not found or on macOS (LD_PRELOAD n/a). + """ + if not sys.platform.startswith("linux"): + return None + cand = os.environ.get("JEMALLOC_PATH") + if cand and os.path.exists(cand): + return cand + import glob as _glob + + for pat in ( + "/usr/lib/x86_64-linux-gnu/libjemalloc.so*", + "/usr/lib/aarch64-linux-gnu/libjemalloc.so*", + "/usr/lib64/libjemalloc.so*", + "/usr/local/lib/libjemalloc.so*", + ): + hits = sorted(_glob.glob(pat)) + if hits: + return hits[0] + return None + + +def _ratio(a, b): + return round(a / b, 2) if a and b else None + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + sub = p.add_subparsers(dest="cmd") + + c = sub.add_parser("case", help="internal: one arm, one long Ray session") + c.add_argument("--shape", choices=list(SHAPE_BINS), required=True) + c.add_argument("--reader", choices=["pa", "rs"], required=True) + c.add_argument("--path", required=True) + c.add_argument("--workers", type=int, default=4) + c.add_argument("--rounds", type=int, default=6) + c.add_argument("--path-repeat", type=int, default=8) + c.add_argument("--settle-s", type=float, default=3.0) + c.add_argument("--sample-s", type=float, default=0.5) + c.add_argument("--workdir", default=".") + c.add_argument("--tag", default="case") + c.add_argument( + "--write-out", + default=None, + help="write-shape output root (s3://... allowed); default /...", + ) + + p.add_argument("--fixture-root", default=None) + p.add_argument("--outdir", default=None) + p.add_argument("--shapes", default="auto,write") + p.add_argument("--arms", default="pa,rs,rs_arena2,rs_trim,rs_jemalloc") + p.add_argument("--workers", type=int, default=4) + p.add_argument( + "--rounds", type=int, default=None, help="override per-shape default" + ) + p.add_argument( + "--path-repeat", type=int, default=None, help="override per-shape default" + ) + p.add_argument("--settle-s", type=float, default=3.0) + p.add_argument("--sample-s", type=float, default=0.5) + p.add_argument( + "--transports", + default="local", + help="comma list of local,s3 — s3 soaks read from AND write to the bucket", + ) + p.add_argument("--s3-bucket", default=os.environ.get("ARROW_RS_S3_BUCKET")) + args = p.parse_args() + + if args.cmd == "case": + run_case(args) + return + + if not args.fixture_root: + p.error("--fixture-root is required (see gen_local_fixtures.py)") + fixture_root = os.path.expanduser(args.fixture_root) + with open(os.path.join(fixture_root, "manifest.json")) as fh: + manifest = json.load(fh) + + shapes = [s.strip() for s in args.shapes.split(",") if s.strip()] + arm_defs = { + "pa": ("pa", {}), + "rs": ("rs", {}), + "rs_arena2": ("rs", {"MALLOC_ARENA_MAX": "2"}), + # Eager glibc trim: return top-of-heap to the OS on every free that + # leaves >=0 bytes free above the break (also freezes the dynamic + # mmap/trim threshold adjustment - fine for an ablation arm). + "rs_trim": ("rs", {"MALLOC_TRIM_THRESHOLD_": "0"}), + } + jemalloc = _find_jemalloc() + if jemalloc: + arm_defs["rs_jemalloc"] = ("rs", {"LD_PRELOAD": jemalloc}) + arms = [a.strip() for a in args.arms.split(",") if a.strip()] + # rs_trim: glibc trim threshold at N MiB (rs_trim alone = 0 = every free). + # Sweeping N asks how much of rs_trim's floor collapse survives at a + # threshold cheap enough to leave wall time intact. + for a in arms: + m = re.fullmatch(r"rs_trim(\d+)", a) + if m and a not in arm_defs: + arm_defs[a] = ("rs", {"MALLOC_TRIM_THRESHOLD_": str(int(m.group(1)) * MiB)}) + if "rs_jemalloc" in arms and not jemalloc: + print( + "WARNING: rs_jemalloc arm skipped - no libjemalloc found " + "(set JEMALLOC_PATH or `apt install libjemalloc2`)", + flush=True, + ) + arms = [a for a in arms if a != "rs_jemalloc"] + + outdir = args.outdir or os.path.join( + HERE, "soak_runs", time.strftime("%Y%m%d_%H%M%S") + ) + os.makedirs(outdir, exist_ok=True) + print(f"shapes={shapes} arms={arms} workers={args.workers} outdir={outdir}") + + transports = [t.strip() for t in args.transports.split(",") if t.strip()] + if "s3" in transports and not args.s3_bucket: + p.error("--transports s3 needs --s3-bucket or ARROW_RS_S3_BUCKET") + + summary = {"shapes": shapes, "arms": arms, "workers": args.workers, "cells": {}} + for shape in shapes: + entry = manifest[SHAPE_FIXTURE[shape]] + local_path = entry["path"] if isinstance(entry, dict) else entry + for transport in transports: + if transport == "s3": + from loss_triage import s3_sync + + bucket = args.s3_bucket.rstrip("/") + path = f"{bucket}/soak/{SHAPE_FIXTURE[shape]}" + s3_sync(local_path, path) + write_out = f"{bucket}/soak_write_out/{shape}" + else: + path, write_out = local_path, None + print( + f"\n=== [{shape}] soak/{transport} ({SHAPE_FIXTURE[shape]}) ===", + flush=True, + ) + cells = {} + for tag in arms: + reader, env_extra = arm_defs[tag] + cells[tag] = run_cell( + outdir, + f"{shape}.{transport}.{tag}", + shape, + reader, + path, + env_extra, + args, + write_out=write_out, + ) + summary["cells"][f"{shape}.{transport}"] = cells + + print("\n\n============ SOAK SUMMARY (R = arrow_rs/pyarrow) ============") + if sys.platform == "darwin": + print("(macOS: USS unavailable - rows are RSS; smoke run only)") + header = ( + f"{'cell':<20} {'end floor MiB':>14} {'climb MiB':>10} " + f"{'tUSS avg last':>14} {'floor R':>8}" + ) + print(header) + print("-" * len(header)) + for shape, cells in summary["cells"].items(): + pa = cells.get("pa") or {} + for tag in arms: + r = cells.get(tag) or {} + floor_r = ( + _ratio(r.get("end_floor_mib"), pa.get("end_floor_mib")) + if tag != "pa" + else None + ) + fmt = lambda v: f"{v}" if v is not None else "-" # noqa: E731 + print( + f"{shape + '.' + tag:<20} {fmt(r.get('end_floor_mib')):>14} " + f"{fmt(r.get('floor_climb_mib')):>10} " + f"{fmt(r.get('task_uss_avg_last_mib')):>14} " + f"{fmt(floor_r):>8}" + ) + print( + "\nRead it as: rs floor climbs while pa stays flat => retention reproduced\n" + "single-node (worker lifetime was the release variable). Whichever rs_* arm\n" + "collapses floor R to ~pa names the mechanism AND the fix (arena cap / trim /\n" + "LD_PRELOAD jemalloc - see the module docstring verdict table);\n" + "rs flat like pa => release-cluster-only (escalate to TODO item 18).\n" + "Climb curves: plot each cell's series.jsonl. Full metrics: " + + os.path.join(outdir, "summary.json") + ) + with open(os.path.join(outdir, "summary.json"), "w") as fh: + json.dump(summary, fh, indent=2) + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/tensors_nbytes_probe.py b/release/nightly_tests/dataset/arrow_rs_probe/tensors_nbytes_probe.py new file mode 100644 index 000000000000..8e0aa649e5c8 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/tensors_nbytes_probe.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +"""M39: is wide_schema tensors' bigger peak batch representation or batch sizing? + +A/B #4 measured, at byte-identical per-task decoded bytes (ratio 1.00), an +arrow-rs largest-yielded-table of 288 MiB vs PyArrow's 236 (peak_batch 1.22x, +findings M39) on the 5000-tensor-column shape - the only loss whose decoder +distributions did NOT match. Two non-defect explanations, distinguishable +standalone with no Ray: + + (a) batch SIZING - the reader's floor (min 2048 rows) or its batch math makes + arrow-rs yield fewer, bigger tables: same bytes/row, more rows per batch; + (b) REPRESENTATION - the realigned tables genuinely carry more bytes per row + (offset widths, non-shared validity buffers, per-chunk overhead), which + would also inflate every downstream block and could explain the shape's + worker-USS loss (M33) where retention doesn't. + +Decodes the tensors_cp fixture through both readers exactly as loss_triage.py +does (same computed batch_size for both; the rs leg realigns to the extension +schema via the reader's own `_cast_table_to`, so tables are what do_read would +see), then compares rows/batch and bytes/row per batch plus a per-arrow-type +breakdown of the first batch. Equal bytes/row with bigger rs batches => (a), +close M39 as benign sizing. A bytes/row gap => (b), and the type table says +which column class carries it. + +Representation is platform-independent - run anywhere the crate imports: + + python tensors_nbytes_probe.py --fixture-root ~/arrow_rs_repl_fixtures +""" +import argparse +import json +import os +from collections import defaultdict + +# Must be set before ray.data is imported (loss_triage SHAPE_ENV, same rule). +os.environ.setdefault("RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA", "1") + +MiB = 1024 * 1024 + + +def _batches(reader, path, batch_size, realign_fields): + from loss_triage import _pa_batches, _reader_knobs, _rs_batches + + if reader == "pa": + return _pa_batches(path, batch_size) + return _rs_batches(path, batch_size, _reader_knobs(), realign_fields) + + +def _describe(tables): + """Per-batch (rows, nbytes) + per-type nbytes breakdown of the 1st batch.""" + batches = [] + by_type = None + for t in tables: + batches.append((t.num_rows, t.nbytes)) + if by_type is None: + by_type = defaultdict(lambda: [0, 0]) # type -> [n_cols, nbytes] + for name, col in zip(t.schema.names, t.columns): + key = str(t.schema.field(name).type) + by_type[key][0] += 1 + by_type[key][1] += col.nbytes + return batches, dict(sorted(by_type.items())) if by_type else {} + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--fixture-root", required=True) + p.add_argument("--max-batches", type=int, default=None) + args = p.parse_args() + + import pyarrow.parquet as pq + import ray.data # noqa: F401 (registers the tensor extension types) + + from loss_triage import _batch_size, _reader_knobs + + root = os.path.expanduser(args.fixture_root) + with open(os.path.join(root, "manifest.json")) as fh: + entry = json.load(fh)["tensors_cp"] + fdir = entry["path"] if isinstance(entry, dict) else entry + path = sorted( + os.path.join(fdir, f) for f in os.listdir(fdir) if f.endswith(".parquet") + )[0] + + md = pq.read_metadata(path) + bs = _batch_size(md, _reader_knobs()) + realign = list(pq.read_schema(path)) + print(f"file={path} rows={md.num_rows} computed batch_size={bs}\n") + + out = {} + for reader in ("pa", "rs"): + it = _batches(reader, path, bs, realign if reader == "rs" else None) + if args.max_batches: + it = (t for i, t in enumerate(it) if i < args.max_batches) + batches, by_type = _describe(it) + rows = sum(r for r, _ in batches) + nbytes = sum(b for _, b in batches) + out[reader] = dict(batches=batches, by_type=by_type, rows=rows, nbytes=nbytes) + print(f"--- {reader} ---") + print( + f" batches: {len(batches)} total rows: {rows} total MiB: {nbytes / MiB:.1f}" + ) + print(f" bytes/row overall: {nbytes / rows:.1f}") + print( + f" per-batch (rows, MiB): {[(r, round(b / MiB, 1)) for r, b in batches]}" + ) + print(f" max batch MiB: {max(b for _, b in batches) / MiB:.1f}\n") + + pa_r, rs_r = out["pa"], out["rs"] + bpr_pa = pa_r["nbytes"] / pa_r["rows"] + bpr_rs = rs_r["nbytes"] / rs_r["rows"] + print("=== VERDICT INPUTS ===") + print(f"bytes/row pa={bpr_pa:.1f} rs={bpr_rs:.1f} R={bpr_rs / bpr_pa:.3f}") + print( + f"max batch MiB pa={max(b for _, b in pa_r['batches']) / MiB:.1f} " + f"rs={max(b for _, b in rs_r['batches']) / MiB:.1f}" + ) + print( + "R~1.00 bytes/row with bigger rs batches => M39 is batch SIZING (benign);\n" + "R>1.05 => REPRESENTATION - see the per-type rows below for the carrier." + ) + print("\n=== PER-TYPE (first batch; type: n_cols, MiB, R) ===") + keys = sorted(set(pa_r["by_type"]) | set(rs_r["by_type"])) + for k in keys: + pn, pb = pa_r["by_type"].get(k, (0, 0)) + rn, rb = rs_r["by_type"].get(k, (0, 0)) + r = round(rb / pb, 3) if pb else None + print( + f" {k[:70]:<70} pa=({pn}, {pb / MiB:.1f}) rs=({rn}, {rb / MiB:.1f}) R={r}" + ) + + print("\n=== RESULT ===") + print( + json.dumps( + dict( + bytes_per_row_pa=round(bpr_pa, 1), + bytes_per_row_rs=round(bpr_rs, 1), + bytes_per_row_ratio=round(bpr_rs / bpr_pa, 4), + max_batch_mib_pa=round(max(b for _, b in pa_r["batches"]) / MiB, 1), + max_batch_mib_rs=round(max(b for _, b in rs_r["batches"]) / MiB, 1), + n_batches_pa=len(pa_r["batches"]), + n_batches_rs=len(rs_r["batches"]), + ) + ) + ) + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/tpch_probe.py b/release/nightly_tests/dataset/arrow_rs_probe/tpch_probe.py new file mode 100644 index 000000000000..7eaf538b73fd --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/tpch_probe.py @@ -0,0 +1,284 @@ +"""The two suspect release TPC-H queries, A/B'd on one box (M46 / T27). + +A/B #4's only real tpch signals (findings M46, T27; everything else is +symmetric autoscaling noise — group medians 0.99-1.00): + + q9 the T-only spiller: arrow-rs spilled 3.9 GB (A/B #4) / 3.4 GB (A/B #3) + on autoscaling while PyArrow spilled 0 — wall 1.65x/1.53x follow the + spill. Multi-join (lineitem x part x supplier x partsupp x orders). + q20 the ONLY fixed_size-replicated tpch wall loss, and it is + hash_shuffle_v2-only: fv2 1.18 / av2 1.15 vs 1.01 on both v1 variants. + Semi-join heavy. Suspect: reader output block granularity interacting + with hash_shuffle_v2's partitioning. + +This probe runs the RELEASE tpch scripts themselves (release/nightly_tests/ +dataset/tpch/tpch_q*.py — same code, same public bucket, smaller --sf) in a +fresh process per cell, over the matrix + + queries x shuffle strategies (hash_shuffle, hash_shuffle_v2) + x readers (RAY_DATA_USE_ARROW_RS_PARQUET_READER=0/1) + +and reports wall + spilled_gb per cell, R per (query, strategy). What it can +and cannot settle: a box-visible q20 gap that follows hash_shuffle_v2 under +the rs reader = the block-granularity suspect is real and local; no gap = the +loss needs the release regime (autoscaling cluster / sf1000) — fold into TODO +items 19/20. q9's spill is object-store pressure, so a single node with a +default object store may not reproduce it; a T-vs-B spill *difference* here +would still be signal. + +Usage: + python tpch_probe.py --outdir DIR [--sf 10] [--repeat 1] + [--queries tpch_q9,tpch_q20] [--strategies hash_shuffle,hash_shuffle_v2] + [--dry-run] +Needs AWS credentials (public bucket s3://ray-benchmark-data/tpch/parquet). +""" +import argparse +import json +import os +import signal +import subprocess +import sys +import time + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) +from release_regression_probe import ARMS, arm_env # noqa: E402 same dir + +DATASET_DIR = os.path.abspath(os.path.join(HERE, "..")) +TPCH_DIR = os.path.join(DATASET_DIR, "tpch") + +# One cell, fresh process: import the release query module, time main(), then +# read this session's spill total. ray.init happens here (the scripts only +# init under __main__). +SNIPPET = r""" +import importlib, json, re, sys, time +from types import SimpleNamespace + +query, sf, dry = sys.argv[1], int(sys.argv[2]), sys.argv[3] == "1" +sched_mem_gb = int(sys.argv[4]) +mod = importlib.import_module(query) +if dry: + print("CELL_JSON " + json.dumps({"dry_run": True})) + raise SystemExit(0) +import ray + +# _memory inflates only the SCHEDULING memory resource (nothing is allocated). +# Without it, one 8-core/30GB box deadlocks on multi-join hash_shuffle queries: +# two JoinOperators each reserve num_partitions x ~450MB of the ~14.4GiB budget +# and upstream shuffle tasks starve forever (seen on q2 sf10, PyArrow arm too). +ray.init(address="local", **({"_memory": sched_mem_gb << 30} if sched_mem_gb else {})) +# The workspace's own Ray (:6379) coexists with this cell's local instance, and +# the state API's address autodetection dies on "multiple active Ray instances" +# — which silently emptied every per-task stats dist. Pin it to this cell. +import os + +os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address +t0 = time.monotonic() +mod.main(SimpleNamespace(sf=sf)) +wall = time.monotonic() - t0 +spilled_gb = None +try: + import ray._private.internal_api as api + + m = re.search(r"Spilled (\d+) MiB", api.memory_summary(stats_only=True)) + spilled_gb = round(int(m.group(1)) / 1024, 3) if m else 0.0 +except Exception: + pass +print("CELL_JSON " + json.dumps({"wall_s": round(wall, 1), "spilled_gb": spilled_gb})) +""" + + +def run_cell( + query, + strategy, + reader, + sf, + outdir, + dry_run, + timeout_s, + sched_mem_gb, + monitor_interval=1.0, +): + tag = f"{query}.{strategy}.{reader}" + env = dict(os.environ) + env["PYTHONPATH"] = ( + TPCH_DIR + os.pathsep + DATASET_DIR + os.pathsep + env.get("PYTHONPATH", "") + ) + env["RAY_DATA_DEFAULT_SHUFFLE_STRATEGY"] = strategy + env["RAY_DATA_BENCH_NODE_MEM_MONITOR"] = "1" + # Anyscale pins RAY_OVERRIDE_RESOURCES (memory=14.4GiB on this box) and it + # beats ray.init(_memory=...): that budget deadlocks multi-join hash_shuffle + # cells (two JoinOperators' aggregator reservations consume all of it). + # Rewrite just the memory field for the cell; scheduling-only, not allocated. + if sched_mem_gb and env.get("RAY_OVERRIDE_RESOURCES"): + ovr = json.loads(env["RAY_OVERRIDE_RESOURCES"]) + ovr["memory"] = sched_mem_gb << 30 + env["RAY_OVERRIDE_RESOURCES"] = json.dumps(ovr) + arm_env(env, reader) + env["RAY_DATA_BENCH_NODE_MEM_INTERVAL"] = str(monitor_interval) + env["TEST_OUTPUT_JSON"] = os.path.join(outdir, f"{tag}.benchmark.json") + cmd = [ + sys.executable, + "-c", + SNIPPET, + query, + str(sf), + "1" if dry_run else "0", + str(sched_mem_gb), + ] + log_path = os.path.join(outdir, f"{tag}.log") + if not dry_run: + # A q9 cell at sf10 runs for MINUTES; stream the query's own output to + # the log live (stdout+stderr interleaved) so `tail -f` shows progress — + # the buffered version looked like a hang. + print(f" -> {tag} running (tail -f {log_path})", flush=True) + t0 = time.perf_counter() + timed_out = False + with open(log_path, "w") as fh: + fh.write(f"# strategy={strategy} reader={reader} sf={sf}\n") + fh.flush() + # start_new_session puts the cell + its local Ray daemons in one process + # group, so a timeout can reap raylet/gcs/workers too (a bare kill of the + # driver leaks them, and idle workers hold RSS the next cells then lack). + proc = subprocess.Popen( + cmd, env=env, stdout=fh, stderr=subprocess.STDOUT, start_new_session=True + ) + try: + proc.wait(timeout=None if dry_run else timeout_s) + except subprocess.TimeoutExpired: + timed_out = True + os.killpg(proc.pid, signal.SIGTERM) + try: + proc.wait(timeout=30) + except subprocess.TimeoutExpired: + pass + try: + os.killpg(proc.pid, signal.SIGKILL) + except ProcessLookupError: + pass + proc.wait() + with open(log_path) as fh: + out = fh.read() + if timed_out: + print(f" !! {tag} TIMEOUT after {timeout_s}s (see {tag}.log)", flush=True) + return {"timeout_s": timeout_s} + line = next((ln for ln in out.splitlines() if ln.startswith("CELL_JSON ")), None) + if line is None: + print(f" !! {tag} FAILED rc={proc.returncode} (see {tag}.log)", flush=True) + print(" " + out.strip()[-400:], flush=True) + return None + rec = json.loads(line[len("CELL_JSON ") :]) + rec["wall_incl_startup_s"] = round(time.perf_counter() - t0, 1) + print(f" {tag:<40} {rec}", flush=True) + # Fold in the query's own benchmark.py result (per-operator wall/cpu, + # per-task dists, node-mem monitor) like release_regression_probe does; + # summary.json otherwise carried walls only (2026-09-04 q17 leg). + try: + with open(env["TEST_OUTPUT_JSON"]) as fh: + rec["bench"] = next(iter(json.load(fh).values())) + except Exception: + rec["bench"] = {} + return rec + + +def main(): + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--outdir", required=True) + p.add_argument("--sf", type=int, default=10) + p.add_argument("--repeat", type=int, default=1) + p.add_argument("--queries", default="tpch_q9,tpch_q20") + p.add_argument("--strategies", default="hash_shuffle,hash_shuffle_v2") + p.add_argument( + "--cell-timeout", + type=int, + default=int(os.environ.get("TPCH_CELL_TIMEOUT", "1200")), + help="kill a cell's whole process group after this many seconds", + ) + p.add_argument( + "--sched-mem-gb", + type=int, + default=int(os.environ.get("PROBE_SCHED_MEM_GB", "64")), + help="scheduling-only memory resource for the local Ray instance " + "(0 = stock; stock deadlocks multi-join hash_shuffle cells on one box)", + ) + p.add_argument( + "--dry-run", + action="store_true", + help="import-and-exit per cell: validates module/env plumbing offline", + ) + p.add_argument( + "--arms", + default="pa,rs,rseos", + help=f"comma list from {sorted(ARMS)}; pa is the denominator", + ) + p.add_argument( + "--monitor-interval", + type=float, + default=float(os.environ.get("RAY_DATA_BENCH_NODE_MEM_INTERVAL", "1.0")), + help="node-memory sampler period in seconds (release: 1.0; 0.1 = 10 Hz " + "for the short q6 sustained-wUSS rows)", + ) + a = p.parse_args() + os.makedirs(a.outdir, exist_ok=True) + arms = [s.strip() for s in a.arms.split(",") if s.strip()] + for arm in arms: + arm_env({}, arm) # validate names up front + + results = {} + for query in a.queries.split(","): + for strategy in a.strategies.split(","): + for reader in arms: + runs = [ + run_cell( + query, + strategy, + reader, + a.sf, + a.outdir, + a.dry_run, + a.cell_timeout, + a.sched_mem_gb, + a.monitor_interval, + ) + for _ in range(a.repeat) + ] + good = sorted( + (r for r in runs if r and "wall_s" in r), + key=lambda r: r["wall_s"], + ) + results[f"{query}.{strategy}.{reader}"] = ( + good[len(good) // 2] if good else (runs[0] if runs else None) + ) + + with open(os.path.join(a.outdir, "summary.json"), "w") as fh: + json.dump(results, fh, indent=2) + if a.dry_run: + print("\ndry run OK — all query modules import") + return + + print("\n================ TPCH PROBE (R = arm/pyarrow) ================") + print( + f"{'cell [arm]':<42} {'wall pa':>8} {'wall arm':>8} {'R':>6} " + f"{'spill pa/arm GB':>15}" + ) + for query in a.queries.split(","): + for strategy in a.strategies.split(","): + pa_r = results.get(f"{query}.{strategy}.pa") or {} + for arm in (m for m in arms if m != "pa"): + rs_r = results.get(f"{query}.{strategy}.{arm}") or {} + wp, wr = pa_r.get("wall_s"), rs_r.get("wall_s") + ratio = f"{wr / wp:.2f}" if wp and wr else "—" + spill = f"{pa_r.get('spilled_gb')}/{rs_r.get('spilled_gb')}" + cell = f"{query}.{strategy} [{arm}]" + print( + f"{cell:<42} {wp or '—':>8} {wr or '—':>8} {ratio:>6} {spill:>15}" + ) + print( + "\nRead it as: q20 gap only under hash_shuffle_v2+rs => block-granularity" + "\nsuspect confirmed locally (M46); no gap => release-regime-only, fold into" + "\nitems 19/20. Any rs-only spill on q9 = T27 reproduced." + ) + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/arrow_rs_probe/worker_constant_probe.py b/release/nightly_tests/dataset/arrow_rs_probe/worker_constant_probe.py new file mode 100644 index 000000000000..4d61c07cd0b2 --- /dev/null +++ b/release/nightly_tests/dataset/arrow_rs_probe/worker_constant_probe.py @@ -0,0 +1,291 @@ +"""worker_constant_probe.py — TODO 31: measure the per-worker resident constant. + +A/B #3-#5 showed arrow-rs read workers carrying a ~90-105 MB higher resident +floor than pyarrow workers (findings M84), which multiplies into the +sustained-wUSS tripper cluster (M77) at high workers-per-node. This probe +measures it directly instead of inferring it from release aggregates: + + - one cell = a fresh subprocess (reader pinned by env BEFORE the first + ray/ray.data import — the DataContext singleton fix from the cliff probe) + starting a fresh 2-CPU local Ray cluster (2, not 1: the V2 listing/footer + actor pool must not starve the read task) reading ONE file at + concurrency=1, so one worker does the read and pre-existing + (workspace/platform) ray processes are excluded by a before-init /proc + census; the sampler keeps DISCOVERING new ray:: workers while it runs + (Ray may fork a fresh worker for the read instead of reusing a prestarted + one — a fixed-pid census would then report an idle worker's numbers as + the read worker's); the read worker is identified POSITIVELY by + proctitle — Ray retitles an executing worker `ray::`, and the + sampler records every title a pid ever shows — because USS growth cannot + identify a worker first seen mid-read (its first-seen baseline already + holds decode buffers, so its settled delta reads ~0 and loses to an idle + pid's noise); growth-based selection remains only as a flagged fallback; + - USS (Private_Clean + Private_Dirty from /proc//smaps_rollup) of that + worker is recorded idle-after-spawn, at 20 Hz through two consecutive + reads (peak), and settled after each read — read #2 separates a one-time + constant (after1 == after2) from per-task growth (retention slope); + - the read is S3 (one file, concurrency=1) so the arrow-rs arm pays its full + release-path setup: crate .so mapping, shared tokio runtime + blocking + pool, object_store client, allocator state. The pa arm is the control; + the rs-pa delta on `after` IS the constant. + +Usage (box): + python worker_constant_probe.py --path s3:///cliff_probe/.parquet +Requires AWS creds in the env (source env.sh). macOS: /proc absent — Linux only. +""" + +import argparse +import json +import os +import subprocess +import sys +import threading +import time + +FLAG = "RAY_DATA_USE_ARROW_RS_PARQUET_READER" + + +def _proc_pids(): + return [int(d) for d in os.listdir("/proc") if d.isdigit()] + + +def _cmdline(pid): + try: + with open(f"/proc/{pid}/cmdline", "rb") as f: + return f.read().replace(b"\0", b" ").decode(errors="replace") + except OSError: + return "" + + +def _uss_kb(pid): + """USS in KiB from smaps_rollup (Private_Clean + Private_Dirty).""" + try: + total = 0 + with open(f"/proc/{pid}/smaps_rollup") as f: + for line in f: + if line.startswith(("Private_Clean:", "Private_Dirty:")): + total += int(line.split()[1]) + return total + except OSError: + return None + + +# Rescan /proc for newly forked workers every N 50 ms ticks (0.5 s: worker +# fork-to-first-task is slower than that, and full-scan cmdline reads at 20 Hz +# would be needless load). +_DISCOVER_EVERY_TICKS = 10 + +# V2 read tasks execute under this proctitle (ray::ReadFilesParquetV2, both +# arms — the reader flag branches inside the task). A pid that ever bore it is +# the read worker, positively; no match falls back to growth selection. +_READ_TASK_TITLE_SUBSTR = "ReadFiles" + + +class _WorkerSampler(threading.Thread): + """20 Hz USS sampler over ray:: workers that discovers new ones as they spawn. + + The read task is not guaranteed to land on a prestarted worker: Ray may + fork one after any one-shot census, and a fixed-pid sampler would then + attribute the read to an idle worker. Each newly seen pid gets its + first-seen USS as baseline; pids first seen after sampling started are + flagged `late` (their baseline already includes whatever the read did + before discovery, ≤0.5 s in). + """ + + def __init__(self, exclude_pids): + super().__init__(daemon=True) + self.exclude = set(exclude_pids) + self.baseline = {} # pid -> first-seen USS KiB + self.peak = {} # pid -> peak USS KiB + self.titles = {} # pid -> set of ray:: cmdlines ever observed + self.late = set() # pids first seen after start() + # NOT named _stop: threading.Thread.join() calls its own internal + # self._stop() method, which an Event attribute would shadow. + self._halt = threading.Event() + + def census(self, late): + for p in _proc_pids(): + if p in self.exclude: + continue + cmd = _cmdline(p).strip() + if not cmd.startswith("ray::"): + continue + if p not in self.baseline: + v = _uss_kb(p) + if v is None: + continue + self.baseline[p] = v + self.peak[p] = v + if late: + self.late.add(p) + # A worker is retitled ray:: while executing, so the + # accumulated title set positively identifies the read worker. + self.titles.setdefault(p, set()).add(cmd) + + def snapshot(self): + return {p: _uss_kb(p) for p in list(self.baseline)} + + def run(self): + tick = 0 + while not self._halt.is_set(): + if tick % _DISCOVER_EVERY_TICKS == 0: + self.census(late=True) + for p in list(self.peak): + v = _uss_kb(p) + if v is not None and v > self.peak[p]: + self.peak[p] = v + tick += 1 + time.sleep(0.05) + + def stop(self): + self._halt.set() + self.join(timeout=2) + + +def run_cell_body(arm, path): + want_rs = arm == "rs" + os.environ[FLAG] = "1" if want_rs else "0" + + before_init = set(_proc_pids()) + + import ray + + # address="local" forces a NEW cluster from THIS venv (plain init would + # auto-discover a workspace raylet via /tmp/ray/ray_current_cluster). + ray.init( + address="local", num_cpus=2, include_dashboard=False, logging_level="ERROR" + ) + import ray.data + + ctx = ray.data.DataContext.get_current() + ctx.use_arrow_rs_parquet_reader = want_rs + if want_rs: + import ray_data_arrow_rs # noqa: F401 fail loudly if crate absent + + # Warm the pool, census OUR cluster's workers (new ray:: pids only), then + # keep the sampler discovering workers forked later. + ray.get(ray.remote(lambda: os.getpid()).remote()) + time.sleep(2) + sampler = _WorkerSampler(before_init) + sampler.census(late=False) + if not sampler.baseline: + raise RuntimeError("no worker found in the fresh local cluster") + sampler.start() + + def one_read(): + ds = ray.data.read_parquet(path, concurrency=1) + n = 0 + for bundle in ds.iter_internal_ref_bundles(): + n += bundle.num_rows() or 0 + return n + + rows1 = one_read() + time.sleep(2) # settle before the after-read floor + after1 = sampler.snapshot() + rows2 = one_read() + time.sleep(2) + sampler.census(late=True) # catch a worker forked at the tail of read 2 + after2 = sampler.snapshot() + sampler.stop() + + idle_uss = sampler.baseline + + # Positive ID first: any pid that ever bore the read-op proctitle. Growth + # selection is only the fallback (a sub-0.5 s read the title poll missed): + # a worker first seen mid-read has an inflated baseline, so its settled + # delta reads ~0 and an idle pid's noise can win the argmax. + matched = [ + p + for p in idle_uss + if any(_READ_TASK_TITLE_SUBSTR in t for t in sampler.titles.get(p, ())) + ] + + def _delta(snap, p): + v = snap.get(p) + # Not yet discovered at that snapshot (or /proc read failed) = no + # observed growth, not negative growth. + return 0 if v is None else v - idle_uss[p] + + if matched: + rw = max(matched, key=lambda p: after2.get(p) or 0) + id_method = "title" + same_worker = len(matched) == 1 + else: + # Attribute each read separately: if read 1 and read 2 landed on + # different pids, one pid's after1/after2 deltas would mix the reads. + rw = max(idle_uss, key=lambda p: _delta(after2, p)) + grower1 = max(idle_uss, key=lambda p: _delta(after1, p)) + grower2 = max(idle_uss, key=lambda p: _delta(after2, p) - _delta(after1, p)) + id_method = "growth" + same_worker = grower1 == grower2 == rw + + result = { + "arm": arm, + "rows_read": rows1, + "rows_read_2": rows2, + "worker_pid": rw, + "idle_uss_mb": round(idle_uss[rw] / 1024, 1), + "peak_uss_mb": round(sampler.peak[rw] / 1024, 1), + "after_read1_uss_mb": round((after1.get(rw) or idle_uss[rw]) / 1024, 1), + "after_read2_uss_mb": round((after2.get(rw) or idle_uss[rw]) / 1024, 1), + "n_workers_seen": len(idle_uss), + # "title" = pid positively bore ray::ReadFiles* while executing; + # "growth" = fallback max-delta pick (weaker — see comment above). + "read_worker_id_method": id_method, + "n_read_task_workers": len(matched), + # True = worker first seen after sampling began; its idle baseline may + # already include read work (constant would read low). + "read_worker_late_spawn": rw in sampler.late, + # False = read tasks ran on more than one pid (title path) or the two + # reads grew different pids (growth path); after1/after2 deltas on rw + # then do NOT mean "read-1 floor / read-2 floor" for one worker. + "reads_on_same_worker": same_worker, + } + print("CELLRESULT " + json.dumps(result), flush=True) + ray.shutdown() + + +def run_cell_subprocess(arm, path): + env = dict(os.environ) + env[FLAG] = "1" if arm == "rs" else "0" + env.pop("RAY_ADDRESS", None) + proc = subprocess.run( + [sys.executable, os.path.abspath(__file__), "--cell", arm, "--path", path], + env=env, + capture_output=True, + text=True, + ) + for line in proc.stdout.splitlines(): + if line.startswith("CELLRESULT "): + return json.loads(line[len("CELLRESULT ") :]) + raise RuntimeError( + f"cell {arm} produced no CELLRESULT (rc={proc.returncode})\n" + f"stdout tail: {proc.stdout[-2000:]}\nstderr tail: {proc.stderr[-2000:]}" + ) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--path", required=True, help="parquet dir/file (S3 or local)") + ap.add_argument("--arms", default="pa,rs") + ap.add_argument("--cell", help="internal: run one arm in-process") + args = ap.parse_args() + + if args.cell: + run_cell_body(args.cell, args.path) + return + + results = [run_cell_subprocess(a.strip(), args.path) for a in args.arms.split(",")] + print(json.dumps(results, indent=2)) + if len(results) == 2: + a, b = results + print( + f"\nper-worker constant (rs - pa, settled after read 2): " + f"{b['after_read2_uss_mb'] - a['after_read2_uss_mb']:+.1f} MB " + f"(idle {b['idle_uss_mb'] - a['idle_uss_mb']:+.1f}, " + f"peak {b['peak_uss_mb'] - a['peak_uss_mb']:+.1f})" + ) + + +if __name__ == "__main__": + main() diff --git a/release/nightly_tests/dataset/backpressure_benchmark.py b/release/nightly_tests/dataset/backpressure_benchmark.py index d86578107a9d..53233c799fa1 100644 --- a/release/nightly_tests/dataset/backpressure_benchmark.py +++ b/release/nightly_tests/dataset/backpressure_benchmark.py @@ -6,7 +6,7 @@ import pyarrow as pa import ray -from benchmark import Benchmark +from benchmark import Benchmark, collect_operator_metrics, consume_ref_bundles def parse_args() -> argparse.Namespace: @@ -70,10 +70,11 @@ def run_fast_producer_slow_consumer(args: argparse.Namespace): .map_batches(producer) .map_batches(consumer, compute=ray.data.TaskPoolStrategy(size=1)) ) - for _ in ds.iter_internal_ref_bundles(): - pass + consume_ref_bundles(ds) - return vars(args) + # Arguments, plus per-operator wall time / output bytes / per-task USS+RSS — + # backpressure is about where memory piles up, so per-operator is the useful grain. + return {**vars(args), **collect_operator_metrics(ds)} def run_training_prefetch(args: argparse.Namespace): diff --git a/release/nightly_tests/dataset/benchmark.py b/release/nightly_tests/dataset/benchmark.py index 28fea379b63b..c272380e7619 100644 --- a/release/nightly_tests/dataset/benchmark.py +++ b/release/nightly_tests/dataset/benchmark.py @@ -6,7 +6,7 @@ import threading import time from enum import Enum -from typing import Any, Callable, Dict, List, Union +from typing import Any, Callable, Dict, List, Tuple, Union import dataclasses import ray from ray._private.internal_api import get_memory_info_reply, get_state_from_address @@ -14,6 +14,20 @@ logger = logging.getLogger(__name__) +# Poll per-task memory (MemoryProfiler in each map worker) at 20 Hz instead of the +# 1 Hz production default. At 1 Hz, a task shorter than ~3 s gets its "max USS" +# from the single synchronous end-of-task sample — end-of-task RESIDENT memory, +# not peak working set — which made 16 of 37 instrumented tests unmeasurable in +# the 2026-08-14 A/B. Each sample is one /proc/self/statm read (microseconds), so +# 20 Hz costs ~0.01% CPU per worker. Set here (release harness only, both arms of +# an A/B identically) rather than upstream, where the conservative default is +# deliberate. The DataContext is captured at dataset creation, so importing +# benchmark.py before building datasets — which every release script does — is +# sufficient for this to reach the workers. +ray.data.DataContext.get_current().memory_usage_poll_interval_s = float( + os.environ.get("RAY_DATA_BENCH_MEMORY_POLL_S", "0.05") +) + def _get_spilled_bytes_total(state) -> float: """Get the total number of spilled bytes across the cluster.""" @@ -120,6 +134,414 @@ def collect_dataset_stats(ds: "ray.data.Dataset") -> Dict[str, Any]: } +class _NullMonitor: + """Stand-in used when the node memory monitor is disabled or unimportable.""" + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + def summary(self) -> Dict[str, Any]: + return {} + + +def _node_memory_monitor(case_name: str): + """Return a started-on-enter node memory monitor, or a no-op stand-in. + + Import is deferred and failure is swallowed: the monitor is diagnostics, and a + benchmark must not fail because a diagnostic could not be imported (release images + and local runs do not always carry the same files). + """ + try: + import node_memory_monitor + + if node_memory_monitor.enabled(): + return node_memory_monitor.NodeMemoryMonitor(case_name) + except Exception: # noqa: BLE001 + logger.warning("node memory monitor unavailable", exc_info=True) + return _NullMonitor() + + +def consume_ref_bundles(ds: "ray.data.Dataset", per_bundle: Callable = None) -> int: + """Consume ``ds`` as RefBundles *without* losing the executor's final stats. + + Use this instead of ``ds.iter_internal_ref_bundles()`` in any benchmark that then + reads stats off ``ds`` (``collect_operator_metrics`` / ``collect_dataset_stats``). + Consumption is identical — the same zero-copy bundle iterator, no blocks fetched — + the only difference is ``capture_executor=True``. + + Why it matters: ``Dataset._execute_to_iterator`` caches ``executor.get_stats()`` + right after the FIRST bundle, and ``iter_internal_ref_bundles`` drops the executor + (``capture_executor=False``, ``dataset.py:7482-7486``), so a later + ``get_stats_summary()`` is pinned to that mid-execution snapshot — taken before the + last read task's ``on_task_finished`` populated ``average_max_uss_per_task``. The + per-task memory metrics then come back ``None`` non-deterministically, depending on + whether the operator happened to finish early: ``ListFiles`` always had values, + ``ReadParquet`` often did not. It cost both arms of two multi-node release A/Bs + their ``read_large_parquet`` memory numbers. + + Args: + ds: the dataset to consume. + per_bundle: optional callback invoked with each ``RefBundle``. Leave unset to + drop each bundle as it arrives (the usual "read it all and throw it away"). + + Returns: + The number of bundles consumed. + """ + bundle_iter, _, _ = ds._execute_to_iterator(capture_executor=True) + # Deliberately NOT calling ds._synchronize_progress_bar() here, though + # iter_internal_ref_bundles does: it is a no-op there (nothing was captured) but + # here it would `shutdown(force=True)` the executor we just captured, truncating + # consumption to the one bundle _execute_to_iterator already forced. Verified + # locally: with the call, an 8-block dataset yields 1 bundle and 1/8 of the rows. + num_bundles = 0 + for bundle in bundle_iter: + num_bundles += 1 + if per_bundle is not None: + per_bundle(bundle) + return num_bundles + + +def collect_operator_metrics(ds: "ray.data.Dataset") -> Dict[str, Any]: + """Per-operator time / output-bytes / worker-memory, for merging into a result dict. + + Surfaces numbers that otherwise live only on the Prometheus dashboard, not in the + release log or databricks: each operator's wall time, output size/rows, and its + per-task peak worker memory — USS (private working set) and RSS (OS-visible + footprint, includes mapped object-store pages), both as the average across tasks + and the single worst task. All four come from ``MemoryProfiler`` sampling inside + the task (Linux-only; ``None`` on macOS). This isolates the read operator's cost + from downstream compute and exposes the decode-memory metrics that the aggregate + object-store peak cannot see. Best-effort: returns a partial/empty dict rather + than failing the benchmark. + + A ``read_*`` top-level convenience is filled from the ``Read*`` operator with + the most tasks (ties: plan order) so the "parquet part" (read wall time + + output bytes + decode USS/RSS) is a first-class field; ``read_operators`` + lists every ``Read*`` operator compactly, since a multi-table plan (TPC-H + q17: ``part`` AND ``lineitem``) has several and the headline used to be + whichever came first. + + NOTE: stats attach to the consumed dataset handle. Consume ``ds`` itself + (``iter_*``/``write_*``/``materialize``) before calling this; ``ds.count()`` + executes a *copy* of the plan and leaves ``ds`` without stats. And consume via + ``consume_ref_bundles(ds)``, not ``ds.iter_internal_ref_bundles()`` — the latter + drops the executor, which pins the stats to a snapshot taken after the first bundle + and silently nulls every per-task memory field below. + """ + from ray.data._internal.stats import DatasetStatsSummary + + def _sum(stat) -> Any: + return stat.sum if stat is not None else None + + # (result-dict key, extra_metrics key) for the per-task memory metrics. + mem_keys = [ + ("avg_max_uss_per_task_bytes", "average_max_uss_per_task"), + ("max_uss_per_task_bytes", "max_uss_per_task"), + ("avg_max_rss_per_task_bytes", "average_max_rss_per_task"), + ("max_rss_per_task_bytes", "max_rss_per_task"), + # Operator TOTALS (seconds summed over the op's tasks) for the part of a + # task the reader-level timers cannot see: block generation is the + # transform's own wall (reader + block build), serialization is the + # worker's plasma put of each output (``object_creation_dur_s`` per + # generated object), and the two backpressure walls are the driver's + # (time it could not launch tasks / time it held outputs). Divide by + # ``num_tasks_finished`` for a per-task mean. Findings M117–M119: a + # read task's duration minus its in-iterator decode leaves ~5 s that + # is not CPU and not the planner's yield — this splits it. + ("block_generation_time_s_total", "block_generation_time"), + ("block_serialization_time_s_total", "block_serialization_time_s"), + ("task_gen_and_ser_time_s_total", "task_block_gen_and_ser_time_s"), + ("task_submission_backpressure_s_total", "task_submission_backpressure_time"), + ("task_output_backpressure_s_total", "task_output_backpressure_time"), + ("num_tasks_finished", "num_tasks_finished"), + ] + # (result-dict key, extra_metrics key) for the full per-task distributions. + # ``max_uss_bytes``/``max_rss_bytes`` are DistributionTracker.as_dict() outputs: + # num_samples/mean/variance/min/max plus p25..p99 (quantiles come from a KLL + # sketch and are None unless ``datasketches`` is importable in the worker — + # it is pinned in requirements_compiled.txt, so release images have it). + dist_keys = [ + ("max_uss_per_task_dist", "max_uss_bytes"), + ("max_rss_per_task_dist", "max_rss_bytes"), + ("task_duration_dist", "op_task_duration_stats"), + # Reader-level per-task aggregates (ReadFilesTaskStats, reported by the + # ReadFiles transform on every V2 file read): what the DECODER did, + # independent of the memory profiler — bytes/decode-seconds per task + # and the largest single table the reader yielded (working-set proxy). + ("decoded_bytes_per_task_dist", "read_task_decoded_bytes"), + ("decode_wall_s_per_task_dist", "read_task_decode_wall_s"), + ("peak_batch_bytes_per_task_dist", "read_task_peak_batch_bytes"), + # Wall seconds of the reader's end-of-stream finalizer (the arrow-rs + # malloc_trim, on by default, RAY_DATA_ARROW_RS_MALLOC_TRIM_EOS=0 turns + # it off); inside decode_wall_s. 0 for pyarrow and with the knob off. + ("trim_wall_s_per_task_dist", "read_task_trim_wall_s"), + ("yield_wall_s_per_task_dist", "read_task_yield_wall_s"), + ("first_table_wall_s_per_task_dist", "read_task_first_table_wall_s"), + ] + + out: Dict[str, Any] = {"operators_detail": []} + try: + summary = ds.get_stats_summary(detail=True) + for node in DatasetStatsSummary._collect_dataset_stats_summaries(summary): + extra = getattr(node, "extra_metrics", {}) or {} + mem = {out_key: extra.get(in_key) for out_key, in_key in mem_keys} + dists = {out_key: extra.get(in_key) for out_key, in_key in dist_keys} + for op in node.operators_stats or []: + # Output-block granularity: StatsSummary is per-block, so its + # count/min/mean/max ARE the block-size distribution — the + # signal for shuffle-feeding-granularity questions. + osb = op.output_size_bytes + out["operators_detail"].append( + { + "operator_name": op.operator_name, + "wall_time_s": _sum(op.wall_time), + "cpu_time_s": _sum(op.cpu_time), + "udf_time_s": _sum(op.udf_time), + "output_num_rows": _sum(op.output_num_rows), + "output_size_bytes": _sum(op.output_size_bytes), + "output_num_blocks": osb.count if osb else None, + "block_size_bytes_min": osb.min if osb else None, + "block_size_bytes_mean": osb.mean if osb else None, + "block_size_bytes_max": osb.max if osb else None, + **mem, + **dists, + } + ) + + def _n_tasks(entry) -> int: + # One sample per finished task that emitted a block. Task duration + # is recorded on every platform; max_uss only where USS is readable + # (Linux), so it is the fallback. + for key in ("task_duration_dist", "max_uss_per_task_dist"): + n = (entry.get(key) or {}).get("num_samples") + if n: + return n + return 0 + + def _q(entry, dist_key, stat): + return (entry.get(dist_key) or {}).get(stat) + + read_entries = [ + e for e in out["operators_detail"] if "Read" in (e["operator_name"] or "") + ] + # One compact row per Read operator, in plan order. A task that emits + # no block carries no worker stats at all (TaskExecWorkerStats rides + # on block metadata), so ``tasks`` can be below the operator's task + # count and ``decoded_*_n`` == ``tasks`` when the reader stats fired. + out["read_operators"] = [ + { + "operator_name": e["operator_name"], + "tasks": _n_tasks(e), + "wall_time_s": e["wall_time_s"], + "output_num_rows": e["output_num_rows"], + "output_size_bytes": e["output_size_bytes"], + "output_num_blocks": e["output_num_blocks"], + "max_uss_per_task_p50": _q(e, "max_uss_per_task_dist", "p50"), + "max_uss_per_task_max": _q(e, "max_uss_per_task_dist", "max"), + "task_duration_p50": _q(e, "task_duration_dist", "p50"), + "decoded_bytes_n": _q(e, "decoded_bytes_per_task_dist", "num_samples"), + "decoded_bytes_p50": _q(e, "decoded_bytes_per_task_dist", "p50"), + "decoded_bytes_max": _q(e, "decoded_bytes_per_task_dist", "max"), + "peak_batch_bytes_p50": _q(e, "peak_batch_bytes_per_task_dist", "p50"), + "peak_batch_bytes_max": _q(e, "peak_batch_bytes_per_task_dist", "max"), + "trim_wall_s_p50": _q(e, "trim_wall_s_per_task_dist", "p50"), + "trim_wall_s_max": _q(e, "trim_wall_s_per_task_dist", "max"), + "yield_wall_s_p50": _q(e, "yield_wall_s_per_task_dist", "p50"), + "yield_wall_s_max": _q(e, "yield_wall_s_per_task_dist", "max"), + "first_table_wall_s_p50": _q( + e, "first_table_wall_s_per_task_dist", "p50" + ), + "first_table_wall_s_max": _q( + e, "first_table_wall_s_per_task_dist", "max" + ), + } + for e in read_entries + ] + if read_entries: + # Headline = the Read with the most tasks; ``max`` keeps the first + # maximal entry, i.e. plan order breaks ties. + entry = max(read_entries, key=_n_tasks) + out["read_operator_name"] = entry["operator_name"] + out["read_wall_time_s"] = entry["wall_time_s"] + out["read_output_size_bytes"] = entry["output_size_bytes"] + out["read_output_num_blocks"] = entry["output_num_blocks"] + out["read_block_size_bytes_mean"] = entry["block_size_bytes_mean"] + for out_key, _ in mem_keys: + out[f"read_{out_key}"] = entry[out_key] + for out_key, _ in dist_keys: + out[f"read_{out_key}"] = entry[out_key] + except Exception: + logger.warning("collect_operator_metrics failed", exc_info=True) + return out + + +def collect_task_distribution(case_start_unix_ms: float) -> Dict[str, Any]: + """How this case's tasks were spread over workers and nodes. + + Answers the placement question the per-task memory distributions cannot: + two arms can run identical tasks yet land them on different worker pools — + a faster arm satisfies autoscaling demand with fewer workers, so each + worker runs MORE tasks and accumulates a higher between-tasks memory + floor. These aggregates make that visible per arm: tasks-per-worker / + per-node distributions, worker lifespan (first task start to last task + end), and worker busy fraction (summed task seconds over lifespan). + + Uses the state API's task list, filtered to tasks that started after this + case began (cases run sequentially in one job). The API caps at 10k tasks; + ``task_dist_truncated`` flags when the cap was hit, in which case the + per-worker numbers undercount instead of erroring. + """ + from ray.util.state.api import list_tasks + + def _quantile(sorted_vals, q): + if not sorted_vals: + return None + i = max(0, min(len(sorted_vals) - 1, round(q * (len(sorted_vals) - 1)))) + return sorted_vals[i] + + out: Dict[str, Any] = {} + try: + tasks = list_tasks(detail=True, limit=10_000, raise_on_missing_output=False) + per_worker: Dict[str, List[Tuple[float, float]]] = {} + nodes = set() + durations = [] + recs = [] + n = 0 + for t in tasks: + if not t.start_time_ms or t.start_time_ms < case_start_unix_ms: + continue + if not t.worker_id or not t.end_time_ms: + continue + n += 1 + per_worker.setdefault(t.worker_id, []).append( + (t.start_time_ms, t.end_time_ms) + ) + if t.node_id: + nodes.add((t.node_id, t.worker_id)) + durations.append((t.end_time_ms - t.start_time_ms) / 1000.0) + recs.append( + ( + t.name, + t.start_time_ms, + t.end_time_ms, + t.node_id, + t.worker_id, + t.worker_pid, + ) + ) + + out["task_dist_num_tasks"] = n + out["task_dist_truncated"] = len(tasks) >= 10_000 + out["task_dist_num_workers"] = len(per_worker) + node_ids = {nid for nid, _ in nodes} + out["task_dist_num_nodes"] = len(node_ids) + + counts = sorted(len(v) for v in per_worker.values()) + out["task_dist_tasks_per_worker_mean"] = ( + round(n / len(per_worker), 2) if per_worker else None + ) + out["task_dist_tasks_per_worker_p50"] = _quantile(counts, 0.5) + out["task_dist_tasks_per_worker_max"] = counts[-1] if counts else None + + per_node: Dict[str, int] = {} + for nid, wid in nodes: + per_node[nid] = per_node.get(nid, 0) + 1 + workers_per_node = sorted(per_node.values()) + out["task_dist_workers_per_node_max"] = ( + workers_per_node[-1] if workers_per_node else None + ) + + durations.sort() + out["task_dist_task_duration_s_p50"] = _quantile(durations, 0.5) + out["task_dist_task_duration_s_p90"] = _quantile(durations, 0.9) + out["task_dist_task_duration_s_max"] = durations[-1] if durations else None + + lifespans = [] + busy_fracs = [] + for spans in per_worker.values(): + start = min(s for s, _ in spans) + end = max(e for _, e in spans) + lifespan_s = (end - start) / 1000.0 + lifespans.append(lifespan_s) + busy_s = sum(e - s for s, e in spans) / 1000.0 + if lifespan_s > 0: + busy_fracs.append(busy_s / lifespan_s) + lifespans.sort() + busy_fracs.sort() + out["task_dist_worker_lifespan_s_p50"] = _quantile(lifespans, 0.5) + out["task_dist_worker_lifespan_s_max"] = lifespans[-1] if lifespans else None + out["task_dist_worker_busy_frac_p50"] = ( + round(_quantile(busy_fracs, 0.5), 4) if busy_fracs else None + ) + out["task_dist_worker_busy_frac_min"] = ( + round(busy_fracs[0], 4) if busy_fracs else None + ) + + # Per-task spawn timeline (the 2x2 topology probe wants WHEN and WHERE + # tasks ran, not just the aggregates above). Compact legend-indexed + # rows keep this a few hundred KB inside result.json -- loose files do + # not survive a release run. RAY_DATA_BENCH_TASK_TIMELINE=0 disables. + if os.environ.get("RAY_DATA_BENCH_TASK_TIMELINE", "1") != "0": + out["task_timeline"] = _build_task_timeline(recs, case_start_unix_ms) + except Exception: + logger.warning("collect_task_distribution failed", exc_info=True) + return out + + +_TASK_TIMELINE_CAP = 6000 + + +def _build_task_timeline(recs, t0_ms): + """Legend-indexed task timeline: each row is + [name_idx, start_ms_rel, dur_ms, node_idx, worker_idx]. + + When the cap bites, read tasks are kept preferentially (they are what the + 2x2 experiment is about), earliest-first within each class, and + ``truncated`` is set. Node/worker ids are truncated to 16 hex chars in the + legends; ``worker_pids`` aligns with ``workers`` so rows can be joined + against node_memory_monitor's per-pid USS samples. + """ + recs.sort(key=lambda r: ("Read" not in (r[0] or ""), r[1])) + truncated = len(recs) > _TASK_TIMELINE_CAP + recs = recs[:_TASK_TIMELINE_CAP] + recs.sort(key=lambda r: r[1]) + names, nodes, workers, worker_pids, rows = [], [], [], [], [] + name_idx, node_idx, worker_idx = {}, {}, {} + for name, start, end, node, worker, pid in recs: + name, node, worker = name or "", (node or "")[:16], (worker or "")[:16] + if name not in name_idx: + name_idx[name] = len(names) + names.append(name) + if node not in node_idx: + node_idx[node] = len(nodes) + nodes.append(node) + if worker not in worker_idx: + worker_idx[worker] = len(workers) + workers.append(worker) + worker_pids.append(pid) + rows.append( + [ + name_idx[name], + int(start - t0_ms), + int(end - start), + node_idx[node], + worker_idx[worker], + ] + ) + return { + "t0_unix_ms": t0_ms, + "names": names, + "nodes": nodes, + "workers": workers, + "worker_pids": worker_pids, + "rows": rows, + "truncated": truncated, + } + + class RuntimeEnvSetupTracker: """Collects runtime environment creation times across the cluster. @@ -167,6 +589,9 @@ def benchmark_py_modules() -> List[str]: return [ os.path.realpath(__file__), os.path.join(dataset_dir, "profiling"), + # The sampler actor class must be importable on every worker node, not just + # wherever the driver ran. + os.path.join(dataset_dir, "node_memory_monitor.py"), ] @@ -232,7 +657,14 @@ def run_fn( print(f"Running case: {name}") state = get_state_from_address(ray.get_runtime_context().gcs_address) - with ObjectStoreMemorySampler(state) as memory_sampler: + # Per-node worker memory with stage provenance. Off unless + # RAY_DATA_BENCH_NODE_MEM_MONITOR=1; a no-op context manager otherwise, so the + # default path is byte-for-byte what it was. + node_mem = _node_memory_monitor(name) + + case_start_unix_ms = time.time() * 1000.0 + + with node_mem, ObjectStoreMemorySampler(state) as memory_sampler: start_time = time.perf_counter() start_spilled_bytes = _get_spilled_bytes_total(state) @@ -256,6 +688,8 @@ def run_fn( memory_sampler.peak_utilization, 4, ), + **node_mem.summary(), + **collect_task_distribution(case_start_unix_ms), } if isinstance(fn_output, dict): for key, value in fn_output.items(): diff --git a/release/nightly_tests/dataset/gpu_batch_inference.py b/release/nightly_tests/dataset/gpu_batch_inference.py index f7502ff9e2a3..0180b38f7364 100644 --- a/release/nightly_tests/dataset/gpu_batch_inference.py +++ b/release/nightly_tests/dataset/gpu_batch_inference.py @@ -9,6 +9,8 @@ BenchmarkMetric, RuntimeEnvSetupTracker, collect_dataset_stats, + collect_operator_metrics, + consume_ref_bundles, benchmark_py_modules, ) from torchvision.models import ResNet50_Weights, resnet50 @@ -117,10 +119,16 @@ def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: total_images = 0 # NOTE: We're iterating over ref-bundles to avoid pulling blocks into the - # driver, therefore making it a factor impacting benchmark performance - for bundle in ds.iter_internal_ref_bundles(): + # driver, therefore making it a factor impacting benchmark performance. + # consume_ref_bundles is that same iteration with capture_executor=True, + # which is what keeps the stats read below from being a snapshot taken + # after the first bundle. + def tally(bundle): + nonlocal total_images total_images += bundle.num_rows() + consume_ref_bundles(ds, tally) + end_time = time.time() total_time = end_time - start_time @@ -142,9 +150,13 @@ def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]: assert dead_nodes print(f"Total chaos killed: {dead_nodes}") - # For structured output integration with internal tooling - results = collect_dataset_stats(ds) + # For structured output integration with internal tooling. + # collect_dataset_stats was previously assigned and then immediately overwritten by + # the literal below, so none of it reached the result JSON; it is merged in now, + # together with the per-operator wall/output/USS+RSS breakdown. results = { + **collect_dataset_stats(ds), + **collect_operator_metrics(ds), BenchmarkMetric.RUNTIME: total_time, BenchmarkMetric.THROUGHPUT: throughput, "data_directory": data_directory, diff --git a/release/nightly_tests/dataset/groupby_benchmark.py b/release/nightly_tests/dataset/groupby_benchmark.py index 7cc0f15215fa..67a267e46bfb 100644 --- a/release/nightly_tests/dataset/groupby_benchmark.py +++ b/release/nightly_tests/dataset/groupby_benchmark.py @@ -5,7 +5,7 @@ import pyarrow.compute as pc import ray -from benchmark import Benchmark +from benchmark import Benchmark, collect_operator_metrics, consume_ref_bundles from ray.data import DataContext from ray.data.context import ShuffleStrategy @@ -62,28 +62,32 @@ def benchmark_fn(): grouped_ds = ray.data.read_parquet( path, override_num_blocks=override_num_blocks ).groupby(args.group_by) - consume_fn(grouped_ds) + consumed_ds = consume_fn(grouped_ds) - # Report arguments for the benchmark. - return vars(args) + # Arguments, plus per-operator wall time / output bytes / per-task USS+RSS so a + # regression can be attributed to the read, the shuffle or the aggregation + # rather than to the job as a whole. + return {**vars(args), **collect_operator_metrics(consumed_ds)} benchmark.run_fn("main", benchmark_fn) benchmark.write_result() def get_consume_fn(args: argparse.Namespace): + # Each consume_fn returns the *consumed* dataset handle: execution stats attach to + # the handle that was consumed, so that is what collect_operator_metrics needs. if args.aggregate: def consume_fn(grouped_ds): # 'column05' is 'l_extendedprice' - grouped_ds.mean("column05").materialize() + return grouped_ds.mean("column05").materialize() elif args.map_groups: def consume_fn(grouped_ds): ds = grouped_ds.map_groups(normalize_table, batch_format="pyarrow") - for _ in ds.iter_internal_ref_bundles(): - pass + consume_ref_bundles(ds) + return ds else: assert False, f"Invalid consume argument: {args}" diff --git a/release/nightly_tests/dataset/join_benchmark.py b/release/nightly_tests/dataset/join_benchmark.py index 87b2f797e953..9eaa2dc2e6ac 100644 --- a/release/nightly_tests/dataset/join_benchmark.py +++ b/release/nightly_tests/dataset/join_benchmark.py @@ -1,7 +1,7 @@ import ray import argparse -from benchmark import Benchmark +from benchmark import Benchmark, collect_operator_metrics, consume_ref_bundles def parse_args() -> argparse.Namespace: @@ -60,8 +60,21 @@ def benchmark_fn(): join_type=args.join_type, ) - # Process joined_ds if needed - print(f"Join completed with {joined_ds.count()} records.") + # Consume the bundles rather than calling count(): count() executes a *copy* of + # the plan, so the stats never attach to `joined_ds` and the per-operator + # numbers below would all be empty. Row count comes from the bundles instead. + total_rows = 0 + + def tally(bundle): + nonlocal total_rows + total_rows += bundle.num_rows() + + consume_ref_bundles(joined_ds, tally) + print(f"Join completed with {total_rows} records.") + + # Per-operator wall time / output bytes / per-task USS+RSS: separates the two + # reads from the shuffle and the join itself. + return {"num_rows": total_rows, **collect_operator_metrics(joined_ds)} benchmark.run_fn(str(vars(args)), benchmark_fn) benchmark.write_result() diff --git a/release/nightly_tests/dataset/map_benchmark.py b/release/nightly_tests/dataset/map_benchmark.py index b31033c92078..541d6cae6360 100644 --- a/release/nightly_tests/dataset/map_benchmark.py +++ b/release/nightly_tests/dataset/map_benchmark.py @@ -7,7 +7,7 @@ import pandas as pd import ray -from benchmark import Benchmark +from benchmark import Benchmark, collect_operator_metrics, consume_ref_bundles def parse_args() -> argparse.Namespace: @@ -111,9 +111,12 @@ def apply_map_batches(ds): ) def benchmark_fn(): - ctx = ray.data.DataContext.get_current() - ctx.use_datasource_v2 = False - # Use V1 for this benchmark, since V2 is currently spilling + # Leave ``num_cpus`` unset: the fusion rule canonicalizes an unspecified + # ``num_cpus`` to 1, so any other value here (0.99 included) makes the + # read incompatible with the downstream ``map_batches`` ops and they + # stop fusing. Unfused, every read output crosses the object store -- + # 1.13 TiB on sf1000 lineitem, ~485s of serialization against ~0.6s + # fused. ds = ray.data.read_parquet(path) # Apply the map transformation. @@ -131,11 +134,12 @@ def dummy_write(batch): ds = ds.map_batches(dummy_write) - for _ in ds.iter_internal_ref_bundles(): - pass + consume_ref_bundles(ds) - # Report arguments for the benchmark. - return vars(args) + # Arguments, plus per-operator wall time / output bytes / per-task USS+RSS. + # The read and the maps fuse into one operator here, so these numbers are the + # fused pipeline's — which is exactly the question when a map case regresses. + return {**vars(args), **collect_operator_metrics(ds)} benchmark.run_fn("main", benchmark_fn) benchmark.write_result() diff --git a/release/nightly_tests/dataset/node_memory_monitor.py b/release/nightly_tests/dataset/node_memory_monitor.py new file mode 100644 index 000000000000..b809336a7e5c --- /dev/null +++ b/release/nightly_tests/dataset/node_memory_monitor.py @@ -0,0 +1,508 @@ +"""Per-node worker-memory sampling for the dataset release benchmarks. + +**What problem this solves.** The release suite's memory evidence came from Prometheus +``query_range`` over node memory. That series is scraped rarely enough that a short job +gets a handful of distinct values, it says nothing about *which* stage held the memory, +and a node's total tells you nothing about whether the decode heap or the object store +grew. Ray's own per-task ``MemoryProfiler`` (surfaced by +``benchmark.collect_operator_metrics``) fixes the attribution but only exists inside map +tasks — shuffle tasks, actors and the raylet are invisible to it. + +This module is the third view: one actor per node, sampling every ``ray::*`` worker +process on that node at ~1 s, keyed by the **proctitle** — which Ray sets to +``ray::`` — plus the node's own used bytes. That yields +peak-with-provenance: not just "the node reached 58 GB" but "the node reached 58 GB +while ``ray::ReadParquet`` workers held 41 GB of it across 8 processes". + +**Off by default.** Set ``RAY_DATA_BENCH_NODE_MEM_MONITOR=1`` to enable; every driver +that runs through ``Benchmark.run_fn`` then gets it with no code change. The sampler +actors take ``num_cpus=0`` so they never displace benchmark work. + +Environment: + RAY_DATA_BENCH_NODE_MEM_MONITOR "1" to enable (default off) + RAY_DATA_BENCH_NODE_MEM_INTERVAL seconds between samples (default 1.0) + RAY_DATA_BENCH_NODE_MEM_DIR where to write the JSONL trace + (default: alongside TEST_OUTPUT_JSON, else CWD) + RAY_DATA_BENCH_NODE_MEM_MAX_SAMPLES per-node sample cap before decimation + (default 20000; ~5.5 h at 1 Hz) + +Result-dict fields it adds (all ``None``/absent if disabled or unavailable): + node_mem_peak_used_gb peak "used" bytes on the worst node + node_mem_peak_used_node that node's IP + node_mem_peak_used_source cgroup | meminfo | psutil — which one it came from + node_mem_peak_worker_uss_gb peak summed worker USS on the worst node (Linux) + node_mem_p50/p90_worker_uss_gb same series' time-percentiles on that node — + peak>>p50 = transient spike, peak~=p50 = sustained + node_mem_peak_worker_rss_gb same in RSS, which includes object-store pages + node_mem_top_workers_uss/_rss {proctitle: peak GB}, biggest first — the provenance + node_mem_trace_path the JSONL trace, one line per (node, sample) + node_mem_nodes / node_mem_samples / node_mem_error +""" +import json +import logging +import os +import sys +import threading +import time +from typing import Any, Dict, List, Optional + +import ray + +logger = logging.getLogger(__name__) + +ENABLE_ENV = "RAY_DATA_BENCH_NODE_MEM_MONITOR" +INTERVAL_ENV = "RAY_DATA_BENCH_NODE_MEM_INTERVAL" +DIR_ENV = "RAY_DATA_BENCH_NODE_MEM_DIR" +MAX_SAMPLES_ENV = "RAY_DATA_BENCH_NODE_MEM_MAX_SAMPLES" + +# Proctitles are ``ray::``; keep the whole thing, it is already the stage label. +_WORKER_PREFIX = "ray::" +_SELF_PREFIX = "ray::_NodeMemorySampler" + + +def enabled() -> bool: + return os.environ.get(ENABLE_ENV, "0") == "1" + + +def _bytes_to_gb(b: Optional[float]) -> Optional[float]: + return round(b / (1024**3), 4) if b else b + + +def _read_node_used_bytes() -> Dict[str, Optional[int]]: + """This node's memory pressure, by the same convention Ray's memory monitor uses. + + Three numbers, because they disagree in ways that matter: + + * ``cgroup_used`` — the container's own accounting (``memory.current`` on cgroup v2, + ``memory.usage_in_bytes`` on v1). This is what the OOM killer acts on, and it + *includes* the object store's ``/dev/shm`` pages. + * ``meminfo_used`` — ``MemTotal - MemAvailable`` from ``/proc/meminfo``: the whole + box, including anything outside our cgroup. + * ``psutil_used`` — ``total - available`` from psutil. Same convention as the second, + but portable, so local (macOS) smoke runs still produce a number. On Linux it is a + cross-check, not a third opinion. + + The first two are ``None`` off Linux. + """ + out: Dict[str, Optional[int]] = { + "cgroup_used": None, + "meminfo_used": None, + "psutil_used": None, + } + for path in ( + "/sys/fs/cgroup/memory.current", # v2 + "/sys/fs/cgroup/memory/memory.usage_in_bytes", # v1 + ): + try: + with open(path) as fh: + out["cgroup_used"] = int(fh.read().strip()) + break + except (OSError, ValueError): + continue + try: + total = avail = None + with open("/proc/meminfo") as fh: + for line in fh: + if line.startswith("MemTotal:"): + total = int(line.split()[1]) * 1024 + elif line.startswith("MemAvailable:"): + avail = int(line.split()[1]) * 1024 + if total is not None and avail is not None: + break + if total is not None and avail is not None: + out["meminfo_used"] = total - avail + except (OSError, ValueError, IndexError): + pass + try: + import psutil + + vm = psutil.virtual_memory() + out["psutil_used"] = vm.total - vm.available + except Exception: # noqa: BLE001 - diagnostics only + pass + return out + + +@ray.remote(num_cpus=0) +class _NodeMemorySampler: + """Samples one node's Ray worker processes, grouped by proctitle. + + ``num_cpus=0`` on purpose: this must not take a slot away from the workload it is + measuring, and it is idle between samples. + """ + + def __init__(self, interval_s: float, max_samples: int): + self._interval_s = interval_s + self._max_samples = max_samples + self._stop = threading.Event() + self._thread: Optional[threading.Thread] = None + self._samples: List[Dict[str, Any]] = [] + # Decimation factor: after the cap, keep every Nth sample and thin what we hold, + # so a long run degrades resolution instead of dying or truncating its tail. + self._keep_every = 1 + self._seen = 0 + self._uss_ok = True + self._error: Optional[str] = None + self._t0 = time.time() + + def node_info(self) -> Dict[str, Any]: + return { + "node_id": ray.get_runtime_context().get_node_id(), + "node_ip": ray.util.get_node_ip_address(), + } + + def _sample_once(self) -> Optional[Dict[str, Any]]: + try: + import psutil + except ImportError: # pragma: no cover - psutil ships with ray[default] + self._error = "psutil not available" + return None + + by_title: Dict[str, Dict[str, float]] = {} + for proc in psutil.process_iter(["cmdline"]): + try: + cmdline = proc.info.get("cmdline") or [] + if not cmdline or not cmdline[0].startswith(_WORKER_PREFIX): + continue + title = cmdline[0] + # Don't measure the measurement: this actor's own process appears as + # ray::_NodeMemorySampler[.method] and would otherwise show up in the + # provenance table alongside the stages we care about. + if title.startswith(_SELF_PREFIX): + continue + if self._uss_ok: + try: + mi = proc.memory_full_info() + uss = getattr(mi, "uss", 0) + rss = mi.rss + except (psutil.AccessDenied, NotImplementedError): + self._uss_ok = False + uss, rss = 0, proc.memory_info().rss + else: + uss, rss = 0, proc.memory_info().rss + except (psutil.NoSuchProcess, psutil.AccessDenied, IndexError): + continue + entry = by_title.setdefault(title, {"n": 0, "rss": 0, "uss": 0}) + entry["n"] += 1 + entry["rss"] += rss + entry["uss"] += uss + + return { + "t": round(time.time() - self._t0, 3), + **_read_node_used_bytes(), + "workers_rss": sum(e["rss"] for e in by_title.values()), + "workers_uss": sum(e["uss"] for e in by_title.values()) + if self._uss_ok + else None, + "by_title": by_title, + } + + def _record(self): + sample = self._sample_once() + if sample is None: + return + self._seen += 1 + if self._seen % self._keep_every: + return + self._samples.append(sample) + if len(self._samples) >= self._max_samples: + # Halve resolution: drop every other sample we already hold, and take half + # as many from here on. Repeats as needed, so memory is bounded. + self._samples = self._samples[::2] + self._keep_every *= 2 + + def _run(self): + while not self._stop.wait(self._interval_s): + self._record() + + def start(self) -> bool: + self._t0 = time.time() + self._record() + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + return True + + def stop(self) -> Dict[str, Any]: + self._stop.set() + if self._thread is not None: + self._thread.join(timeout=self._interval_s * 5) + self._record() + + # Per-stage peak, kept separately for USS and RSS. They are NOT interchangeable: + # RSS counts mapped object-store pages, so a task that merely reads big blocks + # looks huge in RSS and small in USS. Never silently substitute one for the + # other — the caller decides, knowing which it got. + peak_uss_by_title: Dict[str, float] = {} + peak_rss_by_title: Dict[str, float] = {} + for sample in self._samples: + for title, entry in sample["by_title"].items(): + peak_rss_by_title[title] = max( + peak_rss_by_title.get(title, 0), entry["rss"] + ) + if self._uss_ok: + peak_uss_by_title[title] = max( + peak_uss_by_title.get(title, 0), entry["uss"] + ) + + def _peak(key): + vals = [s[key] for s in self._samples if s.get(key) is not None] + return max(vals) if vals else None + + return { + **self.node_info(), + "error": self._error, + "uss_available": self._uss_ok, + "num_samples": len(self._samples), + "keep_every": self._keep_every, + "peak_cgroup_used": _peak("cgroup_used"), + "peak_meminfo_used": _peak("meminfo_used"), + "peak_psutil_used": _peak("psutil_used"), + "peak_workers_rss": _peak("workers_rss"), + "peak_workers_uss": _peak("workers_uss"), + "peak_uss_by_title": peak_uss_by_title, + "peak_rss_by_title": peak_rss_by_title, + "samples": self._samples, + } + + +class NodeMemoryMonitor: + """Driver-side context manager: one sampler actor per alive node. + + Best-effort throughout — a monitor that fails must never fail the benchmark, so + every step degrades to a recorded ``node_mem_error`` instead of raising. Use as:: + + with NodeMemoryMonitor("read_parquet") as mon: + ...run the workload... + result.update(mon.summary()) + """ + + def __init__(self, case_name: str, interval_s: Optional[float] = None): + self._case = "".join(c if c.isalnum() or c in "-_." else "_" for c in case_name) + self._interval_s = interval_s or float(os.environ.get(INTERVAL_ENV, "1.0")) + self._max_samples = int(os.environ.get(MAX_SAMPLES_ENV, "20000")) + self._actors: List[Any] = [] + self._summary: Dict[str, Any] = {} + self._error: Optional[str] = None + + def __enter__(self) -> "NodeMemoryMonitor": + try: + self._start() + except Exception as e: # noqa: BLE001 - never fail the benchmark + self._error = f"start failed: {e!r}" + logger.warning("NodeMemoryMonitor failed to start", exc_info=True) + return self + + def __exit__(self, exc_type, exc_value, traceback): + try: + self._stop() + except Exception as e: # noqa: BLE001 + self._error = f"stop failed: {e!r}" + logger.warning("NodeMemoryMonitor failed to stop", exc_info=True) + + def _start(self): + from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + + # Serialize the actor class BY VALUE. By default cloudpickle stores a class as + # (module, name), so the workers must be able to `import node_memory_monitor` — + # true when the driver's directory is uploaded as the working_dir, false when a + # benchmark is launched from anywhere else, and the failure mode is an + # ActorDiedError raised at stop(), by which point every sample is already gone. + # (A per-actor `runtime_env={"py_modules": [...]}` cannot fix this: local paths + # are only accepted at the job level, i.e. in ray.init.) + try: + import ray.cloudpickle as cloudpickle + + cloudpickle.register_pickle_by_value(sys.modules[__name__]) + except Exception: # noqa: BLE001 - fall back to import-by-reference + logger.warning("could not register by-value pickling", exc_info=True) + + nodes = [n for n in ray.nodes() if n.get("Alive")] + for node in nodes: + actor = _NodeMemorySampler.options( + scheduling_strategy=NodeAffinitySchedulingStrategy( + node_id=node["NodeID"], soft=False + ), + # The workload is what should be scheduled; a sampler that cannot be + # placed (drained node, race with autoscaling) is dropped, not waited on. + max_restarts=0, + ).remote(self._interval_s, self._max_samples) + self._actors.append(actor) + ray.get([a.start.remote() for a in self._actors], timeout=120) + + def _outdir(self) -> str: + explicit = os.environ.get(DIR_ENV) + if explicit: + return explicit + # Default next to the release result JSON so the trace is picked up as an + # artifact wherever the results are. The "./result.json" fallback mirrors + # Benchmark.write_result exactly, so an unset TEST_OUTPUT_JSON puts the trace + # in the CWD beside the result file — dirname(abspath(".")) would give the + # CWD's *parent*. + return os.path.dirname( + os.path.abspath(os.environ.get("TEST_OUTPUT_JSON", "./result.json")) + ) + + def _stop(self): + if not self._actors: + return + per_node = ray.get([a.stop.remote() for a in self._actors], timeout=300) + for actor in self._actors: + ray.kill(actor) + self._actors = [] + + outdir = self._outdir() + trace_path = os.path.join(outdir, f"node_mem_{self._case}.jsonl") + try: + os.makedirs(outdir, exist_ok=True) + with open(trace_path, "w") as fh: + for node in per_node: + for sample in node["samples"]: + fh.write( + json.dumps( + { + "node_ip": node["node_ip"], + "node_id": node["node_id"], + **sample, + } + ) + + "\n" + ) + except OSError as e: + self._error = f"trace write failed: {e!r}" + trace_path = None + + # Peak "used" is per node, so take the worst node rather than a sum: a sum + # across nodes answers no question anyone asks about an OOM. + def _worst(key): + vals = [(n[key], n["node_ip"]) for n in per_node if n.get(key) is not None] + return max(vals) if vals else (None, None) + + # Source order is deliberate: the cgroup number is the one the OOM killer acts + # on, so prefer it and fall back only where it does not exist (macOS). + peak_used, peak_used_node = _worst("peak_cgroup_used") + used_source = "cgroup" + for key, label in ( + ("peak_meminfo_used", "meminfo"), + ("peak_psutil_used", "psutil"), + ): + if peak_used is not None: + break + peak_used, peak_used_node = _worst(key) + used_source = label + peak_worker_uss, _ = _worst("peak_workers_uss") + peak_worker_rss, _ = _worst("peak_workers_rss") + + # p50/p90 of summed worker USS OVER TIME, on the same node that produced + # the peak. The peak answers "how bad did it get"; these answer "how + # loaded was that node typically" - peak >> p50 is a transient spike, + # peak ~= p50 is sustained pressure (retention floors show up here). + def _series_pctile(node, key, q): + vals = sorted( + smp[key] + for smp in (node.get("samples") or []) + if smp.get(key) is not None + ) + if not vals: + return None + return vals[min(len(vals) - 1, int(q * len(vals)))] + + uss_node = ( + next( + (n for n in per_node if n.get("peak_workers_uss") == peak_worker_uss), + None, + ) + if peak_worker_uss is not None + else None + ) + p50_worker_uss = ( + _series_pctile(uss_node, "workers_uss", 0.50) if uss_node else None + ) + p90_worker_uss = ( + _series_pctile(uss_node, "workers_uss", 0.90) if uss_node else None + ) + + # Per-node compact view + worker counts. Sustained (p50-over-time) summed + # USS on a SHORT test tracks how many workers were alive, not what each + # held (arrow_rs_docs TODO 36: wide_schema/q6 p50 rows), so report the + # count and the per-worker mean alongside it. "Workers" = every ray::* + # process the sampler saw (task workers incl. ray::IDLE, actors; the + # sampler itself excluded) — the same set workers_uss sums over. + def _worker_count(smp): + return sum(e.get("n", 0) for e in (smp.get("by_title") or {}).values()) + + def _pctile(vals, q): + vals = sorted(vals) + return vals[min(len(vals) - 1, int(q * len(vals)))] if vals else None + + def _node_compact(node): + samples = node.get("samples") or [] + counts = [_worker_count(s) for s in samples] + means = [ + s["workers_uss"] / n + for s, n in zip(samples, counts) + if n and s.get("workers_uss") is not None + ] + mean_p50 = _pctile(means, 0.50) + return { + "node_ip": node.get("node_ip"), + "peak_worker_uss_gb": _bytes_to_gb(node.get("peak_workers_uss")), + "p50_worker_uss_gb": _bytes_to_gb( + _series_pctile(node, "workers_uss", 0.50) + ), + "workers_p50": _pctile(counts, 0.50), + "workers_max": max(counts) if counts else None, + "mean_worker_uss_p50_mb": ( + round(mean_p50 / (1024**2), 1) if mean_p50 is not None else None + ), + "num_samples": len(samples), + } + + per_node_compact = sorted( + (_node_compact(n) for n in per_node), + key=lambda c: c["peak_worker_uss_gb"] or 0, + reverse=True, + )[:64] + uss_node_compact = _node_compact(uss_node) if uss_node else {} + + # Worst single node per proctitle — the provenance the node total lacks. + # USS where available, RSS as a separate field rather than a silent substitute. + def _top(key): + top: Dict[str, float] = {} + for node in per_node: + for title, val in (node.get(key) or {}).items(): + top[title] = max(top.get(title, 0), val) + ranked = sorted(top.items(), key=lambda kv: kv[1], reverse=True)[:8] + return {t: _bytes_to_gb(v) for t, v in ranked} + + self._summary = { + "node_mem_peak_used_gb": _bytes_to_gb(peak_used), + "node_mem_peak_used_node": peak_used_node, + "node_mem_peak_used_source": used_source if peak_used is not None else None, + "node_mem_peak_worker_uss_gb": _bytes_to_gb(peak_worker_uss), + "node_mem_p50_worker_uss_gb": _bytes_to_gb(p50_worker_uss), + "node_mem_p90_worker_uss_gb": _bytes_to_gb(p90_worker_uss), + # Same (peak-USS) node: how many ray:: workers it typically held and + # what one of them typically held — the sustained row's denominator. + "node_mem_workers_p50": uss_node_compact.get("workers_p50"), + "node_mem_workers_max": uss_node_compact.get("workers_max"), + "node_mem_mean_worker_uss_p50_mb": uss_node_compact.get( + "mean_worker_uss_p50_mb" + ), + "node_mem_per_node": per_node_compact, + "node_mem_interval_s": self._interval_s, + "node_mem_peak_worker_rss_gb": _bytes_to_gb(peak_worker_rss), + "node_mem_top_workers_uss": _top("peak_uss_by_title"), + "node_mem_top_workers_rss": _top("peak_rss_by_title"), + "node_mem_nodes": len(per_node), + "node_mem_samples": sum(n["num_samples"] for n in per_node), + "node_mem_trace_path": trace_path, + } + node_errors = [n["error"] for n in per_node if n.get("error")] + if node_errors: + self._error = "; ".join(sorted(set(node_errors))) + + def summary(self) -> Dict[str, Any]: + out = dict(self._summary) + if self._error: + out["node_mem_error"] = self._error + return out diff --git a/release/nightly_tests/dataset/read_and_consume_benchmark.py b/release/nightly_tests/dataset/read_and_consume_benchmark.py index e66c3a388b01..cbc134b91697 100644 --- a/release/nightly_tests/dataset/read_and_consume_benchmark.py +++ b/release/nightly_tests/dataset/read_and_consume_benchmark.py @@ -3,7 +3,7 @@ import uuid from typing import Callable -from benchmark import Benchmark +from benchmark import Benchmark, collect_operator_metrics, consume_ref_bundles import ray from ray.data import SaveMode @@ -80,8 +80,10 @@ def benchmark_fn(): ds = read_fn(args.path) consume_fn(ds) - # Report arguments for the benchmark. - return vars(args) + # Report arguments for the benchmark, plus per-operator time / output bytes / + # decode-USS (isolates the read from downstream consume; surfaces the decode + # memory the object-store peak can't see). ``ds`` is still in scope here. + return {**vars(args), **collect_operator_metrics(ds)} if args.write_delta and args.write_delta_mode == "overwrite": # Populate the table once first (same source/scale as the timed run @@ -118,8 +120,7 @@ def consume_fn(ds): elif args.iter_bundles: def consume_fn(ds): - for _ in ds.iter_internal_ref_bundles(): - pass + consume_ref_bundles(ds) elif args.iter_batches: diff --git a/release/nightly_tests/dataset/read_from_uris_benchmark.py b/release/nightly_tests/dataset/read_from_uris_benchmark.py index eed591aaf502..91f6a21831fd 100644 --- a/release/nightly_tests/dataset/read_from_uris_benchmark.py +++ b/release/nightly_tests/dataset/read_from_uris_benchmark.py @@ -7,7 +7,7 @@ import ray from ray.data.expressions import download -from benchmark import Benchmark +from benchmark import Benchmark, collect_operator_metrics, consume_ref_bundles BUCKET = "anyscale-imagenet" # This Parquet file contains the keys of images in the 'anyscale-imagenet' bucket. @@ -43,8 +43,12 @@ def convert_key(table): ds = metadata.map_batches(convert_key, batch_format="pyarrow") ds = ds.with_column("image_bytes", download("key")) ds = ds.map_batches(decode_images) - for _ in ds.iter_internal_ref_bundles(): - pass + consume_ref_bundles(ds) + + # Per-operator wall time + per-task USS/RSS (avg and worst task): separates + # the tiny metadata.parquet read from the download/decode maps, so a memory + # regression here can be attributed to the right operator. + return collect_operator_metrics(ds) if __name__ == "__main__": diff --git a/release/nightly_tests/dataset/single_node_all_to_all_compute.yaml b/release/nightly_tests/dataset/single_node_all_to_all_compute.yaml new file mode 100644 index 000000000000..93ea6b1856eb --- /dev/null +++ b/release/nightly_tests/dataset/single_node_all_to_all_compute.yaml @@ -0,0 +1,20 @@ +# Single-fat-node analog of the all-to-all fleets, for the arrow-rs 2x2 +# topology axis (arrow_rs_probe/gen_2x2_release_tests.py). NOTE: unlike the +# CPU analog, one node CANNOT match this fleet's aggregate (32x m5.4xlarge = +# 512 vCPU / 2 TiB; largest m5 is 96 vCPU / 384 GiB). Same story for joins' +# 100x m5.2xlarge. The 2x2 comparison is reader-vs-reader WITHIN a topology, +# so under-provisioning is shared by both arms; expect longer walls and real +# spill (the generated single-node entries set RAYTEST_FAIL_ON_SPILLING=0 +# and double the timeout). 2 TiB disk absorbs the spill that a fleet would +# spread over 32 x 512 GiB volumes. +cloud: {{env["ANYSCALE_CLOUD_NAME"]}} + +advanced_instance_config: + BlockDeviceMappings: + - DeviceName: /dev/sda1 + Ebs: + DeleteOnTermination: true + VolumeSize: 2048 + +head_node: + instance_type: m5.24xlarge diff --git a/release/nightly_tests/dataset/single_node_cpu_compute.yaml b/release/nightly_tests/dataset/single_node_cpu_compute.yaml new file mode 100644 index 000000000000..4bb1251a69e6 --- /dev/null +++ b/release/nightly_tests/dataset/single_node_cpu_compute.yaml @@ -0,0 +1,24 @@ +# Single-fat-node analog of the m5.2xlarge CPU fleets, for the arrow-rs 2x2 +# topology axis (arrow_rs_probe/gen_2x2_release_tests.py). One m5.24xlarge +# (96 vCPU / 384 GiB) meets or exceeds every CPU fleet it stands in for: +# fixed/autoscaling_cpu = 10x m5.2xlarge (80 vCPU / 320 GiB), mix's +# compute_8_cpu = 8x (64 / 256). Head-only: the point is that every task, +# and every long-lived worker, lands on ONE node. +# +# 512 GiB disk: the generated single-node cells run with +# RAYTEST_FAIL_ON_SPILLING=0 because one node absorbs spill a fleet spreads, +# and Ray spills to the root volume -- so provision for it (worst case is the +# whole ~115 GiB object store spilling). Matches the per-node volume the +# all-to-all fleet gives its workers. +cloud: {{env["ANYSCALE_CLOUD_NAME"]}} + +advanced_instance_config: + IamInstanceProfile: {"Name": "ray-autoscaler-v1"} + BlockDeviceMappings: + - DeviceName: /dev/sda1 + Ebs: + DeleteOnTermination: true + VolumeSize: 512 + +head_node: + instance_type: m5.24xlarge diff --git a/release/nightly_tests/dataset/sort_benchmark.py b/release/nightly_tests/dataset/sort_benchmark.py index 4a25e38e8753..301d0abcffa1 100644 --- a/release/nightly_tests/dataset/sort_benchmark.py +++ b/release/nightly_tests/dataset/sort_benchmark.py @@ -6,7 +6,7 @@ import numpy as np import psutil -from benchmark import Benchmark +from benchmark import Benchmark, collect_operator_metrics import ray from ray._private.internal_api import memory_summary from ray.data._internal.util import _check_pyarrow_version, GiB @@ -174,6 +174,10 @@ def run_benchmark(args): "num_partitions": num_partitions, "partition_size": partition_size, "peak_driver_memory": maxrss, + # Per-operator wall time / output bytes / per-task USS+RSS. `peak_driver_ + # memory` above is the *driver*; the sort's memory lives in the workers, + # which only these fields see. + **collect_operator_metrics(ds), } # Wait until after the stats have been printed to raise any exceptions. diff --git a/release/nightly_tests/dataset/tpch/common.py b/release/nightly_tests/dataset/tpch/common.py index 1a23074dff4e..245fbb613e9b 100644 --- a/release/nightly_tests/dataset/tpch/common.py +++ b/release/nightly_tests/dataset/tpch/common.py @@ -121,7 +121,35 @@ def load_table( return ds +# The last dataset a query handed to ``record_dataset``. Module-global because +# the query's final Dataset handle is otherwise local to ``benchmark_fn`` and +# its per-operator stats (shuffle/join walls, output-block granularity) are +# lost — tpch rows had no operator breakdown at all before this hook. +_recorded_dataset = None + + +def record_dataset(ds): + """Stash a query's final materialized dataset so ``run_tpch_benchmark`` can + attach per-operator metrics to the result row. Returns ``ds`` unchanged so + call sites can wrap an existing ``.materialize()`` expression. If called + more than once in a query, the last call wins.""" + global _recorded_dataset + _recorded_dataset = ds + return ds + + def run_tpch_benchmark(name: str, benchmark_fn): + from benchmark import collect_operator_metrics + + def instrumented_fn(): + global _recorded_dataset + _recorded_dataset = None + out = benchmark_fn() or {} + if _recorded_dataset is not None: + out = {**out, **collect_operator_metrics(_recorded_dataset)} + _recorded_dataset = None + return out + benchmark = Benchmark() - benchmark.run_fn(name, benchmark_fn) + benchmark.run_fn(name, instrumented_fn) benchmark.write_result() diff --git a/release/nightly_tests/dataset/tpch/tpch_q1.py b/release/nightly_tests/dataset/tpch/tpch_q1.py index 2fd6d1102193..6f1eed329c57 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q1.py +++ b/release/nightly_tests/dataset/tpch/tpch_q1.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Count, Mean, Sum from ray.data.expressions import col -from common import parse_tpch_args, load_table, to_f64, run_tpch_benchmark +from common import ( + parse_tpch_args, + load_table, + to_f64, + run_tpch_benchmark, + record_dataset, +) def main(args): @@ -41,7 +47,7 @@ def benchmark_fn(): ] ) - _ = ( + _ = record_dataset( ds.groupby(["l_returnflag", "l_linestatus"]) .aggregate( Sum(on="l_quantity_f", alias_name="sum_qty"), diff --git a/release/nightly_tests/dataset/tpch/tpch_q10.py b/release/nightly_tests/dataset/tpch/tpch_q10.py index 56c83d0da645..4a5860e65d41 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q10.py +++ b/release/nightly_tests/dataset/tpch/tpch_q10.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Sum from ray.data.expressions import col -from common import parse_tpch_args, load_table, to_f64, run_tpch_benchmark +from common import ( + parse_tpch_args, + load_table, + to_f64, + run_tpch_benchmark, + record_dataset, +) def main(args): @@ -111,7 +117,7 @@ def benchmark_fn(): ) # Aggregate by customer key, customer name, address, phone, account balance, and nation - _ = ( + _ = record_dataset( ds.groupby( [ "o_custkey", diff --git a/release/nightly_tests/dataset/tpch/tpch_q11.py b/release/nightly_tests/dataset/tpch/tpch_q11.py index c85395ccd937..5a34a04e4756 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q11.py +++ b/release/nightly_tests/dataset/tpch/tpch_q11.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Sum from ray.data.expressions import col -from common import parse_tpch_args, load_table, to_f64, run_tpch_benchmark +from common import ( + parse_tpch_args, + load_table, + to_f64, + run_tpch_benchmark, + record_dataset, +) def main(args): @@ -78,7 +84,7 @@ def benchmark_fn(): total = partsupp_germany.aggregate(Sum(on="value", alias_name="total"))["total"] threshold = total * fraction - _ = ( + _ = record_dataset( partsupp_germany.groupby("ps_partkey") .aggregate(Sum(on="value", alias_name="value")) .filter(expr=col("value") > threshold) diff --git a/release/nightly_tests/dataset/tpch/tpch_q12.py b/release/nightly_tests/dataset/tpch/tpch_q12.py index 0802d8deaae7..3084c878fd7e 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q12.py +++ b/release/nightly_tests/dataset/tpch/tpch_q12.py @@ -2,7 +2,7 @@ from ray.data.aggregate import Sum from ray.data.datatype import DataType from ray.data.expressions import col -from common import parse_tpch_args, load_table, run_tpch_benchmark +from common import parse_tpch_args, load_table, run_tpch_benchmark, record_dataset def main(args): @@ -81,7 +81,7 @@ def benchmark_fn(): col("o_orderpriority").not_in(high_priorities).cast(DataType.int64()), ) - _ = ( + _ = record_dataset( joined.groupby("l_shipmode") .aggregate( Sum(on="high_line_count", alias_name="high_line_count"), diff --git a/release/nightly_tests/dataset/tpch/tpch_q13.py b/release/nightly_tests/dataset/tpch/tpch_q13.py index 0a11facd299a..5896d0665343 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q13.py +++ b/release/nightly_tests/dataset/tpch/tpch_q13.py @@ -1,7 +1,7 @@ import ray from ray.data.aggregate import Count from ray.data.expressions import col -from common import parse_tpch_args, load_table, run_tpch_benchmark +from common import parse_tpch_args, load_table, run_tpch_benchmark, record_dataset def main(args): @@ -61,7 +61,7 @@ def benchmark_fn(): # ... # GROUP BY c_count # ORDER BY custdist DESC, c_count DESC - _ = ( + _ = record_dataset( c_orders.groupby(["c_count"]) .aggregate(Count(alias_name="custdist")) .sort(key=["custdist", "c_count"], descending=[True, True]) diff --git a/release/nightly_tests/dataset/tpch/tpch_q15.py b/release/nightly_tests/dataset/tpch/tpch_q15.py index ba5349639e6e..9d548cc04bc4 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q15.py +++ b/release/nightly_tests/dataset/tpch/tpch_q15.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Max, Sum from ray.data.expressions import col -from common import parse_tpch_args, load_table, to_f64, run_tpch_benchmark +from common import ( + parse_tpch_args, + load_table, + to_f64, + run_tpch_benchmark, + record_dataset, +) def main(args): @@ -62,7 +68,7 @@ def benchmark_fn(): top = revenue.filter(expr=col("total_revenue") == max_revenue) - _ = ( + _ = record_dataset( supplier.join( top, join_type="inner", diff --git a/release/nightly_tests/dataset/tpch/tpch_q17.py b/release/nightly_tests/dataset/tpch/tpch_q17.py index 72df030f0eaa..371d66d26d15 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q17.py +++ b/release/nightly_tests/dataset/tpch/tpch_q17.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Mean, Sum from ray.data.expressions import col -from common import load_table, parse_tpch_args, run_tpch_benchmark, to_f64 +from common import ( + load_table, + parse_tpch_args, + run_tpch_benchmark, + to_f64, + record_dataset, +) def main(args): @@ -50,7 +56,7 @@ def benchmark_fn(): # result for dual consumption (avg_qty groupby + filter pipeline). # This avoids a double S3 read of lineitem and reduces the groupby # from the full lineitem table to only matching rows. - joined = ( + joined = record_dataset( part_filtered.join( lineitem, join_type="inner", diff --git a/release/nightly_tests/dataset/tpch/tpch_q18.py b/release/nightly_tests/dataset/tpch/tpch_q18.py index 40bb5db11fbc..5eea5165d5ef 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q18.py +++ b/release/nightly_tests/dataset/tpch/tpch_q18.py @@ -1,7 +1,7 @@ import ray from ray.data.aggregate import Sum from ray.data.expressions import col -from common import parse_tpch_args, load_table, run_tpch_benchmark +from common import parse_tpch_args, load_table, run_tpch_benchmark, record_dataset def main(args): @@ -83,7 +83,7 @@ def benchmark_fn(): ) # Aggregate by customer name, customer key, order key, and order date - _ = ( + _ = record_dataset( ds.groupby( ["c_name", "o_custkey", "l_orderkey", "o_orderdate", "o_totalprice"] ) diff --git a/release/nightly_tests/dataset/tpch/tpch_q2.py b/release/nightly_tests/dataset/tpch/tpch_q2.py index b4648c88ffc6..59b99633dd24 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q2.py +++ b/release/nightly_tests/dataset/tpch/tpch_q2.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Min from ray.data.expressions import col -from common import parse_tpch_args, load_table, to_f64, run_tpch_benchmark +from common import ( + parse_tpch_args, + load_table, + to_f64, + run_tpch_benchmark, + record_dataset, +) def main(args): @@ -175,7 +181,7 @@ def benchmark_fn(): ds = ds.with_column("s_acctbal", to_f64(col("s_acctbal"))) # Select output columns, sort, and limit - _ = ( + _ = record_dataset( ds.select_columns( [ "s_acctbal", diff --git a/release/nightly_tests/dataset/tpch/tpch_q20.py b/release/nightly_tests/dataset/tpch/tpch_q20.py index e23a687dc759..e30d1f388e14 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q20.py +++ b/release/nightly_tests/dataset/tpch/tpch_q20.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Sum from ray.data.expressions import col -from common import parse_tpch_args, load_table, to_f64, run_tpch_benchmark +from common import ( + parse_tpch_args, + load_table, + to_f64, + run_tpch_benchmark, + record_dataset, +) def main(args): @@ -121,7 +127,7 @@ def benchmark_fn(): right_on=("ps_suppkey",), ) - _ = ( + _ = record_dataset( result.select_columns(["s_name", "s_address"]) .sort(key="s_name") .materialize() diff --git a/release/nightly_tests/dataset/tpch/tpch_q21.py b/release/nightly_tests/dataset/tpch/tpch_q21.py index ec16837118aa..d7d426e797e5 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q21.py +++ b/release/nightly_tests/dataset/tpch/tpch_q21.py @@ -1,7 +1,7 @@ import ray from ray.data.aggregate import Count, CountDistinct from ray.data.expressions import col -from common import parse_tpch_args, load_table, run_tpch_benchmark +from common import parse_tpch_args, load_table, run_tpch_benchmark, record_dataset def main(args): @@ -141,7 +141,7 @@ def benchmark_fn(): ) # Group by supplier name, count, sort, and limit. - _ = ( + _ = record_dataset( ds.groupby("s_name") .aggregate(Count(alias_name="numwait")) .sort(key=["numwait", "s_name"], descending=[True, False]) diff --git a/release/nightly_tests/dataset/tpch/tpch_q22.py b/release/nightly_tests/dataset/tpch/tpch_q22.py index ca78a857f914..50bf2e2a7050 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q22.py +++ b/release/nightly_tests/dataset/tpch/tpch_q22.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Count, Sum from ray.data.expressions import col -from common import parse_tpch_args, load_table, to_f64, run_tpch_benchmark +from common import ( + parse_tpch_args, + load_table, + to_f64, + run_tpch_benchmark, + record_dataset, +) def main(args): @@ -74,7 +80,7 @@ def benchmark_fn(): ) # Group by country code, aggregate count and total balance. - _ = ( + _ = record_dataset( custsale.groupby("cntrycode") .aggregate( Count(alias_name="numcust"), diff --git a/release/nightly_tests/dataset/tpch/tpch_q3.py b/release/nightly_tests/dataset/tpch/tpch_q3.py index 13dac36d31c6..1dbca233d781 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q3.py +++ b/release/nightly_tests/dataset/tpch/tpch_q3.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Sum from ray.data.expressions import col -from common import parse_tpch_args, load_table, to_f64, run_tpch_benchmark +from common import ( + parse_tpch_args, + load_table, + to_f64, + run_tpch_benchmark, + record_dataset, +) def main(args): @@ -77,7 +83,7 @@ def benchmark_fn(): ) # Aggregate by order key, order date, and ship priority - _ = ( + _ = record_dataset( ds.groupby(["o_orderkey", "o_orderdate", "o_shippriority"]) .aggregate(Sum(on="revenue", alias_name="revenue")) .sort(key=["revenue", "o_orderdate"], descending=[True, False]) diff --git a/release/nightly_tests/dataset/tpch/tpch_q4.py b/release/nightly_tests/dataset/tpch/tpch_q4.py index b3992c6a9c78..954992493e30 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q4.py +++ b/release/nightly_tests/dataset/tpch/tpch_q4.py @@ -1,7 +1,7 @@ import ray from ray.data.aggregate import Count from ray.data.expressions import col -from common import load_table, parse_tpch_args, run_tpch_benchmark +from common import load_table, parse_tpch_args, run_tpch_benchmark, record_dataset def main(args): @@ -63,7 +63,7 @@ def benchmark_fn(): ) # Group by order priority and count. - _ = ( + _ = record_dataset( ds.groupby("o_orderpriority") .aggregate(Count(alias_name="order_count")) .sort(key="o_orderpriority") diff --git a/release/nightly_tests/dataset/tpch/tpch_q5.py b/release/nightly_tests/dataset/tpch/tpch_q5.py index 482bccde44a5..7541f1155186 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q5.py +++ b/release/nightly_tests/dataset/tpch/tpch_q5.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Sum from ray.data.expressions import col -from common import load_table, parse_tpch_args, run_tpch_benchmark, to_f64 +from common import ( + load_table, + parse_tpch_args, + run_tpch_benchmark, + to_f64, + record_dataset, +) def main(args): @@ -115,7 +121,7 @@ def benchmark_fn(): to_f64(col("l_extendedprice")) * (1 - to_f64(col("l_discount"))), ) - _ = ( + _ = record_dataset( ds.groupby("n_name") .aggregate(Sum(on="revenue", alias_name="revenue")) .sort(key="revenue", descending=True) diff --git a/release/nightly_tests/dataset/tpch/tpch_q7.py b/release/nightly_tests/dataset/tpch/tpch_q7.py index d0377581263a..30c190772fe8 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q7.py +++ b/release/nightly_tests/dataset/tpch/tpch_q7.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Sum from ray.data.expressions import col -from common import parse_tpch_args, load_table, to_f64, run_tpch_benchmark +from common import ( + parse_tpch_args, + load_table, + to_f64, + run_tpch_benchmark, + record_dataset, +) def main(args): @@ -137,7 +143,7 @@ def benchmark_fn(): ) # Aggregate by supplier nation, customer nation, and year - _ = ( + _ = record_dataset( ds.groupby(["n_name_supp", "n_name_cust", "l_year"]) .aggregate(Sum(on="revenue", alias_name="revenue")) .sort(key=["n_name_supp", "n_name_cust", "l_year"]) diff --git a/release/nightly_tests/dataset/tpch/tpch_q8.py b/release/nightly_tests/dataset/tpch/tpch_q8.py index e77b00886a82..008469405d61 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q8.py +++ b/release/nightly_tests/dataset/tpch/tpch_q8.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Sum from ray.data.expressions import col -from common import parse_tpch_args, load_table, to_f64, run_tpch_benchmark +from common import ( + parse_tpch_args, + load_table, + to_f64, + run_tpch_benchmark, + record_dataset, +) def main(args): @@ -184,7 +190,7 @@ def benchmark_fn(): ) # Select and sort by year - _ = ( + _ = record_dataset( result.select_columns(["o_year", "mkt_share"]) .sort(key="o_year") .materialize() diff --git a/release/nightly_tests/dataset/tpch/tpch_q9.py b/release/nightly_tests/dataset/tpch/tpch_q9.py index 6a9d46e3ff0e..0573bc0dc3d2 100644 --- a/release/nightly_tests/dataset/tpch/tpch_q9.py +++ b/release/nightly_tests/dataset/tpch/tpch_q9.py @@ -1,7 +1,13 @@ import ray from ray.data.aggregate import Sum from ray.data.expressions import col -from common import parse_tpch_args, load_table, to_f64, run_tpch_benchmark +from common import ( + parse_tpch_args, + load_table, + to_f64, + run_tpch_benchmark, + record_dataset, +) def main(args): @@ -157,7 +163,7 @@ def benchmark_fn(): ) # Aggregate by nation and year - _ = ( + _ = record_dataset( ds.groupby(["n_name", "o_year"]) .aggregate(Sum(on="profit", alias_name="profit")) .sort(key=["n_name", "o_year"], descending=[False, True]) diff --git a/release/nightly_tests/dataset/wide_schema_pipeline_benchmark.py b/release/nightly_tests/dataset/wide_schema_pipeline_benchmark.py index 373afb23e55f..3b12198fb853 100644 --- a/release/nightly_tests/dataset/wide_schema_pipeline_benchmark.py +++ b/release/nightly_tests/dataset/wide_schema_pipeline_benchmark.py @@ -2,7 +2,7 @@ from typing import Dict, Any import ray -from benchmark import Benchmark +from benchmark import Benchmark, collect_operator_metrics, consume_ref_bundles def parse_args() -> argparse.Namespace: @@ -35,8 +35,7 @@ def run_pipeline() -> Dict[str, Any]: """Run the data pipeline: read -> map_batches -> write""" ds = ray.data.read_parquet(input_path) - for _ in ds.iter_internal_ref_bundles(): - pass + consume_ref_bundles(ds) # Get dataset stats for reporting actual_num_columns = len(ds.schema().base_schema) @@ -45,6 +44,10 @@ def run_pipeline() -> Dict[str, Any]: "num_columns": actual_num_columns, "data_type": args.data_type, "input_path": input_path, + # Read-operator wall time + per-task USS/RSS (avg and worst task): + # isolates the parquet decode from downstream, and surfaces the + # per-worker memory the aggregate object-store peak can't see. + **collect_operator_metrics(ds), } # Run the timed benchmark diff --git a/release/ray_release/byod/byod_arrow_rs_parquet.sh b/release/ray_release/byod/byod_arrow_rs_parquet.sh new file mode 100755 index 000000000000..51753dda6f72 --- /dev/null +++ b/release/ray_release/byod/byod_arrow_rs_parquet.sh @@ -0,0 +1,47 @@ +#!/bin/bash +# Build + install the ray_data_arrow_rs native Parquet decoder into the +# release-test BYOD image, so tests that read Parquet exercise the Rust reader +# instead of PyArrow. Runs as a Docker RUN layer at IMAGE BUILD time (has +# internet). The crate is NOT in the Ray wheel or on PyPI, so without this the +# reader would raise on `import ray_data_arrow_rs`. +# +# We build from source (rather than shipping a prebuilt wheel): install a +# minimal Rust toolchain, fetch the crate source from the branch, and let +# `pip install ` drive maturin via the pyproject build backend. Only rustc +# needs to pre-exist; pip installs maturin itself in the isolated build env, and +# the committed Cargo.lock makes the dependency set reproducible. +# +# Build cost (a few minutes to compile arrow/parquet/object_store in release) +# is paid once per image build, is cached, and is NOT part of any test's +# measured runtime — it happens before the cluster boots. +set -exo pipefail + +# The crate source MUST match the Ray code under test: the reader passes +# keyword args that only exist in matching crate versions, so a drifted .so +# fails late (or worse, silently). Prefer the build's own commit +# (BUILDKITE_COMMIT, if the release pipeline exports it into this layer) so +# branch pushes during a build can't desync the pair; fall back to the branch +# head otherwise. The echo makes the chosen ref auditable in the image-build +# log — check it on the first run after any branch switch. +BRANCH="arrow-rs-on-64985" +REF="${BUILDKITE_COMMIT:-refs/heads/${BRANCH}}" +CRATE_SUBDIR="python/ray/data/_internal/datasource_v2/native/ray_data_arrow_rs" + +# Minimal stable Rust toolchain (~1 min for rustup itself). +curl -sSf https://sh.rustup.rs -o /tmp/rustup-init.sh +sh /tmp/rustup-init.sh -y --profile minimal --default-toolchain stable +export PATH="$HOME/.cargo/bin:$PATH" + +# Fetch the crate source at the pinned ref and build+install it. pip reads +# the maturin build-backend from pyproject.toml and compiles the extension. +# -f fails the pipe on HTTP errors (404 = bad ref) instead of feeding tar +# an error page; --strip-components drops the ref-dependent top-level dir. +echo "arrow-rs byod: fetching crate source at ${REF}" +mkdir -p /tmp/ray-src +curl -sfL "https://github.com/AarryaSaraf/ray/archive/${REF}.tar.gz" \ + | tar xz -C /tmp/ray-src --strip-components=1 +pip3 install --no-cache-dir "/tmp/ray-src/${CRATE_SUBDIR}" + +# Fail the image build loudly if the crate isn't importable / is a partial +# build, so it surfaces here instead of as a scanner error at test run time. +python3 -c "import ray_data_arrow_rs as m; assert hasattr(m, 'read_row_groups')" diff --git a/release/release_data_tests.yaml b/release/release_data_tests.yaml index aca5680bfe63..b20122751d4d 100644 --- a/release/release_data_tests.yaml +++ b/release/release_data_tests.yaml @@ -6,8 +6,22 @@ team: data cluster: + # WARNING: a test that declares its own ``byod.runtime_env`` REPLACES this + # list rather than extending it -- ``deep_update`` (release/ray_release/ + # util.py) merges dicts but overwrites lists. A test that sets one env var + # here silently loses every RAYTEST_FAIL_ON_* below, i.e. it stops failing + # on OOM / dead nodes / spilling. Re-list the four entries verbatim in any + # test that needs its own runtime_env, or deliberately override them (as + # ``heterogeneous_memory_batch_inference_multitenancy`` does). byod: + # TESTING COMMIT ONLY (revert before opening the PR): build + install the + # Rust ray_data_arrow_rs crate into every data-test image so the flipped + # use_arrow_rs_parquet_reader default (context.py) has a crate to load. + # Inherited by every test that does not define its own byod.post_build_script. + post_build_script: byod_arrow_rs_parquet.sh runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 # Enable verbose stats for resource manager (to troubleshoot autoscaling) - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 # Fail the test if a worker OOMs @@ -29,6 +43,18 @@ python: "3.10" cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: both scaling variants regressed when the fallback default + # moved 64 -> 128 MiB (autoscaling -20.9% -> +3.1%, fixed_size -6.1% -> + # -3.8% vs master). Pinned to hold the 64 MiB behaviour. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 cluster_compute: "{{scaling}}_cpu_compute.yaml" matrix: @@ -42,15 +68,28 @@ s3://ray-benchmark-data-internal-us-west-2/imagenet/parquet --format parquet --iter-bundles -- name: "read_large_parquet_{{scaling}}" +# NOTE: the two scaling variants are spelled out rather than driven by a matrix +# because they want opposite bin budgets, and a matrix would take the cartesian +# product. Files here are ~3.9 GiB uncompressed with ~69 MiB row groups, so the +# default budget can't fit two row groups and every one becomes its own read +# task (~5.8k tasks for 103 files). On a fixed cluster that task explosion +# dominates and a larger budget wins; on an autoscaling cluster the extra tasks +# give the autoscaler demand to act on earlier and the default wins. +- name: "read_large_parquet_fixed_size" python: "3.10" cluster: anyscale_sdk_2026: true - cluster_compute: "{{scaling}}_cpu_compute.yaml" - - matrix: - setup: - scaling: [fixed_size, autoscaling] + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 1 GiB packs ~14 row groups per bin, cutting the task count ~14x. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1073741824 + cluster_compute: "fixed_size_cpu_compute.yaml" run: timeout: 3600 @@ -64,10 +103,44 @@ s3://ray-benchmark-data-internal-us-west-2/large-parquet/ --format parquet --iter-bundles --memory 3650722201 +- name: "read_large_parquet_autoscaling" + python: "3.10" + cluster: + anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: -8.1% vs master at 64 MiB, +32.9% once the fallback default + # moved to 128 MiB. The largest Parquet-path regression from that change. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + # Default budget: -25% vs master here, against +2% at 4 GiB. + cluster_compute: "autoscaling_cpu_compute.yaml" + + run: + timeout: 3600 + # See the fixed_size variant for why --memory is pinned. + script: > + python read_and_consume_benchmark.py + s3://ray-benchmark-data-internal-us-west-2/large-parquet/ --format parquet + --iter-bundles --memory 3650722201 + - name: "read_images_{{scaling}}" python: "3.10" cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 cluster_compute: "{{scaling}}_cpu_compute.yaml" matrix: @@ -84,6 +157,16 @@ python: "3.10" cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: pinned for parity. NOTE: inert -- tfrecords, not Parquet. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 run: timeout: 3600 script: > @@ -91,15 +174,36 @@ s3://ray-benchmark-data-internal-us-west-2/imagenet/tfrecords --format tfrecords --iter-bundles -- name: "read_from_uris_{{scaling}}" +# NOTE: the two scaling variants are spelled out rather than driven by a matrix +# because only the autoscaling one wants a pinned bin budget. Leaving +# fixed_size unpinned keeps it following the fallback default. +- name: "read_from_uris_fixed_size" python: "3.10" cluster: anyscale_sdk_2026: true - cluster_compute: "{{scaling}}_cpu_compute.yaml" + cluster_compute: "fixed_size_cpu_compute.yaml" - matrix: - setup: - scaling: [fixed_size, autoscaling] + run: + timeout: 5400 + script: python read_from_uris_benchmark.py + +- name: "read_from_uris_autoscaling" + python: "3.10" + cluster: + anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: -17.4% vs master at 64 MiB, +0.3% once the fallback default + # moved to 128 MiB. fixed_size was unaffected (-0.4% -> -2.1%), hence + # the split above. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: "autoscaling_cpu_compute.yaml" run: timeout: 5400 @@ -113,6 +217,20 @@ python: "3.10" cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # sf1000 lineitem is read with its full 17-column schema: ~1.16 GiB + # uncompressed per file over 1000 files. 1.25 GiB packs ~one file per + # bin (~943 read tasks vs the ~1000 blocks master's round-robin + # partitioner produced); the default budget + # would produce ~18k tasks. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 run: timeout: 3600 script: > @@ -184,12 +302,36 @@ matrix: setup: - mode: [append, upsert, overwrite] + mode: [append, upsert] run: timeout: 4800 script: python iceberg_benchmark.py --mode {{mode}} +# Split out of the matrix above: only ``overwrite`` wants a pinned bin budget, +# so ``append``/``upsert`` keep following the fallback default. +- name: "iceberg_benchmark_overwrite" + python: "3.10" + cluster: + anyscale_sdk_2026: true + byod: + post_build_script: byod_install_pyiceberg.sh + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: +0.2% vs master at 64 MiB, +5.4% at 128 MiB. append/upsert + # were flat or improved, hence the split. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: iceberg_benchmark_compute.yaml + + run: + timeout: 4800 + script: python iceberg_benchmark.py --mode overwrite + ################### # Aggregation tests ################### @@ -198,6 +340,17 @@ python: "3.10" cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: autoscaling went -15.4% -> -1.7% and fixed_size +2.6% -> +3.9% + # when the fallback default moved 64 -> 128 MiB. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 cluster_compute: "{{scaling}}_cpu_compute.yaml" matrix: @@ -237,21 +390,124 @@ - name: "aggregate_groups_{{scaling}}_{{shuffle_strategy}}_{{columns}}" python: "3.10" + # The bin budget is per-variant rather than per-block: only + # fixed_size/sort_shuffle/84-groups needed retuning, and the other seven were + # measured healthy at 32 MiB, so they stay there rather than move onto an + # unmeasured value. ``setup`` is left empty and every combination is + # enumerated in ``adjustments`` so each can carry its own budget. + # + # Baseline for the eight sort_shuffle/hash_shuffle variants is release build + # #102329 (32 MiB everywhere) against master build #102323. The four + # hash_shuffle_v2 variants came from master after that build and have no + # measurement here yet. + # + # | variant | 32 MiB | master | delta | + # |----------------------------------|--------|--------|------------| + # | fixed_size / sort / 84 groups | 215.7s | 145.0s | +48.8% | + # | autoscaling / sort / 84 groups | 463.6s | 420.7s | +10.2% | + # | autoscaling / hash / 84 groups | 208.8s | 201.9s | +3.4% | + # | autoscaling / hash / 7M groups | 2216.7s| 2208.7s| +0.4% | + # | fixed_size / hash / 84 groups | 19.3s | 19.8s | -2.7% | + # | fixed_size / hash / 7M groups | 663.5s | 771.2s | -14.0% | + # | fixed_size / sort / 7M groups | 550.7s | 651.7s | -15.5% | + # | autoscaling / sort / 7M groups | 716.6s |1033.8s | -30.7% | matrix: setup: - scaling: [fixed_size, autoscaling] - shuffle_strategy: [sort_shuffle_pull_based, hash_shuffle, hash_shuffle_v2] - columns: - - "column08 column13 column14" # 84 groups - - "column02 column14" # 7M groups + scaling: [] + shuffle_strategy: [] + columns: [] + bin_packing_bytes: [] + adjustments: + # 256 MiB (4x the default): the one variant that needed retuning. It sat + # at 211s in #102329 on 32 MiB and 220-259s on smaller budgets, against + # master's ~145s; 256 MiB brings it to 137s, at or just under master. + # + # What's being tuned away is per-block sort-shuffle overhead, not bytes. + # At 32 MiB the scan emitted 4x the ``ListFiles`` outputs and ~1.5x the + # blocks into the shuffle (Sort Sample 3050 vs master's 2000) for an + # identical 20.8GB object-store peak and identical row count. With only + # 84 reduce partitions, block count is what dominates -- so this variant + # wants the same direction as ``map_groups`` (1.25 GiB on this same + # dataset), just less far. + - with: + scaling: fixed_size + shuffle_strategy: sort_shuffle_pull_based + columns: "column08 column13 column14" # 84 groups + bin_packing_bytes: 268435456 + # The remaining seven stay at 32 MiB, where they were measured. None is a + # tuned optimum -- they are a known-acceptable state, and 256 MiB is + # unmeasured for all of them. Worth a sweep given how far it moved the + # variant above, but not worth assuming. + - with: + scaling: fixed_size + shuffle_strategy: sort_shuffle_pull_based + columns: "column02 column14" # 7M groups + bin_packing_bytes: 33554432 + - with: + scaling: fixed_size + shuffle_strategy: hash_shuffle + columns: "column08 column13 column14" + bin_packing_bytes: 33554432 + - with: + scaling: fixed_size + shuffle_strategy: hash_shuffle + columns: "column02 column14" + bin_packing_bytes: 33554432 + - with: + scaling: autoscaling + shuffle_strategy: sort_shuffle_pull_based + columns: "column08 column13 column14" + bin_packing_bytes: 33554432 + - with: + scaling: autoscaling + shuffle_strategy: sort_shuffle_pull_based + columns: "column02 column14" + bin_packing_bytes: 33554432 + - with: + scaling: autoscaling + shuffle_strategy: hash_shuffle + columns: "column08 column13 column14" + bin_packing_bytes: 33554432 + - with: + scaling: autoscaling + shuffle_strategy: hash_shuffle + columns: "column02 column14" + bin_packing_bytes: 33554432 + # ``hash_shuffle_v2`` arrived from master and has never run on this + # branch, so there is no measurement to carry over. Left at the 64 MiB + # branch default rather than inheriting the 32 MiB above, which was a + # deliberate pin for the variants that had been measured at it. + - with: + scaling: fixed_size + shuffle_strategy: hash_shuffle_v2 + columns: "column08 column13 column14" + bin_packing_bytes: 67108864 + - with: + scaling: fixed_size + shuffle_strategy: hash_shuffle_v2 + columns: "column02 column14" + bin_packing_bytes: 67108864 + - with: + scaling: autoscaling + shuffle_strategy: hash_shuffle_v2 + columns: "column08 column13 column14" + bin_packing_bytes: 67108864 + - with: + scaling: autoscaling + shuffle_strategy: hash_shuffle_v2 + columns: "column02 column14" + bin_packing_bytes: 67108864 cluster: anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=1 - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES={{bin_packing_bytes}} cluster_compute: "{{scaling}}_all_to_all_compute.yaml" run: @@ -275,10 +531,15 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=1 - RAYTEST_FAIL_ON_SPILLING=0 + # Same sf100 lineitem full-schema read as the aggregate variants; see + # the note there for why 1.25 GiB. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 cluster_compute: "{{scaling}}_all_to_all_compute.yaml" run: @@ -316,10 +577,16 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=1 - RAYTEST_FAIL_ON_SPILLING=0 + # Both sides are read with their full schema: ~1.17 GiB uncompressed + # per lineitem file. 1.25 GiB packs ~one file per bin (~95 read tasks + # vs the 100 blocks master's round-robin partitioner produced). + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 cluster_compute: fixed_size_100_cpu_compute.yaml matrix: @@ -342,12 +609,25 @@ # Wide Schema tests ############### -- name: wide_schema_pipeline_{{data_type}} +# NOTE: `primitives` and `tensors` are spelled out rather than folded into the +# matrix below because each wants its own bin budget, and a matrix takes the +# cartesian product of its lists (see the read_large_parquet NOTE above). Their +# names match the matrix-generated ones exactly so baseline history is unbroken. +# +# The bin budget is compared against row-group `total_byte_size` -- data-page +# bytes only, post-encoding. For these 5000-column datasets that is a poor proxy +# for the work a read task does: `primitives` files are 21.4 MiB on S3 but +# report only 0.45 MiB, because dictionary + RLE collapses the repeated 500-char +# values and the real weight is dictionary pages plus a 5000-column footer +# carrying per-chunk min/max statistics. +- name: wide_schema_pipeline_primitives python: "3.10" cluster: anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 # Preserve the default verbose stats for resource manager. - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 @@ -355,17 +635,134 @@ - RAYTEST_FAIL_ON_SPILLING=1 # S3 tensor data was written by Ray 2.49-2.54 using cloudpickle. - RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA=1 + # 27 single-row-group files reporting 0.45 MiB each: at the 64 MiB + # default all 27 land in one shared bin, so the whole 579 MiB dataset + # reads as a single task. 256 KiB is below one row group, so each gets + # its own bin (indivisible-oversize relaxation in OnlineBinPacker._place) + # -> 27 tasks, the ceiling while row groups stay atomic. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=262144 + # + # 27 actors x batch size 1 == one footer read per actor, so all 27 + # footers are read in a single concurrent wave. Worth the extra actor + # startup here because these footers are unusually expensive: 5000 + # column chunks each carrying 500-char min/max statistics. + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=27 + - RAY_DATA_PARQUET_FOOTER_BATCH_SIZE=1 + - RAY_DATA_PARQUET_READER_IO_THREAD_COUNT=5000 cluster_compute: fixed_size_cpu_compute.yaml - matrix: - setup: - data_type: [primitives, tensors, objects, nested_structs] + run: + timeout: 300 + script: > + python wide_schema_pipeline_benchmark.py + --data-type primitives + +- name: wide_schema_pipeline_tensors + python: "3.10" + cluster: + anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + # Preserve the default verbose stats for resource manager. + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # S3 tensor data was written by Ray 2.49-2.54 using cloudpickle. + - RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA=1 + # 39 MiB sits just above these 37.25 MiB row groups, so one fills a bin + # and the next cannot join it -- one row group per read task, same + # placement the 64 MiB default already produced via its two-per-bin + # threshold, but pinned so it survives a change to the default. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=40894464 + # + # These datasets are small enough that the footer-reader pool is pure + # startup cost; keep it narrow and spread the few files across it. + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=22 + - RAY_DATA_PARQUET_FOOTER_BATCH_SIZE=1 + - RAY_DATA_PARQUET_READER_IO_THREAD_COUNT=5000 + cluster_compute: fixed_size_cpu_compute.yaml + + run: + timeout: 300 + script: > + python wide_schema_pipeline_benchmark.py + --data-type tensors + +- name: wide_schema_pipeline_nested_structs + python: "3.10" + cluster: + anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + # Preserve the default verbose stats for resource manager. + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # S3 tensor data was written by Ray 2.49-2.54 using cloudpickle. + - RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA=1 + # 21 single-row-group files reporting 1.30 MiB each: at the 64 MiB + # default all 21 land in one shared bin and the dataset reads as a + # single task. 2 MiB is the widest budget that still keeps one file per + # bin -- two would need 2.6 MiB -- so this is 21 tasks, the ceiling + # while row groups stay atomic. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=2097152 + # + # 21 actors x batch size 1 == one footer read per actor, so all 21 + # footers are read in a single concurrent wave -- same shape as + # `primitives`. Worth the extra actor startup here because these + # footers are unusually expensive: 15000 column chunks each carrying + # min/max statistics. + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=21 + - RAY_DATA_PARQUET_FOOTER_BATCH_SIZE=1 + - RAY_DATA_PARQUET_READER_IO_THREAD_COUNT=5000 + cluster_compute: fixed_size_cpu_compute.yaml + + run: + timeout: 300 + script: > + python wide_schema_pipeline_benchmark.py + --data-type nested_structs + +- name: wide_schema_pipeline_objects + python: "3.10" + cluster: + anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + # Preserve the default verbose stats for resource manager. + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # S3 tensor data was written by Ray 2.49-2.54 using cloudpickle. + - RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA=1 + # 9 single-row-group files of 4.23 MiB: at the 64 MiB default all 9 land + # in one shared bin and the dataset reads as a single task. 4.20 MiB sits + # just under one row group, so each is oversized and gets its own bin; + # the smallest possible pair (2 x 2.71 MiB) also exceeds it, so no two + # files ever share -> 9 tasks, the ceiling while row groups stay atomic. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=4404020 + # + # 9 actors x batch size 1 == one footer read per actor, so all 9 footers + # are read in a single concurrent wave -- same shape as the other three. + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=9 + - RAY_DATA_PARQUET_FOOTER_BATCH_SIZE=1 + - RAY_DATA_PARQUET_READER_IO_THREAD_COUNT=5000 + cluster_compute: fixed_size_cpu_compute.yaml run: timeout: 300 script: > python wide_schema_pipeline_benchmark.py - --data-type {{data_type}} + --data-type objects ####################### # Streaming split tests @@ -375,6 +772,29 @@ python: "3.10" cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # imagenet parquet: 7416 single-row-group files of ~8.15 MiB, listed on + # one task. ListFiles dominates this test, and its cost tracks the number + # of results the driver pulls back rather than the footer IO itself: at + # the default result batch of 1 the listing task pays one object-store + # fetch per file (7.4k of them). Batching 50 files per actor call and per + # streamed result cuts that to ~149 fetches. + - RAY_DATA_PARQUET_FOOTER_BATCH_SIZE=50 + - RAY_DATA_PARQUET_FOOTER_RESULT_BATCH_SIZE=50 + # imagenet does not compress (uncompressed/on-disk = 1.00), so master's + # round-robin partitioner over-counted by 5x -- it sizes on-disk bytes + # through a fixed 5x encoding estimate -- and it landed on + # ~4 files (~31 MB) per block. Matching that granularity was tried at + # 36 MiB and made things worse -- it doubles the block count, and block + # generation is ~4 ms per manifest block. Left at the default. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 run: timeout: 300 wait_for_nodes: @@ -406,6 +826,24 @@ python: "3.10" cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # No bin budget override: at the default this dataset (513 single-row- + # group files of ~88 MB) already packs ~one file per bin, matching the + # granularity master's round-robin partitioner produced. + # + # This benchmark opens 8 datasets over the same path, so the footer + # path costs 8 x 513 footer reads and 8 x 32 reader actors -- work + # master never did, and enough to dominate the ~40s `random_mix` + # variants. + # Narrow the pool; the per-actor IO concurrency still covers 513 files. + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=4 cluster_compute: dataset_mixing/compute_8_cpu.yaml run: timeout: 600 @@ -414,12 +852,40 @@ variations: - __suffix__: 8ds_equal + cluster: + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB pin. NOTE: likely inert -- this dataset is 513 + # single-row-group files of ~88 MB, which pack one file per + # bin at both 64 and 128 MiB, so the granularity the block + # comment above describes is unchanged either way. Pinned + # because this variant moved (+17.0pp / +0.5pp) when the + # default changed; the cause is probably elsewhere. + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=4 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 run: script: > python dataset_mixing/mix_benchmark.py --num-datasets 8 --num-workers 16 --max-rows-per-worker 100000 - __suffix__: 8ds_power_law + cluster: + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=4 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 run: script: > python dataset_mixing/mix_benchmark.py --num-datasets 8 @@ -451,6 +917,17 @@ anyscale_sdk_2026: true byod: post_build_script: byod_install_mosaicml.sh + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # This benchmark hands in 200 explicit parquet paths, listed on a single + # task. At the default footer batch size of 10 that is 20 batches, so + # actors past the 20th are provisioned and never dispatched to. + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=20 cluster_compute: dataset/multi_node_train_16_workers.yaml run: @@ -465,10 +942,18 @@ cluster: byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=0 - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: +5.3% vs master at 64 MiB, +23.0% at 128 MiB. The + # ``regular`` variant went the other way, so only chaos is + # pinned. NOTE: this list replaces the parent's, so it also + # drops RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=20 -- pre-existing + # behaviour, left unchanged here. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 run: prepare: > python setup_chaos.py --kill-interval 200 --max-to-kill 1 --task-names @@ -480,6 +965,19 @@ cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: five of the six variants regressed when the fallback + # default moved 64 -> 128 MiB. s3_read_images_gpu improved and is + # exempted below. NOTE: the s3_url_image_* and s3_read_images_* + # variants read images, not Parquet, so the pin is inert for them. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 variations: - __suffix__: s3_parquet_cpu @@ -531,6 +1029,17 @@ - __suffix__: s3_read_images_gpu cluster: + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + # Exempt from the block-level 64 MiB pin: this variant improved + # at the 128 MiB default (-1.5% -> -2.4%), so it keeps following + # the default. + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 cluster_compute: dataset/fixed_size_xlarge_gpu_compute.yaml run: timeout: 4800 @@ -550,6 +1059,8 @@ byod: type: gpu runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION=0.5 # Preserve DEFAULTS' runtime_env (setting runtime_env here replaces, # doesn't merge). @@ -606,7 +1117,7 @@ anyscale_sdk_2026: true matrix: setup: - format: [numpy, pandas, pyarrow] + format: [numpy, pandas] run: timeout: 2400 @@ -615,38 +1126,96 @@ s3://ray-benchmark-data/tpch/parquet/sf10/lineitem --format parquet --iter-batches {{format}} -- name: to_tf +# Split out of the matrix above: only the pyarrow format wants a pinned bin +# budget, so numpy/pandas keep following the fallback default. +- name: "iter_batches_pyarrow" python: "3.10" cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: -2.2% vs master at 64 MiB, -0.0% at 128 MiB. numpy and pandas + # both improved at 128 MiB, hence the split. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + # No cluster_compute: inherit DEFAULTS (fixed_size_cpu_compute.yaml), which + # is what this test ran on inside the matrix. Pinning the same value here + # would only make this variant stop tracking numpy/pandas if DEFAULTS moves. + run: timeout: 2400 script: > python read_and_consume_benchmark.py - s3://air-example-data-2/100G-image-data-synthetic-raw/ --format image - --to-tf image image + s3://ray-benchmark-data/tpch/parquet/sf10/lineitem --format parquet + --iter-batches pyarrow -- name: iter_torch_batches +- name: to_tf python: "3.10" cluster: anyscale_sdk_2026: true - cluster_compute: fixed_size_gpu_head_compute.yaml - + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: +1.1% -> +4.4% vs master when the default moved to 128 MiB. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 run: timeout: 2400 script: > python read_and_consume_benchmark.py s3://air-example-data-2/100G-image-data-synthetic-raw/ --format image - --iter-torch-batches + --to-tf image image -########### -# Map tests -########### +- name: iter_torch_batches + python: "3.10" + cluster: + anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: -0.1% -> +1.1% vs master when the default moved to 128 MiB. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: fixed_size_gpu_head_compute.yaml + + run: + timeout: 2400 + script: > + python read_and_consume_benchmark.py + s3://air-example-data-2/100G-image-data-synthetic-raw/ --format image + --iter-torch-batches + +########### +# Map tests +########### - name: map python: "3.10" cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: -17.1% -> -14.1% vs master when the default moved to 128 MiB. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 run: timeout: 1800 script: python map_benchmark.py --api map --sf 100 @@ -655,6 +1224,16 @@ python: "3.10" cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: -23.2% -> -15.5% vs master when the default moved to 128 MiB. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 run: timeout: 1800 script: python map_benchmark.py --api flat_map --sf 100 @@ -690,6 +1269,23 @@ cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # sf1000 lineitem, full 17-column schema: 1000 files x 49 row groups of + # 24.36 MiB uncompressed = 1.16 GiB per file, and the distribution is + # flat (p0 24.35 MiB, p100 24.78 MiB). 1 GiB packs 42 row groups per bin + # at 99% fill -> ~1170 read tasks, close to the ~1000 one-task-per-file + # that DSv1 runs on master. The 64 MiB default is the worst case here: + # three row groups exceed it, so bins take two and sit 24% empty, giving + # 24.5k read tasks (measured 24,432). + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1073741824 + - RAY_DATA_CLUSTER_SCALING_UP_UTIL_THRESHOLD=0.6 cluster_compute: "{{scaling}}_cpu_compute.yaml" run: @@ -789,10 +1385,14 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=1 - RAYTEST_FAIL_ON_SPILLING=0 + # 64 MiB: -2.4% -> +4.0% vs master when the default moved to 128 MiB. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 cluster_compute: "{{scaling}}_all_to_all_compute.yaml" run: @@ -809,6 +1409,8 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=0 @@ -837,6 +1439,8 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=1 @@ -859,6 +1463,8 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=0 @@ -890,6 +1496,8 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=1 @@ -915,6 +1523,8 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=0 - RAYTEST_FAIL_ON_DEAD_NODES=0 @@ -962,6 +1572,8 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=0 @@ -1050,6 +1662,8 @@ cluster_compute: image_embedding_from_uris/{{cluster_type}}_cluster_compute.yaml byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES={{fail_on_dead_nodes}} @@ -1102,6 +1716,8 @@ cluster_compute: image_embedding_from_jsonl/{{cluster_type}}_cluster_compute.yaml byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES={{fail_on_dead_nodes}} @@ -1144,6 +1760,8 @@ cluster_compute: text_embedding/{{cluster_type}}_cluster_compute.yaml byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES={{fail_on_dead_nodes}} @@ -1185,16 +1803,69 @@ # TPCH Queries ############## -- name: "tpch_q1_{{scaling}}_{{shuffle_strategy}}" +# NOTE: the two scaling variants are spelled out rather than driven by a matrix +# because only the autoscaling one lowers the scale-up threshold, and a matrix +# cannot vary ``byod.runtime_env`` per setup value. The rendered test names are +# unchanged. Both variants keep the same 192 MiB bin budget. +- name: "tpch_q1_fixed_size_{{shuffle_strategy}}" python: "3.10" matrix: setup: - scaling: [fixed_size, autoscaling] shuffle_strategy: [hash_shuffle, hash_shuffle_v2] cluster: anyscale_sdk_2026: true - cluster_compute: "{{scaling}}_all_to_all_compute.yaml" + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # sf1000 lineitem: 1000 files x 49 row groups, 24.36 MiB uncompressed + # per row group at the full 17-column schema. q1 projects 7 of those + # columns and pre-filters on l_shipdate, so the packer sizes on a + # fraction of that -- which is why this lands well below the 1 GiB the + # full-schema readers of the same table (map_batches, write_parquet) + # want. 192 MiB measured best here, matching q20 on the same table. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=201326592 + cluster_compute: "fixed_size_all_to_all_compute.yaml" + + run: + timeout: 5400 + script: RAY_DATA_DEFAULT_SHUFFLE_STRATEGY={{shuffle_strategy}} python tpch/tpch_q1.py --sf 1000 + +- name: "tpch_q1_autoscaling_{{shuffle_strategy}}" + python: "3.10" + matrix: + setup: + shuffle_strategy: [hash_shuffle, hash_shuffle_v2] + + cluster: + anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # Same 192 MiB budget as the fixed_size variants; see the block above + # for how it was sized. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=201326592 + # Scale up at 60% logical utilization instead of the 0.75 default, as + # for map_batches. + # + # hash_shuffle read +53.9% vs master in #102732 and +55.5% in #102650 + # against the same master build. The gap is in the baseline, not the + # branch: both autoscaling variants land at 364-385s here, while master + # runs hash_shuffle in 249s and hash_shuffle_v2 in 367s. A branch floor + # that ignores the shuffle strategy points at scale-up latency rather + # than shuffle cost, which is what this threshold targets. + - RAY_DATA_CLUSTER_SCALING_UP_UTIL_THRESHOLD=0.6 + cluster_compute: "autoscaling_all_to_all_compute.yaml" run: timeout: 5400 @@ -1228,6 +1899,8 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=1 @@ -1332,6 +2005,8 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=1 @@ -1354,6 +2029,8 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 - RAYTEST_FAIL_ON_WORKER_OOM=1 - RAYTEST_FAIL_ON_DEAD_NODES=1 @@ -1496,33 +2173,101 @@ timeout: 5400 script: RAY_DATA_DEFAULT_SHUFFLE_STRATEGY={{shuffle_strategy}} python tpch/tpch_q15.py --sf 100 -- name: "tpch_q17_{{scaling}}_{{shuffle_strategy}}" +# NOTE: the two scaling variants are spelled out rather than driven by a matrix +# because only the autoscaling one pins a bin budget, and a matrix cannot vary +# ``byod.runtime_env`` per setup value. The rendered test names are unchanged. +- name: "tpch_q17_fixed_size_{{shuffle_strategy}}" python: "3.10" frequency: manual matrix: setup: - scaling: [fixed_size, autoscaling] shuffle_strategy: [hash_shuffle, hash_shuffle_v2] cluster: anyscale_sdk_2026: true - cluster_compute: "{{scaling}}_all_to_all_compute.yaml" + cluster_compute: "fixed_size_all_to_all_compute.yaml" run: timeout: 5400 script: RAY_DATA_DEFAULT_SHUFFLE_STRATEGY={{shuffle_strategy}} python tpch/tpch_q17.py --sf 100 -- name: "tpch_q18_{{scaling}}_{{shuffle_strategy}}" +- name: "tpch_q17_autoscaling_{{shuffle_strategy}}" python: "3.10" frequency: manual matrix: setup: - scaling: [fixed_size, autoscaling] shuffle_strategy: [hash_shuffle, hash_shuffle_v2] cluster: anyscale_sdk_2026: true - cluster_compute: "{{scaling}}_all_to_all_compute.yaml" + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB: smaller bins mean more read tasks, which give the autoscaler + # demand to act on earlier. The fixed_size variant runs ~5x faster on + # the same query at the default budget, so this is scaling-specific. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: "autoscaling_all_to_all_compute.yaml" + + run: + timeout: 5400 + script: RAY_DATA_DEFAULT_SHUFFLE_STRATEGY={{shuffle_strategy}} python tpch/tpch_q17.py --sf 100 + +# NOTE: the two scaling variants are spelled out rather than driven by a matrix +# because only the autoscaling one pins a bin budget, and a matrix cannot vary +# ``byod.runtime_env`` per setup value. The rendered test names are unchanged. +- name: "tpch_q18_fixed_size_{{shuffle_strategy}}" + python: "3.10" + frequency: manual + matrix: + setup: + shuffle_strategy: [hash_shuffle, hash_shuffle_v2] + + cluster: + anyscale_sdk_2026: true + cluster_compute: "fixed_size_all_to_all_compute.yaml" + + run: + timeout: 5400 + script: RAY_DATA_DEFAULT_SHUFFLE_STRATEGY={{shuffle_strategy}} python tpch/tpch_q18.py --sf 100 + +- name: "tpch_q18_autoscaling_{{shuffle_strategy}}" + python: "3.10" + frequency: manual + matrix: + setup: + shuffle_strategy: [hash_shuffle, hash_shuffle_v2] + + cluster: + anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 64 MiB, half the default, matching the q17 autoscaling pin. Smaller + # bins mean more and smaller read tasks, which both give the autoscaler + # demand to act on earlier and cut the peak block size feeding the + # shuffle. + # + # hash_shuffle_v2 has failed two builds running -- in #102732 the + # workload completed in 2075.9s and was then failed by + # RAYTEST_FAIL_ON_SPILLING on 4.62 GiB spilled against a 17.75 GiB + # object-store peak (61.6% utilization). That is spill under moderate + # pressure, not a cluster out of memory, which is the shape a smaller + # bin budget is most likely to help. Its sibling + # tpch_q21_autoscaling_hash_shuffle_v2 fails the same way and is left + # at the default for now as a control. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: "autoscaling_all_to_all_compute.yaml" run: timeout: 5400 @@ -1538,6 +2283,22 @@ cluster: anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # q20 projects 4 of lineitem's 17 columns -> ~160 MiB uncompressed per + # file. At the default budget that splits each file ~4 ways (400 blocks + # vs the 100 master's round-robin partitioner produced), and the extra + # blocks become + # extra shuffle-map tasks feeding HashAggregate. 192 MiB packs one file + # per bin. Per-query value: the same table under q21's 2-column + # projection is only ~92 MiB per file. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=201326592 cluster_compute: "{{scaling}}_all_to_all_compute.yaml" run: @@ -1578,6 +2339,42 @@ timeout: 5400 script: RAY_DATA_DEFAULT_SHUFFLE_STRATEGY={{shuffle_strategy}} python tpch/tpch_q21.py --sf 100 +# NOTE: fixed_size/hash_shuffle is spelled out separately because it is the one +# variant that pins a bin budget, and a matrix cannot vary ``byod.runtime_env`` +# per setup value. The other three keep tracking the default. The rendered test +# names are unchanged. +- name: tpch_q22_fixed_size_hash_shuffle + python: "3.10" + frequency: nightly + + cluster: + anyscale_sdk_2026: true + byod: + runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + # 16 MiB, an eighth of the 128 MiB default. This was the worst + # regression in release build #102732 at +76.7% vs master (23.5s vs + # 13.3s), and it reproduced at +72.0% in #102650 against the same + # master build with only +2.7% drift between the two, so the size of + # the gap is measured rather than inferred from one run. + # + # q22 is sf100 customer/orders, far smaller than the sf1000 lineitem + # tables that want 192 MiB (q1, q20) or 1 GiB (map_batches). At the + # 128 MiB default the whole scan collapses into a handful of bins, + # which starves a 23-second query of read parallelism; 16 MiB trades + # per-task overhead for enough tasks to fill the cluster. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=16777216 + cluster_compute: "fixed_size_all_to_all_compute.yaml" + + run: + timeout: 5400 + script: RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle python tpch/tpch_q22.py --sf 100 + - name: "tpch_q22_{{scaling}}_{{shuffle_strategy}}" python: "3.10" frequency: "{{frequency}}" @@ -1587,10 +2384,6 @@ frequency: [] shuffle_strategy: [] adjustments: - - with: - scaling: fixed_size - frequency: nightly - shuffle_strategy: hash_shuffle - with: scaling: fixed_size frequency: nightly @@ -1665,11 +2458,24 @@ anyscale_sdk_2026: true byod: runtime_env: + # Per-node worker-memory sampling with stage provenance (node_memory_monitor.py). + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 - RAY_health_check_period_ms=10000 - RAY_health_check_timeout_ms=100000 - RAY_health_check_failure_threshold=10 - RAY_gcs_rpc_server_connect_timeout_s=60 - RAYTEST_FAIL_ON_SPILLING=1 + # 1 GiB: 16x the default, so the scan yields far fewer, larger read + # tasks. ``map_batches`` wants this direction -- same as ``map_groups`` + # at 1.25 GiB, and the opposite of ``aggregate_groups`` at 32 MiB. + # + # This test read via V1 until the ``use_datasource_v2 = False`` pin came + # out of ``map_benchmark.py``. On V2 at the 64 MiB default it doubled + # object-store peak (310GB on V1 -> 616GB on V2, release build #101536 + # vs #102329) and ran 725s against 304s on V1. The larger budget cuts + # the block count to bring that peak back down; ``RAYTEST_FAIL_ON_ + # SPILLING=1`` above will fail the test if it moves the wrong way. + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1073741824 cluster_compute: dataset/cross_az_250_350_compute_gce.yaml run: @@ -1706,3 +2512,1017 @@ run: timeout: 3600 script: python autoscaling/does_not_over_provision.py + +# === BEGIN arrow-rs 2x2 (generated by arrow_rs_probe/gen_2x2_release_tests.py; do not edit by hand) === +# matrix: alloc +# {pa, rs} on the original fleet over the memory / sustained / control +# targets + 3 wall-fix confirmations; single cells only for the two shapes +# whose single-node reading is the question; two gate cells x3. rs = the +# arrow-rs reader as shipped, which since 2026-09-08 includes the +# end-of-stream malloc_trim (former rseos arm; rstrim retired, M108). +# frequency:manual -- trigger explicitly, all arms in one window (M75: +# readings drift across windows). Regenerate with: +# python nightly_tests/dataset/arrow_rs_probe/gen_2x2_release_tests.py + +# --- tpch_q18_fixed_size_hash_shuffle_v2 [wall] --- +- name: tpch_q18_fixed_size_hash_shuffle_v2_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle_v2 + python tpch/tpch_q18.py --sf 100 + +- name: tpch_q18_fixed_size_hash_shuffle_v2_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle_v2 + python tpch/tpch_q18.py --sf 100 + +# --- read_parquet_autoscaling [wall] --- +- name: read_parquet_autoscaling_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: autoscaling_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python read_and_consume_benchmark.py s3://ray-benchmark-data-internal-us-west-2/imagenet/parquet + --format parquet --iter-bundles + +- name: read_parquet_autoscaling_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: autoscaling_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python read_and_consume_benchmark.py s3://ray-benchmark-data-internal-us-west-2/imagenet/parquet + --format parquet --iter-bundles + +# --- map_groups_autoscaling_hash_shuffle_column02+column14 [wall] --- +- name: map_groups_autoscaling_hash_shuffle_column02+column14_2x2_pa_multi + frequency: manual + repeated_run: 3 + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: autoscaling_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python groupby_benchmark.py --sf 100 --map-groups + --group-by column02 column14 --shuffle-strategy hash_shuffle + +- name: map_groups_autoscaling_hash_shuffle_column02+column14_2x2_pa_single + frequency: manual + repeated_run: 3 + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: single_node_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 7200 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python groupby_benchmark.py --sf 100 --map-groups + --group-by column02 column14 --shuffle-strategy hash_shuffle + +- name: map_groups_autoscaling_hash_shuffle_column02+column14_2x2_rs_multi + frequency: manual + repeated_run: 3 + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: autoscaling_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python groupby_benchmark.py --sf 100 --map-groups + --group-by column02 column14 --shuffle-strategy hash_shuffle + +- name: map_groups_autoscaling_hash_shuffle_column02+column14_2x2_rs_single + frequency: manual + repeated_run: 3 + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: single_node_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 7200 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python groupby_benchmark.py --sf 100 --map-groups + --group-by column02 column14 --shuffle-strategy hash_shuffle + +# --- read_large_parquet_autoscaling [memory] --- +- name: read_large_parquet_autoscaling_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: autoscaling_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python read_and_consume_benchmark.py s3://ray-benchmark-data-internal-us-west-2/large-parquet/ + --format parquet --iter-bundles --memory 3650722201 + +- name: read_large_parquet_autoscaling_2x2_pa_single + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: single_node_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 7200 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python read_and_consume_benchmark.py s3://ray-benchmark-data-internal-us-west-2/large-parquet/ + --format parquet --iter-bundles --memory 3650722201 + +- name: read_large_parquet_autoscaling_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: autoscaling_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python read_and_consume_benchmark.py s3://ray-benchmark-data-internal-us-west-2/large-parquet/ + --format parquet --iter-bundles --memory 3650722201 + +- name: read_large_parquet_autoscaling_2x2_rs_single + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: single_node_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 7200 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python read_and_consume_benchmark.py s3://ray-benchmark-data-internal-us-west-2/large-parquet/ + --format parquet --iter-bundles --memory 3650722201 + +# --- read_large_parquet_fixed_size [memory] --- +- name: read_large_parquet_fixed_size_2x2_pa_multi + frequency: manual + repeated_run: 3 + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1073741824 + cluster_compute: fixed_size_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python read_and_consume_benchmark.py s3://ray-benchmark-data-internal-us-west-2/large-parquet/ + --format parquet --iter-bundles --memory 3650722201 + +- name: read_large_parquet_fixed_size_2x2_rs_multi + frequency: manual + repeated_run: 3 + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1073741824 + cluster_compute: fixed_size_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python read_and_consume_benchmark.py s3://ray-benchmark-data-internal-us-west-2/large-parquet/ + --format parquet --iter-bundles --memory 3650722201 + +# --- write_parquet [memory] --- +- name: write_parquet_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python read_and_consume_benchmark.py s3://ray-benchmark-data/tpch/parquet/sf1000/lineitem + --format parquet --write + +- name: write_parquet_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python read_and_consume_benchmark.py s3://ray-benchmark-data/tpch/parquet/sf1000/lineitem + --format parquet --write + +# --- tpch_q6_fixed_size_hash_shuffle [memory] --- +- name: tpch_q6_fixed_size_hash_shuffle_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle + python tpch/tpch_q6.py --sf 100 + +- name: tpch_q6_fixed_size_hash_shuffle_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle + python tpch/tpch_q6.py --sf 100 + +# --- tpch_q6_fixed_size_hash_shuffle_v2 [memory] --- +- name: tpch_q6_fixed_size_hash_shuffle_v2_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle_v2 + python tpch/tpch_q6.py --sf 100 + +- name: tpch_q6_fixed_size_hash_shuffle_v2_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle_v2 + python tpch/tpch_q6.py --sf 100 + +# --- tpch_q17_fixed_size_hash_shuffle [memory] --- +- name: tpch_q17_fixed_size_hash_shuffle_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle + python tpch/tpch_q17.py --sf 100 + +- name: tpch_q17_fixed_size_hash_shuffle_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle + python tpch/tpch_q17.py --sf 100 + +# --- tpch_q17_fixed_size_hash_shuffle_v2 [memory] --- +- name: tpch_q17_fixed_size_hash_shuffle_v2_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle_v2 + python tpch/tpch_q17.py --sf 100 + +- name: tpch_q17_fixed_size_hash_shuffle_v2_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle_v2 + python tpch/tpch_q17.py --sf 100 + +# --- map_groups_autoscaling_hash_shuffle_column08+column13+column14 [memory] --- +- name: map_groups_autoscaling_hash_shuffle_column08+column13+column14_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: autoscaling_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python groupby_benchmark.py --sf 100 --map-groups + --group-by column08 column13 column14 --shuffle-strategy hash_shuffle + +- name: map_groups_autoscaling_hash_shuffle_column08+column13+column14_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: autoscaling_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python groupby_benchmark.py --sf 100 --map-groups + --group-by column08 column13 column14 --shuffle-strategy hash_shuffle + +# --- wide_schema_pipeline_objects [memory] --- +- name: wide_schema_pipeline_objects_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=4404020 + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=9 + - RAY_DATA_PARQUET_FOOTER_BATCH_SIZE=1 + - RAY_DATA_PARQUET_READER_IO_THREAD_COUNT=5000 + cluster_compute: fixed_size_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 300 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python wide_schema_pipeline_benchmark.py + --data-type objects + +- name: wide_schema_pipeline_objects_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=4404020 + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=9 + - RAY_DATA_PARQUET_FOOTER_BATCH_SIZE=1 + - RAY_DATA_PARQUET_READER_IO_THREAD_COUNT=5000 + cluster_compute: fixed_size_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 300 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python wide_schema_pipeline_benchmark.py + --data-type objects + +# --- map_groups_fixed_size_hash_shuffle_v2_column08+column13+column14 [sustained] --- +- name: map_groups_fixed_size_hash_shuffle_v2_column08+column13+column14_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python groupby_benchmark.py --sf 100 --map-groups + --group-by column08 column13 column14 --shuffle-strategy hash_shuffle_v2 + +- name: map_groups_fixed_size_hash_shuffle_v2_column08+column13+column14_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python groupby_benchmark.py --sf 100 --map-groups + --group-by column08 column13 column14 --shuffle-strategy hash_shuffle_v2 + +# --- map_groups_fixed_size_hash_shuffle_v2_column02+column14 [sustained] --- +- name: map_groups_fixed_size_hash_shuffle_v2_column02+column14_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python groupby_benchmark.py --sf 100 --map-groups + --group-by column02 column14 --shuffle-strategy hash_shuffle_v2 + +- name: map_groups_fixed_size_hash_shuffle_v2_column02+column14_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python groupby_benchmark.py --sf 100 --map-groups + --group-by column02 column14 --shuffle-strategy hash_shuffle_v2 + +# --- map_groups_fixed_size_sort_shuffle_pull_based_column02+column14 [sustained] --- +- name: map_groups_fixed_size_sort_shuffle_pull_based_column02+column14_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python groupby_benchmark.py --sf 100 --map-groups + --group-by column02 column14 --shuffle-strategy sort_shuffle_pull_based + +- name: map_groups_fixed_size_sort_shuffle_pull_based_column02+column14_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python groupby_benchmark.py --sf 100 --map-groups + --group-by column02 column14 --shuffle-strategy sort_shuffle_pull_based + +# --- joins_sf100_inner [sustained] --- +- name: joins_sf100_inner_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_100_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python join_benchmark.py --left_dataset + s3://ray-benchmark-data/tpch/parquet/sf100/lineitem --right_dataset s3://ray-benchmark-data/tpch/parquet/sf100/orders + --left_join_keys column00 --right_join_keys column0 --join_type inner --num_partitions + 50 + +- name: joins_sf100_inner_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_100_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python join_benchmark.py --left_dataset + s3://ray-benchmark-data/tpch/parquet/sf100/lineitem --right_dataset s3://ray-benchmark-data/tpch/parquet/sf100/orders + --left_join_keys column00 --right_join_keys column0 --join_type inner --num_partitions + 50 + +# --- joins_sf100_left_outer [sustained] --- +- name: joins_sf100_left_outer_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_100_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python join_benchmark.py --left_dataset + s3://ray-benchmark-data/tpch/parquet/sf100/lineitem --right_dataset s3://ray-benchmark-data/tpch/parquet/sf100/orders + --left_join_keys column00 --right_join_keys column0 --join_type left_outer --num_partitions + 50 + +- name: joins_sf100_left_outer_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_100_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python join_benchmark.py --left_dataset + s3://ray-benchmark-data/tpch/parquet/sf100/lineitem --right_dataset s3://ray-benchmark-data/tpch/parquet/sf100/orders + --left_join_keys column00 --right_join_keys column0 --join_type left_outer --num_partitions + 50 + +# --- joins_sf100_full_outer [sustained] --- +- name: joins_sf100_full_outer_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_100_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python join_benchmark.py --left_dataset + s3://ray-benchmark-data/tpch/parquet/sf100/lineitem --right_dataset s3://ray-benchmark-data/tpch/parquet/sf100/orders + --left_join_keys column00 --right_join_keys column0 --join_type full_outer --num_partitions + 50 + +- name: joins_sf100_full_outer_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=1342177280 + cluster_compute: fixed_size_100_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python join_benchmark.py --left_dataset + s3://ray-benchmark-data/tpch/parquet/sf100/lineitem --right_dataset s3://ray-benchmark-data/tpch/parquet/sf100/orders + --left_join_keys column00 --right_join_keys column0 --join_type full_outer --num_partitions + 50 + +# --- wide_schema_pipeline_primitives [sustained] --- +- name: wide_schema_pipeline_primitives_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=262144 + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=27 + - RAY_DATA_PARQUET_FOOTER_BATCH_SIZE=1 + - RAY_DATA_PARQUET_READER_IO_THREAD_COUNT=5000 + cluster_compute: fixed_size_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 300 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python wide_schema_pipeline_benchmark.py + --data-type primitives + +- name: wide_schema_pipeline_primitives_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_AUTOLOAD_CLOUDPICKLE_TENSOR_METADATA=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=262144 + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=27 + - RAY_DATA_PARQUET_FOOTER_BATCH_SIZE=1 + - RAY_DATA_PARQUET_READER_IO_THREAD_COUNT=5000 + cluster_compute: fixed_size_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 300 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python wide_schema_pipeline_benchmark.py + --data-type primitives + +# --- mix.8ds_equal [sustained] --- +- name: mix.8ds_equal_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=4 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: dataset_mixing/compute_8_cpu.yaml + anyscale_sdk_2026: true + run: + timeout: 600 + wait_for_nodes: + num_nodes: 8 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python dataset_mixing/mix_benchmark.py + --num-datasets 8 --num-workers 16 --max-rows-per-worker 100000 + +- name: mix.8ds_equal_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_FOOTER_NUM_ACTORS=4 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: dataset_mixing/compute_8_cpu.yaml + anyscale_sdk_2026: true + run: + timeout: 600 + wait_for_nodes: + num_nodes: 8 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python dataset_mixing/mix_benchmark.py + --num-datasets 8 --num-workers 16 --max-rows-per-worker 100000 + +# --- aggregate_groups_autoscaling_hash_shuffle_column08+column13+column14 [control] --- +- name: aggregate_groups_autoscaling_hash_shuffle_column08+column13+column14_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=33554432 + cluster_compute: autoscaling_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python groupby_benchmark.py --sf 100 --aggregate + --group-by column08 column13 column14 --shuffle-strategy hash_shuffle + +- name: aggregate_groups_autoscaling_hash_shuffle_column08+column13+column14_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=0 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=33554432 + cluster_compute: autoscaling_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 3600 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python groupby_benchmark.py --sf 100 --aggregate + --group-by column08 column13 column14 --shuffle-strategy hash_shuffle + +# --- tpch_q12_autoscaling_hash_shuffle [control] --- +- name: tpch_q12_autoscaling_hash_shuffle_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: autoscaling_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle + python tpch/tpch_q12.py --sf 100 + +- name: tpch_q12_autoscaling_hash_shuffle_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + cluster_compute: autoscaling_all_to_all_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 5400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 RAY_DATA_DEFAULT_SHUFFLE_STRATEGY=hash_shuffle + python tpch/tpch_q12.py --sf 100 + +# --- iter_batches_pyarrow [control] --- +- name: iter_batches_pyarrow_2x2_pa_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: fixed_size_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 2400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=0 python read_and_consume_benchmark.py s3://ray-benchmark-data/tpch/parquet/sf10/lineitem + --format parquet --iter-batches pyarrow + +- name: iter_batches_pyarrow_2x2_rs_multi + frequency: manual + python: '3.10' + cluster: + byod: + runtime_env: + - RAY_DATA_BENCH_NODE_MEM_MONITOR=1 + - RAY_DATA_DEBUG_RESOURCE_MANAGER=1 + - RAYTEST_FAIL_ON_WORKER_OOM=1 + - RAYTEST_FAIL_ON_DEAD_NODES=1 + - RAYTEST_FAIL_ON_SPILLING=1 + - RAY_DATA_PARQUET_BIN_PACKING_BYTES=67108864 + cluster_compute: fixed_size_cpu_compute.yaml + anyscale_sdk_2026: true + run: + timeout: 2400 + script: RAY_DATA_USE_ARROW_RS_PARQUET_READER=1 python read_and_consume_benchmark.py s3://ray-benchmark-data/tpch/parquet/sf10/lineitem + --format parquet --iter-batches pyarrow + +# === END arrow-rs 2x2 ===