chore: add duration_col to the sqlglot unit tests (#1998)

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