Introduce traits for node read/write from storage

pgvectorscale/pgvectorscale_derive/src/lib.rs

@@ -43,8 +43,9 @@ fn impl_readable_macro(ast: &syn::DeriveInput) -> TokenStream {
             }
         }
 
-        impl #name {
-            pub unsafe fn read<'a, 'b, S: crate::access_method::stats::StatsNodeRead>(index: &'a PgRelation, index_pointer: ItemPointer, stats: &'b mut S) -> #readable_name<'a> {
+        impl ReadableNode for #name {
+            type Node<'a> = #readable_name<'a>;
+            unsafe fn read<'a, 'b, S: crate::access_method::stats::StatsNodeRead>(index: &'a PgRelation, index_pointer: ItemPointer, stats: &'b mut S) -> Self::Node<'a> {
                 let rb = index_pointer.read_bytes(index);
                 stats.record_read();
                 #readable_name::with_readable_buffer(rb)
@@ -81,21 +82,22 @@ fn impl_writeable_macro(ast: &syn::DeriveInput) -> TokenStream {
             }
         }
 
-        impl #name {
-            pub unsafe fn modify<'a, 'b, S: crate::access_method::stats::StatsNodeModify>(index: &'a PgRelation, index_pointer: ItemPointer, stats: &'b mut S) -> #writeable_name<'a> {
+        impl WriteableNode for #name {
+            type Node<'a> = #writeable_name<'a>;
+            unsafe fn modify<'a, 'b, S: crate::access_method::stats::StatsNodeModify>(index: &'a PgRelation, index_pointer: ItemPointer, stats: &'b mut S) -> Self::Node<'a> {
                 let wb = index_pointer.modify_bytes(index);
                 stats.record_modify();
                 #writeable_name { wb: wb }
             }
 
