enum

Author: Jordan Maples

diskann-benchmark/src/index/bftree/quantizer_util.rs

@@ -13,7 +13,10 @@ use diskann_quantization::{
 use diskann_utils::views::MatrixView;
 use rand::SeedableRng;
 
-use crate::{inputs::bftree::QuantConfig, utils::SimilarityMeasure};
+use crate::{
+    inputs::bftree::{QuantConfig, SphericalBits},
+    utils::SimilarityMeasure,
+};
 
 fn new_quantizer<const NBITS: usize>(
     quantizer: SphericalQuantizer,
@@ -54,11 +57,10 @@ pub(super) fn build_quantizer(
                 GlobalAllocator,
             )?;
 
-            let poly = match num_bits.get() {
-                1 => new_quantizer::<1>(quantizer)?,
-                2 => new_quantizer::<2>(quantizer)?,
-                4 => new_quantizer::<4>(quantizer)?,
-                n => anyhow::bail!("{n} bits not supported for spherical quantization"),
+            let poly = match *num_bits {
+                SphericalBits::One => new_quantizer::<1>(quantizer)?,
+                SphericalBits::Two => new_quantizer::<2>(quantizer)?,
+                SphericalBits::Four => new_quantizer::<4>(quantizer)?,
             };
 
             Ok(Some(poly))

diskann-benchmark/src/inputs/bftree.rs

@@ -2,7 +2,7 @@
  * Copyright (c) Microsoft Corporation.
  * Licensed under the MIT license.
  */
-use std::num::{NonZero, NonZeroUsize};
+use std::num::NonZero;
 
 use crate::inputs::{
     as_input, exhaustive,
@@ -211,6 +211,54 @@ fn bftree_parameters_from(
     }
 }
 
+/// Supported bit widths for spherical quantization.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
+pub(crate) enum SphericalBits {
+    #[serde(rename = "1")]
+    One = 1,
+    #[serde(rename = "2")]
+    Two = 2,
+    #[serde(rename = "4")]
+    Four = 4,
+}
+
+impl SphericalBits {
+    pub(crate) const fn get(self) -> usize {
+        self as usize
+    }
+}
+
+impl std::fmt::Display for SphericalBits {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", self.get())
+    }
+}
+
+impl TryFrom<usize> for SphericalBits {
+    type Error = String;
+
+    fn try_from(value: usize) -> Result<Self, Self::Error> {
+        match value {
+            1 => Ok(Self::One),
+            2 => Ok(Self::Two),
+            4 => Ok(Self::Four),
+            n => Err(format!(
+                "{n} bits not supported for spherical quantization; expected 1, 2, or 4"
+            )),
+        }
+    }
+}
+
+impl<'de> Deserialize<'de> for SphericalBits {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        let n = usize::deserialize(deserializer)?;
+        SphericalBits::try_from(n).map_err(serde::de::Error::custom)
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(tag = "kind")]
 pub(crate) enum QuantConfig {
@@ -220,7 +268,7 @@ pub(crate) enum QuantConfig {
     Spherical {
         seed: u64,
         transform_kind: exhaustive::TransformKind,
-        num_bits: NonZeroUsize,
+        num_bits: SphericalBits,
         #[serde(deserialize_with = "Deserialize::deserialize")]
         pre_scale: Option<exhaustive::PreScale>,
         #[serde(deserialize_with = "Deserialize::deserialize")]
@@ -233,14 +281,8 @@ impl QuantConfig {
         match self {
             Self::None => Ok(()),
             Self::Spherical {
-                num_bits,
-                quant_store_config,
-                ..
+                quant_store_config, ..
             } => {
-                match num_bits.get() {
-                    1 | 2 | 4 => {}
-                    n => anyhow::bail!("{n} bits are not supported for spherical quantization"),
-                }
                 if let Some(cfg) = quant_store_config {
                     cfg.fill_defaults();
                     cfg.validate()?;