Roaring Bitmap Filter in CAGRA

cpp/CMakeLists.txt

@@ -189,6 +189,7 @@ if(NOT BUILD_CPU_ONLY)
   include(cmake/thirdparty/get_raft.cmake)
   include(cmake/thirdparty/get_kvikio.cmake)
   include(cmake/thirdparty/get_cutlass.cmake)
+  rapids_cpm_package_override("${CMAKE_CURRENT_SOURCE_DIR}/cmake/patches/cuco_override.json")
   include(${rapids-cmake-dir}/cpm/cuco.cmake)
   rapids_cpm_cuco()
   include(cmake/thirdparty/get_rtcx.cmake)
@@ -1361,6 +1362,8 @@ if(NOT BUILD_CPU_ONLY)
     src/cluster/spectral.cu
     src/core/bitset.cu
     src/core/bloom_filter.cu
+    src/core/roaring_allowlist.cu
+    src/neighbors/roaring_filter.cu
     src/core/omp_wrapper.cpp
     src/util/file_io.cpp
     src/util/host_memory.cpp

cpp/cmake/patches/cuco_override.json

@@ -0,0 +1,9 @@
+{
+  "packages" : {
+    "cuco" : {
+      "version": "0.0.1",
+      "git_url": "https://github.com/NVIDIA/cuCollections.git",
+      "git_tag": "9d7c9307395c3b8795d93ad65d0751c98471dde6"
+    }
+  }
+}

cpp/include/cuvs/core/roaring_allowlist.hpp

