fix(optimizer): keep projections referenced by SORT BY / DISTRIBUTE BY / CLUSTER BY

sqlglot/optimizer/pushdown_projections.py

@@ -29,6 +29,18 @@
 GROUPING_CONSTRUCTS = (exp.Cube, exp.GroupingSets, exp.Paren, exp.Rollup, exp.Tuple)
 
 
+def _output_column_refs(expression: exp.Expr, scoped: bool) -> set[str]:
+    refs: set[str] = set()
+
+    for arg in ("order", "sort", "distribute", "cluster"):
+        node = expression.args.get(arg)
+        if node:
+            columns = find_all_in_scope(node, exp.Column) if scoped else node.find_all(exp.Column)
+            refs.update(c.name for c in columns if not c.table)
+
+    return refs
+
+
 def _is_self_referencing_cte(scope: Scope) -> bool:
     cte = scope.expression.parent
     return (
@@ -113,12 +125,11 @@ def pushdown_projections(
                 scope_sql = scope_expression.sql(dialect=dialect)
                 raise OptimizeError(f"Invalid set operation due to column mismatch: {scope_sql}.")
 
-            # Columns in ORDER BY need to be kept too
-            order = scope_expression.args.get("order")
-            if order and SELECT_ALL not in parent_selections:
-                order_refs = {c.name for c in find_all_in_scope(order, exp.Column) if not c.table}
-                if order_refs:
-                    parent_selections = parent_selections | order_refs
+            # Columns referenced by ORDER BY and friends need to be kept too
+            if SELECT_ALL not in parent_selections:
+                output_refs = _output_column_refs(scope_expression, scoped=True)
+                if output_refs:
+                    parent_selections = parent_selections | output_refs
 
             referenced_columns[left] = parent_selections
 
@@ -168,13 +179,7 @@ def pushdown_projections(
 
 def _remove_unused_selections(scope, parent_selections, schema, alias_count, journal=None):
     expression = scope.expression
-    order = expression.args.get("order")
-
-    if order:
-        # Assume columns without a qualified table are references to output columns
-        order_refs = {c.name for c in order.find_all(exp.Column) if not c.table}
-    else:
-        order_refs = set()
+    output_refs = _output_column_refs(expression, scoped=False)
 
     # Resolve GROUP BY ordinals before pruning
     ordinal_refs = _group_by_ordinal_refs(expression)
@@ -200,7 +205,7 @@ def _remove_unused_selections(scope, parent_selections, schema, alias_count, jou
         if (
             select_all
             or name in parent_selections
-            or name in order_refs
+            or name in output_refs
             or alias_count > 0
             or id(selection) in group_ordinal_selection_ids
             or (implicit_group_by_all and not is_agg_selection)

sqlglot/optimizer/scope.py

@@ -381,6 +381,7 @@ def columns(self) -> list[exp.Column]:
                     exp.Select,
                     exp.Qualify,
                     exp.Order,
+                    exp.Cluster,
                     exp.Having,
                     exp.Hint,
                     exp.Table,
@@ -393,7 +394,7 @@ def columns(self) -> list[exp.Column]:
                     or isinstance(ancestor, exp.Select)
                     or (isinstance(ancestor, exp.Table) and not isinstance(ancestor.this, exp.Func))
                     or (
-                        isinstance(ancestor, (exp.Order, exp.Distinct))
+                        isinstance(ancestor, (exp.Order, exp.Cluster, exp.Distinct))
                         and (
                             isinstance(ancestor.parent, (exp.Window, exp.WithinGroup))
                             or not isinstance(ancestor.parent, exp.Select)

sqlglot/parser.py

@@ -1866,7 +1866,7 @@ def _parse_partitioned_by_bucket_or_truncate(self) -> exp.Expr | None:
 
     # Whether query modifiers such as LIMIT are attached to the UNION node (vs its right operand)
     MODIFIERS_ATTACHED_TO_SET_OP: t.ClassVar = True
-    SET_OP_MODIFIERS: t.ClassVar = {"order", "limit", "offset"}
+    SET_OP_MODIFIERS: t.ClassVar = {"order", "limit", "offset", "sort", "distribute", "cluster"}
 
     # Whether to parse IF statements that aren't followed by a left parenthesis as commands
     NO_PAREN_IF_COMMANDS: t.ClassVar = True