chore: add duration_col to the sqlglot unit tests

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/conftest.py

@@ -89,6 +89,7 @@ def scalar_types_table_schema() -> typing.Sequence[bigquery.SchemaField]:
         bigquery.SchemaField("string_col", "STRING"),
         bigquery.SchemaField("time_col", "TIME"),
         bigquery.SchemaField("timestamp_col", "TIMESTAMP"),
+        bigquery.SchemaField("duration_col", "INTEGER"),
     ]
 
 

tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_mul_timedelta/out.sql

@@ -2,30 +2,42 @@ WITH `bfcte_0` AS (
   SELECT
     `int64_col` AS `bfcol_0`,
     `rowindex` AS `bfcol_1`,
-    `timestamp_col` AS `bfcol_2`
+    `timestamp_col` AS `bfcol_2`,
+    `duration_col` AS `bfcol_3`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
     *,
-    `bfcol_1` AS `bfcol_6`,
-    `bfcol_2` AS `bfcol_7`,
-    `bfcol_0` AS `bfcol_8`,
-    CAST(FLOOR(86400000000 * `bfcol_0`) AS INT64) AS `bfcol_9`
+    `bfcol_1` AS `bfcol_8`,
+    `bfcol_2` AS `bfcol_9`,
+    `bfcol_0` AS `bfcol_10`,
+    INTERVAL `bfcol_3` MICROSECOND AS `bfcol_11`
   FROM `bfcte_0`
 ), `bfcte_2` AS (
   SELECT
     *,
-    `bfcol_6` AS `bfcol_14`,
-    `bfcol_7` AS `bfcol_15`,
     `bfcol_8` AS `bfcol_16`,
     `bfcol_9` AS `bfcol_17`,
-    CAST(FLOOR(`bfcol_8` * 86400000000) AS INT64) AS `bfcol_18`
+    `bfcol_10` AS `bfcol_18`,
+    `bfcol_11` AS `bfcol_19`,
+    CAST(FLOOR(`bfcol_11` * `bfcol_10`) AS INT64) AS `bfcol_20`
   FROM `bfcte_1`
+), `bfcte_3` AS (
+  SELECT
+    *,
+    `bfcol_16` AS `bfcol_26`,
+    `bfcol_17` AS `bfcol_27`,
+    `bfcol_18` AS `bfcol_28`,
+    `bfcol_19` AS `bfcol_29`,
+    `bfcol_20` AS `bfcol_30`,
+    CAST(FLOOR(`bfcol_18` * `bfcol_19`) AS INT64) AS `bfcol_31`
+  FROM `bfcte_2`
 )
 SELECT
-  `bfcol_14` AS `rowindex`,
-  `bfcol_15` AS `timestamp_col`,
-  `bfcol_16` AS `int64_col`,
-  `bfcol_17` AS `timedelta_mul_numeric`,
-  `bfcol_18` AS `numeric_mul_timedelta`
-FROM `bfcte_2`
+  `bfcol_26` AS `rowindex`,
+  `bfcol_27` AS `timestamp_col`,
+  `bfcol_28` AS `int64_col`,
+  `bfcol_29` AS `duration_col`,
+  `bfcol_30` AS `timedelta_mul_numeric`,
+  `bfcol_31` AS `numeric_mul_timedelta`
+FROM `bfcte_3`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_sub_timedelta/out.sql

