fix: preserve universal-label fallback in unified indexes

docs/unified_index_format.md

@@ -193,7 +193,7 @@ Replaces today's `_labels_map.txt` + `_labels_to_medoids.txt`. One row per disti
 
 `label_integer` is always written as a 4-byte little-endian unsigned integer, independent of the build-time `LabelT` template parameter (`uint16_t` values are zero-extended). This makes the on-disk dictionary self-describing and uniform across writer instantiations. Row count is implicit: read rows until `label_dictionary_len` bytes are consumed.
 
-If `header.universal_label != 0`, the dictionary MAY contain a row whose `label_integer` matches it; otherwise the universal label has no explicit dictionary entry.
+If `header.universal_label != 0`, the dictionary MUST contain a row whose `label_integer` matches it. That row supplies the universal label's filtered-search medoid.
 
 #### 2.4.2 Per-point labels
 
@@ -268,6 +268,16 @@ This change is encapsulated in a single helper (`node_read_window(N)`) so the bu
 
 ---
 
+### 3.3 Filtered search compatibility
+
+Unified-index search preserves the legacy invalid-label behavior:
+
+- Unknown label strings are ignored when at least one requested label resolves.
+- If no requested label resolves and no universal label is configured, search completes successfully with an empty result and reports `QueryStats::label_valid == false`.
+- If no requested label resolves and a universal label is configured with its required dictionary medoid, search reports `label_valid == false` but performs an ANN search using that universal label and medoid.
+- A filtered index missing the resolved label medoid, including the universal medoid, is treated as invalid data and returns an empty successful result.
+- Supplying filter labels to an unfiltered index remains an error.
+
 ## 4. Build Path (informative)
 
 `build_unified_index` reuses the existing pipeline (preprocess → optional PQ training → `build_merged_vamana_index`) up to the point where the legacy code would write separate files or call `create_disk_layout`. From there:

include/percentile_stats.h

@@ -33,6 +33,7 @@ struct QueryStats
     unsigned n_cmps = 0;       // # cmps
     unsigned n_cache_hits = 0; // # cache_hits
     unsigned n_hops = 0;       // # search hops
+    bool label_valid = true;   // whether at least one requested filter label was valid
 };
 
 struct TableStats

include/unified_index.h