@@ -0,0 +1,135 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cuvs/core/export.hpp>
+
+#include <raft/core/device_mdspan.hpp>
+#include <raft/core/host_mdspan.hpp>
+#include <raft/core/resources.hpp>
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+namespace CUVS_EXPORT cuvs {
+namespace core {
+
+/**
+ * @brief Non-owning device view of one immutable Roaring allowlist.
+ *
+ * The view contains an opaque pointer to an already initialized device-side cuCollections
+ * reference plus immutable shape and cardinality metadata. Creating or copying it is O(1) and
+ * performs no allocation, parsing, kernel launch, or synchronization. The owning
+ * @ref roaring_allowlist must outlive the view and every operation that uses it.
+ */
+class CUVS_EXPORT roaring_allowlist_view {
+ public:
+  roaring_allowlist_view() = default;
+
+  [[nodiscard]] std::size_t dataset_rows() const noexcept { return dataset_rows_; }
+  [[nodiscard]] std::size_t cardinality() const noexcept { return cardinality_; }
+  [[nodiscard]] bool empty() const noexcept { return cardinality_ == 0; }
+  [[nodiscard]] bool valid() const noexcept { return valid_; }
+
+  /** @brief Opaque device pointer to the initialized cuCollections reference, or null if empty. */
+  [[nodiscard]] void const* device_reference() const noexcept { return device_reference_; }
+
+ private:
+  friend class roaring_allowlist;
+
+  roaring_allowlist_view(void const* device_reference,
+                         std::size_t dataset_rows,
+                         std::size_t cardinality) noexcept
+    : device_reference_(device_reference),
+      dataset_rows_(dataset_rows),
+      cardinality_(cardinality),
+      valid_(true)
+  {
+  }
+
+  void const* device_reference_{};
+  std::size_t dataset_rows_{};
+  std::size_t cardinality_{};
+  bool valid_{};
+};
+
+/**
+ * @brief Owning immutable exact Roaring allowlist over CAGRA dataset-row IDs.
+ *
+ * Build an allowlist through @ref from_ids, then pass its zero-copy @ref view to a
+ * cuvs::neighbors::filtering::roaring_filter. A filter maps one such view to each query; owners
+ * remain independent and can therefore be reused across filters and queries.
+ *
+ * Construction sorts IDs on the GPU unless @p pre_sorted is true. Setting @p pre_sorted promises
+ * that IDs are already in strictly increasing order; this promise is not verified. IDs must be
+ * unique and smaller than @p dataset_rows. The encoded bytes and the initialized
+ * cuco::experimental::roaring_bitmap_ref<uint32_t> are retained on the device. Creating a view
+ * never copies or reparses them, and CAGRA search performs no Roaring initialization.
+ *
+ * ID-based construction emits the standard portable 32-bit Roaring array and bitmap container
+ * forms. Each ID is partitioned by its high 16 bits; the low 16 bits are stored as an array for at
+ * most 4,096 values in a partition and as an 8 KiB bitmap otherwise.
+ *
+ * @see https://github.com/RoaringBitmap/RoaringFormatSpec
+ * @see https://github.com/NVIDIA/cuCollections/pull/839
+ */
+class CUVS_EXPORT roaring_allowlist {
+ private:
+  struct impl;
+
+ public:
+  using key_type = std::uint32_t;
+
+  /**
+   * @brief Build one allowlist from host IDs.
+   *
+   * Host IDs are copied to the construction stream and then use the same device builder as the
+   * device overload. IDs must be smaller than dataset_rows; this precondition is not checked.
+   * Empty input is valid and rejects every candidate.
+   */
+  static roaring_allowlist from_ids(raft::resources const& res,
+                                    std::size_t dataset_rows,
+                                    raft::host_vector_view<const key_type, std::int64_t> ids,
+                                    bool pre_sorted = false);
+
+  /**
+   * @brief Build one allowlist from device IDs.
+   *
+   * The input must remain valid until the construction stream reaches the enqueued work. IDs must
+   * be smaller than dataset_rows; this precondition is not checked.
+   * Temporary memory is O(cardinality + container count); no dataset-sized dense bitmap is used.
+   */
+  static roaring_allowlist from_ids(raft::resources const& res,
+                                    std::size_t dataset_rows,
+                                    raft::device_vector_view<const key_type, std::int64_t> ids,
+                                    bool pre_sorted = false);
+
+  ~roaring_allowlist();
+
+  roaring_allowlist(roaring_allowlist const&)            = delete;
+  roaring_allowlist& operator=(roaring_allowlist const&) = delete;
+  roaring_allowlist(roaring_allowlist&&) noexcept;
+  roaring_allowlist& operator=(roaring_allowlist&&) noexcept;
+
+  [[nodiscard]] std::size_t dataset_rows() const noexcept;
+  [[nodiscard]] std::size_t cardinality() const noexcept;
+  [[nodiscard]] bool empty() const noexcept;
+
+  /** @brief Total device bytes retained by the encoded allowlist and initialized reference. */
+  [[nodiscard]] std::size_t size_bytes() const noexcept;
+
+  /** @brief Return a zero-copy view. */
+  [[nodiscard]] roaring_allowlist_view view() const noexcept;
+
+ private:
+  explicit roaring_allowlist(std::unique_ptr<impl> impl) noexcept;
+
+  std::unique_ptr<impl> impl_;
+};
+
+}  // namespace core
+}  // namespace CUVS_EXPORT cuvs

cpp/include/cuvs/detail/jit_lto/common_fragments.hpp

@@ -15,6 +15,7 @@ struct tag_u8 {};
 struct tag_filter_none {};
 struct tag_filter_bitset {};
 struct tag_filter_bloom_filter {};
+struct tag_filter_roaring {};
 struct tag_filter_udf {};
 
 struct tag_bitset_u32 {};

cpp/include/cuvs/neighbors/common.hpp

@@ -33,6 +33,7 @@
 #include <cstring>
 #include <memory>
 #include <numeric>
+#include <span>
 #include <string>
 #include <type_traits>
 #include <utility>
