Fix(optimizer)!: fix datetime coercion in the canonicalize rule (#4431)

georgesittas · web-flow · commit fe3539464a15 · 2024-11-22T21:13:53.000+02:00
* Fix(optimizer): fix datetime coercion in the canonicalize rule

* Refactor: infer literal type and promote as needed (mysql only atm)

* Reduce nesting level
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
@@ -397,6 +397,13 @@ class Dialect(metaclass=_Dialect):
     ARRAY_AGG_INCLUDES_NULLS: t.Optional[bool] = True
     """Whether ArrayAgg needs to filter NULL values."""
 
+    PROMOTE_TO_INFERRED_DATETIME_TYPE = False
+    """
+    This flag is used in the optimizer's canonicalize rule and determines whether x will be promoted
+    to the literal's type in x::DATE < '2020-01-01 12:05:03' (i.e., DATETIME). When false, the literal
+    is cast to x's type to match it instead.
+    """
+
     REGEXP_EXTRACT_DEFAULT_GROUP = 0
     """The default value for the capturing group."""
 
diff --git a/sqlglot/dialects/mysql.py b/sqlglot/dialects/mysql.py
@@ -145,6 +145,8 @@ def func(self: MySQL.Generator, expression: exp.Func) -> str:
 
 
 class MySQL(Dialect):
+    PROMOTE_TO_INFERRED_DATETIME_TYPE = True
+
     # https://dev.mysql.com/doc/refman/8.0/en/identifiers.html
     IDENTIFIERS_CAN_START_WITH_DIGIT = True
 
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -7866,6 +7866,7 @@ def cast(
         types_are_equivalent = type_mapping.get(
             existing_cast_type, existing_cast_type.value
         ) == type_mapping.get(new_cast_type, new_cast_type.value)
+
         if expr.is_type(data_type) or types_are_equivalent:
             return expr
 
diff --git a/sqlglot/optimizer/canonicalize.py b/sqlglot/optimizer/canonicalize.py
@@ -4,10 +4,12 @@
 import typing as t
 
 from sqlglot import exp
+from sqlglot.dialects.dialect import Dialect, DialectType
 from sqlglot.helper import is_date_unit, is_iso_date, is_iso_datetime
+from sqlglot.optimizer.annotate_types import TypeAnnotator
 
 
-def canonicalize(expression: exp.Expression) -> exp.Expression:
+def canonicalize(expression: exp.Expression, dialect: DialectType = None) -> exp.Expression:
     """Converts a sql expression into a standard form.
 
     This method relies on annotate_types because many of the
@@ -17,10 +19,12 @@ def canonicalize(expression: exp.Expression) -> exp.Expression:
         expression: The expression to canonicalize.
     """
 
+    dialect = Dialect.get_or_raise(dialect)
+
     def _canonicalize(expression: exp.Expression) -> exp.Expression:
         expression = add_text_to_concat(expression)
         expression = replace_date_funcs(expression)
-        expression = coerce_type(expression)
+        expression = coerce_type(expression, dialect.PROMOTE_TO_INFERRED_DATETIME_TYPE)
         expression = remove_redundant_casts(expression)
         expression = ensure_bools(expression, _replace_int_predicate)
         expression = remove_ascending_order(expression)
@@ -68,11 +72,11 @@ def replace_date_funcs(node: exp.Expression) -> exp.Expression:
 )
 
 
-def coerce_type(node: exp.Expression) -> exp.Expression:
+def coerce_type(node: exp.Expression, promote_to_inferred_datetime_type: bool) -> exp.Expression:
     if isinstance(node, COERCIBLE_DATE_OPS):
-        _coerce_date(node.left, node.right)
+        _coerce_date(node.left, node.right, promote_to_inferred_datetime_type)
     elif isinstance(node, exp.Between):
-        _coerce_date(node.this, node.args["low"])
+        _coerce_date(node.this, node.args["low"], promote_to_inferred_datetime_type)
     elif isinstance(node, exp.Extract) and not node.expression.type.is_type(
         *exp.DataType.TEMPORAL_TYPES
     ):
@@ -128,17 +132,48 @@ 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
-            and b.type
-            and b.type.this in exp.DataType.TEXT_TYPES
+            not a_type
+            or a_type.this not in exp.DataType.TEMPORAL_TYPES
+            or not b.type
+            or b.type.this not in exp.DataType.TEXT_TYPES
         ):
