Merge branch 'main' into shuowei-anywidget-col

Author: shuoweil

bigframes/core/compile/polars/compiler.py

@@ -328,6 +328,26 @@ def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
             assert isinstance(op, string_ops.StrContainsRegexOp)
             return input.str.contains(pattern=op.pat, literal=False)
 
+        @compile_op.register(string_ops.UpperOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.UpperOp)
+            return input.str.to_uppercase()
+
+        @compile_op.register(string_ops.LowerOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.LowerOp)
+            return input.str.to_lowercase()
+
+        @compile_op.register(string_ops.ArrayLenOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.ArrayLenOp)
+            return input.list.len()
+
+        @compile_op.register(string_ops.StrLenOp)
+        def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
+            assert isinstance(op, string_ops.StrLenOp)
+            return input.str.len_chars()
+
         @compile_op.register(string_ops.StartsWithOp)
         def _(self, op: ops.ScalarOp, input: pl.Expr) -> pl.Expr:
             assert isinstance(op, string_ops.StartsWithOp)

bigframes/core/compile/polars/lowering.py

@@ -27,6 +27,7 @@
     generic_ops,
     json_ops,
     numeric_ops,
+    string_ops,
 )
 import bigframes.operations as ops
 
@@ -347,11 +348,28 @@ def lower(self, expr: expression.OpExpression) -> expression.Expression:
             return ops.coalesce_op.as_expr(new_isin, expression.const(False))
 
 
+class LowerLenOp(op_lowering.OpLoweringRule):
+    @property
+    def op(self) -> type[ops.ScalarOp]:
+        return string_ops.LenOp
+
+    def lower(self, expr: expression.OpExpression) -> expression.Expression:
+        assert isinstance(expr.op, string_ops.LenOp)
+        arg = expr.children[0]
+
+        if dtypes.is_string_like(arg.output_type):
+            return string_ops.StrLenOp().as_expr(arg)
+        elif dtypes.is_array_like(arg.output_type):
+            return string_ops.ArrayLenOp().as_expr(arg)
+        else:
+            raise ValueError(f"Unexpected type: {arg.output_type}")
+
+
 def _coerce_comparables(
     expr1: expression.Expression,
     expr2: expression.Expression,
     *,
-    bools_only: bool = False
+    bools_only: bool = False,
 ):
     if bools_only:
         if (
@@ -446,6 +464,7 @@ def _lower_cast(cast_op: ops.AsTypeOp, arg: expression.Expression):
     LowerAsTypeRule(),
     LowerInvertOp(),
     LowerIsinOp(),
+    LowerLenOp(),
 )
 
 

bigframes/core/compile/polars/operations/numeric_ops.py

@@ -29,15 +29,6 @@
     import polars as pl
 
 
