wehub/ray-project--ray
1
0
Fork 0
Code Issues Pull Requests Actions 1 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:17:40 +08:00
commit f1825c8ceb
10096 changed files with 2364182 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ray/data/_internal/datasource_v2/readers/__init__.py
View File
python/ray/data/_internal/datasource_v2/readers/base_reader.py
+34
View File
@@ -0,0 +1,34 @@
from abc import ABC, abstractmethod
from typing import Generic, Iterator
import pyarrow as pa
from ray.data._internal.datasource_v2 import InputSplit
from ray.util.annotations import DeveloperAPI
@DeveloperAPI
class Reader(ABC, Generic[InputSplit]):
"""Abstract base class for reading data from input buckets.
Readers execute on workers to actually read data. They receive an InputSplit
(e.g., FileManifest for file-based sources) and yield Arrow tables.
The Reader is created by Scanner.create_reader() and is configured with all
pushdown optimizations (columns, predicates, limits) that were applied.
"""
@abstractmethod
def read(self, input_split: InputSplit) -> Iterator[pa.Table]:
"""Read data from the input bucket and yield Arrow tables.
This method is called on workers to perform the actual read operation.
It should respect all pushdowns configured on this reader.
Args:
input_split: Work unit describing what data to read.
Returns:
Iterator[pa.Table]: Iterator of PyArrow Tables containing the read data.
"""
...
python/ray/data/_internal/datasource_v2/readers/file_reader.py
+513
View File
@@ -0,0 +1,513 @@
from enum import Enum
from functools import cached_property, partial
from typing import Any, Iterator, List, Optional, Set, Tuple
import pyarrow as pa
import pyarrow.dataset as pds
from pyarrow.fs import FileSystem, LocalFileSystem
from ray._common.utils import env_integer
from ray.data._internal.arrow_block import _BATCH_SIZE_PRESERVING_STUB_COL_NAME
from ray.data._internal.datasource.parquet_datasource import _compute_row_hashes
from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
from ray.data._internal.datasource_v2.readers.base_reader import Reader
from ray.data._internal.util import iterate_with_retry, make_async_gen
from ray.data.context import DataContext
from ray.data.datasource.partitioning import Partitioning, PathPartitionParser
from ray.data.expressions import Expr
from ray.util.annotations import DeveloperAPI
# Synthetic column name produced when ``include_paths=True``. Shared with
# the V2 datasource and scanner layers so all references are spelled the
# same way.
INCLUDE_PATHS_COLUMN_NAME = "path"
# https://arrow.apache.org/docs/python/generated/pyarrow.dataset.Scanner.html#pyarrow.dataset.Scanner.from_batches
# Default is specified by PyArrow.
_ARROW_DEFAULT_BATCH_SIZE = 131_072
# Number of batches read ahead per scanner. PyArrow's default is 16,
# which can retain a multi-GB working set when scanning jumbo tensor
# columns. 8 keeps I/O pipelined on remote filesystems for typical
# Parquet workloads without doubling memory peak. Drop to 1 via the
# env var when reading wide tensor columns.
_ARROW_SCANNER_BATCH_READAHEAD = env_integer(
"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"
class FileFormat(str, Enum):
PARQUET = "parquet"
CSV = "csv"
FEATHER = "feather"
JSON = "json"
ARROW = "arrow"
IPC = "ipc"
@DeveloperAPI
class FileReader(Reader[FileManifest]):
"""Reader for file-based sources.
This reader uses PyArrow's Dataset API which automatically handles:
- Column pruning
- Filter pushdown (row group pruning)
- Batch-level filtering
"""
def __init__(
self,
format: FileFormat,
batch_size: int = _ARROW_DEFAULT_BATCH_SIZE,
columns: Optional[List[str]] = None,
predicate: Optional[Expr] = None,
limit: Optional[int] = None,
filesystem: Optional[FileSystem] = None,
partitioning: Optional[Partitioning] = None,
ignore_prefixes: Optional[List[str]] = None,
include_paths: bool = False,
include_row_hash: bool = False,
schema: Optional[pa.Schema] = None,
):
"""Initialize the reader.
Refer to https://arrow.apache.org/docs/python/generated/pyarrow.dataset.dataset.html for more details.
Args:
format: Format of the files to read.
batch_size: Number of rows per batch.
columns: Columns to read. None means all columns.
predicate: Ray Data expression for filtering. Converted to a
PyArrow expression at the scanner-kwargs boundary.
limit: Maximum number of rows to read.
filesystem: Filesystem for reading files.
partitioning: Ray ``Partitioning`` object. Partition columns are
synthesized per-path via ``PathPartitionParser`` after each
batch is read, producing string-typed columns (V1 parity).
ignore_prefixes: Prefixes to ignore when reading files. Default is ['.', '_'] set by PyArrow.
include_paths: If True, include the source file path in a
``'path'`` column for each row.
include_row_hash: If True, include a deterministic uint64 hash
per row in a ``'row_hash'`` column. The hash is derived from
the source file path and the row's post-filter output
position within the fragment, matching V1 semantics. If a
``'row_hash'`` column already exists in the file, it is
overwritten.
schema: Caller-supplied unified schema used both to override
pyarrow's per-fragment inference (so a file whose column
is all-null doesn't pin the type to ``null``) and to cast
path-derived partition values to their target types when
``Partitioning(field_types=...)`` is set.
"""
self._format = format
self._columns = columns
self._predicate = predicate
self._batch_size = batch_size
self._limit = limit
self._filesystem = filesystem
self._partition_parser: Optional[PathPartitionParser] = (
PathPartitionParser(partitioning) if partitioning is not None else None
)
self._ignore_prefixes = ignore_prefixes
self._include_paths = include_paths
self._include_row_hash = include_row_hash
self._schema = schema
@cached_property
def _file_dataset_schema(self) -> Optional[pa.Schema]:
"""Schema passed to ``pds.dataset`` — partition keys and ``path``
stripped out since those are synthesized post-read.
Pinning the caller-supplied schema at the pyarrow layer is how
we cover the "first file has an all-null column, later files
have the real type" case (e.g.
``test_read_null_data_in_first_file``): without the pin,
pyarrow locks column X to ``null`` across the fragment group
and the later string-typed file fails the cast.
But pyarrow refuses extension-to-extension casts (e.g.
``ArrowTensorTypeV2(shape=X)`` → ``ArrowVariableShapedTensor``),
and files with different per-file tensor shapes only unify
through ``ArrowVariableShapedTensor``. When the caller schema
contains *any* extension column we skip the pin entirely and
let pyarrow infer per-file — downstream concat handles the
heterogeneous blocks. Losing the all-null promotion in this
narrow case is acceptable; the combination of an all-null
first file *and* an extension column is uncommon, whereas
reading multiple files with variable-shape tensors is a
supported V1 feature.
"""
if self._schema is None:
return None
if any(isinstance(f.type, pa.ExtensionType) for f in self._schema):
return None
partition_keys = (
set(self._partition_parser._scheme.field_names or [])
if self._partition_parser is not None
else set()
)
synthesized = {INCLUDE_PATHS_COLUMN_NAME}
if self._include_row_hash:
# ``row_hash`` is synthesized post-read, and the schema's type
# (``uint64``) may not match the on-disk column's type when a
# file already carries a ``row_hash`` column. Strip it from the
# dataset schema so pyarrow doesn't try to cast.
synthesized.add(ROW_HASH_COLUMN_NAME)
fields = [
f
for f in self._schema
if f.name not in partition_keys and f.name not in synthesized
]
return pa.schema(fields) if fields else None
def _broadcast_partition_value(
self, name: str, value: Any, num_rows: int
) -> pa.Array:
"""Broadcast a single path-derived partition value to ``num_rows``,
casting to the caller-supplied schema's field type if set.
Values are stringified first (``PathPartitionParser`` in
``explicit`` mode can return arrow-scalar-like non-strings) and
then cast to the target type, so ``Partitioning(field_types=
{"year": int})`` still promotes them correctly.
"""
str_val = None if value is None else str(value)
arr = pa.repeat(pa.scalar(str_val, type=pa.string()), num_rows)
if self._schema is not None:
idx = self._schema.get_field_index(name)
if idx != -1 and self._schema.field(idx).type != pa.string():
arr = arr.cast(self._schema.field(idx).type)
return arr
def read(self, input_split: FileManifest) -> Iterator[pa.Table]:
"""Read data from the input bucket and yield Arrow tables.
This method is called on workers to perform the actual read operation.
It should respect all pushdowns configured on this reader.
Args:
input_split: Work unit describing what data to read.
Yields:
pa.Table: PyArrow Tables containing the read data.
"""
if len(input_split) == 0:
return
# Dedupe paths before handing them to pyarrow. When chunking is on,
# a manifest can carry multiple rows per file (each describing a
# different row-group slice); pyarrow only needs one fragment per
# file, and ``_get_fragments_to_read`` then fans out chunk-level
# sub-fragments using the per-row chunk metadata.
paths = list(dict.fromkeys(list(input_split.paths)))
filesystem = self._filesystem or LocalFileSystem()
# Build a ``pds.Dataset`` over *all* manifest paths so pyarrow's
# listing + column metadata is shared, but then iterate its
# fragments one at a time. ``dataset.scanner(fragments=...)``
# at the aggregate level would force a cross-fragment cast —
# which breaks variable-shape tensor extensions where each
# file has its own ``ArrowTensorTypeV2(shape=...)``. Per-
# fragment scanners let pyarrow use the native per-file type,
# and downstream concat handles unification.
dataset = pds.dataset(
source=paths,
format=self._make_format(),
filesystem=filesystem,
schema=self._file_dataset_schema,
ignore_prefixes=self._ignore_prefixes,
)
# 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.
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
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_readahead": _ARROW_SCANNER_BATCH_READAHEAD,
}
scanner_kwargs.update(self._arrow_scanner_kwargs())
rows_read = 0
for table, fragment_path, fragment_row_offset in self._read_fragment_batches(
dataset, scanner_kwargs, input_split
):
if self._limit is not None:
if rows_read >= self._limit:
break
if len(table) > self._limit - rows_read:
table = table.slice(0, self._limit - rows_read)
# Build the list of (name, value) pairs to synthesize from
# the fragment path: hive partitions + optional ``path``.
derived_items: List[Tuple[str, Any]] = []
if self._partition_parser is not None:
derived_items.extend(self._partition_parser(fragment_path).items())
if self._include_paths:
derived_items.append((INCLUDE_PATHS_COLUMN_NAME, fragment_path))
for name, value in derived_items:
if (
columns_to_synthesize is not None
and name not in columns_to_synthesize
):
continue
if name in table.column_names:
# When the caller schema names a partition key, pyarrow
# expects it in every file and fills it with nulls when
# absent (the hive-typical case). Drop that placeholder
# so the path-derived value below replaces it.
table = table.drop([name])
table = table.append_column(
name,
self._broadcast_partition_value(name, value, table.num_rows),
)
# Skip when projection pushdown has narrowed ``columns`` to
# exclude ``row_hash`` — the projection below would just drop it.
if self._include_row_hash and (
columns_to_synthesize is None
or ROW_HASH_COLUMN_NAME in columns_to_synthesize
):
hashes = _compute_row_hashes(
fragment_path, fragment_row_offset, table.num_rows
)
if ROW_HASH_COLUMN_NAME in table.column_names:
table = table.drop([ROW_HASH_COLUMN_NAME])
table = table.append_column(
ROW_HASH_COLUMN_NAME, pa.array(hashes, type=pa.uint64())
)
if self._columns is not None:
# Project/reorder to the caller's requested column order;
# drop any that weren't produced (matches V1's lenient
# behavior). Always select — an empty projection must
# narrow the table to zero columns so the stub-column
# guard below handles row preservation.
produced = set(table.column_names)
projected = [c for c in self._columns if c in produced]
table = table.select(projected)
if table.num_columns == 0 and table.num_rows > 0:
# Guards against ``pa.concat_tables`` collapsing rows
# when a batch has zero columns (e.g., empty projection
# for a count query). The stub column is dropped by
# downstream projections.
table = table.append_column(
_BATCH_SIZE_PRESERVING_STUB_COL_NAME,
pa.nulls(table.num_rows),
)
self._on_batch_read(table)
rows_read += len(table)
yield table
def _resolve_batch_size(self, dataset: pds.Dataset) -> int:
"""Return the batch size to use for scanning.
Subclasses can override this to implement adaptive batch sizing.
"""
return self._batch_size
def _on_batch_read(self, table: pa.Table) -> None:
"""Hook called after each batch is read.
Subclasses can override this to update internal state (e.g., refine
batch size estimates from actual data).
"""
pass
def _arrow_scanner_kwargs(self) -> dict:
"""Additional keyword arguments passed to ``pds.Dataset.scanner()``.
Subclasses override this to inject format-specific options.
"""
return {}
def _make_format(self) -> Any:
"""Format passed to ``pds.dataset(format=...)``.
Defaults to the format string (e.g. ``"parquet"``); subclasses
override to return a configured ``pds.FileFormat`` instance when
format-specific options (read options, fragment scan options) need
to be threaded through.
"""
return self._format.value
def _get_fragments_to_read(
self,
dataset: pds.Dataset,
manifest: FileManifest,
) -> List[Tuple[pds.Fragment, int]]:
"""Return ``(fragment, file_row_offset)`` pairs to scan for this
manifest.
``file_row_offset`` is the cumulative pre-filter row count of all
rows in the underlying file that precede this fragment. It seeds
the per-fragment hashing offset so chunked sub-fragments of the
same file produce unique ``_compute_row_hashes`` keys instead of
colliding on ``(path, 0, n)``.
Default impl returns one ``(fragment, 0)`` per file in the dataset
(paths are deduped in :meth:`read` before the dataset is built).
Subclasses that support per-row chunk metadata
(e.g. :class:`ParquetFileReader`) override this to fan a single
file fragment out into N sub-fragments — one per row-group slice —
based on :attr:`FileManifest.file_chunk_metadatas`, each paired
with its starting row offset in the file.
"""
return [(fragment, 0) for fragment in dataset.get_fragments()]
def _read_fragment_batches(
self,
dataset: pds.Dataset,
scanner_kwargs: dict,
manifest: FileManifest,
) -> Iterator[Tuple[pa.Table, str, int]]:
"""Yield non-empty (table, fragment_path, fragment_row_offset) triples.
``fragment_row_offset`` is the post-filter row position of the first
row of ``table`` within its fragment. ``iterate_with_retry`` skips
already-yielded items on retry, so ``offset`` reflects only the
rows that actually surface to the caller — matching V1 row-hash
semantics even when a fragment fails partway through.
Retry is scoped per-fragment: if a fragment fails mid-read, only
that fragment is re-read (skipping batches already yielded).
Wrapping the whole manifest in a single retry would re-iterate
fragments that already succeeded and double-emit their batches.
Each fragment gets its own scanner so pyarrow uses the native
per-file schema. A cross-fragment scanner would force a unified
schema cast, which refuses extension-to-extension conversion
(e.g. variable-shape tensors). V1 ``ParquetDatasource`` follows
the same per-fragment pattern via ``fragment.to_batches``.
When ``RAY_DATA_READ_FILES_NUM_THREADS > 1`` and
``execution_options.preserve_order`` is False, fragments are
read concurrently via :func:`make_async_gen`. We still pass
``preserve_ordering=True`` so concurrent reads emit blocks in
fragment order; otherwise Ray Data task retries (block
reconstruction) could produce a different block sequence.
``make_async_gen`` consumes the whole input iterator up front
when preserving order. That is acceptable here because the input
is the finite fragment manifest from ``_get_fragments_to_read``,
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
# an early-fallback when the manifest has a single fragment.
# Subclasses (e.g. ``ParquetFileReader``) override
# ``_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)
if not fragments_with_offsets:
return
num_workers = min(_DEFAULT_NUM_THREADS, 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
)
return
# Set `preserve_ordering=True` to ensure deterministic output ordering.
# This is required so that Ray Data task retries (block reconstruction)
yield from make_async_gen(
base_iterator=iter(fragments_with_offsets),
fn=partial(self._read_fragments_sequential, scanner_kwargs=scanner_kwargs),
preserve_ordering=True,
num_workers=num_workers,
)
def _read_fragments_sequential(
self,
fragments_with_offsets: Iterator[Tuple[pds.Fragment, int]],
scanner_kwargs: dict,
) -> Iterator[Tuple[pa.Table, str, int]]:
"""Read each fragment in ``fragments_with_offsets`` in order, yielding
``(table, fragment_path, fragment_row_offset)`` triples.
Each input pair is ``(fragment, file_row_offset)``. The yielded
``fragment_row_offset`` starts at ``file_row_offset`` (the row
position of the fragment's first row within its underlying file)
and accumulates per yielded batch, so the per-fragment row-hash
math in :meth:`read` keys off the right window even when chunking
fans one file into multiple sub-fragments sharing ``fragment.path``.
``iterate_with_retry`` is scoped to a single fragment so a
transient I/O failure only re-reads the failing file (skipping
batches already yielded), not the whole input.
This is the per-worker body for the threaded path in
:meth:`_read_fragment_batches` (one thread per call, each
consuming a disjoint slice of fragments via ``make_async_gen``)
and is also the entire read loop for the sequential path.
"""
ctx = DataContext.get_current()
for fragment, file_row_offset in fragments_with_offsets:
offset = file_row_offset
for table in iterate_with_retry(
partial(self._iter_fragment_tables, fragment, scanner_kwargs),
f"read fragment {fragment.path}",
match=ctx.retried_io_errors,
):
if table.num_rows > 0:
yield table, fragment.path, offset
offset += table.num_rows
def _iter_fragment_tables(
self,
fragment: pds.Fragment,
scanner_kwargs: dict,
) -> Iterator[pa.Table]:
"""Yield Arrow tables for a single fragment.
Subclasses override this to swap in a format-specific reader for
fragments that don't fit the default scanner-based path (e.g.
Parquet's ARROW-5030 nested-type fallback).
When a non-extension caller schema is available we pin it at the
scanner so pyarrow null-fills any column the unified schema names
but the fragment lacks (V1 parity — ``ParquetDatasource`` passes
``read_schema`` to ``fragment.to_batches``). Falling back to the
per-fragment ``physical_schema`` preserves the variable-shape
tensor escape hatch already encoded in ``_file_dataset_schema``.
"""
fragment_schema = (
self._file_dataset_schema
if self._file_dataset_schema is not None
else fragment.physical_schema
)
scanner = fragment.scanner(**scanner_kwargs, schema=fragment_schema)
for tagged in scanner.scan_batches():
yield pa.Table.from_batches(batches=[tagged.record_batch])
python/ray/data/_internal/datasource_v2/readers/in_memory_size_estimator.py
+146
View File
@@ -0,0 +1,146 @@
from abc import ABC, abstractmethod
from typing import Optional
import numpy as np
from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
from ray.data._internal.datasource_v2.readers.file_reader import FileReader
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
from ray.data.block import BlockAccessor
from ray.util.annotations import DeveloperAPI
@DeveloperAPI
class InMemorySizeEstimator(ABC):
@abstractmethod
def estimate_in_memory_sizes(self, manifest: FileManifest) -> np.ndarray:
"""Estimate the in-memory sizes of the paths in the given manifest.
Some `FilePartitioner` implementations use this method to ensure that each
read task receives an appropriate amount of data. To ensure that file listing
is efficient, this method must be cheap to call, on average.
Args:
manifest: A manifest containing the paths and on-disk sizes of the files.
Returns:
The estimated in-memory sizes of the data in bytes.
"""
...
@DeveloperAPI
class SamplingInMemorySizeEstimator(InMemorySizeEstimator):
"""Estimates in-memory sizes by reading files.
This class estimates the in-memory size of files by multiplying the on-disk
size by an estimated encoding ratio. If an instance hasn't estimated an encoding
ratio yet, it'll read a file to estimate it. Otherwise, it'll use the previously
estimated encoding ratio.
TODO: This approach doesn't work well for formats that produce multiple batches
(because we assume a 1:1 encoding ratio) or for formats that vary in encoding
ratios (e.g. videos).
"""
def __init__(self, reader: "FileReader"):
self._reader = reader
self._encoding_ratio = None
def estimate_in_memory_sizes(self, manifest: FileManifest) -> np.ndarray:
assert np.all(manifest.file_sizes >= 0)
for path, file_size in zip(manifest.paths, manifest.file_sizes):
if self._encoding_ratio is None:
# Estimating the encoding ratio can be expensive since it requires
# reading the file. So, we only estimate the encoding ratio if we don't
# already have one.
self._encoding_ratio = self._estimate_encoding_ratio(path, file_size)
break
if self._encoding_ratio is None:
# If we couldn't estimate the encoding ratio, assume a 1:1 encoding ratio.
return manifest.file_sizes
else:
return manifest.file_sizes * self._encoding_ratio
def _estimate_encoding_ratio(
self,
path: str,
file_size: int,
) -> Optional[float]:
"""
Estimate the encoding ratio (in-memory size / on-disk size) for a file.
Args:
path: The path to the file.
file_size: The on-disk size of the file/chunk in bytes.
Returns:
The estimated encoding ratio of the file, or `None` if the ratio can't
be estimated.
"""
# If the file is empty, we can't estimate the encoding ratio.
if not file_size:
return None
# Use ``None`` chunk metadata: the size estimator reads the file whole
# to estimate the encoding ratio; chunk-level splitting is irrelevant here.
manifest = FileManifest.construct_manifest(
[path],
[file_size],
[None],
)
batches = self._reader.read(manifest)
try:
first_batch = next(batches)
except StopIteration:
# If there's no data, we can't estimate the encoding ratio.
return None
try:
# Try to read a second batch. If it succeeds, it means the file contains
# multiple batches.
next(batches)
except StopIteration:
# Each file contains exactly one batch.
builder = DelegatingBlockBuilder()
builder.add_batch(first_batch)
block = builder.build()
in_memory_size = BlockAccessor.for_block(block).size_bytes()
else:
# Each file contains multiple batches.
#
# NOTE: To avoid reading the entire file to estimate the encoding ratio,
# we assume the file is 1:1 encoded. We can't return `None` because if
# all files contain multiple batches, then we'd try to re-estimate the
# encoding ratio for every file, and that'd be very expensive.
in_memory_size = file_size
return in_memory_size / file_size
# Default Parquet encoding ratio: in-memory is ~5x on-disk size.
# Parquet uses columnar compression and encoding, so Arrow in-memory
# representation is significantly larger than the on-disk format.
PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT = 5
@DeveloperAPI
class ParquetInMemorySizeEstimator(InMemorySizeEstimator):
"""Estimates in-memory sizes for Parquet files using a fixed encoding ratio.
Parquet files are typically much smaller on disk than in memory due to
columnar compression and encoding. This estimator applies a constant
ratio (default 5x) to avoid the overhead of reading file metadata or
sampling data, which can be slow for Parquet files and hurt startup time.
"""
def __init__(self, encoding_ratio: float = PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT):
self._encoding_ratio = encoding_ratio
def estimate_in_memory_sizes(self, manifest: FileManifest) -> np.ndarray:
return self._encoding_ratio * manifest.file_sizes
python/ray/data/_internal/datasource_v2/readers/parquet_file_reader.py
+516
View File
@@ -0,0 +1,516 @@
import logging
import math
from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple
import pyarrow as pa
import pyarrow.dataset as pds
import pyarrow.parquet as pq
from pyarrow.fs import FileSystem
from typing_extensions import override
if TYPE_CHECKING:
from ray.data.datasource.partitioning import Partitioning
from ray._common.utils import env_bool, env_integer
from ray.data._internal.datasource.parquet_datasource import (
AUTOLOAD_PICKLE_OBJECT_SCALAR_ENV_VAR,
_check_for_pickle_object_columns,
)
from ray.data._internal.datasource_v2.chunkers.parquet_file_chunking_utils import (
_fragments_from_chunk_metadata,
)
from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
from ray.data._internal.datasource_v2.readers.file_reader import (
_ARROW_DEFAULT_BATCH_SIZE,
FileFormat,
FileReader,
)
from ray.data._internal.datasource_v2.readers.in_memory_size_estimator import (
PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT,
)
from ray.data._internal.util import MiB
from ray.data.expressions import Expr
from ray.util.annotations import DeveloperAPI
from ray.util.debug import log_once
logger = logging.getLogger(__name__)
_UNSET = object()
# 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.
_PARQUET_FRAGMENT_BUFFER_SIZE = env_integer(
"RAY_DATA_PARQUET_FRAGMENT_BUFFER_SIZE", 8 * MiB
)
def _estimate_batch_size_from_metadata(
fragment: pds.ParquetFileFragment,
columns: Optional[List[str]],
target_block_size: int,
) -> Optional[int]:
"""Estimate batch size from Parquet row group metadata without reading data.
Uses uncompressed column sizes from row group metadata and the encoding
ratio to estimate in-memory row size, then computes how many rows fit
in ``target_block_size``.
Args:
fragment: A PyArrow Parquet fragment with accessible metadata.
columns: Columns being read, or None for all columns.
target_block_size: Target in-memory size per batch in bytes.
Returns:
Estimated batch size in rows, or None if metadata is unavailable.
"""
try:
# Accessing `metadata` triggers I/O via `EnsureCompleteMetadata()` to
# read the Parquet footer. `check_status` maps C++ Status codes to
# Python exceptions: IOError (OSError) for I/O failures,
# ArrowInvalid for corrupt footers.
# https://github.com/apache/arrow/blob/apache-arrow-23.0.0/python/pyarrow/error.pxi#L110-L126
metadata: pq.FileMetaData = fragment.metadata
except (OSError, pa.ArrowInvalid) as e:
logger.debug("Failed to read Parquet metadata for batch size estimation: %s", e)
return None
if metadata is None or metadata.num_row_groups == 0:
return None
row_group_idx: int = (
fragment.row_groups[0].id if fragment.row_groups is not None else 0
)
row_group_meta: pq.RowGroupMetaData = metadata.row_group(row_group_idx)
row_group_num_rows: int = row_group_meta.num_rows
if row_group_num_rows == 0:
return None
if columns is not None:
projected_columns = tuple(columns)
target_column_indices = []
for col_idx in range(row_group_meta.num_columns):
leaf_path = row_group_meta.column(col_idx).path_in_schema
# Account for nested columns
if any(
leaf_path == col_name or leaf_path.startswith(f"{col_name}.")
for col_name in projected_columns
):
target_column_indices.append(col_idx)
row_group_uncompressed_size = sum(
row_group_meta.column(col_idx).total_uncompressed_size
for col_idx in target_column_indices
)
else:
# Sum per-column uncompressed sizes instead of using
# row_group_meta.total_byte_size, which can return the *compressed* size
# for some files (apache/arrow#48138).
row_group_uncompressed_size = sum(
row_group_meta.column(col_idx).total_uncompressed_size
for col_idx in range(row_group_meta.num_columns)
)
# Estimate the in-memory size of the row group
estimated_in_mem_row_group_size = (
row_group_uncompressed_size * PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT
)
estimated_in_mem_row_size = estimated_in_mem_row_group_size / row_group_num_rows
if estimated_in_mem_row_size == 0:
return None
# Never request more rows than the row group actually contains.
target_batch_size = min(
math.ceil(target_block_size / estimated_in_mem_row_size),
row_group_num_rows,
)
logger.debug(
f"Estimated target batch size to be: {target_batch_size} (with "
f"{target_block_size=} bytes and {estimated_in_mem_row_size=} bytes)"
)
return target_batch_size
@DeveloperAPI
class ParquetFileReader(FileReader):
"""Parquet-specific file reader with adaptive batch sizing.
Extends :class:`FileReader` with:
- **Metadata-based batch size estimation**: Uses Parquet row group metadata
(uncompressed column sizes) to estimate an optimal batch size before
reading any data.
- **Adaptive refinement**: After reading each batch, refines the batch size
estimate from actual in-memory sizes for subsequent reads.
For non-Parquet formats, use :class:`FileReader` directly.
"""
def __init__(
self,
batch_size: Optional[int] = None,
columns: Optional[List[str]] = None,
predicate: Optional[Expr] = None,
limit: Optional[int] = None,
filesystem: Optional[FileSystem] = None,
partitioning: "Optional[Partitioning]" = None,
ignore_prefixes: Optional[List[str]] = None,
target_block_size: Optional[int] = None,
include_paths: bool = False,
include_row_hash: bool = False,
schema: Optional[pa.Schema] = None,
parquet_format_kwargs: Optional[Dict[str, Any]] = None,
):
"""Initialize the Parquet reader.
Args:
batch_size: Explicit batch size override. If provided, disables
adaptive batch sizing.
columns: Columns to read. None means all columns.
predicate: Ray Data expression for filtering.
limit: Maximum number of rows to read.
filesystem: Filesystem for reading files.
partitioning: Ray ``Partitioning`` for synthesizing partition
columns from file paths.
ignore_prefixes: Prefixes to ignore when reading files.
target_block_size: Target in-memory size per batch in bytes.
Used for adaptive batch sizing when ``batch_size`` is not set.
include_paths: If True, include the source file path in a
``'path'`` column for each row.
include_row_hash: If True, include a deterministic uint64 hash
per row in a ``'row_hash'`` column.
schema: Caller-supplied unified schema forwarded to the base
:class:`FileReader` for per-fragment inference override
and partition-column type casting.
parquet_format_kwargs: Extra kwargs spread into
:class:`pyarrow.dataset.ParquetFileFormat` (e.g.
``coerce_int96_timestamp_unit``, ``pre_buffer``,
``dictionary_columns``). Used to forward the deprecated
``dataset_kwargs`` arg on the V2 path.
"""
super().__init__(
format=FileFormat.PARQUET,
batch_size=batch_size or _ARROW_DEFAULT_BATCH_SIZE,
columns=columns,
predicate=predicate,
limit=limit,
filesystem=filesystem,
partitioning=partitioning,
ignore_prefixes=ignore_prefixes,
include_paths=include_paths,
include_row_hash=include_row_hash,
schema=schema,
)
self._explicit_batch_size = batch_size
self._allow_pickle_object_columns = env_bool(
AUTOLOAD_PICKLE_OBJECT_SCALAR_ENV_VAR, False
)
self._target_block_size = target_block_size
self._parquet_format_kwargs: Dict[str, Any] = parquet_format_kwargs or {}
self._sampled_batch_size: int | object = (
_UNSET # pyrefly: ignore[bad-assignment]
)
@override
def _make_format(self) -> pds.ParquetFileFormat:
return pds.ParquetFileFormat(**self._parquet_format_kwargs)
@override
def _resolve_batch_size(self, dataset: pds.Dataset) -> int:
"""Determine batch size from explicit setting, metadata, or default.
Priority: explicit batch_size > sampled estimate > metadata estimate > default.
On the first call ``_sampled_batch_size`` is ``_UNSET``, so we fall
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.
"""
if self._explicit_batch_size is not None:
return self._explicit_batch_size
if self._sampled_batch_size is not _UNSET:
return self._sampled_batch_size # pyrefly: ignore[bad-return]
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
)
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
self._sampled_batch_size = batch_size
return batch_size
@override
def _on_batch_read(self, table: pa.Table) -> None:
"""Refine batch size estimate from actual in-memory data."""
if self._target_block_size is None or table.nbytes == 0 or table.num_rows == 0:
return
row_size = table.nbytes / table.num_rows
self._sampled_batch_size = max(math.ceil(self._target_block_size / row_size), 1)
@override
def _get_fragments_to_read(
self,
dataset: pds.Dataset,
manifest: FileManifest,
) -> List[Tuple[pds.Fragment, int]]:
"""Fan file fragments into chunk-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`;
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``.
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.
"""
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]
if chunk_metadata is None:
fragments.append((fragment, 0))
else:
fragments.extend(
_fragments_from_chunk_metadata(fragment, chunk_metadata)
)
return fragments
@override
def _iter_fragment_tables(
self,
fragment: pds.Fragment,
scanner_kwargs: dict,
) -> "Iterator[pa.Table]":
for table in self._iter_fragment_tables_without_pickle_check(
fragment, scanner_kwargs
):
if not self._allow_pickle_object_columns:
_check_for_pickle_object_columns(table)
yield table
def _iter_fragment_tables_without_pickle_check(
self,
fragment: pds.Fragment,
scanner_kwargs: dict,
) -> "Iterator[pa.Table]":
"""Use V1's nested-type fallback path when the fragment has nested
columns whose row-group size exceeds Arrow's ~2GB chunking limit
(ARROW-5030).
"""
import pyarrow.compute as pc
from ray.data._internal.arrow_ops.transform_pyarrow import (
_align_struct_fields,
)
from ray.data._internal.datasource.parquet_datasource import (
_get_safe_batch_size_for_nested_types,
_needs_nested_type_fallback,
_resolve_leaf_column_indices,
_resolve_read_columns,
)
from ray.data._internal.planner.plan_expression.expression_visitors import (
get_column_references,
)
columns = scanner_kwargs.get("columns")
filter_expr: pc.Expression = scanner_kwargs.get("filter")
# Include filter-referenced columns in the fallback check: a filter
# that touches a large nested column outside the projection still
# forces row-level decoding of that column, which would otherwise
# hit ARROW-5030 in the normal scanner path.
filter_columns = (
get_column_references(self._predicate)
if self._predicate is not None
else None
)
read_columns = _resolve_read_columns(columns, filter_expr, filter_columns)
if not _needs_nested_type_fallback(fragment, read_columns):
yield from super()._iter_fragment_tables(fragment, scanner_kwargs)
return
if log_once(f"parquet_nested_fallback_v2:{fragment.path}"):
logger.warning(
"Using pyarrow.parquet row-level batched reader for '%s' due "
"to Arrow nested type chunking limitation (ARROW-5030). "
"Consider writing Parquet files with smaller row group sizes "
"to avoid this.",
fragment.path,
)
batch_size = scanner_kwargs.get("batch_size")
pf = pq.ParquetFile(
fragment.path,
filesystem=fragment.filesystem, # pyrefly: ignore[unexpected-keyword]
)
# Scope the safe batch-size calculation to the columns actually being
# decoded so we don't shrink batches based on columns we won't read.
leaf_indices = (
_resolve_leaf_column_indices(pf.metadata, read_columns)
if read_columns is not None and pf.metadata.num_row_groups > 0
else None
)
safe_batch_size = _get_safe_batch_size_for_nested_types(pf, leaf_indices)
fallback_batch_size = (
min(batch_size, safe_batch_size) if batch_size else safe_batch_size
)
# Apply row-group-level predicate pushdown via fragment.subset; the
# row-level filter is applied per-batch below since iter_batches
# doesn't accept a filter expression. Under schema evolution the
# filter may reference a column absent from this fragment's
# physical schema — fragment.subset uses that schema (not the
# unified one) and raises ArrowInvalid, so skip row-group pruning
# in that case and let the per-batch filter (post null-fill) do
# all the row-dropping.
fragment_physical_columns = set(fragment.physical_schema.names)
filter_touches_missing_column = filter_columns is not None and any(
c not in fragment_physical_columns for c in filter_columns
)
if filter_expr is not None and not filter_touches_missing_column:
subset = fragment.subset(filter=filter_expr)
else:
subset = fragment
row_groups = (
[rg.id for rg in subset.row_groups]
if subset.row_groups is not None
else None
)
if row_groups is not None and len(row_groups) == 0:
return
# ``pq.ParquetFile.iter_batches`` returns batches with the fragment's
# physical schema, so the fallback path would otherwise emit tables
# that differ from the scanner path (which pins
# ``_file_dataset_schema``) in struct field order, integer width,
# or missing columns. Align + cast to the same unified schema so
# fallback and non-fallback fragments concat cleanly downstream.
# Scoped to ``columns`` (not ``read_columns``) since filter-only
# columns are projected away before alignment.
file_dataset_schema = self._file_dataset_schema
if file_dataset_schema is not None and columns is not None:
align_schema = pa.schema(
[
file_dataset_schema.field(c)
for c in columns
if file_dataset_schema.get_field_index(c) != -1
]
)
else:
align_schema = file_dataset_schema
# Under schema evolution a filter-referenced column may live in
# the unified dataset schema but be absent from this fragment.
# The scanner path null-fills such columns via dataset-level
# schema pinning; ``pq.ParquetFile.iter_batches`` silently drops
# them and then ``table.filter(filter_expr)`` raises
# ``ArrowInvalid: No match for FieldRef.Name``. Mirror the
# scanner: append a null column of the unified type before the
# filter evaluates, so ``null > 15`` resolves to false and the
# fragment contributes 0 rows.
columns_to_null_fill: List[str] = (
[c for c in read_columns if c not in fragment_physical_columns]
if read_columns is not None
else []
)
null_fill_type_by_column = {
column_name: (
file_dataset_schema.field(column_name).type
if file_dataset_schema is not None
and file_dataset_schema.get_field_index(column_name) != -1
else pa.null()
)
for column_name in columns_to_null_fill
}
for batch in pf.iter_batches(
batch_size=fallback_batch_size,
columns=read_columns,
use_threads=False,
row_groups=row_groups,
):
table = pa.Table.from_batches([batch])
for column_name in columns_to_null_fill:
if column_name not in table.column_names:
table = table.append_column(
column_name,
pa.nulls(
table.num_rows,
type=null_fill_type_by_column[column_name],
),
)
if filter_expr is not None:
table = table.filter(filter_expr)
# Skip downstream select/align/cast on fully-filtered
# batches — the caller discards empty tables anyway.
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])
if align_schema is not None:
table = _align_struct_fields([table], align_schema)[0].cast(
align_schema
)
yield table
@override
def _arrow_scanner_kwargs(self) -> dict:
# ``pre_buffer`` is left at pyarrow's default (``True``). With
# ``pre_buffer=True`` pyarrow plans a single coalesced range
# request covering all needed column chunks for a fragment and
# issues it in one I/O burst, then decodes from memory. With
# ``pre_buffer=False`` pyarrow opens a per-column buffered
# stream and fetches lazily — fine on narrow schemas (few large
# columns) but catastrophic on wide schemas (thousands of small
# columns become thousands of range requests). V1
# ``ParquetDatasource`` also relies on the default. The
# cross-fragment memory accumulation that originally motivated
# disabling ``pre_buffer`` (apache/arrow#39808) is already
# addressed by V2's per-fragment scanners.
#
# ``buffer_size`` controls the per-stream read-ahead buffer
# pyarrow issues against the filesystem when ``use_buffered_stream``
# is on. The default is small (8 KiB), which produces many tiny
# range requests on S3. 8 MiB amortizes S3 latency across
# meaningful bytes per round-trip. Tunable via env var for
# workloads that need a different point on the latency/memory-
# peak curve.
kwargs: dict = {
"fragment_scan_options": pds.ParquetFragmentScanOptions(
use_buffered_stream=True,
buffer_size=_PARQUET_FRAGMENT_BUFFER_SIZE,
),
"fragment_readahead": 1,
}
return kwargs