Feat/n param for bench (#244)

Author: matthewmturner

README.md

@@ -105,6 +105,8 @@ Optionally you can use the `--run-before` param to run a query before the benchm
 
 To save benchmark results to a file use the `--save` parameter with a file path.  Further, you can use the `--append` parameter to append to the file instead of overwriting it.
 
+The number of benchmark iterations is defined in your configuration (default is 10) and can be configured per benchmark run with `-n` parameter.
+
 #### Analyze Queries
 
 The output from `EXPLAIN ANALYZE` provides a wealth of information on a queries execution - however, the amount of information and connecting the dots can be difficult and manual.  Further, there is detail in the `MetricSet`'s of the underlying `ExecutionPlan`'s that is lost in the output.

src/args.rs

@@ -87,6 +87,9 @@ pub struct DftArgs {
 
     #[clap(long, help = "Append the benchmark results to an existing file")]
     pub append: bool,
+
+    #[clap(short = 'n', help = "Set the number of benchmark iterations to run")]
+    pub benchmark_iterations: Option<usize>,
 }
 
 impl DftArgs {

src/cli/mod.rs

@@ -386,7 +386,7 @@ impl CliApp {
         let stats = self
             .app_execution
             .execution_ctx()
-            .benchmark_query(sql)
+            .benchmark_query(sql, self.args.benchmark_iterations)
             .await?;
         Ok(stats)
     }
@@ -407,7 +407,7 @@ impl CliApp {
         let stats = self
             .app_execution
             .flightsql_ctx()
-            .benchmark_query(sql)
+            .benchmark_query(sql, self.args.benchmark_iterations)
             .await?;
         Ok(stats)
     }

src/execution/flightsql.rs

@@ -67,8 +67,12 @@ impl FlightSQLContext {
         }
     }
 
-    pub async fn benchmark_query(&self, query: &str) -> Result<FlightSQLBenchmarkStats> {
-        let iterations = self.config.benchmark_iterations;
+    pub async fn benchmark_query(
+        &self,
+        query: &str,
+        cli_iterations: Option<usize>,
+    ) -> Result<FlightSQLBenchmarkStats> {
+        let iterations = cli_iterations.unwrap_or(self.config.benchmark_iterations);
         let mut rows_returned = Vec::with_capacity(iterations);
         let mut get_flight_info_durations = Vec::with_capacity(iterations);
         let mut ttfb_durations = Vec::with_capacity(iterations);

src/execution/local.rs

@@ -243,8 +243,12 @@ impl ExecutionContext {
     }
 
     /// Benchmark the provided query.  Currently, only a single statement can be benchmarked
-    pub async fn benchmark_query(&self, query: &str) -> Result<LocalBenchmarkStats> {
-        let iterations = self.config.benchmark_iterations;
+    pub async fn benchmark_query(
+        &self,
+        query: &str,
+        cli_iterations: Option<usize>,
+    ) -> Result<LocalBenchmarkStats> {
+        let iterations = cli_iterations.unwrap_or(self.config.benchmark_iterations);
         info!("Benchmarking query with {} iterations", iterations);
         let mut rows_returned = Vec::with_capacity(iterations);
         let mut logical_planning_durations = Vec::with_capacity(iterations);

tests/cli_cases/bench.rs

@@ -179,3 +179,24 @@ SELECT 1
     let lines: Vec<&str> = contents.lines().collect();
     assert_eq!(lines.len(), 3);
 }
+
+#[test]
+fn test_bench_command_with_custom_iterations() {
+    let assert = Command::cargo_bin("dft")
+        .unwrap()
+        .arg("-c")
+        .arg("SELECT 1")
+        .arg("--bench")
+        .arg("-n")
+        .arg("3")
+        .assert()
+        .success();
+
+    let expected = r##"
+----------------------------
+Benchmark Stats (3 runs)
+----------------------------
+SELECT 1
+----------------------------"##;
+    assert.stdout(contains_str(expected));
+}

tests/extension_cases/flightsql.rs

@@ -465,3 +465,33 @@ SELECT 1
 
     fixture.shutdown_and_wait().await;
 }
+
+#[tokio::test]
+pub async fn test_bench_command_customer_iterations() {
+    let test_server = TestFlightSqlServiceImpl::new();
+    let fixture = TestFixture::new(test_server.service(), "127.0.0.1:50051").await;
+
+    let assert = tokio::task::spawn_blocking(move || {
+        Command::cargo_bin("dft")
+            .unwrap()
+            .arg("-c")
+            .arg("SELECT 1")
+            .arg("--bench")
+            .arg("--flightsql")
+            .arg("-n")
+            .arg("3")
+            .assert()
+            .success()
+    })
+    .await
+    .unwrap();
+
+    let expected = r##"
+----------------------------
+Benchmark Stats (3 runs)
+----------------------------
+SELECT 1
+----------------------------"##;
+    assert.stdout(contains_str(expected));
+    fixture.shutdown_and_wait().await;
+}