Merge branch 'main' into shuowei-anywidget-html-repr

Author: shuoweil

.github/.OwlBot.lock.yaml

@@ -0,0 +1,10 @@
+image: us-central1-docker.pkg.dev/cloud-sdk-librarian-prod/images-prod/python-librarian-generator@sha256:c8612d3fffb3f6a32353b2d1abd16b61e87811866f7ec9d65b59b02eb452a620
+libraries:
+  - id: bigframes
+    version: 2.28.0
+    apis: []
+    source_roots:
+      - .
+    preserve_regex: []
+    remove_regex: []
+    tag_format: v{version}

.github/.OwlBot.yaml

@@ -19,7 +19,7 @@
 import functools
 import itertools
 import typing
-from typing import Callable, Mapping, TypeVar
+from typing import Callable, Mapping, Tuple, TypeVar
 
 from bigframes import dtypes
 from bigframes.core import expression, window_spec
@@ -63,6 +63,10 @@ def inputs(
     ) -> typing.Tuple[expression.Expression, ...]:
         ...
 
+    @property
+    def children(self) -> Tuple[expression.Expression, ...]:
+        return self.inputs
+
     @property
     def free_variables(self) -> typing.Tuple[str, ...]:
         return tuple(
@@ -73,6 +77,10 @@ def free_variables(self) -> typing.Tuple[str, ...]:
     def is_const(self) -> bool:
         return all(child.is_const for child in self.inputs)
 
+    @functools.cached_property
+    def is_scalar_expr(self) -> bool:
+        return False
+
     @abc.abstractmethod
     def replace_args(self: TExpression, *arg) -> TExpression:
         ...
@@ -176,8 +184,13 @@ def output_type(self) -> dtypes.ExpressionType:
     def inputs(
         self,
     ) -> typing.Tuple[expression.Expression, ...]:
+        # TODO: Maybe make the window spec itself an expression?
         return (self.analytic_expr, *self.window.expressions)
 
+    @property
+    def children(self) -> Tuple[expression.Expression, ...]:
+        return self.inputs
+
     @property
     def free_variables(self) -> typing.Tuple[str, ...]:
         return tuple(
@@ -188,12 +201,16 @@ def free_variables(self) -> typing.Tuple[str, ...]:
     def is_const(self) -> bool:
         return all(child.is_const for child in self.inputs)
 
+    @functools.cached_property
+    def is_scalar_expr(self) -> bool:
+        return False
+
     def transform_children(
         self: WindowExpression,
         t: Callable[[expression.Expression], expression.Expression],
     ) -> WindowExpression:
         return WindowExpression(
-            self.analytic_expr.transform_children(t),
+            t(self.analytic_expr),  # type: ignore
             self.window.transform_exprs(t),
         )
 

.github/auto-approve.yml

@@ -16,19 +16,24 @@
 from dataclasses import dataclass
 import datetime
 import functools
+import itertools
 import typing
 from typing import Iterable, List, Mapping, Optional, Sequence, Tuple
 
 import google.cloud.bigquery
 import pandas
 import pyarrow as pa
 
-from bigframes.core import agg_expressions, bq_data
+from bigframes.core import (
+    agg_expressions,
+    bq_data,
+    expression_factoring,
+    join_def,
+    local_data,
+)
 import bigframes.core.expression as ex
 import bigframes.core.guid
 import bigframes.core.identifiers as ids
-import bigframes.core.join_def as join_def
-import bigframes.core.local_data as local_data
 import bigframes.core.nodes as nodes
 from bigframes.core.ordering import OrderingExpression
 import bigframes.core.ordering as orderings
@@ -261,6 +266,23 @@ def compute_values(self, assignments: Sequence[ex.Expression]):
             col_ids,
         )
 
+    def compute_general_expression(self, assignments: Sequence[ex.Expression]):
+        named_exprs = [
+            expression_factoring.NamedExpression(expr, ids.ColumnId.unique())
+            for expr in assignments
+        ]
+        # TODO: Push this to rewrite later to go from block expression to planning form
+        # TODO: Jointly fragmentize expressions to more efficiently reuse common sub-expressions
+        fragments = tuple(
+            itertools.chain.from_iterable(
+                expression_factoring.fragmentize_expression(expr)
+                for expr in named_exprs
+            )
+        )
+        target_ids = tuple(named_expr.name for named_expr in named_exprs)
+        new_root = expression_factoring.push_into_tree(self.node, fragments, target_ids)
+        return (ArrayValue(new_root), target_ids)
+
     def project_to_id(self, expression: ex.Expression):
         array_val, ids = self.compute_values(
             [expression],
@@ -428,13 +450,15 @@ def project_window_expr(
         )
 
     def isin(
-        self, other: ArrayValue, lcol: str, rcol: str
+        self,
+        other: ArrayValue,
+        lcol: str,
     ) -> typing.Tuple[ArrayValue, str]:
+        assert len(other.column_ids) == 1
         node = nodes.InNode(
             self.node,
             other.node,
             ex.deref(lcol),
-            ex.deref(rcol),
             indicator_col=ids.ColumnId.unique(),
         )
         return ArrayValue(node), node.indicator_col.name