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

Author: shuoweil

bigframes/core/array_value.py

@@ -401,37 +401,10 @@ def aggregate(
             )
         )
 
-    def project_window_op(
-        self,
-        column_name: str,
-        op: agg_ops.UnaryWindowOp,
-        window_spec: WindowSpec,
-        *,
-        never_skip_nulls=False,
-        skip_reproject_unsafe: bool = False,
-    ) -> Tuple[ArrayValue, str]:
-        """
-        Creates a new expression based on this expression with unary operation applied to one column.
-        column_name: the id of the input column present in the expression
-        op: the windowable operator to apply to the input column
-        window_spec: a specification of the window over which to apply the operator
-        output_name: the id to assign to the output of the operator, by default will replace input col if distinct output id not provided
-        never_skip_nulls: will disable null skipping for operators that would otherwise do so
-        skip_reproject_unsafe: skips the reprojection step, can be used when performing many non-dependent window operations, user responsible for not nesting window expressions, or using outputs as join, filter or aggregation keys before a reprojection
-        """
-
-        return self.project_window_expr(
-            agg_expressions.UnaryAggregation(op, ex.deref(column_name)),
-            window_spec,
-            never_skip_nulls,
-            skip_reproject_unsafe,
-        )
-
     def project_window_expr(
         self,
         expression: agg_expressions.Aggregation,
         window: WindowSpec,
-        never_skip_nulls=False,
         skip_reproject_unsafe: bool = False,
     ):
         output_name = self._gen_namespaced_uid()
@@ -442,7 +415,6 @@ def project_window_expr(
                     expression=expression,
                     window_spec=window,
                     output_name=ids.ColumnId(output_name),
-                    never_skip_nulls=never_skip_nulls,
                     skip_reproject_unsafe=skip_reproject_unsafe,
                 )
             ),

bigframes/core/blocks.py

@@ -1090,7 +1090,6 @@ def multi_apply_window_op(
         window_spec: windows.WindowSpec,
         *,
         skip_null_groups: bool = False,
-        never_skip_nulls: bool = False,
     ) -> typing.Tuple[Block, typing.Sequence[str]]:
         block = self
         result_ids = []
@@ -1103,7 +1102,6 @@ def multi_apply_window_op(
                 skip_reproject_unsafe=(i + 1) < len(columns),
                 result_label=label,
                 skip_null_groups=skip_null_groups,
-                never_skip_nulls=never_skip_nulls,
             )
             result_ids.append(result_id)
         return block, result_ids
@@ -1184,15 +1182,13 @@ def apply_window_op(
         result_label: Label = None,
         skip_null_groups: bool = False,
         skip_reproject_unsafe: bool = False,
-        never_skip_nulls: bool = False,
     ) -> typing.Tuple[Block, str]:
         agg_expr = agg_expressions.UnaryAggregation(op, ex.deref(column))
         return self.apply_analytic(
             agg_expr,
             window_spec,
             result_label,
             skip_reproject_unsafe=skip_reproject_unsafe,
-            never_skip_nulls=never_skip_nulls,
             skip_null_groups=skip_null_groups,
         )
 
@@ -1203,7 +1199,6 @@ def apply_analytic(
         result_label: Label,
         *,
         skip_reproject_unsafe: bool = False,
-        never_skip_nulls: bool = False,
         skip_null_groups: bool = False,
     ) -> typing.Tuple[Block, str]:
         block = self
@@ -1214,7 +1209,6 @@ def apply_analytic(
             agg_expr,
             window,
             skip_reproject_unsafe=skip_reproject_unsafe,
-            never_skip_nulls=never_skip_nulls,
         )
         block = Block(
             expr,

bigframes/core/compile/compiled.py

@@ -394,16 +394,13 @@ def project_window_op(
         expression: ex_types.Aggregation,
         window_spec: WindowSpec,
         output_name: str,
-        *,
-        never_skip_nulls=False,
     ) -> UnorderedIR:
         """
         Creates a new expression based on this expression with unary operation applied to one column.
         column_name: the id of the input column present in the expression
         op: the windowable operator to apply to the input column
         window_spec: a specification of the window over which to apply the operator
         output_name: the id to assign to the output of the operator
-        never_skip_nulls: will disable null skipping for operators that would otherwise do so
         """
         # Cannot nest analytic expressions, so reproject to cte first if needed.
         # Also ibis cannot window literals, so need to reproject those (even though this is legal in googlesql)
@@ -425,7 +422,6 @@ def project_window_op(
                 expression,
                 window_spec,
                 output_name,
-                never_skip_nulls=never_skip_nulls,
             )
 
         if expression.op.order_independent and window_spec.is_unbounded:
@@ -437,9 +433,6 @@ def project_window_op(
             expression, window_spec
         )
         clauses: list[tuple[ex.Expression, ex.Expression]] = []
-        if expression.op.skips_nulls and not never_skip_nulls:
-            for input in expression.inputs:
-                clauses.append((ops.isnull_op.as_expr(input), ex.const(None)))
         if window_spec.min_periods and len(expression.inputs) > 0:
             if not expression.op.nulls_count_for_min_values:
                 is_observation = ops.notnull_op.as_expr()

bigframes/core/compile/ibis_compiler/ibis_compiler.py

@@ -269,7 +269,6 @@ def compile_window(node: nodes.WindowOpNode, child: compiled.UnorderedIR):
         node.expression,
         node.window_spec,
         node.output_name.sql,
-        never_skip_nulls=node.never_skip_nulls,
     )
     return result
 

