Roaring Bitmap Filter in CAGRA

cpp/CMakeLists.txt

@@ -1358,6 +1358,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/include/cuvs/core/roaring_allowlist.hpp

@@ -0,0 +1,174 @@
+/*
+ * 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 row in a batched 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 pre-parsed 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 batch of exact per-query Roaring allowlists.
+ *
+ * Logically, the owner is a sparse matrix with one allowlist row per query and one possible column
+ * per dataset row. @ref from_ids accepts one contiguous ID vector plus an indptr vector that
+ * delimits independently sized query rows. Every row is sorted and encoded independently. All
+ * variable-length portable Roaring streams and their initialized
+ * `cuco::experimental::roaring_bitmap_ref<uint32_t>` objects share one packed device allocation.
+ *
+ * For multiple rows, the builder uses indptr directly for segmented device radix sort and
+ * schedules analysis/encoding over all containers in all rows. A one-row input retains the tuned
+ * single-allowlist builder. Final encoding and reference initialization remain stream ordered.
+ *
+ * IDs must be unique within each row. Setting @p pre_sorted skips sorting and promises that every
+ * row is strictly increasing; ordering and uniqueness are not checked. Every ID must be smaller
+ * than  dataset_rows.
+ *
+ * @see https://github.com/RoaringBitmap/RoaringFormatSpec
+ * @see
+ * https://github.com/NVIDIA/cuCollections/blob/6001618aaa7f17ea2bbcd444650e9573c4f3d6c5/include/cuco/roaring_bitmap_ref.cuh
+ */
+class CUVS_EXPORT roaring_allowlist {
+ private:
+  struct impl;
+
+ public:
+  using key_type    = std::uint32_t;
+  using indptr_type = std::int64_t;
+
+  /**
+   * @brief Build ragged allowlist rows from contiguous host IDs and row offsets.
+   *
+   * `indptr` contains `num_allowlists + 1` entries, starts at zero, is nondecreasing,
+   * and ends at `ids.extent(0)`. Empty slices are valid allowlists.
+   */
+  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,
+                                    raft::host_vector_view<const indptr_type, std::int64_t> indptr,
+                                    bool pre_sorted = false);
+
+  /**
+   * @brief Build ragged allowlist rows from contiguous device IDs and row offsets.
+   *
+   * The same indptr invariants as the host overload apply. The row offsets are copied to the host
+   * once for validation, shape-aware dispatch, and exact packed allocation.
+   *
+   * The input must remain valid until the construction stream reaches the work enqueued by this
+   * call. Temporary memory is O(total input IDs + total containers); no dense dataset-sized bitmap
+   * is materialized.
+   */
+  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,
+    raft::device_vector_view<const indptr_type, std::int64_t> indptr,
+    bool pre_sorted = false);
+
+  /**
+   * @brief Import packed standard 32-bit portable Roaring rows.
+   *
+   * `byte_offsets` has `num_allowlists + 1` entries, starts at zero, is nondecreasing, and ends at
+   * `bytes.extent(0)`. Empty slices represent empty allowlists. Every row is strictly validated on
+   * the host before its bytes are copied. The host buffers must remain valid until the construction
+   * stream completes; pinned bytes are recommended when overlap matters.
+   */
+  static roaring_allowlist from_serialized(
+    raft::resources const& res,
+    std::size_t dataset_rows,
+    raft::host_vector_view<const std::byte, std::int64_t> bytes,
+    raft::host_vector_view<const std::uint64_t, std::int64_t> byte_offsets);
+
+  ~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 num_allowlists() const noexcept;
+  [[nodiscard]] std::size_t dataset_rows() const noexcept;
+  [[nodiscard]] std::size_t cardinality(std::size_t allowlist_id) const;
+  [[nodiscard]] bool empty(std::size_t allowlist_id) const;
+  [[nodiscard]] std::size_t total_cardinality() const noexcept;
+
+  /** @brief Total device bytes retained by packed rows, references, and row pointer metadata. */
+  [[nodiscard]] std::size_t size_bytes() const noexcept;
+
+  /** @brief Return a zero-copy view of one row. */
+  [[nodiscard]] roaring_allowlist_view view(std::size_t allowlist_id) const;
+
+  /**
+   * @brief Test a matrix of row IDs and synchronize the resource stream.
+   *
+   * `row_ids[q][i]` is tested against allowlist row `q`. Input and output shapes must match, and
+   * their first extent must equal @ref num_allowlists.
+   */
+  void contains(raft::resources const& res,
+                raft::device_matrix_view<const key_type, std::int64_t, raft::row_major> row_ids,
+                raft::device_matrix_view<std::uint8_t, std::int64_t, raft::row_major> output) const;
+
+  /** @brief Stream-ordered asynchronous version of @ref contains. */
+  void contains_async(
+    raft::resources const& res,
+    raft::device_matrix_view<const key_type, std::int64_t, raft::row_major> row_ids,
+    raft::device_matrix_view<std::uint8_t, std::int64_t, raft::row_major> output) const;
+
+ 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,91 @@ struct bloom_filter : public base_filter {
   FilterType get_filter_type() const override { return FilterType::Bloom; }
 };
 
+/**
+ * @brief Reusable per-query mapping to an immutable batch of exact Roaring allowlists.
+ *
+ * Entry @c q selects row @c q of the owner. CAGRA retains candidate dataset row @c r when that
+ * allowlist contains @c r. Constructing from a @c cuvs::core::roaring_allowlist copies only its
+ * 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, parsing,
+ * initialization, synchronization, or per-query preprocessing.
+ *
+ * @code{.cpp}
+ * // Flat IDs plus num_queries + 1 row offsets.
+ * auto allowlists = cuvs::core::roaring_allowlist::from_ids(
+ *   res, dataset_rows,
+ *   raft::make_host_vector_view<const std::uint32_t, std::int64_t>(allowed_ids.data(),
+ *                                                                    allowed_ids.size()),
+ *   raft::make_host_vector_view<const std::int64_t, std::int64_t>(indptr.data(),
+ *                                                                   indptr.size()));
+ * std::vector<cuvs::core::roaring_allowlist_view> views;
+ * for (std::size_t q = 0; q < allowlists.num_allowlists(); ++q) {
+ *   views.push_back(allowlists.view(q));
+ * }
+ * auto filter = cuvs::neighbors::filtering::roaring_filter(res, views);
+ * @endcode
+ *
+ * The span overload remains useful when queries reuse rows from several owners or when one query's
+ * mapping must be replaced without rebuilding encoded allowlists. This filter owns its mapping
+ * tables and device payload, but not the referenced owner(s), 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.
  *