make vq codebook optional

Author: tarang-jain

c/src/preprocessing/quantize/pq.cpp

@@ -272,10 +272,12 @@ extern "C" cuvsError_t cuvsProductQuantizerGetVqCodebook(cuvsProductQuantizer_t
     if (quantizer != nullptr) {
       auto quant_addr = quantizer->addr;
       if (quantizer->dtype.code == kDLFloat && quantizer->dtype.bits == 32) {
-        auto pq_mdspan =
+        auto vq_opt =
           (reinterpret_cast<cuvs::preprocessing::quantize::pq::quantizer<float>*>(quant_addr))
             ->codebooks.vq_code_book();
-        cuvs::core::to_dlpack(pq_mdspan, vq_codebook);
+        RAFT_EXPECTS(vq_opt.has_value(),
+                     "quantizer has no VQ codebook (build with use_vq=true to enable)");
+        cuvs::core::to_dlpack(vq_opt.value(), vq_codebook);
       } else {
         RAFT_FAIL("Unsupported quantizer dtype: %d and bits: %d",
                   quantizer->dtype.code,

cpp/include/cuvs/preprocessing/quantize/vpq_dataset.hpp

@@ -7,6 +7,8 @@
 
 #include <cuvs/neighbors/common.hpp>
 
+#include <optional>
+
 namespace cuvs::preprocessing::quantize::pq {
 
 /**
@@ -21,18 +23,29 @@ class vpq_codebooks_iface {
 
   virtual ~vpq_codebooks_iface() = default;
 
-  /** VQ codebook [vq_n_centers, dim]. */
+  /**
+   * VQ codebook [vq_n_centers, dim].
+   *
+   * Returns std::nullopt when no VQ codebook is configured (i.e. PQ-only,
+   * use_vq=false). Callers that need to forward a `device_matrix_view`
+   * downstream should materialize an empty 0x0 view themselves on nullopt.
+   */
   [[nodiscard]] virtual auto vq_code_book() const noexcept
-    -> raft::device_matrix_view<const math_type, uint32_t, raft::row_major> = 0;
+    -> std::optional<raft::device_matrix_view<const math_type, uint32_t, raft::row_major>> = 0;
 
   /** PQ codebook [pq_n_centers (× pq_dim for subspaces), pq_len]. */
   [[nodiscard]] virtual auto pq_code_book() const noexcept
     -> raft::device_matrix_view<const math_type, uint32_t, raft::row_major> = 0;
 
-  [[nodiscard]] virtual auto dim() const noexcept -> uint32_t { return vq_code_book().extent(1); }
+  [[nodiscard]] virtual auto dim() const noexcept -> uint32_t
+  {
+    auto vq = vq_code_book();
+    return vq.has_value() ? vq->extent(1) : 0;
+  }
   [[nodiscard]] virtual auto vq_n_centers() const noexcept -> uint32_t
   {
-    return vq_code_book().extent(0);
+    auto vq = vq_code_book();
+    return vq.has_value() ? vq->extent(0) : 0;
   }
   [[nodiscard]] virtual auto pq_len() const noexcept -> uint32_t
   {
@@ -89,8 +102,11 @@ class vpq_codebooks {
   vpq_codebooks& operator=(vpq_codebooks&&)      = default;
   ~vpq_codebooks()                               = default;
 
+  /**
+   * VQ codebook view, or std::nullopt when no VQ codebook is configured.
+   */
   [[nodiscard]] auto vq_code_book() const noexcept
-    -> raft::device_matrix_view<const math_type, uint32_t, raft::row_major>
+    -> std::optional<raft::device_matrix_view<const math_type, uint32_t, raft::row_major>>
   {
     return impl_->vq_code_book();
   }
@@ -150,8 +166,11 @@ class vpq_dataset : public cuvs::neighbors::dataset<IdxT> {
   [[nodiscard]] uint32_t dim() const noexcept override { return codebooks.dim(); }
   [[nodiscard]] bool is_owning() const noexcept override { return true; }
 
+  /**
+   * VQ codebook view, or std::nullopt when no VQ codebook is configured.
+   */
   [[nodiscard]] auto vq_code_book() const noexcept
-    -> raft::device_matrix_view<const math_type, uint32_t, raft::row_major>
+    -> std::optional<raft::device_matrix_view<const math_type, uint32_t, raft::row_major>>
   {
     return codebooks.vq_code_book();
   }

cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp

@@ -50,10 +50,11 @@ struct vpq_descriptor_spec : public instance_spec<DataT, IndexT, DistanceT> {
                    cuvs::distance::DistanceType metric,
                    const DistanceT* dataset_norms = nullptr) -> host_type
   {
+    auto vq_view = dataset.vq_code_book();
     return init_(params,
                  dataset.data.data_handle(),
                  dataset.encoded_row_length(),
-                 dataset.vq_code_book().data_handle(),
+                 vq_view.has_value() ? vq_view->data_handle() : nullptr,
                  dataset.pq_code_book().data_handle(),
                  IndexT(dataset.n_rows()),
                  dataset.dim());

cpp/src/neighbors/detail/dataset_serialize.hpp

@@ -69,7 +69,9 @@ void serialize(const raft::resources& res,
   raft::serialize_scalar(res, os, dataset.pq_n_centers());
   raft::serialize_scalar(res, os, dataset.pq_len());
   raft::serialize_scalar(res, os, dataset.encoded_row_length());
-  raft::serialize_mdspan(res, os, dataset.vq_code_book());
+  auto vq_view = dataset.vq_code_book().value_or(
+    raft::make_device_matrix_view<const MathT, uint32_t, raft::row_major>(nullptr, 0, 0));
+  raft::serialize_mdspan(res, os, vq_view);
   raft::serialize_mdspan(res, os, dataset.pq_code_book());
   raft::serialize_mdspan(res, os, dataset.data.view());
 }
@@ -160,10 +162,12 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is)
   raft::deserialize_mdspan(res, is, pq_code_book.view());
   raft::deserialize_mdspan(res, is, data.view());
 
+  std::optional<raft::device_matrix<MathT, uint32_t, raft::row_major>> vq_code_book_opt;
+  if (vq_n_centers > 0) { vq_code_book_opt = std::move(vq_code_book); }
   return std::make_unique<cuvs::preprocessing::quantize::pq::vpq_dataset<MathT, IdxT>>(
     cuvs::preprocessing::quantize::pq::vpq_codebooks<MathT>{
       std::make_unique<cuvs::preprocessing::quantize::pq::vpq_codebooks_owning<MathT>>(
-        std::move(vq_code_book), std::move(pq_code_book))},
+        std::move(pq_code_book), std::move(vq_code_book_opt))},
     std::move(data));
 }
 

cpp/src/preprocessing/quantize/detail/pq.cuh

@@ -154,9 +154,11 @@ quantizer<MathT> build(
     pq_code_book = cuvs::neighbors::detail::train_pq<MathT>(
       res, vpq_params, dataset, raft::make_const_mdspan(vq_code_book.view()));
   }
+  std::optional<raft::device_matrix<MathT, uint32_t, raft::row_major>> vq_code_book_opt;
+  if (filled_params.use_vq) { vq_code_book_opt = std::move(vq_code_book); }
   return {filled_params,
           vpq_codebooks<MathT>{std::make_unique<vpq_codebooks_owning<MathT>>(
-            std::move(vq_code_book), std::move(pq_code_book))}};
+            std::move(pq_code_book), std::move(vq_code_book_opt))}};
 }
 
 template <typename MathT>
@@ -195,12 +197,9 @@ quantizer<MathT> build_view(
                  vq_centers.value().extent(0));
   }
 