@@ -43,7 +44,9 @@
 namespace CUVS_EXPORT cuvs {
 namespace core {
 class bloom_filter;
-}
+class roaring_allowlist_view;
+class roaring_allowlist;
+}  // namespace core
 namespace neighbors {
 /**
  * @addtogroup cagra_cpp_index_params
@@ -1346,7 +1349,7 @@ namespace filtering {
  * @{
  */
 
-enum class FilterType : int { None = 0, Bitmap = 1, Bitset = 2, Bloom = 3, UDF = 100 };
+enum class FilterType : int { None = 0, Bitmap = 1, Bitset = 2, Bloom = 3, Roaring = 4, UDF = 100 };
 
 struct base_filter {
   ~base_filter()                             = default;
@@ -1501,6 +1504,87 @@ struct bloom_filter : public base_filter {
   FilterType get_filter_type() const override { return FilterType::Bloom; }
 };
 
+/**
+ * @brief Reusable per-query mapping to immutable exact Roaring allowlists.
+ *
+ * Entry @c q selects view @c q. CAGRA retains candidate dataset row @c r when the selected
+ * allowlist contains @c r. Construction copies only already initialized device-reference pointers
+ * and empty flags into the filter payload; encoded bytes are neither copied nor parsed. Search
+ * therefore performs no Roaring allocation, initialization, synchronization, or preprocessing.
+ *
+ * @code{.cpp}
+ * auto first = cuvs::core::roaring_allowlist::from_ids(
+ *   res, dataset_rows,
+ *   raft::make_host_vector_view<const std::uint32_t, std::int64_t>(first_ids.data(),
+ *                                                                    first_ids.size()));
+ * auto second = cuvs::core::roaring_allowlist::from_ids(
+ *   res, dataset_rows,
+ *   raft::make_host_vector_view<const std::uint32_t, std::int64_t>(second_ids.data(),
+ *                                                                    second_ids.size()));
+ * std::array views{first.view(), second.view()};
+ * auto filter = cuvs::neighbors::filtering::roaring_filter(res, views);
+ * @endcode
+ *
+ * Owners and views can be reused across filters and queries. This filter owns its mapping tables
+ * and device payload, but not the referenced owners, which must outlive the filter and all searches
+ * using it. Copies are cheap shared handles required by CAGRA query-offset wrappers.
+ *
+ * @see cuvs::core::roaring_allowlist
+ * @see https://github.com/RoaringBitmap/RoaringFormatSpec
+ */
+struct roaring_filter : public base_filter {
+ private:
+  struct impl;
+
+ public:
+  /** @brief Construct an invalid handle. It cannot be passed to CAGRA search. */
+  roaring_filter() = default;
+
+  /**
+   * @brief Materialize the query-to-allowlist device pointer table.
+   *
+   * @p allowlists must be nonempty, every view must be valid, and every view must have the same
+   * `dataset_rows()`. Query count is inferred from the span length.
+   */
+  explicit roaring_filter(raft::resources const& res,
+                          std::span<const cuvs::core::roaring_allowlist_view> allowlists);
+
+  [[nodiscard]] bool valid() const noexcept;
+  [[nodiscard]] std::size_t num_queries() const noexcept;
+  [[nodiscard]] std::size_t dataset_rows() const noexcept;
+  [[nodiscard]] std::size_t cardinality(std::size_t query_id) const;
+  [[nodiscard]] bool empty(std::size_t query_id) const;
+
+  /**
+   * @brief Maximum rejected fraction among all query allowlists.
+   *
+   * CAGRA uses this precomputed value when `search_params::filtering_rate` is unset.
+   */
+  [[nodiscard]] float filtering_rate() const noexcept;
+
+  /** @brief Device bytes owned by this mapping, excluding the referenced allowlists. */
+  [[nodiscard]] std::size_t size_bytes() const noexcept;
+
+  /**
+   * @brief Replace one query's allowlist pointer outside the search path.
+   *
+   * The replacement must have the same `dataset_rows()`. The method copies one pointer and one
+   * empty flag to the device and synchronizes @p res before returning. Do not call it concurrently
+   * with a search, and keep the replacement owner alive for all subsequent searches.
+   */
+  void set_allowlist(raft::resources const& res,
+                     std::size_t query_id,
+                     cuvs::core::roaring_allowlist_view replacement);
+
+  /** @brief Internal device payload already prepared for the linked CAGRA predicate. */
+  [[nodiscard]] void* device_payload() const noexcept;
+
+  FilterType get_filter_type() const override { return FilterType::Roaring; }
+
+ private:
+  std::shared_ptr<impl> impl_;
+};
+
 /**
  * @brief JIT-LTO user-defined filter predicate.
  *