Repository: tobymao/sqlglot
Pull request: Fix(optimizer)!: fix datetime coercion in the canonicalize rule
Fixes #4429

Problem statement:
Bigquery annotate_types incorrectly sets TIMESTAMP as DATETIME even when schema type is provided
```python
import sqlglot
import sqlglot.errors
from sqlglot.optimizer import optimize

dialect = "bigquery"
schema = {"test": {"products": {"created_at": "TIMESTAMP"}}}
expression = sqlglot.parse_one("SELECT * FROM test.products WHERE created_at > '2024-10-01'", read=dialect)
optimized_expression = optimize(
    expression,
    dialect=dialect,
    schema=schema,
)
formatted_sql = optimized_expression.sql(dialect=dialect, pretty=True)
print("@formatted_sql", formatted_sql)
```
formatted_sql is 
```SQL
SELECT
  `products`.`created_at` AS `created_at`
FROM `test`.`products` AS `products`
WHERE
  `products`.`created_at` > CAST('2024-10-01' AS DATETIME)
```

which cannot be ran due to 
No matching signature for operator < for argument types: TIMESTAMP, DATETIME.

Hints from issue discussion:
```
diff --git a/sqlglot/optimizer/canonicalize.py b/sqlglot/optimizer/canonicalize.py
index fd002622..4382dc70 100644
--- a/sqlglot/optimizer/canonicalize.py
+++ b/sqlglot/optimizer/canonicalize.py
@@ -138,7 +138,7 @@ def _coerce_date(a: exp.Expression, b: exp.Expression) -> None:
             and b.type
             and b.type.this in exp.DataType.TEXT_TYPES
         ):
-            _replace_cast(b, exp.DataType.Type.DATETIME)
+            _replace_cast(b, a.type)
```

cc @barakalon

Commit message: Fix(optimizer): fix datetime coercion in the canonicalize rule

Patch to review:
diff --git a/sqlglot/optimizer/canonicalize.py b/sqlglot/optimizer/canonicalize.py
index fd0026228b..f1c4ba573b 100644
--- a/sqlglot/optimizer/canonicalize.py
+++ b/sqlglot/optimizer/canonicalize.py
@@ -138,7 +138,7 @@ def _coerce_date(a: exp.Expression, b: exp.Expression) -> None:
             and b.type
             and b.type.this in exp.DataType.TEXT_TYPES
         ):
