refactor: Limit access to ArrayValue node field. (#977)

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -36,6 +36,7 @@
 import bigframes.core.ordering as orderings
 import bigframes.core.rewrite
 import bigframes.core.schema as schemata
+import bigframes.core.tree_properties
 import bigframes.core.utils
 from bigframes.core.window_spec import WindowSpec
 import bigframes.dtypes
@@ -124,6 +125,20 @@ def schema(self) -> schemata.ArraySchema:
     def _compiled_schema(self) -> schemata.ArraySchema:
         return bigframes.core.compile.test_only_ibis_inferred_schema(self.node)
 
+    @property
+    def explicitly_ordered(self) -> bool:
+        # see BigFrameNode.explicitly_ordered
+        return self.node.explicitly_ordered
+
+    @property
+    def order_ambiguous(self) -> bool:
+        # see BigFrameNode.order_ambiguous
+        return self.node.order_ambiguous
+
+    @property
+    def supports_fast_peek(self) -> bool:
+        return bigframes.core.tree_properties.can_fast_peek(self.node)
+
     def as_cached(
         self: ArrayValue,
         cache_table: google.cloud.bigquery.Table,

bigframes/core/blocks.py

@@ -49,7 +49,6 @@
 import bigframes.core.ordering as ordering
 import bigframes.core.schema as bf_schema
 import bigframes.core.sql as sql
-import bigframes.core.tree_properties as tree_properties
 import bigframes.core.utils as utils
 import bigframes.core.window_spec as window_specs
 import bigframes.dtypes
@@ -205,7 +204,7 @@ def shape(self) -> typing.Tuple[int, int]:
         row_count_expr = self.expr.row_count()
 
         # Support in-memory engines for hermetic unit tests.
-        if self.expr.node.session is None:
+        if self.expr.session is None:
             try:
                 row_count = row_count_expr._try_evaluate_local().squeeze()
                 return (row_count, len(self.value_columns))
@@ -283,7 +282,7 @@ def index_name_to_col_id(self) -> typing.Mapping[Label, typing.Sequence[str]]:
 
     @property
     def explicitly_ordered(self) -> bool:
-        return self.expr.node.explicitly_ordered
+        return self.expr.explicitly_ordered
 
     def cols_matching_label(self, partial_label: Label) -> typing.Sequence[str]:
         """
@@ -466,7 +465,7 @@ def _validate_result_schema(
     ):
         actual_schema = tuple(bq_result_schema)
         ibis_schema = self.expr._compiled_schema
-        internal_schema = self.expr.node.schema
+        internal_schema = self.expr.schema
         if not bigframes.features.PANDAS_VERSIONS.is_arrow_list_dtype_usable:
             return
         if internal_schema.to_bigquery() != actual_schema:
@@ -561,7 +560,7 @@ def to_pandas(
     def try_peek(
         self, n: int = 20, force: bool = False
     ) -> typing.Optional[pd.DataFrame]:
-        if force or tree_properties.can_fast_peek(self.expr.node):
+        if force or self.expr.supports_fast_peek:
             iterator, _ = self.session._peek(self.expr, n)
             df = self._to_dataframe(iterator)
             self._copy_index_to_pandas(df)
@@ -2365,10 +2364,7 @@ def cached(self, *, force: bool = False, session_aware: bool = False) -> None:
         if (not force) and self.session._executor._is_trivially_executable(self.expr):
             return
         elif session_aware:
-            bfet_roots = [obj._block._expr.node for obj in self.session.objects]
-            self.session._executor._cache_with_session_awareness(
-                self.expr, session_forest=bfet_roots
-            )
+            self.session._executor._cache_with_session_awareness(self.expr)
         else:
             self.session._executor._cache_with_cluster_cols(
                 self.expr, cluster_cols=self.index_columns

bigframes/core/validations.py

@@ -66,7 +66,7 @@ def enforce_ordered(
     object: HasSession, opname: str, suggestion: Optional[str] = None
 ) -> None:
     session = object._session
-    if session._strictly_ordered or not object._block.expr.node.order_ambiguous:
+    if session._strictly_ordered or not object._block.expr.order_ambiguous:
         # No ambiguity for how to calculate ordering, so no error or warning
         return None
     if not session._allows_ambiguity:

bigframes/session/executor.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import math
-from typing import cast, Iterable, Literal, Mapping, Optional, Sequence, Tuple, Union
+from typing import cast, Literal, Mapping, Optional, Sequence, Tuple, Union
 import warnings
 import weakref
 
@@ -381,8 +381,8 @@ def _cache_with_offsets(self, array_value: bigframes.core.ArrayValue):
     def _cache_with_session_awareness(
         self,
         array_value: bigframes.core.ArrayValue,
-        session_forest: Iterable[nodes.BigFrameNode],
     ) -> None:
+        session_forest = [obj._block._expr.node for obj in array_value.session.objects]
         # These node types are cheap to re-compute
         target, cluster_cols = bigframes.session.planner.session_aware_cache_plan(
             array_value.node, list(session_forest)