Merge branch 'main' into shuowei-blob-null

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())
@@ -148,6 +159,30 @@ def _(*cases_and_outputs: TypedExpr) -> sge.Expression:
     )
 
 
+@register_nary_op(ops.RowKey)
+def _(*values: TypedExpr) -> sge.Expression:
+    # All inputs into hash must be non-null or resulting hash will be null
+    str_values = [_convert_to_nonnull_string_sqlglot(value) for value in values]
+
+    full_row_hash_p1 = sge.func("FARM_FINGERPRINT", sge.Concat(expressions=str_values))
+
+    # By modifying value slightly, we get another hash uncorrelated with the first
+    full_row_hash_p2 = sge.func(
+        "FARM_FINGERPRINT", sge.Concat(expressions=[*str_values, sge.convert("_")])
+    )
+
+    # Used to disambiguate between identical rows (which will have identical hash)
+    random_hash_p3 = sge.func("RAND")
+
+    return sge.Concat(
+        expressions=[
+            sge.Cast(this=full_row_hash_p1, to="STRING"),
+            sge.Cast(this=full_row_hash_p2, to="STRING"),
+            sge.Cast(this=random_hash_p3, to="STRING"),
+        ]
+    )
+
+
 # Helper functions
 def _cast_to_json(expr: TypedExpr, op: ops.AsTypeOp) -> sge.Expression:
     from_type = expr.dtype
@@ -207,3 +242,32 @@ def _cast(expr: sge.Expression, to: str, safe: bool):
         return sge.TryCast(this=expr, to=to)
     else:
         return sge.Cast(this=expr, to=to)
+
+
+def _convert_to_nonnull_string_sqlglot(expr: TypedExpr) -> sge.Expression:
+    col_type = expr.dtype
+    sg_expr = expr.expr
+
+    if col_type == dtypes.STRING_DTYPE:
+        result = sg_expr
+    elif (
+        dtypes.is_numeric(col_type)
+        or dtypes.is_time_or_date_like(col_type)
+        or col_type == dtypes.BYTES_DTYPE
+    ):
+        result = sge.Cast(this=sg_expr, to="STRING")
+    elif col_type == dtypes.GEO_DTYPE:
+        result = sge.func("ST_ASTEXT", sg_expr)
+    else:
+        # TO_JSON_STRING works with all data types, but isn't the most efficient
+        # Needed for JSON, STRUCT and ARRAY datatypes
+        result = sge.func("TO_JSON_STRING", sg_expr)
+
+    # Escape backslashes and use backslash as delineator
+    escaped = sge.func(
+        "REPLACE",
+        sge.func("COALESCE", result, sge.convert("")),
+        sge.convert("\\"),
+        sge.convert("\\\\"),
+    )
+    return sge.Concat(expressions=[sge.convert("\\"), escaped])

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)
+        ]
+    )