test: stop checking ml large tests exact numbers (#690)

* test: stop checking ml large tests exact numbers

* clean up

Author: GarrettWu

tests/system/large/ml/test_cluster.py

@@ -13,13 +13,11 @@
 # limitations under the License.
 
 import pandas as pd
-import pytest
 
 from bigframes.ml import cluster
-from tests.system.utils import assert_pandas_df_equal
+from tests.system import utils
 
 
-@pytest.mark.flaky(retries=2)
 def test_cluster_configure_fit_score_predict(
     session, penguins_df_default_index, dataset_id
 ):
@@ -88,26 +86,18 @@ def test_cluster_configure_fit_score_predict(
 
     # Check score to ensure the model was fitted
     score_result = model.score(new_penguins).to_pandas()
-    score_expected = pd.DataFrame(
-        {"davies_bouldin_index": [1.502182], "mean_squared_distance": [1.953408]},
-        dtype="Float64",
-    )
-    score_expected = score_expected.reindex(index=score_expected.index.astype("Int64"))
 
-    pd.testing.assert_frame_equal(
-        score_result, score_expected, check_exact=False, rtol=0.1
-    )
+    eval_metrics = ["davies_bouldin_index", "mean_squared_distance"]
+    utils.check_pandas_df_schema_and_index(score_result, columns=eval_metrics, index=1)
 
     predictions = model.predict(new_penguins).to_pandas()
     assert predictions.shape == (4, 9)
-    result = predictions[["CENTROID_ID"]]
-    expected = pd.DataFrame(
-        {"CENTROID_ID": [2, 3, 1, 2]},
-        dtype="Int64",
-        index=pd.Index(["test1", "test2", "test3", "test4"], dtype="string[pyarrow]"),
+    utils.check_pandas_df_schema_and_index(
+        predictions,
+        columns=["CENTROID_ID"],
+        index=["test1", "test2", "test3", "test4"],
+        col_exact=False,
     )
-    expected.index.name = "observation"
-    assert_pandas_df_equal(result, expected, ignore_order=True)
 
     # save, load, check n_clusters to ensure configuration was kept
     reloaded_model = model.to_gbq(

tests/system/large/ml/test_compose.py

@@ -12,9 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pandas
-
 from bigframes.ml import compose, preprocessing
+from tests.system import utils
 
 
 def test_columntransformer_standalone_fit_and_transform(
@@ -45,26 +44,18 @@ def test_columntransformer_standalone_fit_and_transform(
     )
     result = transformer.transform(new_penguins_df).to_pandas()
 
-    expected = pandas.DataFrame(
-        {
-            "onehotencoded_species": [
-                [{"index": 1, "value": 1.0}],
-                [{"index": 1, "value": 1.0}],
-                [{"index": 2, "value": 1.0}],
-            ],
-            "standard_scaled_culmen_length_mm": [
-                -0.811119671289163,
-                -0.9945520581113803,
-                -1.104611490204711,
-            ],
-            "min_max_scaled_culmen_length_mm": [0.269, 0.232, 0.210],
-            "standard_scaled_flipper_length_mm": [-0.350044, -1.418336, -0.9198],
-        },
-        index=pandas.Index([1633, 1672, 1690], dtype="Int64", name="tag_number"),
+    utils.check_pandas_df_schema_and_index(
+        result,
+        columns=[
+            "onehotencoded_species",
+            "standard_scaled_culmen_length_mm",
+            "min_max_scaled_culmen_length_mm",
+            "standard_scaled_flipper_length_mm",
+        ],
+        index=[1633, 1672, 1690],
+        col_exact=False,
     )
 
-    pandas.testing.assert_frame_equal(result, expected, rtol=0.1, check_dtype=False)
-
 
 def test_columntransformer_standalone_fit_transform(new_penguins_df):
     transformer = compose.ColumnTransformer(
@@ -86,25 +77,17 @@ def test_columntransformer_standalone_fit_transform(new_penguins_df):
         new_penguins_df[["species", "culmen_length_mm", "flipper_length_mm"]]
     ).to_pandas()
 
-    expected = pandas.DataFrame(
-        {
-            "onehotencoded_species": [
-                [{"index": 1, "value": 1.0}],
-                [{"index": 1, "value": 1.0}],
-                [{"index": 2, "value": 1.0}],
-            ],
-            "standard_scaled_culmen_length_mm": [
-                1.313249,
-                -0.20198,
-                -1.111118,
-            ],
-            "standard_scaled_flipper_length_mm": [1.251098, -1.196588, -0.054338],
-        },
-        index=pandas.Index([1633, 1672, 1690], dtype="Int64", name="tag_number"),
+    utils.check_pandas_df_schema_and_index(
+        result,
+        columns=[
+            "onehotencoded_species",
+            "standard_scaled_culmen_length_mm",
+            "standard_scaled_flipper_length_mm",
+        ],
+        index=[1633, 1672, 1690],
+        col_exact=False,
     )
 
-    pandas.testing.assert_frame_equal(result, expected, rtol=0.1, check_dtype=False)
-
 
 def test_columntransformer_save_load(new_penguins_df, dataset_id):
     transformer = compose.ColumnTransformer(
@@ -147,23 +130,13 @@ def test_columntransformer_save_load(new_penguins_df, dataset_id):
         new_penguins_df[["species", "culmen_length_mm", "flipper_length_mm"]]
     ).to_pandas()
 
-    # TODO(b/340888429): fix type error
-    expected = pandas.DataFrame(  # type: ignore
-        {
-            "onehotencoded_species": [
-                [{"index": 1, "value": 1.0}],
-                [{"index": 1, "value": 1.0}],
-                [{"index": 2, "value": 1.0}],
-            ],
-            "standard_scaled_culmen_length_mm": [
-                1.313249,
-                -0.20198,
-                -1.111118,
-            ],
-            "standard_scaled_flipper_length_mm": [1.251098, -1.196588, -0.054338],
-        },
-        index=pandas.Index([1633, 1672, 1690], dtype="Int64", name="tag_number"),
+    utils.check_pandas_df_schema_and_index(
+        result,
+        columns=[
+            "onehotencoded_species",
+            "standard_scaled_culmen_length_mm",
+            "standard_scaled_flipper_length_mm",
+        ],
+        index=[1633, 1672, 1690],
+        col_exact=False,
     )
-
-    # TODO(b/340888429): fix type error
-    pandas.testing.assert_frame_equal(result, expected, rtol=0.1, check_dtype=False)  # type: ignore

tests/system/large/ml/test_core.py

@@ -12,14 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pandas
-import pytest
-
 from bigframes.ml import globals
+from tests.system import utils
 
 
-# TODO(garrettwu): Re-enable or not check exact numbers.
-@pytest.mark.skip(reason="bqml regression")
 def test_bqml_e2e(session, dataset_id, penguins_df_default_index, new_penguins_df):
     df = penguins_df_default_index.dropna()
     X_train = df[
@@ -38,41 +34,33 @@ def test_bqml_e2e(session, dataset_id, penguins_df_default_index, new_penguins_d
         X_train, y_train, options={"model_type": "linear_reg"}
     )
 
+    eval_metrics = [
+        "mean_absolute_error",
+        "mean_squared_error",
+        "mean_squared_log_error",
+        "median_absolute_error",
+        "r2_score",
+        "explained_variance",
+    ]
     # no data - report evaluation from the automatic data split
     evaluate_result = model.evaluate().to_pandas()
-    evaluate_expected = pandas.DataFrame(
-        {
-            "mean_absolute_error": [225.817334],
-            "mean_squared_error": [80540.705944],
-            "mean_squared_log_error": [0.004972],
-            "median_absolute_error": [173.080816],
-            "r2_score": [0.87529],
-            "explained_variance": [0.87529],
-        },
-        dtype="Float64",
-    )
-    evaluate_expected = evaluate_expected.reindex(
-        index=evaluate_expected.index.astype("Int64")
-    )
-    pandas.testing.assert_frame_equal(
-        evaluate_result, evaluate_expected, check_exact=False, rtol=0.1
+    utils.check_pandas_df_schema_and_index(
+        evaluate_result, columns=eval_metrics, index=1
     )
 
     # evaluate on all training data
     evaluate_result = model.evaluate(df).to_pandas()
-    pandas.testing.assert_frame_equal(
-        evaluate_result, evaluate_expected, check_exact=False, rtol=0.1
+    utils.check_pandas_df_schema_and_index(
+        evaluate_result, columns=eval_metrics, index=1
     )
 
     # predict new labels
     predictions = model.predict(new_penguins_df).to_pandas()
-    expected = pandas.DataFrame(
-        {"predicted_body_mass_g": [4030.1, 3280.8, 3177.9]},
-        dtype="Float64",
-        index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"),
-    )
-    pandas.testing.assert_frame_equal(
-        predictions[["predicted_body_mass_g"]], expected, check_exact=False, rtol=0.1
+    utils.check_pandas_df_schema_and_index(
+        predictions,
+        columns=["predicted_body_mass_g"],
+        index=[1633, 1672, 1690],
+        col_exact=False,
     )
 
     new_name = f"{dataset_id}.my_model"
@@ -108,42 +96,34 @@ def test_bqml_manual_preprocessing_e2e(
         X_train, y_train, transforms=transforms, options=options
     )
 
+    eval_metrics = [
+        "mean_absolute_error",
+        "mean_squared_error",
+        "mean_squared_log_error",
+        "median_absolute_error",
+        "r2_score",
+        "explained_variance",
+    ]
+
     # no data - report evaluation from the automatic data split
     evaluate_result = model.evaluate().to_pandas()
-    evaluate_expected = pandas.DataFrame(
-        {
-            "mean_absolute_error": [309.477334],
-            "mean_squared_error": [152184.227218],
-            "mean_squared_log_error": [0.009524],
-            "median_absolute_error": [257.727777],
-            "r2_score": [0.764356],
-            "explained_variance": [0.764356],
-        },
-        dtype="Float64",
-    )
-    evaluate_expected = evaluate_expected.reindex(
-        index=evaluate_expected.index.astype("Int64")
-    )
-
-    pandas.testing.assert_frame_equal(
-        evaluate_result, evaluate_expected, check_exact=False, rtol=0.1
+    utils.check_pandas_df_schema_and_index(
+        evaluate_result, columns=eval_metrics, index=1
     )
 
     # evaluate on all training data
     evaluate_result = model.evaluate(df).to_pandas()
-    pandas.testing.assert_frame_equal(
-        evaluate_result, evaluate_expected, check_exact=False, rtol=0.1
+    utils.check_pandas_df_schema_and_index(
+        evaluate_result, columns=eval_metrics, index=1
     )
 
     # predict new labels
     predictions = model.predict(new_penguins_df).to_pandas()
-    expected = pandas.DataFrame(
-        {"predicted_body_mass_g": [3968.8, 3176.3, 3545.2]},
-        dtype="Float64",
-        index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"),
-    )
-    pandas.testing.assert_frame_equal(
-        predictions[["predicted_body_mass_g"]], expected, check_exact=False, rtol=0.1
+    utils.check_pandas_df_schema_and_index(
+        predictions,
+        columns=["predicted_body_mass_g"],
+        index=[1633, 1672, 1690],
+        col_exact=False,
     )
 
     new_name = f"{dataset_id}.my_model"
@@ -168,24 +148,9 @@ def test_bqml_standalone_transform(penguins_df_default_index, new_penguins_df):
     )
 
     transformed = model.transform(new_penguins_df).to_pandas()
-    expected = pandas.DataFrame(
-        {
-            "scaled_culmen_length_mm": [-0.8099, -0.9931, -1.103],
-            "onehotencoded_species": [
-                [{"index": 1, "value": 1.0}],
-                [{"index": 1, "value": 1.0}],
-                [{"index": 2, "value": 1.0}],
-            ],
-        },
-        index=pandas.Index([1633, 1672, 1690], name="tag_number", dtype="Int64"),
-    )
-    expected["scaled_culmen_length_mm"] = expected["scaled_culmen_length_mm"].astype(
-        "Float64"
-    )
-    pandas.testing.assert_frame_equal(
-        transformed[["scaled_culmen_length_mm", "onehotencoded_species"]],
-        expected,
-        check_exact=False,
-        rtol=0.1,
-        check_dtype=False,
+    utils.check_pandas_df_schema_and_index(
+        transformed,
+        columns=["scaled_culmen_length_mm", "onehotencoded_species"],
+        index=[1633, 1672, 1690],
+        col_exact=False,
     )