@@ -40,16 +40,18 @@ struct UnifiedLoadContext
 struct UnifiedSearchContext
 {
     // ---- Inputs ----
-    const void *query = nullptr;          // typed by caller as const T*
+    const void *query = nullptr; // typed by caller as const T*
     size_t K = 10;
     uint32_t L = 100;
-    // Filter labels as user-facing strings. Required non-empty if the loaded
-    // index has labels; required empty otherwise. The index converts strings
-    // to internal label ints per its encoding.
+    // Filter labels as user-facing strings. Required empty for an unfiltered
+    // index. A filtered index ignores unknown labels; if none resolve, it
+    // returns an empty successful result unless a universal label is available,
+    // in which case it performs a universal-label search. QueryStats::label_valid
+    // reports whether at least one requested label resolved.
     std::vector<std::string> filter_labels;
-    std::optional<uint32_t> beam_width;                            // SSD-only
-    std::optional<uint32_t> io_limit;                              // SSD-only
-    std::function<float(const std::uint8_t *, size_t)> rerank_fn;  // SSD-only
+    std::optional<uint32_t> beam_width;                           // SSD-only
+    std::optional<uint32_t> io_limit;                             // SSD-only
+    std::function<float(const std::uint8_t *, size_t)> rerank_fn; // SSD-only
 
     // ---- Outputs (caller-allocated, length >= K) ----
     uint64_t *indices = nullptr;
@@ -90,7 +92,7 @@ DISKANN_DLLEXPORT std::unique_ptr<unified_index> make_unified_index_memory(const
 
 // Factory: open a unified file in disk-resident (SSD) mode. The supplied
 // AlignedFileReader is handed to the constructed unified_index_ssd<T>.
-DISKANN_DLLEXPORT std::unique_ptr<unified_index> make_unified_index_ssd(
-    std::shared_ptr<AlignedFileReader> reader, const UnifiedLoadContext &ctx);
+DISKANN_DLLEXPORT std::unique_ptr<unified_index> make_unified_index_ssd(std::shared_ptr<AlignedFileReader> reader,
+                                                                        const UnifiedLoadContext &ctx);
 
 } // namespace diskann

include/unified_label_data.h

@@ -20,6 +20,17 @@ namespace diskann
 
 class UnifiedIndexReader;
 
+struct filter_resolution_result
+{
+    // True when at least one requested label string resolved through the
+    // dictionary. Callers use this to report invalid-label telemetry.
+    bool label_valid = false;
+
+    // True when search can proceed through a resolved label or the universal
+    // label. False means the caller should return a successful empty result.
+    bool should_search = false;
+};
+
 // ---------------------------------------------------------------------------
 // Abstract base for the label-data trio.
 // Owns the shared, encoding-independent state (label dictionary, universal
@@ -68,16 +79,14 @@ class unified_label_data_base
     bool get_converted_label(const std::string &s, uint32_t &out) const;
 
     // Resolve filter label strings to their internal label ints AND per-label
-    // medoids in a single dictionary probe per string. out_label_ints[i] and
-    // out_medoids[i] both correspond to filter_label_strings[i]; both vectors
-    // are caller-owned, cleared, then filled in lockstep. Throws ANNException
-    // on an unknown label string. The unified format stores exactly one medoid
-    // per label, packed in the same dictionary row as the label int, so the
-    // search path gets the proxy input (label int) and the init-id seed
-    // (medoid) from one map lookup instead of two.
-    void resolve_filters(const std::vector<std::string> &filter_label_strings,
-                                           std::vector<uint32_t> &out_label_ints,
-                                           std::vector<uint32_t> &out_medoids) const;
+    // medoids in a single dictionary probe per string. Unknown labels are
+    // ignored, matching the legacy filtered-search path. If no requested label
+    // resolves and a universal label with a dictionary medoid exists, both are
+    // inserted into the outputs. Otherwise the caller returns a successful
+    // empty result.
+    filter_resolution_result resolve_filters(const std::vector<std::string> &filter_label_strings,
+                                             std::vector<uint32_t> &out_label_ints,
+                                             std::vector<uint32_t> &out_medoids) const;
 
     // Append every per-label entry-point medoid (the unified format stores
     // exactly one per label) to `out`. Used to seed SSD cache priming so that
@@ -92,20 +101,21 @@ class unified_label_data_base
     // and is shared with init-id seeding). The returned proxy borrows internal
     // storage of `this` -- lifetime must not exceed `this`. No external scratch
     // is needed; the concrete proxy owns any per-query scratch it requires.
-    virtual std::unique_ptr<filter_match_proxy> make_match_proxy(
-        const std::vector<uint32_t> &filter_label_ints) = 0;
+    virtual std::unique_ptr<filter_match_proxy> make_match_proxy(const std::vector<uint32_t> &filter_label_ints) = 0;
 
   protected:
     // Derived classes load their encoding-specific region(s) after the base
     // has parsed the shared dictionary.
     virtual void load_encoding(UnifiedIndexReader &r, const UnifiedIndexHeader &h, uint64_t npts) = 0;
 
-    // Helper: convert strings -> uint32 label ints via dictionary; throws on unknown.
+    // Helper: parse the string -> {label int, medoid} dictionary.
     void parse_dictionary(const std::vector<uint8_t> &dict_bytes);
 
     bool _has_labels = false;
     bool _use_universal_label = false;
     uint32_t _universal_label = 0;
+    bool _has_universal_medoid = false;
+    uint32_t _universal_medoid = 0;
 
     // Dictionary row: label string -> {internal label int, per-label medoid}.
     // Both fields come from the same on-disk dictionary entry (see
@@ -135,8 +145,7 @@ class unified_label_data_bitmask final : public unified_label_data_base
         return _bitmask_buf._buf.size() * sizeof(std::uint64_t);
     }
 
-    std::unique_ptr<filter_match_proxy> make_match_proxy(
-        const std::vector<uint32_t> &filter_label_ints) override;
+    std::unique_ptr<filter_match_proxy> make_match_proxy(const std::vector<uint32_t> &filter_label_ints) override;
 
     simple_bitmask_buf &bitmask_buf()
     {
@@ -169,8 +178,7 @@ class unified_label_data_integer final : public unified_label_data_base
         return _label_vector.get_memory_usage();
     }
 
-    std::unique_ptr<filter_match_proxy> make_match_proxy(
-        const std::vector<uint32_t> &filter_label_ints) override;
+    std::unique_ptr<filter_match_proxy> make_match_proxy(const std::vector<uint32_t> &filter_label_ints) override;
 
     integer_label_vector &label_vector()
     {
@@ -191,8 +199,7 @@ class unified_label_data_integer final : public unified_label_data_base
 // Factory: peeks at `h.label_encoding`, constructs the correct derived class,
 // runs load(), and returns it. Returns nullptr when the header carries no
 // labels (HAS_LABELS flag unset or encoding == None).
-std::unique_ptr<unified_label_data_base> make_unified_label_data(UnifiedIndexReader &r,
-                                                                                    const UnifiedIndexHeader &h,
-                                                                                    uint64_t npts);
+std::unique_ptr<unified_label_data_base> make_unified_label_data(UnifiedIndexReader &r, const UnifiedIndexHeader &h,
+                                                                 uint64_t npts);
 
 } // namespace diskann

