123 lines
3.8 KiB
Python
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.
|
|
"""
|
|
...
|