Repository: mabel-dev/opteryx

Issue:
🪲 UNNEST constructor returns invalid column types
InvalidInternalStateError: Internal Error - function visit_function_dataset put unexpected items in 'columns' attribute


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
@@ -356,15 +356,24 @@ def visit_function_dataset(
             context.schemas[relation_name] = schema
             node.columns = columns
         elif node.function == "UNNEST":
+            column_to_unnest = node.args[0].value
+            if node.args[0].node_type == NodeType.LITERAL:
+                column_to_unnest = "<value>"
             if not node.alias:
-                if node.args[0].node_type == NodeType.IDENTIFIER:
-                    node.alias = f"UNNEST({node.args[0].value})"
+                node.alias = f"UNNEST({column_to_unnest})"
             relation_name = f"$unnest-{random_string()}"
-            schema = RelationSchema(
-                name=relation_name, columns=[FlatColumn(name=node.alias or "unnest", type=0)]
-            )
+
+            columns = [
+                LogicalColumn(
+                    node_type=NodeType.IDENTIFIER,
+                    source_column=node.alias,
+                    source=relation_name,
+                    schema_column=FlatColumn(name=node.alias, type=0),
+                )
+            ]
+            schema = RelationSchema(name=relation_name, columns=[c.schema_column for c in columns])
             context.schemas[relation_name] = schema
-            node.columns = [schema.columns[0].identity]
+            node.columns = columns
         elif node.function == "GENERATE_SERIES":
             schema = RelationSchema(
                 name=node.alias,
diff --git a/opteryx/operators/function_dataset_node.py b/opteryx/operators/function_dataset_node.py
--- a/opteryx/operators/function_dataset_node.py
+++ b/opteryx/operators/function_dataset_node.py
@@ -41,7 +41,7 @@ def _unnest(**kwargs):
         list_items = [kwargs["args"][0].centre.value]
     else:
         list_items = kwargs["args"][0].value
-    column_name = kwargs["columns"][0]
+    column_name = kwargs["columns"][0].schema_column.identity
     return [{column_name: row} for row in list_items]
 
 


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
@@ -555,14 +555,14 @@
 
         ("SELECT * FROM (VALUES ('High', 3),('Medium', 2),('Low', 1)) AS ratings(name, rating)", 3, 2, None),
         ("SELECT * FROM (VALUES ('High', 3),('Medium', 2),('Low', 1)) AS ratings(name, rating) WHERE rating = 3", 1, 2, None),
-]
-A = [
+
         ("SELECT * FROM UNNEST(('foo', 'bar', 'baz', 'qux', 'corge', 'garply', 'waldo', 'fred'))", 8, 1, UnnamedColumnError),
         ("SELECT * FROM UNNEST(('foo', 'bar', 'baz', 'qux', 'corge', 'garply', 'waldo', 'fred')) AS element", 8, 1, None),
         ("SELECT * FROM UNNEST(('foo', 'bar', 'baz', 'qux', 'corge', 'garply', 'waldo', 'fred')) AS element WHERE element LIKE '%e%'", 2, 1, None),
         ("SELECT * FROM UNNEST(('foo', 'bar', 'baz', 'qux', 'corge', 'garply', 'waldo', 'fred')) AS UN", 8, 1, None),
         ("SELECT * FROM UNNEST(('foo', 'bar', 'baz', 'qux', 'corge', 'garply', 'waldo', 'fred')) AS UN WHERE UN LIKE '%e%'", 2, 1, None),
-
+]
+A = [
         ("SELECT * FROM generate_series(1, 10)", 10, 1, UnnamedColumnError),
         ("SELECT * FROM generate_series(1, 10) AS GS", 10, 1, None),
         ("SELECT * FROM generate_series(-10,10) AS GS", 21, 1, None),