Merge branch 'main' into issue687-numpy

Author: tswast

noxfile.py

@@ -195,6 +195,7 @@ def run_unit(session, install_test_extra):
     install_unittest_dependencies(session, install_test_extra, "-c", constraints_path)
 
     # Run py.test against the unit tests.
+    scripts_path = "scripts"
     tests_path = os.path.join("tests", "unit")
     third_party_tests_path = os.path.join("third_party", "bigframes_vendored")
     session.run(
@@ -209,6 +210,7 @@ def run_unit(session, install_test_extra):
         "--cov-fail-under=0",
         tests_path,
         third_party_tests_path,
+        scripts_path,
         *session.posargs,
     )
 

scripts/__init__.py

@@ -0,0 +1,13 @@
+# 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.

scripts/publish_api_coverage.py

@@ -107,8 +107,15 @@ def names_from_signature(signature):
 
 
 def calculate_missing_parameters(bigframes_function, target_function):
-    bigframes_params = names_from_signature(inspect.signature(bigframes_function))
-    target_params = names_from_signature(inspect.signature(target_function))
+    # Some built-in functions can't be inspected. These raise a ValueError.
+    try:
+        bigframes_signature = inspect.signature(bigframes_function)
+        target_signature = inspect.signature(target_function)
+    except ValueError:
+        return {}
+
+    bigframes_params = names_from_signature(bigframes_signature)
+    target_params = names_from_signature(target_signature)
     return target_params - bigframes_params
 
 
@@ -164,13 +171,20 @@ def generate_pandas_api_coverage():
                 token_type = "property"
 
             is_in_bigframes = hasattr(bigframes_obj, member)
-            requires_index = False
-            requires_ordering = False
+            requires_index = ""
+            requires_ordering = ""
 
             if is_in_bigframes:
                 attr = getattr(bigframes_obj, member)
-                requires_index = hasattr(attr, "_validations_requires_index")
-                requires_ordering = hasattr(attr, "_validations_requires_ordering")
+
+                # TODO(b/361101138): Add check/documentation for partial
+                # support (e.g. with some parameters).
+                requires_index = (
+                    "Y" if hasattr(attr, "_validations_requires_index") else ""
+                )
+                requires_ordering = (
+                    "Y" if hasattr(attr, "_validations_requires_ordering") else ""
+                )
 
             api_patterns.append(
                 [
@@ -279,9 +293,12 @@ def build_api_coverage_table(bigframes_version: str, release_version: str):
     sklearn_cov_df["module"] = "bigframes.ml"
     combined_df = pd.concat([pandas_cov_df, sklearn_cov_df])
     combined_df["timestamp"] = pd.Timestamp.now()
+    # BigQuery only supports microsecond precision timestamps.
+    combined_df["timestamp"] = combined_df["timestamp"].astype("datetime64[us]")
     combined_df["bigframes_version"] = bigframes_version
     combined_df["release_version"] = release_version
-    return combined_df.infer_objects().convert_dtypes()
+    combined_df = combined_df.infer_objects().convert_dtypes()
+    return combined_df
 
 
 def format_api(api_names, is_in_bigframes, api_prefix):
@@ -313,16 +330,14 @@ def generate_api_coverage(df, api_prefix):
                 api_prefix,
             ),
             "Implemented": "",
-            "Requires index": "",
-            "Requires ordering": "",
+            "Requires index": dataframe_apis["requires_index"],
+            "Requires ordering": dataframe_apis["requires_ordering"],
             "Missing parameters": dataframe_apis["missing_parameters"],
         }
     )
     dataframe_table.loc[fully_implemented, "Implemented"] = "Y"
     dataframe_table.loc[partial_implemented, "Implemented"] = "P"
     dataframe_table.loc[not_implemented, "Implemented"] = "N"
-    dataframe_table.loc[dataframe_apis["requires_index"], "Requires index"] = "Y"
-    dataframe_table.loc[dataframe_apis["requires_ordering"], "Requires ordering"] = "Y"
     return dataframe_table
 
 

scripts/test_publish_api_coverage.py

@@ -12,36 +12,48 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import sys
+
 import pandas
