Ports tests and removes deprecated APIs

Author: rchowell

Cargo.lock

@@ -12,7 +12,6 @@ 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}
@@ -50,15 +49,17 @@ sysinfo = {workspace = true}
 [features]
 # maturin will turn this on
 python = [
+  "dep:pyo3",
+  "dep:pyo3-log",
   "common-daft-config/python",
   "common-display/python",
   "common-partitioning/python",
   "common-resource-request/python",
   "common-file-formats/python",
   "common-scan-info/python",
   "common-system-info/python",
-  "daft-catalog-python-catalog/python",
   "daft-catalog/python",
+  "daft-connect/python",
   "daft-context/python",
   "daft-core/python",
   "daft-csv/python",
@@ -80,12 +81,7 @@ python = [
   "daft-session/python",
   "daft-stats/python",
   "daft-recordbatch/python",
-  "daft-writers/python",
-  "dep:daft-catalog-python-catalog",
-  "dep:daft-connect",
-  "daft-connect/python",
-  "dep:pyo3",
-  "dep:pyo3-log"
+  "daft-writers/python"
 ]
 
 [lib]

Cargo.toml

@@ -2,7 +2,6 @@
 common-error = {path = "../common/error", default-features = false}
 daft-core = {path = "../daft-core", default-features = false}
 daft-logical-plan = {path = "../daft-logical-plan", default-features = false}
-lazy_static = {workspace = true}
 pyo3 = {workspace = true, optional = true}
 sqlparser = {workspace = true}
 snafu.workspace = true

src/daft-catalog/Cargo.toml

@@ -32,11 +32,8 @@ pub enum Error {
     Unsupported { message: String },
 
     #[cfg(feature = "python")]
-    #[snafu(display("Python error during {}: {}", context, source))]
-    PythonError {
-        source: pyo3::PyErr,
-        context: String,
-    },
+    #[snafu(display("Python error: {}", source))]
+    PythonError { source: pyo3::PyErr },
 }
 
 impl Error {
@@ -81,3 +78,10 @@ impl From<Error> for PyErr {
         daft_error.into()
     }
 }
+
+#[cfg(feature = "python")]
+impl From<PyErr> for Error {
+    fn from(value: PyErr) -> Self {
+        Error::PythonError { source: value }
+    }
+}

src/daft-catalog/python-catalog/Cargo.toml

@@ -16,227 +16,3 @@ pub use python::register_modules;
 
 // TODO audit daft-catalog and daft-session errors.
 pub mod error;
-
-// ----------------------------------
-// TODO deprecated catalog APIs #3819
-// ----------------------------------
-
-mod data_catalog;
-mod data_catalog_table;
-
-// Export public-facing traits
-use std::{collections::HashMap, default, sync::Arc};
-
-use daft_logical_plan::LogicalPlanBuilder;
-pub use data_catalog::DataCatalog;
-pub use data_catalog_table::DataCatalogTable;
-use error::{Error, Result};
-
-pub mod global_catalog {
-    use std::sync::{Arc, RwLock};
-
-    use lazy_static::lazy_static;
-
-    use crate::{DaftCatalog, DataCatalog};
-
-    lazy_static! {
-        pub(crate) static ref GLOBAL_DAFT_META_CATALOG: RwLock<DaftCatalog> =
-            RwLock::new(DaftCatalog::new_from_env());
-    }
-
-    /// Register a DataCatalog with the global DaftMetaCatalog
-    pub fn register_catalog(catalog: Arc<dyn DataCatalog>, name: Option<&str>) {
-        GLOBAL_DAFT_META_CATALOG
-            .write()
-            .unwrap()
-            .register_catalog(catalog, name);
-    }
-
-    /// Unregisters a catalog with the global DaftMetaCatalog
-    pub fn unregister_catalog(name: Option<&str>) -> bool {
-        GLOBAL_DAFT_META_CATALOG
-            .write()
-            .unwrap()
-            .unregister_catalog(name)
-    }
-}
-
-/// Name of the default catalog
-static DEFAULT_CATALOG_NAME: &str = "default";
-
-/// The [`DaftMetaCatalog`] is a catalog of [`DataCatalog`] implementations
-///
-/// Users of Daft can register various [`DataCatalog`] with Daft, enabling
-/// discovery of tables across various [`DataCatalog`] implementations.
-#[derive(Debug, Clone, Default)]
-pub struct DaftCatalog {
-    /// Map of catalog names to the DataCatalog impls.
-    ///
-    /// NOTE: The default catalog is always named "default"
-    data_catalogs: HashMap<String, Arc<dyn DataCatalog>>,
-
-    /// LogicalPlans that were "named" and registered with Daft
-    named_tables: HashMap<String, LogicalPlanBuilder>,
-}
-
-impl DaftCatalog {
-    /// Create a `DaftMetaCatalog` from the current environment
-    pub fn new_from_env() -> Self {
-        // TODO: Parse a YAML file to produce the catalog
-        DaftCatalog {
-            data_catalogs: default::Default::default(),
-            named_tables: default::Default::default(),
-        }
-    }
-
-    /// Register a new [`DataCatalog`] with the `DaftMetaCatalog`.
-    ///
-    /// # Arguments
-    ///
-    /// * `catalog` - The [`DataCatalog`] to register.
-    pub fn register_catalog(&mut self, catalog: Arc<dyn DataCatalog>, name: Option<&str>) {
-        let name = name.unwrap_or(DEFAULT_CATALOG_NAME);
-        self.data_catalogs.insert(name.to_string(), catalog);
-    }
-
-    /// Unregister a [`DataCatalog`] from the `DaftMetaCatalog`.
-    ///
-    /// # Arguments
-    ///
-    /// * `name` - The name of the catalog to unregister. If None, the default catalog will be unregistered.
-    ///
-    /// # Returns
-    ///
-    /// Returns `true` if a catalog was successfully unregistered, `false` otherwise.
-    pub fn unregister_catalog(&mut self, name: Option<&str>) -> bool {
-        let name = name.unwrap_or(DEFAULT_CATALOG_NAME);
-        self.data_catalogs.remove(name).is_some()
-    }
-
-    /// Registers a LogicalPlan with a name in the DaftMetaCatalog
-    pub fn register_table(
-        &mut self,
-        name: &str,
-        view: impl Into<LogicalPlanBuilder>,
-    ) -> Result<()> {
-        // TODO this API is being removed, for now preserve the exact name as if it were delimited.
-        self.named_tables.insert(name.into(), view.into());
-        Ok(())
-    }
-
-    /// Check if a named table is registered in the DaftCatalog
-    pub fn contains_table(&self, name: &str) -> bool {
-        self.named_tables.contains_key(name)
-    }
-
-    /// Provides high-level functionality for reading a table of data against a [`DaftMetaCatalog`]
-    ///
-    /// Resolves the provided table_identifier against the catalog:
-    ///
-    /// 1. If there is an exact match for the provided `table_identifier` in the catalog's registered named tables, immediately return the exact match
-    /// 2. If the [`DaftMetaCatalog`] has a default catalog, we will attempt to resolve the `table_identifier` against the default catalog
-    /// 3. If the `table_identifier` is hierarchical (delimited by "."), use the first component as the Data Catalog name and resolve the rest of the components against
-    ///     the selected Data Catalog
-    pub fn read_table(&self, table_identifier: &str) -> error::Result<LogicalPlanBuilder> {
-        // If the name is an exact match with a registered view, return it.
-        if let Some(view) = self.named_tables.get(table_identifier) {
-            return Ok(view.clone());
-        }
-
-        let mut searched_catalog_name = "default";
-        let mut searched_table_name = table_identifier;
-
-        // Check the default catalog for a match
-        if let Some(default_data_catalog) = self.data_catalogs.get(DEFAULT_CATALOG_NAME) {
-            if let Some(tbl) = default_data_catalog.get_table(table_identifier)? {
-                return tbl.as_ref().to_logical_plan_builder();
-            }
-        }
-
-        // Try to parse the catalog name from the provided table identifier by taking the first segment, split by '.'
-        if let Some((catalog_name, table_name)) = table_identifier.split_once('.') {
-            if let Some(data_catalog) = self.data_catalogs.get(catalog_name) {
-                searched_catalog_name = catalog_name;
-                searched_table_name = table_name;
-                if let Some(tbl) = data_catalog.get_table(table_name)? {
-                    return tbl.as_ref().to_logical_plan_builder();
-                }
-            }
-        }
-
-        // Return the error containing the last catalog/table pairing that we attempted to search on
-        Err(Error::TableNotFound {
-            catalog_name: searched_catalog_name.to_string(),
-            table_id: searched_table_name.to_string(),
-        })
-    }
-
-    /// Copy from another catalog, using tables from other in case of conflict
-    pub fn copy_from(&mut self, other: &Self) {
-        for (name, plan) in &other.named_tables {
-            self.named_tables.insert(name.clone(), plan.clone());
-        }
-        for (name, catalog) in &other.data_catalogs {
-            self.data_catalogs.insert(name.clone(), catalog.clone());
-        }
-    }
-
-    /// TODO remove py register and read methods are moved to session
-    /// I cannot remove DaftMetaCatalog until I invert the dependency
-    /// so that the current register_ methods use the session rather than the catalog.
-    pub fn into_catalog_map(self) -> HashMap<String, Arc<dyn DataCatalog>> {
-        self.data_catalogs
-    }
-}
-
-#[cfg(test)]
-mod tests {
-    use std::sync::Arc;
-
-    use daft_core::prelude::*;
-    use daft_logical_plan::{
-        ops::Source, source_info::PlaceHolderInfo, ClusteringSpec, LogicalPlan, LogicalPlanRef,
-        SourceInfo,
-    };
-
-    use super::*;
-
-    fn mock_plan() -> LogicalPlanRef {
-        let schema = Arc::new(
-            Schema::new(vec![
-                Field::new("text", DataType::Utf8),
-                Field::new("id", DataType::Int32),
-            ])
-            .unwrap(),
-        );
-        LogicalPlan::Source(Source::new(
-            schema.clone(),
-            Arc::new(SourceInfo::PlaceHolder(PlaceHolderInfo {
-                source_schema: schema,
-                clustering_spec: Arc::new(ClusteringSpec::unknown()),
-                source_id: 0,
-            })),
-        ))
-        .arced()
-    }
-
-    #[test]
-    fn test_register_and_unregister_named_table() {
-        let mut catalog = DaftCatalog::new_from_env();
-        let plan = LogicalPlanBuilder::from(mock_plan());
-
-        // Register a table
-        assert!(catalog.register_table("test_table", plan.clone()).is_ok());
-    }
-
-    #[test]
-    fn test_read_registered_table() {
-        let mut catalog = DaftCatalog::new_from_env();
-        let plan = LogicalPlanBuilder::from(mock_plan());
-
-        catalog.register_table("test_table", plan).unwrap();
-
-        assert!(catalog.read_table("test_table").is_ok());
-        assert!(catalog.read_table("non_existent_table").is_err());
-    }
-}

src/daft-catalog/python-catalog/src/lib.rs

@@ -4,100 +4,7 @@ use daft_core::prelude::SchemaRef;
 use daft_logical_plan::{LogicalPlanRef, PyLogicalPlanBuilder};
 use pyo3::{exceptions::PyIndexError, intern, prelude::*};
 
-use crate::{
-    error::Result, global_catalog, Catalog, CatalogRef, Identifier, Table, TableRef, View,
-};
-
-/// Read a table from the specified `DaftMetaCatalog`.
-///
-/// TODO deprecated catalog APIs #3819
-///
-/// This function reads a table from a `DaftMetaCatalog` and returns a PyLogicalPlanBuilder
-/// object representing the plan required to read the table.
-///
-/// The provided `table_identifier` can be:
-///
-/// 1. Name of a registered dataframe/SQL view (manually registered using `DaftMetaCatalog.register_view`)
-/// 2. Name of a table within the default catalog (without inputting the catalog name) for example: `"my.table.name"`
-/// 3. Name of a fully-qualified table path with the catalog name for example: `"my_catalog.my.table.name"`
-///
-/// Args:
-///     table_identifier (str): The identifier of the table to read.
-///
-/// Returns:
-///     PyLogicalPlanBuilder: A PyLogicalPlanBuilder object representing the table's data.
-///
-/// Raises:
-///     DaftError: If the table cannot be read or the specified table identifier is not found.
-///
-/// Example:
-///     >>> import daft
-///     >>> df = daft.read_table("foo")
-#[pyfunction]
-#[pyo3(name = "read_table")]
-fn py_read_table(table_identifier: &str) -> PyResult<PyLogicalPlanBuilder> {
-    let logical_plan_builder = global_catalog::GLOBAL_DAFT_META_CATALOG
-        .read()
-        .unwrap()
-        .read_table(table_identifier)?;
-    Ok(PyLogicalPlanBuilder::new(logical_plan_builder))
-}
-
-/// Register a table with the global catalog.
-///
-/// TODO deprecated catalog APIs #3819
-///
-/// This function registers a table with the global `DaftMetaCatalog` using the provided
-/// table identifier and logical plan.
-///
-/// Args:
-///     table_identifier (str): The identifier to use for the registered table.
-///     logical_plan (PyLogicalPlanBuilder): The logical plan representing the table's data.
-///
-/// Returns:
-///     str: The table identifier used for registration.
-///
-/// Example:
-///     >>> import daft
-///     >>> df = daft.read_csv("data.csv")
-///     >>> daft.register_table("my_table", df)
-#[pyfunction]
-#[pyo3(name = "register_table")]
-fn py_register_table(
-    table_identifier: &str,
-    logical_plan: &PyLogicalPlanBuilder,
-) -> PyResult<String> {
-    global_catalog::GLOBAL_DAFT_META_CATALOG
-        .write()
-        .unwrap()
-        .register_table(table_identifier, logical_plan.builder.clone())?;
-    Ok(table_identifier.to_string())
-}
-
-/// Unregisters a catalog from the Daft catalog system
-///
-/// TODO deprecated catalog APIs #3819
-///
-/// This function removes a previously registered catalog from the Daft catalog system.
-///
-/// Args:
-///     catalog_name (Optional[str]): The name of the catalog to unregister. If None, the default catalog will be unregistered.
-///
-/// Returns:
-///     bool: True if a catalog was successfully unregistered, False otherwise.
-///
-/// Example:
-///     >>> import daft
-///     >>> daft.unregister_catalog("my_catalog")
-///     True
-#[pyfunction]
-#[pyo3(
-    name = "unregister_catalog",
-    signature = (catalog_name=None)
-)]
-pub fn py_unregister_catalog(catalog_name: Option<&str>) -> bool {
-    crate::global_catalog::unregister_catalog(catalog_name)
-}
+use crate::{error::Result, Catalog, CatalogRef, Identifier, Table, TableRef, View};
 
 /// PyCatalog implements the Catalog ABC for some Catalog trait impl (rust->py).
 #[pyclass]
