Merge remote-tracking branch 'origin/main' into search-function-18023861678525870072

Author: tswast

bigframes/core/block_transforms.py

@@ -625,21 +625,7 @@ def skew(
     # counts, moment3 for each column
     aggregations = []
     for col in original_columns:
-        delta3_expr = _mean_delta_to_power(3, col)
-        count_agg = agg_expressions.UnaryAggregation(
-            agg_ops.count_op,
-            ex.deref(col),
-        )
-        moment3_agg = agg_expressions.UnaryAggregation(
-            agg_ops.mean_op,
-            delta3_expr,
-        )
-        variance_agg = agg_expressions.UnaryAggregation(
-            agg_ops.PopVarOp(),
-            ex.deref(col),
-        )
-        skew_expr = _skew_from_moments_and_count(count_agg, moment3_agg, variance_agg)
-        aggregations.append(skew_expr)
+        aggregations.append(skew_expr(ex.deref(col)))
 
     block = block.aggregate(
         aggregations, grouping_column_ids, column_labels=column_labels
@@ -663,16 +649,7 @@ def kurt(
     # counts, moment4 for each column
     kurt_exprs = []
     for col in original_columns:
-        delta_4_expr = _mean_delta_to_power(4, col)
-        count_agg = agg_expressions.UnaryAggregation(agg_ops.count_op, ex.deref(col))
-        moment4_agg = agg_expressions.UnaryAggregation(agg_ops.mean_op, delta_4_expr)
-        variance_agg = agg_expressions.UnaryAggregation(
-            agg_ops.PopVarOp(), ex.deref(col)
-        )
-
-        # Corresponds to order of aggregations in preceding loop
-        kurt_expr = _kurt_from_moments_and_count(count_agg, moment4_agg, variance_agg)
-        kurt_exprs.append(kurt_expr)
+        kurt_exprs.append(kurt_expr(ex.deref(col)))
 
     block = block.aggregate(
         kurt_exprs, grouping_column_ids, column_labels=column_labels
@@ -686,13 +663,38 @@ def kurt(
     return block
 
 
+def skew_expr(expr: ex.Expression) -> ex.Expression:
+    delta3_expr = _mean_delta_to_power(3, expr)
+    count_agg = agg_expressions.UnaryAggregation(
+        agg_ops.count_op,
+        expr,
+    )
+    moment3_agg = agg_expressions.UnaryAggregation(
+        agg_ops.mean_op,
+        delta3_expr,
+    )
+    variance_agg = agg_expressions.UnaryAggregation(
+        agg_ops.PopVarOp(),
+        expr,
+    )
+    return _skew_from_moments_and_count(count_agg, moment3_agg, variance_agg)
+
+
+def kurt_expr(expr: ex.Expression) -> ex.Expression:
+    delta_4_expr = _mean_delta_to_power(4, expr)
+    count_agg = agg_expressions.UnaryAggregation(agg_ops.count_op, expr)
+    moment4_agg = agg_expressions.UnaryAggregation(agg_ops.mean_op, delta_4_expr)
+    variance_agg = agg_expressions.UnaryAggregation(agg_ops.PopVarOp(), expr)
+    return _kurt_from_moments_and_count(count_agg, moment4_agg, variance_agg)
+
+
 def _mean_delta_to_power(
     n_power: int,
-    val_id: str,
+    col_expr: ex.Expression,
 ) -> ex.Expression:
     """Calculate (x-mean(x))^n. Useful for calculating moment statistics such as skew and kurtosis."""
-    mean_expr = agg_expressions.UnaryAggregation(agg_ops.mean_op, ex.deref(val_id))
-    delta = ops.sub_op.as_expr(val_id, mean_expr)
+    mean_expr = agg_expressions.UnaryAggregation(agg_ops.mean_op, col_expr)
+    delta = ops.sub_op.as_expr(col_expr, mean_expr)
     return ops.pow_op.as_expr(delta, ex.const(n_power))
 
 

bigframes/core/bq_data.py

@@ -64,6 +64,21 @@ def from_table(table: bq.Table, columns: Sequence[str] = ()) -> GbqTable:
             else tuple(table.clustering_fields),
         )
 
+    @staticmethod
+    def from_ref_and_schema(
+        table_ref: bq.TableReference,
+        schema: Sequence[bq.SchemaField],
+        cluster_cols: Optional[Sequence[str]] = None,
+    ) -> GbqTable:
+        return GbqTable(
+            project_id=table_ref.project,
+            dataset_id=table_ref.dataset_id,
+            table_id=table_ref.table_id,
+            physical_schema=tuple(schema),
+            is_physically_stored=True,
+            cluster_cols=tuple(cluster_cols) if cluster_cols else None,
+        )
+
     def get_table_ref(self) -> bq.TableReference:
         return bq.TableReference(
             bq.DatasetReference(self.project_id, self.dataset_id), self.table_id

bigframes/core/compile/sqlglot/compiler.py

@@ -378,7 +378,7 @@ def compile_window(node: nodes.WindowOpNode, child: ir.SQLGlotIR) -> ir.SQLGlotI
             window_op = sge.Case(ifs=when_expressions, default=window_op)
 
         # TODO: check if we can directly window the expression.
-        result = child.window(
+        result = result.window(
             window_op=window_op,
             output_column_id=cdef.id.sql,
         )

bigframes/core/compile/sqlglot/expressions/ai_ops.py

@@ -93,6 +93,7 @@ def _construct_prompt(
     for elem in prompt_context:
         if elem is None:
             prompt.append(exprs[column_ref_idx].expr)
+            column_ref_idx += 1
         else:
             prompt.append(sge.Literal.string(elem))
 

bigframes/core/compile/sqlglot/expressions/json_ops.py

@@ -69,6 +69,11 @@ def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("PARSE_JSON", expr.expr)
 
 
+@register_unary_op(ops.ToJSON)
+def _(expr: TypedExpr) -> sge.Expression:
+    return sge.func("TO_JSON", expr.expr)
+
+
 @register_unary_op(ops.ToJSONString)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("TO_JSON_STRING", expr.expr)

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -21,14 +21,15 @@
 
 from google.cloud import bigquery
 import numpy as np
+import pandas as pd
 import pyarrow as pa
 import sqlglot as sg
 import sqlglot.dialects.bigquery
 import sqlglot.expressions as sge
 
 from bigframes import dtypes
 from bigframes.core import guid, local_data, schema, utils
-from bigframes.core.compile.sqlglot.expressions import typed_expr
+from bigframes.core.compile.sqlglot.expressions import constants, typed_expr
 import bigframes.core.compile.sqlglot.sqlglot_types as sgt
 
 # shapely.wkt.dumps was moved to shapely.io.to_wkt in 2.0.
@@ -639,12 +640,30 @@ def _select_to_cte(expr: sge.Select, cte_name: sge.Identifier) -> sge.Select:
 def _literal(value: typing.Any, dtype: dtypes.Dtype) -> sge.Expression:
     sqlglot_type = sgt.from_bigframes_dtype(dtype) if dtype else None
     if sqlglot_type is None:
-        if value is not None:
-            raise ValueError("Cannot infer SQLGlot type from None dtype.")
+        if not pd.isna(value):
+            raise ValueError(f"Cannot infer SQLGlot type from None dtype: {value}")
         return sge.Null()
 
     if value is None:
         return _cast(sge.Null(), sqlglot_type)
+    if dtypes.is_struct_like(dtype):
+        items = [
+            _literal(value=value[field_name], dtype=field_dtype).as_(
+                field_name, quoted=True
+            )
+            for field_name, field_dtype in dtypes.get_struct_fields(dtype).items()
+        ]
+        return sge.Struct.from_arg_list(items)
+    elif dtypes.is_array_like(dtype):
+        value_type = dtypes.get_array_inner_type(dtype)
+        values = sge.Array(
+            expressions=[_literal(value=v, dtype=value_type) for v in value]
+        )
+        return values if len(value) > 0 else _cast(values, sqlglot_type)
+    elif pd.isna(value):
+        return _cast(sge.Null(), sqlglot_type)
+    elif dtype == dtypes.JSON_DTYPE:
+        return sge.ParseJSON(this=sge.convert(str(value)))
     elif dtype == dtypes.BYTES_DTYPE:
         return _cast(str(value), sqlglot_type)
     elif dtypes.is_time_like(dtype):
@@ -658,24 +677,12 @@ def _literal(value: typing.Any, dtype: dtypes.Dtype) -> sge.Expression:
     elif dtypes.is_geo_like(dtype):
         wkt = value if isinstance(value, str) else to_wkt(value)
         return sge.func("ST_GEOGFROMTEXT", sge.convert(wkt))
-    elif dtype == dtypes.JSON_DTYPE:
-        return sge.ParseJSON(this=sge.convert(str(value)))
     elif dtype == dtypes.TIMEDELTA_DTYPE:
         return sge.convert(utils.timedelta_to_micros(value))
-    elif dtypes.is_struct_like(dtype):
-        items = [
-            _literal(value=value[field_name], dtype=field_dtype).as_(
-                field_name, quoted=True
-            )
-            for field_name, field_dtype in dtypes.get_struct_fields(dtype).items()
-        ]
-        return sge.Struct.from_arg_list(items)
-    elif dtypes.is_array_like(dtype):
-        value_type = dtypes.get_array_inner_type(dtype)
-        values = sge.Array(
-            expressions=[_literal(value=v, dtype=value_type) for v in value]
-        )
-        return values if len(value) > 0 else _cast(values, sqlglot_type)
+    elif dtype == dtypes.FLOAT_DTYPE:
+        if np.isinf(value):
+            return constants._INF if value > 0 else constants._NEG_INF
+        return sge.convert(value)
     else:
         if isinstance(value, np.generic):
             value = value.item()

bigframes/core/expression.py

@@ -15,12 +15,11 @@
 from __future__ import annotations
 
 import abc
-import collections
 import dataclasses
 import functools
 import itertools
 import typing
-from typing import Callable, Dict, Generator, Mapping, Tuple, TypeVar, Union
+from typing import Callable, Generator, Mapping, TypeVar, Union
 
 import pandas as pd
 
@@ -162,57 +161,6 @@ def walk(self) -> Generator[Expression, None, None]:
         for child in self.children:
             yield from child.children
 
-    def unique_nodes(
-        self: Expression,
-    ) -> Generator[Expression, None, None]:
-        """Walks the tree for unique nodes"""
-        seen = set()
-        stack: list[Expression] = [self]
-        while stack:
-            item = stack.pop()
-            if item not in seen:
-                yield item
-                seen.add(item)
-                stack.extend(item.children)
-
-    def iter_nodes_topo(
-        self: Expression,
-    ) -> Generator[Expression, None, None]:
-        """Returns nodes in reverse topological order, using Kahn's algorithm."""
-        child_to_parents: Dict[Expression, list[Expression]] = collections.defaultdict(
-            list
-        )
-        out_degree: Dict[Expression, int] = collections.defaultdict(int)
-
-        queue: collections.deque["Expression"] = collections.deque()
-        for node in list(self.unique_nodes()):
-            num_children = len(node.children)
-            out_degree[node] = num_children
-            if num_children == 0:
-                queue.append(node)
-            for child in node.children:
-                child_to_parents[child].append(node)
-
-        while queue:
-            item = queue.popleft()
-            yield item
-            parents = child_to_parents.get(item, [])
-            for parent in parents:
-                out_degree[parent] -= 1
-                if out_degree[parent] == 0:
-                    queue.append(parent)
-
-    def reduce_up(self, reduction: Callable[[Expression, Tuple[T, ...]], T]) -> T:
-        """Apply a bottom-up reduction to the tree."""
-        results: dict[Expression, T] = {}
-        for node in list(self.iter_nodes_topo()):
-            # child nodes have already been transformed
-            child_results = tuple(results[child] for child in node.children)
-            result = reduction(node, child_results)
-            results[node] = result
-
-        return results[self]
-
 
 @dataclasses.dataclass(frozen=True)
 class ScalarConstantExpression(Expression):