SNOW-1805840: Augment telemetry with method_call_count

CHANGELOG.md

@@ -54,7 +54,8 @@
 - Updated integration testing for `session.lineage.trace` to exclude deleted objects
 - Added documentation for `DataFrame.map`.
 - Improve performance of `DataFrame.apply` by mapping numpy functions to snowpark functions if possible.
-- Added documentation on the extent of Snowpark pandas interoperability with scikit-learn
+- Added documentation on the extent of Snowpark pandas interoperability with scikit-learn.
+- Added `method_call_count` and `interchange_call_count` to telemetry.
 
 ## 1.26.0 (2024-12-05)
 

src/snowflake/snowpark/modin/plugin/_internal/telemetry.py

@@ -18,6 +18,7 @@
 from snowflake.snowpark.modin.plugin._internal.utils import (
     is_snowpark_pandas_dataframe_or_series_type,
 )
+from collections import Counter
 from snowflake.snowpark.query_history import QueryHistory
 from snowflake.snowpark.session import Session
 
@@ -36,6 +37,8 @@ class SnowparkPandasTelemetryField(Enum):
     ARGS = "argument"
     # fallback flag
     IS_FALLBACK = "is_fallback"
+    CALL_COUNT = "call_count"
+    INTERCHANGE_CALL_COUNT = "interchange_call_count"
 
 
 # Argument truncating size after converted to str. Size amount can be later specified after analysis and needs.
@@ -58,6 +61,8 @@ def _send_snowpark_pandas_telemetry_helper(
     func_name: str,
     query_history: Optional[QueryHistory],
     api_calls: Union[str, list[dict[str, Any]]],
+    method_call_count: Counter[str],
+    interchange_call_count: Counter[str],
 ) -> None:
     """
     A helper function that sends Snowpark pandas API telemetry data.
@@ -71,14 +76,30 @@ def _send_snowpark_pandas_telemetry_helper(
         query_history: The query history context manager to record queries that are pushed down to the Snowflake
         database in the session.
         api_calls: Optional list of Snowpark pandas API calls made during the function execution.
+        method_call_count: Number of times a method has been called.
+        interchange_call_count: Number of times __dataframe__ has been called.
 
     Returns:
         None
     """
