leaves12138 commented on code in PR #62:
URL:
https://github.com/apache/paimon-vector-index/pull/62#discussion_r3651319456
##########
core/src/index.rs:
##########
@@ -671,36 +1248,118 @@ impl<R: SeekRead> VectorIndexReader<R> {
self.metadata().total_vectors
}
+ pub fn diskann_search_stats(&self) -> Option<DiskAnnSearchStats> {
+ match self {
+ Self::DiskAnn(reader) => Some(reader.last_search_stats()),
+ _ => None,
+ }
+ }
+
+ pub fn effective_storage_profile(&self) -> Option<StorageProfile> {
+ match self {
+ Self::DiskAnn(reader) => Some(reader.effective_storage_profile()),
+ _ => None,
+ }
+ }
+
pub fn optimize_for_search(&mut self) -> io::Result<()> {
match self {
Self::IvfFlat(reader) => reader.ensure_loaded(),
+ Self::IvfSq(reader) => reader.optimize_for_search(),
Self::IvfPq(reader) => reader.optimize_for_search(),
Self::IvfRq(reader) => reader.ensure_loaded(),
- Self::IvfHnswFlat(reader) => reader.ensure_loaded(),
- // IVF_HNSW_SQ warms SQ scan/fallback structures used by filtered
- // searches; normal unfiltered search primarily uses the HNSW
graph.
- Self::IvfHnswSq(reader) => reader.optimize_for_search(),
+ Self::DiskAnn(reader) => reader.optimize_for_search(),
}
}
- pub fn search(
+ /// Warm query-dependent caches with representative queries. DiskANN runs
+ /// the graph and rerank path; other index types perform their normal
+ /// resident optimization because they do not expose a paged query cache.
+ pub fn warmup_queries(
&mut self,
- query: &[f32],
- params: VectorSearchParams,
- ) -> io::Result<(Vec<i64>, Vec<f32>)> {
- validate_query(query, self.dimension())?;
- validate_query_bits_for_index(self.index_type(), params.query_bits)?;
- match self {
- Self::IvfFlat(reader) => reader.search(query, params.top_k,
params.nprobe),
- Self::IvfPq(reader) => search_with_reader(reader, query,
params.top_k, params.nprobe),
- Self::IvfRq(reader) => {
- reader.search_with_query_bits(query, params.top_k,
params.nprobe, params.query_bits)
- }
- Self::IvfHnswFlat(reader) => {
- reader.search(query, params.top_k, params.nprobe,
params.hnsw_ef_search())
+ queries: &[f32],
+ query_count: usize,
+ l_search: usize,
+ ) -> io::Result<()> {
+ let expected_len = query_count
+ .checked_mul(self.dimension())
+ .ok_or_else(|| invalid_input("warmup query count * dimension
overflows usize"))?;
+ if queries.len() != expected_len {
+ return Err(invalid_input(format!(
+ "warmup queries length {} does not match query count *
dimension {}",
+ queries.len(),
+ expected_len
+ )));
+ }
+ validate_finite_values(queries, expected_len, "warmup queries")?;
+ match self {
+ Self::DiskAnn(reader) => reader.warmup_queries(queries, l_search),
+ _ => self.optimize_for_search(),
+ }
+ }
+
+ pub fn calibrate_search_width(
+ &mut self,
+ queries: &[f32],
+ query_count: usize,
+ top_k: usize,
+ ) -> io::Result<usize> {
+ validate_queries(queries, query_count, self.dimension())?;
+ validate_positive(top_k, "top_k")?;
+ match self {
+ Self::DiskAnn(reader) => reader.calibrate_l_search(queries, top_k),
+ _ => Err(invalid_input(
+ "search-width calibration is currently only available for
DiskANN",
+ )),
+ }
+ }
+
+ pub fn search(
+ &mut self,
+ query: &[f32],
+ params: VectorSearchParams,
+ ) -> io::Result<(Vec<i64>, Vec<f32>)> {
+ validate_query(query, self.dimension())?;
+ match self {
+ Self::IvfFlat(reader) => {
+ let nprobe = params.resolve_ivf_nprobe(
+ reader.nlist,
+ usize::try_from(reader.total_vectors)
+ .map_err(|_| invalid_input("negative IVF vector
count"))?,
+ None,
+ )?;
+ reader.search(query, params.top_k, nprobe)
Review Comment:
Automatic IVF search stops after the single inferred `nprobe` on the
unfiltered paths, unlike the filtered paths below that call
`progressive_ivf_search`. This can return padded `-1` entries even though the
index contains enough rows. I reproduced it with 64 lists: the eight closest
lists contain one row each and another list contains 1,000 rows; with
`total_vectors=1008` and `top_k=10`, auto infers `nprobe=8` and returns only
eight IDs plus two paddings, while `nprobe=64` returns ten IDs. Could we apply
progressive expansion for `SearchWidth::Auto` here as well, including batch
search and the other IVF variants?
--
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]
