feat: Add .cache() method to persist intermediate dataframe (#626)

Author: TrevorBergeron

bigframes/dataframe.py

@@ -3397,6 +3397,17 @@ def _set_block(self, block: blocks.Block):
  def _get_block(self) -> blocks.Block:
  return self._block
 
+ def cache(self):
+ """
+ Materializes the DataFrame to a temporary table.
+
+ Useful if the dataframe will be used multiple times, as this will avoid recomputating the shared intermediate value.
+
+ Returns:
+ DataFrame: Self
+ """
+ return self._cached(force=True)
+
  def _cached(self, *, force: bool = False) -> DataFrame:
  """Materialize dataframe to a temporary table.
  No-op if the dataframe represents a trivial transformation of an existing materialization.

bigframes/ml/core.py

@@ -83,7 +83,7 @@ def distance(
  """
  assert len(x.columns) == 1 and len(y.columns) == 1
 
- input_data = x._cached().join(y._cached(), how="outer")
+ input_data = x.cache().join(y.cache(), how="outer")
  x_column_id, y_column_id = x._block.value_columns[0], y._block.value_columns[0]
 
  return self._apply_sql(
@@ -310,11 +310,9 @@ def create_model(
  # Cache dataframes to make sure base table is not a snapshot
  # cached dataframe creates a full copy, never uses snapshot
  if y_train is None:
- input_data = X_train._cached(force=True)
+ input_data = X_train.cache()
  else:
- input_data = X_train._cached(force=True).join(
- y_train._cached(force=True), how="outer"
- )
+ input_data = X_train.cache().join(y_train.cache(), how="outer")
  options.update({"INPUT_LABEL_COLS": y_train.columns.tolist()})
 
  session = X_train._session
@@ -354,9 +352,7 @@ def create_llm_remote_model(
  options = dict(options)
  # Cache dataframes to make sure base table is not a snapshot
  # cached dataframe creates a full copy, never uses snapshot
- input_data = X_train._cached(force=True).join(
- y_train._cached(force=True), how="outer"
- )
+ input_data = X_train.cache().join(y_train.cache(), how="outer")
  options.update({"INPUT_LABEL_COLS": y_train.columns.tolist()})
 
  session = X_train._session
@@ -389,9 +385,7 @@ def create_time_series_model(
  options = dict(options)
  # Cache dataframes to make sure base table is not a snapshot
  # cached dataframe creates a full copy, never uses snapshot
- input_data = X_train._cached(force=True).join(
- y_train._cached(force=True), how="outer"
- )
+ input_data = X_train.cache().join(y_train.cache(), how="outer")
  options.update({"TIME_SERIES_TIMESTAMP_COL": X_train.columns.tolist()[0]})
  options.update({"TIME_SERIES_DATA_COL": y_train.columns.tolist()[0]})
 

bigframes/series.py

@@ -1682,6 +1682,17 @@ def _slice(
  ),
  )
 
+ def cache(self):
+ """
+ Materializes the Series to a temporary table.
+
+ Useful if the series will be used multiple times, as this will avoid recomputating the shared intermediate value.
+
+ Returns:
+ Series: Self
+ """
+ return self._cached(force=True)
+
  def _cached(self, *, force: bool = True) -> Series:
  self._set_block(self._block.cached(force=force))
  return self

tests/system/small/test_dataframe.py

@@ -4204,7 +4204,7 @@ def test_df_cached(scalars_df_index):
  )
  df = df[df["rowindex_2"] % 2 == 0]
 
- df_cached_copy = df._cached()
+ df_cached_copy = df.cache()
  pandas.testing.assert_frame_equal(df.to_pandas(), df_cached_copy.to_pandas())
 
 

tests/unit/ml/test_golden_sql.py

@@ -63,7 +63,7 @@ def bqml_model_factory(mocker: pytest_mock.MockerFixture):
 def mock_y():
  mock_y = mock.create_autospec(spec=bpd.DataFrame)
  mock_y.columns = pd.Index(["input_column_label"])
- mock_y._cached.return_value = mock_y
+ mock_y.cache.return_value = mock_y
 
  return mock_y
 
@@ -83,7 +83,7 @@ def mock_X(mock_y, mock_session):
  ["index_column_id"],
  ["index_column_label"],
  )
- mock_X._cached.return_value = mock_X
+ mock_X.cache.return_value = mock_X
 
  return mock_X