python-bigquery-dataframes/bigframes/testing/utils.py at col_numpy · googleapis/python-bigquery-dataframes

This repository was archived by the owner on Apr 1, 2026. It is now read-only.
550 lines (462 loc) · 19.7 KB
# Copyright 2023 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.
import base64
import decimal
from typing import Iterable, Optional, Sequence, Set, TypeVar, Union
import geopandas as gpd  # type: ignore
import google.api_core.operation
from google.cloud import bigquery, functions_v2
from google.cloud.functions_v2.types import functions
import numpy as np
import pandas as pd
import pandas.api.types as pd_types
import pyarrow as pa  # type: ignore
import pytest
from bigframes import operations as ops
from bigframes.core import expression as ex
import bigframes.functions._utils as bff_utils
import bigframes.pandas as bpd
ML_REGRESSION_METRICS = [
    "mean_absolute_error",
    "mean_squared_error",
    "mean_squared_log_error",
    "median_absolute_error",
    "r2_score",
    "explained_variance",
ML_CLASSFICATION_METRICS = [
    "precision",
    "recall",
    "accuracy",
    "f1_score",
    "log_loss",
    "roc_auc",
ML_GENERATE_TEXT_OUTPUT = [
    "ml_generate_text_llm_result",
    "ml_generate_text_status",
    "prompt",
ML_GENERATE_EMBEDDING_OUTPUT = [
    "ml_generate_embedding_result",
    "ml_generate_embedding_statistics",
    "ml_generate_embedding_status",
    "content",
ML_MULTIMODAL_GENERATE_EMBEDDING_OUTPUT = [
    "ml_generate_embedding_result",
    "ml_generate_embedding_status",
    # start and end sec depend on input format. Images and videos input will contain these 2.
    "ml_generate_embedding_start_sec",
    "ml_generate_embedding_end_sec",
    "content",
SeriesOrIndexT = TypeVar("SeriesOrIndexT", pd.Series, pd.Index)
def pandas_major_version() -> int:
    match = re.search(r"^v?(\d+)", pd.__version__.strip())
    assert match is not None
    return int(match.group(1))
# Prefer this function for tests that run in both ordered and unordered mode
def assert_dfs_equivalent(pd_df: pd.DataFrame, bf_df: bpd.DataFrame, **kwargs):
    bf_df_local = bf_df.to_pandas()
    ignore_order = not bf_df._session._strictly_ordered
    assert_frame_equal(bf_df_local, pd_df, ignore_order=ignore_order, **kwargs)
def assert_series_equivalent(pd_series: pd.Series, bf_series: bpd.Series, **kwargs):
    bf_df_local = bf_series.to_pandas()
    ignore_order = not bf_series._session._strictly_ordered
    assert_series_equal(bf_df_local, pd_series, ignore_order=ignore_order, **kwargs)
def _normalize_all_nulls(col: pd.Series) -> pd.Series:
    if pd_types.is_float_dtype(col.dtype):
        col = col.astype("float64").astype("Float64")
    return col
def _normalize_index_nulls(idx: pd.Index) -> pd.Index:
    if isinstance(idx, pd.MultiIndex):
        new_levels = [
            _normalize_index_nulls(idx.get_level_values(i)) for i in range(idx.nlevels)
        return pd.MultiIndex.from_arrays(new_levels, names=idx.names)
    if idx.hasnans:
        if pd_types.is_float_dtype(idx.dtype):
            idx = idx.astype("float64").astype("Float64")
    return idx
def assert_frame_equal(
    left: pd.DataFrame,
    right: pd.DataFrame,
    ignore_order: bool = False,
    nulls_are_nan: bool = True,
    downcast_object: bool = True,
    **kwargs,
    if ignore_order:
        # Sort by a column to get consistent results.
        if left.index.name != "rowindex":
            left = left.sort_values(
                list(left.columns.drop("geography_col", errors="ignore"))
            ).reset_index(drop=True)
            right = right.sort_values(
                list(right.columns.drop("geography_col", errors="ignore"))
            ).reset_index(drop=True)
        else:
            left = left.sort_index()
            right = right.sort_index()
    # Pandas sometimes likes to produce object dtype columns
    # However, nan/None/Null inconsistency makes comparison futile, convert to typed column
    if downcast_object:
        left = left.apply(lambda x: x.infer_objects())
        right = right.apply(lambda x: x.infer_objects())
    if nulls_are_nan:
        left = left.apply(_normalize_all_nulls)
        right = right.apply(_normalize_all_nulls)
        left.index = _normalize_index_nulls(left.index)
        right.index = _normalize_index_nulls(right.index)
    pd.testing.assert_frame_equal(left, right, **kwargs)
def assert_series_equal(
    left: pd.Series,
    right: pd.Series,
    ignore_order: bool = False,
    nulls_are_nan: bool = True,
    **kwargs,
    if ignore_order:
        if left.index.name is None:
            left = left.sort_values().reset_index(drop=True)
            right = right.sort_values().reset_index(drop=True)
        else:
            left = left.sort_index()
            right = right.sort_index()
    if isinstance(left.index, pd.RangeIndex) or pd_types.is_integer_dtype(
        left.index.dtype,
        left.index = left.index.astype("Int64")
    if isinstance(right.index, pd.RangeIndex) or pd_types.is_integer_dtype(
        right.index.dtype,
        right.index = right.index.astype("Int64")
    if nulls_are_nan:
        left = _normalize_all_nulls(left.infer_objects())
        right = _normalize_all_nulls(right.infer_objects())
        left.index = _normalize_index_nulls(left.index)
        right.index = _normalize_index_nulls(right.index)
        left.name = pd.NA if pd.isna(left.name) else left.name  # type: ignore
        right.name = pd.NA if pd.isna(right.name) else right.name  # type: ignore
    pd.testing.assert_series_equal(left, right, **kwargs)
def assert_index_equal(left, right, **kwargs):
    pd.testing.assert_index_equal(left, right, **kwargs)
def _standardize_index(idx):
    return pd.Index(list(idx), name=idx.name)
def assert_pandas_index_equal_ignore_index_type(idx0, idx1):
    idx0 = _standardize_index(idx0)
    idx1 = _standardize_index(idx1)
    pd.testing.assert_index_equal(idx0, idx1)
def convert_pandas_dtypes(df: pd.DataFrame, bytes_col: bool):
    """Convert pandas dataframe dtypes compatible with bigframes dataframe."""
    # TODO(chelsealin): updates the function to accept dtypes as input rather than
    # hard-code the column names here.
    # Convert basic types columns
    df["bool_col"] = df["bool_col"].astype(pd.BooleanDtype())
    df["int64_col"] = df["int64_col"].astype(pd.Int64Dtype())
    df["int64_too"] = df["int64_too"].astype(pd.Int64Dtype())
    df["float64_col"] = df["float64_col"].astype(pd.Float64Dtype())
    df["string_col"] = df["string_col"].astype(pd.StringDtype(storage="pyarrow"))
    if "rowindex" in df.columns:
        df["rowindex"] = df["rowindex"].astype(pd.Int64Dtype())
    if "rowindex_2" in df.columns:
        df["rowindex_2"] = df["rowindex_2"].astype(pd.Int64Dtype())
    # Convert time types columns. The `astype` works for Pandas 2.0 but hits an assert
    # error at Pandas 1.5. Hence, we have to convert to arrow table and convert back
    # to pandas dataframe.
    if not isinstance(df["date_col"].dtype, pd.ArrowDtype):
        df["date_col"] = pd.to_datetime(df["date_col"], format="%Y-%m-%d")
        arrow_table = pa.Table.from_pandas(
            pd.DataFrame(df, columns=["date_col"]),
            schema=pa.schema([("date_col", pa.date32())]),
        df["date_col"] = arrow_table.to_pandas(types_mapper=pd.ArrowDtype)["date_col"]
    if not isinstance(df["datetime_col"].dtype, pd.ArrowDtype):
        df["datetime_col"] = pd.to_datetime(
            df["datetime_col"], format="%Y-%m-%d %H:%M:%S"
        arrow_table = pa.Table.from_pandas(
            pd.DataFrame(df, columns=["datetime_col"]),
            schema=pa.schema([("datetime_col", pa.timestamp("us"))]),
        df["datetime_col"] = arrow_table.to_pandas(types_mapper=pd.ArrowDtype)[
            "datetime_col"
    if not isinstance(df["time_col"].dtype, pd.ArrowDtype):
        df["time_col"] = pd.to_datetime(df["time_col"], format="%H:%M:%S.%f")
        arrow_table = pa.Table.from_pandas(
            pd.DataFrame(df, columns=["time_col"]),
            schema=pa.schema([("time_col", pa.time64("us"))]),
        df["time_col"] = arrow_table.to_pandas(types_mapper=pd.ArrowDtype)["time_col"]
    if not isinstance(df["timestamp_col"].dtype, pd.ArrowDtype):
        df["timestamp_col"] = pd.to_datetime(
            df["timestamp_col"], format="%Y-%m-%d %H:%M:%S.%f%Z"
        arrow_table = pa.Table.from_pandas(
            pd.DataFrame(df, columns=["timestamp_col"]),
            schema=pa.schema([("timestamp_col", pa.timestamp("us", tz="UTC"))]),
        df["timestamp_col"] = arrow_table.to_pandas(types_mapper=pd.ArrowDtype)[
            "timestamp_col"
    if not isinstance(df["duration_col"].dtype, pd.ArrowDtype):
        df["duration_col"] = df["duration_col"].astype(pd.Int64Dtype())
        arrow_table = pa.Table.from_pandas(
            pd.DataFrame(df, columns=["duration_col"]),
            schema=pa.schema([("duration_col", pa.duration("us"))]),
        df["duration_col"] = arrow_table.to_pandas(types_mapper=pd.ArrowDtype)[
            "duration_col"
    # Convert geography types columns.
    if "geography_col" in df.columns:
        df["geography_col"] = df["geography_col"].astype(
            pd.StringDtype(storage="pyarrow")
        df["geography_col"] = gpd.GeoSeries.from_wkt(
            df["geography_col"].replace({np.nan: None})
    if bytes_col and not isinstance(df["bytes_col"].dtype, pd.ArrowDtype):
        df["bytes_col"] = df["bytes_col"].apply(
            lambda value: base64.b64decode(value) if not pd.isnull(value) else value
        arrow_table = pa.Table.from_pandas(
            pd.DataFrame(df, columns=["bytes_col"]),
            schema=pa.schema([("bytes_col", pa.binary())]),
        df["bytes_col"] = arrow_table.to_pandas(types_mapper=pd.ArrowDtype)["bytes_col"]
    if not isinstance(df["numeric_col"].dtype, pd.ArrowDtype):
        # Convert numeric types column.
        df["numeric_col"] = df["numeric_col"].apply(
            lambda value: decimal.Decimal(str(value)) if value else None  # type: ignore
        arrow_table = pa.Table.from_pandas(
            pd.DataFrame(df, columns=["numeric_col"]),
            schema=pa.schema([("numeric_col", pa.decimal128(38, 9))]),
        df["numeric_col"] = arrow_table.to_pandas(types_mapper=pd.ArrowDtype)[
            "numeric_col"
def assert_pandas_df_equal_pca_components(actual, expected, **kwargs):
    """Compare two pandas dataframes representing PCA components. The columns
    required to be present in the dataframes are:
        numerical_value: numeric,
        categorical_value: List[object(category, value)]
    The index types of `actual` and `expected` are ignored in the comparison.
        actual: Actual Pandas DataFrame
        expected: Expected Pandas DataFrame
        kwargs: kwargs to use in `pandas.testing.assert_series_equal` per column
    # Compare the index, columns and values separately, as the polarity of the
    # PCA vectors can be arbitrary
    pd.testing.assert_index_equal(
        actual.index, expected.index.astype(actual.index.dtype)
    )  # dtype agnostic index comparison
    pd.testing.assert_index_equal(actual.columns, expected.columns)
    for column in expected.columns:
        try:
            pd.testing.assert_series_equal(actual[column], expected[column], **kwargs)
        except AssertionError:
            if column not in {"numerical_value", "categorical_value"}:
                raise
            # Allow for sign difference per numeric/categorical column
            if column == "numerical_value":
                actual_ = -actual[column]
                expected_ = expected[column]
            else:
                # In this column each element is an array of objects, where the
                # object has attributes "category" and "value". For the sake of
                # comparison let's normalize by flipping the polarity of "value".
                def normalize_array_of_objects(arr, reverse_polarity=False):
                    newarr = []
                    for element in arr:
                        newelement = dict(element)
                        if reverse_polarity:
                            newelement["value"] = -newelement["value"]
                        newarr.append(newelement)
                    return sorted(newarr, key=lambda d: d["category"])
                actual_ = actual[column].apply(normalize_array_of_objects, args=(True,))
                expected_ = expected[column].apply(normalize_array_of_objects)
            pd.testing.assert_series_equal(actual_, expected_, **kwargs)
def assert_pandas_df_equal_pca(actual, expected, **kwargs):
    """Compare two pandas dataframes representing PCA predictions. The columns
    in the dataframes are expected to be numeric.
        actual: Actual Pandas DataFrame
        expected: Expected Pandas DataFrame
        kwargs: kwargs to use in `pandas.testing.assert_series_equal` per column
    # Compare the index, columns and values separately, as the polarity of the
    # PCA vector can be arbitrary
    pd.testing.assert_index_equal(actual.index, expected.index)
    pd.testing.assert_index_equal(actual.columns, expected.columns)
    for column in expected.columns:
        try:
            pd.testing.assert_series_equal(actual[column], expected[column], **kwargs)
        except AssertionError:
            # Allow for sign difference per column
            pd.testing.assert_series_equal(-actual[column], expected[column], **kwargs)
def check_pandas_df_schema_and_index(
    pd_df: pd.DataFrame,
    columns: Iterable,
    index: Optional[Union[int, Iterable]] = None,
    col_exact: bool = True,
    """Check pandas df schema and index. But not the values.
        pd_df: the input pandas df
        columns: target columns to check with
        index: int or Iterable or None, default None. If int, only check the length (index size) of the df. If Iterable, check index values match. If None, skip checking index.
        col_exact: If True, check the columns param are exact match. Otherwise only check the df contains all of those columns
    if col_exact:
        assert list(pd_df.columns) == list(columns)
        assert set(columns) <= set(pd_df.columns)
    if index is None:
        pass
    elif isinstance(index, int):
        assert len(pd_df) == index
    elif isinstance(index, Iterable):
        assert list(pd_df.index) == list(index)
        raise ValueError("Unsupported index type.")
def get_remote_function_endpoints(
    bigquery_client: bigquery.Client, dataset_id: str
) -> Set[str]:
    """Get endpoints used by the remote functions in a datset"""
    endpoints = set()
    routines = bigquery_client.list_routines(dataset=dataset_id)
    for routine in routines:
        rf_options = routine._properties.get("remoteFunctionOptions")
        if not rf_options:
            continue
        rf_endpoint = rf_options.get("endpoint")
        if rf_endpoint:
            endpoints.add(rf_endpoint)
    return endpoints
def get_cloud_functions(
    functions_client: functions_v2.FunctionServiceClient,
    project: str,
    location: str,
    name: Optional[str] = None,
    name_prefix: Optional[str] = None,
) -> Iterable[functions.ListFunctionsResponse]:
    """Get the cloud functions in the given project and location."""
    assert (
        not name or not name_prefix
    ), "Either 'name' or 'name_prefix' can be passed but not both."
    _, location = bff_utils.get_remote_function_locations(location)
    parent = f"projects/{project}/locations/{location}"
    request = functions_v2.ListFunctionsRequest(parent=parent)
    page_result = functions_client.list_functions(request=request)
    for response in page_result:
        # If name is provided and it does not match then skip
        if bool(name):
            full_name = parent + f"/functions/{name}"
            if response.name != full_name:
                continue
        # If name prefix is provided and it does not match then skip
        elif bool(name_prefix):
            full_name_prefix = parent + f"/functions/{name_prefix}"
            if not response.name.startswith(full_name_prefix):
                continue
        yield response
def delete_cloud_function(
    functions_client: functions_v2.FunctionServiceClient, full_name: str
) -> google.api_core.operation.Operation:
    """Delete a cloud function with the given fully qualified name."""
    request = functions_v2.DeleteFunctionRequest(name=full_name)
    operation = functions_client.delete_function(request=request)
    return operation
def get_first_file_from_wildcard(path):
    return path.replace("*", "000000000000")
def cleanup_function_assets(
    bigframes_func,
    bigquery_client,
    cloudfunctions_client=None,
    ignore_failures=True,
    """Clean up the GCP assets behind a bigframess function."""
    # Clean up bigframes bigquery function.
        bigquery_client.delete_routine(bigframes_func.bigframes_bigquery_function)
    except Exception:
        # By default don't raise exception in cleanup.
        if not ignore_failures:
            raise
    if not ignore_failures:
        # Make sure that the BQ routins is actually deleted
        with pytest.raises(google.api_core.exceptions.NotFound):
            bigquery_client.get_routine(bigframes_func.bigframes_bigquery_function)
    # Clean up bigframes cloud run function
    if cloudfunctions_client:
        # Clean up cloud function
        try:
            delete_cloud_function(
                cloudfunctions_client, bigframes_func.bigframes_cloud_function
        except Exception:
            # By default don't raise exception in cleanup.
            if not ignore_failures:
                raise
        if not ignore_failures:
            # Make sure the cloud run function is actually deleted
            try:
                gcf = cloudfunctions_client.get_function(
                    name=bigframes_func.bigframes_cloud_function
                assert gcf.state is functions_v2.Function.State.DELETING
            except google.cloud.exceptions.NotFound:
def _apply_ops_to_sql(
    obj: bpd.DataFrame,
    ops_list: Sequence[ex.Expression],
    new_names: Sequence[str],
    """Applies a list of ops to the given DataFrame and returns the SQL
    representing the resulting DataFrame."""
    array_value = obj._block.expr
    result, old_names = array_value.compute_values(ops_list)
    # Rename columns for deterministic golden SQL results.
    assert len(old_names) == len(new_names)
    col_ids = {old_name: new_name for old_name, new_name in zip(old_names, new_names)}
    result = result.rename_columns(col_ids).select_columns(new_names)
    sql = result.session._executor.to_sql(result, enable_cache=False)
    return sql
def _apply_binary_op(
    obj: bpd.DataFrame,
    op: ops.BinaryOp,
    l_arg: str,
    r_arg: Union[str, ex.Expression],
    """Applies a binary op to the given DataFrame and return the SQL representing
    the resulting DataFrame."""
    return _apply_nary_op(obj, op, l_arg, r_arg)
def _apply_nary_op(
    obj: bpd.DataFrame,
    op: Union[ops.BinaryOp, ops.NaryOp],
    *args: Union[str, ex.Expression],
    """Applies a nary op to the given DataFrame and return the SQL representing
    the resulting DataFrame."""
    op_expr = op.as_expr(*args)
    sql = _apply_ops_to_sql(obj, [op_expr], [args[0]])  # type: ignore
    return sql
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

utils.py

Latest commit

History

utils.py

File metadata and controls
