Class Project

python/ray/data/_internal/logical/operators/map_operator.py:345–453

Logical operator for all Projection Operations.

@dataclass(frozen=True, repr=False, eq=False)
class Project(AbstractMap, LogicalOperatorSupportsPredicatePassThrough):
    """Logical operator for all Projection Operations."""

    exprs: list["Expr"]
    input_dependencies: list[LogicalOperator] = field(repr=False, kw_only=True)
    compute: Optional[ComputeStrategy] = None
    ray_remote_args: Dict[str, Any] = field(default_factory=dict)
    ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None
    can_modify_num_rows: bool = field(init=False, default=False)
    min_rows_per_bundled_input: Optional[int] = field(init=False, default=None)
    batch_size: Optional[int] = field(init=False, default=None)
    batch_format: str = field(init=False, default="pyarrow")
    zero_copy_batch: bool = field(init=False, default=True)
    per_block_limit: Optional[int] = None
    _num_outputs: Optional[int] = field(init=False, default=None, repr=False)

    def __post_init__(self):
        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
        # Eagerly expand ``StarExpr`` when the input schema is known. By the time
        # optimizer rules see this op, the projection list contains only
        # explicit ``col()`` and computed expressions, no ``StarExpr``.
        # When the input schema is opaque (e.g., upstream UDF map), the
        # ``StarExpr`` is preserved and runtime ``eval_projection``
        # expands it on a per-block basis.
        import pyarrow as pa

        input_schema = self.input_dependencies[0].infer_schema()
        if isinstance(input_schema, pa.Schema):
            object.__setattr__(
                self, "exprs", expand_star_exprs(self.exprs, input_schema)
            )
        if self.compute is None:
            object.__setattr__(
                self, "compute", self._detect_and_get_compute_strategy(self.exprs)
            )
        for expr in self.exprs:
            if expr.name is None and not isinstance(expr, StarExpr):
                raise TypeError(
                    "All Project expressions must be named (use .alias(name) or col(name)), "
                    "or be a star() expression."
                )
        object.__setattr__(self, "_num_outputs", None)

    def _detect_and_get_compute_strategy(self, exprs: list["Expr"]) -> ComputeStrategy:
        """Detect if expressions contain callable class UDFs and return appropriate compute strategy.

        If any expression contains a callable class UDF, returns ActorPoolStrategy.
        Otherwise returns TaskPoolStrategy.
        """
        from ray.data._internal.planner.plan_expression.expression_visitors import (
            _CallableClassUDFCollector,
        )

        # Check all expressions for callable class UDFs
        for expr in exprs:
            collector = _CallableClassUDFCollector()
            collector.visit(expr)
            if collector.get_callable_class_udfs():