pq opt

timescale · Mar 25, 2024 · 5cd70ad · 5cd70ad
1 parent 578965f
commit 5cd70ad
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 2 deletions.
diff --git a/timescale_vector/src/access_method/pq_quantizer.rs b/timescale_vector/src/access_method/pq_quantizer.rs
@@ -241,9 +241,30 @@ impl PqQuantizer {
         PqDistanceTable::new_for_full_query(&self.pq.as_ref().unwrap(), distance_fn, query)
     }
 
-    pub fn get_distance_table_pq_query(&self, pq_vector: &[u8]) -> PqDistanceTable {
+    pub fn get_distance_table_pq_query(&self, pq_vector: &[PqVectorElement]) -> PqDistanceTable {
         PqDistanceTable::new_for_pq_query(&self.pq.as_ref().unwrap(), pq_vector)
     }
+
+    pub fn get_distance_directly(
+        &self,
+        left: &[PqVectorElement],
+        right: &[PqVectorElement],
+    ) -> f32 {
+        let pq = self.pq.as_ref().unwrap();
+        let sq = pq.subquantizers();
+
+        let mut dist = 0.0;
+        for (subquantizer_index, subquantizer) in sq.outer_iter().enumerate() {
+            let left_slice = subquantizer.index_axis(Axis(0), left[subquantizer_index] as usize);
+            let right_slice = subquantizer.index_axis(Axis(0), right[subquantizer_index] as usize);
+            assert!(left_slice.len() == right_slice.len());
+            dist += distance_l2_optimized_for_few_dimensions(
+                left_slice.as_slice().unwrap(),
+                right_slice.as_slice().unwrap(),
+            );
+        }
+        dist
+    }
 }
 
 /// DistanceCalculator encapsulates the code to generate distances between a PQ vector and a query.

diff --git a/timescale_vector/src/access_method/pq_storage.rs b/timescale_vector/src/access_method/pq_storage.rs
@@ -22,7 +22,7 @@ use crate::util::{
 
 use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance};
 
-pub struct PqNodeDistanceMeasure<'a> {
+/*pub struct PqNodeDistanceMeasure<'a> {
     storage: &'a PqCompressionStorage<'a>,
     table: PqDistanceTable,
 }
@@ -57,6 +57,42 @@ impl<'a> NodeDistanceMeasure for PqNodeDistanceMeasure<'a> {
         let node1 = rn1.get_archived_node();
         self.table.distance(node1.pq_vector.as_slice())
     }
+}*/
+
+pub struct PqNodeDistanceMeasure<'a> {
+    storage: &'a PqCompressionStorage<'a>,
+    vector: Vec<PqVectorElement>,
+}
+
+impl<'a> PqNodeDistanceMeasure<'a> {
+    pub unsafe fn with_index_pointer<T: StatsNodeRead>(
+        storage: &'a PqCompressionStorage,
+        index_pointer: IndexPointer,
+        stats: &mut T,
+    ) -> Self {
+        let rn = unsafe { Node::read(storage.index, index_pointer, stats) };
+        let node = rn.get_archived_node();
+        assert!(node.pq_vector.len() > 0);
+        let vector = node.pq_vector.as_slice().to_vec();
+        Self {
+            storage: storage,
+            vector: vector,
+        }
+    }
+}
+
+impl<'a> NodeDistanceMeasure for PqNodeDistanceMeasure<'a> {
+    unsafe fn get_distance<T: StatsNodeRead + StatsDistanceComparison>(
+        &self,
+        index_pointer: IndexPointer,
+        stats: &mut T,
+    ) -> f32 {
+        let rn1 = Node::read(self.storage.index, index_pointer, stats);
+        let node1 = rn1.get_archived_node();
+        self.storage
+            .quantizer
+            .get_distance_directly(self.vector.as_slice(), node1.pq_vector.as_slice())
+    }
 }
 
 pub struct PqCompressionStorage<'a> {