Repository: ray-project/ray
File: python/ray/data/expressions.py (Python)

Code before:
    Example:
        >>> from ray.data.expressions import download
        >>> import ray
        >>> # Create dataset with URIs
        >>> ds = ray.data.from_items([
        ...     {"uri": "s3://bucket/file1.jpg", "id": "1"},
        ...     {"uri": "s3://bucket/file2.jpg", "id": "2"}
        ... ])
        >>> # Add downloaded bytes column
        >>> ds_with_bytes = ds.with_column("bytes", download("uri"))
    """
    return DownloadExpr(uri_column_name=uri_column_name)


@PublicAPI(stability="alpha")
def monotonically_increasing_id() -> MonotonicallyIncreasingIdExpr:
    """
    Create an expression that generates monotonically increasing IDs.

    The generated IDs are guaranteed to be monotonically increasing and unique,
    but not consecutive. The current implementation puts the block ID in the
    upper 31 bits, and the lower 33 bits represent the record number within
    each block. The assumption is that the dataset has less than 1 billion blocks,
    and each block has less than 8 billion records.

    The function is non-deterministic because its result depends on block IDs.

    Returns:
        A MonotonicallyIncreasingIdExpr that generates unique IDs.

    Example:
        >>> from ray.data.expressions import monotonically_increasing_id
        >>> import ray
        >>> ds = ray.data.range(10)
        >>> ds = ds.with_column("uid", monotonically_increasing_id())
    """
    return MonotonicallyIncreasingIdExpr()


# ──────────────────────────────────────
# Public API for evaluation
# ──────────────────────────────────────
# Note: Implementation details are in _expression_evaluator.py

# Re-export eval_expr for public use

__all__ = [
    "Operation",
    "Expr",
    "ColumnExpr",
    "LiteralExpr",

Review comment (suggestion):
Dumb question -- if this is non-deterministic, when would this implementation be better than just doing random IDs? I think this is monotonically increasing with respect to task indices, but task indices are non-deterministic. Worried this might confuse or mislead users who expect this to be to be monotonically increasing with respect to the input data.

Code after:
            and all(a.structurally_equals(b) for a, b in zip(self.args, other.args))
            and self.kwargs.keys() == other.kwargs.keys()
            and all(
                self.kwargs[k].structurally_equals(other.kwargs[k])
                for k in self.kwargs.keys()
            )
        )


def _create_udf_callable(
    fn: Callable[..., BatchColumn],
    return_dtype: DataType,
) -> Callable[..., UDFExpr]:
    """Create a callable that generates UDFExpr when called with expressions.

    Args:
        fn: The user-defined function to wrap. Can be a regular function
            or an _CallableClassUDF instance (for callable classes).
        return_dtype: The return data type of the UDF

    Returns:
        A callable that creates UDFExpr instances when called with expressions
    """

    def udf_callable(*args, **kwargs) -> UDFExpr:
        # Convert arguments to expressions if they aren't already
        expr_args = []
        for arg in args:
            if isinstance(arg, Expr):
                expr_args.append(arg)
            else:
                expr_args.append(LiteralExpr(arg))

        expr_kwargs = {}
        for k, v in kwargs.items():
            if isinstance(v, Expr):
                expr_kwargs[k] = v
            else:
                expr_kwargs[k] = LiteralExpr(v)

        return UDFExpr(
            fn=fn,
            args=expr_args,
            kwargs=expr_kwargs,
            data_type=return_dtype,
        )

    # Preserve original function metadata
    functools.update_wrapper(udf_callable, fn)

    # Store the original function for access if needed