Repository: ray-project/ray
File: python/ray/data/_internal/planner/plan_expression/expression_evaluator.py (Python)

Code before:

        Args:
            expr: The monotonically_increasing_id expression.

        Returns:
            The result of the monotonically_increasing_id expression as a BlockColumn.
        """
        ctx = TaskContext.get_current()
        if ctx is None:
            raise RuntimeError(
                "monotonically_increasing_id() must be evaluated within a Ray task"
            )

        # Key the counter by expression instance ID so that multiple expressions
        # in the same projection will have isolated row count state
        counter_key = f"_mono_id_{expr._instance_id}_counter"

        start_idx = ctx.kwargs.get(counter_key, 0)
        num_rows = self.block_accessor.num_rows()
        ctx.kwargs[counter_key] = start_idx + num_rows

        # Upper 31 bits = task/partition ID, lower 33 bits = row number
        partition_mask = ctx.task_idx << 33
        ids = partition_mask + np.arange(
            start_idx, start_idx + num_rows, dtype=np.int64
        )

        block_type = self.block_accessor.block_type()
        if block_type == BlockType.PANDAS:
            return pd.Series(ids)
        elif block_type == BlockType.ARROW:
            return pa.array(ids)
        else:
            return ids


def eval_expr(expr: Expr, block: Block) -> Union[BlockColumn, ScalarType]:
    """Evaluate an expression against a block using the visitor pattern.

    Args:
        expr: The expression to evaluate.
        block: The Block to evaluate against.

    Returns:
        The evaluated result as a BlockColumn or a scalar value.
    """
    evaluator = NativeExpressionEvaluator(block)
    return evaluator.visit(expr)


def eval_projection(projection_exprs: List[Expr], block: Block) -> Block:

Review comment (suggestion):
In what situations will this be non-zero (i.e., when would we visit the same `expr` multiple times in a task?)

I think the logic makes sense. Mostly wondering for my own understanding

Code after:
            "DownloadExpr evaluation is not yet implemented in NativeExpressionEvaluator."
        )

    def visit_monotonically_increasing_id(
        self, expr: MonotonicallyIncreasingIdExpr
    ) -> Union[BlockColumn, ScalarType]:
        """Visit a monotonically_increasing_id expression.

        Args:
            expr: The monotonically_increasing_id expression.

        Returns:
            The result of the monotonically_increasing_id expression as a BlockColumn.
        """
        ctx = TaskContext.get_current()
        assert (
            ctx is not None
        ), "TaskContext is required for monotonically_increasing_id()"

        # Key the counter by expression instance ID so that multiple expressions
        # in the same projection will have isolated row count state.
        # This is required because a single task may process multiple blocks if
        # the upstream data source does not compress the data into a single block.
        counter_key = f"_mono_id_{expr._instance_id}_counter"

        start_idx = ctx.kwargs.get(counter_key, 0)
        num_rows = self.block_accessor.num_rows()
        end_idx = start_idx + num_rows
        ctx.kwargs[counter_key] = end_idx

        # int64 (signed): upper 30 bits = task ID, lower 33 bits = row number.
        # Note end_idx is an exclusive upper bound, as the max row ID is end_idx - 1.
        ROW_BITS = 33
        TASK_BITS = 30
        if end_idx > (1 << ROW_BITS):
            raise ValueError(
                f"Cannot generate monotonically increasing IDs: row count for this task exceeds the maximum allowed value of {(1<<ROW_BITS)-1}"
            )
        if ctx.task_idx >= (1 << TASK_BITS):
            raise ValueError(
                f"Cannot generate monotonically increasing IDs: number of tasks exceeds the maximum allowed value of {(1<<TASK_BITS)-1}"
            )

        partition_mask = ctx.task_idx << ROW_BITS
        ids = partition_mask + np.arange(start_idx, end_idx, dtype=np.int64)

        block_type = self.block_accessor.block_type()
        if block_type == BlockType.PANDAS:
            return pd.Series(ids)
        elif block_type == BlockType.ARROW:
            return pa.array(ids)