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

Code before:
            child₁ -> Limit ->│
                               │
            child₂ -> Limit ->┤ Union ──► Limit   (original)
                               │
            …    -> Limit  ->│
        """
        union_op = limit_op.input_dependency
        assert isinstance(union_op, Union)

        def _branch_has_limit(op: LogicalOperator, limit: int) -> bool:
            current = op
            while True:
                if isinstance(current, Limit) and current._limit == limit:
                    return True
                if not (
                    isinstance(current, AbstractOneToOne)
                    and not current.can_modify_num_rows()
                ):
                    return False
                if not current.input_dependencies:
                    return False
                current = current.input_dependency

        # Insert a branch-local Limit and push it further upstream.
        branch_tails: List[LogicalOperator] = []
        for child in union_op.input_dependencies:
            if _branch_has_limit(child, limit_op._limit):
                branch_tails.append(child)
                continue
            raw_limit = Limit(child, limit_op._limit)  # child → limit

            if isinstance(raw_limit.input_dependency, Union):
                # This represents the limit operator appended after the union.
                pushed_tail = self._push_limit_into_union(raw_limit)
            else:
                # This represents the operator that takes place of the original limit position.
                pushed_tail = self._push_limit_down(raw_limit)
            branch_tails.append(pushed_tail)

        new_union = Union(*branch_tails)
        return Limit(new_union, limit_op._limit)

    def _push_limit_down(self, limit_op: Limit) -> LogicalOperator:
        """Push a single limit down through compatible operators conservatively.

        Creates entirely new operators instead of mutating existing ones.
        """
        # Traverse up the DAG until we reach the first operator that meets
        # one of the stopping conditions
        current_op = limit_op.input_dependency
        num_rows_preserving_ops: List[LogicalOperator] = []

Review comment (refactor):
I felt confused because `input_dependency` isn't an attribute on all logical operators. After reading through the code more carefully, I think it's because it's guaranteed to be a `AbstractOneToOne` by this point in the code.

Code after:
            child₁ -> Limit ->│
                               │
            child₂ -> Limit ->┤ Union ──► Limit   (original)
                               │
            …    -> Limit  ->│

        Example (skip duplicate limit on a branch that already has it):
            before:
                child -> Limit(n) -> Union -> Limit(n)
            after:
                child -> Limit(n) -> Union -> Limit(n)   (no extra branch limit inserted)
        """
        union_op = limit_op.input_dependency
        assert isinstance(union_op, Union)

        def _branch_has_limit(op: LogicalOperator, limit: int) -> bool:
            current = op
            while (
                isinstance(current, AbstractOneToOne)
                and not current.can_modify_num_rows()
                and current.input_dependencies
            ):
                if isinstance(current, Limit):
                    return current._limit == limit
                # Safe to use input_dependency: current is an AbstractOneToOne here.
                current = current.input_dependency

            return isinstance(current, Limit) and current._limit == limit

        # Insert a branch-local Limit and push it further upstream.
        branch_tails: List[LogicalOperator] = []
        for child in union_op.input_dependencies:
            # Avoid inserting a duplicate Limit on a branch that already has the same
            # limit upstream of row-preserving ops.
            if _branch_has_limit(child, limit_op._limit):
                branch_tails.append(child)
                continue
            raw_limit = Limit(child, limit_op._limit)  # child → limit

            if isinstance(raw_limit.input_dependency, Union):
                # This represents the limit operator appended after the union.
                pushed_tail = self._push_limit_into_union(raw_limit)
            else:
                # This represents the operator that takes place of the original limit position.
                pushed_tail = self._push_limit_down(raw_limit)
            branch_tails.append(pushed_tail)

        new_union = Union(*branch_tails)
        return Limit(new_union, limit_op._limit)

    def _push_limit_down(self, limit_op: Limit) -> LogicalOperator: