add duration column

Author: chelsea-lin

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/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

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import typing
-
 import pytest
 
 from bigframes import operations as ops
@@ -23,21 +21,14 @@
 pytest.importorskip("pytest_snapshot")
 
 
-def _apply_unary_op(
-    obj: bpd.DataFrame, ops: ops.UnaryOp | typing.Sequence[ops.UnaryOp], arg: str
-) -> str:
+def _apply_unary_op(obj: bpd.DataFrame, op: ops.UnaryOp, arg: str) -> str:
     array_value = obj._block.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)
+    op_expr = op.as_expr(arg)
+    result, col_ids = array_value.compute_values([op_expr])
 
     # Rename columns for deterministic golden SQL results.
-    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])
+    assert len(col_ids) == 1
+    result = result.rename_columns({col_ids[0]: arg}).select_columns([arg])
 
     sql = result.session._executor.to_sql(result, enable_cache=False)
     return sql

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`