finish benchmark calls

paradedb · Apr 15, 2024 · 9b8671c · 9b8671c
1 parent 3cd152c
commit 9b8671c
Show file tree

Hide file tree

Showing 5 changed files with 343 additions and 49 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/cargo-paradedb/Cargo.toml b/cargo-paradedb/Cargo.toml
@@ -11,6 +11,8 @@ async-std = "1.12.0"
 chrono = { version = "0.4.34", features = ["clock", "alloc", "serde"] }
 clap = { version = "4.5.4", features = ["derive", "env"] }
 cmd_lib = "1.9.3"
+criterion = { version = "0.5.1", features = ["async_std"] }
+dotenvy = "0.15.7"
 glob = "0.3.1"
 itertools = "0.12.1"
 serde = "1.0.197"

diff --git a/cargo-paradedb/src/cli.rs b/cargo-paradedb/src/cli.rs
@@ -65,18 +65,37 @@ pub enum EsLogsCommand {
         table: String,
         /// Postgres table name to index.
         #[arg(short, long, default_value = DEFAULT_BENCH_ESLOGS_INDEX_NAME)]
-        index_name: String,
+        index: String,
         /// Postgres database url to connect to.
         #[arg(short, long, env = "DATABASE_URL")]
         url: String,
     },
     QuerySearchIndex {
-        /// Postgres table name to query.
+        /// Postgres index name to query.
+        #[arg(short, long, default_value = DEFAULT_BENCH_ESLOGS_INDEX_NAME)]
+        index: String,
+        /// Query to run.
+        #[arg(short, long, default_value = "message:flame")]
+        query: String,
+        /// Limit results to return.
+        #[arg(short, long, default_value_t = 1)]
+        limit: u64,
+        /// Postgres database url to connect to.
+        #[arg(short, long, env = "DATABASE_URL")]
+        url: String,
+    },
+    BuildParquetTable {
+        /// Postgres table name to build from.
+        #[arg(short, long, default_value = DEFAULT_BENCH_ESLOGS_TABLE)]
+        table: String,
+        /// Postgres database url to connect to.
+        #[arg(short, long, env = "DATABASE_URL")]
+        url: String,
+    },
+    CountParquetTable {
+        /// Postgres table name to build from.
         #[arg(short, long, default_value = DEFAULT_BENCH_ESLOGS_TABLE)]
         table: String,
-        /// Postgres table name to query.
-        #[arg(short, long, default_value = DEFAULT_BENCH_ESLOGS_INDEX_NAME)]
-        index_name: String,
         /// Postgres database url to connect to.
         #[arg(short, long, env = "DATABASE_URL")]
         url: String,

diff --git a/cargo-paradedb/src/main.rs b/cargo-paradedb/src/main.rs
@@ -5,13 +5,17 @@ mod tables;
 use anyhow::Result;
 use async_std::task::block_on;
 use cli::{Cli, Corpus, EsLogsCommand, Subcommand};
+use dotenvy::dotenv;
 use tracing_subscriber::EnvFilter;
 
 fn main() -> Result<()> {
     tracing_subscriber::fmt()
         .with_env_filter(EnvFilter::from_default_env())
         .init();
 
+    // Load env vars from a .env in any parent folder.
+    dotenv().ok();
+
     let cli = Cli::default();
     match cli.subcommand {
         Subcommand::Install => subcommand::install(),
@@ -23,14 +27,23 @@ fn main() -> Result<()> {
                     table,
                     url,
                 } => block_on(subcommand::bench_eslogs_generate(seed, events, table, url)),
-                EsLogsCommand::BuildSearchIndex {
-                    table,
-                    index_name,
+                EsLogsCommand::BuildSearchIndex { table, index, url } => block_on(
+                    subcommand::bench_eslogs_build_search_index(table, index, url),
+                ),
+                EsLogsCommand::QuerySearchIndex {
+                    index,
+                    query,
+                    limit,
                     url,
-                } => block_on(subcommand::bench_eslogs_build_search_index(
-                    table, index_name, url,
+                } => block_on(subcommand::bench_eslogs_query_search_index(
+                    index, query, limit, url,
                 )),
-                EsLogsCommand::QuerySearchIndex { .. } => Ok(()),
+                EsLogsCommand::BuildParquetTable { table, url } => {
+                    block_on(subcommand::bench_eslogs_build_parquet_table(table, url))
+                }
+                EsLogsCommand::CountParquetTable { table, url } => {
+                    block_on(subcommand::bench_eslogs_count_parquet_table(table, url))
+                }
             },
         },
     }