diff --git a/sqlglot-integration-tests b/sqlglot-integration-tests index 8d9d5e1430..bfd6be8f22 160000 --- a/sqlglot-integration-tests +++ b/sqlglot-integration-tests @@ -1 +1 @@ -Subproject commit 8d9d5e14302330cecf5fb784f983862d1f893db9 +Subproject commit bfd6be8f221dfa821b33893095be264ab31445ab diff --git a/sqlglot/generators/databricks.py b/sqlglot/generators/databricks.py index 5ec44b1c22..819c17bd1c 100644 --- a/sqlglot/generators/databricks.py +++ b/sqlglot/generators/databricks.py @@ -109,13 +109,21 @@ def uniform_sql(self, expression: exp.Uniform) -> str: return self.func("UNIFORM", expression.this, expression.expression, seed) - def regravgx_sql(self, expression: exp.RegrAvgx) -> str: + def _regr_sql(self, expression: exp.RegrAvgx | exp.RegrSxx | exp.RegrSyy) -> str: + name = expression.sql_name() x = expression.expression if isinstance(x, exp.Distinct): - return self.func( - "REGR_AVGX", exp.Distinct(expressions=[expression.this]), *x.expressions - ) - return self.func("REGR_AVGX", expression.this, x) + return self.func(name, exp.Distinct(expressions=[expression.this]), *x.expressions) + return self.func(name, expression.this, x) + + def regravgx_sql(self, expression: exp.RegrAvgx) -> str: + return self._regr_sql(expression) + + def regrsxx_sql(self, expression: exp.RegrSxx) -> str: + return self._regr_sql(expression) + + def regrsyy_sql(self, expression: exp.RegrSyy) -> str: + return self._regr_sql(expression) def clusterproperty_sql(self, expression): this = self.sql(expression, "this") or f"({self.expressions(expression, flat=True)})" diff --git a/sqlglot/parsers/databricks.py b/sqlglot/parsers/databricks.py index 6f5b4b202e..a5f8ce4865 100644 --- a/sqlglot/parsers/databricks.py +++ b/sqlglot/parsers/databricks.py @@ -35,6 +35,9 @@ class DatabricksParser(SparkParser): **SparkParser.FUNCTION_PARSERS, "REGR_AVGX": lambda self: self._parse_distinct_arg_function(exp.RegrAvgx, distinct_index=1), "REGR_AVGY": lambda self: self._parse_distinct_arg_function(exp.RegrAvgy), + "REGR_SXX": lambda self: self._parse_distinct_arg_function(exp.RegrSxx, distinct_index=1), + "REGR_SXY": lambda self: self._parse_distinct_arg_function(exp.RegrSxy), + "REGR_SYY": lambda self: self._parse_distinct_arg_function(exp.RegrSyy, distinct_index=1), } FACTOR = { diff --git a/sqlglot/typing/databricks.py b/sqlglot/typing/databricks.py index bf2d4faf56..ee898d33f4 100644 --- a/sqlglot/typing/databricks.py +++ b/sqlglot/typing/databricks.py @@ -13,6 +13,9 @@ exp.RegrIntercept, exp.RegrR2, exp.RegrSlope, + exp.RegrSxx, + exp.RegrSxy, + exp.RegrSyy, } }, **{ diff --git a/tests/fixtures/optimizer/annotate_functions.sql b/tests/fixtures/optimizer/annotate_functions.sql index 5b09906732..817eee42b6 100644 --- a/tests/fixtures/optimizer/annotate_functions.sql +++ b/tests/fixtures/optimizer/annotate_functions.sql @@ -4439,6 +4439,66 @@ DOUBLE; REGR_SYY(tbl.decfloat_col, tbl.decfloat_col); DECFLOAT; +# dialect: databricks +REGR_SXX(tbl.double_col, tbl.double_col); +DOUBLE; + +# dialect: databricks +REGR_SXX(tbl.int_col, tbl.int_col); +DOUBLE; + +# dialect: databricks +REGR_SXX(ALL tbl.double_col, tbl.double_col); +DOUBLE; + +# dialect: databricks +REGR_SXX(DISTINCT tbl.double_col, tbl.double_col); +DOUBLE; + +# dialect: databricks +REGR_SXX(tbl.double_col, tbl.double_col) OVER (PARTITION BY 1); +DOUBLE; + +# dialect: databricks +REGR_SXY(tbl.double_col, tbl.double_col); +DOUBLE; + +# dialect: databricks +REGR_SXY(tbl.int_col, tbl.int_col); +DOUBLE; + +# dialect: databricks +REGR_SXY(ALL tbl.double_col, tbl.double_col); +DOUBLE; + +# dialect: databricks +REGR_SXY(DISTINCT tbl.double_col, tbl.double_col); +DOUBLE; + +# dialect: databricks +REGR_SXY(tbl.double_col, tbl.double_col) OVER (PARTITION BY 1); +DOUBLE; + +# dialect: databricks +REGR_SYY(tbl.double_col, tbl.double_col); +DOUBLE; + +# dialect: databricks +REGR_SYY(tbl.int_col, tbl.int_col); +DOUBLE; + +# dialect: databricks +REGR_SYY(ALL tbl.double_col, tbl.double_col); +DOUBLE; + +# dialect: databricks +REGR_SYY(DISTINCT tbl.double_col, tbl.double_col); +DOUBLE; + +# dialect: databricks +REGR_SYY(tbl.double_col, tbl.double_col) OVER (PARTITION BY 1); +DOUBLE; + # dialect: snowflake REGR_SLOPE(tbl.double_col, tbl.double_col); DOUBLE;