Adds create_temp_table to allow using dataframes in sessions

Author: rchowell

daft/catalog/__init__.py

@@ -42,14 +42,16 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
-from daft.daft import catalog as native_catalog
+from daft.daft import PyTableSource, catalog as native_catalog
 from daft.daft import PyIdentifier, PyTable
 from daft.logical.builder import LogicalPlanBuilder
 
 from daft.dataframe import DataFrame
 
 from typing import TYPE_CHECKING
 
+from daft.logical.schema import Schema
+
 if TYPE_CHECKING:
     from daft.dataframe.dataframe import ColumnInputType
 
@@ -286,9 +288,11 @@ class Table(ABC):
     """Interface for python table implementations."""
 
     @staticmethod
-    def from_df(dataframe: DataFrame) -> Table:
+    def from_df(name: str, dataframe: DataFrame) -> Table:
         """Returns a read-only table backed by the DataFrame."""
-        return PyTable.from_builder(dataframe._builder._builder)
+        from daft.catalog.__memory import MemoryTable
+
+        return MemoryTable(name, dataframe)
 
     @staticmethod
     def from_iceberg(obj: object) -> Table:
@@ -314,10 +318,24 @@ def from_unity(obj: object) -> Table:
     @staticmethod
     def _from_obj(obj: object) -> Table:
         """Returns a Daft Table from a supported object type or raises an error."""
-        if isinstance(obj, DataFrame):
-            return Table.from_df(obj)
         raise ValueError(f"Unsupported table type: {type(obj)}")
 
+    # TODO catalog APIs part 3
+    # @property
+    # @abstractmethod
+    # def name(self) -> str:
+    #     """Returns the table name."""
+
+    # TODO catalog APIs part 3
+    # @property
+    # @abstractmethod
+    # def inner(self) -> object | None:
+    #     """Returns the inner table object if this is an adapter."""
+
+    @abstractmethod
+    def read(self) -> DataFrame:
+        """Returns a DataFrame from this table."""
+
     # TODO deprecated catalog APIs #3819
     def to_dataframe(self) -> DataFrame:
         """DEPRECATED: Please use `read` instead; version 0.5.0!"""
@@ -327,14 +345,62 @@ def to_dataframe(self) -> DataFrame:
         )
         return self.read()
 
-    @abstractmethod
-    def read(self) -> DataFrame:
-        """Returns a DataFrame from this table."""
-
     def select(self, *columns: ColumnInputType) -> DataFrame:
         """Returns a DataFrame from this table with the selected columns."""
         return self.read().select(*columns)
 
     def show(self, n: int = 8) -> None:
         """Shows the first n rows from this table."""
         return self.read().show(n)
+
+
+class TableSource:
+    """A TableSource is used to create a new table; this could be a Schema or DataFrame."""
+
+    _source: PyTableSource
+
+    def __init__(self) -> None:
+        raise ValueError("We do not support creating a TableSource via __init__")
+
+    @staticmethod
+    def from_df(df: DataFrame) -> TableSource:
+        s = TableSource.__new__(TableSource)
+        s._source = PyTableSource.from_builder(df._builder._builder)
+        return s
+
+    @staticmethod
+    def _from_obj(obj: object = None) -> TableSource:
+        # TODO for future sources, consider https://github.com/Eventual-Inc/Daft/pull/2864
+        if obj is None:
+            return TableSource._from_none()
+        elif isinstance(obj, DataFrame):
+            return TableSource.from_df(obj)
+        elif isinstance(obj, str):
+            return TableSource._from_path(obj)
+        elif isinstance(obj, Schema):
+            return TableSource._from_schema(obj)
+        else:
+            raise Exception(f"Unknown table source: {obj}")
+
+    @staticmethod
+    def _from_none() -> TableSource:
+        # for creating temp mutable tables, but we don't have those yet
+        # s = TableSource.__new__(TableSource)
+        # s._source = PyTableSource.empty()
+        # return s
+        # todo temp workaround just use an empty schema
+        return TableSource._from_schema(Schema._from_fields([]))
+
+    @staticmethod
+    def _from_schema(schema: Schema) -> TableSource:
+        # we don't have mutable temp tables, so just make an empty view
+        # s = TableSource.__new__(TableSource)
+        # s._source = PyTableSource.from_schema(schema._schema)
+        # return s
+        # todo temp workaround until create_table is wired
+        return TableSource.from_df(DataFrame._from_pylist([]))
+
+    @staticmethod
+    def _from_path(path: str) -> TableSource:
+        # for supporting daft.create_table("t", "/path/to/data") <-> CREATE TABLE t AS '/path/to/my.data'
+        raise NotImplementedError("creating a table source from a path is not yet supported.")

daft/catalog/__memory.py

@@ -41,9 +41,11 @@ def get_table(self, name: str | Identifier) -> Table:
 class MemoryTable(Table):
     """An in-memory table holds a reference to an existing dataframe."""
 