bigframes/core/compile/polars/compiler.py

@@ -16,7 +16,6 @@
 import dataclasses
 import functools
 import itertools
-import operator
 from typing import cast, Literal, Optional, Sequence, Tuple, Type, TYPE_CHECKING
 
 import pandas as pd
@@ -868,26 +867,6 @@ def compile_window(self, node: nodes.WindowOpNode):
                     df, node.expression, node.window_spec, node.output_name.sql
                 )
                 result = pl.concat([df, window_result], how="horizontal")
-
-            # Probably easier just to pull this out as a rewriter
-            if (
-                node.expression.op.skips_nulls
-                and not node.never_skip_nulls
-                and node.expression.column_references
-            ):
-                nullity_expr = functools.reduce(
-                    operator.or_,
-                    (
-                        pl.col(column.sql).is_null()
-                        for column in node.expression.column_references
-                    ),
-                )
-                result = result.with_columns(
-                    pl.when(nullity_expr)
-                    .then(None)
-                    .otherwise(pl.col(node.output_name.sql))
-                    .alias(node.output_name.sql)
-                )
             return result
 
         def _calc_row_analytic_func(

bigframes/core/compile/sqlglot/aggregations/unary_compiler.py

@@ -111,6 +111,140 @@ def _(
     return apply_window_if_present(sge.func("COUNT", column.expr), window)
 
 
+@UNARY_OP_REGISTRATION.register(agg_ops.CutOp)
+def _(
+    op: agg_ops.CutOp,
+    column: typed_expr.TypedExpr,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Expression:
+    if isinstance(op.bins, int):
+        case_expr = _cut_ops_w_int_bins(op, column, op.bins, window)
+    else:  # Interpret as intervals
+        case_expr = _cut_ops_w_intervals(op, column, op.bins, window)
+    return case_expr
+
+
+def _cut_ops_w_int_bins(
+    op: agg_ops.CutOp,
+    column: typed_expr.TypedExpr,
+    bins: int,
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Case:
+    case_expr = sge.Case()
+    col_min = apply_window_if_present(
+        sge.func("MIN", column.expr), window or window_spec.WindowSpec()
+    )
+    col_max = apply_window_if_present(
+        sge.func("MAX", column.expr), window or window_spec.WindowSpec()
+    )
+    adj: sge.Expression = sge.Sub(this=col_max, expression=col_min) * sge.convert(0.001)
+    bin_width: sge.Expression = sge.func(
+        "IEEE_DIVIDE",
+        sge.Sub(this=col_max, expression=col_min),
+        sge.convert(bins),
+    )
+
+    for this_bin in range(bins):
+        value: sge.Expression
+        if op.labels is False:
+            value = ir._literal(this_bin, dtypes.INT_DTYPE)
+        elif isinstance(op.labels, typing.Iterable):
+            value = ir._literal(list(op.labels)[this_bin], dtypes.STRING_DTYPE)
+        else:
+            left_adj: sge.Expression = (
+                adj if this_bin == 0 and op.right else sge.convert(0)
+            )
+            right_adj: sge.Expression = (
+                adj if this_bin == bins - 1 and not op.right else sge.convert(0)
+            )
+
+            left: sge.Expression = (
+                col_min + sge.convert(this_bin) * bin_width - left_adj
+            )
+            right: sge.Expression = (
+                col_min + sge.convert(this_bin + 1) * bin_width + right_adj
+            )
+            if op.right:
+                left_identifier = sge.Identifier(this="left_exclusive", quoted=True)
+                right_identifier = sge.Identifier(this="right_inclusive", quoted=True)
+            else:
+                left_identifier = sge.Identifier(this="left_inclusive", quoted=True)
+                right_identifier = sge.Identifier(this="right_exclusive", quoted=True)
+
+            value = sge.Struct(
+                expressions=[
+                    sge.PropertyEQ(this=left_identifier, expression=left),
+                    sge.PropertyEQ(this=right_identifier, expression=right),
+                ]
+            )
+
+        condition: sge.Expression
+        if this_bin == bins - 1:
+            condition = sge.Is(this=column.expr, expression=sge.Not(this=sge.Null()))
+        else:
+            if op.right:
+                condition = sge.LTE(
+                    this=column.expr,
+                    expression=(col_min + sge.convert(this_bin + 1) * bin_width),
+                )
+            else:
+                condition = sge.LT(
+                    this=column.expr,
+                    expression=(col_min + sge.convert(this_bin + 1) * bin_width),
+                )
+        case_expr = case_expr.when(condition, value)
+    return case_expr
+
+
+def _cut_ops_w_intervals(
+    op: agg_ops.CutOp,
+    column: typed_expr.TypedExpr,
+    bins: typing.Iterable[typing.Tuple[typing.Any, typing.Any]],
+    window: typing.Optional[window_spec.WindowSpec] = None,
+) -> sge.Case:
+    case_expr = sge.Case()
+    for this_bin, interval in enumerate(bins):
+        left: sge.Expression = ir._literal(
+            interval[0], dtypes.infer_literal_type(interval[0])
+        )
+        right: sge.Expression = ir._literal(
+            interval[1], dtypes.infer_literal_type(interval[1])
+        )
+        condition: sge.Expression
+        if op.right:
+            condition = sge.And(
+                this=sge.GT(this=column.expr, expression=left),
+                expression=sge.LTE(this=column.expr, expression=right),
+            )
+        else:
+            condition = sge.And(
+                this=sge.GTE(this=column.expr, expression=left),
+                expression=sge.LT(this=column.expr, expression=right),
+            )
+
+        value: sge.Expression
+        if op.labels is False:
+            value = ir._literal(this_bin, dtypes.INT_DTYPE)
+        elif isinstance(op.labels, typing.Iterable):
+            value = ir._literal(list(op.labels)[this_bin], dtypes.STRING_DTYPE)
+        else:
+            if op.right:
+                left_identifier = sge.Identifier(this="left_exclusive", quoted=True)
+                right_identifier = sge.Identifier(this="right_inclusive", quoted=True)
+            else:
+                left_identifier = sge.Identifier(this="left_inclusive", quoted=True)
+                right_identifier = sge.Identifier(this="right_exclusive", quoted=True)
+
+            value = sge.Struct(
+                expressions=[
+                    sge.PropertyEQ(this=left_identifier, expression=left),
+                    sge.PropertyEQ(this=right_identifier, expression=right),
+                ]
+            )
+        case_expr = case_expr.when(condition, value)
+    return case_expr
+
+
 @UNARY_OP_REGISTRATION.register(agg_ops.DateSeriesDiffOp)
 def _(
     op: agg_ops.DateSeriesDiffOp,

bigframes/core/compile/sqlglot/compiler.py

@@ -324,12 +324,8 @@ def compile_window(node: nodes.WindowOpNode, child: ir.SQLGlotIR) -> ir.SQLGlotI
     )
 
     clauses: list[tuple[sge.Expression, sge.Expression]] = []
-    if node.expression.op.skips_nulls and not node.never_skip_nulls:
-        for column in inputs:
-            clauses.append((sge.Is(this=column, expression=sge.Null()), sge.Null()))
-
     if window_spec.min_periods and len(inputs) > 0:
-        if node.expression.op.skips_nulls:
+        if not node.expression.op.nulls_count_for_min_values:
             # Most operations do not count NULL values towards min_periods
             not_null_columns = [
                 sge.Not(this=sge.Is(this=column, expression=sge.Null()))

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

@@ -116,6 +116,11 @@ def _(expr: TypedExpr) -> sge.Expression:
     return sge.func("SAFE.ST_Y", expr.expr)
 
 
+@register_binary_op(ops.GeoStDistanceOp, pass_op=True)
+def _(left: TypedExpr, right: TypedExpr, op: ops.GeoStDistanceOp) -> sge.Expression:
+    return sge.func("ST_DISTANCE", left.expr, right.expr, sge.convert(op.use_spheroid))
+
+
 @register_binary_op(ops.geo_st_difference_op)
 def _(left: TypedExpr, right: TypedExpr) -> sge.Expression:
     return sge.func("ST_DIFFERENCE", left.expr, right.expr)