@@ -2,59 +2,81 @@ WITH `bfcte_0` AS (
   SELECT
     `date_col` AS `bfcol_0`,
     `rowindex` AS `bfcol_1`,
-    `timestamp_col` AS `bfcol_2`
+    `timestamp_col` AS `bfcol_2`,
+    `duration_col` AS `bfcol_3`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
     *,
-    `bfcol_1` AS `bfcol_6`,
-    `bfcol_2` AS `bfcol_7`,
-    `bfcol_0` AS `bfcol_8`,
-    TIMESTAMP_SUB(CAST(`bfcol_0` AS DATETIME), INTERVAL 86400000000 MICROSECOND) AS `bfcol_9`
+    `bfcol_1` AS `bfcol_8`,
+    `bfcol_2` AS `bfcol_9`,
+    `bfcol_0` AS `bfcol_10`,
+    INTERVAL `bfcol_3` MICROSECOND AS `bfcol_11`
   FROM `bfcte_0`
 ), `bfcte_2` AS (
   SELECT
     *,
-    `bfcol_6` AS `bfcol_14`,
-    `bfcol_7` AS `bfcol_15`,
     `bfcol_8` AS `bfcol_16`,
     `bfcol_9` AS `bfcol_17`,
-    TIMESTAMP_SUB(`bfcol_7`, INTERVAL 86400000000 MICROSECOND) AS `bfcol_18`
+    `bfcol_11` AS `bfcol_18`,
+    `bfcol_10` AS `bfcol_19`,
+    TIMESTAMP_SUB(CAST(`bfcol_10` AS DATETIME), INTERVAL `bfcol_11` MICROSECOND) AS `bfcol_20`
   FROM `bfcte_1`
 ), `bfcte_3` AS (
   SELECT
     *,
-    `bfcol_14` AS `bfcol_24`,
-    `bfcol_15` AS `bfcol_25`,
     `bfcol_16` AS `bfcol_26`,
     `bfcol_17` AS `bfcol_27`,
     `bfcol_18` AS `bfcol_28`,
-    TIMESTAMP_DIFF(CAST(`bfcol_16` AS DATETIME), CAST(`bfcol_16` AS DATETIME), MICROSECOND) AS `bfcol_29`
+    `bfcol_19` AS `bfcol_29`,
+    `bfcol_20` AS `bfcol_30`,
+    TIMESTAMP_SUB(`bfcol_17`, INTERVAL `bfcol_18` MICROSECOND) AS `bfcol_31`
   FROM `bfcte_2`
 ), `bfcte_4` AS (
   SELECT
     *,
-    `bfcol_24` AS `bfcol_36`,
-    `bfcol_25` AS `bfcol_37`,
     `bfcol_26` AS `bfcol_38`,
     `bfcol_27` AS `bfcol_39`,
     `bfcol_28` AS `bfcol_40`,
     `bfcol_29` AS `bfcol_41`,
-    TIMESTAMP_DIFF(`bfcol_25`, `bfcol_25`, MICROSECOND) AS `bfcol_42`
+    `bfcol_30` AS `bfcol_42`,
+    `bfcol_31` AS `bfcol_43`,
+    TIMESTAMP_DIFF(CAST(`bfcol_29` AS DATETIME), CAST(`bfcol_29` AS DATETIME), MICROSECOND) AS `bfcol_44`
   FROM `bfcte_3`
 ), `bfcte_5` AS (
   SELECT
     *,
-    0 AS `bfcol_50`
+    `bfcol_38` AS `bfcol_52`,
+    `bfcol_39` AS `bfcol_53`,
+    `bfcol_40` AS `bfcol_54`,
+    `bfcol_41` AS `bfcol_55`,
+    `bfcol_42` AS `bfcol_56`,
+    `bfcol_43` AS `bfcol_57`,
+    `bfcol_44` AS `bfcol_58`,
+    TIMESTAMP_DIFF(`bfcol_39`, `bfcol_39`, MICROSECOND) AS `bfcol_59`
   FROM `bfcte_4`
+), `bfcte_6` AS (
+  SELECT
+    *,
+    `bfcol_52` AS `bfcol_68`,
+    `bfcol_53` AS `bfcol_69`,
+    `bfcol_54` AS `bfcol_70`,
+    `bfcol_55` AS `bfcol_71`,
+    `bfcol_56` AS `bfcol_72`,
+    `bfcol_57` AS `bfcol_73`,
+    `bfcol_58` AS `bfcol_74`,
+    `bfcol_59` AS `bfcol_75`,
+    `bfcol_54` - `bfcol_54` AS `bfcol_76`
+  FROM `bfcte_5`
 )
 SELECT
