wehub-resource-sync
2140 lines
80 KiB
Python
2140 lines
80 KiB
Python
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
|