test_ml.py

# Copyright 2024 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 unittest import mock

import pandas as pd
import pytest

import bigframes.bigquery._operations.ml as ml_ops
import bigframes.session


@pytest.fixture
def mock_session():
    return mock.create_autospec(spec=bigframes.session.Session)


MODEL_SERIES = pd.Series(
    {
        "modelReference": {
            "projectId": "test-project",
            "datasetId": "test-dataset",
            "modelId": "test-model",
        }
    }
)

MODEL_NAME = "test-project.test-dataset.test-model"


def test_get_model_name_and_session_with_pandas_series_model_input():
    model_name, _ = ml_ops._get_model_name_and_session(MODEL_SERIES)
    assert model_name == MODEL_NAME


def test_get_model_name_and_session_with_pandas_series_model_input_missing_model_reference():
    model_series = pd.Series({"some_other_key": "value"})
    with pytest.raises(
        ValueError, match="modelReference must be present in the pandas Series"
    ):
        ml_ops._get_model_name_and_session(model_series)


@mock.patch("bigframes.pandas.read_pandas")
def test_to_sql_with_pandas_dataframe(read_pandas_mock):
    df = pd.DataFrame({"col1": [1, 2, 3]})
    read_pandas_mock.return_value._to_sql_query.return_value = (
        "SELECT * FROM `pandas_df`",
        [],
        [],
    )
    ml_ops._to_sql(df)
    read_pandas_mock.assert_called_once()


@mock.patch("bigframes.bigquery._operations.ml._get_model_metadata")
@mock.patch("bigframes.pandas.read_pandas")
def test_create_model_with_pandas_dataframe(
    read_pandas_mock, _get_model_metadata_mock, mock_session
):
    df = pd.DataFrame({"col1": [1, 2, 3]})
    read_pandas_mock.return_value._to_sql_query.return_value = (
        "SELECT * FROM `pandas_df`",
        [],
        [],
    )
    ml_ops.create_model("model_name", training_data=df, session=mock_session)
    read_pandas_mock.assert_called_once()
    mock_session.read_gbq_query.assert_called_once()
    generated_sql = mock_session.read_gbq_query.call_args[0][0]
    assert "CREATE MODEL `model_name`" in generated_sql
    assert "AS SELECT * FROM `pandas_df`" in generated_sql


@mock.patch("bigframes.pandas.read_gbq_query")
@mock.patch("bigframes.pandas.read_pandas")
def test_evaluate_with_pandas_dataframe(read_pandas_mock, read_gbq_query_mock):
    df = pd.DataFrame({"col1": [1, 2, 3]})
    read_pandas_mock.return_value._to_sql_query.return_value = (
        "SELECT * FROM `pandas_df`",
        [],
        [],
    )
    ml_ops.evaluate(MODEL_SERIES, input_=df)
    read_pandas_mock.assert_called_once()
    read_gbq_query_mock.assert_called_once()
    generated_sql = read_gbq_query_mock.call_args[0][0]
    assert "ML.EVALUATE" in generated_sql
    assert f"MODEL `{MODEL_NAME}`" in generated_sql
    assert "(SELECT * FROM `pandas_df`)" in generated_sql


@mock.patch("bigframes.pandas.read_gbq_query")
@mock.patch("bigframes.pandas.read_pandas")
def test_predict_with_pandas_dataframe(read_pandas_mock, read_gbq_query_mock):
    df = pd.DataFrame({"col1": [1, 2, 3]})
    read_pandas_mock.return_value._to_sql_query.return_value = (
        "SELECT * FROM `pandas_df`",
        [],
        [],
    )
    ml_ops.predict(MODEL_SERIES, input_=df)
    read_pandas_mock.assert_called_once()
    read_gbq_query_mock.assert_called_once()
    generated_sql = read_gbq_query_mock.call_args[0][0]
    assert "ML.PREDICT" in generated_sql
    assert f"MODEL `{MODEL_NAME}`" in generated_sql
    assert "(SELECT * FROM `pandas_df`)" in generated_sql


@mock.patch("bigframes.pandas.read_gbq_query")
@mock.patch("bigframes.pandas.read_pandas")
def test_explain_predict_with_pandas_dataframe(read_pandas_mock, read_gbq_query_mock):
    df = pd.DataFrame({"col1": [1, 2, 3]})
    read_pandas_mock.return_value._to_sql_query.return_value = (
        "SELECT * FROM `pandas_df`",
        [],
        [],
    )
    ml_ops.explain_predict(MODEL_SERIES, input_=df)
    read_pandas_mock.assert_called_once()
    read_gbq_query_mock.assert_called_once()
    generated_sql = read_gbq_query_mock.call_args[0][0]
    assert "ML.EXPLAIN_PREDICT" in generated_sql
    assert f"MODEL `{MODEL_NAME}`" in generated_sql
    assert "(SELECT * FROM `pandas_df`)" in generated_sql


@mock.patch("bigframes.pandas.read_gbq_query")
def test_global_explain_with_pandas_series_model(read_gbq_query_mock):
    ml_ops.global_explain(MODEL_SERIES)
    read_gbq_query_mock.assert_called_once()
    generated_sql = read_gbq_query_mock.call_args[0][0]
    assert "ML.GLOBAL_EXPLAIN" in generated_sql
    assert f"MODEL `{MODEL_NAME}`" in generated_sql


@mock.patch("bigframes.pandas.read_gbq_query")
@mock.patch("bigframes.pandas.read_pandas")
def test_transform_with_pandas_dataframe(read_pandas_mock, read_gbq_query_mock):
    df = pd.DataFrame({"col1": [1, 2, 3]})
    read_pandas_mock.return_value._to_sql_query.return_value = (
        "SELECT * FROM `pandas_df`",
        [],
        [],
    )
    ml_ops.transform(MODEL_SERIES, input_=df)
    read_pandas_mock.assert_called_once()
    read_gbq_query_mock.assert_called_once()
    generated_sql = read_gbq_query_mock.call_args[0][0]
    assert "ML.TRANSFORM" in generated_sql
    assert f"MODEL `{MODEL_NAME}`" in generated_sql
    assert "(SELECT * FROM `pandas_df`)" in generated_sql