refactor: Add function to make all column ids in a tree unique and sequential (#1094)

Author: TrevorBergeron

bigframes/core/__init__.py

@@ -268,7 +268,13 @@ def promote_offsets(self) -> Tuple[ArrayValue, str]:
  def concat(self, other: typing.Sequence[ArrayValue]) -> ArrayValue:
  """Append together multiple ArrayValue objects."""
  return ArrayValue(
- nodes.ConcatNode(children=tuple([self.node, *[val.node for val in other]]))
+ nodes.ConcatNode(
+ children=tuple([self.node, *[val.node for val in other]]),
+ output_ids=tuple(
+ ids.ColumnId(bigframes.core.guid.generate_guid())
+ for id in self.column_ids
+ ),
+ )
  )
 
  def compute_values(self, assignments: Sequence[ex.Expression]):

bigframes/core/blocks.py

@@ -3139,7 +3139,7 @@ def _pd_index_to_array_value(
  rows = []
  labels_as_tuples = utils.index_as_tuples(index)
  for row_offset in range(len(index)):
- id_gen = bigframes.core.identifiers.standard_identifiers()
+ id_gen = bigframes.core.identifiers.standard_id_strings()
  row_label = labels_as_tuples[row_offset]
  row_label = (row_label,) if not isinstance(row_label, tuple) else row_label
  row = {}

bigframes/core/compile/api.py

@@ -18,14 +18,15 @@
 import google.cloud.bigquery as bigquery
 
 import bigframes.core.compile.compiler as compiler
-import bigframes.core.rewrite as rewrites
 
 if TYPE_CHECKING:
  import bigframes.core.nodes
  import bigframes.core.ordering
  import bigframes.core.schema
 
-_STRICT_COMPILER = compiler.Compiler(strict=True)
+_STRICT_COMPILER = compiler.Compiler(
+ strict=True, enable_pruning=True, enable_densify_ids=True
+)
 
 
 class SQLCompiler:
@@ -34,7 +35,7 @@ def __init__(self, strict: bool = True):
 
  def compile_peek(self, node: bigframes.core.nodes.BigFrameNode, n_rows: int) -> str:
  """Compile node into sql that selects N arbitrary rows, may not execute deterministically."""
- return self._compiler.compile_unordered_ir(node).peek_sql(n_rows)
+ return self._compiler.compile_peek_sql(node, n_rows)
 
  def compile_unordered(
  self,
@@ -44,9 +45,8 @@ def compile_unordered(
  ) -> str:
  """Compile node into sql where rows are unsorted, and no ordering information is preserved."""
  # TODO: Enable limit pullup, but only if not being used to write to clustered table.
- return self._compiler.compile_unordered_ir(node).to_sql(
- col_id_overrides=col_id_overrides
- )
+ output_ids = [col_id_overrides.get(id, id) for id in node.schema.names]
+ return self._compiler.compile_sql(node, ordered=False, output_ids=output_ids)
 
  def compile_ordered(
  self,
@@ -56,10 +56,8 @@ def compile_ordered(
  ) -> str:
  """Compile node into sql where rows are sorted with ORDER BY."""
  # If we are ordering the query anyways, compiling the slice as a limit is probably a good idea.
- new_node, limit = rewrites.pullup_limit_from_slice(node)
- return self._compiler.compile_ordered_ir(new_node).to_sql(
- col_id_overrides=col_id_overrides, ordered=True, limit=limit
- )
+ output_ids = [col_id_overrides.get(id, id) for id in node.schema.names]
+ return self._compiler.compile_sql(node, ordered=True, output_ids=output_ids)
 
  def compile_raw(
  self,
@@ -68,13 +66,12 @@ def compile_raw(
  str, Sequence[bigquery.SchemaField], bigframes.core.ordering.RowOrdering
  ]:
  """Compile node into sql that exposes all columns, including hidden ordering-only columns."""
- ir = self._compiler.compile_ordered_ir(node)
- sql, schema = ir.raw_sql_and_schema()
- return sql, schema, ir._ordering
+ return self._compiler.compile_raw(node)
 
 
 def test_only_try_evaluate(node: bigframes.core.nodes.BigFrameNode):
  """Use only for unit testing paths - not fully featured. Will throw exception if fails."""
+ node = _STRICT_COMPILER._preprocess(node)
  ibis = _STRICT_COMPILER.compile_ordered_ir(node)._to_ibis_expr(
  ordering_mode="unordered"
  )
@@ -85,9 +82,10 @@ def test_only_ibis_inferred_schema(node: bigframes.core.nodes.BigFrameNode):
  """Use only for testing paths to ensure ibis inferred schema does not diverge from bigframes inferred schema."""
  import bigframes.core.schema
 
+ node = _STRICT_COMPILER._preprocess(node)
  compiled = _STRICT_COMPILER.compile_unordered_ir(node)
  items = tuple(
- bigframes.core.schema.SchemaItem(id, compiled.get_column_type(id))
- for id in compiled.column_ids
+ bigframes.core.schema.SchemaItem(name, compiled.get_column_type(ibis_id))
+ for name, ibis_id in zip(node.schema.names, compiled.column_ids)
  )
  return bigframes.core.schema.ArraySchema(items)

bigframes/core/compile/compiled.py

@@ -202,7 +202,12 @@ def _aggregate_base(
  )
  # Must have deterministic ordering, so order by the unique "by" column
  ordering = TotalOrdering(
- tuple([OrderingExpression(column_id) for column_id in by_column_ids]),
+ tuple(
+ [
+ OrderingExpression(ex.DerefOp(ref.id.local_normalized))
+ for ref in by_column_ids
+ ]
+ ),
  total_ordering_columns=frozenset(
  [ex.DerefOp(ref.id.local_normalized) for ref in by_column_ids]
  ),
@@ -266,31 +271,26 @@ def peek_sql(self, n: int):
  def to_sql(
  self,
  offset_column: typing.Optional[str] = None,
- col_id_overrides: typing.Mapping[str, str] = {},
  ordered: bool = False,
  ) -> str:
  if offset_column or ordered:
  raise ValueError("Cannot produce sorted sql in partial ordering mode")
- sql = ibis_bigquery.Backend().compile(
- self._to_ibis_expr(
- col_id_overrides=col_id_overrides,
- )
- )
+ sql = ibis_bigquery.Backend().compile(self._to_ibis_expr())
  return typing.cast(str, sql)
 
- def row_count(self) -> OrderedIR:
+ def row_count(self, name: str) -> OrderedIR:
  original_table = self._to_ibis_expr()
  ibis_table = original_table.agg(
  [
- original_table.count().name("count"),
+ original_table.count().name(name),
  ]
  )
  return OrderedIR(
  ibis_table,
- (ibis_table["count"],),
+ (ibis_table[name],),
  ordering=TotalOrdering(
- ordering_value_columns=(ascending_over("count"),),
- total_ordering_columns=frozenset([ex.deref("count")]),
+ ordering_value_columns=(ascending_over(name),),
+ total_ordering_columns=frozenset([ex.deref(name)]),
  ),
  )
 
@@ -299,7 +299,6 @@ def _to_ibis_expr(
  *,
  expose_hidden_cols: bool = False,
  fraction: Optional[float] = None,
- col_id_overrides: typing.Mapping[str, str] = {},
  ):
  """
  Creates an Ibis table expression representing the DataFrame.
@@ -320,8 +319,6 @@ def _to_ibis_expr(
  If True, include the hidden ordering columns in the results.
  Only compatible with `order_by` and `unordered`
  ``ordering_mode``.
- col_id_overrides:
- overrides the column ids for the result
  Returns:
  An ibis expression representing the data help by the ArrayValue object.
  """
@@ -346,10 +343,6 @@ def _to_ibis_expr(
  if self._reduced_predicate is not None:
  table = table.filter(base_table[PREDICATE_COLUMN])
  table = table.drop(*columns_to_drop)
- if col_id_overrides:
- table = table.rename(
- {value: key for key, value in col_id_overrides.items()}
- )
  if fraction is not None:
  table = table.filter(ibis.random() < ibis.literal(fraction))
  return table
@@ -941,7 +934,6 @@ def _reproject_to_table(self) -> OrderedIR:
 
  def to_sql(
  self,
- col_id_overrides: typing.Mapping[str, str] = {},
  ordered: bool = False,
  limit: Optional[int] = None,
  ) -> str:
@@ -951,17 +943,13 @@ def to_sql(
  sql = ibis_bigquery.Backend().compile(
  baked_ir._to_ibis_expr(
  ordering_mode="unordered",
- col_id_overrides=col_id_overrides,
  expose_hidden_cols=True,
  )
  )
- output_columns = [
- col_id_overrides.get(col, col) for col in baked_ir.column_ids
- ]
  sql = (
  bigframes.core.compile.googlesql.Select()
  .from_(sql)
- .select(output_columns)
+ .select(self.column_ids)
  .sql()
  )
 
@@ -979,24 +967,26 @@ def to_sql(
  sql = ibis_bigquery.Backend().compile(
  self._to_ibis_expr(
  ordering_mode="unordered",
- col_id_overrides=col_id_overrides,
  expose_hidden_cols=False,
  )
  )
  return typing.cast(str, sql)
 
  def raw_sql_and_schema(
  self,
+ column_ids: typing.Sequence[str],
  ) -> typing.Tuple[str, typing.Sequence[google.cloud.bigquery.SchemaField]]:
  """Return sql with all hidden columns. Used to cache with ordering information.
 
  Also returns schema, as the extra ordering columns are determined compile-time.
  """
+ col_id_overrides = dict(zip(self.column_ids, column_ids))
  all_columns = (*self.column_ids, *self._hidden_ordering_column_names.keys())
  as_ibis = self._to_ibis_expr(
  ordering_mode="unordered",
  expose_hidden_cols=True,
- ).select(all_columns)
+ )
+ as_ibis = as_ibis.select(all_columns).rename(col_id_overrides)
 
  # Ibis will produce non-nullable schema types, but bigframes should always be nullable
  fixed_ibis_schema = ibis_schema.Schema.from_tuples(
@@ -1013,7 +1003,6 @@ def _to_ibis_expr(
  *,
  expose_hidden_cols: bool = False,
  fraction: Optional[float] = None,
- col_id_overrides: typing.Mapping[str, str] = {},
  ordering_mode: Literal["string_encoded", "unordered"],
  order_col_name: Optional[str] = ORDER_ID_COLUMN,
  ):
@@ -1043,8 +1032,6 @@ def _to_ibis_expr(
  order_col_name:
  If the ordering mode outputs a single ordering or offsets
  column, use this as the column name.
- col_id_overrides:
- overrides the column ids for the result
  Returns:
  An ibis expression representing the data help by the ArrayValue object.
  """
@@ -1086,10 +1073,6 @@ def _to_ibis_expr(
  if self._reduced_predicate is not None:
  table = table.filter(base_table[PREDICATE_COLUMN])
  table = table.drop(*columns_to_drop)
- if col_id_overrides:
- table = table.rename(
- {value: key for key, value in col_id_overrides.items()}
- )
  if fraction is not None:
  table = table.filter(ibis.random() < ibis.literal(fraction))
  return table

bigframes/core/compile/compiler.py

@@ -18,6 +18,7 @@
 import io
 import typing
 
+import google.cloud.bigquery
 import ibis
 import ibis.backends
 import ibis.backends.bigquery
@@ -32,6 +33,7 @@
 import bigframes.core.compile.scalar_op_compiler as compile_scalar
 import bigframes.core.compile.schema_translator
 import bigframes.core.compile.single_column
+import bigframes.core.expression as ex
 import bigframes.core.guid as guids
 import bigframes.core.identifiers as ids
 import bigframes.core.nodes as nodes
@@ -50,31 +52,66 @@ class Compiler:
  strict: bool = True
  scalar_op_compiler = compile_scalar.ScalarOpCompiler()
  enable_pruning: bool = False
+ enable_densify_ids: bool = False
+
+ def compile_sql(
+ self, node: nodes.BigFrameNode, ordered: bool, output_ids: typing.Sequence[str]
+ ) -> str:
+ node = self.set_output_names(node, output_ids)
+ if ordered:
+ node, limit = rewrites.pullup_limit_from_slice(node)
+ return self.compile_ordered_ir(self._preprocess(node)).to_sql(
+ ordered=True, limit=limit
+ )
+ else:
+ return self.compile_unordered_ir(self._preprocess(node)).to_sql()
+
+ def compile_peek_sql(self, node: nodes.BigFrameNode, n_rows: int) -> str:
+ return self.compile_unordered_ir(self._preprocess(node)).peek_sql(n_rows)
+
+ def compile_raw(
+ self,
+ node: bigframes.core.nodes.BigFrameNode,
+ ) -> typing.Tuple[
+ str, typing.Sequence[google.cloud.bigquery.SchemaField], bf_ordering.RowOrdering
+ ]:
+ ir = self.compile_ordered_ir(self._preprocess(node))
+ sql, schema = ir.raw_sql_and_schema(column_ids=node.schema.names)
+ return sql, schema, ir._ordering
 
  def _preprocess(self, node: nodes.BigFrameNode):
  if self.enable_pruning:
  used_fields = frozenset(field.id for field in node.fields)
  node = node.prune(used_fields)
  node = functools.cache(rewrites.replace_slice_ops)(node)
+ if self.enable_densify_ids:
+ original_names = [id.name for id in node.ids]
+ node, _ = rewrites.remap_variables(
+ node, id_generator=ids.anonymous_serial_ids()
+ )
+ node = self.set_output_names(node, original_names)
  return node
 
- def compile_ordered_ir(self, node: nodes.BigFrameNode) -> compiled.OrderedIR:
- ir = typing.cast(
- compiled.OrderedIR, self.compile_node(self._preprocess(node), True)
+ def set_output_names(
+ self, node: bigframes.core.nodes.BigFrameNode, output_ids: typing.Sequence[str]
+ ):
+ # TODO: Create specialized output operators that will handle final names
+ return nodes.SelectionNode(
+ node,
+ tuple(
+ (ex.DerefOp(old_id), ids.ColumnId(out_id))
+ for old_id, out_id in zip(node.ids, output_ids)
+ ),
  )
+
+ def compile_ordered_ir(self, node: nodes.BigFrameNode) -> compiled.OrderedIR:
+ ir = typing.cast(compiled.OrderedIR, self.compile_node(node, True))
  if self.strict:
  assert ir.has_total_order
  return ir
 
  def compile_unordered_ir(self, node: nodes.BigFrameNode) -> compiled.UnorderedIR:
- return typing.cast(
- compiled.UnorderedIR, self.compile_node(self._preprocess(node), False)
- )
-
- def compile_peak_sql(
- self, node: nodes.BigFrameNode, n_rows: int
- ) -> typing.Optional[str]:
- return self.compile_unordered_ir(self._preprocess(node)).peek_sql(n_rows)
+ return typing.cast(compiled.UnorderedIR, self.compile_node(node, False))
 
  # TODO: Remove cache when schema no longer requires compilation to derive schema (and therefor only compiles for execution)
  @functools.lru_cache(maxsize=5000)
@@ -144,11 +181,11 @@ def compile_fromrange(self, node: nodes.FromRangeNode, ordered: bool = True):
 
  labels_array_table = ibis.range(
  joined_table[start_column], joined_table[end_column] + node.step, node.step
- ).name("labels")
+ ).name(node.output_id.sql)
  labels = (
  typing.cast(ibis.expr.types.ArrayValue, labels_array_table)
  .as_table()
- .unnest(["labels"])
+ .unnest([node.output_id.sql])
  )
  if ordered:
  return compiled.OrderedIR(
@@ -307,18 +344,19 @@ def compile_projection(self, node: nodes.ProjectionNode, ordered: bool = True):
 
  @_compile_node.register
  def compile_concat(self, node: nodes.ConcatNode, ordered: bool = True):
+ output_ids = [id.sql for id in node.output_ids]
  if ordered:
  compiled_ordered = [self.compile_ordered_ir(node) for node in node.children]
- return concat_impl.concat_ordered(compiled_ordered)
+ return concat_impl.concat_ordered(compiled_ordered, output_ids)
  else:
  compiled_unordered = [
  self.compile_unordered_ir(node) for node in node.children
  ]
- return concat_impl.concat_unordered(compiled_unordered)
+ return concat_impl.concat_unordered(compiled_unordered, output_ids)
 
  @_compile_node.register
  def compile_rowcount(self, node: nodes.RowCountNode, ordered: bool = True):
- result = self.compile_unordered_ir(node.child).row_count()
+ result = self.compile_unordered_ir(node.child).row_count(name=node.col_id.sql)
  return result if ordered else result.to_unordered()
 
  @_compile_node.register