Executor doesn't support "ORDER BY AVG(b)"

[stonyw]

sqlglot.errors.ExecuteError: Step 'Sort: x (4361584080)' failed: 'b'
for 'SELECT a, AVG(b) FROM x GROUP BY a ORDER BY AVG(b)'

import unittest

from sqlglot.executor import execute


class TestExecutor(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        pass

    @classmethod
    def tearDownClass(cls):
        pass

    def test_execute_callable(self):
        tables = {
            "x": [
                {"a": "a", "b": 10},
                {"a": "b", "b": 20},
                {"a": "c", "b": 28},
                {"a": "b", "b": 25},
                {"a": "a", "b": 40},
            ],
        }
        schema = {
            "x": {
                "a": "VARCHAR",
                "b": "NUMERIC",
            },
        }

        for sql, cols, rows in [
            ('SELECT a, AVG(b) FROM x GROUP BY a', ["a", "_col_1"], [("a", 25), ("b", 22.5), ("c", 28)]),  # PASS
            ('SELECT a, AVG(b) FROM x GROUP BY a ORDER BY a', ["a", "_col_1"], [("a", 25), ("b", 22.5), ("c", 28)]),  # PASS
            ('SELECT a, AVG(b) FROM x GROUP BY a ORDER BY AVG(b)', ["a", "_col_1"], [("b", 22.5), ("a", 25), ("c", 28)]),  # FAILED
        ]:
            with self.subTest(sql):
                result = execute(sql, schema=schema, tables=tables)
                self.assertEqual(result.columns, tuple(cols))
                self.assertEqual(result.rows, rows)

Test logs with error message:

Testing started at 13:55 ...
Launching unittests with arguments python -m unittest test_executor.TestExecutor.test_execute_callable in /Users/..app../src/tr/dbtool

SubTest error: Traceback (most recent call last):
  File "/Users/.../lib/python3.11/site-packages/sqlglot/executor/python.py", line 46, in execute
    contexts[node] = self.sort(node, context)
                     ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../lib/python3.11/site-packages/sqlglot/executor/python.py", line 325, in sort
    sort_ctx.sort(self.generate_tuple(step.key))
  File "/Users/.../lib/python3.11/site-packages/sqlglot/executor/context.py", line 83, in sort
    self.table.rows.sort(key=sort_key)
  File "/Users/.../lib/python3.11/site-packages/sqlglot/executor/context.py", line 81, in sort_key
    return self.eval_tuple(key)
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../lib/python3.11/site-packages/sqlglot/executor/context.py", line 38, in eval_tuple
    return tuple(self.eval(code) for code in codes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/ ../lib/python3.11/site-packages/sqlglot/executor/context.py", line 38, in <genexpr>
    return tuple(self.eval(code) for code in codes)
                 ^^^^^^^^^^^^^^^
  File "/Users/.../lib/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/.../lib/python3.11/site-packages/sqlglot/executor/table.py", line 104, in __getitem__
    return self.row[self.columns[column]]
                    ~~~~~~~~~~~~^^^^^^^^
KeyError: 'b'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/local/Cellar/[email protected]/3.11.3/Frameworks/Python.framework/Versions/3.11/lib/python3.11/unittest/case.py", line 57, in testPartExecutor
    yield
  File "/usr/local/Cellar/[email protected]/3.11.3/Frameworks/Python.framework/Versions/3.11/lib/python3.11/unittest/case.py", line 538, in subTest
    yield
  File "/Users/..app../src/tr/dbtool/test_executor.py", line 44, in test_execute_callable
    result = execute(sql, schema=schema, tables=tables)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../lib/python3.11/site-packages/sqlglot/executor/__init__.py", line 82, in execute
    result = PythonExecutor(tables=tables_).execute(plan)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/.../lib/python3.11/site-packages/sqlglot/executor/python.py", line 63, in execute
    raise ExecuteError(f"Step '{node.id}' failed: {e}") from e
sqlglot.errors.ExecuteError: Step 'Sort: x (4361584080)' failed: 'b'



One or more subtests failed
Failed subtests list: [SELECT a, AVG(b) FROM x GROUP BY a ORDER BY AVG(b)]


Ran 1 test in 0.025s

FAILED (errors=1)

Possible cause:
ORDER clause uses a temp middle table including new column for AVG(b) (named "_col_1") while no column b.