From 7508043b62d2d4500fce223cca949ad21e34e4a8 Mon Sep 17 00:00:00 2001
From: pedrocarlo <pedro.muniz.carlo@gmail.com>
Date: Sun, 4 May 2025 01:32:09 -0300
Subject: [PATCH] add bench for select count

---
 core/benches/benchmark.rs | 55 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 54 insertions(+), 1 deletion(-)

diff --git a/core/benches/benchmark.rs b/core/benches/benchmark.rs
index 1fed5308b..388fcf933 100644
--- a/core/benches/benchmark.rs
+++ b/core/benches/benchmark.rs
@@ -180,9 +180,62 @@ fn bench_execute_select_1(criterion: &mut Criterion) {
     group.finish();
 }
 
+fn bench_execute_select_count(criterion: &mut Criterion) {
+    // https://github.com/tursodatabase/limbo/issues/174
+    // The rusqlite benchmark crashes on Mac M1 when using the flamegraph features
+    let enable_rusqlite = std::env::var("DISABLE_RUSQLITE_BENCHMARK").is_err();
+
+    #[allow(clippy::arc_with_non_send_sync)]
+    let io = Arc::new(PlatformIO::new().unwrap());
+    let db = Database::open_file(io.clone(), "../testing/testing.db", false).unwrap();
+    let limbo_conn = db.connect().unwrap();
+
+    let mut group = criterion.benchmark_group("Execute `SELECT count() FROM users`");
+
+    group.bench_function("limbo_execute_select_count", |b| {
+        let mut stmt = limbo_conn.prepare("SELECT count() FROM users").unwrap();
+        let io = io.clone();
+        b.iter(|| {
+            loop {
+                match stmt.step().unwrap() {
+                    limbo_core::StepResult::Row => {
+                        black_box(stmt.row());
+                    }
+                    limbo_core::StepResult::IO => {
+                        let _ = io.run_once();
+                    }
+                    limbo_core::StepResult::Done => {
+                        break;
+                    }
+                    limbo_core::StepResult::Interrupt | limbo_core::StepResult::Busy => {
+                        unreachable!();
+                    }
+                }
+            }
+            stmt.reset();
+        });
+    });
+
+    if enable_rusqlite {
+        let sqlite_conn = rusqlite_open();
+
+        group.bench_function("sqlite_execute_select_count", |b| {
+            let mut stmt = sqlite_conn.prepare("SELECT count() FROM users").unwrap();
+            b.iter(|| {
+                let mut rows = stmt.raw_query();
+                while let Some(row) = rows.next().unwrap() {
+                    black_box(row);
+                }
+            });
+        });
+    }
+
+    group.finish();
+}
+
 criterion_group! {
     name = benches;
     config = Criterion::default().with_profiler(PProfProfiler::new(100, Output::Flamegraph(None)));
-    targets = bench_prepare_query, bench_execute_select_1, bench_execute_select_rows
+    targets = bench_prepare_query, bench_execute_select_1, bench_execute_select_rows, bench_execute_select_count
 }
 criterion_main!(benches);