feat: add support for add to async python API (#1037)

In order to add support for `add` we needed to migrate the rust `Table`
trait to a `Table` struct and `TableInternal` trait (similar to the way
the connection is designed).

While doing this we also cleaned up some inconsistencies between the
SDKs:

* Python and Node are garbage collected languages and it can be
difficult to trigger something to be freed. The convention for these
languages is to have some kind of close method. I added a close method
to both the table and connection which will drop the underlying rust
object.
* We made significant improvements to table creation in
https://github.com/lancedb/lancedb/commit/cc5f2136a61234a9cdbe20e19c56459db8ca1051
for the `node` SDK. I copied these changes to the `nodejs` SDK.
* The nodejs tables were using fs to create tmp directories and these
were not getting cleaned up. This is mostly harmless but annoying and so
I changed it up a bit to ensure we cleanup tmp directories.
* ~~countRows in the node SDK was returning `bigint`. I changed it to
return `number`~~ (this actually happened in a previous PR)
* Tables and connections now implement `std::fmt::Display` which is
hooked into python's `__repr__`. Node has no concept of a regular "to
string" function and so I added a `display` method.
* Python method signatures are changing so that optional parameters are
always `Optional[foo] = None` instead of something like `foo = False`.
This is because we want those defaults to be in rust whenever possible
(though we still need to mention the default in documentation).
* I changed the python `AsyncConnection/AsyncTable` classes from
abstract classes with a single implementation to just classes because we
no longer have the remote implementation in python.

Note: this does NOT add the `add` function to the remote table. This PR
was already large enough, and the remote implementation is unique
enough, that I am going to do all the remote stuff at a later date (we
should have the structure in place and correct so there shouldn't be any
refactor concerns)

---------

Co-authored-by: Will Jones <willjones127@gmail.com>

python/src/connection.rs

@@ -17,7 +17,8 @@ use std::{sync::Arc, time::Duration};
 use arrow::{datatypes::Schema, ffi_stream::ArrowArrayStreamReader, pyarrow::FromPyArrow};
 use lancedb::connection::{Connection as LanceConnection, CreateTableMode};
 use pyo3::{
-    exceptions::PyValueError, pyclass, pyfunction, pymethods, PyAny, PyRef, PyResult, Python,
+    exceptions::{PyRuntimeError, PyValueError},
+    pyclass, pyfunction, pymethods, PyAny, PyRef, PyResult, Python,
 };
 use pyo3_asyncio::tokio::future_into_py;
 
@@ -25,7 +26,19 @@ use crate::{error::PythonErrorExt, table::Table};
 
 #[pyclass]
 pub struct Connection {
-    inner: LanceConnection,
+    inner: Option<LanceConnection>,
+}
+
+impl Connection {
+    pub(crate) fn new(inner: LanceConnection) -> Self {
+        Self { inner: Some(inner) }
+    }
+
+    fn get_inner(&self) -> PyResult<&LanceConnection> {
+        self.inner
+            .as_ref()
+            .ok_or_else(|| PyRuntimeError::new_err("Connection is closed"))
+    }
 }
 
 impl Connection {
@@ -41,8 +54,23 @@ impl Connection {
 
 #[pymethods]
 impl Connection {
+    fn __repr__(&self) -> String {
+        match &self.inner {
+            Some(inner) => inner.to_string(),
+            None => "ClosedConnection".to_string(),
+        }
+    }
+
+    fn is_open(&self) -> bool {
+        self.inner.is_some()
+    }
+
+    fn close(&mut self) {
+        self.inner.take();
+    }
+
     pub fn table_names(self_: PyRef<'_, Self>) -> PyResult<&PyAny> {
-        let inner = self_.inner.clone();
+        let inner = self_.get_inner()?.clone();
         future_into_py(self_.py(), async move {
             inner.table_names().await.infer_error()
         })
@@ -54,7 +82,7 @@ impl Connection {
         mode: &str,
         data: &PyAny,
     ) -> PyResult<&'a PyAny> {
-        let inner = self_.inner.clone();
+        let inner = self_.get_inner()?.clone();
 
         let mode = Self::parse_create_mode_str(mode)?;
 
@@ -76,7 +104,7 @@ impl Connection {
         mode: &str,
         schema: &PyAny,
     ) -> PyResult<&'a PyAny> {
-        let inner = self_.inner.clone();
+        let inner = self_.get_inner()?.clone();
 
         let mode = Self::parse_create_mode_str(mode)?;
 
@@ -92,6 +120,14 @@ impl Connection {
             Ok(Table::new(table))
         })
     }
+
+    pub fn open_table(self_: PyRef<'_, Self>, name: String) -> PyResult<&PyAny> {
+        let inner = self_.get_inner()?.clone();
+        future_into_py(self_.py(), async move {
+            let table = inner.open_table(&name).execute().await.infer_error()?;
+            Ok(Table::new(table))
+        })
+    }
 }
 
 #[pyfunction]
@@ -118,8 +154,6 @@ pub fn connect(
             let read_consistency_interval = Duration::from_secs_f64(read_consistency_interval);
             builder = builder.read_consistency_interval(read_consistency_interval);
         }
-        Ok(Connection {
-            inner: builder.execute().await.infer_error()?,
-        })
+        Ok(Connection::new(builder.execute().await.infer_error()?))
     })
 }

