feat(search): support vector index controls in hybrid search (#190)

Author: Xuanwo

docs/sql.md

@@ -124,7 +124,10 @@ Positional arguments:
 
 Named parameters:
 - `k` (BIGINT, default `10`): Number of results to return. Must be > 0.
+- `nprobs` (BIGINT, optional): Number of IVF partitions to probe when using a vector index. Must be > 0. Only affects IVF-based vector indices.
+- `refine_factor` (BIGINT, optional): Refine factor for the vector branch. Must be > 0.
 - `prefilter` (BOOLEAN, default `false`): If `true`, filters are applied before top-k selection.
+- `use_index` (BOOLEAN, default `true`): If `true`, allow ANN index usage for the vector branch when available. If `false`, the vector branch runs exact KNN.
 - `alpha` (FLOAT, default `0.5`): Vector/text mixing weight. Larger values weigh vector similarity more heavily.
 - `oversample_factor` (INTEGER, default `4`): Oversample factor for candidate generation. If provided, must be > 0.
 

rust/ffi/search.rs

@@ -191,9 +191,12 @@ pub unsafe extern "C" fn lance_create_hybrid_stream_ir(
     text_column: *const c_char,
     text_query: *const c_char,
     k: u64,
+    nprobes: u64,
+    refine_factor: u64,
     filter_ir: *const u8,
     filter_ir_len: usize,
     prefilter: u8,
+    use_index: u8,
     alpha: f32,
     oversample_factor: u32,
 ) -> *mut c_void {
@@ -205,9 +208,12 @@ pub unsafe extern "C" fn lance_create_hybrid_stream_ir(
         text_column,
         text_query,
         k,
+        nprobes,
+        refine_factor,
         filter_ir,
         filter_ir_len,
         prefilter,
+        use_index,
         alpha,
         oversample_factor,
     ) {
@@ -231,9 +237,12 @@ fn create_hybrid_stream_ir_inner(
     text_column: *const c_char,
     text_query: *const c_char,
     k: u64,
+    nprobes: u64,
+    refine_factor: u64,
     filter_ir: *const u8,
     filter_ir_len: usize,
     prefilter: u8,
+    use_index: u8,
     alpha: f32,
     oversample_factor: u32,
 ) -> FfiResult<StreamHandle> {
@@ -265,8 +274,11 @@ fn create_hybrid_stream_ir_inner(
         text_column,
         text_query,
         k_usize,
+        nprobes,
+        refine_factor,
         filter,
         prefilter,
+        use_index,
         alpha,
         oversample_factor,
     )?;
@@ -282,8 +294,11 @@ fn create_hybrid_batch(
     text_column: &str,
     text_query: &str,
     k: usize,
+    nprobes: u64,
+    refine_factor: u64,
     filter: Option<datafusion_expr::Expr>,
     prefilter: u8,
+    use_index: u8,
     alpha: f32,
     oversample_factor: u32,
 ) -> FfiResult<RecordBatch> {
@@ -303,7 +318,17 @@ fn create_hybrid_batch(
                 format!("hybrid vector nearest: {err}"),
             )
         })?;
-    vector_scan.use_index(false);
+    if nprobes != 0 {
+        let nprobes_usize = nonzero_u64_to_usize(nprobes, "nprobes")?;
+        vector_scan.nprobes(nprobes_usize);
+    }
+    if refine_factor != 0 {
+        let refine_factor_u32: u32 = refine_factor.try_into().map_err(|_| {
+            FfiError::new(ErrorCode::InvalidArgument, "refine_factor must fit in u32")
+        })?;
+        vector_scan.refine(refine_factor_u32);
+    }
+    vector_scan.use_index(use_index != 0);
     vector_scan.with_row_id();
     vector_scan.disable_scoring_autoprojection();
     vector_scan

src/include/lance_ffi.hpp

@@ -273,13 +273,12 @@ void *lance_create_fts_stream_ir(void *dataset, const char *text_column,
                                  uint8_t prefilter);
 
 void *lance_get_hybrid_schema(void *dataset);
-void *lance_create_hybrid_stream_ir(void *dataset, const char *vector_column,
-                                    const float *query_values, size_t query_len,
-                                    const char *text_column,
-                                    const char *text_query, uint64_t k,
-                                    const uint8_t *filter_ir,
-                                    size_t filter_ir_len, uint8_t prefilter,
-                                    float alpha, uint32_t oversample_factor);
+void *lance_create_hybrid_stream_ir(
+    void *dataset, const char *vector_column, const float *query_values,
+    size_t query_len, const char *text_column, const char *text_query,
+    uint64_t k, uint64_t nprobes, uint64_t refine_factor,
+    const uint8_t *filter_ir, size_t filter_ir_len, uint8_t prefilter,
+    uint8_t use_index, float alpha, uint32_t oversample_factor);
 
 // Index DDL / metadata
 int32_t lance_dataset_create_index(void *dataset, const char *index_name,

src/lance_search.cpp

@@ -814,6 +814,9 @@ struct LanceSearchBindData : public TableFunctionData {
   string vector_column;
   vector<float> vector_query;
   string text_query;
+  uint64_t nprobes = 0;
+  uint64_t refine_factor = 0;
+  bool use_index = true;
   float alpha = 0.5F;
   uint32_t oversample_factor = 4;
 
@@ -882,8 +885,9 @@ static bool LanceSearchLoadNextBatch(LanceSearchLocalState &local_state,
           bind_data.dataset, bind_data.vector_column.c_str(),
           bind_data.vector_query.data(), bind_data.vector_query.size(),
           bind_data.text_column.c_str(), bind_data.text_query.c_str(),
-          bind_data.k, ir, NumericCast<size_t>(ir_len),
-          bind_data.prefilter ? 1 : 0, bind_data.alpha,
+          bind_data.k, bind_data.nprobes, bind_data.refine_factor, ir,
+          NumericCast<size_t>(ir_len), bind_data.prefilter ? 1 : 0,
+          bind_data.use_index ? 1 : 0, bind_data.alpha,
           bind_data.oversample_factor);
     };
 
@@ -1073,13 +1077,51 @@ LanceHybridBind(ClientContext &context, TableFunctionBindInput &input,
   }
   result->k = NumericCast<uint64_t>(k_val);
 
+  bool has_nprobes = false;
+  int64_t nprobes_val = 0;
+  auto nprobes_named = input.named_parameters.find("nprobs");
+  if (nprobes_named != input.named_parameters.end() &&
+      !nprobes_named->second.IsNull()) {
+    has_nprobes = true;
+    nprobes_val = nprobes_named->second.DefaultCastAs(LogicalType::BIGINT)
+                      .GetValue<int64_t>();
+  }
+  if (has_nprobes && nprobes_val <= 0) {
+    throw InvalidInputException("lance_hybrid_search requires nprobs > 0");
+  }
+  result->nprobes = has_nprobes ? NumericCast<uint64_t>(nprobes_val) : 0;
+
+  bool has_refine_factor = false;
+  int64_t refine_factor_val = 0;
+  auto refine_factor_named = input.named_parameters.find("refine_factor");
+  if (refine_factor_named != input.named_parameters.end() &&
+      !refine_factor_named->second.IsNull()) {
+    has_refine_factor = true;
+    refine_factor_val =
+        refine_factor_named->second.DefaultCastAs(LogicalType::BIGINT)
+            .GetValue<int64_t>();
+  }
+  if (has_refine_factor && refine_factor_val <= 0) {
+    throw InvalidInputException(
+        "lance_hybrid_search requires refine_factor > 0");
+  }
+  result->refine_factor =
+      has_refine_factor ? NumericCast<uint64_t>(refine_factor_val) : 0;
+
   auto prefilter_named = input.named_parameters.find("prefilter");
   if (prefilter_named != input.named_parameters.end() &&
       !prefilter_named->second.IsNull()) {
     result->prefilter =
         prefilter_named->second.DefaultCastAs(LogicalType::BOOLEAN)
             .GetValue<bool>();
   }
+  auto use_index_named = input.named_parameters.find("use_index");
+  if (use_index_named != input.named_parameters.end() &&
+      !use_index_named->second.IsNull()) {
+    result->use_index =
+        use_index_named->second.DefaultCastAs(LogicalType::BOOLEAN)
+            .GetValue<bool>();
+  }
 
   auto alpha_named = input.named_parameters.find("alpha");
   if (alpha_named != input.named_parameters.end() &&
@@ -1266,6 +1308,9 @@ LanceSearchBindToString(const LanceSearchBindData &bind_data) {
     result["Lance Text Column"] = bind_data.text_column;
     result["Lance Vector Query Dim"] = to_string(bind_data.vector_query.size());
     result["Lance Text Query"] = bind_data.text_query;
+    result["Lance Nprobes"] = to_string(bind_data.nprobes);
+    result["Lance Refine Factor"] = to_string(bind_data.refine_factor);
+    result["Lance Use Index"] = bind_data.use_index ? "true" : "false";
     result["Lance Alpha"] = to_string(bind_data.alpha);
     result["Lance Oversample Factor"] = to_string(bind_data.oversample_factor);
   }
@@ -1320,7 +1365,10 @@ static void RegisterLanceFtsSearch(ExtensionLoader &loader) {
 static void RegisterLanceHybridSearch(ExtensionLoader &loader) {
   auto configure = [](TableFunction &fun) {
     fun.named_parameters["k"] = LogicalType::BIGINT;
+    fun.named_parameters["nprobs"] = LogicalType::BIGINT;
+    fun.named_parameters["refine_factor"] = LogicalType::BIGINT;
     fun.named_parameters["prefilter"] = LogicalType::BOOLEAN;
+    fun.named_parameters["use_index"] = LogicalType::BOOLEAN;
     fun.named_parameters["alpha"] = LogicalType::FLOAT;
     fun.named_parameters["oversample_factor"] = LogicalType::INTEGER;
     fun.projection_pushdown = true;

test/sql/search_functions.test

@@ -34,6 +34,32 @@ SELECT * FROM lance_vector_search('test/data/test_data.lance', 'vec', [1.0]::FLO
 ----
 Invalid Input Error: lance_vector_search requires refine_factor > 0
 
+# Hybrid search rejects non-positive nprobs
+statement error
+SELECT * FROM lance_hybrid_search(
+  'test/data/search_test_data.lance',
+  'vec',
+  [0.0, 0.0, 0.0, 0.0]::FLOAT[4],
+  'text',
+  'puppy',
+  nprobs = 0
+)
+----
+Invalid Input Error: lance_hybrid_search requires nprobs > 0
+
+# Hybrid search rejects non-positive refine_factor
+statement error
+SELECT * FROM lance_hybrid_search(
+  'test/data/search_test_data.lance',
+  'vec',
+  [0.0, 0.0, 0.0, 0.0]::FLOAT[4],
+  'text',
+  'puppy',
+  refine_factor = 0
+)
+----
+Invalid Input Error: lance_hybrid_search requires refine_factor > 0
+
 # Sanity: dataset is readable
 query I
 SELECT count(*) FROM 'test/data/search_test_data.lance'
@@ -308,6 +334,100 @@ ORDER BY _hybrid_score DESC
 3
 2
 
+# Hybrid search accepts vector search controls in flat mode
+query I
+SELECT id
+FROM lance_hybrid_search(
+  'test/data/search_test_data.lance',
+  'vec',
+  [0.0, 0.0, 0.0, 0.0]::FLOAT[4],
+  'text',
+  'puppy',
+  k = 3,
+  nprobs = 2,
+  refine_factor = 2,
+  use_index = false,
+  alpha = 0.5,
+  oversample_factor = 4
+)
+ORDER BY _hybrid_score DESC
+----
+1
+3
+2
+
+statement ok
+COPY (SELECT * FROM 'test/data/search_test_data.lance')
+TO 'test/.tmp/search_hybrid_indexed.lance' (FORMAT lance, mode 'overwrite');
+
+statement ok
+CREATE INDEX vec_idx ON 'test/.tmp/search_hybrid_indexed.lance' (vec)
+USING IVF_FLAT WITH (num_partitions=1, metric_type='l2');
+
+statement ok
+CREATE INDEX text_idx ON 'test/.tmp/search_hybrid_indexed.lance' (text)
+USING INVERTED;
+
+# Hybrid search exposes vector search controls in EXPLAIN when using an index
+query II
+EXPLAIN (FORMAT JSON)
+SELECT id
+FROM lance_hybrid_search(
+  'test/.tmp/search_hybrid_indexed.lance',
+  'vec',
+  [0.0, 0.0, 0.0, 0.0]::FLOAT[4],
+  'text',
+  'puppy',
+  k = 3,
+  nprobs = 1,
+  refine_factor = 2,
+  use_index = true,
+  alpha = 0.5,
+  oversample_factor = 4
+);
+----
+physical_plan	<REGEX>:[\s\S]*"Lance Search Mode": "hybrid"[\s\S]*"Lance Nprobes": "1"[\s\S]*"Lance Refine Factor": "2"[\s\S]*"Lance Use Index": "true"[\s\S]*
+
+query II
+EXPLAIN (FORMAT JSON)
+SELECT id
+FROM lance_hybrid_search(
+  'test/.tmp/search_hybrid_indexed.lance',
+  'vec',
+  [0.0, 0.0, 0.0, 0.0]::FLOAT[4],
+  'text',
+  'puppy',
+  k = 3,
+  nprobs = 1,
+  refine_factor = 2,
+  use_index = false,
+  alpha = 0.5,
+  oversample_factor = 4
+);
+----
+physical_plan	<REGEX>:[\s\S]*"Lance Search Mode": "hybrid"[\s\S]*"Lance Use Index": "false"[\s\S]*
+
+query I
+SELECT id
+FROM lance_hybrid_search(
+  'test/.tmp/search_hybrid_indexed.lance',
+  'vec',
+  [0.0, 0.0, 0.0, 0.0]::FLOAT[4],
+  'text',
+  'puppy',
+  k = 3,
+  nprobs = 1,
+  refine_factor = 2,
+  use_index = true,
+  alpha = 0.5,
+  oversample_factor = 4
+)
+ORDER BY _hybrid_score DESC
+----
+1
+3
+2
+
 # ===================================================================
 # Table name resolution: search functions support catalog table names
 # (Users can ATTACH a Lance directory and use table names instead of