-  auto vq_view =
-    vq_centers.has_value()
-      ? vq_centers.value()
-      : raft::make_device_matrix_view<const MathT, uint32_t, raft::row_major>(nullptr, 0, 0);
-  return {params,
-          vpq_codebooks<MathT>{std::make_unique<vpq_codebooks_view<MathT>>(vq_view, pq_centers)}};
+  return {
+    params,
+    vpq_codebooks<MathT>{std::make_unique<vpq_codebooks_view<MathT>>(pq_centers, vq_centers)}};
 }
 
 template <typename T, typename QuantI, typename AccessorType>
@@ -223,7 +222,9 @@ void transform(
   RAFT_EXPECTS(quantizer.params_quantizer.pq_bits >= 4 && quantizer.params_quantizer.pq_bits <= 16,
                "PQ bits must be within [4, 16]");
 
-  auto vq_centers     = quantizer.codebooks.vq_code_book();
+  auto vq_centers_opt = quantizer.codebooks.vq_code_book();
+  auto vq_centers     = vq_centers_opt.value_or(
+    raft::make_device_matrix_view<const T, uint32_t, raft::row_major>(nullptr, 0, 0));
   auto pq_centers     = quantizer.codebooks.pq_code_book();
   auto vq_labels_view = raft::make_device_vector_view<uint32_t, int64_t>(nullptr, 0);
   if (vq_labels.has_value()) { vq_labels_view = vq_labels.value(); }
@@ -360,11 +361,14 @@ void inverse_transform(
   RAFT_EXPECTS(quantizer.params_quantizer.pq_bits >= 4 && quantizer.params_quantizer.pq_bits <= 16,
                "PQ bits must be within [4, 16]");
 
+  auto vq_centers_opt = quantizer.codebooks.vq_code_book();
+  auto vq_centers     = vq_centers_opt.value_or(
+    raft::make_device_matrix_view<const T, uint32_t, raft::row_major>(nullptr, 0, 0));
   reconstruct_vectors<T, T, idx_t, label_t>(res,
                                             quantizer.params_quantizer,
                                             codes,
                                             quantizer.codebooks.pq_code_book(),
-                                            quantizer.codebooks.vq_code_book(),
+                                            vq_centers,
                                             vq_labels,
                                             out,
                                             quantizer.params_quantizer.use_subspaces);
@@ -374,21 +378,26 @@ template <typename NewMathT, typename OldMathT>
 auto vpq_convert_math_type(const raft::resources& res, const vpq_codebooks<OldMathT>& src)
   -> vpq_codebooks<NewMathT>
 {
-  auto vq_src = src.vq_code_book();
-  auto pq_src = src.pq_code_book();
+  auto vq_src_opt = src.vq_code_book();
+  auto pq_src     = src.pq_code_book();
 
-  auto vq_new = raft::make_device_matrix<NewMathT, uint32_t, raft::row_major>(
-    res, vq_src.extent(0), vq_src.extent(1));
   auto pq_new = raft::make_device_matrix<NewMathT, uint32_t, raft::row_major>(
     res, pq_src.extent(0), pq_src.extent(1));
-
-  raft::linalg::map(
-    res, vq_new.view(), cuvs::spatial::knn::detail::utils::mapping<NewMathT>{}, vq_src);
   raft::linalg::map(
     res, pq_new.view(), cuvs::spatial::knn::detail::utils::mapping<NewMathT>{}, pq_src);
 
+  std::optional<raft::device_matrix<NewMathT, uint32_t, raft::row_major>> vq_new_opt;
+  if (vq_src_opt.has_value()) {
+    auto vq_src = vq_src_opt.value();
+    auto vq_new = raft::make_device_matrix<NewMathT, uint32_t, raft::row_major>(
+      res, vq_src.extent(0), vq_src.extent(1));
+    raft::linalg::map(
+      res, vq_new.view(), cuvs::spatial::knn::detail::utils::mapping<NewMathT>{}, vq_src);
+    vq_new_opt = std::move(vq_new);
+  }
+
   return vpq_codebooks<NewMathT>{
-    std::make_unique<vpq_codebooks_owning<NewMathT>>(std::move(vq_new), std::move(pq_new))};
+    std::make_unique<vpq_codebooks_owning<NewMathT>>(std::move(pq_new), std::move(vq_new_opt))};
 }
 
 template <typename DatasetT, typename MathT, typename IdxT>
@@ -419,8 +428,8 @@ auto vpq_build(const raft::resources& res,
     true);
 
   return vpq_dataset<MathT, IdxT>{
-    vpq_codebooks<MathT>{std::make_unique<vpq_codebooks_owning<MathT>>(std::move(vq_code_book),
-                                                                       std::move(pq_code_book))},
+    vpq_codebooks<MathT>{std::make_unique<vpq_codebooks_owning<MathT>>(std::move(pq_code_book),
+                                                                       std::move(vq_code_book))},
     std::move(codes)};
 }
 

