Repository: mabel-dev/opteryx

Issue:
🪲 GROUP BY is losing column source information
### Thank you for taking the time to report a problem with Opteryx.
_To help us to respond to your request we ask that you try to provide the below detail about the bug._

**Describe the bug** _A clear and specific description of what the bug is. What the error, incorrect or unexpected behaviour was._


**Expected behaviour** _A clear and concise description of what you expected to happen._


**Sample Code/Statement** _If you can, please submit the SQL statement or Python code snippet, or a representative example using the sample datasets._

~~~sql

~~~

**Additional context** _Add any other context about the problem here, for example what you have done to try to diagnose or workaround the problem._



Discussion:

Resolving patch:
diff --git a/opteryx/components/binder/binder_visitor.py b/opteryx/components/binder/binder_visitor.py
--- a/opteryx/components/binder/binder_visitor.py
+++ b/opteryx/components/binder/binder_visitor.py
@@ -27,6 +27,7 @@
 from opteryx.exceptions import ColumnNotFoundError
 from opteryx.exceptions import InvalidInternalStateError
 from opteryx.managers.expression import NodeType
+from opteryx.managers.expression import get_all_nodes_of_type
 from opteryx.models import LogicalColumn
 from opteryx.models import Node
 from opteryx.third_party.travers import Graph