+    _name: str
     _inner: DataFrame
 
-    def __init__(self, inner: DataFrame):
+    def __init__(self, name: str, inner: DataFrame):
+        self._name = name
         self._inner = inner
 
     ###

daft/daft/__init__.pyi

@@ -1867,11 +1867,13 @@ class PyIdentifier:
     def __repr__(self) -> str: ...
 
 class PyTable(Table):
+    def read(self):
+        LogicalPlanBuilder
+
+class PyTableSource:
     @staticmethod
     def from_builder(builder: LogicalPlanBuilder):
         PyTable
-    def read(self):
-        LogicalPlanBuilder
 
 ###
 # daft-session
@@ -1885,6 +1887,7 @@ class PySession:
     def attach_table(self, table: Table, alias: str): ...
     def detach_catalog(self, alias: str): ...
     def detach_table(self, alias: str): ...
+    def create_temp_table(self, name: str, source: PyTableSource, replace: bool): ...
     def current_catalog(self) -> Catalog: ...
     def get_catalog(self, name: str) -> Catalog: ...
     def get_table(self, name: PyIdentifier) -> Table: ...

daft/session.py

@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING
 
-from daft.catalog import Catalog, Identifier, Table
+from daft.catalog import Catalog, Identifier, Table, TableSource
 from daft.daft import PySession
 
 if TYPE_CHECKING:
@@ -15,19 +15,12 @@ class Session:
     _session: PySession
 
     def __init__(self):
-        raise NotImplementedError("We do not support creating a Session via __init__ ")
+        self._session = PySession.empty()
 
     ###
     # factory methods
     ###
 
-    @staticmethod
-    def empty() -> Session:
-        """Creates an empty session."""
-        s = Session.__new__(Session)
-        s._session = PySession.empty()
-        return s
-
     @staticmethod
     def _from_pysession(session: PySession) -> Session:
         """Creates a session from a rust session wrapper."""
@@ -39,7 +32,7 @@ def _from_pysession(session: PySession) -> Session:
     def _from_env() -> Session:
         """Creates a session from the environment's configuration."""
         # todo session builders, raise if DAFT_SESSION=0
-        return Session.empty()
+        return Session()
 
     ###
     # attach & detach
@@ -63,6 +56,15 @@ def detach_table(self, alias: str):
         """Detaches the table from this session."""
         return self._session.detach_table(alias)
 
+    ###
+    # create_*
+    ###
+
+    def create_temp_table(self, name: str, source: object | TableSource = None) -> Table:
+        """Creates a temp table scoped to this session's lifetime."""
+        s = source if isinstance(source, TableSource) else TableSource._from_obj(source)
+        return self._session.create_temp_table(name, s._source, replace=True)
+
     ###
     # session state
     ###

src/daft-catalog/src/python.rs

@@ -1,11 +1,12 @@
 use std::sync::Arc;
 
-use daft_core::prelude::SchemaRef;
+use daft_core::{prelude::SchemaRef, python::PySchema};
 use daft_logical_plan::{LogicalPlanRef, PyLogicalPlanBuilder};
 use pyo3::{exceptions::PyIndexError, intern, prelude::*};
 
 use crate::{
-    error::Result, global_catalog, Catalog, CatalogRef, Identifier, Table, TableRef, View,
+    error::Result, global_catalog, Catalog, CatalogRef, Identifier, Table, TableRef, TableSource,
+    View,
 };
 
 /// Read a table from the specified `DaftMetaCatalog`.
@@ -277,10 +278,41 @@ impl Table for PyTableWrapper {
     }
 }
 
+/// PyTableSource wraps either a schema or dataframe.
+#[pyclass]
+pub struct PyTableSource(TableSource);
+
+impl From<TableSource> for PyTableSource {
+    fn from(source: TableSource) -> Self {
+        Self(source)
+    }
+}
+
+/// PyTableSource -> TableSource
+impl AsRef<TableSource> for PyTableSource {
+    fn as_ref(&self) -> &TableSource {
+        &self.0
+    }
+}
+
+#[pymethods]
+impl PyTableSource {
+    #[staticmethod]
+    pub fn from_schema(schema: PySchema) -> PyTableSource {
+        Self(TableSource::Schema(schema.schema))
+    }
+
+    #[staticmethod]
+    pub fn from_builder(view: &PyLogicalPlanBuilder) -> PyTableSource {
+        Self(TableSource::View(view.builder.build()))
+    }
+}
+
 pub fn register_modules<'py>(parent: &Bound<'py, PyModule>) -> PyResult<Bound<'py, PyModule>> {
     parent.add_class::<PyCatalog>()?;
     parent.add_class::<PyIdentifier>()?;
     parent.add_class::<PyTable>()?;
+    parent.add_class::<PyTableSource>()?;
     // TODO deprecated catalog APIs #3819
     let module = PyModule::new(parent.py(), "catalog")?;
     module.add_wrapped(wrap_pyfunction!(py_read_table))?;

