Defines the catalog trait/abc across both sides

Author: rchowell

Cargo.lock

@@ -58,7 +58,12 @@ def refresh_logger() -> None:
 # Daft top-level imports
 ###
 
-from daft.catalog import Catalog, Identifier, load_catalog
+from daft.catalog import (
+    Catalog,
+    Identifier,
+    Table,
+    load_catalog,
+)
 from daft.convert import (
     from_arrow,
     from_dask_dataframe,
@@ -113,6 +118,7 @@ def refresh_logger() -> None:
     "Schema",
     "Series",
     "Session",
+    "Table",
     "TimeUnit",
     "coalesce",
     "col",

daft/__init__.py

@@ -1,12 +1,15 @@
-"""Add catalog documentation..
-"""
+"""The daft-catalog moduel documentation..."""
 
 from __future__ import annotations
 from abc import ABC, abstractmethod
 from typing import Sequence
 from collections.abc import Sequence
 from daft.daft import PyCatalog, PyIdentifier
-from daft.table import Source, Table
+from daft.dataframe import DataFrame
+from daft.expressions import Expression
+from daft.logical.schema import Schema
+from daft.recordbatch import MicroPartition
+
 
 def load_catalog(name: str, options: object | None = None) -> Catalog:
     """Loads a new catalog from the configuration options or creates an in-memory catalog if none given."""
@@ -15,15 +18,18 @@ def load_catalog(name: str, options: object | None = None) -> Catalog:
     else:
         return Catalog._from_some(name, options)
 
+
 class Catalog(ABC):
+    """Catalog documentation..."""
 
     def __repr__(self) -> str:
         return f"Catalog('{self.name()}')"
 
     @staticmethod
     def _from_none(name: str):
-        from daft.catalog.__memory import MemoryCatalog
-        return MemoryCatalog(name)
+        from daft.catalog.__temp import TempCatalog
+
+        return TempCatalog(name)
 
     @staticmethod
     def _from_some(name: str, options: object) -> Catalog:
@@ -140,14 +146,57 @@ def __len__(self) -> int:
     def __repr__(self) -> str:
         return f"Identifier('{self._identifier.__repr__()}')"
 
+
 # TODO make a sequence
 Namespace = tuple[str]
 
+
+# TODO for future sources, consider https://github.com/Eventual-Inc/Daft/pull/2864
+# pandas/arrow/arrow_record_batches/pydict
+TableSource = Schema | DataFrame | str | None
+
+
+class Table(ABC):
+    """Table documentation..."""
+
+    def __repr__(self) -> str:
+        return f"Table('{self._name}')"
+
+    @abstractmethod
+    def name(self) -> str:
+        """Returns the table name."""
+
+    @abstractmethod
+    def schema(self) -> Schema:
+        """Returns the table schema."""
+
+    ###
+    # Creation Methods
+    ###
+
+    @staticmethod
+    def _from_source(name: str, source: TableSource = None) -> Table:
+        from daft.catalog.__temp import TempTable
+        return TempTable._from_source(name, source)
+
+    ###
+    # DataFrame Methods
+    ###
+
+    @abstractmethod
+    def read(self) -> DataFrame:
+        """Returns a DataFrame from this table."""
+
+    @abstractmethod
+    def show(self, n: int = 8) -> None:
+        """Shows the first n rows from this table."""
+
+
 __all__ = [
     "Catalog",
     "Identifier",
     "Namespace",
-    "Source",
     "Table",
+    "TableSource",
     "load_catalog",
 ]

daft/catalog/__init__.py

@@ -1,18 +1,29 @@
 from __future__ import annotations
 
-from daft.dataframe import DataFrame
-from daft.expressions import Expression
+from daft.catalog import Catalog, Table, TableSource
+from daft.dataframe.dataframe import DataFrame
 from daft.logical.schema import Schema
-from daft.recordbatch import MicroPartition
+from daft.recordbatch.micropartition import MicroPartition
 
-# for future sources, consider https://github.com/Eventual-Inc/Daft/pull/2864
-# pandas/arrow/arrow_record_batches/pydict
-Source = Schema | DataFrame | str | None
 
-class Table:
+class TempCatalog(Catalog):
+    """A temporary catalog scoped to a given session."""
+
+    def __init__(self, name: str):
+        self._name: str = name
+
+    def __repr__(self) -> str:
+        return f"TempCatalog('{self._name}')"
+
+    def name(self) -> str:
+        return self._name
+
+
+class TempTable(Table):
+    """A temp table holds a reference to an existing dataframe."""
 
     def __init__(self) -> Table:
-        raise NotImplementedError("Creating a Table via __init__ is not supported")
+        raise NotImplementedError("Creating a TempTable via __init__ is not supported")
 
     def name(self) -> str:
         return self._name
@@ -23,12 +34,8 @@ def schema(self) -> Schema:
     def __repr__(self) -> str:
         return f"table('{self._name}')"
 
-    ###
-    # Creation Methods
-    ###
-
     @staticmethod
-    def _from_source(name: str, source: Source = None) -> Table:
+    def _from_source(name: str, source: TableSource = None) -> Table:
         if source is None:
             return Table._from_none(name)
         elif isinstance(source, DataFrame):
@@ -74,8 +81,3 @@ def read(self) -> DataFrame:
 
     def show(self, n: int = 8) -> None:
         return self._inner.show(n)
-
-
-__all__ = [
-    "Table",
-]

daft/catalog/__memory.py

@@ -2,9 +2,8 @@
 
 from daft.context import DaftContext, get_context
 from daft.dataframe import DataFrame
-from daft.catalog import Catalog, Namespace
+from daft.catalog import Catalog, Identifier, Namespace, Table, TableSource
 from daft.daft import PySession
-from daft.table import Source, Table
 
 class Session:
 
@@ -57,14 +56,20 @@ def create_catalog(self, name: str) -> Catalog:
         """Create a new catalog scoped to this session."""
         return self._session.create_catalog(name)
 
-    # def create_namespace(self, name: str) -> Namespace:
-    #     """Create a new namespace scope to this session's current catalog."""
-    #     return self._session.create_namespace(name)
+    def create_namespace(self, name: str) -> Namespace:
+        """Create a new namespace scope to this session's current catalog."""
+        return self._session.create_namespace(name)
 
-    def create_table(self, name: str, source: Source = None) -> Table:
+    def create_table(self, name: str, source: TableSource = None) -> Table:
         """Creates a new table scoped to this session's current catalog and namespace."""
         return self._session.create_table(name, source)
 
+    def create_temp_table(self, name: str, source: TableSource = None) -> Table:
+        """Creates a temp table scoped to this session's lifetime."""
+        # TODO implement TableSource on the rust side.
+        source = Table._from_source(name, source)
+        return self._session.create_temp_table(name, source)
+
     ###
     # session state
     ###
@@ -73,9 +78,9 @@ def current_catalog(self) -> Catalog:
         """Returns the session's current catalog."""
         return self._session.current_catalog()
 
-    # def current_namespace(self) -> Namespace:
-    #     """Returns the session's current namespace."""
-    #     return self._session.current_namespace()
+    def current_namespace(self) -> Namespace:
+        """Returns the session's current namespace."""
+        return self._session.current_namespace()
 
     ###
     # get_*
@@ -89,8 +94,10 @@ def get_namespace(self, name: str) -> Namespace:
         """Returns the namespace or raises an exception if it does not exist."""
         return self._session.get_namespace(name)
 
-    def get_table(self, name: str) -> Table:
-        """Returns the table ."""
+    def get_table(self, name: str | Identifier) -> Table:
+        """Returns the table or raises an exception if it does not exist."""
+        if isinstance(name, str):
+            name = Identifier(*name.split("."))
         return self._session.get_table(name)
 
     ###

daft/table/__init__.py daft/catalog/__temp.py

@@ -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

daft/session.py

@@ -0,0 +1,55 @@
+use std::collections::HashMap;
+
+/// Bindings
+///
+/// Notes:
+///  - Currently using Clone because all references are arcs.
+///  - This is intentionally lightweight and everything is exact-case.
+///  - All APIs are non-fallible because callers determine what is an error.
+///  - It does not necessarily have map semantics.
+///  - Callers are responsible for case-normalization hence String, &str.
+///  - Intentionally using String and &str rather than Into and AsRef.
+#[derive(Debug)]
+pub struct Bindings<T>(HashMap<String, T>);
+
+impl<T> Bindings<T> {
+    /// Creates an empty catalog provider.
+    pub fn empty() -> Self {
+        Self(HashMap::new())
+    }
+
+    /// Inserts a new binding with ownership.
+    pub fn insert(&mut self, name: String, object: T) {
+        self.0.insert(name, object);
+    }
+
+    /// Removes the binding if it exists.
+    pub fn remove(&mut self, name: &str) {
+        self.0.remove(name);
+    }
+
+    /// Returns true if the binding exists.
+    pub fn exists(&self, name: &str) -> bool {
+        self.0.contains_key(name)
+    }
+
+    /// Get an object reference by name.
+    pub fn get(&self, name: &str) -> Option<&T> {
+        self.0.get(name)
+    }
+
+    /// List all objects matching the pattern.
+    pub fn list(&self, pattern: Option<&str>) -> Vec<String> {
+        self.0
+            .keys()
+            .map(|k| k.as_str())
+            .filter(|k| pattern.is_none() || k.contains(pattern.unwrap_or("")))
+            .map(|k| k.to_string())
+            .collect()
+    }
+
+    /// Returns true iff there are no bindings.
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+}