@@ -263,27 +264,45 @@ def visit_aggregate_and_group(
             Tuple[Node, Dict[str, Any]]
             The modified node and the updated context.
         """
-        if node.groups:
-            tmp_groups, _ = zip(
-                *(inner_binder(group, context, node.identity) for group in node.groups)
-            )
-            node.groups = list(tmp_groups)
-
         if node.aggregates:
             tmp_aggregates, _ = zip(
                 *(inner_binder(aggregate, context, node.identity) for aggregate in node.aggregates)
             )
             node.aggregates = list(tmp_aggregates)
 
-        columns = [
-            i.schema_column
-            for i in (node.aggregates or []) + (node.groups or [])
-            if i.schema_column
+        # We're going to trim down the schemas to just the columns used in the GROUP BY.
+        # 1) the easy one - the columns explictly in the GROUP BY
+        columns_to_keep = set()
+        if node.groups:
+            tmp_groups, _ = zip(
+                *(inner_binder(group, context, node.identity) for group in node.groups)
+            )
+            columns_to_keep = {col.schema_column.identity for col in tmp_groups}
+        # 2) the columns referenced in the SELECT
+        all_identifiers = [
+            node.schema_column.identity
+            for node in get_all_nodes_of_type(
+                node.aggregates + node.groups, select_nodes=(NodeType.IDENTIFIER,)
+            )
         ]
+        columns_to_keep = columns_to_keep.union(all_identifiers)
+        node.all_identifiers = columns_to_keep
+
+        for name, schema in list(context.schemas.items()):
+            schema_columns = [
+                column for column in schema.columns if column.identity in columns_to_keep
+            ]
+            if schema_columns:
+                context.schemas[name].columns = schema_columns
+            else:
+                context.schemas.pop(name)
+
+        # we should always have a derived schema
+        if not "$derived" in context.schemas:
+            context.schemas["$derived"] = derived.schema()
 
-        # although this is now the only dataset, we use $derived as there is logic
-        # specific to that dataset with regards to reevaluating expressions
-        context.schemas = {"$derived": RelationSchema(name="$derived", columns=columns)}
+        # the aggregates and any calculated expressions in the SELECT should be in $derived
+        context.schemas["$derived"].columns.extend(col.schema_column for col in node.aggregates)
 
         return node, context
 
@@ -309,6 +328,8 @@ def name_column(qualifier, column):
 
         def keep_column(column, identities):
             if len(node.columns) == 1 and node.columns[0].node_type == NodeType.WILDCARD:
+                if node.columns[0].value:
+                    return node.columns[0].value[0] == column.origin
                 return True
             return column.identity in identities
 
diff --git a/opteryx/exceptions.py b/opteryx/exceptions.py
--- a/opteryx/exceptions.py
+++ b/opteryx/exceptions.py
@@ -40,7 +40,7 @@
                  ├── AmbiguousDatasetError
                  ├── AmbiguousIdentifierError
                  ├── ColumnNotFoundError
-                 ├── ColumnReferencedBeforeEvaluation
+                 ├── ColumnReferencedBeforeEvaluationError
                  ├── DatasetNotFoundError
                  ├── FunctionNotFoundError
                  ├── IncorrectTypeError
@@ -151,7 +151,7 @@ def __init__(
         super().__init__(message)
 
 
-class ColumnReferencedBeforeEvaluation(SqlError):
+class ColumnReferencedBeforeEvaluationError(SqlError):
     """
     Return an error message when the column reference order is incorrect
     """
diff --git a/opteryx/managers/expression/__init__.py b/opteryx/managers/expression/__init__.py
--- a/opteryx/managers/expression/__init__.py
+++ b/opteryx/managers/expression/__init__.py
@@ -27,7 +27,7 @@
 from orso.types import OrsoTypes
 from pyarrow import Table
 
-from opteryx.exceptions import ColumnReferencedBeforeEvaluation
+from opteryx.exceptions import ColumnReferencedBeforeEvaluationError
 from opteryx.exceptions import UnsupportedSyntaxError
 from opteryx.functions.binary_operators import binary_operations
 from opteryx.functions.unary_operations import UNARY_OPERATIONS
@@ -201,7 +201,7 @@ def _inner_evaluate(root: Node, table: Table, context: ExecutionContext):
         if node_type == NodeType.EVALUATED:
             column = root.schema_column
             if not root.schema_column.identity in table.column_names:
-                raise ColumnReferencedBeforeEvaluation(column=root.schema_column.name)
+                raise ColumnReferencedBeforeEvaluationError(column=root.schema_column.name)
             return table[root.schema_column.identity].to_numpy()
         if node_type == NodeType.COMPARISON_OPERATOR:
             left = _inner_evaluate(root.left, table, context)


Test patch:
diff --git a/tests/sql_battery/test_shapes_and_errors_battery.py b/tests/sql_battery/test_shapes_and_errors_battery.py
--- a/tests/sql_battery/test_shapes_and_errors_battery.py
+++ b/tests/sql_battery/test_shapes_and_errors_battery.py
@@ -50,7 +50,7 @@
     AmbiguousIdentifierError,
     AmbiguousDatasetError,
     ColumnNotFoundError,
-    ColumnReferencedBeforeEvaluation,
+    ColumnReferencedBeforeEvaluationError,
     DatasetNotFoundError,
     EmptyDatasetError,
     IncompatibleTypesError,
@@ -268,7 +268,7 @@
         ("SELECT UPPER(name) as NAME FROM $satellites", 177, 1, None),
         ("SELECT name as n FROM $satellites WHERE n = 'Titan'", None, 1, ColumnNotFoundError),
         ("SELECT id as Identifier FROM $satellites ORDER BY Identifier", 177, 1, None),
-        ("SELECT name as n FROM $satellites GROUP BY name HAVING COUNT(n) > 1", None, 1, ColumnReferencedBeforeEvaluation),  # TEMP
+        ("SELECT name as n FROM $satellites GROUP BY name HAVING COUNT(n) > 1", None, 1, ColumnReferencedBeforeEvaluationError),  # TEMP
         ("SELECT name as Name, name as NAME FROM $satellites", 177, 2, AmbiguousIdentifierError),
         ("SELECT COUNT(id) as countID, MIN(id) as minID FROM $satellites", 1, 2, None),
         ("SELECT s.id as satelliteID, p.id as planetID FROM $satellites s JOIN $planets p ON s.planetId = p.id", 177, 2, None),
@@ -276,7 +276,7 @@
         ("SELECT name as n as m FROM $satellites", None, 1, SqlError),
         ("SELECT id*2 as doubleID FROM $satellites", 177, 1, None),
         ("SELECT id as Identifier FROM $satellites ORDER BY Identifier", 177, 1, None),
-        ("SELECT name as n FROM $satellites GROUP BY name HAVING COUNT(n) > 1", None, 1, ColumnReferencedBeforeEvaluation),
+        ("SELECT name as n FROM $satellites GROUP BY name HAVING COUNT(n) > 1", None, 1, ColumnReferencedBeforeEvaluationError),
         ("SELECT name as n FROM $satellites WHERE n = 'Calypso'", None, 1, ColumnNotFoundError),
         ("SELECT id * 2 as DoubleID FROM $satellites", 177, 1, None),
         ("SELECT LEFT(name, 3) as newName FROM $satellites", 177, 1, None),
@@ -437,7 +437,7 @@
         ("SELECT LEFT(name, 1), COUNT(*) FROM $satellites GROUP BY LEFT(name, 1) ORDER BY 2 DESC", 21, 2, UnsupportedSyntaxError),
         ("SELECT LEFT(name, 1), COUNT(*) FROM $satellites GROUP BY name ORDER BY 2 DESC", 177, 2, UnsupportedSyntaxError),
         ("SELECT LEFT(name, 2) as le, COUNT(*) FROM $satellites GROUP BY LEFT(name, 2) ORDER BY 2 DESC", 87, 2, UnsupportedSyntaxError),
-        ("SELECT LEFT(name, 2) as le, COUNT(*) FROM $satellites GROUP BY LEFT(name, 1)", 87, 2, ColumnNotFoundError),
+        ("SELECT LEFT(name, 2) as le, COUNT(*) FROM $satellites GROUP BY LEFT(name, 1)", 87, 2, TypeError),
         ("SELECT RIGHT(name, 10), COUNT(*) FROM $satellites GROUP BY RIGHT(name, 10) ORDER BY 2 DESC", 177, 2, UnsupportedSyntaxError),
         ("SELECT RIGHT(name, 2) as le, COUNT(*) FROM $satellites GROUP BY RIGHT(name, 2) ORDER BY 2 DESC", 91, 2, UnsupportedSyntaxError),
         ("SELECT RIGHT(name, 2) as le, COUNT(*) FROM $satellites GROUP BY le ORDER BY 2 DESC", 91, 2, UnsupportedSyntaxError),
@@ -747,13 +747,12 @@
         ("SELECT P_1.* FROM $planets AS P_1 INNER JOIN $satellites AS P_2 ON P_1.id = P_2.planet_id ORDER BY P_2.id", 177, 20, None),
         ("SELECT COUNT(*) FROM $planets AS P_1 INNER JOIN $planets AS P_2 ON P_1.id = P_2.id", 1, 1, None),
         ("SELECT COUNT(*) FROM $planets AS P_1 INNER JOIN $satellites AS P_2 ON P_1.id = P_2.planet_id", 1, 1, None),
-        ("SELECT P_1.id, COUNT(P_2.id) FROM $planets AS P_1 INNER JOIN $satellites AS P_2 ON P_1.id = P_2.planet_id GROUP BY P_1.id", 9, 2, None),
+        ("SELECT P_1.id, COUNT(P_2.id) FROM $planets AS P_1 INNER JOIN $satellites AS P_2 ON P_1.id = P_2.planet_id GROUP BY P_1.id", 7, 2, None),
 #        ("SELECT P_1.* FROM $planets AS P_1 INNER JOIN $satellites AS P_2 ON P_1.id = P_2.planet_id INNER JOIN $planets AS P_3 ON P_2.planet_id = P_3.id", 9, 20, None),
 #        ("SELECT P_1.* FROM $planets AS P_1 INNER JOIN $planets AS P_2 ON P_1.id = P_2.id AND P_1.id > 2", 7, 20, None),
-#        ("SELECT P_1.* FROM $planets AS P_1 INNER JOIN $satellites AS P_2 ON P_1.id = P_2.planet_id AND P_2.id != 5", 9, 20, None),
+#        ("SELECT P_1.* FROM $planets AS P_1 INNER JOIN $satellites AS P_2 ON P_1.id = P_2.planet_id AND P_2.id != 5", 177, 20, None),
         ("SELECT P_1.* FROM $planets AS P_1 INNER JOIN $planets AS P_2 ON P_1.id = P_2.id AND P_1.name = P_2.name AND P_1.mass = P_2.mass", 9, 20, None),
-]
-A = [
+
         ("SELECT * FROM $planets AS P_1 INNER JOIN $planets AS P_2 ON P_1.id = P_2.id AND P_2.name = P_1.name", 9, 40, None),
         ("SELECT * FROM $planets NATURAL JOIN generate_series(1, 5) as id", 5, 20, None),
         ("SELECT * FROM $planets AS P1 NATURAL JOIN $planets AS P2", 5, 20, None),
@@ -770,7 +769,8 @@
         ("SELECT DISTINCT planetId FROM $satellites FULL OUTER JOIN $planets ON $satellites.planetId = $planets.id", 8, 1, None),
         ("SELECT DISTINCT planetId FROM $satellites FULL JOIN $planets ON $satellites.planetId = $planets.id", 8, 1, None),
         ("SELECT planetId FROM $satellites FULL JOIN $planets ON $satellites.planetId = $planets.id", 179, 1, None),
-
+]
+A = [
         ("SELECT pid FROM ( SELECT id AS pid FROM $planets) AS SQ WHERE pid > 5", 4, 1, None),
         ("SELECT * FROM ( SELECT id AS pid FROM $planets) AS SQ WHERE pid > 5", 4, 1, None),
         ("SELECT * FROM ( SELECT COUNT(planetId) AS moons, planetId FROM $satellites GROUP BY planetId ) AS SQ WHERE moons > 10", 4, 2, None),