Prevent InsertBuilder from taking ownership of record batches (#518)

Author: nuno-faria

core/src/mysql.rs

@@ -382,8 +382,9 @@ impl MySQL {
         batch: RecordBatch,
         on_conflict: Option<OnConflict>,
     ) -> Result<()> {
+        let batches = vec![batch];
         let insert_table_builder =
-            InsertBuilder::new(&TableReference::bare(self.table_name.clone()), vec![batch]);
+            InsertBuilder::new(&TableReference::bare(self.table_name.clone()), &batches);
 
         let sea_query_on_conflict =
             on_conflict.map(|oc| oc.build_sea_query_on_conflict(&self.schema));

core/src/postgres.rs

@@ -419,7 +419,8 @@ impl Postgres {
         batch: RecordBatch,
         on_conflict: Option<OnConflict>,
     ) -> Result<()> {
-        let insert_table_builder = InsertBuilder::new(&self.table, vec![batch]);
+        let batches = vec![batch];
+        let insert_table_builder = InsertBuilder::new(&self.table, &batches);
 
         let sea_query_on_conflict =
             on_conflict.map(|oc| oc.build_sea_query_on_conflict(&self.schema));

core/src/sql/arrow_sql_gen/statement.rs

@@ -190,9 +190,9 @@ macro_rules! push_list_values {
     }};
 }
 
-pub struct InsertBuilder {
+pub struct InsertBuilder<'a> {
     table: TableReference,
-    record_batches: Vec<RecordBatch>,
+    record_batches: &'a Vec<RecordBatch>,
 }
 
 pub fn use_json_insert_for_type<T: QueryBuilder + 'static>(
@@ -218,9 +218,9 @@ pub fn use_json_insert_for_type<T: QueryBuilder + 'static>(
     false
 }
 
-impl InsertBuilder {
+impl<'a> InsertBuilder<'a> {
     #[must_use]
-    pub fn new(table: &TableReference, record_batches: Vec<RecordBatch>) -> Self {
+    pub fn new(table: &TableReference, record_batches: &'a Vec<RecordBatch>) -> Self {
         Self {
             table: table.clone(),
             record_batches,
@@ -1082,7 +1082,7 @@ impl InsertBuilder {
             .columns(columns)
             .to_owned();
 
-        for record_batch in &self.record_batches {
+        for record_batch in self.record_batches {
             self.construct_insert_stmt(&mut insert_stmt, record_batch, &query_builder)?;
         }
         if let Some(on_conflict) = on_conflict {
@@ -1500,7 +1500,7 @@ mod tests {
         .expect("Unable to build record batch");
         let record_batches = vec![batch1, batch2];
 
-        let sql = InsertBuilder::new(&TableReference::from("users"), record_batches)
+        let sql = InsertBuilder::new(&TableReference::from("users"), &record_batches)
             .build_postgres(None)
             .expect("Failed to build insert statement");
         assert_eq!(sql, "INSERT INTO \"users\" (\"id\", \"name\", \"age\") VALUES (1, 'a', 10), (2, 'b', 20), (3, 'c', 30), (1, 'a', 10), (2, 'b', 20), (3, 'c', 30)");
@@ -1544,7 +1544,7 @@ mod tests {
         .expect("Unable to build record batch");
         let record_batches = vec![batch1, batch2];
 
-        let sql = InsertBuilder::new(&TableReference::from("schema.users"), record_batches)
+        let sql = InsertBuilder::new(&TableReference::from("schema.users"), &record_batches)
             .build_postgres(None)
             .expect("Failed to build insert statement");
         assert_eq!(sql, "INSERT INTO \"schema\".\"users\" (\"id\", \"name\", \"age\") VALUES (1, 'a', 10), (2, 'b', 20), (3, 'c', 30), (1, 'a', 10), (2, 'b', 20), (3, 'c', 30)");
@@ -1595,7 +1595,7 @@ mod tests {
         let batch = RecordBatch::try_new(Arc::new(schema1.clone()), vec![Arc::new(list_array)])
             .expect("Unable to build record batch");
 
-        let sql = InsertBuilder::new(&TableReference::from("arrays"), vec![batch])
+        let sql = InsertBuilder::new(&TableReference::from("arrays"), &vec![batch])
             .build_postgres(None)
             .expect("Failed to build insert statement");
         assert_eq!(

core/src/sqlite.rs

@@ -71,13 +71,17 @@ pub enum Error {
     },
 
     #[snafu(display("Unable to create table in Sqlite: {source}"))]
-    UnableToCreateTable { source: tokio_rusqlite::Error<rusqlite::Error> },
+    UnableToCreateTable {
+        source: tokio_rusqlite::Error<rusqlite::Error>,
+    },
 
     #[snafu(display("Unable to insert data into the Sqlite table: {source}"))]
     UnableToInsertIntoTable { source: rusqlite::Error },
 
     #[snafu(display("Unable to insert data into the Sqlite table: {source}"))]
-    UnableToInsertIntoTableAsync { source: tokio_rusqlite::Error<rusqlite::Error> },
+    UnableToInsertIntoTableAsync {
+        source: tokio_rusqlite::Error<rusqlite::Error>,
+    },
 
     #[snafu(display("Unable to insert data into the Sqlite table. The disk is full."))]
     DiskFull {},
@@ -549,7 +553,8 @@ impl Sqlite {
         batch: RecordBatch,
         on_conflict: Option<&OnConflict>,
     ) -> rusqlite::Result<()> {
-        let insert_table_builder = InsertBuilder::new(&self.table, vec![batch]);
+        let batches = vec![batch];
+        let insert_table_builder = InsertBuilder::new(&self.table, &batches);
 
         let sea_query_on_conflict =
             on_conflict.map(|oc| oc.build_sea_query_on_conflict(&self.schema));

core/tests/sqlite/mod.rs

@@ -52,12 +52,10 @@ async fn arrow_sqlite_round_trip(
     // Create sqlite table from arrow records and insert arrow records
     let schema = Arc::clone(&arrow_record.schema());
     let create_table_stmts = CreateTableBuilder::new(schema, table_name).build_sqlite();
-    let insert_table_stmt = InsertBuilder::new(
-        &TableReference::from(table_name),
-        vec![arrow_record.clone()],
-    )
-    .build_sqlite(None)
-    .expect("SQLite insert statement should be constructed");
+    let batches = vec![arrow_record.clone()];
+    let insert_table_stmt = InsertBuilder::new(&TableReference::from(table_name), &batches)
+        .build_sqlite(None)
+        .expect("SQLite insert statement should be constructed");
 
     // Test arrow -> Sqlite row coverage
     let _ = conn
@@ -211,9 +209,18 @@ fn create_comprehensive_test_data() -> (RecordBatch, SchemaRef) {
         None,
         Some(500000u64),
     ]);
-    let col_float32 = Float32Array::from(vec![Some(1.5), None, Some(-std::f32::consts::PI), Some(2.71)]);
-    let col_float64 =
-        Float64Array::from(vec![None, Some(std::f64::consts::E), Some(-1.414), Some(std::f64::consts::PI)]);
+    let col_float32 = Float32Array::from(vec![
+        Some(1.5),
+        None,
+        Some(-std::f32::consts::PI),
+        Some(2.71),
+    ]);
+    let col_float64 = Float64Array::from(vec![
+        None,
+        Some(std::f64::consts::E),
+        Some(-1.414),
+        Some(std::f64::consts::PI),
+    ]);
     let col_utf8 = StringArray::from(vec![Some("hello"), Some("world"), None, Some("test")]);
     let col_large_utf8 = LargeStringArray::from(vec![
         None,