Fix iceberg naming issues

opteryx/__version__.py

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

opteryx/connectors/__init__.py

@@ -181,6 +181,11 @@ def register_store(prefix, connector, *, remove_prefix: bool = False, **kwargs):
         # uninstantiated classes aren't a type
         raise ValueError("connectors registered with `register_store` must be uninstantiated.")
 
+    if connector.__name__ == "IcebergConnector" and not remove_prefix:
+        raise ValueError(
+            "IcebergConnector requires remove_prefix=True so catalog prefixes don't leak into table names."
+        )
+
     # Store connector class directly (not as a string)
     _storage_prefixes[prefix] = {
         "connector": connector,  # type: ignore
@@ -289,6 +294,9 @@ def connector_factory(dataset, statistics, **config):
         dataset = dataset[len(prefix) :]
         if dataset.startswith(".") or dataset.startswith("//"):
             dataset = dataset[1:] if dataset.startswith(".") else dataset[2:]
+        if connector.__name__ == "IcebergConnector" and dataset and "." not in dataset:
+            # Default to a namespace matching the prefix when one isn't provided
+            dataset = f"{prefix}.{dataset}"
 
     return connector(dataset=dataset, statistics=statistics, **connector_entry)
 

opteryx/connectors/aws_s3_connector.py

@@ -176,7 +176,7 @@ def get_dataset_schema(self) -> RelationSchema:
         self.schema = next(self.read_dataset(just_schema=True), None)
 
         if self.schema is None:
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         return self.schema
 

opteryx/connectors/disk_connector.py

@@ -376,6 +376,6 @@ def get_dataset_schema(self) -> RelationSchema:
         if self.schema is None:
             if os.path.isdir(self.dataset):
                 raise EmptyDatasetError(dataset=self.dataset.replace(OS_SEP, "."))
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         return self.schema

opteryx/connectors/file_connector.py

@@ -135,7 +135,7 @@ 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)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         # Check if dataset contains wildcards
         self.has_wildcards = any(char in self.dataset for char in ["*", "?", "["])
@@ -144,7 +144,7 @@ def __init__(self, *args, **kwargs):
             # Expand wildcards to get list of files
             self.files = self._expand_wildcards(self.dataset)
             if not self.files:
-                raise DatasetNotFoundError(dataset=self.dataset)
+                raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
             # Use the first file to determine the decoder
             self.decoder = get_decoder(self.files[0])
         else:
