feat: adaptive vector search

[killme2008]

I hereby agree to the terms of the GreptimeDB CLA.

Refer to a related PR or issue link (optional)

#7668

What's changed and what's your intention?

PR Checklist

Please convert it to a draft if some of the following conditions are not met.

• I have written the necessary rustdoc comments.
• I have added the necessary unit tests and integration tests.
• This PR requires documentation updates.
• API changes are backward compatible.
• Schema or data changes are backward compatible.

[gemini-code-assist]

Summary of Changes

Hello @killme2008, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant enhancement to vector search capabilities by implementing an adaptive top-k mechanism. This system intelligently refines vector search queries by iteratively adjusting the number of candidates considered, ensuring more accurate results, especially when complex filters are applied. The changes span across query optimization, physical execution, and storage layer interactions, providing a robust and observable solution for vector similarity search.

Highlights

• Adaptive Vector Search Implementation: Introduced a new adaptive vector search mechanism to dynamically adjust the 'k' value (number of nearest neighbors) during queries, especially when combined with filtering predicates, to improve accuracy and efficiency.
• Optimizer Rule and Execution Plan: Added the AdaptiveVectorTopKRule optimizer rule to transform Sort + Limit logical plans into an AdaptiveVectorTopKLogicalPlan. This is then executed by the new AdaptiveVectorTopKExec physical plan, which manages the adaptive retry loop.
• Predicate Pushdown for Vector Indexes: Enhanced VectorIndexEngine and IndexApplier traits with a search_with_predicate method, allowing filtering predicates to be pushed down directly to the vector index engine for more efficient pruning.
• Configurable Adaptive Behavior: Added vector_topk_max_rounds and vector_topk_max_k options to QueryOptions and FlownodeOptions, enabling configuration of the adaptive search's retry limits and maximum 'k' value.
• Detailed Vector Search Metrics: Integrated new metrics (vector_index_requested_k, vector_index_returned_k, retry_rounds, last_k, desired_rows, result_len, kth_distance_micros) to monitor the performance and behavior of the adaptive vector search process.
• Enhanced Scan Hinting: Updated the ScanHintRule to capture additional vector search parameters like limit, offset, and tie_breakers from the query plan, providing more context for the adaptive search.

Changelog

• src/cmd/tests/load_config_test.rs
  • Added vector_topk_max_rounds and vector_topk_max_k to FlownodeOptions in test configuration.
• src/flow/src/adapter.rs
  • Added vector_topk_max_rounds and vector_topk_max_k to default FlownodeOptions.
• src/index/src/vector/apply.rs
  • Imported VectorSearchPredicate trait.
  • Added search_with_predicate method to IndexApplier trait with a default implementation.
  • Implemented search_with_predicate for HnswVectorIndexApplier, including an adapter for key mapping.
  • Added a new test case test_hnsw_vector_index_applier_search_with_predicate.
• src/index/src/vector/engine/usearch.rs
  • Imported VectorSearchPredicate trait and Key type.
  • Implemented search_with_predicate for UsearchEngine using filtered_search.
• src/mito2/src/read/pruner.rs
  • Added a comment regarding vector index k counters and cache hits.
  • Modified error handling for prune_file_range.
• src/mito2/src/read/scan_region.rs
  • Removed VECTOR_INDEX_OVERFETCH_MULTIPLIER constant.
  • Simplified vector_index_k calculation.
• src/mito2/src/read/scan_util.rs
  • Imported Gauge from DataFusion metrics.
  • Added vector_index_requested_k and vector_index_returned_k fields to ScanMetricsSet and PartitionMetricsInner.
  • Updated Debug implementation for ScanMetricsSet to include new vector index k metrics.
  • Updated merge_reader_metrics to aggregate new vector index k metrics.
  • Updated PartitionMetrics::new and PartitionMetrics::merge_reader_metrics to handle new vector index k metrics.
• src/mito2/src/sst/index/vector_index/applier.rs
  • Imported VectorSearchPredicate trait.
  • Modified apply_with_k to accept an optional predicate argument.
  • Updated calls to applier.search to applier.search_with_predicate.
  • Updated test calls to apply_with_k to pass None for the new predicate argument.
• src/mito2/src/sst/parquet/reader.rs
  • Updated apply_with_k call to pass None for the new predicate argument.
  • Added logic to record vector_index_requested_k and vector_index_returned_k metrics.
  • Added vector_index_requested_k and vector_index_returned_k fields to ReaderFilterMetrics.
  • Updated merge method of ReaderFilterMetrics to aggregate new vector index k metrics.
