Files
2026-07-13 13:17:40 +08:00

123 lines
3.8 KiB
Python

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, List, Optional, Set, Tuple
from ray.data.expressions import Expr
from ray.util.annotations import DeveloperAPI
if TYPE_CHECKING:
from ray.data._internal.datasource_v2.scanners.scanner import Scanner
@DeveloperAPI
class SupportsFilterPushdown(ABC):
"""Mixin for scanners that support filter/predicate pushdown.
Filter pushdown allows predicates to be evaluated at the data source level,
reducing the amount of data that needs to be read and transferred.
"""
@abstractmethod
def push_filters(self, predicate: "Expr") -> Tuple["Scanner", Optional["Expr"]]:
"""Push a filter predicate down to the scanner.
Args:
predicate: Expression representing the filter condition.
Returns:
Tuple of (new_scanner, residual_predicate) where:
- new_scanner: New Scanner instance with the filter applied
- residual_predicate: Any part of the predicate that couldn't be
pushed down and must be applied post-scan. None if fully pushed.
"""
...
@DeveloperAPI
class SupportsColumnPruning(ABC):
"""Mixin for scanners that support column pruning/projection pushdown.
Column pruning allows reading only the columns needed by the query,
which is especially beneficial for columnar formats like Parquet.
"""
@abstractmethod
def prune_columns(self, columns: List[str]) -> "Scanner":
"""Prune the scanner to only read the specified columns.
Args:
columns: List of column names to read.
Returns:
New Scanner instance configured to read only the specified columns.
"""
...
@abstractmethod
def pruned_column_names(self) -> Optional[Tuple[str, ...]]:
"""Physical column names selected after pruning, if any.
Returns:
``None`` when no pruning has been applied (read all columns).
A tuple (possibly empty) after :meth:`prune_columns` has been
applied, listing on-disk / reader column names in read order.
"""
...
@DeveloperAPI
class SupportsLimitPushdown(ABC):
"""Mixin for scanners that support limit pushdown.
Limit pushdown allows the scanner to stop early once the required number
of rows has been read.
"""
@abstractmethod
def push_limit(self, limit: int) -> "Scanner":
"""Push a row limit down to the scanner.
Args:
limit: Maximum number of rows to read.
Returns:
New Scanner instance with the limit applied.
"""
...
@DeveloperAPI
class SupportsPartitionPruning(ABC):
"""Mixin for scanners that support partition pruning.
Partition pruning allows skipping entire files/partitions based on
predicates that reference partition columns.
"""
@property
@abstractmethod
def partition_columns(self) -> Set[str]:
"""Names of columns that are partition keys.
Callers (e.g. the predicate-pushdown rule) use this to decide
whether a predicate should be routed through :meth:`push_filters`
(data columns) or :meth:`prune_partitions` (partition columns).
Must be fully populated by schema inference at planning time.
"""
...
@abstractmethod
def prune_partitions(self, predicate: "Expr") -> "Scanner":
"""Prune partitions based on a predicate.
The scanner determines its partition columns from its
``Partitioning`` configuration, which is fully populated
by schema inference at planning time.
Args:
predicate: Expression to evaluate against partition values.
Returns:
New Scanner instance with partition pruning applied.
"""
...