-            pub fn write<S: crate::access_method::stats::StatsNodeWrite>(&self, tape: &mut crate::util::tape::Tape, stats: &mut S) -> ItemPointer {
+            fn write<S: crate::access_method::stats::StatsNodeWrite>(&self, tape: &mut crate::util::tape::Tape, stats: &mut S) -> ItemPointer {
                 //TODO 256 probably too small
                 let bytes = self.serialize_to_vec();
                 stats.record_write();
                 unsafe { tape.write(&bytes) }
             }
 
-            pub fn serialize_to_vec(&self) -> rkyv::util::AlignedVec {
+            fn serialize_to_vec(&self) -> rkyv::util::AlignedVec {
                 //TODO 256 probably too small
                 rkyv::to_bytes::<_, 256>(self).unwrap()
             }

pgvectorscale/src/access_method/debugging.rs

@@ -4,9 +4,10 @@ use std::collections::HashMap;
 
 use pgrx::PgRelation;
 
+use crate::access_method::node::ReadableNode;
 use crate::util::ItemPointer;
 
-use super::{plain_node::Node, stats::GreedySearchStats};
+use super::{plain_node::PlainNode, stats::GreedySearchStats};
 
 #[allow(dead_code)]
 pub fn print_graph_from_disk(index: &PgRelation, init_id: ItemPointer) {
@@ -25,7 +26,7 @@ unsafe fn print_graph_from_disk_visitor(
     sb: &mut String,
 ) {
     let mut stats = GreedySearchStats::new();
-    let data_node = Node::read(index, index_pointer, &mut stats);
+    let data_node = PlainNode::read(index, index_pointer, &mut stats);
     let node = data_node.get_archived_node();
     let v = node.vector.as_slice();
     let copy: Vec<f32> = v.to_vec();

pgvectorscale/src/access_method/meta_page.rs

@@ -4,17 +4,17 @@ use pgvectorscale_derive::{Readable, Writeable};
 use rkyv::{Archive, Deserialize, Serialize};
 use semver::Version;
 
-use crate::access_method::options::TSVIndexOptions;
-use crate::util::page;
-use crate::util::*;
-
 use super::distance::{DistanceFn, DistanceType};
 use super::options::{
     NUM_DIMENSIONS_DEFAULT_SENTINEL, NUM_NEIGHBORS_DEFAULT_SENTINEL,
     SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL,
 };
 use super::stats::StatsNodeModify;
 use super::storage::StorageType;
+use crate::access_method::node::{ReadableNode, WriteableNode};
+use crate::access_method::options::TSVIndexOptions;
+use crate::util::page;
+use crate::util::*;
 
 const TSV_MAGIC_NUMBER: u32 = 768756476; //Magic number, random
 const TSV_VERSION: u32 = 2;

pgvectorscale/src/access_method/mod.rs

@@ -7,6 +7,7 @@ mod graph_neighbor_store;
 pub mod guc;
 mod meta_page;
 mod neighbor_with_distance;
+mod node;
 pub mod options;
 pub mod pg_vector;
 mod plain_node;

pgvectorscale/src/access_method/node.rs

@@ -0,0 +1,29 @@
+use pgrx::PgRelation;
+use rkyv::AlignedVec;
+
+use crate::util::{tape::Tape, ItemPointer};
+
+use super::stats::{StatsNodeModify, StatsNodeRead, StatsNodeWrite};
+
+pub trait ReadableNode {
+    type Node<'a>;
+    unsafe fn read<'a, S: StatsNodeRead>(
+        index: &'a PgRelation,
+        index_pointer: ItemPointer,
+        stats: &mut S,
+    ) -> Self::Node<'a>;
+}
+
+pub trait WriteableNode {
+    type Node<'a>;
+
+    unsafe fn modify<'a, S: StatsNodeModify>(
+        index: &'a PgRelation,
+        index_pointer: ItemPointer,
+        stats: &mut S,
+    ) -> Self::Node<'a>;
+
+    fn write<S: StatsNodeWrite>(&self, tape: &mut Tape, stats: &mut S) -> ItemPointer;
+
+    fn serialize_to_vec(&self) -> AlignedVec;
+}

pgvectorscale/src/access_method/plain_node.rs

@@ -6,22 +6,22 @@ use pgvectorscale_derive::{Readable, Writeable};
 use rkyv::vec::ArchivedVec;
 use rkyv::{Archive, Deserialize, Serialize};
 
+use super::meta_page::MetaPage;
 use super::neighbor_with_distance::NeighborWithDistance;
 use super::storage::ArchivedData;
+use crate::access_method::node::{ReadableNode, WriteableNode};
 use crate::util::{ArchivedItemPointer, HeapPointer, ItemPointer, ReadableBuffer, WritableBuffer};
 
-use super::meta_page::MetaPage;
-
 #[derive(Archive, Deserialize, Serialize, Readable, Writeable)]
 #[archive(check_bytes)]
-pub struct Node {
+pub struct PlainNode {
     pub vector: Vec<f32>,
     pub pq_vector: Vec<u8>,
     neighbor_index_pointers: Vec<ItemPointer>,
     pub heap_item_pointer: HeapPointer,
 }
 
-impl Node {
+impl PlainNode {
     fn new_internal(
         vector: Vec<f32>,
         pq_vector: Vec<u8>,
@@ -52,7 +52,7 @@ impl Node {
 }
 
 /// contains helpers for mutate-in-place. See struct_mutable_refs in test_alloc.rs in rkyv
-impl ArchivedNode {
+impl ArchivedPlainNode {
     pub fn is_deleted(&self) -> bool {
         self.heap_item_pointer.offset == InvalidOffsetNumber
     }
@@ -106,9 +106,9 @@ impl ArchivedNode {
     }
 }
 
-impl ArchivedData for ArchivedNode {
-    fn with_data(data: &mut [u8]) -> Pin<&mut ArchivedNode> {
-        ArchivedNode::with_data(data)
+impl ArchivedData for ArchivedPlainNode {
+    fn with_data(data: &mut [u8]) -> Pin<&mut ArchivedPlainNode> {
+        ArchivedPlainNode::with_data(data)
     }
 
     fn get_index_pointer_to_neighbors(&self) -> Vec<ItemPointer> {

pgvectorscale/src/access_method/plain_storage.rs

@@ -4,7 +4,7 @@ use super::{
     graph_neighbor_store::GraphNeighborStore,
     neighbor_with_distance::DistanceWithTieBreak,
     pg_vector::PgVector,
-    plain_node::{ArchivedNode, Node, ReadableNode},
+    plain_node::{ArchivedPlainNode, PlainNode, ReadablePlainNode},
     stats::{
         GreedySearchStats, StatsDistanceComparison, StatsHeapNodeRead, StatsNodeModify,
         StatsNodeRead, StatsNodeWrite, WriteStats,
@@ -15,12 +15,12 @@ use super::{
 
 use pgrx::{PgBox, PgRelation};
 
+use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance};
+use crate::access_method::node::{ReadableNode, WriteableNode};
 use crate::util::{
     page::PageType, table_slot::TableSlot, tape::Tape, HeapPointer, IndexPointer, ItemPointer,
 };
 
-use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance};
-
 pub struct PlainStorage<'a> {
     pub index: &'a PgRelation,
     pub distance_fn: DistanceFn,
@@ -89,7 +89,7 @@ impl PlainDistanceMeasure {
 
 /* This is only applicable to plain, so keep here not in storage_common */
 pub struct IndexFullDistanceMeasure<'a> {
-    readable_node: ReadableNode<'a>,
+    readable_node: ReadablePlainNode<'a>,
     storage: &'a PlainStorage<'a>,
 }
 
@@ -99,7 +99,7 @@ impl<'a> IndexFullDistanceMeasure<'a> {
         index_pointer: IndexPointer,
         stats: &mut T,
     ) -> Self {
-        let rn = unsafe { Node::read(storage.index, index_pointer, stats) };
+        let rn = unsafe { PlainNode::read(storage.index, index_pointer, stats) };
         Self {
             readable_node: rn,
             storage,
@@ -108,7 +108,7 @@ impl<'a> IndexFullDistanceMeasure<'a> {
 
     pub unsafe fn with_readable_node(
         storage: &'a PlainStorage<'a>,
-        readable_node: ReadableNode<'a>,
+        readable_node: ReadablePlainNode<'a>,
     ) -> Self {
         Self {
             readable_node,
@@ -123,7 +123,7 @@ impl NodeDistanceMeasure for IndexFullDistanceMeasure<'_> {
         index_pointer: IndexPointer,
         stats: &mut T,
     ) -> f32 {
-        let rn1 = Node::read(self.storage.index, index_pointer, stats);
+        let rn1 = PlainNode::read(self.storage.index, index_pointer, stats);
         let rn2 = &self.readable_node;
         let node1 = rn1.get_archived_node();
         let node2 = rn2.get_archived_node();
@@ -144,7 +144,7 @@ pub struct PlainStorageLsnPrivateData {
 impl PlainStorageLsnPrivateData {
     pub fn new(
         index_pointer_to_node: IndexPointer,
-        node: &ArchivedNode,
+        node: &ArchivedPlainNode,
         gns: &GraphNeighborStore,
     ) -> Self {
         let heap_pointer = node.heap_item_pointer.deserialize_item_pointer();
@@ -165,7 +165,7 @@ impl Storage for PlainStorage<'_> {
         = IndexFullDistanceMeasure<'b>
     where
         Self: 'b;
-    type ArchivedType = ArchivedNode;
+    type ArchivedType = ArchivedPlainNode;
     type LSNPrivateData = PlainStorageLsnPrivateData;
 
     fn page_type() -> PageType {
@@ -181,7 +181,7 @@ impl Storage for PlainStorage<'_> {
         stats: &mut S,
     ) -> ItemPointer {
         //OPT: avoid the clone?
-        let node = Node::new_for_full_vector(full_vector.to_vec(), heap_pointer, meta_page);
+        let node = PlainNode::new_for_full_vector(full_vector.to_vec(), heap_pointer, meta_page);
         let index_pointer: IndexPointer = node.write(tape, stats);
         index_pointer
     }
@@ -199,7 +199,7 @@ impl Storage for PlainStorage<'_> {
         neighbors: &[NeighborWithDistance],
         stats: &mut S,
     ) {
-        let mut node = unsafe { Node::modify(self.index, index_pointer, stats) };
+        let mut node = unsafe { PlainNode::modify(self.index, index_pointer, stats) };
         let mut archived = node.get_archived_node();
         archived.as_mut().set_neighbors(neighbors, meta);
         node.commit();
@@ -252,7 +252,7 @@ impl Storage for PlainStorage<'_> {
         result: &mut Vec<NeighborWithDistance>,
         stats: &mut S,
     ) {
-        let rn = unsafe { Node::read(self.index, neighbors_of, stats) };
+        let rn = unsafe { PlainNode::read(self.index, neighbors_of, stats) };
         //get neighbors copy before givining ownership of rn to the distance state
         let neighbors: Vec<_> = rn.get_archived_node().iter_neighbors().collect();
         let dist_state = unsafe { IndexFullDistanceMeasure::with_readable_node(self, rn) };
@@ -277,7 +277,7 @@ impl Storage for PlainStorage<'_> {
             panic!("should not have had an init id already inserted");
         }
 
-        let rn = unsafe { Node::read(self.index, index_pointer, &mut lsr.stats) };
+        let rn = unsafe { PlainNode::read(self.index, index_pointer, &mut lsr.stats) };
         let node = rn.get_archived_node();
 
         let distance = match lsr.sdm.as_ref().unwrap() {
@@ -312,7 +312,7 @@ impl Storage for PlainStorage<'_> {
             }
 
             let rn_neighbor =
-                unsafe { Node::read(self.index, neighbor_index_pointer, &mut lsr.stats) };
+                unsafe { PlainNode::read(self.index, neighbor_index_pointer, &mut lsr.stats) };
             let node_neighbor = rn_neighbor.get_archived_node();
 
             let distance = match lsr.sdm.as_ref().unwrap() {
@@ -348,7 +348,7 @@ impl Storage for PlainStorage<'_> {
         neighbors: &[NeighborWithDistance],
         stats: &mut S,
     ) {
-        let mut node = unsafe { Node::modify(self.index, index_pointer, stats) };
+        let mut node = unsafe { PlainNode::modify(self.index, index_pointer, stats) };
         let mut archived = node.get_archived_node();
         archived.as_mut().set_neighbors(neighbors, meta);
         node.commit();

pgvectorscale/src/access_method/sbq.rs

@@ -17,15 +17,15 @@ use std::{cell::RefCell, collections::HashMap, iter::once, marker::PhantomData};
 use pgrx::{PgBox, PgRelation};
 use rkyv::{Archive, Deserialize, Serialize};
 
+use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance};
+use crate::access_method::node::{ReadableNode, WriteableNode};
 use crate::util::{
     chain::{ChainItemReader, ChainTapeWriter},
     page::{PageType, ReadablePage},
     table_slot::TableSlot,
     tape::Tape,
     HeapPointer, IndexPointer, ItemPointer, ReadableBuffer, WritableBuffer,
 };
-
-use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance};
 use pgvectorscale_derive::{Readable, Writeable};
 
 pub type SbqVectorElement = u64;

pgvectorscale/src/access_method/sbq_node.rs

@@ -1,3 +1,4 @@
+use crate::access_method::node::{ReadableNode, WriteableNode};
 use crate::access_method::PgRelation;
 use crate::util::{ArchivedItemPointer, HeapPointer, ItemPointer, ReadableBuffer, WritableBuffer};
 use pgrx::pg_sys::{InvalidBlockNumber, InvalidOffsetNumber, BLCKSZ};