feat(optimizer)!: annotate compress for mysql

sqlglot/typing/mysql.py

@@ -38,6 +38,54 @@ def _annotate_regexp_replace(self: TypeAnnotator, expression: exp.RegexpReplace)
     return self._set_type(expression, exp.DType.LONGBLOB if has_binary else exp.DType.LONGTEXT)
 
 
+COMPRESS_VARBINARY_TYPES = {
+    exp.DType.BINARY,
+    exp.DType.VARBINARY,
+    exp.DType.TINYBLOB,
+    exp.DType.ENUM,
+    exp.DType.INT,
+    exp.DType.BIGINT,
+    exp.DType.DECIMAL,
+    exp.DType.DOUBLE,
+    exp.DType.DATE,
+    exp.DType.DATETIME,
+}
+
+
+def _annotate_compress(self: TypeAnnotator, expression: exp.Compress) -> exp.Expr:
+    this = expression.this
+
+    if this.is_type(exp.DType.TINYTEXT):
+        return self._set_type(expression, exp.DType.BLOB)
+
+    if this.is_type(*exp.DataType.TEXT_TYPES):
+        return self._set_type(
+            expression,
+            exp.DType.VARBINARY
+            if this.is_type(
+                exp.DType.CHAR,
+                exp.DType.NCHAR,
+                exp.DType.VARCHAR,
+                exp.DType.NVARCHAR,
+                exp.DType.NAME,
+            )
+            else exp.DType.LONGBLOB,
+        )
+
+    if this.is_type(*COMPRESS_VARBINARY_TYPES):
+        return self._set_type(expression, exp.DType.VARBINARY)
+
+    if this.is_type(
+        exp.DType.BLOB,
+        exp.DType.MEDIUMBLOB,
+        exp.DType.LONGBLOB,
+        exp.DType.JSON,
+    ):
+        return self._set_type(expression, exp.DType.LONGBLOB)
+
+    return self._set_type(expression, exp.DType.UNKNOWN)
+
+
 EXPRESSION_METADATA = {
     **EXPRESSION_METADATA,
     **{
@@ -114,4 +162,5 @@ def _annotate_regexp_replace(self: TypeAnnotator, expression: exp.RegexpReplace)
     exp.Reverse: {"annotator": _annotate_reverse},
     exp.Trunc: {"annotator": _annotate_truncate},
     exp.RegexpReplace: {"annotator": _annotate_regexp_replace},
+    exp.Compress: {"annotator": _annotate_compress},
 }

tests/fixtures/optimizer/annotate_functions.sql

@@ -6739,6 +6739,81 @@ DOUBLE;
 MATCH(tbl.str_col) AGAINST ('database' WITH QUERY EXPANSION);
 DOUBLE;
 
+# dialect: mysql
+COMPRESS(CAST('test' AS TEXT));
+LONGBLOB;
+
+# dialect: mysql
+COMPRESS(CAST('test' AS MEDIUMTEXT));
+LONGBLOB;
+
+# dialect: mysql
+COMPRESS(CAST('test' AS LONGTEXT));
+LONGBLOB;
+
+# dialect: mysql
+COMPRESS(CAST('test' AS BLOB));
+LONGBLOB;
+
+# dialect: mysql
+COMPRESS(CAST('test' AS MEDIUMBLOB));
+LONGBLOB;
+
+# dialect: mysql
+COMPRESS(CAST('test' AS LONGBLOB));
+LONGBLOB;
+
+# dialect: mysql
+COMPRESS(CAST('{}' AS JSON));
+LONGBLOB;
+
+# dialect: mysql
+COMPRESS(CAST('test' AS TINYTEXT));
+BLOB;
+
+# dialect: mysql
+COMPRESS(CAST('test' AS CHAR));
+VARBINARY;
+
+# dialect: mysql
+COMPRESS(CAST('test' AS VARCHAR));
+VARBINARY;
+
+# dialect: mysql
+COMPRESS(CAST('test' AS BINARY));
+VARBINARY;
+
+# dialect: mysql
+COMPRESS(CAST('test' AS VARBINARY));
+VARBINARY;
+
+# dialect: mysql
+COMPRESS(CAST('test' AS TINYBLOB));
+VARBINARY;
+
+# dialect: mysql
+COMPRESS(CAST(1 AS INT));
+VARBINARY;
+
+# dialect: mysql
+COMPRESS(CAST(1 AS BIGINT));
+VARBINARY;
+
+# dialect: mysql
+COMPRESS(CAST(1.5 AS DECIMAL));
+VARBINARY;
+
+# dialect: mysql
+COMPRESS(CAST(1.5 AS DOUBLE));
+VARBINARY;
+
+# dialect: mysql
+COMPRESS(CAST('2024-01-01' AS DATE));
+VARBINARY;
+
+# dialect: mysql
+COMPRESS(CAST('2024-01-01 12:00:00' AS DATETIME));
+VARBINARY;
 --------------------------------------
 -- DuckDB
 --------------------------------------