feat(snowflake): transpile CORR with NaN-->NULL (#6619)

* Transpile CORR with NaN-->NULL

* Cleanup and address PR feedback

Author: treysp

sqlglot/dialects/duckdb.py

@@ -871,6 +871,21 @@ def _regr_val_sql(
     )
 
 
+def _maybe_corr_null_to_false(
+    expression: t.Union[exp.Filter, exp.Window, exp.Corr],
+) -> t.Optional[t.Union[exp.Filter, exp.Window, exp.Corr]]:
+    expr = expression.copy()
+    corr = expr
+    while isinstance(corr, (exp.Window, exp.Filter)):
+        corr = corr.this
+
+    if not isinstance(corr, exp.Corr) or not corr.args.get("null_on_zero_variance"):
+        return None
+
+    corr.set("null_on_zero_variance", False)
+    return expr
+
+
 def _date_from_parts_sql(self, expression: exp.DateFromParts) -> str:
     """
     Snowflake's DATE_FROM_PARTS allows out-of-range values for the month and day input.
@@ -1348,6 +1363,7 @@ class Generator(generator.Generator):
             exp.BitwiseOrAgg: _bitwise_agg_sql,
             exp.BitwiseXorAgg: _bitwise_agg_sql,
             exp.CommentColumnConstraint: no_comment_column_constraint_sql,
+            exp.Corr: lambda self, e: self._corr_sql(e),
             exp.CosineDistance: rename_func("LIST_COSINE_DISTANCE"),
             exp.CurrentTime: lambda *_: "CURRENT_TIME",
             exp.CurrentTimestamp: lambda *_: "CURRENT_TIMESTAMP",
@@ -2542,3 +2558,37 @@ def bitwisenot_sql(self, expression: exp.BitwiseNot) -> str:
             result_sql = f"~{self.sql(expression, 'this')}"
 
             return _gen_with_cast_to_blob(self, expression, result_sql)
+
+        def window_sql(self, expression: exp.Window) -> str:
+            this = expression.this
+            if isinstance(this, exp.Corr) or (
+                isinstance(this, exp.Filter) and isinstance(this.this, exp.Corr)
+            ):
+                return self._corr_sql(expression)
+
+            return super().window_sql(expression)
+
+        def filter_sql(self, expression: exp.Filter) -> str:
+            if isinstance(expression.this, exp.Corr):
+                return self._corr_sql(expression)
+
+            return super().filter_sql(expression)
+
+        def _corr_sql(
+            self,
+            expression: t.Union[exp.Filter, exp.Window, exp.Corr],
+        ) -> str:
+            if isinstance(expression, exp.Corr) and not expression.args.get(
+                "null_on_zero_variance"
+            ):
+                return self.func("CORR", expression.this, expression.expression)
+
+            corr_expr = _maybe_corr_null_to_false(expression)
+            if corr_expr is None:
+                if isinstance(expression, exp.Window):
+                    return super().window_sql(expression)
+                if isinstance(expression, exp.Filter):
+                    return super().filter_sql(expression)
+                corr_expr = expression  # make mypy happy
+
+            return self.sql(exp.case().when(exp.IsNan(this=corr_expr), exp.null()).else_(corr_expr))

sqlglot/dialects/snowflake.py

@@ -769,6 +769,11 @@ class Parser(parser.Parser):
             "BIT_XORAGG": exp.BitwiseXorAgg.from_arg_list,
             "BITMAP_OR_AGG": exp.BitmapOrAgg.from_arg_list,
             "BOOLXOR": _build_bitwise(exp.Xor, "BOOLXOR"),
+            "CORR": lambda args: exp.Corr(
+                this=seq_get(args, 0),
+                expression=seq_get(args, 1),
+                null_on_zero_variance=True,
+            ),
             "DATE": _build_datetime("DATE", exp.DataType.Type.DATE),
             "DATEFROMPARTS": _build_date_from_parts,
             "DATE_FROM_PARTS": _build_date_from_parts,

sqlglot/expressions.py

@@ -8300,8 +8300,11 @@ class Upper(Func):
     _sql_names = ["UPPER", "UCASE"]
 
 
-class Corr(AggFunc):
-    arg_types = {"this": True, "expression": True}
+class Corr(Binary, AggFunc):
+    # Correlation divides by variance(column). If a column has 0 variance, the denominator
+    # is 0 - some dialects return NaN (DuckDB) while others return NULL (Snowflake).
+    # `null_on_zero_variance` is set to True at parse time for dialects that return NULL.
+    arg_types = {"this": True, "expression": True, "null_on_zero_variance": False}
 
 
 # https://docs.oracle.com/en/database/oracle/oracle-database/19/sqlrf/CUME_DIST.html

tests/dialects/test_postgres.py

@@ -1398,6 +1398,36 @@ def test_variance(self):
             },
         )
 
+    def test_corr(self):
+        self.validate_all(
+            "SELECT CORR(a, b)",
+            write={
+                "duckdb": "SELECT CORR(a, b)",
+                "postgres": "SELECT CORR(a, b)",
+            },
+        )
+        self.validate_all(
+            "SELECT CORR(a, b) OVER (PARTITION BY c)",
+            write={
+                "duckdb": "SELECT CORR(a, b) OVER (PARTITION BY c)",
+                "postgres": "SELECT CORR(a, b) OVER (PARTITION BY c)",
+            },
+        )
+        self.validate_all(
+            "SELECT CORR(a, b) FILTER(WHERE c > 0)",
+            write={
+                "duckdb": "SELECT CORR(a, b) FILTER(WHERE c > 0)",
+                "postgres": "SELECT CORR(a, b) FILTER(WHERE c > 0)",
+            },
+        )
+        self.validate_all(
+            "SELECT CORR(a, b) FILTER(WHERE c > 0) OVER (PARTITION BY d)",
+            write={
+                "duckdb": "SELECT CORR(a, b) FILTER(WHERE c > 0) OVER (PARTITION BY d)",
+                "postgres": "SELECT CORR(a, b) FILTER(WHERE c > 0) OVER (PARTITION BY d)",
+            },
+        )
+
     def test_regexp_binary(self):
         """See https://github.com/tobymao/sqlglot/pull/2404 for details."""
         self.assertIsInstance(self.parse_one("'thomas' ~ '.*thomas.*'"), exp.Binary)

tests/dialects/test_snowflake.py

@@ -1170,7 +1170,6 @@ def test_snowflake(self):
                 },
             )
             for func in (
-                "CORR",
                 "COVAR_POP",
                 "COVAR_SAMP",
             ):