-  `bfcol_36` AS `rowindex`,
-  `bfcol_37` AS `timestamp_col`,
-  `bfcol_38` AS `date_col`,
-  `bfcol_39` AS `date_sub_timedelta`,
-  `bfcol_40` AS `timestamp_sub_timedelta`,
-  `bfcol_41` AS `timestamp_sub_date`,
-  `bfcol_42` AS `date_sub_timestamp`,
-  `bfcol_50` AS `timedelta_sub_timedelta`
-FROM `bfcte_5`
+  `bfcol_68` AS `rowindex`,
+  `bfcol_69` AS `timestamp_col`,
+  `bfcol_70` AS `duration_col`,
+  `bfcol_71` AS `date_col`,
+  `bfcol_72` AS `date_sub_timedelta`,
+  `bfcol_73` AS `timestamp_sub_timedelta`,
+  `bfcol_74` AS `timestamp_sub_date`,
+  `bfcol_75` AS `date_sub_timestamp`,
+  `bfcol_76` AS `timedelta_sub_timedelta`
+FROM `bfcte_6`

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_binary_compiler.py

@@ -124,14 +124,14 @@ def test_sub_numeric(scalar_types_df: bpd.DataFrame, snapshot):
 
 
 def test_sub_timedelta(scalar_types_df: bpd.DataFrame, snapshot):
-    bf_df = scalar_types_df[["timestamp_col", "date_col"]]
-    timedelta = pd.Timedelta(1, unit="d")
+    bf_df = scalar_types_df[["timestamp_col", "duration_col", "date_col"]]
+    bf_df["duration_col"] = bpd.to_timedelta(bf_df["duration_col"], unit="us")
 
-    bf_df["date_sub_timedelta"] = bf_df["date_col"] - timedelta
-    bf_df["timestamp_sub_timedelta"] = bf_df["timestamp_col"] - timedelta
+    bf_df["date_sub_timedelta"] = bf_df["date_col"] - bf_df["duration_col"]
+    bf_df["timestamp_sub_timedelta"] = bf_df["timestamp_col"] - bf_df["duration_col"]
     bf_df["timestamp_sub_date"] = bf_df["date_col"] - bf_df["date_col"]
     bf_df["date_sub_timestamp"] = bf_df["timestamp_col"] - bf_df["timestamp_col"]
-    bf_df["timedelta_sub_timedelta"] = timedelta - timedelta
+    bf_df["timedelta_sub_timedelta"] = bf_df["duration_col"] - bf_df["duration_col"]
 
     snapshot.assert_match(bf_df.sql, "out.sql")
 
@@ -157,11 +157,11 @@ def test_mul_numeric(scalar_types_df: bpd.DataFrame, snapshot):
 
 
 def test_mul_timedelta(scalar_types_df: bpd.DataFrame, snapshot):
-    bf_df = scalar_types_df[["timestamp_col", "int64_col"]]
-    timedelta = pd.Timedelta(1, unit="d")
+    bf_df = scalar_types_df[["timestamp_col", "int64_col", "duration_col"]]
+    bf_df["duration_col"] = bpd.to_timedelta(bf_df["duration_col"], unit="us")
 
-    bf_df["timedelta_mul_numeric"] = timedelta * bf_df["int64_col"]
-    bf_df["numeric_mul_timedelta"] = bf_df["int64_col"] * timedelta
+    bf_df["timedelta_mul_numeric"] = bf_df["duration_col"] * bf_df["int64_col"]
+    bf_df["numeric_mul_timedelta"] = bf_df["int64_col"] * bf_df["duration_col"]
 
     snapshot.assert_match(bf_df.sql, "out.sql")
 

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

@@ -599,9 +599,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):

tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable/out.sql

@@ -13,7 +13,8 @@ WITH `bfcte_0` AS (
     `rowindex_2` AS `bfcol_10`,
     `string_col` AS `bfcol_11`,
     `time_col` AS `bfcol_12`,
-    `timestamp_col` AS `bfcol_13`
+    `timestamp_col` AS `bfcol_13`,
+    `duration_col` AS `bfcol_14`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 )
 SELECT
@@ -31,5 +32,6 @@ SELECT
   `bfcol_10` AS `rowindex_2`,
   `bfcol_11` AS `string_col`,
   `bfcol_12` AS `time_col`,
-  `bfcol_13` AS `timestamp_col`
+  `bfcol_13` AS `timestamp_col`,
+  `bfcol_14` AS `duration_col`
 FROM `bfcte_0`