Adds back register_python_catalog and rebase on main

Author: rchowell

Cargo.lock

@@ -12,6 +12,7 @@ common-tracing = {path = "src/common/tracing", default-features = false}
 common-version = {path = "src/common/version", default-features = false}
 daft-algebra = {path = "src/daft-algebra", default-features = false}
 daft-catalog = {path = "src/daft-catalog", default-features = false}
+daft-catalog-python-catalog = {path = "src/daft-catalog/python-catalog", optional = true}
 daft-compression = {path = "src/daft-compression", default-features = false}
 daft-connect = {path = "src/daft-connect", optional = true}
 daft-context = {path = "src/daft-context", default-features = false}
@@ -56,6 +57,7 @@ python = [
   "common-file-formats/python",
   "common-scan-info/python",
   "common-system-info/python",
+  "daft-catalog-python-catalog/python",
   "daft-catalog/python",
   "daft-context/python",
   "daft-core/python",

Cargo.toml

@@ -34,11 +34,6 @@ def _from_obj(obj: object) -> IcebergCatalog:
             return c
         raise ValueError(f"Unsupported iceberg catalog type: {type(obj)}")
 
-    @property
-    def inner(self) -> InnerCatalog:
-        """Returns the inner iceberg catalog."""
-        return self._inner
-
     ###
     # get_*
     ###
@@ -72,11 +67,6 @@ def _from_obj(obj: object) -> IcebergTable | None:
             return t
         raise ValueError(f"Unsupported iceberg table type: {type(obj)}")
 
-    @property
-    def inner(self) -> InnerTable:
-        """Returns the inner iceberg table."""
-        return self._inner
-
     @staticmethod
     def _try_from(obj: object) -> IcebergTable | None:
         """Returns an IcebergTable if the given object can be adapted so."""

daft/catalog/__iceberg.py

@@ -43,7 +43,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from daft.daft import PyTableSource, catalog as native_catalog
-from daft.daft import PyIdentifier, PyTable
+from daft.daft import PyIdentifier
 from daft.logical.builder import LogicalPlanBuilder
 
 from daft.dataframe import DataFrame
@@ -185,10 +185,6 @@ def _from_obj(obj: object) -> Catalog:
             f"Unsupported catalog type: {type(obj)}; please ensure all required extra dependencies are installed."
         )
 
-    @property
-    def inner(self) -> object | None:
-        """Returns the inner catalog object if this is an adapter."""
-
     ###
     # list_*
     ###
@@ -314,7 +310,6 @@ def from_unity(obj: object) -> Table:
         except ImportError:
             raise ImportError("Unity support not installed: pip install -U 'getdaft[unity]'")
 
-
     @staticmethod
     def _from_obj(obj: object) -> Table:
         """Returns a Daft Table from a supported object type or raises an error."""

daft/catalog/__init__.py

@@ -82,11 +82,6 @@ def _from_obj(obj: object) -> UnityTable | None:
             return t
         raise ValueError(f"Unsupported unity table type: {type(obj)}")
 
-    @property
-    def inner(self) -> InnerTable:
-        """Returns the inner unity table."""
-        return self._inner
-
     @staticmethod
     def _try_from(obj: object) -> UnityTable | None:
         """Returns an UnityTable if the given object can be adapted so."""

daft/catalog/__unity.py

@@ -8,7 +8,12 @@ sqlparser = {workspace = true}
 snafu.workspace = true
 
 [features]
-python = ["dep:pyo3", "common-error/python", "daft-logical-plan/python", "daft-core/python"]
+python = [
+  "dep:pyo3",
+  "common-error/python",
+  "daft-core/python",
+  "daft-logical-plan/python"
+]
 
 [package]
 name = "daft-catalog"

src/daft-catalog/Cargo.toml

@@ -168,6 +168,19 @@ impl<'a> SQLPlanner<'a> {
         Ref::map(self.context.borrow(), |i| &i.bound_ctes)
     }
 
+    /// Lookup a table by identifier in the current session
+    fn get_table(&self, ident: &Identifier) -> Option<LogicalPlanBuilder> {
+        self.session()
+            .get_table(ident)
+            .map(|table| {
+                table
+                    .get_logical_plan()
+                    .expect("could not create a logical plan from the table")
+                    .into()
+            })
+            .ok()
+    }
+
     /// Borrow the planning session
     fn session(&self) -> Ref<'_, Rc<Session>> {
         Ref::map(self.context.borrow(), |i| &i.session)
@@ -887,13 +900,14 @@ impl<'a> SQLPlanner<'a> {
         let ident = normalize(name);
         let table = if ident.has_namespace() {
             // qualified search of sesison metadata
-            self.session().get_table(&ident).ok()
+            self.get_table(&ident)
         } else {
             // search bindings then session metadata
             self.bound_ctes()
                 .get(&ident.name)
                 .cloned()
-                .or_else(|| self.session().get_table(&ident).ok())
+                .or_else(|| self.bound_ctes().get(&ident.name).cloned())
+                .or_else(|| self.get_table(&ident))
         };
         table.ok_or_else(|| PlannerError::table_not_found(ident.to_string()))
     }

src/daft-sql/src/planner.rs

@@ -124,6 +124,11 @@ pub mod pylib {
         daft_session::register_modules(m)?;
         daft_sql::register_modules(m)?;
 
+        // Register `py_register_python_catalog` to daft.catalog
+        let catalog_attr = m.getattr("catalog")?;
+        let catalog_module = catalog_attr.downcast()?;
+        daft_catalog_python_catalog::python::register_modules(catalog_module)?;
+
         // Register testing module
         let testing_module = PyModule::new(m.py(), "testing")?;
         m.add_submodule(&testing_module)?;

src/lib.rs

@@ -23,3 +23,12 @@ def test_try_from_unity():
     # assert doesn't throw!
     assert Catalog._from_obj(unity_catalog) is not None
     assert Catalog.from_unity(unity_catalog) is not None
+
+
+def test_register_python_catalog():
+    import daft.unity_catalog
+
+    # sanity check for backwards compatibility
+    cat1 = daft.unity_catalog.UnityCatalog("", "")
+    daft.catalog.register_python_catalog(cat1, "test")
+    daft.catalog.unregister_catalog("test")