src/unified_index_base.cpp

@@ -9,8 +9,7 @@
 namespace diskann
 {
 
-template <typename T>
-unified_index_base<T>::unified_index_base(diskann::Metric metric) : _metric(metric)
+template <typename T> unified_index_base<T>::unified_index_base(diskann::Metric metric) : _metric(metric)
 {
 }
 
@@ -26,8 +25,7 @@ template <typename T> void unified_index_base<T>::validate_header(const UnifiedI
         throw ANNException("unified_index_base: data_type mismatch with T", -1, __FUNCSIG__, __FILE__, __LINE__);
 }
 
-template <typename T>
-void unified_index_base<T>::validate_search_context(const UnifiedSearchContext &ctx) const
+template <typename T> void unified_index_base<T>::validate_search_context(const UnifiedSearchContext &ctx) const
 {
     if (ctx.query == nullptr)
         throw ANNException("UnifiedSearchContext: query == nullptr", -1, __FUNCSIG__, __FILE__, __LINE__);
@@ -41,9 +39,6 @@ void unified_index_base<T>::validate_search_context(const UnifiedSearchContext &
 
     const bool filtered = has_labels();
     const bool has_filters = !ctx.filter_labels.empty();
-    if (filtered && !has_filters)
-        throw ANNException("UnifiedSearchContext: filter_labels must be non-empty for a filtered index", -1,
-                           __FUNCSIG__, __FILE__, __LINE__);
     if (!filtered && has_filters)
         throw ANNException("UnifiedSearchContext: filter_labels must be empty for a non-filtered index", -1,
                            __FUNCSIG__, __FILE__, __LINE__);

src/unified_index_memory.cpp

@@ -41,8 +41,7 @@ template <typename T> unified_index_memory<T>::~unified_index_memory()
     }
 }
 
-template <typename T>
-void unified_index_memory<T>::load_storage(UnifiedIndexReader &r, const UnifiedLoadContext &ctx)
+template <typename T> void unified_index_memory<T>::load_storage(UnifiedIndexReader &r, const UnifiedLoadContext &ctx)
 {
     // Build the resident node store.
     auto store = std::make_unique<unified_node_store_memory<T>>();
@@ -91,17 +90,18 @@ template <typename T> void unified_index_memory<T>::init_scratch_pool(uint32_t n
     std::vector<uint32_t> empty_sellers;
     for (uint32_t i = 0; i < num_threads; ++i)
     {
-        auto *s = new InMemQueryScratch<T>(search_l, search_l, R, maxc, dim, aligned_dim, alignment_factor,
-                                           empty_sellers,
-                                           /*init_pq_scratch=*/false, bitmask_size);
+        auto *s =
+            new InMemQueryScratch<T>(search_l, search_l, R, maxc, dim, aligned_dim, alignment_factor, empty_sellers,
+                                     /*init_pq_scratch=*/false, bitmask_size);
         _query_scratch.push(s);
     }
 }
 
 template <typename T>
-std::pair<uint32_t, uint32_t> unified_index_memory<T>::iterate_to_fixed_point(
-    InMemQueryScratch<T> *scratch, uint32_t L, const T *query, const std::vector<uint32_t> &init_ids,
-    filter_match_proxy *match_proxy)
+std::pair<uint32_t, uint32_t> unified_index_memory<T>::iterate_to_fixed_point(InMemQueryScratch<T> *scratch, uint32_t L,
+                                                                              const T *query,
+                                                                              const std::vector<uint32_t> &init_ids,
+                                                                              filter_match_proxy *match_proxy)
 {
     auto *store = static_cast<unified_node_store_memory<T> *>(this->_store.get());
     const uint64_t aligned_dim = this->_header.aligned_dim;
@@ -124,9 +124,8 @@ std::pair<uint32_t, uint32_t> unified_index_memory<T>::iterate_to_fixed_point(
     {
         if (inserted_into_pool_bs.size() < total_num_points)
         {
-            auto resize_size = 2 * total_num_points > MAX_POINTS_FOR_USING_BITSET
-                                   ? MAX_POINTS_FOR_USING_BITSET
-                                   : 2 * total_num_points;
+            auto resize_size =
+                2 * total_num_points > MAX_POINTS_FOR_USING_BITSET ? MAX_POINTS_FOR_USING_BITSET : 2 * total_num_points;
             inserted_into_pool_bs.resize(resize_size);
         }
     }
@@ -206,6 +205,38 @@ std::pair<uint32_t, uint32_t> unified_index_memory<T>::iterate_to_fixed_point(
 
 template <typename T> void unified_index_memory<T>::search_impl(UnifiedSearchContext &ctx)
 {
+    // Resolve filters before borrowing scratch so an all-invalid request
+    // without a universal label can return a successful empty result cheaply.
+    std::unique_ptr<filter_match_proxy> proxy;
+    thread_local std::vector<uint32_t> init_ids;
+    init_ids.clear();
+    if (ctx.stats != nullptr)
+    {
+        ctx.stats->label_valid = true;
+    }
+    if (this->_labels && this->_labels->has_labels())
+    {
+        thread_local std::vector<uint32_t> filter_label_ints;
+        const auto resolution = this->_labels->resolve_filters(ctx.filter_labels, filter_label_ints, init_ids);
+        if (ctx.stats != nullptr)
+        {
+            ctx.stats->label_valid = resolution.label_valid;
+        }
+        if (!resolution.should_search)
+        {
+            std::fill_n(ctx.indices, ctx.K, std::numeric_limits<uint64_t>::max());
+            std::fill_n(ctx.distances, ctx.K, std::numeric_limits<float>::max());
+            if (ctx.stats != nullptr)
+            {
+                ctx.stats->n_hops = 0;
+                ctx.stats->n_cmps = 0;
+                ctx.stats->n_ios = 0;
+            }
+            return;
+        }
+        proxy = this->_labels->make_match_proxy(filter_label_ints);
+    }
+
     // Borrow per-thread scratch. ScratchStoreManager blocks until a scratch
     // object is available, so when more queries arrive concurrently than the
     // pool holds, callers wait for a free slot rather than failing. The pool is
@@ -231,26 +262,6 @@ template <typename T> void unified_index_memory<T>::search_impl(UnifiedSearchCon
         std::memcpy(aligned_query, ctx.query, dim * sizeof(T));
     }
 
-    // Build the label match proxy if the index is filtered, resolving the
-    // filter label strings once: resolve_filters yields both the internal
-    // label ints (for the proxy) and the per-label medoid seed ids (init_ids
-    // below) from a single dictionary probe per label.
-    //
-    // init_ids / filter_label_ints are thread_local: search_impl runs once per
-    // query on a pooled thread, so reusing these buffers across queries avoids a
-    // per-call heap allocation (mirrors unified_index_ssd::cached_beam_search).
-    // init_ids must be cleared up front because the unfiltered branch below
-    // relies on it being empty when no filter is applied.
-    std::unique_ptr<filter_match_proxy> proxy;
-    thread_local std::vector<uint32_t> init_ids;
-    init_ids.clear();
-    if (this->_labels && this->_labels->has_labels())
-    {
-        thread_local std::vector<uint32_t> filter_label_ints;
-        this->_labels->resolve_filters(ctx.filter_labels, filter_label_ints, init_ids);
-        proxy = this->_labels->make_match_proxy(filter_label_ints);
-    }
-
     // Seed init_ids. Aligned with unified_index_ssd::cached_beam_search:
     // - Unfiltered: pick the single closest medoid from _medoids by
     //   full-vector L2 (memory has all coords resident, so we don't need the
@@ -260,9 +271,9 @@ template <typename T> void unified_index_memory<T>::search_impl(UnifiedSearchCon
     //   seeding dramatically improves recall on filtered search because the
     //   global start node may not lie within any filter-label cluster.
     auto *store = static_cast<unified_node_store_memory<T> *>(this->_store.get());
-    if (init_ids.empty())
+    if (proxy == nullptr)
     {
-        // Unfiltered path -- pick closest of the global medoid set.
+        // Unfiltered path -- pick the closest global medoid.
         uint32_t best_id = _medoids.empty() ? _start : _medoids[0];
         float best_dist = std::numeric_limits<float>::max();
         for (uint32_t mid : _medoids)
@@ -277,6 +288,20 @@ template <typename T> void unified_index_memory<T>::search_impl(UnifiedSearchCon
         }
         init_ids.push_back(best_id);
     }
+    else if (init_ids.empty())
+    {
+        // A valid filtered index supplies a medoid for every label, including
+        // the universal label. Missing init IDs indicate invalid index data.
+        std::fill_n(ctx.indices, ctx.K, std::numeric_limits<uint64_t>::max());
+        std::fill_n(ctx.distances, ctx.K, std::numeric_limits<float>::max());
+        if (ctx.stats != nullptr)
+        {
+            ctx.stats->n_hops = 0;
+            ctx.stats->n_cmps = 0;
+            ctx.stats->n_ios = 0;
+        }
+        return;
+    }
 
     auto [hops, cmps] = iterate_to_fixed_point(scratch, ctx.L, aligned_query, init_ids, proxy.get());