gqlite 1.5.0 → 1.7.0

data/ext/db-index/src/hnsw/store.rs

@@ -0,0 +1,472 @@
+use super::error::{Error, Result};
+
+#[derive(Debug, Clone, Copy, PartialEq, serde::Serialize, serde::Deserialize)]
+pub enum VectorEncoding
+{
+  F32,
+  F16,
+  Qi8
+  {
+    scale: f32,
+    zero_point: i8,
+  },
+}
+
+#[allow(dead_code)]
+pub fn encode_adjacency(level: u8, layers: &[Vec<u32>]) -> Vec<u8>
+{
+  assert!(
+    level as usize + 1 == layers.len(),
+    "layers must have level+1 entries"
+  );
+  let mut out = Vec::with_capacity(1 + layers.iter().map(|l| 2 + 4 * l.len()).sum::<usize>());
+  out.push(level);
+  for layer in layers
+  {
+    let count: u16 = layer.len().try_into().expect("neighbour count exceeds u16");
+    out.extend_from_slice(&count.to_le_bytes());
+    for &n in layer
+    {
+      out.extend_from_slice(&n.to_le_bytes());
+    }
+  }
+  out
+}
+
+/// Decode a live adjacency record into its level and per-layer neighbour
+/// lists. Passing empty bytes returns `Err(EmptyDecodedAdjency)`;
+/// structurally truncated or padded records return `Err(TruncatedDecodedAdjency)`
+/// or `Err(TrailingBytesDecodedAdjency)` respectively.
+#[allow(dead_code)]
+pub fn decode_adjacency_layers(bytes: &[u8]) -> Result<(u8, Vec<Vec<u32>>)>
+{
+  if bytes.is_empty()
+  {
+    return Err(Error::EmptyDecodedAdjency);
+  }
+  let level = bytes[0];
+  let mut pos = 1usize;
+  let mut layers = Vec::with_capacity(level as usize + 1);
+  for _ in 0..=level
+  {
+    if pos + 2 > bytes.len()
+    {
+      return Err(Error::TruncatedDecodedAdjency);
+    }
+    let count = u16::from_le_bytes([bytes[pos], bytes[pos + 1]]) as usize;
+    pos += 2;
+    let mut layer = Vec::with_capacity(count);
+    for _ in 0..count
+    {
+      if pos + 4 > bytes.len()
+      {
+        return Err(Error::TruncatedDecodedAdjency);
+      }
+      layer.push(u32::from_le_bytes([
+        bytes[pos],
+        bytes[pos + 1],
+        bytes[pos + 2],
+        bytes[pos + 3],
+      ]));
+      pos += 4;
+    }
+    layers.push(layer);
+  }
+  if pos != bytes.len()
+  {
+    return Err(Error::TrailingBytesDecodedAdjency);
+  }
+  Ok((level, layers))
+}
+
+#[allow(dead_code)]
+pub(crate) fn validate_qi8_scale(encoding: &VectorEncoding) -> Result<()>
+{
+  if let VectorEncoding::Qi8 { scale, .. } = encoding
+  {
+    if *scale <= 0.0 || !scale.is_finite()
+    {
+      return Err(Error::InvalidQi8Scale { scale: *scale });
+    }
+  }
+  Ok(())
+}
+
+#[allow(dead_code)]
+pub fn encode_vector(encoding: &VectorEncoding, floats: &[f32]) -> Result<Vec<u8>>
+{
+  validate_qi8_scale(encoding)?;
+  Ok(match encoding
+  {
+    VectorEncoding::F32 => floats.iter().flat_map(|f| f.to_le_bytes()).collect(),
+    VectorEncoding::F16 => floats
+      .iter()
+      .flat_map(|f| half::f16::from_f32(*f).to_le_bytes())
+      .collect(),
+    VectorEncoding::Qi8 { scale, zero_point } => floats
+      .iter()
+      .map(|f| {
+        let q = (*f / scale) + (*zero_point as f32);
+        q.round().clamp(i8::MIN as f32, i8::MAX as f32) as i8 as u8
+      })
+      .collect(),
+  })
+}
+
+#[allow(dead_code)]
+pub fn decode_vector(encoding: &VectorEncoding, bytes: &[u8], dim: usize) -> Result<Vec<f32>>
+{
+  validate_qi8_scale(encoding)?;
+  let unit_size = match encoding
+  {
+    VectorEncoding::F32 => 4,
+    VectorEncoding::F16 => 2,
+    VectorEncoding::Qi8 { .. } => 1,
+  };
+  let expected = dim * unit_size;
+  if bytes.len() != expected
+  {
+    return Err(Error::DimensionMismatch {
+      expected,
+      actual: bytes.len(),
+    });
+  }
+  Ok(match encoding
+  {
+    VectorEncoding::F32 => bytes
+      .as_chunks::<4>()
+      .0
+      .iter()
+      .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
+      .collect(),
+    VectorEncoding::F16 => bytes
+      .as_chunks::<2>()
+      .0
+      .iter()
+      .map(|c| half::f16::from_le_bytes([c[0], c[1]]).to_f32())
+      .collect(),
+    VectorEncoding::Qi8 { scale, zero_point } => bytes
+      .iter()
+      .map(|&b| (b as i8 as f32 - *zero_point as f32) * scale)
+      .collect(),
+  })
+}
+
+#[allow(dead_code)]
+const HNSW_META_SCHEMA_VERSION: u8 = 1;
+
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+#[allow(dead_code)]
+struct VersionOnly
+{
+  schema_version: u8,
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
+#[allow(dead_code)]
+pub(crate) struct HnswMetaBlob
+{
+  pub schema_version: u8,
+  pub entry_point: Option<u32>,
+  pub max_level: i32,
+  pub next_node_id: u32,
+  pub dim: u32,
+  pub m: u32,
+  pub ef_construction: u32,
+  pub ef_search: u32,
+  pub vector_encoding: VectorEncoding,
+}
+
+#[allow(dead_code)]
+pub(crate) fn encode_meta(meta: &HnswMetaBlob) -> Vec<u8>
+{
+  let mut data = Vec::<u8>::new();
+  ciborium::into_writer(meta, &mut data).unwrap();
+  data
+}
+
+#[allow(dead_code)]
+pub(crate) fn decode_meta(bytes: &[u8]) -> Result<HnswMetaBlob>
+{
+  let version_probe: VersionOnly =
+    ciborium::from_reader(bytes).map_err(|e| Error::MalformedMetadata(e.to_string()))?;
+  if version_probe.schema_version > HNSW_META_SCHEMA_VERSION
+  {
+    return Err(Error::SchemaVersionTooNew {
+      found: version_probe.schema_version,
+    });
+  }
+  if version_probe.schema_version < HNSW_META_SCHEMA_VERSION
+  {
+    return Err(Error::SchemaVersionUnsupported {
+      found: version_probe.schema_version,
+    });
+  }
+  let meta: HnswMetaBlob =
+    ciborium::from_reader(bytes).map_err(|e| Error::MalformedMetadata(e.to_string()))?;
+  validate_qi8_scale(&meta.vector_encoding)?;
+  Ok(meta)
+}
+
+#[cfg(test)]
+mod testing
+{
+  use super::*;
+  use crate::hnsw::id::NodeId;
+  use std::assert_matches;
+
+  #[test]
+  fn node_id_from_u32_roundtrips_through_as_u32()
+  {
+    assert_eq!(NodeId::from_u32(12345).as_u32(), 12345);
+  }
+
+  #[test]
+  fn adjacency_roundtrip_level0_m16()
+  {
+    let layer0: Vec<u32> = (0..16).collect();
+    let bytes = encode_adjacency(0, std::slice::from_ref(&layer0));
+    assert_eq!(bytes.len(), 1 + 2 + 16 * 4);
+    let (level, layers) = decode_adjacency_layers(&bytes).unwrap();
+    assert_eq!(level, 0);
+    assert_eq!(layers, vec![layer0]);
+  }
+
+  #[test]
+  fn adjacency_count_field_does_not_overflow_past_u8_neighbour_counts()
+  {
+    let layer0: Vec<u32> = (0..200).collect();
+    let bytes = encode_adjacency(0, std::slice::from_ref(&layer0));
+    let (_, layers) = decode_adjacency_layers(&bytes).unwrap();
+    assert_eq!(layers[0].len(), 200);
+  }
+
+  #[test]
+  #[should_panic(expected = "neighbour count exceeds u16")]
+  fn adjacency_encode_panics_rather_than_silently_truncating_above_u16_max()
+  {
+    let huge: Vec<u32> = (0..70_000).collect();
+    encode_adjacency(0, &[huge]);
+  }
+
+  #[test]
+  fn vector_f32_roundtrip_exact()
+  {
+    let v = vec![1.0f32, -2.5, 0.0, std::f32::consts::PI];
+    let bytes = encode_vector(&VectorEncoding::F32, &v).unwrap();
+    let back = decode_vector(&VectorEncoding::F32, &bytes, v.len()).unwrap();
+    assert_eq!(back, v);
+  }
+
+  #[test]
+  fn vector_f16_roundtrip_within_tolerance()
+  {
+    let v = vec![1.0f32, -2.5, 100.25, 0.001];
+    let bytes = encode_vector(&VectorEncoding::F16, &v).unwrap();
+    let back = decode_vector(&VectorEncoding::F16, &bytes, v.len()).unwrap();
+    for (a, b) in v.iter().zip(back.iter())
+    {
+      let rel_err = if *a != 0.0
+      {
+        (a - b).abs() / a.abs()
+      }
+      else
+      {
+        (a - b).abs()
+      };
+      assert!(rel_err < 0.001, "a={a} b={b} rel_err={rel_err}");
+    }
+  }
+
+  #[test]
+  fn vector_qi8_roundtrip_within_scale()
+  {
+    let scale = 0.05f32;
+    let zero_point = 0i8;
+    let v = vec![1.0f32, -1.0, 0.0, 2.5];
+    let enc = VectorEncoding::Qi8 { scale, zero_point };
+    let bytes = encode_vector(&enc, &v).unwrap();
+    let back = decode_vector(&enc, &bytes, v.len()).unwrap();
+    for (a, b) in v.iter().zip(back.iter())
+    {
+      assert!((a - b).abs() <= scale, "a={a} b={b}");
+    }
+  }
+
+  #[test]
+  fn vector_qi8_clamps_out_of_range_values_instead_of_wrapping()
+  {
+    let enc = VectorEncoding::Qi8 {
+      scale: 1.0,
+      zero_point: 0,
+    };
+    let bytes = encode_vector(&enc, &[1000.0, -1000.0]).unwrap();
+    assert_eq!(bytes[0] as i8, i8::MAX);
+    assert_eq!(bytes[1] as i8, i8::MIN);
+  }
+
+  #[test]
+  fn vector_decode_rejects_wrong_length()
+  {
+    let bytes = encode_vector(&VectorEncoding::F32, &[1.0, 2.0]).unwrap();
+    assert_matches!(
+      decode_vector(&VectorEncoding::F32, &bytes, 3),
+      Err(Error::DimensionMismatch {
+        expected: 12,
+        actual: 8
+      })
+    );
+  }
+
+  #[test]
+  fn vector_decode_rejects_extra_trailing_bytes_even_when_a_naive_take_would_hide_it()
+  {
+    let bytes = encode_vector(&VectorEncoding::F32, &[1.0, 2.0, 3.0, 4.0, 5.0]).unwrap();
+    assert_matches!(
+      decode_vector(&VectorEncoding::F32, &bytes, 4),
+      Err(Error::DimensionMismatch {
+        expected: 16,
+        actual: 20
+      })
+    );
+  }
+
+  #[test]
+  fn vector_qi8_encode_rejects_zero_scale()
+  {
+    let enc = VectorEncoding::Qi8 {
+      scale: 0.0,
+      zero_point: 0,
+    };
+    assert_matches!(
+      encode_vector(&enc, &[1.0, 2.0]),
+      Err(Error::InvalidQi8Scale { scale: 0.0 })
+    );
+  }
+
+  #[test]
+  fn vector_qi8_encode_rejects_negative_scale()
+  {
+    let enc = VectorEncoding::Qi8 {
+      scale: -1.0,
+      zero_point: 0,
+    };
+    assert_matches!(
+      encode_vector(&enc, &[1.0]),
+      Err(Error::InvalidQi8Scale { scale: -1.0 })
+    );
+  }
+
+  #[test]
+  fn vector_qi8_decode_also_rejects_invalid_scale_not_just_encode()
+  {
+    let enc = VectorEncoding::Qi8 {
+      scale: 0.0,
+      zero_point: 0,
+    };
+    assert_matches!(
+      decode_vector(&enc, &[0, 0], 2),
+      Err(Error::InvalidQi8Scale { scale: 0.0 })
+    );
+  }
+
+  #[test]
+  fn vector_encoding_sizes_are_expected()
+  {
+    let v = vec![1.0f32; 10];
+    assert_eq!(encode_vector(&VectorEncoding::F32, &v).unwrap().len(), 40);
+    assert_eq!(encode_vector(&VectorEncoding::F16, &v).unwrap().len(), 20);
+    assert_eq!(
+      encode_vector(
+        &VectorEncoding::Qi8 {
+          scale: 1.0,
+          zero_point: 0
+        },
+        &v
+      )
+      .unwrap()
+      .len(),
+      10
+    );
+  }
+
+  fn sample_meta() -> HnswMetaBlob
+  {
+    HnswMetaBlob {
+      schema_version: HNSW_META_SCHEMA_VERSION,
+      entry_point: Some(7),
+      max_level: 2,
+      next_node_id: 103,
+      dim: 384,
+      m: 16,
+      ef_construction: 200,
+      ef_search: 50,
+      vector_encoding: VectorEncoding::F16,
+    }
+  }
+
+  #[test]
+  fn meta_roundtrip()
+  {
+    let meta = sample_meta();
+    let bytes = encode_meta(&meta);
+    let back = decode_meta(&bytes).unwrap();
+    assert_eq!(meta, back);
+  }
+
+  #[test]
+  fn meta_empty_graph_has_no_entry_point()
+  {
+    let mut meta = sample_meta();
+    meta.entry_point = None;
+    meta.max_level = -1;
+    let bytes = encode_meta(&meta);
+    let back = decode_meta(&bytes).unwrap();
+    assert_eq!(back.entry_point, None);
+    assert_eq!(back.max_level, -1);
+  }
+
+  #[test]
+  fn meta_rejects_newer_schema_version()
+  {
+    let mut meta = sample_meta();
+    meta.schema_version = HNSW_META_SCHEMA_VERSION + 1;
+    let bytes = encode_meta(&meta);
+
+    const HNSW_META_SCHEMA_VERSION_PLUS_1: u8 = HNSW_META_SCHEMA_VERSION + 1;
+
+    assert_matches!(
+      decode_meta(&bytes),
+      Err(Error::SchemaVersionTooNew {
+        found: HNSW_META_SCHEMA_VERSION_PLUS_1,
+      })
+    );
+  }
+
+  #[test]
+  fn meta_rejects_older_schema_version_no_downgrade_path()
+  {
+    let mut meta = sample_meta();
+    meta.schema_version = 0;
+    let bytes = encode_meta(&meta);
+    assert_matches!(
+      decode_meta(&bytes),
+      Err(Error::SchemaVersionUnsupported { found: 0 })
+    );
+  }
+
+  #[test]
+  fn meta_rejects_corrupted_qi8_scale_at_load_boundary()
+  {
+    let mut meta = sample_meta();
+    meta.vector_encoding = VectorEncoding::Qi8 {
+      scale: 0.0,
+      zero_point: 0,
+    };
+    let bytes = encode_meta(&meta);
+    assert_matches!(
+      decode_meta(&bytes),
+      Err(Error::InvalidQi8Scale { scale: 0.0 })
+    );
+  }
+}

data/ext/db-index/src/hnsw/vector.rs

@@ -0,0 +1,105 @@
+use super::scalar::Dtype;
+use super::scalar::Scalar;
+use std::marker::PhantomData;
+
+pub trait VectorRef: Copy
+{
+  fn len(self) -> usize;
+  fn is_empty(self) -> bool;
+}
+
+impl<T> VectorRef for &[T]
+{
+  #[inline]
+  fn len(self) -> usize
+  {
+    self.len()
+  }
+  #[inline]
+  fn is_empty(self) -> bool
+  {
+    self.is_empty()
+  }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct Qi8Ref<'a>
+{
+  pub data: &'a [i8],
+  pub scale: f32,
+  pub zero_point: i8,
+}
+
+impl VectorRef for Qi8Ref<'_>
+{
+  #[inline]
+  fn len(self) -> usize
+  {
+    self.data.len()
+  }
+  #[inline]
+  fn is_empty(self) -> bool
+  {
+    self.data.is_empty()
+  }
+}
+
+pub trait VectorFamily: Copy + Clone + Send + Sync + 'static
+{
+  type Ref<'a>: VectorRef
+  where
+    Self: 'a;
+
+  const DTYPE: Dtype;
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct Dense<S: Scalar>(PhantomData<S>);
+
+impl<S: Scalar> VectorFamily for Dense<S>
+{
+  type Ref<'a>
+    = &'a [S]
+  where
+    Self: 'a;
+
+  const DTYPE: Dtype = S::DTYPE;
+}
+
+#[derive(Clone, Copy, Debug, Default)]
+pub struct Qi8;
+
+impl VectorFamily for Qi8
+{
+  type Ref<'a>
+    = Qi8Ref<'a>
+  where
+    Self: 'a;
+
+  const DTYPE: Dtype = Dtype::QI8;
+}
+
+pub trait VectorView<F: VectorFamily>: Clone + Send + Sync
+{
+  fn view<'a>(&'a self) -> <F as VectorFamily>::Ref<'a>;
+}
+
+impl<S: Scalar> VectorView<Dense<S>> for &[S]
+{
+  fn view(&self) -> &[S]
+  {
+    self
+  }
+}
+
+impl<'v> VectorView<Qi8> for Qi8Ref<'v>
+{
+  fn view<'a>(&'a self) -> Qi8Ref<'a>
+  {
+    Qi8Ref {
+      data: self.data,
+      scale: self.scale,
+      zero_point: self.zero_point,
+    }
+  }
+}