Re-open DuckDB context between iterations (take 2) (#4362)

Reverts #4355 and re-attempts #4346 while fixing the segfault issue

---------

Signed-off-by: Will Manning <[email protected]>

Author: lwwmanning

.github/workflows/ci.yml

@@ -305,7 +305,7 @@ jobs:
       - name: Run TPC-H
         if: ${{ matrix.suite == 'tpc-h' }}
         run: |
-          cargo run --bin query_bench -- tpch -i1 --scale-factor 0.1
+          cargo run --bin query_bench -- tpch -i2 --scale-factor 0.1
       - name: Run FFI Example
         if: ${{ matrix.suite == 'ffi' }}
         run: |

bench-vortex/src/benchmark_driver.rs

@@ -188,6 +188,9 @@ fn execute_queries<B: Benchmark>(
                 let mut row_count = None;
 
                 for _ in 0..iterations {
+                    // Ensure we reopen the database to clear caches between runs.
+                    ctx.reopen()?;
+
                     let (duration, current_row_count) =
                         ctx.execute_query(query_string).unwrap_or_else(|err| {
                             vortex_panic!("query: {query_idx} failed with: {err}")

bench-vortex/src/engines/ddb/mod.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::path::PathBuf;
 use std::time::{Duration, Instant};
 
 use anyhow::Result;
@@ -31,6 +32,7 @@ impl DuckDBObject {
 pub struct DuckDBCtx {
     pub db: Database,
     pub connection: Connection,
+    pub db_path: Option<PathBuf>,
 }
 
 impl DuckDBCtx {
@@ -55,18 +57,51 @@ impl DuckDBCtx {
         if delete_database {
             std::fs::remove_file(&db_path)?;
         }
-        let db = Database::open(db_path)?;
+        let db = Database::open(&db_path)?;
         let connection = db.connect()?;
         vortex_duckdb::register_table_functions(&connection)?;
-        Ok(Self { db, connection })
+
+        Ok(Self {
+            db,
+            connection,
+            db_path: Some(db_path),
+        })
+    }
+
+    pub fn reopen(&mut self) -> Result<()> {
+        // take ownership of the connection & database
+        let mut connection = unsafe { Connection::borrow(self.connection.as_ptr()) };
+        std::mem::swap(&mut self.connection, &mut connection);
+        let mut db = unsafe { Database::borrow(self.db.as_ptr()) };
+        std::mem::swap(&mut self.db, &mut db);
+
+        // drop the connection, then the database (order might be important?)
+        // NB: self.db and self.connection will be dangling pointers, which we'll fix below
+        drop(connection);
+        drop(db);
+
+        let mut db = match &self.db_path {
+            Some(path) => Database::open(path),
+            None => Database::open_in_memory(),
+        }?;
+        let mut connection = db.connect()?;
+        vortex_duckdb::register_table_functions(&connection)?;
+
+        std::mem::swap(&mut self.connection, &mut connection);
+        std::mem::swap(&mut self.db, &mut db);
+
+        Ok(())
     }
 
     pub fn new_in_memory() -> Result<Self> {
         let db = Database::open_in_memory()?;
         let connection = db.connect()?;
         vortex_duckdb::register_table_functions(&connection)?;
-
-        Ok(Self { db, connection })
+        Ok(Self {
+            db,
+            connection,
+            db_path: None,
+        })
     }
 
     /// Execute DuckDB queries for benchmarks using the internal connection