feat: log most recent API calls as recent-bigframes-api-xx labels on BigQuery jobs

bigframes/core/groupby/__init__.py

@@ -30,9 +30,11 @@
 import bigframes.operations as ops
 import bigframes.operations.aggregations as agg_ops
 import bigframes.series as series
+from bigframes.utils import log_adapter
 import third_party.bigframes_vendored.pandas.core.groupby as vendored_pandas_groupby
 
 
+@log_adapter.class_logger
 class DataFrameGroupBy(vendored_pandas_groupby.DataFrameGroupBy):
     __doc__ = vendored_pandas_groupby.GroupBy.__doc__
 
@@ -408,6 +410,7 @@ def _resolve_label(self, label: blocks.Label) -> str:
         return col_ids[0]
 
 
+@log_adapter.class_logger
 class SeriesGroupBy(vendored_pandas_groupby.SeriesGroupBy):
     __doc__ = vendored_pandas_groupby.GroupBy.__doc__
 

bigframes/core/window/__init__.py

@@ -19,9 +19,11 @@
 import bigframes.core as core
 import bigframes.core.blocks as blocks
 import bigframes.operations.aggregations as agg_ops
+from bigframes.utils import log_adapter
 import third_party.bigframes_vendored.pandas.core.window.rolling as vendored_pandas_rolling
 
 
+@log_adapter.class_logger
 class Window(vendored_pandas_rolling.Window):
     __doc__ = vendored_pandas_rolling.Window.__doc__
 

bigframes/dataframe.py

@@ -56,6 +56,7 @@
 import bigframes.series
 import bigframes.series as bf_series
 import bigframes.session._io.bigquery
+from bigframes.utils import log_adapter
 import third_party.bigframes_vendored.pandas.core.frame as vendored_pandas_frame
 import third_party.bigframes_vendored.pandas.pandas._typing as vendored_pandas_typing
 
@@ -80,6 +81,7 @@
 
 
 # Inherits from pandas DataFrame so that we can use the same docstrings.
+@log_adapter.class_logger
 class DataFrame(vendored_pandas_frame.DataFrame):
     __doc__ = vendored_pandas_frame.DataFrame.__doc__
 

bigframes/operations/datetimes.py

@@ -17,9 +17,11 @@
 import bigframes.operations as ops
 import bigframes.operations.base
 import bigframes.series as series
+from bigframes.utils import log_adapter
 import third_party.bigframes_vendored.pandas.core.indexes.accessor as vendordt
 
 
+@log_adapter.class_logger
 class DatetimeMethods(
     bigframes.operations.base.SeriesMethods, vendordt.DatetimeProperties
 ):

bigframes/operations/strings.py

@@ -22,6 +22,7 @@
 import bigframes.operations as ops
 import bigframes.operations.base
 import bigframes.series as series
+from bigframes.utils import log_adapter
 import third_party.bigframes_vendored.pandas.core.strings.accessor as vendorstr
 
 # Maps from python to re2
@@ -32,6 +33,7 @@
 }
 
 
+@log_adapter.class_logger
 class StringMethods(bigframes.operations.base.SeriesMethods, vendorstr.StringMethods):
     __doc__ = vendorstr.StringMethods.__doc__
 

bigframes/operations/structs.py

@@ -22,6 +22,7 @@
 import bigframes.operations
 import bigframes.operations.base
 import bigframes.series
+from bigframes.utils import log_adapter
 import third_party.bigframes_vendored.pandas.core.arrays.arrow.accessors as vendoracessors
 
 
@@ -38,6 +39,7 @@ def _as_ibis(self, x: ibis_types.Value):
         return struct_value[name].name(name)
 
 
+@log_adapter.class_logger
 class StructAccessor(
     bigframes.operations.base.SeriesMethods, vendoracessors.StructAccessor
 ):

bigframes/series.py

@@ -52,12 +52,14 @@
 import bigframes.operations.datetimes as dt
 import bigframes.operations.strings as strings
 import bigframes.operations.structs as structs
+from bigframes.utils import log_adapter
 import third_party.bigframes_vendored.pandas.core.series as vendored_pandas_series
 
 LevelType = typing.Union[str, int]
 LevelsType = typing.Union[LevelType, typing.Sequence[LevelType]]
 
 
+@log_adapter.class_logger
 class Series(bigframes.operations.base.SeriesMethods, vendored_pandas_series.Series):
     def __init__(self, *args, **kwargs):
         self._query_job: Optional[bigquery.QueryJob] = None

bigframes/session/__init__.py

@@ -75,6 +75,7 @@
 from bigframes.remote_function import remote_function as bigframes_rf
 import bigframes.session._io.bigquery as bigframes_io
 import bigframes.session.clients
+from bigframes.utils import log_adapter
 import bigframes.version
 
 # Even though the ibis.backends.bigquery.registry import is unused, it's needed
