1use git_vdb::{
2 CollectionConfig, Condition, Database, Filter, Point, PointId, Query, QueryParams, Range,
3 SnapshotEngine, SnapshotMutation,
4};
5use serde::Deserialize;
6use serde_json::{json, Map, Value};
7use std::env;
8use std::fs;
9use std::path::{Path, PathBuf};
10use std::thread;
11use std::time::Instant;
12
13#[derive(Deserialize)]
14struct RunSpec {
15 schema_version: u32,
16 case_name: String,
17 dimension: usize,
18 point_count: usize,
19 query_count: usize,
20 points_path: PathBuf,
21 queries_path: PathBuf,
22 k: Vec<usize>,
23 mutation_fractions: Vec<f64>,
24 filter_selectivities: Vec<f64>,
25 concurrency: Vec<usize>,
26}
27
28fn main() -> Result<(), Box<dyn std::error::Error>> {
29 let mut args = env::args_os().skip(1);
30 let input = args.next().ok_or("usage: runner INPUT.json OUTPUT.json")?;
31 let output = args.next().ok_or("usage: runner INPUT.json OUTPUT.json")?;
32 if args.next().is_some() {
33 return Err("usage: runner INPUT.json OUTPUT.json".into());
34 }
35 let spec: RunSpec = serde_json::from_slice(&fs::read(input)?)?;
36 if spec.schema_version != 1 {
37 return Err(format!("unsupported harness schema version {}", spec.schema_version).into());
38 }
39 let vectors = read_vectors(&spec.points_path, spec.point_count, spec.dimension)?;
40 let queries = read_vectors(&spec.queries_path, spec.query_count, spec.dimension)?;
41 let points = make_points(&vectors);
42 let maximum_k = *spec.k.iter().max().ok_or("k must not be empty")?;
43 let config = CollectionConfig {
44 dimension: spec.dimension,
45 ..CollectionConfig::default()
46 };
47
48 let setup_started = Instant::now();
49 let core_dir = tempfile::TempDir::new()?;
50 let engine = SnapshotEngine::init(core_dir.path())?;
51 let setup_us = micros(setup_started);
52 let build_started = Instant::now();
53 let snapshot = engine.build(config.clone(), points.clone())?;
54 let build_us = micros(build_started);
55 let baseline_on_disk_bytes = directory_bytes(core_dir.path())?;
56
57 let (exact_query_us, exact_results, exact_vectors_scored) =
58 query_all(&snapshot, &queries, maximum_k, true, None)?;
59 let (approximate_query_us, approximate_results, approximate_vectors_scored) =
60 query_all(&snapshot, &queries, maximum_k, false, None)?;
61 let snapshot_throughput = json!({
62 "exact": query_snapshot_throughput(
63 &snapshot,
64 &queries,
65 maximum_k,
66 true,
67 &spec.concurrency,
68 )?,
69 "approximate": query_snapshot_throughput(
70 &snapshot,
71 &queries,
72 maximum_k,
73 false,
74 &spec.concurrency,
75 )?,
76 });
77
78 let mut filtered = Map::new();
79 for selectivity in &spec.filter_selectivities {
80 let filter = selectivity_filter(*selectivity);
81 let (exact_us, exact, exact_scored) =
82 query_all(&snapshot, &queries, maximum_k, true, Some(filter.clone()))?;
83 let (approximate_us, approximate, approximate_scored) =
84 query_all(&snapshot, &queries, maximum_k, false, Some(filter))?;
85 filtered.insert(
86 selectivity.to_string(),
87 json!({
88 "exact_query_us": exact_us,
89 "approximate_query_us": approximate_us,
90 "exact_results": exact,
91 "approximate_results": approximate,
92 "exact_vectors_scored": exact_scored,
93 "approximate_vectors_scored": approximate_scored,
94 }),
95 );
96 }
97
98 let mut mutations = Map::new();
99 for fraction in &spec.mutation_fractions {
100 let count = fraction_count(spec.point_count, *fraction);
101 let changed = changed_points(&points[..count]);
102 let upsert_started = Instant::now();
103 let upserted = engine.apply(
104 snapshot.root(),
105 changed.into_iter().map(SnapshotMutation::upsert).collect(),
106 )?;
107 let upsert_us = micros(upsert_started);
108 let delete_started = Instant::now();
109 let deleted = engine.apply(
110 snapshot.root(),
111 vec![SnapshotMutation::delete_ids(
112 (0..count).map(|id| PointId::from(id as u64)),
113 )],
114 )?;
115 let delete_us = micros(delete_started);
116 mutations.insert(
117 fraction.to_string(),
118 json!({
119 "points": count,
120 "upsert_us": upsert_us,
121 "delete_us": delete_us,
122 "upsert_root": upserted.root(),
123 "delete_root": deleted.root(),
124 }),
125 );
126 }
127
128 let adapter_setup_started = Instant::now();
129 let adapter_dir = tempfile::TempDir::new()?;
130 let database = Database::init_bare(adapter_dir.path())?;
131 let collection = database.create_collection("benchmark", config)?;
132 let adapter_setup_us = micros(adapter_setup_started);
133 let adapter_build_started = Instant::now();
134 let adapter_root = collection.upsert(points)?.root;
135 let adapter_build_us = micros(adapter_build_started);
136 if adapter_root != snapshot.root() {
137 return Err("snapshot-core and named-adapter roots differ".into());
138 }
139 query_collection_all(&collection, &queries[..1], maximum_k, true, None)?;
140 query_collection_all(&collection, &queries[..1], maximum_k, false, None)?;
141 let (adapter_exact_query_us, adapter_exact_results, adapter_exact_vectors_scored) =
142 query_collection_all(&collection, &queries, maximum_k, true, None)?;
143 let (
144 adapter_approximate_query_us,
145 adapter_approximate_results,
146 adapter_approximate_vectors_scored,
147 ) = query_collection_all(&collection, &queries, maximum_k, false, None)?;
148 let adapter_throughput = json!({
149 "exact": query_collection_throughput(
150 &collection,
151 &queries,
152 maximum_k,
153 true,
154 &spec.concurrency,
155 )?,
156 "approximate": query_collection_throughput(
157 &collection,
158 &queries,
159 maximum_k,
160 false,
161 &spec.concurrency,
162 )?,
163 });
164 let historical_started = Instant::now();
165 let historical = collection.at(&adapter_root)?;
166 let historical_count = historical.count(None)?.count;
167 let historical_read_us = micros(historical_started);
168
169 let report = json!({
170 "schema_version": 1,
171 "engine": "git-vdb",
172 "case_name": spec.case_name,
173 "point_count": spec.point_count,
174 "dimension": spec.dimension,
175 "query_count": spec.query_count,
176 "k": spec.k,
177 "root": snapshot.root(),
178 "setup_us": setup_us,
179 "snapshot_core": {
180 "build_us": build_us,
181 "exact_query_us": exact_query_us,
182 "approximate_query_us": approximate_query_us,
183 "exact_results": exact_results,
184 "approximate_results": approximate_results,
185 "exact_vectors_scored": exact_vectors_scored,
186 "approximate_vectors_scored": approximate_vectors_scored,
187 "throughput": snapshot_throughput,
188 "filtered": filtered,
189 "mutations": mutations,
190 "on_disk_bytes": baseline_on_disk_bytes,
191 },
192 "named_adapter": {
193 "setup_us": adapter_setup_us,
194 "build_us": adapter_build_us,
195 "exact_query_us": adapter_exact_query_us,
196 "approximate_query_us": adapter_approximate_query_us,
197 "exact_results": adapter_exact_results,
198 "approximate_results": adapter_approximate_results,
199 "exact_vectors_scored": adapter_exact_vectors_scored,
200 "approximate_vectors_scored": adapter_approximate_vectors_scored,
201 "throughput": adapter_throughput,
202 "historical_read_us": historical_read_us,
203 "historical_count": historical_count,
204 "on_disk_bytes": directory_bytes(adapter_dir.path())?,
205 }
206 });
207 fs::write(output, serde_json::to_vec_pretty(&report)?)?;
208 Ok(())
209}
210
211fn read_vectors(
212 path: &Path,
213 count: usize,
214 dimension: usize,
215) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>> {
216 let bytes = fs::read(path)?;
217 let expected = count
218 .checked_mul(dimension)
219 .and_then(|components| components.checked_mul(4))
220 .ok_or("dataset size overflow")?;
221 if bytes.len() != expected {
222 return Err(format!(
223 "{} has {} bytes, expected {expected}",
224 path.display(),
225 bytes.len()
226 )
227 .into());
228 }
229 Ok(bytes
230 .chunks_exact(dimension * 4)
231 .map(|row| {
232 row.chunks_exact(4)
233 .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
234 .collect()
235 })
236 .collect())
237}
238
239fn make_points(vectors: &[Vec<f32>]) -> Vec<Point> {
240 vectors
241 .iter()
242 .enumerate()
243 .map(|(id, vector)| {
244 let mut payload = Map::new();
245 payload.insert("selectivity_bucket".into(), json!(id % 1000));
246 Point {
247 id: (id as u64).into(),
248 vector: vector.clone(),
249 payload,
250 }
251 })
252 .collect()
253}
254
255type QueryBatch = (Vec<u64>, Vec<Value>, Vec<usize>);
256
257fn query_all(
258 snapshot: &git_vdb::Snapshot,
259 queries: &[Vec<f32>],
260 limit: usize,
261 exact: bool,
262 filter: Option<Filter>,
263) -> git_vdb::Result<QueryBatch> {
264 let mut durations = Vec::with_capacity(queries.len());
265 let mut results = Vec::with_capacity(queries.len());
266 let mut vectors_scored = Vec::with_capacity(queries.len());
267 for vector in queries {
268 let started = Instant::now();
269 let result = snapshot.query(Query {
270 vector: vector.clone(),
271 limit,
272 filter: filter.clone(),
273 params: QueryParams {
274 exact: Some(exact),
275 ..QueryParams::default()
276 },
277 ..Query::default()
278 })?;
279 durations.push(micros(started));
280 vectors_scored.push(result.stats.vectors_scored);
281 results.push(Value::Array(
282 result
283 .points
284 .into_iter()
285 .map(|point| json!({"id": point.id, "score": point.score}))
286 .collect(),
287 ));
288 }
289 Ok((durations, results, vectors_scored))
290}
291
292fn query_collection_all(
293 collection: &git_vdb::Collection,
294 queries: &[Vec<f32>],
295 limit: usize,
296 exact: bool,
297 filter: Option<Filter>,
298) -> git_vdb::Result<QueryBatch> {
299 let mut durations = Vec::with_capacity(queries.len());
300 let mut results = Vec::with_capacity(queries.len());
301 let mut vectors_scored = Vec::with_capacity(queries.len());
302 for vector in queries {
303 let started = Instant::now();
304 let result = collection.query(Query {
305 vector: vector.clone(),
306 limit,
307 filter: filter.clone(),
308 params: QueryParams {
309 exact: Some(exact),
310 ..QueryParams::default()
311 },
312 ..Query::default()
313 })?;
314 durations.push(micros(started));
315 vectors_scored.push(result.stats.vectors_scored);
316 results.push(Value::Array(
317 result
318 .points
319 .into_iter()
320 .map(|point| json!({"id": point.id, "score": point.score}))
321 .collect(),
322 ));
323 }
324 Ok((durations, results, vectors_scored))
325}
326
327fn query_snapshot_throughput(
328 snapshot: &git_vdb::Snapshot,
329 queries: &[Vec<f32>],
330 limit: usize,
331 exact: bool,
332 concurrencies: &[usize],
333) -> git_vdb::Result<Map<String, Value>> {
334 let mut measurements = Map::new();
335 for &workers in concurrencies {
336 if workers == 0 {
337 return Err(git_vdb::Error::Invalid(
338 "benchmark concurrency must be positive".into(),
339 ));
340 }
341 let started = Instant::now();
342 let outcome: git_vdb::Result<()> = thread::scope(|scope| {
343 let mut handles = Vec::with_capacity(workers);
344 for worker in 0..workers {
345 let snapshot = snapshot.clone();
346 handles.push(scope.spawn(move || -> git_vdb::Result<()> {
347 for vector in queries.iter().skip(worker).step_by(workers) {
348 snapshot.query(Query {
349 vector: vector.clone(),
350 limit,
351 params: QueryParams {
352 exact: Some(exact),
353 ..QueryParams::default()
354 },
355 ..Query::default()
356 })?;
357 }
358 Ok(())
359 }));
360 }
361 join_query_workers(handles)
362 });
363 outcome?;
364 measurements.insert(
365 workers.to_string(),
366 throughput_measurement(queries.len(), micros(started)),
367 );
368 }
369 Ok(measurements)
370}
371
372fn query_collection_throughput(
373 collection: &git_vdb::Collection,
374 queries: &[Vec<f32>],
375 limit: usize,
376 exact: bool,
377 concurrencies: &[usize],
378) -> git_vdb::Result<Map<String, Value>> {
379 let mut measurements = Map::new();
380 for &workers in concurrencies {
381 if workers == 0 {
382 return Err(git_vdb::Error::Invalid(
383 "benchmark concurrency must be positive".into(),
384 ));
385 }
386 let started = Instant::now();
387 let outcome: git_vdb::Result<()> = thread::scope(|scope| {
388 let mut handles = Vec::with_capacity(workers);
389 for worker in 0..workers {
390 let collection = collection.clone();
391 handles.push(scope.spawn(move || -> git_vdb::Result<()> {
392 for vector in queries.iter().skip(worker).step_by(workers) {
393 collection.query(Query {
394 vector: vector.clone(),
395 limit,
396 params: QueryParams {
397 exact: Some(exact),
398 ..QueryParams::default()
399 },
400 ..Query::default()
401 })?;
402 }
403 Ok(())
404 }));
405 }
406 join_query_workers(handles)
407 });
408 outcome?;
409 measurements.insert(
410 workers.to_string(),
411 throughput_measurement(queries.len(), micros(started)),
412 );
413 }
414 Ok(measurements)
415}
416
417fn join_query_workers(
418 handles: Vec<thread::ScopedJoinHandle<'_, git_vdb::Result<()>>>,
419) -> git_vdb::Result<()> {
420 for handle in handles {
421 match handle.join() {
422 Ok(result) => result?,
423 Err(_) => {
424 return Err(git_vdb::Error::Invalid(
425 "benchmark query worker panicked".into(),
426 ));
427 }
428 }
429 }
430 Ok(())
431}
432
433fn throughput_measurement(queries: usize, wall_us: u64) -> Value {
434 json!({
435 "queries": queries,
436 "wall_us": wall_us,
437 "queries_per_second": queries as f64 * 1_000_000.0 / wall_us as f64,
438 })
439}
440
441fn selectivity_filter(selectivity: f64) -> Filter {
442 Filter::must([Condition::range(
443 "selectivity_bucket",
444 Range {
445 lt: Some((selectivity * 1000.0).round()),
446 ..Range::default()
447 },
448 )])
449}
450
451fn fraction_count(total: usize, fraction: f64) -> usize {
452 ((total as f64 * fraction).round() as usize).clamp(1, total)
453}
454
455fn changed_points(points: &[Point]) -> Vec<Point> {
456 points
457 .iter()
458 .cloned()
459 .map(|mut point| {
460 point.vector[0] += 0.001;
461 point
462 })
463 .collect()
464}
465
466fn micros(started: Instant) -> u64 {
467 started.elapsed().as_micros().try_into().unwrap_or(u64::MAX)
468}
469
470fn directory_bytes(path: &Path) -> std::io::Result<u64> {
471 let mut total = 0;
472 for entry in fs::read_dir(path)? {
473 let entry = entry?;
474 let metadata = entry.metadata()?;
475 if metadata.is_dir() {
476 total += directory_bytes(&entry.path())?;
477 } else {
478 total += metadata.len();
479 }
480 }
481 Ok(total)
482}