Repository: tobymao/sqlglot

Issue:
Executor doesn't work when GROUP BY..ORDER BY COUNT(*)
sqlglot==16.8.1



```python
    def test_execute_group_by_order_by_count(self):
        tables = {
            "x": [
                {"t": "a", "a": 6, "y": 2019},
                {"t": "b", "a": 12, "y": 2019},
                {"t": "c", "a": 15, "y": 2019},
                {"t": "a", "a": 10, "y": 2018},
                {"t": "b", "a": 12, "y": 2018},
                {"t": "c", "a": 2, "y": 2018},
                {"t": "a", "a": 9, "y": 2017},
            ],
        }
        schema = {
            "x": {
                "t": "text",
                "a": "number",
                "y": "number"
            },
        }
        gen = SqlGenerator("mysql")
        creates = gen.get_sql_tables_schema(schema)
        for c in creates:
            print(c.sql(pretty=True, dialect="mysql"), ";", sep="")
        inserts = gen.get_sql_tables_content(tables)
        for i in inserts:
            print(i.sql(pretty=True, dialect="mysql"), ";", sep="")

        for sql, cols, rows in [
            ('select t , sum ( a ) from x group by t order by avg(a)  asc limit 1;', ["t", "_col_1"], [("a", 25)]), # PASS
            ('select t , sum ( a ) from x group by t order by count ( * ) asc limit 1;', ["t", "_col_1"], [("b", 24)]),  # FAIL
        ]:
            with self.subTest(sql):
                result = execute(sql, read="mysql", schema=schema, tables=tables)
                self.assertEqual(result.columns, tuple(cols))
                self.assertEqual(result.rows, rows)
```

```
SubTest error: Traceback (most recent call last):
  File "/Users/.../python3.11/site-packages/sqlglot/executor/python.py", line 40, in execute
    contexts[node] = self.aggregate(node, context)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../python3.11/site-packages/sqlglot/executor/python.py", line 288, in aggregate
    add_row()
  File "/Users/.../python3.11/site-packages/sqlglot/executor/python.py", line 278, in add_row
    table.append(group + context.eval_tuple(aggregations))
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  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 "COUNT(scope[None]["_a_1"])", line 1, in <module>
KeyError: None


The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/.../python3.11/unittest/case.py", line 57, in testPartExecutor
    yield
  File "/usr/.../python3.11/unittest/case.py", line 538, in subTest
    yield
  File "/Users/.../src/tr/dbtool/test_executor.py", line 237, in test_execute_group_by_order_by_count
    result = execute(sql, read="mysql", schema=schema, tables=tables)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/..../python3.11/site-packages/sqlglot/executor/__init__.py", line 82, in execute
    result = PythonExecutor(tables=tables_).execute(plan)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../python3.11/site-packages/sqlglot/executor/python.py", line 60, in execute
    raise ExecuteError(f"Step '{node.id}' failed: {e}") from e
sqlglot.errors.ExecuteError: Step 'Aggregate: x (4552043856)' failed: None

```

Comments:

Patch:
diff --git a/sqlglot/planner.py b/sqlglot/planner.py
index 232c3b94..07ee7399 100644
--- a/sqlglot/planner.py
+++ b/sqlglot/planner.py
@@ -142,6 +142,10 @@ class Step:
 
             return bool(agg_funcs)
 
+        def set_ops_and_aggs(step):
+            step.operands = tuple(alias(operand, alias_) for operand, alias_ in operands.items())
+            step.aggregations = list(aggregations)
+
         for e in expression.expressions:
             if e.find(exp.AggFunc):
                 projections.append(exp.column(e.alias_or_name, step.name, quoted=True))
@@ -169,10 +173,7 @@ class Step:
                 else:
                     aggregate.condition = having.this
 
-            aggregate.operands = tuple(
-                alias(operand, alias_) for operand, alias_ in operands.items()
-            )
-            aggregate.aggregations = list(aggregations)
+            set_ops_and_aggs(aggregate)
 
             # give aggregates names and replace projections with references to them
             aggregate.group = {
@@ -204,13 +205,11 @@ class Step:
 
         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))
+                for i, ordered in enumerate(order.expressions):
+                    if extract_agg_operands(exp.alias_(ordered.this, f"_o_{i}", quoted=True)):
+                        ordered.this.replace(exp.column(f"_o_{i}", step.name, quoted=True))
 
-                step.aggregations = list(aggregations)
+                set_ops_and_aggs(aggregate)
 
             sort = Sort()
             sort.name = step.name
@@ -355,7 +354,10 @@ class Join(Step):
     def _to_s(self, indent: str) -> t.List[str]:
         lines = []
         for name, join in self.joins.items():
-            lines.append(f"{indent}{name}: {join['side']}")
+            lines.append(f"{indent}{name}: {join['side'] or 'INNER'}")
+            join_key = ", ".join(str(key) for key in t.cast(list, join.get("join_key") or []))
+            if join_key:
+                lines.append(f"{indent}Key: {join_key}")
             if join.get("condition"):
                 lines.append(f"{indent}On: {join['condition'].sql()}")  # type: ignore
         return lines


Test patch:
diff --git a/tests/test_executor.py b/tests/test_executor.py
index 006cf9df..9dacbbf3 100644
--- a/tests/test_executor.py
+++ b/tests/test_executor.py
@@ -708,6 +708,16 @@ class TestExecutor(unittest.TestCase):
                 [(2,), (1,), (3,)],
                 ("a",),
             ),
+            (
+                "SELECT a, SUM(b) FROM x GROUP BY a ORDER BY COUNT(*)",
+                [(3, 28), (1, 50), (2, 45)],
+                ("a", "_col_1"),
+            ),
+            (
+                "SELECT a, SUM(b) FROM x GROUP BY a ORDER BY COUNT(*) DESC",
+                [(1, 50), (2, 45), (3, 28)],
+                ("a", "_col_1"),
+            ),
         ):
             with self.subTest(sql):
                 result = execute(sql, tables=tables)