chore: import upstream snapshot with attribution

python/ray/data/_internal/datasource_v2/scanners/arrow_file_scanner.py

@@ -0,0 +1,201 @@
+import logging
+from dataclasses import dataclass, replace
+from typing import List, Optional, Set, Tuple
+
+import pyarrow as pa
+from pyarrow.fs import FileSystem
+from typing_extensions import override
+
+from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
+from ray.data._internal.datasource_v2.logical_optimizers import (
+    SupportsColumnPruning,
+    SupportsFilterPushdown,
+    SupportsLimitPushdown,
+    SupportsPartitionPruning,
+)
+from ray.data._internal.datasource_v2.scanners.file_scanner import FileScanner
+from ray.data.datasource.partitioning import Partitioning, PathPartitionParser
+from ray.data.expressions import Expr
+from ray.util.annotations import DeveloperAPI
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+@dataclass(frozen=True)
+class ArrowFileScanner(
+    FileScanner,
+    SupportsFilterPushdown,
+    SupportsColumnPruning,
+    SupportsLimitPushdown,
+    SupportsPartitionPruning,
+):
+    """Base scanner for file-based datasources that use PyArrow's Dataset API.
+
+    Holds shared Arrow types and options (schema, projection, filesystem,
+    partitioning, etc.). Subclasses set the file format in :meth:`create_reader`.
+
+    Provides default implementations of filter pushdown, column pruning,
+    limit pushdown, and partition pruning that work for all Arrow-backed
+    formats.
+
+    Non-Arrow file formats should subclass :class:`FileScanner` directly.
+    """
+
+    schema: pa.Schema
+    batch_size: Optional[int] = None
+    columns: Optional[Tuple[str, ...]] = None
+    predicate: Optional[Expr] = None
+    partition_predicate: Optional[Expr] = None
+    limit: Optional[int] = None
+    filesystem: Optional[FileSystem] = None
+    partitioning: Optional[Partitioning] = None
+    ignore_prefixes: Optional[List[str]] = None
+
+    @property
+    def partition_columns(self) -> Set[str]:
+        """Return the set of partition column names, or empty if unpartitioned."""
+        if self.partitioning is None:
+            return set()
+        return set(self.partitioning.field_names or [])
+
+    def read_schema(self) -> pa.Schema:
+        """Return the logical schema after column pruning.
+
+        ``columns is None`` → no projection applied, return the full schema.
+        ``columns = ()`` → empty projection (``ds.select_columns([])``),
+        return an empty schema.
+
+        The physical read may still inject a stub column (see
+        ``_BATCH_SIZE_PRESERVING_STUB_COL_NAME``) so that row counts
+        survive a zero-column scan; that stub is an execution-layer detail
+        and is deliberately not reflected in this logical schema.
+        """
+        if self.columns is None:
+            return self.schema
+        fields = []
+        for name in self.columns:
+            idx = self.schema.get_field_index(name)
+            assert idx >= 0, f"Column {name} not found in schema"
+            fields.append(self.schema.field(idx))
+        return pa.schema(fields)
+
+    @override
+    def push_filters(
+        self, predicate: "Expr"
+    ) -> Tuple["ArrowFileScanner", Optional["Expr"]]:
+        """Push filter predicate down to the scanner.
+
+        ANDs the predicate with any existing predicate. The Ray ``Expr`` is
+        retained as the source of truth so the reader can introspect filter
+        columns; conversion to a PyArrow expression happens at the
+        scanner-kwargs boundary in :class:`FileReader`.
+
+        This method handles data-column predicates only. Partition predicates
+        should be pushed via :meth:`prune_partitions` instead; the optimizer
+        is responsible for splitting them before calling either method.
+
+        Args:
+            predicate: Ray Data expression to push down.
+
+        Returns:
+            A pair ``(scanner, residual)`` where ``scanner`` has the predicate
+            merged into its PyArrow filter. ``residual`` is ``None`` because
+            PyArrow handles the full filter at scan time.
+        """
+        if self.predicate is not None:
+            combined = self.predicate & predicate
+        else:
+            combined = predicate
+
+        return replace(self, predicate=combined), None
+
+    @override
+    def prune_columns(self, columns: List[str]) -> "ArrowFileScanner":
+        """Prune to only the specified columns.
+
+        Args:
+            columns: List of column names to keep.
+
+        Returns:
+            New scanner with column pruning applied.
+        """
+        if self.columns:
+            existing = set(self.columns)
+            columns = [c for c in columns if c in existing]
+
+        return replace(self, columns=tuple(columns))
+
+    @override
+    def pruned_column_names(self) -> Optional[Tuple[str, ...]]:
+        return self.columns
+
+    @override
+    def push_limit(self, limit: int) -> "ArrowFileScanner":
+        """Push row limit down to the scanner.
+
+        Args:
+            limit: Maximum number of rows to read.
+
+        Returns:
+            New scanner with limit applied.
+        """
+        current = self.limit
+        new_limit = min(current, limit) if current is not None else limit
+        return replace(self, limit=new_limit)
+
+    @override
+    def prune_partitions(self, predicate: "Expr") -> "ArrowFileScanner":
+        """Store a partition predicate for file-level pruning during plan().
+
+        The predicate is ANDed with any existing partition predicate. Actual
+        file pruning happens in :meth:`plan` when the manifest is available,
+        using :class:`PathPartitionParser` to evaluate partition values from
+        file paths.
+
+        Args:
+            predicate: Expression referencing only partition columns.
+
+        Returns:
+            New scanner with partition predicate stored.
+        """
+        if self.partition_predicate is not None:
+            combined = self.partition_predicate & predicate
+        else:
+            combined = predicate
+
+        return replace(self, partition_predicate=combined)
+
+    @override
+    def prune_manifest(self, manifest: FileManifest) -> FileManifest:
+        """Filter manifest to only files matching ``self.partition_predicate``.
+
+        Called by :func:`plan_read_files_op.do_read` for every incoming
+        manifest block. No-op when either the predicate or the
+        partitioning spec is absent. Uses
+        :class:`PathPartitionParser` to parse partition values from
+        each file path and evaluate the predicate.
+        """
+        if self.partition_predicate is None or self.partitioning is None:
+            return manifest
+
+        parser = PathPartitionParser(self.partitioning)
+        keep_indices = []
+
+        for i, path in enumerate(manifest.paths):
+            if parser.evaluate_predicate_on_partition(path, self.partition_predicate):
+                keep_indices.append(i)
+
+        if len(keep_indices) == len(manifest):
+            return manifest
+
+        pruned_count = len(manifest) - len(keep_indices)
+        logger.debug(
+            "Partition pruning removed %d of %d files",
+            pruned_count,
+            len(manifest),
+        )
+
+        block = manifest.as_block()
+        pruned_block = block.take(keep_indices)
+        return FileManifest(pruned_block)

