From d701947f0bce1672c8902db4c3f40b2aec348e27 Mon Sep 17 00:00:00 2001
From: Lei Xu <lei@lancedb.com>
Date: Wed, 5 Jul 2023 18:20:31 -0700
Subject: [PATCH] [Rust] Re-export WriteMode from lancedb instead of lance
 (#257)

`Table::add(.., mode: WriteMode)`, which is a public API, currently uses
the WriteMode exported from `lance`. Re-export it to lancedb so that the
pub API looks better.
---
 rust/ffi/node/src/lib.rs   |  6 ++++--
 rust/vectordb/src/lib.rs   |  3 +++
 rust/vectordb/src/table.rs | 13 +++++++++----
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/rust/ffi/node/src/lib.rs b/rust/ffi/node/src/lib.rs
index 608de1b7..40cf3743 100644
--- a/rust/ffi/node/src/lib.rs
+++ b/rust/ffi/node/src/lib.rs
@@ -239,7 +239,7 @@ fn table_create(mut cx: FunctionContext) -> JsResult<JsPromise> {
         "overwrite" => WriteMode::Overwrite,
         "append" => WriteMode::Append,
         "create" => WriteMode::Create,
-        _ => return cx.throw_error("Table::create only supports 'overwrite' and 'create' modes")
+        _ => return cx.throw_error("Table::create only supports 'overwrite' and 'create' modes"),
     };
     let mut params = WriteParams::default();
     params.mode = mode;
@@ -255,7 +255,9 @@ fn table_create(mut cx: FunctionContext) -> JsResult<JsPromise> {
             batches.into_iter().map(Ok),
             schema,
         ));
-        let table_rst = database.create_table(&table_name, batch_reader, Some(params)).await;
+        let table_rst = database
+            .create_table(&table_name, batch_reader, Some(params))
+            .await;
 
         deferred.settle_with(&channel, move |mut cx| {
             let table = Arc::new(Mutex::new(
diff --git a/rust/vectordb/src/lib.rs b/rust/vectordb/src/lib.rs
index c191784b..2a89851b 100644
--- a/rust/vectordb/src/lib.rs
+++ b/rust/vectordb/src/lib.rs
@@ -19,3 +19,6 @@ pub mod query;
 pub mod table;
 
 pub use database::Database;
+pub use table::Table;
+
+pub use lance::dataset::WriteMode;
diff --git a/rust/vectordb/src/table.rs b/rust/vectordb/src/table.rs
index 1967ce67..ac9544a6 100644
--- a/rust/vectordb/src/table.rs
+++ b/rust/vectordb/src/table.rs
@@ -16,12 +16,13 @@ use std::path::Path;
 use std::sync::Arc;
 
 use arrow_array::{Float32Array, RecordBatchReader};
-use lance::dataset::{Dataset, ReadParams, WriteMode, WriteParams};
+use lance::dataset::{Dataset, ReadParams, WriteParams};
 use lance::index::IndexType;
 use snafu::prelude::*;
 
 use crate::error::{Error, InvalidTableNameSnafu, Result};
 use crate::index::vector::VectorIndexBuilder;
+use crate::WriteMode;
 use crate::query::Query;
 
 pub const VECTOR_COLUMN_NAME: &str = "vector";
@@ -305,12 +306,16 @@ mod tests {
         let mut table = Table::create(&uri, "test", batches, None).await.unwrap();
         assert_eq!(table.count_rows().await.unwrap(), 10);
 
-        let new_batches: Box<dyn RecordBatchReader> =
-            Box::new(RecordBatchIterator::new(vec![RecordBatch::try_new(
+        let new_batches: Box<dyn RecordBatchReader> = Box::new(RecordBatchIterator::new(
+            vec![RecordBatch::try_new(
                 schema.clone(),
                 vec![Arc::new(Int32Array::from_iter_values(100..110))],
             )
-            .unwrap()].into_iter().map(Ok), schema.clone()));
+            .unwrap()]
+            .into_iter()
+            .map(Ok),
+            schema.clone(),
+        ));
 
         table.add(new_batches, None).await.unwrap();
         assert_eq!(table.count_rows().await.unwrap(), 20);