Repository: tobymao/sqlglot

Issue:
Executor cannot find the aggr column in order-by clause if aggr column is not in the select clause.
sqlglot==16.6.0

Related to #1822

```python
('SELECT a, AVG(b) FROM x GROUP BY a ORDER BY AVG(b)', ["a", "_col_1"], [("b", 22.5), ("a", 25), ("c", 28)]), # PASS
('SELECT a FROM x GROUP BY a ORDER BY AVG(b)', ["a"], [("b",), ("a",), ("c",)]),  # FAILED
```
if avg not in select clause , i have the same issue.


Here is the error log.
```
Traceback (most recent call last):
  File "/Users/.../python3.11/site-packages/sqlglot/executor/python.py", line 46, in execute
    contexts[node] = self.sort(node, context)
                     ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../python3.11/site-packages/sqlglot/executor/python.py", line 323, in sort
    sort_ctx.sort(self.generate_tuple(step.key))
  File "/Users/.../python3.11/site-packages/sqlglot/executor/context.py", line 81, in sort
    self.table.rows.sort(key=sort_key)
  File "/Users/.../python3.11/site-packages/sqlglot/executor/context.py", line 79, in sort_key
    return self.eval_tuple(key)
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../python3.11/site-packages/sqlglot/executor/context.py", line 38, in eval_tuple
    return tuple(self.eval(code) for code in codes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../python3.11/site-packages/sqlglot/executor/context.py", line 38, in <genexpr>
    return tuple(self.eval(code) for code in codes)
                 ^^^^^^^^^^^^^^^
  File "/Users/.../python3.11/site-packages/sqlglot/executor/context.py", line 35, in eval
    return eval(code, self.env)
           ^^^^^^^^^^^^^^^^^^^^
  File "ORDERED(AVG(scope["x"]["b"]), False, True)", line 1, in <module>
  File "/Users/.../python3.11/site-packages/sqlglot/executor/table.py", line 104, in __getitem__
    return self.row[self.columns[column]]
                    ~~~~~~~~~~~~^^^^^^^^
KeyError: 'b'
```

Comments:

Patch:
diff --git a/sqlglot/executor/context.py b/sqlglot/executor/context.py
index 630cb657..d7952c15 100644
--- a/sqlglot/executor/context.py
+++ b/sqlglot/executor/context.py
@@ -41,11 +41,13 @@ class Context:
     def table(self) -> Table:
         if self._table is None:
             self._table = list(self.tables.values())[0]
+
             for other in self.tables.values():
                 if self._table.columns != other.columns:
                     raise Exception(f"Columns are different.")
                 if len(self._table.rows) != len(other.rows):
                     raise Exception(f"Rows are different.")
+
         return self._table
 
     def add_columns(self, *columns: str) -> None:
diff --git a/sqlglot/planner.py b/sqlglot/planner.py
index f2467021..232c3b94 100644
--- a/sqlglot/planner.py
+++ b/sqlglot/planner.py
@@ -23,9 +23,11 @@ class Plan:
             while nodes:
                 node = nodes.pop()
                 dag[node] = set()
+
                 for dep in node.dependencies:
                     dag[node].add(dep)
                     nodes.add(dep)
+
             self._dag = dag
 
         return self._dag
@@ -128,13 +130,16 @@ class Step:
             agg_funcs = tuple(expression.find_all(exp.AggFunc))
             if agg_funcs:
                 aggregations.add(expression)
+
             for agg in agg_funcs:
                 for operand in agg.unnest_operands():
                     if isinstance(operand, exp.Column):
                         continue
                     if operand not in operands:
                         operands[operand] = next_operand_name()
+
                     operand.replace(exp.column(operands[operand], quoted=True))
+
             return bool(agg_funcs)
 
         for e in expression.expressions:
@@ -178,13 +183,14 @@ class Step:
             for k, v in aggregate.group.items():
                 intermediate[v] = k
                 if isinstance(v, exp.Column):
-                    intermediate[v.alias_or_name] = k
+                    intermediate[v.name] = k
 
             for projection in projections:
                 for node, *_ in projection.walk():
                     name = intermediate.get(node)
                     if name:
                         node.replace(exp.column(name, step.name))
+
             if aggregate.condition:
                 for node, *_ in aggregate.condition.walk():
                     name = intermediate.get(node) or intermediate.get(node.name)
@@ -197,6 +203,15 @@ class Step:
         order = expression.args.get("order")
 
         if order:
+            if isinstance(step, Aggregate):
+                for ordered in order.expressions:
+                    if ordered.find(exp.AggFunc):
+                        operand_name = next_operand_name()
+                        extract_agg_operands(exp.alias_(ordered.this, operand_name, quoted=True))
+                        ordered.this.replace(exp.column(operand_name, quoted=True))
+
+                step.aggregations = list(aggregations)
+
             sort = Sort()
             sort.name = step.name
             sort.key = order.expressions


Test patch:
diff --git a/tests/test_executor.py b/tests/test_executor.py
index f257bd14..13376ce5 100644
--- a/tests/test_executor.py
+++ b/tests/test_executor.py
@@ -700,6 +700,11 @@ class TestExecutor(unittest.TestCase):
                 [(2, 25.0)],
                 ("_col_0", "_col_1"),
             ),
+            (
+                "SELECT a FROM x GROUP BY a ORDER BY AVG(b)",
+                [(2,), (1,), (3,)],
+                ("a",),
+            ),
         ):
             with self.subTest(sql):
                 result = execute(sql, tables=tables)