Support appending large Arrow record batches (#530)

Based on the test from #271 by @plaflamme, this PR enhances/fixes
`append_record_batch` to automatically chunk the passed record batch
into smaller pieces (using Arrow's zero-copy slicing) that fit within
DuckDB's vector size limit.

Author: mlafeldt

crates/duckdb/src/appender/arrow.rs

@@ -6,10 +6,10 @@ use crate::{
     Error,
 };
 use arrow::record_batch::RecordBatch;
-use ffi::duckdb_append_data_chunk;
+use ffi::{duckdb_append_data_chunk, duckdb_vector_size};
 
 impl Appender<'_> {
-    /// Append one record_batch
+    /// Append one record batch
     ///
     /// ## Example
     ///
@@ -28,27 +28,43 @@ impl Appender<'_> {
     /// 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<LogicalTypeHandle> = 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 logical_types: Vec<LogicalTypeHandle> = record_batch
+            .schema()
+            .fields()
+            .iter()
+            .map(|field| {
+                to_duckdb_logical_type(field.data_type())
+                    .map_err(|_op| Error::ArrowTypeToDuckdbType(field.to_string(), field.data_type().clone()))
+            })
+            .collect::<Result<Vec<_>, _>>()?;
+
+        let vector_size = unsafe { duckdb_vector_size() } as usize;
+        let num_rows = record_batch.num_rows();
+
+        // Process record batch in chunks that fit within DuckDB's vector size
+        let mut offset = 0;
+        while offset < num_rows {
+            let slice_len = std::cmp::min(vector_size, num_rows - offset);
+            let slice = record_batch.slice(offset, slice_len);
 
-        let mut data_chunk = DataChunkHandle::new(&logical_type);
-        record_batch_to_duckdb_data_chunk(&record_batch, &mut data_chunk).map_err(|_op| Error::AppendError)?;
+            let mut data_chunk = DataChunkHandle::new(&logical_types);
+            record_batch_to_duckdb_data_chunk(&slice, &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, &mut self.app)
+            let rc = unsafe { duckdb_append_data_chunk(self.app, data_chunk.get_ptr()) };
+            result_from_duckdb_appender(rc, &mut self.app)?;
+
+            offset += slice_len;
+        }
+
+        Ok(())
     }
 }
 
 #[cfg(test)]
 mod test {
     use crate::{Connection, Result};
     use arrow::{
-        array::{Int8Array, StringArray},
+        array::{Int32Array, Int8Array, StringArray},
         datatypes::{DataType, Field, Schema},
         record_batch::RecordBatch,
     };
@@ -80,4 +96,27 @@ mod test {
         assert_eq!(rbs.iter().map(|op| op.num_rows()).sum::<usize>(), 5);
         Ok(())
     }
+
+    #[test]
+    fn test_append_record_batch_large() -> Result<()> {
+        let record_count = usize::pow(2, 16) + 1;
+        let db = Connection::open_in_memory()?;
+        db.execute_batch("CREATE TABLE foo(id INT)")?;
+        {
+            let id_array = Int32Array::from((0..record_count as i32).collect::<Vec<_>>());
+            let schema = Schema::new(vec![Field::new("id", DataType::Int32, true)]);
+            let record_batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(id_array)]).unwrap();
+            let mut app = db.appender("foo")?;
+            app.append_record_batch(record_batch)?;
+        }
+        let count: usize = db.query_row("SELECT COUNT(*) FROM foo", [], |row| row.get(0))?;
+        assert_eq!(count, record_count);
+
+        // Verify the data is correct
+        let sum: i64 = db.query_row("SELECT SUM(id) FROM foo", [], |row| row.get(0))?;
+        let expected_sum: i64 = (0..record_count as i64).sum();
+        assert_eq!(sum, expected_sum);
+
+        Ok(())
+    }
 }