@@ -168,7 +168,7 @@ def _expand_wildcards(self, pattern: str) -> List[str]:
         """
         # Additional path traversal check after expansion
         if ".." in pattern:
-            raise DatasetNotFoundError(dataset=pattern)
+            raise DatasetNotFoundError(dataset=pattern, connector=self.__type__)
 
         # Use glob to expand the pattern
         matched_files = glob.glob(pattern, recursive=False)

opteryx/connectors/gcp_cloudstorage_connector.py

@@ -318,7 +318,7 @@ def get_dataset_schema(self) -> RelationSchema:
             break
 
         if self.schema is None:
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         # if we have more than one blob we need to estimate the row count
         if self.schema.row_count_metric and number_of_blobs > 1:

opteryx/connectors/gcp_firestore_connector.py

@@ -126,7 +126,7 @@ def get_dataset_schema(self) -> RelationSchema:
         record = next(self.read_dataset(chunk_size=10), None)
 
         if record is None:
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         arrow_schema = record.schema
 

opteryx/connectors/iceberg_connector.py

@@ -148,12 +148,11 @@ def __init__(self, *args, catalog=None, io=DiskConnector, **kwargs):
 
         import pyiceberg
 
-        self.dataset = self.dataset.lower()
         try:
             self.table = catalog.load_table(self.dataset)
             self.io_connector = io(**kwargs)
         except pyiceberg.exceptions.NoSuchTableError:
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__) from None
 
     def get_dataset_schema(self) -> RelationSchema:
         iceberg_schema = self.table.schema()

opteryx/connectors/mongodb_connector.py

@@ -76,7 +76,7 @@ def get_dataset_schema(self) -> RelationSchema:
         record = next(self.read_dataset(chunk_size=25), None)
 
         if record is None:
-            raise DatasetNotFoundError(dataset=self.dataset)
+            raise DatasetNotFoundError(dataset=self.dataset, connector=self.__type__)
 
         arrow_schema = record.schema
 

opteryx/connectors/virtual_data.py

@@ -93,7 +93,9 @@ def read_dataset(self, columns: list = None, **kwargs) -> "DatasetReader":
     def get_dataset_schema(self) -> RelationSchema:
         if self.dataset not in WELL_KNOWN_DATASETS:
             suggestion = suggest(self.dataset)
-            raise DatasetNotFoundError(suggestion=suggestion, dataset=self.dataset)
+            raise DatasetNotFoundError(
+                suggestion=suggestion, dataset=self.dataset, connector=self.__type__
+            )
         data_provider, _ = _load_provider(self.dataset)
         self.relation_statistics = data_provider.statistics()
         return data_provider.schema()
@@ -138,6 +140,8 @@ def __next__(self) -> "pyarrow.Table":
         data_provider, _ = _load_provider(self.dataset_name)
         if data_provider is None:
             suggestion = suggest(self.dataset_name.lower())
-            raise DatasetNotFoundError(suggestion=suggestion, dataset=self.dataset_name)
+            raise DatasetNotFoundError(
+                suggestion=suggestion, dataset=self.dataset_name, connector="SAMPLE"
+            )
         table = data_provider.read(self.date, self.variables)
         return arrow.post_read_projector(table, self.columns)

opteryx/draken/morsels/morsel.pxd

@@ -11,6 +11,7 @@ cdef class Morsel:
     cpdef Vector column(self, bytes name)
     cpdef uint64_t[::1] hash(self, object columns=*)
     cdef void _take_inplace(self, object indices)
+    cdef void _empty_inplace(self)
     cdef void _select_inplace(self, object columns)
     cdef Morsel _full_copy(self)
     cdef inline void _rebuild_name_to_index(self)

opteryx/exceptions.py

@@ -172,9 +172,10 @@ def __init__(self, column: str):
 class DatasetNotFoundError(SqlError):
     """Exception raised when a dataset is not found."""
 
-    def __init__(self, dataset: str = None, suggestion: Optional[str] = None):
+    def __init__(self, connector: str, dataset: str = None, suggestion: Optional[str] = None):
         self.dataset = dataset
-        message = f"The requested dataset, '{dataset}', could not be found."
+        self.connector = connector
+        message = f"The requested dataset, '{dataset}', could not be found by '{connector}'."
         if suggestion is not None:
             message += f" Did you mean '{suggestion}'?"
         super().__init__(message)

opteryx/operators/bench/#information_schema_node.py

@@ -175,5 +175,5 @@ def execute(self) -> Iterable:
         elif self._dataset == "information_schema.routines":
             yield information_schema_routines()
         else:
-            raise DatasetNotFoundError(dataset=self._dataset)
+            raise DatasetNotFoundError(dataset=self._dataset, connector=self.__type__)
         return

opteryx/operators/distinct_node.py

@@ -39,26 +39,22 @@ def name(self):  # pragma: no cover
         return "Distinction"
 
     def execute(self, morsel: Table, **kwargs) -> Table:
-        morsel = self.ensure_arrow_table(morsel)
+        morsel = self.ensure_draken_morsel(morsel)
 
-        import opteryx.draken as draken
         from opteryx.compiled.table_ops.distinct import distinct
 
         if morsel == EOS:
             yield EOS
             return
 
-        # Convert Arrow table to Draken morsel
-        draken_morsel = draken.Morsel.from_arrow(morsel) if isinstance(morsel, Table) else morsel
-
         # Use Draken-based distinct with column names as bytes
         unique_indexes, self.hash_set = distinct(
-            draken_morsel, columns=self._distinct_on, seen_hashes=self.hash_set
+            morsel, columns=self._distinct_on, seen_hashes=self.hash_set
         )
 
         if len(unique_indexes) > 0:
-            distinct_table = morsel.take(unique_indexes)
-            yield distinct_table
+            morsel.take(unique_indexes)
+            yield morsel
         else:
-            distinct_table = morsel.slice(0, 0)
-            yield distinct_table
+            morsel.empty()
+            yield morsel

opteryx/operators/non_equi_join_node.py

@@ -117,4 +117,4 @@ def execute(self, morsel: Table, join_leg: str) -> Table:
             result_morsel = align_tables(
                 self.left_morsel, right_morsel, left_indexes, right_indexes
             )
-            yield result_morsel.to_arrow()
+            yield result_morsel

opteryx/operators/show_create_node.py

@@ -45,6 +45,6 @@ def execute(self, morsel: pyarrow.Table, **kwargs) -> pyarrow.Table:
                 yield table
                 return
 
-            raise DatasetNotFoundError(self.object_name)
+            raise DatasetNotFoundError(dataset=self.object_name, connector="VIEW")
 
         raise UnsupportedSyntaxError("Invalid SHOW statement")

opteryx/planner/optimizer/strategies/projection_pushdown.py

@@ -88,7 +88,7 @@ def visit(self, node: LogicalPlanNode, context: OptimizerContext) -> OptimizerCo
                 LogicalColumn(
                     node_type=NodeType.IDENTIFIER,
                     source_column=col.name,
-                    source=col.origin[0],
+                    source=(col.origin[0] if col.origin else None),
                     schema_column=col,
                 )
                 for col in node.schema.columns
@@ -106,7 +106,7 @@ def visit(self, node: LogicalPlanNode, context: OptimizerContext) -> OptimizerCo
                         LogicalColumn(
                             node_type=NodeType.IDENTIFIER,
                             source_column=col.name,
-                            source=col.origin[0],
+                            source=(col.origin[0] if col.origin else None),
                             schema_column=col,
                         )
                         for col in schema.columns

opteryx/planner/views/__init__.py

@@ -34,7 +34,7 @@ def view_as_plan(view_name: str) -> dict:
     from opteryx.utils.sql import remove_comments
 
     if not is_view(view_name):
-        raise DatasetNotFoundError(view_name)
+        raise DatasetNotFoundError(dataset=view_name, connector="VIEW")
 
     operation = view_as_sql(view_name)
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "opteryx"
-version = "0.26.2-beta.1862"
+version = "0.26.2-beta.1871"
 description = "Query your data, where it lives"
 requires-python = '>=3.11'
 readme = {file = "README.md", content-type = "text/markdown"}

tests/__init__.py

@@ -44,7 +44,7 @@ def test_example():
     if __name__ == "__main__":
         run_tests()
 """