@@ -265,27 +172,32 @@ impl PyTableWrapper {
 
 impl Table for PyTableWrapper {
     fn get_schema(&self) -> SchemaRef {
-        todo!()
+        todo!("Table.schema not yet defined")
     }
 
     fn get_logical_plan(&self) -> Result<LogicalPlanRef> {
-        todo!()
+        Python::with_gil(|py| {
+            let builder = self
+                .0
+                .bind(py)
+                .call_method0("read")?
+                .getattr("_builder")?
+                .getattr("_builder")?;
+            let builder = builder
+                .downcast::<PyLogicalPlanBuilder>()
+                .expect("downcast PyAny to PyLogicalPlanBuilder");
+            Ok(builder.borrow().builder.plan.clone())
+        })
     }
 
     fn to_py(&self, py: Python<'_>) -> PyResult<PyObject> {
         self.0.extract(py)
     }
 }
 
-pub fn register_modules<'py>(parent: &Bound<'py, PyModule>) -> PyResult<Bound<'py, PyModule>> {
+pub fn register_modules(parent: &Bound<'_, PyModule>) -> PyResult<()> {
     parent.add_class::<PyCatalog>()?;
     parent.add_class::<PyIdentifier>()?;
     parent.add_class::<PyTable>()?;
-    // TODO deprecated catalog APIs #3819
-    let module = PyModule::new(parent.py(), "catalog")?;
-    module.add_wrapped(wrap_pyfunction!(py_read_table))?;
-    module.add_wrapped(wrap_pyfunction!(py_register_table))?;
-    module.add_wrapped(wrap_pyfunction!(py_unregister_catalog))?;
-    parent.add_submodule(&module)?;
-    Ok(module)
+    Ok(())
 }