• src/query/src/dist_plan/merge_scan.rs
  • Added collect_vector_index_k_from_region_metrics function to extract vector index k metrics.
  • Added vector_index_requested_k and vector_index_returned_k gauges to MergeScanMetric.
  • Updated MergeScanMetric::new to initialize new gauges.
  • Added record_vector_index_k method to MergeScanMetric.
  • Modified MergeScanExec to record vector index k metrics from sub-stages.
  • Added unit tests for vector index k metric collection and recording.
• src/query/src/lib.rs
  • Added pub mod vector_search;.
• src/query/src/optimizer.rs
  • Added pub mod adaptive_vector_topk;.
• src/query/src/optimizer/adaptive_vector_topk.rs
  • New file: Introduces AdaptiveVectorTopKRule optimizer rule to rewrite Sort + Limit plans into AdaptiveVectorTopKLogicalPlan for vector distance sorts.
  • Includes a SKIP_REWRITE guard to prevent recursive rewrites during adaptive execution.
• src/query/src/optimizer/scan_hint/vector_search.rs
  • Imported SortOptions and OrderOption.
  • Added tie_breakers field to VectorDistanceInfo.
  • Modified VectorSearchState::new to capture limit, offset, and tie_breakers.
  • Added build_tie_breakers function to extract tie-breaker columns from sort expressions.
  • Added a new test case test_vector_hint_with_tie_breaker_and_offset.
• src/query/src/options.rs
  • Added vector_topk_max_rounds and vector_topk_max_k fields to QueryOptions with default values.
• src/query/src/query_engine/state.rs
  • Imported AdaptiveVectorTopKRule, AdaptiveVectorTopKOptions, and AdaptiveVectorTopKPlanner.
  • Configured SessionConfig to insert AdaptiveVectorTopKOptions.
  • Added AdaptiveVectorTopKRule to the optimizer rules, ensuring it runs after ScanHintRule.
  • Added AdaptiveVectorTopKPlanner to the list of extension planners.
• src/query/src/vector_search/exec.rs
  • New file: Defines AdaptiveVectorTopKExec for adaptive vector top-k search execution.
  • Implements the adaptive loop, re-executing sub-plans with increasing k values.
  • Includes logic for sorting, limiting, and collecting detailed vector search metrics.
  • Provides helper functions for batch processing, metric collection, and adaptive loop control.
• src/query/src/vector_search/mod.rs
  • New file: Declares exec, options, plan, and planner modules for vector search.
• src/query/src/vector_search/options.rs
  • New file: Defines AdaptiveVectorTopKOptions for configuring adaptive vector top-k behavior (max rounds, max k, max rows per round).
  • Implements ConfigExtension and ExtensionOptions for DataFusion integration.
• src/query/src/vector_search/plan.rs
  • New file: Defines AdaptiveVectorTopKLogicalPlan as a DataFusion UserDefinedLogicalNodeCore.
• src/query/src/vector_search/planner.rs
  • New file: Implements AdaptiveVectorTopKPlanner to convert AdaptiveVectorTopKLogicalPlan into AdaptiveVectorTopKExec.
• src/store-api/src/storage.rs
  • Exported VectorSearchPredicate trait.
• src/store-api/src/storage/requests.rs
  • Added limit, offset, and tie_breakers fields to VectorSearchRequest.
  • Added VectorSearchPredicate trait for filtering vector search candidates.
  • Added search_with_predicate method to VectorIndexEngine trait with a default implementation.
  • Updated Display implementation for ScanRequest to include new vector search fields.
• tests-integration/tests/http.rs
  • Added vector_topk_max_rounds to the [query] section of the example configuration.
• tests/cases/standalone/common/function/vector/dynamic_overfetch.result
  • New file: Added expected results for dynamic overfetch SQL queries.
• tests/cases/standalone/common/function/vector/dynamic_overfetch.sql
  • New file: Added SQL test cases for dynamic overfetch functionality, including LIMIT and OFFSET with vector search and predicates.
• tests/cases/standalone/common/function/vector/vector_index_explain.result
  • Updated expected EXPLAIN output to reflect the new AdaptiveVectorTopKExec.
• tests/cases/standalone/common/tql-explain-analyze/explain.result
  • Updated EXPLAIN VERBOSE output to include logical_plan after AdaptiveVectorTopKRule.

Activity

