From 6fc5418c2cba5ac999c93e78269c8cf420a1492d Mon Sep 17 00:00:00 2001 From: geooo109 Date: Wed, 22 Apr 2026 12:55:32 +0300 Subject: [PATCH 01/13] fix(optimizer): unpivot annotate types for bq --- sqlglot-integration-tests | 2 +- sqlglot/optimizer/annotate_types.py | 53 ++++++++++++++++++++++++++++- 2 files changed, 53 insertions(+), 2 deletions(-) diff --git a/sqlglot-integration-tests b/sqlglot-integration-tests index 7281197267..ccfe6a98df 160000 --- a/sqlglot-integration-tests +++ b/sqlglot-integration-tests @@ -1 +1 @@ -Subproject commit 7281197267141c2802f8fc603fe3e1c2a27b9509 +Subproject commit ccfe6a98df29d2c811270861379d6e2f852a5241 diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py index 0e6b1c407c..852fb4b10f 100644 --- a/sqlglot/optimizer/annotate_types.py +++ b/sqlglot/optimizer/annotate_types.py @@ -330,6 +330,50 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: elif isinstance(expression, exp.Selectable): selects[name] = {s.alias_or_name: s.type for s in expression.selects if s.type} + for pivot in scope.pivots: + pivot_source = scope.sources.get(pivot.alias) + if not pivot_source: + continue + + inner_name = ( + pivot_source.name if isinstance(pivot_source, exp.Table) else pivot.alias + ) + col_types = dict(selects.get(inner_name, {})) + + if pivot.unpivot: + for field in pivot.fields: + field_col = field.this + if not isinstance(field_col, exp.Column) or not field.expressions: + continue + + first = field.expressions[0] + is_pivot_alias = isinstance(first, exp.PivotAlias) + + # FOR column type from the alias literal, or VARCHAR if no alias + if is_pivot_alias: + alias_node = first.args.get("alias") + if alias_node: + col_types[field_col.name] = alias_node.type + else: + col_types[field_col.name] = exp.DType.VARCHAR + + # Value column types from the IN source columns + src = first.this if is_pivot_alias else first + src_cols = src.expressions if isinstance(src, exp.Tuple) else [src] + for val_expr in pivot.expressions: + val_cols = ( + val_expr.expressions + if isinstance(val_expr, exp.Tuple) + else [val_expr] + ) + for val_col, src_col in zip(val_cols, src_cols): + src_type = col_types.get(src_col.output_name) + if src_type: + col_types[val_col.output_name] = src_type + + if col_types: + selects[pivot.alias] = col_types + self._scope_selects[scope] = selects return self._scope_selects[scope] @@ -419,7 +463,14 @@ def _annotate_expression( source_scope = source_scope.parent if isinstance(source, exp.Table): - self._set_type(expr, self.schema.get_column_type(source, expr)) + schema_type = self.schema.get_column_type(source, expr) + if schema_type.is_type(exp.DType.UNKNOWN) and source.args.get("pivots"): + pivot_type = ( + self._get_scope_selects(scope).get(expr.table, {}).get(expr.name) + ) + if pivot_type: + schema_type = pivot_type + self._set_type(expr, schema_type) elif source and source_scope: col_type = ( self._get_scope_selects(source_scope).get(expr.table, {}).get(expr.name) From cf10ba7070ecbde8432ef3f7e10ee2f674c145ec Mon Sep 17 00:00:00 2001 From: geooo109 Date: Thu, 23 Apr 2026 14:35:47 +0300 Subject: [PATCH 02/13] fix varchar --- sqlglot-integration-tests | 2 +- sqlglot/optimizer/annotate_types.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/sqlglot-integration-tests b/sqlglot-integration-tests index ccfe6a98df..ca99759fc5 160000 --- a/sqlglot-integration-tests +++ b/sqlglot-integration-tests @@ -1 +1 @@ -Subproject commit ccfe6a98df29d2c811270861379d6e2f852a5241 +Subproject commit ca99759fc52d5b307438b42cb43bdd4a818a2e38 diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py index 852fb4b10f..4c050cb654 100644 --- a/sqlglot/optimizer/annotate_types.py +++ b/sqlglot/optimizer/annotate_types.py @@ -355,7 +355,9 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: if alias_node: col_types[field_col.name] = alias_node.type else: - col_types[field_col.name] = exp.DType.VARCHAR + col_types[field_col.name] = exp.DataType.build( + "VARCHAR", dialect=self.dialect + ) # Value column types from the IN source columns src = first.this if is_pivot_alias else first From 43dfe68e967c99a37bde9ece52c1afe72a526b5c Mon Sep 17 00:00:00 2001 From: geooo109 Date: Thu, 23 Apr 2026 17:01:06 +0300 Subject: [PATCH 03/13] fix --- sqlglot/optimizer/annotate_types.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py index 4c050cb654..13f3be66ff 100644 --- a/sqlglot/optimizer/annotate_types.py +++ b/sqlglot/optimizer/annotate_types.py @@ -369,8 +369,10 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: else [val_expr] ) for val_col, src_col in zip(val_cols, src_cols): - src_type = col_types.get(src_col.output_name) - if src_type: + src_type = col_types.get(src_col.output_name) or src_col.type + if isinstance(src_type, exp.DataType) and not src_type.is_type( + exp.DType.UNKNOWN + ): col_types[val_col.output_name] = src_type if col_types: From 3fe5e29dc341b72327564eabd2751e975aa9a81e Mon Sep 17 00:00:00 2001 From: geooo109 Date: Thu, 23 Apr 2026 19:27:12 +0300 Subject: [PATCH 04/13] simplify --- sqlglot/optimizer/annotate_types.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py index 13f3be66ff..3dbdebeb19 100644 --- a/sqlglot/optimizer/annotate_types.py +++ b/sqlglot/optimizer/annotate_types.py @@ -343,10 +343,11 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: if pivot.unpivot: for field in pivot.fields: field_col = field.this - if not isinstance(field_col, exp.Column) or not field.expressions: + + first = seq_get(field.expressions, 0) + if not first: continue - first = field.expressions[0] is_pivot_alias = isinstance(first, exp.PivotAlias) # FOR column type from the alias literal, or VARCHAR if no alias From e4d315db93a0ab792209127a06cb41ad3bc8c064 Mon Sep 17 00:00:00 2001 From: geooo109 Date: Thu, 23 Apr 2026 19:43:37 +0300 Subject: [PATCH 05/13] refactor --- sqlglot/optimizer/annotate_types.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py index 3dbdebeb19..f3fe2cb076 100644 --- a/sqlglot/optimizer/annotate_types.py +++ b/sqlglot/optimizer/annotate_types.py @@ -345,23 +345,20 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: field_col = field.this first = seq_get(field.expressions, 0) - if not first: - continue - - is_pivot_alias = isinstance(first, exp.PivotAlias) # FOR column type from the alias literal, or VARCHAR if no alias - if is_pivot_alias: + if isinstance(first, exp.PivotAlias): alias_node = first.args.get("alias") if alias_node: col_types[field_col.name] = alias_node.type + src = first.this else: col_types[field_col.name] = exp.DataType.build( "VARCHAR", dialect=self.dialect ) + src = first # Value column types from the IN source columns - src = first.this if is_pivot_alias else first src_cols = src.expressions if isinstance(src, exp.Tuple) else [src] for val_expr in pivot.expressions: val_cols = ( From d8e8959cacbc8adade434d59168085266524bb5e Mon Sep 17 00:00:00 2001 From: geooo109 Date: Thu, 23 Apr 2026 20:14:14 +0300 Subject: [PATCH 06/13] remove copy --- sqlglot/optimizer/annotate_types.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py index f3fe2cb076..cc724a1e18 100644 --- a/sqlglot/optimizer/annotate_types.py +++ b/sqlglot/optimizer/annotate_types.py @@ -338,8 +338,10 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: inner_name = ( pivot_source.name if isinstance(pivot_source, exp.Table) else pivot.alias ) - col_types = dict(selects.get(inner_name, {})) + source_types = selects.get(inner_name, {}) + + col_types = {} if pivot.unpivot: for field in pivot.fields: field_col = field.this @@ -367,7 +369,7 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: else [val_expr] ) for val_col, src_col in zip(val_cols, src_cols): - src_type = col_types.get(src_col.output_name) or src_col.type + src_type = source_types.get(src_col.output_name) or src_col.type if isinstance(src_type, exp.DataType) and not src_type.is_type( exp.DType.UNKNOWN ): From dee01a94fa453162ec8994d72bbe40d15881c67a Mon Sep 17 00:00:00 2001 From: geooo109 Date: Fri, 24 Apr 2026 02:57:45 +0300 Subject: [PATCH 07/13] revert copy --- sqlglot/optimizer/annotate_types.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py index cc724a1e18..278e292759 100644 --- a/sqlglot/optimizer/annotate_types.py +++ b/sqlglot/optimizer/annotate_types.py @@ -339,9 +339,8 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: pivot_source.name if isinstance(pivot_source, exp.Table) else pivot.alias ) - source_types = selects.get(inner_name, {}) + col_types = selects.get(inner_name, {}).copy() - col_types = {} if pivot.unpivot: for field in pivot.fields: field_col = field.this @@ -369,7 +368,7 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: else [val_expr] ) for val_col, src_col in zip(val_cols, src_cols): - src_type = source_types.get(src_col.output_name) or src_col.type + src_type = col_types.get(src_col.output_name) or src_col.type if isinstance(src_type, exp.DataType) and not src_type.is_type( exp.DType.UNKNOWN ): From 0f27e514201261e64872a6426c8906dd36c1aef0 Mon Sep 17 00:00:00 2001 From: geooo109 Date: Fri, 24 Apr 2026 02:59:14 +0300 Subject: [PATCH 08/13] Sync w/ integration tests --- sqlglot-integration-tests | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sqlglot-integration-tests b/sqlglot-integration-tests index ca99759fc5..cd6302ffd3 160000 --- a/sqlglot-integration-tests +++ b/sqlglot-integration-tests @@ -1 +1 @@ -Subproject commit ca99759fc52d5b307438b42cb43bdd4a818a2e38 +Subproject commit cd6302ffd3934e081e6e87fe020de3e6fff6bcbc From c50a089796e6414e37d1750695720ca9ee210b47 Mon Sep 17 00:00:00 2001 From: geooo109 Date: Fri, 24 Apr 2026 17:52:44 +0300 Subject: [PATCH 09/13] refactor impl --- sqlglot-integration-tests | 2 +- sqlglot/optimizer/annotate_types.py | 23 ++++++++--------------- 2 files changed, 9 insertions(+), 16 deletions(-) diff --git a/sqlglot-integration-tests b/sqlglot-integration-tests index cd6302ffd3..a153b5ad97 160000 --- a/sqlglot-integration-tests +++ b/sqlglot-integration-tests @@ -1 +1 @@ -Subproject commit cd6302ffd3934e081e6e87fe020de3e6fff6bcbc +Subproject commit a153b5ad975db8412faf938ff47f2fd5c5ff8e68 diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py index 278e292759..f36319939b 100644 --- a/sqlglot/optimizer/annotate_types.py +++ b/sqlglot/optimizer/annotate_types.py @@ -331,12 +331,8 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: selects[name] = {s.alias_or_name: s.type for s in expression.selects if s.type} for pivot in scope.pivots: - pivot_source = scope.sources.get(pivot.alias) - if not pivot_source: - continue - inner_name = ( - pivot_source.name if isinstance(pivot_source, exp.Table) else pivot.alias + pivot.parent.name if isinstance(pivot.parent, exp.Table) else pivot.alias ) col_types = selects.get(inner_name, {}).copy() @@ -355,7 +351,7 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: src = first.this else: col_types[field_col.name] = exp.DataType.build( - "VARCHAR", dialect=self.dialect + "TEXT", dialect=self.dialect ) src = first @@ -465,15 +461,12 @@ def _annotate_expression( if not source: source_scope = source_scope.parent - if isinstance(source, exp.Table): - schema_type = self.schema.get_column_type(source, expr) - if schema_type.is_type(exp.DType.UNKNOWN) and source.args.get("pivots"): - pivot_type = ( - self._get_scope_selects(scope).get(expr.table, {}).get(expr.name) - ) - if pivot_type: - schema_type = pivot_type - self._set_type(expr, schema_type) + # Pivot-indexed selects win first: they capture UNPIVOT outputs whether + # or not the pivot alias made it into scope.sources. + if pivot_type := self._get_scope_selects(scope).get(expr.table, {}).get(expr.name): + self._set_type(expr, pivot_type) + elif isinstance(source, exp.Table): + self._set_type(expr, self.schema.get_column_type(source, expr)) elif source and source_scope: col_type = ( self._get_scope_selects(source_scope).get(expr.table, {}).get(expr.name) From f704e0c68f9d35e379ba26fbc67652789262593f Mon Sep 17 00:00:00 2001 From: geooo109 Date: Fri, 24 Apr 2026 17:57:29 +0300 Subject: [PATCH 10/13] Sync w/ integration tests --- sqlglot-integration-tests | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sqlglot-integration-tests b/sqlglot-integration-tests index a153b5ad97..ccfe6a98df 160000 --- a/sqlglot-integration-tests +++ b/sqlglot-integration-tests @@ -1 +1 @@ -Subproject commit a153b5ad975db8412faf938ff47f2fd5c5ff8e68 +Subproject commit ccfe6a98df29d2c811270861379d6e2f852a5241 From 6c507d80e835a9f66204ad8001a63a8fb1cd37a9 Mon Sep 17 00:00:00 2001 From: geooo109 Date: Fri, 24 Apr 2026 18:54:54 +0300 Subject: [PATCH 11/13] varchar default --- sqlglot/optimizer/annotate_types.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py index f36319939b..740cd69db7 100644 --- a/sqlglot/optimizer/annotate_types.py +++ b/sqlglot/optimizer/annotate_types.py @@ -350,9 +350,7 @@ def _get_scope_selects(self, scope: Scope) -> dict[str, dict[str, t.Any]]: col_types[field_col.name] = alias_node.type src = first.this else: - col_types[field_col.name] = exp.DataType.build( - "TEXT", dialect=self.dialect - ) + col_types[field_col.name] = exp.DType.VARCHAR.into_expr() src = first # Value column types from the IN source columns From 7238fb5164b1efa2e6ca846e8addce11faf8c220 Mon Sep 17 00:00:00 2001 From: geooo109 Date: Fri, 24 Apr 2026 21:12:53 +0300 Subject: [PATCH 12/13] Sync w/ integration tests --- sqlglot-integration-tests | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sqlglot-integration-tests b/sqlglot-integration-tests index ccfe6a98df..ca99759fc5 160000 --- a/sqlglot-integration-tests +++ b/sqlglot-integration-tests @@ -1 +1 @@ -Subproject commit ccfe6a98df29d2c811270861379d6e2f852a5241 +Subproject commit ca99759fc52d5b307438b42cb43bdd4a818a2e38 From a409090bc7445a73d806ad669f3d13765569f3bb Mon Sep 17 00:00:00 2001 From: geooo109 Date: Sat, 25 Apr 2026 00:55:11 +0300 Subject: [PATCH 13/13] Sync w/ integration tests