python/ray/data/_internal/datasource_v2/scanners/file_scanner.py

@@ -0,0 +1,46 @@
+from dataclasses import dataclass, field
+from typing import Literal, Union
+
+from ray.data._internal.datasource_v2.listing.file_manifest import FileManifest
+from ray.data._internal.datasource_v2.scanners.scanner import Scanner
+from ray.data.datasource.file_based_datasource import (
+    FileShuffleConfig,
+    _validate_shuffle_arg,
+)
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+@dataclass(frozen=True)
+class FileScanner(Scanner[FileManifest]):
+    """Base scanner for file-based datasources.
+
+    Subclasses implement format-specific ``read_schema()`` and
+    ``create_reader()``. Shuffling and parallel bucketing are handled
+    upstream in the ``ListFiles`` transform chain (``shuffle_files`` +
+    ``RoundRobinPartitioner`` via ``plan_list_files_op``), not here.
+
+    PyArrow Dataset-based scanners should subclass ``ArrowFileScanner``;
+    use ``FileScanner`` directly for non-Arrow file formats.
+    """
+
+    # kw_only so subclass dataclasses can declare their own required fields
+    # (like ``ArrowFileScanner.schema``) without running into the "non-default
+    # argument follows default argument" dataclass inheritance rule.
+    shuffle: Union[Literal["files"], FileShuffleConfig, None] = field(
+        default=None, kw_only=True
+    )
+
+    def __post_init__(self) -> None:
+        _validate_shuffle_arg(self.shuffle)
+
+    def prune_manifest(self, manifest: FileManifest) -> FileManifest:
+        """Return a filtered view of ``manifest``.
+
+        Default: identity. Subclasses that support file-level predicate
+        pruning (e.g. :class:`ArrowFileScanner`'s ``partition_predicate``)
+        override this to drop rows whose partition values fail the
+        predicate. Invoked per-block from
+        :func:`plan_read_files_op.do_read`.
+        """
+        return manifest

python/ray/data/_internal/datasource_v2/scanners/parquet_scanner.py

