add append_record_batch (#194)

* add append_record_batch

* reset default features

* add to_duckdb_type_id

* refactor code

---------

Co-authored-by: simon <[email protected]>

Cargo.toml

@@ -30,7 +30,8 @@ vtab = []
 vtab-loadable = ["vtab", "duckdb-loadable-macros"]
 vtab-excel = ["vtab", "calamine"]
 vtab-arrow = ["vtab", "num"]
-vtab-full = ["vtab-excel", "vtab-arrow"]
+appender-arrow = ["vtab-arrow"]
+vtab-full = ["vtab-excel", "vtab-arrow", "appender-arrow"]
 extensions-full = ["httpfs", "json", "parquet", "vtab-full"]
 buildtime_bindgen = ["libduckdb-sys/buildtime_bindgen"]
 modern-full = ["chrono", "serde_json", "url", "r2d2", "uuid", "polars"]

README.md

@@ -129,7 +129,7 @@ See to [Contributing.md](CONTRIBUTING.md)
 
 ### Checklist
 
-- Run `cargo fmt` to ensure your Rust code is correctly formatted.
+- Run `cargo +nightly fmt` to ensure your Rust code is correctly formatted.
 - Run `cargo clippy --fix --allow-dirty --all-targets --workspace --all-features -- -D warnings` to fix all clippy issues.
 - Ensure `cargo test --all-targets --workspace --features "modern-full extensions-full"` reports no failures.
 

src/appender/arrow.rs

@@ -0,0 +1,82 @@
+use super::{ffi, Appender, Result};
+use crate::{
+    error::result_from_duckdb_appender,
+    vtab::{record_batch_to_duckdb_data_chunk, to_duckdb_logical_type, DataChunk, LogicalType},
+    Error,
+};
+use arrow::record_batch::RecordBatch;
+use ffi::duckdb_append_data_chunk;
+
+impl Appender<'_> {
+    /// Append one record_batch
+    ///
+    /// ## Example
+    ///
+    /// ```rust,no_run
+    /// # use duckdb::{Connection, Result, params};
+    ///   use arrow::record_batch::RecordBatch;
+    /// fn insert_record_batch(conn: &Connection,record_batch:RecordBatch) -> Result<()> {
+    ///     let mut app = conn.appender("foo")?;
+    ///     app.append_record_batch(record_batch)?;
+    ///     Ok(())
+    /// }
+    /// ```
+    ///
+    /// # Failure
+    ///
+    /// Will return `Err` if append column count not the same with the table schema
+    #[inline]
+    pub fn append_record_batch(&mut self, record_batch: RecordBatch) -> Result<()> {
+        let schema = record_batch.schema();
+        let mut logical_type: Vec<LogicalType> = vec![];
+        for field in schema.fields() {
+            let logical_t = to_duckdb_logical_type(field.data_type())
+                .map_err(|_op| Error::ArrowTypeToDuckdbType(field.to_string(), field.data_type().clone()))?;
+            logical_type.push(logical_t);
+        }
+
+        let mut data_chunk = DataChunk::new(&logical_type);
+        record_batch_to_duckdb_data_chunk(&record_batch, &mut data_chunk).map_err(|_op| Error::AppendError)?;
+
+        let rc = unsafe { duckdb_append_data_chunk(self.app, data_chunk.get_ptr()) };
+        result_from_duckdb_appender(rc, self.app)
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::{Connection, Result};
+    use arrow::{
+        array::{Int8Array, StringArray},
+        datatypes::{DataType, Field, Schema},
+        record_batch::RecordBatch,
+    };
+    use std::sync::Arc;
+
+    #[test]
+    fn test_append_record_batch() -> Result<()> {
+        let db = Connection::open_in_memory()?;
+        db.execute_batch("CREATE TABLE foo(id TINYINT not null,area TINYINT not null,name Varchar)")?;
+        {
+            let id_array = Int8Array::from(vec![1, 2, 3, 4, 5]);
+            let area_array = Int8Array::from(vec![11, 22, 33, 44, 55]);
+            let name_array = StringArray::from(vec![Some("11"), None, None, Some("44"), None]);
+            let schema = Schema::new(vec![
+                Field::new("id", DataType::Int8, true),
+                Field::new("area", DataType::Int8, true),
+                Field::new("area", DataType::Utf8, true),
+            ]);
+            let record_batch = RecordBatch::try_new(
+                Arc::new(schema),
+                vec![Arc::new(id_array), Arc::new(area_array), Arc::new(name_array)],
+            )
+            .unwrap();
+            let mut app = db.appender("foo")?;
+            app.append_record_batch(record_batch)?;
+        }
+        let mut stmt = db.prepare("SELECT id, area,name  FROM foo")?;
+        let rbs: Vec<RecordBatch> = stmt.query_arrow([])?.collect();
+        assert_eq!(rbs.iter().map(|op| op.num_rows()).sum::<usize>(), 5);
+        Ok(())
+    }
+}

src/appender.rs → src/appender/mod.rs

@@ -13,6 +13,9 @@ pub struct Appender<'conn> {
     app: ffi::duckdb_appender,
 }
 