src/daft-connect/src/execute.rs

@@ -2,7 +2,7 @@ use std::{future::ready, rc::Rc, sync::Arc};
 
 use common_error::DaftResult;
 use common_file_formats::FileFormat;
-use daft_catalog::View;
+use daft_catalog::TableSource;
 use daft_context::get_context;
 use daft_dsl::LiteralValue;
 use daft_logical_plan::LogicalPlanBuilder;
@@ -234,19 +234,16 @@ impl ConnectSession {
                 )
             })?;
 
-        {
-            // TODO session should handle the pre-existence error
-            if !replace && self.session().has_table(&name.clone().into()) {
-                return Err(Status::internal("Dataframe view already exists"));
-            }
-        }
-
         let session = self.session_mut();
-        let view = View::from(input.build()).arced();
+        let source = TableSource::from(input);
 
-        session.attach_table(view, name).map_err(|e| {
-            Status::internal(textwrap::wrap(&format!("Error in Daft server: {e}"), 120).join("\n"))
-        })?;
+        session
+            .create_temp_table(name, &source, replace)
+            .map_err(|e| {
+                Status::internal(
+                    textwrap::wrap(&format!("Error in Daft server: {e}"), 120).join("\n"),
+                )
+            })?;
 
         let response = rb.result_complete_response();
         let stream = stream::once(ready(Ok(response)));

src/daft-session/src/python.rs

@@ -1,4 +1,6 @@
-use daft_catalog::python::{PyCatalogWrapper, PyIdentifier, PyTableWrapper};
+use daft_catalog::python::{
+    PyCatalogWrapper, PyIdentifier, PyTable, PyTableSource, PyTableWrapper,
+};
 use pyo3::prelude::*;
 
 use crate::Session;
@@ -19,8 +21,8 @@ impl PySession {
             .attach_catalog(PyCatalogWrapper::wrap(catalog), alias)?)
     }
 
-    pub fn current_catalog(&self, py: Python<'_>) -> PyResult<PyObject> {
-        self.0.current_catalog()?.to_py(py)
+    pub fn attach_table(&self, table: PyObject, alias: String) -> PyResult<()> {
+        Ok(self.0.attach_table(PyTableWrapper::wrap(table), alias)?)
     }
 
     pub fn detach_catalog(&self, alias: &str) -> PyResult<()> {
@@ -31,8 +33,19 @@ impl PySession {
         Ok(self.0.detach_table(alias)?)
     }
 
-    pub fn attach_table(&self, table: PyObject, alias: String) -> PyResult<()> {
-        Ok(self.0.attach_table(PyTableWrapper::wrap(table), alias)?)
+    pub fn create_temp_table(
+        &self,
+        name: String,
+        source: &PyTableSource,
+        replace: bool,
+    ) -> PyResult<PyTable> {
+        let table = self.0.create_temp_table(name, source.as_ref(), replace)?;
+        let table = PyTable::new(table);
+        Ok(table)
+    }
+
+    pub fn current_catalog(&self, py: Python<'_>) -> PyResult<PyObject> {
+        self.0.current_catalog()?.to_py(py)
     }
 
     pub fn get_catalog(&self, py: Python<'_>, name: &str) -> PyResult<PyObject> {

src/daft-session/src/session.rs

@@ -1,6 +1,6 @@
 use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
 
-use daft_catalog::{Bindings, CatalogRef, Identifier, TableRef};
+use daft_catalog::{Bindings, CatalogRef, Identifier, TableRef, TableSource, View};
 use uuid::Uuid;
 
 use crate::{
@@ -72,6 +72,32 @@ impl Session {
         Ok(())
     }
 
+    /// Creates a temp table scoped to this session from an existing view.
+    ///
+    /// TODO feat: consider making a CreateTableSource object for more complicated options.
+    ///
+    /// ```
+    /// CREATE [OR REPLACE] TEMP TABLE [IF NOT EXISTS] <name> <source>;
+    /// ```
+    pub fn create_temp_table(
+        &self,
+        name: impl Into<String>,
+        source: &TableSource,
+        replace: bool,
+    ) -> Result<TableRef> {
+        let name = name.into();
+        if !replace && self.state().tables.exists(&name) {
+            obj_already_exists_err!("Temporary table", &name.into())
+        }
+        // we don't have mutable temporary tables, only immutable views over dataframes.
+        let table = match source {
+            TableSource::Schema(_) => unsupported_err!("temporary table with schema"),
+            TableSource::View(plan) => View::from(plan.clone()).arced(),
+        };
+        self.state_mut().tables.insert(name, table.clone());
+        Ok(table)
+    }
+
     /// Returns the session's current catalog.
     pub fn current_catalog(&self) -> Result<CatalogRef> {
         self.get_catalog(&self.state().options.curr_catalog)