diff --git a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py b/bigframes/core/compile/sqlglot/expressions/unary_compiler.py index 60a8baea96..1551c555a7 100644 --- a/bigframes/core/compile/sqlglot/expressions/unary_compiler.py +++ b/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) diff --git a/tests/unit/core/compile/sqlglot/conftest.py b/tests/unit/core/compile/sqlglot/conftest.py index 754c19ac90..f65343fd66 100644 --- a/tests/unit/core/compile/sqlglot/conftest.py +++ b/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"), ] diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_mul_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_mul_timedelta/out.sql index 082a714845..c8a8cf6cbf 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_mul_timedelta/out.sql +++ b/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` \ No newline at end of file + `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` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_sub_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_sub_timedelta/out.sql index 41e45d3333..460f941d1b 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_sub_timedelta/out.sql +++ b/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` \ No newline at end of file + `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` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_timedelta/out.sql index b89056d65f..01ebebc455 100644 --- a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_unary_compiler/test_to_timedelta/out.sql +++ b/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` \ No newline at end of file + `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` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py index e4e1f93259..0d3fd42607 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py +++ b/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") diff --git a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py b/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py index e5a76b067f..4a5b586c77 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_unary_compiler.py +++ b/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): diff --git a/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable/out.sql b/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable/out.sql index 34fc8e3c49..10c2a2088a 100644 --- a/tests/unit/core/compile/sqlglot/snapshots/test_compile_readtable/test_compile_readtable/out.sql +++ b/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` \ No newline at end of file