perf: Speedup internal tree comparisons (#1060)

Author: TrevorBergeron

bigframes/core/nodes.py

@@ -51,8 +51,8 @@ class Field:
  dtype: bigframes.dtypes.Dtype
 
 
-@dataclass(frozen=True)
-class BigFrameNode:
+@dataclass(eq=False, frozen=True)
+class BigFrameNode(abc.ABC):
  """
  Immutable node for representing 2D typed array as a tree of operators.
 
@@ -95,12 +95,30 @@ def session(self):
  return sessions[0]
  return None
 
+ def _as_tuple(self) -> Tuple:
+ """Get all fields as tuple."""
+ return tuple(getattr(self, field.name) for field in fields(self))
+
+ def __hash__(self) -> int:
+ # Custom hash that uses cache to avoid costly recomputation
+ return self._cached_hash
+
+ def __eq__(self, other) -> bool:
+ # Custom eq that tries to short-circuit full structural comparison
+ if not isinstance(other, self.__class__):
+ return False
+ if self is other:
+ return True
+ if hash(self) != hash(other):
+ return False
+ return self._as_tuple() == other._as_tuple()
+
  # BigFrameNode trees can be very deep so its important avoid recalculating the hash from scratch
  # Each subclass of BigFrameNode should use this property to implement __hash__
  # The default dataclass-generated __hash__ method is not cached
  @functools.cached_property
- def _node_hash(self):
- return hash(tuple(hash(getattr(self, field.name)) for field in fields(self)))
+ def _cached_hash(self):
+ return hash(self._as_tuple())
 
  @property
  def roots(self) -> typing.Set[BigFrameNode]:
@@ -226,7 +244,7 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
  return self.transform_children(lambda x: x.prune(used_cols))
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class UnaryNode(BigFrameNode):
  child: BigFrameNode
 
@@ -252,7 +270,7 @@ def order_ambiguous(self) -> bool:
  return self.child.order_ambiguous
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class JoinNode(BigFrameNode):
  left_child: BigFrameNode
  right_child: BigFrameNode
@@ -285,9 +303,6 @@ def explicitly_ordered(self) -> bool:
  # Do not consider user pre-join ordering intent - they need to re-order post-join in unordered mode.
  return False
 
- def __hash__(self):
- return self._node_hash
-
  @functools.cached_property
  def fields(self) -> Tuple[Field, ...]:
  return tuple(itertools.chain(self.left_child.fields, self.right_child.fields))
@@ -320,7 +335,7 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
  return self.transform_children(lambda x: x.prune(new_used))
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class ConcatNode(BigFrameNode):
  # TODO: Explcitly map column ids from each child
  children: Tuple[BigFrameNode, ...]
@@ -345,9 +360,6 @@ def explicitly_ordered(self) -> bool:
  # Consider concat as an ordered operations (even though input frames may not be ordered)
  return True
 
- def __hash__(self):
- return self._node_hash
-
  @functools.cached_property
  def fields(self) -> Tuple[Field, ...]:
  # TODO: Output names should probably be aligned beforehand or be part of concat definition
@@ -371,16 +383,13 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
  return self
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class FromRangeNode(BigFrameNode):
  # TODO: Enforce single-row, single column constraint
  start: BigFrameNode
  end: BigFrameNode
  step: int
 
- def __hash__(self):
- return self._node_hash
-
  @property
  def roots(self) -> typing.Set[BigFrameNode]:
  return {self}
@@ -419,7 +428,7 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
 # Input Nodex
 # TODO: Most leaf nodes produce fixed column names based on the datasource
 # They should support renaming
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class LeafNode(BigFrameNode):
  @property
  def roots(self) -> typing.Set[BigFrameNode]:
@@ -451,7 +460,7 @@ class ScanList:
  items: typing.Tuple[ScanItem, ...]
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class ReadLocalNode(LeafNode):
  feather_bytes: bytes
  data_schema: schemata.ArraySchema
@@ -460,9 +469,6 @@ class ReadLocalNode(LeafNode):
  scan_list: ScanList
  session: typing.Optional[bigframes.session.Session] = None
 
- def __hash__(self):
- return self._node_hash
-
  @functools.cached_property
  def fields(self) -> Tuple[Field, ...]:
  return tuple(Field(col_id, dtype) for col_id, dtype, _ in self.scan_list.items)
@@ -547,7 +553,7 @@ class BigqueryDataSource:
 
 
 ## Put ordering in here or just add order_by node above?
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class ReadTableNode(LeafNode):
  source: BigqueryDataSource
  # Subset of physical schema column
@@ -570,9 +576,6 @@ def __post_init__(self):
  def session(self):
  return self.table_session
 
- def __hash__(self):
- return self._node_hash
-
  @functools.cached_property
  def fields(self) -> Tuple[Field, ...]:
  return tuple(Field(col_id, dtype) for col_id, dtype, _ in self.scan_list.items)
@@ -616,15 +619,12 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
  return ReadTableNode(self.source, new_scan_list, self.table_session)
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class CachedTableNode(ReadTableNode):
  # The original BFET subtree that was cached
  # note: this isn't a "child" node.
  original_node: BigFrameNode = field()
 
- def __hash__(self):
- return self._node_hash
-
  def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
  new_scan_list = ScanList(
  tuple(item for item in self.scan_list.items if item.id in used_cols)
@@ -635,13 +635,10 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
 
 
 # Unary nodes
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class PromoteOffsetsNode(UnaryNode):
  col_id: bigframes.core.identifiers.ColumnId
 
- def __hash__(self):
- return self._node_hash
-
  @property
  def non_local(self) -> bool:
  return True
@@ -666,17 +663,14 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
  return self.transform_children(lambda x: x.prune(new_used))
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class FilterNode(UnaryNode):
  predicate: ex.Expression
 
  @property
  def row_preserving(self) -> bool:
  return False
 
- def __hash__(self):
- return self._node_hash
-
  @property
  def variables_introduced(self) -> int:
  return 1
@@ -687,13 +681,10 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
  return FilterNode(pruned_child, self.predicate)
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class OrderByNode(UnaryNode):
  by: Tuple[OrderingExpression, ...]
 
- def __hash__(self):
- return self._node_hash
-
  @property
  def variables_introduced(self) -> int:
  return 0
@@ -716,14 +707,11 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
  return OrderByNode(pruned_child, self.by)
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class ReversedNode(UnaryNode):
  # useless field to make sure has distinct hash
  reversed: bool = True
 
- def __hash__(self):
- return self._node_hash
-
  @property
  def variables_introduced(self) -> int:
  return 0
@@ -734,15 +722,12 @@ def relation_ops_created(self) -> int:
  return 0
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class SelectionNode(UnaryNode):
  input_output_pairs: typing.Tuple[
  typing.Tuple[ex.DerefOp, bigframes.core.identifiers.ColumnId], ...
  ]
 
- def __hash__(self):
- return self._node_hash
-
  @functools.cached_property
  def fields(self) -> Tuple[Field, ...]:
  return tuple(
@@ -772,7 +757,7 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
  return SelectionNode(pruned_child, pruned_selections)
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class ProjectionNode(UnaryNode):
  """Assigns new variables (without modifying existing ones)"""
 
@@ -788,9 +773,6 @@ def __post_init__(self):
  # Cannot assign to existing variables - append only!
  assert all(name not in self.child.schema.names for _, name in self.assignments)
 
- def __hash__(self):
- return self._node_hash
-
  @functools.cached_property
  def fields(self) -> Tuple[Field, ...]:
  input_types = self.child._dtype_lookup
@@ -819,7 +801,7 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
 
 # TODO: Merge RowCount into Aggregate Node?
 # Row count can be compute from table metadata sometimes, so it is a bit special.
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class RowCountNode(UnaryNode):
  @property
  def row_preserving(self) -> bool:
@@ -842,7 +824,7 @@ def defines_namespace(self) -> bool:
  return True
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class AggregateNode(UnaryNode):
  aggregations: typing.Tuple[
  typing.Tuple[ex.Aggregation, bigframes.core.identifiers.ColumnId], ...
@@ -854,9 +836,6 @@ class AggregateNode(UnaryNode):
  def row_preserving(self) -> bool:
  return False
 
- def __hash__(self):
- return self._node_hash
-
  @property
  def non_local(self) -> bool:
  return True
@@ -904,7 +883,7 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
  return AggregateNode(pruned_child, pruned_aggs, self.by_column_ids, self.dropna)
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class WindowOpNode(UnaryNode):
  column_name: ex.DerefOp
  op: agg_ops.UnaryWindowOp
@@ -913,9 +892,6 @@ class WindowOpNode(UnaryNode):
  never_skip_nulls: bool = False
  skip_reproject_unsafe: bool = False
 
- def __hash__(self):
- return self._node_hash
-
  @property
  def non_local(self) -> bool:
  return True
@@ -945,11 +921,8 @@ def prune(self, used_cols: COLUMN_SET) -> BigFrameNode:
 
 
 # TODO: Remove this op
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class ReprojectOpNode(UnaryNode):
- def __hash__(self):
- return self._node_hash
-
  @property
  def variables_introduced(self) -> int:
  return 0
@@ -960,7 +933,7 @@ def relation_ops_created(self) -> int:
  return 0
 
 
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class RandomSampleNode(UnaryNode):
  fraction: float
 
@@ -972,26 +945,20 @@ def deterministic(self) -> bool:
  def row_preserving(self) -> bool:
  return False
 
- def __hash__(self):
- return self._node_hash
-
  @property
  def variables_introduced(self) -> int:
  return 1
 
 
 # TODO: Explode should create a new column instead of overriding the existing one
-@dataclass(frozen=True)
+@dataclass(frozen=True, eq=False)
 class ExplodeNode(UnaryNode):
  column_ids: typing.Tuple[ex.DerefOp, ...]
 
  @property
  def row_preserving(self) -> bool:
  return False
 
- def __hash__(self):
- return self._node_hash
-
  @functools.cached_property
  def fields(self) -> Tuple[Field, ...]:
  return tuple(