-@polars_compiler.register_op(numeric_ops.CosOp)
-def cos_op_impl(
-    compiler: polars_compiler.PolarsExpressionCompiler,
-    op: numeric_ops.CosOp,  # type: ignore
-    input: pl.Expr,
-) -> pl.Expr:
-    return input.cos()
-
-
 @polars_compiler.register_op(numeric_ops.LnOp)
 def ln_op_impl(
     compiler: polars_compiler.PolarsExpressionCompiler,
@@ -80,6 +71,78 @@ def sin_op_impl(
     return input.sin()
 
 
+@polars_compiler.register_op(numeric_ops.CosOp)
+def cos_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.CosOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.cos()
+
+
+@polars_compiler.register_op(numeric_ops.TanOp)
+def tan_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.SinOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.tan()
+
+
+@polars_compiler.register_op(numeric_ops.SinhOp)
+def sinh_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.SinOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.sinh()
+
+
+@polars_compiler.register_op(numeric_ops.CoshOp)
+def cosh_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.CosOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.cosh()
+
+
+@polars_compiler.register_op(numeric_ops.TanhOp)
+def tanh_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.SinOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.tanh()
+
+
+@polars_compiler.register_op(numeric_ops.ArcsinOp)
+def asin_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.ArcsinOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.arcsin()
+
+
+@polars_compiler.register_op(numeric_ops.ArccosOp)
+def acos_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.ArccosOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.arccos()
+
+
+@polars_compiler.register_op(numeric_ops.ArctanOp)
+def atan_op_impl(
+    compiler: polars_compiler.PolarsExpressionCompiler,
+    op: numeric_ops.ArctanOp,  # type: ignore
+    input: pl.Expr,
+) -> pl.Expr:
+    return input.arctan()
+
+
 @polars_compiler.register_op(numeric_ops.SqrtOp)
 def sqrt_op_impl(
     compiler: polars_compiler.PolarsExpressionCompiler,

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

@@ -35,10 +35,7 @@ def _(expr: TypedExpr) -> sge.Expression:
 
 @register_unary_op(ops.dayofweek_op)
 def _(expr: TypedExpr) -> sge.Expression:
-    # Adjust the 1-based day-of-week index (from SQL) to a 0-based index.
-    return sge.Extract(
-        this=sge.Identifier(this="DAYOFWEEK"), expression=expr.expr
-    ) - sge.convert(1)
+    return dayofweek_op_impl(expr)
 
 
 @register_unary_op(ops.dayofyear_op)
@@ -48,7 +45,8 @@ def _(expr: TypedExpr) -> sge.Expression:
 
 @register_unary_op(ops.iso_day_op)
 def _(expr: TypedExpr) -> sge.Expression:
-    return sge.Extract(this=sge.Identifier(this="DAYOFWEEK"), expression=expr.expr)
+    # Plus 1 because iso day of week uses 1-based indexing
+    return dayofweek_op_impl(expr) + sge.convert(1)
 
 
 @register_unary_op(ops.iso_week_op)
@@ -59,3 +57,16 @@ def _(expr: TypedExpr) -> sge.Expression:
 @register_unary_op(ops.iso_year_op)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.Extract(this=sge.Identifier(this="ISOYEAR"), expression=expr.expr)
+
+
+# Helpers
+def dayofweek_op_impl(expr: TypedExpr) -> sge.Expression:
+    # BigQuery SQL Extract(DAYOFWEEK) returns 1 for Sunday through 7 for Saturday.
+    # We want 0 for Monday through 6 for Sunday to be compatible with Pandas.
+    extract_expr = sge.Extract(
+        this=sge.Identifier(this="DAYOFWEEK"), expression=expr.expr
+    )
+    return sge.Cast(
+        this=sge.Mod(this=extract_expr + sge.convert(5), expression=sge.convert(7)),
+        to="INT64",
+    )

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

@@ -25,8 +25,28 @@
 
 @register_unary_op(ops.FloorDtOp, pass_op=True)
 def _(expr: TypedExpr, op: ops.FloorDtOp) -> sge.Expression:
-    # TODO: Remove this method when it is covered by ops.FloorOp
-    return sge.TimestampTrunc(this=expr.expr, unit=sge.Identifier(this=op.freq))
+    pandas_to_bq_freq_map = {
+        "Y": "YEAR",
+        "Q": "QUARTER",
+        "M": "MONTH",
+        "W": "WEEK(MONDAY)",
+        "D": "DAY",
+        "h": "HOUR",
+        "min": "MINUTE",
+        "s": "SECOND",
+        "ms": "MILLISECOND",
+        "us": "MICROSECOND",
+        "ns": "NANOSECOND",
+    }
+    if op.freq not in pandas_to_bq_freq_map.keys():
+        raise NotImplementedError(
+            f"Unsupported freq paramater: {op.freq}"
+            + " Supported freq parameters are: "
+            + ",".join(pandas_to_bq_freq_map.keys())
+        )
+
+    bq_freq = pandas_to_bq_freq_map[op.freq]
+    return sge.TimestampTrunc(this=expr.expr, unit=sge.Identifier(this=bq_freq))
 
 
 @register_unary_op(ops.hour_op)

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

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import sqlglot as sg
 import sqlglot.expressions as sge
 
 from bigframes import dtypes
@@ -80,6 +81,16 @@ def _(expr: TypedExpr) -> sge.Expression:
     return sge.BitwiseNot(this=sge.paren(expr.expr))
 
 
+@register_nary_op(ops.SqlScalarOp, pass_op=True)
+def _(*operands: TypedExpr, op: ops.SqlScalarOp) -> sge.Expression:
+    return sg.parse_one(
+        op.sql_template.format(
+            *[operand.expr.sql(dialect="bigquery") for operand in operands]
+        ),
+        dialect="bigquery",
+    )
+
+
 @register_unary_op(ops.isnull_op)
 def _(expr: TypedExpr) -> sge.Expression:
     return sge.Is(this=expr.expr, expression=sge.Null())

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

@@ -24,6 +24,7 @@
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 import bigframes.core.compile.sqlglot.scalar_compiler as scalar_compiler
 
+register_nary_op = scalar_compiler.scalar_op_compiler.register_nary_op
 register_unary_op = scalar_compiler.scalar_op_compiler.register_unary_op
 
 
@@ -40,3 +41,13 @@ def _(expr: TypedExpr, op: ops.StructFieldOp) -> sge.Expression:
         this=sge.to_identifier(name, quoted=True),
         catalog=expr.expr,
     )
+
+
+@register_nary_op(ops.StructOp, pass_op=True)
+def _(*exprs: TypedExpr, op: ops.StructOp) -> sge.Struct:
+    return sge.Struct(
+        expressions=[
+            sge.PropertyEQ(this=sge.to_identifier(col), expression=expr.expr)
+            for col, expr in zip(op.column_names, exprs)
+        ]
+    )

bigframes/core/compile/sqlglot/scalar_compiler.py

@@ -31,6 +31,37 @@ class ScalarOpCompiler:
         typing.Callable[[typing.Sequence[TypedExpr], ops.RowOp], sge.Expression],
     ] = {}
 
+    # A set of SQLGlot classes that may need to be parenthesized
+    SQLGLOT_NEEDS_PARENS = {
+        # Numeric operations
+        sge.Add,
+        sge.Sub,
+        sge.Mul,
+        sge.Div,
+        sge.Mod,
+        sge.Pow,
+        # Comparison operations
+        sge.GTE,
+        sge.GT,
+        sge.LTE,
+        sge.LT,
+        sge.EQ,
+        sge.NEQ,
+        # Logical operations
+        sge.And,
+        sge.Or,
+        sge.Xor,
+        # Bitwise operations
+        sge.BitwiseAnd,
+        sge.BitwiseOr,
+        sge.BitwiseXor,
+        sge.BitwiseLeftShift,
+        sge.BitwiseRightShift,
+        sge.BitwiseNot,
+        # Other operations
+        sge.Is,
+    }
+
     @functools.singledispatchmethod
     def compile_expression(
         self,
@@ -110,10 +141,12 @@ def register_binary_op(
 
         def decorator(impl: typing.Callable[..., sge.Expression]):
             def normalized_impl(args: typing.Sequence[TypedExpr], op: ops.RowOp):
+                left = self._add_parentheses(args[0])
+                right = self._add_parentheses(args[1])
                 if pass_op:
-                    return impl(args[0], args[1], op)
+                    return impl(left, right, op)
                 else:
-                    return impl(args[0], args[1])
+                    return impl(left, right)
 
             self._register(key, normalized_impl)
             return impl
@@ -177,6 +210,12 @@ def _register(
             raise ValueError(f"Operation name {op_name} already registered")
         self._registry[op_name] = impl
 
+    @classmethod
+    def _add_parentheses(cls, expr: TypedExpr) -> TypedExpr:
+        if type(expr.expr) in cls.SQLGLOT_NEEDS_PARENS:
+            return TypedExpr(sge.paren(expr.expr, copy=False), expr.dtype)
+        return expr
+
 
 # Singleton compiler
 scalar_op_compiler = ScalarOpCompiler()