@@ -0,0 +1,88 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+import pyarrow as pa
+
+from ray.data._internal.datasource.parquet_datasource import (
+    check_for_legacy_tensor_type,
+)
+from ray.data._internal.datasource_v2.readers.file_reader import (
+    INCLUDE_PATHS_COLUMN_NAME,
+    ROW_HASH_COLUMN_NAME,
+)
+from ray.data._internal.datasource_v2.readers.parquet_file_reader import (
+    ParquetFileReader,
+)
+from ray.data._internal.datasource_v2.scanners.arrow_file_scanner import (
+    ArrowFileScanner,
+)
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+@dataclass(frozen=True)
+class ParquetScanner(ArrowFileScanner):
+    """Parquet-specific scanner implementation.
+
+    Inherits filter pushdown, column pruning, limit pushdown, partition
+    pruning, and file shuffle from ArrowFileScanner. Adds Parquet-specific
+    reader creation with adaptive batch sizing and the Parquet-only
+    legacy-tensor-type schema check.
+    """
+
+    target_block_size: Optional[int] = None
+    include_paths: bool = False
+    include_row_hash: bool = False
+    # Extra kwargs forwarded to ``pds.ParquetFileFormat(**kwargs)`` inside
+    # the per-task ``ParquetFileReader`` (e.g. ``coerce_int96_timestamp_unit``,
+    # ``pre_buffer``, ``dictionary_columns``). Carries the deprecated
+    # ``dataset_kwargs`` payload from ``read_parquet`` to the worker.
+    parquet_format_kwargs: Dict[str, Any] = field(default_factory=dict)
+
+    def read_schema(self) -> pa.Schema:
+        """Return schema after column pruning and tensor check.
+
+        ``path`` and ``row_hash`` are synthesized post-read by the file
+        reader, but only for columns listed in ``self.columns`` (see
+        ``file_reader.read``'s ``columns_to_synthesize`` filter). When a
+        projection has pruned a synthesized column away, advertising it
+        here would put the schema out of sync with the actual blocks — so
+        only append when no projection is active or when it survives.
+        """
+        schema = super().read_schema()
+        synthesized = (
+            (self.include_paths, INCLUDE_PATHS_COLUMN_NAME, pa.string()),
+            (self.include_row_hash, ROW_HASH_COLUMN_NAME, pa.uint64()),
+        )
+        for enabled, name, dtype in synthesized:
+            if not enabled:
+                continue
+            if self.columns is not None and name not in self.columns:
+                continue
+            if schema.get_field_index(name) != -1:
+                continue
+            schema = schema.append(pa.field(name, dtype))
+
+        check_for_legacy_tensor_type(schema)
+        return schema
+
+    def create_reader(self) -> ParquetFileReader:
+        """Create a ParquetFileReader configured for this scanner.
+
+        Returns:
+            ParquetFileReader with all pushdowns and adaptive batch sizing.
+        """
+        return ParquetFileReader(
+            batch_size=self.batch_size,
+            columns=list(self.columns) if self.columns is not None else None,
+            predicate=self.predicate,
+            limit=self.limit,
+            filesystem=self.filesystem,
+            partitioning=self.partitioning,
+            ignore_prefixes=self.ignore_prefixes,
+            target_block_size=self.target_block_size,
+            include_paths=self.include_paths,
+            include_row_hash=self.include_row_hash,
+            schema=self.schema,
+            parquet_format_kwargs=dict(self.parquet_format_kwargs),
+        )

python/ray/data/_internal/datasource_v2/scanners/scanner.py

@@ -0,0 +1,51 @@
+from abc import ABC, abstractmethod
+from typing import Generic
+
+import pyarrow as pa
+
+from ray.data._internal.datasource_v2 import InputSplit
+from ray.data._internal.datasource_v2.readers.base_reader import Reader
+from ray.util.annotations import DeveloperAPI
+
+
+@DeveloperAPI
+class Scanner(ABC, Generic[InputSplit]):
+    """Abstract base class for configured scanners.
+
+    A Scanner represents the logical result of reading data, including applied
+    filters, projections, limits, and other pushdown operations. It is an
+    immutable abstraction: each push operation returns a new Scanner instance
+    via cloning rather than mutation.
+
+    The Scanner is responsible for:
+    1. Determining the output schema after all projections
+    2. Creating Reader instances configured with all pushdowns
+
+    Splitting the input into parallel work units used to live here as a
+    ``plan()`` method. That responsibility now belongs to the listing-side
+    pipeline (``ListFiles`` + ``FilePartitioner``); scanners only
+    need to answer "what schema?" and "give me a reader."
+    """
+
+    @abstractmethod
+    def read_schema(self) -> pa.Schema:
+        """Return the schema that will be produced by this scanner.
+
+        This reflects the schema after all column pruning has been applied.
+
+        Returns:
+            PyArrow Schema describing the output data.
+        """
+        ...
+
+    @abstractmethod
+    def create_reader(self) -> Reader[InputSplit]:
+        """Create a Reader configured for this scanner.
+
+        The returned Reader will have all pushdowns (columns, predicates, limits)
+        applied and is ready to execute on workers.
+
+        Returns:
+            Configured Reader instance.
+        """
+        ...