/python/ray/data/_internal/datasource/parquet_datasource.py

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