Feat: improve transpilation of Doris' MONTHS_ADD (#3012)

tobymao · Feb 22, 2024 · d2e15ed · d2e15ed
1 parent 2a88e40
commit d2e15ed
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 1 deletion.
diff --git a/sqlglot/dataframe/sql/functions.py b/sqlglot/dataframe/sql/functions.py
@@ -592,7 +592,7 @@ def date_diff(end: ColumnOrName, start: ColumnOrName) -> Column:
 
 
 def add_months(start: ColumnOrName, months: t.Union[ColumnOrName, int]) -> Column:
-    return Column.invoke_anonymous_function(start, "ADD_MONTHS", months)
+    return Column.invoke_expression_over_column(start, expression.AddMonths, expression=months)
 
 
 def months_between(

diff --git a/sqlglot/dialects/doris.py b/sqlglot/dialects/doris.py
@@ -21,6 +21,7 @@ class Parser(MySQL.Parser):
             **MySQL.Parser.FUNCTIONS,
             "COLLECT_SET": exp.ArrayUniqueAgg.from_arg_list,
             "DATE_TRUNC": build_timestamp_trunc,
+            "MONTHS_ADD": exp.AddMonths.from_arg_list,
             "REGEXP": exp.RegexpLike.from_arg_list,
             "TO_DATE": exp.TsOrDsToDate.from_arg_list,
         }
@@ -41,6 +42,7 @@ class Generator(MySQL.Generator):
 
         TRANSFORMS = {
             **MySQL.Generator.TRANSFORMS,
+            exp.AddMonths: rename_func("MONTHS_ADD"),
             exp.ApproxDistinct: approx_count_distinct_sql,
             exp.ArgMax: rename_func("MAX_BY"),
             exp.ArgMin: rename_func("MIN_BY"),

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -5197,6 +5197,10 @@ class Month(Func):
     pass
 
 
+class AddMonths(Func):
+    arg_types = {"this": True, "expression": True}
+
+
 class Nvl2(Func):
     arg_types = {"this": True, "true": True, "false": False}
 

diff --git a/tests/dialects/test_doris.py b/tests/dialects/test_doris.py
@@ -26,6 +26,16 @@ def test_doris(self):
                 "doris": "SELECT ARRAY_SUM(x -> x * x, ARRAY(2, 3))",
             },
         )
+        self.validate_all(
+            "MONTHS_ADD(d, n)",
+            read={
+                "oracle": "ADD_MONTHS(d, n)",
+            },
+            write={
+                "doris": "MONTHS_ADD(d, n)",
+                "oracle": "ADD_MONTHS(d, n)",
+            },
+        )
 
     def test_identity(self):
         self.validate_identity("COALECSE(a, b, c, d)")

diff --git a/tests/test_expressions.py b/tests/test_expressions.py
@@ -645,6 +645,7 @@ def test_functions(self):
         self.assertIsInstance(parse_one("TO_HEX(foo)", read="bigquery"), exp.Hex)
         self.assertIsInstance(parse_one("TO_HEX(MD5(foo))", read="bigquery"), exp.MD5)
         self.assertIsInstance(parse_one("TRANSFORM(a, b)", read="spark"), exp.Transform)
+        self.assertIsInstance(parse_one("ADD_MONTHS(a, b)"), exp.AddMonths)
 
     def test_column(self):
         column = parse_one("a.b.c.d")