add: WIP Phase 3 first pass; Python API

Author: brooksomics

apis/python/src/tiledb/vector_search/type_erased_module.cc

@@ -436,13 +436,16 @@ void init_type_erased_module(py::module_& m) {
           [](IndexVamana& index,
              const FeatureVectorArray& vectors,
              size_t k,
-             uint32_t l_search) {
-            auto r = index.query(vectors, k, l_search);
+             uint32_t l_search,
+             std::optional<std::unordered_set<uint32_t>> query_filter =
+                 std::nullopt) {
+            auto r = index.query(vectors, k, l_search, query_filter);
             return make_python_pair(std::move(r));
           },
           py::arg("vectors"),
           py::arg("k"),
-          py::arg("l_search"))
+          py::arg("l_search"),
+          py::arg("query_filter") = std::nullopt)
       .def(
           "write_index",
           [](IndexVamana& index,

apis/python/src/tiledb/vector_search/vamana_index.py

@@ -7,11 +7,13 @@
 
   Singh, Aditi, et al. FreshDiskANN: A Fast and Accurate Graph-Based ANN Index for Streaming Similarity Search. arXiv:2105.09613, arXiv, 20 May 2021, http://arxiv.org/abs/2105.09613.
 
-  Gollapudi, Siddharth, et al. “Filtered-DiskANN: Graph Algorithms for Approximate Nearest Neighbor Search with Filters.” Proceedings of the ACM Web Conference 2023, ACM, 2023, pp. 3406-16, https://doi.org/10.1145/3543507.3583552.
+  Gollapudi, Siddharth, et al. "Filtered-DiskANN: Graph Algorithms for Approximate Nearest Neighbor Search with Filters." Proceedings of the ACM Web Conference 2023, ACM, 2023, pp. 3406-16, https://doi.org/10.1145/3543507.3583552.
 ```
 """
+import json
+import re
 import warnings
-from typing import Any, Mapping
+from typing import Any, Mapping, Optional, Set
 
 import numpy as np
 
@@ -25,6 +27,55 @@
 from tiledb.vector_search.utils import MAX_UINT64
 from tiledb.vector_search.utils import to_temporal_policy
 
+
+def _parse_where_clause(where: str, label_enumeration: dict) -> Set[int]:
+    """
+    Parse a simple where clause and return a set of label IDs.
+
+    Supports basic equality conditions like: "label_col == 'value'"
+
+    Parameters
+    ----------
+    where : str
+        The where clause string to parse
+    label_enumeration : dict
+        Mapping from label strings to enumeration IDs
+
+    Returns
+    -------
+    Set[int]
+        Set of label IDs matching the where clause
+
+    Raises
+    ------
+    ValueError
+        If the where clause is invalid or references non-existent labels
+    """
+    # Simple pattern for: column_name == 'value'
+    # We support single or double quotes
+    pattern = r"\s*\w+\s*==\s*['\"]([^'\"]+)['\"]\s*"
+    match = re.match(pattern, where.strip())
+
+    if not match:
+        raise ValueError(
+            f"Invalid where clause: '{where}'. "
+            "Expected format: \"label_col == 'value'\""
+        )
+
+    label_value = match.group(1)
+
+    # Check if the label exists in the enumeration
+    if label_value not in label_enumeration:
+        available_labels = ", ".join(sorted(label_enumeration.keys()))
+        raise ValueError(
+            f"Label '{label_value}' not found in index. "
+            f"Available labels: {available_labels}"
+        )
+
+    # Return the enumeration ID for this label
+    label_id = label_enumeration[label_value]
+    return {label_id}
+
 INDEX_TYPE = "VAMANA"
 
 L_BUILD_DEFAULT = 100
@@ -94,6 +145,7 @@ def query_internal(
         queries: np.ndarray,
         k: int = 10,
         l_search: Optional[int] = L_SEARCH_DEFAULT,
+        where: Optional[str] = None,
         **kwargs,
     ):
         """
@@ -108,6 +160,9 @@ def query_internal(
         l_search: int
             How deep to search. Larger parameters will result in slower latencies, but higher accuracies.
             Should be >= k, and if it's not, we will set it to k.
+        where: Optional[str]
+            Optional filter condition for filtered queries.
+            Example: "label_col == 'dataset_1'"
         """
         if self.size == 0:
             return np.full((queries.shape[0], k), MAX_FLOAT32), np.full(
@@ -125,7 +180,26 @@ def query_internal(
             queries = queries.copy(order="F")
         queries_feature_vector_array = vspy.FeatureVectorArray(queries)
 
-        distances, ids = self.index.query(queries_feature_vector_array, k, l_search)
+        # NEW: Handle filtered queries
+        query_filter = None
+        if where is not None:
+            # Get label enumeration from metadata
+            label_enum_str = self.group.meta.get("label_enumeration", None)
+            if label_enum_str is None:
+                raise ValueError(
+                    "Cannot use 'where' parameter: index does not have filter metadata. "
+                    "This index was not created with filter support."
+                )
+
+            # Parse JSON string to get label enumeration
+            label_enumeration = json.loads(label_enum_str)
+
+            # Parse where clause and get filter label IDs
+            query_filter = _parse_where_clause(where, label_enumeration)
+
+        distances, ids = self.index.query(
+            queries_feature_vector_array, k, l_search, query_filter
+        )
 
         return np.array(distances, copy=False), np.array(ids, copy=False)
 

src/include/api/vamana_index.h

@@ -225,11 +225,12 @@ class IndexVamana {
   [[nodiscard]] auto query(
       const QueryVectorArray& vectors,
       size_t top_k,
-      std::optional<uint32_t> l_search = std::nullopt) {
+      std::optional<uint32_t> l_search = std::nullopt,
+      std::optional<std::unordered_set<uint32_t>> query_filter = std::nullopt) {
     if (!index_) {
       throw std::runtime_error("Cannot query() because there is no index.");
     }
-    return index_->query(vectors, top_k, l_search);
+    return index_->query(vectors, top_k, l_search, query_filter);
   }
 
   void write_index(
@@ -348,7 +349,8 @@ class IndexVamana {
     query(
         const QueryVectorArray& vectors,
         size_t top_k,
-        std::optional<uint32_t> l_search) = 0;
+        std::optional<uint32_t> l_search,
+        std::optional<std::unordered_set<uint32_t>> query_filter) = 0;
 
     virtual void write_index(
         const tiledb::Context& ctx,
@@ -436,7 +438,8 @@ class IndexVamana {
     [[nodiscard]] std::tuple<FeatureVectorArray, FeatureVectorArray> query(
         const QueryVectorArray& vectors,
         size_t top_k,
-        std::optional<uint32_t> l_search) override {
+        std::optional<uint32_t> l_search,
+        std::optional<std::unordered_set<uint32_t>> query_filter) override {
       // @todo using index_type = size_t;
       auto dtype = vectors.feature_type();
 
@@ -448,7 +451,7 @@ class IndexVamana {
               (float*)vectors.data(),
               extents(vectors)[0],
               extents(vectors)[1]};  // @todo ??
-          auto [s, t] = impl_index_.query(qspan, top_k, l_search);
+          auto [s, t] = impl_index_.query(qspan, top_k, l_search, query_filter);
           auto x = FeatureVectorArray{std::move(s)};
           auto y = FeatureVectorArray{std::move(t)};
           return {std::move(x), std::move(y)};
@@ -458,7 +461,7 @@ class IndexVamana {
               (uint8_t*)vectors.data(),
               extents(vectors)[0],
               extents(vectors)[1]};  // @todo ??
-          auto [s, t] = impl_index_.query(qspan, top_k, l_search);
+          auto [s, t] = impl_index_.query(qspan, top_k, l_search, query_filter);
           auto x = FeatureVectorArray{std::move(s)};
           auto y = FeatureVectorArray{std::move(t)};
           return {std::move(x), std::move(y)};

src/include/index/vamana_index.h

@@ -811,13 +811,15 @@ class vamana_index {
    * @param query_set Container of query vectors
    * @param k How many nearest neighbors to return
    * @param l_search How deep to search
+   * @param query_filter Optional filter labels for filtered search
    * @return Tuple of top k scores and top k ids
    */
   template <query_vector_array Q>
   auto query(
       const Q& query_set,
       size_t k,
       std::optional<uint32_t> l_search = std::nullopt,
+      std::optional<std::unordered_set<filter_label_type>> query_filter = std::nullopt,
       Distance distance = Distance{}) {
     scoped_timer _("vamana_index@query");
 
@@ -833,18 +835,50 @@ class vamana_index {
 
     stdx::range_for_each(
         std::move(par), query_set, [&](auto&& query_vec, auto n, auto i) {
-          auto&& [tk_scores, tk, V] = greedy_search(
-              graph_,
-              feature_vectors_,
-              medoid_,
-              query_vec,
-              k,
-              L,
-              distance_function_,
-              true);
-          std::copy(
-              tk_scores.data(), tk_scores.data() + k, top_k_scores[i].data());
-          std::copy(tk.data(), tk.data() + k, top_k[i].data());
+          // NEW: Use filtered or unfiltered search based on query_filter
+          if (filter_enabled_ && query_filter.has_value()) {
+            // Determine start nodes for ALL labels in query filter (multi-start)
+            std::vector<id_type> start_nodes_for_query;
+            for (uint32_t label : *query_filter) {
+              if (start_nodes_.find(label) != start_nodes_.end()) {
+                start_nodes_for_query.push_back(start_nodes_.at(label));
+              }
+            }
+
+            if (start_nodes_for_query.empty()) {
+              throw std::runtime_error(
+                  "No start nodes found for query filter labels");
+            }
+
+            auto&& [tk_scores, tk, V] = filtered_greedy_search_multi_start(
+                graph_,
+                feature_vectors_,
+                filter_labels_,
+                start_nodes_for_query,
+                query_vec,
+                *query_filter,
+                k,
+                L,
+                distance_function_,
+                true);
+            std::copy(
+                tk_scores.data(), tk_scores.data() + k, top_k_scores[i].data());
+            std::copy(tk.data(), tk.data() + k, top_k[i].data());
+          } else {
+            // Unfiltered search
+            auto&& [tk_scores, tk, V] = greedy_search(
+                graph_,
+                feature_vectors_,
+                medoid_,
+                query_vec,
+                k,
+                L,
+                distance_function_,
+                true);
+            std::copy(
+                tk_scores.data(), tk_scores.data() + k, top_k_scores[i].data());
+            std::copy(tk.data(), tk.data() + k, top_k[i].data());
+          }
         });
 
     return std::make_tuple(std::move(top_k_scores), std::move(top_k));
@@ -857,26 +891,60 @@ class vamana_index {
    * @param query_vec The vector to query
    * @param k How many nearest neighbors to return
    * @param l_search How deep to search
+   * @param query_filter Optional filter labels for filtered search
    * @return Top k scores and top k ids
    */
   template <query_vector Q>
   auto query(
       const Q& query_vec,
       size_t k,
       std::optional<uint32_t> l_search = std::nullopt,
+      std::optional<std::unordered_set<filter_label_type>> query_filter = std::nullopt,
       Distance distance = Distance{}) {
     uint32_t L = l_search ? *l_search : l_build_;
-    auto&& [top_k_scores, top_k, V] = greedy_search(
-        graph_,
-        feature_vectors_,
-        medoid_,
-        query_vec,
-        k,
-        L,
-        distance_function_,
-        true);
 
-    return std::make_tuple(std::move(top_k_scores), std::move(top_k));
+    // NEW: Use filtered or unfiltered search based on query_filter
+    if (filter_enabled_ && query_filter.has_value()) {
+      // Determine start nodes for ALL labels in query filter (multi-start)
+      std::vector<id_type> start_nodes_for_query;
+      for (uint32_t label : *query_filter) {
+        if (start_nodes_.find(label) != start_nodes_.end()) {
+          start_nodes_for_query.push_back(start_nodes_.at(label));
+        }
+      }
+
+      if (start_nodes_for_query.empty()) {
+        throw std::runtime_error(
+            "No start nodes found for query filter labels");
+      }
+
+      auto&& [top_k_scores, top_k, V] = filtered_greedy_search_multi_start(
+          graph_,
+          feature_vectors_,
+          filter_labels_,
+          start_nodes_for_query,
+          query_vec,
+          *query_filter,
+          k,
+          L,
+          distance_function_,
+          true);
+
+      return std::make_tuple(std::move(top_k_scores), std::move(top_k));
+    } else {
+      // Unfiltered search
+      auto&& [top_k_scores, top_k, V] = greedy_search(
+          graph_,
+          feature_vectors_,
+          medoid_,
+          query_vec,
+          k,
+          L,
+          distance_function_,
+          true);
+
+      return std::make_tuple(std::move(top_k_scores), std::move(top_k));
+    }
   }
 
   constexpr uint64_t dimensions() const {