Merge branch 'main' into copilot/create-fast-json-lines-decoder

Author: joocer

opteryx/compiled/joins/cross_join.pyx

@@ -161,11 +161,11 @@ cpdef tuple numpy_build_filtered_rows_indices_and_column(numpy.ndarray column_da
 
     # Handle set initialization based on element dtype
     if numpy.issubdtype(element_dtype, numpy.integer):
-        valid_values_typed = set([int(v) for v in valid_values])
+        valid_values_typed = {int(v) for v in valid_values}
     elif numpy.issubdtype(element_dtype, numpy.floating):
-        valid_values_typed = set([parse_fast_float(v) for v in valid_values])
+        valid_values_typed = {parse_fast_float(v) for v in valid_values}
     elif numpy.issubdtype(element_dtype, numpy.str_):
-        valid_values_typed = set([unicode(v) for v in valid_values])
+        valid_values_typed = {unicode(v) for v in valid_values}
     else:
         valid_values_typed = valid_values  # Fallback to generic Python set
 

opteryx/cursor.py

@@ -336,7 +336,25 @@ def execute_to_arrow(
         if isinstance(result_data, pyarrow.Table):
             return result_data
         try:
-            return pyarrow.concat_tables(result_data, promote_options="permissive")
+            # arrow allows duplicate column names, but not when concatting
+            from itertools import chain
+
+            first_table = next(result_data, None)
+            if first_table is not None:
+                column_names = first_table.column_names
+                if len(column_names) != len(set(column_names)):
+                    temporary_names = [f"col_{i}" for i in range(len(column_names))]
+                    first_table = first_table.rename_columns(temporary_names)
+                    return_table = pyarrow.concat_tables(
+                        chain(
+                            [first_table], (t.rename_columns(temporary_names) for t in result_data)
+                        ),
+                        promote_options="permissive",
+                    )
+                    return return_table.rename_columns(column_names)
+            return pyarrow.concat_tables(
+                chain([first_table], result_data), promote_options="permissive"
+            )
         except (
             pyarrow.ArrowInvalid,
             pyarrow.ArrowTypeError,

opteryx/operators/exit_node.py

@@ -40,7 +40,7 @@ def __init__(self, properties: QueryProperties, **parameters):
         final_names = []
         for column in self.columns:
             final_columns.append(column.schema_column.identity)
-            final_names.append(column.current_name)
+            final_names.append(column.alias)
 
         if len(final_columns) != len(set(final_columns)):  # pragma: no cover
             from collections import Counter
@@ -57,7 +57,7 @@ def __init__(self, properties: QueryProperties, **parameters):
                 # if column.schema_column.origin:
                 #    final_names.append(f"{column.schema_column.origin[0]}.{column.current_name}")
                 # else:
-                final_names.append(column.qualified_name)
+                final_names.append(column.alias)
 
         self.final_columns = final_columns
         self.final_names = final_names

opteryx/planner/binder/binder_visitor.py

@@ -380,14 +380,7 @@ def visit_exit(self, node: Node, context: BindingContext) -> Tuple[Node, Binding
         # clear the derived schema
         context.schemas.pop("$derived", None)
 
-        seen = set()
-        needs_qualifier = len(context.schemas) > 1 or any(
-            column.name in seen or seen.add(column.name) is not None  # type: ignore
-            for schema in context.schemas.values()
-            for column in schema.columns
-        )
-
-        def name_column(qualifier, column):
+        def name_column(column):
             for projection_column in node.columns:
                 if (
                     projection_column.schema_column
@@ -396,20 +389,11 @@ def name_column(qualifier, column):
                     if projection_column.alias:
                         return projection_column.alias
 
-                    if len(context.relations) > 1 or needs_qualifier:
-                        if isinstance(projection_column, LogicalColumn):
-                            if qualifier:
-                                projection_column.source = qualifier
-                            return projection_column.qualified_name
-                        return f"{qualifier}.{column.name}"
-
                     if projection_column.query_column:
                         return str(projection_column.query_column)
                     if projection_column.current_name:
                         return projection_column.current_name
 
-            if needs_qualifier:
-                return f"{qualifier}.{column.name}"
             return column.name
 
         def keep_column(column, identities):
@@ -441,15 +425,15 @@ def keep_column(column, identities):
                 identities.append(column.identity)
 
         columns = []
-        for qualifier, schema in context.schemas.items():
+        for _, schema in context.schemas.items():
             for column in schema.columns:
                 if keep_column(column, identities):
-                    column_name = name_column(qualifier=qualifier, column=column)
+                    column_name = name_column(column=column)
                     column_reference = LogicalColumn(
                         node_type=NodeType.IDENTIFIER,
                         source_column=column_name,
                         source=None,
-                        alias=None,
+                        alias=column_name,
                         schema_column=column,
                     )
                     columns.append(column_reference)

opteryx/planner/logical_planner/logical_planner_builders.py

@@ -271,12 +271,21 @@ def ceiling(value, alias: Optional[List[str]] = None, key=None):
 
 
 def compound_identifier(branch, alias: Optional[List[str]] = None, key=None):
-    return LogicalColumn(
+    column = LogicalColumn(
         node_type=NodeType.IDENTIFIER,  # column type
         alias=alias,  # type: ignore
         source_column=branch[-1]["value"],  # the source column
         source=".".join(p["value"] for p in branch[:-1]),  # the source relation
     )
+    alias_name = alias[0] if isinstance(alias, list) and alias else alias
+    if alias_name:
+        column.query_column = alias_name
+    else:
+        qualifier = column.source
+        column.query_column = (
+            f"{qualifier}.{column.source_column}" if qualifier else column.source_column
+        )
+    return column
 
 
 def expression_with_alias(branch, alias: Optional[List[str]] = None, key=None):
@@ -424,11 +433,14 @@ def identifier(branch, alias: Optional[List[str]] = None, key=None):
     """idenitifier doesn't have a qualifier (recorded in source)"""
     if "Identifier" in branch:
         return build(branch["Identifier"], alias=alias)
-    return LogicalColumn(
+    column = LogicalColumn(
         node_type=NodeType.IDENTIFIER,  # column type
         alias=alias,  # type: ignore
         source_column=branch["value"],  # the source column
     )
+    alias_name = alias[0] if isinstance(alias, list) and alias else alias
+    column.query_column = alias_name or column.source_column
+    return column
 
 
 def in_list(branch, alias: Optional[List[str]] = None, key=None):