From 4ed4cb28dfe74d038a53ccfaae01e5bfcc83400a Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 14 Aug 2026 14:06:48 +0800 Subject: [PATCH 1/2] feat: add vector search timing logs --- crates/paimon/src/arrow/format/parquet.rs | 99 ++++++ .../paimon/src/table/vector_search_builder.rs | 308 ++++++++++++++++-- crates/paimon/src/vindex/mod.rs | 7 + crates/paimon/src/vindex/range_reader.rs | 145 +++++++++ crates/paimon/src/vindex/reader.rs | 108 +++++- 5 files changed, 632 insertions(+), 35 deletions(-) diff --git a/crates/paimon/src/arrow/format/parquet.rs b/crates/paimon/src/arrow/format/parquet.rs index 960828ec..96369d68 100644 --- a/crates/paimon/src/arrow/format/parquet.rs +++ b/crates/paimon/src/arrow/format/parquet.rs @@ -3291,11 +3291,20 @@ mod tests { /// pages. `id` runs 0..total_rows so page `p` covers ids /// `[p*page_row_limit, (p+1)*page_row_limit)`. async fn write_multi_page_parquet(page_row_limit: usize, total_rows: i32) -> Vec { + write_multi_page_parquet_with_offset_index(page_row_limit, total_rows, true).await + } + + async fn write_multi_page_parquet_with_offset_index( + page_row_limit: usize, + total_rows: i32, + write_offset_index: bool, + ) -> Vec { let schema = writer_arrow_schema(); let props = parquet::file::properties::WriterProperties::builder() .set_data_page_row_count_limit(page_row_limit) .set_write_batch_size(page_row_limit) .set_max_row_group_row_count(Some(total_rows as usize)) + .set_offset_index_disabled(!write_offset_index) .build(); let mut buf: Vec = Vec::new(); { @@ -3335,6 +3344,10 @@ mod tests { .sum() } + fn read_calls(&self) -> usize { + self.ranges.lock().unwrap().len() + } + fn reset(&self) { self.ranges.lock().unwrap().clear(); } @@ -3431,6 +3444,92 @@ mod tests { ); } + async fn read_row_ranges(data: Bytes, row_ranges: Vec) -> (usize, usize, u64) { + let file_size = data.len() as u64; + let file_read = TrackingFileRead::new(data); + let tracker = file_read.clone(); + let fields = vec![int_field("id"), int_field("value")]; + let stream = ParquetFormatReader::default() + .read_batch_stream( + Box::new(file_read), + file_size, + &fields, + None, + Some(32), + Some(row_ranges), + ) + .await + .unwrap(); + tracker.reset(); + let rows = stream + .try_fold( + 0usize, + |rows, batch| async move { Ok(rows + batch.num_rows()) }, + ) + .await + .unwrap(); + (rows, tracker.read_calls(), tracker.bytes_read()) + } + + #[tokio::test] + #[ignore = "controlled refine I/O diagnostic"] + async fn refine_io_amplification_clustered() { + let data = Bytes::from(write_multi_page_parquet(10, 80).await); + let requested_payload_bytes = 4 * std::mem::size_of::(); + let (rows, read_calls, actual_bytes) = + read_row_ranges(data, vec![RowRange::new(30, 33)]).await; + + eprintln!( + "event=refine_io_amplification scenario=clustered rows={} read_calls={} actual_bytes={} requested_payload_bytes={}", + rows, read_calls, actual_bytes, requested_payload_bytes + ); + assert_eq!(rows, 4); + assert!(actual_bytes >= requested_payload_bytes as u64); + } + + #[tokio::test] + #[ignore = "controlled refine I/O diagnostic"] + async fn refine_io_amplification_scattered() { + let data = Bytes::from(write_multi_page_parquet(10, 80).await); + let (_, clustered_calls, clustered_bytes) = + read_row_ranges(data.clone(), vec![RowRange::new(30, 33)]).await; + let ranges = vec![ + RowRange::new(1, 1), + RowRange::new(21, 21), + RowRange::new(41, 41), + RowRange::new(61, 61), + ]; + let requested_payload_bytes = ranges.len() * std::mem::size_of::(); + let (rows, read_calls, actual_bytes) = read_row_ranges(data, ranges).await; + + eprintln!( + "event=refine_io_amplification scenario=scattered rows={} read_calls={} actual_bytes={} requested_payload_bytes={}", + rows, read_calls, actual_bytes, requested_payload_bytes + ); + assert_eq!(rows, 4); + assert!(read_calls >= clustered_calls); + assert!(actual_bytes >= clustered_bytes); + } + + #[tokio::test] + #[ignore = "controlled refine I/O diagnostic"] + async fn refine_io_amplification_without_offset_index() { + let data = Bytes::from(write_multi_page_parquet_with_offset_index(10, 80, false).await); + assert!(load_metadata_with_page_index(&data, true) + .offset_index() + .is_none()); + let requested_payload_bytes = std::mem::size_of::(); + let (rows, read_calls, actual_bytes) = + read_row_ranges(data, vec![RowRange::new(35, 35)]).await; + + eprintln!( + "event=refine_io_amplification scenario=without_offset_index rows={} read_calls={} actual_bytes={} requested_payload_bytes={}", + rows, read_calls, actual_bytes, requested_payload_bytes + ); + assert_eq!(rows, 1); + assert!(actual_bytes > requested_payload_bytes as u64); + } + /// Parse metadata from in-memory parquet bytes, optionally loading the page /// index — mirrors what the reader does via `with_page_index_policy`. fn load_metadata_with_page_index( diff --git a/crates/paimon/src/table/vector_search_builder.rs b/crates/paimon/src/table/vector_search_builder.rs index 1295db93..b772a28d 100644 --- a/crates/paimon/src/table/vector_search_builder.rs +++ b/crates/paimon/src/table/vector_search_builder.rs @@ -60,7 +60,7 @@ use crate::vindex::pkvector::exact::validate_query; use crate::vindex::pkvector::metric::VectorSearchMetric; use crate::vindex::range_reader::VindexFileReader; use crate::vindex::reader::VindexVectorGlobalIndexReader; -use crate::vindex::{is_vindex_index_type, VindexVectorIndexOptions}; +use crate::vindex::{is_vindex_index_type, vector_search_timing_enabled, VindexVectorIndexOptions}; use arrow_array::{Array, FixedSizeListArray, Float32Array, Int64Array, ListArray, RecordBatch}; use arrow_select::interleave::interleave_record_batch; use futures::{stream, TryStreamExt}; @@ -72,6 +72,7 @@ use std::cmp::Ordering; use std::collections::{BinaryHeap, HashMap, HashSet}; use std::io::Cursor; use std::sync::Arc; +use std::time::{Duration, Instant}; const INDEX_DIR: &str = "index"; @@ -1152,6 +1153,8 @@ impl<'a> BatchVectorSearchBuilder<'a> { } pub async fn execute(&self) -> crate::Result> { + let timing_enabled = vector_search_timing_enabled(); + let total_start = timing_enabled.then(Instant::now); // Fail closed: like `execute_read` and the single-query builder, this // returns data-derived row ids/scores outside `TableScan`/`TableRead`, // so it must refuse a `query-auth.enabled` table before any fast path @@ -1227,12 +1230,32 @@ impl<'a> BatchVectorSearchBuilder<'a> { .collect::>>()?; let snapshot_manager = self.table.snapshot_manager(); + let setup = total_start.map_or(Duration::ZERO, |start| start.elapsed()); + let snapshot_start = timing_enabled.then(Instant::now); let snapshot = match crate::table::time_travel::resolve_snapshot(self.table).await? { Some(s) => s, - None => return Ok(vec![SearchResult::empty(); vector_searches.len()]), + None => { + let snapshot = snapshot_start.map_or(Duration::ZERO, |start| start.elapsed()); + let results = vec![SearchResult::empty(); vector_searches.len()]; + if let Some(total_start) = total_start { + let total = total_start.elapsed(); + let unattributed = total.saturating_sub(setup.saturating_add(snapshot)); + eprintln!( + "event=paimon_vector_search_api nq={} index_entries=0 result_count=0 total_ms={:.3} setup_ms={:.3} snapshot_ms={:.3} manifest_ms=0.000 evaluate_ms=0.000 unattributed_ms={:.3}", + vector_searches.len(), + total.as_secs_f64() * 1000.0, + setup.as_secs_f64() * 1000.0, + snapshot.as_secs_f64() * 1000.0, + unattributed.as_secs_f64() * 1000.0, + ); + } + return Ok(results); + } }; + let snapshot_elapsed = snapshot_start.map_or(Duration::ZERO, |start| start.elapsed()); + let manifest_start = timing_enabled.then(Instant::now); let index_entries = match snapshot.index_manifest() { Some(index_manifest_name) => { let manifest_path = snapshot_manager.manifest_path(index_manifest_name); @@ -1240,8 +1263,10 @@ impl<'a> BatchVectorSearchBuilder<'a> { } None => Vec::new(), }; + let manifest = manifest_start.map_or(Duration::ZERO, |start| start.elapsed()); - evaluate_batch_vector_search( + let evaluate_start = timing_enabled.then(Instant::now); + let results = evaluate_batch_vector_search( VectorSearchEvaluation { table: Some(self.table), file_io: self.table.file_io(), @@ -1253,7 +1278,32 @@ impl<'a> BatchVectorSearchBuilder<'a> { &index_entries, &vector_searches, ) - .await + .await?; + if let (Some(total_start), Some(evaluate_start)) = (total_start, evaluate_start) { + let total = total_start.elapsed(); + let evaluate = evaluate_start.elapsed(); + let children = setup + .saturating_add(snapshot_elapsed) + .saturating_add(manifest) + .saturating_add(evaluate); + let result_count = results + .iter() + .map(|result| result.row_ids.len()) + .sum::(); + eprintln!( + "event=paimon_vector_search_api nq={} index_entries={} result_count={} total_ms={:.3} setup_ms={:.3} snapshot_ms={:.3} manifest_ms={:.3} evaluate_ms={:.3} unattributed_ms={:.3}", + vector_searches.len(), + index_entries.len(), + result_count, + total.as_secs_f64() * 1000.0, + setup.as_secs_f64() * 1000.0, + snapshot_elapsed.as_secs_f64() * 1000.0, + manifest.as_secs_f64() * 1000.0, + evaluate.as_secs_f64() * 1000.0, + total.saturating_sub(children).as_secs_f64() * 1000.0, + ); + } + Ok(results) } /// Run a batch of vector searches and materialize each query's matching rows as @@ -1408,6 +1458,12 @@ struct VectorSearchEvaluation<'a> { next_row_id: Option, } +#[derive(Default)] +struct IndexSearchTiming { + permit_wait: Duration, + file_reader_open: Duration, +} + #[cfg(test)] async fn evaluate_vector_search( evaluation: VectorSearchEvaluation<'_>, @@ -1429,6 +1485,8 @@ async fn evaluate_batch_vector_search( index_entries: &[IndexManifestEntry], vector_searches: &[VectorSearch], ) -> crate::Result> { + let timing_enabled = vector_search_timing_enabled(); + let total_start = timing_enabled.then(Instant::now); if vector_searches.is_empty() { return Ok(Vec::new()); } @@ -1480,6 +1538,7 @@ async fn evaluate_batch_vector_search( return Ok(vec![SearchResult::empty(); vector_searches.len()]); } + let deletion_vector_start = timing_enabled.then(Instant::now); let deleted_row_index = if core_options.data_evolution_enabled() { match evaluation.table { Some(table) => { @@ -1492,6 +1551,7 @@ async fn evaluate_batch_vector_search( } else { None }; + let deletion_vector = deletion_vector_start.map_or(Duration::ZERO, |start| start.elapsed()); let max_limit = vector_searches .iter() @@ -1509,8 +1569,16 @@ async fn evaluate_batch_vector_search( }; let index_search_limit = indexed_search_limit(max_limit, refine_factor)?; + let vector_entry_count = vector_entries.len(); + let mut permit_wait = Duration::ZERO; + let mut file_reader_open = Duration::ZERO; + let mut index_search = Duration::ZERO; + let mut merge = Duration::ZERO; + let mut refine = Duration::ZERO; + let mut raw_fallback = Duration::ZERO; let mut merged = vec![SearchResult::empty(); vector_searches.len()]; if !vector_entries.is_empty() { + let index_search_start = timing_enabled.then(Instant::now); let concurrency = core_options.global_index_thread_num()?; if concurrency > tokio::sync::Semaphore::MAX_PERMITS { return Err(crate::Error::DataInvalid { @@ -1551,9 +1619,13 @@ async fn evaluate_batch_vector_search( options.extend(search_options.clone()); let input = evaluation.file_io.new_input(&path); async move { + let permit_start = timing_enabled.then(Instant::now); let permit = acquire_process_global_search_permit(concurrency).await?; + let permit_wait = + permit_start.map_or(Duration::ZERO, |start| start.elapsed()); let input = input?; let query_count = vector_searches.len(); + let mut file_reader_open = Duration::ZERO; let io_meta = GlobalIndexIOMeta::new(file_name.clone(), file_size, index_meta_bytes); let results = match backend { @@ -1585,6 +1657,8 @@ async fn evaluate_batch_vector_search( VectorIndexBackend::Vindex => { match tokio::runtime::Handle::try_current() { Ok(runtime) => { + let file_reader_open_start = + timing_enabled.then(Instant::now); let file_reader = input.reader().await.map_err(|e| { crate::Error::DataInvalid { message: format!( @@ -1594,6 +1668,8 @@ async fn evaluate_batch_vector_search( source: None, } })?; + file_reader_open = file_reader_open_start + .map_or(Duration::ZERO, |start| start.elapsed()); let source = VindexFileReader::new_with_permits( Arc::new(file_reader), runtime, @@ -1601,16 +1677,31 @@ async fn evaluate_batch_vector_search( file_size, file_name.clone(), ); - execute_vindex_searches( + let range_io_stats = source.range_io_stats(); + let results = execute_vindex_searches( io_meta, options, vector_searches, source, - file_name, + file_name.clone(), concurrency, permit, ) - .await? + .await?; + if let Some(stats) = range_io_stats { + let stats = stats.snapshot(); + eprintln!( + "event=paimon_vector_range_io file={} nq={} logical_ranges={} requested_bytes={} file_read_calls={} returned_bytes={} read_ahead_hits={}", + file_name, + query_count, + stats.logical_ranges, + stats.requested_bytes, + stats.file_read_calls, + stats.returned_bytes, + stats.read_ahead_hits, + ); + } + results } Err(_) if query_count > 1 => { let data = input.read().await.map_err(|e| { @@ -1655,7 +1746,7 @@ async fn evaluate_batch_vector_search( }); } - Ok::<_, crate::Error>( + Ok::<_, crate::Error>(( results .into_iter() .map(|result| match result { @@ -1664,20 +1755,30 @@ async fn evaluate_batch_vector_search( None => SearchResult::empty(), }) .collect::>(), - ) + IndexSearchTiming { + permit_wait, + file_reader_open, + }, + )) } }) .collect(); let results = drain_indexed_jobs(futures.into_iter(), concurrency).await?; - for per_entry in &results { + index_search = index_search_start.map_or(Duration::ZERO, |start| start.elapsed()); + let merge_start = timing_enabled.then(Instant::now); + for (per_entry, entry_timing) in &results { + permit_wait = permit_wait.saturating_add(entry_timing.permit_wait); + file_reader_open = file_reader_open.saturating_add(entry_timing.file_reader_open); for (query_index, result) in per_entry.iter().enumerate() { merged[query_index] = merged[query_index].or(result); } } + merge = merge_start.map_or(Duration::ZERO, |start| start.elapsed()); } if refine_factor != 0 { + let refine_start = timing_enabled.then(Instant::now); merged = maybe_rerank_indexed_batch_results( evaluation, index_entries, @@ -1688,9 +1789,11 @@ async fn evaluate_batch_vector_search( index_search_limit, ) .await?; + refine = refine_start.map_or(Duration::ZERO, |start| start.elapsed()); } if search_mode != GlobalIndexSearchMode::Fast { + let raw_fallback_start = timing_enabled.then(Instant::now); let detail_ranges = if search_mode == GlobalIndexSearchMode::Detail { let table = evaluation.table.ok_or_else(|| crate::Error::DataInvalid { message: "Vector raw search in detail mode requires table context".to_string(), @@ -1714,7 +1817,8 @@ async fn evaluate_batch_vector_search( message: "Vector raw search requires table context".to_string(), source: None, })?; - let metric = resolve_raw_vector_metric( + let metric_start = timing_enabled.then(Instant::now); + let (metric, metric_bytes) = resolve_raw_vector_metric( evaluation.file_io, table_path, evaluation.table_options, @@ -1723,15 +1827,35 @@ async fn evaluate_batch_vector_search( field_name, ) .await?; - let raw_results = + let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); + let (raw_results, raw_timing) = read_raw_batch_vector_search(table, vector_searches, &raw_ranges, metric).await?; + if let Some(raw_timing) = raw_timing { + eprintln!( + "event=paimon_vector_raw_fallback nq={} row_ranges={} metric_resolve_ms={:.3} metric_index_bytes={} raw_plan_ms={:.3} split_count={} file_count={} raw_stream_wait_ms={:.3} raw_score_cpu_ms={:.3} arrow_batches={} arrow_rows={} total_raw_read_ms={:.3}", + vector_searches.len(), + raw_ranges.len(), + metric_resolve.as_secs_f64() * 1000.0, + metric_bytes, + raw_timing.plan.as_secs_f64() * 1000.0, + raw_timing.split_count, + raw_timing.file_count, + raw_timing.stream_wait.as_secs_f64() * 1000.0, + raw_timing.score_cpu.as_secs_f64() * 1000.0, + raw_timing.batch_count, + raw_timing.row_count, + raw_timing.total.as_secs_f64() * 1000.0, + ); + } for (query_index, result) in raw_results.iter().enumerate() { merged[query_index] = merged[query_index].or(result); } } + raw_fallback = raw_fallback_start.map_or(Duration::ZERO, |start| start.elapsed()); } - merged + let finalize_start = timing_enabled.then(Instant::now); + let results = merged .into_iter() .zip(vector_searches) .map(|(result, vector_search)| { @@ -1739,7 +1863,40 @@ async fn evaluate_batch_vector_search( .without_deleted_row_ranges(deleted_row_index.as_ref())? .top_k(vector_search.limit)) }) - .collect() + .collect::>>()?; + let finalize = finalize_start.map_or(Duration::ZERO, |start| start.elapsed()); + if let Some(total_start) = total_start { + let total = total_start.elapsed(); + let children = deletion_vector + .saturating_add(index_search) + .saturating_add(merge) + .saturating_add(refine) + .saturating_add(raw_fallback) + .saturating_add(finalize); + let result_count = results + .iter() + .map(|result| result.row_ids.len()) + .sum::(); + eprintln!( + "event=paimon_vector_search_evaluate nq={} index_entries={} index_files={} result_count={} refine_factor={} total_ms={:.3} deletion_vector_ms={:.3} index_search_ms={:.3} global_permit_wait_ms={:.3} file_reader_open_ms={:.3} merge_ms={:.3} refine_ms={:.3} raw_fallback_ms={:.3} finalize_ms={:.3} unattributed_ms={:.3}", + vector_searches.len(), + index_entries.len(), + vector_entry_count, + result_count, + refine_factor, + total.as_secs_f64() * 1000.0, + deletion_vector.as_secs_f64() * 1000.0, + index_search.as_secs_f64() * 1000.0, + permit_wait.as_secs_f64() * 1000.0, + file_reader_open.as_secs_f64() * 1000.0, + merge.as_secs_f64() * 1000.0, + refine.as_secs_f64() * 1000.0, + raw_fallback.as_secs_f64() * 1000.0, + finalize.as_secs_f64() * 1000.0, + total.saturating_sub(children).as_secs_f64() * 1000.0, + ); + } + Ok(results) } fn is_vector_global_index_file(index_file: &IndexFileMeta) -> bool { @@ -2292,12 +2449,16 @@ async fn maybe_rerank_indexed_batch_results( results: Vec, index_search_limit: usize, ) -> crate::Result> { + let timing_enabled = vector_search_timing_enabled(); + let total_start = timing_enabled.then(Instant::now); let mut candidate_searches = Vec::with_capacity(vector_searches.len()); let mut candidate_results = Vec::with_capacity(vector_searches.len()); let mut union_candidates = RoaringTreemap::new(); + let mut candidate_references = 0usize; for (result, vector_search) in results.into_iter().zip(vector_searches) { let candidates = result.top_k(index_search_limit); + candidate_references = candidate_references.saturating_add(candidates.row_ids.len()); let mut include_row_ids = RoaringTreemap::new(); for &row_id in &candidates.row_ids { include_row_ids.insert(row_id); @@ -2318,8 +2479,10 @@ async fn maybe_rerank_indexed_batch_results( message: "Vector index rerank requires table context".to_string(), source: None, })?; + let unique_candidates = union_candidates.len(); let raw_ranges = sorted_row_ids_to_row_ranges(union_candidates.iter())?; - let metric = resolve_raw_vector_metric( + let metric_start = timing_enabled.then(Instant::now); + let (metric, metric_index_bytes) = resolve_raw_vector_metric( evaluation.file_io, evaluation.table_path.trim_end_matches('/'), evaluation.table_options, @@ -2328,8 +2491,30 @@ async fn maybe_rerank_indexed_batch_results( field_name, ) .await?; - - read_raw_batch_vector_search(table, &candidate_searches, &raw_ranges, metric).await + let metric_resolve = metric_start.map_or(Duration::ZERO, |start| start.elapsed()); + + let (results, raw_timing) = + read_raw_batch_vector_search(table, &candidate_searches, &raw_ranges, metric).await?; + if let (Some(total_start), Some(raw_timing)) = (total_start, raw_timing) { + eprintln!( + "event=paimon_vector_refine nq={} candidate_references={} unique_candidates={} row_ranges={} metric_resolve_ms={:.3} metric_index_bytes={} raw_plan_ms={:.3} split_count={} file_count={} raw_stream_wait_ms={:.3} raw_score_cpu_ms={:.3} arrow_batches={} arrow_rows={} total_refine_ms={:.3}", + vector_searches.len(), + candidate_references, + unique_candidates, + raw_ranges.len(), + metric_resolve.as_secs_f64() * 1000.0, + metric_index_bytes, + raw_timing.plan.as_secs_f64() * 1000.0, + raw_timing.split_count, + raw_timing.file_count, + raw_timing.stream_wait.as_secs_f64() * 1000.0, + raw_timing.score_cpu.as_secs_f64() * 1000.0, + raw_timing.batch_count, + raw_timing.row_count, + total_start.elapsed().as_secs_f64() * 1000.0, + ); + } + Ok(results) } fn sorted_row_ids_to_row_ranges( @@ -2509,7 +2694,7 @@ async fn resolve_raw_vector_metric( index_entries: &[IndexManifestEntry], field_id: i32, field_name: &str, -) -> crate::Result { +) -> crate::Result<(RawVectorMetric, usize)> { for entry in index_entries { if entry.kind != FileKind::Add { continue; @@ -2529,7 +2714,7 @@ async fn resolve_raw_vector_metric( if let Some(index_meta) = global_meta.index_meta.as_ref() { if !index_meta.is_empty() { let metric = LuminaIndexMeta::deserialize(index_meta)?.metric()?; - return Ok(RawVectorMetric::from_lumina(metric)); + return Ok((RawVectorMetric::from_lumina(metric), 0)); } } } @@ -2543,6 +2728,7 @@ async fn resolve_raw_vector_metric( ), source: None, })?; + let index_bytes = bytes.len(); let reader = VIndexReader::open(Cursor::new(bytes.to_vec())).map_err(|e| { crate::Error::DataInvalid { message: format!( @@ -2552,12 +2738,15 @@ async fn resolve_raw_vector_metric( source: Some(Box::new(e)), } })?; - return Ok(RawVectorMetric::from_vindex(reader.metadata().metric)); + return Ok(( + RawVectorMetric::from_vindex(reader.metadata().metric), + index_bytes, + )); } } } - configured_raw_vector_metric(table_options, field_name) + Ok((configured_raw_vector_metric(table_options, field_name)?, 0)) } fn configured_raw_vector_metric( @@ -2598,17 +2787,31 @@ fn configured_raw_vector_metric( Ok(inferred.unwrap_or(RawVectorMetric::L2)) } +#[derive(Default)] +struct RawVectorReadTiming { + plan: Duration, + stream_wait: Duration, + score_cpu: Duration, + total: Duration, + split_count: usize, + file_count: usize, + batch_count: usize, + row_count: usize, +} + async fn read_raw_batch_vector_search( table: &Table, vector_searches: &[VectorSearch], raw_ranges: &[RowRange], metric: RawVectorMetric, -) -> crate::Result> { +) -> crate::Result<(Vec, Option)> { + let timing_enabled = vector_search_timing_enabled(); + let total_start = timing_enabled.then(Instant::now); if vector_searches.is_empty() { - return Ok(Vec::new()); + return Ok((Vec::new(), None)); } if raw_ranges.is_empty() { - return Ok(vec![SearchResult::empty(); vector_searches.len()]); + return Ok((vec![SearchResult::empty(); vector_searches.len()], None)); } let field_name = &vector_searches[0].field_name; @@ -2623,13 +2826,28 @@ async fn read_raw_batch_vector_search( }); } + let plan_start = timing_enabled.then(Instant::now); let mut read_builder = table.new_read_builder(); read_builder .with_projection(&[field_name.as_str(), ROW_ID_FIELD_NAME])? .with_row_ranges(raw_ranges.to_vec()); let plan = read_builder.new_scan().plan().await?; + let plan_elapsed = plan_start.map_or(Duration::ZERO, |start| start.elapsed()); + let split_count = plan.splits().len(); + let file_count = plan + .splits() + .iter() + .map(|split| split.data_files().len()) + .sum(); if plan.splits().is_empty() { - return Ok(vec![SearchResult::empty(); vector_searches.len()]); + return Ok(( + vec![SearchResult::empty(); vector_searches.len()], + total_start.map(|start| RawVectorReadTiming { + plan: plan_elapsed, + total: start.elapsed(), + ..RawVectorReadTiming::default() + }), + )); } let read = read_builder.new_read()?; let mut stream = read.to_arrow(plan.splits())?; @@ -2639,14 +2857,44 @@ async fn read_raw_batch_vector_search( .iter() .map(|vector_search| RawScoreTopK::new(vector_search.limit)) .collect::>(); - while let Some(batch) = stream.try_next().await? { + let mut timing = timing_enabled.then(|| RawVectorReadTiming { + plan: plan_elapsed, + split_count, + file_count, + ..RawVectorReadTiming::default() + }); + loop { + let stream_wait_start = timing_enabled.then(Instant::now); + let batch = stream.try_next().await?; + if let (Some(timing), Some(stream_wait_start)) = (&mut timing, stream_wait_start) { + timing.stream_wait = timing + .stream_wait + .saturating_add(stream_wait_start.elapsed()); + } + let Some(batch) = batch else { + break; + }; + if let Some(timing) = &mut timing { + timing.batch_count += 1; + timing.row_count = timing.row_count.saturating_add(batch.num_rows()); + } + let score_start = timing_enabled.then(Instant::now); collect_raw_batch_vector_batch(&batch, vector_searches, metric, &scoring_plan, &mut top_k)?; + if let (Some(timing), Some(score_start)) = (&mut timing, score_start) { + timing.score_cpu = timing.score_cpu.saturating_add(score_start.elapsed()); + } } - Ok(top_k - .into_iter() - .map(RawScoreTopK::into_search_result) - .collect()) + if let (Some(timing), Some(total_start)) = (&mut timing, total_start) { + timing.total = total_start.elapsed(); + } + Ok(( + top_k + .into_iter() + .map(RawScoreTopK::into_search_result) + .collect(), + timing, + )) } struct RawScoringPlan { diff --git a/crates/paimon/src/vindex/mod.rs b/crates/paimon/src/vindex/mod.rs index c6db3adf..c059c3e2 100644 --- a/crates/paimon/src/vindex/mod.rs +++ b/crates/paimon/src/vindex/mod.rs @@ -24,6 +24,7 @@ pub mod pkvector; use crate::spec::{DataField, DataType}; use paimon_vindex_core::index::VectorIndexConfig; use std::collections::HashMap; +use std::sync::OnceLock; pub const IVF_FLAT_IDENTIFIER: &str = "ivf-flat"; pub const IVF_PQ_IDENTIFIER: &str = "ivf-pq"; @@ -33,6 +34,12 @@ const DEFAULT_METRIC: &str = "inner_product"; const DEFAULT_NLIST: &str = "256"; const DEFAULT_PQ_M: &str = "16"; const DEFAULT_PQ_USE_OPQ: &str = "false"; +const VECTOR_SEARCH_TIMING_ENV: &str = "PAIMON_LOG_VECTOR_SEARCH_TIMING"; + +pub(crate) fn vector_search_timing_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| std::env::var_os(VECTOR_SEARCH_TIMING_ENV).is_some_and(|v| v == "1")) +} pub fn is_vindex_index_type(index_type: &str) -> bool { matches!(index_type, IVF_FLAT_IDENTIFIER | IVF_PQ_IDENTIFIER) diff --git a/crates/paimon/src/vindex/range_reader.rs b/crates/paimon/src/vindex/range_reader.rs index 07431749..e54b9bc7 100644 --- a/crates/paimon/src/vindex/range_reader.rs +++ b/crates/paimon/src/vindex/range_reader.rs @@ -16,11 +16,13 @@ // under the License. use crate::io::FileRead; +use crate::vindex::vector_search_timing_enabled; use bytes::Bytes; use futures::future::try_join_all; use paimon_vindex_core::io::{ReadRequest, SeekRead, SeekReadCapabilities}; use std::io; use std::ops::Range; +use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{mpsc, Arc}; const SCALAR_READ_MAX: usize = 64; @@ -54,6 +56,36 @@ struct MergedRange { requested_bytes: u64, } +#[derive(Debug, Default)] +pub(crate) struct RangeIoStats { + logical_ranges: AtomicU64, + requested_bytes: AtomicU64, + file_read_calls: AtomicU64, + returned_bytes: AtomicU64, + read_ahead_hits: AtomicU64, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub(crate) struct RangeIoStatsSnapshot { + pub(crate) logical_ranges: u64, + pub(crate) requested_bytes: u64, + pub(crate) file_read_calls: u64, + pub(crate) returned_bytes: u64, + pub(crate) read_ahead_hits: u64, +} + +impl RangeIoStats { + pub(crate) fn snapshot(&self) -> RangeIoStatsSnapshot { + RangeIoStatsSnapshot { + logical_ranges: self.logical_ranges.load(Ordering::Relaxed), + requested_bytes: self.requested_bytes.load(Ordering::Relaxed), + file_read_calls: self.file_read_calls.load(Ordering::Relaxed), + returned_bytes: self.returned_bytes.load(Ordering::Relaxed), + read_ahead_hits: self.read_ahead_hits.load(Ordering::Relaxed), + } + } +} + /// Bridges vindex-core's synchronous positional reads to Paimon's asynchronous /// range reader. This type is consumed from a blocking search task; it captures /// the surrounding Tokio runtime so remote storage reads still run asynchronously. @@ -64,6 +96,7 @@ pub(crate) struct VindexFileReader { file_size: u64, path: String, scalar_cache: Option, + stats: Option>, } impl VindexFileReader { @@ -97,9 +130,14 @@ impl VindexFileReader { file_size, path, scalar_cache: None, + stats: vector_search_timing_enabled().then(|| Arc::new(RangeIoStats::default())), } } + pub(crate) fn range_io_stats(&self) -> Option> { + self.stats.clone() + } + fn validate_range(&self, pos: u64, len: usize) -> io::Result> { let end = pos.checked_add(len as u64).ok_or_else(|| { io::Error::new( @@ -127,6 +165,9 @@ impl VindexFileReader { if buf.len() <= SCALAR_READ_MAX { if let Some(cache) = &self.scalar_cache { if cache.contains(&range) { + if let Some(stats) = &self.stats { + stats.read_ahead_hits.fetch_add(1, Ordering::Relaxed); + } let start = (range.start - cache.start) as usize; buf.copy_from_slice(&cache.data[start..start + buf.len()]); return Ok(()); @@ -163,17 +204,22 @@ impl VindexFileReader { let permits = Arc::clone(&self.permits); let path = self.path.clone(); let requested = ranges.to_vec(); + let stats = self.stats.clone(); let (sender, receiver) = mpsc::sync_channel(1); self.runtime.spawn(async move { let fetched = try_join_all(requested.iter().cloned().map(|range| { let reader = Arc::clone(&reader); let permits = Arc::clone(&permits); let path = path.clone(); + let stats = stats.clone(); async move { let _permit = permits.acquire_owned().await.map_err(|_| { io::Error::other("vindex range read concurrency limiter closed") })?; let expected = (range.end - range.start) as usize; + if let Some(stats) = &stats { + stats.file_read_calls.fetch_add(1, Ordering::Relaxed); + } let data = reader.read(range.clone()).await.map_err(|error| { io::Error::other(format!( "failed to read vindex file '{path}' range {}..{}: {error}", @@ -191,6 +237,11 @@ impl VindexFileReader { ), )); } + if let Some(stats) = &stats { + stats + .returned_bytes + .fetch_add(data.len() as u64, Ordering::Relaxed); + } Ok(data) } })) @@ -266,6 +317,20 @@ impl VindexFileReader { impl SeekRead for VindexFileReader { fn pread(&mut self, requests: &mut [ReadRequest<'_>]) -> io::Result<()> { + if let Some(stats) = &self.stats { + let (logical_ranges, requested_bytes) = requests + .iter() + .filter(|request| !request.buf.is_empty()) + .fold((0u64, 0u64), |(ranges, bytes), request| { + (ranges + 1, bytes.saturating_add(request.buf.len() as u64)) + }); + stats + .logical_ranges + .fetch_add(logical_ranges, Ordering::Relaxed); + stats + .requested_bytes + .fetch_add(requested_bytes, Ordering::Relaxed); + } let non_empty = requests .iter() .filter(|request| !request.buf.is_empty()) @@ -289,6 +354,7 @@ impl SeekRead for VindexFileReader { file_size: self.file_size, path: self.path.clone(), scalar_cache: None, + stats: self.stats.clone(), })) } @@ -592,6 +658,85 @@ mod tests { assert_eq!(cloned.read_capabilities(), SeekReadCapabilities::default()); } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn range_io_stats_are_shared_across_clones() { + let data = Bytes::from(vec![7u8; 1024]); + let source: Arc = TrackingRead::new(data.clone()); + let mut reader = VindexFileReader::new( + source, + tokio::runtime::Handle::current(), + data.len() as u64, + "index".to_string(), + ); + let stats = Arc::new(RangeIoStats::default()); + reader.stats = Some(Arc::clone(&stats)); + let mut cloned = reader.try_clone_reader().unwrap().unwrap(); + + tokio::task::spawn_blocking(move || { + let mut first = [0u8; 128]; + reader + .pread(&mut [ReadRequest::new(0, &mut first)]) + .unwrap(); + let mut second = [0u8; 128]; + cloned + .pread(&mut [ReadRequest::new(128, &mut second)]) + .unwrap(); + }) + .await + .unwrap(); + + assert_eq!( + stats.snapshot(), + RangeIoStatsSnapshot { + logical_ranges: 2, + requested_bytes: 256, + file_read_calls: 2, + returned_bytes: 256, + read_ahead_hits: 0, + } + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn range_io_stats_count_coalesced_reads() { + let data = Bytes::from(vec![7u8; 100_000]); + let source: Arc = TrackingRead::new(data.clone()); + let mut reader = VindexFileReader::new( + source, + tokio::runtime::Handle::current(), + data.len() as u64, + "index".to_string(), + ); + let stats = Arc::new(RangeIoStats::default()); + reader.stats = Some(Arc::clone(&stats)); + + tokio::task::spawn_blocking(move || { + let mut first = [0u8; 4]; + let mut second = [0u8; 4]; + let mut third = [0u8; 4]; + reader + .pread(&mut [ + ReadRequest::new(0, &mut first), + ReadRequest::new(8, &mut second), + ReadRequest::new(20_000, &mut third), + ]) + .unwrap(); + }) + .await + .unwrap(); + + assert_eq!( + stats.snapshot(), + RangeIoStatsSnapshot { + logical_ranges: 3, + requested_bytes: 12, + file_read_calls: 2, + returned_bytes: 16, + read_ahead_hits: 0, + } + ); + } + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn shared_permits_bound_reads_across_independent_readers() { let data = Bytes::from(vec![8u8; 1024]); diff --git a/crates/paimon/src/vindex/reader.rs b/crates/paimon/src/vindex/reader.rs index 21603f17..e231b6f1 100644 --- a/crates/paimon/src/vindex/reader.rs +++ b/crates/paimon/src/vindex/reader.rs @@ -17,6 +17,7 @@ use crate::spec::CoreOptions; use crate::vector_search::{GlobalIndexIOMeta, VectorSearch}; +use crate::vindex::vector_search_timing_enabled; use paimon_vindex_core::distance::MetricType; use paimon_vindex_core::index::{ VectorIndexMetadata, VectorIndexReader as VIndexReader, VectorSearchParams, @@ -25,11 +26,28 @@ use paimon_vindex_core::io::{ReadRequest, SeekRead, SeekReadCapabilities}; use std::collections::BinaryHeap; use std::collections::HashMap; use std::io; +use std::time::{Duration, Instant}; const DEFAULT_NPROBE: usize = 16; const NPROBE_PARAMETER: &str = "ivf.nprobe"; const NATIVE_BATCH_OPERATION_WORKING_SET_BYTES: usize = 64 * 1024 * 1024; +#[derive(Clone, Copy, Default)] +struct VindexLoadTiming { + vindex_open: Duration, + metadata: Duration, + optimize: Duration, +} + +#[derive(Default)] +struct VindexBatchStats { + native_chunk_queries: Vec, + scalar_chunk_count: usize, + max_chunk_size: usize, + memory_budget_bytes: usize, + batch_index_parallelism: usize, +} + trait ErasedSeekRead: Send { fn pread_erased(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()>; @@ -81,6 +99,9 @@ pub struct VindexVectorGlobalIndexReader { batch_shard_concurrency: Option, reader: Option>, metadata: Option, + timing_enabled: bool, + load_timing: VindexLoadTiming, + batch_stats: Option, } impl VindexVectorGlobalIndexReader { @@ -91,6 +112,9 @@ impl VindexVectorGlobalIndexReader { batch_shard_concurrency: None, reader: None, metadata: None, + timing_enabled: vector_search_timing_enabled(), + load_timing: VindexLoadTiming::default(), + batch_stats: None, } } @@ -113,8 +137,52 @@ impl VindexVectorGlobalIndexReader { vector_searches: &[VectorSearch], stream_fn: impl FnOnce(&str) -> crate::Result, ) -> crate::Result>>> { + let total_start = self.timing_enabled.then(Instant::now); self.ensure_loaded(stream_fn, |_| Ok(()))?; - self.search_batch(vector_searches) + let search_start = self.timing_enabled.then(Instant::now); + let results = self.search_batch(vector_searches)?; + if let (Some(total_start), Some(search_start), Some(stats)) = + (total_start, search_start, self.batch_stats.as_ref()) + { + let total = total_start.elapsed(); + let native_search = search_start.elapsed(); + let load = self + .load_timing + .vindex_open + .saturating_add(self.load_timing.metadata) + .saturating_add(self.load_timing.optimize); + let unattributed = total.saturating_sub(load.saturating_add(native_search)); + let chunk_queries = stats + .native_chunk_queries + .iter() + .map(usize::to_string) + .collect::>() + .join(","); + let nprobe = self + .options + .get(NPROBE_PARAMETER) + .map(String::as_str) + .unwrap_or("16"); + eprintln!( + "event=paimon_vindex_reader file={} nq={} nprobe={} batch_index_parallelism={} memory_budget_bytes={} max_chunk_size={} native_chunk_count={} native_chunk_queries={} scalar_chunk_count={} total_ms={:.3} vindex_open_ms={:.3} metadata_ms={:.3} optimize_ms={:.3} native_search_wall_ms={:.3} unattributed_ms={:.3}", + self.io_meta.file_path, + vector_searches.len(), + nprobe, + stats.batch_index_parallelism, + stats.memory_budget_bytes, + stats.max_chunk_size, + stats.native_chunk_queries.len(), + chunk_queries, + stats.scalar_chunk_count, + total.as_secs_f64() * 1000.0, + self.load_timing.vindex_open.as_secs_f64() * 1000.0, + self.load_timing.metadata.as_secs_f64() * 1000.0, + self.load_timing.optimize.as_secs_f64() * 1000.0, + native_search.as_secs_f64() * 1000.0, + unattributed.as_secs_f64() * 1000.0, + ); + } + Ok(results) } #[cfg(test)] @@ -168,13 +236,16 @@ impl VindexVectorGlobalIndexReader { message: "vindex metadata not initialized".to_string(), source: None, })?; - search_batch_vindex( + let (results, batch_stats) = search_batch_vindex( reader, metadata, &self.options, vector_searches, shard_concurrency, - ) + self.timing_enabled, + )?; + self.batch_stats = batch_stats; + Ok(results) } fn search(&mut self, vector_search: &VectorSearch) -> crate::Result>> { @@ -227,9 +298,11 @@ impl VindexVectorGlobalIndexReader { O: FnOnce(&mut VIndexReader) -> crate::Result<()>, { if self.reader.is_some() { + self.load_timing = VindexLoadTiming::default(); return validate(self.metadata()?); } + let open_start = self.timing_enabled.then(Instant::now); let source = stream_fn(&self.io_meta.file_path)?; let mut reader = VIndexReader::open(VindexInput::new(source)).map_err(|e| { crate::Error::DataInvalid { @@ -237,12 +310,22 @@ impl VindexVectorGlobalIndexReader { source: Some(Box::new(e)), } })?; + let vindex_open = open_start.map_or(Duration::ZERO, |start| start.elapsed()); + let metadata_start = self.timing_enabled.then(Instant::now); let metadata = reader.metadata(); + let metadata_elapsed = metadata_start.map_or(Duration::ZERO, |start| start.elapsed()); validate(&metadata)?; + let optimize_start = self.timing_enabled.then(Instant::now); optimize(&mut reader)?; + let optimize_elapsed = optimize_start.map_or(Duration::ZERO, |start| start.elapsed()); self.reader = Some(reader); self.metadata = Some(metadata); + self.load_timing = VindexLoadTiming { + vindex_open, + metadata: metadata_elapsed, + optimize: optimize_elapsed, + }; Ok(()) } } @@ -349,10 +432,16 @@ fn search_batch_vindex( options: &HashMap, vector_searches: &[VectorSearch], shard_concurrency: usize, -) -> crate::Result>>> { + timing_enabled: bool, +) -> crate::Result<(Vec>>, Option)> { let mut results: Vec>> = (0..vector_searches.len()).map(|_| None).collect(); let mut groups: Vec<(PreparedSearch, Vec)> = Vec::new(); + let mut batch_stats = timing_enabled.then(|| VindexBatchStats { + memory_budget_bytes: NATIVE_BATCH_OPERATION_WORKING_SET_BYTES / shard_concurrency.max(1), + batch_index_parallelism: shard_concurrency, + ..VindexBatchStats::default() + }); for (index, search) in vector_searches.iter().enumerate() { let Some(prepared) = prepare_search(metadata, options, search)? else { @@ -367,8 +456,14 @@ fn search_batch_vindex( for (prepared, indices) in groups { let chunk_size = native_batch_chunk_size(metadata, &prepared, shard_concurrency); + if let Some(stats) = &mut batch_stats { + stats.max_chunk_size = stats.max_chunk_size.max(chunk_size); + } for indices in indices.chunks(chunk_size) { if indices.len() == 1 { + if let Some(stats) = &mut batch_stats { + stats.scalar_chunk_count += 1; + } let index = indices[0]; let (labels, distances) = execute_scalar_search(reader, &vector_searches[index], &prepared)?; @@ -378,6 +473,9 @@ fn search_batch_vindex( } continue; } + if let Some(stats) = &mut batch_stats { + stats.native_chunk_queries.push(indices.len()); + } let mut queries = Vec::with_capacity(indices.len() * metadata.dimension); for &index in indices { @@ -425,7 +523,7 @@ fn search_batch_vindex( } } - Ok(results) + Ok((results, batch_stats)) } fn native_batch_chunk_size( From 40f9fc995a9deb927f29d8cc882cf2436a4e5688 Mon Sep 17 00:00:00 2001 From: yantian Date: Fri, 14 Aug 2026 14:23:38 +0800 Subject: [PATCH 2/2] fix: simplify vindex batch result type --- crates/paimon/src/vindex/reader.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/paimon/src/vindex/reader.rs b/crates/paimon/src/vindex/reader.rs index e231b6f1..6f049ff1 100644 --- a/crates/paimon/src/vindex/reader.rs +++ b/crates/paimon/src/vindex/reader.rs @@ -48,6 +48,8 @@ struct VindexBatchStats { batch_index_parallelism: usize, } +type VindexBatchSearchResult = (Vec>>, Option); + trait ErasedSeekRead: Send { fn pread_erased(&mut self, ranges: &mut [ReadRequest<'_>]) -> io::Result<()>; @@ -433,7 +435,7 @@ fn search_batch_vindex( vector_searches: &[VectorSearch], shard_concurrency: usize, timing_enabled: bool, -) -> crate::Result<(Vec>>, Option)> { +) -> crate::Result { let mut results: Vec>> = (0..vector_searches.len()).map(|_| None).collect(); let mut groups: Vec<(PreparedSearch, Vec)> = Vec::new();