python/src/table.rs

@@ -1,34 +1,90 @@
-use std::sync::Arc;
-
-use arrow::pyarrow::ToPyArrow;
-use lancedb::table::Table as LanceTable;
-use pyo3::{pyclass, pymethods, PyAny, PyRef, PyResult, Python};
+use arrow::{
+    ffi_stream::ArrowArrayStreamReader,
+    pyarrow::{FromPyArrow, ToPyArrow},
+};
+use lancedb::table::{AddDataMode, Table as LanceDbTable};
+use pyo3::{
+    exceptions::{PyRuntimeError, PyValueError},
+    pyclass, pymethods, PyAny, PyRef, PyResult, Python,
+};
 use pyo3_asyncio::tokio::future_into_py;
 
 use crate::error::PythonErrorExt;
 
 #[pyclass]
 pub struct Table {
-    inner: Arc<dyn LanceTable>,
+    // We keep a copy of the name to use if the inner table is dropped
+    name: String,
+    inner: Option<LanceDbTable>,
 }
 
 impl Table {
-    pub(crate) fn new(inner: Arc<dyn LanceTable>) -> Self {
-        Self { inner }
+    pub(crate) fn new(inner: LanceDbTable) -> Self {
+        Self {
+            name: inner.name().to_string(),
+            inner: Some(inner),
+        }
+    }
+}
+
+impl Table {
+    fn inner_ref(&self) -> PyResult<&LanceDbTable> {
+        self.inner
+            .as_ref()
+            .ok_or_else(|| PyRuntimeError::new_err(format!("Table {} is closed", self.name)))
     }
 }
 
 #[pymethods]
 impl Table {
     pub fn name(&self) -> String {
-        self.inner.name().to_string()
+        self.name.clone()
+    }
+
+    pub fn is_open(&self) -> bool {
+        self.inner.is_some()
+    }
+
+    pub fn close(&mut self) {
+        self.inner.take();
     }
 
     pub fn schema(self_: PyRef<'_, Self>) -> PyResult<&PyAny> {
-        let inner = self_.inner.clone();
+        let inner = self_.inner_ref()?.clone();
         future_into_py(self_.py(), async move {
             let schema = inner.schema().await.infer_error()?;
             Python::with_gil(|py| schema.to_pyarrow(py))
         })
     }
+
+    pub fn add<'a>(self_: PyRef<'a, Self>, data: &PyAny, mode: String) -> PyResult<&'a PyAny> {
+        let batches = Box::new(ArrowArrayStreamReader::from_pyarrow(data)?);
+        let mut op = self_.inner_ref()?.add(batches);
+        if mode == "append" {
+            op = op.mode(AddDataMode::Append);
+        } else if mode == "overwrite" {
+            op = op.mode(AddDataMode::Overwrite);
+        } else {
+            return Err(PyValueError::new_err(format!("Invalid mode: {}", mode)));
+        }
+
+        future_into_py(self_.py(), async move {
+            op.execute().await.infer_error()?;
+            Ok(())
+        })
+    }
+
+    pub fn count_rows(self_: PyRef<'_, Self>, filter: Option<String>) -> PyResult<&PyAny> {
+        let inner = self_.inner_ref()?.clone();
+        future_into_py(self_.py(), async move {
+            inner.count_rows(filter).await.infer_error()
+        })
+    }
+
+    pub fn __repr__(&self) -> String {
+        match &self.inner {
+            None => format!("ClosedTable({})", self.name),
+            Some(inner) => inner.to_string(),
+        }
+    }
 }