fix: cache DataFrames to temp tables in bigframes.bigquery.ml methods to avoid time travel (#2318)

tswast · web-flow · commit d99383195ac3 · 2025-12-09T12:50:41.000-06:00
See internal issue b/310266666 &#129429;
diff --git a/bigframes/bigquery/_operations/ml.py b/bigframes/bigquery/_operations/ml.py
@@ -39,6 +39,10 @@ def _to_sql(df_or_sql: Union[pd.DataFrame, dataframe.DataFrame, str]) -> str:
     else:
         bf_df = cast(dataframe.DataFrame, df_or_sql)
 
+    # Cache dataframes to make sure base table is not a snapshot.
+    # Cached dataframe creates a full copy, never uses snapshot.
+    # This is a workaround for internal issue b/310266666.
+    bf_df.cache()
     sql, _, _ = bf_df._to_sql_query(include_index=False)
     return sql
 
diff --git a/bigframes/ml/core.py b/bigframes/ml/core.py
@@ -436,8 +436,9 @@ def create_model(
         Returns: a BqmlModel, wrapping a trained model in BigQuery
         """
         options = dict(options)
-        # Cache dataframes to make sure base table is not a snapshot
-        # cached dataframe creates a full copy, never uses snapshot
+        # Cache dataframes to make sure base table is not a snapshot.
+        # Cached dataframe creates a full copy, never uses snapshot.
+        # This is a workaround for internal issue b/310266666.
         if y_train is None:
             input_data = X_train.reset_index(drop=True).cache()
         else:
