JingsongLi commented on code in PR #646:
URL: https://github.com/apache/paimon-rust/pull/646#discussion_r3709511654
##########
crates/paimon/src/table/vector_search_builder.rs:
##########
@@ -1482,65 +1482,40 @@ async fn evaluate_batch_vector_search(
.await?
}
VectorIndexBackend::Vindex => {
- if vector_searches.len() > 1 {
- let data = input.read().await.map_err(|e| {
- crate::Error::DataInvalid {
- message: format!(
- "Failed to read vindex index file
'{}': {}",
- file_name, e
- ),
- source: None,
- }
- })?;
- execute_global_index(
- "vindex global-index batch search task
failed",
- move || {
- let mut reader =
VindexVectorGlobalIndexReader::new(
- io_meta, options,
- );
-
reader.visit_batch_vector_search(&vector_searches, |_| {
- Ok(Cursor::new(data))
- })
- },
- )
- .await?
+ let file_reader = input.reader().await.map_err(|e|
{
+ crate::Error::DataInvalid {
+ message: format!(
+ "Failed to open vindex file '{}' for
range reads: {}",
+ file_name, e
+ ),
+ source: None,
+ }
+ })?;
+ let source = VindexFileReader::new(
+ Arc::new(file_reader),
+ current_tokio_runtime_handle()?,
Review Comment:
[P2] Preserve the non-Tokio batch fallback
This now unconditionally requires a Tokio handle, but before this change the
`vector_searches.len() > 1` branch read the file into a `Cursor` and therefore
worked when polled by async-std, smol, or a plain `futures` executor. A public
batch with two or more queries will now fail before searching with `Vector
index range reader requires a Tokio runtime`.
Please retain the whole-file `Cursor` fallback when `Handle::try_current()`
fails, or make the range bridge executor-agnostic, and add a non-Tokio batch
regression test. The existing test only covers the pre-existing single-query
limitation.
##########
crates/paimon/src/vindex/reader.rs:
##########
@@ -165,61 +222,175 @@ fn search_vindex(
options: &HashMap<String, String>,
vector_search: &VectorSearch,
) -> crate::Result<Option<HashMap<u64, f32>>> {
- let expected_dim = metadata.dimension;
- if vector_search.vector.len() != expected_dim {
+ let Some(prepared) = prepare_search(metadata, options, vector_search)?
else {
+ return Ok(None);
+ };
+ let (labels, distances) = execute_scalar_search(reader, vector_search,
&prepared)?;
+ let id_to_scores = collect_results(&labels, &distances, prepared.top_k,
metadata.metric);
+ if id_to_scores.is_empty() {
+ return Ok(None);
+ }
+
+ Ok(Some(id_to_scores))
+}
+
+#[derive(Clone, PartialEq, Eq)]
+struct PreparedSearch {
+ top_k: usize,
+ nprobe: usize,
+ filter_bytes: Option<Vec<u8>>,
+}
+
+fn prepare_search(
+ metadata: &VectorIndexMetadata,
+ options: &HashMap<String, String>,
+ vector_search: &VectorSearch,
+) -> crate::Result<Option<PreparedSearch>> {
+ if vector_search.vector.len() != metadata.dimension {
return Err(crate::Error::DataInvalid {
message: format!(
"Query vector dimension mismatch: index expects {}, but got
{}",
- expected_dim,
+ metadata.dimension,
vector_search.vector.len()
),
source: None,
});
}
let count = usize::try_from(metadata.total_vectors).unwrap_or(0);
- let effective_k = std::cmp::min(vector_search.limit, count);
- if effective_k == 0 {
+ let mut top_k = vector_search.limit.min(count);
+ if top_k == 0 {
return Ok(None);
}
-
let nprobe = int_parameter(options, NPROBE_PARAMETER, DEFAULT_NPROBE)?;
- let params = VectorSearchParams::new(effective_k, nprobe);
- let (labels, distances) = if let Some(include_ids) =
&vector_search.include_row_ids {
+ let filter_bytes = if let Some(include_ids) =
&vector_search.include_row_ids {
if include_ids.is_empty() {
return Ok(None);
}
- let ek = std::cmp::min(effective_k, include_ids.len() as usize);
- let params = VectorSearchParams::new(params.top_k.min(ek), nprobe);
- let mut filter_bytes = Vec::new();
+ top_k = top_k.min(include_ids.len() as usize);
+ let mut bytes = Vec::new();
include_ids
- .serialize_into(&mut filter_bytes)
+ .serialize_into(&mut bytes)
.map_err(|e| crate::Error::DataInvalid {
message: format!("Failed to serialize vector search row-id
filter: {}", e),
source: Some(Box::new(e)),
})?;
- reader
- .search_with_roaring_filter(&vector_search.vector, params,
&filter_bytes)
- .map_err(|e| crate::Error::DataInvalid {
- message: format!("paimon-vindex-core filtered search failed:
{}", e),
- source: Some(Box::new(e)),
- })?
+ Some(bytes)
} else {
- reader
- .search(&vector_search.vector, params)
+ None
+ };
+
+ Ok(Some(PreparedSearch {
+ top_k,
+ nprobe,
+ filter_bytes,
+ }))
+}
+
+fn execute_scalar_search(
+ reader: &mut VIndexReader<impl SeekRead>,
+ vector_search: &VectorSearch,
+ prepared: &PreparedSearch,
+) -> crate::Result<(Vec<i64>, Vec<f32>)> {
+ let params = VectorSearchParams::new(prepared.top_k, prepared.nprobe);
+ match &prepared.filter_bytes {
+ Some(filter) => reader
+ .search_with_roaring_filter(&vector_search.vector, params, filter)
.map_err(|e| crate::Error::DataInvalid {
- message: format!("paimon-vindex-core search failed: {}", e),
+ message: format!("paimon-vindex-core filtered search failed:
{}", e),
source: Some(Box::new(e)),
- })?
- };
+ }),
+ None => {
+ reader
+ .search(&vector_search.vector, params)
+ .map_err(|e| crate::Error::DataInvalid {
+ message: format!("paimon-vindex-core search failed: {}",
e),
+ source: Some(Box::new(e)),
+ })
+ }
+ }
+}
- let id_to_scores = collect_results(&labels, &distances, effective_k,
metadata.metric);
- if id_to_scores.is_empty() {
- return Ok(None);
+fn search_batch_vindex(
+ reader: &mut VIndexReader<impl SeekRead>,
+ metadata: &VectorIndexMetadata,
+ options: &HashMap<String, String>,
+ vector_searches: &[VectorSearch],
+) -> crate::Result<Vec<Option<HashMap<u64, f32>>>> {
+ let mut results: Vec<Option<HashMap<u64, f32>>> =
+ (0..vector_searches.len()).map(|_| None).collect();
+ let mut groups: Vec<(PreparedSearch, Vec<usize>)> = Vec::new();
+
+ for (index, search) in vector_searches.iter().enumerate() {
+ let Some(prepared) = prepare_search(metadata, options, search)? else {
+ continue;
+ };
+ if let Some((_, indices)) = groups.iter_mut().find(|(key, _)| key ==
&prepared) {
+ indices.push(index);
+ } else {
+ groups.push((prepared, vec![index]));
+ }
}
- Ok(Some(id_to_scores))
+ for (prepared, indices) in groups {
+ if indices.len() == 1 {
+ let index = indices[0];
+ let (labels, distances) =
+ execute_scalar_search(reader, &vector_searches[index],
&prepared)?;
+ let map = collect_results(&labels, &distances, prepared.top_k,
metadata.metric);
+ if !map.is_empty() {
+ results[index] = Some(map);
+ }
+ continue;
+ }
+
+ let mut queries = Vec::with_capacity(indices.len() *
metadata.dimension);
+ for &index in &indices {
+ queries.extend_from_slice(&vector_searches[index].vector);
+ }
+ let params = VectorSearchParams::new(prepared.top_k, prepared.nprobe);
+ let (labels, distances) = match &prepared.filter_bytes {
+ Some(filter) => reader
+ .search_batch_with_roaring_filter(&queries, indices.len(),
params, filter)
+ .map_err(|e| crate::Error::DataInvalid {
+ message: format!("paimon-vindex-core filtered batch search
failed: {}", e),
+ source: Some(Box::new(e)),
+ })?,
+ None => reader
+ .search_batch(&queries, indices.len(), params)
Review Comment:
[P2] Bound the native batch by its full working set
Every compatible group is flattened and passed to one native batch call. In
`paimon-vindex-core` 0.3.0, this allocates a `query_count * nlist` `f32`
centroid-product matrix, plus query copies, per-query heaps, and `query_count *
top_k` result arrays; L2 IVF-PQ also retains per-query `m * ksub` tables.
Posting-list reads are byte-bounded, but this query-side scratch is not.
For example, 10,000 queries with `nlist=4096` consume about 164 MB per
active index-file job for the centroid matrix alone. With the default 32
concurrent shard jobs, those matrices alone can exceed 5 GB. The previous
scalar loop bounded this scratch to one query at a time. Please chunk
compatible groups using a working-set budget derived from at least `nlist`,
dimension, `top_k`, and shard concurrency, then restore results to their
original positions.
##########
crates/paimon/src/table/vector_search_builder.rs:
##########
@@ -1482,65 +1482,40 @@ async fn evaluate_batch_vector_search(
.await?
}
VectorIndexBackend::Vindex => {
- if vector_searches.len() > 1 {
- let data = input.read().await.map_err(|e| {
- crate::Error::DataInvalid {
- message: format!(
- "Failed to read vindex index file
'{}': {}",
- file_name, e
- ),
- source: None,
- }
- })?;
- execute_global_index(
- "vindex global-index batch search task
failed",
- move || {
- let mut reader =
VindexVectorGlobalIndexReader::new(
- io_meta, options,
- );
-
reader.visit_batch_vector_search(&vector_searches, |_| {
- Ok(Cursor::new(data))
- })
- },
- )
- .await?
+ let file_reader = input.reader().await.map_err(|e|
{
+ crate::Error::DataInvalid {
+ message: format!(
+ "Failed to open vindex file '{}' for
range reads: {}",
+ file_name, e
+ ),
+ source: None,
+ }
+ })?;
+ let source = VindexFileReader::new(
Review Comment:
[P2] Apply `global-index.thread-num` to aggregate range I/O
Each admitted shard now creates an independent `VindexFileReader`, and each
reader owns a private 32-permit range-read semaphore. The outer scheduler
already admits up to `global-index.thread-num` shard jobs (default 32), so
native multi-range reads can drive up to `32 * 32 = 1024` concurrent
`FileRead::read` calls. This defeats the option's documented per-operation
global-index I/O fan-out limit; in particular, setting it to `1` no longer
reproduces strictly sequential I/O. The removed batch path issued one
whole-file read per admitted shard.
Please share one operation-wide semaphore across all readers, or otherwise
divide the range-read budget across admitted shard jobs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
