diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_add_string/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_add_string/out.sql new file mode 100644 index 0000000000..de5129a6a3 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_add_string/out.sql @@ -0,0 +1,13 @@ +WITH `bfcte_0` AS ( + SELECT + `string_col` AS `bfcol_0` + FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` +), `bfcte_1` AS ( + SELECT + *, + CONCAT(`bfcol_0`, 'a') AS `bfcol_1` + FROM `bfcte_0` +) +SELECT + `bfcol_1` AS `string_col` +FROM `bfcte_1` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_add_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_add_timedelta/out.sql new file mode 100644 index 0000000000..a47531999b --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_add_timedelta/out.sql @@ -0,0 +1,60 @@ +WITH `bfcte_0` AS ( + SELECT + `date_col` AS `bfcol_0`, + `rowindex` AS `bfcol_1`, + `timestamp_col` AS `bfcol_2` + 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_ADD(CAST(`bfcol_0` AS DATETIME), INTERVAL 86400000000 MICROSECOND) AS `bfcol_9` + 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_ADD(`bfcol_7`, INTERVAL 86400000000 MICROSECOND) AS `bfcol_18` + 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_ADD(CAST(`bfcol_16` AS DATETIME), INTERVAL 86400000000 MICROSECOND) AS `bfcol_29` + 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_ADD(`bfcol_25`, INTERVAL 86400000000 MICROSECOND) AS `bfcol_42` + FROM `bfcte_3` +), `bfcte_5` AS ( + SELECT + *, + 172800000000 AS `bfcol_50` + FROM `bfcte_4` +) +SELECT + `bfcol_36` AS `rowindex`, + `bfcol_37` AS `timestamp_col`, + `bfcol_38` AS `date_col`, + `bfcol_39` AS `date_add_timedelta`, + `bfcol_40` AS `timestamp_add_timedelta`, + `bfcol_41` AS `timedelta_add_date`, + `bfcol_42` AS `timedelta_add_timestamp`, + `bfcol_50` AS `timedelta_add_timedelta` +FROM `bfcte_5` \ No newline at end of file diff --git a/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_mul_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_mul_timedelta/out.sql new file mode 100644 index 0000000000..f8752d0a60 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_mul_timedelta/out.sql @@ -0,0 +1,43 @@ +WITH `bfcte_0` AS ( + SELECT + `int64_col` AS `bfcol_0`, + `rowindex` AS `bfcol_1`, + `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_8`, + `bfcol_2` AS `bfcol_9`, + `bfcol_0` AS `bfcol_10`, + `bfcol_3` AS `bfcol_11` + FROM `bfcte_0` +), `bfcte_2` AS ( + SELECT + *, + `bfcol_8` AS `bfcol_16`, + `bfcol_9` AS `bfcol_17`, + `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_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_numeric_ops/test_sub_timedelta/out.sql b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_sub_timedelta/out.sql new file mode 100644 index 0000000000..2d615fcca6 --- /dev/null +++ b/tests/unit/core/compile/sqlglot/expressions/snapshots/test_numeric_ops/test_sub_timedelta/out.sql @@ -0,0 +1,82 @@ +WITH `bfcte_0` AS ( + SELECT + `date_col` AS `bfcol_0`, + `rowindex` AS `bfcol_1`, + `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_8`, + `bfcol_2` AS `bfcol_9`, + `bfcol_0` AS `bfcol_10`, + `bfcol_3` AS `bfcol_11` + FROM `bfcte_0` +), `bfcte_2` AS ( + SELECT + *, + `bfcol_8` AS `bfcol_16`, + `bfcol_9` AS `bfcol_17`, + `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_16` AS `bfcol_26`, + `bfcol_17` AS `bfcol_27`, + `bfcol_18` AS `bfcol_28`, + `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_26` AS `bfcol_38`, + `bfcol_27` AS `bfcol_39`, + `bfcol_28` AS `bfcol_40`, + `bfcol_29` AS `bfcol_41`, + `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 + *, + `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_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/test_numeric_ops.py b/tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py index 59726da73b..fe9a53a558 100644 --- a/tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py +++ b/tests/unit/core/compile/sqlglot/expressions/test_numeric_ops.py @@ -16,6 +16,7 @@ import pytest from bigframes import operations as ops +import bigframes.core.expression as ex import bigframes.pandas as bpd from bigframes.testing import utils @@ -218,6 +219,34 @@ def test_add_numeric(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(bf_df.sql, "out.sql") +def test_add_string(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["string_col"]] + sql = utils._apply_binary_op(bf_df, ops.add_op, "string_col", ex.const("a")) + + snapshot.assert_match(sql, "out.sql") + + +def test_add_timedelta(scalar_types_df: bpd.DataFrame, snapshot): + bf_df = scalar_types_df[["timestamp_col", "date_col"]] + timedelta = pd.Timedelta(1, unit="d") + + bf_df["date_add_timedelta"] = bf_df["date_col"] + timedelta + bf_df["timestamp_add_timedelta"] = bf_df["timestamp_col"] + timedelta + bf_df["timedelta_add_date"] = timedelta + bf_df["date_col"] + bf_df["timedelta_add_timestamp"] = timedelta + bf_df["timestamp_col"] + bf_df["timedelta_add_timedelta"] = timedelta + timedelta + + snapshot.assert_match(bf_df.sql, "out.sql") + + +def test_add_unsupported_raises(scalar_types_df: bpd.DataFrame): + with pytest.raises(TypeError): + utils._apply_binary_op(scalar_types_df, ops.add_op, "timestamp_col", "date_col") + + with pytest.raises(TypeError): + utils._apply_binary_op(scalar_types_df, ops.add_op, "int64_col", "string_col") + + def test_div_numeric(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["int64_col", "bool_col", "float64_col"]] @@ -279,6 +308,16 @@ def test_mul_numeric(scalar_types_df: bpd.DataFrame, snapshot): snapshot.assert_match(bf_df.sql, "out.sql") +def test_mul_timedelta(scalar_types_df: bpd.DataFrame, snapshot): + 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"] = 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") + + def test_mod_numeric(scalar_types_df: bpd.DataFrame, snapshot): bf_df = scalar_types_df[["int64_col", "float64_col"]] @@ -305,3 +344,24 @@ def test_sub_numeric(scalar_types_df: bpd.DataFrame, snapshot): bf_df["bool_add_int"] = bf_df["bool_col"] - bf_df["int64_col"] snapshot.assert_match(bf_df.sql, "out.sql") + + +def test_sub_timedelta(scalar_types_df: bpd.DataFrame, snapshot): + 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"] - 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"] = bf_df["duration_col"] - bf_df["duration_col"] + + snapshot.assert_match(bf_df.sql, "out.sql") + + +def test_sub_unsupported_raises(scalar_types_df: bpd.DataFrame): + with pytest.raises(TypeError): + utils._apply_binary_op(scalar_types_df, ops.sub_op, "string_col", "string_col") + + with pytest.raises(TypeError): + utils._apply_binary_op(scalar_types_df, ops.sub_op, "int64_col", "string_col")