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

Code before:
        Existing topology:
            child₁ , child₂ , …  ->  Union  ->  Limit

        New topology:
            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.
        """

Review comment (refactor):
```
while (
    isinstance(current, AbstractOneToOne)
    and not current.can_modify_num_rows()
    and current.input_dependencies
):
    if isinstance(current, Limit) and current._limit == limit:
        return True
    current = current.input_dependency
return False
```
is a little clearer to understand

Code after:
        Existing topology:
            child₁ , child₂ , …  ->  Union  ->  Limit

        New topology:
            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)