@@ -4559,6 +4558,45 @@ def test_ceil(self):
             },
         )
 
+    def test_corr(self):
+        self.validate_all(
+            "SELECT CORR(a, b)",
+            read={
+                "snowflake": "SELECT CORR(a, b)",
+                "postgres": "SELECT CORR(a, b)",
+            },
+            write={
+                "snowflake": "SELECT CORR(a, b)",
+                "postgres": "SELECT CORR(a, b)",
+                "duckdb": "SELECT CASE WHEN ISNAN(CORR(a, b)) THEN NULL ELSE CORR(a, b) END",
+            },
+        )
+        self.validate_all(
+            "SELECT CORR(a, b) OVER (PARTITION BY c)",
+            read={
+                "snowflake": "SELECT CORR(a, b) OVER (PARTITION BY c)",
+                "postgres": "SELECT CORR(a, b) OVER (PARTITION BY c)",
+            },
+            write={
+                "snowflake": "SELECT CORR(a, b) OVER (PARTITION BY c)",
+                "postgres": "SELECT CORR(a, b) OVER (PARTITION BY c)",
+                "duckdb": "SELECT CASE WHEN ISNAN(CORR(a, b) OVER (PARTITION BY c)) THEN NULL ELSE CORR(a, b) OVER (PARTITION BY c) END",
+            },
+        )
+
+        self.validate_all(
+            "SELECT CORR(a, b) FILTER(WHERE c > 0)",
+            write={
+                "duckdb": "SELECT CASE WHEN ISNAN(CORR(a, b) FILTER(WHERE c > 0)) THEN NULL ELSE CORR(a, b) FILTER(WHERE c > 0) END",
+            },
+        )
+        self.validate_all(
+            "SELECT CORR(a, b) FILTER(WHERE c > 0) OVER (PARTITION BY d)",
+            write={
+                "duckdb": "SELECT CASE WHEN ISNAN(CORR(a, b) FILTER(WHERE c > 0) OVER (PARTITION BY d)) THEN NULL ELSE CORR(a, b) FILTER(WHERE c > 0) OVER (PARTITION BY d) END",
+            },
+        )
+
     def test_encryption_functions(self):
         # ENCRYPT
         self.validate_identity("ENCRYPT(value, 'passphrase')")