cpp/src/preprocessing/quantize/detail/vpq_dataset_impl.hpp

@@ -7,16 +7,21 @@
 
 #include <cuvs/preprocessing/quantize/vpq_dataset.hpp>
 
+#include <optional>
+
 namespace cuvs::preprocessing::quantize::pq {
 
 template <typename MathT>
 class vpq_codebooks_owning : public vpq_codebooks_iface<MathT> {
  public:
-  using math_type = MathT;
-
-  vpq_codebooks_owning(raft::device_matrix<math_type, uint32_t, raft::row_major>&& vq_code_book,
-                       raft::device_matrix<math_type, uint32_t, raft::row_major>&& pq_code_book)
-    : vq_code_book_{std::move(vq_code_book)}, pq_code_book_{std::move(pq_code_book)}
+  using math_type   = MathT;
+  using matrix_type = raft::device_matrix<math_type, uint32_t, raft::row_major>;
+
+  // PQ codebook is required; VQ codebook is optional and defaults to absent
+  // When VQ is not provided, vq_code_book() returns std::nullopt.
+  explicit vpq_codebooks_owning(matrix_type&& pq_code_book,
+                                std::optional<matrix_type>&& vq_code_book = std::nullopt)
+    : pq_code_book_{std::move(pq_code_book)}, vq_code_book_{std::move(vq_code_book)}
   {
   }
 
@@ -27,15 +32,10 @@ class vpq_codebooks_owning : public vpq_codebooks_iface<MathT> {
   ~vpq_codebooks_owning() override                             = default;
 
   [[nodiscard]] auto vq_code_book() const noexcept
-    -> raft::device_matrix_view<const math_type, uint32_t, raft::row_major> override
-  {
-    return vq_code_book_.view();
-  }
-
-  [[nodiscard]] auto vq_code_book() noexcept
-    -> raft::device_matrix_view<math_type, uint32_t, raft::row_major>
+    -> std::optional<raft::device_matrix_view<const math_type, uint32_t, raft::row_major>> override
   {
-    return vq_code_book_.view();
+    if (!vq_code_book_.has_value()) { return std::nullopt; }
+    return vq_code_book_.value().view();
   }
 
   [[nodiscard]] auto pq_code_book() const noexcept
@@ -51,19 +51,21 @@ class vpq_codebooks_owning : public vpq_codebooks_iface<MathT> {
   }
 
  private:
-  raft::device_matrix<math_type, uint32_t, raft::row_major> vq_code_book_;
-  raft::device_matrix<math_type, uint32_t, raft::row_major> pq_code_book_;
+  matrix_type pq_code_book_;
+  std::optional<matrix_type> vq_code_book_;
 };
 
 template <typename MathT>
 class vpq_codebooks_view : public vpq_codebooks_iface<MathT> {
  public:
   using math_type = MathT;
+  using view_type = raft::device_matrix_view<const math_type, uint32_t, raft::row_major>;
 
-  vpq_codebooks_view(
-    raft::device_matrix_view<const math_type, uint32_t, raft::row_major> vq_code_book_view,
-    raft::device_matrix_view<const math_type, uint32_t, raft::row_major> pq_code_book_view)
-    : vq_code_book_view_{vq_code_book_view}, pq_code_book_view_{pq_code_book_view}
+  // PQ codebook view is required; VQ codebook view is optional and defaults to absent.
+  // When VQ is not provided, vq_code_book() returns std::nullopt.
+  explicit vpq_codebooks_view(view_type pq_code_book_view,
+                              std::optional<view_type> vq_code_book_view = std::nullopt)
+    : pq_code_book_view_{pq_code_book_view}, vq_code_book_view_{vq_code_book_view}
   {
   }
 
@@ -73,20 +75,18 @@ class vpq_codebooks_view : public vpq_codebooks_iface<MathT> {
   vpq_codebooks_view& operator=(vpq_codebooks_view&&)      = default;
   ~vpq_codebooks_view() override                           = default;
 
-  [[nodiscard]] auto vq_code_book() const noexcept
-    -> raft::device_matrix_view<const math_type, uint32_t, raft::row_major> override
+  [[nodiscard]] auto vq_code_book() const noexcept -> std::optional<view_type> override
   {
     return vq_code_book_view_;
   }
-  [[nodiscard]] auto pq_code_book() const noexcept
-    -> raft::device_matrix_view<const math_type, uint32_t, raft::row_major> override
+  [[nodiscard]] auto pq_code_book() const noexcept -> view_type override
   {
     return pq_code_book_view_;
   }
 
  private:
-  raft::device_matrix_view<const math_type, uint32_t, raft::row_major> vq_code_book_view_;
-  raft::device_matrix_view<const math_type, uint32_t, raft::row_major> pq_code_book_view_;
+  view_type pq_code_book_view_;
+  std::optional<view_type> vq_code_book_view_;
 };
 
 }  // namespace cuvs::preprocessing::quantize::pq