import hashlib import logging import math import os from dataclasses import dataclass from typing import ( TYPE_CHECKING, Any, Callable, Dict, Iterable, Iterator, List, Literal, Optional, Tuple, Union, ) import numpy as np from packaging.version import parse as parse_version import ray from ray._common.utils import env_bool, env_integer from ray.data._internal.arrow_block import ( _BATCH_SIZE_PRESERVING_STUB_COL_NAME, ArrowBlockAccessor, ) from ray.data._internal.execution.util import merge_label_selector from ray.data._internal.object_extensions.arrow import ArrowPythonObjectType from ray.data._internal.planner.plan_expression.expression_visitors import ( get_column_references, ) from ray.data._internal.progress.progress_bar import ProgressBar from ray.data._internal.remote_fn import cached_remote_fn from ray.data._internal.util import ( MiB, RetryingPyFileSystem, _check_pyarrow_version, _is_local_scheme, iterate_with_retry, ) from ray.data._internal.utils.arrow_utils import get_pyarrow_version from ray.data.block import Block, BlockAccessor, BlockMetadata from ray.data.context import DataContext from ray.data.datasource import Datasource from ray.data.datasource.datasource import ReadTask from ray.data.datasource.file_based_datasource import ( _shuffle_file_metadata, _validate_shuffle_arg, ) from ray.data.datasource.file_meta_provider import ( FileMetadataProvider, _handle_read_os_error, _list_files, ) from ray.data.datasource.partitioning import ( PartitionDataType, Partitioning, PathPartitionFilter, PathPartitionParser, ) from ray.data.datasource.path_util import ( _resolve_paths_and_filesystem, ) from ray.data.expressions import BinaryExpr, Expr, Operation from ray.util.debug import log_once if TYPE_CHECKING: import pyarrow from pyarrow.dataset import ParquetFileFragment from ray.data.datasource.file_based_datasource import FileShuffleConfig # Type aliases for tensor column schema ColumnName = str # Shape of the tensor Shape = Tuple[int, ...] TensorColumnSchema = Dict[ColumnName, Tuple[np.dtype, Shape]] logger = logging.getLogger(__name__) MIN_PYARROW_TO_BATCHES_READAHEAD = parse_version("10.0.0") _MIN_PYARROW_VERSION_FS_FACTORY_INSPECT_PROMOTE_OPTIONS = parse_version("22.0.0") # The `num_cpus` for each metadata prefetching task. # Default to 0.5 instead of 1 because it is cheaper than normal read task. NUM_CPUS_FOR_META_FETCH_TASK = 0.5 # The number of rows to read per batch. This is sized to generate 10MiB batches # for rows about 1KiB in size. DEFAULT_PARQUET_READER_ROW_BATCH_SIZE = 10_000 FILE_READING_RETRY = 8 # `ParquetFileFragment.to_batches` passes `batch_size` through PyArrow's Cython # layer as a C ``int`` (32-bit). Larger values raise # `OverflowError: value too large to convert to int` (e.g. when estimated batch # size from bytes-per-row blows up for sparse or highly compressed batches). _MAX_PYARROW_TO_BATCHES_BATCH_SIZE = 2**31 - 1 # The default size multiplier for reading Parquet data source in Arrow. # Parquet data format is encoded with various encoding techniques (such as # dictionary, RLE, delta), so Arrow in-memory representation uses much more memory # compared to Parquet encoded representation. Parquet file statistics only record # encoded (i.e. uncompressed) data size information. # # To estimate real-time in-memory data size, Datasets will try to estimate the # correct inflation ratio from Parquet to Arrow, using this constant as the default # value for safety. See https://github.com/ray-project/ray/pull/26516 for more context. PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT = 5 # The lower bound size to estimate Parquet encoding ratio. PARQUET_ENCODING_RATIO_ESTIMATE_LOWER_BOUND = 1 # The percentage of files (1% by default) to be sampled from the dataset to estimate # Parquet encoding ratio. PARQUET_ENCODING_RATIO_ESTIMATE_SAMPLING_RATIO = 0.01 # The minimal and maximal number of file samples to take from the dataset to estimate # Parquet encoding ratio. # This is to restrict `PARQUET_ENCODING_RATIO_ESTIMATE_SAMPLING_RATIO` within the # proper boundary. PARQUET_ENCODING_RATIO_ESTIMATE_MIN_NUM_SAMPLES = 2 PARQUET_ENCODING_RATIO_ESTIMATE_MAX_NUM_SAMPLES = 10 # The number of rows to read from each file for sampling. Try to keep it low to avoid # reading too much data into memory. PARQUET_ENCODING_RATIO_ESTIMATE_NUM_ROWS = 1024 # Arrow's nested type chunking limit # See: https://github.com/apache/arrow/issues/21526 (ARROW-5030) _ARROW_CHUNK_LIMIT = 2 * 1024**3 # 2GB _MIN_PYARROW_VERSION_FOR_SCANNER_DEFAULTS = parse_version("12.0.1") # Opt-in env var to allow reading Parquet files that contain # ray.data.arrow_pickled_object columns. Disabled by default because # pickle.load on attacker-controlled data enables arbitrary code execution. AUTOLOAD_PICKLE_OBJECT_SCALAR_ENV_VAR = "RAY_DATA_AUTOLOAD_PICKLE_OBJECT_SCALAR" class _ParquetFragment: """This wrapper class is created to avoid utilizing `ParquetFileFragment` original serialization protocol that actually does network RPCs during serialization (to fetch actual parquet metadata)""" def __init__(self, f: "ParquetFileFragment", file_size: int): self._fragment = f self._file_size = file_size @property def file_size(self) -> int: return self._file_size @property def original(self) -> "ParquetFileFragment": return self._fragment def __reduce__(self): return _ParquetFragment.make_fragment, ( self._fragment.format, self._fragment.path, self._fragment.filesystem, self._fragment.partition_expression, self._file_size, ) @staticmethod def make_fragment(format, path, filesystem, partition_expression, file_size): fragment = format.make_fragment(path, filesystem, partition_expression) return _ParquetFragment(fragment, file_size) def check_for_legacy_tensor_type(schema): """Check for the legacy tensor extension type and raise an error if found. Ray Data uses an extension type to represent tensors in Arrow tables. Previously, the extension type extended `PyExtensionType`. However, this base type can expose users to arbitrary code execution. To prevent this, we don't load the type by default. """ import pyarrow as pa for name, type in zip(schema.names, schema.types): if isinstance(type, pa.UnknownExtensionType) and isinstance( type, pa.PyExtensionType ): raise RuntimeError( f"Ray Data couldn't infer the type of column '{name}' (got " f"`UnknownExtensionType` with pickled class ref " f"'{type.__arrow_ext_serialize__()}'). This might mean you're trying " f"to read data written with an older version of Ray. Reading data " f"written with older versions of Ray might expose you to arbitrary code " f"execution. To try reading the data anyway, " f"preset `RAY_DATA_AUTOLOAD_PYEXTENSIONTYPE=1` on *all* nodes." "To learn more, see https://github.com/ray-project/ray/issues/41314." ) @dataclass class _SplitPredicateResult: """Result of splitting a predicate by column type. Attributes: data_predicate: Conjuncts referencing only data columns (for PyArrow pushdown), or None if none could be extracted. partition_predicate: Conjuncts referencing only partition columns (for partition pruning), or None if none could be extracted. residual_predicate: Conjuncts that mix partition and data columns and can't be split safely (e.g. an ``OR`` straddling both kinds). The caller must keep these as a ``Filter`` above the read; dropping them would over-include rows. """ data_predicate: Optional[Expr] partition_predicate: Optional[Expr] residual_predicate: Optional[Expr] def _split_predicate_by_columns( predicate: Expr, partition_columns: set, ) -> _SplitPredicateResult: """Split a predicate into data, partition, and residual parts. This function walks the top-level ``AND`` chain and classifies each conjunct by the columns it references: - References only data columns (or none) → data bucket; pyarrow can evaluate it at scan time. - References only partition columns → partition bucket; the partition parser can evaluate it from file paths. - References both kinds (i.e. a non-``AND`` whose column set spans both) → residual bucket; semantics-preserving splitting is impossible (e.g. ``data > 5 OR partition == "US"``), so the caller must keep these as a ``Filter`` above the read. Args: predicate: The predicate expression to analyze. partition_columns: Set of partition column names. Returns: :class:`_SplitPredicateResult` with the three buckets. Combining ``data_predicate``, ``partition_predicate``, and ``residual_predicate`` with ``AND`` reproduces the original predicate exactly. Examples: >>> from ray.data.expressions import col >>> # Pure data predicate: >>> result = _split_predicate_by_columns(col("data1") > 5, {"partition_col"}) >>> result.data_predicate is not None True >>> result.partition_predicate is None and result.residual_predicate is None True >>> # Pure partition predicate: >>> result = _split_predicate_by_columns(col("partition_col") == "US", {"partition_col"}) >>> result.partition_predicate is not None True >>> result.data_predicate is None and result.residual_predicate is None True >>> # Mixed AND - can split into data and partition parts: >>> result = _split_predicate_by_columns( ... (col("data1") > 5) & (col("partition_col") == "US"), ... {"partition_col"} ... ) >>> result.data_predicate is not None and result.partition_predicate is not None True >>> result.residual_predicate is None True >>> # Mixed OR - kept as residual; caller wraps it in a Filter above: >>> result = _split_predicate_by_columns( ... (col("data1") > 5) | (col("partition_col") == "US"), ... {"partition_col"} ... ) >>> result.data_predicate is None and result.partition_predicate is None True >>> result.residual_predicate is not None True """ referenced_cols = set(get_column_references(predicate)) data_cols = referenced_cols - partition_columns partition_cols_in_predicate = referenced_cols & partition_columns if not partition_cols_in_predicate: # Pure data predicate (or no column refs). return _SplitPredicateResult( data_predicate=predicate, partition_predicate=None, residual_predicate=None, ) if not data_cols: # Pure partition predicate. return _SplitPredicateResult( data_predicate=None, partition_predicate=predicate, residual_predicate=None, ) # Mixed predicate - keep splitting if it's an AND chain. if isinstance(predicate, BinaryExpr) and predicate.op == Operation.AND: left_result = _split_predicate_by_columns(predicate.left, partition_columns) right_result = _split_predicate_by_columns(predicate.right, partition_columns) def combine_predicates( left: Optional[Expr], right: Optional[Expr] ) -> Optional[Expr]: if left and right: return left & right return left or right return _SplitPredicateResult( data_predicate=combine_predicates( left_result.data_predicate, right_result.data_predicate ), partition_predicate=combine_predicates( left_result.partition_predicate, right_result.partition_predicate ), residual_predicate=combine_predicates( left_result.residual_predicate, right_result.residual_predicate ), ) # ``OR``/``NOT``/etc. straddling both column kinds — not safely # splittable. Surface as residual so the caller doesn't silently drop # it (the prior version returned ``(None, None)`` here, which let the # surrounding ``AND`` chain push partial conjuncts and over-include # rows that should have been filtered by this one). return _SplitPredicateResult( data_predicate=None, partition_predicate=None, residual_predicate=predicate, ) class ParquetDatasource(Datasource): """Parquet datasource, for reading and writing Parquet files. This implementation uses PyArrow's `ParquetDataset` abstraction for dataset reads, and thus offers automatic Arrow dataset schema inference and row count collection at the cost of some potential performance and/or compatibility penalties. """ # Number of fragments that Pyarrow should look at to determine schema. # # NOTE: Default is 1 for backwards compatibility, which means only 1 fragment # (in no particular order) will be inspected. To inspect all fragments # set this to None. _DEFAULT_NUM_FRAGMENTS_TO_INSPECT_FOR_SCHEMA: Optional[int] = 1 _FILE_EXTENSIONS = ["parquet"] # Denotes number of batches to read ahead in a fragment. Default is 16 # as per https://arrow.apache.org/docs/python/generated/pyarrow.dataset.Dataset.html#pyarrow.dataset.Dataset.to_batches # Chose 8 based on past experiments. _DEFAULT_BATCH_READAHEAD = env_integer("RAY_DATA_PARQUET_READER_BATCH_READAHEAD", 8) # NOTE: We're essentially stubbing out this value as currently # ParquetDatasource reads individual fragments independently # Default is 4. This refers to the number of files to readahead, which was chosen based on past experiments. _DEFAULT_FRAGMENT_READAHEAD = env_integer( "RAY_DATA_PARQUET_READER_FRAGMENT_READAHEAD", 1 ) # Default is False as per https://arrow.apache.org/docs/python/generated/pyarrow.dataset.ParquetFragmentScanOptions.html # This parameter when set to True, reads files through buffered input streams rather than loading entire row groups at once. _DEFAULT_FRAGMENT_USE_BUFFERED_STREAM = env_bool( "RAY_DATA_PARQUET_READER_FRAGMENT_USE_BUFFERED_STREAM", True ) # Default is 8 KiB as per https://arrow.apache.org/docs/python/generated/pyarrow.dataset.ParquetFragmentScanOptions.html # Based on experiments, this was a good default. _DEFAULT_FRAGMENT_SCAN_BUFFER_SIZE = env_integer( "RAY_DATA_PARQUET_READER_FRAGMENT_SCAN_BUFFER_SIZE", 8 * MiB ) def __init__( self, paths: Union[str, List[str]], *, columns: Optional[List[str]] = None, dataset_kwargs: Optional[Dict[str, Any]] = None, to_batch_kwargs: Optional[Dict[str, Any]] = None, _block_udf: Optional[Callable[[Block], Block]] = None, filesystem: Optional["pyarrow.fs.FileSystem"] = None, schema: Optional[Union["pyarrow.lib.Schema"]] = None, meta_provider: Optional[FileMetadataProvider] = None, partition_filter: Optional[PathPartitionFilter] = None, partitioning: Optional[Partitioning] = Partitioning("hive"), shuffle: "FileShuffleConfig" | Literal["files"] | None = None, include_paths: bool = False, include_row_hash: bool = False, file_extensions: Optional[List[str]] = None, ): super().__init__() _check_pyarrow_version() supports_distributed_reads = not _is_local_scheme(paths) if not supports_distributed_reads and ray.util.client.ray.is_connected(): raise ValueError( "Because you're using Ray Client, read tasks scheduled on the Ray " "cluster can't access your local files. To fix this issue, store " "files in cloud storage or a distributed filesystem like NFS." ) local_scheduling = None if not supports_distributed_reads: local_scheduling = { ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id() } # Need this property for lineage tracking. We should not directly assign paths # to self since it is captured every read_task_fn during serialization and # causing this data being duplicated and excessive object store spilling. source_paths_ref = ray.put(paths) paths, resolved_filesystem = _resolve_paths_and_filesystem(paths, filesystem) filesystem = RetryingPyFileSystem.wrap( resolved_filesystem, retryable_errors=DataContext.get_current().retried_io_errors, ) listed_files = _list_files( paths, filesystem, partition_filter=partition_filter, file_extensions=file_extensions, ) if listed_files: paths, file_sizes = zip(*listed_files) else: paths, file_sizes = [], [] if dataset_kwargs is not None: logger.warning( "Please note that `ParquetDatasource.__init__`s `dataset_kwargs` " "is a deprecated parameter and will be removed in the future." ) else: dataset_kwargs = {} if "partitioning" in dataset_kwargs: raise ValueError( "The 'partitioning' parameter isn't supported in 'dataset_kwargs'. " "Use the top-level 'partitioning' parameter instead." ) # This datasource manually adds partition data at the Ray Data-level. To avoid # duplicating the partition data, we disable PyArrow's partitioning. dataset_kwargs["partitioning"] = None # NOTE: ParquetDataset only accepts list of paths, hence we need to convert # it to a list pq_ds = get_parquet_dataset( list(paths), filesystem=filesystem, schema=schema, inspect_num_fragments=self._DEFAULT_NUM_FRAGMENTS_TO_INSPECT_FOR_SCHEMA, dataset_kwargs=dataset_kwargs, ) fragments = list(pq_ds.fragments) # Users can pass both data columns and partition columns in the 'columns' # argument. To prevent PyArrow from complaining about missing columns, we # separate the partition columns from the data columns. When we read the # fragments, we pass the data columns to PyArrow and add the partition # columns manually. data_columns, partition_columns = None, None if columns is not None: if fragments: data_columns, partition_columns = _infer_data_and_partition_columns( columns, fragments[0], partitioning ) else: # Empty dataset - can't infer columns without fragments data_columns, partition_columns = [], [] if to_batch_kwargs is None: to_batch_kwargs = {} # Store as projection_map (identity mapping if columns specified, None otherwise) # Note: Empty list [] means no columns, None means all columns # Include partition columns in projection_map if they were requested, so that # projection pushdown can properly track them if data_columns is None and partition_columns is None: projection_map = None else: projection_map = {} if data_columns is not None: projection_map.update({col: col for col in data_columns}) if partition_columns is not None: projection_map.update({col: col for col in partition_columns}) # Eagerly compute the actual partition columns for _partition_columns. # This ensures _partition_columns is always a list (never None). actual_partition_columns = partition_columns if partition_columns is None and partitioning is not None and fragments: parse = PathPartitionParser(partitioning) parsed_partitions = parse(fragments[0].path) if parsed_partitions: actual_partition_columns = list(parsed_partitions.keys()) # Store selected partition columns. Always a list (never None) representing # the actual partition columns to include. actual_partition_columns = ( actual_partition_columns if actual_partition_columns is not None else [] ) # Track whether partition columns were explicitly part of the user's column selection partition_columns_selected = ( partition_columns is not None and len(actual_partition_columns) > 0 ) self._init_state( supports_distributed_reads=supports_distributed_reads, local_scheduling=local_scheduling, source_paths_ref=source_paths_ref, filesystem=resolved_filesystem, fragments=fragments, file_sizes=list(file_sizes), file_schema=pq_ds.schema, read_schema=schema, partition_columns=actual_partition_columns, partition_columns_selected=partition_columns_selected, partition_schema=_get_partition_columns_schema( partitioning, [p.path for p in fragments] ), partitioning=partitioning, projection_map=projection_map, to_batch_kwargs=to_batch_kwargs, _block_udf=_block_udf, shuffle=shuffle, include_paths=include_paths, include_row_hash=include_row_hash, ) def _init_state( self, *, supports_distributed_reads: bool, local_scheduling: Optional[Dict[str, str]], source_paths_ref: "ray.ObjectRef", filesystem: "pyarrow.fs.FileSystem", fragments: List["ParquetFileFragment"], file_sizes: List[int], file_schema: "pyarrow.Schema", read_schema: Optional["pyarrow.lib.Schema"], partition_columns: List[str], partition_columns_selected: bool, partition_schema: "pyarrow.Schema", partitioning: Optional[Partitioning], projection_map: Optional[Dict[str, str]], to_batch_kwargs: Optional[Dict[str, Any]], _block_udf: Optional[Callable[[Block], Block]], shuffle: Union["FileShuffleConfig", Literal["files"], None], include_paths: bool, include_row_hash: bool = False, ): """Shared initialization for all instance state and sampling estimates. Called by both ``__init__`` (after path resolution and file listing) and ``from_state`` (used by alternate constructors like ``from_pyarrow_dataset``). """ self._allow_pickle_object_columns = env_bool( AUTOLOAD_PICKLE_OBJECT_SCALAR_ENV_VAR, False ) self._supports_distributed_reads = supports_distributed_reads self._local_scheduling = local_scheduling self._source_paths_ref = source_paths_ref self._filesystem = filesystem # NOTE: Store the custom serialized `ParquetFileFragment` to avoid unexpected # network calls when `_ParquetDatasourceReader` is serialized. See # `_SerializedFragment()` implementation for more details. self._pq_fragments = [ _ParquetFragment(fragment, file_size) for fragment, file_size in zip(fragments, file_sizes) ] self._pq_paths = [f.path for f in fragments] self._block_udf = _block_udf self._projection_map = projection_map self._partition_columns = partition_columns self._partition_columns_selected = partition_columns_selected self._read_schema = read_schema self._file_schema = file_schema self._partition_schema = partition_schema self._file_metadata_shuffler = None self._include_paths = include_paths self._include_row_hash = include_row_hash self._partitioning = partitioning _validate_shuffle_arg(shuffle) self._shuffle = shuffle # Sample small number of parquet files to estimate # - Encoding ratio: ratio of file size on disk to approximate expected # size of the corresponding block in memory # - Default batch-size: number of rows to be read from a file at a time, # used to limit amount of memory pressure sampled_fragments = _sample_fragments( self._pq_fragments, ) sampled_file_infos = _fetch_file_infos( sampled_fragments, columns=self._get_data_columns(), schema=read_schema, local_scheduling=self._local_scheduling, ) self._encoding_ratio = _estimate_files_encoding_ratio( sampled_fragments, sampled_file_infos, ) estimated_batch_size = _estimate_reader_batch_size( sampled_file_infos, DataContext.get_current().target_max_block_size ) self._scanner_kwargs = self._get_scanner_kwargs( to_batch_kwargs, estimated_batch_size ) @classmethod def from_state(cls, **kwargs) -> "ParquetDatasource": """Create a fully-initialized instance from pre-computed state. This is the preferred entry point for alternate constructors (e.g. ``from_pyarrow_dataset``) that bypass the normal ``__init__`` path. All keyword arguments are forwarded to ``_init_state``. """ instance = cls.__new__(cls) Datasource.__init__(instance) _check_pyarrow_version() instance._init_state(**kwargs) return instance @classmethod def from_pyarrow_dataset( cls, pa_dataset: "pyarrow.dataset.Dataset", *, columns: Optional[List[str]] = None, to_batch_kwargs: Optional[Dict[str, Any]] = None, _block_udf: Optional[Callable[[Block], Block]] = None, schema: Optional["pyarrow.lib.Schema"] = None, shuffle: "FileShuffleConfig" | Literal["files"] | None = None, include_paths: bool = False, include_row_hash: bool = False, ) -> "ParquetDatasource": """Create a ParquetDatasource from a pre-built PyArrow dataset. This bypasses path resolution, file listing, and dataset construction — useful when another system (e.g. delta-rs) has already built a PyArrow dataset with the correct filesystem, schema, and fragment metadata. """ import pyarrow as pa from pyarrow.fs import LocalFileSystem fragments = list(pa_dataset.get_fragments()) filesystem = pa_dataset.filesystem pq_paths = [f.path for f in fragments] # Determine locality from the filesystem type rather than path parsing. # PyArrow fragment paths are filesystem-relative (e.g. "bucket/key/file.parquet" # instead of "s3://bucket/key/..."), so _is_local_scheme cannot distinguish # local from remote paths here. is_local = isinstance(filesystem, LocalFileSystem) supports_distributed_reads = not is_local local_scheduling = None if pq_paths: if is_local and ray.util.client.ray.is_connected(): raise ValueError( "Because you're using Ray Client, read tasks scheduled on " "the Ray cluster can't access your local files. To fix " "this issue, store files in cloud storage or a distributed " "filesystem like NFS." ) if is_local: local_scheduling = { ray._raylet.RAY_NODE_ID_KEY: ray.get_runtime_context().get_node_id() } infos = filesystem.get_file_info(pq_paths) file_sizes = [info.size if info.size is not None else 0 for info in infos] else: file_sizes = [] # Partition columns are intentionally left empty. For PyArrow datasets # produced by systems like delta-rs, partition columns are materialized # by PyArrow itself (via fragment.to_batches(schema=...)) rather than # parsed from file paths. Setting these to empty prevents Ray Data's # path-based partition parsing from duplicating columns that PyArrow # already provides. return cls.from_state( supports_distributed_reads=supports_distributed_reads, local_scheduling=local_scheduling, source_paths_ref=ray.put(pq_paths), filesystem=filesystem, fragments=fragments, file_sizes=file_sizes, file_schema=pa_dataset.schema, read_schema=schema if schema is not None else pa_dataset.schema, partition_columns=[], partition_columns_selected=False, partition_schema=pa.schema([]), partitioning=None, projection_map={col: col for col in columns} if columns is not None else None, to_batch_kwargs=to_batch_kwargs, _block_udf=_block_udf, shuffle=shuffle, include_paths=include_paths, include_row_hash=include_row_hash, ) @property def _source_paths(self) -> List[str]: return ray.get(self._source_paths_ref) def estimate_inmemory_data_size(self) -> int: # In case of empty projections no data will be read if self._projection_map == {}: return 0 return self._estimate_in_mem_size(self._pq_fragments) def _get_scanner_kwargs( self, to_batch_kwargs: Optional[Dict[str, Any]], batch_size: Optional[int], ) -> dict[str, Any]: import pyarrow.dataset as pds scanner_kwargs: Dict[str, Any] = (to_batch_kwargs or {}).copy() # NOTE: We inject ``batch_size`` via kwargs, since Scanner doesn't accept # nulls. Use setdefault so a user-provided ``batch_size`` in # ``to_batch_kwargs`` wins. if batch_size is not None: scanner_kwargs.setdefault("batch_size", batch_size) # Override default `batch_readahead` value to reduce amount of data prefetched # by Pyarrow's Parquet reader pyarrow_version = get_pyarrow_version() if ( pyarrow_version is not None and pyarrow_version >= _MIN_PYARROW_VERSION_FOR_SCANNER_DEFAULTS ): scanner_kwargs.setdefault("batch_readahead", self._DEFAULT_BATCH_READAHEAD) scanner_kwargs.setdefault( "fragment_readahead", self._DEFAULT_FRAGMENT_READAHEAD ) # Refer https://arrow.apache.org/docs/python/generated/pyarrow.dataset.ParquetFragmentScanOptions.html # Read files through buffered input streams rather than loading # entire row groups at once. if self._DEFAULT_FRAGMENT_USE_BUFFERED_STREAM: scanner_kwargs.setdefault( "fragment_scan_options", pds.ParquetFragmentScanOptions( use_buffered_stream=True, buffer_size=self._DEFAULT_FRAGMENT_SCAN_BUFFER_SIZE, ), ) return scanner_kwargs def get_read_tasks( self, parallelism: int, per_task_row_limit: Optional[int] = None, data_context: Optional["DataContext"] = None, ) -> List[ReadTask]: # NOTE: We override the base class FileBasedDatasource.get_read_tasks() # method in order to leverage pyarrow's ParquetDataset abstraction, # which simplifies partitioning logic. We still use # FileBasedDatasource's write side, however. execution_idx = data_context._execution_idx if data_context is not None else 0 pq_fragments, pq_paths = _shuffle_file_metadata( self._pq_fragments, self._pq_paths, self._shuffle, execution_idx ) # Derive expected target schema of the blocks being read target_schema = self._derive_schema( self._read_schema, file_schema=self._file_schema, partition_schema=self._partition_schema, projected_columns=self.get_current_projection(), _block_udf=self._block_udf, include_paths=self._include_paths, include_row_hash=self._include_row_hash, ) read_tasks = [] filter_expr = ( self._predicate_expr.to_pyarrow() if self._predicate_expr is not None else None ) filter_columns = None if self._predicate_expr is not None: filter_columns = get_column_references(self._predicate_expr) for fragments, paths in zip( np.array_split(pq_fragments, parallelism), np.array_split(pq_paths, parallelism), ): if len(fragments) <= 0: continue meta = BlockMetadata( num_rows=None, size_bytes=self._estimate_in_mem_size(fragments), input_files=paths, exec_stats=None, ) ( block_udf, to_batches_kwargs, data_columns, partition_columns, read_schema, include_paths, include_row_hash, partitioning, ) = ( self._block_udf, self._scanner_kwargs, self._get_data_columns(), self._get_partition_columns(), self._read_schema, self._include_paths, self._include_row_hash, self._partitioning, ) allow_pickle = self._allow_pickle_object_columns read_tasks.append( ReadTask( lambda f=fragments: read_fragments( block_udf, to_batches_kwargs, data_columns, partition_columns, read_schema, f, include_paths, include_row_hash, partitioning, filter_expr, filter_columns, allow_pickle, ), meta, schema=target_schema, per_task_row_limit=per_task_row_limit, ) ) return read_tasks def get_name(self): """Return a human-readable name for this datasource. This will be used as the names of the read tasks. """ return "Parquet" @property def supports_distributed_reads(self) -> bool: return self._supports_distributed_reads def supports_projection_pushdown(self) -> bool: return True def supports_predicate_pushdown(self) -> bool: return True def get_current_projection(self) -> Optional[List[str]]: """Override to include partition columns in addition to data columns.""" # NOTE: In case there's no projection both file and partition columns # will be none data_columns = self._get_data_columns() partition_columns = self._get_partition_columns() if data_columns is None and partition_columns is None: result = None else: result = (data_columns or []) + (partition_columns or []) # If include_paths is True, make sure to include the path column in the projection # NOTE: When result is None (no projection), the path column is already added # via _derive_schema, so we only need to add it when there is a projection. if self._include_paths and "path" not in result: result = result + ["path"] if self._include_row_hash and "row_hash" not in result: result = result + ["row_hash"] return result def _get_partition_columns(self) -> Optional[List[str]]: """Extract partition columns from projection map. This method extracts partition columns from _projection_map, which is the source of truth after projection pushdown. Since partition columns are now included in _projection_map during initialization when requested, we can reliably extract them from the map. Returns: List of partition column names in the projection, None if there's no projection (meaning include all partition columns), or [] if partition columns aren't in the projection map (meaning include no partition columns). """ if self._projection_map is None: return None if not self._partition_columns: # If a projection is active but the dataset has no partition columns, # then no partition columns should be included in the output. # Returning [] ensures that no partition columns are added, # `None` is interpreted as including all partition columns. return [] # Extract partition columns that are in the projection map partition_cols = [ col for col in self._projection_map.keys() if col in self._partition_columns ] # If partition columns are found in projection map, return them if partition_cols: return partition_cols # No partition columns in projection map. # Since the projection map exists and is the source of truth after # projection pushdown, return [] (no partition columns to include). return [] def _get_data_columns(self) -> Optional[List[str]]: """Extract data columns from projection map, excluding partition columns. Partition columns aren't in the physical file schema, so they must be filtered out before passing to PyArrow's to_batches(). Similarly, the synthetic "path" column (when include_paths=True) isn't in the physical file schema, so it must also be filtered out. Returns: List of data column names to read from files, or None if no projection. Can return empty list if only partition columns are projected. """ if self._projection_map is None: return None # Get partition columns and filter them out from the projection partition_cols = self._partition_columns # Also filter out synthetic columns (path, row_hash) as they are # added after reading from the file cols_to_filter = set(partition_cols) if self._include_paths: cols_to_filter.add("path") if self._include_row_hash: cols_to_filter.add("row_hash") data_cols = [ col for col in self._projection_map.keys() if col not in cols_to_filter ] return data_cols def apply_predicate( self, predicate_expr: Expr, ) -> "ParquetDatasource": """Apply a predicate with data pushdown and partition pruning. This method optimizes predicates in three ways: 1. Data predicates → pushed to PyArrow (row-level filtering) 2. Partition predicates → used for partition pruning (file-level filtering) 3. Mixed predicates → both optimizations applied together """ partition_cols = set(self._partition_columns) if not partition_cols: # No partition columns - can push down everything normally return super().apply_predicate(predicate_expr) # Split predicate into data and partition parts split_result = _split_predicate_by_columns(predicate_expr, partition_cols) # If a mixed-column conjunct can't be safely split (e.g. ``data > 5 # OR partition == "US"``), the V1 ``Read.apply_predicate`` wrapper # has no way to keep it as a ``Filter`` above the read — its return # type is ``Read``, not ``LogicalOperator``. Pushing the splittable # parts and silently dropping the residual would over-include rows. # Return ``self`` so ``PredicatePushdown`` keeps the original # ``Filter`` above intact. if split_result.residual_predicate is not None: return self # Apply partition pruning if we have a partition predicate if ( split_result.partition_predicate is not None and self._partitioning is not None ): parser = PathPartitionParser(self._partitioning) pruned_fragments = [] pruned_paths = [] for fragment, path in zip(self._pq_fragments, self._pq_paths): # Evaluate partition predicate - skip if it doesn't match if parser.evaluate_predicate_on_partition( path, split_result.partition_predicate ): pruned_fragments.append(fragment) pruned_paths.append(path) # Apply partition pruning directly to self self._pq_fragments = pruned_fragments self._pq_paths = pruned_paths # Push down data predicate to PyArrow if present # Create a copy and push down the data predicate to PyArrow import copy datasource = copy.copy(self) # Only call apply_predicate if there's a data predicate to push down # If data_predicate is None (pure partition predicate), skip it to avoid # creating invalid expressions like existing_expr & None if split_result.data_predicate is not None: return super(ParquetDatasource, datasource).apply_predicate( split_result.data_predicate ) return datasource def _estimate_in_mem_size(self, fragments: List[_ParquetFragment]) -> int: in_mem_size = sum([f.file_size for f in fragments]) * self._encoding_ratio return round(in_mem_size) @staticmethod def _derive_schema( read_schema: Optional["pyarrow.Schema"], *, file_schema: "pyarrow.Schema", partition_schema: Optional["pyarrow.Schema"], projected_columns: Optional[List[str]], _block_udf, include_paths: bool = False, include_row_hash: bool = False, ) -> "pyarrow.Schema": """Derives target schema for read operation""" import pyarrow as pa # Use target read schema if provided if read_schema is not None: target_schema = read_schema else: file_schema_fields = list(file_schema) partition_schema_fields = ( list(partition_schema) if partition_schema is not None else [] ) # Otherwise, fallback to file + partitioning schema by default target_schema = pa.schema( fields=( file_schema_fields + [ f for f in partition_schema_fields # Ignore fields from partition schema overlapping with # file's schema if file_schema.get_field_index(f.name) == -1 ] ), metadata=file_schema.metadata, ) # Add path column if include_paths is True and path column is not already present if include_paths and target_schema.get_field_index("path") == -1: target_schema = target_schema.append(pa.field("path", pa.string())) if include_row_hash: idx = target_schema.get_field_index("row_hash") if idx == -1: target_schema = target_schema.append(pa.field("row_hash", pa.uint64())) elif target_schema.field(idx).type != pa.uint64(): target_schema = target_schema.set( idx, pa.field("row_hash", pa.uint64()) ) # Project schema if necessary if projected_columns is not None: target_schema = pa.schema( [target_schema.field(column) for column in projected_columns], target_schema.metadata, ) if _block_udf is not None: # Try to infer dataset schema by passing dummy table through UDF. try: # An empty table with extensions will fail for pyarrow==9.0.0 dummy_table = target_schema.empty_table() target_schema = _block_udf(dummy_table).schema.with_metadata( target_schema.metadata ) except Exception: logger.debug( "Failed to infer schema of dataset by passing dummy table " "through UDF due to the following exception:", exc_info=True, ) check_for_legacy_tensor_type(target_schema) return target_schema def _check_for_pickle_object_columns(table: "pyarrow.Table") -> None: pickle_cols = [ field.name for field in table.schema if isinstance(field.type, ArrowPythonObjectType) ] if pickle_cols: raise ValueError( f"This Parquet file contains columns stored as " f"'ray.data.arrow_pickled_object': {pickle_cols}. Reading these " f"columns requires unpickling, which can execute arbitrary code " f"and is unsafe with untrusted files.\n\n" f"If you trust the source of this data, set the environment " f"variable {AUTOLOAD_PICKLE_OBJECT_SCALAR_ENV_VAR}=1 to allow " f"reading these columns. In a Ray cluster, this variable must " f"be set on all worker nodes (e.g. via 'runtime_env')." ) def read_fragments( block_udf: Callable[[Block], Optional[Block]], to_batches_kwargs: Dict[str, Any], data_columns: Optional[List[str]], partition_columns: Optional[List[str]], schema: Optional[Union[type, "pyarrow.lib.Schema"]], fragments: List[_ParquetFragment], include_paths: bool, include_row_hash: bool, partitioning: Partitioning, filter_expr: Optional["pyarrow.dataset.Expression"] = None, filter_columns: Optional[List[str]] = None, allow_pickle: bool = False, ) -> Iterator["pyarrow.Table"]: """Yield Arrow tables from Parquet fragments via ``to_batches_kwargs``.""" # This import is necessary to load the tensor extension type. from ray.data.extensions.tensor_extension import ArrowTensorType # noqa # Ensure that we're reading at least one dataset fragment. assert len(fragments) > 0 logger.debug(f"Reading {len(fragments)} parquet fragments") for fragment in fragments: # S3 can raise transient errors during iteration, and PyArrow doesn't expose a # way to retry specific batches. ctx = ray.data.DataContext.get_current() for table in iterate_with_retry( lambda: _read_batches_from( fragment.original, schema=schema, data_columns=data_columns, partition_columns=partition_columns, partitioning=partitioning, include_path=include_paths, include_row_hash=include_row_hash, filter_expr=filter_expr, filter_columns=filter_columns, to_batches_kwargs=to_batches_kwargs, ), "reading batches", match=ctx.retried_io_errors, ): # If the table is empty, drop it. if table.num_rows > 0: if not allow_pickle: _check_for_pickle_object_columns(table) if block_udf is not None: yield block_udf(table) else: yield table def _coerce_pyarrow_fragment_batch_size(batch_size: int) -> int: """Clamp batch size for ``ParquetFileFragment.to_batches`` to PyArrow's C int range. Expects a value already converted with :func:`int` (callers reading from untyped kwargs should do ``int(...)`` before calling). Values outside ``[1, _MAX_PYARROW_TO_BATCHES_BATCH_SIZE]`` are clamped. """ if batch_size <= 0: raise ValueError(f"Batch size must be > 0, got {batch_size}") coerced = min(batch_size, _MAX_PYARROW_TO_BATCHES_BATCH_SIZE) if coerced != batch_size: logger.debug( "Clamping Parquet fragment read batch_size from %s to %s " "(PyArrow ``to_batches`` requires batch_size in [1, %s]).", batch_size, coerced, _MAX_PYARROW_TO_BATCHES_BATCH_SIZE, ) return coerced def _has_susceptible_nested_types(schema: "pyarrow.Schema") -> bool: """Check if a schema contains nested column types wrapping variable-length leaves that are susceptible to Arrow's chunked array limitation (ARROW-5030). The error only occurs when a nested container (list, struct, map) contains a variable-length leaf (string, binary, and their large/view variants) whose data exceeds ~2GB in a single row group. Fixed-width leaves (int, float, bool, etc.) never trigger chunking. """ import pyarrow as pa # is_string_view / is_binary_view only exist in PyArrow >= 16.0 _has_view_types = hasattr(pa.types, "is_string_view") def _is_variable_length(t): return ( pa.types.is_string(t) or pa.types.is_binary(t) or pa.types.is_large_string(t) or pa.types.is_large_binary(t) or (_has_view_types and pa.types.is_string_view(t)) or (_has_view_types and pa.types.is_binary_view(t)) ) def _is_nested(t): return ( pa.types.is_list(t) or pa.types.is_large_list(t) or pa.types.is_struct(t) or pa.types.is_map(t) or pa.types.is_fixed_size_list(t) ) def _nested_contains_variable_length(t): """Recursively check if a nested type contains a variable-length leaf.""" if _is_variable_length(t): return True if ( pa.types.is_list(t) or pa.types.is_large_list(t) or pa.types.is_fixed_size_list(t) ): return _nested_contains_variable_length(t.value_type) if pa.types.is_struct(t): return any(_nested_contains_variable_length(f.type) for f in t) if pa.types.is_map(t): return _nested_contains_variable_length( t.key_type ) or _nested_contains_variable_length(t.item_type) return False return any( _is_nested(field.type) and _nested_contains_variable_length(field.type) for field in schema ) def _row_group_uncompressed_size( rg_meta: "pyarrow.parquet.RowGroupMetaData", column_indices: Optional[List[int]] = None, ) -> int: """Total uncompressed byte size of columns in a row group. When *column_indices* is ``None`` all columns are summed, otherwise only the listed (leaf-level) column indices are included. NOTE: We intentionally avoid ``rg_meta.total_byte_size`` because it can return the *compressed* size for some files (apache/arrow#48138). """ indices = range(rg_meta.num_columns) if column_indices is None else column_indices return sum(rg_meta.column(i).total_uncompressed_size for i in indices) def _resolve_leaf_column_indices( metadata: "pyarrow.parquet.FileMetaData", columns: List[str], ) -> List[int]: """Map top-level column names to Parquet metadata leaf column indices. Parquet metadata enumerates *leaf* columns (nested types are flattened), and each leaf's ``path_in_schema`` starts with the top-level field name. """ col_set = set(columns) return [ i for i in range(metadata.num_columns) if metadata.row_group(0).column(i).path_in_schema.split(".")[0] in col_set ] def _get_safe_batch_size_for_nested_types( pf: "pyarrow.parquet.ParquetFile", column_indices: Optional[List[int]] = None, ) -> int: """Compute a batch size that keeps each batch under Arrow's ~2GB nested-type chunking threshold. Uses Parquet row group metadata (uncompressed column sizes) to estimate bytes per row, then picks a batch size with a 50% safety margin. """ safe_batch_size = pf.metadata.num_rows for rg_idx in range(pf.metadata.num_row_groups): rg_meta = pf.metadata.row_group(rg_idx) if rg_meta.num_rows == 0: continue uncompressed = _row_group_uncompressed_size(rg_meta, column_indices) if uncompressed == 0: continue bytes_per_row = uncompressed / rg_meta.num_rows rg_safe = max(int(_ARROW_CHUNK_LIMIT // bytes_per_row // 2), 1) safe_batch_size = min(safe_batch_size, rg_safe) return safe_batch_size def _needs_nested_type_fallback( fragment: "ParquetFileFragment", columns: Optional[List[str]] = None, ) -> bool: """Check if a fragment requires the fallback reader for nested types. Returns True if the *requested* columns (or all columns when ``columns`` is ``None``) contain nested types AND any row group has uncompressed data exceeding Arrow's ~2GB chunking threshold. This is a metadata-only check (no data read). """ import pyarrow as pa physical_schema = fragment.physical_schema if columns is not None: physical_schema = pa.schema( [ physical_schema.field(c) for c in columns if physical_schema.get_field_index(c) != -1 ] ) if not _has_susceptible_nested_types(physical_schema): return False metadata = fragment.metadata column_indices = ( _resolve_leaf_column_indices(metadata, columns) if columns is not None and metadata.num_row_groups > 0 else None ) # fragment.row_groups is non-None when the fragment is a subset of the # file (e.g. only row group 0). Only inspect those row groups to avoid # falsely triggering the fallback because of a *different* large row # group elsewhere in the same file. if fragment.row_groups is not None: rg_indices = [rg.id for rg in fragment.row_groups] else: rg_indices = range(metadata.num_row_groups) return any( _row_group_uncompressed_size(metadata.row_group(rg_idx), column_indices) >= _ARROW_CHUNK_LIMIT for rg_idx in rg_indices ) def _resolve_read_columns( columns: Optional[List[str]], filter_expr: Optional["pyarrow.dataset.Expression"], filter_columns: Optional[List[str]], ) -> Optional[List[str]]: """Compute the union of projected and filter-referenced columns. When a filter references columns outside the projection, we must read the union so the filter can evaluate. Returns ``None`` (meaning "all columns") when filter_columns is unknown. """ if filter_expr is not None and columns is not None: if filter_columns is not None: return list(dict.fromkeys(columns + filter_columns)) return None return columns def _iter_batches_with_nested_fallback( fragment: "ParquetFileFragment", *, columns: Optional[List[str]] = None, schema: Optional["pyarrow.Schema"] = None, to_batches_kwargs: Optional[Dict[str, Any]] = None, use_threads: bool = False, filter_expr: Optional["pyarrow.dataset.Expression"] = None, filter_columns: Optional[List[str]] = None, ) -> Iterable["pyarrow.RecordBatch"]: """Iterate over batches from a fragment, using the fallback reader when the fragment has nested column types with row groups exceeding Arrow's ~2GB chunking threshold (ARROW-5030). The reading strategy is chosen upfront based on schema and metadata to avoid mid-stream fallback that could duplicate already-yielded batches. In the normal path, ``filter_expr`` is pushed down to the scanner. In the fallback path, row-group-level predicate pushdown is applied via ``fragment.subset(filter=)`` and row-level filtering is done post-read. """ to_batches_kwargs = dict(to_batches_kwargs or {}) read_columns = _resolve_read_columns(columns, filter_expr, filter_columns) if not _needs_nested_type_fallback(fragment, read_columns): yield from fragment.to_batches( columns=columns, filter=filter_expr, schema=schema, use_threads=use_threads, **to_batches_kwargs, ) return yield from _iter_batches_fallback( fragment, columns=columns, read_columns=read_columns, schema=schema, batch_size=to_batches_kwargs.get("batch_size"), use_threads=use_threads, filter_expr=filter_expr, ) def _iter_batches_fallback( fragment: "ParquetFileFragment", *, columns: Optional[List[str]], read_columns: Optional[List[str]], schema: Optional["pyarrow.Schema"], batch_size: Optional[int], use_threads: bool, filter_expr: Optional["pyarrow.dataset.Expression"], ) -> Iterable["pyarrow.RecordBatch"]: """Row-level batched reader for fragments with nested types that exceed Arrow's ~2GB chunking threshold (ARROW-5030). """ import pyarrow as pa import pyarrow.parquet as pq from ray.data._internal.arrow_ops.transform_pyarrow import ( _align_struct_fields, ) if log_once("parquet_nested_fallback"): 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, ) pf = pq.ParquetFile(fragment.path, filesystem=fragment.filesystem) # Scope batch-size calculation to only the columns being decoded. 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 = _get_safe_batch_size_for_nested_types(pf, leaf_indices) fallback_batch_size = min(batch_size, safe) if batch_size else safe # fragment.subset(filter=) respects the fragment's existing row-group # subset, so a single call handles both the original constraint and # filter-based predicate pushdown. subset = ( fragment.subset(filter=filter_expr) if filter_expr is not None else fragment ) row_groups = ( [rg.id for rg in subset.row_groups] if subset.row_groups is not None else None ) # Filter pruned every row group — nothing to read. if row_groups is not None and len(row_groups) == 0: return # Build a sub-schema covering only the output columns so alignment # doesn't pad with unneeded null columns. Scoped to ``columns`` # (not ``read_columns``) because filter-referenced columns may not # appear in the target schema. if schema is not None and columns is not None: align_schema = pa.schema( [schema.field(c) for c in columns if schema.get_field_index(c) != -1] ) else: align_schema = schema for batch in pf.iter_batches( batch_size=fallback_batch_size, columns=read_columns, use_threads=use_threads, row_groups=row_groups, ): table = pa.Table.from_batches([batch]) # Row-level filter runs on the physical schema before alignment # so that filter-referenced columns outside the output projection # (and possibly outside the target schema) are still present. if filter_expr is not None: table = table.filter(filter_expr) # Project down to the requested columns, dropping any # filter-only columns before the (potentially expensive) align. if columns is not None: table = table.select(columns) # pq.ParquetFile.iter_batches() doesn't accept a schema arg, so # batches come back with the file's physical schema. Align to the # unified dataset schema to match the normal (scanner) path which # handles type promotion and missing-column filling automatically. if align_schema is not None: table = _align_struct_fields([table], align_schema)[0].cast(align_schema) yield from table.to_batches() def _read_batches_from( fragment: "ParquetFileFragment", *, schema: "pyarrow.Schema", data_columns: Optional[List[str]], partition_columns: Optional[List[str]], partitioning: Partitioning, filter_expr: Optional["pyarrow.dataset.Expression"] = None, filter_columns: Optional[List[str]] = None, include_path: bool = False, include_row_hash: bool = False, use_threads: bool = False, to_batches_kwargs: Optional[Dict[str, Any]] = None, ) -> Iterable["pyarrow.Table"]: """Get an iterable of batches from a parquet fragment. Row batching is controlled via ``to_batches_kwargs["batch_size"]`` (when present), which is coerced to values PyArrow accepts as a C ``int``. """ import pyarrow as pa # Copy to avoid modifying passed in arg to_batches_kwargs = dict(to_batches_kwargs or {}) # NOTE: Passed in kwargs overrides always take precedence # TODO deprecate to_batches_kwargs use_threads = to_batches_kwargs.pop("use_threads", use_threads) # TODO: We should deprecate filter through the read_parquet API and only allow through dataset.filter() filter_from_kwargs = to_batches_kwargs.pop("filter", None) if filter_from_kwargs is not None: filter_expr = ( filter_from_kwargs if filter_expr is None else filter_expr & filter_from_kwargs ) # Cannot determine columns from an opaque PyArrow filter expression, # so invalidate filter_columns to fall back to reading all columns. filter_columns = None if to_batches_kwargs.get("batch_size") is not None: to_batches_kwargs["batch_size"] = _coerce_pyarrow_fragment_batch_size( int(to_batches_kwargs["batch_size"]) ) partition_col_values = _parse_partition_column_values( fragment, partition_columns, partitioning ) row_offset = 0 def _generate_tables() -> "pa.Table": nonlocal row_offset def _postprocess_table(table): if partition_col_values: table = _add_partitions_to_table(partition_col_values, table) if include_path: table = ArrowBlockAccessor.for_block(table).fill_column( "path", fragment.path ) # ``ParquetFileFragment.to_batches`` returns ``RecordBatch``, # which could have empty projection (ie ``num_columns`` == 0) # while having non-empty rows (ie ``num_rows`` > 0), which # could occur when list of requested columns is empty. # # However, when ``RecordBatches`` are concatenated using # ``pyarrow.concat_tables`` it will return a single ``Table`` # with 0 columns and therefore 0 rows (since ``Table``s number of # rows is determined as the length of its columns). # # To avoid running into this pitfall, we introduce a stub column # holding just nulls to maintain invariance of the number of rows. # # NOTE: There's no impact from this as the binary size of the # extra column is basically 0 if table.num_columns == 0 and table.num_rows > 0: table = table.append_column( _BATCH_SIZE_PRESERVING_STUB_COL_NAME, pa.nulls(table.num_rows) ) return table try: for batch in _iter_batches_with_nested_fallback( fragment, columns=data_columns, filter_expr=filter_expr, filter_columns=filter_columns, schema=schema, use_threads=use_threads, to_batches_kwargs=to_batches_kwargs, ): table = _postprocess_table(pa.Table.from_batches([batch])) if include_row_hash: hashes = _compute_row_hashes( fragment.path, row_offset, table.num_rows ) table = ArrowBlockAccessor.for_block(table).fill_column( "row_hash", pa.array(hashes, type=pa.uint64()) ) row_offset += table.num_rows yield table except pa.lib.ArrowInvalid as e: error_message = str(e) if ( "No match for FieldRef.Name" in error_message and filter_expr is not None ): filename = os.path.basename(fragment.path) file_columns = set(fragment.physical_schema.names) raise RuntimeError( f"Filter expression: '{filter_expr}' failed on parquet " f"file: '{filename}' with columns: {file_columns}" ) raise yield from _generate_tables() def _compute_row_hashes(file_path: str, start_row: int, num_rows: int) -> np.ndarray: """Compute deterministic uint64 hashes from file path and output row position. ``start_row`` is the position within the output stream (post-filter), not the physical file offset. This means hashes are reproducible for a given pipeline configuration (same file + same filter) but will differ across reads with different filters. Hashes the file path with MD5 to obtain a 64-bit seed, adds the row indices, then applies the splitmix64 finalizer (a bijective 64-bit mixing function) to produce well-distributed, reproducible hashes. Fully vectorized via numpy. """ path_seed = np.uint64( int.from_bytes( hashlib.md5(file_path.encode("utf-8")).digest()[:8], byteorder="little" ) ) keys = path_seed + np.arange(start_row, start_row + num_rows, dtype=np.uint64) # splitmix64 finalizer – a bijective 64-bit mixing function from # Steele, Lea & Flood, "Fast Splittable Pseudorandom Number Generators", # OOPSLA 2014. Also used in Java's SplittableRandom. # Reference: https://xorshift.di.unimi.it/splitmix64.c keys ^= keys >> np.uint64(30) keys *= np.uint64(0xBF58476D1CE4E5B9) keys ^= keys >> np.uint64(27) keys *= np.uint64(0x94D049BB133111EB) keys ^= keys >> np.uint64(31) return keys def _parse_partition_column_values( fragment: "ParquetFileFragment", partition_columns: Optional[List[str]], partitioning: Partitioning, ) -> Dict[str, PartitionDataType]: partitions = {} if partitioning is not None: parse = PathPartitionParser(partitioning) partitions = parse(fragment.path) # Filter out partitions that aren't in the user-specified columns list. if partition_columns is not None: partitions = { field_name: value for field_name, value in partitions.items() if field_name in partition_columns } return partitions def _fetch_parquet_file_info( fragment: _ParquetFragment, *, columns: Optional[List[str]], schema: Optional["pyarrow.Schema"], ) -> Optional["_ParquetFileInfo"]: # If the fragment has no row groups, it's an empty or metadata-only file. # Skip it by returning empty sample info. # # NOTE: Accessing `ParquetFileFragment.metadata` does fetch a parquet footer # from storage metadata = fragment.original.metadata if metadata.num_row_groups == 0: return None # Only sample the first row group. row_group_fragment = fragment.original.subset(row_group_ids=[0]) batch_size = max( min( row_group_fragment.metadata.num_rows, PARQUET_ENCODING_RATIO_ESTIMATE_NUM_ROWS, ), 1, ) to_batches_kwargs = {} if get_pyarrow_version() >= MIN_PYARROW_TO_BATCHES_READAHEAD: # Limit prefetching to just 1 batch to_batches_kwargs["batch_readahead"] = 1 to_batches_kwargs["batch_size"] = batch_size avg_row_size: Optional[int] = None # Use first non-empty batch to estimate the avg size of the row in-memory for batch in _iter_batches_with_nested_fallback( row_group_fragment, columns=columns, schema=schema, to_batches_kwargs=to_batches_kwargs, ): if batch.num_rows > 0: avg_row_size = math.ceil(batch.nbytes / batch.num_rows) break return _ParquetFileInfo( avg_row_in_mem_bytes=avg_row_size, metadata=metadata, ) @dataclass class _ParquetFileInfo: # Estimated avg byte size of a row (in-memory) avg_row_in_mem_bytes: Optional[int] # Corresponding file metadata metadata: "pyarrow._parquet.FileMetaData" def estimate_in_memory_bytes(self) -> Optional[int]: if self.avg_row_in_mem_bytes is None: return None return self.avg_row_in_mem_bytes * self.metadata.num_rows def _estimate_files_encoding_ratio( fragments: List[_ParquetFragment], file_infos: List[_ParquetFileInfo], ) -> float: """Return an estimate of the Parquet files encoding ratio. To avoid OOMs, it is safer to return an over-estimate than an underestimate. """ if not DataContext.get_current().decoding_size_estimation: return PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT assert len(file_infos) == len(fragments) # Estimate size of the rows in a file in memory estimated_in_mem_size_arr = [ fi.estimate_in_memory_bytes() if fi is not None else None for fi in file_infos ] file_size_arr = [f.file_size for f in fragments] estimated_encoding_ratios = [ float(in_mem_size) / file_size for in_mem_size, file_size in zip(estimated_in_mem_size_arr, file_size_arr) if file_size > 0 and in_mem_size is not None ] # Return default estimate of 5 if all sampled files turned out to be empty if not estimated_encoding_ratios: return PARQUET_ENCODING_RATIO_ESTIMATE_DEFAULT estimated_ratio = np.mean(estimated_encoding_ratios) logger.info(f"Estimated parquet encoding ratio is {estimated_ratio:.3f}.") return max(estimated_ratio, PARQUET_ENCODING_RATIO_ESTIMATE_LOWER_BOUND) def _fetch_file_infos( sampled_fragments: List[_ParquetFragment], *, columns: Optional[List[str]], schema: Optional["pyarrow.Schema"], local_scheduling: Optional[Dict[str, str]], ) -> List[Optional[_ParquetFileInfo]]: fetch_file_info = cached_remote_fn(_fetch_parquet_file_info) futures = [] # Retry in case of transient errors during sampling. # Cap retries to avoid hanging indefinitely on permanent errors # (e.g., permission denied, invalid credentials). task_options = {"retry_exceptions": [OSError], "max_retries": 3} ctx = DataContext.get_current() if local_scheduling: task_options["label_selector"] = local_scheduling else: task_options["scheduling_strategy"] = ctx.scheduling_strategy task_options = merge_label_selector( task_options, ctx.execution_options.label_selector ) for fragment in sampled_fragments: # Sample the first rows batch in i-th file. # Use SPREAD scheduling strategy to avoid packing many sampling tasks on # same machine to cause OOM issue, as sampling can be memory-intensive. futures.append( fetch_file_info.options(**task_options).remote( fragment, columns=columns, schema=schema, ) ) sample_bar = ProgressBar("Parquet dataset sampling", len(futures), unit="file") try: file_infos = sample_bar.fetch_until_complete(futures) except ray.exceptions.RayTaskError as e: logger.warning( "Parquet dataset sampling failed. " "If this is a credentials or permissions issue, " "check your cloud storage access configuration. " f"Underlying error: {e.cause}" ) raise finally: sample_bar.close() return file_infos def _estimate_reader_batch_size( file_infos: List[Optional[_ParquetFileInfo]], target_block_size: Optional[int] ) -> Optional[int]: if target_block_size is None: return None avg_num_rows_per_block = [ target_block_size / fi.avg_row_in_mem_bytes for fi in file_infos if ( fi is not None and fi.avg_row_in_mem_bytes is not None and fi.avg_row_in_mem_bytes > 0 ) ] if not avg_num_rows_per_block: return DEFAULT_PARQUET_READER_ROW_BATCH_SIZE estimated_batch_size: int = max(math.ceil(np.mean(avg_num_rows_per_block)), 1) logger.info(f"Estimated parquet reader batch size at {estimated_batch_size} rows") return estimated_batch_size def get_parquet_dataset( paths: List[str], schema: Optional["pyarrow.Schema"] = None, filesystem: Optional["pyarrow.fs.FileSystem"] = None, inspect_num_fragments: Optional[int] = 0, dataset_kwargs: Optional[Dict[str, Any]] = None, ): assert inspect_num_fragments is None or inspect_num_fragments >= 0, ( f"`inspect_num_fragments` could either be null (inspect all fragments) " f"or >= 0 (got {inspect_num_fragments})" ) paths = paths if isinstance(paths, list) else [paths] # For linter dataset = None try: dataset = _get_parquet_dataset_internal( paths, schema, filesystem, inspect_num_fragments, dataset_kwargs, ) except TypeError: from ray.data.datasource.path_util import _resolve_paths_and_filesystem try: # Fallback: resolve filesystem locally in the worker resolved_paths, resolved_filesystem = _resolve_paths_and_filesystem( paths, filesystem=None ) resolved_filesystem = RetryingPyFileSystem.wrap( resolved_filesystem, retryable_errors=DataContext.get_current().retried_io_errors, ) dataset = _get_parquet_dataset_internal( resolved_paths, schema, resolved_filesystem, inspect_num_fragments, dataset_kwargs, ) except OSError as os_e: _handle_read_os_error(os_e, paths) except OSError as e: _handle_read_os_error(e, paths) return dataset def _get_parquet_dataset_internal( paths: List[str], schema: Optional["pyarrow.Schema"], filesystem: Optional["pyarrow.fs.FileSystem"], inspect_num_fragments: Optional[int], dataset_kwargs: Optional[Dict[str, Any]] = None, ) -> "pyarrow.parquet.ParquetDataset": import pyarrow.parquet as pq should_inspect = inspect_num_fragments != 0 if schema is None and should_inspect: # NOTE: In case no schema is provided we must infer schema = _infer_schema( paths, inspect_num_fragments, filesystem, ) dataset_kwargs = dataset_kwargs or {} return pq.ParquetDataset( # When passing directories, Pyarrow expects single items and not # a list (otherwise erroring out) paths[0] if len(paths) == 1 else paths, schema=schema, filesystem=filesystem, **dataset_kwargs, ) def _infer_schema( paths: List[str], inspect_num_fragments: Optional[int], filesystem: Optional["pyarrow.fs.FileSystem"], ) -> "pyarrow.Schema": import pyarrow as pa import pyarrow.dataset as pds factory = pds.FileSystemDatasetFactory( filesystem, paths, format=pds.ParquetFileFormat(), ) # NOTE: By default we're inspecting all the fragments. # The ``fragments`` kwarg was added in PyArrow 21.0 (previously # all fragments were inspected unconditionally). # PyArrow 22.0 added ``promote_options`` for proper null→concrete # type promotion across fragments (GH-46629). pa_version = get_pyarrow_version() if pa_version >= _MIN_PYARROW_VERSION_FS_FACTORY_INSPECT_PROMOTE_OPTIONS: inspect_kwargs = { "fragments": inspect_num_fragments, "promote_options": "permissive", } else: inspect_kwargs = {} schema = factory.inspect(**inspect_kwargs) # Before Pyarrow 22.0, ``factory.inspect`` doesn't promote ``null`` types # to concrete types when unifying schemas across fragments (which # happens when some files have all-null values for a column). # # In that case we manually collect physical schemas from all fragments and # call ``pa.unify_schemas`` to correctly promote the types. if pa_version < _MIN_PYARROW_VERSION_FS_FACTORY_INSPECT_PROMOTE_OPTIONS and any( field.type == pa.null() for field in schema ): dataset = factory.finish(schema) fragment_schemas = [f.physical_schema for f in dataset.get_fragments()] schema = pa.unify_schemas([schema] + fragment_schemas) return schema def _sample_fragments( fragments: List[_ParquetFragment], ) -> List[_ParquetFragment]: if not fragments: return [] target_num_samples = math.ceil( len(fragments) * PARQUET_ENCODING_RATIO_ESTIMATE_SAMPLING_RATIO ) target_num_samples = max( min(target_num_samples, PARQUET_ENCODING_RATIO_ESTIMATE_MAX_NUM_SAMPLES), PARQUET_ENCODING_RATIO_ESTIMATE_MIN_NUM_SAMPLES, ) # Make sure number of samples doesn't exceed total # of files target_num_samples = min(target_num_samples, len(fragments)) # Evenly distributed to choose which file to sample, to avoid biased prediction # if data is skewed. pivots = np.linspace(0, len(fragments) - 1, target_num_samples).astype(int) return [fragments[idx] for idx in pivots.tolist()] def _add_partitions_to_table( partition_col_values: Dict[str, PartitionDataType], table: "pyarrow.Table" ) -> "pyarrow.Table": for partition_col, value in partition_col_values.items(): field_index = table.schema.get_field_index(partition_col) if field_index == -1: table = BlockAccessor.for_block(table).fill_column(partition_col, value) elif log_once(f"duplicate_partition_field_{partition_col}"): logger.warning( f"The partition field '{partition_col}' also exists in the Parquet " f"file. Ray Data will default to using the value in the Parquet file." ) return table def _get_partition_columns_schema( partitioning: Partitioning, file_paths: List[str], ) -> "pyarrow.Schema": """Return a new schema with partition fields added. This function infers the partition fields from the first file path in the dataset. """ import pyarrow as pa # If the dataset is empty, we can't infer the partitioning if len(file_paths) == 0: return pa.schema([]) # If the dataset isn't partitioned, there's no partition schema elif partitioning is None: return pa.schema([]) first_path = file_paths[0] fields = [] parser = PathPartitionParser(partitioning) partitions = parser(first_path) for field_name in partitions: if field_name in partitioning.field_types: field_type = pa.from_numpy_dtype(partitioning.field_types[field_name]) else: field_type = pa.string() # Without this check, we would add the same partition field multiple times, # which silently fails when asking for `pa.field()`. fields.append(pa.field(field_name, field_type)) return pa.schema(fields) def _infer_data_and_partition_columns( user_specified_columns: List[str], fragment: "ParquetFileFragment", partitioning: Optional[Partitioning], ) -> Tuple[List[str], List[str]]: """Infer which columns are in the files and which columns are partition columns. This function uses the schema and path of the first file to infer what columns represent. Args: user_specified_columns: A list of column names that the user specified. fragment: The first fragment in the dataset. partitioning: The partitioning scheme used to partition the data. Returns: A tuple of lists of column names. The first list contains the columns that are in the file, and the second list contains the columns that are partition columns. """ data_columns = [ column for column in user_specified_columns if column in fragment.physical_schema.names ] if partitioning is not None: parse = PathPartitionParser(partitioning) partitions = parse(fragment.path) partition_columns = [ column for column in user_specified_columns if column in partitions ] else: partition_columns = [] return data_columns, partition_columns