• The pull request was created by killme2008.
• The PR introduces a new feature: adaptive vector search.
• The author has agreed to the GreptimeDB CLA.
• The PR references a related issue or pull request: #7668.
• The PR checklist indicates that unit and integration tests are included, and documentation updates may be required.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces an adaptive vector search mechanism, a significant feature for more accurate top-k vector searches when combined with filters. It works by adaptively fetching more candidates from the vector index. The implementation includes a new optimizer rule (AdaptiveVectorTopKRule), a custom logical plan node (AdaptiveVectorTopKLogicalPlan), and a physical execution node (AdaptiveVectorTopKExec) that contains the adaptive loop logic. This loop progressively increases the number of candidates (k) until enough results are found to satisfy the query's LIMIT and OFFSET. The PR also adds new configuration options and metrics for control and observability, extends vector search interfaces with predicate support, and includes new integration tests. The implementation appears solid and well-tested. I have one minor suggestion for code simplification.

[Copilot]

Pull request overview

Adds adaptive (dynamic overfetch) execution for vector Top-K queries so results remain correct with filters / OFFSET / tie-breakers, while improving observability via new vector-index k metrics.

Changes:

• Introduces AdaptiveVectorTopKRule + DataFusion extension logical/physical nodes to iteratively re-plan/retry vector Top-K with increasing k.
• Extends vector search plumbing with optional limit/offset/tie-breakers and adds predicate-capable vector index search APIs.
• Adds/updates golden tests and metrics reporting for vector-index requested/returned k (including distributed MergeScan aggregation).

Reviewed changes

Copilot reviewed 28 out of 28 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/cases/standalone/common/tql-explain-analyze/explain.result	Updates explain output to include the new adaptive optimizer rule stage.
tests/cases/standalone/common/function/vector/vector_index_explain.result	Updates vector explain plans to show AdaptiveVectorTopKExec.
tests/cases/standalone/common/function/vector/dynamic_overfetch.sql	Adds SQL test coverage for dynamic overfetch correctness (filters, offsets, partitions).
tests/cases/standalone/common/function/vector/dynamic_overfetch.result	Adds expected outputs for dynamic overfetch test cases.
tests-integration/tests/http.rs	Sets vector_topk_max_rounds in integration test config.
src/store-api/src/storage/requests.rs	Extends VectorSearchRequest; introduces VectorSearchPredicate and default predicate-aware engine search.
src/store-api/src/storage.rs	Re-exports VectorSearchPredicate.
src/query/src/vector_search/planner.rs	Adds extension planner to build AdaptiveVectorTopKExec from the logical node.
src/query/src/vector_search/plan.rs	Adds AdaptiveVectorTopKLogicalPlan (DataFusion extension logical node).
src/query/src/vector_search/options.rs	Adds DataFusion config extension for adaptive vector Top-K runtime bounds.
src/query/src/vector_search/mod.rs	Wires new vector_search module exports.
src/query/src/vector_search/exec.rs	Implements adaptive retry loop executor + metrics and unit tests.
src/query/src/query_engine/state.rs	Registers optimizer rule, extension planner, and session extension options.
src/query/src/options.rs	Adds query config knobs for adaptive vector Top-K.
src/query/src/optimizer/scan_hint/vector_search.rs	Captures tie-breakers + limit/offset into vector search hints.
src/query/src/optimizer/adaptive_vector_topk.rs	Adds optimizer rewrite from Sort+Limit to AdaptiveVectorTopK extension node.
src/query/src/optimizer.rs	Exposes adaptive vector top-k optimizer module.
src/query/src/lib.rs	Exports vector_search module.
src/query/src/dist_plan/merge_scan.rs	Fixes vector-index k metric aggregation (avoid over-counting) and adds tests.
src/mito2/src/sst/parquet/reader.rs	Passes predicate placeholder to vector index applier; records requested/returned k metrics.
src/mito2/src/sst/index/vector_index/applier.rs	Adds predicate parameter and uses predicate-aware search API.
src/mito2/src/read/scan_util.rs	Plumbs and exposes vector-index requested/returned k scan metrics via DataFusion gauges.
src/mito2/src/read/scan_region.rs	Removes fixed overfetch multiplier and uses request k directly.
src/mito2/src/read/pruner.rs	Clarifies cache-hit metric behavior and keeps error path explicit.
src/index/src/vector/engine/usearch.rs	Implements predicate-aware search using USearch filtered_search.
src/index/src/vector/apply.rs	Adds predicate-aware search at applier level + mapping adapter + unit test.
src/flow/src/adapter.rs	Updates flownode defaults to include new query options.
src/cmd/tests/load_config_test.rs	Updates config loading test expectations for new query options.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.