-            _replace_cast(b, exp.DataType.Type.DATETIME)
+            continue
+
+        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
+                else:
+                    b_type = a_type.this
+            else:
+                # If b is not a datetime string, we conservatively promote it to a DATETIME,
+                # in order to ensure there are no surprising truncations due to downcasting
+                b_type = exp.DataType.Type.DATETIME
+
+            target_type = (
+                b_type if b_type in TypeAnnotator.COERCES_TO.get(a_type.this, {}) else a_type
+            )
+        else:
+            target_type = a_type
+
+        if target_type != a_type:
+            _replace_cast(a, target_type)
+
+        _replace_cast(b, target_type)
 
 
 def _coerce_timeunit_arg(arg: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.Expression:
@@ -168,7 +203,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
@@ -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,14 +102,25 @@ 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";
 
 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";
+
+# dialect: mysql
+SELECT `t`.`d` < '2023-01-01 00:00:01' AS `col` FROM `temporal` AS `t`;
+SELECT CAST(`t`.`d` AS DATETIME) < CAST('2023-01-01 00:00:01' AS DATETIME) AS `col` FROM `temporal` AS `t`;
+
+# dialect: mysql
+SELECT CAST(`t`.`some_col` AS DATE) < CAST(`t`.`other_col` AS CHAR) AS `col` FROM `other_table` AS `t`;
+SELECT CAST(CAST(`t`.`some_col` AS DATE) AS DATETIME) < CAST(CAST(`t`.`other_col` AS CHAR) AS DATETIME) AS `col` FROM `other_table` 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
@@ -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"
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
@@ -132,7 +132,6 @@ def check_file(
         func,
         pretty=False,
         execute=False,
-        set_dialect=False,
         only=None,
         **kwargs,
     ):
@@ -158,7 +157,7 @@ def check_file(
                         validate_qualify_columns
                     )
 
-                if set_dialect and dialect:
+                if dialect:
                     func_kwargs["dialect"] = dialect
 
                 future = pool.submit(parse_and_optimize, func, sql, dialect, **func_kwargs)
@@ -207,7 +206,6 @@ def test_optimize(self, logger):
             pretty=True,
             execute=True,
             schema=schema,
-            set_dialect=True,
         )
 
     def test_isolate_table_selects(self):
@@ -235,7 +233,6 @@ def test_qualify_tables(self):
             optimizer.qualify_tables.qualify_tables,
             db="db",
             catalog="c",
-            set_dialect=True,
         )
 
     def test_normalize(self):
@@ -446,11 +443,8 @@ def test_qualify_columns(self, logger):
             qualify_columns,
             execute=True,
             schema=self.schema,
-            set_dialect=True,
-        )
-        self.check_file(
-            "qualify_columns_ddl", qualify_columns, schema=self.schema, set_dialect=True
         )
+        self.check_file("qualify_columns_ddl", qualify_columns, schema=self.schema)
 
     def test_qualify_columns__with_invisible(self):
         schema = MappingSchema(self.schema, {"x": {"a"}, "y": {"b"}, "z": {"b"}})
@@ -475,7 +469,6 @@ def test_normalize_identifiers(self):
         self.check_file(
             "normalize_identifiers",
             optimizer.normalize_identifiers.normalize_identifiers,
-            set_dialect=True,
         )
 
         self.assertEqual(optimizer.normalize_identifiers.normalize_identifiers("a%").sql(), '"a%"')
@@ -484,14 +477,13 @@ def test_quote_identifiers(self):
         self.check_file(
             "quote_identifiers",
             optimizer.qualify_columns.quote_identifiers,
-            set_dialect=True,
         )
 
     def test_pushdown_projection(self):
         self.check_file("pushdown_projections", pushdown_projections, schema=self.schema)
 
     def test_simplify(self):
-        self.check_file("simplify", simplify, set_dialect=True)
+        self.check_file("simplify", simplify)
 
         expression = parse_one("SELECT a, c, b FROM table1 WHERE 1 = 1")
         self.assertEqual(simplify(simplify(expression.find(exp.Where))).sql(), "WHERE TRUE")