+#[cfg(feature = "appender-arrow")]
+mod arrow;
+
 impl Appender<'_> {
     /// Append multiple rows from Iterator
     ///

src/error.rs

@@ -1,3 +1,5 @@
+use arrow::datatypes::DataType;
+
 use super::Result;
 use crate::{
     ffi,
@@ -58,6 +60,9 @@ pub enum Error {
     /// Rust type.
     InvalidColumnType(usize, String, Type),
 
+    /// Error when datatype to duckdb type
+    ArrowTypeToDuckdbType(String, DataType),
+
     /// Error when a query that was expected to insert one row did not insert
     /// any or insert many.
     StatementChangedRows(usize),
@@ -170,6 +175,9 @@ impl fmt::Display for Error {
             Error::InvalidColumnType(i, ref name, ref t) => {
                 write!(f, "Invalid column type {t} at index: {i}, name: {name}")
             }
+            Error::ArrowTypeToDuckdbType(ref name, ref t) => {
+                write!(f, "Invalid column type {t} , name: {name}")
+            }
             Error::InvalidParameterCount(i1, n1) => {
                 write!(f, "Wrong number of parameters passed to query. Got {i1}, needed {n1}")
             }
@@ -201,6 +209,7 @@ impl error::Error for Error {
             | Error::StatementChangedRows(_)
             | Error::InvalidQuery
             | Error::AppendError
+            | Error::ArrowTypeToDuckdbType(..)
             | Error::MultipleStatement => None,
             Error::FromSqlConversionFailure(_, _, ref err) | Error::ToSqlConversionFailure(ref err) => Some(&**err),
         }

src/vtab/arrow.rs

@@ -119,6 +119,7 @@ impl VTab for ArrowVTab {
     }
 }
 
+/// Convert arrow DataType to duckdb type id
 pub fn to_duckdb_type_id(data_type: &DataType) -> Result<LogicalTypeId, Box<dyn std::error::Error>> {
     use LogicalTypeId::*;
 
@@ -160,6 +161,7 @@ pub fn to_duckdb_type_id(data_type: &DataType) -> Result<LogicalTypeId, Box<dyn
     Ok(type_id)
 }
 
+/// Convert arrow DataType to duckdb logical type
 pub fn to_duckdb_logical_type(data_type: &DataType) -> Result<LogicalType, Box<dyn std::error::Error>> {
     if data_type.is_primitive()
         || matches!(

src/vtab/data_chunk.rs

@@ -59,6 +59,11 @@ impl DataChunk {
     pub fn num_columns(&self) -> usize {
         unsafe { duckdb_data_chunk_get_column_count(self.ptr) as usize }
     }
+
+    /// Get the ptr of duckdb_data_chunk in this [DataChunk].
+    pub fn get_ptr(&self) -> duckdb_data_chunk {
+        self.ptr
+    }
 }
 
 impl From<duckdb_data_chunk> for DataChunk {

src/vtab/mod.rs

@@ -14,7 +14,7 @@ mod arrow;
 #[cfg(feature = "vtab-arrow")]
 pub use self::arrow::{
     arrow_arraydata_to_query_params, arrow_ffi_to_query_params, arrow_recordbatch_to_query_params,
-    record_batch_to_duckdb_data_chunk,
+    record_batch_to_duckdb_data_chunk, to_duckdb_logical_type, to_duckdb_type_id,
 };
 #[cfg(feature = "vtab-excel")]
 mod excel;