Skip to content

fix(optimizer)!: annotate type for databricks REGR_AVGY, REGR_COUNT, REGR_INTERCEPT, REGR_R2, REGR_SLOPE #7820

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
fivetran-amrutabhimsenayachit merged 4 commits into from
Jul 8, 2026
Merged

fix(optimizer)!: annotate type for databricks REGR_AVGY, REGR_COUNT, REGR_INTERCEPT, REGR_R2, REGR_SLOPE #7820

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
627692c
perf(optimizer): memoize UNNEST/EXPLODE column-type resolution [CLAUD…
dvryaboy Jul 3, 2026
5c05391
feat(typing): add DISTINCT/ALL/OVER fixtures and parser fix for REGR_…
fivetran-amrutabhimsenayachit Jul 2, 2026
69d59e1
fix: update aggregate expressions and databricks parser
fivetran-amrutabhimsenayachit Jul 7, 2026
02cbca0
fix: update databricks parser to have a generic function for regr_avgy/x
fivetran-amrutabhimsenayachit Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion sqlglot-integration-tests
Open in desktop
Submodule sqlglot-integration-tests updated from 1bcd84 to eaf74a
8 changes: 4 additions & 4 deletions sqlglot/expressions/aggregate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -256,19 +256,19 @@ class RegrAvgy(Expression, AggFunc):


class RegrCount(Expression, AggFunc):
arg_types = {"this": True, "expression": True}
arg_types = {"this": True, "expression": False}


class RegrIntercept(Expression, AggFunc):
arg_types = {"this": True, "expression": True}
arg_types = {"this": True, "expression": False}


class RegrR2(Expression, AggFunc):
arg_types = {"this": True, "expression": True}
arg_types = {"this": True, "expression": False}


class RegrSlope(Expression, AggFunc):
arg_types = {"this": True, "expression": True}
arg_types = {"this": True, "expression": False}


class RegrSxx(Expression, AggFunc):
Expand Down
8 changes: 8 additions & 0 deletions sqlglot/generators/databricks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,14 @@ 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:
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)

def clusterproperty_sql(self, expression):
this = self.sql(expression, "this") or f"({self.expressions(expression, flat=True)})"
return f"CLUSTER BY {this}"
6 changes: 6 additions & 0 deletions sqlglot/parsers/databricks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,12 @@ class DatabricksParser(SparkParser):
"CURDATE": lambda self: self._parse_curdate(),
}

FUNCTION_PARSERS = {
**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),
}

FACTOR = {
**SparkParser.FACTOR,
TokenType.COLON: exp.JSONExtract,
Expand Down
20 changes: 10 additions & 10 deletions sqlglot/parsers/hive.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,8 +63,8 @@ class HiveParser(parser.Parser):

FUNCTION_PARSERS = {
**parser.Parser.FUNCTION_PARSERS,
"PERCENTILE": lambda self: self._parse_quantile_function(exp.Quantile),
"PERCENTILE_APPROX": lambda self: self._parse_quantile_function(exp.ApproxQuantile),
"PERCENTILE": lambda self: self._parse_distinct_arg_function(exp.Quantile),
"PERCENTILE_APPROX": lambda self: self._parse_distinct_arg_function(exp.ApproxQuantile),
}

FUNCTIONS = {
Expand Down Expand Up @@ -178,19 +178,19 @@ def _parse_transform(self) -> exp.Transform | exp.QueryTransform | None:
)
)

def _parse_quantile_function(self, func: type[F]) -> F:
if self._match(TokenType.DISTINCT):
first_arg: exp.Expr | None = self.expression(
exp.Distinct(expressions=[self._parse_lambda()])
)
else:
def _parse_distinct_arg_function(self, func: type[F], distinct_index: int = 0) -> F:
is_distinct = self._match(TokenType.DISTINCT)
if not is_distinct:
self._match(TokenType.ALL)
first_arg = self._parse_lambda()

args = [first_arg]
args = [self._parse_lambda()]
if self._match(TokenType.COMMA):
args.extend(self._parse_function_args())

target = seq_get(args, distinct_index)
if is_distinct and target:
args[distinct_index] = self.expression(exp.Distinct(expressions=[target]))

return func.from_arg_list(args)