-            _replace_cast(b, exp.DataType.Type.DATETIME)
+            _replace_cast(b, a.type)
 
 
 def _coerce_timeunit_arg(arg: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.Expression:
@@ -168,7 +168,7 @@ def _coerce_datediff_args(node: exp.DateDiff) -> None:
             e.replace(exp.cast(e.copy(), to=exp.DataType.Type.DATETIME))
 
 
-def _replace_cast(node: exp.Expression, to: exp.DataType.Type) -> None:
+def _replace_cast(node: exp.Expression, to: exp.DATA_TYPE) -> None:
     node.replace(exp.cast(node.copy(), to=to))
 
 
diff --git a/tests/fixtures/optimizer/canonicalize.sql b/tests/fixtures/optimizer/canonicalize.sql
index 66c6c95c93..aa3eae988e 100644
--- a/tests/fixtures/optimizer/canonicalize.sql
+++ b/tests/fixtures/optimizer/canonicalize.sql
@@ -2,7 +2,7 @@ SELECT w.d + w.e AS c FROM w AS w;
 SELECT CONCAT("w"."d", "w"."e") AS "c" FROM "w" AS "w";
 
 SELECT CAST(w.d AS DATE) > w.e AS a FROM w AS w;
-SELECT CAST("w"."d" AS DATE) > CAST("w"."e" AS DATETIME) AS "a" FROM "w" AS "w";
+SELECT CAST("w"."d" AS DATE) > CAST("w"."e" AS DATE) AS "a" FROM "w" AS "w";
 
 SELECT CAST(1 AS VARCHAR) AS a FROM w AS w;
 SELECT CAST(1 AS VARCHAR) AS "a" FROM "w" AS "w";
@@ -102,7 +102,7 @@ DATEDIFF('2023-01-01', '2023-01-02', DAY);
 DATEDIFF(CAST('2023-01-01' AS DATETIME), CAST('2023-01-02' AS DATETIME), DAY);
 
 SELECT "t"."d" > '2023-01-01' AS "d" FROM "temporal" AS "t";
-SELECT "t"."d" > CAST('2023-01-01' AS DATETIME) AS "d" FROM "temporal" AS "t";
+SELECT "t"."d" > CAST('2023-01-01' AS DATE) AS "d" FROM "temporal" AS "t";
 
 SELECT "t"."d" > CAST('2023-01-01' AS DATETIME) AS "d" FROM "temporal" AS "t";
 SELECT "t"."d" > CAST('2023-01-01' AS DATETIME) AS "d" FROM "temporal" AS "t";
@@ -110,6 +110,9 @@ SELECT "t"."d" > CAST('2023-01-01' AS DATETIME) AS "d" FROM "temporal" AS "t";
 SELECT "t"."t" > '2023-01-01 00:00:01' AS "t" FROM "temporal" AS "t";
 SELECT "t"."t" > CAST('2023-01-01 00:00:01' AS DATETIME) AS "t" FROM "temporal" AS "t";
 
+WITH "t" AS (SELECT CAST("ext"."created_at" AS TIMESTAMP) AS "created_at" FROM "ext" AS "ext") SELECT "t"."created_at" > '2024-10-01 12:05:02' AS "col" FROM "t" AS "t";
+WITH "t" AS (SELECT CAST("ext"."created_at" AS TIMESTAMP) AS "created_at" FROM "ext" AS "ext") SELECT "t"."created_at" > CAST('2024-10-01 12:05:02' AS TIMESTAMP) AS "col" FROM "t" AS "t";
+
 --------------------------------------
 -- Remove redundant casts
 --------------------------------------
diff --git a/tests/fixtures/optimizer/tpc-ds/tpc-ds.sql b/tests/fixtures/optimizer/tpc-ds/tpc-ds.sql
index 290d2760b8..59bc432fdd 100644
--- a/tests/fixtures/optimizer/tpc-ds/tpc-ds.sql
+++ b/tests/fixtures/optimizer/tpc-ds/tpc-ds.sql
@@ -736,8 +736,8 @@ WITH "salesreturns" AS (
     "date_dim"."d_date" AS "d_date"
   FROM "date_dim" AS "date_dim"
   WHERE
-    CAST("date_dim"."d_date" AS DATETIME) <= CAST('2002-09-05' AS DATE)
-    AND CAST("date_dim"."d_date" AS DATETIME) >= CAST('2002-08-22' AS DATE)
+    CAST("date_dim"."d_date" AS DATE) <= CAST('2002-09-05' AS DATE)
+    AND CAST("date_dim"."d_date" AS DATE) >= CAST('2002-08-22' AS DATE)
 ), "ssr" AS (
   SELECT
     "store"."s_store_id" AS "s_store_id",
@@ -1853,8 +1853,8 @@ SELECT
 FROM "web_sales" AS "web_sales"
 JOIN "date_dim" AS "date_dim"
   ON "date_dim"."d_date_sk" = "web_sales"."ws_sold_date_sk"
-  AND CAST("date_dim"."d_date" AS DATETIME) <= CAST('2000-06-10' AS DATE)
-  AND CAST("date_dim"."d_date" AS DATETIME) >= CAST('2000-05-11' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) <= CAST('2000-06-10' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) >= CAST('2000-05-11' AS DATE)
 JOIN "item" AS "item"
   ON "item"."i_category" IN ('Home', 'Men', 'Women')
   AND "item"."i_item_sk" = "web_sales"."ws_item_sk"
@@ -2422,7 +2422,7 @@ JOIN "date_dim" AS "date_dim"
   AND "date_dim"."d_date" >= '2002-3-01'
   AND (
     CAST('2002-3-01' AS DATE) + INTERVAL '60' DAY
-  ) >= CAST("date_dim"."d_date" AS DATETIME)
+  ) >= CAST("date_dim"."d_date" AS DATE)
 WHERE
   "_u_3"."_u_4" IS NULL
   AND ARRAY_ANY("_u_0"."_u_2", "_x" -> "cs1"."cs_warehouse_sk" <> "_x")
@@ -2731,8 +2731,8 @@ SELECT
 FROM "catalog_sales" AS "catalog_sales"
 JOIN "date_dim" AS "date_dim"
   ON "catalog_sales"."cs_sold_date_sk" = "date_dim"."d_date_sk"
-  AND CAST("date_dim"."d_date" AS DATETIME) <= CAST('2001-03-05' AS DATE)
-  AND CAST("date_dim"."d_date" AS DATETIME) >= CAST('2001-02-03' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) <= CAST('2001-03-05' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) >= CAST('2001-02-03' AS DATE)
 JOIN "item" AS "item"
   ON "catalog_sales"."cs_item_sk" = "item"."i_item_sk"
   AND "item"."i_category" IN ('Children', 'Women', 'Electronics')
@@ -2811,8 +2811,8 @@ WITH "x" AS (
   FROM "inventory" AS "inventory"
   JOIN "date_dim" AS "date_dim"
     ON "date_dim"."d_date_sk" = "inventory"."inv_date_sk"
-    AND CAST("date_dim"."d_date" AS DATETIME) <= CAST('2000-06-12' AS DATE)
-    AND CAST("date_dim"."d_date" AS DATETIME) >= CAST('2000-04-13' AS DATE)
+    AND CAST("date_dim"."d_date" AS DATE) <= CAST('2000-06-12' AS DATE)
+    AND CAST("date_dim"."d_date" AS DATE) >= CAST('2000-04-13' AS DATE)
   JOIN "item" AS "item"
     ON "inventory"."inv_item_sk" = "item"."i_item_sk"
     AND "item"."i_current_price" <= 1.49
@@ -3944,7 +3944,7 @@ WITH "catalog_sales_2" AS (
   FROM "date_dim" AS "date_dim"
   WHERE
     "date_dim"."d_date" >= '2001-03-04'
-    AND CAST("date_dim"."d_date" AS DATETIME) <= CAST('2001-06-02' AS DATE)
+    AND CAST("date_dim"."d_date" AS DATE) <= CAST('2001-06-02' AS DATE)
 ), "_u_0" AS (
   SELECT
     1.3 * AVG("catalog_sales"."cs_ext_discount_amt") AS "_col_0",
@@ -4510,8 +4510,8 @@ JOIN "inventory" AS "inventory"
   AND "inventory"."inv_quantity_on_hand" >= 100
 JOIN "date_dim" AS "date_dim"
   ON "date_dim"."d_date_sk" = "inventory"."inv_date_sk"
-  AND CAST("date_dim"."d_date" AS DATETIME) <= CAST('1999-05-05' AS DATE)
-  AND CAST("date_dim"."d_date" AS DATETIME) >= CAST('1999-03-06' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) <= CAST('1999-05-05' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) >= CAST('1999-03-06' AS DATE)
 WHERE
   "item"."i_current_price" <= 50
   AND "item"."i_current_price" >= 20
@@ -4787,8 +4787,8 @@ LEFT JOIN "catalog_returns" AS "catalog_returns"
   AND "catalog_returns"."cr_order_number" = "catalog_sales"."cs_order_number"
 JOIN "date_dim" AS "date_dim"
   ON "catalog_sales"."cs_sold_date_sk" = "date_dim"."d_date_sk"
-  AND CAST("date_dim"."d_date" AS DATETIME) <= CAST('2002-07-01' AS DATE)
-  AND CAST("date_dim"."d_date" AS DATETIME) >= CAST('2002-05-02' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) <= CAST('2002-07-01' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) >= CAST('2002-05-02' AS DATE)
 JOIN "item" AS "item"
   ON "catalog_sales"."cs_item_sk" = "item"."i_item_sk"
   AND "item"."i_current_price" <= 1.49
@@ -10318,8 +10318,8 @@ WITH "date_dim_2" AS (
     "date_dim"."d_date" AS "d_date"
   FROM "date_dim" AS "date_dim"
   WHERE
-    CAST("date_dim"."d_date" AS DATETIME) <= CAST('2001-09-15' AS DATE)
-    AND CAST("date_dim"."d_date" AS DATETIME) >= CAST('2001-08-16' AS DATE)
+    CAST("date_dim"."d_date" AS DATE) <= CAST('2001-09-15' AS DATE)
+    AND CAST("date_dim"."d_date" AS DATE) >= CAST('2001-08-16' AS DATE)
 ), "store_2" AS (
   SELECT
     "store"."s_store_sk" AS "s_store_sk"
@@ -10828,8 +10828,8 @@ WITH "date_dim_2" AS (
     "date_dim"."d_date" AS "d_date"
   FROM "date_dim" AS "date_dim"
   WHERE
-    CAST("date_dim"."d_date" AS DATETIME) <= CAST('2000-09-25' AS DATE)
-    AND CAST("date_dim"."d_date" AS DATETIME) >= CAST('2000-08-26' AS DATE)
+    CAST("date_dim"."d_date" AS DATE) <= CAST('2000-09-25' AS DATE)
+    AND CAST("date_dim"."d_date" AS DATE) >= CAST('2000-08-26' AS DATE)
 ), "item_2" AS (
   SELECT
     "item"."i_item_sk" AS "i_item_sk",
@@ -11109,8 +11109,8 @@ JOIN "store_sales" AS "store_sales"
   ON "item"."i_item_sk" = "store_sales"."ss_item_sk"
 JOIN "date_dim" AS "date_dim"
   ON "date_dim"."d_date_sk" = "inventory"."inv_date_sk"
-  AND CAST("date_dim"."d_date" AS DATETIME) <= CAST('1998-06-26' AS DATE)
-  AND CAST("date_dim"."d_date" AS DATETIME) >= CAST('1998-04-27' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) <= CAST('1998-06-26' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) >= CAST('1998-04-27' AS DATE)
 WHERE
   "item"."i_current_price" <= 93
   AND "item"."i_current_price" >= 63
@@ -12180,7 +12180,7 @@ WITH "web_sales_2" AS (
   FROM "date_dim" AS "date_dim"
   WHERE
     "date_dim"."d_date" >= '2002-03-29'
-    AND CAST("date_dim"."d_date" AS DATETIME) <= CAST('2002-06-27' AS DATE)
+    AND CAST("date_dim"."d_date" AS DATE) <= CAST('2002-06-27' AS DATE)
 ), "_u_0" AS (
   SELECT
     1.3 * AVG("web_sales"."ws_ext_discount_amt") AS "_col_0",
@@ -12321,7 +12321,7 @@ JOIN "date_dim" AS "date_dim"
   AND "date_dim"."d_date_sk" = "ws1"."ws_ship_date_sk"
   AND (
     CAST('2000-3-01' AS DATE) + INTERVAL '60' DAY
-  ) >= CAST("date_dim"."d_date" AS DATETIME)
+  ) >= CAST("date_dim"."d_date" AS DATE)
 JOIN "web_site" AS "web_site"
   ON "web_site"."web_company_name" = 'pri'
   AND "web_site"."web_site_sk" = "ws1"."ws_web_site_sk"
@@ -12411,7 +12411,7 @@ JOIN "date_dim" AS "date_dim"
   AND "date_dim"."d_date_sk" = "ws1"."ws_ship_date_sk"
   AND (
     CAST('2000-4-01' AS DATE) + INTERVAL '60' DAY
-  ) >= CAST("date_dim"."d_date" AS DATETIME)
+  ) >= CAST("date_dim"."d_date" AS DATE)
 JOIN "web_site" AS "web_site"
   ON "web_site"."web_company_name" = 'pri'
   AND "web_site"."web_site_sk" = "ws1"."ws_web_site_sk"
@@ -12595,8 +12595,8 @@ SELECT
 FROM "store_sales" AS "store_sales"
 JOIN "date_dim" AS "date_dim"
   ON "date_dim"."d_date_sk" = "store_sales"."ss_sold_date_sk"
-  AND CAST("date_dim"."d_date" AS DATETIME) <= CAST('2000-06-17' AS DATE)
-  AND CAST("date_dim"."d_date" AS DATETIME) >= CAST('2000-05-18' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) <= CAST('2000-06-17' AS DATE)
+  AND CAST("date_dim"."d_date" AS DATE) >= CAST('2000-05-18' AS DATE)
 JOIN "item" AS "item"
   ON "item"."i_category" IN ('Men', 'Home', 'Electronics')
   AND "item"."i_item_sk" = "store_sales"."ss_item_sk"


Code review comments:
diff_hunk: @@ -138,7 +138,7 @@ def _coerce_date(a: exp.Expression, b: exp.Expression) -> None:
             and b.type
             and b.type.this in exp.DataType.TEXT_TYPES
         ):
-            _replace_cast(b, exp.DataType.Type.DATETIME)
+            _replace_cast(b, a.type)
path: sqlglot/optimizer/canonicalize.py
text: @user1:
Do we need to worry about type promotion?

```sql
date_col < '2024-01-01 00:00:01'
```



@author:
That's a fair concern, I think we do... for MySQL:

```
mysql> SELECT CAST('2024-01-01' AS DATE) < '2024-01-01 12:00:00';
+----------------------------------------------------+
| CAST('2024-01-01' AS DATE) < '2024-01-01 12:00:00' |
+----------------------------------------------------+
|                                                  1 |
+----------------------------------------------------+
1 row in set (0.00 sec)
```

But not for Postgres?!

```
georgesittas=# SELECT '2024-01-01'::DATE < '2024-01-01 12:00:00';
 ?column?
----------
 f
(1 row)
```

@user1:
Oy vey

@user2:
this canoncicalization is really only necessary for presto i think. 

@user1:
Yeah, that’s why I added it. But it’s got to match the behavior of one or the other (or both, configurably)

@user1:
Were extending MySQL in my application, so I’m generally biased. But Im ok matching Postgres here, since that’s simpler and this case feels pretty rare.

@user1:
Another thing to think about: I believe cast(‘2024-01-01 00:00:01’ as date) would fail in presto

@author:
Should we keep it Presto/Trino-specific then by controlling it with a dialect argument? FWIW, for the above check:

- Databricks, DuckDB, Postgres, Redshift, Snowflake, Spark, T-SQL - return false
- MySQL - return true
- BigQuery, ClickHouse, Presto, Trino - fails

So this looks like MySQL weirdness? Based on these observations, I'd assume the others promote string -> whatever datetime type of the other operand.

@user2:
presto doesn't allow x < '2020-01-01' which is why it was added in the first place

@user2:
yes, i'd keep it specific to presto/trino.

@author:
So, double-checking: keep the old logic but only have it kick in for Presto/Trino?

@user1:
Should the dialect arg be the parsed dialect or the target dialect?

I feel like the dialect arg should be avoided when possible, since it violates the idea that our AST is a superset of dialects 

@user2:
it's the parser dialect. there's no notion of a target dialect.

i guess given that, this code isn't valid because x < text is not legal in presto

@user2:
another option is to just disable this for the ones that fail on type mismatch

@user1:
Right but we use the optimizer to convert MySQL queries to Trino where normal transpilation typically fails. I thought that was the point of canonicalize in the first place?



@user2:
the point of canocalizie is to make a canonical sql to make two different sqls identical

@user2:
if the goal is to parse mysql queries, you need to respect it,

thus x < '2022-01-01 01:00:00' should return false

@user1:
Right. So the dialect config should be the automatic coercion behavior of the parsed dialect, which is either to match types or promote

@author:
FYI I decided to investigate this a little further to make sure I understand MySQL's semantics well before moving on.

@author:
I checked this [SO thread](https://stackoverflow.com/a/47561792) which helped, it seems like we're interested in this doc. [page](https://dev.mysql.com/doc/refman/8.4/en/date-and-time-type-conversion.html) in MySQL:

> Conversion of [DATE](https://dev.mysql.com/doc/refman/8.4/en/datetime.html) values:
>
>    Conversion to a [DATETIME](https://dev.mysql.com/doc/refman/8.4/en/datetime.html) or [TIMESTAMP](https://dev.mysql.com/doc/refman/8.4/en/datetime.html) value adds a time part of '00:00:00' because the [DATE](https://dev.mysql.com/doc/refman/8.4/en/datetime.html) value contains no time information.
>
>    Conversion to a [TIME](https://dev.mysql.com/doc/refman/8.4/en/time.html) value is not useful; the result is '00:00:00'.

diff_hunk: @@ -128,17 +132,46 @@ def remove_ascending_order(expression: exp.Expression) -> exp.Expression:
     return expression
 
 
-def _coerce_date(a: exp.Expression, b: exp.Expression) -> None:
+def _coerce_date(
+    a: exp.Expression,
+    b: exp.Expression,
+    promote_to_inferred_datetime_type: bool,
+) -> None:
     for a, b in itertools.permutations([a, b]):
         if isinstance(b, exp.Interval):
             a = _coerce_timeunit_arg(a, b.unit)
+
+        a_type = a.type
         if (
-            a.type
-            and a.type.this in exp.DataType.TEMPORAL_TYPES
+            a_type
+            and a_type.this in exp.DataType.TEMPORAL_TYPES
             and b.type
             and b.type.this in exp.DataType.TEXT_TYPES
         ):
-            _replace_cast(b, exp.DataType.Type.DATETIME)
+            if promote_to_inferred_datetime_type:
+                if b.is_string:
+                    date_text = b.name
+                    if is_iso_date(date_text):
+                        b_type = exp.DataType.Type.DATE
+                    elif is_iso_datetime(date_text):
+                        b_type = exp.DataType.Type.DATETIME
path: sqlglot/optimizer/canonicalize.py
text: @author:
@user1 it feels like this `is_iso*` logic is somewhat repeated now. I didn't refactor yet because the other sections where we do it take a `unit` argument as well into account which looks somewhat different, but curious to see if there's room for DRYing it up.

Other references:
- https://github.com/tobymao/sqlglot/blob/main/sqlglot/optimizer/annotate_types.py#L63-L74
- https://github.com/tobymao/sqlglot/blob/main/sqlglot/optimizer/canonicalize.py#L144-L162