-
+import contextlib
 import datetime
 import os
 import platform
@@ -678,6 +678,7 @@ def set_up_iceberg():
     import pyarrow
     import opteryx
     from pyiceberg.catalog.sql import SqlCatalog
+    from pyiceberg.exceptions import NamespaceAlreadyExistsError, NoSuchTableError
     from opteryx.connectors.iceberg_connector import IcebergConnector
 
     worker_id = os.environ.get('PYTEST_XDIST_WORKER', 'gw0')
@@ -724,18 +725,29 @@ def cast_dataset(dataset):
         },
     )
 
+    needs_setup = True
+    if existing:
+        try:
+            catalog.load_table("opteryx.planets")
+            catalog.load_table("opteryx.tweets")
+            needs_setup = False
+        except NoSuchTableError:
+            needs_setup = True
 
-    if existing: 
+    if not needs_setup:
         return catalog
 
-    catalog.create_namespace("iceberg")
+    with contextlib.suppress(NamespaceAlreadyExistsError):
+        catalog.create_namespace("opteryx")
 
     data = opteryx.query_to_arrow("SELECT tweet_id, text, timestamp, user_id, user_verified, user_name, hash_tags, followers, following, tweets_by_user, is_quoting, is_reply_to, is_retweeting FROM testdata.flat.formats.parquet")
-    table = catalog.create_table("iceberg.tweets", schema=data.schema)
+    table = catalog.create_table("opteryx.tweets", schema=data.schema)
     table.append(data.slice(0, 50000))
     table.append(data.slice(50000, 50000))
 
-    opteryx.register_store("iceberg", IcebergConnector, catalog=catalog)
+    opteryx.register_store(
+        "iceberg", IcebergConnector, catalog=catalog, remove_prefix=True
+    )
 
     for dataset in ('planets', 'satellites', 'missions', 'astronauts'):
         if dataset == 'planets':
@@ -762,7 +774,7 @@ def load_planet_snapshot(cutoff):
 
             snapshot_label, cutoff = snapshots[0]
             snapshot_data = load_planet_snapshot(cutoff)
-            table = catalog.create_table("iceberg.planets", schema=snapshot_data.schema)
+            table = catalog.create_table("opteryx.planets", schema=snapshot_data.schema)
             table.append(snapshot_data, snapshot_properties=snapshot_props(snapshot_label, cutoff))
 
             latest_snapshot = snapshot_data
@@ -775,22 +787,22 @@ def load_planet_snapshot(cutoff):
             del latest_snapshot
             del snapshot_data
 
-            iceberged = opteryx.query("SELECT * FROM iceberg.planets")
+            iceberged = opteryx.query("SELECT * FROM iceberg.opteryx.planets")
             assert iceberged.rowcount == expected_rows
             del iceberged  # Free memory immediately
             continue
 
         data = opteryx.query_to_arrow(f"SELECT * FROM ${dataset}")
         data = cast_dataset(data)
 
-        table = catalog.create_table(f"iceberg.{dataset}", schema=data.schema)
+        table = catalog.create_table(f"opteryx.{dataset}", schema=data.schema)
         table.append(data)
 
         # Verify row count without loading full result set into memory
         expected_rows = data.num_rows
         del data  # Free memory immediately
 
-        iceberged = opteryx.query(f"SELECT * FROM iceberg.{dataset}")
+        iceberged = opteryx.query(f"SELECT * FROM iceberg.opteryx.{dataset}")
         assert iceberged.rowcount == expected_rows
         del iceberged  # Free memory immediately