mathematical.py

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from typing import Union

from bigframes import dataframe, dtypes
from bigframes import operations as ops
from bigframes import series


def rand(input_data: Union[series.Series, dataframe.DataFrame]) -> series.Series:
    """
    Generates a pseudo-random value of type FLOAT64 in the range of [0, 1),
    inclusive of 0 and exclusive of 1.

    .. warning::
        This method introduces non-determinism to the expression. Reading the
        same column twice may result in different results. This value might
        change and not to use this value or any value derived from it as a join
        key.

    **Examples:**

        >>> import bigframes.pandas as bpd
        >>> import bigframes.bigquery as bbq
        >>> df = bpd.DataFrame({"a": [1, 2, 3]})
        >>> df['random'] = bbq.rand(df)
        >>> # Resulting column 'random' will contain random floats between 0 and 1.

    Args:
        input_data (bigframes.pandas.Series or bigframes.pandas.DataFrame):
            A Series or DataFrame to determine the number of rows and the index
            of the result. The actual values in this input are ignored.

    Returns:
        bigframes.pandas.Series: A new Series of random float values.
    """
    if isinstance(input_data, dataframe.DataFrame):
        if len(input_data.columns) == 0:
            raise ValueError("Input DataFrame must have at least one column.")
        # Use the first column as anchor
        anchor = input_data.iloc[:, 0]
    elif isinstance(input_data, series.Series):
        anchor = input_data
    else:
        raise TypeError(
            f"Unsupported type {type(input_data)}. "
            "Expected bigframes.pandas.Series or bigframes.pandas.DataFrame."
        )

    op = ops.SqlScalarOp(
        _output_type=dtypes.FLOAT_DTYPE,
        sql_template="RAND()",
        is_deterministic=False,
    )
    return anchor._apply_nary_op(op, [])