address comments

Author: chelsea-lin

bigframes/core/compile/sqlglot/expressions/binary_compiler.py

@@ -83,7 +83,7 @@ def _(op, left: TypedExpr, right: TypedExpr) -> sge.Expression:
         right_expr = sge.Cast(this=right_expr, to="INT64")
 
     result = sge.func("IEEE_DIVIDE", left_expr, right_expr)
-    if dtypes.is_numeric(right.dtype) and left.dtype == dtypes.TIMEDELTA_DTYPE:
+    if left.dtype == dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype):
         return sge.Cast(this=sge.Floor(this=result), to="INT64")
     else:
         return result
@@ -111,7 +111,7 @@ def _(op, left: TypedExpr, right: TypedExpr) -> sge.Expression:
     result = sge.Mul(this=left_expr, expression=right_expr)
 
     if (dtypes.is_numeric(left.dtype) and right.dtype == dtypes.TIMEDELTA_DTYPE) or (
-        dtypes.is_numeric(right.dtype) and left.dtype == dtypes.TIMEDELTA_DTYPE
+        left.dtype == dtypes.TIMEDELTA_DTYPE and dtypes.is_numeric(right.dtype)
     ):
         return sge.Cast(this=sge.Floor(this=result), to="INT64")
     else:

tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_div_timedelta/out.sql

@@ -1,18 +1,21 @@
 WITH `bfcte_0` AS (
   SELECT
-    `date_col` AS `bfcol_0`,
+    `int64_col` AS `bfcol_0`,
     `rowindex` AS `bfcol_1`,
     `timestamp_col` AS `bfcol_2`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
     *,
-    43200000000 AS `bfcol_6`
+    `bfcol_1` AS `bfcol_6`,
+    `bfcol_2` AS `bfcol_7`,
+    `bfcol_0` AS `bfcol_8`,
+    CAST(FLOOR(IEEE_DIVIDE(86400000000, `bfcol_0`)) AS INT64) AS `bfcol_9`
   FROM `bfcte_0`
 )
 SELECT
-  `bfcol_1` AS `rowindex`,
-  `bfcol_2` AS `timestamp_col`,
-  `bfcol_0` AS `date_col`,
-  `bfcol_6` AS `timedelta_div_numeric`
+  `bfcol_6` AS `rowindex`,
+  `bfcol_7` AS `timestamp_col`,
+  `bfcol_8` AS `int64_col`,
+  `bfcol_9` AS `timedelta_div_numeric`
 FROM `bfcte_1`

tests/unit/core/compile/sqlglot/expressions/snapshots/test_binary_compiler/test_mul_timedelta/out.sql

@@ -1,24 +1,31 @@
 WITH `bfcte_0` AS (
   SELECT
-    `date_col` AS `bfcol_0`,
+    `int64_col` AS `bfcol_0`,
     `rowindex` AS `bfcol_1`,
     `timestamp_col` AS `bfcol_2`
   FROM `bigframes-dev`.`sqlglot_test`.`scalar_types`
 ), `bfcte_1` AS (
   SELECT
     *,
-    172800000000 AS `bfcol_6`
+    `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`
   FROM `bfcte_0`
 ), `bfcte_2` AS (
   SELECT
     *,
-    172800000000 AS `bfcol_7`
+    `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`
   FROM `bfcte_1`
 )
 SELECT
-  `bfcol_1` AS `rowindex`,
-  `bfcol_2` AS `timestamp_col`,
-  `bfcol_0` AS `date_col`,
-  `bfcol_6` AS `timedelta_mul_numeric`,
-  `bfcol_7` AS `numeric_mul_timedelta`
+  `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`

tests/unit/core/compile/sqlglot/expressions/test_binary_compiler.py

@@ -95,10 +95,9 @@ def test_div_numeric(scalar_types_df: bpd.DataFrame, snapshot):
 
 
 def test_div_timedelta(scalar_types_df: bpd.DataFrame, snapshot):
-    bf_df = scalar_types_df[["timestamp_col", "date_col"]]
+    bf_df = scalar_types_df[["timestamp_col", "int64_col"]]
     timedelta = pd.Timedelta(1, unit="d")
-
-    bf_df["timedelta_div_numeric"] = timedelta / 2
+    bf_df["timedelta_div_numeric"] = timedelta / bf_df["int64_col"]
 
     snapshot.assert_match(bf_df.sql, "out.sql")
 
@@ -158,10 +157,10 @@ 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", "date_col"]]
+    bf_df = scalar_types_df[["timestamp_col", "int64_col"]]
     timedelta = pd.Timedelta(1, unit="d")
 
-    bf_df["timedelta_mul_numeric"] = timedelta * 2
-    bf_df["numeric_mul_timedelta"] = 2 * timedelta
+    bf_df["timedelta_mul_numeric"] = timedelta * bf_df["int64_col"]
+    bf_df["numeric_mul_timedelta"] = bf_df["int64_col"] * timedelta
 
     snapshot.assert_match(bf_df.sql, "out.sql")