src/daft-catalog/python-catalog/src/python.rs

@@ -63,6 +63,12 @@ impl From<LogicalPlanRef> for View {
     }
 }
 
+impl From<LogicalPlanBuilder> for View {
+    fn from(value: LogicalPlanBuilder) -> Self {
+        Self(value.plan)
+    }
+}
+
 impl View {
     pub fn arced(self) -> Arc<View> {
         Arc::new(self)

src/daft-catalog/src/data_catalog.rs

@@ -1,5 +1,6 @@
 [dependencies]
 daft-catalog = {path = "../daft-catalog"}
+daft-core = {path = "../daft-core", default-features = false}
 daft-logical-plan = {path = "../daft-logical-plan"}
 pyo3 = {workspace = true, optional = true}
 uuid = {version = "1.10.0", features = ["v4"]}
@@ -8,6 +9,7 @@ uuid = {version = "1.10.0", features = ["v4"]}
 python = [
   "dep:pyo3",
   "daft-catalog/python",
+  "daft-core/python",
   "daft-logical-plan/python"
 ]
 

src/daft-catalog/src/data_catalog_table.rs

@@ -153,3 +153,62 @@ impl Default for Session {
         Self::empty()
     }
 }
+
+/// Migrated from daft-catalog DaftMetaCatalog tests
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use daft_catalog::View;
+    use daft_core::prelude::*;
+    use daft_logical_plan::{
+        ops::Source, source_info::PlaceHolderInfo, ClusteringSpec, LogicalPlan, LogicalPlanBuilder,
+        LogicalPlanRef, SourceInfo,
+    };
+
+    use super::*;
+
+    fn mock_plan() -> LogicalPlanRef {
+        let schema = Arc::new(
+            Schema::new(vec![
+                Field::new("text", DataType::Utf8),
+                Field::new("id", DataType::Int32),
+            ])
+            .unwrap(),
+        );
+        LogicalPlan::Source(Source::new(
+            schema.clone(),
+            Arc::new(SourceInfo::PlaceHolder(PlaceHolderInfo {
+                source_schema: schema,
+                clustering_spec: Arc::new(ClusteringSpec::unknown()),
+                source_id: 0,
+            })),
+        ))
+        .arced()
+    }
+
+    #[test]
+    fn test_attach_table() {
+        let sess = Session::empty();
+        let plan = LogicalPlanBuilder::from(mock_plan());
+        let view = View::from(plan).arced();
+
+        // Register a table
+        assert!(sess.attach_table(view, "test_table").is_ok());
+    }
+
+    #[test]
+    fn test_get_table() {
+        let sess = Session::empty();
+        let plan = LogicalPlanBuilder::from(mock_plan());
+        let view = View::from(plan).arced();
+
+        sess.attach_table(view, "test_table")
+            .expect("failed to attach table");
+
+        assert!(sess.get_table(&Identifier::simple("test_table")).is_ok());
+        assert!(sess
+            .get_table(&Identifier::simple("non_existent_table"))
+            .is_err());
+    }
+}