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

Open
fivetran-amrutabhimsenayachit wants to merge 2 commits into
base: main
Choose a base branch
from
+127 −1
Open

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
2 commits
Select commit Hold shift + click to select a range
50c6dda
perf(optimizer): memoize UNNEST/EXPLODE column-type resolution [CLAUD…
dvryaboy Jul 3, 2026
7a265fb
feat(typing): add DISTINCT/ALL/OVER fixtures and parser fix for REGR_…
fivetran-amrutabhimsenayachit Jul 2, 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 6b0ede to af84e3
17 changes: 17 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,16 @@ class DatabricksParser(SparkParser):
"CURDATE": lambda self: self._parse_curdate(),
}

FUNCTION_PARSERS = {
**SparkParser.FUNCTION_PARSERS,
"REGR_AVGX": lambda self: self._parse_quantile_function(exp.RegrAvgx),

@geooo109 geooo109 Jul 6, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did you verify where the DISTINCT is applied to? The default _parse_quantile_function always applies it to the first arg, but here we also have the second arg (x) as a candidate, right?

"REGR_AVGY": lambda self: self._parse_quantile_function(exp.RegrAvgy),
"REGR_COUNT": lambda self: self._parse_regr_count(),
"REGR_INTERCEPT": lambda self: self._parse_quantile_function(exp.RegrIntercept),
"REGR_R2": lambda self: self._parse_quantile_function(exp.RegrR2),
"REGR_SLOPE": lambda self: self._parse_quantile_function(exp.RegrSlope),
Comment on lines +39 to +41

@geooo109 geooo109 Jul 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did you check for these 3 functions for databricks where the DISTINCT is applied to?

SELECT
  regr_intercept(DISTINCT y, x) AS distinct_result,
  regr_intercept(y, x)          AS all_result
FROM VALUES
  (1, 10),
  (1, 20), <- duplicate y
  (2, 30),
  (3, 40)
AS T(y, x);
> 
distinct_result	all_result
-2.220446049250313e-16	-2.220446049250313e-16

SELECT
  regr_intercept(DISTINCT y, x) AS distinct_result,
  regr_intercept(y, x)          AS all_result
FROM VALUES
  (1, 10),
  (1, 10), <- duplicate pair
  (2, 30),
  (3, 40)
AS T(y, x);
>
distinct_result	all_result
0.28571428571428537	0.3333333333333335

SELECT
  regr_intercept(DISTINCT y, x) AS distinct_result,
  regr_intercept(y, x)          AS all_result
FROM VALUES
  (1, 10),
  (5, 10), <- duplicate x
  (2, 30),
  (3, 40)
AS T(y, x);
>
distinct_result	all_result
3	3

As it seems, databricks deduplicates based on the pair of (y, x), which is different from REGR_AVGX and REGR_AVGY.

Let's verify for the rest of the REGR_* funcs

}

FACTOR = {
**SparkParser.FACTOR,
TokenType.COLON: exp.JSONExtract,
Expand All @@ -55,6 +65,13 @@ def _parse_curdate(self) -> exp.CurrentDate:
self._match_r_paren()
return self.expression(exp.CurrentDate())

def _parse_regr_count(self) -> exp.RegrCount:
if self._match(TokenType.DISTINCT):
return exp.RegrCount(this=exp.Distinct(expressions=self._parse_function_args()))
self._match(TokenType.ALL)
args = self._parse_function_args()
return self.expression(exp.RegrCount(this=seq_get(args, 0), expression=seq_get(args, 1)))
Comment on lines +68 to +73

@geooo109 geooo109 Jul 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can make the expression arg false to cover the DISTINCT case and remove this method, the parsing will be automated, we follow this logic for other functions such as PercentileCont.


def _parse_cluster_property(self):
if self._match_texts(("AUTO", "NONE")):
return self.expression(exp.ClusterProperty(this=self._prev.text.upper()))
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 @@ -4247,6 +4247,10 @@ DOUBLE;
REGR_AVGX(tbl.decfloat_col, tbl.decfloat_col);
DECFLOAT;

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

# dialect: snowflake
REGR_AVGY(tbl.double_col, tbl.double_col);
DOUBLE;
Expand All @@ -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