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

Code before:
        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] = []
        while (
            isinstance(current_op, AbstractOneToOne)
            and not current_op.can_modify_num_rows()
        ):
            if isinstance(current_op, AbstractMap):
                min_rows = current_op._min_rows_per_bundled_input
                if min_rows is not None and min_rows > limit_op._limit:
                    # Avoid pushing the limit past batch-based maps that require more
                    # rows than the limit to produce stable outputs (e.g. schema).

Review comment (suggestion):
Can you add a comment about the motivation for `_branch_has_limit`? I think the logic is correct, but also new, and I'm not sure if it'll be obvious to users *why* we do this.

Code after:
            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:
        """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] = []
        while (