@@ -110,6 +111,7 @@ def _is_query(query_or_table: str) -> bool:
     return re.search(r"\s", query_or_table.strip(), re.MULTILINE) is not None
 
 
+@log_adapter.class_logger
 class Session(
     third_party_pandas_gbq.GBQIOMixin,
     third_party_pandas_parquet.ParquetIOMixin,
@@ -1496,10 +1498,19 @@ def _start_query(
         """
         Starts query job and waits for results
         """
+        api_methods = log_adapter._api_methods
         if job_config is not None:
-            query_job = self.bqclient.query(sql, job_config=job_config)
+            job_config.labels = bigframes_io.create_job_configs_labels(
+                job_configs_labels=job_config.labels, api_methods=api_methods
+            )
         else:
-            query_job = self.bqclient.query(sql)
+            job_config = bigquery.QueryJobConfig()
+            job_config.labels = bigframes_io.create_job_configs_labels(
+                job_configs_labels=None, api_methods=api_methods
+            )
+        query_job = self.bqclient.query(sql, job_config=job_config)
+        # Clear out the global api logger
+        log_adapter._api_methods = []
 
         opts = bigframes.options.display
         if opts.progress_bar is not None and not query_job.configuration.dry_run:

bigframes/session/_io/bigquery.py

@@ -17,11 +17,41 @@
 import datetime
 import textwrap
 import types
-from typing import Dict, Iterable, Union
+from typing import Dict, Iterable, Optional, Sequence, Union
 
 import google.cloud.bigquery as bigquery
 
 IO_ORDERING_ID = "bqdf_row_nums"
+MAX_LABELS_COUNT = 64
+
+
+def create_job_configs_labels(
+    job_configs_labels: Optional[Dict[str, str]],
+    api_methods: Sequence[str],
+) -> Dict[str, str]:
+    # If there is no label set
+    if job_configs_labels is None:
+        labels = {}
+        label_values = list(api_methods)
+    else:
+        labels = job_configs_labels.copy()
+        cur_labels_len = len(job_configs_labels)
+        api_methods_len = len(api_methods)
+        # If the total number of labels is under the limit of labels count
+        if cur_labels_len + api_methods_len <= MAX_LABELS_COUNT:
+            label_values = list(api_methods)
+        # We capture the latest label if it is out of the length limit of labels count
+        else:
+            added_api_len = cur_labels_len + api_methods_len - MAX_LABELS_COUNT
+            label_values = list(api_methods)[-added_api_len:]
+
+    for i, label_value in enumerate(label_values):
+        if job_configs_labels is not None:
+            label_key = "bigframes-api-" + str(i + len(job_configs_labels))
+        else:
+            label_key = "bigframes-api-" + str(i)
+        labels[label_key] = label_value
+    return labels
 
 
 def create_export_csv_statement(

bigframes/utils/__init__.py

@@ -0,0 +1,17 @@
+# 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.
+
+from bigframes.utils.log_adapter import class_logger, method_logger
+
+__all__ = ["class_logger", "method_logger"]

bigframes/utils/log_adapter.py

@@ -0,0 +1,32 @@
+import functools
+
+from loguru import logger
+
+_log_file_path = None
+_logger = logger
+_api_methods = []
+
+
+def class_logger(decorated_cls):
+    """Decorator that adds logging functionality to each method of the class."""
+    for attr_name, attr_value in decorated_cls.__dict__.items():
+        if callable(attr_value):
+            setattr(decorated_cls, attr_name, method_logger(attr_value))
+    return decorated_cls
+
+
+def method_logger(method):
+    """Decorator that adds logging functionality to a method."""
+
+    @functools.wraps(method)
+    def wrapper(*args, **kwargs):
+        api_method_name = str(method.__name__)
+        if not api_method_name.startswith("__"):
+            add_api_method(api_method_name)
+
+    return wrapper
+
+
+def add_api_method(method: str):
+    global _api_methods
+    _api_methods.append(method)

setup.py

@@ -52,6 +52,7 @@
     "sqlalchemy >=1.4,<3.0dev",
     "ipywidgets >=7.7.1",
     "humanize >= 4.6.0",
+    "logruru >=0.6.0",
 ]
 extras = {
     # Optional test dependencies packages. If they're missed, may skip some tests.

tests/unit/session/test_io_bigquery.py

@@ -18,15 +18,122 @@
 import google.cloud.bigquery as bigquery
 import pytest
 
-import bigframes.session._io.bigquery
+import bigframes.pandas as bpd
+import bigframes.session._io.bigquery as io_bq
+from bigframes.utils import log_adapter
+
+
+def test_create_job_configs_labels_is_none():
+    api_methods = ["df-agg", "series-mode"]
+    labels = io_bq.create_job_configs_labels(
+        job_configs_labels=None, api_methods=api_methods
+    )
+    expected_dict = {"bigframes-api-0": "df-agg", "bigframes-api-1": "series-mode"}
+    assert labels is not None
+    assert labels == expected_dict
+
+
+def test_create_job_configs_labels_length_limit_not_met():
+    cur_labels = {
+        "bigframes-api": "read_pandas",
+        "source": "bigquery-dataframes-temp",
+    }
+    api_methods = ["df-agg", "series-mode"]
+    labels = io_bq.create_job_configs_labels(
+        job_configs_labels=cur_labels, api_methods=api_methods
+    )
+    expected_dict = {
+        "bigframes-api": "read_pandas",
+        "source": "bigquery-dataframes-temp",
+        "bigframes-api-2": "df-agg",
+        "bigframes-api-3": "series-mode",
+    }
+    assert labels is not None
+    assert len(labels) == 4
+    assert labels == expected_dict
+
+
+def test_create_job_configs_labels_log_adaptor_under_length_limit():
+    log_adapter._api_methods = ["df-agg", "series-mode"]
+    cur_labels = {
+        "bigframes-api": "read_pandas",
+        "source": "bigquery-dataframes-temp",
+    }
+    api_methods = log_adapter._api_methods
+    labels = io_bq.create_job_configs_labels(
+        job_configs_labels=cur_labels, api_methods=api_methods
+    )
+    expected_dict = {
+        "bigframes-api": "read_pandas",
+        "source": "bigquery-dataframes-temp",
+        "bigframes-api-2": "df-agg",
+        "bigframes-api-3": "series-mode",
+    }
+    assert labels is not None
+    assert len(labels) == 4
+    assert labels == expected_dict
+
+
+def test_create_job_configs_labels_log_adaptor_call_method_under_length_limit():
+    cur_labels = {
+        "bigframes-api": "read_pandas",
+        "source": "bigquery-dataframes-temp",
+    }
+    log_adapter._api_methods = []
+    df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+    # Test running two methods
+    df.head()
+    df.max()
+    api_methods = log_adapter._api_methods
+
+    labels = io_bq.create_job_configs_labels(
+        job_configs_labels=cur_labels, api_methods=api_methods
+    )
+    expected_dict = {
+        "bigframes-api": "read_pandas",
+        "source": "bigquery-dataframes-temp",
+        "bigframes-api-2": "head",
+        "bigframes-api-3": "max",
+    }
+    assert labels is not None
+    assert len(labels) == 4
+    assert labels == expected_dict
+
+
+def test_create_job_configs_labels_length_limit_met():
+    cur_labels = {
+        "bigframes-api": "read_pandas",
+        "source": "bigquery-dataframes-temp",
+    }
+    for i in range(61):
+        key = f"bigframes-api-{i}"
+        value = f"test{i}"
+        cur_labels[key] = value
+    # If cur_labels length is 63, we can only add one label from api_methods
+    log_adapter._api_methods = []
+    df = bpd.DataFrame({"col1": [1, 2], "col2": [3, 4]})
+    # Test running two methods
+    df.head()
+    df.max()
+    api_methods = log_adapter._api_methods
+
+    labels = io_bq.create_job_configs_labels(
+        job_configs_labels=cur_labels, api_methods=api_methods
+    )
+    assert labels is not None
+    assert len(labels) == 64
+    assert "head" not in labels.values()
+    assert "max" in labels.values()
+    assert "bigframes-api" in labels.keys()
+    assert "source" in labels.keys()
 
 
 def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets():
     table_ref = bigquery.TableReference.from_string(
         "my-test-project._e8166e0cdb.anonbb92cd"
     )
 
-    sql = bigframes.session._io.bigquery.create_snapshot_sql(
+    sql = io_bq.create_snapshot_sql(
         table_ref, datetime.datetime.now(datetime.timezone.utc)
     )
 
@@ -40,7 +147,7 @@ def test_create_snapshot_sql_doesnt_timetravel_anonymous_datasets():
 def test_create_snapshot_sql_doesnt_timetravel_session_datasets():
     table_ref = bigquery.TableReference.from_string("my-test-project._session.abcdefg")
 
-    sql = bigframes.session._io.bigquery.create_snapshot_sql(
+    sql = io_bq.create_snapshot_sql(
         table_ref, datetime.datetime.now(datetime.timezone.utc)
     )
 
@@ -101,5 +208,5 @@ def test_create_snapshot_sql_doesnt_timetravel_session_datasets():
     ),
 )
 def test_bq_schema_to_sql(schema: Iterable[bigquery.SchemaField], expected: str):
-    sql = bigframes.session._io.bigquery.bq_schema_to_sql(schema)
+    sql = io_bq.bq_schema_to_sql(schema)
     assert sql == expected

third_party/bigframes_vendored/sklearn/cluster/_kmeans.py

@@ -12,7 +12,6 @@
 # License: BSD 3 clause
 
 from abc import ABC
-from typing import List, Optional
 
 from bigframes import constants
 from third_party.bigframes_vendored.sklearn.base import BaseEstimator