python/sqlglot
1
0
Fork 0
Code Issues Activity

Merging upstream version 25.29.0.

Browse source 
Signed-off-by: Daniel Baumann <daniel@debian.org>
This commit is contained in:
Daniel Baumann 2025-02-13 21:56:19 +01:00
parent de8c8a17d0
commit 1e53504dfc
Signed by: daniel
GPG key ID: FBB4F0E80A80222F
97 changed files with 64720 additions and 61752 deletions
Download patch file Download diff file Expand all files Collapse all files

57
sqlglot/optimizer/annotate_types.py
View file

@ -192,6 +192,11 @@ class TypeAnnotator(metaclass=_TypeAnnotator):
# Caches the ids of annotated sub-Expressions, to ensure we only visit them once
self._visited: t.Set[int] = set()
# Maps an exp.SetOperation's id (e.g. UNION) to its projection types. This is computed if the
# exp.SetOperation is the expression of a scope source, as selecting from it multiple times
# would reprocess the entire subtree to coerce the types of its operands' projections
self._setop_column_types: t.Dict[int, t.Dict[str, exp.DataType.Type]] = {}
def _set_type(
self, expression: exp.Expression, target_type: t.Optional[exp.DataType | exp.DataType.Type]
) -> None:
@ -231,22 +236,42 @@ class TypeAnnotator(metaclass=_TypeAnnotator):
elif isinstance(expression, exp.SetOperation) and len(expression.left.selects) == len(
expression.right.selects
):
if expression.args.get("by_name"):
r_type_by_select = {s.alias_or_name: s.type for s in expression.right.selects}
selects[name] = {
s.alias_or_name: self._maybe_coerce(
t.cast(exp.DataType, s.type),
r_type_by_select.get(s.alias_or_name) or exp.DataType.Type.UNKNOWN,
)
for s in expression.left.selects
}
else:
selects[name] = {
ls.alias_or_name: self._maybe_coerce(
t.cast(exp.DataType, ls.type), t.cast(exp.DataType, rs.type)
)
for ls, rs in zip(expression.left.selects, expression.right.selects)
}
selects[name] = col_types = self._setop_column_types.setdefault(id(expression), {})
if not col_types:
# Process a chain / sub-tree of set operations
for set_op in expression.walk(
prune=lambda n: not isinstance(n, (exp.SetOperation, exp.Subquery))
):
if not isinstance(set_op, exp.SetOperation):
continue
if set_op.args.get("by_name"):
r_type_by_select = {
s.alias_or_name: s.type for s in set_op.right.selects
}
setop_cols = {
s.alias_or_name: self._maybe_coerce(
t.cast(exp.DataType, s.type),
r_type_by_select.get(s.alias_or_name)
or exp.DataType.Type.UNKNOWN,
)
for s in set_op.left.selects
}
else:
setop_cols = {
ls.alias_or_name: self._maybe_coerce(
t.cast(exp.DataType, ls.type), t.cast(exp.DataType, rs.type)
)
for ls, rs in zip(set_op.left.selects, set_op.right.selects)
}
# Coerce intermediate results with the previously registered types, if they exist
for col_name, col_type in setop_cols.items():
col_types[col_name] = self._maybe_coerce(
col_type, col_types.get(col_name, exp.DataType.Type.NULL)
)
else:
selects[name] = {s.alias_or_name: s.type for s in expression.selects}