refactor(optimizer): continuation on typing coverage improvements (#7572)

* refactor (optimizer): add type annotations, raise `OptimizeError` if parent is None in `reorder_joins`

* refactor (optimizer): go back to original func body in `_distribute`

* refactor (optimizer): added a few other type annotations

* fix: added from future import annotations for <3.10 python compatibility

Co-authored-by: Copilot <copilot@github.com>

* fix: narrow pushdown and pushdown_cnf `join_index` mapping key type

* fix: narrow `replace_aliases()` return type, as well as the closure argument

Co-authored-by: Copilot <copilot@github.com>

* fix: delete incorrect `None` return type for `_distribute`, and improve `_predicate_lengths` typing

Co-authored-by: Copilot <copilot@github.com>

* chore: ruff formatting + linting

* fix: annotate `_is_limit_1` function, modified None check to make it pass

Co-authored-by: Copilot <copilot@github.com>

* refactor:
  - improve `_unique_outputs` internal and signature typing
  - added a few annotations to `_join_is_used`,  `eliminate_joins`, and `_is_limit_1` functions

Co-authored-by: Copilot <copilot@github.com>

* fix: prefer parent class `Query` instead of `Select` for `_unique_outputs` internal cast

---------

Co-authored-by: Copilot <copilot@github.com>

Author: OutSquareCapital

sqlglot/optimizer/eliminate_joins.py

@@ -30,7 +30,7 @@ def eliminate_joins(expression: E) -> E:
         The optimized expression
     """
     for scope in traverse_scope(expression):
-        joins = scope.expression.args.get("joins", [])
+        joins: list[exp.Join] = scope.expression.args.get("joins", [])
         if not joins:
             continue
 
@@ -53,7 +53,7 @@ def eliminate_joins(expression: E) -> E:
     return expression
 
 
-def _should_eliminate_join(scope, join, alias):
+def _should_eliminate_join(scope: Scope, join: exp.Join, alias: str) -> bool:
     inner_source = scope.sources.get(alias)
     return (
         isinstance(inner_source, Scope)
@@ -65,20 +65,20 @@ def _should_eliminate_join(scope, join, alias):
     )
 
 
-def _join_is_used(scope, join, alias):
+def _join_is_used(scope: Scope, join: exp.Join, alias: str) -> bool:
     # We need to find all columns that reference this join.
     # But columns in the ON clause shouldn't count.
-    on = join.args.get("on")
-    if on:
-        on_clause_columns = {id(column) for column in on.find_all(exp.Column)}
+    on: exp.Expr | None = join.args.get("on")
+    if on is not None:
+        on_clause_columns: set[int] = {id(column) for column in on.find_all(exp.Column)}
     else:
         on_clause_columns = set()
     return any(
         column for column in scope.source_columns(alias) if id(column) not in on_clause_columns
     )
 
 
-def _is_joined_on_all_unique_outputs(scope, join):
+def _is_joined_on_all_unique_outputs(scope: Scope, join: exp.Join) -> bool:
     unique_outputs = _unique_outputs(scope)
     if not unique_outputs:
         return False
@@ -88,18 +88,19 @@ def _is_joined_on_all_unique_outputs(scope, join):
     return not remaining_unique_outputs
 
 
-def _unique_outputs(scope):
+def _unique_outputs(scope: Scope) -> set[str]:
     """Determine output columns of `scope` that must have a unique combination per row"""
-    if scope.expression.args.get("distinct"):
-        return set(scope.expression.named_selects)
+    expr = t.cast(exp.Query, scope.expression)
+    if expr.args.get("distinct") is not None:
+        return set(expr.named_selects)
 
-    group = scope.expression.args.get("group")
-    if group:
-        grouped_expressions = set(group.expressions)
-        grouped_outputs = set()
+    group: exp.Group | None = expr.args.get("group")
+    if group is not None:
+        grouped_expressions: set[exp.Expr] = set(group.expressions)
+        grouped_outputs: set[exp.Expr] = set()
 
-        unique_outputs = set()
-        for select in scope.expression.selects:
+        unique_outputs: set[str] = set()
+        for select in expr.selects:
             output = select.unalias()
             if output in grouped_expressions:
                 grouped_outputs.add(output)
@@ -112,22 +113,22 @@ def _unique_outputs(scope):
             return set()
 
     if _has_single_output_row(scope):
-        return set(scope.expression.named_selects)
+        return set(expr.named_selects)
 
     return set()
 
 
-def _has_single_output_row(scope):
+def _has_single_output_row(scope: Scope) -> bool:
     return isinstance(scope.expression, exp.Select) and (
         all(isinstance(e.unalias(), exp.AggFunc) for e in scope.expression.selects)
         or _is_limit_1(scope)
         or not scope.expression.args.get("from_")
     )
 
 
-def _is_limit_1(scope):
-    limit = scope.expression.args.get("limit")
-    return limit and limit.expression.this == "1"
+def _is_limit_1(scope: Scope) -> bool:
+    limit: exp.Limit | None = scope.expression.args.get("limit")
+    return limit is not None and limit.expression.this == "1"
 
 
 def join_condition(join):

sqlglot/optimizer/normalize.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Callable, Iterator
 import logging
 
 from sqlglot import exp
@@ -123,7 +124,9 @@ def normalization_distance(
     return total
 
 
-def _predicate_lengths(expression, dnf, max_=float("inf"), depth=0):
+def _predicate_lengths(
+    expression: exp.Expr, dnf: bool, max_: float = float("inf"), depth: int = 0
+) -> Iterator[int]:
     """
     Returns a list of predicate lengths when expanded to normalized form.
 
@@ -151,7 +154,9 @@ def _predicate_lengths(expression, dnf, max_=float("inf"), depth=0):
         yield from _predicate_lengths(right, dnf, max_, depth)
 
 
-def distributive_law(expression, dnf, max_distance, simplifier=None):
+def distributive_law(
+    expression: exp.Expr, dnf: bool, max_distance: float, simplifier: Simplifier | None = None
+):
     """
     x OR (y AND z) -> (x OR y) AND (x OR z)
     (x AND y) OR (y AND z) -> (x OR y) AND (x OR z) AND (y OR y) AND (y OR z)
@@ -187,7 +192,13 @@ def distributive_law(expression, dnf, max_distance, simplifier=None):
     return expression
 
 
-def _distribute(a, b, from_func, to_func, simplifier):
+def _distribute(
+    a,
+    b,
+    from_func: Callable[..., exp.Condition],
+    to_func: Callable[..., exp.Condition],
+    simplifier: Simplifier,
+):
     if isinstance(a, exp.Connector):
         exp.replace_children(
             a,

sqlglot/optimizer/optimize_joins.py

@@ -2,6 +2,7 @@
 from collections.abc import Iterable
 
 from sqlglot import exp
+from sqlglot.errors import OptimizeError
 from sqlglot.helper import tsort
 
 JOIN_ATTRS = ("on", "side", "kind", "using", "method")
@@ -18,7 +19,7 @@ def optimize_joins(expression: exp.Expr) -> exp.Expr:
     """
 
     for select in expression.find_all(exp.Select):
-        joins = select.args.get("joins", [])
+        joins: list[exp.Join] = select.args.get("joins", [])
 
         if not _is_reorderable(joins):
             continue
@@ -57,13 +58,15 @@ def optimize_joins(expression: exp.Expr) -> exp.Expr:
     return expression
 
 
-def reorder_joins(expression) -> exp.Expr:
+def reorder_joins(expression: exp.Expr) -> exp.Expr:
     """
     Reorder joins by topological sort order based on predicate references.
     """
     for from_ in expression.find_all(exp.From):
         parent = from_.parent
-        joins = parent.args.get("joins", [])
+        if parent is None:
+            raise OptimizeError("FROM clause without parent expression")
+        joins: list[exp.Join] = parent.args.get("joins", [])
 
         if not _is_reorderable(joins):
             continue

sqlglot/optimizer/pushdown_predicates.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
 import typing as t
 
 from sqlglot import exp
@@ -80,7 +81,13 @@ def pushdown_predicates(expression: E, dialect: DialectType = None) -> E:
     return expression
 
 
-def pushdown(condition, sources, scope_ref_count, dialect, join_index=None):
+def pushdown(
+    condition: exp.Expr | None,
+    sources,
+    scope_ref_count,
+    dialect: DialectType,
+    join_index: Mapping[str, int] | None = None,
+):
     if not condition:
         return
 
@@ -99,7 +106,7 @@ def pushdown(condition, sources, scope_ref_count, dialect, join_index=None):
         pushdown_dnf(predicates, sources, scope_ref_count, join_index=join_index)
 
 
-def pushdown_cnf(predicates, sources, scope_ref_count, join_index=None):
+def pushdown_cnf(predicates, sources, scope_ref_count, join_index: Mapping[str, int] | None = None):
     """
     If the predicates are in CNF like form, we can simply replace each block in the parent.
     """

sqlglot/optimizer/resolver.py

@@ -21,14 +21,14 @@ class Resolver:
     This is a class so we can lazily load some things and easily share them across functions.
     """
 
-    def __init__(self, scope: Scope, schema: Schema, infer_schema: bool = True):
-        self.scope = scope
-        self.schema = schema
-        self.dialect = schema.dialect or Dialect()
+    def __init__(self, scope: Scope, schema: Schema, infer_schema: bool = True) -> None:
+        self.scope: Scope = scope
+        self.schema: Schema = schema
+        self.dialect: Dialect = schema.dialect or Dialect()
         self._source_columns: dict[str, Sequence[str]] | None = None
         self._unambiguous_columns: Mapping[str, str] | None = None
         self._all_columns: set[str] | None = None
-        self._infer_schema = infer_schema
+        self._infer_schema: bool = infer_schema
         self._get_source_columns_cache: dict[tuple[str, bool], Sequence[str]] = {}
 
     def get_table(self, column: str | exp.Column) -> exp.Identifier | None:

sqlglot/optimizer/unnest_subqueries.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 from sqlglot import exp
 from sqlglot.helper import name_sequence
 from sqlglot.optimizer.scope import ScopeType, find_in_scope, traverse_scope
@@ -303,11 +304,11 @@ def remove_aggs(node):
     )
 
 
-def _replace(expression, condition):
+def _replace(expression: exp.Expr, condition: exp.ExpOrStr) -> exp.Expr:
     return expression.replace(exp.condition(condition))
 
 
-def _other_operand(expression):
+def _other_operand(expression: object) -> exp.Expr | None:
     if isinstance(expression, exp.In):
         return expression.this