refactor: Reduce executor cache interface to one method (#1133)

Author: TrevorBergeron

bigframes/core/blocks.py

@@ -2426,14 +2426,12 @@ def to_sql_query(
  def cached(self, *, force: bool = False, session_aware: bool = False) -> None:
  """Write the block to a session table."""
  # use a heuristic for whether something needs to be cached
- if (not force) and self.session._executor._is_trivially_executable(self.expr):
- return
- elif session_aware:
- self.session._executor._cache_with_session_awareness(self.expr)
- else:
- self.session._executor._cache_with_cluster_cols(
- self.expr, cluster_cols=self.index_columns
- )
+ self.session._executor.cached(
+ self.expr,
+ force=force,
+ use_session=session_aware,
+ cluster_cols=self.index_columns,
+ )
 
  def _is_monotonic(
  self, column_ids: typing.Union[str, Sequence[str]], increasing: bool

bigframes/session/executor.py

@@ -94,6 +94,7 @@ def __init__(
  bqclient: bigquery.Client,
  storage_manager: bigframes.session.temp_storage.TemporaryGbqStorageManager,
  bqstoragereadclient: google.cloud.bigquery_storage_v1.BigQueryReadClient,
+ *,
  strictly_ordered: bool = True,
  metrics: Optional[bigframes.session.metrics.ExecutionMetrics] = None,
  ):
@@ -354,6 +355,23 @@ def get_row_count(self, array_value: bigframes.core.ArrayValue) -> int:
  iter, _ = self._run_execute_query(sql)
  return next(iter)[0]
 
+ def cached(
+ self,
+ array_value: bigframes.core.ArrayValue,
+ *,
+ force: bool = False,
+ use_session: bool = False,
+ cluster_cols: Sequence[str] = (),
+ ) -> None:
+ """Write the block to a session table."""
+ # use a heuristic for whether something needs to be cached
+ if (not force) and self._is_trivially_executable(array_value):
+ return
+ elif use_session:
+ self._cache_with_session_awareness(array_value)
+ else:
+ self._cache_with_cluster_cols(array_value, cluster_cols=cluster_cols)
+
  def _local_get_row_count(
  self, array_value: bigframes.core.ArrayValue
  ) -> Optional[int]: