refactor: add some numeric tests in the golden sql (#2195)

Author: chelsea-lin

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`

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`

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`

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`

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")