fix to_timedelta compiler

Author: chelsea-lin

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

@@ -22,6 +22,7 @@
 import sqlglot.expressions as sge
 
 from bigframes import operations as ops
+from bigframes.core.compile.constants import UNIT_TO_US_CONVERSION_FACTORS
 from bigframes.core.compile.sqlglot.expressions.op_registration import OpRegistration
 from bigframes.core.compile.sqlglot.expressions.typed_expr import TypedExpr
 
@@ -618,7 +619,11 @@ def _(op: ops.ToTimestampOp, expr: TypedExpr) -> sge.Expression:
 
 @UNARY_OP_REGISTRATION.register(ops.ToTimedeltaOp)
 def _(op: ops.ToTimedeltaOp, expr: TypedExpr) -> sge.Expression:
-    return sge.Interval(this=expr.expr, unit=sge.Identifier(this="SECOND"))
+    value = expr.expr
+    factor = UNIT_TO_US_CONVERSION_FACTORS[op.unit]
+    if factor != 1:
+        value = sge.Mul(this=value, expression=sge.convert(factor))
+    return sge.Interval(this=value, unit=sge.Identifier(this="MICROSECOND"))
 
 
 @UNARY_OP_REGISTRATION.register(ops.UnixMicros)

tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_timedelta/out.sql

@@ -1,13 +1,37 @@
 WITH `bfcte_0` AS (
   SELECT
-    `int64_col` AS `bfcol_0`
+    `int64_col` AS `bfcol_0`,
+    `rowindex` AS `bfcol_1`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
     *,
-    INTERVAL `bfcol_0` SECOND AS `bfcol_1`
+    `bfcol_1` AS `bfcol_4`,
+    `bfcol_0` AS `bfcol_5`,
+    INTERVAL `bfcol_0` MICROSECOND AS `bfcol_6`
   FROM `bfcte_0`
+), `bfcte_2` AS (
+  SELECT
+    *,
+    `bfcol_4` AS `bfcol_10`,
+    `bfcol_5` AS `bfcol_11`,
+    `bfcol_6` AS `bfcol_12`,
+    INTERVAL (`bfcol_5` * 1000000) MICROSECOND AS `bfcol_13`
+  FROM `bfcte_1`
+), `bfcte_3` AS (
+  SELECT
+    *,
+    `bfcol_10` AS `bfcol_18`,
+    `bfcol_11` AS `bfcol_19`,
+    `bfcol_12` AS `bfcol_20`,
+    `bfcol_13` AS `bfcol_21`,
+    INTERVAL (`bfcol_11` * 604800000000) MICROSECOND AS `bfcol_22`
+  FROM `bfcte_2`
 )
 SELECT
-  `bfcol_1` AS `int64_col`
-FROM `bfcte_1`
+  `bfcol_18` AS `rowindex`,
+  `bfcol_19` AS `int64_col`,
+  `bfcol_20` AS `duration_us`,
+  `bfcol_21` AS `duration_s`,
+  `bfcol_22` AS `duration_w`
+FROM `bfcte_3`

tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import typing
+
 import pytest
 
 from bigframes import operations as ops
@@ -21,14 +23,21 @@
 pytest.importorskip("pytest_snapshot")
 
 
-def _apply_unary_op(obj: bpd.DataFrame, op: ops.UnaryOp, arg: str) -> str:
+def _apply_unary_op(
+    obj: bpd.DataFrame, ops: ops.UnaryOp | typing.Sequence[ops.UnaryOp], arg: str
+) -> str:
     array_value = obj._block.expr
-    op_expr = op.as_expr(arg)
-    result, col_ids = array_value.compute_values([op_expr])
+    if not isinstance(ops, typing.Sequence):
+        ops = [ops]
+    op_expr = [op.as_expr(arg) for op in ops]
+    result, col_ids = array_value.compute_values(op_expr)
 
     # Rename columns for deterministic golden SQL results.
-    assert len(col_ids) == 1
-    result = result.rename_columns({col_ids[0]: arg}).select_columns([arg])
+    rename_cols = {
+        col_id: arg if index == 0 else f"{arg}_{index}"
+        for index, col_id in enumerate(col_ids)
+    }
+    result = result.rename_columns(rename_cols).select_columns([arg])
 
     sql = result.session._executor.to_sql(result, enable_cache=False)
     return sql
@@ -599,9 +608,11 @@ def test_to_timestamp(scalar_types_df: bpd.DataFrame, snapshot):
 
 def test_to_timedelta(scalar_types_df: bpd.DataFrame, snapshot):
     bf_df = scalar_types_df[["int64_col"]]
-    sql = _apply_unary_op(bf_df, ops.ToTimedeltaOp("s"), "int64_col")
+    bf_df["duration_us"] = bpd.to_timedelta(bf_df["int64_col"], "us")
+    bf_df["duration_s"] = bpd.to_timedelta(bf_df["int64_col"], "s")
+    bf_df["duration_w"] = bpd.to_timedelta(bf_df["int64_col"], "W")
 
-    snapshot.assert_match(sql, "out.sql")
+    snapshot.assert_match(bf_df.sql, "out.sql")
 
 
 def test_unix_micros(scalar_types_df: bpd.DataFrame, snapshot):