-import publish_api_coverage
+import pytest
+
+from . import publish_api_coverage
+
+
+@pytest.fixture
+def api_coverage_df():
+    return publish_api_coverage.build_api_coverage_table("my_bf_ver", "my_release_ver")
 
 
-def test_api_coverage_produces_expected_schema():
-    df = publish_api_coverage.build_api_coverage_table("my_bf_ver", "my_release_ver")
+def test_api_coverage_produces_expected_schema(api_coverage_df):
+    if sys.version.split(".")[:2] == ["3", "9"]:
+        pytest.skip(
+            "Python 3.9 uses older pandas without good microsecond timestamp support."
+        )
+
     pandas.testing.assert_series_equal(
-        df.dtypes,
+        api_coverage_df.dtypes,
         pandas.Series(
-            data=[
-                "string",
-                "string",
-                "string",
-                "boolean",
-                "string",
-                "string",
-                "datetime64[ns]",
-                "string",
-                "string",
-            ],
-            index=[
-                "api",
-                "pattern",
-                "kind",
-                "is_in_bigframes",
-                "missing_parameters",
-                "module",
-                "timestamp",
-                "bigframes_version",
-                "release_version",
-            ],
+            data={
+                # Note to developer: if you update this test, you will also
+                # need to update schema of the API coverage BigQuery table in
+                # the bigframes-metrics project.
+                "api": "string",
+                "pattern": "string",
+                "kind": "string",
+                "is_in_bigframes": "boolean",
+                "missing_parameters": "string",
+                "requires_index": "string",
+                "requires_ordering": "string",
+                "module": "string",
+                "timestamp": "datetime64[us]",
+                "bigframes_version": "string",
+                "release_version": "string",
+            },
         ),
     )
+
+
+def test_api_coverage_produces_missing_parameters(api_coverage_df):
+    """Make sure at least some functions have reported missing parameters."""
+    assert (api_coverage_df["missing_parameters"].str.len() > 0).any()

tests/system/small/ml/test_llm.py

@@ -409,6 +409,7 @@ def test_gemini_text_generator_predict_with_params_success(
     "model_name",
     ("claude-3-sonnet", "claude-3-haiku", "claude-3-5-sonnet", "claude-3-opus"),
 )
+@pytest.mark.flaky(retries=3, delay=120)
 def test_claude3_text_generator_create_load(
     dataset_id, model_name, session, session_us_east5, bq_connection
 ):
@@ -433,7 +434,7 @@ def test_claude3_text_generator_create_load(
     "model_name",
     ("claude-3-sonnet", "claude-3-haiku", "claude-3-5-sonnet", "claude-3-opus"),
 )
-@pytest.mark.flaky(retries=2)
+@pytest.mark.flaky(retries=3, delay=120)
 def test_claude3_text_generator_predict_default_params_success(
     llm_text_df, model_name, session, session_us_east5, bq_connection
 ):
@@ -453,7 +454,7 @@ def test_claude3_text_generator_predict_default_params_success(
     "model_name",
     ("claude-3-sonnet", "claude-3-haiku", "claude-3-5-sonnet", "claude-3-opus"),
 )
-@pytest.mark.flaky(retries=2)
+@pytest.mark.flaky(retries=3, delay=120)
 def test_claude3_text_generator_predict_with_params_success(
     llm_text_df, model_name, session, session_us_east5, bq_connection
 ):

tests/system/small/test_encryption.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 
+import random
+
 from google.cloud import bigquery
 import pandas
 import pytest
@@ -153,14 +155,19 @@ def test_read_csv_gcs(
     if not bq_cmek:  # pragma: NO COVER
         pytest.skip("no cmek set for testing")  # pragma: NO COVER
 
-    # Create a csv in gcs
+    # Let's make the source data non-deterministic so that the test doesn't run
+    # into a BQ caching path
+    df = scalars_df_index.copy()
+    df["int_random"] = random.randint(0, 1_000_000_000)
+
+    # Export the dataframe to a csv in gcs
     write_path = gcs_folder + "test_read_csv_gcs_bigquery_engine*.csv"
     read_path = (
         utils.get_first_file_from_wildcard(write_path) if engine is None else write_path
     )
-    scalars_df_index.to_csv(write_path)
+    df.to_csv(write_path)
 
-    # Read the BQ table
+    # Read the gcs csv
     df = session_with_bq_cmek.read_csv(read_path, engine=engine)
 
     # Assert encryption
@@ -205,7 +212,9 @@ def test_read_pandas(bq_cmek, session_with_bq_cmek):
         pytest.skip("no cmek set for testing")  # pragma: NO COVER
 
     # Read a pandas dataframe
-    df = session_with_bq_cmek.read_pandas(pandas.DataFrame([1]))
+    df = session_with_bq_cmek.read_pandas(
+        pandas.DataFrame([random.randint(0, 1_000_000_000)])
+    )
 
     # Assert encryption
     _assert_bq_table_is_encrypted(df, bq_cmek, session_with_bq_cmek)