22// SPDX-FileCopyrightText: Copyright the Vortex contributors
33
44use std:: fs:: File ;
5+ use std:: future:: Future ;
56use std:: io:: { Write , stdout} ;
67use std:: path:: PathBuf ;
78
8- use bench_vortex:: bench_run:: run_with_setup ;
9+ use bench_vortex:: bench_run:: run_timed_with_setup ;
910use bench_vortex:: datasets:: taxi_data:: * ;
1011use bench_vortex:: display:: { DisplayFormat , print_measurements_json, render_table} ;
1112use bench_vortex:: measurements:: TimingMeasurement ;
@@ -34,8 +35,9 @@ struct Args {
3435 default_values_t = vec![ Format :: Parquet , Format :: OnDiskVortex ]
3536 ) ]
3637 formats : Vec < Format > ,
37- #[ arg( short, long, default_value_t = 10 ) ]
38- iterations : usize ,
38+ /// Time limit in seconds for each benchmark target (e.g., 10 for 10 seconds).
39+ #[ arg( long, default_value_t = 10 ) ]
40+ time_limit : u64 ,
3941 #[ arg( short, long) ]
4042 threads : Option < usize > ,
4143 #[ arg( short, long) ]
@@ -61,41 +63,48 @@ fn main() -> anyhow::Result<()> {
6163 random_access (
6264 args. formats ,
6365 runtime,
64- args. iterations ,
66+ args. time_limit ,
6567 args. display_format ,
6668 indices,
6769 & args. output_path ,
6870 )
6971}
7072
71- /// Given a benchmark future, runs it and returns a [`TimingMeasurement`].
72- fn create_timing_measurement < O , B , F > (
73- benchmark : B ,
73+ /// Configuration for timing measurements
74+ struct TimingConfig < ' a > {
7475 name : String ,
7576 storage : String ,
76- runtime : & Runtime ,
77- indices : & Buffer < u64 > ,
78- iterations : usize ,
77+ runtime : & ' a Runtime ,
78+ indices : & ' a Buffer < u64 > ,
79+ time_limit : u64 ,
7980 target : Target ,
80- ) -> TimingMeasurement
81+ }
82+
83+ /// Given a benchmark future, runs it and returns a [`TimingMeasurement`].
84+ fn create_timing_measurement < O , B , F > ( benchmark : B , config : TimingConfig ) -> TimingMeasurement
8185where
8286 B : FnMut ( Buffer < u64 > ) -> F ,
8387 F : Future < Output = O > ,
8488{
85- let benchmark_duration = run_with_setup ( runtime, iterations, || indices. clone ( ) , benchmark) ;
89+ let runs = run_timed_with_setup (
90+ config. runtime ,
91+ config. time_limit ,
92+ || config. indices . clone ( ) ,
93+ benchmark,
94+ ) ;
8695
8796 TimingMeasurement {
88- name,
89- storage,
90- target,
91- time : benchmark_duration ,
97+ name : config . name ,
98+ storage : config . storage ,
99+ target : config . target ,
100+ runs ,
92101 }
93102}
94103
95104fn random_access (
96105 formats : Vec < Format > ,
97106 runtime : Runtime ,
98- iterations : usize ,
107+ time_limit : u64 ,
99108 display_format : DisplayFormat ,
100109 indices : Buffer < u64 > ,
101110 output_path : & Option < PathBuf > ,
@@ -123,12 +132,14 @@ fn random_access(
123132 |indices| async {
124133 take_vortex_tokio ( & taxi_vortex, indices, validate_vortex_array) . await
125134 } ,
126- "random-access/vortex-tokio-local-disk" . to_string ( ) ,
127- STORAGE_NVME . to_owned ( ) ,
128- & runtime,
129- & indices,
130- iterations,
131- target,
135+ TimingConfig {
136+ name : "random-access/vortex-tokio-local-disk" . to_string ( ) ,
137+ storage : STORAGE_NVME . to_owned ( ) ,
138+ runtime : & runtime,
139+ indices : & indices,
140+ time_limit,
141+ target,
142+ } ,
132143 )
133144 }
134145 Format :: VortexCompact => {
@@ -139,25 +150,29 @@ fn random_access(
139150 take_vortex_tokio ( & taxi_vortex_compact, indices, validate_vortex_array)
140151 . await
141152 } ,
142- "random-access/vortex-compact-tokio-local-disk" . to_string ( ) ,
143- STORAGE_NVME . to_owned ( ) ,
144- & runtime,
145- & indices,
146- iterations,
147- target,
153+ TimingConfig {
154+ name : "random-access/vortex-compact-tokio-local-disk" . to_string ( ) ,
155+ storage : STORAGE_NVME . to_owned ( ) ,
156+ runtime : & runtime,
157+ indices : & indices,
158+ time_limit,
159+ target,
160+ } ,
148161 )
149162 }
150163 Format :: Parquet => {
151164 let taxi_parquet = runtime. block_on ( taxi_data_parquet ( ) ) ?;
152165
153166 create_timing_measurement (
154167 |indices| async { take_parquet ( & taxi_parquet, indices) . await } ,
155- "random-access/parquet-tokio-local-disk" . to_string ( ) ,
156- STORAGE_NVME . to_owned ( ) ,
157- & runtime,
158- & indices,
159- iterations,
160- target,
168+ TimingConfig {
169+ name : "random-access/parquet-tokio-local-disk" . to_string ( ) ,
170+ storage : STORAGE_NVME . to_owned ( ) ,
171+ runtime : & runtime,
172+ indices : & indices,
173+ time_limit,
174+ target,
175+ } ,
161176 )
162177 }
163178 #[ cfg( feature = "lance" ) ]
@@ -166,12 +181,14 @@ fn random_access(
166181
167182 create_timing_measurement (
168183 |indices| async { take_lance ( & taxi_lance, indices) . await } ,
169- "random-access/lance-tokio-local-disk" . to_string ( ) ,
170- STORAGE_NVME . to_owned ( ) ,
171- & runtime,
172- & indices,
173- iterations,
174- target,
184+ TimingConfig {
185+ name : "random-access/lance-tokio-local-disk" . to_string ( ) ,
186+ storage : STORAGE_NVME . to_owned ( ) ,
187+ runtime : & runtime,
188+ indices : & indices,
189+ time_limit,
190+ target,
191+ } ,
175192 )
176193 }
177194 Format :: Csv | Format :: Arrow | Format :: OnDiskDuckDB => unimplemented ! ( ) ,
0 commit comments