-    data: dict[str, Union[str, list[dict[str, Any]], list[str], Optional[str]]] = {
+    data: dict[
+        str, Union[str, list[dict[str, Any]], list[str], Optional[str], Counter[str]]
+    ] = {
         TelemetryField.KEY_FUNC_NAME.value: func_name,
         TelemetryField.KEY_CATEGORY.value: SnowparkPandasTelemetryField.FUNC_CATEGORY_SNOWPARK_PANDAS.value,
         TelemetryField.KEY_ERROR_MSG.value: error_msg,
+        **(
+            {SnowparkPandasTelemetryField.CALL_COUNT.value: method_call_count}
+            if method_call_count is not None
+            else {}
+        ),
+        **(
+            {
+                SnowparkPandasTelemetryField.INTERCHANGE_CALL_COUNT.value: interchange_call_count
+            }
+            if interchange_call_count is not None
+            else {}
+        ),
     }
     if len(api_calls) > 0:
         data[TelemetryField.KEY_API_CALLS.value] = api_calls
@@ -274,6 +295,8 @@ def _telemetry_helper(
     # Moving existing api call out first can avoid to generate duplicates.
     existing_api_calls = []
     need_to_restore_args0_api_calls = False
+    method_call_count = None
+    interchange_call_count = None
 
     # If the decorated func is a class method or a standalone function, we need to get an active session:
     if is_standalone_function or (len(args) > 0 and isinstance(args[0], type)):
@@ -295,6 +318,13 @@ def _telemetry_helper(
             need_to_restore_args0_api_calls = True
             session = args[0]._query_compiler._modin_frame.ordered_dataframe.session
             class_prefix = args[0].__class__.__name__
+            args[0]._query_compiler._method_call_counts[func.__qualname__] += 1
+            method_call_count = args[0]._query_compiler._method_call_counts[
+                func.__qualname__
+            ]
+            interchange_call_count = args[0]._query_compiler._method_call_counts[
+                "__dataframe__"
+            ]
         except (TypeError, IndexError, AttributeError):
             # TypeError: args might not support indexing; IndexError: args is empty; AttributeError: args[0] might not
             # have _query_compiler attribute.
@@ -337,6 +367,8 @@ def _telemetry_helper(
             func_name=func_name,
             query_history=query_history,
             api_calls=existing_api_calls + [curr_api_call],
+            method_call_count=method_call_count,
+            interchange_call_count=interchange_call_count,
         )
         raise e
 
@@ -371,6 +403,8 @@ def _telemetry_helper(
             func_name=func_name,
             query_history=query_history,
             api_calls=existing_api_calls + [curr_api_call],
+            method_call_count=method_call_count,
+            interchange_call_count=interchange_call_count,
         )
         if need_to_restore_args0_api_calls:
             args[0]._query_compiler.snowpark_pandas_api_calls = existing_api_calls

src/snowflake/snowpark/modin/plugin/compiler/snowflake_query_compiler.py

@@ -10,6 +10,7 @@
 import json
 import logging
 import re
+from collections import Counter
 import typing
 import uuid
 from collections.abc import Hashable, Iterable, Mapping, Sequence
@@ -533,6 +534,7 @@ def __init__(self, frame: InternalFrame) -> None:
         # Copying and modifying self.snowpark_pandas_api_calls is taken care of in telemetry decorators
         self.snowpark_pandas_api_calls: list = []
         self._attrs: dict[Any, Any] = {}
+        self._method_call_counts: Counter[str] = Counter[str]()
 
     def _raise_not_implemented_error_for_timedelta(
         self, frame: InternalFrame = None

tests/integ/modin/test_telemetry.py

@@ -143,6 +143,8 @@ def test_snowpark_pandas_telemetry_method_decorator(test_table_name):
         "sfqids",
         "func_name",
         "error_msg",
+        "call_count",
+        "interchange_call_count",
     }
     assert data["category"] == "snowpark_pandas"
     assert data["api_calls"] == df1_expected_api_calls + [
@@ -178,6 +180,8 @@ def test_send_snowpark_pandas_telemetry_helper(send_mock):
         func_name="test_send_func",
         query_history=None,
         api_calls=[],
+        method_call_count=None,
+        interchange_call_count=None,
     )
     send_mock.assert_called_with(
         {
@@ -559,6 +563,87 @@ def test_telemetry_repr():
     ]
 
 
+@sql_count_checker(query_count=6, join_count=4)
+def test_telemetry_interchange_call_count():
+    s = pd.DataFrame([1, 2, 3, 4])
+    t = pd.DataFrame([5])
+    s.__dataframe__()
+    s.__dataframe__()
+    t.__dataframe__()
+
+    s.iloc[0, 0] = 7
+    s.__dataframe__()
+    s.__dataframe__()
+    t.__dataframe__()
+
+    def _get_data(call):
+        try:
+            return call.to_dict()["message"][TelemetryField.KEY_DATA.value]
+        except Exception:
+            return None
+
+    telemetry_data = [
+        _get_data(call)
+        for call in pd.session._conn._telemetry_client.telemetry._log_batch
+        if _get_data(call) is not None
+        and "func_name" in _get_data(call)
+        and _get_data(call)["func_name"] == "DataFrame.__dataframe__"
+    ]
+    assert len(telemetry_data) == 6
+    # s calls __dataframe__() for the first time.
+    assert telemetry_data[0]["call_count"] == 1
+    assert telemetry_data[0]["interchange_call_count"] == 1
+    # s calls __dataframe__() for the second time.
+    assert telemetry_data[1]["call_count"] == 2
+    assert telemetry_data[1]["interchange_call_count"] == 2
+    # t calls __dataframe__() for the first time.
+    assert telemetry_data[2]["call_count"] == 1
+    assert telemetry_data[2]["interchange_call_count"] == 1
+    # the new version of s calls __dataframe__() for the first time.
+    assert telemetry_data[3]["call_count"] == 1
+    assert telemetry_data[3]["interchange_call_count"] == 1
+    # the new version of s calls __dataframe__() for the second time.
+    assert telemetry_data[4]["call_count"] == 2
+    assert telemetry_data[4]["interchange_call_count"] == 2
+    # t calls __dataframe__() for the second time.
+    assert telemetry_data[5]["call_count"] == 2
+    assert telemetry_data[5]["interchange_call_count"] == 2
+
+
+@sql_count_checker(query_count=4)
+def test_telemetry_func_call_count():
+    s = pd.DataFrame([1, 2, np.nan, 4])
+    t = pd.DataFrame([5])
+
+    s.__repr__()
+    s.__repr__()
+    s.__repr__()
+
+    t.__repr__()
+
+    def _get_data(call):
+        try:
+            return call.to_dict()["message"][TelemetryField.KEY_DATA.value]
+        except Exception:
+            return None
+
+    telemetry_data = [
+        _get_data(call)
+        for call in pd.session._conn._telemetry_client.telemetry._log_batch
+        if _get_data(call) is not None
+        and "func_name" in _get_data(call)
+        and _get_data(call)["func_name"] == "DataFrame.__repr__"
+    ]
+
+    # second to last call from telemetry data
+    assert telemetry_data[-2]["call_count"] == 3
+    assert telemetry_data[-2]["interchange_call_count"] == 0
+
+    # last call from telemetry data
+    assert telemetry_data[-1]["call_count"] == 1
+    assert telemetry_data[-1]["interchange_call_count"] == 0
+
+
 @sql_count_checker(query_count=0)
 def test_telemetry_copy():
     # copy() is defined in upstream Modin's BasePandasDataset class, and not overridden by any

tests/unit/modin/test_telemetry.py

@@ -54,6 +54,8 @@ def snowpark_pandas_error_test_helper(
         query_history=ANY,
         telemetry_type=telemetry_type,
         error_msg=error_msg,
+        method_call_count=ANY,
+        interchange_call_count=ANY,
     )
 
 
@@ -115,6 +117,8 @@ def raise_real_type_error(_):
         query_history=ANY,
         telemetry_type="snowpark_pandas_type_error",
         error_msg=None,
+        method_call_count=ANY,
+        interchange_call_count=ANY,
     )
     assert len(mock_arg2._query_compiler.snowpark_pandas_api_calls) == 0
 
@@ -133,6 +137,8 @@ def raise_real_type_error(_):
         query_history=ANY,
         telemetry_type="snowpark_pandas_type_error",
         error_msg=None,
+        method_call_count=ANY,
+        interchange_call_count=ANY,
     )