Repository: ray-project/ray
File: python/ray/data/_internal/logical/operators/one_to_one_operator.py (Python)

Code before:
    @property
    def num_outputs(self) -> Optional[int]:
        return self._num_outputs

    @property
    def input_dependency(self) -> LogicalOperator:
        return self.input_dependencies[0]


class Limit(AbstractOneToOne, LogicalOperatorSupportsPredicatePassThrough):
    """Logical operator for limit."""

    def __init__(
        self,
        input_op: LogicalOperator,
        limit: int,
    ):
        super().__init__(
            input_op=input_op,
            can_modify_num_rows=True,
            name=f"limit={limit}",
        )
        self.limit = limit

    def infer_metadata(self) -> BlockMetadata:
        return BlockMetadata(
            num_rows=self._num_rows(),
            size_bytes=None,
            input_files=self._input_files(),
            exec_stats=None,
        )

    def infer_schema(
        self,
    ) -> Optional["Schema"]:
        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
        assert isinstance(self.input_dependencies[0], LogicalOperator)
        return self.input_dependencies[0].infer_schema()

    def _num_rows(self):
        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
        assert isinstance(self.input_dependencies[0], LogicalOperator)
        input_rows = self.input_dependencies[0].infer_metadata().num_rows
        if input_rows is not None:
            return min(input_rows, self.limit)
        else:
            return None

    def _input_files(self):
        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
        assert isinstance(self.input_dependencies[0], LogicalOperator)

Review comment (refactor):
When `LogicalOperator` is a pure abstract class (no constructor parameters), I think we should work towards avoiding `__setattr__`. Maybe it could look like this:
```python
num_outputs: Optional[int] = field(init=False, default=None, repr=False)

@property
def name(self) -> str:
    return f"limit={self.limit"}"

@property
def input_dependencies(self):
    return [self.input_op]
```

Okay as-is because this is a transient state and I trust it'll all look good in the end

Code after:

    @property
    def num_outputs(self) -> Optional[int]:
        return self._num_outputs

    @property
    def input_dependency(self) -> LogicalOperator:
        return self.input_dependencies[0]


@dataclass(frozen=True, repr=False, eq=False)
class Limit(AbstractOneToOne, LogicalOperatorSupportsPredicatePassThrough):
    """Logical operator for limit."""

    input_op: InitVar[LogicalOperator]
    limit: int
    can_modify_num_rows: bool = field(init=False, default=True)
    _name: str = field(init=False, repr=False)
    _input_dependencies: List[LogicalOperator] = field(init=False, repr=False)
    _num_outputs: Optional[int] = field(init=False, default=None, repr=False)

    def __post_init__(self, input_op: LogicalOperator):
        assert isinstance(input_op, LogicalOperator), input_op
        object.__setattr__(self, "_name", f"limit={self.limit}")
        object.__setattr__(self, "_input_dependencies", [input_op])
        object.__setattr__(self, "_num_outputs", None)

    def _apply_transform(
        self, transform: Callable[[LogicalOperator], LogicalOperator]
    ) -> LogicalOperator:
        transformed_input = self.input_dependency._apply_transform(transform)
        target: LogicalOperator
        if transformed_input is self.input_dependency:
            target = self
        else:
            target = Limit(transformed_input, self.limit)
        return transform(target)

    def infer_metadata(self) -> BlockMetadata:
        return BlockMetadata(
            num_rows=self._num_rows(),
            size_bytes=None,
            input_files=self._input_files(),
            exec_stats=None,
        )

    def infer_schema(
        self,
    ) -> Optional["Schema"]:
        assert len(self.input_dependencies) == 1, len(self.input_dependencies)
        assert isinstance(self.input_dependencies[0], LogicalOperator)