def _parse_types(
Expand Down
2 changes: 1 addition & 1 deletion sqlglot/parsers/spark2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ class Spark2Parser(HiveParser):

FUNCTION_PARSERS = {
**HiveParser.FUNCTION_PARSERS,
"APPROX_PERCENTILE": lambda self: self._parse_quantile_function(exp.ApproxQuantile),
"APPROX_PERCENTILE": lambda self: self._parse_distinct_arg_function(exp.ApproxQuantile),
"BROADCAST": lambda self: self._parse_join_hint("BROADCAST"),
"BROADCASTJOIN": lambda self: self._parse_join_hint("BROADCASTJOIN"),
"MAPJOIN": lambda self: self._parse_join_hint("MAPJOIN"),
Expand Down
5 changes: 5 additions & 0 deletions sqlglot/typing/databricks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,10 @@
exp_type: {"returns": exp.DType.DOUBLE}
for exp_type in {
exp.RegrAvgx,
exp.RegrAvgy,
exp.RegrIntercept,
exp.RegrR2,
exp.RegrSlope,
}
},
**{
Expand All @@ -19,6 +23,7 @@
}
},
exp.RegexpSubstr: {"returns": exp.DType.VARCHAR},
exp.RegrCount: {"returns": exp.DType.BIGINT},
exp.RegexpExtractAll: {
"annotator": lambda self, e: self._set_type(
e, exp.DataType.from_str("ARRAY<STRING>", dialect="databricks")
Expand Down
104 changes: 104 additions & 0 deletions tests/fixtures/optimizer/annotate_functions.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1051,6 +1051,10 @@ DOUBLE;
REGR_AVGX(ALL tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_AVGX(DISTINCT tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_AVGX(tbl.double_col, tbl.double_col) OVER (PARTITION BY 1);
DOUBLE;
Expand Down Expand Up @@ -4259,6 +4263,26 @@ DOUBLE;
REGR_AVGY(tbl.decfloat_col, tbl.decfloat_col);
DECFLOAT;

# dialect: databricks
REGR_AVGY(tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_AVGY(tbl.int_col, tbl.int_col);
DOUBLE;
Comment thread
geooo109 marked this conversation as resolved.

# dialect: databricks
REGR_AVGY(ALL tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_AVGY(DISTINCT tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_AVGY(tbl.double_col, tbl.double_col) OVER (PARTITION BY 1);
DOUBLE;

# dialect: snowflake
REGR_COUNT(tbl.double_col, tbl.double_col);
DOUBLE;
Expand All @@ -4275,6 +4299,26 @@ DOUBLE;
REGR_COUNT(tbl.decfloat_col, tbl.decfloat_col);
DECFLOAT;

# dialect: databricks
REGR_COUNT(tbl.double_col, tbl.double_col);
BIGINT;

# dialect: databricks
REGR_COUNT(tbl.int_col, tbl.int_col);
BIGINT;
Comment thread
geooo109 marked this conversation as resolved.

# dialect: databricks
REGR_COUNT(ALL tbl.double_col, tbl.double_col);
BIGINT;

# dialect: databricks
REGR_COUNT(DISTINCT tbl.double_col, tbl.double_col);
BIGINT;

# dialect: databricks
REGR_COUNT(tbl.double_col, tbl.double_col) OVER (PARTITION BY 1);
BIGINT;

# dialect: snowflake
REGR_INTERCEPT(tbl.double_col, tbl.double_col);
DOUBLE;
Expand All @@ -4291,6 +4335,26 @@ DOUBLE;
REGR_INTERCEPT(tbl.decfloat_col, tbl.decfloat_col);
DECFLOAT;

# dialect: databricks
REGR_INTERCEPT(tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_INTERCEPT(tbl.int_col, tbl.int_col);
DOUBLE;
Comment thread
geooo109 marked this conversation as resolved.

# dialect: databricks
REGR_INTERCEPT(ALL tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_INTERCEPT(DISTINCT tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_INTERCEPT(tbl.double_col, tbl.double_col) OVER (PARTITION BY 1);
DOUBLE;

# dialect: snowflake
REGR_R2(tbl.double_col, tbl.double_col);
DOUBLE;
Expand All @@ -4307,6 +4371,26 @@ DOUBLE;
REGR_R2(tbl.decfloat_col, tbl.decfloat_col);
DECFLOAT;

# dialect: databricks
REGR_R2(tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_R2(tbl.int_col, tbl.int_col);
DOUBLE;
Comment thread
geooo109 marked this conversation as resolved.

# dialect: databricks
REGR_R2(ALL tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_R2(DISTINCT tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_R2(tbl.double_col, tbl.double_col) OVER (PARTITION BY 1);
DOUBLE;

# dialect: snowflake
REGR_SXX(tbl.double_col, tbl.double_col);
DOUBLE;
Expand Down Expand Up @@ -4371,6 +4455,26 @@ DOUBLE;
REGR_SLOPE(tbl.decfloat_col, tbl.decfloat_col);
DECFLOAT;

# dialect: databricks
REGR_SLOPE(tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_SLOPE(tbl.int_col, tbl.int_col);
DOUBLE;
Comment thread
geooo109 marked this conversation as resolved.

# dialect: databricks
REGR_SLOPE(ALL tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_SLOPE(DISTINCT tbl.double_col, tbl.double_col);
DOUBLE;

# dialect: databricks
REGR_SLOPE(tbl.double_col, tbl.double_col) OVER (PARTITION BY 1);
DOUBLE;

# dialect: snowflake
REGR_VALX(NULL, 2.0);
DOUBLE;
Expand Down
Loading