Merge pull request #2850 from mabel-dev/#2849

Return column names as per SELECT #2849

Author: joocer

opteryx/__version__.py

@@ -1,9 +1,9 @@
 # THIS FILE IS AUTOMATICALLY UPDATED DURING THE BUILD PROCESS
 # DO NOT EDIT THIS FILE DIRECTLY
 
-__build__ = 1652
+__build__ = 1654
 __author__ = "@joocer"
-__version__ = "0.26.0-beta.1652"
+__version__ = "0.26.0-beta.1654"
 
 # Store the version here so:
 # 1) we don't load dependencies by storing it in __init__.py

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/connectors/__init__.py

@@ -270,7 +270,7 @@ def connector_factory(dataset, statistics, **config):
             break
     else:
         # Check if dataset is a file or contains wildcards
-        has_wildcards = any(char in dataset for char in ['*', '?', '['])
+        has_wildcards = any(char in dataset for char in ["*", "?", "["])
         if os.path.isfile(dataset) or has_wildcards:
             from opteryx.connectors import file_connector
 
@@ -286,7 +286,7 @@ def connector_factory(dataset, statistics, **config):
     remove_prefix = connector_entry.pop("remove_prefix", False)
     if prefix and remove_prefix and dataset.startswith(prefix):
         # Remove the prefix. If there's a separator (. or //) after the prefix, skip it too
-        dataset = dataset[len(prefix):]
+        dataset = dataset[len(prefix) :]
         if dataset.startswith(".") or dataset.startswith("//"):
             dataset = dataset[1:] if dataset.startswith(".") else dataset[2:]
 

opteryx/connectors/aws_s3_connector.py

@@ -86,13 +86,13 @@ def __init__(self, credentials=None, **kwargs):
             )
 
         self.minio = Minio(end_point, access_key, secret_key, secure=secure)
-        
+
         # Only convert dots to path separators if the dataset doesn't already contain slashes
         # Dataset references like "my.dataset.table" use dots as separators
         # File paths like "bucket/path/file.parquet" already have slashes and should not be converted
         if OS_SEP not in self.dataset and "/" not in self.dataset:
             self.dataset = self.dataset.replace(".", OS_SEP)
-        
+
         # Check if dataset contains wildcards
         self.has_wildcards = paths.has_wildcards(self.dataset)
         if self.has_wildcards:
@@ -111,28 +111,28 @@ def get_list_of_blob_names(self, *, prefix: str) -> List[str]:
         else:
             list_prefix = prefix
             filter_pattern = None
-            
+
         bucket, object_path, _, _ = paths.get_parts(list_prefix)
         blobs = self.minio.list_objects(bucket_name=bucket, prefix=object_path, recursive=True)
-        
+
         blob_list = []
         for blob in blobs:
             if blob.object_name.endswith("/"):
                 continue
-                
+
             full_path = bucket + "/" + blob.object_name
-            
+
             # Check if blob has valid extension
             if ("." + full_path.split(".")[-1].lower()) not in VALID_EXTENSIONS:
                 continue
-            
+
             # If we have a wildcard pattern, filter by it
             if filter_pattern:
                 if paths.match_wildcard(filter_pattern, full_path):
                     blob_list.append(full_path)
             else:
                 blob_list.append(full_path)
-        
+
         return sorted(blob_list)
 
     def read_dataset(

opteryx/connectors/file_connector.py

@@ -136,10 +136,10 @@ def __init__(self, *args, **kwargs):
         if ".." in self.dataset or self.dataset[0] in ("\\", "/", "~"):
             # Don't find any datasets which look like path traversal
             raise DatasetNotFoundError(dataset=self.dataset)
-        
+
         # Check if dataset contains wildcards
-        self.has_wildcards = any(char in self.dataset for char in ['*', '?', '['])
-        
+        self.has_wildcards = any(char in self.dataset for char in ["*", "?", "["])
+
         if self.has_wildcards:
             # Expand wildcards to get list of files
             self.files = self._expand_wildcards(self.dataset)
@@ -150,43 +150,43 @@ def __init__(self, *args, **kwargs):
         else:
             self.files = [self.dataset]
             self.decoder = get_decoder(self.dataset)
-    
+
     def _expand_wildcards(self, pattern: str) -> List[str]:
         """
         Expand wildcard patterns in file paths while preventing path traversal.
-        
+
         Supports wildcards:
         - * matches any number of characters
-        - ? matches a single character  
+        - ? matches a single character
         - [range] matches a range of characters (e.g., [0-9], [a-z])
-        
+
         Args:
             pattern: File path pattern with wildcards
-            
+
         Returns:
             List of matching file paths
         """
         # Additional path traversal check after expansion
         if ".." in pattern:
             raise DatasetNotFoundError(dataset=pattern)
-        
+
         # Use glob to expand the pattern
         matched_files = glob.glob(pattern, recursive=False)
-        
+
         # Filter out any results that might have path traversal
         # This is an extra safety check
         safe_files = []
         for file_path in matched_files:
             if ".." not in file_path and os.path.isfile(file_path):
                 safe_files.append(file_path)
-        
+
         return sorted(safe_files)
 
     def read_dataset(
         self, columns: list = None, predicates: list = None, limit: int = None, **kwargs
     ) -> pyarrow.Table:
         rows_read = 0
-        
+
         # Iterate over all matched files
         for file_path in self.files:
             morsel = read_blob(
@@ -221,7 +221,7 @@ def get_dataset_schema(self) -> RelationSchema:
 
         # Use the first file to get the schema
         first_file = self.files[0]
-        
+
         try:
             file_descriptor = os.open(first_file, os.O_RDONLY | os.O_BINARY)
             size = os.path.getsize(first_file)

opteryx/connectors/gcp_cloudstorage_connector.py

@@ -97,7 +97,7 @@ def __init__(self, credentials=None, **kwargs):
         if OS_SEP not in self.dataset and "/" not in self.dataset:
             self.dataset = self.dataset.replace(".", OS_SEP)
         self.credentials = credentials
-        
+
         # Check if dataset contains wildcards
         self.has_wildcards = paths.has_wildcards(self.dataset)
         if self.has_wildcards:
@@ -231,9 +231,9 @@ def get_list_of_blob_names(self, *, prefix: str) -> List[str]:
                 name = blob["name"]
                 if not name.endswith(TUPLE_OF_VALID_EXTENSIONS):
                     continue
-                    
+
                 full_path = f"{bucket}/{name}"
-                
+
                 # If we have a wildcard pattern, filter by it
                 if filter_pattern:
                     if paths.match_wildcard(filter_pattern, full_path):

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):