From 8955cedcb707f4649bad34c1f551a770db9345bd Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Wed, 13 Dec 2023 17:19:27 -0500 Subject: [PATCH 01/44] Refactor to isolate quantization code --- timescale_vector/src/access_method/build.rs | 105 +++++++------ .../src/access_method/builder_graph.rs | 42 +++-- .../src/access_method/disk_index_graph.rs | 12 +- timescale_vector/src/access_method/graph.rs | 101 +++++------- .../src/access_method/meta_page.rs | 13 +- timescale_vector/src/access_method/mod.rs | 1 + timescale_vector/src/access_method/model.rs | 111 +++++++------- timescale_vector/src/access_method/pq.rs | 144 ++++++++++++------ .../src/access_method/quantizer.rs | 22 +++ timescale_vector/src/access_method/scan.rs | 24 ++- 10 files changed, 329 insertions(+), 246 deletions(-) create mode 100644 timescale_vector/src/access_method/quantizer.rs diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index 0408118f..4ef76622 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -1,7 +1,6 @@ use std::time::Instant; use pgrx::*; -use reductive::pq::Pq; use crate::access_method::disk_index_graph::DiskIndexGraph; use crate::access_method::graph::Graph; @@ -9,7 +8,6 @@ use crate::access_method::graph::InsertStats; use crate::access_method::graph::VectorProvider; use crate::access_method::model::PgVector; use crate::access_method::options::TSVIndexOptions; -use crate::access_method::pq::{PgPq, PqTrainer}; use crate::util::page; use crate::util::tape::Tape; use crate::util::*; @@ -17,6 +15,7 @@ use crate::util::*; use super::builder_graph::BuilderGraph; use super::meta_page::MetaPage; use super::model::{self}; +use super::quantizer::Quantizer; struct BuildState<'a, 'b> { memcxt: PgMemoryContexts, @@ -26,17 +25,24 @@ struct BuildState<'a, 'b> { node_builder: BuilderGraph<'b>, started: Instant, stats: InsertStats, - pq_trainer: Option, + quantizer: Quantizer, } impl<'a, 'b> BuildState<'a, 'b> { - fn new(index_relation: &'a PgRelation, meta_page: MetaPage, bg: BuilderGraph<'b>) -> Self { + fn new( + index_relation: &'a PgRelation, + meta_page: MetaPage, + bg: BuilderGraph<'b>, + mut quantizer: Quantizer, + ) -> Self { let tape = unsafe { Tape::new(index_relation, page::PageType::Node) }; - let pq = if meta_page.get_use_pq() { - Some(PqTrainer::new(&meta_page)) - } else { - None - }; + + match &mut quantizer { + Quantizer::None => {} + Quantizer::PQ(pq) => { + pq.start_training(&meta_page); + } + } //TODO: some ways to get rid of meta_page.clone? BuildState { memcxt: PgMemoryContexts::new("tsv build context"), @@ -46,7 +52,7 @@ impl<'a, 'b> BuildState<'a, 'b> { node_builder: bg, started: Instant::now(), stats: InsertStats::new(), - pq_trainer: pq, + quantizer: quantizer, } } } @@ -82,16 +88,7 @@ pub extern "C" fn ambuild( } assert!(dimensions > 0 && dimensions < 2000); let meta_page = unsafe { MetaPage::create(&index_relation, dimensions as _, opt.clone()) }; - let (ntuples, pq_opt) = do_heap_scan(index_info, &heap_relation, &index_relation, meta_page); - - // When using PQ, we initialize a node to store the model we use to quantize the vectors. - unsafe { - if opt.use_pq { - let pq = pq_opt.unwrap(); - let index_pointer: IndexPointer = model::write_pq(pq, &index_relation); - super::meta_page::MetaPage::update_pq_pointer(&index_relation, index_pointer) - } - } + let ntuples = do_heap_scan(index_info, &heap_relation, &index_relation, meta_page); let mut result = unsafe { PgBox::::alloc0() }; result.heap_tuples = ntuples as f64; @@ -122,24 +119,24 @@ pub unsafe extern "C" fn aminsert( let vector = (*vec).to_slice(); let heap_pointer = ItemPointer::with_item_pointer_data(*heap_tid); let meta_page = MetaPage::read(&index_relation); + + let mut quantizer = meta_page.get_quantizer(); + match &mut quantizer { + Quantizer::None => {} + Quantizer::PQ(pq) => { + pq.load(&index_relation, &meta_page); + } + } + let vp = VectorProvider::new( Some(&heap_relation), Some(get_attribute_number(index_info)), - meta_page.get_use_pq(), + &quantizer, false, ); let mut graph = DiskIndexGraph::new(&index_relation, vp); - let mut node = model::Node::new(vector.to_vec(), heap_pointer, &meta_page); - - // Populate the PQ version of the vector if it exists. - let pq = PgPq::new(&meta_page, &index_relation); - match pq { - None => {} - Some(pq) => { - node.pq_vector = pq.quantize(vector.to_vec()); - } - } + let node = model::Node::new(vector.to_vec(), heap_pointer, &meta_page, &quantizer); let mut tape = unsafe { Tape::new(&index_relation, page::PageType::Node) }; let index_pointer: IndexPointer = node.write(&mut tape); @@ -163,15 +160,18 @@ fn do_heap_scan<'a>( heap_relation: &'a PgRelation, index_relation: &'a PgRelation, meta_page: MetaPage, -) -> (usize, Option>) { +) -> usize { + let quantizer = meta_page.get_quantizer(); let vp = VectorProvider::new( Some(heap_relation), Some(get_attribute_number(index_info)), - meta_page.get_use_pq(), + &quantizer, false, ); + let bg = BuilderGraph::new(meta_page.clone(), vp); - let mut state = BuildState::new(index_relation, meta_page.clone(), bg); + let quantizer = meta_page.get_quantizer(); + let mut state = BuildState::new(index_relation, meta_page.clone(), bg, quantizer); unsafe { pg_sys::IndexBuildHeapScan( heap_relation.as_ptr(), @@ -183,9 +183,14 @@ fn do_heap_scan<'a>( } // we train the quantizer and add prepare to write quantized values to the nodes. - let pq = state.pq_trainer.map(|pq| pq.train_pq()); + match &mut state.quantizer { + Quantizer::None => {} + Quantizer::PQ(pq) => { + pq.finish_training(); + } + } - let write_stats = unsafe { state.node_builder.write(index_relation, &pq) }; + let write_stats = unsafe { state.node_builder.write(index_relation, &state.quantizer) }; assert_eq!(write_stats.num_nodes, state.ntuples); let writing_took = Instant::now() @@ -210,7 +215,15 @@ fn do_heap_scan<'a>( let ntuples = state.ntuples; warning!("Indexed {} tuples", ntuples); - (ntuples, pq) + + match state.quantizer { + Quantizer::None => {} + Quantizer::PQ(pq) => { + pq.write_metadata(index_relation); + } + } + + ntuples } #[pg_guard] @@ -261,16 +274,20 @@ fn build_callback_internal( ); } - match &state.pq_trainer { - Some(_) => { - let pqt = state.pq_trainer.as_mut(); - pqt.expect("error adding sample") - .add_sample(vector.to_vec()) + match &mut state.quantizer { + Quantizer::None => {} + Quantizer::PQ(pq) => { + pq.add_sample(vector); } - None => {} } - let node = model::Node::new(vector.to_vec(), heap_pointer, &state.meta_page); + let node = model::Node::new( + vector.to_vec(), + heap_pointer, + &state.meta_page, + &state.quantizer, + ); + let index_pointer: IndexPointer = node.write(&mut state.tape); let new_stats = state.node_builder.insert(&index, index_pointer, vector); state.stats.combine(new_stats); diff --git a/timescale_vector/src/access_method/builder_graph.rs b/timescale_vector/src/access_method/builder_graph.rs index 8a050b03..6d1770d6 100644 --- a/timescale_vector/src/access_method/builder_graph.rs +++ b/timescale_vector/src/access_method/builder_graph.rs @@ -1,15 +1,14 @@ use std::collections::HashMap; use std::time::Instant; -use ndarray::Array1; use pgrx::*; -use reductive::pq::{Pq, QuantizeVector}; use crate::util::{IndexPointer, ItemPointer}; use super::graph::{Graph, VectorProvider}; use super::meta_page::MetaPage; use super::model::*; +use super::quantizer::Quantizer; /// A builderGraph is a graph that keep the neighbors in-memory in the neighbor_map below /// The idea is that during the index build, you don't want to update the actual Postgres @@ -32,20 +31,14 @@ impl<'a> BuilderGraph<'a> { } } - unsafe fn get_pq_vector( - &self, - index: &PgRelation, - index_pointer: ItemPointer, - pq: &Pq, - ) -> Vec { + unsafe fn get_full_vector(&self, heap_pointer: ItemPointer) -> Vec { let vp = self.get_vector_provider(); - let copy = vp.get_full_vector_copy_from_heap(index, index_pointer); - let og_vec = Array1::from(copy); - pq.quantize_vector(og_vec).to_vec() + vp.get_full_vector_copy_from_heap_pointer(heap_pointer) } - pub unsafe fn write(&self, index: &PgRelation, pq: &Option>) -> WriteStats { + pub unsafe fn write(&self, index: &PgRelation, quantizer: &Quantizer) -> WriteStats { let mut stats = WriteStats::new(); + let meta = self.get_meta_page(index); //TODO: OPT: do this in order of item pointers for (index_pointer, neighbors) in &self.neighbor_map { @@ -62,17 +55,22 @@ impl<'a> BuilderGraph<'a> { }; stats.num_neighbors += neighbors.len(); - let pqv = match pq { - Some(pq) => Some(self.get_pq_vector(index, *index_pointer, pq)), - None => None, + let node = Node::modify(index, *index_pointer); + let mut archived = node.get_archived_node(); + archived.as_mut().set_neighbors(neighbors, &meta); + + match quantizer { + Quantizer::None => {} + Quantizer::PQ(pq) => { + let heap_pointer = node + .get_archived_node() + .heap_item_pointer + .deserialize_item_pointer(); + let full_vector = self.get_full_vector(heap_pointer); + pq.update_node_after_traing(&mut archived, full_vector); + } }; - Node::update_neighbors_and_pq( - index, - *index_pointer, - neighbors, - self.get_meta_page(index), - pqv, - ); + node.commit(); } stats } diff --git a/timescale_vector/src/access_method/disk_index_graph.rs b/timescale_vector/src/access_method/disk_index_graph.rs index c3bd6eea..ac63b54a 100644 --- a/timescale_vector/src/access_method/disk_index_graph.rs +++ b/timescale_vector/src/access_method/disk_index_graph.rs @@ -86,14 +86,12 @@ impl<'h> Graph for DiskIndexGraph<'h> { MetaPage::update_init_ids(index, vec![neighbors_of]); self.meta_page = MetaPage::read(index); } + unsafe { - Node::update_neighbors_and_pq( - index, - neighbors_of, - &new_neighbors, - self.get_meta_page(index), - None, - ); + let node = Node::modify(index, neighbors_of); + let archived = node.get_archived_node(); + archived.set_neighbors(&new_neighbors, &self.meta_page); + node.commit(); } } } diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index 05894324..190dc3c4 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -4,11 +4,11 @@ use pgrx::pg_sys::{Datum, TupleTableSlot}; use pgrx::{pg_sys, PgBox, PgRelation}; use crate::access_method::model::Node; -use crate::access_method::pq::{DistanceCalculator, PgPq}; use crate::util::ports::slot_getattr; use crate::util::{HeapPointer, IndexPointer, ItemPointer}; use super::model::PgVector; +use super::quantizer::Quantizer; use super::{ meta_page::MetaPage, model::{NeighborWithDistance, ReadableNode}, @@ -59,8 +59,8 @@ impl Drop for TableSlot { #[derive(Clone)] pub struct VectorProvider<'a> { - pq_enabled: bool, - calc_distance_with_pq: bool, + quantizer: &'a Quantizer, + calc_distance_with_quantizer: bool, heap_rel: Option<&'a PgRelation>, heap_attr_number: Option, distance_fn: fn(&[f32], &[f32]) -> f32, @@ -70,24 +70,22 @@ impl<'a> VectorProvider<'a> { pub fn new( heap_rel: Option<&'a PgRelation>, heap_attr_number: Option, - pq_enabled: bool, - calc_distance_with_pq: bool, + quantizer: &'a Quantizer, + calc_distance_with_quantizer: bool, ) -> Self { Self { - pq_enabled, - calc_distance_with_pq, + quantizer, + calc_distance_with_quantizer, heap_rel, heap_attr_number, distance_fn: distance, } } - pub unsafe fn get_full_vector_copy_from_heap( + pub unsafe fn get_full_vector_copy_from_heap_pointer( &self, - index: &PgRelation, - index_pointer: ItemPointer, + heap_pointer: ItemPointer, ) -> Vec { - let heap_pointer = self.get_heap_pointer(index, index_pointer); let slot = TableSlot::new(self.heap_rel.unwrap()); self.init_slot(&slot, heap_pointer); let slice = self.get_slice(&slot); @@ -124,6 +122,10 @@ impl<'a> VectorProvider<'a> { heap_pointer } + fn get_distance_measure(&self, query: &[f32]) -> DistanceMeasure { + return DistanceMeasure::new(self.quantizer, query, self.calc_distance_with_quantizer); + } + unsafe fn get_distance( &self, index: &PgRelation, @@ -132,20 +134,20 @@ impl<'a> VectorProvider<'a> { dm: &DistanceMeasure, stats: &mut GreedySearchStats, ) -> (f32, HeapPointer) { - if self.calc_distance_with_pq { + if self.calc_distance_with_quantizer { let rn = unsafe { Node::read(index, index_pointer) }; stats.node_reads += 1; let node = rn.get_archived_node(); assert!(node.pq_vector.len() > 0); let vec = node.pq_vector.as_slice(); - let distance = dm.get_pq_distance(vec); + let distance = dm.get_quantized_distance(vec); stats.pq_distance_comparisons += 1; stats.distance_comparisons += 1; return (distance, node.heap_item_pointer.deserialize_item_pointer()); } //now we know we're doing a distance calc on the full-sized vector - if self.pq_enabled { + if self.quantizer.is_some() { //have to get it from the heap let heap_pointer = self.get_heap_pointer(index, index_pointer); stats.node_reads += 1; @@ -172,7 +174,7 @@ impl<'a> VectorProvider<'a> { index: &'i PgRelation, index_pointer: IndexPointer, ) -> FullVectorDistanceState<'i> { - if self.pq_enabled { + if self.quantizer.is_some() { let heap_pointer = self.get_heap_pointer(index, index_pointer); let slot = TableSlot::new(self.heap_rel.unwrap()); self.init_slot(&slot, heap_pointer); @@ -195,7 +197,7 @@ impl<'a> VectorProvider<'a> { index: &PgRelation, index_pointer: IndexPointer, ) -> f32 { - if self.pq_enabled { + if self.quantizer.is_some() { let heap_pointer = self.get_heap_pointer(index, index_pointer); let slot = TableSlot::new(self.heap_rel.unwrap()); self.init_slot(&slot, heap_pointer); @@ -222,40 +224,37 @@ pub struct FullVectorDistanceState<'a> { } pub struct DistanceMeasure { - distance_calculator: Option, - //query: Option<&[f32]> + pq_distance_table: Option, } impl DistanceMeasure { - pub fn new( - index: &PgRelation, - meta_page: &MetaPage, - query: &[f32], - calc_distance_with_pq: bool, - ) -> Self { - let use_pq = meta_page.get_use_pq(); - if calc_distance_with_pq { - assert!(use_pq); - let pq = PgPq::new(meta_page, index); - let dc = pq.unwrap().distance_calculator(query, distance); + pub fn new(quantizer: &Quantizer, query: &[f32], calc_distance_with_quantizer: bool) -> Self { + if !calc_distance_with_quantizer { return Self { - distance_calculator: Some(dc), + pq_distance_table: None, }; } - - return Self { - distance_calculator: None, - }; + match quantizer { + Quantizer::None => Self { + pq_distance_table: None, + }, + Quantizer::PQ(pq) => { + let dc = pq.get_distance_table(query, distance); + Self { + pq_distance_table: Some(dc), + } + } + } } - fn get_pq_distance(&self, vec: &[u8]) -> f32 { - let dc = self.distance_calculator.as_ref().unwrap(); + fn get_quantized_distance(&self, vec: &[u8]) -> f32 { + let dc = self.pq_distance_table.as_ref().unwrap(); let distance = dc.distance(vec); distance } fn get_full_vector_distance(&self, vec: &[f32], query: &[f32]) -> f32 { - assert!(self.distance_calculator.is_none()); + assert!(self.pq_distance_table.is_none()); distance(vec, query) } } @@ -305,7 +304,7 @@ impl ListSearchResult { inserted: HashSet::new(), max_history_size: None, dm: DistanceMeasure { - distance_calculator: None, + pq_distance_table: None, }, stats: GreedySearchStats::new(), } @@ -422,16 +421,6 @@ pub trait Graph { fn get_vector_provider(&self) -> VectorProvider; - fn get_distance_measure( - &self, - index: &PgRelation, - query: &[f32], - calc_distance_with_pq: bool, - ) -> DistanceMeasure { - let meta_page = self.get_meta_page(index); - return DistanceMeasure::new(index, meta_page, query, calc_distance_with_pq); - } - fn set_neighbors( &mut self, index: &PgRelation, @@ -467,13 +456,8 @@ pub trait Graph { //no nodes in the graph return (ListSearchResult::empty(), None); } - let dm = { - self.get_distance_measure( - index, - query, - self.get_vector_provider().calc_distance_with_pq, - ) - }; + let dm = self.get_vector_provider().get_distance_measure(query); + let mut l = ListSearchResult::new( index, Some(search_list_size), @@ -498,11 +482,8 @@ pub trait Graph { //no nodes in the graph return ListSearchResult::empty(); } - let dm = self.get_distance_measure( - index, - query, - self.get_vector_provider().calc_distance_with_pq, - ); + let dm = self.get_vector_provider().get_distance_measure(query); + ListSearchResult::new(index, None, self, init_ids.unwrap(), query, dm) } diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index 2d68f8dc..256cf440 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -5,6 +5,9 @@ use crate::access_method::options::TSVIndexOptions; use crate::util::page; use crate::util::*; +use super::pq::PqQuantizer; +use super::quantizer::Quantizer; + const TSV_MAGIC_NUMBER: u32 = 768756476; //Magic number, random const TSV_VERSION: u32 = 1; const GRAPH_SLACK_FACTOR: f64 = 1.3_f64; @@ -55,10 +58,18 @@ impl MetaPage { self.max_alpha } - pub fn get_use_pq(&self) -> bool { + fn get_use_pq(&self) -> bool { self.use_pq } + pub fn get_quantizer(&self) -> Quantizer { + if self.get_use_pq() { + Quantizer::PQ(PqQuantizer::new()) + } else { + Quantizer::None + } + } + pub fn get_max_neighbors_during_build(&self) -> usize { return ((self.get_num_neighbors() as f64) * GRAPH_SLACK_FACTOR).ceil() as usize; } diff --git a/timescale_vector/src/access_method/mod.rs b/timescale_vector/src/access_method/mod.rs index 77d8e607..83dea8fa 100644 --- a/timescale_vector/src/access_method/mod.rs +++ b/timescale_vector/src/access_method/mod.rs @@ -9,6 +9,7 @@ pub mod guc; mod meta_page; mod model; pub mod options; +mod quantizer; mod scan; mod vacuum; diff --git a/timescale_vector/src/access_method/model.rs b/timescale_vector/src/access_method/model.rs index 4de0b607..7c168061 100644 --- a/timescale_vector/src/access_method/model.rs +++ b/timescale_vector/src/access_method/model.rs @@ -16,6 +16,7 @@ use crate::util::{ }; use super::meta_page::MetaPage; +use super::quantizer::Quantizer; //Ported from pg_vector code #[repr(C)] @@ -93,26 +94,35 @@ impl<'a> WritableNode<'a> { } impl Node { - pub fn new(vector: Vec, heap_item_pointer: ItemPointer, meta_page: &MetaPage) -> Self { + pub fn new( + full_vector: Vec, + heap_item_pointer: ItemPointer, + meta_page: &MetaPage, + quantizer: &Quantizer, + ) -> Self { let num_neighbors = meta_page.get_num_neighbors(); - let (vector, pq_vector) = if meta_page.get_use_pq() { - let pq_vec_len = meta_page.get_pq_vector_length(); - ( - Vec::with_capacity(0), - (0..pq_vec_len).map(|_| 0u8).collect(), - ) - } else { - (vector, Vec::with_capacity(0)) - }; - Self { - vector, - // always use vectors of num_clusters on length because we never want the serialized size of a Node to change - pq_vector, - // always use vectors of num_neighbors on length because we never want the serialized size of a Node to change - neighbor_index_pointers: (0..num_neighbors) - .map(|_| ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber)) - .collect(), - heap_item_pointer, + // always use vectors of num_neighbors in length because we never want the serialized size of a Node to change + let neighbor_index_pointers: Vec<_> = (0..num_neighbors) + .map(|_| ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber)) + .collect(); + + match quantizer { + Quantizer::None => Self { + vector: full_vector, + pq_vector: Vec::with_capacity(0), + neighbor_index_pointers: neighbor_index_pointers, + heap_item_pointer, + }, + Quantizer::PQ(pq) => { + let mut node = Self { + vector: Vec::with_capacity(0), + pq_vector: Vec::with_capacity(0), + neighbor_index_pointers: neighbor_index_pointers, + heap_item_pointer, + }; + pq.initialize_node(&mut node, meta_page, full_vector); + node + } } } @@ -126,46 +136,6 @@ impl Node { WritableNode { wb: wb } } - pub unsafe fn update_neighbors_and_pq( - index: &PgRelation, - index_pointer: ItemPointer, - neighbors: &Vec, - meta_page: &MetaPage, - vector: Option>, - ) { - let node = Node::modify(index, index_pointer); - let mut archived = node.get_archived_node(); - for (i, new_neighbor) in neighbors.iter().enumerate() { - //TODO: why do we need to recreate the archive? - let mut a_index_pointer = archived.as_mut().neighbor_index_pointer().index_pin(i); - //TODO hate that we have to set each field like this - a_index_pointer.block_number = - new_neighbor.get_index_pointer_to_neighbor().block_number; - a_index_pointer.offset = new_neighbor.get_index_pointer_to_neighbor().offset; - } - //set the marker that the list ended - if neighbors.len() < meta_page.get_num_neighbors() as _ { - //TODO: why do we need to recreate the archive? - let archived = node.get_archived_node(); - let mut past_last_index_pointers = - archived.neighbor_index_pointer().index_pin(neighbors.len()); - past_last_index_pointers.block_number = InvalidBlockNumber; - past_last_index_pointers.offset = InvalidOffsetNumber; - } - - match vector { - Some(v) => { - assert!(v.len() == archived.pq_vector.len()); - for i in 0..=v.len() - 1 { - let mut pgv = archived.as_mut().pq_vectors().index_pin(i); - *pgv = v[i]; - } - } - None => {} - } - - node.commit() - } pub fn write(&self, tape: &mut Tape) -> ItemPointer { let bytes = rkyv::to_bytes::<_, 256>(self).unwrap(); unsafe { tape.write(&bytes) } @@ -216,6 +186,27 @@ impl ArchivedNode { f(neighbor); } } + + pub fn set_neighbors( + mut self: Pin<&mut Self>, + neighbors: &Vec, + meta_page: &MetaPage, + ) { + for (i, new_neighbor) in neighbors.iter().enumerate() { + let mut a_index_pointer = self.as_mut().neighbor_index_pointer().index_pin(i); + //TODO hate that we have to set each field like this + a_index_pointer.block_number = + new_neighbor.get_index_pointer_to_neighbor().block_number; + a_index_pointer.offset = new_neighbor.get_index_pointer_to_neighbor().offset; + } + //set the marker that the list ended + if neighbors.len() < meta_page.get_num_neighbors() as _ { + let mut past_last_index_pointers = + self.neighbor_index_pointer().index_pin(neighbors.len()); + past_last_index_pointers.block_number = InvalidBlockNumber; + past_last_index_pointers.offset = InvalidOffsetNumber; + } + } } //TODO is this right? @@ -381,7 +372,7 @@ pub unsafe fn read_pq(index: &PgRelation, index_pointer: &IndexPointer) -> Pq, index: &PgRelation) -> ItemPointer { +pub unsafe fn write_pq(pq: &Pq, index: &PgRelation) -> ItemPointer { let vec = pq.subquantizers().to_slice_memory_order().unwrap().to_vec(); let shape = pq.subquantizers().dim(); let mut pq_node = PqQuantizerDef::new(shape.0, shape.1, shape.2, vec.len()); diff --git a/timescale_vector/src/access_method/pq.rs b/timescale_vector/src/access_method/pq.rs index 9e5cb8fc..d0dac299 100644 --- a/timescale_vector/src/access_method/pq.rs +++ b/timescale_vector/src/access_method/pq.rs @@ -1,9 +1,16 @@ +use std::pin::Pin; + use ndarray::{Array1, Array2, Axis}; use pgrx::{error, notice, PgRelation}; use rand::Rng; use reductive::pq::{Pq, QuantizeVector, TrainPq}; -use crate::access_method::model::read_pq; +use crate::{ + access_method::model::{self, read_pq}, + util::IndexPointer, +}; + +use super::meta_page::MetaPage; /// pq aka Product quantization (PQ) is one of the most widely used algorithms for memory-efficient approximated nearest neighbor search, /// This module encapsulates a vanilla implementation of PQ that we use for the vector index. @@ -50,16 +57,16 @@ impl PqTrainer { /// add_sample adds vectors to the training set via uniform reservoir sampling to keep the /// number of vectors within a reasonable memory limit. - pub fn add_sample(&mut self, sample: Vec) { + pub fn add_sample(&mut self, sample: &[f32]) { if self.training_set.len() >= NUM_TRAINING_SET_SIZE { // TODO: Cache this somehow. let mut rng = rand::thread_rng(); let index = rng.gen_range(0..self.considered_samples + 1); if index < NUM_TRAINING_SET_SIZE { - self.training_set[index] = sample; + self.training_set[index] = sample.to_vec(); } } else { - self.training_set.push(sample); + self.training_set.push(sample.to_vec()); } self.considered_samples += 1; } @@ -91,42 +98,6 @@ impl PqTrainer { } } -/// PgPq encapsulates functions to work with PQ. -pub struct PgPq { - pq: Pq, -} - -impl PgPq { - pub fn new( - meta_page: &super::meta_page::MetaPage, - index_relation: &PgRelation, - ) -> Option { - if !meta_page.get_use_pq() { - return None; - } - let pq_id = meta_page.get_pq_pointer(); - match pq_id { - None => None, - Some(pq_id) => { - let pq = unsafe { read_pq(&index_relation, &pq_id) }; - Some(PgPq { pq }) - } - } - } - /// quantize produces a quantized vector from the raw pg vector. - pub fn quantize(self, vector: Vec) -> Vec { - let og_vec = Array1::from(vector.to_vec()); - self.pq.quantize_vector(og_vec).to_vec() - } - pub fn distance_calculator( - self, - query: &[f32], - distance_fn: fn(&[f32], &[f32]) -> f32, - ) -> DistanceCalculator { - DistanceCalculator::new(&self.pq, distance_fn, query) - } -} - /// build_distance_table produces an Asymmetric Distance Table to quickly compute distances. /// We compute the distance from every centroid and cache that so actual distance calculations /// can be fast. @@ -158,18 +129,103 @@ fn build_distance_table( distance_table } +pub struct PqQuantizer { + pq_trainer: Option, + pq: Option>, +} + +impl PqQuantizer { + pub fn new() -> PqQuantizer { + Self { + pq_trainer: None, + pq: None, + } + } + + pub fn load(&mut self, index_relation: &PgRelation, meta_page: &super::meta_page::MetaPage) { + assert!(self.pq_trainer.is_none()); + let pq_item_pointer = meta_page.get_pq_pointer().unwrap(); + self.pq = unsafe { Some(read_pq(&index_relation, &pq_item_pointer)) }; + } + + pub fn initialize_node( + &self, + node: &mut super::model::Node, + meta_page: &MetaPage, + full_vector: Vec, + ) { + if self.pq_trainer.is_some() { + let pq_vec_len = meta_page.get_pq_vector_length(); + node.pq_vector = (0..pq_vec_len).map(|_| 0u8).collect(); + } else { + assert!(self.pq.is_some()); + let pq_vec_len = meta_page.get_pq_vector_length(); + node.pq_vector = self.quantize(full_vector); + assert!(node.pq_vector.len() == pq_vec_len); + } + } + + pub fn update_node_after_traing( + &self, + archived: &mut Pin<&mut super::model::ArchivedNode>, + full_vector: Vec, + ) { + let pq_vector = self.quantize(full_vector); + + assert!(pq_vector.len() == archived.pq_vector.len()); + for i in 0..=pq_vector.len() - 1 { + let mut pgv = archived.as_mut().pq_vectors().index_pin(i); + *pgv = pq_vector[i]; + } + } + + pub fn start_training(&mut self, meta_page: &super::meta_page::MetaPage) { + self.pq_trainer = Some(PqTrainer::new(meta_page)); + } + + pub fn add_sample(&mut self, sample: &[f32]) { + self.pq_trainer.as_mut().unwrap().add_sample(sample); + } + + pub fn finish_training(&mut self) { + self.pq = Some(self.pq_trainer.take().unwrap().train_pq()); + } + + pub fn write_metadata(&self, index: &PgRelation) { + assert!(self.pq.is_some()); + let index_pointer: IndexPointer = + unsafe { model::write_pq(self.pq.as_ref().unwrap(), &index) }; + super::meta_page::MetaPage::update_pq_pointer(&index, index_pointer); + } + + pub fn quantize(&self, full_vector: Vec) -> Vec { + assert!(self.pq.is_some()); + let pq = self.pq.as_ref().unwrap(); + let array_vec = Array1::from(full_vector); + pq.quantize_vector(array_vec).to_vec() + } + + pub fn get_distance_table( + &self, + query: &[f32], + distance_fn: fn(&[f32], &[f32]) -> f32, + ) -> PqDistanceTable { + PqDistanceTable::new(&self.pq.as_ref().unwrap(), distance_fn, query) + } +} + /// DistanceCalculator encapsulates the code to generate distances between a PQ vector and a query. -pub struct DistanceCalculator { +pub struct PqDistanceTable { distance_table: Vec, } -impl DistanceCalculator { +impl PqDistanceTable { pub fn new( pq: &Pq, distance_fn: fn(&[f32], &[f32]) -> f32, query: &[f32], - ) -> DistanceCalculator { - DistanceCalculator { + ) -> PqDistanceTable { + PqDistanceTable { distance_table: build_distance_table(pq, query, distance_fn), } } diff --git a/timescale_vector/src/access_method/quantizer.rs b/timescale_vector/src/access_method/quantizer.rs new file mode 100644 index 00000000..5dee6053 --- /dev/null +++ b/timescale_vector/src/access_method/quantizer.rs @@ -0,0 +1,22 @@ +use super::pq::PqQuantizer; + +/*pub trait Quantizer { + fn initialize_node(&self, node: &mut Node, meta_page: &MetaPage); + fn start_training(&mut self, meta_page: &super::meta_page::MetaPage); + fn add_sample(&mut self, sample: Vec); + fn finish_training(&mut self); +}*/ + +pub enum Quantizer { + PQ(PqQuantizer), + None, +} + +impl Quantizer { + pub fn is_some(&self) -> bool { + match self { + Quantizer::None => false, + _ => true, + } + } +} diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 58e6ed77..5acf2832 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -8,7 +8,7 @@ use crate::{ util::{buffer::PinnedBufferShare, HeapPointer}, }; -use super::graph::ListSearchResult; +use super::{graph::ListSearchResult, quantizer::Quantizer}; struct TSVResponseIterator<'a> { query: Vec, @@ -16,14 +16,21 @@ struct TSVResponseIterator<'a> { search_list_size: usize, current: usize, last_buffer: Option>, + quantizer: Quantizer, } impl<'a> TSVResponseIterator<'a> { fn new(index: &PgRelation, query: &[f32], search_list_size: usize) -> Self { let meta_page = MetaPage::read(&index); - let use_pq = meta_page.get_use_pq(); - let mut graph = - DiskIndexGraph::new(&index, VectorProvider::new(None, None, use_pq, use_pq)); + let mut quantizer = meta_page.get_quantizer(); + match &mut quantizer { + Quantizer::None => {} + Quantizer::PQ(pq) => pq.load(index, &meta_page), + } + let mut graph = DiskIndexGraph::new( + &index, + VectorProvider::new(None, None, &quantizer, quantizer.is_some()), + ); use super::graph::Graph; let lsr = graph.greedy_search_streaming_init(&index, query); Self { @@ -32,16 +39,17 @@ impl<'a> TSVResponseIterator<'a> { lsr, current: 0, last_buffer: None, + quantizer, } } } impl<'a> TSVResponseIterator<'a> { fn next(&mut self, index: &'a PgRelation) -> Option { - let meta_page = MetaPage::read(&index); - let use_pq = meta_page.get_use_pq(); - let mut graph = - DiskIndexGraph::new(&index, VectorProvider::new(None, None, use_pq, use_pq)); + let mut graph = DiskIndexGraph::new( + &index, + VectorProvider::new(None, None, &self.quantizer, self.quantizer.is_some()), + ); use super::graph::Graph; /* Iterate until we find a non-deleted tuple */ From 041db69d72abd7377391cbc50292f4c8f52b7897 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Thu, 14 Dec 2023 09:42:39 -0500 Subject: [PATCH 02/44] Optimize: don't reread node for neighbor list --- .../src/access_method/builder_graph.rs | 19 ++--- .../src/access_method/disk_index_graph.rs | 17 ++--- timescale_vector/src/access_method/graph.rs | 74 +++++++++---------- 3 files changed, 49 insertions(+), 61 deletions(-) diff --git a/timescale_vector/src/access_method/builder_graph.rs b/timescale_vector/src/access_method/builder_graph.rs index 6d1770d6..2c31a8fe 100644 --- a/timescale_vector/src/access_method/builder_graph.rs +++ b/timescale_vector/src/access_method/builder_graph.rs @@ -86,21 +86,14 @@ impl<'a> Graph for BuilderGraph<'a> { self.meta_page.get_init_ids() } - fn get_neighbors( - &self, - _index: &PgRelation, - neighbors_of: ItemPointer, - result: &mut Vec, - ) -> bool { + fn get_neighbors(&self, _node: &ArchivedNode, neighbors_of: ItemPointer) -> Vec { let neighbors = self.neighbor_map.get(&neighbors_of); match neighbors { - Some(n) => { - for nwd in n { - result.push(nwd.get_index_pointer_to_neighbor()); - } - true - } - None => false, + Some(n) => n + .iter() + .map(|n| n.get_index_pointer_to_neighbor()) + .collect(), + None => vec![], } } diff --git a/timescale_vector/src/access_method/disk_index_graph.rs b/timescale_vector/src/access_method/disk_index_graph.rs index ac63b54a..20b241ad 100644 --- a/timescale_vector/src/access_method/disk_index_graph.rs +++ b/timescale_vector/src/access_method/disk_index_graph.rs @@ -1,11 +1,11 @@ use pgrx::PgRelation; -use crate::util::{IndexPointer, ItemPointer}; +use crate::util::ItemPointer; use super::{ graph::{Graph, VectorProvider}, meta_page::MetaPage, - model::{NeighborWithDistance, Node, ReadableNode}, + model::{ArchivedNode, NeighborWithDistance, Node, ReadableNode}, }; pub struct DiskIndexGraph<'a> { @@ -36,18 +36,13 @@ impl<'h> Graph for DiskIndexGraph<'h> { self.meta_page.get_init_ids() } - fn get_neighbors( - &self, - index: &PgRelation, - neighbors_of: ItemPointer, - result: &mut Vec, - ) -> bool { - let rn = self.read(index, neighbors_of); - rn.get_archived_node().apply_to_neighbors(|n| { + fn get_neighbors(&self, node: &ArchivedNode, _neighbors_of: ItemPointer) -> Vec { + let mut result = Vec::with_capacity(node.num_neighbors()); + node.apply_to_neighbors(|n| { let n = n.deserialize_item_pointer(); result.push(n) }); - true + result } fn get_neighbors_with_distances( diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index 190dc3c4..81b61e00 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -7,7 +7,7 @@ use crate::access_method::model::Node; use crate::util::ports::slot_getattr; use crate::util::{HeapPointer, IndexPointer, ItemPointer}; -use super::model::PgVector; +use super::model::{ArchivedNode, PgVector}; use super::quantizer::Quantizer; use super::{ meta_page::MetaPage, @@ -128,44 +128,37 @@ impl<'a> VectorProvider<'a> { unsafe fn get_distance( &self, - index: &PgRelation, - index_pointer: IndexPointer, + node: &ArchivedNode, query: &[f32], dm: &DistanceMeasure, stats: &mut GreedySearchStats, - ) -> (f32, HeapPointer) { + ) -> f32 { if self.calc_distance_with_quantizer { - let rn = unsafe { Node::read(index, index_pointer) }; - stats.node_reads += 1; - let node = rn.get_archived_node(); assert!(node.pq_vector.len() > 0); let vec = node.pq_vector.as_slice(); let distance = dm.get_quantized_distance(vec); stats.pq_distance_comparisons += 1; stats.distance_comparisons += 1; - return (distance, node.heap_item_pointer.deserialize_item_pointer()); + return distance; } //now we know we're doing a distance calc on the full-sized vector if self.quantizer.is_some() { //have to get it from the heap - let heap_pointer = self.get_heap_pointer(index, index_pointer); - stats.node_reads += 1; + let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); let slot = TableSlot::new(self.heap_rel.unwrap()); self.init_slot(&slot, heap_pointer); let slice = self.get_slice(&slot); + let distance = dm.get_full_vector_distance(slice, query); stats.distance_comparisons += 1; - return (dm.get_full_vector_distance(slice, query), heap_pointer); + return distance; } else { //have to get it from the index - let rn = unsafe { Node::read(index, index_pointer) }; - stats.node_reads += 1; - let node = rn.get_archived_node(); assert!(node.vector.len() > 0); let vec = node.vector.as_slice(); let distance = dm.get_full_vector_distance(vec, query); stats.distance_comparisons += 1; - return (distance, node.heap_item_pointer.deserialize_item_pointer()); + return distance; } } @@ -262,6 +255,7 @@ impl DistanceMeasure { struct ListSearchNeighbor { index_pointer: IndexPointer, heap_pointer: HeapPointer, + neighbor_index_pointers: Vec, distance: f32, visited: bool, } @@ -279,10 +273,16 @@ impl PartialEq for ListSearchNeighbor { } impl ListSearchNeighbor { - pub fn new(index_pointer: IndexPointer, heap_pointer: HeapPointer, distance: f32) -> Self { + pub fn new( + index_pointer: IndexPointer, + heap_pointer: HeapPointer, + distance: f32, + neighbor_index_pointers: Vec, + ) -> Self { Self { index_pointer, heap_pointer, + neighbor_index_pointers, distance, visited: false, } @@ -349,12 +349,21 @@ impl ListSearchResult { return; } + let rn = unsafe { Node::read(index, index_pointer) }; + self.stats.node_reads += 1; + let node = rn.get_archived_node(); + let vp = graph.get_vector_provider(); - let (dist, heap_pointer) = - unsafe { vp.get_distance(index, index_pointer, query, &self.dm, &mut self.stats) }; + let distance = unsafe { vp.get_distance(node, query, &self.dm, &mut self.stats) }; - let neighbor = ListSearchNeighbor::new(index_pointer, heap_pointer, dist); - self._insert_neighbor(neighbor); + let neighbors = graph.get_neighbors(node, index_pointer); + let lsn = ListSearchNeighbor::new( + index_pointer, + node.heap_item_pointer.deserialize_item_pointer(), + distance, + neighbors, + ); + self._insert_neighbor(lsn); } /// Internal function @@ -374,7 +383,7 @@ impl ListSearchResult { self.best_candidate.insert(idx, n) } - fn visit_closest(&mut self, pos_limit: usize) -> Option<(ItemPointer, f32)> { + fn visit_closest(&mut self, pos_limit: usize) -> Option<&ListSearchNeighbor> { //OPT: should we optimize this not to do a linear search each time? let neighbor_position = self.best_candidate.iter().position(|n| !n.visited); match neighbor_position { @@ -384,7 +393,7 @@ impl ListSearchResult { } let n = &mut self.best_candidate[pos]; n.visited = true; - Some((n.index_pointer, n.distance)) + Some(n) } None => None, } @@ -404,12 +413,7 @@ impl ListSearchResult { pub trait Graph { fn read<'a>(&self, index: &'a PgRelation, index_pointer: ItemPointer) -> ReadableNode<'a>; fn get_init_ids(&mut self) -> Option>; - fn get_neighbors( - &self, - index: &PgRelation, - neighbors_of: ItemPointer, - result: &mut Vec, - ) -> bool; + fn get_neighbors(&self, node: &ArchivedNode, neighbors_of: ItemPointer) -> Vec; fn get_neighbors_with_distances( &self, index: &PgRelation, @@ -499,20 +503,16 @@ pub trait Graph { Self: Graph, { //OPT: Only build v when needed. - let mut v: HashSet<_> = HashSet::::with_capacity(visit_n_closest); let mut neighbors = Vec::::with_capacity(self.get_meta_page(index).get_num_neighbors() as _); - while let Some((index_pointer, distance)) = lsr.visit_closest(visit_n_closest) { + let mut v: HashSet<_> = HashSet::::with_capacity(visit_n_closest); + while let Some(node) = lsr.visit_closest(visit_n_closest) { neighbors.clear(); - let neighbors_existed = self.get_neighbors(index, index_pointer, &mut neighbors); - if !neighbors_existed { - panic!("Nodes in the list search results that aren't in the builder"); - } - - for neighbor_index_pointer in &neighbors { + v.insert(NeighborWithDistance::new(node.index_pointer, node.distance)); + neighbors.extend_from_slice(node.neighbor_index_pointers.as_slice()); + for neighbor_index_pointer in neighbors.iter() { lsr.insert(index, self, *neighbor_index_pointer, query) } - v.insert(NeighborWithDistance::new(index_pointer, distance)); } Some(v) From 0daa77b1eca9e618ef558a2181307adecf9bef5d Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Thu, 14 Dec 2023 10:44:54 -0500 Subject: [PATCH 03/44] Optimization: memory allocations --- .../src/access_method/builder_graph.rs | 2 +- .../src/access_method/disk_index_graph.rs | 2 +- timescale_vector/src/access_method/graph.rs | 74 +++++++++++++------ timescale_vector/src/access_method/scan.rs | 14 +++- 4 files changed, 63 insertions(+), 29 deletions(-) diff --git a/timescale_vector/src/access_method/builder_graph.rs b/timescale_vector/src/access_method/builder_graph.rs index 2c31a8fe..f14c8e0c 100644 --- a/timescale_vector/src/access_method/builder_graph.rs +++ b/timescale_vector/src/access_method/builder_graph.rs @@ -81,7 +81,7 @@ impl<'a> Graph for BuilderGraph<'a> { unsafe { Node::read(index, index_pointer) } } - fn get_init_ids(&mut self) -> Option> { + fn get_init_ids(&self) -> Option> { //returns a vector for generality self.meta_page.get_init_ids() } diff --git a/timescale_vector/src/access_method/disk_index_graph.rs b/timescale_vector/src/access_method/disk_index_graph.rs index 20b241ad..7ae548f4 100644 --- a/timescale_vector/src/access_method/disk_index_graph.rs +++ b/timescale_vector/src/access_method/disk_index_graph.rs @@ -32,7 +32,7 @@ impl<'h> Graph for DiskIndexGraph<'h> { unsafe { Node::read(index, index_pointer) } } - fn get_init_ids(&mut self) -> Option> { + fn get_init_ids(&self) -> Option> { self.meta_page.get_init_ids() } diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index 81b61e00..6ce19b47 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -317,13 +317,16 @@ impl ListSearchResult { init_ids: Vec, query: &[f32], dm: DistanceMeasure, + search_list_size: usize, + meta_page: &MetaPage, ) -> Self where G: Graph + ?Sized, { + let neigbors = meta_page.get_num_neighbors() as usize; let mut res = Self { - best_candidate: Vec::new(), - inserted: HashSet::new(), + best_candidate: Vec::with_capacity(search_list_size * neigbors), + inserted: HashSet::with_capacity(search_list_size * neigbors), max_history_size, stats: GreedySearchStats::new(), dm: dm, @@ -412,7 +415,7 @@ impl ListSearchResult { pub trait Graph { fn read<'a>(&self, index: &'a PgRelation, index_pointer: ItemPointer) -> ReadableNode<'a>; - fn get_init_ids(&mut self) -> Option>; + fn get_init_ids(&self) -> Option>; fn get_neighbors(&self, node: &ArchivedNode, neighbors_of: ItemPointer) -> Vec; fn get_neighbors_with_distances( &self, @@ -447,20 +450,21 @@ pub trait Graph { /// Note this is the one-shot implementation that keeps only the closest `search_list_size` results in /// the returned ListSearchResult elements. It shouldn't be used with self.greedy_search_iterate fn greedy_search( - &mut self, + &self, index: &PgRelation, query: &[f32], - search_list_size: usize, - ) -> (ListSearchResult, Option>) + meta_page: &MetaPage, + ) -> (ListSearchResult, HashSet) where Self: Graph, { let init_ids = self.get_init_ids(); if let None = init_ids { //no nodes in the graph - return (ListSearchResult::empty(), None); + return (ListSearchResult::empty(), HashSet::with_capacity(0)); } let dm = self.get_vector_provider().get_distance_measure(query); + let search_list_size = meta_page.get_search_list_size_for_build() as usize; let mut l = ListSearchResult::new( index, @@ -469,17 +473,28 @@ pub trait Graph { init_ids.unwrap(), query, dm, + search_list_size, + meta_page, + ); + let mut visited_nodes = HashSet::with_capacity(search_list_size); + self.greedy_search_iterate( + &mut l, + index, + query, + search_list_size, + Some(&mut visited_nodes), ); - let v = self.greedy_search_iterate(&mut l, index, query, search_list_size); - return (l, v); + return (l, visited_nodes); } /// Returns a ListSearchResult initialized for streaming. The output should be used with greedy_search_iterate to obtain /// the next elements. fn greedy_search_streaming_init( - &mut self, + &self, index: &PgRelation, query: &[f32], + search_list_size: usize, + meta_page: &MetaPage, ) -> ListSearchResult { let init_ids = self.get_init_ids(); if let None = init_ids { @@ -488,34 +503,48 @@ pub trait Graph { } let dm = self.get_vector_provider().get_distance_measure(query); - ListSearchResult::new(index, None, self, init_ids.unwrap(), query, dm) + ListSearchResult::new( + index, + None, + self, + init_ids.unwrap(), + query, + dm, + search_list_size, + meta_page, + ) } /// Advance the state of the lsr until the closest `visit_n_closest` elements have been visited. fn greedy_search_iterate( - &mut self, + &self, lsr: &mut ListSearchResult, index: &PgRelation, query: &[f32], visit_n_closest: usize, - ) -> Option> - where + mut visited_nodes: Option<&mut HashSet>, + ) where Self: Graph, { //OPT: Only build v when needed. let mut neighbors = Vec::::with_capacity(self.get_meta_page(index).get_num_neighbors() as _); - let mut v: HashSet<_> = HashSet::::with_capacity(visit_n_closest); - while let Some(node) = lsr.visit_closest(visit_n_closest) { + while let Some(list_search_entry) = lsr.visit_closest(visit_n_closest) { neighbors.clear(); - v.insert(NeighborWithDistance::new(node.index_pointer, node.distance)); - neighbors.extend_from_slice(node.neighbor_index_pointers.as_slice()); + match visited_nodes { + None => {} + Some(ref mut visited_nodes) => { + visited_nodes.insert(NeighborWithDistance::new( + list_search_entry.index_pointer, + list_search_entry.distance, + )); + } + } + neighbors.extend_from_slice(list_search_entry.neighbor_index_pointers.as_slice()); for neighbor_index_pointer in neighbors.iter() { lsr.insert(index, self, *neighbor_index_pointer, query) } } - - Some(v) } /// Prune neigbors by prefering neighbors closer to the point in question @@ -656,11 +685,10 @@ pub trait Graph { } //TODO: make configurable? - let (l, v) = - self.greedy_search(index, vec, meta_page.get_search_list_size_for_build() as _); + let (l, v) = self.greedy_search(index, vec, meta_page); greedy_search_stats.combine(l.stats); let (neighbor_list, forward_stats) = - self.prune_neighbors(index, index_pointer, v.unwrap().into_iter().collect()); + self.prune_neighbors(index, index_pointer, v.into_iter().collect()); prune_neighbor_stats.combine(forward_stats); //set forward pointers diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 5acf2832..660ba2a2 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -27,12 +27,12 @@ impl<'a> TSVResponseIterator<'a> { Quantizer::None => {} Quantizer::PQ(pq) => pq.load(index, &meta_page), } - let mut graph = DiskIndexGraph::new( + let graph = DiskIndexGraph::new( &index, VectorProvider::new(None, None, &quantizer, quantizer.is_some()), ); use super::graph::Graph; - let lsr = graph.greedy_search_streaming_init(&index, query); + let lsr = graph.greedy_search_streaming_init(&index, query, search_list_size, &meta_page); Self { query: query.to_vec(), search_list_size, @@ -46,7 +46,7 @@ impl<'a> TSVResponseIterator<'a> { impl<'a> TSVResponseIterator<'a> { fn next(&mut self, index: &'a PgRelation) -> Option { - let mut graph = DiskIndexGraph::new( + let graph = DiskIndexGraph::new( &index, VectorProvider::new(None, None, &self.quantizer, self.quantizer.is_some()), ); @@ -54,7 +54,13 @@ impl<'a> TSVResponseIterator<'a> { /* Iterate until we find a non-deleted tuple */ loop { - graph.greedy_search_iterate(&mut self.lsr, index, &self.query, self.search_list_size); + graph.greedy_search_iterate( + &mut self.lsr, + index, + &self.query, + self.search_list_size, + None, + ); let item = self.lsr.consume(); From dc39a0fe8f4a383c0d3be4d538e8277632c3ea2a Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Thu, 14 Dec 2023 13:53:19 -0500 Subject: [PATCH 04/44] Optimization: optimize lsr with separate storage --- timescale_vector/src/access_method/graph.rs | 24 +++++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index 6ce19b47..77b54256 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -290,7 +290,8 @@ impl ListSearchNeighbor { } pub struct ListSearchResult { - best_candidate: Vec, //keep sorted by distanced + candidate_storage: Vec, //plain storage + best_candidate: Vec, //pos in candidate storage, sorted by distance inserted: HashSet, max_history_size: Option, dm: DistanceMeasure, @@ -300,6 +301,7 @@ pub struct ListSearchResult { impl ListSearchResult { fn empty() -> Self { Self { + candidate_storage: vec![], best_candidate: vec![], inserted: HashSet::new(), max_history_size: None, @@ -325,6 +327,7 @@ impl ListSearchResult { { let neigbors = meta_page.get_num_neighbors() as usize; let mut res = Self { + candidate_storage: Vec::with_capacity(search_list_size * neigbors), best_candidate: Vec::with_capacity(search_list_size * neigbors), inserted: HashSet::with_capacity(search_list_size * neigbors), max_history_size, @@ -374,7 +377,7 @@ impl ListSearchResult { if let Some(max_size) = self.max_history_size { if self.best_candidate.len() >= max_size { let last = self.best_candidate.last().unwrap(); - if n >= *last { + if n >= self.candidate_storage[*last] { //n is too far in the list to be the best candidate. return; } @@ -382,19 +385,26 @@ impl ListSearchResult { } } //insert while preserving sort order. - let idx = self.best_candidate.partition_point(|x| *x < n); - self.best_candidate.insert(idx, n) + let idx = self + .best_candidate + .partition_point(|x| self.candidate_storage[*x] < n); + self.candidate_storage.push(n); + let pos = self.candidate_storage.len() - 1; + self.best_candidate.insert(idx, pos) } fn visit_closest(&mut self, pos_limit: usize) -> Option<&ListSearchNeighbor> { //OPT: should we optimize this not to do a linear search each time? - let neighbor_position = self.best_candidate.iter().position(|n| !n.visited); + let neighbor_position = self + .best_candidate + .iter() + .position(|n| !self.candidate_storage[*n].visited); match neighbor_position { Some(pos) => { if pos > pos_limit { return None; } - let n = &mut self.best_candidate[pos]; + let n = &mut self.candidate_storage[self.best_candidate[pos]]; n.visited = true; Some(n) } @@ -408,7 +418,7 @@ impl ListSearchResult { if self.best_candidate.is_empty() { return None; } - let f = self.best_candidate.remove(0); + let f = &self.candidate_storage[self.best_candidate.remove(0)]; return Some((f.heap_pointer, f.index_pointer)); } } From 6e6ccd14c0c9fa05cfa18fd17c2a5f47b7d83c93 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Mon, 11 Dec 2023 10:39:08 -0500 Subject: [PATCH 05/44] preallocate vec capacity --- timescale_vector/src/access_method/model.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/timescale_vector/src/access_method/model.rs b/timescale_vector/src/access_method/model.rs index 7c168061..e750b757 100644 --- a/timescale_vector/src/access_method/model.rs +++ b/timescale_vector/src/access_method/model.rs @@ -352,7 +352,8 @@ impl<'a> ReadablePqVectorNode<'a> { pub unsafe fn read_pq(index: &PgRelation, index_pointer: &IndexPointer) -> Pq { let rpq = PqQuantizerDef::read(index, &index_pointer); let rpn = rpq.get_archived_node(); - let mut result: Vec = Vec::new(); + let size = rpn.dim_0 * rpn.dim_1 * rpn.dim_2; + let mut result: Vec = Vec::with_capacity(size as usize); let mut next = rpn.next_vector_pointer.deserialize_item_pointer(); loop { if next.offset == 0 && next.block_number == 0 { From fe8a516e55a12ef1009bbfa4fed93dbeeb8abe3d Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Thu, 14 Dec 2023 15:03:00 -0500 Subject: [PATCH 06/44] optimize pq load --- timescale_vector/src/access_method/model.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/timescale_vector/src/access_method/model.rs b/timescale_vector/src/access_method/model.rs index e750b757..20e52397 100644 --- a/timescale_vector/src/access_method/model.rs +++ b/timescale_vector/src/access_method/model.rs @@ -361,8 +361,7 @@ pub unsafe fn read_pq(index: &PgRelation, index_pointer: &IndexPointer) -> Pq Date: Thu, 14 Dec 2023 21:23:23 -0500 Subject: [PATCH 07/44] Make changes to distance functions --- timescale_vector/Cargo.toml | 14 +- timescale_vector/benches/distance.rs | 164 +++++++++++++++++ .../src/access_method/distance.rs | 54 ++++++ .../src/access_method/distance_x86.rs | 166 ++++++++++++++++-- timescale_vector/src/access_method/graph.rs | 26 +-- timescale_vector/src/access_method/mod.rs | 1 + timescale_vector/src/lib.rs | 2 +- 7 files changed, 385 insertions(+), 42 deletions(-) create mode 100644 timescale_vector/benches/distance.rs create mode 100644 timescale_vector/src/access_method/distance.rs diff --git a/timescale_vector/Cargo.toml b/timescale_vector/Cargo.toml index c43b4803..625362e5 100644 --- a/timescale_vector/Cargo.toml +++ b/timescale_vector/Cargo.toml @@ -4,7 +4,7 @@ version = "0.0.2" edition = "2021" [lib] -crate-type = ["cdylib"] +crate-type = ["cdylib", "rlib"] [features] default = ["pg16"] @@ -14,9 +14,9 @@ pg_test = [] [dependencies] memoffset = "0.9.0" -pgrx = "=0.11.1" +pgrx = "=0.11.2" rkyv = { version="0.7.42", features=["validation"]} -simdeez = {version = "1.0"} +simdeez = {version = "1.0.8"} reductive = { version = "0.9.0"} ndarray = { version = "0.15.0", features = ["blas"] } blas-src = { version = "0.8", features = ["openblas"] } @@ -29,7 +29,8 @@ rayon = "1" [dev-dependencies] -pgrx-tests = "=0.11.1" +pgrx-tests = "=0.11.2" +criterion = "0.5.1" [profile.dev] panic = "unwind" @@ -39,3 +40,8 @@ panic = "unwind" opt-level = 3 lto = "fat" codegen-units = 1 +#debug = true + +[[bench]] +name = "distance" +harness = false diff --git a/timescale_vector/benches/distance.rs b/timescale_vector/benches/distance.rs new file mode 100644 index 00000000..dfce659e --- /dev/null +++ b/timescale_vector/benches/distance.rs @@ -0,0 +1,164 @@ +use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use timescale_vector::access_method::distance::{distance_cosine, distance_l2}; + +//copy and use qdrants simd code, purely for benchmarking purposes +//not used in the actual extension +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +use std::arch::x86_64::*; + +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +#[target_feature(enable = "avx")] +#[target_feature(enable = "fma")] +unsafe fn hsum256_ps_avx(x: __m256) -> f32 { + let x128: __m128 = _mm_add_ps(_mm256_extractf128_ps(x, 1), _mm256_castps256_ps128(x)); + let x64: __m128 = _mm_add_ps(x128, _mm_movehl_ps(x128, x128)); + let x32: __m128 = _mm_add_ss(x64, _mm_shuffle_ps(x64, x64, 0x55)); + _mm_cvtss_f32(x32) +} + +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +#[target_feature(enable = "avx")] +#[target_feature(enable = "fma")] +pub unsafe fn dot_similarity_avx_qdrant(v1: &[f32], v2: &[f32]) -> f32 { + let n = v1.len(); + let m = n - (n % 32); + let mut ptr1: *const f32 = v1.as_ptr(); + let mut ptr2: *const f32 = v2.as_ptr(); + let mut sum256_1: __m256 = _mm256_setzero_ps(); + let mut sum256_2: __m256 = _mm256_setzero_ps(); + let mut sum256_3: __m256 = _mm256_setzero_ps(); + let mut sum256_4: __m256 = _mm256_setzero_ps(); + let mut i: usize = 0; + while i < m { + sum256_1 = _mm256_fmadd_ps(_mm256_loadu_ps(ptr1), _mm256_loadu_ps(ptr2), sum256_1); + sum256_2 = _mm256_fmadd_ps( + _mm256_loadu_ps(ptr1.add(8)), + _mm256_loadu_ps(ptr2.add(8)), + sum256_2, + ); + sum256_3 = _mm256_fmadd_ps( + _mm256_loadu_ps(ptr1.add(16)), + _mm256_loadu_ps(ptr2.add(16)), + sum256_3, + ); + sum256_4 = _mm256_fmadd_ps( + _mm256_loadu_ps(ptr1.add(24)), + _mm256_loadu_ps(ptr2.add(24)), + sum256_4, + ); + + ptr1 = ptr1.add(32); + ptr2 = ptr2.add(32); + i += 32; + } + + let mut result = hsum256_ps_avx(sum256_1) + + hsum256_ps_avx(sum256_2) + + hsum256_ps_avx(sum256_3) + + hsum256_ps_avx(sum256_4); + + for i in 0..n - m { + result += (*ptr1.add(i)) * (*ptr2.add(i)); + } + result +} + +/// Copy of Diskann's distance function. again just for benchmarking +/// not used in the actual extension +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +#[inline(never)] +pub unsafe fn distance_l2_vector_f32(a: &[f32], b: &[f32]) -> f32 { + let n = a.len(); + + // make sure the addresses are bytes aligned + debug_assert_eq!(a.as_ptr().align_offset(32), 0); + debug_assert_eq!(b.as_ptr().align_offset(32), 0); + + unsafe { + let mut sum = _mm256_setzero_ps(); + + // Iterate over the elements in steps of 8 + for i in (0..n).step_by(8) { + let a_vec = _mm256_load_ps(&a[i]); + let b_vec = _mm256_load_ps(&b[i]); + let diff = _mm256_sub_ps(a_vec, b_vec); + sum = _mm256_fmadd_ps(diff, diff, sum); + } + + let x128: __m128 = _mm_add_ps(_mm256_extractf128_ps(sum, 1), _mm256_castps256_ps128(sum)); + /* ( -, -, x1+x3+x5+x7, x0+x2+x4+x6 ) */ + let x64: __m128 = _mm_add_ps(x128, _mm_movehl_ps(x128, x128)); + /* ( -, -, -, x0+x1+x2+x3+x4+x5+x6+x7 ) */ + let x32: __m128 = _mm_add_ss(x64, _mm_shuffle_ps(x64, x64, 0x55)); + /* Conversion to float is a no-op on x86-64 */ + _mm_cvtss_f32(x32) + } +} + +//only used for alignment +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +#[repr(C, align(32))] +struct Vector32ByteAligned { + v: [f32; 2000], +} + +//the diskann version requires alignment so run benchmarks with aligned vectors +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +fn benchmark_distance_x86_aligned_vectors(c: &mut Criterion) { + let a = Box::new(Vector32ByteAligned { + v: [(); 2000].map(|_| 100.1), + }); + + let b = Box::new(Vector32ByteAligned { + v: [(); 2000].map(|_| 22.1), + }); + + let l = a.v; + let r = b.v; + + assert_eq!(r.as_ptr().align_offset(32), 0); + assert_eq!(l.as_ptr().align_offset(32), 0); + + c.bench_function("distance comparison qdrant (aligned)", |b| { + b.iter(|| unsafe { dot_similarity_avx_qdrant(black_box(&r), black_box(&l)) }) + }); + c.bench_function("distance comparison diskann (aligned)", |b| { + b.iter(|| unsafe { distance_l2_vector_f32(black_box(&r), black_box(&l)) }) + }); +} + +//compare qdrant on unaligned vectors (we don't have alignment so this is apples to apples with us) +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +fn benchmark_distance_x86_unaligned_vectors(c: &mut Criterion) { + let r: Vec = (0..2000).map(|v| v as f32 + 1000.1).collect(); + let l: Vec = (0..2000).map(|v| v as f32 + 2000.2).collect(); + + c.bench_function("distance comparison qdrant (unaligned)", |b| { + b.iter(|| unsafe { dot_similarity_avx_qdrant(black_box(&r), black_box(&l)) }) + }); +} + +fn benchmark_distance(c: &mut Criterion) { + let r: Vec = (0..2000).map(|v| v as f32 + 1000.1).collect(); + let l: Vec = (0..2000).map(|v| v as f32 + 2000.2).collect(); + + let mut group = c.benchmark_group("Distance"); + group.bench_function("distance l2", |b| { + b.iter(|| distance_l2(black_box(&r), black_box(&l))) + }); + group.bench_function("distance cosine", |b| { + b.iter(|| distance_cosine(black_box(&r), black_box(&l))) + }); +} + +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +criterion_group!( + benches, + benchmark_distance, + benchmark_distance_x86_unaligned_vectors, + benchmark_distance_x86_aligned_vectors +); +#[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))] +criterion_group!(benches, benchmark_distance); + +criterion_main!(benches); diff --git a/timescale_vector/src/access_method/distance.rs b/timescale_vector/src/access_method/distance.rs new file mode 100644 index 00000000..f96ccdd7 --- /dev/null +++ b/timescale_vector/src/access_method/distance.rs @@ -0,0 +1,54 @@ +/* we use the avx2 version of x86 functions. This verifies that's kosher */ +#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] +#[cfg(not(target_feature = "avx2"))] +compile_error!( + "On x86, the AVX2 feature must be enabled. Set RUSTFLAGS=\"-C target-feature=+avx2,+fma\"" +); + +#[inline] +pub fn distance_l2(a: &[f32], b: &[f32]) -> f32 { + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + //note safety is guraranteed by compile_error above + unsafe { + return super::distance_x86::distance_l2_x86_avx2(a, b); + } + + #[allow(unreachable_code)] + { + return distance_l2_unoptimized(a, b); + } +} + +#[inline(always)] +pub fn distance_l2_unoptimized(a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len()); + let norm: f32 = a + .iter() + .zip(b.iter()) + .map(|t| (*t.0 as f32 - *t.1 as f32) * (*t.0 as f32 - *t.1 as f32)) + .sum(); + assert!(norm >= 0.); + //don't sqrt for performance. These are only used for ordering so sqrt not needed + norm +} + +#[inline] +pub fn distance_cosine(a: &[f32], b: &[f32]) -> f32 { + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] + //note safety is guraranteed by compile_error above + unsafe { + return super::distance_x86::distance_cosine_x86_avx2(a, b); + } + + #[allow(unreachable_code)] + { + return distance_cosine_unoptimized(a, b); + } +} + +#[inline(always)] +pub fn distance_cosine_unoptimized(a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len()); + let res: f32 = a.iter().zip(b).map(|(a, b)| *a * *b).sum(); + -res +} diff --git a/timescale_vector/src/access_method/distance_x86.rs b/timescale_vector/src/access_method/distance_x86.rs index d12de428..46e580cf 100644 --- a/timescale_vector/src/access_method/distance_x86.rs +++ b/timescale_vector/src/access_method/distance_x86.rs @@ -6,11 +6,25 @@ use simdeez::sse41::*; //use simdeez::avx::*; use simdeez::avx2::*; +#[cfg(not(target_feature = "avx2"))] +compile_error!( + "On x86, the AVX2 feature must be enabled. Set RUSTFLAGS=\"-C target-feature=+avx2,+fma\"" +); + +//note: without fmadd, the performance degrades pretty badly. Benchmark before disbaling +#[cfg(not(target_feature = "fma"))] +compile_error!( + "On x86, the fma feature must be enabled. Set RUSTFLAGS=\"-C target-feature=+avx2,+fma\"" +); + simdeez::simd_runtime_generate!( - pub fn distance_opt(x: &[f32], y: &[f32]) -> f32 { - let mut res = S::setzero_ps(); + pub fn distance_l2_x86(x: &[f32], y: &[f32]) -> f32 { + let mut accum0 = S::setzero_ps(); + let mut accum1 = S::setzero_ps(); + let mut accum2 = S::setzero_ps(); + let mut accum3 = S::setzero_ps(); - assert!(x.len() == y.len()); + //assert!(x.len() == y.len()); let mut x = &x[..]; let mut y = &y[..]; @@ -19,22 +33,30 @@ simdeez::simd_runtime_generate!( // the width of a vector type is provided as a constant // so the compiler is free to optimize it more. // S::VF32_WIDTH is a constant, 4 when using SSE, 8 when using AVX2, etc - while x.len() >= S::VF32_WIDTH { + while x.len() >= S::VF32_WIDTH * 4 { //load data from your vec into an SIMD value - let xv = S::loadu_ps(&x[0]); - let yv = S::loadu_ps(&y[0]); - - let mut diff = S::sub_ps(xv, yv); - diff *= diff; - - res = res + diff; + accum0 = accum0 + + ((S::loadu_ps(&x[S::VF32_WIDTH * 0]) - S::loadu_ps(&y[S::VF32_WIDTH * 0])) + * (S::loadu_ps(&x[S::VF32_WIDTH * 0]) - S::loadu_ps(&y[S::VF32_WIDTH * 0]))); + accum1 = accum1 + + ((S::loadu_ps(&x[S::VF32_WIDTH * 1]) - S::loadu_ps(&y[S::VF32_WIDTH * 1])) + * (S::loadu_ps(&x[S::VF32_WIDTH * 1]) - S::loadu_ps(&y[S::VF32_WIDTH * 1]))); + accum2 = accum2 + + ((S::loadu_ps(&x[S::VF32_WIDTH * 2]) - S::loadu_ps(&y[S::VF32_WIDTH * 2])) + * (S::loadu_ps(&x[S::VF32_WIDTH * 2]) - S::loadu_ps(&y[S::VF32_WIDTH * 2]))); + accum3 = accum3 + + ((S::loadu_ps(&x[S::VF32_WIDTH * 3]) - S::loadu_ps(&y[S::VF32_WIDTH * 3])) + * (S::loadu_ps(&x[S::VF32_WIDTH * 3]) - S::loadu_ps(&y[S::VF32_WIDTH * 3]))); // Move each slice to the next position - x = &x[S::VF32_WIDTH..]; - y = &y[S::VF32_WIDTH..]; + x = &x[S::VF32_WIDTH * 4..]; + y = &y[S::VF32_WIDTH * 4..]; } - let mut dist = S::horizontal_add_ps(res); + let mut dist = S::horizontal_add_ps(accum0) + + S::horizontal_add_ps(accum1) + + S::horizontal_add_ps(accum2) + + S::horizontal_add_ps(accum3); // compute for the remaining elements for i in 0..x.len() { @@ -43,6 +65,120 @@ simdeez::simd_runtime_generate!( } assert!(dist >= 0.); - dist.sqrt() + //dist.sqrt() + dist } ); + +simdeez::simd_runtime_generate!( + pub fn distance_cosine_x86(x: &[f32], y: &[f32]) -> f32 { + let mut accum0 = S::setzero_ps(); + let mut accum1 = S::setzero_ps(); + let mut accum2 = S::setzero_ps(); + let mut accum3 = S::setzero_ps(); + + let mut x = &x[..]; + let mut y = &y[..]; + + //assert!(x.len() == y.len()); + + // Operations have to be done in terms of the vector width + // so that it will work with any size vector. + // the width of a vector type is provided as a constant + // so the compiler is free to optimize it more. + // S::VF32_WIDTH is a constant, 4 when using SSE, 8 when using AVX2, etc + while x.len() >= S::VF32_WIDTH * 4 { + accum0 = S::fmadd_ps( + S::loadu_ps(&x[S::VF32_WIDTH * 0]), + S::loadu_ps(&y[S::VF32_WIDTH * 0]), + accum0, + ); + accum1 = S::fmadd_ps( + S::loadu_ps(&x[S::VF32_WIDTH * 1]), + S::loadu_ps(&y[S::VF32_WIDTH * 1]), + accum1, + ); + accum2 = S::fmadd_ps( + S::loadu_ps(&x[S::VF32_WIDTH * 2]), + S::loadu_ps(&y[S::VF32_WIDTH * 2]), + accum2, + ); + accum3 = S::fmadd_ps( + S::loadu_ps(&x[S::VF32_WIDTH * 3]), + S::loadu_ps(&y[S::VF32_WIDTH * 3]), + accum3, + ); + + // Move each slice to the next position + x = &x[S::VF32_WIDTH * 4..]; + y = &y[S::VF32_WIDTH * 4..]; + } + + let mut dist = S::horizontal_add_ps(accum0) + + S::horizontal_add_ps(accum1) + + S::horizontal_add_ps(accum2) + + S::horizontal_add_ps(accum3); + + // compute for the remaining elements + for i in 0..x.len() { + dist += x[i] * y[i]; + } + + -dist + } +); + +#[cfg(test)] +mod tests { + #[test] + fn distances_equal() { + let r: Vec = (0..2000).map(|_| 1.0).collect(); + let l: Vec = (0..2000).map(|_| 2.0).collect(); + + assert_eq!( + unsafe { super::distance_cosine_x86_avx2(&r, &l) }, + super::super::distance::distance_cosine_unoptimized(&r, &l) + ); + + assert_eq!( + unsafe { super::distance_l2_x86_avx2(&r, &l) }, + super::super::distance::distance_l2_unoptimized(&r, &l) + ); + + //don't use too many dimensions to avoid overflow + let r: Vec = (0..20).map(|v| v as f32).collect(); + let l: Vec = (0..20).map(|v| v as f32).collect(); + + assert_eq!( + unsafe { super::distance_cosine_x86_avx2(&r, &l) }, + super::super::distance::distance_cosine_unoptimized(&r, &l) + ); + assert_eq!( + unsafe { super::distance_l2_x86_avx2(&r, &l) }, + super::super::distance::distance_l2_unoptimized(&r, &l) + ); + + //many dimensions but normalized + let r: Vec = (0..2000).map(|v| v as f32 + 1.0).collect(); + let l: Vec = (0..2000).map(|v| v as f32 + 2.0).collect(); + + let r_size = r.iter().map(|v| v * v).sum::().sqrt(); + let l_size = l.iter().map(|v| v * v).sum::().sqrt(); + + let r: Vec = r.iter().map(|v| v / r_size).collect(); + let l: Vec = l.iter().map(|v| v / l_size).collect(); + + assert!( + (unsafe { super::distance_cosine_x86_avx2(&r, &l) } + - super::super::distance::distance_cosine_unoptimized(&r, &l)) + .abs() + < 0.000001 + ); + assert!( + (unsafe { super::distance_l2_x86_avx2(&r, &l) } + - super::super::distance::distance_l2_unoptimized(&r, &l)) + .abs() + < 0.000001 + ); + } +} diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index 77b54256..13c4d4db 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -7,6 +7,7 @@ use crate::access_method::model::Node; use crate::util::ports::slot_getattr; use crate::util::{HeapPointer, IndexPointer, ItemPointer}; +use super::distance::distance_l2 as default_distance; use super::model::{ArchivedNode, PgVector}; use super::quantizer::Quantizer; use super::{ @@ -14,25 +15,6 @@ use super::{ model::{NeighborWithDistance, ReadableNode}, }; -#[cfg(any(target_arch = "x86", target_arch = "x86_64"))] -fn distance(a: &[f32], b: &[f32]) -> f32 { - super::distance_x86::distance_opt_runtime_select(a, b) -} - -//TODO: use slow L2 for now. Make pluggable and simd -#[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))] -fn distance(a: &[f32], b: &[f32]) -> f32 { - assert_eq!(a.len(), b.len()); - - let norm: f32 = a - .iter() - .zip(b.iter()) - .map(|t| (*t.0 as f32 - *t.1 as f32) * (*t.0 as f32 - *t.1 as f32)) - .sum(); - assert!(norm >= 0.); - norm.sqrt() -} - struct TableSlot { slot: PgBox, } @@ -78,7 +60,7 @@ impl<'a> VectorProvider<'a> { calc_distance_with_quantizer, heap_rel, heap_attr_number, - distance_fn: distance, + distance_fn: default_distance, } } @@ -232,7 +214,7 @@ impl DistanceMeasure { pq_distance_table: None, }, Quantizer::PQ(pq) => { - let dc = pq.get_distance_table(query, distance); + let dc = pq.get_distance_table(query, default_distance); Self { pq_distance_table: Some(dc), } @@ -248,7 +230,7 @@ impl DistanceMeasure { fn get_full_vector_distance(&self, vec: &[f32], query: &[f32]) -> f32 { assert!(self.pq_distance_table.is_none()); - distance(vec, query) + default_distance(vec, query) } } diff --git a/timescale_vector/src/access_method/mod.rs b/timescale_vector/src/access_method/mod.rs index 83dea8fa..9db1bcd9 100644 --- a/timescale_vector/src/access_method/mod.rs +++ b/timescale_vector/src/access_method/mod.rs @@ -15,6 +15,7 @@ mod vacuum; extern crate blas_src; +pub mod distance; #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] mod distance_x86; mod pq; diff --git a/timescale_vector/src/lib.rs b/timescale_vector/src/lib.rs index 9753ddf4..72edc3c6 100644 --- a/timescale_vector/src/lib.rs +++ b/timescale_vector/src/lib.rs @@ -2,7 +2,7 @@ use pgrx::prelude::*; pgrx::pg_module_magic!(); -mod access_method; +pub mod access_method; mod util; #[allow(non_snake_case)] From 2ff844ad9ce6b6f752c2e3becaf1c0725f83d03a Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Mon, 18 Dec 2023 14:03:00 -0500 Subject: [PATCH 08/44] Make cosine distance always positive We need positive distances for the ratio in the prune function. This also requires us to make sure the vectors are normalized since cosine on non-normalized vectors could return a distance less than 1.0, which would result in negative distances. --- .../src/access_method/distance.rs | 13 +++++- .../src/access_method/distance_x86.rs | 2 +- timescale_vector/src/access_method/graph.rs | 44 +++++++++++++++---- timescale_vector/src/access_method/model.rs | 8 ++-- timescale_vector/src/access_method/scan.rs | 4 +- 5 files changed, 56 insertions(+), 15 deletions(-) diff --git a/timescale_vector/src/access_method/distance.rs b/timescale_vector/src/access_method/distance.rs index f96ccdd7..f30dccab 100644 --- a/timescale_vector/src/access_method/distance.rs +++ b/timescale_vector/src/access_method/distance.rs @@ -50,5 +50,16 @@ pub fn distance_cosine(a: &[f32], b: &[f32]) -> f32 { pub fn distance_cosine_unoptimized(a: &[f32], b: &[f32]) -> f32 { assert_eq!(a.len(), b.len()); let res: f32 = a.iter().zip(b).map(|(a, b)| *a * *b).sum(); - -res + 1.0 - res +} + +pub fn preprocess_cosine(a: &mut [f32]) { + let norm = a.iter().map(|v| v * v).sum::(); + if norm < f32::EPSILON { + return; + } + let norm = norm.sqrt(); + if norm > 1.0 + f32::EPSILON || norm < 1.0 - f32::EPSILON { + a.iter_mut().for_each(|v| *v /= norm); + } } diff --git a/timescale_vector/src/access_method/distance_x86.rs b/timescale_vector/src/access_method/distance_x86.rs index 46e580cf..9ea6e71e 100644 --- a/timescale_vector/src/access_method/distance_x86.rs +++ b/timescale_vector/src/access_method/distance_x86.rs @@ -124,7 +124,7 @@ simdeez::simd_runtime_generate!( dist += x[i] * y[i]; } - -dist + 1.0 - dist } ); diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index 13c4d4db..95055836 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -7,7 +7,7 @@ use crate::access_method::model::Node; use crate::util::ports::slot_getattr; use crate::util::{HeapPointer, IndexPointer, ItemPointer}; -use super::distance::distance_l2 as default_distance; +use super::distance::distance_cosine as default_distance; use super::model::{ArchivedNode, PgVector}; use super::quantizer::Quantizer; use super::{ @@ -628,7 +628,7 @@ pub trait Graph { } //todo handle the non-pq case - let distance_between_candidate_and_existing_neighbor = unsafe { + let mut distance_between_candidate_and_existing_neighbor = unsafe { vp.get_distance_pair_for_full_vectors_from_state( &dist_state, index, @@ -637,14 +637,39 @@ pub trait Graph { }; stats.node_reads += 2; stats.distance_comparisons += 1; - let distance_between_candidate_and_point = candidate_neighbor.get_distance(); + let mut distance_between_candidate_and_point = + candidate_neighbor.get_distance(); + + //We need both values to be positive. + //Otherwise, the case where distance_between_candidate_and_point > 0 and distance_between_candidate_and_existing_neighbor < 0 is totally wrong. + //If we implement inner product distance we'll have to figure something else out. + if distance_between_candidate_and_point < 0.0 + && distance_between_candidate_and_point >= 0.0 - f32::EPSILON + { + distance_between_candidate_and_point = 0.0; + } + + if distance_between_candidate_and_existing_neighbor < 0.0 + && distance_between_candidate_and_existing_neighbor >= 0.0 - f32::EPSILON + { + distance_between_candidate_and_existing_neighbor = 0.0; + } + + debug_assert!(distance_between_candidate_and_point >= 0.0); + debug_assert!(distance_between_candidate_and_existing_neighbor >= 0.0); + //factor is high if the candidate is closer to an existing neighbor than the point it's being considered for - let factor = if distance_between_candidate_and_existing_neighbor == 0.0 { - f64::MAX //avoid division by 0 - } else { - distance_between_candidate_and_point as f64 - / distance_between_candidate_and_existing_neighbor as f64 - }; + let factor = + if distance_between_candidate_and_existing_neighbor < 0.0 + f32::EPSILON { + if distance_between_candidate_and_point < 0.0 + f32::EPSILON { + 1.0 + } else { + f64::MAX + } + } else { + distance_between_candidate_and_point as f64 + / distance_between_candidate_and_existing_neighbor as f64 + }; max_factors[j] = max_factors[j].max(factor) } } @@ -662,6 +687,7 @@ pub trait Graph { let mut prune_neighbor_stats: PruneNeighborStats = PruneNeighborStats::new(); let mut greedy_search_stats = GreedySearchStats::new(); let meta_page = self.get_meta_page(index); + if self.is_empty() { self.set_neighbors( index, diff --git a/timescale_vector/src/access_method/model.rs b/timescale_vector/src/access_method/model.rs index 20e52397..a17a6d6f 100644 --- a/timescale_vector/src/access_method/model.rs +++ b/timescale_vector/src/access_method/model.rs @@ -15,6 +15,7 @@ use crate::util::{ ArchivedItemPointer, HeapPointer, IndexPointer, ItemPointer, ReadableBuffer, WritableBuffer, }; +use super::distance::preprocess_cosine; use super::meta_page::MetaPage; use super::quantizer::Quantizer; @@ -48,10 +49,11 @@ impl PgVector { casted } - pub fn to_slice(&self) -> &[f32] { + pub fn to_slice(&mut self) -> &[f32] { let dim = (*self).dim; - unsafe { (*self).x.as_slice(dim as _) } - // unsafe { std::slice::from_raw_parts((*self).x, (*self).dim as _) } + let raw_slice = unsafe { (*self).x.as_mut_slice(dim as _) }; + preprocess_cosine(raw_slice); + raw_slice } } diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 660ba2a2..bb3fb5f7 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -205,6 +205,8 @@ pub extern "C" fn amendscan(scan: pg_sys::IndexScanDesc) { mod tests { use pgrx::*; + //TODO: add test where inserting and querying with vectors that are all the same. + #[pg_test] unsafe fn test_index_scan() -> spi::Result<()> { Spi::run(&format!( @@ -212,7 +214,7 @@ mod tests { INSERT INTO test(embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[7,8,10]'); - INSERT INTO test(embedding) SELECT ('[' || g::text ||', 0, 0]')::vector FROM generate_series(0, 100) g; + INSERT INTO test(embedding) SELECT ('[' || g::text ||', 1.0, 1.0]')::vector FROM generate_series(0, 100) g; CREATE INDEX idxtest ON test From 30e807779ac949c967d3f8e548e0b97476a34b9e Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Mon, 18 Dec 2023 21:03:48 -0500 Subject: [PATCH 09/44] Switch to using euclidean distance for PQ always --- timescale_vector/src/access_method/pq.rs | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/timescale_vector/src/access_method/pq.rs b/timescale_vector/src/access_method/pq.rs index d0dac299..4ebff5c4 100644 --- a/timescale_vector/src/access_method/pq.rs +++ b/timescale_vector/src/access_method/pq.rs @@ -6,7 +6,10 @@ use rand::Rng; use reductive::pq::{Pq, QuantizeVector, TrainPq}; use crate::{ - access_method::model::{self, read_pq}, + access_method::{ + distance::distance_l2, + model::{self, read_pq}, + }, util::IndexPointer, }; @@ -105,7 +108,7 @@ impl PqTrainer { fn build_distance_table( pq: &Pq, query: &[f32], - distance_fn: fn(&[f32], &[f32]) -> f32, + _distance_fn: fn(&[f32], &[f32]) -> f32, ) -> Vec { let sq = pq.subquantizers(); let num_centroids = pq.n_quantizer_centroids(); @@ -118,7 +121,10 @@ fn build_distance_table( for (subquantizer_index, subquantizer) in sq.outer_iter().enumerate() { let sl = &query[subquantizer_index * ds..(subquantizer_index + 1) * ds]; for (centroid_index, c) in subquantizer.outer_iter().enumerate() { - let dist = distance_fn(sl, c.to_slice().unwrap()); + /* always use l2 for pq measurements since centeroids use k-means (which uses euclidean/l2 distance) + * The quantization also uses euclidean distance too. In the future we can experiment with k-mediods + * using a different distance measure, but this may make little difference. */ + let dist = distance_l2(sl, c.to_slice().unwrap()); assert!(subquantizer_index < num_subquantizers); assert!(centroid_index * num_subquantizers + subquantizer_index < dt_size); distance_table[centroid_index * num_subquantizers + subquantizer_index] = dist; From 08df297f835c1d5f8bd2da86685d75579c5f99fc Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Wed, 20 Dec 2023 14:40:44 -0500 Subject: [PATCH 10/44] Optimize distance calc for PQ PQ table creation calculates distance on low-dimensional arrays. This needs to be optimized separately. --- timescale_vector/benches/distance.rs | 46 +++++++++++++++- .../src/access_method/distance.rs | 54 +++++++++++++++++++ timescale_vector/src/access_method/pq.rs | 4 +- 3 files changed, 100 insertions(+), 4 deletions(-) diff --git a/timescale_vector/benches/distance.rs b/timescale_vector/benches/distance.rs index dfce659e..836c4b0b 100644 --- a/timescale_vector/benches/distance.rs +++ b/timescale_vector/benches/distance.rs @@ -1,5 +1,7 @@ use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use timescale_vector::access_method::distance::{distance_cosine, distance_l2}; +use timescale_vector::access_method::distance::{ + distance_cosine, distance_l2, distance_l2_optimized_for_few_dimensions, distance_l2_unoptimized, +}; //copy and use qdrants simd code, purely for benchmarking purposes //not used in the actual extension @@ -151,14 +153,54 @@ fn benchmark_distance(c: &mut Criterion) { }); } +#[inline(always)] +pub fn distance_l2_fixed_size_opt(a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), 6); + let norm: f32 = a[..6] + .iter() + .zip(b[..6].iter()) + .map(|t| (*t.0 as f32 - *t.1 as f32) * (*t.0 as f32 - *t.1 as f32)) + .sum(); + assert!(norm >= 0.); + //don't sqrt for performance. These are only used for ordering so sqrt not needed + norm +} + +//PQ uses l2 distance on small vectors (6 dims or so). Benchmark that. +fn benchmark_distance_few_dimensions(c: &mut Criterion) { + let r: Vec = (0..6).map(|v| v as f32 + 1000.1).collect(); + let l: Vec = (0..6).map(|v| v as f32 + 2000.2).collect(); + + let mut group = c.benchmark_group("Distance"); + group.bench_function("pq distance l2 optimized for many dimensions", |b| { + b.iter(|| distance_l2(black_box(&r), black_box(&l))) + }); + group.bench_function("pq distance l2 unoptimized", |b| { + b.iter(|| distance_l2_unoptimized(black_box(&r), black_box(&l))) + }); + group.bench_function( + "pq distance l2 auto-vectorized for 6 dimensionl arrays", + |b| b.iter(|| distance_l2_fixed_size_opt(black_box(&r), black_box(&l))), + ); + group.bench_function( + "pq distance l2 optimized for few dimensions (what's used in the code now)", + |b| b.iter(|| distance_l2_optimized_for_few_dimensions(black_box(&r), black_box(&l))), + ); +} + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] criterion_group!( benches, benchmark_distance, + benchmark_distance_few_dimensions, benchmark_distance_x86_unaligned_vectors, benchmark_distance_x86_aligned_vectors ); #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))] -criterion_group!(benches, benchmark_distance); +criterion_group!( + benches, + benchmark_distance, + benchmark_distance_few_dimensions +); criterion_main!(benches); diff --git a/timescale_vector/src/access_method/distance.rs b/timescale_vector/src/access_method/distance.rs index f30dccab..e66b7c53 100644 --- a/timescale_vector/src/access_method/distance.rs +++ b/timescale_vector/src/access_method/distance.rs @@ -32,6 +32,60 @@ pub fn distance_l2_unoptimized(a: &[f32], b: &[f32]) -> f32 { norm } +/* PQ computes distances on subsegments that have few dimensions (e.g. 6). This function optimizes that. +* We optimize by telling the compiler exactly how long the slices are. This allows the compiler to figure +* out SIMD optimizations. Look at the benchmark results. */ +#[inline] +pub fn distance_l2_optimized_for_few_dimensions(a: &[f32], b: &[f32]) -> f32 { + let norm: f32 = match a.len() { + 0 => 0., + 1 => a[..1] + .iter() + .zip(b[..1].iter()) + .map(|t| (*t.0 as f32 - *t.1 as f32) * (*t.0 as f32 - *t.1 as f32)) + .sum(), + 2 => a[..2] + .iter() + .zip(b[..2].iter()) + .map(|t| (*t.0 as f32 - *t.1 as f32) * (*t.0 as f32 - *t.1 as f32)) + .sum(), + 3 => a[..3] + .iter() + .zip(b[..3].iter()) + .map(|t| (*t.0 as f32 - *t.1 as f32) * (*t.0 as f32 - *t.1 as f32)) + .sum(), + 4 => a[..4] + .iter() + .zip(b[..4].iter()) + .map(|t| (*t.0 as f32 - *t.1 as f32) * (*t.0 as f32 - *t.1 as f32)) + .sum(), + 5 => a[..5] + .iter() + .zip(b[..5].iter()) + .map(|t| (*t.0 as f32 - *t.1 as f32) * (*t.0 as f32 - *t.1 as f32)) + .sum(), + 6 => a[..6] + .iter() + .zip(b[..6].iter()) + .map(|t| (*t.0 as f32 - *t.1 as f32) * (*t.0 as f32 - *t.1 as f32)) + .sum(), + 7 => a[..7] + .iter() + .zip(b[..7].iter()) + .map(|t| (*t.0 as f32 - *t.1 as f32) * (*t.0 as f32 - *t.1 as f32)) + .sum(), + 8 => a[..8] + .iter() + .zip(b[..8].iter()) + .map(|t| (*t.0 as f32 - *t.1 as f32) * (*t.0 as f32 - *t.1 as f32)) + .sum(), + _ => distance_l2(a, b), + }; + assert!(norm >= 0.); + //don't sqrt for performance. These are only used for ordering so sqrt not needed + norm +} + #[inline] pub fn distance_cosine(a: &[f32], b: &[f32]) -> f32 { #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] diff --git a/timescale_vector/src/access_method/pq.rs b/timescale_vector/src/access_method/pq.rs index 4ebff5c4..1a202143 100644 --- a/timescale_vector/src/access_method/pq.rs +++ b/timescale_vector/src/access_method/pq.rs @@ -7,7 +7,7 @@ use reductive::pq::{Pq, QuantizeVector, TrainPq}; use crate::{ access_method::{ - distance::distance_l2, + distance::distance_l2_optimized_for_few_dimensions, model::{self, read_pq}, }, util::IndexPointer, @@ -124,7 +124,7 @@ fn build_distance_table( /* always use l2 for pq measurements since centeroids use k-means (which uses euclidean/l2 distance) * The quantization also uses euclidean distance too. In the future we can experiment with k-mediods * using a different distance measure, but this may make little difference. */ - let dist = distance_l2(sl, c.to_slice().unwrap()); + let dist = distance_l2_optimized_for_few_dimensions(sl, c.to_slice().unwrap()); assert!(subquantizer_index < num_subquantizers); assert!(centroid_index * num_subquantizers + subquantizer_index < dt_size); distance_table[centroid_index * num_subquantizers + subquantizer_index] = dist; From 3e793a8a7923ef2379d7abdaaa66f4d162980b0b Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Thu, 21 Dec 2023 12:22:45 -0500 Subject: [PATCH 11/44] Initial bq implementation --- timescale_vector/src/access_method/bq.rs | 216 ++++++++++++++++++ timescale_vector/src/access_method/build.rs | 67 ++++++ .../src/access_method/builder_graph.rs | 9 + timescale_vector/src/access_method/graph.rs | 24 +- .../src/access_method/meta_page.rs | 9 +- timescale_vector/src/access_method/mod.rs | 1 + timescale_vector/src/access_method/model.rs | 10 + timescale_vector/src/access_method/options.rs | 48 +++- .../src/access_method/quantizer.rs | 3 +- timescale_vector/src/access_method/scan.rs | 1 + timescale_vector/src/util/page.rs | 2 + 11 files changed, 380 insertions(+), 10 deletions(-) create mode 100644 timescale_vector/src/access_method/bq.rs diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs new file mode 100644 index 00000000..dc83e485 --- /dev/null +++ b/timescale_vector/src/access_method/bq.rs @@ -0,0 +1,216 @@ +use std::pin::Pin; + +use pgrx::PgRelation; +use rkyv::{Archive, Deserialize, Serialize}; + +use crate::util::{page::PageType, tape::Tape, IndexPointer, ItemPointer, ReadableBuffer}; + +use super::meta_page::MetaPage; + +const BITS_STORE_TYPE_SIZE: usize = 8; + +#[derive(Archive, Deserialize, Serialize)] +#[archive(check_bytes)] +#[repr(C)] +pub struct BqMeans { + count: u64, + means: Vec, +} + +impl BqMeans { + pub unsafe fn write(&self, tape: &mut Tape) -> ItemPointer { + let bytes = rkyv::to_bytes::<_, 8192>(self).unwrap(); + tape.write(&bytes) + } + pub unsafe fn read<'a>( + index: &'a PgRelation, + index_pointer: &ItemPointer, + ) -> ReadableBqMeans<'a> { + let rb = index_pointer.read_bytes(index); + ReadableBqMeans { _rb: rb } + } +} + +//ReadablePqNode ties an archive node to it's underlying buffer +pub struct ReadableBqMeans<'a> { + _rb: ReadableBuffer<'a>, +} + +impl<'a> ReadableBqMeans<'a> { + pub fn get_archived_node(&self) -> &ArchivedBqMeans { + // checking the code here is expensive during build, so skip it. + // TODO: should we check the data during queries? + //rkyv::check_archived_root::(self._rb.get_data_slice()).unwrap() + unsafe { rkyv::archived_root::(self._rb.get_data_slice()) } + } +} + +pub unsafe fn read_bq(index: &PgRelation, index_pointer: &IndexPointer) -> (u64, Vec) { + let rpq = BqMeans::read(index, &index_pointer); + let rpn = rpq.get_archived_node(); + (rpn.count, rpn.means.as_slice().to_vec()) +} + +pub unsafe fn write_bq(index: &PgRelation, count: u64, means: &[f32]) -> ItemPointer { + let mut tape = Tape::new(index, PageType::BqMeans); + let node = BqMeans { + count, + means: means.to_vec(), + }; + let ptr = node.write(&mut tape); + tape.close(); + ptr +} + +pub struct BqQuantizer { + use_mean: bool, + count: u64, + mean: Vec, +} + +impl BqQuantizer { + pub fn new() -> BqQuantizer { + Self { + use_mean: true, + count: 0, + mean: vec![], + } + } + + pub fn load(&mut self, index_relation: &PgRelation, meta_page: &super::meta_page::MetaPage) { + if self.use_mean { + if meta_page.get_pq_pointer().is_none() { + pgrx::error!("No PQ pointer found in meta page"); + } + let pq_item_pointer = meta_page.get_pq_pointer().unwrap(); + (self.count, self.mean) = unsafe { read_bq(&index_relation, &pq_item_pointer) }; + } + } + + pub fn initialize_node( + &self, + node: &mut super::model::Node, + meta_page: &MetaPage, + full_vector: Vec, + ) { + if self.use_mean { + node.pq_vector = vec![0; Self::quantized_size(meta_page.get_num_dimensions() as _)]; + } else { + node.pq_vector = self.quantize(&full_vector); + } + } + + pub fn update_node_after_traing( + &self, + archived: &mut Pin<&mut super::model::ArchivedNode>, + full_vector: Vec, + ) { + if self.use_mean { + let bq_vector = self.quantize(&full_vector); + + assert!(bq_vector.len() == archived.pq_vector.len()); + for i in 0..=bq_vector.len() - 1 { + let mut pgv = archived.as_mut().pq_vectors().index_pin(i); + *pgv = bq_vector[i]; + } + } + } + + pub fn start_training(&mut self, meta_page: &super::meta_page::MetaPage) { + if self.use_mean { + self.count = 0; + self.mean = vec![0.0; meta_page.get_num_dimensions() as _]; + } + } + + pub fn add_sample(&mut self, sample: &[f32]) { + if self.use_mean { + self.count += 1; + assert!(self.mean.len() == sample.len()); + + self.mean + .iter_mut() + .zip(sample.iter()) + .for_each(|(m, s)| *m += ((s - *m) / self.count as f32)); + } + } + + pub fn finish_training(&mut self) {} + + pub fn write_metadata(&self, index: &PgRelation) { + if self.use_mean { + let index_pointer = unsafe { write_bq(&index, self.count, &self.mean) }; + super::meta_page::MetaPage::update_pq_pointer(&index, index_pointer); + } + } + + fn quantized_size(full_vector_size: usize) -> usize { + if full_vector_size % BITS_STORE_TYPE_SIZE == 0 { + full_vector_size / BITS_STORE_TYPE_SIZE + } else { + (full_vector_size / BITS_STORE_TYPE_SIZE) + 1 + } + } + + pub fn quantize(&self, full_vector: &[f32]) -> Vec { + if self.use_mean { + let mut res_vector = vec![0; Self::quantized_size(full_vector.len())]; + + for (i, &v) in full_vector.iter().enumerate() { + if v > self.mean[i] { + res_vector[i / BITS_STORE_TYPE_SIZE] |= 1 << (i % BITS_STORE_TYPE_SIZE); + } + } + + res_vector + } else { + let mut res_vector = vec![0; Self::quantized_size(full_vector.len())]; + + for (i, &v) in full_vector.iter().enumerate() { + if v > 0.0 { + res_vector[i / BITS_STORE_TYPE_SIZE] |= 1 << (i % BITS_STORE_TYPE_SIZE); + } + } + + res_vector + } + } + + pub fn get_distance_table( + &self, + query: &[f32], + distance_fn: fn(&[f32], &[f32]) -> f32, + ) -> BqDistanceTable { + BqDistanceTable::new(self.quantize(query)) + } +} + +/// DistanceCalculator encapsulates the code to generate distances between a PQ vector and a query. +pub struct BqDistanceTable { + quantized_vector: Vec, +} + +fn xor_unoptimized(v1: &[u8], v2: &[u8]) -> usize { + let mut result = 0; + for (b1, b2) in v1.iter().zip(v2.iter()) { + result += (b1 ^ b2).count_ones() as usize; + } + result +} + +impl BqDistanceTable { + pub fn new(query: Vec) -> BqDistanceTable { + BqDistanceTable { + quantized_vector: query, + } + } + + /// distance emits the sum of distances between each centroid in the quantized vector. + pub fn distance(&self, bq_vector: &[u8]) -> f32 { + let count_ones = xor_unoptimized(&self.quantized_vector, bq_vector); + //dot product is LOWER the more xors that lead to 1 becaues that means a negative times a positive = negative component + //but the distance is 1 - dot product, so the more count_ones the higher the distance. + // one other check for distance(a,a), xor=0, count_ones=0, distance=0 + count_ones as f32 + } +} diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index 4ef76622..02277b2e 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -42,6 +42,9 @@ impl<'a, 'b> BuildState<'a, 'b> { Quantizer::PQ(pq) => { pq.start_training(&meta_page); } + Quantizer::BQ(bq) => { + bq.start_training(&meta_page); + } } //TODO: some ways to get rid of meta_page.clone? BuildState { @@ -126,6 +129,9 @@ pub unsafe extern "C" fn aminsert( Quantizer::PQ(pq) => { pq.load(&index_relation, &meta_page); } + Quantizer::BQ(bq) => { + bq.load(&index_relation, &meta_page); + } } let vp = VectorProvider::new( @@ -188,6 +194,9 @@ fn do_heap_scan<'a>( Quantizer::PQ(pq) => { pq.finish_training(); } + Quantizer::BQ(bq) => { + bq.finish_training(); + } } let write_stats = unsafe { state.node_builder.write(index_relation, &state.quantizer) }; @@ -221,6 +230,9 @@ fn do_heap_scan<'a>( Quantizer::PQ(pq) => { pq.write_metadata(index_relation); } + Quantizer::BQ(bq) => { + bq.write_metadata(index_relation); + } } ntuples @@ -279,6 +291,9 @@ fn build_callback_internal( Quantizer::PQ(pq) => { pq.add_sample(vector); } + Quantizer::BQ(bq) => { + bq.add_sample(vector); + } } let node = model::Node::new( @@ -371,6 +386,58 @@ mod tests { Ok(()) } + #[pg_test] + unsafe fn test_bq_index_creation() -> spi::Result<()> { + Spi::run(&format!( + "CREATE TABLE test_bq ( + embedding vector (1536) + ); + + -- generate 300 vectors + INSERT INTO test_bq (embedding) + SELECT + * + FROM ( + SELECT + ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding + FROM + generate_series(1, 1536 * 300) i + GROUP BY + i % 300) g; + + CREATE INDEX idx_tsv_pq ON test_bq USING tsv (embedding) WITH (num_neighbors = 64, search_list_size = 125, max_alpha = 1.0, use_bq = TRUE); + + ; + + SET enable_seqscan = 0; + -- perform index scans on the vectors + SELECT + * + FROM + test_bq + ORDER BY + embedding <=> ( + SELECT + ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding + FROM generate_series(1, 1536)); + + EXPLAIN ANALYZE + SELECT + * + FROM + test_bq + ORDER BY + embedding <=> ( + SELECT + ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding + FROM generate_series(1, 1536)); + + DROP INDEX idx_tsv_pq; + ", + ))?; + Ok(()) + } + #[pg_test] unsafe fn test_insert() -> spi::Result<()> { Spi::run(&format!( diff --git a/timescale_vector/src/access_method/builder_graph.rs b/timescale_vector/src/access_method/builder_graph.rs index f14c8e0c..be95950f 100644 --- a/timescale_vector/src/access_method/builder_graph.rs +++ b/timescale_vector/src/access_method/builder_graph.rs @@ -69,6 +69,15 @@ impl<'a> BuilderGraph<'a> { let full_vector = self.get_full_vector(heap_pointer); pq.update_node_after_traing(&mut archived, full_vector); } + Quantizer::BQ(bq) => { + //TODO: OPT: this may not be needed + let heap_pointer = node + .get_archived_node() + .heap_item_pointer + .deserialize_item_pointer(); + let full_vector = self.get_full_vector(heap_pointer); + bq.update_node_after_traing(&mut archived, full_vector); + } }; node.commit(); } diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index 95055836..53662c58 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -200,6 +200,7 @@ pub struct FullVectorDistanceState<'a> { pub struct DistanceMeasure { pq_distance_table: Option, + bq_distance_table: Option, } impl DistanceMeasure { @@ -207,25 +208,41 @@ impl DistanceMeasure { if !calc_distance_with_quantizer { return Self { pq_distance_table: None, + bq_distance_table: None, }; } match quantizer { Quantizer::None => Self { pq_distance_table: None, + bq_distance_table: None, }, Quantizer::PQ(pq) => { let dc = pq.get_distance_table(query, default_distance); Self { pq_distance_table: Some(dc), + bq_distance_table: None, + } + } + Quantizer::BQ(bq) => { + let dc = bq.get_distance_table(query, default_distance); + Self { + pq_distance_table: None, + bq_distance_table: Some(dc), } } } } fn get_quantized_distance(&self, vec: &[u8]) -> f32 { - let dc = self.pq_distance_table.as_ref().unwrap(); - let distance = dc.distance(vec); - distance + if self.pq_distance_table.is_some() { + let dc = self.pq_distance_table.as_ref().unwrap(); + let distance = dc.distance(vec); + distance + } else { + let dc = self.bq_distance_table.as_ref().unwrap(); + let distance = dc.distance(vec); + distance + } } fn get_full_vector_distance(&self, vec: &[f32], query: &[f32]) -> f32 { @@ -289,6 +306,7 @@ impl ListSearchResult { max_history_size: None, dm: DistanceMeasure { pq_distance_table: None, + bq_distance_table: None, }, stats: GreedySearchStats::new(), } diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index 256cf440..83a2ef54 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -5,6 +5,7 @@ use crate::access_method::options::TSVIndexOptions; use crate::util::page; use crate::util::*; +use super::bq::BqQuantizer; use super::pq::PqQuantizer; use super::quantizer::Quantizer; @@ -31,6 +32,7 @@ pub struct MetaPage { pq_vector_length: usize, pq_block_number: pg_sys::BlockNumber, pq_block_offset: pg_sys::OffsetNumber, + use_bq: bool, } impl MetaPage { @@ -65,6 +67,8 @@ impl MetaPage { pub fn get_quantizer(&self) -> Quantizer { if self.get_use_pq() { Quantizer::PQ(PqQuantizer::new()) + } else if self.use_bq { + Quantizer::BQ(BqQuantizer::new()) } else { Quantizer::None } @@ -84,7 +88,9 @@ impl MetaPage { } pub fn get_pq_pointer(&self) -> Option { - if !self.use_pq || (self.pq_block_number == 0 && self.pq_block_offset == 0) { + if (!self.use_pq && !self.use_bq) + || (self.pq_block_number == 0 && self.pq_block_offset == 0) + { return None; } @@ -128,6 +134,7 @@ impl MetaPage { (*meta).pq_block_offset = 0; (*meta).init_ids_block_number = 0; (*meta).init_ids_offset = 0; + (*meta).use_bq = (*opt).use_bq; let header = page.cast::(); let meta_end = (meta as Pointer).add(std::mem::size_of::()); diff --git a/timescale_vector/src/access_method/mod.rs b/timescale_vector/src/access_method/mod.rs index 9db1bcd9..6ae883fa 100644 --- a/timescale_vector/src/access_method/mod.rs +++ b/timescale_vector/src/access_method/mod.rs @@ -15,6 +15,7 @@ mod vacuum; extern crate blas_src; +mod bq; pub mod distance; #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] mod distance_x86; diff --git a/timescale_vector/src/access_method/model.rs b/timescale_vector/src/access_method/model.rs index a17a6d6f..b10d0bfa 100644 --- a/timescale_vector/src/access_method/model.rs +++ b/timescale_vector/src/access_method/model.rs @@ -125,6 +125,16 @@ impl Node { pq.initialize_node(&mut node, meta_page, full_vector); node } + Quantizer::BQ(bq) => { + let mut node = Self { + vector: Vec::with_capacity(0), + pq_vector: Vec::with_capacity(0), + neighbor_index_pointers: neighbor_index_pointers, + heap_item_pointer, + }; + bq.initialize_node(&mut node, meta_page, full_vector); + node + } } } diff --git a/timescale_vector/src/access_method/options.rs b/timescale_vector/src/access_method/options.rs index b225295d..45b7f8ab 100644 --- a/timescale_vector/src/access_method/options.rs +++ b/timescale_vector/src/access_method/options.rs @@ -13,6 +13,7 @@ pub struct TSVIndexOptions { pub search_list_size: u32, pub max_alpha: f64, pub use_pq: bool, + pub use_bq: bool, pub pq_vector_length: usize, } @@ -24,10 +25,11 @@ impl TSVIndexOptions { // use defaults let mut ops = unsafe { PgBox::::alloc0() }; ops.num_neighbors = 50; - ops.search_list_size = 65; + ops.search_list_size = 100; ops.max_alpha = 1.0; ops.use_pq = false; - ops.pq_vector_length = 64; + ops.use_bq = false; + ops.pq_vector_length = 256; unsafe { set_varsize( ops.as_ptr().cast(), @@ -41,7 +43,7 @@ impl TSVIndexOptions { } } -const NUM_REL_OPTS: usize = 5; +const NUM_REL_OPTS: usize = 6; static mut RELOPT_KIND_TSV: pg_sys::relopt_kind = 0; #[allow(clippy::unneeded_field_pattern)] // b/c of offset_of!() @@ -72,6 +74,11 @@ pub unsafe extern "C" fn amoptions( opttype: pg_sys::relopt_type_RELOPT_TYPE_BOOL, offset: offset_of!(TSVIndexOptions, use_pq) as i32, }, + pg_sys::relopt_parse_elt { + optname: "use_bq".as_pg_cstr(), + opttype: pg_sys::relopt_type_RELOPT_TYPE_BOOL, + offset: offset_of!(TSVIndexOptions, use_bq) as i32, + }, pg_sys::relopt_parse_elt { optname: "pq_vector_length".as_pg_cstr(), opttype: pg_sys::relopt_type_RELOPT_TYPE_INT, @@ -142,6 +149,13 @@ pub unsafe fn init() { false, pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, ); + pg_sys::add_bool_reloption( + RELOPT_KIND_TSV, + "use_bq".as_pg_cstr(), + "Enable binary quantization".as_pg_cstr(), + false, + pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, + ); pg_sys::add_int_reloption( RELOPT_KIND_TSV, "pq_vector_length".as_pg_cstr(), @@ -191,10 +205,34 @@ mod tests { let indexrel = PgRelation::from_pg(pg_sys::RelationIdGetRelation(index_oid)); let options = TSVIndexOptions::from_relation(&indexrel); assert_eq!(options.num_neighbors, 50); - assert_eq!(options.search_list_size, 65); + assert_eq!(options.search_list_size, 100); + assert_eq!(options.max_alpha, 1.0); + assert_eq!(options.use_pq, false); + assert_eq!(options.use_bq, false); + assert_eq!(options.pq_vector_length, 256); + Ok(()) + } + + #[pg_test] + unsafe fn test_index_options_bq() -> spi::Result<()> { + Spi::run(&format!( + "CREATE TABLE test(encoding vector(3)); + CREATE INDEX idxtest + ON test + USING tsv(encoding) + WITH (use_bq = TRUE);", + ))?; + + let index_oid = + Spi::get_one::("SELECT 'idxtest'::regclass::oid")?.expect("oid was null"); + let indexrel = PgRelation::from_pg(pg_sys::RelationIdGetRelation(index_oid)); + let options = TSVIndexOptions::from_relation(&indexrel); + assert_eq!(options.num_neighbors, 50); + assert_eq!(options.search_list_size, 100); assert_eq!(options.max_alpha, 1.0); assert_eq!(options.use_pq, false); - assert_eq!(options.pq_vector_length, 64); + assert_eq!(options.use_bq, true); + assert_eq!(options.pq_vector_length, 256); Ok(()) } } diff --git a/timescale_vector/src/access_method/quantizer.rs b/timescale_vector/src/access_method/quantizer.rs index 5dee6053..c11e6098 100644 --- a/timescale_vector/src/access_method/quantizer.rs +++ b/timescale_vector/src/access_method/quantizer.rs @@ -1,4 +1,4 @@ -use super::pq::PqQuantizer; +use super::{bq::BqQuantizer, pq::PqQuantizer}; /*pub trait Quantizer { fn initialize_node(&self, node: &mut Node, meta_page: &MetaPage); @@ -8,6 +8,7 @@ use super::pq::PqQuantizer; }*/ pub enum Quantizer { + BQ(BqQuantizer), PQ(PqQuantizer), None, } diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index bb3fb5f7..61125bc7 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -26,6 +26,7 @@ impl<'a> TSVResponseIterator<'a> { match &mut quantizer { Quantizer::None => {} Quantizer::PQ(pq) => pq.load(index, &meta_page), + Quantizer::BQ(bq) => bq.load(index, &meta_page), } let graph = DiskIndexGraph::new( &index, diff --git a/timescale_vector/src/util/page.rs b/timescale_vector/src/util/page.rs index bfe09d34..64a758ab 100644 --- a/timescale_vector/src/util/page.rs +++ b/timescale_vector/src/util/page.rs @@ -26,6 +26,7 @@ pub enum PageType { Node = 1, PqQuantizerDef = 2, PqQuantizerVector = 3, + BqMeans = 4, } impl PageType { @@ -35,6 +36,7 @@ impl PageType { 1 => PageType::Node, 2 => PageType::PqQuantizerDef, 3 => PageType::PqQuantizerVector, + 4 => PageType::BqMeans, _ => panic!("Unknown PageType number {}", value), } } From 5c30fcf2fe2e1d89fb4b41ac78e9b21f299cb50a Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Wed, 3 Jan 2024 22:15:38 -0500 Subject: [PATCH 12/44] Big refactor to break out storage properly and fix architecture --- README.md | 2 +- timescale_vector/Cargo.toml | 2 +- timescale_vector/src/access_method/bq.rs | 856 +++++++++++++++--- timescale_vector/src/access_method/build.rs | 611 ++++++++----- .../src/access_method/builder_graph.rs | 174 ---- .../src/access_method/debugging.rs | 22 +- .../src/access_method/disk_index_graph.rs | 92 -- .../src/access_method/distance.rs | 33 +- .../src/access_method/distance_x86.rs | 2 +- timescale_vector/src/access_method/graph.rs | 767 +++++----------- .../src/access_method/graph_neighbor_store.rs | 124 +++ .../src/access_method/meta_page.rs | 12 +- timescale_vector/src/access_method/mod.rs | 16 +- timescale_vector/src/access_method/model.rs | 423 --------- .../access_method/neighbor_with_distance.rs | 56 ++ .../src/access_method/pg_vector.rs | 76 ++ .../src/access_method/plain_node.rs | 156 ++++ .../src/access_method/plain_storage.rs | 353 ++++++++ .../access_method/{pq.rs => pq_quantizer.rs} | 114 +-- .../src/access_method/pq_quantizer_storage.rs | 123 +++ .../src/access_method/pq_storage.rs | 446 +++++++++ .../src/access_method/quantizer.rs | 23 - timescale_vector/src/access_method/scan.rs | 315 ++++--- timescale_vector/src/access_method/stats.rs | 238 +++++ timescale_vector/src/access_method/storage.rs | 141 +++ .../src/access_method/storage_common.rs | 74 ++ timescale_vector/src/access_method/vacuum.rs | 225 +++-- timescale_vector/src/util/buffer.rs | 17 +- timescale_vector/src/util/mod.rs | 1 + timescale_vector/src/util/page.rs | 2 + timescale_vector/src/util/table_slot.rs | 60 ++ .../timescale_vector_derive/Cargo.toml | 11 + .../timescale_vector_derive/src/lib.rs | 91 ++ 33 files changed, 3739 insertions(+), 1919 deletions(-) delete mode 100644 timescale_vector/src/access_method/builder_graph.rs delete mode 100644 timescale_vector/src/access_method/disk_index_graph.rs create mode 100644 timescale_vector/src/access_method/graph_neighbor_store.rs delete mode 100644 timescale_vector/src/access_method/model.rs create mode 100644 timescale_vector/src/access_method/neighbor_with_distance.rs create mode 100644 timescale_vector/src/access_method/pg_vector.rs create mode 100644 timescale_vector/src/access_method/plain_node.rs create mode 100644 timescale_vector/src/access_method/plain_storage.rs rename timescale_vector/src/access_method/{pq.rs => pq_quantizer.rs} (83%) create mode 100644 timescale_vector/src/access_method/pq_quantizer_storage.rs create mode 100644 timescale_vector/src/access_method/pq_storage.rs delete mode 100644 timescale_vector/src/access_method/quantizer.rs create mode 100644 timescale_vector/src/access_method/stats.rs create mode 100644 timescale_vector/src/access_method/storage.rs create mode 100644 timescale_vector/src/access_method/storage_common.rs create mode 100644 timescale_vector/src/util/table_slot.rs create mode 100644 timescale_vector/timescale_vector_derive/Cargo.toml create mode 100644 timescale_vector/timescale_vector_derive/src/lib.rs diff --git a/README.md b/README.md index b34ecee9..33524b31 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ Next you need cargo-pgx, which can be installed with cargo install --locked cargo-pgrx ``` -You must reinstall cargo-pgx whenever you update your Rust compiler, since cargo-pgx needs to be built with the same compiler as Toolkit. +You must reinstall cargo-pgx whenever you update your Rust compiler, since cargo-pgx needs to be built with the same compiler as Timescale Vector. Finally, setup the pgx development environment with ```shell diff --git a/timescale_vector/Cargo.toml b/timescale_vector/Cargo.toml index 625362e5..9dc6952c 100644 --- a/timescale_vector/Cargo.toml +++ b/timescale_vector/Cargo.toml @@ -26,7 +26,7 @@ rand_chacha = "0.3" rand_core = "0.6" rand_xorshift = "0.3" rayon = "1" - +timescale_vector_derive = { path = "timescale_vector_derive" } [dev-dependencies] pgrx-tests = "=0.11.2" diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index dc83e485..ec652cca 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -1,15 +1,35 @@ -use std::pin::Pin; - -use pgrx::PgRelation; -use rkyv::{Archive, Deserialize, Serialize}; - -use crate::util::{page::PageType, tape::Tape, IndexPointer, ItemPointer, ReadableBuffer}; - -use super::meta_page::MetaPage; - +use super::{ + distance::distance_cosine as default_distance, + graph::{ListSearchNeighbor, ListSearchResult}, + graph_neighbor_store::GraphNeighborStore, + pg_vector::PgVector, + stats::{ + GreedySearchStats, StatsDistanceComparison, StatsNodeModify, StatsNodeRead, StatsNodeWrite, + WriteStats, + }, + storage::{ArchivedData, NodeFullDistanceMeasure, Storage, StorageFullDistanceFromHeap}, + storage_common::{calculate_full_distance, HeapFullDistanceMeasure}, +}; +use std::{collections::HashMap, iter::once, marker::PhantomData, pin::Pin}; + +use pgrx::{ + pg_sys::{InvalidBlockNumber, InvalidOffsetNumber}, + PgRelation, +}; +use rkyv::{vec::ArchivedVec, Archive, Archived, Deserialize, Serialize}; + +use crate::util::{ + page::PageType, table_slot::TableSlot, tape::Tape, ArchivedItemPointer, HeapPointer, + IndexPointer, ItemPointer, ReadableBuffer, +}; + +use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance}; +use crate::util::WritableBuffer; + +type BqVectorElement = u8; const BITS_STORE_TYPE_SIZE: usize = 8; -#[derive(Archive, Deserialize, Serialize)] +#[derive(Archive, Deserialize, Serialize, Readable, Writeable)] #[archive(check_bytes)] #[repr(C)] pub struct BqMeans { @@ -18,112 +38,105 @@ pub struct BqMeans { } impl BqMeans { - pub unsafe fn write(&self, tape: &mut Tape) -> ItemPointer { - let bytes = rkyv::to_bytes::<_, 8192>(self).unwrap(); - tape.write(&bytes) - } - pub unsafe fn read<'a>( - index: &'a PgRelation, - index_pointer: &ItemPointer, - ) -> ReadableBqMeans<'a> { - let rb = index_pointer.read_bytes(index); - ReadableBqMeans { _rb: rb } - } -} - -//ReadablePqNode ties an archive node to it's underlying buffer -pub struct ReadableBqMeans<'a> { - _rb: ReadableBuffer<'a>, -} + pub unsafe fn load( + index: &PgRelation, + meta_page: &super::meta_page::MetaPage, + stats: &mut S, + ) -> BqQuantizer { + let mut quantizer = BqQuantizer::new(); + if quantizer.use_mean { + if meta_page.get_pq_pointer().is_none() { + pgrx::error!("No BQ pointer found in meta page"); + } + let pq_item_pointer = meta_page.get_pq_pointer().unwrap(); + let rpq = BqMeans::read(index, pq_item_pointer, stats); + let rpn = rpq.get_archived_node(); -impl<'a> ReadableBqMeans<'a> { - pub fn get_archived_node(&self) -> &ArchivedBqMeans { - // checking the code here is expensive during build, so skip it. - // TODO: should we check the data during queries? - //rkyv::check_archived_root::(self._rb.get_data_slice()).unwrap() - unsafe { rkyv::archived_root::(self._rb.get_data_slice()) } + quantizer.load(rpn.count, rpn.means.to_vec()); + } + quantizer } -} - -pub unsafe fn read_bq(index: &PgRelation, index_pointer: &IndexPointer) -> (u64, Vec) { - let rpq = BqMeans::read(index, &index_pointer); - let rpn = rpq.get_archived_node(); - (rpn.count, rpn.means.as_slice().to_vec()) -} -pub unsafe fn write_bq(index: &PgRelation, count: u64, means: &[f32]) -> ItemPointer { - let mut tape = Tape::new(index, PageType::BqMeans); - let node = BqMeans { - count, - means: means.to_vec(), - }; - let ptr = node.write(&mut tape); - tape.close(); - ptr + pub unsafe fn store( + index: &PgRelation, + quantizer: &BqQuantizer, + stats: &mut S, + ) -> ItemPointer { + let mut tape = Tape::new(index, PageType::BqMeans); + let node = BqMeans { + count: quantizer.count, + means: quantizer.mean.to_vec(), + }; + let ptr = node.write(&mut tape, stats); + tape.close(); + ptr + } } +#[derive(Clone)] pub struct BqQuantizer { - use_mean: bool, - count: u64, - mean: Vec, + pub use_mean: bool, + training: bool, + pub count: u64, + pub mean: Vec, } impl BqQuantizer { - pub fn new() -> BqQuantizer { + fn new() -> BqQuantizer { Self { use_mean: true, + training: false, count: 0, mean: vec![], } } - pub fn load(&mut self, index_relation: &PgRelation, meta_page: &super::meta_page::MetaPage) { - if self.use_mean { - if meta_page.get_pq_pointer().is_none() { - pgrx::error!("No PQ pointer found in meta page"); - } - let pq_item_pointer = meta_page.get_pq_pointer().unwrap(); - (self.count, self.mean) = unsafe { read_bq(&index_relation, &pq_item_pointer) }; - } + fn load(&mut self, count: u64, mean: Vec) { + self.count = count; + self.mean = mean; } - pub fn initialize_node( - &self, - node: &mut super::model::Node, - meta_page: &MetaPage, - full_vector: Vec, - ) { - if self.use_mean { - node.pq_vector = vec![0; Self::quantized_size(meta_page.get_num_dimensions() as _)]; + fn quantized_size(full_vector_size: usize) -> usize { + if full_vector_size % BITS_STORE_TYPE_SIZE == 0 { + full_vector_size / BITS_STORE_TYPE_SIZE } else { - node.pq_vector = self.quantize(&full_vector); + (full_vector_size / BITS_STORE_TYPE_SIZE) + 1 } } - pub fn update_node_after_traing( - &self, - archived: &mut Pin<&mut super::model::ArchivedNode>, - full_vector: Vec, - ) { + fn quantize(&self, full_vector: &[f32]) -> Vec { if self.use_mean { - let bq_vector = self.quantize(&full_vector); + let mut res_vector = vec![0; Self::quantized_size(full_vector.len())]; + + for (i, &v) in full_vector.iter().enumerate() { + if v > self.mean[i] { + res_vector[i / BITS_STORE_TYPE_SIZE] |= 1 << (i % BITS_STORE_TYPE_SIZE); + } + } + + res_vector + } else { + let mut res_vector = vec![0; Self::quantized_size(full_vector.len())]; - assert!(bq_vector.len() == archived.pq_vector.len()); - for i in 0..=bq_vector.len() - 1 { - let mut pgv = archived.as_mut().pq_vectors().index_pin(i); - *pgv = bq_vector[i]; + for (i, &v) in full_vector.iter().enumerate() { + if v > 0.0 { + res_vector[i / BITS_STORE_TYPE_SIZE] |= 1 << (i % BITS_STORE_TYPE_SIZE); + } } + + res_vector } } - pub fn start_training(&mut self, meta_page: &super::meta_page::MetaPage) { + fn start_training(&mut self, meta_page: &super::meta_page::MetaPage) { + self.training = true; if self.use_mean { self.count = 0; self.mean = vec![0.0; meta_page.get_num_dimensions() as _]; } } - pub fn add_sample(&mut self, sample: &[f32]) { + fn add_sample(&mut self, sample: &[f32]) { if self.use_mean { self.count += 1; assert!(self.mean.len() == sample.len()); @@ -131,66 +144,45 @@ impl BqQuantizer { self.mean .iter_mut() .zip(sample.iter()) - .for_each(|(m, s)| *m += ((s - *m) / self.count as f32)); + .for_each(|(m, s)| *m += (s - *m) / self.count as f32); } } - pub fn finish_training(&mut self) {} - - pub fn write_metadata(&self, index: &PgRelation) { - if self.use_mean { - let index_pointer = unsafe { write_bq(&index, self.count, &self.mean) }; - super::meta_page::MetaPage::update_pq_pointer(&index, index_pointer); - } + fn finish_training(&mut self) { + self.training = false; } - fn quantized_size(full_vector_size: usize) -> usize { - if full_vector_size % BITS_STORE_TYPE_SIZE == 0 { - full_vector_size / BITS_STORE_TYPE_SIZE + fn vector_for_new_node( + &self, + meta_page: &super::meta_page::MetaPage, + full_vector: &[f32], + ) -> Vec { + if self.use_mean && self.training { + vec![0; BqQuantizer::quantized_size(meta_page.get_num_dimensions() as _)] } else { - (full_vector_size / BITS_STORE_TYPE_SIZE) + 1 + self.quantize(&full_vector) } } - pub fn quantize(&self, full_vector: &[f32]) -> Vec { - if self.use_mean { - let mut res_vector = vec![0; Self::quantized_size(full_vector.len())]; - - for (i, &v) in full_vector.iter().enumerate() { - if v > self.mean[i] { - res_vector[i / BITS_STORE_TYPE_SIZE] |= 1 << (i % BITS_STORE_TYPE_SIZE); - } - } - - res_vector - } else { - let mut res_vector = vec![0; Self::quantized_size(full_vector.len())]; - - for (i, &v) in full_vector.iter().enumerate() { - if v > 0.0 { - res_vector[i / BITS_STORE_TYPE_SIZE] |= 1 << (i % BITS_STORE_TYPE_SIZE); - } - } - - res_vector - } + fn vector_needs_update_after_training(&self) -> bool { + self.use_mean } - pub fn get_distance_table( + fn get_distance_table( &self, query: &[f32], - distance_fn: fn(&[f32], &[f32]) -> f32, + _distance_fn: fn(&[f32], &[f32]) -> f32, ) -> BqDistanceTable { BqDistanceTable::new(self.quantize(query)) } } -/// DistanceCalculator encapsulates the code to generate distances between a PQ vector and a query. +/// DistanceCalculator encapsulates the code to generate distances between a BQ vector and a query. pub struct BqDistanceTable { - quantized_vector: Vec, + quantized_vector: Vec, } -fn xor_unoptimized(v1: &[u8], v2: &[u8]) -> usize { +fn xor_unoptimized(v1: &[BqVectorElement], v2: &[BqVectorElement]) -> usize { let mut result = 0; for (b1, b2) in v1.iter().zip(v2.iter()) { result += (b1 ^ b2).count_ones() as usize; @@ -199,14 +191,14 @@ fn xor_unoptimized(v1: &[u8], v2: &[u8]) -> usize { } impl BqDistanceTable { - pub fn new(query: Vec) -> BqDistanceTable { + pub fn new(query: Vec) -> BqDistanceTable { BqDistanceTable { quantized_vector: query, } } /// distance emits the sum of distances between each centroid in the quantized vector. - pub fn distance(&self, bq_vector: &[u8]) -> f32 { + pub fn distance(&self, bq_vector: &[BqVectorElement]) -> f32 { let count_ones = xor_unoptimized(&self.quantized_vector, bq_vector); //dot product is LOWER the more xors that lead to 1 becaues that means a negative times a positive = negative component //but the distance is 1 - dot product, so the more count_ones the higher the distance. @@ -214,3 +206,621 @@ impl BqDistanceTable { count_ones as f32 } } + +pub enum BqSearchDistanceMeasure { + Full(PgVector), + Bq(BqDistanceTable), +} + +impl BqSearchDistanceMeasure { + pub fn calculate_bq_distance( + table: &BqDistanceTable, + bq_vector: &[BqVectorElement], + stats: &mut S, + ) -> f32 { + assert!(bq_vector.len() > 0); + let vec = bq_vector; + stats.record_quantized_distance_comparison(); + table.distance(vec) + } +} + +struct QuantizedVectorCache { + quantized_vector_map: HashMap>, +} + +/* should be a LRU cache for quantized vector. For now cheat and never evict + TODO: implement LRU cache +*/ +impl QuantizedVectorCache { + fn new(capacity: usize) -> Self { + Self { + quantized_vector_map: HashMap::with_capacity(capacity), + } + } + + fn get( + &mut self, + index_pointer: IndexPointer, + storage: &BqSpeedupStorage, + stats: &mut S, + ) -> &[BqVectorElement] { + self.quantized_vector_map + .entry(index_pointer) + .or_insert_with(|| { + storage.get_quantized_vector_from_index_pointer(index_pointer, stats) + }) + } + + fn must_get(&self, index_pointer: IndexPointer) -> &[BqVectorElement] { + self.quantized_vector_map.get(&index_pointer).unwrap() + } + + /* Ensure that all these elements are in the cache. If the capacity isn't big enough throw an error. + must_get must succeed on all the elements after this call prior to another get or preload call */ + + fn preload, S: StatsNodeRead>( + &mut self, + index_pointers: I, + storage: &BqSpeedupStorage, + stats: &mut S, + ) { + for index_pointer in index_pointers { + self.get(index_pointer, storage, stats); + } + } +} + +pub struct BqSpeedupStorage<'a> { + pub index: &'a PgRelation, + pub distance_fn: fn(&[f32], &[f32]) -> f32, + quantizer: BqQuantizer, + heap_rel: Option<&'a PgRelation>, + heap_attr: Option, + qv_cache: Option, +} + +impl<'a> BqSpeedupStorage<'a> { + pub fn new_for_build( + index: &'a PgRelation, + heap_rel: &'a PgRelation, + heap_attr: pgrx::pg_sys::AttrNumber, + ) -> BqSpeedupStorage<'a> { + Self { + index: index, + distance_fn: default_distance, + quantizer: BqQuantizer::new(), + heap_rel: Some(heap_rel), + heap_attr: Some(heap_attr), + qv_cache: None, + } + } + + fn load_quantizer( + index_relation: &PgRelation, + meta_page: &super::meta_page::MetaPage, + stats: &mut S, + ) -> BqQuantizer { + unsafe { BqMeans::load(&index_relation, meta_page, stats) } + } + + pub fn load_for_insert( + heap_rel: &'a PgRelation, + heap_attr: pgrx::pg_sys::AttrNumber, + index_relation: &'a PgRelation, + meta_page: &super::meta_page::MetaPage, + stats: &mut S, + ) -> BqSpeedupStorage<'a> { + Self { + index: index_relation, + distance_fn: default_distance, + quantizer: Self::load_quantizer(index_relation, meta_page, stats), + heap_rel: Some(heap_rel), + heap_attr: Some(heap_attr), + qv_cache: None, + } + } + + pub fn load_for_search( + index_relation: &'a PgRelation, + quantizer: &BqQuantizer, + ) -> BqSpeedupStorage<'a> { + Self { + index: index_relation, + distance_fn: default_distance, + //OPT: get rid of clone + quantizer: quantizer.clone(), + heap_rel: None, + heap_attr: None, + qv_cache: None, + } + } + + fn get_quantized_vector_from_index_pointer( + &self, + index_pointer: IndexPointer, + stats: &mut S, + ) -> Vec { + let slot = unsafe { self.get_heap_table_slot_from_index_pointer(index_pointer, stats) }; + let slice = unsafe { slot.get_pg_vector() }; + self.quantizer.quantize(slice.to_slice()) + } + + fn write_quantizer_metadata(&self, stats: &mut S) { + if self.quantizer.use_mean { + let index_pointer = unsafe { BqMeans::store(&self.index, &self.quantizer, stats) }; + super::meta_page::MetaPage::update_pq_pointer(&self.index, index_pointer); + } + } + + fn visit_lsn_internal( + &self, + lsr: &mut ListSearchResult< + as Storage>::QueryDistanceMeasure, + as Storage>::LSNPrivateData, + >, + lsn_index_pointer: IndexPointer, + gns: &GraphNeighborStore, + ) { + //Opt shouldn't need to read the node in the builder graph case. + let rn_visiting = unsafe { BqNode::read(self.index, lsn_index_pointer, &mut lsr.stats) }; + let node_visiting = rn_visiting.get_archived_node(); + + let neighbors = match gns { + GraphNeighborStore::Disk => node_visiting.get_index_pointer_to_neighbors(), + GraphNeighborStore::Builder(b) => b.get_neighbors(lsn_index_pointer), + }; + + for (i, &neighbor_index_pointer) in neighbors.iter().enumerate() { + if !lsr.prepare_insert(neighbor_index_pointer) { + continue; + } + + let distance = match lsr.sdm.as_ref().unwrap() { + BqSearchDistanceMeasure::Full(query) => { + let rn_neighbor = + unsafe { BqNode::read(self.index, neighbor_index_pointer, &mut lsr.stats) }; + let node_neighbor = rn_neighbor.get_archived_node(); + if node_neighbor.is_deleted() { + self.visit_lsn_internal(lsr, neighbor_index_pointer, gns); + continue; + } + let heap_pointer_neighbor = + node_neighbor.heap_item_pointer.deserialize_item_pointer(); + unsafe { + calculate_full_distance( + self, + heap_pointer_neighbor, + query.to_slice(), + &mut lsr.stats, + ) + } + } + BqSearchDistanceMeasure::Bq(table) => { + /* Note: there is no additional node reads here. We get all of our info from node_visiting + * This is what gives us a speedup in BQ Speedup */ + if let GraphNeighborStore::Builder(_) = gns { + assert!( + false, + "BQ distance should not be used with the builder graph store" + ) + } + let bq_vector = node_visiting.neighbor_vectors[i].as_slice(); + BqSearchDistanceMeasure::calculate_bq_distance(table, bq_vector, &mut lsr.stats) + } + }; + let lsn = + ListSearchNeighbor::new(neighbor_index_pointer, distance, PhantomData::); + + lsr.insert_neighbor(lsn); + } + } +} + +pub type BqSpeedupStorageLsnPrivateData = PhantomData; //no data stored + +impl<'a> Storage for BqSpeedupStorage<'a> { + type QueryDistanceMeasure = BqSearchDistanceMeasure; + type NodeFullDistanceMeasure<'b> = HeapFullDistanceMeasure<'b, BqSpeedupStorage<'b>> where Self: 'b; + type ArchivedType = ArchivedBqNode; + type LSNPrivateData = BqSpeedupStorageLsnPrivateData; //no data stored + + fn page_type() -> PageType { + PageType::BqNode + } + + fn create_node( + &self, + full_vector: &[f32], + heap_pointer: HeapPointer, + meta_page: &MetaPage, + tape: &mut Tape, + stats: &mut S, + ) -> ItemPointer { + let bq_vector = self.quantizer.vector_for_new_node(meta_page, full_vector); + + let node = BqNode::new(heap_pointer, &meta_page, bq_vector.as_slice()); + + let index_pointer: IndexPointer = node.write(tape, stats); + index_pointer + } + + fn start_training(&mut self, meta_page: &super::meta_page::MetaPage) { + self.quantizer.start_training(meta_page); + } + + fn add_sample(&mut self, sample: &[f32]) { + self.quantizer.add_sample(sample); + } + + fn finish_training(&mut self, stats: &mut WriteStats) { + self.quantizer.finish_training(); + self.write_quantizer_metadata(stats); + self.qv_cache = Some(QuantizedVectorCache::new(1000)); + } + + fn finalize_node_at_end_of_build( + &mut self, + meta: &MetaPage, + index_pointer: IndexPointer, + neighbors: &Vec, + stats: &mut S, + ) { + let mut cache = self.qv_cache.take().unwrap(); + /* It's important to preload cache with all the items since you can run into deadlocks + if you try to fetch a quantized vector while holding the BqNode::modify lock */ + let iter = neighbors + .iter() + .map(|n| n.get_index_pointer_to_neighbor()) + .chain(once(index_pointer)); + cache.preload(iter, self, stats); + + let node = unsafe { BqNode::modify(self.index, index_pointer, stats) }; + let mut archived = node.get_archived_node(); + archived.as_mut().set_neighbors(neighbors, &meta, &cache); + + if self.quantizer.vector_needs_update_after_training() { + let bq_vector = cache.must_get(index_pointer); + archived.as_mut().set_bq_vector(bq_vector); + } + + node.commit(); + self.qv_cache = Some(cache); + } + + unsafe fn get_full_vector_distance_state<'b, S: StatsNodeRead>( + &'b self, + index_pointer: IndexPointer, + stats: &mut S, + ) -> HeapFullDistanceMeasure<'b, BqSpeedupStorage<'b>> { + HeapFullDistanceMeasure::with_index_pointer(self, index_pointer, stats) + } + + fn get_search_distance_measure( + &self, + query: PgVector, + calc_distance_with_quantizer: bool, + ) -> BqSearchDistanceMeasure { + if !calc_distance_with_quantizer { + return BqSearchDistanceMeasure::Full(query); + } else { + return BqSearchDistanceMeasure::Bq( + self.quantizer + .get_distance_table(query.to_slice(), self.distance_fn), + ); + } + } + + fn get_neighbors_with_full_vector_distances_from_disk< + S: StatsNodeRead + StatsDistanceComparison, + >( + &self, + neighbors_of: ItemPointer, + result: &mut Vec, + stats: &mut S, + ) { + let rn = unsafe { BqNode::read(self.index, neighbors_of, stats) }; + let heap_pointer = rn + .get_archived_node() + .heap_item_pointer + .deserialize_item_pointer(); + let dist_state = + unsafe { HeapFullDistanceMeasure::with_heap_pointer(self, heap_pointer, stats) }; + for n in rn.get_archived_node().iter_neighbors() { + let dist = unsafe { dist_state.get_distance(n, stats) }; + result.push(NeighborWithDistance::new(n, dist)) + } + } + + /* get_lsn and visit_lsn are different because the distance + comparisons for BQ get the vector from different places */ + fn create_lsn_for_init_id( + &self, + lsr: &mut ListSearchResult, + index_pointer: ItemPointer, + _gns: &GraphNeighborStore, + ) -> ListSearchNeighbor { + if !lsr.prepare_insert(index_pointer) { + panic!("should not have had an init id already inserted"); + } + + let rn = unsafe { BqNode::read(self.index, index_pointer, &mut lsr.stats) }; + let node = rn.get_archived_node(); + + let distance = match lsr.sdm.as_ref().unwrap() { + BqSearchDistanceMeasure::Full(query) => { + if node.is_deleted() { + //FIXME: handle deleted init ids + panic!("don't handle deleted init ids yet"); + } + let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); + unsafe { + calculate_full_distance(self, heap_pointer, query.to_slice(), &mut lsr.stats) + } + } + BqSearchDistanceMeasure::Bq(table) => BqSearchDistanceMeasure::calculate_bq_distance( + table, + node.bq_vector.as_slice(), + &mut lsr.stats, + ), + }; + + ListSearchNeighbor::new(index_pointer, distance, PhantomData::) + } + + fn visit_lsn( + &self, + lsr: &mut ListSearchResult, + lsn_idx: usize, + gns: &GraphNeighborStore, + ) { + let lsn_index_pointer = lsr.get_lsn_by_idx(lsn_idx).index_pointer; + self.visit_lsn_internal(lsr, lsn_index_pointer, gns); + } + + fn return_lsn( + &self, + lsn: &ListSearchNeighbor, + stats: &mut GreedySearchStats, + ) -> HeapPointer { + let lsn_index_pointer = lsn.index_pointer; + let rn = unsafe { BqNode::read(self.index, lsn_index_pointer, stats) }; + let node = rn.get_archived_node(); + let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); + heap_pointer + } + + fn set_neighbors_on_disk( + &self, + meta: &MetaPage, + index_pointer: IndexPointer, + neighbors: &[NeighborWithDistance], + stats: &mut S, + ) { + let mut cache = QuantizedVectorCache::new(neighbors.len() + 1); + + /* It's important to preload cache with all the items since you can run into deadlocks + if you try to fetch a quantized vector while holding the BqNode::modify lock */ + let iter = neighbors + .iter() + .map(|n| n.get_index_pointer_to_neighbor()) + .chain(once(index_pointer)); + cache.preload(iter, self, stats); + + let node = unsafe { BqNode::modify(self.index, index_pointer, stats) }; + let mut archived = node.get_archived_node(); + archived.as_mut().set_neighbors(neighbors, &meta, &cache); + node.commit(); + } + + fn get_distance_function(&self) -> fn(&[f32], &[f32]) -> f32 { + self.distance_fn + } +} + +impl<'a> StorageFullDistanceFromHeap for BqSpeedupStorage<'a> { + unsafe fn get_heap_table_slot_from_index_pointer( + &self, + index_pointer: IndexPointer, + stats: &mut S, + ) -> TableSlot { + let rn = unsafe { BqNode::read(self.index, index_pointer, stats) }; + let node = rn.get_archived_node(); + let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); + + self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) + } + + unsafe fn get_heap_table_slot_from_heap_pointer( + &self, + heap_pointer: HeapPointer, + stats: &mut T, + ) -> TableSlot { + TableSlot::new( + self.heap_rel.unwrap(), + heap_pointer, + self.heap_attr.unwrap(), + stats, + ) + } +} + +use timescale_vector_derive::{Readable, Writeable}; + +#[derive(Archive, Deserialize, Serialize, Readable, Writeable)] +#[archive(check_bytes)] +pub struct BqNode { + pub heap_item_pointer: HeapPointer, + pub bq_vector: Vec, + neighbor_index_pointers: Vec, + neighbor_vectors: Vec>, +} + +impl BqNode { + pub fn new( + heap_pointer: HeapPointer, + meta_page: &MetaPage, + bq_vector: &[BqVectorElement], + ) -> Self { + let num_neighbors = meta_page.get_num_neighbors(); + // always use vectors of num_neighbors in length because we never want the serialized size of a Node to change + let neighbor_index_pointers: Vec<_> = (0..num_neighbors) + .map(|_| ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber)) + .collect(); + + let neighbor_vectors: Vec<_> = (0..num_neighbors) + .map(|_| vec![0; BqQuantizer::quantized_size(meta_page.get_num_dimensions() as _)]) + .collect(); + + Self { + heap_item_pointer: heap_pointer, + bq_vector: bq_vector.to_vec(), + neighbor_index_pointers: neighbor_index_pointers, + neighbor_vectors: neighbor_vectors, + } + } +} + +impl ArchivedBqNode { + pub fn neighbor_index_pointer( + self: Pin<&mut Self>, + ) -> Pin<&mut ArchivedVec> { + unsafe { self.map_unchecked_mut(|s| &mut s.neighbor_index_pointers) } + } + + pub fn neighbor_vector(self: Pin<&mut Self>) -> Pin<&mut ArchivedVec>> { + unsafe { self.map_unchecked_mut(|s| &mut s.neighbor_vectors) } + } + + pub fn bq_vector(self: Pin<&mut Self>) -> Pin<&mut Archived>> { + unsafe { self.map_unchecked_mut(|s| &mut s.bq_vector) } + } + + fn set_bq_vector(mut self: Pin<&mut Self>, bq_vector: &[BqVectorElement]) { + assert!(bq_vector.len() == self.bq_vector.len()); + for i in 0..=bq_vector.len() - 1 { + let mut pgv = self.as_mut().bq_vector().index_pin(i); + *pgv = bq_vector[i]; + } + } + + fn set_neighbors( + mut self: Pin<&mut Self>, + neighbors: &[NeighborWithDistance], + meta_page: &MetaPage, + cache: &QuantizedVectorCache, + ) { + for (i, new_neighbor) in neighbors.iter().enumerate() { + let mut a_index_pointer = self.as_mut().neighbor_index_pointer().index_pin(i); + let ip = new_neighbor.get_index_pointer_to_neighbor(); + //TODO hate that we have to set each field like this + a_index_pointer.block_number = ip.block_number; + a_index_pointer.offset = ip.offset; + + let quantized = cache.must_get(ip); + + let mut neighbor_vector = self.as_mut().neighbor_vector().index_pin(i); + for (index_in_q_vec, val) in quantized.iter().enumerate() { + let mut x = neighbor_vector.as_mut().index_pin(index_in_q_vec); + *x = *val; + } + } + //set the marker that the list ended + if neighbors.len() < meta_page.get_num_neighbors() as _ { + let mut past_last_index_pointers = + self.neighbor_index_pointer().index_pin(neighbors.len()); + past_last_index_pointers.block_number = InvalidBlockNumber; + past_last_index_pointers.offset = InvalidOffsetNumber; + } + } + + pub fn num_neighbors(&self) -> usize { + self.neighbor_index_pointers + .iter() + .position(|f| f.block_number == InvalidBlockNumber) + .unwrap_or(self.neighbor_index_pointers.len()) + } + + pub fn iter_neighbors(&self) -> impl Iterator + '_ { + self.neighbor_index_pointers + .iter() + .take(self.num_neighbors()) + .map(|ip| ip.deserialize_item_pointer()) + } +} + +impl ArchivedData for ArchivedBqNode { + fn with_data(data: &mut [u8]) -> Pin<&mut ArchivedBqNode> { + ArchivedBqNode::with_data(data) + } + + fn get_index_pointer_to_neighbors(&self) -> Vec { + self.iter_neighbors().collect() + } + + fn is_deleted(&self) -> bool { + self.heap_item_pointer.offset == InvalidOffsetNumber + } + + fn delete(self: Pin<&mut Self>) { + //TODO: actually optimize the deletes by removing index tuples. For now just mark it. + let mut heap_pointer = unsafe { self.map_unchecked_mut(|s| &mut s.heap_item_pointer) }; + heap_pointer.offset = InvalidOffsetNumber; + heap_pointer.block_number = InvalidBlockNumber; + } + + fn get_heap_item_pointer(&self) -> HeapPointer { + self.heap_item_pointer.deserialize_item_pointer() + } +} + +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + use pgrx::*; + + #[pg_test] + unsafe fn test_bq_storage_index_creation() -> spi::Result<()> { + crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( + "num_neighbors=38, USE_BQ = TRUE", + )?; + Ok(()) + } + + #[pg_test] + unsafe fn test_bq_storage_index_creation_few_neighbors() -> spi::Result<()> { + //a test with few neighbors tests the case that nodes share a page, which has caused deadlocks in the past. + crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( + "num_neighbors=10, USE_BQ = TRUE", + )?; + Ok(()) + } + + #[test] + fn test_bq_storage_delete_vacuum_plain() { + crate::access_method::vacuum::tests::test_delete_vacuum_plain_scaffold( + "num_neighbors = 10, use_bq = TRUE", + ); + } + + #[test] + fn test_bq_storage_delete_vacuum_full() { + crate::access_method::vacuum::tests::test_delete_vacuum_full_scaffold( + "num_neighbors = 38, use_bq = TRUE", + ); + } + + #[pg_test] + unsafe fn test_bq_storage_empty_table_insert() -> spi::Result<()> { + crate::access_method::build::tests::test_empty_table_insert_scaffold( + "num_neighbors=38, use_bq = TRUE", + ) + } + + #[pg_test] + unsafe fn test_bq_storage_insert_empty_insert() -> spi::Result<()> { + crate::access_method::build::tests::test_insert_empty_insert_scaffold( + "num_neighbors=38, use_bq = TRUE", + ) + } +} diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index 02277b2e..898a5e7d 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -2,60 +2,59 @@ use std::time::Instant; use pgrx::*; -use crate::access_method::disk_index_graph::DiskIndexGraph; use crate::access_method::graph::Graph; -use crate::access_method::graph::InsertStats; -use crate::access_method::graph::VectorProvider; -use crate::access_method::model::PgVector; +use crate::access_method::graph_neighbor_store::GraphNeighborStore; use crate::access_method::options::TSVIndexOptions; -use crate::util::page; +use crate::access_method::pg_vector::PgVector; +use crate::access_method::stats::{InsertStats, WriteStats}; + +use crate::util::page::PageType; use crate::util::tape::Tape; use crate::util::*; -use super::builder_graph::BuilderGraph; +use super::bq::BqSpeedupStorage; +use super::graph_neighbor_store::BuilderNeighborCache; + use super::meta_page::MetaPage; -use super::model::{self}; -use super::quantizer::Quantizer; + +use super::plain_storage::PlainStorage; +use super::pq_storage::PqCompressionStorage; +use super::storage::{Storage, StorageType}; + +enum StorageBuildState<'a, 'b, 'c, 'd, 'e> { + BqSpeedup(&'a mut BqSpeedupStorage<'b>, &'c mut BuildState<'d, 'e>), + PqCompression(&'a mut PqCompressionStorage<'b>, &'c mut BuildState<'d, 'e>), + Plain(&'a mut PlainStorage<'b>, &'c mut BuildState<'d, 'e>), +} struct BuildState<'a, 'b> { memcxt: PgMemoryContexts, meta_page: MetaPage, ntuples: usize, tape: Tape<'a>, //The tape is a memory abstraction over Postgres pages for writing data. - node_builder: BuilderGraph<'b>, + graph: Graph<'b>, started: Instant, stats: InsertStats, - quantizer: Quantizer, } impl<'a, 'b> BuildState<'a, 'b> { fn new( index_relation: &'a PgRelation, meta_page: MetaPage, - bg: BuilderGraph<'b>, - mut quantizer: Quantizer, + graph: Graph<'b>, + page_type: PageType, ) -> Self { - let tape = unsafe { Tape::new(index_relation, page::PageType::Node) }; + let tape = unsafe { Tape::new(index_relation, page_type) }; - match &mut quantizer { - Quantizer::None => {} - Quantizer::PQ(pq) => { - pq.start_training(&meta_page); - } - Quantizer::BQ(bq) => { - bq.start_training(&meta_page); - } - } //TODO: some ways to get rid of meta_page.clone? BuildState { memcxt: PgMemoryContexts::new("tsv build context"), ntuples: 0, - meta_page: meta_page.clone(), + meta_page: meta_page, tape, - node_builder: bg, + graph: graph, started: Instant::now(), stats: InsertStats::new(), - quantizer: quantizer, } } } @@ -91,6 +90,7 @@ pub extern "C" fn ambuild( } assert!(dimensions > 0 && dimensions < 2000); let meta_page = unsafe { MetaPage::create(&index_relation, dimensions as _, opt.clone()) }; + let ntuples = do_heap_scan(index_info, &heap_relation, &index_relation, meta_page); let mut result = unsafe { PgBox::::alloc0() }; @@ -119,36 +119,80 @@ pub unsafe extern "C" fn aminsert( return false; } let vec = vec.unwrap(); - let vector = (*vec).to_slice(); let heap_pointer = ItemPointer::with_item_pointer_data(*heap_tid); - let meta_page = MetaPage::read(&index_relation); - - let mut quantizer = meta_page.get_quantizer(); - match &mut quantizer { - Quantizer::None => {} - Quantizer::PQ(pq) => { - pq.load(&index_relation, &meta_page); + let mut meta_page = MetaPage::read(&index_relation); + + let mut storage = meta_page.get_storage_type(); + let mut stats = InsertStats::new(); + match &mut storage { + StorageType::Plain => { + let plain = PlainStorage::load_for_insert(&index_relation); + insert_storage( + &plain, + &index_relation, + vec, + heap_pointer, + &mut meta_page, + &mut stats, + ); } - Quantizer::BQ(bq) => { - bq.load(&index_relation, &meta_page); + StorageType::PqCompression => { + let pq = PqCompressionStorage::load_for_insert( + &heap_relation, + get_attribute_number(index_info), + &index_relation, + &meta_page, + &mut stats.quantizer_stats, + ); + insert_storage( + &pq, + &index_relation, + vec, + heap_pointer, + &mut meta_page, + &mut stats, + ); + } + StorageType::BqSpeedup => { + let bq = BqSpeedupStorage::load_for_insert( + &heap_relation, + get_attribute_number(index_info), + &index_relation, + &meta_page, + &mut stats.quantizer_stats, + ); + insert_storage( + &bq, + &index_relation, + vec, + heap_pointer, + &mut meta_page, + &mut stats, + ); } } + false +} - let vp = VectorProvider::new( - Some(&heap_relation), - Some(get_attribute_number(index_info)), - &quantizer, - false, +unsafe fn insert_storage( + storage: &S, + index_relation: &PgRelation, + vector: PgVector, + heap_pointer: ItemPointer, + meta_page: &mut MetaPage, + stats: &mut InsertStats, +) { + let mut tape = Tape::new(&index_relation, S::page_type()); + let index_pointer = storage.create_node( + vector.to_slice(), + heap_pointer, + &meta_page, + &mut tape, + stats, ); - let mut graph = DiskIndexGraph::new(&index_relation, vp); - - let node = model::Node::new(vector.to_vec(), heap_pointer, &meta_page, &quantizer); - let mut tape = unsafe { Tape::new(&index_relation, page::PageType::Node) }; - let index_pointer: IndexPointer = node.write(&mut tape); - - let _stats = graph.insert(&index_relation, index_pointer, vector); - false + let mut graph = Graph::new(GraphNeighborStore::Disk, meta_page); + graph.insert(&index_relation, index_pointer, vector, storage, stats) } #[pg_guard] @@ -167,39 +211,120 @@ fn do_heap_scan<'a>( index_relation: &'a PgRelation, meta_page: MetaPage, ) -> usize { - let quantizer = meta_page.get_quantizer(); - let vp = VectorProvider::new( - Some(heap_relation), - Some(get_attribute_number(index_info)), - &quantizer, - false, + let storage = meta_page.get_storage_type(); + + let mut mp2 = meta_page.clone(); + let graph = Graph::new( + GraphNeighborStore::Builder(BuilderNeighborCache::new()), + &mut mp2, ); + match storage { + StorageType::Plain => { + let mut plain = PlainStorage::new_for_build(index_relation); + plain.start_training(&meta_page); + let page_type = PlainStorage::page_type(); + let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); + let mut state = StorageBuildState::Plain(&mut plain, &mut bs); + + unsafe { + pg_sys::IndexBuildHeapScan( + heap_relation.as_ptr(), + index_relation.as_ptr(), + index_info, + Some(build_callback), + &mut state, + ); + } - let bg = BuilderGraph::new(meta_page.clone(), vp); - let quantizer = meta_page.get_quantizer(); - let mut state = BuildState::new(index_relation, meta_page.clone(), bg, quantizer); - unsafe { - pg_sys::IndexBuildHeapScan( - heap_relation.as_ptr(), - index_relation.as_ptr(), - index_info, - Some(build_callback), - &mut state, - ); + do_heap_scan_with_state(&mut plain, &mut bs) + } + StorageType::PqCompression => { + let mut pq = PqCompressionStorage::new_for_build( + index_relation, + heap_relation, + get_attribute_number(index_info), + ); + pq.start_training(&meta_page); + let page_type = PqCompressionStorage::page_type(); + let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); + let mut state = StorageBuildState::PqCompression(&mut pq, &mut bs); + + unsafe { + pg_sys::IndexBuildHeapScan( + heap_relation.as_ptr(), + index_relation.as_ptr(), + index_info, + Some(build_callback), + &mut state, + ); + } + + do_heap_scan_with_state(&mut pq, &mut bs) + } + StorageType::BqSpeedup => { + let mut bq = BqSpeedupStorage::new_for_build( + index_relation, + heap_relation, + get_attribute_number(index_info), + ); + bq.start_training(&meta_page); + let page_type = BqSpeedupStorage::page_type(); + let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); + let mut state = StorageBuildState::BqSpeedup(&mut bq, &mut bs); + + unsafe { + pg_sys::IndexBuildHeapScan( + heap_relation.as_ptr(), + index_relation.as_ptr(), + index_info, + Some(build_callback), + &mut state, + ); + } + + do_heap_scan_with_state(&mut bq, &mut bs) + } } +} - // we train the quantizer and add prepare to write quantized values to the nodes. - match &mut state.quantizer { - Quantizer::None => {} - Quantizer::PQ(pq) => { - pq.finish_training(); +fn do_heap_scan_with_state(storage: &mut S, state: &mut BuildState) -> usize { + // we train the quantizer and add prepare to write quantized values to the nodes.\ + let mut write_stats = WriteStats::new(); + storage.finish_training(&mut write_stats); + + match state.graph.get_neighbor_store() { + GraphNeighborStore::Builder(builder) => { + for (&index_pointer, neighbors) in builder.iter() { + write_stats.num_nodes += 1; + let prune_neighbors; + let neighbors = + if neighbors.len() > state.graph.get_meta_page().get_num_neighbors() as _ { + //OPT: get rid of this clone + prune_neighbors = state.graph.prune_neighbors( + neighbors.clone(), + storage, + &mut write_stats.prune_stats, + ); + &prune_neighbors + } else { + neighbors + }; + write_stats.num_neighbors += neighbors.len(); + + storage.finalize_node_at_end_of_build( + &state.meta_page, + index_pointer, + neighbors, + &mut write_stats, + ); + } } - Quantizer::BQ(bq) => { - bq.finish_training(); + GraphNeighborStore::Disk => { + panic!("Should not be using the disk neighbor store during build"); } } - let write_stats = unsafe { state.node_builder.write(index_relation, &state.quantizer) }; + info!("write done"); assert_eq!(write_stats.num_nodes, state.ntuples); let writing_took = Instant::now() @@ -213,28 +338,18 @@ fn do_heap_scan<'a>( write_stats.num_neighbors / write_stats.num_nodes ); } - if write_stats.num_prunes > 0 { + if write_stats.prune_stats.calls > 0 { info!( "When pruned for cleanup: avg neighbors before/after {}/{} of {} prunes", - write_stats.num_neighbors_before_prune / write_stats.num_prunes, - write_stats.num_neighbors_after_prune / write_stats.num_prunes, - write_stats.num_prunes + write_stats.prune_stats.num_neighbors_before_prune / write_stats.prune_stats.calls, + write_stats.prune_stats.num_neighbors_after_prune / write_stats.prune_stats.calls, + write_stats.prune_stats.calls ); } let ntuples = state.ntuples; warning!("Indexed {} tuples", ntuples); - match state.quantizer { - Quantizer::None => {} - Quantizer::PQ(pq) => { - pq.write_metadata(index_relation); - } - Quantizer::BQ(bq) => { - bq.write_metadata(index_relation); - } - } - ntuples } @@ -250,25 +365,47 @@ unsafe extern "C" fn build_callback( let index_relation = unsafe { PgRelation::from_pg(index) }; let vec = PgVector::from_pg_parts(values, isnull, 0); if let Some(vec) = vec { - let state = (state as *mut BuildState).as_mut().unwrap(); - - let mut old_context = state.memcxt.set_as_current(); + let state = (state as *mut StorageBuildState).as_mut().unwrap(); let heap_pointer = ItemPointer::with_item_pointer_data(*ctid); - build_callback_internal(index_relation, heap_pointer, (*vec).to_slice(), state); - - old_context.set_as_current(); - state.memcxt.reset(); + match state { + StorageBuildState::BqSpeedup(bq, state) => { + build_callback_memory_wrapper(index_relation, heap_pointer, vec, state, *bq); + } + StorageBuildState::PqCompression(pq, state) => { + build_callback_memory_wrapper(index_relation, heap_pointer, vec, state, *pq); + } + StorageBuildState::Plain(plain, state) => { + build_callback_memory_wrapper(index_relation, heap_pointer, vec, state, *plain); + } + } } //todo: what do we do with nulls? } #[inline(always)] -fn build_callback_internal( +unsafe fn build_callback_memory_wrapper( + index: PgRelation, + heap_pointer: ItemPointer, + vector: PgVector, + state: &mut BuildState, + storage: &mut S, +) { + let mut old_context = state.memcxt.set_as_current(); + + build_callback_internal(index, heap_pointer, vector, state, storage); + + old_context.set_as_current(); + state.memcxt.reset(); +} + +#[inline(always)] +fn build_callback_internal( index: PgRelation, heap_pointer: ItemPointer, - vector: &[f32], + vector: PgVector, state: &mut BuildState, + storage: &mut S, ) { check_for_interrupts!(); @@ -281,68 +418,49 @@ fn build_callback_internal( Instant::now().duration_since(state.started).as_secs_f64(), (Instant::now().duration_since(state.started) / state.ntuples as u32).as_secs_f64(), state.stats.prune_neighbor_stats.distance_comparisons / state.ntuples, - state.stats.greedy_search_stats.distance_comparisons / state.ntuples, + state.stats.greedy_search_stats.get_total_distance_comparisons() / state.ntuples, state.stats, ); } - match &mut state.quantizer { - Quantizer::None => {} - Quantizer::PQ(pq) => { - pq.add_sample(vector); - } - Quantizer::BQ(bq) => { - bq.add_sample(vector); - } - } + storage.add_sample(vector.to_slice()); - let node = model::Node::new( - vector.to_vec(), + let index_pointer = storage.create_node( + vector.to_slice(), heap_pointer, &state.meta_page, - &state.quantizer, + &mut state.tape, + &mut state.stats, ); - let index_pointer: IndexPointer = node.write(&mut state.tape); - let new_stats = state.node_builder.insert(&index, index_pointer, vector); - state.stats.combine(new_stats); + state + .graph + .insert(&index, index_pointer, vector, storage, &mut state.stats); } #[cfg(any(test, feature = "pg_test"))] #[pgrx::pg_schema] -mod tests { - use pgrx::*; +pub mod tests { + use std::collections::HashSet; - #[pg_test] - unsafe fn test_index_creation() -> spi::Result<()> { - Spi::run(&format!( - "CREATE TABLE test(embedding vector(3)); - - INSERT INTO test(embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[7,8,10]'); + use pgrx::*; - CREATE INDEX idxtest - ON test - USING tsv(embedding) - WITH (num_neighbors=30); + use crate::util::ItemPointer; - set enable_seqscan =0; - select * from test order by embedding <=> '[0,0,0]'; - explain analyze select * from test order by embedding <=> '[0,0,0]'; - drop index idxtest; - ", - ))?; - Ok(()) - } + //TODO: add test where inserting and querying with vectors that are all the same. - #[pg_test] - unsafe fn test_pq_index_creation() -> spi::Result<()> { + #[cfg(any(test, feature = "pg_test"))] + pub unsafe fn test_index_creation_and_accuracy_scaffold( + index_options: &str, + ) -> spi::Result<()> { Spi::run(&format!( - "CREATE TABLE test_pq ( + "CREATE TABLE test_data ( embedding vector (1536) ); + select setseed(0.5); -- generate 300 vectors - INSERT INTO test_pq (embedding) + INSERT INTO test_data (embedding) SELECT * FROM ( @@ -353,136 +471,182 @@ mod tests { GROUP BY i % 300) g; - CREATE INDEX idx_tsv_pq ON test_pq USING tsv (embedding) WITH (num_neighbors = 64, search_list_size = 125, max_alpha = 1.0, use_pq = TRUE, pq_vector_length = 64); + CREATE INDEX idx_tsv_bq ON test_data USING tsv (embedding) WITH ({index_options}); - ; SET enable_seqscan = 0; -- perform index scans on the vectors SELECT * FROM - test_pq - ORDER BY - embedding <=> ( - SELECT - ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding - FROM generate_series(1, 1536)); - - EXPLAIN ANALYZE - SELECT - * - FROM - test_pq + test_data ORDER BY embedding <=> ( SELECT ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding FROM generate_series(1, 1536)); - DROP INDEX idx_tsv_pq; - ", - ))?; - Ok(()) - } - - #[pg_test] - unsafe fn test_bq_index_creation() -> spi::Result<()> { - Spi::run(&format!( - "CREATE TABLE test_bq ( - embedding vector (1536) - ); - - -- generate 300 vectors - INSERT INTO test_bq (embedding) + -- test insert 2 vectors + INSERT INTO test_data (embedding) SELECT * FROM ( SELECT ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding FROM - generate_series(1, 1536 * 300) i + generate_series(1, 1536 * 2) i GROUP BY - i % 300) g; - - CREATE INDEX idx_tsv_pq ON test_bq USING tsv (embedding) WITH (num_neighbors = 64, search_list_size = 125, max_alpha = 1.0, use_bq = TRUE); + i % 2) g; - ; - SET enable_seqscan = 0; - -- perform index scans on the vectors + EXPLAIN ANALYZE SELECT * FROM - test_bq + test_data ORDER BY embedding <=> ( SELECT ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding FROM generate_series(1, 1536)); - EXPLAIN ANALYZE + -- test insert 10 vectors to search for that aren't random + INSERT INTO test_data (embedding) SELECT * - FROM - test_bq - ORDER BY - embedding <=> ( - SELECT - ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding - FROM generate_series(1, 1536)); - - DROP INDEX idx_tsv_pq; - ", - ))?; - Ok(()) - } - - #[pg_test] - unsafe fn test_insert() -> spi::Result<()> { - Spi::run(&format!( - "CREATE TABLE test(embedding vector(3)); - - INSERT INTO test(embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[7,8,10]'); - - CREATE INDEX idxtest - ON test - USING tsv(embedding) - WITH (num_neighbors=30); + FROM ( + SELECT + ('[' || array_to_string(array_agg(1.0), ',', '0') || ']')::vector AS embedding + FROM + generate_series(1, 1536 * 10) i + GROUP BY + i % 10) g; - INSERT INTO test(embedding) VALUES ('[11,12,13]'); ", ))?; - let res: Option = Spi::get_one(&format!( - " set enable_seqscan = 0; - WITH cte as (select * from test order by embedding <=> '[0,0,0]') SELECT count(*) from cte;", + let test_vec: Option> = Spi::get_one(&format!( + "SELECT('{{' || array_to_string(array_agg(1.0), ',', '0') || '}}')::real[] AS embedding +FROM generate_series(1, 1536)" ))?; - assert_eq!(4, res.unwrap()); - Spi::run(&format!( - "INSERT INTO test(embedding) VALUES ('[11,12,13]'), ('[14,15,16]');", - ))?; - let res: Option = Spi::get_one(&format!( - " set enable_seqscan = 0; - WITH cte as (select * from test order by embedding <=> '[0,0,0]') SELECT count(*) from cte;", - ))?; - assert_eq!(6, res.unwrap()); - - Spi::run(&format!("drop index idxtest;",))?; + let with_index: Option> = Spi::get_one_with_args( + &format!( + " + SET enable_seqscan = 0; + SET enable_indexscan = 1; + SET tsv.query_search_list_size = 25; + WITH cte AS ( + SELECT + ctid + FROM + test_data + ORDER BY + embedding <=> $1::vector + LIMIT 10 + ) + SELECT array_agg(ctid) from cte;" + ), + vec![( + pgrx::PgOid::Custom(pgrx::pg_sys::FLOAT4ARRAYOID), + test_vec.clone().into_datum(), + )], + )?; + + /* Test that the explain plan is generated ok */ + let explain: Option = Spi::get_one_with_args( + &format!( + " + SET enable_seqscan = 0; + SET enable_indexscan = 1; + EXPLAIN (format json) WITH cte AS ( + SELECT + ctid + FROM + test_data + ORDER BY + embedding <=> $1::vector + LIMIT 10 + ) + SELECT array_agg(ctid) from cte;" + ), + vec![( + pgrx::PgOid::Custom(pgrx::pg_sys::FLOAT4ARRAYOID), + test_vec.clone().into_datum(), + )], + )?; + assert!(explain.is_some()); + //warning!("explain: {}", explain.unwrap().0); + + let without_index: Option> = Spi::get_one_with_args( + &format!( + " + SET enable_seqscan = 1; + SET enable_indexscan = 0; + WITH cte AS ( + SELECT + ctid + FROM + test_data + ORDER BY + embedding <=> $1::vector + LIMIT 10 + ) + SELECT array_agg(ctid) from cte;" + ), + vec![( + pgrx::PgOid::Custom(pgrx::pg_sys::FLOAT4ARRAYOID), + test_vec.clone().into_datum(), + )], + )?; + + let set: HashSet<_> = without_index + .unwrap() + .iter() + .map(|&ctid| ItemPointer::with_item_pointer_data(ctid)) + .collect(); + + let mut matches = 0; + for ctid in with_index.unwrap() { + if set.contains(&ItemPointer::with_item_pointer_data(ctid)) { + matches += 1; + } + } + assert!(matches > 9, "Low number of matches: {}", matches); + + //FIXME: should work in all cases + if !index_options.contains("num_neighbors=10") { + //make sure you can scan entire table with index + let cnt: Option = Spi::get_one_with_args( + &format!( + " + SET enable_seqscan = 0; + SET enable_indexscan = 1; + SET tsv.query_search_list_size = 2; + WITH cte as (select * from test_data order by embedding <=> $1::vector) SELECT count(*) from cte; + ", + ), + vec![( + pgrx::PgOid::Custom(pgrx::pg_sys::FLOAT4ARRAYOID), + test_vec.into_datum(), + )], + )?; + + assert_eq!(cnt.unwrap(), 312); + } Ok(()) } - #[pg_test] - unsafe fn test_empty_table_insert() -> spi::Result<()> { + #[cfg(any(test, feature = "pg_test"))] + pub unsafe fn test_empty_table_insert_scaffold(index_options: &str) -> spi::Result<()> { Spi::run(&format!( "CREATE TABLE test(embedding vector(3)); CREATE INDEX idxtest ON test USING tsv(embedding) - WITH (num_neighbors=30); + WITH ({index_options}); INSERT INTO test(embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[7,8,10]'); ", @@ -494,20 +658,27 @@ mod tests { ))?; assert_eq!(3, res.unwrap()); + Spi::run(&format!( + " + set enable_seqscan = 0; + explain analyze select * from test order by embedding <=> '[0,0,0]'; + ", + ))?; + Spi::run(&format!("drop index idxtest;",))?; Ok(()) } - #[pg_test] - unsafe fn test_insert_empty_insert() -> spi::Result<()> { + #[cfg(any(test, feature = "pg_test"))] + pub unsafe fn test_insert_empty_insert_scaffold(index_options: &str) -> spi::Result<()> { Spi::run(&format!( "CREATE TABLE test(embedding vector(3)); CREATE INDEX idxtest ON test USING tsv(embedding) - WITH (num_neighbors=30); + WITH ({index_options}); INSERT INTO test(embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[7,8,10]'); DELETE FROM test; diff --git a/timescale_vector/src/access_method/builder_graph.rs b/timescale_vector/src/access_method/builder_graph.rs deleted file mode 100644 index be95950f..00000000 --- a/timescale_vector/src/access_method/builder_graph.rs +++ /dev/null @@ -1,174 +0,0 @@ -use std::collections::HashMap; -use std::time::Instant; - -use pgrx::*; - -use crate::util::{IndexPointer, ItemPointer}; - -use super::graph::{Graph, VectorProvider}; -use super::meta_page::MetaPage; -use super::model::*; -use super::quantizer::Quantizer; - -/// A builderGraph is a graph that keep the neighbors in-memory in the neighbor_map below -/// The idea is that during the index build, you don't want to update the actual Postgres -/// pages every time you change the neighbors. Instead you change the neighbors in memory -/// until the build is done. Afterwards, calling the `write` method, will write out all -/// the neighbors to the right pages. -pub struct BuilderGraph<'a> { - //maps node's pointer to the representation on disk - neighbor_map: HashMap>, - meta_page: MetaPage, - vector_provider: VectorProvider<'a>, -} - -impl<'a> BuilderGraph<'a> { - pub fn new(meta_page: MetaPage, vp: VectorProvider<'a>) -> Self { - Self { - neighbor_map: HashMap::with_capacity(200), - meta_page, - vector_provider: vp, - } - } - - unsafe fn get_full_vector(&self, heap_pointer: ItemPointer) -> Vec { - let vp = self.get_vector_provider(); - vp.get_full_vector_copy_from_heap_pointer(heap_pointer) - } - - pub unsafe fn write(&self, index: &PgRelation, quantizer: &Quantizer) -> WriteStats { - let mut stats = WriteStats::new(); - let meta = self.get_meta_page(index); - - //TODO: OPT: do this in order of item pointers - for (index_pointer, neighbors) in &self.neighbor_map { - stats.num_nodes += 1; - let prune_neighbors; - let neighbors = if neighbors.len() > self.meta_page.get_num_neighbors() as _ { - stats.num_prunes += 1; - stats.num_neighbors_before_prune += neighbors.len(); - (prune_neighbors, _) = self.prune_neighbors(index, *index_pointer, vec![]); - stats.num_neighbors_after_prune += prune_neighbors.len(); - &prune_neighbors - } else { - neighbors - }; - stats.num_neighbors += neighbors.len(); - - let node = Node::modify(index, *index_pointer); - let mut archived = node.get_archived_node(); - archived.as_mut().set_neighbors(neighbors, &meta); - - match quantizer { - Quantizer::None => {} - Quantizer::PQ(pq) => { - let heap_pointer = node - .get_archived_node() - .heap_item_pointer - .deserialize_item_pointer(); - let full_vector = self.get_full_vector(heap_pointer); - pq.update_node_after_traing(&mut archived, full_vector); - } - Quantizer::BQ(bq) => { - //TODO: OPT: this may not be needed - let heap_pointer = node - .get_archived_node() - .heap_item_pointer - .deserialize_item_pointer(); - let full_vector = self.get_full_vector(heap_pointer); - bq.update_node_after_traing(&mut archived, full_vector); - } - }; - node.commit(); - } - stats - } -} - -impl<'a> Graph for BuilderGraph<'a> { - fn read<'b>(&self, index: &'b PgRelation, index_pointer: ItemPointer) -> ReadableNode<'b> { - unsafe { Node::read(index, index_pointer) } - } - - fn get_init_ids(&self) -> Option> { - //returns a vector for generality - self.meta_page.get_init_ids() - } - - fn get_neighbors(&self, _node: &ArchivedNode, neighbors_of: ItemPointer) -> Vec { - let neighbors = self.neighbor_map.get(&neighbors_of); - match neighbors { - Some(n) => n - .iter() - .map(|n| n.get_index_pointer_to_neighbor()) - .collect(), - None => vec![], - } - } - - fn get_neighbors_with_distances( - &self, - _index: &PgRelation, - neighbors_of: ItemPointer, - result: &mut Vec, - ) -> bool { - let neighbors = self.neighbor_map.get(&neighbors_of); - match neighbors { - Some(n) => { - for nwd in n { - result.push(nwd.clone()); - } - true - } - None => false, - } - } - - fn is_empty(&self) -> bool { - self.neighbor_map.len() == 0 - } - - fn get_vector_provider(&self) -> VectorProvider { - return self.vector_provider.clone(); - } - - fn get_meta_page(&self, _index: &PgRelation) -> &MetaPage { - &self.meta_page - } - - fn set_neighbors( - &mut self, - index: &PgRelation, - neighbors_of: ItemPointer, - new_neighbors: Vec, - ) { - if self.meta_page.get_init_ids().is_none() { - //TODO probably better set off of centeroids - MetaPage::update_init_ids(index, vec![neighbors_of]); - self.meta_page = MetaPage::read(index); - } - self.neighbor_map.insert(neighbors_of, new_neighbors); - } -} - -pub struct WriteStats { - pub started: Instant, - pub num_nodes: usize, - pub num_prunes: usize, - pub num_neighbors_before_prune: usize, - pub num_neighbors_after_prune: usize, - pub num_neighbors: usize, -} - -impl WriteStats { - pub fn new() -> Self { - Self { - started: Instant::now(), - num_nodes: 0, - num_prunes: 0, - num_neighbors_before_prune: 0, - num_neighbors_after_prune: 0, - num_neighbors: 0, - } - } -} diff --git a/timescale_vector/src/access_method/debugging.rs b/timescale_vector/src/access_method/debugging.rs index 5be1674e..40a86392 100644 --- a/timescale_vector/src/access_method/debugging.rs +++ b/timescale_vector/src/access_method/debugging.rs @@ -3,11 +3,10 @@ use std::collections::HashMap; use pgrx::PgRelation; -use rkyv::Deserialize; use crate::util::ItemPointer; -use super::model::Node; +use super::{plain_node::Node, stats::GreedySearchStats}; #[allow(dead_code)] pub fn print_graph_from_disk(index: &PgRelation, init_id: ItemPointer) { @@ -26,7 +25,8 @@ unsafe fn print_graph_from_disk_visitor( sb: &mut String, level: usize, ) { - let data_node = Node::read(&index, index_pointer); + let mut stats = GreedySearchStats::new(); + let data_node = Node::read(&index, index_pointer, &mut stats); let node = data_node.get_archived_node(); let v = node.vector.as_slice(); let copy: Vec = v.iter().map(|f| *f).collect(); @@ -34,20 +34,18 @@ unsafe fn print_graph_from_disk_visitor( map.insert(index_pointer, copy); - node.apply_to_neighbors(|neighbor_pointer| { - let p = neighbor_pointer.deserialize_item_pointer(); + for neighbor_pointer in node.iter_neighbors() { + let p = neighbor_pointer; if !map.contains_key(&p) { print_graph_from_disk_visitor(index, p, map, sb, level + 1); } - }); + } sb.push_str(&name); sb.push_str("\n"); - node.apply_to_neighbors(|neighbor_pointer| { - let ip: ItemPointer = (neighbor_pointer) - .deserialize(&mut rkyv::Infallible) - .unwrap(); - let neighbor = map.get(&ip).unwrap(); + + for neighbor_pointer in node.iter_neighbors() { + let neighbor = map.get(&neighbor_pointer).unwrap(); sb.push_str(&format!("->{:?}\n", neighbor)) - }); + } sb.push_str("\n") } diff --git a/timescale_vector/src/access_method/disk_index_graph.rs b/timescale_vector/src/access_method/disk_index_graph.rs deleted file mode 100644 index 7ae548f4..00000000 --- a/timescale_vector/src/access_method/disk_index_graph.rs +++ /dev/null @@ -1,92 +0,0 @@ -use pgrx::PgRelation; - -use crate::util::ItemPointer; - -use super::{ - graph::{Graph, VectorProvider}, - meta_page::MetaPage, - model::{ArchivedNode, NeighborWithDistance, Node, ReadableNode}, -}; - -pub struct DiskIndexGraph<'a> { - meta_page: MetaPage, - vector_provider: VectorProvider<'a>, -} - -impl<'a> DiskIndexGraph<'a> { - pub fn new(index: &PgRelation, vp: VectorProvider<'a>) -> Self { - let meta = MetaPage::read(index); - Self { - meta_page: meta, - vector_provider: vp, - } - } -} - -impl<'h> Graph for DiskIndexGraph<'h> { - fn get_vector_provider(&self) -> VectorProvider { - return self.vector_provider.clone(); - } - - fn read<'a>(&self, index: &'a PgRelation, index_pointer: ItemPointer) -> ReadableNode<'a> { - unsafe { Node::read(index, index_pointer) } - } - - fn get_init_ids(&self) -> Option> { - self.meta_page.get_init_ids() - } - - fn get_neighbors(&self, node: &ArchivedNode, _neighbors_of: ItemPointer) -> Vec { - let mut result = Vec::with_capacity(node.num_neighbors()); - node.apply_to_neighbors(|n| { - let n = n.deserialize_item_pointer(); - result.push(n) - }); - result - } - - fn get_neighbors_with_distances( - &self, - index: &PgRelation, - neighbors_of: ItemPointer, - result: &mut Vec, - ) -> bool { - let rn = self.read(index, neighbors_of); - let vp = self.get_vector_provider(); - let dist_state = unsafe { vp.get_full_vector_distance_state(index, neighbors_of) }; - rn.get_archived_node().apply_to_neighbors(|n| { - let n = n.deserialize_item_pointer(); - let dist = - unsafe { vp.get_distance_pair_for_full_vectors_from_state(&dist_state, index, n) }; - result.push(NeighborWithDistance::new(n, dist)) - }); - true - } - - fn is_empty(&self) -> bool { - self.meta_page.get_init_ids().is_none() - } - - fn get_meta_page(&self, _index: &PgRelation) -> &MetaPage { - &self.meta_page - } - - fn set_neighbors( - &mut self, - index: &PgRelation, - neighbors_of: ItemPointer, - new_neighbors: Vec, - ) { - if self.meta_page.get_init_ids().is_none() { - MetaPage::update_init_ids(index, vec![neighbors_of]); - self.meta_page = MetaPage::read(index); - } - - unsafe { - let node = Node::modify(index, neighbors_of); - let archived = node.get_archived_node(); - archived.set_neighbors(&new_neighbors, &self.meta_page); - node.commit(); - } - } -} diff --git a/timescale_vector/src/access_method/distance.rs b/timescale_vector/src/access_method/distance.rs index e66b7c53..21396761 100644 --- a/timescale_vector/src/access_method/distance.rs +++ b/timescale_vector/src/access_method/distance.rs @@ -103,17 +103,38 @@ pub fn distance_cosine(a: &[f32], b: &[f32]) -> f32 { #[inline(always)] pub fn distance_cosine_unoptimized(a: &[f32], b: &[f32]) -> f32 { assert_eq!(a.len(), b.len()); + debug_assert!(preprocess_cosine_get_norm(a).is_none()); + debug_assert!(preprocess_cosine_get_norm(b).is_none()); let res: f32 = a.iter().zip(b).map(|(a, b)| *a * *b).sum(); - 1.0 - res + (1.0 - res).max(0.0) } -pub fn preprocess_cosine(a: &mut [f32]) { +pub fn preprocess_cosine_get_norm(a: &[f32]) -> Option { let norm = a.iter().map(|v| v * v).sum::(); + //adjust the epsilon to the length of the vector + let adj_epsilon = f32::EPSILON * a.len() as f32; + + /* this mainly handles the zero-vector case */ if norm < f32::EPSILON { - return; + return None; + } + /* no need to renormalize if norm around 1.0 */ + if norm >= 1.0 - adj_epsilon && norm <= 1.0 + adj_epsilon { + return None; } - let norm = norm.sqrt(); - if norm > 1.0 + f32::EPSILON || norm < 1.0 - f32::EPSILON { - a.iter_mut().for_each(|v| *v /= norm); + return Some(norm.sqrt()); +} + +pub fn preprocess_cosine(a: &mut [f32]) { + let norm = preprocess_cosine_get_norm(a); + match norm { + None => (), + Some(norm) => { + a.iter_mut().for_each(|v| *v /= norm); + debug_assert!( + preprocess_cosine_get_norm(a).is_none(), + "preprocess_cosine isn't idempotent", + ); + } } } diff --git a/timescale_vector/src/access_method/distance_x86.rs b/timescale_vector/src/access_method/distance_x86.rs index 9ea6e71e..36f78537 100644 --- a/timescale_vector/src/access_method/distance_x86.rs +++ b/timescale_vector/src/access_method/distance_x86.rs @@ -124,7 +124,7 @@ simdeez::simd_runtime_generate!( dist += x[i] * y[i]; } - 1.0 - dist + (1.0 - dist).max(0.0) } ); diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index 53662c58..5c318109 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -1,330 +1,84 @@ use std::{cmp::Ordering, collections::HashSet}; -use pgrx::pg_sys::{Datum, TupleTableSlot}; -use pgrx::{pg_sys, PgBox, PgRelation}; +use pgrx::PgRelation; -use crate::access_method::model::Node; -use crate::util::ports::slot_getattr; -use crate::util::{HeapPointer, IndexPointer, ItemPointer}; - -use super::distance::distance_cosine as default_distance; -use super::model::{ArchivedNode, PgVector}; -use super::quantizer::Quantizer; -use super::{ - meta_page::MetaPage, - model::{NeighborWithDistance, ReadableNode}, -}; - -struct TableSlot { - slot: PgBox, -} - -impl TableSlot { - unsafe fn new(relation: &PgRelation) -> Self { - let slot = PgBox::from_pg(pg_sys::table_slot_create( - relation.as_ptr(), - std::ptr::null_mut(), - )); - Self { slot } - } - - unsafe fn get_attribute(&self, attribute_number: pg_sys::AttrNumber) -> Option { - slot_getattr(&self.slot, attribute_number) - } -} - -impl Drop for TableSlot { - fn drop(&mut self) { - unsafe { pg_sys::ExecDropSingleTupleTableSlot(self.slot.as_ptr()) }; - } -} - -#[derive(Clone)] -pub struct VectorProvider<'a> { - quantizer: &'a Quantizer, - calc_distance_with_quantizer: bool, - heap_rel: Option<&'a PgRelation>, - heap_attr_number: Option, - distance_fn: fn(&[f32], &[f32]) -> f32, -} - -impl<'a> VectorProvider<'a> { - pub fn new( - heap_rel: Option<&'a PgRelation>, - heap_attr_number: Option, - quantizer: &'a Quantizer, - calc_distance_with_quantizer: bool, - ) -> Self { - Self { - quantizer, - calc_distance_with_quantizer, - heap_rel, - heap_attr_number, - distance_fn: default_distance, - } - } - - pub unsafe fn get_full_vector_copy_from_heap_pointer( - &self, - heap_pointer: ItemPointer, - ) -> Vec { - let slot = TableSlot::new(self.heap_rel.unwrap()); - self.init_slot(&slot, heap_pointer); - let slice = self.get_slice(&slot); - slice.to_vec() - } - - unsafe fn init_slot(&self, slot: &TableSlot, heap_pointer: HeapPointer) { - let table_am = self.heap_rel.unwrap().rd_tableam; - let fetch_row_version = (*table_am).tuple_fetch_row_version.unwrap(); - let mut ctid: pg_sys::ItemPointerData = pg_sys::ItemPointerData { - ..Default::default() - }; - heap_pointer.to_item_pointer_data(&mut ctid); - fetch_row_version( - self.heap_rel.unwrap().as_ptr(), - &mut ctid, - &mut pg_sys::SnapshotAnyData, - slot.slot.as_ptr(), - ); - } - - unsafe fn get_slice<'s>(&self, slot: &'s TableSlot) -> &'s [f32] { - let vector = - PgVector::from_datum(slot.get_attribute(self.heap_attr_number.unwrap()).unwrap()); - - //note pgvector slice is only valid as long as the slot is valid that's why the lifetime is tied to it. - (*vector).to_slice() - } - - fn get_heap_pointer(&self, index: &PgRelation, index_pointer: IndexPointer) -> HeapPointer { - let rn = unsafe { Node::read(index, index_pointer) }; - let node = rn.get_archived_node(); - let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); - heap_pointer - } - - fn get_distance_measure(&self, query: &[f32]) -> DistanceMeasure { - return DistanceMeasure::new(self.quantizer, query, self.calc_distance_with_quantizer); - } - - unsafe fn get_distance( - &self, - node: &ArchivedNode, - query: &[f32], - dm: &DistanceMeasure, - stats: &mut GreedySearchStats, - ) -> f32 { - if self.calc_distance_with_quantizer { - assert!(node.pq_vector.len() > 0); - let vec = node.pq_vector.as_slice(); - let distance = dm.get_quantized_distance(vec); - stats.pq_distance_comparisons += 1; - stats.distance_comparisons += 1; - return distance; - } - - //now we know we're doing a distance calc on the full-sized vector - if self.quantizer.is_some() { - //have to get it from the heap - let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); - let slot = TableSlot::new(self.heap_rel.unwrap()); - self.init_slot(&slot, heap_pointer); - let slice = self.get_slice(&slot); - let distance = dm.get_full_vector_distance(slice, query); - stats.distance_comparisons += 1; - return distance; - } else { - //have to get it from the index - assert!(node.vector.len() > 0); - let vec = node.vector.as_slice(); - let distance = dm.get_full_vector_distance(vec, query); - stats.distance_comparisons += 1; - return distance; - } - } - - pub unsafe fn get_full_vector_distance_state<'i>( - &self, - index: &'i PgRelation, - index_pointer: IndexPointer, - ) -> FullVectorDistanceState<'i> { - if self.quantizer.is_some() { - let heap_pointer = self.get_heap_pointer(index, index_pointer); - let slot = TableSlot::new(self.heap_rel.unwrap()); - self.init_slot(&slot, heap_pointer); - FullVectorDistanceState { - table_slot: Some(slot), - readable_node: None, - } - } else { - let rn = Node::read(index, index_pointer); - FullVectorDistanceState { - table_slot: None, - readable_node: Some(rn), - } - } - } +use crate::access_method::storage::NodeFullDistanceMeasure; - pub unsafe fn get_distance_pair_for_full_vectors_from_state( - &self, - state: &FullVectorDistanceState, - index: &PgRelation, - index_pointer: IndexPointer, - ) -> f32 { - if self.quantizer.is_some() { - let heap_pointer = self.get_heap_pointer(index, index_pointer); - let slot = TableSlot::new(self.heap_rel.unwrap()); - self.init_slot(&slot, heap_pointer); - let slice1 = self.get_slice(&slot); - let slice2 = self.get_slice(state.table_slot.as_ref().unwrap()); - (self.distance_fn)(slice1, slice2) - } else { - let rn1 = Node::read(index, index_pointer); - let rn2 = state.readable_node.as_ref().unwrap(); - let node1 = rn1.get_archived_node(); - let node2 = rn2.get_archived_node(); - assert!(node1.vector.len() > 0); - assert!(node1.vector.len() == node2.vector.len()); - let vec1 = node1.vector.as_slice(); - let vec2 = node2.vector.as_slice(); - (self.distance_fn)(vec1, vec2) - } - } -} - -pub struct FullVectorDistanceState<'a> { - table_slot: Option, - readable_node: Option>, -} - -pub struct DistanceMeasure { - pq_distance_table: Option, - bq_distance_table: Option, -} - -impl DistanceMeasure { - pub fn new(quantizer: &Quantizer, query: &[f32], calc_distance_with_quantizer: bool) -> Self { - if !calc_distance_with_quantizer { - return Self { - pq_distance_table: None, - bq_distance_table: None, - }; - } - match quantizer { - Quantizer::None => Self { - pq_distance_table: None, - bq_distance_table: None, - }, - Quantizer::PQ(pq) => { - let dc = pq.get_distance_table(query, default_distance); - Self { - pq_distance_table: Some(dc), - bq_distance_table: None, - } - } - Quantizer::BQ(bq) => { - let dc = bq.get_distance_table(query, default_distance); - Self { - pq_distance_table: None, - bq_distance_table: Some(dc), - } - } - } - } +use crate::util::{HeapPointer, IndexPointer, ItemPointer}; - fn get_quantized_distance(&self, vec: &[u8]) -> f32 { - if self.pq_distance_table.is_some() { - let dc = self.pq_distance_table.as_ref().unwrap(); - let distance = dc.distance(vec); - distance - } else { - let dc = self.bq_distance_table.as_ref().unwrap(); - let distance = dc.distance(vec); - distance - } - } +use super::graph_neighbor_store::GraphNeighborStore; - fn get_full_vector_distance(&self, vec: &[f32], query: &[f32]) -> f32 { - assert!(self.pq_distance_table.is_none()); - default_distance(vec, query) - } -} +use super::pg_vector::PgVector; +use super::stats::{GreedySearchStats, InsertStats, PruneNeighborStats}; +use super::storage::Storage; +use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance}; -struct ListSearchNeighbor { - index_pointer: IndexPointer, - heap_pointer: HeapPointer, - neighbor_index_pointers: Vec, +pub struct ListSearchNeighbor { + pub index_pointer: IndexPointer, distance: f32, visited: bool, + private_data: PD, } -impl PartialOrd for ListSearchNeighbor { +impl PartialOrd for ListSearchNeighbor { fn partial_cmp(&self, other: &Self) -> Option { self.distance.partial_cmp(&other.distance) } } -impl PartialEq for ListSearchNeighbor { +impl PartialEq for ListSearchNeighbor { fn eq(&self, other: &Self) -> bool { self.index_pointer == other.index_pointer } } -impl ListSearchNeighbor { - pub fn new( - index_pointer: IndexPointer, - heap_pointer: HeapPointer, - distance: f32, - neighbor_index_pointers: Vec, - ) -> Self { +impl ListSearchNeighbor { + pub fn new(index_pointer: IndexPointer, distance: f32, private_data: PD) -> Self { + assert!(!distance.is_nan()); + debug_assert!(distance >= 0.0); Self { index_pointer, - heap_pointer, - neighbor_index_pointers, + private_data, distance, visited: false, } } + + pub fn get_private_data(&self) -> &PD { + &self.private_data + } } -pub struct ListSearchResult { - candidate_storage: Vec, //plain storage - best_candidate: Vec, //pos in candidate storage, sorted by distance +pub struct ListSearchResult { + candidate_storage: Vec>, //plain storage + best_candidate: Vec, //pos in candidate storage, sorted by distance inserted: HashSet, max_history_size: Option, - dm: DistanceMeasure, + pub sdm: Option, pub stats: GreedySearchStats, } -impl ListSearchResult { +impl ListSearchResult { fn empty() -> Self { Self { candidate_storage: vec![], best_candidate: vec![], inserted: HashSet::new(), max_history_size: None, - dm: DistanceMeasure { - pq_distance_table: None, - bq_distance_table: None, - }, + sdm: None, stats: GreedySearchStats::new(), } } - fn new( - index: &PgRelation, + fn new>( max_history_size: Option, - graph: &G, init_ids: Vec, - query: &[f32], - dm: DistanceMeasure, + sdm: S::QueryDistanceMeasure, search_list_size: usize, meta_page: &MetaPage, - ) -> Self - where - G: Graph + ?Sized, - { + gns: &GraphNeighborStore, + storage: &S, + ) -> Self { let neigbors = meta_page.get_num_neighbors() as usize; let mut res = Self { candidate_storage: Vec::with_capacity(search_list_size * neigbors), @@ -332,48 +86,22 @@ impl ListSearchResult { inserted: HashSet::with_capacity(search_list_size * neigbors), max_history_size, stats: GreedySearchStats::new(), - dm: dm, + sdm: Some(sdm), }; - res.stats.calls += 1; + res.stats.record_call(); for index_pointer in init_ids { - res.insert(index, graph, index_pointer, query); + let lsn = storage.create_lsn_for_init_id(&mut res, index_pointer, gns); + res.insert_neighbor(lsn); } res } - fn insert( - &mut self, - index: &PgRelation, - graph: &G, - index_pointer: ItemPointer, - query: &[f32], - ) where - G: Graph + ?Sized, - { - //no point reprocessing a point. Distance calcs are expensive. - if !self.inserted.insert(index_pointer) { - return; - } - - let rn = unsafe { Node::read(index, index_pointer) }; - self.stats.node_reads += 1; - let node = rn.get_archived_node(); - - let vp = graph.get_vector_provider(); - let distance = unsafe { vp.get_distance(node, query, &self.dm, &mut self.stats) }; - - let neighbors = graph.get_neighbors(node, index_pointer); - let lsn = ListSearchNeighbor::new( - index_pointer, - node.heap_item_pointer.deserialize_item_pointer(), - distance, - neighbors, - ); - self._insert_neighbor(lsn); + pub fn prepare_insert(&mut self, ip: ItemPointer) -> bool { + return self.inserted.insert(ip); } /// Internal function - fn _insert_neighbor(&mut self, n: ListSearchNeighbor) { + pub fn insert_neighbor(&mut self, n: ListSearchNeighbor) { if let Some(max_size) = self.max_history_size { if self.best_candidate.len() >= max_size { let last = self.best_candidate.last().unwrap(); @@ -393,7 +121,11 @@ impl ListSearchResult { self.best_candidate.insert(idx, pos) } - fn visit_closest(&mut self, pos_limit: usize) -> Option<&ListSearchNeighbor> { + pub fn get_lsn_by_idx(&self, idx: usize) -> &ListSearchNeighbor { + &self.candidate_storage[idx] + } + + fn visit_closest(&mut self, pos_limit: usize) -> Option { //OPT: should we optimize this not to do a linear search each time? let neighbor_position = self .best_candidate @@ -406,7 +138,7 @@ impl ListSearchResult { } let n = &mut self.candidate_storage[self.best_candidate[pos]]; n.visited = true; - Some(n) + Some(self.best_candidate[pos]) } None => None, } @@ -414,38 +146,101 @@ impl ListSearchResult { //removes and returns the first element. Given that the element remains in self.inserted, that means the element will never again be insereted //into the best_candidate list, so it will never again be returned. - pub fn consume(&mut self) -> Option<(HeapPointer, IndexPointer)> { + pub fn consume>( + &mut self, + storage: &S, + ) -> Option<(HeapPointer, IndexPointer)> { if self.best_candidate.is_empty() { return None; } - let f = &self.candidate_storage[self.best_candidate.remove(0)]; - return Some((f.heap_pointer, f.index_pointer)); + let idx = self.best_candidate.remove(0); + let lsn = &self.candidate_storage[idx]; + let heap_pointer = storage.return_lsn(lsn, &mut self.stats); + return Some((heap_pointer, lsn.index_pointer)); } } -pub trait Graph { - fn read<'a>(&self, index: &'a PgRelation, index_pointer: ItemPointer) -> ReadableNode<'a>; - fn get_init_ids(&self) -> Option>; - fn get_neighbors(&self, node: &ArchivedNode, neighbors_of: ItemPointer) -> Vec; - fn get_neighbors_with_distances( - &self, - index: &PgRelation, - neighbors_of: ItemPointer, - result: &mut Vec, - ) -> bool; +pub struct Graph<'a> { + neighbor_store: GraphNeighborStore, + meta_page: &'a mut MetaPage, +} + +impl<'a> Graph<'a> { + pub fn new(neighbor_store: GraphNeighborStore, meta_page: &'a mut MetaPage) -> Self { + Self { + neighbor_store, + meta_page, + } + } - fn is_empty(&self) -> bool; + pub fn get_neighbor_store(&self) -> &GraphNeighborStore { + &self.neighbor_store + } - fn get_vector_provider(&self) -> VectorProvider; + fn get_init_ids(&self) -> Option> { + self.meta_page.get_init_ids() + } - fn set_neighbors( + fn add_neighbors( &mut self, - index: &PgRelation, + storage: &S, neighbors_of: ItemPointer, - new_neighbors: Vec, - ); + additional_neighbors: Vec, + stats: &mut PruneNeighborStats, + ) -> (bool, Vec) { + let mut candidates = Vec::::with_capacity( + (self.neighbor_store.max_neighbors(self.get_meta_page()) as usize) + + additional_neighbors.len(), + ); + self.neighbor_store + .get_neighbors_with_full_vector_distances( + neighbors_of, + storage, + &mut candidates, + stats, + ); + + let mut hash: HashSet = candidates + .iter() + .map(|c| c.get_index_pointer_to_neighbor()) + .collect(); + for n in additional_neighbors { + if hash.insert(n.get_index_pointer_to_neighbor()) { + candidates.push(n); + } + } + //remove myself + if !hash.insert(neighbors_of) { + //prevent self-loops + let index = candidates + .iter() + .position(|x| x.get_index_pointer_to_neighbor() == neighbors_of) + .unwrap(); + candidates.remove(index); + } + + let (pruned, new_neighbors) = + if candidates.len() > self.neighbor_store.max_neighbors(self.get_meta_page()) { + let new_list = self.prune_neighbors(candidates, storage, stats); + (true, new_list) + } else { + (false, candidates) + }; + + //OPT: remove clone + self.neighbor_store.set_neighbors( + storage, + self.meta_page, + neighbors_of, + new_neighbors.clone(), + stats, + ); + (pruned, new_neighbors) + } - fn get_meta_page(&self, index: &PgRelation) -> &MetaPage; + pub fn get_meta_page(&self) -> &MetaPage { + &self.meta_page + } /// greedy search looks for the closest neighbors to a query vector /// You may think that this needs the "K" parameter but it does not, @@ -459,101 +254,82 @@ pub trait Graph { /// /// Note this is the one-shot implementation that keeps only the closest `search_list_size` results in /// the returned ListSearchResult elements. It shouldn't be used with self.greedy_search_iterate - fn greedy_search( + fn greedy_search_for_build( &self, - index: &PgRelation, - query: &[f32], + query: PgVector, meta_page: &MetaPage, - ) -> (ListSearchResult, HashSet) - where - Self: Graph, - { + storage: &S, + stats: &mut GreedySearchStats, + ) -> HashSet { let init_ids = self.get_init_ids(); if let None = init_ids { //no nodes in the graph - return (ListSearchResult::empty(), HashSet::with_capacity(0)); + return HashSet::with_capacity(0); } - let dm = self.get_vector_provider().get_distance_measure(query); + let dm = storage.get_search_distance_measure(query, false); let search_list_size = meta_page.get_search_list_size_for_build() as usize; let mut l = ListSearchResult::new( - index, Some(search_list_size), - self, init_ids.unwrap(), - query, dm, search_list_size, meta_page, + self.get_neighbor_store(), + storage, ); let mut visited_nodes = HashSet::with_capacity(search_list_size); - self.greedy_search_iterate( - &mut l, - index, - query, - search_list_size, - Some(&mut visited_nodes), - ); - return (l, visited_nodes); + self.greedy_search_iterate(&mut l, search_list_size, Some(&mut visited_nodes), storage); + stats.combine(&l.stats); + return visited_nodes; } /// Returns a ListSearchResult initialized for streaming. The output should be used with greedy_search_iterate to obtain /// the next elements. - fn greedy_search_streaming_init( + pub fn greedy_search_streaming_init( &self, - index: &PgRelation, - query: &[f32], + query: PgVector, search_list_size: usize, - meta_page: &MetaPage, - ) -> ListSearchResult { + storage: &S, + ) -> ListSearchResult { let init_ids = self.get_init_ids(); if let None = init_ids { //no nodes in the graph return ListSearchResult::empty(); } - let dm = self.get_vector_provider().get_distance_measure(query); + let dm = storage.get_search_distance_measure(query, true); ListSearchResult::new( - index, None, - self, init_ids.unwrap(), - query, dm, search_list_size, - meta_page, + &self.meta_page, + self.get_neighbor_store(), + storage, ) } /// Advance the state of the lsr until the closest `visit_n_closest` elements have been visited. - fn greedy_search_iterate( + pub fn greedy_search_iterate( &self, - lsr: &mut ListSearchResult, - index: &PgRelation, - query: &[f32], + lsr: &mut ListSearchResult, visit_n_closest: usize, mut visited_nodes: Option<&mut HashSet>, - ) where - Self: Graph, - { - //OPT: Only build v when needed. - let mut neighbors = - Vec::::with_capacity(self.get_meta_page(index).get_num_neighbors() as _); - while let Some(list_search_entry) = lsr.visit_closest(visit_n_closest) { - neighbors.clear(); + storage: &S, + ) { + while let Some(list_search_entry_idx) = lsr.visit_closest(visit_n_closest) { match visited_nodes { None => {} Some(ref mut visited_nodes) => { + let list_search_entry = &lsr.candidate_storage[list_search_entry_idx]; visited_nodes.insert(NeighborWithDistance::new( list_search_entry.index_pointer, list_search_entry.distance, )); } } - neighbors.extend_from_slice(list_search_entry.neighbor_index_pointers.as_slice()); - for neighbor_index_pointer in neighbors.iter() { - lsr.insert(index, self, *neighbor_index_pointer, query) - } + storage.visit_lsn(lsr, list_search_entry_idx, &self.neighbor_store); } } @@ -562,40 +338,17 @@ pub trait Graph { /// /// TODO: this is the ann-disk implementation. There may be better implementations /// if we save the factors or the distances and add incrementally. Not sure. - fn prune_neighbors( + pub fn prune_neighbors( &self, - index: &PgRelation, - index_pointer: ItemPointer, - new_neigbors: Vec, - ) -> (Vec, PruneNeighborStats) { - let mut stats = PruneNeighborStats::new(); + mut candidates: Vec, + storage: &S, + stats: &mut PruneNeighborStats, + ) -> Vec { stats.calls += 1; //TODO make configurable? - let max_alpha = self.get_meta_page(index).get_max_alpha(); - //get a unique candidate pool - let mut candidates = Vec::::with_capacity( - (self.get_meta_page(index).get_num_neighbors() as usize) + new_neigbors.len(), - ); - self.get_neighbors_with_distances(index, index_pointer, &mut candidates); + let max_alpha = self.get_meta_page().get_max_alpha(); - let mut hash: HashSet = candidates - .iter() - .map(|c| c.get_index_pointer_to_neighbor()) - .collect(); - for n in new_neigbors { - if hash.insert(n.get_index_pointer_to_neighbor()) { - candidates.push(n); - } - } - //remove myself - if !hash.insert(index_pointer) { - //prevent self-loops - let index = candidates - .iter() - .position(|x| x.get_index_pointer_to_neighbor() == index_pointer) - .unwrap(); - candidates.remove(index); - } + stats.num_neighbors_before_prune += candidates.len(); //TODO remove deleted nodes //TODO diskann has something called max_occlusion_size/max_candidate_size(default:750). Do we need to implement? @@ -603,20 +356,19 @@ pub trait Graph { //sort by distance candidates.sort(); let mut results = Vec::::with_capacity( - self.get_meta_page(index).get_max_neighbors_during_build(), + self.get_meta_page().get_num_neighbors() as _, ); let mut max_factors: Vec = vec![0.0; candidates.len()]; let mut alpha = 1.0; + let dimension_epsilon = self.get_meta_page().get_num_dimensions() as f32 * f32::EPSILON; //first we add nodes that "pass" a small alpha. Then, if there //is still room we loop again with a larger alpha. - while alpha <= max_alpha - && results.len() < self.get_meta_page(index).get_num_neighbors() as _ - { + while alpha <= max_alpha && results.len() < self.get_meta_page().get_num_neighbors() as _ { for (i, neighbor) in candidates.iter().enumerate() { - if results.len() >= self.get_meta_page(index).get_num_neighbors() as _ { - return (results, stats); + if results.len() >= self.get_meta_page().get_num_neighbors() as _ { + return results; } if max_factors[i] > alpha { continue; @@ -630,11 +382,10 @@ pub trait Graph { //rename for clarity. let existing_neighbor = neighbor; - let vp = self.get_vector_provider(); let dist_state = unsafe { - vp.get_full_vector_distance_state( - index, + storage.get_full_vector_distance_state( existing_neighbor.get_index_pointer_to_neighbor(), + stats, ) }; @@ -647,14 +398,9 @@ pub trait Graph { //todo handle the non-pq case let mut distance_between_candidate_and_existing_neighbor = unsafe { - vp.get_distance_pair_for_full_vectors_from_state( - &dist_state, - index, - candidate_neighbor.get_index_pointer_to_neighbor(), - ) + dist_state + .get_distance(candidate_neighbor.get_index_pointer_to_neighbor(), stats) }; - stats.node_reads += 2; - stats.distance_comparisons += 1; let mut distance_between_candidate_and_point = candidate_neighbor.get_distance(); @@ -662,18 +408,24 @@ pub trait Graph { //Otherwise, the case where distance_between_candidate_and_point > 0 and distance_between_candidate_and_existing_neighbor < 0 is totally wrong. //If we implement inner product distance we'll have to figure something else out. if distance_between_candidate_and_point < 0.0 - && distance_between_candidate_and_point >= 0.0 - f32::EPSILON + && distance_between_candidate_and_point >= 0.0 - dimension_epsilon { distance_between_candidate_and_point = 0.0; } if distance_between_candidate_and_existing_neighbor < 0.0 - && distance_between_candidate_and_existing_neighbor >= 0.0 - f32::EPSILON + && distance_between_candidate_and_existing_neighbor + >= 0.0 - dimension_epsilon { distance_between_candidate_and_existing_neighbor = 0.0; } - debug_assert!(distance_between_candidate_and_point >= 0.0); + debug_assert!( + distance_between_candidate_and_point >= 0.0, + "distance_between_candidate_and_point is negative: {}, {}", + distance_between_candidate_and_point, + f32::EPSILON + ); debug_assert!(distance_between_candidate_and_existing_neighbor >= 0.0); //factor is high if the candidate is closer to an existing neighbor than the point it's being considered for @@ -693,156 +445,73 @@ pub trait Graph { } alpha = alpha * 1.2 } - (results, stats) + stats.num_neighbors_after_prune += results.len(); + results } - fn insert( + pub fn insert( &mut self, index: &PgRelation, index_pointer: IndexPointer, - vec: &[f32], - ) -> InsertStats { - let mut prune_neighbor_stats: PruneNeighborStats = PruneNeighborStats::new(); - let mut greedy_search_stats = GreedySearchStats::new(); - let meta_page = self.get_meta_page(index); - - if self.is_empty() { - self.set_neighbors( - index, + vec: PgVector, + storage: &S, + stats: &mut InsertStats, + ) { + if self.meta_page.get_init_ids().is_none() { + //TODO probably better set off of centeroids + MetaPage::update_init_ids(index, vec![index_pointer]); + *self.meta_page = MetaPage::read(index); + + self.neighbor_store.set_neighbors( + storage, + self.meta_page, index_pointer, Vec::::with_capacity( - meta_page.get_max_neighbors_during_build() as _, + self.neighbor_store.max_neighbors(self.meta_page) as _, ), + stats, ); - return InsertStats { - prune_neighbor_stats: prune_neighbor_stats, - greedy_search_stats: greedy_search_stats, - }; } + let meta_page = self.get_meta_page(); + //TODO: make configurable? - let (l, v) = self.greedy_search(index, vec, meta_page); - greedy_search_stats.combine(l.stats); - let (neighbor_list, forward_stats) = - self.prune_neighbors(index, index_pointer, v.into_iter().collect()); - prune_neighbor_stats.combine(forward_stats); + let v = + self.greedy_search_for_build(vec, meta_page, storage, &mut stats.greedy_search_stats); - //set forward pointers - self.set_neighbors(index, index_pointer, neighbor_list.clone()); + let (_, neighbor_list) = self.add_neighbors( + storage, + index_pointer, + v.into_iter().collect(), + &mut stats.prune_neighbor_stats, + ); //update back pointers let mut cnt = 0; for neighbor in neighbor_list { - let (needed_prune, backpointer_stats) = self.update_back_pointer( - index, + let needed_prune = self.update_back_pointer( neighbor.get_index_pointer_to_neighbor(), index_pointer, neighbor.get_distance(), + storage, + &mut stats.prune_neighbor_stats, ); if needed_prune { cnt = cnt + 1; } - prune_neighbor_stats.combine(backpointer_stats); } - //info!("pruned {} neighbors", cnt); - return InsertStats { - prune_neighbor_stats, - greedy_search_stats, - }; } - fn update_back_pointer( + fn update_back_pointer( &mut self, - index: &PgRelation, from: IndexPointer, to: IndexPointer, distance: f32, - ) -> (bool, PruneNeighborStats) { - let mut current_links = Vec::::new(); - self.get_neighbors_with_distances(index, from, &mut current_links); - - if current_links.len() < current_links.capacity() as _ { - current_links.push(NeighborWithDistance::new(to, distance)); - self.set_neighbors(index, from, current_links); - (false, PruneNeighborStats::new()) - } else { - //info!("sizes {} {} {}", current_links.len() + 1, current_links.capacity(), self.meta_page.get_max_neighbors_during_build()); - //Note prune_neighbors will reduce to current_links.len() to num_neighbors while capacity is num_neighbors * 1.3 - //thus we are avoiding prunning every time - let (new_list, stats) = - self.prune_neighbors(index, from, vec![NeighborWithDistance::new(to, distance)]); - self.set_neighbors(index, from, new_list); - (true, stats) - } - } -} - -#[derive(Debug)] -pub struct PruneNeighborStats { - pub calls: usize, - pub distance_comparisons: usize, - pub node_reads: usize, -} - -impl PruneNeighborStats { - pub fn new() -> Self { - PruneNeighborStats { - calls: 0, - distance_comparisons: 0, - node_reads: 0, - } - } - - pub fn combine(&mut self, other: Self) { - self.calls += other.calls; - self.distance_comparisons += other.distance_comparisons; - self.node_reads += other.node_reads; - } -} - -#[derive(Debug)] -pub struct GreedySearchStats { - pub calls: usize, - pub distance_comparisons: usize, - pub node_reads: usize, - pub pq_distance_comparisons: usize, -} - -impl GreedySearchStats { - pub fn new() -> Self { - GreedySearchStats { - calls: 0, - distance_comparisons: 0, - node_reads: 0, - pq_distance_comparisons: 0, - } - } - - pub fn combine(&mut self, other: Self) { - self.calls += other.calls; - self.distance_comparisons += other.distance_comparisons; - self.node_reads += other.node_reads; - self.pq_distance_comparisons += other.pq_distance_comparisons; - } -} - -#[derive(Debug)] -pub struct InsertStats { - pub prune_neighbor_stats: PruneNeighborStats, - pub greedy_search_stats: GreedySearchStats, -} - -impl InsertStats { - pub fn new() -> Self { - return InsertStats { - prune_neighbor_stats: PruneNeighborStats::new(), - greedy_search_stats: GreedySearchStats::new(), - }; - } - - pub fn combine(&mut self, other: InsertStats) { - self.prune_neighbor_stats - .combine(other.prune_neighbor_stats); - self.greedy_search_stats.combine(other.greedy_search_stats); + storage: &S, + prune_stats: &mut PruneNeighborStats, + ) -> bool { + let new = vec![NeighborWithDistance::new(to, distance)]; + let (pruned, _) = self.add_neighbors(storage, from, new, prune_stats); + pruned } } diff --git a/timescale_vector/src/access_method/graph_neighbor_store.rs b/timescale_vector/src/access_method/graph_neighbor_store.rs new file mode 100644 index 00000000..2e185fd3 --- /dev/null +++ b/timescale_vector/src/access_method/graph_neighbor_store.rs @@ -0,0 +1,124 @@ +use std::collections::HashMap; + +use crate::util::{IndexPointer, ItemPointer}; + +use super::stats::{StatsDistanceComparison, StatsNodeModify, StatsNodeRead}; + +use super::meta_page::MetaPage; +use super::neighbor_with_distance::*; +use super::storage::Storage; + +/// A builderGraph is a graph that keep the neighbors in-memory in the neighbor_map below +/// The idea is that during the index build, you don't want to update the actual Postgres +/// pages every time you change the neighbors. Instead you change the neighbors in memory +/// until the build is done. Afterwards, calling the `write` method, will write out all +/// the neighbors to the right pages. +pub struct BuilderNeighborCache { + //maps node's pointer to the representation on disk + neighbor_map: HashMap>, +} + +impl BuilderNeighborCache { + pub fn new() -> Self { + Self { + neighbor_map: HashMap::with_capacity(200), + } + } + pub fn iter(&self) -> impl Iterator)> { + self.neighbor_map.iter() + } + + pub fn get_neighbors(&self, neighbors_of: ItemPointer) -> Vec { + let neighbors = self.neighbor_map.get(&neighbors_of); + match neighbors { + Some(n) => n + .iter() + .map(|n| n.get_index_pointer_to_neighbor()) + .collect(), + None => vec![], + } + } + + pub fn get_neighbors_with_full_vector_distances( + &self, + neighbors_of: ItemPointer, + result: &mut Vec, + ) { + let neighbors = self.neighbor_map.get(&neighbors_of); + match neighbors { + Some(n) => { + for nwd in n { + result.push(nwd.clone()); + } + } + None => (), + } + } + + pub fn set_neighbors( + &mut self, + neighbors_of: ItemPointer, + new_neighbors: Vec, + ) { + self.neighbor_map.insert(neighbors_of, new_neighbors); + } + + pub fn max_neighbors(&self, meta_page: &MetaPage) -> usize { + meta_page.get_max_neighbors_during_build() + } +} + +pub enum GraphNeighborStore { + Builder(BuilderNeighborCache), + Disk, +} + +impl GraphNeighborStore { + pub fn get_neighbors_with_full_vector_distances< + S: Storage, + T: StatsNodeRead + StatsDistanceComparison, + >( + &self, + neighbors_of: ItemPointer, + storage: &S, + result: &mut Vec, + stats: &mut T, + ) { + match self { + GraphNeighborStore::Builder(b) => { + b.get_neighbors_with_full_vector_distances(neighbors_of, result) + } + GraphNeighborStore::Disk => storage.get_neighbors_with_full_vector_distances_from_disk( + neighbors_of, + result, + stats, + ), + }; + } + + pub fn set_neighbors( + &mut self, + storage: &S, + meta_page: &MetaPage, + neighbors_of: ItemPointer, + new_neighbors: Vec, + stats: &mut T, + ) { + match self { + GraphNeighborStore::Builder(b) => b.set_neighbors(neighbors_of, new_neighbors), + GraphNeighborStore::Disk => storage.set_neighbors_on_disk( + meta_page, + neighbors_of, + new_neighbors.as_slice(), + stats, + ), + } + } + + pub fn max_neighbors(&self, meta_page: &MetaPage) -> usize { + match self { + GraphNeighborStore::Builder(b) => b.max_neighbors(meta_page), + GraphNeighborStore::Disk => meta_page.get_num_neighbors() as _, + } + } +} diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index 83a2ef54..ebc0d9ee 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -5,9 +5,7 @@ use crate::access_method::options::TSVIndexOptions; use crate::util::page; use crate::util::*; -use super::bq::BqQuantizer; -use super::pq::PqQuantizer; -use super::quantizer::Quantizer; +use super::storage::StorageType; const TSV_MAGIC_NUMBER: u32 = 768756476; //Magic number, random const TSV_VERSION: u32 = 1; @@ -64,13 +62,13 @@ impl MetaPage { self.use_pq } - pub fn get_quantizer(&self) -> Quantizer { + pub fn get_storage_type(&self) -> StorageType { if self.get_use_pq() { - Quantizer::PQ(PqQuantizer::new()) + StorageType::PqCompression } else if self.use_bq { - Quantizer::BQ(BqQuantizer::new()) + StorageType::BqSpeedup } else { - Quantizer::None + StorageType::Plain } } diff --git a/timescale_vector/src/access_method/mod.rs b/timescale_vector/src/access_method/mod.rs index 6ae883fa..1f9620ab 100644 --- a/timescale_vector/src/access_method/mod.rs +++ b/timescale_vector/src/access_method/mod.rs @@ -1,16 +1,20 @@ use pgrx::*; mod build; -mod builder_graph; mod cost_estimate; mod debugging; -mod disk_index_graph; mod graph; +mod graph_neighbor_store; pub mod guc; mod meta_page; -mod model; +mod neighbor_with_distance; pub mod options; -mod quantizer; +pub mod pg_vector; +mod plain_node; +mod plain_storage; mod scan; +pub mod stats; +mod storage; +mod storage_common; mod vacuum; extern crate blas_src; @@ -19,7 +23,9 @@ mod bq; pub mod distance; #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] mod distance_x86; -mod pq; +mod pq_quantizer; +mod pq_quantizer_storage; +mod pq_storage; #[pg_extern(sql = " CREATE OR REPLACE FUNCTION tsv_amhandler(internal) RETURNS index_am_handler PARALLEL SAFE IMMUTABLE STRICT COST 0.0001 LANGUAGE c AS 'MODULE_PATHNAME', '@FUNCTION_NAME@'; diff --git a/timescale_vector/src/access_method/model.rs b/timescale_vector/src/access_method/model.rs deleted file mode 100644 index b10d0bfa..00000000 --- a/timescale_vector/src/access_method/model.rs +++ /dev/null @@ -1,423 +0,0 @@ -use std::cmp::Ordering; -use std::mem::size_of; -use std::pin::Pin; - -use ndarray::Array3; -use pgrx::pg_sys::{InvalidBlockNumber, InvalidOffsetNumber, BLCKSZ}; -use pgrx::*; -use reductive::pq::Pq; -use rkyv::vec::ArchivedVec; -use rkyv::{Archive, Archived, Deserialize, Serialize}; - -use crate::util::page::PageType; -use crate::util::tape::Tape; -use crate::util::{ - ArchivedItemPointer, HeapPointer, IndexPointer, ItemPointer, ReadableBuffer, WritableBuffer, -}; - -use super::distance::preprocess_cosine; -use super::meta_page::MetaPage; -use super::quantizer::Quantizer; - -//Ported from pg_vector code -#[repr(C)] -#[derive(Debug)] -pub struct PgVector { - vl_len_: i32, /* varlena header (do not touch directly!) */ - pub dim: i16, /* number of dimensions */ - unused: i16, - pub x: pg_sys::__IncompleteArrayField, -} - -impl PgVector { - pub unsafe fn from_pg_parts( - datum_parts: *mut pg_sys::Datum, - isnull_parts: *mut bool, - index: usize, - ) -> Option<*mut PgVector> { - let isnulls = std::slice::from_raw_parts(isnull_parts, index + 1); - if isnulls[index] { - return None; - } - let datums = std::slice::from_raw_parts(datum_parts, index + 1); - Some(Self::from_datum(datums[index])) - } - - pub unsafe fn from_datum(datum: pg_sys::Datum) -> *mut PgVector { - let detoasted = pg_sys::pg_detoast_datum(datum.cast_mut_ptr()); - let casted = detoasted.cast::(); - casted - } - - pub fn to_slice(&mut self) -> &[f32] { - let dim = (*self).dim; - let raw_slice = unsafe { (*self).x.as_mut_slice(dim as _) }; - preprocess_cosine(raw_slice); - raw_slice - } -} - -#[derive(Archive, Deserialize, Serialize)] -#[archive(check_bytes)] -pub struct Node { - pub vector: Vec, - pub pq_vector: Vec, - neighbor_index_pointers: Vec, - pub heap_item_pointer: HeapPointer, -} - -//ReadableNode ties an archive node to it's underlying buffer -pub struct ReadableNode<'a> { - _rb: ReadableBuffer<'a>, -} - -impl<'a> ReadableNode<'a> { - pub fn get_archived_node(&self) -> &ArchivedNode { - // checking the code here is expensive during build, so skip it. - // TODO: should we check the data during queries? - //rkyv::check_archived_root::(self._rb.get_data_slice()).unwrap() - unsafe { rkyv::archived_root::(self._rb.get_data_slice()) } - } -} - -//WritableNode ties an archive node to it's underlying buffer that can be modified -pub struct WritableNode<'a> { - wb: WritableBuffer<'a>, -} - -impl<'a> WritableNode<'a> { - pub fn get_archived_node(&self) -> Pin<&mut ArchivedNode> { - ArchivedNode::with_data(self.wb.get_data_slice()) - } - - pub fn commit(self) { - self.wb.commit() - } -} - -impl Node { - pub fn new( - full_vector: Vec, - heap_item_pointer: ItemPointer, - meta_page: &MetaPage, - quantizer: &Quantizer, - ) -> Self { - let num_neighbors = meta_page.get_num_neighbors(); - // always use vectors of num_neighbors in length because we never want the serialized size of a Node to change - let neighbor_index_pointers: Vec<_> = (0..num_neighbors) - .map(|_| ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber)) - .collect(); - - match quantizer { - Quantizer::None => Self { - vector: full_vector, - pq_vector: Vec::with_capacity(0), - neighbor_index_pointers: neighbor_index_pointers, - heap_item_pointer, - }, - Quantizer::PQ(pq) => { - let mut node = Self { - vector: Vec::with_capacity(0), - pq_vector: Vec::with_capacity(0), - neighbor_index_pointers: neighbor_index_pointers, - heap_item_pointer, - }; - pq.initialize_node(&mut node, meta_page, full_vector); - node - } - Quantizer::BQ(bq) => { - let mut node = Self { - vector: Vec::with_capacity(0), - pq_vector: Vec::with_capacity(0), - neighbor_index_pointers: neighbor_index_pointers, - heap_item_pointer, - }; - bq.initialize_node(&mut node, meta_page, full_vector); - node - } - } - } - - pub unsafe fn read<'a>(index: &'a PgRelation, index_pointer: ItemPointer) -> ReadableNode<'a> { - let rb = index_pointer.read_bytes(index); - ReadableNode { _rb: rb } - } - - pub unsafe fn modify(index: &PgRelation, index_pointer: ItemPointer) -> WritableNode { - let wb = index_pointer.modify_bytes(index); - WritableNode { wb: wb } - } - - pub fn write(&self, tape: &mut Tape) -> ItemPointer { - let bytes = rkyv::to_bytes::<_, 256>(self).unwrap(); - unsafe { tape.write(&bytes) } - } -} - -/// contains helpers for mutate-in-place. See struct_mutable_refs in test_alloc.rs in rkyv -impl ArchivedNode { - pub fn with_data(data: &mut [u8]) -> Pin<&mut ArchivedNode> { - let pinned_bytes = Pin::new(data); - unsafe { rkyv::archived_root_mut::(pinned_bytes) } - } - - pub fn is_deleted(&self) -> bool { - self.heap_item_pointer.offset == InvalidOffsetNumber - } - - pub fn delete(self: Pin<&mut Self>) { - //TODO: actually optimize the deletes by removing index tuples. For now just mark it. - let mut heap_pointer = unsafe { self.map_unchecked_mut(|s| &mut s.heap_item_pointer) }; - heap_pointer.offset = InvalidOffsetNumber; - heap_pointer.block_number = InvalidBlockNumber; - } - - pub fn neighbor_index_pointer( - self: Pin<&mut Self>, - ) -> Pin<&mut ArchivedVec> { - unsafe { self.map_unchecked_mut(|s| &mut s.neighbor_index_pointers) } - } - - pub fn pq_vectors(self: Pin<&mut Self>) -> Pin<&mut Archived>> { - unsafe { self.map_unchecked_mut(|s| &mut s.pq_vector) } - } - - pub fn num_neighbors(&self) -> usize { - self.neighbor_index_pointers - .iter() - .position(|f| f.block_number == InvalidBlockNumber) - .unwrap_or(self.neighbor_index_pointers.len()) - } - - pub fn apply_to_neighbors(&self, mut f: F) - where - F: FnMut(&ArchivedItemPointer), - { - for i in 0..self.num_neighbors() { - let neighbor = &self.neighbor_index_pointers[i]; - f(neighbor); - } - } - - pub fn set_neighbors( - mut self: Pin<&mut Self>, - neighbors: &Vec, - meta_page: &MetaPage, - ) { - for (i, new_neighbor) in neighbors.iter().enumerate() { - let mut a_index_pointer = self.as_mut().neighbor_index_pointer().index_pin(i); - //TODO hate that we have to set each field like this - a_index_pointer.block_number = - new_neighbor.get_index_pointer_to_neighbor().block_number; - a_index_pointer.offset = new_neighbor.get_index_pointer_to_neighbor().offset; - } - //set the marker that the list ended - if neighbors.len() < meta_page.get_num_neighbors() as _ { - let mut past_last_index_pointers = - self.neighbor_index_pointer().index_pin(neighbors.len()); - past_last_index_pointers.block_number = InvalidBlockNumber; - past_last_index_pointers.offset = InvalidOffsetNumber; - } - } -} - -//TODO is this right? -pub type Distance = f32; -#[derive(Clone)] -pub struct NeighborWithDistance { - index_pointer: IndexPointer, - distance: Distance, -} - -impl NeighborWithDistance { - pub fn new(neighbor_index_pointer: ItemPointer, distance: Distance) -> Self { - Self { - index_pointer: neighbor_index_pointer, - distance, - } - } - - pub fn get_index_pointer_to_neighbor(&self) -> ItemPointer { - return self.index_pointer; - } - pub fn get_distance(&self) -> Distance { - return self.distance; - } -} - -impl PartialOrd for NeighborWithDistance { - fn partial_cmp(&self, other: &Self) -> Option { - self.distance.partial_cmp(&other.distance) - } -} - -impl Ord for NeighborWithDistance { - fn cmp(&self, other: &Self) -> Ordering { - self.distance.total_cmp(&other.distance) - } -} - -impl PartialEq for NeighborWithDistance { - fn eq(&self, other: &Self) -> bool { - self.index_pointer == other.index_pointer - } -} - -//promise that PartialEq is reflexive -impl Eq for NeighborWithDistance {} - -impl std::hash::Hash for NeighborWithDistance { - fn hash(&self, state: &mut H) { - self.index_pointer.hash(state); - } -} - -#[derive(Archive, Deserialize, Serialize)] -#[archive(check_bytes)] -#[repr(C)] -pub struct PqQuantizerDef { - dim_0: usize, - dim_1: usize, - dim_2: usize, - vec_len: usize, - next_vector_pointer: ItemPointer, -} - -impl PqQuantizerDef { - pub fn new(dim_0: usize, dim_1: usize, dim_2: usize, vec_len: usize) -> PqQuantizerDef { - { - Self { - dim_0, - dim_1, - dim_2, - vec_len, - next_vector_pointer: ItemPointer { - block_number: 0, - offset: 0, - }, - } - } - } - - pub unsafe fn write(&self, tape: &mut Tape) -> ItemPointer { - let bytes = rkyv::to_bytes::<_, 256>(self).unwrap(); - tape.write(&bytes) - } - pub unsafe fn read<'a>( - index: &'a PgRelation, - index_pointer: &ItemPointer, - ) -> ReadablePqQuantizerDef<'a> { - let rb = index_pointer.read_bytes(index); - ReadablePqQuantizerDef { _rb: rb } - } -} - -pub struct ReadablePqQuantizerDef<'a> { - _rb: ReadableBuffer<'a>, -} - -impl<'a> ReadablePqQuantizerDef<'a> { - pub fn get_archived_node(&self) -> &ArchivedPqQuantizerDef { - // checking the code here is expensive during build, so skip it. - // TODO: should we check the data during queries? - //rkyv::check_archived_root::(self._rb.get_data_slice()).unwrap() - unsafe { rkyv::archived_root::(self._rb.get_data_slice()) } - } -} - -#[derive(Archive, Deserialize, Serialize)] -#[archive(check_bytes)] -#[repr(C)] -pub struct PqQuantizerVector { - vec: Vec, - next_vector_pointer: ItemPointer, -} - -impl PqQuantizerVector { - pub unsafe fn write(&self, tape: &mut Tape) -> ItemPointer { - let bytes = rkyv::to_bytes::<_, 8192>(self).unwrap(); - tape.write(&bytes) - } - pub unsafe fn read<'a>( - index: &'a PgRelation, - index_pointer: &ItemPointer, - ) -> ReadablePqVectorNode<'a> { - let rb = index_pointer.read_bytes(index); - ReadablePqVectorNode { _rb: rb } - } -} - -//ReadablePqNode ties an archive node to it's underlying buffer -pub struct ReadablePqVectorNode<'a> { - _rb: ReadableBuffer<'a>, -} - -impl<'a> ReadablePqVectorNode<'a> { - pub fn get_archived_node(&self) -> &ArchivedPqQuantizerVector { - // checking the code here is expensive during build, so skip it. - // TODO: should we check the data during queries? - //rkyv::check_archived_root::(self._rb.get_data_slice()).unwrap() - unsafe { rkyv::archived_root::(self._rb.get_data_slice()) } - } -} - -pub unsafe fn read_pq(index: &PgRelation, index_pointer: &IndexPointer) -> Pq { - let rpq = PqQuantizerDef::read(index, &index_pointer); - let rpn = rpq.get_archived_node(); - let size = rpn.dim_0 * rpn.dim_1 * rpn.dim_2; - let mut result: Vec = Vec::with_capacity(size as usize); - let mut next = rpn.next_vector_pointer.deserialize_item_pointer(); - loop { - if next.offset == 0 && next.block_number == 0 { - break; - } - let qvn = PqQuantizerVector::read(index, &next); - let vn = qvn.get_archived_node(); - result.extend(vn.vec.iter()); - next = vn.next_vector_pointer.deserialize_item_pointer(); - } - let sq = Array3::from_shape_vec( - (rpn.dim_0 as usize, rpn.dim_1 as usize, rpn.dim_2 as usize), - result, - ) - .unwrap(); - Pq::new(None, sq) -} - -pub unsafe fn write_pq(pq: &Pq, index: &PgRelation) -> ItemPointer { - let vec = pq.subquantizers().to_slice_memory_order().unwrap().to_vec(); - let shape = pq.subquantizers().dim(); - let mut pq_node = PqQuantizerDef::new(shape.0, shape.1, shape.2, vec.len()); - - let mut pqt = Tape::new(index, PageType::PqQuantizerDef); - - // write out the large vector bits. - // we write "from the back" - let mut prev: IndexPointer = ItemPointer { - block_number: 0, - offset: 0, - }; - let mut prev_vec = vec; - - // get numbers that can fit in a page by subtracting the item pointer. - let block_fit = (BLCKSZ as usize / size_of::()) - size_of::() - 64; - let mut tape = Tape::new(index, PageType::PqQuantizerVector); - loop { - let l = prev_vec.len(); - if l == 0 { - pq_node.next_vector_pointer = prev; - return pq_node.write(&mut pqt); - } - let lv = prev_vec; - let ni = if l > block_fit { l - block_fit } else { 0 }; - let (b, a) = lv.split_at(ni); - - let pqv_node = PqQuantizerVector { - vec: a.to_vec(), - next_vector_pointer: prev, - }; - let index_pointer: IndexPointer = pqv_node.write(&mut tape); - prev = index_pointer; - prev_vec = b.clone().to_vec(); - } -} diff --git a/timescale_vector/src/access_method/neighbor_with_distance.rs b/timescale_vector/src/access_method/neighbor_with_distance.rs new file mode 100644 index 00000000..0a17c1ab --- /dev/null +++ b/timescale_vector/src/access_method/neighbor_with_distance.rs @@ -0,0 +1,56 @@ +use std::cmp::Ordering; + +use crate::util::{IndexPointer, ItemPointer}; + +//TODO is this right? +pub type Distance = f32; +#[derive(Clone, Debug)] +pub struct NeighborWithDistance { + index_pointer: IndexPointer, + distance: Distance, +} + +impl NeighborWithDistance { + pub fn new(neighbor_index_pointer: ItemPointer, distance: Distance) -> Self { + assert!(!distance.is_nan()); + assert!(distance >= 0.0); + Self { + index_pointer: neighbor_index_pointer, + distance, + } + } + + pub fn get_index_pointer_to_neighbor(&self) -> ItemPointer { + return self.index_pointer; + } + pub fn get_distance(&self) -> Distance { + return self.distance; + } +} + +impl PartialOrd for NeighborWithDistance { + fn partial_cmp(&self, other: &Self) -> Option { + self.distance.partial_cmp(&other.distance) + } +} + +impl Ord for NeighborWithDistance { + fn cmp(&self, other: &Self) -> Ordering { + self.distance.total_cmp(&other.distance) + } +} + +impl PartialEq for NeighborWithDistance { + fn eq(&self, other: &Self) -> bool { + self.index_pointer == other.index_pointer + } +} + +//promise that PartialEq is reflexive +impl Eq for NeighborWithDistance {} + +impl std::hash::Hash for NeighborWithDistance { + fn hash(&self, state: &mut H) { + self.index_pointer.hash(state); + } +} diff --git a/timescale_vector/src/access_method/pg_vector.rs b/timescale_vector/src/access_method/pg_vector.rs new file mode 100644 index 00000000..2ab4fe8d --- /dev/null +++ b/timescale_vector/src/access_method/pg_vector.rs @@ -0,0 +1,76 @@ +use pgrx::*; + +use super::distance::preprocess_cosine; + +//Ported from pg_vector code +#[repr(C)] +#[derive(Debug)] +pub struct PgVectorInternal { + vl_len_: i32, /* varlena header (do not touch directly!) */ + pub dim: i16, /* number of dimensions */ + unused: i16, + pub x: pg_sys::__IncompleteArrayField, +} + +impl PgVectorInternal { + pub fn to_slice(&self) -> &[f32] { + let dim = (*self).dim; + let raw_slice = unsafe { (*self).x.as_slice(dim as _) }; + raw_slice + } +} + +pub struct PgVector { + inner: *mut PgVectorInternal, + need_pfree: bool, +} + +impl Drop for PgVector { + fn drop(&mut self) { + if self.need_pfree { + unsafe { + pg_sys::pfree(self.inner.cast()); + } + } + } +} + +impl PgVector { + pub unsafe fn from_pg_parts( + datum_parts: *mut pg_sys::Datum, + isnull_parts: *mut bool, + index: usize, + ) -> Option { + let isnulls = std::slice::from_raw_parts(isnull_parts, index + 1); + if isnulls[index] { + return None; + } + let datums = std::slice::from_raw_parts(datum_parts, index + 1); + Some(Self::from_datum(datums[index])) + } + + pub unsafe fn from_datum(datum: pg_sys::Datum) -> PgVector { + //FIXME: we are using a copy here to avoid lifetime issues and because in some cases we have to + //modify the datum in preprocess_cosine. We should find a way to avoid the copy if the vector is + //normalized and preprocess_cosine is a noop; + let detoasted = pg_sys::pg_detoast_datum_copy(datum.cast_mut_ptr()); + let is_copy = !std::ptr::eq( + detoasted.cast::(), + datum.cast_mut_ptr::(), + ); + let casted = detoasted.cast::(); + + let dim = (*casted).dim; + let raw_slice = unsafe { (*casted).x.as_mut_slice(dim as _) }; + preprocess_cosine(raw_slice); + + PgVector { + inner: casted, + need_pfree: is_copy, + } + } + + pub fn to_slice(&self) -> &[f32] { + unsafe { (*self.inner).to_slice() } + } +} diff --git a/timescale_vector/src/access_method/plain_node.rs b/timescale_vector/src/access_method/plain_node.rs new file mode 100644 index 00000000..29ba9659 --- /dev/null +++ b/timescale_vector/src/access_method/plain_node.rs @@ -0,0 +1,156 @@ +use std::pin::Pin; + +use pgrx::pg_sys::{InvalidBlockNumber, InvalidOffsetNumber}; +use pgrx::*; +use rkyv::vec::ArchivedVec; +use rkyv::{Archive, Archived, Deserialize, Serialize}; +use timescale_vector_derive::{Readable, Writeable}; + +use super::neighbor_with_distance::NeighborWithDistance; +use super::pq_quantizer::PqVectorElement; +use super::stats::{StatsNodeModify, StatsNodeRead, StatsNodeWrite}; +use super::storage::ArchivedData; +use crate::util::tape::Tape; +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 vector: Vec, + pub pq_vector: Vec, + neighbor_index_pointers: Vec, + pub heap_item_pointer: HeapPointer, +} + +impl Node { + fn new_internal( + vector: Vec, + pq_vector: Vec, + heap_item_pointer: ItemPointer, + meta_page: &MetaPage, + ) -> Self { + let num_neighbors = meta_page.get_num_neighbors(); + Self { + vector, + // always use vectors of num_clusters on length because we never want the serialized size of a Node to change + pq_vector, + // always use vectors of num_neighbors on length because we never want the serialized size of a Node to change + neighbor_index_pointers: (0..num_neighbors) + .map(|_| ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber)) + .collect(), + heap_item_pointer, + } + } + + pub fn new_for_full_vector( + vector: Vec, + heap_item_pointer: ItemPointer, + meta_page: &MetaPage, + ) -> Self { + let pq_vector = Vec::with_capacity(0); + Self::new_internal(vector, pq_vector, heap_item_pointer, meta_page) + } + + pub fn new_for_pq( + heap_item_pointer: ItemPointer, + pq_vector: Vec, + meta_page: &MetaPage, + ) -> Self { + let vector = Vec::with_capacity(0); + + Self::new_internal(vector, pq_vector, heap_item_pointer, meta_page) + } +} + +/// contains helpers for mutate-in-place. See struct_mutable_refs in test_alloc.rs in rkyv +impl ArchivedNode { + pub fn is_deleted(&self) -> bool { + self.heap_item_pointer.offset == InvalidOffsetNumber + } + + pub fn delete(self: Pin<&mut Self>) { + //TODO: actually optimize the deletes by removing index tuples. For now just mark it. + let mut heap_pointer = unsafe { self.map_unchecked_mut(|s| &mut s.heap_item_pointer) }; + heap_pointer.offset = InvalidOffsetNumber; + heap_pointer.block_number = InvalidBlockNumber; + } + + pub fn neighbor_index_pointer( + self: Pin<&mut Self>, + ) -> Pin<&mut ArchivedVec> { + unsafe { self.map_unchecked_mut(|s| &mut s.neighbor_index_pointers) } + } + + pub fn pq_vectors(self: Pin<&mut Self>) -> Pin<&mut Archived>> { + unsafe { self.map_unchecked_mut(|s| &mut s.pq_vector) } + } + + pub fn num_neighbors(&self) -> usize { + self.neighbor_index_pointers + .iter() + .position(|f| f.block_number == InvalidBlockNumber) + .unwrap_or(self.neighbor_index_pointers.len()) + } + + pub fn iter_neighbors(&self) -> impl Iterator + '_ { + self.neighbor_index_pointers + .iter() + .take(self.num_neighbors()) + .map(|ip| ip.deserialize_item_pointer()) + } + + pub fn set_neighbors( + mut self: Pin<&mut Self>, + neighbors: &[NeighborWithDistance], + meta_page: &MetaPage, + ) { + for (i, new_neighbor) in neighbors.iter().enumerate() { + let mut a_index_pointer = self.as_mut().neighbor_index_pointer().index_pin(i); + //TODO hate that we have to set each field like this + a_index_pointer.block_number = + new_neighbor.get_index_pointer_to_neighbor().block_number; + a_index_pointer.offset = new_neighbor.get_index_pointer_to_neighbor().offset; + } + //set the marker that the list ended + if neighbors.len() < meta_page.get_num_neighbors() as _ { + let mut past_last_index_pointers = + self.neighbor_index_pointer().index_pin(neighbors.len()); + past_last_index_pointers.block_number = InvalidBlockNumber; + past_last_index_pointers.offset = InvalidOffsetNumber; + } + } + + pub fn set_pq_vector(mut self: Pin<&mut Self>, pq_vector: &[u8]) { + for i in 0..=pq_vector.len() - 1 { + let mut pgv = self.as_mut().pq_vectors().index_pin(i); + *pgv = pq_vector[i]; + } + } +} + +impl ArchivedData for ArchivedNode { + fn with_data(data: &mut [u8]) -> Pin<&mut ArchivedNode> { + ArchivedNode::with_data(data) + } + + fn get_index_pointer_to_neighbors(&self) -> Vec { + self.iter_neighbors().collect() + } + + fn is_deleted(&self) -> bool { + self.heap_item_pointer.offset == InvalidOffsetNumber + } + + fn delete(self: Pin<&mut Self>) { + //TODO: actually optimize the deletes by removing index tuples. For now just mark it. + let mut heap_pointer = unsafe { self.map_unchecked_mut(|s| &mut s.heap_item_pointer) }; + heap_pointer.offset = InvalidOffsetNumber; + heap_pointer.block_number = InvalidBlockNumber; + } + + fn get_heap_item_pointer(&self) -> HeapPointer { + self.heap_item_pointer.deserialize_item_pointer() + } +} diff --git a/timescale_vector/src/access_method/plain_storage.rs b/timescale_vector/src/access_method/plain_storage.rs new file mode 100644 index 00000000..a5bbfcd2 --- /dev/null +++ b/timescale_vector/src/access_method/plain_storage.rs @@ -0,0 +1,353 @@ +use super::{ + distance::distance_cosine as default_distance, + graph::{ListSearchNeighbor, ListSearchResult}, + graph_neighbor_store::GraphNeighborStore, + pg_vector::PgVector, + plain_node::{ArchivedNode, Node, ReadableNode}, + stats::{ + GreedySearchStats, StatsDistanceComparison, StatsNodeModify, StatsNodeRead, StatsNodeWrite, + WriteStats, + }, + storage::{ArchivedData, NodeFullDistanceMeasure, Storage}, +}; + +use pgrx::PgRelation; + +use crate::util::{page::PageType, 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: fn(&[f32], &[f32]) -> f32, +} + +impl<'a> PlainStorage<'a> { + pub fn new_for_build(index: &'a PgRelation) -> PlainStorage<'a> { + Self { + index: index, + distance_fn: default_distance, + } + } + + pub fn load_for_insert(index_relation: &'a PgRelation) -> PlainStorage<'a> { + Self { + index: index_relation, + distance_fn: default_distance, + } + } + + pub fn load_for_search(index_relation: &'a PgRelation) -> PlainStorage<'a> { + Self { + index: index_relation, + distance_fn: default_distance, + } + } +} + +pub enum PlainDistanceMeasure { + Full(PgVector), +} + +impl PlainDistanceMeasure { + pub fn calculate_distance( + distance_fn: fn(&[f32], &[f32]) -> f32, + query: &[f32], + vector: &[f32], + stats: &mut S, + ) -> f32 { + assert!(vector.len() > 0); + assert!(vector.len() == query.len()); + stats.record_full_distance_comparison(); + (distance_fn)(query, vector) + } +} + +/* This is only applicable to plain, so keep here not in storage_common */ +pub struct IndexFullDistanceMeasure<'a> { + readable_node: ReadableNode<'a>, + storage: &'a PlainStorage<'a>, +} + +impl<'a> IndexFullDistanceMeasure<'a> { + pub unsafe fn with_index_pointer( + storage: &'a PlainStorage<'a>, + index_pointer: IndexPointer, + stats: &mut T, + ) -> Self { + let rn = unsafe { Node::read(storage.index, index_pointer, stats) }; + Self { + readable_node: rn, + storage: storage, + } + } + + pub unsafe fn with_readable_node( + storage: &'a PlainStorage<'a>, + readable_node: ReadableNode<'a>, + ) -> Self { + Self { + readable_node: readable_node, + storage: storage, + } + } +} + +impl<'a> NodeFullDistanceMeasure for IndexFullDistanceMeasure<'a> { + unsafe fn get_distance( + &self, + index_pointer: IndexPointer, + stats: &mut T, + ) -> f32 { + let rn1 = Node::read(self.storage.index, index_pointer, stats); + let rn2 = &self.readable_node; + let node1 = rn1.get_archived_node(); + let node2 = rn2.get_archived_node(); + assert!(node1.vector.len() > 0); + assert!(node1.vector.len() == node2.vector.len()); + let vec1 = node1.vector.as_slice(); + let vec2 = node2.vector.as_slice(); + (self.storage.get_distance_function())(vec1, vec2) + } +} + +//todo move to storage_common +pub struct PlainStorageLsnPrivateData { + pub heap_pointer: HeapPointer, + pub neighbors: Vec, +} + +impl PlainStorageLsnPrivateData { + pub fn new( + index_pointer_to_node: IndexPointer, + node: &ArchivedNode, + gns: &GraphNeighborStore, + ) -> Self { + let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); + let neighbors = match gns { + GraphNeighborStore::Disk => node.get_index_pointer_to_neighbors(), + GraphNeighborStore::Builder(b) => b.get_neighbors(index_pointer_to_node), + }; + Self { + heap_pointer: heap_pointer, + neighbors: neighbors, + } + } +} + +impl<'a> Storage for PlainStorage<'a> { + type QueryDistanceMeasure = PlainDistanceMeasure; + type NodeFullDistanceMeasure<'b> = IndexFullDistanceMeasure<'b> where Self: 'b; + type ArchivedType = ArchivedNode; + type LSNPrivateData = PlainStorageLsnPrivateData; + + fn page_type() -> PageType { + PageType::Node + } + + fn create_node( + &self, + full_vector: &[f32], + heap_pointer: HeapPointer, + meta_page: &MetaPage, + tape: &mut Tape, + stats: &mut S, + ) -> ItemPointer { + //OPT: avoid the clone? + let node = Node::new_for_full_vector(full_vector.to_vec(), heap_pointer, meta_page); + let index_pointer: IndexPointer = node.write(tape, stats); + index_pointer + } + + fn start_training(&mut self, _meta_page: &super::meta_page::MetaPage) {} + + fn add_sample(&mut self, _sample: &[f32]) {} + + fn finish_training(&mut self, _stats: &mut WriteStats) {} + + fn finalize_node_at_end_of_build( + &mut self, + meta: &MetaPage, + index_pointer: IndexPointer, + neighbors: &Vec, + stats: &mut S, + ) { + let node = unsafe { Node::modify(self.index, index_pointer, stats) }; + let mut archived = node.get_archived_node(); + archived.as_mut().set_neighbors(neighbors, &meta); + node.commit(); + } + + unsafe fn get_full_vector_distance_state<'b, S: StatsNodeRead>( + &'b self, + index_pointer: IndexPointer, + stats: &mut S, + ) -> Self::NodeFullDistanceMeasure<'b> { + IndexFullDistanceMeasure::with_index_pointer(self, index_pointer, stats) + } + + fn get_search_distance_measure( + &self, + query: PgVector, + _calc_distance_with_quantizer: bool, + ) -> PlainDistanceMeasure { + return PlainDistanceMeasure::Full(query); + } + + fn get_neighbors_with_full_vector_distances_from_disk< + S: StatsNodeRead + StatsDistanceComparison, + >( + &self, + neighbors_of: ItemPointer, + result: &mut Vec, + stats: &mut S, + ) { + let rn = unsafe { Node::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) }; + for n in neighbors { + let dist = unsafe { dist_state.get_distance(n, stats) }; + result.push(NeighborWithDistance::new(n, dist)) + } + } + + /* get_lsn and visit_lsn are different because the distance + comparisons for BQ get the vector from different places */ + fn create_lsn_for_init_id( + &self, + lsr: &mut ListSearchResult, + index_pointer: ItemPointer, + gns: &GraphNeighborStore, + ) -> ListSearchNeighbor { + if !lsr.prepare_insert(index_pointer) { + panic!("should not have had an init id already inserted"); + } + + let rn = unsafe { Node::read(self.index, index_pointer, &mut lsr.stats) }; + let node = rn.get_archived_node(); + + let distance = match lsr.sdm.as_ref().unwrap() { + PlainDistanceMeasure::Full(query) => PlainDistanceMeasure::calculate_distance( + self.distance_fn, + query.to_slice(), + node.vector.as_slice(), + &mut lsr.stats, + ), + }; + + ListSearchNeighbor::new( + index_pointer, + distance, + PlainStorageLsnPrivateData::new(index_pointer, node, gns), + ) + } + + fn visit_lsn( + &self, + lsr: &mut ListSearchResult, + lsn_idx: usize, + gns: &GraphNeighborStore, + ) { + let lsn = lsr.get_lsn_by_idx(lsn_idx); + //clone needed so we don't continue to borrow lsr + let neighbors = lsn.get_private_data().neighbors.clone(); + + for &neighbor_index_pointer in neighbors.iter() { + if !lsr.prepare_insert(neighbor_index_pointer) { + continue; + } + + let rn_neighbor = + unsafe { Node::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() { + PlainDistanceMeasure::Full(query) => PlainDistanceMeasure::calculate_distance( + self.distance_fn, + query.to_slice(), + node_neighbor.vector.as_slice(), + &mut lsr.stats, + ), + }; + let lsn = ListSearchNeighbor::new( + neighbor_index_pointer, + distance, + PlainStorageLsnPrivateData::new(neighbor_index_pointer, node_neighbor, gns), + ); + + lsr.insert_neighbor(lsn); + } + } + + fn return_lsn( + &self, + lsn: &ListSearchNeighbor, + _stats: &mut GreedySearchStats, + ) -> HeapPointer { + lsn.get_private_data().heap_pointer + } + + fn set_neighbors_on_disk( + &self, + meta: &MetaPage, + index_pointer: IndexPointer, + neighbors: &[NeighborWithDistance], + stats: &mut S, + ) { + let node = unsafe { Node::modify(self.index, index_pointer, stats) }; + let mut archived = node.get_archived_node(); + archived.as_mut().set_neighbors(neighbors, &meta); + node.commit(); + } + + fn get_distance_function(&self) -> fn(&[f32], &[f32]) -> f32 { + self.distance_fn + } +} + +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + + use pgrx::*; + + #[pg_test] + unsafe fn test_plain_storage_index_creation() -> spi::Result<()> { + crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( + "num_neighbors=38", + )?; + Ok(()) + } + + #[pg_test] + unsafe fn test_plain_storage_index_creation_few_neighbors() -> spi::Result<()> { + //a test with few neighbors tests the case that nodes share a page, which has caused deadlocks in the past. + crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( + "num_neighbors=10", + )?; + Ok(()) + } + + #[test] + fn test_plain_storage_delete_vacuum_plain() { + crate::access_method::vacuum::tests::test_delete_vacuum_plain_scaffold( + "num_neighbors = 10", + ); + } + + #[test] + fn test_plain_storage_delete_vacuum_full() { + crate::access_method::vacuum::tests::test_delete_vacuum_full_scaffold("num_neighbors = 38"); + } + + #[pg_test] + unsafe fn test_plain_storage_empty_table_insert() -> spi::Result<()> { + crate::access_method::build::tests::test_empty_table_insert_scaffold("num_neighbors=38") + } + + #[pg_test] + unsafe fn test_plain_storage_insert_empty_insert() -> spi::Result<()> { + crate::access_method::build::tests::test_insert_empty_insert_scaffold("num_neighbors=38") + } +} diff --git a/timescale_vector/src/access_method/pq.rs b/timescale_vector/src/access_method/pq_quantizer.rs similarity index 83% rename from timescale_vector/src/access_method/pq.rs rename to timescale_vector/src/access_method/pq_quantizer.rs index 1a202143..c1d757d5 100644 --- a/timescale_vector/src/access_method/pq.rs +++ b/timescale_vector/src/access_method/pq_quantizer.rs @@ -1,19 +1,17 @@ -use std::pin::Pin; - use ndarray::{Array1, Array2, Axis}; use pgrx::{error, notice, PgRelation}; use rand::Rng; use reductive::pq::{Pq, QuantizeVector, TrainPq}; -use crate::{ - access_method::{ - distance::distance_l2_optimized_for_few_dimensions, - model::{self, read_pq}, - }, - util::IndexPointer, +use crate::access_method::{ + distance::distance_l2_optimized_for_few_dimensions, pq_quantizer_storage::read_pq, }; -use super::meta_page::MetaPage; +use super::{ + meta_page::MetaPage, + pg_vector::PgVector, + stats::{StatsDistanceComparison, StatsNodeRead}, +}; /// pq aka Product quantization (PQ) is one of the most widely used algorithms for memory-efficient approximated nearest neighbor search, /// This module encapsulates a vanilla implementation of PQ that we use for the vector index. @@ -35,7 +33,10 @@ const NUM_TRAINING_ATTEMPTS: usize = 1; /// We pick a value used by DiskANN implementations. const NUM_TRAINING_SET_SIZE: usize = 256000; +pub type PqVectorElement = u8; /// PqTrainer is a utility that produces a product quantizer from training with sample vectors. + +#[derive(Clone)] pub struct PqTrainer { /// training_set contains the vectors we'll use to train PQ. training_set: Vec>, @@ -135,6 +136,7 @@ fn build_distance_table( distance_table } +#[derive(Clone)] pub struct PqQuantizer { pq_trainer: Option, pq: Option>, @@ -148,41 +150,30 @@ impl PqQuantizer { } } - pub fn load(&mut self, index_relation: &PgRelation, meta_page: &super::meta_page::MetaPage) { - assert!(self.pq_trainer.is_none()); + pub fn load( + index_relation: &PgRelation, + meta_page: &super::meta_page::MetaPage, + stats: &mut S, + ) -> Self { let pq_item_pointer = meta_page.get_pq_pointer().unwrap(); - self.pq = unsafe { Some(read_pq(&index_relation, &pq_item_pointer)) }; - } + let pq = unsafe { Some(read_pq(&index_relation, pq_item_pointer, stats)) }; - pub fn initialize_node( - &self, - node: &mut super::model::Node, - meta_page: &MetaPage, - full_vector: Vec, - ) { - if self.pq_trainer.is_some() { - let pq_vec_len = meta_page.get_pq_vector_length(); - node.pq_vector = (0..pq_vec_len).map(|_| 0u8).collect(); - } else { - assert!(self.pq.is_some()); - let pq_vec_len = meta_page.get_pq_vector_length(); - node.pq_vector = self.quantize(full_vector); - assert!(node.pq_vector.len() == pq_vec_len); + Self { + pq_trainer: None, + pq: pq, } } - pub fn update_node_after_traing( - &self, - archived: &mut Pin<&mut super::model::ArchivedNode>, - full_vector: Vec, - ) { - let pq_vector = self.quantize(full_vector); - - assert!(pq_vector.len() == archived.pq_vector.len()); - for i in 0..=pq_vector.len() - 1 { - let mut pgv = archived.as_mut().pq_vectors().index_pin(i); - *pgv = pq_vector[i]; - } + pub fn must_get_pq(&self) -> &Pq { + self.pq.as_ref().unwrap() + } + + pub fn quantize(&self, full_vector: &[f32]) -> Vec { + assert!(self.pq.is_some()); + let pq = self.pq.as_ref().unwrap(); + //OPT is the copy really necessary? + let array_vec = Array1::from(full_vector.to_vec()); + pq.quantize_vector(array_vec).to_vec() } pub fn start_training(&mut self, meta_page: &super::meta_page::MetaPage) { @@ -197,18 +188,21 @@ impl PqQuantizer { self.pq = Some(self.pq_trainer.take().unwrap().train_pq()); } - pub fn write_metadata(&self, index: &PgRelation) { - assert!(self.pq.is_some()); - let index_pointer: IndexPointer = - unsafe { model::write_pq(self.pq.as_ref().unwrap(), &index) }; - super::meta_page::MetaPage::update_pq_pointer(&index, index_pointer); - } - - pub fn quantize(&self, full_vector: Vec) -> Vec { - assert!(self.pq.is_some()); - let pq = self.pq.as_ref().unwrap(); - let array_vec = Array1::from(full_vector); - pq.quantize_vector(array_vec).to_vec() + pub fn vector_for_new_node( + &self, + meta_page: &MetaPage, + full_vector: &[f32], + ) -> Vec { + if self.pq_trainer.is_some() { + let pq_vec_len = meta_page.get_pq_vector_length(); + vec![0; pq_vec_len] + } else { + assert!(self.pq.is_some()); + let pq_vec_len = meta_page.get_pq_vector_length(); + let res = self.quantize(full_vector); + assert!(res.len() == pq_vec_len); + res + } } pub fn get_distance_table( @@ -249,3 +243,21 @@ impl PqDistanceTable { d } } + +pub enum PqSearchDistanceMeasure { + Full(PgVector), + Pq(PqDistanceTable), +} + +impl PqSearchDistanceMeasure { + pub fn calculate_pq_distance( + table: &PqDistanceTable, + pq_vector: &[PqVectorElement], + stats: &mut S, + ) -> f32 { + assert!(pq_vector.len() > 0); + let vec = pq_vector; + stats.record_quantized_distance_comparison(); + table.distance(vec) + } +} diff --git a/timescale_vector/src/access_method/pq_quantizer_storage.rs b/timescale_vector/src/access_method/pq_quantizer_storage.rs new file mode 100644 index 00000000..f76668a3 --- /dev/null +++ b/timescale_vector/src/access_method/pq_quantizer_storage.rs @@ -0,0 +1,123 @@ +use std::mem::size_of; +use std::pin::Pin; + +use ndarray::Array3; +use pgrx::pg_sys::BLCKSZ; +use pgrx::*; +use reductive::pq::Pq; + +use rkyv::{Archive, Deserialize, Serialize}; +use timescale_vector_derive::{Readable, Writeable}; + +use crate::util::page::PageType; +use crate::util::tape::Tape; +use crate::util::{IndexPointer, ItemPointer, ReadableBuffer, WritableBuffer}; + +use super::stats::{StatsNodeModify, StatsNodeRead, StatsNodeWrite}; + +#[derive(Archive, Deserialize, Serialize, Readable, Writeable)] +#[archive(check_bytes)] +#[repr(C)] +pub struct PqQuantizerDef { + dim_0: usize, + dim_1: usize, + dim_2: usize, + vec_len: usize, + next_vector_pointer: ItemPointer, +} + +impl PqQuantizerDef { + pub fn new(dim_0: usize, dim_1: usize, dim_2: usize, vec_len: usize) -> PqQuantizerDef { + { + Self { + dim_0, + dim_1, + dim_2, + vec_len, + next_vector_pointer: ItemPointer { + block_number: 0, + offset: 0, + }, + } + } + } +} + +#[derive(Archive, Deserialize, Serialize, Readable, Writeable)] +#[archive(check_bytes)] +#[repr(C)] +pub struct PqQuantizerVector { + vec: Vec, + next_vector_pointer: ItemPointer, +} + +pub unsafe fn read_pq( + index: &PgRelation, + index_pointer: IndexPointer, + stats: &mut S, +) -> Pq { + //TODO: handle stats better + let rpq = PqQuantizerDef::read(index, index_pointer, stats); + stats.record_read(); + let rpn = rpq.get_archived_node(); + let size = rpn.dim_0 * rpn.dim_1 * rpn.dim_2; + let mut result: Vec = Vec::with_capacity(size as usize); + let mut next = rpn.next_vector_pointer.deserialize_item_pointer(); + loop { + if next.offset == 0 && next.block_number == 0 { + break; + } + let qvn = PqQuantizerVector::read(index, next, stats); + let vn = qvn.get_archived_node(); + result.extend(vn.vec.iter()); + next = vn.next_vector_pointer.deserialize_item_pointer(); + } + let sq = Array3::from_shape_vec( + (rpn.dim_0 as usize, rpn.dim_1 as usize, rpn.dim_2 as usize), + result, + ) + .unwrap(); + Pq::new(None, sq) +} + +pub unsafe fn write_pq( + pq: &Pq, + index: &PgRelation, + stats: &mut S, +) -> ItemPointer { + let vec = pq.subquantizers().to_slice_memory_order().unwrap().to_vec(); + let shape = pq.subquantizers().dim(); + let mut pq_node = PqQuantizerDef::new(shape.0, shape.1, shape.2, vec.len()); + + let mut pqt = Tape::new(index, PageType::PqQuantizerDef); + + // write out the large vector bits. + // we write "from the back" + let mut prev: IndexPointer = ItemPointer { + block_number: 0, + offset: 0, + }; + let mut prev_vec = vec; + + // get numbers that can fit in a page by subtracting the item pointer. + let block_fit = (BLCKSZ as usize / size_of::()) - size_of::() - 64; + let mut tape = Tape::new(index, PageType::PqQuantizerVector); + loop { + let l = prev_vec.len(); + if l == 0 { + pq_node.next_vector_pointer = prev; + return pq_node.write(&mut pqt, stats); + } + let lv = prev_vec; + let ni = if l > block_fit { l - block_fit } else { 0 }; + let (b, a) = lv.split_at(ni); + + let pqv_node = PqQuantizerVector { + vec: a.to_vec(), + next_vector_pointer: prev, + }; + let index_pointer: IndexPointer = pqv_node.write(&mut tape, stats); + prev = index_pointer; + prev_vec = b.to_vec(); + } +} diff --git a/timescale_vector/src/access_method/pq_storage.rs b/timescale_vector/src/access_method/pq_storage.rs new file mode 100644 index 00000000..366dd26e --- /dev/null +++ b/timescale_vector/src/access_method/pq_storage.rs @@ -0,0 +1,446 @@ +use super::{ + distance::distance_cosine as default_distance, + graph::{ListSearchNeighbor, ListSearchResult}, + graph_neighbor_store::GraphNeighborStore, + pg_vector::PgVector, + plain_node::{ArchivedNode, Node}, + plain_storage::PlainStorageLsnPrivateData, + pq_quantizer::{PqQuantizer, PqSearchDistanceMeasure, PqVectorElement}, + pq_quantizer_storage::write_pq, + stats::{ + GreedySearchStats, StatsDistanceComparison, StatsNodeModify, StatsNodeRead, StatsNodeWrite, + WriteStats, + }, + storage::{NodeFullDistanceMeasure, Storage, StorageFullDistanceFromHeap}, + storage_common::{calculate_full_distance, HeapFullDistanceMeasure}, +}; + +use pgrx::PgRelation; + +use crate::util::{ + page::PageType, table_slot::TableSlot, tape::Tape, HeapPointer, IndexPointer, ItemPointer, +}; + +use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance}; + +pub struct PqCompressionStorage<'a> { + pub index: &'a PgRelation, + pub distance_fn: fn(&[f32], &[f32]) -> f32, + quantizer: PqQuantizer, + heap_rel: Option<&'a PgRelation>, + heap_attr: Option, +} + +impl<'a> PqCompressionStorage<'a> { + pub fn new_for_build( + index: &'a PgRelation, + heap_rel: &'a PgRelation, + heap_attr: pgrx::pg_sys::AttrNumber, + ) -> PqCompressionStorage<'a> { + Self { + index: index, + distance_fn: default_distance, + quantizer: PqQuantizer::new(), + heap_rel: Some(heap_rel), + heap_attr: Some(heap_attr), + } + } + + fn load_quantizer( + index_relation: &PgRelation, + meta_page: &super::meta_page::MetaPage, + stats: &mut S, + ) -> PqQuantizer { + PqQuantizer::load(&index_relation, meta_page, stats) + } + + pub fn load_for_insert( + heap_rel: &'a PgRelation, + heap_attr: pgrx::pg_sys::AttrNumber, + index_relation: &'a PgRelation, + meta_page: &super::meta_page::MetaPage, + stats: &mut S, + ) -> PqCompressionStorage<'a> { + Self { + index: index_relation, + distance_fn: default_distance, + quantizer: Self::load_quantizer(index_relation, meta_page, stats), + heap_rel: Some(heap_rel), + heap_attr: Some(heap_attr), + } + } + + pub fn load_for_search( + index_relation: &'a PgRelation, + quantizer: &PqQuantizer, + ) -> PqCompressionStorage<'a> { + Self { + index: index_relation, + distance_fn: default_distance, + //OPT: get rid of clone + quantizer: quantizer.clone(), + heap_rel: None, + heap_attr: None, + } + } + + fn get_quantized_vector_from_heap_pointer( + &self, + heap_pointer: HeapPointer, + stats: &mut S, + ) -> Vec { + let slot = unsafe { self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) }; + let slice = unsafe { slot.get_pg_vector() }; + self.quantizer.quantize(slice.to_slice()) + } + + fn write_quantizer_metadata(&self, stats: &mut S) { + let pq = self.quantizer.must_get_pq(); + let index_pointer: IndexPointer = unsafe { write_pq(pq, &self.index, stats) }; + super::meta_page::MetaPage::update_pq_pointer(&self.index, index_pointer); + } + + fn visit_lsn_internal( + &self, + lsr: &mut ListSearchResult< + as Storage>::QueryDistanceMeasure, + as Storage>::LSNPrivateData, + >, + neighbors: &[ItemPointer], + gns: &GraphNeighborStore, + ) { + for &neighbor_index_pointer in neighbors.iter() { + if !lsr.prepare_insert(neighbor_index_pointer) { + continue; + } + + let rn_neighbor = + unsafe { Node::read(self.index, neighbor_index_pointer, &mut lsr.stats) }; + let node_neighbor = rn_neighbor.get_archived_node(); + let deleted = node_neighbor.is_deleted(); + + let distance = match lsr.sdm.as_ref().unwrap() { + PqSearchDistanceMeasure::Full(query) => { + let heap_pointer = node_neighbor.heap_item_pointer.deserialize_item_pointer(); + if deleted { + let pvt_data = PlainStorageLsnPrivateData::new( + neighbor_index_pointer, + node_neighbor, + gns, + ); + self.visit_lsn_internal(lsr, &pvt_data.neighbors, gns); + continue; + //for deleted nodes, we can't get the distance because we don't know the full vector + //so pretend it's the same distance as the parent + } else { + unsafe { + calculate_full_distance( + self, + heap_pointer, + query.to_slice(), + &mut lsr.stats, + ) + } + } + } + PqSearchDistanceMeasure::Pq(table) => { + PqSearchDistanceMeasure::calculate_pq_distance( + table, + node_neighbor.pq_vector.as_slice(), + &mut lsr.stats, + ) + } + }; + let lsn = ListSearchNeighbor::new( + neighbor_index_pointer, + distance, + PlainStorageLsnPrivateData::new(neighbor_index_pointer, node_neighbor, gns), + ); + + lsr.insert_neighbor(lsn); + } + } +} + +impl<'a> Storage for PqCompressionStorage<'a> { + type QueryDistanceMeasure = PqSearchDistanceMeasure; + type NodeFullDistanceMeasure<'b> = HeapFullDistanceMeasure<'b, PqCompressionStorage<'b>> where Self: 'b; + type ArchivedType = ArchivedNode; + type LSNPrivateData = PlainStorageLsnPrivateData; //no data stored + + fn page_type() -> PageType { + PageType::Node + } + + fn create_node( + &self, + full_vector: &[f32], + heap_pointer: HeapPointer, + meta_page: &MetaPage, + tape: &mut Tape, + stats: &mut S, + ) -> ItemPointer { + let pq_vector = self.quantizer.vector_for_new_node(meta_page, full_vector); + let node = Node::new_for_pq(heap_pointer, pq_vector, meta_page); + let index_pointer: IndexPointer = node.write(tape, stats); + index_pointer + } + + fn start_training(&mut self, meta_page: &super::meta_page::MetaPage) { + self.quantizer.start_training(meta_page); + } + + fn add_sample(&mut self, sample: &[f32]) { + self.quantizer.add_sample(sample); + } + + fn finish_training(&mut self, stats: &mut WriteStats) { + self.quantizer.finish_training(); + self.write_quantizer_metadata(stats); + } + + fn finalize_node_at_end_of_build( + &mut self, + meta: &MetaPage, + index_pointer: IndexPointer, + neighbors: &Vec, + stats: &mut S, + ) { + let node = unsafe { Node::modify(self.index, index_pointer, stats) }; + let mut archived = node.get_archived_node(); + archived.as_mut().set_neighbors(neighbors, &meta); + + let quantized = self.get_quantized_vector_from_heap_pointer( + archived.heap_item_pointer.deserialize_item_pointer(), + stats, + ); + + archived.as_mut().set_pq_vector(quantized.as_slice()); + + node.commit(); + } + + unsafe fn get_full_vector_distance_state<'b, S: StatsNodeRead>( + &'b self, + index_pointer: IndexPointer, + stats: &mut S, + ) -> HeapFullDistanceMeasure<'b, PqCompressionStorage<'b>> { + HeapFullDistanceMeasure::with_index_pointer(self, index_pointer, stats) + } + + fn get_search_distance_measure( + &self, + query: PgVector, + calc_distance_with_quantizer: bool, + ) -> PqSearchDistanceMeasure { + if !calc_distance_with_quantizer { + return PqSearchDistanceMeasure::Full(query); + } else { + return PqSearchDistanceMeasure::Pq( + self.quantizer + .get_distance_table(query.to_slice(), self.distance_fn), + ); + } + } + + //todo: same as Bq code? + fn get_neighbors_with_full_vector_distances_from_disk< + S: StatsNodeRead + StatsDistanceComparison, + >( + &self, + neighbors_of: ItemPointer, + result: &mut Vec, + stats: &mut S, + ) { + let rn = unsafe { Node::read(self.index, neighbors_of, stats) }; + let heap_pointer = rn + .get_archived_node() + .heap_item_pointer + .deserialize_item_pointer(); + let dist_state = + unsafe { HeapFullDistanceMeasure::with_heap_pointer(self, heap_pointer, stats) }; + for n in rn.get_archived_node().iter_neighbors() { + let dist = unsafe { dist_state.get_distance(n, stats) }; + result.push(NeighborWithDistance::new(n, dist)) + } + } + + /* get_lsn and visit_lsn are different because the distance + comparisons for BQ get the vector from different places */ + fn create_lsn_for_init_id( + &self, + lsr: &mut ListSearchResult, + index_pointer: ItemPointer, + gns: &GraphNeighborStore, + ) -> ListSearchNeighbor { + if !lsr.prepare_insert(index_pointer) { + panic!("should not have had an init id already inserted"); + } + + let rn = unsafe { Node::read(self.index, index_pointer, &mut lsr.stats) }; + let node = rn.get_archived_node(); + + let distance = match lsr.sdm.as_ref().unwrap() { + PqSearchDistanceMeasure::Full(query) => { + if node.is_deleted() { + //FIXME: need to handle this case + panic!("can't handle the case where the init_id node is deleted"); + } + let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); + unsafe { + calculate_full_distance(self, heap_pointer, query.to_slice(), &mut lsr.stats) + } + } + PqSearchDistanceMeasure::Pq(table) => PqSearchDistanceMeasure::calculate_pq_distance( + table, + node.pq_vector.as_slice(), + &mut lsr.stats, + ), + }; + + ListSearchNeighbor::new( + index_pointer, + distance, + PlainStorageLsnPrivateData::new(index_pointer, node, gns), + ) + } + + fn visit_lsn( + &self, + lsr: &mut ListSearchResult, + lsn_idx: usize, + gns: &GraphNeighborStore, + ) { + let lsn = lsr.get_lsn_by_idx(lsn_idx); + //clone needed so we don't continue to borrow lsr + self.visit_lsn_internal(lsr, &lsn.get_private_data().neighbors.clone(), gns); + } + + fn return_lsn( + &self, + lsn: &ListSearchNeighbor, + _stats: &mut GreedySearchStats, + ) -> HeapPointer { + lsn.get_private_data().heap_pointer + } + + fn set_neighbors_on_disk( + &self, + meta: &MetaPage, + index_pointer: IndexPointer, + neighbors: &[NeighborWithDistance], + stats: &mut S, + ) { + let node = unsafe { Node::modify(self.index, index_pointer, stats) }; + let mut archived = node.get_archived_node(); + archived.as_mut().set_neighbors(neighbors, &meta); + node.commit(); + } + + fn get_distance_function(&self) -> fn(&[f32], &[f32]) -> f32 { + self.distance_fn + } +} + +impl<'a> StorageFullDistanceFromHeap for PqCompressionStorage<'a> { + unsafe fn get_heap_table_slot_from_index_pointer( + &self, + index_pointer: IndexPointer, + stats: &mut S, + ) -> TableSlot { + let rn = unsafe { Node::read(self.index, index_pointer, stats) }; + let node = rn.get_archived_node(); + let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); + + self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) + } + + unsafe fn get_heap_table_slot_from_heap_pointer( + &self, + heap_pointer: HeapPointer, + stats: &mut T, + ) -> TableSlot { + TableSlot::new( + self.heap_rel.unwrap(), + heap_pointer, + self.heap_attr.unwrap(), + stats, + ) + } +} + +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +mod tests { + use pgrx::*; + + #[pg_test] + unsafe fn test_pq_storage_index_creation() -> spi::Result<()> { + crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( + "num_neighbors=38, USE_PQ = TRUE", + )?; + Ok(()) + } + + #[pg_test] + unsafe fn test_pq_storage_index_creation_few_neighbors() -> spi::Result<()> { + //a test with few neighbors tests the case that nodes share a page, which has caused deadlocks in the past. + crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( + "num_neighbors=10, USE_PQ = TRUE", + )?; + Ok(()) + } + + #[test] + fn test_pq_storage_delete_vacuum_plain() { + crate::access_method::vacuum::tests::test_delete_vacuum_plain_scaffold( + "num_neighbors = 10, use_pq = TRUE", + ); + } + + #[test] + fn test_pq_storage_delete_vacuum_full() { + crate::access_method::vacuum::tests::test_delete_vacuum_full_scaffold( + "num_neighbors = 38, use_pq = TRUE", + ); + } + + /* can't run test_pq_storage_empty_table_insert because can't create pq index on pq table */ + + #[pg_test] + unsafe fn test_pq_storage_insert_empty_insert() -> spi::Result<()> { + let suffix = (1..=253) + .map(|i| format!("{}", i)) + .collect::>() + .join(", "); + + Spi::run(&format!( + "CREATE TABLE test(embedding vector(256)); + + INSERT INTO test (embedding) + SELECT + ('[' || i || ',2,3,{suffix}]')::vector + FROM generate_series(1, 300) i; + + CREATE INDEX idxtest + ON test + USING tsv(embedding) + WITH (num_neighbors = 10, use_pq = TRUE); + + DELETE FROM test; + + INSERT INTO test(embedding) VALUES ('[1,2,3,{suffix}]'), ('[14,15,16,{suffix}]'); + ", + ))?; + + let res: Option = Spi::get_one(&format!( + " set enable_seqscan = 0; + WITH cte as (select * from test order by embedding <=> '[0,0,0,{suffix}]') SELECT count(*) from cte;", + ))?; + assert_eq!(2, res.unwrap()); + + Spi::run(&format!("drop index idxtest;",))?; + + Ok(()) + } +} diff --git a/timescale_vector/src/access_method/quantizer.rs b/timescale_vector/src/access_method/quantizer.rs deleted file mode 100644 index c11e6098..00000000 --- a/timescale_vector/src/access_method/quantizer.rs +++ /dev/null @@ -1,23 +0,0 @@ -use super::{bq::BqQuantizer, pq::PqQuantizer}; - -/*pub trait Quantizer { - fn initialize_node(&self, node: &mut Node, meta_page: &MetaPage); - fn start_training(&mut self, meta_page: &super::meta_page::MetaPage); - fn add_sample(&mut self, sample: Vec); - fn finish_training(&mut self); -}*/ - -pub enum Quantizer { - BQ(BqQuantizer), - PQ(PqQuantizer), - None, -} - -impl Quantizer { - pub fn is_some(&self) -> bool { - match self { - Quantizer::None => false, - _ => true, - } - } -} diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 61125bc7..39b69a92 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -2,68 +2,129 @@ use pgrx::{pg_sys::InvalidOffsetNumber, *}; use crate::{ access_method::{ - disk_index_graph::DiskIndexGraph, graph::VectorProvider, meta_page::MetaPage, - model::PgVector, + bq::BqSpeedupStorage, graph_neighbor_store::GraphNeighborStore, meta_page::MetaPage, + pg_vector::PgVector, }, util::{buffer::PinnedBufferShare, HeapPointer}, }; -use super::{graph::ListSearchResult, quantizer::Quantizer}; +use super::{ + bq::{BqMeans, BqQuantizer, BqSearchDistanceMeasure, BqSpeedupStorageLsnPrivateData}, + graph::{Graph, ListSearchResult}, + plain_storage::{PlainDistanceMeasure, PlainStorage, PlainStorageLsnPrivateData}, + pq_quantizer::{PqQuantizer, PqSearchDistanceMeasure}, + pq_storage::PqCompressionStorage, + stats::QuantizerStats, + storage::{Storage, StorageType}, +}; + +/* Be very careful not to transfer PgRelations in the state, as they can change between calls. That means we shouldn't be +using lifetimes here. Everything should be owned */ +enum StorageState { + BqSpeedup( + BqQuantizer, + TSVResponseIterator, + ), + PqCompression( + PqQuantizer, + TSVResponseIterator, + ), + Plain(TSVResponseIterator), +} + +/* no lifetime usage here. */ +struct TSVScanState { + storage: *mut StorageState, +} + +impl TSVScanState { + fn new() -> Self { + Self { + storage: std::ptr::null_mut(), + } + } + + fn initialize(&mut self, index: &PgRelation, query: PgVector, search_list_size: usize) { + let meta_page = MetaPage::read(&index); + let storage = meta_page.get_storage_type(); + + let store_type = match storage { + StorageType::Plain => { + let stats = QuantizerStats::new(); + let bq = PlainStorage::load_for_search(index); + let it = + TSVResponseIterator::new(&bq, index, query, search_list_size, meta_page, stats); + StorageState::Plain(it) + } + StorageType::PqCompression => { + let mut stats = QuantizerStats::new(); + let quantizer = PqQuantizer::load(index, &meta_page, &mut stats); + let pq = PqCompressionStorage::load_for_search(index, &quantizer); + let it = + TSVResponseIterator::new(&pq, index, query, search_list_size, meta_page, stats); + StorageState::PqCompression(quantizer, it) + } + StorageType::BqSpeedup => { + let mut stats = QuantizerStats::new(); + let quantizer = unsafe { BqMeans::load(index, &meta_page, &mut stats) }; + let bq = BqSpeedupStorage::load_for_search(index, &quantizer); + let it = + TSVResponseIterator::new(&bq, index, query, search_list_size, meta_page, stats); + StorageState::BqSpeedup(quantizer, it) + } + }; + + self.storage = PgMemoryContexts::CurrentMemoryContext.leak_and_drop_on_delete(store_type); + } +} -struct TSVResponseIterator<'a> { - query: Vec, - lsr: ListSearchResult, +struct TSVResponseIterator { + lsr: ListSearchResult, search_list_size: usize, current: usize, - last_buffer: Option>, - quantizer: Quantizer, + last_buffer: Option, + meta_page: MetaPage, + quantizer_stats: QuantizerStats, } -impl<'a> TSVResponseIterator<'a> { - fn new(index: &PgRelation, query: &[f32], search_list_size: usize) -> Self { - let meta_page = MetaPage::read(&index); - let mut quantizer = meta_page.get_quantizer(); - match &mut quantizer { - Quantizer::None => {} - Quantizer::PQ(pq) => pq.load(index, &meta_page), - Quantizer::BQ(bq) => bq.load(index, &meta_page), - } - let graph = DiskIndexGraph::new( - &index, - VectorProvider::new(None, None, &quantizer, quantizer.is_some()), - ); - use super::graph::Graph; - let lsr = graph.greedy_search_streaming_init(&index, query, search_list_size, &meta_page); +impl TSVResponseIterator { + fn new>( + storage: &S, + index: &PgRelation, + query: PgVector, + search_list_size: usize, + _meta_page: MetaPage, + quantizer_stats: QuantizerStats, + ) -> Self { + let mut meta_page = MetaPage::read(&index); + let graph = Graph::new(GraphNeighborStore::Disk, &mut meta_page); + + let lsr = graph.greedy_search_streaming_init(query, search_list_size, storage); + Self { - query: query.to_vec(), search_list_size, lsr, current: 0, last_buffer: None, - quantizer, + meta_page, + quantizer_stats, } } } -impl<'a> TSVResponseIterator<'a> { - fn next(&mut self, index: &'a PgRelation) -> Option { - let graph = DiskIndexGraph::new( - &index, - VectorProvider::new(None, None, &self.quantizer, self.quantizer.is_some()), - ); - use super::graph::Graph; +impl TSVResponseIterator { + fn next>( + &mut self, + index: &PgRelation, + storage: &S, + ) -> Option { + let graph = Graph::new(GraphNeighborStore::Disk, &mut self.meta_page); /* Iterate until we find a non-deleted tuple */ loop { - graph.greedy_search_iterate( - &mut self.lsr, - index, - &self.query, - self.search_list_size, - None, - ); + graph.greedy_search_iterate(&mut self.lsr, self.search_list_size, None, storage); - let item = self.lsr.consume(); + let item = self.lsr.consume(storage); match item { Some((heap_pointer, index_pointer)) => { @@ -92,10 +153,11 @@ impl<'a> TSVResponseIterator<'a> { } } -struct TSVScanState<'a> { - iterator: *mut TSVResponseIterator<'a>, +/* +struct TSVScanState<'a, 'b> { + iterator: *mut TSVResponseIterator<'a, 'b>, } - +*/ #[pg_guard] pub extern "C" fn ambeginscan( index_relation: pg_sys::Relation, @@ -109,10 +171,8 @@ pub extern "C" fn ambeginscan( norderbys, )) }; - let state = TSVScanState { - iterator: std::ptr::null_mut(), - }; + let state: TSVScanState = TSVScanState::new(); scandesc.opaque = PgMemoryContexts::CurrentMemoryContext.leak_and_drop_on_delete(state) as void_mut_ptr; @@ -135,7 +195,8 @@ pub extern "C" fn amrescan( } let mut scan: PgBox = unsafe { PgBox::from_pg(scan) }; let indexrel = unsafe { PgRelation::from_pg(scan.indexRelation) }; - let state = unsafe { (scan.opaque as *mut TSVScanState).as_mut() }.expect("no scandesc state"); + let meta_page = MetaPage::read(&indexrel); + let _storage = meta_page.get_storage_type(); if nkeys > 0 { scan.xs_recheck = true; @@ -144,16 +205,25 @@ pub extern "C" fn amrescan( let orderby_keys = unsafe { std::slice::from_raw_parts(orderbys as *const pg_sys::ScanKeyData, norderbys as _) }; - let vec = unsafe { PgVector::from_datum(orderby_keys[0].sk_argument) }; - let query = unsafe { (*vec).to_slice() }; + let query = unsafe { PgVector::from_datum(orderby_keys[0].sk_argument) }; //TODO need to set search_list_size correctly //TODO right now doesn't handle more than LIMIT 100; let search_list_size = super::guc::TSV_QUERY_SEARCH_LIST_SIZE.get() as usize; - let res = TSVResponseIterator::new(&indexrel, query, search_list_size); - - state.iterator = PgMemoryContexts::CurrentMemoryContext.leak_and_drop_on_delete(res); + let state = unsafe { (scan.opaque as *mut TSVScanState).as_mut() }.expect("no scandesc state"); + state.initialize(&indexrel, query, search_list_size); + /*match &mut storage { + Storage::None => pgrx::error!("not implemented"), + Storage::PQ(_pq) => pgrx::error!("not implemented"), + Storage::BQ(_bq) => { + let state = + unsafe { (scan.opaque as *mut TSVScanState).as_mut() }.expect("no scandesc state"); + + let res = TSVResponseIterator::new(&indexrel, query, search_list_size); + state.iterator = PgMemoryContexts::CurrentMemoryContext.leak_and_drop_on_delete(res); + } + }*/ } #[pg_guard] @@ -161,15 +231,37 @@ pub extern "C" fn amgettuple( scan: pg_sys::IndexScanDesc, _direction: pg_sys::ScanDirection, ) -> bool { - let mut scan: PgBox = unsafe { PgBox::from_pg(scan) }; + let scan: PgBox = unsafe { PgBox::from_pg(scan) }; let state = unsafe { (scan.opaque as *mut TSVScanState).as_mut() }.expect("no scandesc state"); - let iter = unsafe { state.iterator.as_mut() }.expect("no iterator in state"); + //let iter = unsafe { state.iterator.as_mut() }.expect("no iterator in state"); let indexrel = unsafe { PgRelation::from_pg(scan.indexRelation) }; - /* no need to recheck stuff for now */ + let mut storage = unsafe { state.storage.as_mut() }.expect("no storage in state"); + match &mut storage { + StorageState::BqSpeedup(quantizer, iter) => { + let bq = BqSpeedupStorage::load_for_search(&indexrel, quantizer); + get_tuple(&bq, &indexrel, iter, scan) + } + StorageState::PqCompression(quantizer, iter) => { + let pq = PqCompressionStorage::load_for_search(&indexrel, quantizer); + get_tuple(&pq, &indexrel, iter, scan) + } + StorageState::Plain(iter) => { + let bq = PlainStorage::load_for_search(&indexrel); + get_tuple(&bq, &indexrel, iter, scan) + } + } +} + +fn get_tuple<'a, S: Storage>( + storage: &S, + index: &'a PgRelation, + iter: &'a mut TSVResponseIterator, + mut scan: PgBox, +) -> bool { scan.xs_recheckorderby = false; - match iter.next(&indexrel) { + match iter.next(&index, storage) { Some(heap_pointer) => { let tid_to_set = &mut scan.xs_heaptid; heap_pointer.to_item_pointer_data(tid_to_set); @@ -190,99 +282,26 @@ pub extern "C" fn amendscan(scan: pg_sys::IndexScanDesc) { let scan: PgBox = unsafe { PgBox::from_pg(scan) }; let state = unsafe { (scan.opaque as *mut TSVScanState).as_mut() }.expect("no scandesc state"); - let iter = unsafe { state.iterator.as_mut() }.expect("no iterator in state"); - debug1!( - "Query stats - node reads:{}, calls: {}, distance comparisons: {}, pq distance comparisons: {}", - iter.lsr.stats.node_reads, - iter.lsr.stats.calls, - iter.lsr.stats.distance_comparisons, - iter.lsr.stats.pq_distance_comparisons, - ); - } -} -#[cfg(any(test, feature = "pg_test"))] -#[pgrx::pg_schema] -mod tests { - use pgrx::*; - - //TODO: add test where inserting and querying with vectors that are all the same. - - #[pg_test] - unsafe fn test_index_scan() -> spi::Result<()> { - Spi::run(&format!( - "CREATE TABLE test(embedding vector(3)); - - INSERT INTO test(embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[7,8,10]'); - - INSERT INTO test(embedding) SELECT ('[' || g::text ||', 1.0, 1.0]')::vector FROM generate_series(0, 100) g; - - CREATE INDEX idxtest - ON test - USING tsv(embedding) - WITH (num_neighbors=30);" - ))?; - - Spi::run(&format!( - " - set enable_seqscan = 0; - select * from test order by embedding <=> '[0,0,0]'; - explain analyze select * from test order by embedding <=> '[0,0,0]'; - ", - ))?; - - Spi::run(&format!( - " - set enable_seqscan = 0; - set tsv.query_search_list_size = 2; - select * from test order by embedding <=> '[0,0,0]'; - ", - ))?; - - let res: Option = Spi::get_one(&format!( - " - set enable_seqscan = 0; - set tsv.query_search_list_size = 2; - WITH cte as (select * from test order by embedding <=> '[0,0,0]') SELECT count(*) from cte; - ", - ))?; - - assert_eq!(104, res.unwrap(), "Testing query over entire table"); - - Spi::run(&format!( - " - drop index idxtest; - ", - ))?; - - Ok(()) + let mut storage = unsafe { state.storage.as_mut() }.expect("no storage in state"); + match &mut storage { + StorageState::BqSpeedup(_bq, iter) => end_scan::(iter), + StorageState::PqCompression(_pq, iter) => end_scan::(iter), + StorageState::Plain(iter) => end_scan::(iter), + } } +} - #[pg_test] - unsafe fn test_index_scan_on_empty_table() -> spi::Result<()> { - Spi::run(&format!( - "CREATE TABLE test(embedding vector(3)); - - CREATE INDEX idxtest - ON test - USING tsv(embedding) - WITH (num_neighbors=30);" - ))?; - - Spi::run(&format!( - " - set enable_seqscan = 0; - select * from test order by embedding <=> '[0,0,0]'; - explain analyze select * from test order by embedding <=> '[0,0,0]'; - ", - ))?; - - Spi::run(&format!( - " - drop index idxtest; - ", - ))?; - - Ok(()) - } +fn end_scan( + iter: &mut TSVResponseIterator, +) { + debug1!( + "Query stats - node reads:{}, calls: {}, total distance comparisons: {}, quantized distance comparisons: {}, quantizer r/w: {}/{}", + iter.lsr.stats.get_node_reads(), + iter.lsr.stats.get_calls(), + iter.lsr.stats.get_total_distance_comparisons(), + iter.lsr.stats.get_quantized_distance_comparisons(), + iter.quantizer_stats.node_reads, + iter.quantizer_stats.node_writes, + ); } diff --git a/timescale_vector/src/access_method/stats.rs b/timescale_vector/src/access_method/stats.rs new file mode 100644 index 00000000..5e780cdb --- /dev/null +++ b/timescale_vector/src/access_method/stats.rs @@ -0,0 +1,238 @@ +use std::time::Instant; + +pub trait StatsNodeRead { + fn record_read(&mut self); +} + +pub trait StatsNodeModify { + fn record_modify(&mut self); +} + +pub trait StatsNodeWrite { + fn record_write(&mut self); +} + +pub trait StatsDistanceComparison { + fn record_full_distance_comparison(&mut self); + fn record_quantized_distance_comparison(&mut self); +} + +#[derive(Debug)] +pub struct PruneNeighborStats { + pub calls: usize, + pub distance_comparisons: usize, + pub node_reads: usize, + pub node_modify: usize, + pub num_neighbors_before_prune: usize, + pub num_neighbors_after_prune: usize, +} + +impl PruneNeighborStats { + pub fn new() -> Self { + PruneNeighborStats { + calls: 0, + distance_comparisons: 0, + node_reads: 0, + node_modify: 0, + num_neighbors_before_prune: 0, + num_neighbors_after_prune: 0, + } + } +} + +impl StatsDistanceComparison for PruneNeighborStats { + fn record_full_distance_comparison(&mut self) { + self.distance_comparisons += 1; + } + + fn record_quantized_distance_comparison(&mut self) { + pgrx::error!("Should not use quantized distance comparisons during pruning"); + } +} + +impl StatsNodeRead for PruneNeighborStats { + fn record_read(&mut self) { + self.node_reads += 1; + } +} + +impl StatsNodeModify for PruneNeighborStats { + fn record_modify(&mut self) { + self.node_modify += 1; + } +} + +#[derive(Debug)] +pub struct GreedySearchStats { + calls: usize, + full_distance_comparisons: usize, + node_reads: usize, + quantized_distance_comparisons: usize, +} + +impl GreedySearchStats { + pub fn new() -> Self { + GreedySearchStats { + calls: 0, + full_distance_comparisons: 0, + node_reads: 0, + quantized_distance_comparisons: 0, + } + } + + pub fn combine(&mut self, other: &Self) { + self.calls += other.calls; + self.full_distance_comparisons += other.full_distance_comparisons; + self.node_reads += other.node_reads; + self.quantized_distance_comparisons += other.quantized_distance_comparisons; + } + + pub fn get_calls(&self) -> usize { + self.calls + } + + pub fn get_node_reads(&self) -> usize { + self.node_reads + } + + pub fn get_total_distance_comparisons(&self) -> usize { + self.full_distance_comparisons + self.quantized_distance_comparisons + } + + pub fn get_quantized_distance_comparisons(&self) -> usize { + self.quantized_distance_comparisons + } + + pub fn get_full_distance_comparisons(&self) -> usize { + self.full_distance_comparisons + } + + pub fn record_call(&mut self) { + self.calls += 1; + } +} + +impl StatsNodeRead for GreedySearchStats { + fn record_read(&mut self) { + self.node_reads += 1; + } +} + +impl StatsDistanceComparison for GreedySearchStats { + fn record_full_distance_comparison(&mut self) { + self.full_distance_comparisons += 1; + } + + fn record_quantized_distance_comparison(&mut self) { + self.quantized_distance_comparisons += 1; + } +} + +#[derive(Debug)] +pub struct QuantizerStats { + pub node_reads: usize, + pub node_writes: usize, +} + +impl QuantizerStats { + pub fn new() -> Self { + QuantizerStats { + node_reads: 0, + node_writes: 0, + } + } +} + +impl StatsNodeRead for QuantizerStats { + fn record_read(&mut self) { + self.node_reads += 1; + } +} + +impl StatsNodeWrite for QuantizerStats { + fn record_write(&mut self) { + self.node_writes += 1; + } +} +#[derive(Debug)] +pub struct InsertStats { + pub prune_neighbor_stats: PruneNeighborStats, + pub greedy_search_stats: GreedySearchStats, + pub quantizer_stats: QuantizerStats, + pub node_reads: usize, + pub node_modify: usize, + pub node_writes: usize, +} + +impl InsertStats { + pub fn new() -> Self { + return InsertStats { + prune_neighbor_stats: PruneNeighborStats::new(), + greedy_search_stats: GreedySearchStats::new(), + quantizer_stats: QuantizerStats::new(), + node_reads: 0, + node_modify: 0, + node_writes: 0, + }; + } +} + +impl StatsNodeRead for InsertStats { + fn record_read(&mut self) { + self.node_reads += 1; + } +} + +impl StatsNodeModify for InsertStats { + fn record_modify(&mut self) { + self.node_modify += 1; + } +} + +impl StatsNodeWrite for InsertStats { + fn record_write(&mut self) { + self.node_writes += 1; + } +} + +pub struct WriteStats { + pub started: Instant, + pub num_nodes: usize, + pub nodes_read: usize, + pub nodes_modified: usize, + pub nodes_written: usize, + pub prune_stats: PruneNeighborStats, + pub num_neighbors: usize, +} + +impl WriteStats { + pub fn new() -> Self { + Self { + started: Instant::now(), + num_nodes: 0, + prune_stats: PruneNeighborStats::new(), + num_neighbors: 0, + nodes_read: 0, + nodes_modified: 0, + nodes_written: 0, + } + } +} + +impl StatsNodeRead for WriteStats { + fn record_read(&mut self) { + self.nodes_read += 1; + } +} + +impl StatsNodeModify for WriteStats { + fn record_modify(&mut self) { + self.nodes_modified += 1; + } +} + +impl StatsNodeWrite for WriteStats { + fn record_write(&mut self) { + self.nodes_written += 1; + } +} diff --git a/timescale_vector/src/access_method/storage.rs b/timescale_vector/src/access_method/storage.rs new file mode 100644 index 00000000..eb502db1 --- /dev/null +++ b/timescale_vector/src/access_method/storage.rs @@ -0,0 +1,141 @@ +use std::pin::Pin; + +use crate::util::{ + page::PageType, table_slot::TableSlot, tape::Tape, HeapPointer, IndexPointer, ItemPointer, +}; + +use super::{ + graph::{ListSearchNeighbor, ListSearchResult}, + graph_neighbor_store::GraphNeighborStore, + meta_page::MetaPage, + neighbor_with_distance::NeighborWithDistance, + pg_vector::PgVector, + stats::{ + GreedySearchStats, StatsDistanceComparison, StatsNodeModify, StatsNodeRead, StatsNodeWrite, + WriteStats, + }, +}; + +pub trait NodeFullDistanceMeasure { + unsafe fn get_distance( + &self, + index_pointer: IndexPointer, + stats: &mut S, + ) -> f32; +} + +pub trait ArchivedData { + fn with_data(data: &mut [u8]) -> Pin<&mut Self>; + fn is_deleted(&self) -> bool; + fn delete(self: Pin<&mut Self>); + fn get_heap_item_pointer(&self) -> HeapPointer; + fn get_index_pointer_to_neighbors(&self) -> Vec; +} + +pub trait Storage { + type QueryDistanceMeasure; + type NodeFullDistanceMeasure<'a>: NodeFullDistanceMeasure + where + Self: 'a; + type ArchivedType: ArchivedData; + type LSNPrivateData; + + fn page_type() -> PageType; + + fn create_node( + &self, + full_vector: &[f32], + heap_pointer: HeapPointer, + meta_page: &MetaPage, + tape: &mut Tape, + stats: &mut S, + ) -> ItemPointer; + + fn start_training(&mut self, meta_page: &super::meta_page::MetaPage); + fn add_sample(&mut self, sample: &[f32]); + fn finish_training(&mut self, stats: &mut WriteStats); + + fn finalize_node_at_end_of_build( + &mut self, + meta: &MetaPage, + index_pointer: IndexPointer, + neighbors: &Vec, + stats: &mut S, + ); + + unsafe fn get_full_vector_distance_state<'a, S: StatsNodeRead>( + &'a self, + index_pointer: IndexPointer, + stats: &mut S, + ) -> Self::NodeFullDistanceMeasure<'a>; + + fn get_search_distance_measure( + &self, + query: PgVector, + calc_distance_with_quantizer: bool, + ) -> Self::QueryDistanceMeasure; + + fn visit_lsn( + &self, + lsr: &mut ListSearchResult, + lsn_idx: usize, + gns: &GraphNeighborStore, + ) where + Self: Sized; + + fn create_lsn_for_init_id( + &self, + lsr: &mut ListSearchResult, + index_pointer: ItemPointer, + gns: &GraphNeighborStore, + ) -> ListSearchNeighbor + where + Self: Sized; + + fn return_lsn( + &self, + lsn: &ListSearchNeighbor, + stats: &mut GreedySearchStats, + ) -> HeapPointer + where + Self: Sized; + + fn get_neighbors_with_full_vector_distances_from_disk< + S: StatsNodeRead + StatsDistanceComparison, + >( + &self, + neighbors_of: ItemPointer, + result: &mut Vec, + stats: &mut S, + ); + + fn set_neighbors_on_disk( + &self, + meta: &MetaPage, + index_pointer: IndexPointer, + neighbors: &[NeighborWithDistance], + stats: &mut S, + ); + + fn get_distance_function(&self) -> fn(&[f32], &[f32]) -> f32; +} + +pub trait StorageFullDistanceFromHeap { + unsafe fn get_heap_table_slot_from_index_pointer( + &self, + index_pointer: IndexPointer, + stats: &mut T, + ) -> TableSlot; + + unsafe fn get_heap_table_slot_from_heap_pointer( + &self, + heap_pointer: HeapPointer, + stats: &mut T, + ) -> TableSlot; +} + +pub enum StorageType { + BqSpeedup, + PqCompression, + Plain, +} diff --git a/timescale_vector/src/access_method/storage_common.rs b/timescale_vector/src/access_method/storage_common.rs new file mode 100644 index 00000000..c55d06c2 --- /dev/null +++ b/timescale_vector/src/access_method/storage_common.rs @@ -0,0 +1,74 @@ +use crate::util::{HeapPointer, IndexPointer}; + +use super::{ + pg_vector::PgVector, + stats::{StatsDistanceComparison, StatsNodeRead}, + storage::{NodeFullDistanceMeasure, Storage, StorageFullDistanceFromHeap}, +}; + +pub struct HeapFullDistanceMeasure<'a, S: Storage + StorageFullDistanceFromHeap> { + storage: &'a S, + vector: PgVector, +} + +impl<'a, S: Storage + StorageFullDistanceFromHeap> HeapFullDistanceMeasure<'a, S> { + pub unsafe fn with_index_pointer( + storage: &'a S, + index_pointer: IndexPointer, + stats: &mut T, + ) -> Self { + let slot = storage.get_heap_table_slot_from_index_pointer(index_pointer, stats); + Self { + storage: storage, + vector: slot.get_pg_vector(), + } + } + + pub unsafe fn with_heap_pointer( + storage: &'a S, + heap_pointer: HeapPointer, + stats: &mut T, + ) -> Self { + let slot = storage.get_heap_table_slot_from_heap_pointer(heap_pointer, stats); + Self { + storage: storage, + vector: slot.get_pg_vector(), + } + } +} + +impl<'a, S: Storage + StorageFullDistanceFromHeap> NodeFullDistanceMeasure + for HeapFullDistanceMeasure<'a, S> +{ + unsafe fn get_distance( + &self, + index_pointer: IndexPointer, + stats: &mut T, + ) -> f32 { + let slot = self + .storage + .get_heap_table_slot_from_index_pointer(index_pointer, stats); + stats.record_full_distance_comparison(); + let slice1 = slot.get_pg_vector(); + let slice2 = &self.vector; + (self.storage.get_distance_function())(slice1.to_slice(), slice2.to_slice()) + } +} + +pub unsafe fn calculate_full_distance< + S: Storage + StorageFullDistanceFromHeap, + T: StatsNodeRead + StatsDistanceComparison, +>( + storage: &S, + heap_pointer: HeapPointer, + query: &[f32], + stats: &mut T, +) -> f32 { + let slot = storage.get_heap_table_slot_from_heap_pointer(heap_pointer, stats); + let slice = unsafe { slot.get_pg_vector() }; + + stats.record_full_distance_comparison(); + let dist = (storage.get_distance_function())(slice.to_slice(), query); + debug_assert!(!dist.is_nan()); + dist +} diff --git a/timescale_vector/src/access_method/vacuum.rs b/timescale_vector/src/access_method/vacuum.rs index c723f2e2..c542daab 100644 --- a/timescale_vector/src/access_method/vacuum.rs +++ b/timescale_vector/src/access_method/vacuum.rs @@ -1,14 +1,24 @@ -use pgrx::{pg_sys::FirstOffsetNumber, *}; +use pgrx::{ + pg_sys::{FirstOffsetNumber, IndexBulkDeleteResult}, + *, +}; use crate::{ - access_method::model::ArchivedNode, + access_method::{ + bq::BqSpeedupStorage, meta_page::MetaPage, plain_storage::PlainStorage, + pq_storage::PqCompressionStorage, + }, util::{ - page::{PageType, WritablePage}, + page::WritablePage, ports::{PageGetItem, PageGetItemId, PageGetMaxOffsetNumber}, ItemPointer, }, }; +use crate::access_method::storage::ArchivedData; + +use super::storage::{Storage, StorageType}; + #[pg_guard] pub extern "C" fn ambulkdelete( info: *mut pg_sys::IndexVacuumInfo, @@ -29,9 +39,51 @@ pub extern "C" fn ambulkdelete( pg_sys::ForkNumber_MAIN_FORKNUM, ) }; + + let meta_page = MetaPage::read(&index_relation); + let storage = meta_page.get_storage_type(); + match storage { + StorageType::BqSpeedup => { + bulk_delete_for_storage::( + &index_relation, + nblocks, + results, + callback, + callback_state, + ); + } + StorageType::PqCompression => { + bulk_delete_for_storage::( + &index_relation, + nblocks, + results, + callback, + callback_state, + ); + } + StorageType::Plain => { + bulk_delete_for_storage::( + &index_relation, + nblocks, + results, + callback, + callback_state, + ); + } + } + results +} + +fn bulk_delete_for_storage( + index: &PgRelation, + nblocks: u32, + results: *mut IndexBulkDeleteResult, + callback: pg_sys::IndexBulkDeleteCallback, + callback_state: *mut ::std::os::raw::c_void, +) { for block_number in 0..nblocks { - let page = unsafe { WritablePage::cleanup(&index_relation, block_number) }; - if page.get_type() != PageType::Node { + let page = unsafe { WritablePage::cleanup(&index, block_number) }; + if page.get_type() != S::page_type() { continue; } let mut modified = false; @@ -45,13 +97,13 @@ pub extern "C" fn ambulkdelete( let item = PageGetItem(*page, item_id) as *mut u8; let len = (*item_id).lp_len(); let data = std::slice::from_raw_parts_mut(item, len as _); - let node = ArchivedNode::with_data(data); + let node = S::ArchivedType::with_data(data); if node.is_deleted() { continue; } - let heap_pointer: ItemPointer = node.heap_item_pointer.deserialize_item_pointer(); + let heap_pointer: ItemPointer = node.get_heap_item_pointer(); let mut ctid: pg_sys::ItemPointerData = pg_sys::ItemPointerData { ..Default::default() }; @@ -71,7 +123,6 @@ pub extern "C" fn ambulkdelete( page.commit(); } } - results } #[pg_guard] @@ -97,11 +148,18 @@ pub extern "C" fn amvacuumcleanup( #[cfg(any(test, feature = "pg_test"))] #[pgrx::pg_schema] -mod tests { +pub mod tests { + use once_cell::sync::Lazy; use pgrx::*; + use std::sync::Mutex; + + static VAC_PLAIN_MUTEX: Lazy> = Lazy::new(Mutex::default); + + #[cfg(test)] + pub fn test_delete_vacuum_plain_scaffold(index_options: &str) { + //do not run this test in parallel. (pgrx tests run in a txn rolled back after each test, but we do not have that luxury here). + let _lock = VAC_PLAIN_MUTEX.lock().unwrap(); - #[test] - fn test_delete_vacuum_plain() { //we need to run vacuum in this test which cannot be run from SPI. //so we cannot use the pg_test framework here. Thus we do a bit of //hackery to bring up the test db and then use a client to run queries against it. @@ -114,29 +172,43 @@ mod tests { ) .unwrap(); + let suffix = (1..=253) + .map(|i| format!("{}", i)) + .collect::>() + .join(", "); + let (mut client, _) = pgrx_tests::client().unwrap(); client - .batch_execute( - "CREATE TABLE test_vac(embedding vector(3)); + .batch_execute(&format!( + "CREATE TABLE test_vac(embedding vector(256)); - INSERT INTO test_vac(embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[7,8,10]'); + INSERT INTO test_vac (embedding) + SELECT + ('[' || i || ',1,1,{suffix}]')::vector + FROM generate_series(1, 300) i; + + + INSERT INTO test_vac(embedding) VALUES ('[1,2,3,{suffix}]'), ('[4,5,6,{suffix}]'), ('[7,8,10,{suffix}]'); CREATE INDEX idxtest_vac ON test_vac USING tsv(embedding) - WITH (num_neighbors=30); - ", - ) + WITH ({index_options}); + " + )) .unwrap(); client.execute("set enable_seqscan = 0;", &[]).unwrap(); - let cnt: i64 = client.query_one("WITH cte as (select * from test_vac order by embedding <=> '[1,1,1]') SELECT count(*) from cte;", &[]).unwrap().get(0); + let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); - assert_eq!(cnt, 3); + assert_eq!(cnt, 303); client - .execute("DELETE FROM test_vac WHERE embedding = '[1,2,3]';", &[]) + .execute( + &format!("DELETE FROM test_vac WHERE embedding = '[1,2,3,{suffix}]';"), + &[], + ) .unwrap(); client.close().unwrap(); @@ -148,22 +220,40 @@ mod tests { //inserts into the previous 1,2,3 spot that was deleted client .execute( - "INSERT INTO test_vac(embedding) VALUES ('[10,12,13]');", + &format!("INSERT INTO test_vac(embedding) VALUES ('[10,12,13,{suffix}]');"), + &[], + ) + .unwrap(); + + client.execute("set enable_seqscan = 0;", &[]).unwrap(); + let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); + //if the old index is still used the count is 304 + assert_eq!(cnt, 303); + + //do another delete for same items (noop) + client + .execute( + &format!("DELETE FROM test_vac WHERE embedding = '[1,2,3,{suffix}]';"), &[], ) .unwrap(); client.execute("set enable_seqscan = 0;", &[]).unwrap(); - let cnt: i64 = client.query_one("WITH cte as (select * from test_vac order by embedding <=> '[1,1,1]') SELECT count(*) from cte;", &[]).unwrap().get(0); - //if the old index is still used the count is 4 - assert_eq!(cnt, 3); + let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); + //if the old index is still used the count is 304 + assert_eq!(cnt, 303); client.execute("DROP INDEX idxtest_vac", &[]).unwrap(); client.execute("DROP TABLE test_vac", &[]).unwrap(); } - #[test] - fn test_delete_vacuum_full() { + static VAC_FULL_MUTEX: Lazy> = Lazy::new(Mutex::default); + + #[cfg(test)] + pub fn test_delete_vacuum_full_scaffold(index_options: &str) { + //do not run this test in parallel + let _lock = VAC_FULL_MUTEX.lock().unwrap(); + //we need to run vacuum in this test which cannot be run from SPI. //so we cannot use the pg_test framework here. Thus we do a bit of //hackery to bring up the test db and then use a client to run queries against it. @@ -178,27 +268,50 @@ mod tests { let (mut client, _) = pgrx_tests::client().unwrap(); + let suffix = (1..=253) + .map(|i| format!("{}", i)) + .collect::>() + .join(", "); + client - .batch_execute( - "CREATE TABLE test_vac_full(embedding vector(3)); + .batch_execute(&format!( + "CREATE TABLE test_vac_full(embedding vector(256)); - INSERT INTO test_vac_full(embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[7,8,10]'); + -- generate 300 vectors + INSERT INTO test_vac_full (embedding) + SELECT + ('[' || i || ',2,3,{suffix}]')::vector + FROM generate_series(1, 300) i; + + INSERT INTO test_vac_full(embedding) VALUES ('[1,2,3,{suffix}]'), ('[4,5,6,{suffix}]'), ('[7,8,10,{suffix}]'); CREATE INDEX idxtest_vac_full ON test_vac_full USING tsv(embedding) - WITH (num_neighbors=30); - ", - ) + WITH ({index_options}); + " + )) .unwrap(); client.execute("set enable_seqscan = 0;", &[]).unwrap(); - let cnt: i64 = client.query_one("WITH cte as (select * from test_vac_full order by embedding <=> '[1,1,1]') SELECT count(*) from cte;", &[]).unwrap().get(0); - - assert_eq!(cnt, 3); + let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac_full order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); + std::thread::sleep(std::time::Duration::from_millis(10000)); + assert_eq!(cnt, 303); client.execute("DELETE FROM test_vac_full", &[]).unwrap(); + client + .execute( + &format!( + "INSERT INTO test_vac_full (embedding) + SELECT + ('[' || i || ',2,3,{suffix}]')::vector + FROM generate_series(1, 300) i;" + ), + &[], + ) + .unwrap(); + client.close().unwrap(); let (mut client, _) = pgrx_tests::client().unwrap(); @@ -206,56 +319,22 @@ mod tests { client .execute( - "INSERT INTO test_vac_full(embedding) VALUES ('[1,2,3]');", + &format!("INSERT INTO test_vac_full(embedding) VALUES ('[1,2,3,{suffix}]');"), &[], ) .unwrap(); client.execute("set enable_seqscan = 0;", &[]).unwrap(); - let cnt: i64 = client.query_one("WITH cte as (select * from test_vac_full order by embedding <=> '[1,1,1]') SELECT count(*) from cte;", &[]).unwrap().get(0); - assert_eq!(cnt, 1); + let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac_full order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); + assert_eq!(cnt, 301); client.execute("DROP INDEX idxtest_vac_full", &[]).unwrap(); client.execute("DROP TABLE test_vac_full", &[]).unwrap(); } + #[pg_test] ///This function is only a mock to bring up the test framewokr in test_delete_vacuum fn test_delete_mock_fn() -> spi::Result<()> { Ok(()) } - - #[pg_test] - unsafe fn test_delete() -> spi::Result<()> { - Spi::run(&format!( - "CREATE TABLE test(embedding vector(3)); - - INSERT INTO test(embedding) VALUES ('[1,2,3]'), ('[4,5,6]'), ('[7,8,10]'); - - CREATE INDEX idxtest - ON test - USING tsv(embedding) - WITH (num_neighbors=30); - - DELETE FROM test WHERE embedding = '[1,2,3]'; - ", - ))?; - - let res: Option = Spi::get_one(&format!( - " set enable_seqscan = 0; - WITH cte as (select * from test order by embedding <=> '[1,1,1]') SELECT count(*) from cte;", - ))?; - assert_eq!(2, res.unwrap()); - - //delete same thing again -- should be a no-op; - Spi::run(&format!("DELETE FROM test WHERE embedding = '[1,2,3]';",))?; - let res: Option = Spi::get_one(&format!( - " set enable_seqscan = 0; - WITH cte as (select * from test order by embedding <=> '[1,1,1]') SELECT count(*) from cte;", - ))?; - assert_eq!(2, res.unwrap()); - - Spi::run(&format!("drop index idxtest;",))?; - - Ok(()) - } } diff --git a/timescale_vector/src/util/buffer.rs b/timescale_vector/src/util/buffer.rs index 611d961d..da71291a 100644 --- a/timescale_vector/src/util/buffer.rs +++ b/timescale_vector/src/util/buffer.rs @@ -197,19 +197,19 @@ impl<'a> Deref for LockedBufferShare<'a> { /// has been pinned but not locked. /// /// It is probably not a good idea to hold on to this too long except during an index scan. -/// Does not use a LWLock. -pub struct PinnedBufferShare<'a> { - _relation: &'a PgRelation, +/// Does not use a LWLock. Note a pinned buffer is valid whether or not the relation that read it +/// is still open. +pub struct PinnedBufferShare { buffer: Buffer, } -impl<'a> PinnedBufferShare<'a> { +impl PinnedBufferShare { /// read return buffer for the given blockNumber in a relation. /// /// The returned block will be pinned /// /// Safety: Safe because it checks the block number doesn't overflow. ReadBufferExtended will throw an error if the block number is out of range for the relation - pub fn read(index: &'a PgRelation, block: BlockNumber) -> Self { + pub fn read(index: &PgRelation, block: BlockNumber) -> Self { let fork_number = ForkNumber_MAIN_FORKNUM; unsafe { @@ -220,15 +220,12 @@ impl<'a> PinnedBufferShare<'a> { ReadBufferMode_RBM_NORMAL, std::ptr::null_mut(), ); - PinnedBufferShare { - _relation: index, - buffer: buf, - } + PinnedBufferShare { buffer: buf } } } } -impl<'a> Drop for PinnedBufferShare<'a> { +impl Drop for PinnedBufferShare { /// drop both unlock and unpins the buffer. fn drop(&mut self) { unsafe { diff --git a/timescale_vector/src/util/mod.rs b/timescale_vector/src/util/mod.rs index 9bf47f60..0f23fc38 100644 --- a/timescale_vector/src/util/mod.rs +++ b/timescale_vector/src/util/mod.rs @@ -1,6 +1,7 @@ pub mod buffer; pub mod page; pub mod ports; +pub mod table_slot; pub mod tape; use pgrx::PgRelation; diff --git a/timescale_vector/src/util/page.rs b/timescale_vector/src/util/page.rs index 64a758ab..823a0259 100644 --- a/timescale_vector/src/util/page.rs +++ b/timescale_vector/src/util/page.rs @@ -27,6 +27,7 @@ pub enum PageType { PqQuantizerDef = 2, PqQuantizerVector = 3, BqMeans = 4, + BqNode = 5, } impl PageType { @@ -37,6 +38,7 @@ impl PageType { 2 => PageType::PqQuantizerDef, 3 => PageType::PqQuantizerVector, 4 => PageType::BqMeans, + 5 => PageType::BqNode, _ => panic!("Unknown PageType number {}", value), } } diff --git a/timescale_vector/src/util/table_slot.rs b/timescale_vector/src/util/table_slot.rs new file mode 100644 index 00000000..147c4752 --- /dev/null +++ b/timescale_vector/src/util/table_slot.rs @@ -0,0 +1,60 @@ +use pgrx::pg_sys::{Datum, TupleTableSlot}; +use pgrx::{pg_sys, PgBox, PgRelation}; + +use crate::access_method::pg_vector::PgVector; +use crate::access_method::stats::StatsNodeRead; +use crate::util::ports::slot_getattr; +use crate::util::HeapPointer; + +pub struct TableSlot { + slot: PgBox, + attribute_number: pg_sys::AttrNumber, +} + +impl TableSlot { + pub unsafe fn new( + heap_rel: &PgRelation, + heap_pointer: HeapPointer, + attribute_number: pg_sys::AttrNumber, + stats: &mut S, + ) -> Self { + let slot = PgBox::from_pg(pg_sys::table_slot_create( + heap_rel.as_ptr(), + std::ptr::null_mut(), + )); + + let table_am = heap_rel.rd_tableam; + let fetch_row_version = (*table_am).tuple_fetch_row_version.unwrap(); + let mut ctid: pg_sys::ItemPointerData = pg_sys::ItemPointerData { + ..Default::default() + }; + heap_pointer.to_item_pointer_data(&mut ctid); + fetch_row_version( + heap_rel.as_ptr(), + &mut ctid, + &mut pg_sys::SnapshotAnyData, + slot.as_ptr(), + ); + stats.record_read(); + + Self { + slot, + attribute_number, + } + } + + unsafe fn get_attribute(&self, attribute_number: pg_sys::AttrNumber) -> Option { + slot_getattr(&self.slot, attribute_number) + } + + pub unsafe fn get_pg_vector(&self) -> PgVector { + let vector = PgVector::from_datum(self.get_attribute(self.attribute_number).unwrap()); + vector + } +} + +impl Drop for TableSlot { + fn drop(&mut self) { + unsafe { pg_sys::ExecDropSingleTupleTableSlot(self.slot.as_ptr()) }; + } +} diff --git a/timescale_vector/timescale_vector_derive/Cargo.toml b/timescale_vector/timescale_vector_derive/Cargo.toml new file mode 100644 index 00000000..3f62da20 --- /dev/null +++ b/timescale_vector/timescale_vector_derive/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "timescale_vector_derive" +version = "0.1.0" +edition = "2021" + +[lib] +proc-macro = true + +[dependencies] +syn = "1.0" +quote = "1.0" \ No newline at end of file diff --git a/timescale_vector/timescale_vector_derive/src/lib.rs b/timescale_vector/timescale_vector_derive/src/lib.rs new file mode 100644 index 00000000..2843b75a --- /dev/null +++ b/timescale_vector/timescale_vector_derive/src/lib.rs @@ -0,0 +1,91 @@ +use proc_macro::TokenStream; +use quote::{format_ident, quote}; + +#[proc_macro_derive(Readable)] +pub fn readable_macro_derive(input: TokenStream) -> TokenStream { + // Construct a representation of Rust code as a syntax tree + // that we can manipulate + let ast = syn::parse(input).unwrap(); + + // Build the trait implementation + impl_readable_macro(&ast) +} + +#[proc_macro_derive(Writeable)] +pub fn writeable_macro_derive(input: TokenStream) -> TokenStream { + let ast = syn::parse(input).unwrap(); + impl_writeable_macro(&ast) +} + +fn impl_readable_macro(ast: &syn::DeriveInput) -> TokenStream { + let name = &ast.ident; + let readable_name = format_ident!("Readable{}", name); + let archived_name = format_ident!("Archived{}", name); + let gen = quote! { + pub struct #readable_name<'a> { + _rb: ReadableBuffer<'a>, + } + + impl<'a> #readable_name<'a> { + pub fn get_archived_node(&self) -> & #archived_name { + // checking the code here is expensive during build, so skip it. + // TODO: should we check the data during queries? + //rkyv::check_archived_root::(self._rb.get_data_slice()).unwrap() + unsafe { rkyv::archived_root::<#name>(self._rb.get_data_slice()) } + } + } + + impl #name { + pub unsafe fn read<'a, 'b, S: StatsNodeRead>(index: &'a PgRelation, index_pointer: ItemPointer, stats: &'b mut S) -> #readable_name<'a> { + let rb = index_pointer.read_bytes(index); + stats.record_read(); + #readable_name { _rb: rb } + } + } + }; + gen.into() +} + +fn impl_writeable_macro(ast: &syn::DeriveInput) -> TokenStream { + let name = &ast.ident; + let writeable_name = format_ident!("Writable{}", name); + let archived_name = format_ident!("Archived{}", name); + let gen = quote! { + pub struct #writeable_name<'a> { + wb: WritableBuffer<'a>, + } + + impl #archived_name { + pub fn with_data(data: &mut [u8]) -> Pin<&mut #archived_name> { + let pinned_bytes = Pin::new(data); + unsafe { rkyv::archived_root_mut::<#name>(pinned_bytes) } + } + } + + impl<'a> #writeable_name<'a> { + pub fn get_archived_node(&self) -> Pin<&mut #archived_name> { + #archived_name::with_data(self.wb.get_data_slice()) + } + + pub fn commit(self) { + self.wb.commit() + } + } + + impl #name { + pub unsafe fn modify<'a, 'b, S: StatsNodeModify>(index: &'a PgRelation, index_pointer: ItemPointer, stats: &'b mut S) -> #writeable_name<'a> { + let wb = index_pointer.modify_bytes(index); + stats.record_modify(); + #writeable_name { wb: wb } + } + + pub fn write(&self, tape: &mut Tape, stats: &mut S) -> ItemPointer { + //TODO 256 probably too small + let bytes = rkyv::to_bytes::<_, 256>(self).unwrap(); + stats.record_write(); + unsafe { tape.write(&bytes) } + } + } + }; + gen.into() +} From ca385d5d55402c25d94ad7d6cbd46b56c299a4ab Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 23 Jan 2024 16:31:44 -0500 Subject: [PATCH 13/44] Optimize lsr --- timescale_vector/Cargo.toml | 4 + timescale_vector/benches/lsr.rs | 206 ++++++++++++++++++++ timescale_vector/src/access_method/graph.rs | 75 ++++--- 3 files changed, 262 insertions(+), 23 deletions(-) create mode 100644 timescale_vector/benches/lsr.rs diff --git a/timescale_vector/Cargo.toml b/timescale_vector/Cargo.toml index 9dc6952c..a2b9d622 100644 --- a/timescale_vector/Cargo.toml +++ b/timescale_vector/Cargo.toml @@ -45,3 +45,7 @@ codegen-units = 1 [[bench]] name = "distance" harness = false + +[[bench]] +name = "lsr" +harness = false diff --git a/timescale_vector/benches/lsr.rs b/timescale_vector/benches/lsr.rs new file mode 100644 index 00000000..53e31c9f --- /dev/null +++ b/timescale_vector/benches/lsr.rs @@ -0,0 +1,206 @@ +use std::{ + cmp::{Ordering, Reverse}, + collections::BinaryHeap, +}; + +use criterion::{black_box, criterion_group, criterion_main, Criterion}; +use rand::Rng; + +pub struct ListSearchNeighbor { + pub index_pointer: u64, + distance: f32, + visited: bool, + _private_data: u64, +} + +impl PartialOrd for ListSearchNeighbor { + fn partial_cmp(&self, other: &Self) -> Option { + self.distance.partial_cmp(&other.distance) + } +} + +impl PartialEq for ListSearchNeighbor { + fn eq(&self, other: &Self) -> bool { + self.index_pointer == other.index_pointer + } +} + +impl Eq for ListSearchNeighbor {} + +impl Ord for ListSearchNeighbor { + fn cmp(&self, other: &Self) -> Ordering { + self.distance.partial_cmp(&other.distance).unwrap() + } +} + +pub struct ListSearchResult { + candidate_storage: Vec, //plain storage + best_candidate: Vec, //pos in candidate storage, sorted by distance +} + +impl ListSearchResult { + pub fn get_lsn_by_idx(&self, idx: usize) -> &ListSearchNeighbor { + &self.candidate_storage[idx] + } + + pub fn insert_neighbor(&mut self, n: ListSearchNeighbor) { + //insert while preserving sort order. + let idx = self + .best_candidate + .partition_point(|x| self.candidate_storage[*x] < n); + self.candidate_storage.push(n); + let pos = self.candidate_storage.len() - 1; + self.best_candidate.insert(idx, pos) + } + + fn visit_closest(&mut self, pos_limit: usize) -> Option { + //OPT: should we optimize this not to do a linear search each time? + let neighbor_position = self + .best_candidate + .iter() + .position(|n| !self.candidate_storage[*n].visited); + match neighbor_position { + Some(pos) => { + if pos > pos_limit { + return None; + } + let n = &mut self.candidate_storage[self.best_candidate[pos]]; + n.visited = true; + Some(self.best_candidate[pos]) + } + None => None, + } + } +} + +pub struct ListSearchResultMinHeap { + candidates: BinaryHeap>, + visited: Vec, +} + +impl ListSearchResultMinHeap { + pub fn insert_neighbor(&mut self, n: ListSearchNeighbor) { + //insert while preserving sort order. + // self.candidate_storage.push(n); + // let pos = self.candidate_storage.len() - 1; + self.candidates.push(Reverse(n)); + + /*let idx = self + .best_candidate + .partition_point(|x| self.candidate_storage[*x].distance < n.distance); + self.candidate_storage.push(n); + let pos = self.candidate_storage.len() - 1; + self.best_candidate.insert(idx, pos)*/ + } + + fn visit_closest(&mut self, pos_limit: usize) -> Option<&ListSearchNeighbor> { + //OPT: should we optimize this not to do a linear search each time? + if self.candidates.len() == 0 { + panic!("no candidates left"); + //return None; + } + + if self.visited.len() > pos_limit { + let node_at_pos = &self.visited[pos_limit - 1]; + let head = self.candidates.peek().unwrap(); + if head.0.distance >= node_at_pos.distance { + return None; + } + } + + let head = self.candidates.pop().unwrap(); + let idx = self + .visited + .partition_point(|x| x.distance < head.0.distance); + self.visited.insert(idx, head.0); + Some(&self.visited[idx]) + } +} + +fn run_lsr_min_heap(lsr: &mut ListSearchResultMinHeap) { + let item = lsr.visit_closest(100000000); + let lsn = item.unwrap(); + + let mut rng = rand::thread_rng(); + let delta: f64 = rng.gen(); // generates a float between 0 and 1 + let distance = lsn.distance + ((delta * 5.0) as f32); + + for _ in 0..20 { + lsr.insert_neighbor(ListSearchNeighbor { + index_pointer: 0, + distance: distance, + visited: false, + _private_data: 2, + }) + } +} + +fn run_lsr(lsr: &mut ListSearchResult) { + let item_idx = lsr.visit_closest(1000000); + let lsn = lsr.get_lsn_by_idx(item_idx.unwrap()); + + let mut rng = rand::thread_rng(); + let delta: f64 = rng.gen(); // generates a float between 0 and 1 + let distance = lsn.distance + ((delta * 5.0) as f32); + + for _ in 0..20 { + lsr.insert_neighbor(ListSearchNeighbor { + index_pointer: 0, + distance: distance, + visited: false, + _private_data: 2, + }) + } +} + +pub fn benchmark_lsr(c: &mut Criterion) { + let mut lsr = ListSearchResult { + candidate_storage: Vec::new(), + best_candidate: Vec::new(), + }; + + lsr.insert_neighbor(ListSearchNeighbor { + index_pointer: 0, + distance: 100.0, + visited: false, + _private_data: 1, + }); + + c.bench_function("lsr OG", |b| b.iter(|| run_lsr(black_box(&mut lsr)))); +} + +pub fn benchmark_lsr_min_heap(c: &mut Criterion) { + let mut lsr = ListSearchResultMinHeap { + candidates: BinaryHeap::new(), + visited: Vec::new(), + }; + + lsr.insert_neighbor(ListSearchNeighbor { + index_pointer: 0, + distance: 100.0, + visited: false, + _private_data: 1, + }); + + c.bench_function("lsr min heap", |b| { + b.iter(|| run_lsr_min_heap(black_box(&mut lsr))) + }); +} + +criterion_group!(benches_lsr, benchmark_lsr, benchmark_lsr_min_heap); + +criterion_main!(benches_lsr); +/* +fn fibonacci(n: u64) -> u64 { + match n { + 0 => 1, + 1 => 1, + n => fibonacci(n - 1) + fibonacci(n - 2), + } +} +pub fn criterion_benchmark(c: &mut Criterion) { + c.bench_function("fib 20", |b| b.iter(|| fibonacci(black_box(20)))); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches);*/ diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index 5c318109..a273e311 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -1,3 +1,5 @@ +use std::cmp::Reverse; +use std::collections::BinaryHeap; use std::{cmp::Ordering, collections::HashSet}; use pgrx::PgRelation; @@ -32,6 +34,14 @@ impl PartialEq for ListSearchNeighbor { } } +impl Eq for ListSearchNeighbor {} + +impl Ord for ListSearchNeighbor { + fn cmp(&self, other: &Self) -> Ordering { + self.distance.partial_cmp(&other.distance).unwrap() + } +} + impl ListSearchNeighbor { pub fn new(index_pointer: IndexPointer, distance: f32, private_data: PD) -> Self { assert!(!distance.is_nan()); @@ -50,8 +60,10 @@ impl ListSearchNeighbor { } pub struct ListSearchResult { - candidate_storage: Vec>, //plain storage - best_candidate: Vec, //pos in candidate storage, sorted by distance + candidates: BinaryHeap>>, + visited: Vec>, + //candidate_storage: Vec>, //plain storage + //best_candidate: Vec, //pos in candidate storage, sorted by distance inserted: HashSet, max_history_size: Option, pub sdm: Option, @@ -61,8 +73,8 @@ pub struct ListSearchResult { impl ListSearchResult { fn empty() -> Self { Self { - candidate_storage: vec![], - best_candidate: vec![], + candidates: BinaryHeap::new(), + visited: vec![], inserted: HashSet::new(), max_history_size: None, sdm: None, @@ -81,8 +93,10 @@ impl ListSearchResult { ) -> Self { let neigbors = meta_page.get_num_neighbors() as usize; let mut res = Self { - candidate_storage: Vec::with_capacity(search_list_size * neigbors), - best_candidate: Vec::with_capacity(search_list_size * neigbors), + candidates: BinaryHeap::with_capacity(search_list_size * neigbors), + visited: Vec::with_capacity(search_list_size * 2), + //candidate_storage: Vec::with_capacity(search_list_size * neigbors), + //best_candidate: Vec::with_capacity(search_list_size * neigbors), inserted: HashSet::with_capacity(search_list_size * neigbors), max_history_size, stats: GreedySearchStats::new(), @@ -102,7 +116,7 @@ impl ListSearchResult { /// Internal function pub fn insert_neighbor(&mut self, n: ListSearchNeighbor) { - if let Some(max_size) = self.max_history_size { + /*if let Some(max_size) = self.max_history_size { if self.best_candidate.len() >= max_size { let last = self.best_candidate.last().unwrap(); if n >= self.candidate_storage[*last] { @@ -111,23 +125,37 @@ impl ListSearchResult { } self.best_candidate.pop(); } - } - //insert while preserving sort order. - let idx = self - .best_candidate - .partition_point(|x| self.candidate_storage[*x] < n); - self.candidate_storage.push(n); - let pos = self.candidate_storage.len() - 1; - self.best_candidate.insert(idx, pos) + }*/ + + self.candidates.push(Reverse(n)); } pub fn get_lsn_by_idx(&self, idx: usize) -> &ListSearchNeighbor { - &self.candidate_storage[idx] + &self.visited[idx] } fn visit_closest(&mut self, pos_limit: usize) -> Option { + if self.candidates.len() == 0 { + return None; + } + + if self.visited.len() > pos_limit { + let node_at_pos = &self.visited[pos_limit - 1]; + let head = self.candidates.peek().unwrap(); + if head.0.distance >= node_at_pos.distance { + return None; + } + } + + let head = self.candidates.pop().unwrap(); + let idx = self + .visited + .partition_point(|x| x.distance < head.0.distance); + self.visited.insert(idx, head.0); + Some(idx) + //OPT: should we optimize this not to do a linear search each time? - let neighbor_position = self + /*let neighbor_position = self .best_candidate .iter() .position(|n| !self.candidate_storage[*n].visited); @@ -141,7 +169,7 @@ impl ListSearchResult { Some(self.best_candidate[pos]) } None => None, - } + }*/ } //removes and returns the first element. Given that the element remains in self.inserted, that means the element will never again be insereted @@ -150,12 +178,13 @@ impl ListSearchResult { &mut self, storage: &S, ) -> Option<(HeapPointer, IndexPointer)> { - if self.best_candidate.is_empty() { + if self.visited.len() == 0 { return None; } - let idx = self.best_candidate.remove(0); - let lsn = &self.candidate_storage[idx]; - let heap_pointer = storage.return_lsn(lsn, &mut self.stats); + let lsn = self.visited.remove(0); + //let idx = self.best_candidate.remove(0); + //let lsn = &self.candidate_storage[idx]; + let heap_pointer = storage.return_lsn(&lsn, &mut self.stats); return Some((heap_pointer, lsn.index_pointer)); } } @@ -322,7 +351,7 @@ impl<'a> Graph<'a> { match visited_nodes { None => {} Some(ref mut visited_nodes) => { - let list_search_entry = &lsr.candidate_storage[list_search_entry_idx]; + let list_search_entry = &lsr.visited[list_search_entry_idx]; visited_nodes.insert(NeighborWithDistance::new( list_search_entry.index_pointer, list_search_entry.distance, From db9ad1dffb03e00b1364197b3b8b7af8840b9025 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 23 Jan 2024 17:41:01 -0500 Subject: [PATCH 14/44] add first xor benchmarks --- timescale_vector/benches/distance.rs | 145 ++++++++++++++++++++++++++- 1 file changed, 143 insertions(+), 2 deletions(-) diff --git a/timescale_vector/benches/distance.rs b/timescale_vector/benches/distance.rs index 836c4b0b..2cf4c7ed 100644 --- a/timescale_vector/benches/distance.rs +++ b/timescale_vector/benches/distance.rs @@ -188,19 +188,160 @@ fn benchmark_distance_few_dimensions(c: &mut Criterion) { ); } +fn pack_bools_to_u8(bools: Vec) -> Vec { + let mut bytes = vec![0u8; (bools.len() + 7) / 8]; + + for (i, &b) in bools.iter().enumerate() { + let byte_index = i / 8; + let bit_index = i % 8; + + if b { + bytes[byte_index] |= 1 << bit_index; + } + } + + bytes +} + +fn pack_bools_to_u64(bools: Vec) -> Vec { + let mut u64s = vec![0u64; (bools.len() + 63) / 64]; + + for (i, &b) in bools.iter().enumerate() { + let u64_index = i / 64; + let bit_index = i % 64; + + if b { + u64s[u64_index] |= 1 << bit_index; + } + } + + u64s +} + +fn pack_bools_to_u128(bools: Vec) -> Vec { + let mut u128s = vec![0u128; (bools.len() + 127) / 128]; + + for (i, &b) in bools.iter().enumerate() { + let u128_index = i / 128; + let bit_index = i % 128; + + if b { + u128s[u128_index] |= 1 << bit_index; + } + } + + u128s +} + +fn xor_unoptimized_u8(v1: &[u8], v2: &[u8]) -> usize { + let mut result = 0; + for (b1, b2) in v1.iter().zip(v2.iter()) { + result += (b1 ^ b2).count_ones() as usize; + } + result +} + +fn xor_unoptimized_u8_fixed_size(v1: &[u8], v2: &[u8]) -> usize { + let mut result = 0; + for (b1, b2) in v1[..192].iter().zip(v2[..192].iter()) { + result += (b1 ^ b2).count_ones() as usize; + } + result +} + +fn xor_unoptimized_u64(v1: &[u64], v2: &[u64]) -> usize { + let mut result = 0; + for (b1, b2) in v1.iter().zip(v2.iter()) { + result += (b1 ^ b2).count_ones() as usize; + } + result +} + +fn xor_unoptimized_u64_fixed_size(v1: &[u64], v2: &[u64]) -> usize { + let mut result = 0; + for (b1, b2) in v1[..24].iter().zip(v2[..24].iter()) { + result += (b1 ^ b2).count_ones() as usize; + } + result +} + +fn xor_unoptimized_u64_fixed_size_map(v1: &[u64], v2: &[u64]) -> usize { + v1[..24] + .iter() + .zip(v2[..24].iter()) + .map(|(&l, &r)| (l ^ r).count_ones() as usize) + .sum() +} + +fn xor_unoptimized_u128(v1: &[u128], v2: &[u128]) -> usize { + let mut result = 0; + for (b1, b2) in v1.iter().zip(v2.iter()) { + result += (b1 ^ b2).count_ones() as usize; + } + result +} + +fn xor_unoptimized_u128_fixed_size(v1: &[u128], v2: &[u128]) -> usize { + let mut result = 0; + for (b1, b2) in v1[..12].iter().zip(v2[..12].iter()) { + result += (b1 ^ b2).count_ones() as usize; + } + result +} + +fn benchmark_distance_xor(c: &mut Criterion) { + let r: Vec = (0..1536).map(|v| v as u64 % 2 == 0).collect(); + let l: Vec = (0..1536).map(|v| v as u64 % 3 == 0).collect(); + let r_u8 = pack_bools_to_u8(r.clone()); + let l_u8 = pack_bools_to_u8(l.clone()); + let r_u64 = pack_bools_to_u64(r.clone()); + let l_u64 = pack_bools_to_u64(l.clone()); + let r_u128 = pack_bools_to_u128(r.clone()); + let l_u128 = pack_bools_to_u128(l.clone()); + + let mut group = c.benchmark_group("Distance xor"); + group.bench_function("xor unoptimized u8", |b| { + b.iter(|| xor_unoptimized_u8(black_box(&r_u8), black_box(&l_u8))) + }); + group.bench_function("xor unoptimized u64", |b| { + b.iter(|| xor_unoptimized_u64(black_box(&r_u64), black_box(&l_u64))) + }); + group.bench_function("xor unoptimized u128", |b| { + b.iter(|| xor_unoptimized_u128(black_box(&r_u128), black_box(&l_u128))) + }); + + assert!(r_u8.len() == 192); + group.bench_function("xor unoptimized u8 fixed size", |b| { + b.iter(|| xor_unoptimized_u8_fixed_size(black_box(&r_u8), black_box(&l_u8))) + }); + assert!(r_u64.len() == 24); + group.bench_function("xor unoptimized u64 fixed size", |b| { + b.iter(|| xor_unoptimized_u64_fixed_size(black_box(&r_u64), black_box(&l_u64))) + }); + group.bench_function("xor unoptimized u64 fixed size_map", |b| { + b.iter(|| xor_unoptimized_u64_fixed_size_map(black_box(&r_u64), black_box(&l_u64))) + }); + assert!(r_u128.len() == 12); + group.bench_function("xor unoptimized u128 fixed size", |b| { + b.iter(|| xor_unoptimized_u128_fixed_size(black_box(&r_u128), black_box(&l_u128))) + }); +} + #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] criterion_group!( benches, benchmark_distance, benchmark_distance_few_dimensions, benchmark_distance_x86_unaligned_vectors, - benchmark_distance_x86_aligned_vectors + benchmark_distance_x86_aligned_vectors, + benchmark_distance_xor, ); #[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))] criterion_group!( benches, benchmark_distance, - benchmark_distance_few_dimensions + benchmark_distance_few_dimensions, + benchmark_distance_xor, ); criterion_main!(benches); From 3af296c2a24998dc07bc22aa6ceb6f3b3dc582fe Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 23 Jan 2024 21:27:04 -0500 Subject: [PATCH 15/44] write optimized xor func --- timescale_vector/benches/distance.rs | 6 ++- .../src/access_method/distance.rs | 51 +++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/timescale_vector/benches/distance.rs b/timescale_vector/benches/distance.rs index 2cf4c7ed..e823df94 100644 --- a/timescale_vector/benches/distance.rs +++ b/timescale_vector/benches/distance.rs @@ -1,6 +1,7 @@ use criterion::{black_box, criterion_group, criterion_main, Criterion}; use timescale_vector::access_method::distance::{ - distance_cosine, distance_l2, distance_l2_optimized_for_few_dimensions, distance_l2_unoptimized, + distance_cosine, distance_l2, distance_l2_optimized_for_few_dimensions, + distance_l2_unoptimized, distance_xor_optimized, }; //copy and use qdrants simd code, purely for benchmarking purposes @@ -321,6 +322,9 @@ fn benchmark_distance_xor(c: &mut Criterion) { group.bench_function("xor unoptimized u64 fixed size_map", |b| { b.iter(|| xor_unoptimized_u64_fixed_size_map(black_box(&r_u64), black_box(&l_u64))) }); + group.bench_function("xor optimized version we use in code", |b| { + b.iter(|| distance_xor_optimized(black_box(&r_u64), black_box(&l_u64))) + }); assert!(r_u128.len() == 12); group.bench_function("xor unoptimized u128 fixed size", |b| { b.iter(|| xor_unoptimized_u128_fixed_size(black_box(&r_u128), black_box(&l_u128))) diff --git a/timescale_vector/src/access_method/distance.rs b/timescale_vector/src/access_method/distance.rs index 21396761..f8a053ec 100644 --- a/timescale_vector/src/access_method/distance.rs +++ b/timescale_vector/src/access_method/distance.rs @@ -138,3 +138,54 @@ pub fn preprocess_cosine(a: &mut [f32]) { } } } + +macro_rules! xor_arm { + ($a: expr, $b: expr, $sz: expr) => { + $a[..$sz] + .iter() + .zip($b[..$sz].iter()) + .map(|(&l, &r)| (l ^ r).count_ones() as usize) + .sum() + }; +} + +#[inline] +pub fn distance_xor_optimized(a: &[u64], b: &[u64]) -> usize { + match a.len() { + 1 => xor_arm!(a, b, 1), + 2 => xor_arm!(a, b, 2), + 3 => xor_arm!(a, b, 3), + 4 => xor_arm!(a, b, 4), + 5 => xor_arm!(a, b, 5), + 6 => xor_arm!(a, b, 6), + 7 => xor_arm!(a, b, 7), + 8 => xor_arm!(a, b, 8), + 9 => xor_arm!(a, b, 9), + 10 => xor_arm!(a, b, 10), + 11 => xor_arm!(a, b, 11), + 12 => xor_arm!(a, b, 12), + 13 => xor_arm!(a, b, 13), + 14 => xor_arm!(a, b, 14), + 15 => xor_arm!(a, b, 15), + 16 => xor_arm!(a, b, 16), + 17 => xor_arm!(a, b, 17), + 18 => xor_arm!(a, b, 18), + 19 => xor_arm!(a, b, 19), + 20 => xor_arm!(a, b, 20), + 21 => xor_arm!(a, b, 21), + 22 => xor_arm!(a, b, 22), + 23 => xor_arm!(a, b, 23), + 24 => xor_arm!(a, b, 24), + 25 => xor_arm!(a, b, 25), + 26 => xor_arm!(a, b, 26), + 27 => xor_arm!(a, b, 27), + 28 => xor_arm!(a, b, 28), + 29 => xor_arm!(a, b, 29), + 30 => xor_arm!(a, b, 30), + _ => a + .iter() + .zip(b.iter()) + .map(|(&l, &r)| (l ^ r).count_ones() as usize) + .sum(), + } +} From f9cb57774430e7007576728fef212f64dbeeb925 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 23 Jan 2024 21:40:19 -0500 Subject: [PATCH 16/44] Change bq to use u64 --- timescale_vector/src/access_method/bq.rs | 14 +++++++------- timescale_vector/src/access_method/distance.rs | 2 +- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index ec652cca..48b03d2a 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -1,5 +1,5 @@ use super::{ - distance::distance_cosine as default_distance, + distance::{distance_cosine as default_distance, distance_xor_optimized}, graph::{ListSearchNeighbor, ListSearchResult}, graph_neighbor_store::GraphNeighborStore, pg_vector::PgVector, @@ -26,8 +26,8 @@ use crate::util::{ use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance}; use crate::util::WritableBuffer; -type BqVectorElement = u8; -const BITS_STORE_TYPE_SIZE: usize = 8; +type BqVectorElement = u64; +const BITS_STORE_TYPE_SIZE: usize = 64; #[derive(Archive, Deserialize, Serialize, Readable, Writeable)] #[archive(check_bytes)] @@ -199,7 +199,7 @@ impl BqDistanceTable { /// distance emits the sum of distances between each centroid in the quantized vector. pub fn distance(&self, bq_vector: &[BqVectorElement]) -> f32 { - let count_ones = xor_unoptimized(&self.quantized_vector, bq_vector); + let count_ones = distance_xor_optimized(&self.quantized_vector, bq_vector); //dot product is LOWER the more xors that lead to 1 becaues that means a negative times a positive = negative component //but the distance is 1 - dot product, so the more count_ones the higher the distance. // one other check for distance(a,a), xor=0, count_ones=0, distance=0 @@ -651,9 +651,9 @@ use timescale_vector_derive::{Readable, Writeable}; #[archive(check_bytes)] pub struct BqNode { pub heap_item_pointer: HeapPointer, - pub bq_vector: Vec, + pub bq_vector: Vec, //don't use BqVectorElement because we don't want to change the size in on-disk format by accident neighbor_index_pointers: Vec, - neighbor_vectors: Vec>, + neighbor_vectors: Vec>, //don't use BqVectorElement because we don't want to change the size in on-disk format by accident } impl BqNode { @@ -688,7 +688,7 @@ impl ArchivedBqNode { unsafe { self.map_unchecked_mut(|s| &mut s.neighbor_index_pointers) } } - pub fn neighbor_vector(self: Pin<&mut Self>) -> Pin<&mut ArchivedVec>> { + pub fn neighbor_vector(self: Pin<&mut Self>) -> Pin<&mut ArchivedVec>> { unsafe { self.map_unchecked_mut(|s| &mut s.neighbor_vectors) } } diff --git a/timescale_vector/src/access_method/distance.rs b/timescale_vector/src/access_method/distance.rs index f8a053ec..ad552f2f 100644 --- a/timescale_vector/src/access_method/distance.rs +++ b/timescale_vector/src/access_method/distance.rs @@ -149,7 +149,7 @@ macro_rules! xor_arm { }; } -#[inline] +#[inline(always)] pub fn distance_xor_optimized(a: &[u64], b: &[u64]) -> usize { match a.len() { 1 => xor_arm!(a, b, 1), From df4ca43f532a8b9bee7d76d65a012899e1a9365d Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 23 Jan 2024 22:02:59 -0500 Subject: [PATCH 17/44] cleanup --- timescale_vector/src/access_method/bq.rs | 8 ----- timescale_vector/src/access_method/graph.rs | 40 --------------------- 2 files changed, 48 deletions(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index 48b03d2a..4a94a7b0 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -182,14 +182,6 @@ pub struct BqDistanceTable { quantized_vector: Vec, } -fn xor_unoptimized(v1: &[BqVectorElement], v2: &[BqVectorElement]) -> usize { - let mut result = 0; - for (b1, b2) in v1.iter().zip(v2.iter()) { - result += (b1 ^ b2).count_ones() as usize; - } - result -} - impl BqDistanceTable { pub fn new(query: Vec) -> BqDistanceTable { BqDistanceTable { diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index a273e311..f420f60f 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -18,7 +18,6 @@ use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance}; pub struct ListSearchNeighbor { pub index_pointer: IndexPointer, distance: f32, - visited: bool, private_data: PD, } @@ -50,7 +49,6 @@ impl ListSearchNeighbor { index_pointer, private_data, distance, - visited: false, } } @@ -62,10 +60,7 @@ impl ListSearchNeighbor { pub struct ListSearchResult { candidates: BinaryHeap>>, visited: Vec>, - //candidate_storage: Vec>, //plain storage - //best_candidate: Vec, //pos in candidate storage, sorted by distance inserted: HashSet, - max_history_size: Option, pub sdm: Option, pub stats: GreedySearchStats, } @@ -76,14 +71,12 @@ impl ListSearchResult { candidates: BinaryHeap::new(), visited: vec![], inserted: HashSet::new(), - max_history_size: None, sdm: None, stats: GreedySearchStats::new(), } } fn new>( - max_history_size: Option, init_ids: Vec, sdm: S::QueryDistanceMeasure, search_list_size: usize, @@ -98,7 +91,6 @@ impl ListSearchResult { //candidate_storage: Vec::with_capacity(search_list_size * neigbors), //best_candidate: Vec::with_capacity(search_list_size * neigbors), inserted: HashSet::with_capacity(search_list_size * neigbors), - max_history_size, stats: GreedySearchStats::new(), sdm: Some(sdm), }; @@ -116,17 +108,6 @@ impl ListSearchResult { /// Internal function pub fn insert_neighbor(&mut self, n: ListSearchNeighbor) { - /*if let Some(max_size) = self.max_history_size { - if self.best_candidate.len() >= max_size { - let last = self.best_candidate.last().unwrap(); - if n >= self.candidate_storage[*last] { - //n is too far in the list to be the best candidate. - return; - } - self.best_candidate.pop(); - } - }*/ - self.candidates.push(Reverse(n)); } @@ -153,23 +134,6 @@ impl ListSearchResult { .partition_point(|x| x.distance < head.0.distance); self.visited.insert(idx, head.0); Some(idx) - - //OPT: should we optimize this not to do a linear search each time? - /*let neighbor_position = self - .best_candidate - .iter() - .position(|n| !self.candidate_storage[*n].visited); - match neighbor_position { - Some(pos) => { - if pos > pos_limit { - return None; - } - let n = &mut self.candidate_storage[self.best_candidate[pos]]; - n.visited = true; - Some(self.best_candidate[pos]) - } - None => None, - }*/ } //removes and returns the first element. Given that the element remains in self.inserted, that means the element will never again be insereted @@ -182,8 +146,6 @@ impl ListSearchResult { return None; } let lsn = self.visited.remove(0); - //let idx = self.best_candidate.remove(0); - //let lsn = &self.candidate_storage[idx]; let heap_pointer = storage.return_lsn(&lsn, &mut self.stats); return Some((heap_pointer, lsn.index_pointer)); } @@ -299,7 +261,6 @@ impl<'a> Graph<'a> { let search_list_size = meta_page.get_search_list_size_for_build() as usize; let mut l = ListSearchResult::new( - Some(search_list_size), init_ids.unwrap(), dm, search_list_size, @@ -329,7 +290,6 @@ impl<'a> Graph<'a> { let dm = storage.get_search_distance_measure(query, true); ListSearchResult::new( - None, init_ids.unwrap(), dm, search_list_size, From 48da6db8636340c19ee30e00e791339dc02f2303 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Mon, 12 Feb 2024 19:38:59 -0500 Subject: [PATCH 18/44] Building in quantized distances instead of full distances --- timescale_vector/src/access_method/bq.rs | 195 +++++++++--------- timescale_vector/src/access_method/build.rs | 107 ++++++++-- timescale_vector/src/access_method/graph.rs | 8 +- .../src/access_method/graph_neighbor_store.rs | 8 +- .../src/access_method/plain_node.rs | 7 - .../src/access_method/plain_storage.rs | 22 +- .../src/access_method/pq_quantizer.rs | 82 ++++++-- .../src/access_method/pq_storage.rs | 158 +++++++------- timescale_vector/src/access_method/stats.rs | 2 +- timescale_vector/src/access_method/storage.rs | 21 +- .../src/access_method/storage_common.rs | 22 +- timescale_vector/src/access_method/vacuum.rs | 49 +++-- 12 files changed, 400 insertions(+), 281 deletions(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index 4a94a7b0..a8270531 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -7,10 +7,9 @@ use super::{ GreedySearchStats, StatsDistanceComparison, StatsNodeModify, StatsNodeRead, StatsNodeWrite, WriteStats, }, - storage::{ArchivedData, NodeFullDistanceMeasure, Storage, StorageFullDistanceFromHeap}, - storage_common::{calculate_full_distance, HeapFullDistanceMeasure}, + storage::{ArchivedData, NodeDistanceMeasure, Storage, StorageFullDistanceFromHeap}, }; -use std::{collections::HashMap, iter::once, marker::PhantomData, pin::Pin}; +use std::{cell::RefCell, collections::HashMap, iter::once, marker::PhantomData, pin::Pin}; use pgrx::{ pg_sys::{InvalidBlockNumber, InvalidOffsetNumber}, @@ -105,6 +104,7 @@ impl BqQuantizer { } fn quantize(&self, full_vector: &[f32]) -> Vec { + assert!(!self.training); if self.use_mean { let mut res_vector = vec![0; Self::quantized_size(full_vector.len())]; @@ -154,18 +154,10 @@ impl BqQuantizer { fn vector_for_new_node( &self, - meta_page: &super::meta_page::MetaPage, + _meta_page: &super::meta_page::MetaPage, full_vector: &[f32], ) -> Vec { - if self.use_mean && self.training { - vec![0; BqQuantizer::quantized_size(meta_page.get_num_dimensions() as _)] - } else { - self.quantize(&full_vector) - } - } - - fn vector_needs_update_after_training(&self) -> bool { - self.use_mean + self.quantize(&full_vector) } fn get_distance_table( @@ -199,8 +191,8 @@ impl BqDistanceTable { } } +//FIXME: cleanup make this into a struct pub enum BqSearchDistanceMeasure { - Full(PgVector), Bq(BqDistanceTable), } @@ -217,6 +209,44 @@ impl BqSearchDistanceMeasure { } } +pub struct BqNodeDistanceMeasure<'a> { + readable_node: ReadableBqNode<'a>, + storage: &'a BqSpeedupStorage<'a>, +} + +impl<'a> BqNodeDistanceMeasure<'a> { + pub unsafe fn with_index_pointer( + storage: &'a BqSpeedupStorage<'a>, + index_pointer: IndexPointer, + stats: &mut T, + ) -> Self { + let rn = unsafe { BqNode::read(storage.index, index_pointer, stats) }; + Self { + readable_node: rn, + storage: storage, + } + } +} + +impl<'a> NodeDistanceMeasure for BqNodeDistanceMeasure<'a> { + unsafe fn get_distance( + &self, + index_pointer: IndexPointer, + stats: &mut T, + ) -> f32 { + //OPT: should I get and memoize the vector from self.readable_node in with_index_pointer above? + let rn1 = BqNode::read(self.storage.index, index_pointer, stats); + let rn2 = &self.readable_node; + let node1 = rn1.get_archived_node(); + let node2 = rn2.get_archived_node(); + assert!(node1.bq_vector.len() > 0); + assert!(node1.bq_vector.len() == node2.bq_vector.len()); + let vec1 = node1.bq_vector.as_slice(); + let vec2 = node2.bq_vector.as_slice(); + distance_xor_optimized(vec1, vec2) as f32 + } +} + struct QuantizedVectorCache { quantized_vector_map: HashMap>, } @@ -269,7 +299,7 @@ pub struct BqSpeedupStorage<'a> { quantizer: BqQuantizer, heap_rel: Option<&'a PgRelation>, heap_attr: Option, - qv_cache: Option, + qv_cache: RefCell, } impl<'a> BqSpeedupStorage<'a> { @@ -284,7 +314,7 @@ impl<'a> BqSpeedupStorage<'a> { quantizer: BqQuantizer::new(), heap_rel: Some(heap_rel), heap_attr: Some(heap_attr), - qv_cache: None, + qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), } } @@ -309,7 +339,7 @@ impl<'a> BqSpeedupStorage<'a> { quantizer: Self::load_quantizer(index_relation, meta_page, stats), heap_rel: Some(heap_rel), heap_attr: Some(heap_attr), - qv_cache: None, + qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), } } @@ -324,7 +354,7 @@ impl<'a> BqSpeedupStorage<'a> { quantizer: quantizer.clone(), heap_rel: None, heap_attr: None, - qv_cache: None, + qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), } } @@ -333,9 +363,9 @@ impl<'a> BqSpeedupStorage<'a> { index_pointer: IndexPointer, stats: &mut S, ) -> Vec { - let slot = unsafe { self.get_heap_table_slot_from_index_pointer(index_pointer, stats) }; - let slice = unsafe { slot.get_pg_vector() }; - self.quantizer.quantize(slice.to_slice()) + let rn = unsafe { BqNode::read(self.index, index_pointer, stats) }; + let node = rn.get_archived_node(); + node.bq_vector.as_slice().to_vec() } fn write_quantizer_metadata(&self, stats: &mut S) { @@ -369,36 +399,31 @@ impl<'a> BqSpeedupStorage<'a> { } let distance = match lsr.sdm.as_ref().unwrap() { - BqSearchDistanceMeasure::Full(query) => { - let rn_neighbor = - unsafe { BqNode::read(self.index, neighbor_index_pointer, &mut lsr.stats) }; - let node_neighbor = rn_neighbor.get_archived_node(); - if node_neighbor.is_deleted() { - self.visit_lsn_internal(lsr, neighbor_index_pointer, gns); - continue; - } - let heap_pointer_neighbor = - node_neighbor.heap_item_pointer.deserialize_item_pointer(); - unsafe { - calculate_full_distance( - self, - heap_pointer_neighbor, - query.to_slice(), - &mut lsr.stats, - ) - } - } BqSearchDistanceMeasure::Bq(table) => { /* Note: there is no additional node reads here. We get all of our info from node_visiting * This is what gives us a speedup in BQ Speedup */ - if let GraphNeighborStore::Builder(_) = gns { - assert!( - false, - "BQ distance should not be used with the builder graph store" - ) + match gns { + GraphNeighborStore::Disk => { + let bq_vector = node_visiting.neighbor_vectors[i].as_slice(); + BqSearchDistanceMeasure::calculate_bq_distance( + table, + bq_vector, + &mut lsr.stats, + ) + } + GraphNeighborStore::Builder(b) => { + let mut cache = self.qv_cache.borrow_mut(); + let bq_vector = cache.get(neighbor_index_pointer, self, &mut lsr.stats); + let dist = BqSearchDistanceMeasure::calculate_bq_distance( + table, + bq_vector, + &mut lsr.stats, + ); + dist + } } - let bq_vector = node_visiting.neighbor_vectors[i].as_slice(); - BqSearchDistanceMeasure::calculate_bq_distance(table, bq_vector, &mut lsr.stats) + //let bq_vector = node_visiting.neighbor_vectors[i].as_slice(); + //BqSearchDistanceMeasure::calculate_bq_distance(table, bq_vector, &mut lsr.stats) } }; let lsn = @@ -413,7 +438,7 @@ pub type BqSpeedupStorageLsnPrivateData = PhantomData; //no data stored impl<'a> Storage for BqSpeedupStorage<'a> { type QueryDistanceMeasure = BqSearchDistanceMeasure; - type NodeFullDistanceMeasure<'b> = HeapFullDistanceMeasure<'b, BqSpeedupStorage<'b>> where Self: 'b; + type NodeDistanceMeasure<'b> = BqNodeDistanceMeasure<'b> where Self: 'b; type ArchivedType = ArchivedBqNode; type LSNPrivateData = BqSpeedupStorageLsnPrivateData; //no data stored @@ -448,7 +473,6 @@ impl<'a> Storage for BqSpeedupStorage<'a> { fn finish_training(&mut self, stats: &mut WriteStats) { self.quantizer.finish_training(); self.write_quantizer_metadata(stats); - self.qv_cache = Some(QuantizedVectorCache::new(1000)); } fn finalize_node_at_end_of_build( @@ -458,7 +482,7 @@ impl<'a> Storage for BqSpeedupStorage<'a> { neighbors: &Vec, stats: &mut S, ) { - let mut cache = self.qv_cache.take().unwrap(); + let mut cache = self.qv_cache.borrow_mut(); /* It's important to preload cache with all the items since you can run into deadlocks if you try to fetch a quantized vector while holding the BqNode::modify lock */ let iter = neighbors @@ -471,56 +495,41 @@ impl<'a> Storage for BqSpeedupStorage<'a> { let mut archived = node.get_archived_node(); archived.as_mut().set_neighbors(neighbors, &meta, &cache); - if self.quantizer.vector_needs_update_after_training() { - let bq_vector = cache.must_get(index_pointer); - archived.as_mut().set_bq_vector(bq_vector); - } - node.commit(); - self.qv_cache = Some(cache); } - unsafe fn get_full_vector_distance_state<'b, S: StatsNodeRead>( + unsafe fn get_node_distance_measure<'b, S: StatsNodeRead>( &'b self, index_pointer: IndexPointer, stats: &mut S, - ) -> HeapFullDistanceMeasure<'b, BqSpeedupStorage<'b>> { - HeapFullDistanceMeasure::with_index_pointer(self, index_pointer, stats) + ) -> BqNodeDistanceMeasure<'b> { + BqNodeDistanceMeasure::with_index_pointer(self, index_pointer, stats) } - fn get_search_distance_measure( - &self, - query: PgVector, - calc_distance_with_quantizer: bool, - ) -> BqSearchDistanceMeasure { - if !calc_distance_with_quantizer { - return BqSearchDistanceMeasure::Full(query); - } else { - return BqSearchDistanceMeasure::Bq( - self.quantizer - .get_distance_table(query.to_slice(), self.distance_fn), - ); - } + fn get_query_distance_measure(&self, query: PgVector) -> BqSearchDistanceMeasure { + return BqSearchDistanceMeasure::Bq( + self.quantizer + .get_distance_table(query.to_slice(), self.distance_fn), + ); } - fn get_neighbors_with_full_vector_distances_from_disk< - S: StatsNodeRead + StatsDistanceComparison, - >( + fn get_neighbors_with_distances_from_disk( &self, neighbors_of: ItemPointer, result: &mut Vec, stats: &mut S, ) { let rn = unsafe { BqNode::read(self.index, neighbors_of, stats) }; - let heap_pointer = rn - .get_archived_node() - .heap_item_pointer - .deserialize_item_pointer(); - let dist_state = - unsafe { HeapFullDistanceMeasure::with_heap_pointer(self, heap_pointer, stats) }; - for n in rn.get_archived_node().iter_neighbors() { - let dist = unsafe { dist_state.get_distance(n, stats) }; - result.push(NeighborWithDistance::new(n, dist)) + let archived = rn.get_archived_node(); + let q = archived.bq_vector.as_slice(); + + for (i, n) in rn.get_archived_node().iter_neighbors().enumerate() { + //let dist = unsafe { dist_state.get_distance(n, stats) }; + assert!(i < archived.neighbor_vectors.len()); + let neighbor_q = archived.neighbor_vectors[i].as_slice(); + stats.record_quantized_distance_comparison(); + let dist = distance_xor_optimized(q, neighbor_q); + result.push(NeighborWithDistance::new(n, dist as f32)) } } @@ -540,16 +549,6 @@ impl<'a> Storage for BqSpeedupStorage<'a> { let node = rn.get_archived_node(); let distance = match lsr.sdm.as_ref().unwrap() { - BqSearchDistanceMeasure::Full(query) => { - if node.is_deleted() { - //FIXME: handle deleted init ids - panic!("don't handle deleted init ids yet"); - } - let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); - unsafe { - calculate_full_distance(self, heap_pointer, query.to_slice(), &mut lsr.stats) - } - } BqSearchDistanceMeasure::Bq(table) => BqSearchDistanceMeasure::calculate_bq_distance( table, node.bq_vector.as_slice(), @@ -688,14 +687,6 @@ impl ArchivedBqNode { unsafe { self.map_unchecked_mut(|s| &mut s.bq_vector) } } - fn set_bq_vector(mut self: Pin<&mut Self>, bq_vector: &[BqVectorElement]) { - assert!(bq_vector.len() == self.bq_vector.len()); - for i in 0..=bq_vector.len() - 1 { - let mut pgv = self.as_mut().bq_vector().index_pin(i); - *pgv = bq_vector[i]; - } - } - fn set_neighbors( mut self: Pin<&mut Self>, neighbors: &[NeighborWithDistance], diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index 898a5e7d..dedd1337 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -218,6 +218,7 @@ fn do_heap_scan<'a>( GraphNeighborStore::Builder(BuilderNeighborCache::new()), &mut mp2, ); + let mut write_stats = WriteStats::new(); match storage { StorageType::Plain => { let mut plain = PlainStorage::new_for_build(index_relation); @@ -236,7 +237,7 @@ fn do_heap_scan<'a>( ); } - do_heap_scan_with_state(&mut plain, &mut bs) + finalize_index_build(&mut plain, &mut bs, write_stats) } StorageType::PqCompression => { let mut pq = PqCompressionStorage::new_for_build( @@ -245,6 +246,17 @@ fn do_heap_scan<'a>( get_attribute_number(index_info), ); pq.start_training(&meta_page); + unsafe { + pg_sys::IndexBuildHeapScan( + heap_relation.as_ptr(), + index_relation.as_ptr(), + index_info, + Some(build_callback_pq_train), + &mut pq, + ); + } + pq.finish_training(&mut write_stats); + let page_type = PqCompressionStorage::page_type(); let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); let mut state = StorageBuildState::PqCompression(&mut pq, &mut bs); @@ -259,7 +271,7 @@ fn do_heap_scan<'a>( ); } - do_heap_scan_with_state(&mut pq, &mut bs) + finalize_index_build(&mut pq, &mut bs, write_stats) } StorageType::BqSpeedup => { let mut bq = BqSpeedupStorage::new_for_build( @@ -268,6 +280,17 @@ fn do_heap_scan<'a>( get_attribute_number(index_info), ); bq.start_training(&meta_page); + unsafe { + pg_sys::IndexBuildHeapScan( + heap_relation.as_ptr(), + index_relation.as_ptr(), + index_info, + Some(build_callback_bq_train), + &mut bq, + ); + } + bq.finish_training(&mut write_stats); + let page_type = BqSpeedupStorage::page_type(); let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); let mut state = StorageBuildState::BqSpeedup(&mut bq, &mut bs); @@ -282,16 +305,16 @@ fn do_heap_scan<'a>( ); } - do_heap_scan_with_state(&mut bq, &mut bs) + finalize_index_build(&mut bq, &mut bs, write_stats) } } } -fn do_heap_scan_with_state(storage: &mut S, state: &mut BuildState) -> usize { - // we train the quantizer and add prepare to write quantized values to the nodes.\ - let mut write_stats = WriteStats::new(); - storage.finish_training(&mut write_stats); - +fn finalize_index_build( + storage: &mut S, + state: &mut BuildState, + mut write_stats: WriteStats, +) -> usize { match state.graph.get_neighbor_store() { GraphNeighborStore::Builder(builder) => { for (&index_pointer, neighbors) in builder.iter() { @@ -353,6 +376,38 @@ fn do_heap_scan_with_state(storage: &mut S, state: &mut BuildState) ntuples } +#[pg_guard] +unsafe extern "C" fn build_callback_bq_train( + _index: pg_sys::Relation, + _ctid: pg_sys::ItemPointer, + values: *mut pg_sys::Datum, + isnull: *mut bool, + _tuple_is_alive: bool, + state: *mut std::os::raw::c_void, +) { + let vec = PgVector::from_pg_parts(values, isnull, 0); + if let Some(vec) = vec { + let bq = (state as *mut BqSpeedupStorage).as_mut().unwrap(); + bq.add_sample(vec.to_slice()); + } +} + +#[pg_guard] +unsafe extern "C" fn build_callback_pq_train( + _index: pg_sys::Relation, + _ctid: pg_sys::ItemPointer, + values: *mut pg_sys::Datum, + isnull: *mut bool, + _tuple_is_alive: bool, + state: *mut std::os::raw::c_void, +) { + let vec = PgVector::from_pg_parts(values, isnull, 0); + if let Some(vec) = vec { + let pq = (state as *mut PqCompressionStorage).as_mut().unwrap(); + pq.add_sample(vec.to_slice()); + } +} + #[pg_guard] unsafe extern "C" fn build_callback( index: pg_sys::Relation, @@ -380,7 +435,6 @@ unsafe extern "C" fn build_callback( } } } - //todo: what do we do with nulls? } #[inline(always)] @@ -423,8 +477,6 @@ fn build_callback_internal( ); } - storage.add_sample(vector.to_slice()); - let index_pointer = storage.create_node( vector.to_slice(), heap_pointer, @@ -484,8 +536,34 @@ pub mod tests { embedding <=> ( SELECT ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding - FROM generate_series(1, 1536)); + FROM generate_series(1, 1536));"))?; + let test_vec: Option> = Spi::get_one(&format!( + "SELECT('{{' || array_to_string(array_agg(1.0), ',', '0') || '}}')::real[] AS embedding + FROM generate_series(1, 1536)" + ))?; + + let cnt: Option = Spi::get_one_with_args( + &format!( + " + SET enable_seqscan = 0; + SET enable_indexscan = 1; + SET tsv.query_search_list_size = 2; + WITH cte as (select * from test_data order by embedding <=> $1::vector) SELECT count(*) from cte; + ", + ), + vec![( + pgrx::PgOid::Custom(pgrx::pg_sys::FLOAT4ARRAYOID), + test_vec.clone().into_datum(), + )], + )?; + + //FIXME: should work in all cases + if !index_options.contains("num_neighbors=10") { + assert_eq!(cnt.unwrap(), 300, "initial count"); + } + + Spi::run(&format!(" -- test insert 2 vectors INSERT INTO test_data (embedding) SELECT @@ -525,11 +603,6 @@ pub mod tests { ", ))?; - let test_vec: Option> = Spi::get_one(&format!( - "SELECT('{{' || array_to_string(array_agg(1.0), ',', '0') || '}}')::real[] AS embedding -FROM generate_series(1, 1536)" - ))?; - let with_index: Option> = Spi::get_one_with_args( &format!( " diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index f420f60f..95a2dffd 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -4,7 +4,7 @@ use std::{cmp::Ordering, collections::HashSet}; use pgrx::PgRelation; -use crate::access_method::storage::NodeFullDistanceMeasure; +use crate::access_method::storage::NodeDistanceMeasure; use crate::util::{HeapPointer, IndexPointer, ItemPointer}; @@ -257,7 +257,7 @@ impl<'a> Graph<'a> { //no nodes in the graph return HashSet::with_capacity(0); } - let dm = storage.get_search_distance_measure(query, false); + let dm = storage.get_query_distance_measure(query); let search_list_size = meta_page.get_search_list_size_for_build() as usize; let mut l = ListSearchResult::new( @@ -287,7 +287,7 @@ impl<'a> Graph<'a> { //no nodes in the graph return ListSearchResult::empty(); } - let dm = storage.get_search_distance_measure(query, true); + let dm = storage.get_query_distance_measure(query); ListSearchResult::new( init_ids.unwrap(), @@ -372,7 +372,7 @@ impl<'a> Graph<'a> { let existing_neighbor = neighbor; let dist_state = unsafe { - storage.get_full_vector_distance_state( + storage.get_node_distance_measure( existing_neighbor.get_index_pointer_to_neighbor(), stats, ) diff --git a/timescale_vector/src/access_method/graph_neighbor_store.rs b/timescale_vector/src/access_method/graph_neighbor_store.rs index 2e185fd3..e0e4e596 100644 --- a/timescale_vector/src/access_method/graph_neighbor_store.rs +++ b/timescale_vector/src/access_method/graph_neighbor_store.rs @@ -88,11 +88,9 @@ impl GraphNeighborStore { GraphNeighborStore::Builder(b) => { b.get_neighbors_with_full_vector_distances(neighbors_of, result) } - GraphNeighborStore::Disk => storage.get_neighbors_with_full_vector_distances_from_disk( - neighbors_of, - result, - stats, - ), + GraphNeighborStore::Disk => { + storage.get_neighbors_with_distances_from_disk(neighbors_of, result, stats) + } }; } diff --git a/timescale_vector/src/access_method/plain_node.rs b/timescale_vector/src/access_method/plain_node.rs index 29ba9659..1e660eb3 100644 --- a/timescale_vector/src/access_method/plain_node.rs +++ b/timescale_vector/src/access_method/plain_node.rs @@ -121,13 +121,6 @@ impl ArchivedNode { past_last_index_pointers.offset = InvalidOffsetNumber; } } - - pub fn set_pq_vector(mut self: Pin<&mut Self>, pq_vector: &[u8]) { - for i in 0..=pq_vector.len() - 1 { - let mut pgv = self.as_mut().pq_vectors().index_pin(i); - *pgv = pq_vector[i]; - } - } } impl ArchivedData for ArchivedNode { diff --git a/timescale_vector/src/access_method/plain_storage.rs b/timescale_vector/src/access_method/plain_storage.rs index a5bbfcd2..ff55b367 100644 --- a/timescale_vector/src/access_method/plain_storage.rs +++ b/timescale_vector/src/access_method/plain_storage.rs @@ -8,7 +8,7 @@ use super::{ GreedySearchStats, StatsDistanceComparison, StatsNodeModify, StatsNodeRead, StatsNodeWrite, WriteStats, }, - storage::{ArchivedData, NodeFullDistanceMeasure, Storage}, + storage::{ArchivedData, NodeDistanceMeasure, Storage}, }; use pgrx::PgRelation; @@ -93,7 +93,7 @@ impl<'a> IndexFullDistanceMeasure<'a> { } } -impl<'a> NodeFullDistanceMeasure for IndexFullDistanceMeasure<'a> { +impl<'a> NodeDistanceMeasure for IndexFullDistanceMeasure<'a> { unsafe fn get_distance( &self, index_pointer: IndexPointer, @@ -137,7 +137,7 @@ impl PlainStorageLsnPrivateData { impl<'a> Storage for PlainStorage<'a> { type QueryDistanceMeasure = PlainDistanceMeasure; - type NodeFullDistanceMeasure<'b> = IndexFullDistanceMeasure<'b> where Self: 'b; + type NodeDistanceMeasure<'b> = IndexFullDistanceMeasure<'b> where Self: 'b; type ArchivedType = ArchivedNode; type LSNPrivateData = PlainStorageLsnPrivateData; @@ -178,25 +178,19 @@ impl<'a> Storage for PlainStorage<'a> { node.commit(); } - unsafe fn get_full_vector_distance_state<'b, S: StatsNodeRead>( + unsafe fn get_node_distance_measure<'b, S: StatsNodeRead>( &'b self, index_pointer: IndexPointer, stats: &mut S, - ) -> Self::NodeFullDistanceMeasure<'b> { + ) -> Self::NodeDistanceMeasure<'b> { IndexFullDistanceMeasure::with_index_pointer(self, index_pointer, stats) } - fn get_search_distance_measure( - &self, - query: PgVector, - _calc_distance_with_quantizer: bool, - ) -> PlainDistanceMeasure { + fn get_query_distance_measure(&self, query: PgVector) -> PlainDistanceMeasure { return PlainDistanceMeasure::Full(query); } - fn get_neighbors_with_full_vector_distances_from_disk< - S: StatsNodeRead + StatsDistanceComparison, - >( + fn get_neighbors_with_distances_from_disk( &self, neighbors_of: ItemPointer, result: &mut Vec, @@ -332,7 +326,7 @@ mod tests { #[test] fn test_plain_storage_delete_vacuum_plain() { crate::access_method::vacuum::tests::test_delete_vacuum_plain_scaffold( - "num_neighbors = 10", + "num_neighbors = 38", ); } diff --git a/timescale_vector/src/access_method/pq_quantizer.rs b/timescale_vector/src/access_method/pq_quantizer.rs index c1d757d5..13a2662b 100644 --- a/timescale_vector/src/access_method/pq_quantizer.rs +++ b/timescale_vector/src/access_method/pq_quantizer.rs @@ -9,7 +9,6 @@ use crate::access_method::{ use super::{ meta_page::MetaPage, - pg_vector::PgVector, stats::{StatsDistanceComparison, StatsNodeRead}, }; @@ -136,6 +135,44 @@ fn build_distance_table( distance_table } +/* +It seems that the node-node comparisons don't benefit from a table. remove this later if still true. +The node comparisons are done in the prune step where we do too few comparisons to benefit from a table. +fn build_distance_table_pq_query(pq: &Pq, query: &[u8]) -> Vec { + let sq = pq.subquantizers(); + let num_centroids = pq.n_quantizer_centroids(); + let num_subquantizers = sq.len_of(Axis(0)); + let dt_size = num_subquantizers * num_centroids; + let mut distance_table = vec![0.0; dt_size]; + + let mut elements_for_assert = 0; + for (subquantizer_index, subquantizer) in sq.outer_iter().enumerate() { + let query_centroid_index = query[subquantizer_index] as usize; + let query_slice = subquantizer.index_axis(Axis(0), query_centroid_index); + for (centroid_index, c) in subquantizer.outer_iter().enumerate() { + /* always use l2 for pq measurements since centeroids use k-means (which uses euclidean/l2 distance) + * The quantization also uses euclidean distance too. In the future we can experiment with k-mediods + * using a different distance measure, but this may make little difference. */ + let dist = if query_centroid_index == centroid_index { + debug_assert!(query_slice.as_slice().unwrap() == c.as_slice().unwrap()); + 0.0 + } else { + distance_l2_optimized_for_few_dimensions( + query_slice.as_slice().unwrap(), + c.as_slice().unwrap(), + ) + }; + assert!(subquantizer_index < num_subquantizers); + assert!(centroid_index * num_subquantizers + subquantizer_index < dt_size); + distance_table[centroid_index * num_subquantizers + subquantizer_index] = dist; + elements_for_assert += 1; + } + } + assert_eq!(dt_size, elements_for_assert); + distance_table +} +*/ + #[derive(Clone)] pub struct PqQuantizer { pq_trainer: Option, @@ -193,24 +230,40 @@ impl PqQuantizer { meta_page: &MetaPage, full_vector: &[f32], ) -> Vec { - if self.pq_trainer.is_some() { - let pq_vec_len = meta_page.get_pq_vector_length(); - vec![0; pq_vec_len] - } else { - assert!(self.pq.is_some()); - let pq_vec_len = meta_page.get_pq_vector_length(); - let res = self.quantize(full_vector); - assert!(res.len() == pq_vec_len); - res - } + assert!(self.pq_trainer.is_none() && self.pq.is_some()); + let pq_vec_len = meta_page.get_pq_vector_length(); + let res = self.quantize(full_vector); + assert!(res.len() == pq_vec_len); + res } - pub fn get_distance_table( + pub fn get_distance_table_full_query( &self, query: &[f32], distance_fn: fn(&[f32], &[f32]) -> f32, ) -> PqDistanceTable { - PqDistanceTable::new(&self.pq.as_ref().unwrap(), distance_fn, query) + PqDistanceTable::new_for_full_query(&self.pq.as_ref().unwrap(), distance_fn, query) + } + + 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 } } @@ -220,7 +273,7 @@ pub struct PqDistanceTable { } impl PqDistanceTable { - pub fn new( + pub fn new_for_full_query( pq: &Pq, distance_fn: fn(&[f32], &[f32]) -> f32, query: &[f32], @@ -245,7 +298,6 @@ impl PqDistanceTable { } pub enum PqSearchDistanceMeasure { - Full(PgVector), Pq(PqDistanceTable), } diff --git a/timescale_vector/src/access_method/pq_storage.rs b/timescale_vector/src/access_method/pq_storage.rs index 366dd26e..b3eb1321 100644 --- a/timescale_vector/src/access_method/pq_storage.rs +++ b/timescale_vector/src/access_method/pq_storage.rs @@ -11,8 +11,7 @@ use super::{ GreedySearchStats, StatsDistanceComparison, StatsNodeModify, StatsNodeRead, StatsNodeWrite, WriteStats, }, - storage::{NodeFullDistanceMeasure, Storage, StorageFullDistanceFromHeap}, - storage_common::{calculate_full_distance, HeapFullDistanceMeasure}, + storage::{NodeDistanceMeasure, Storage, StorageFullDistanceFromHeap}, }; use pgrx::PgRelation; @@ -23,6 +22,79 @@ use crate::util::{ use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance}; +/*pub struct PqNodeDistanceMeasure<'a> { + storage: &'a PqCompressionStorage<'a>, + table: PqDistanceTable, +} + +impl<'a> PqNodeDistanceMeasure<'a> { + pub unsafe fn with_index_pointer( + 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 table = storage + .quantizer + .get_distance_table_pq_query(node.pq_vector.as_slice()); + + Self { + storage: storage, + table: table, + } + } +} + +impl<'a> NodeDistanceMeasure for PqNodeDistanceMeasure<'a> { + unsafe fn get_distance( + &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.table.distance(node1.pq_vector.as_slice()) + } +}*/ + +pub struct PqNodeDistanceMeasure<'a> { + storage: &'a PqCompressionStorage<'a>, + vector: Vec, +} + +impl<'a> PqNodeDistanceMeasure<'a> { + pub unsafe fn with_index_pointer( + 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( + &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> { pub index: &'a PgRelation, pub distance_fn: fn(&[f32], &[f32]) -> f32, @@ -117,32 +189,8 @@ impl<'a> PqCompressionStorage<'a> { let rn_neighbor = unsafe { Node::read(self.index, neighbor_index_pointer, &mut lsr.stats) }; let node_neighbor = rn_neighbor.get_archived_node(); - let deleted = node_neighbor.is_deleted(); let distance = match lsr.sdm.as_ref().unwrap() { - PqSearchDistanceMeasure::Full(query) => { - let heap_pointer = node_neighbor.heap_item_pointer.deserialize_item_pointer(); - if deleted { - let pvt_data = PlainStorageLsnPrivateData::new( - neighbor_index_pointer, - node_neighbor, - gns, - ); - self.visit_lsn_internal(lsr, &pvt_data.neighbors, gns); - continue; - //for deleted nodes, we can't get the distance because we don't know the full vector - //so pretend it's the same distance as the parent - } else { - unsafe { - calculate_full_distance( - self, - heap_pointer, - query.to_slice(), - &mut lsr.stats, - ) - } - } - } PqSearchDistanceMeasure::Pq(table) => { PqSearchDistanceMeasure::calculate_pq_distance( table, @@ -164,7 +212,7 @@ impl<'a> PqCompressionStorage<'a> { impl<'a> Storage for PqCompressionStorage<'a> { type QueryDistanceMeasure = PqSearchDistanceMeasure; - type NodeFullDistanceMeasure<'b> = HeapFullDistanceMeasure<'b, PqCompressionStorage<'b>> where Self: 'b; + type NodeDistanceMeasure<'b> = PqNodeDistanceMeasure<'b> where Self: 'b; type ArchivedType = ArchivedNode; type LSNPrivateData = PlainStorageLsnPrivateData; //no data stored @@ -210,55 +258,33 @@ impl<'a> Storage for PqCompressionStorage<'a> { let mut archived = node.get_archived_node(); archived.as_mut().set_neighbors(neighbors, &meta); - let quantized = self.get_quantized_vector_from_heap_pointer( - archived.heap_item_pointer.deserialize_item_pointer(), - stats, - ); - - archived.as_mut().set_pq_vector(quantized.as_slice()); - node.commit(); } - unsafe fn get_full_vector_distance_state<'b, S: StatsNodeRead>( + unsafe fn get_node_distance_measure<'b, S: StatsNodeRead>( &'b self, index_pointer: IndexPointer, stats: &mut S, - ) -> HeapFullDistanceMeasure<'b, PqCompressionStorage<'b>> { - HeapFullDistanceMeasure::with_index_pointer(self, index_pointer, stats) + ) -> PqNodeDistanceMeasure<'b> { + PqNodeDistanceMeasure::with_index_pointer(self, index_pointer, stats) } - fn get_search_distance_measure( - &self, - query: PgVector, - calc_distance_with_quantizer: bool, - ) -> PqSearchDistanceMeasure { - if !calc_distance_with_quantizer { - return PqSearchDistanceMeasure::Full(query); - } else { - return PqSearchDistanceMeasure::Pq( - self.quantizer - .get_distance_table(query.to_slice(), self.distance_fn), - ); - } + fn get_query_distance_measure(&self, query: PgVector) -> PqSearchDistanceMeasure { + return PqSearchDistanceMeasure::Pq( + self.quantizer + .get_distance_table_full_query(query.to_slice(), self.distance_fn), + ); } //todo: same as Bq code? - fn get_neighbors_with_full_vector_distances_from_disk< - S: StatsNodeRead + StatsDistanceComparison, - >( + fn get_neighbors_with_distances_from_disk( &self, neighbors_of: ItemPointer, result: &mut Vec, stats: &mut S, ) { let rn = unsafe { Node::read(self.index, neighbors_of, stats) }; - let heap_pointer = rn - .get_archived_node() - .heap_item_pointer - .deserialize_item_pointer(); - let dist_state = - unsafe { HeapFullDistanceMeasure::with_heap_pointer(self, heap_pointer, stats) }; + let dist_state = unsafe { self.get_node_distance_measure(neighbors_of, stats) }; for n in rn.get_archived_node().iter_neighbors() { let dist = unsafe { dist_state.get_distance(n, stats) }; result.push(NeighborWithDistance::new(n, dist)) @@ -281,16 +307,6 @@ impl<'a> Storage for PqCompressionStorage<'a> { let node = rn.get_archived_node(); let distance = match lsr.sdm.as_ref().unwrap() { - PqSearchDistanceMeasure::Full(query) => { - if node.is_deleted() { - //FIXME: need to handle this case - panic!("can't handle the case where the init_id node is deleted"); - } - let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); - unsafe { - calculate_full_distance(self, heap_pointer, query.to_slice(), &mut lsr.stats) - } - } PqSearchDistanceMeasure::Pq(table) => PqSearchDistanceMeasure::calculate_pq_distance( table, node.pq_vector.as_slice(), @@ -375,7 +391,7 @@ mod tests { use pgrx::*; #[pg_test] - unsafe fn test_pq_storage_index_creation() -> spi::Result<()> { + unsafe fn test_pq_storage_index_creation_default() -> spi::Result<()> { crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( "num_neighbors=38, USE_PQ = TRUE", )?; @@ -394,7 +410,7 @@ mod tests { #[test] fn test_pq_storage_delete_vacuum_plain() { crate::access_method::vacuum::tests::test_delete_vacuum_plain_scaffold( - "num_neighbors = 10, use_pq = TRUE", + "num_neighbors = 38, use_pq = TRUE", ); } diff --git a/timescale_vector/src/access_method/stats.rs b/timescale_vector/src/access_method/stats.rs index 5e780cdb..f237a2fe 100644 --- a/timescale_vector/src/access_method/stats.rs +++ b/timescale_vector/src/access_method/stats.rs @@ -46,7 +46,7 @@ impl StatsDistanceComparison for PruneNeighborStats { } fn record_quantized_distance_comparison(&mut self) { - pgrx::error!("Should not use quantized distance comparisons during pruning"); + self.distance_comparisons += 1; } } diff --git a/timescale_vector/src/access_method/storage.rs b/timescale_vector/src/access_method/storage.rs index eb502db1..24152b19 100644 --- a/timescale_vector/src/access_method/storage.rs +++ b/timescale_vector/src/access_method/storage.rs @@ -16,7 +16,8 @@ use super::{ }, }; -pub trait NodeFullDistanceMeasure { +/// NodeDistanceMeasure keeps the state to make distance comparison between two nodes. +pub trait NodeDistanceMeasure { unsafe fn get_distance( &self, index_pointer: IndexPointer, @@ -33,8 +34,10 @@ pub trait ArchivedData { } pub trait Storage { + /// A QueryDistanceMeasure keeps the state to make distance comparison between a query given at initialization and a node. type QueryDistanceMeasure; - type NodeFullDistanceMeasure<'a>: NodeFullDistanceMeasure + /// A NodeDistanceMeasure keeps the state to make distance comparison between a node given at initialization and another node. + type NodeDistanceMeasure<'a>: NodeDistanceMeasure where Self: 'a; type ArchivedType: ArchivedData; @@ -63,17 +66,13 @@ pub trait Storage { stats: &mut S, ); - unsafe fn get_full_vector_distance_state<'a, S: StatsNodeRead>( + unsafe fn get_node_distance_measure<'a, S: StatsNodeRead>( &'a self, index_pointer: IndexPointer, stats: &mut S, - ) -> Self::NodeFullDistanceMeasure<'a>; + ) -> Self::NodeDistanceMeasure<'a>; - fn get_search_distance_measure( - &self, - query: PgVector, - calc_distance_with_quantizer: bool, - ) -> Self::QueryDistanceMeasure; + fn get_query_distance_measure(&self, query: PgVector) -> Self::QueryDistanceMeasure; fn visit_lsn( &self, @@ -100,9 +99,7 @@ pub trait Storage { where Self: Sized; - fn get_neighbors_with_full_vector_distances_from_disk< - S: StatsNodeRead + StatsDistanceComparison, - >( + fn get_neighbors_with_distances_from_disk( &self, neighbors_of: ItemPointer, result: &mut Vec, diff --git a/timescale_vector/src/access_method/storage_common.rs b/timescale_vector/src/access_method/storage_common.rs index c55d06c2..43ba256d 100644 --- a/timescale_vector/src/access_method/storage_common.rs +++ b/timescale_vector/src/access_method/storage_common.rs @@ -3,7 +3,7 @@ use crate::util::{HeapPointer, IndexPointer}; use super::{ pg_vector::PgVector, stats::{StatsDistanceComparison, StatsNodeRead}, - storage::{NodeFullDistanceMeasure, Storage, StorageFullDistanceFromHeap}, + storage::{NodeDistanceMeasure, Storage, StorageFullDistanceFromHeap}, }; pub struct HeapFullDistanceMeasure<'a, S: Storage + StorageFullDistanceFromHeap> { @@ -37,7 +37,7 @@ impl<'a, S: Storage + StorageFullDistanceFromHeap> HeapFullDistanceMeasure<'a, S } } -impl<'a, S: Storage + StorageFullDistanceFromHeap> NodeFullDistanceMeasure +impl<'a, S: Storage + StorageFullDistanceFromHeap> NodeDistanceMeasure for HeapFullDistanceMeasure<'a, S> { unsafe fn get_distance( @@ -54,21 +54,3 @@ impl<'a, S: Storage + StorageFullDistanceFromHeap> NodeFullDistanceMeasure (self.storage.get_distance_function())(slice1.to_slice(), slice2.to_slice()) } } - -pub unsafe fn calculate_full_distance< - S: Storage + StorageFullDistanceFromHeap, - T: StatsNodeRead + StatsDistanceComparison, ->( - storage: &S, - heap_pointer: HeapPointer, - query: &[f32], - stats: &mut T, -) -> f32 { - let slot = storage.get_heap_table_slot_from_heap_pointer(heap_pointer, stats); - let slice = unsafe { slot.get_pg_vector() }; - - stats.record_full_distance_comparison(); - let dist = (storage.get_distance_function())(slice.to_slice(), query); - debug_assert!(!dist.is_nan()); - dist -} diff --git a/timescale_vector/src/access_method/vacuum.rs b/timescale_vector/src/access_method/vacuum.rs index c542daab..20049aa9 100644 --- a/timescale_vector/src/access_method/vacuum.rs +++ b/timescale_vector/src/access_method/vacuum.rs @@ -183,10 +183,18 @@ pub mod tests { .batch_execute(&format!( "CREATE TABLE test_vac(embedding vector(256)); + select setseed(0.5); + -- generate 300 vectors INSERT INTO test_vac (embedding) SELECT - ('[' || i || ',1,1,{suffix}]')::vector - FROM generate_series(1, 300) i; + * + FROM ( + SELECT + ('[ 0 , ' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding + FROM + generate_series(1, 255 * 300) i + GROUP BY + i % 300) g; INSERT INTO test_vac(embedding) VALUES ('[1,2,3,{suffix}]'), ('[4,5,6,{suffix}]'), ('[7,8,10,{suffix}]'); @@ -202,7 +210,7 @@ pub mod tests { client.execute("set enable_seqscan = 0;", &[]).unwrap(); let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); - assert_eq!(cnt, 303); + assert_eq!(cnt, 303, "initial count"); client .execute( @@ -228,7 +236,7 @@ pub mod tests { client.execute("set enable_seqscan = 0;", &[]).unwrap(); let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); //if the old index is still used the count is 304 - assert_eq!(cnt, 303); + assert_eq!(cnt, 303, "count after vacuum"); //do another delete for same items (noop) client @@ -241,7 +249,7 @@ pub mod tests { client.execute("set enable_seqscan = 0;", &[]).unwrap(); let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); //if the old index is still used the count is 304 - assert_eq!(cnt, 303); + assert_eq!(cnt, 303, "count after delete"); client.execute("DROP INDEX idxtest_vac", &[]).unwrap(); client.execute("DROP TABLE test_vac", &[]).unwrap(); @@ -277,11 +285,18 @@ pub mod tests { .batch_execute(&format!( "CREATE TABLE test_vac_full(embedding vector(256)); + select setseed(0.5); -- generate 300 vectors INSERT INTO test_vac_full (embedding) SELECT - ('[' || i || ',2,3,{suffix}]')::vector - FROM generate_series(1, 300) i; + * + FROM ( + SELECT + ('[ 0 , ' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding + FROM + generate_series(1, 255 * 300) i + GROUP BY + i % 300) g; INSERT INTO test_vac_full(embedding) VALUES ('[1,2,3,{suffix}]'), ('[4,5,6,{suffix}]'), ('[7,8,10,{suffix}]'); @@ -296,17 +311,25 @@ pub mod tests { client.execute("set enable_seqscan = 0;", &[]).unwrap(); let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac_full order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); std::thread::sleep(std::time::Duration::from_millis(10000)); - assert_eq!(cnt, 303); + assert_eq!(cnt, 303, "initial count"); client.execute("DELETE FROM test_vac_full", &[]).unwrap(); client .execute( &format!( - "INSERT INTO test_vac_full (embedding) - SELECT - ('[' || i || ',2,3,{suffix}]')::vector - FROM generate_series(1, 300) i;" + " + INSERT INTO test_vac_full (embedding) + SELECT + * + FROM ( + SELECT + ('[ 0 , ' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding + FROM + generate_series(1, 255 * 300) i + GROUP BY + i % 300) g; + " ), &[], ) @@ -326,7 +349,7 @@ pub mod tests { client.execute("set enable_seqscan = 0;", &[]).unwrap(); let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac_full order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); - assert_eq!(cnt, 301); + assert_eq!(cnt, 301, "count after full vacuum"); client.execute("DROP INDEX idxtest_vac_full", &[]).unwrap(); client.execute("DROP TABLE test_vac_full", &[]).unwrap(); From f3dc51b921a6947f450ea7ed1d39c02a9df2ca88 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 26 Mar 2024 12:21:17 -0400 Subject: [PATCH 19/44] basic resort --- timescale_vector/src/access_method/bq.rs | 87 +++++++------ timescale_vector/src/access_method/build.rs | 2 +- timescale_vector/src/access_method/guc.rs | 12 ++ .../src/access_method/plain_storage.rs | 14 ++- .../src/access_method/pq_quantizer.rs | 3 +- .../src/access_method/pq_storage.rs | 93 +++++++------- timescale_vector/src/access_method/scan.rs | 117 +++++++++++++++--- timescale_vector/src/access_method/stats.rs | 13 ++ timescale_vector/src/access_method/storage.rs | 30 ++--- .../src/access_method/storage_common.rs | 62 ++-------- timescale_vector/src/util/table_slot.rs | 6 +- 11 files changed, 251 insertions(+), 188 deletions(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index a8270531..336eff1c 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -4,10 +4,11 @@ use super::{ graph_neighbor_store::GraphNeighborStore, pg_vector::PgVector, stats::{ - GreedySearchStats, StatsDistanceComparison, StatsNodeModify, StatsNodeRead, StatsNodeWrite, - WriteStats, + GreedySearchStats, StatsDistanceComparison, StatsHeapNodeRead, StatsNodeModify, + StatsNodeRead, StatsNodeWrite, WriteStats, }, - storage::{ArchivedData, NodeDistanceMeasure, Storage, StorageFullDistanceFromHeap}, + storage::{ArchivedData, NodeDistanceMeasure, Storage}, + storage_common::get_attribute_number_from_index, }; use std::{cell::RefCell, collections::HashMap, iter::once, marker::PhantomData, pin::Pin}; @@ -193,7 +194,7 @@ impl BqDistanceTable { //FIXME: cleanup make this into a struct pub enum BqSearchDistanceMeasure { - Bq(BqDistanceTable), + Bq(BqDistanceTable, PgVector), } impl BqSearchDistanceMeasure { @@ -345,6 +346,7 @@ impl<'a> BqSpeedupStorage<'a> { pub fn load_for_search( index_relation: &'a PgRelation, + heap_relation: &'a PgRelation, quantizer: &BqQuantizer, ) -> BqSpeedupStorage<'a> { Self { @@ -352,8 +354,8 @@ impl<'a> BqSpeedupStorage<'a> { distance_fn: default_distance, //OPT: get rid of clone quantizer: quantizer.clone(), - heap_rel: None, - heap_attr: None, + heap_rel: Some(heap_relation), + heap_attr: Some(get_attribute_number_from_index(index_relation)), qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), } } @@ -399,7 +401,7 @@ impl<'a> BqSpeedupStorage<'a> { } let distance = match lsr.sdm.as_ref().unwrap() { - BqSearchDistanceMeasure::Bq(table) => { + BqSearchDistanceMeasure::Bq(table, _) => { /* Note: there is no additional node reads here. We get all of our info from node_visiting * This is what gives us a speedup in BQ Speedup */ match gns { @@ -411,7 +413,7 @@ impl<'a> BqSpeedupStorage<'a> { &mut lsr.stats, ) } - GraphNeighborStore::Builder(b) => { + GraphNeighborStore::Builder(_) => { let mut cache = self.qv_cache.borrow_mut(); let bq_vector = cache.get(neighbor_index_pointer, self, &mut lsr.stats); let dist = BqSearchDistanceMeasure::calculate_bq_distance( @@ -432,6 +434,19 @@ impl<'a> BqSpeedupStorage<'a> { lsr.insert_neighbor(lsn); } } + + unsafe fn get_heap_table_slot_from_heap_pointer( + &self, + heap_pointer: HeapPointer, + stats: &mut T, + ) -> TableSlot { + TableSlot::new( + self.heap_rel.unwrap(), + heap_pointer, + self.heap_attr.unwrap(), + stats, + ) + } } pub type BqSpeedupStorageLsnPrivateData = PhantomData; //no data stored @@ -510,9 +525,26 @@ impl<'a> Storage for BqSpeedupStorage<'a> { return BqSearchDistanceMeasure::Bq( self.quantizer .get_distance_table(query.to_slice(), self.distance_fn), + query, ); } + fn get_full_distance_for_resort( + &self, + qdm: &Self::QueryDistanceMeasure, + _index_pointer: IndexPointer, + heap_pointer: HeapPointer, + stats: &mut S, + ) -> f32 { + let slot = unsafe { self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) }; + match qdm { + BqSearchDistanceMeasure::Bq(_, query) => self.get_distance_function()( + unsafe { slot.get_pg_vector().to_slice() }, + query.to_slice(), + ), + } + } + fn get_neighbors_with_distances_from_disk( &self, neighbors_of: ItemPointer, @@ -549,11 +581,13 @@ impl<'a> Storage for BqSpeedupStorage<'a> { let node = rn.get_archived_node(); let distance = match lsr.sdm.as_ref().unwrap() { - BqSearchDistanceMeasure::Bq(table) => BqSearchDistanceMeasure::calculate_bq_distance( - table, - node.bq_vector.as_slice(), - &mut lsr.stats, - ), + BqSearchDistanceMeasure::Bq(table, _) => { + BqSearchDistanceMeasure::calculate_bq_distance( + table, + node.bq_vector.as_slice(), + &mut lsr.stats, + ) + } }; ListSearchNeighbor::new(index_pointer, distance, PhantomData::) @@ -609,33 +643,6 @@ impl<'a> Storage for BqSpeedupStorage<'a> { } } -impl<'a> StorageFullDistanceFromHeap for BqSpeedupStorage<'a> { - unsafe fn get_heap_table_slot_from_index_pointer( - &self, - index_pointer: IndexPointer, - stats: &mut S, - ) -> TableSlot { - let rn = unsafe { BqNode::read(self.index, index_pointer, stats) }; - let node = rn.get_archived_node(); - let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); - - self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) - } - - unsafe fn get_heap_table_slot_from_heap_pointer( - &self, - heap_pointer: HeapPointer, - stats: &mut T, - ) -> TableSlot { - TableSlot::new( - self.heap_rel.unwrap(), - heap_pointer, - self.heap_attr.unwrap(), - stats, - ) - } -} - use timescale_vector_derive::{Readable, Writeable}; #[derive(Archive, Deserialize, Serialize, Readable, Writeable)] diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index dedd1337..68e21d32 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -200,7 +200,7 @@ pub extern "C" fn ambuildempty(_index_relation: pg_sys::Relation) { panic!("ambuildempty: not yet implemented") } -fn get_attribute_number(index_info: *mut pg_sys::IndexInfo) -> pg_sys::AttrNumber { +pub fn get_attribute_number(index_info: *mut pg_sys::IndexInfo) -> pg_sys::AttrNumber { unsafe { assert!((*index_info).ii_NumIndexAttrs == 1) }; unsafe { (*index_info).ii_IndexAttrNumbers[0] } } diff --git a/timescale_vector/src/access_method/guc.rs b/timescale_vector/src/access_method/guc.rs index 5e9d8a54..6a83b1bb 100644 --- a/timescale_vector/src/access_method/guc.rs +++ b/timescale_vector/src/access_method/guc.rs @@ -1,6 +1,7 @@ use pgrx::*; pub static TSV_QUERY_SEARCH_LIST_SIZE: GucSetting = GucSetting::::new(100); +pub static TSV_RESORT_SIZE: GucSetting = GucSetting::::new(10); pub fn init() { GucRegistry::define_int_guc( @@ -13,4 +14,15 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + + GucRegistry::define_int_guc( + "tsv.query_resort", + "The resort size used in queries", + "Resort size.", + &TSV_RESORT_SIZE, + 1, + 1000, + GucContext::Userset, + GucFlags::default(), + ); } diff --git a/timescale_vector/src/access_method/plain_storage.rs b/timescale_vector/src/access_method/plain_storage.rs index ff55b367..b7a1ec56 100644 --- a/timescale_vector/src/access_method/plain_storage.rs +++ b/timescale_vector/src/access_method/plain_storage.rs @@ -5,8 +5,8 @@ use super::{ pg_vector::PgVector, plain_node::{ArchivedNode, Node, ReadableNode}, stats::{ - GreedySearchStats, StatsDistanceComparison, StatsNodeModify, StatsNodeRead, StatsNodeWrite, - WriteStats, + GreedySearchStats, StatsDistanceComparison, StatsHeapNodeRead, StatsNodeModify, + StatsNodeRead, StatsNodeWrite, WriteStats, }, storage::{ArchivedData, NodeDistanceMeasure, Storage}, }; @@ -189,7 +189,15 @@ impl<'a> Storage for PlainStorage<'a> { fn get_query_distance_measure(&self, query: PgVector) -> PlainDistanceMeasure { return PlainDistanceMeasure::Full(query); } - + fn get_full_distance_for_resort( + &self, + _qdm: &Self::QueryDistanceMeasure, + _index_pointer: IndexPointer, + _heap_pointer: HeapPointer, + _stats: &mut S, + ) -> f32 { + pgrx::error!("Plain node should never be resorted"); + } fn get_neighbors_with_distances_from_disk( &self, neighbors_of: ItemPointer, diff --git a/timescale_vector/src/access_method/pq_quantizer.rs b/timescale_vector/src/access_method/pq_quantizer.rs index 13a2662b..dfb44005 100644 --- a/timescale_vector/src/access_method/pq_quantizer.rs +++ b/timescale_vector/src/access_method/pq_quantizer.rs @@ -9,6 +9,7 @@ use crate::access_method::{ use super::{ meta_page::MetaPage, + pg_vector::PgVector, stats::{StatsDistanceComparison, StatsNodeRead}, }; @@ -298,7 +299,7 @@ impl PqDistanceTable { } pub enum PqSearchDistanceMeasure { - Pq(PqDistanceTable), + Pq(PqDistanceTable, PgVector), } impl PqSearchDistanceMeasure { diff --git a/timescale_vector/src/access_method/pq_storage.rs b/timescale_vector/src/access_method/pq_storage.rs index b3eb1321..974ceda7 100644 --- a/timescale_vector/src/access_method/pq_storage.rs +++ b/timescale_vector/src/access_method/pq_storage.rs @@ -8,10 +8,11 @@ use super::{ pq_quantizer::{PqQuantizer, PqSearchDistanceMeasure, PqVectorElement}, pq_quantizer_storage::write_pq, stats::{ - GreedySearchStats, StatsDistanceComparison, StatsNodeModify, StatsNodeRead, StatsNodeWrite, - WriteStats, + GreedySearchStats, StatsDistanceComparison, StatsHeapNodeRead, StatsNodeModify, + StatsNodeRead, StatsNodeWrite, WriteStats, }, - storage::{NodeDistanceMeasure, Storage, StorageFullDistanceFromHeap}, + storage::{NodeDistanceMeasure, Storage}, + storage_common::get_attribute_number_from_index, }; use pgrx::PgRelation; @@ -144,6 +145,7 @@ impl<'a> PqCompressionStorage<'a> { pub fn load_for_search( index_relation: &'a PgRelation, + heap_relation: &'a PgRelation, quantizer: &PqQuantizer, ) -> PqCompressionStorage<'a> { Self { @@ -151,21 +153,11 @@ impl<'a> PqCompressionStorage<'a> { distance_fn: default_distance, //OPT: get rid of clone quantizer: quantizer.clone(), - heap_rel: None, - heap_attr: None, + heap_rel: Some(heap_relation), + heap_attr: Some(get_attribute_number_from_index(heap_relation)), } } - fn get_quantized_vector_from_heap_pointer( - &self, - heap_pointer: HeapPointer, - stats: &mut S, - ) -> Vec { - let slot = unsafe { self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) }; - let slice = unsafe { slot.get_pg_vector() }; - self.quantizer.quantize(slice.to_slice()) - } - fn write_quantizer_metadata(&self, stats: &mut S) { let pq = self.quantizer.must_get_pq(); let index_pointer: IndexPointer = unsafe { write_pq(pq, &self.index, stats) }; @@ -191,7 +183,7 @@ impl<'a> PqCompressionStorage<'a> { let node_neighbor = rn_neighbor.get_archived_node(); let distance = match lsr.sdm.as_ref().unwrap() { - PqSearchDistanceMeasure::Pq(table) => { + PqSearchDistanceMeasure::Pq(table, _) => { PqSearchDistanceMeasure::calculate_pq_distance( table, node_neighbor.pq_vector.as_slice(), @@ -208,6 +200,19 @@ impl<'a> PqCompressionStorage<'a> { lsr.insert_neighbor(lsn); } } + + unsafe fn get_heap_table_slot_from_heap_pointer( + &self, + heap_pointer: HeapPointer, + stats: &mut T, + ) -> TableSlot { + TableSlot::new( + self.heap_rel.unwrap(), + heap_pointer, + self.heap_attr.unwrap(), + stats, + ) + } } impl<'a> Storage for PqCompressionStorage<'a> { @@ -273,9 +278,26 @@ impl<'a> Storage for PqCompressionStorage<'a> { return PqSearchDistanceMeasure::Pq( self.quantizer .get_distance_table_full_query(query.to_slice(), self.distance_fn), + query, ); } + fn get_full_distance_for_resort( + &self, + qdm: &Self::QueryDistanceMeasure, + _index_pointer: IndexPointer, + heap_pointer: HeapPointer, + stats: &mut S, + ) -> f32 { + let slot = unsafe { self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) }; + match qdm { + PqSearchDistanceMeasure::Pq(_, query) => self.get_distance_function()( + unsafe { slot.get_pg_vector().to_slice() }, + query.to_slice(), + ), + } + } + //todo: same as Bq code? fn get_neighbors_with_distances_from_disk( &self, @@ -307,11 +329,13 @@ impl<'a> Storage for PqCompressionStorage<'a> { let node = rn.get_archived_node(); let distance = match lsr.sdm.as_ref().unwrap() { - PqSearchDistanceMeasure::Pq(table) => PqSearchDistanceMeasure::calculate_pq_distance( - table, - node.pq_vector.as_slice(), - &mut lsr.stats, - ), + PqSearchDistanceMeasure::Pq(table, _) => { + PqSearchDistanceMeasure::calculate_pq_distance( + table, + node.pq_vector.as_slice(), + &mut lsr.stats, + ) + } }; ListSearchNeighbor::new( @@ -358,33 +382,6 @@ impl<'a> Storage for PqCompressionStorage<'a> { } } -impl<'a> StorageFullDistanceFromHeap for PqCompressionStorage<'a> { - unsafe fn get_heap_table_slot_from_index_pointer( - &self, - index_pointer: IndexPointer, - stats: &mut S, - ) -> TableSlot { - let rn = unsafe { Node::read(self.index, index_pointer, stats) }; - let node = rn.get_archived_node(); - let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); - - self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) - } - - unsafe fn get_heap_table_slot_from_heap_pointer( - &self, - heap_pointer: HeapPointer, - stats: &mut T, - ) -> TableSlot { - TableSlot::new( - self.heap_rel.unwrap(), - heap_pointer, - self.heap_attr.unwrap(), - stats, - ) - } -} - #[cfg(any(test, feature = "pg_test"))] #[pgrx::pg_schema] mod tests { diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 39b69a92..c6d3a126 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -1,3 +1,5 @@ +use std::collections::BinaryHeap; + use pgrx::{pg_sys::InvalidOffsetNumber, *}; use crate::{ @@ -5,7 +7,7 @@ use crate::{ bq::BqSpeedupStorage, graph_neighbor_store::GraphNeighborStore, meta_page::MetaPage, pg_vector::PgVector, }, - util::{buffer::PinnedBufferShare, HeapPointer}, + util::{buffer::PinnedBufferShare, HeapPointer, IndexPointer}, }; use super::{ @@ -44,7 +46,13 @@ impl TSVScanState { } } - fn initialize(&mut self, index: &PgRelation, query: PgVector, search_list_size: usize) { + fn initialize( + &mut self, + index: &PgRelation, + heap: &PgRelation, + query: PgVector, + search_list_size: usize, + ) { let meta_page = MetaPage::read(&index); let storage = meta_page.get_storage_type(); @@ -59,7 +67,7 @@ impl TSVScanState { StorageType::PqCompression => { let mut stats = QuantizerStats::new(); let quantizer = PqQuantizer::load(index, &meta_page, &mut stats); - let pq = PqCompressionStorage::load_for_search(index, &quantizer); + let pq = PqCompressionStorage::load_for_search(index, heap, &quantizer); let it = TSVResponseIterator::new(&pq, index, query, search_list_size, meta_page, stats); StorageState::PqCompression(quantizer, it) @@ -67,7 +75,7 @@ impl TSVScanState { StorageType::BqSpeedup => { let mut stats = QuantizerStats::new(); let quantizer = unsafe { BqMeans::load(index, &meta_page, &mut stats) }; - let bq = BqSpeedupStorage::load_for_search(index, &quantizer); + let bq = BqSpeedupStorage::load_for_search(index, heap, &quantizer); let it = TSVResponseIterator::new(&bq, index, query, search_list_size, meta_page, stats); StorageState::BqSpeedup(quantizer, it) @@ -78,6 +86,34 @@ impl TSVScanState { } } +struct ResortData { + heap_pointer: HeapPointer, + index_pointer: IndexPointer, + distance: f32, +} + +impl PartialEq for ResortData { + fn eq(&self, other: &Self) -> bool { + self.heap_pointer == other.heap_pointer + } +} + +impl PartialOrd for ResortData { + fn partial_cmp(&self, other: &Self) -> Option { + //notice the reverse here. Other is the one that is being compared to self + //this allows us to have a min heap + other.distance.partial_cmp(&self.distance) + } +} + +impl Eq for ResortData {} + +impl Ord for ResortData { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.partial_cmp(other).unwrap() + } +} + struct TSVResponseIterator { lsr: ListSearchResult, search_list_size: usize, @@ -85,6 +121,7 @@ struct TSVResponseIterator { last_buffer: Option, meta_page: MetaPage, quantizer_stats: QuantizerStats, + resort_buffer: BinaryHeap, } impl TSVResponseIterator { @@ -100,6 +137,7 @@ impl TSVResponseIterator { let graph = Graph::new(GraphNeighborStore::Disk, &mut meta_page); let lsr = graph.greedy_search_streaming_init(query, search_list_size, storage); + let resort_size = super::guc::TSV_RESORT_SIZE.get() as usize; Self { search_list_size, @@ -108,6 +146,7 @@ impl TSVResponseIterator { last_buffer: None, meta_page, quantizer_stats, + resort_buffer: BinaryHeap::with_capacity(resort_size), } } } @@ -117,7 +156,7 @@ impl TSVResponseIterator { &mut self, index: &PgRelation, storage: &S, - ) -> Option { + ) -> Option<(HeapPointer, IndexPointer)> { let graph = Graph::new(GraphNeighborStore::Disk, &mut self.meta_page); /* Iterate until we find a non-deleted tuple */ @@ -142,7 +181,7 @@ impl TSVResponseIterator { /* deleted tuple */ continue; } - return Some(heap_pointer); + return Some((heap_pointer, index_pointer)); } None => { self.last_buffer = None; @@ -151,6 +190,43 @@ impl TSVResponseIterator { } } } + + fn next_with_resort>( + &mut self, + index: &PgRelation, + storage: &S, + ) -> Option<(HeapPointer, IndexPointer)> { + if self.resort_buffer.capacity() == 0 { + return self.next(index, storage); + } + + while self.resort_buffer.len() < self.resort_buffer.capacity() { + match self.next(index, storage) { + Some((heap_pointer, index_pointer)) => { + let distance = storage.get_full_distance_for_resort( + self.lsr.sdm.as_ref().unwrap(), + index_pointer, + heap_pointer, + &mut self.lsr.stats, + ); + + self.resort_buffer.push(ResortData { + heap_pointer, + index_pointer, + distance, + }); + } + None => { + break; + } + } + } + + match self.resort_buffer.pop() { + Some(rd) => Some((rd.heap_pointer, rd.index_pointer)), + None => None, + } + } } /* @@ -195,6 +271,7 @@ pub extern "C" fn amrescan( } let mut scan: PgBox = unsafe { PgBox::from_pg(scan) }; let indexrel = unsafe { PgRelation::from_pg(scan.indexRelation) }; + let heaprel = unsafe { PgRelation::from_pg(scan.heapRelation) }; let meta_page = MetaPage::read(&indexrel); let _storage = meta_page.get_storage_type(); @@ -212,7 +289,7 @@ pub extern "C" fn amrescan( let search_list_size = super::guc::TSV_QUERY_SEARCH_LIST_SIZE.get() as usize; let state = unsafe { (scan.opaque as *mut TSVScanState).as_mut() }.expect("no scandesc state"); - state.initialize(&indexrel, query, search_list_size); + state.initialize(&indexrel, &heaprel, query, search_list_size); /*match &mut storage { Storage::None => pgrx::error!("not implemented"), Storage::PQ(_pq) => pgrx::error!("not implemented"), @@ -236,33 +313,35 @@ pub extern "C" fn amgettuple( //let iter = unsafe { state.iterator.as_mut() }.expect("no iterator in state"); let indexrel = unsafe { PgRelation::from_pg(scan.indexRelation) }; + let heaprel = unsafe { PgRelation::from_pg(scan.heapRelation) }; let mut storage = unsafe { state.storage.as_mut() }.expect("no storage in state"); match &mut storage { StorageState::BqSpeedup(quantizer, iter) => { - let bq = BqSpeedupStorage::load_for_search(&indexrel, quantizer); - get_tuple(&bq, &indexrel, iter, scan) + let bq = BqSpeedupStorage::load_for_search(&indexrel, &heaprel, quantizer); + let next = iter.next_with_resort(&indexrel, &bq); + get_tuple(next, scan) } StorageState::PqCompression(quantizer, iter) => { - let pq = PqCompressionStorage::load_for_search(&indexrel, quantizer); - get_tuple(&pq, &indexrel, iter, scan) + let pq = PqCompressionStorage::load_for_search(&indexrel, &heaprel, quantizer); + let next = iter.next_with_resort(&indexrel, &pq); + get_tuple(next, scan) } StorageState::Plain(iter) => { - let bq = PlainStorage::load_for_search(&indexrel); - get_tuple(&bq, &indexrel, iter, scan) + let storage = PlainStorage::load_for_search(&indexrel); + let next = iter.next(&indexrel, &storage); + get_tuple(next, scan) } } } -fn get_tuple<'a, S: Storage>( - storage: &S, - index: &'a PgRelation, - iter: &'a mut TSVResponseIterator, +fn get_tuple( + next: Option<(HeapPointer, IndexPointer)>, mut scan: PgBox, ) -> bool { scan.xs_recheckorderby = false; - match iter.next(&index, storage) { - Some(heap_pointer) => { + match next { + Some((heap_pointer, _)) => { let tid_to_set = &mut scan.xs_heaptid; heap_pointer.to_item_pointer_data(tid_to_set); true diff --git a/timescale_vector/src/access_method/stats.rs b/timescale_vector/src/access_method/stats.rs index f237a2fe..ab2c6ceb 100644 --- a/timescale_vector/src/access_method/stats.rs +++ b/timescale_vector/src/access_method/stats.rs @@ -4,6 +4,10 @@ pub trait StatsNodeRead { fn record_read(&mut self); } +pub trait StatsHeapNodeRead { + fn record_heap_read(&mut self); +} + pub trait StatsNodeModify { fn record_modify(&mut self); } @@ -67,6 +71,7 @@ pub struct GreedySearchStats { calls: usize, full_distance_comparisons: usize, node_reads: usize, + node_heap_reads: usize, quantized_distance_comparisons: usize, } @@ -76,6 +81,7 @@ impl GreedySearchStats { calls: 0, full_distance_comparisons: 0, node_reads: 0, + node_heap_reads: 0, quantized_distance_comparisons: 0, } } @@ -84,6 +90,7 @@ impl GreedySearchStats { self.calls += other.calls; self.full_distance_comparisons += other.full_distance_comparisons; self.node_reads += other.node_reads; + self.node_heap_reads += other.node_heap_reads; self.quantized_distance_comparisons += other.quantized_distance_comparisons; } @@ -118,6 +125,12 @@ impl StatsNodeRead for GreedySearchStats { } } +impl StatsHeapNodeRead for GreedySearchStats { + fn record_heap_read(&mut self) { + self.node_heap_reads += 1; + } +} + impl StatsDistanceComparison for GreedySearchStats { fn record_full_distance_comparison(&mut self) { self.full_distance_comparisons += 1; diff --git a/timescale_vector/src/access_method/storage.rs b/timescale_vector/src/access_method/storage.rs index 24152b19..0bfcbd2f 100644 --- a/timescale_vector/src/access_method/storage.rs +++ b/timescale_vector/src/access_method/storage.rs @@ -1,8 +1,6 @@ use std::pin::Pin; -use crate::util::{ - page::PageType, table_slot::TableSlot, tape::Tape, HeapPointer, IndexPointer, ItemPointer, -}; +use crate::util::{page::PageType, tape::Tape, HeapPointer, IndexPointer, ItemPointer}; use super::{ graph::{ListSearchNeighbor, ListSearchResult}, @@ -11,8 +9,8 @@ use super::{ neighbor_with_distance::NeighborWithDistance, pg_vector::PgVector, stats::{ - GreedySearchStats, StatsDistanceComparison, StatsNodeModify, StatsNodeRead, StatsNodeWrite, - WriteStats, + GreedySearchStats, StatsDistanceComparison, StatsHeapNodeRead, StatsNodeModify, + StatsNodeRead, StatsNodeWrite, WriteStats, }, }; @@ -74,6 +72,14 @@ pub trait Storage { fn get_query_distance_measure(&self, query: PgVector) -> Self::QueryDistanceMeasure; + fn get_full_distance_for_resort( + &self, + query: &Self::QueryDistanceMeasure, + index_pointer: IndexPointer, + heap_pointer: HeapPointer, + stats: &mut S, + ) -> f32; + fn visit_lsn( &self, lsr: &mut ListSearchResult, @@ -117,20 +123,6 @@ pub trait Storage { fn get_distance_function(&self) -> fn(&[f32], &[f32]) -> f32; } -pub trait StorageFullDistanceFromHeap { - unsafe fn get_heap_table_slot_from_index_pointer( - &self, - index_pointer: IndexPointer, - stats: &mut T, - ) -> TableSlot; - - unsafe fn get_heap_table_slot_from_heap_pointer( - &self, - heap_pointer: HeapPointer, - stats: &mut T, - ) -> TableSlot; -} - pub enum StorageType { BqSpeedup, PqCompression, diff --git a/timescale_vector/src/access_method/storage_common.rs b/timescale_vector/src/access_method/storage_common.rs index 43ba256d..38484363 100644 --- a/timescale_vector/src/access_method/storage_common.rs +++ b/timescale_vector/src/access_method/storage_common.rs @@ -1,56 +1,10 @@ -use crate::util::{HeapPointer, IndexPointer}; - -use super::{ - pg_vector::PgVector, - stats::{StatsDistanceComparison, StatsNodeRead}, - storage::{NodeDistanceMeasure, Storage, StorageFullDistanceFromHeap}, -}; - -pub struct HeapFullDistanceMeasure<'a, S: Storage + StorageFullDistanceFromHeap> { - storage: &'a S, - vector: PgVector, -} - -impl<'a, S: Storage + StorageFullDistanceFromHeap> HeapFullDistanceMeasure<'a, S> { - pub unsafe fn with_index_pointer( - storage: &'a S, - index_pointer: IndexPointer, - stats: &mut T, - ) -> Self { - let slot = storage.get_heap_table_slot_from_index_pointer(index_pointer, stats); - Self { - storage: storage, - vector: slot.get_pg_vector(), - } - } - - pub unsafe fn with_heap_pointer( - storage: &'a S, - heap_pointer: HeapPointer, - stats: &mut T, - ) -> Self { - let slot = storage.get_heap_table_slot_from_heap_pointer(heap_pointer, stats); - Self { - storage: storage, - vector: slot.get_pg_vector(), - } - } -} - -impl<'a, S: Storage + StorageFullDistanceFromHeap> NodeDistanceMeasure - for HeapFullDistanceMeasure<'a, S> -{ - unsafe fn get_distance( - &self, - index_pointer: IndexPointer, - stats: &mut T, - ) -> f32 { - let slot = self - .storage - .get_heap_table_slot_from_index_pointer(index_pointer, stats); - stats.record_full_distance_comparison(); - let slice1 = slot.get_pg_vector(); - let slice2 = &self.vector; - (self.storage.get_distance_function())(slice1.to_slice(), slice2.to_slice()) +use pgrx::{pg_sys, PgRelation}; + +pub fn get_attribute_number_from_index(index: &PgRelation) -> pg_sys::AttrNumber { + unsafe { + let a = index.rd_index; + let natts = (*a).indnatts; + assert!(natts == 1); + (*a).indkey.values.as_slice(natts as _)[0] } } diff --git a/timescale_vector/src/util/table_slot.rs b/timescale_vector/src/util/table_slot.rs index 147c4752..dc27d8bb 100644 --- a/timescale_vector/src/util/table_slot.rs +++ b/timescale_vector/src/util/table_slot.rs @@ -2,7 +2,7 @@ use pgrx::pg_sys::{Datum, TupleTableSlot}; use pgrx::{pg_sys, PgBox, PgRelation}; use crate::access_method::pg_vector::PgVector; -use crate::access_method::stats::StatsNodeRead; +use crate::access_method::stats::StatsHeapNodeRead; use crate::util::ports::slot_getattr; use crate::util::HeapPointer; @@ -12,7 +12,7 @@ pub struct TableSlot { } impl TableSlot { - pub unsafe fn new( + pub unsafe fn new( heap_rel: &PgRelation, heap_pointer: HeapPointer, attribute_number: pg_sys::AttrNumber, @@ -35,7 +35,7 @@ impl TableSlot { &mut pg_sys::SnapshotAnyData, slot.as_ptr(), ); - stats.record_read(); + stats.record_heap_read(); Self { slot, From 54e8000ddc8a684cf0cd474a73fff1526ddc1da1 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 26 Mar 2024 13:11:53 -0400 Subject: [PATCH 20/44] cleanup+pq fix --- timescale_vector/src/access_method/bq.rs | 23 +++---- .../src/access_method/pq_storage.rs | 60 +++---------------- 2 files changed, 18 insertions(+), 65 deletions(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index 336eff1c..92c447b4 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -298,8 +298,8 @@ pub struct BqSpeedupStorage<'a> { pub index: &'a PgRelation, pub distance_fn: fn(&[f32], &[f32]) -> f32, quantizer: BqQuantizer, - heap_rel: Option<&'a PgRelation>, - heap_attr: Option, + heap_rel: &'a PgRelation, + heap_attr: pgrx::pg_sys::AttrNumber, qv_cache: RefCell, } @@ -313,8 +313,8 @@ impl<'a> BqSpeedupStorage<'a> { index: index, distance_fn: default_distance, quantizer: BqQuantizer::new(), - heap_rel: Some(heap_rel), - heap_attr: Some(heap_attr), + heap_rel: heap_rel, + heap_attr: heap_attr, qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), } } @@ -338,8 +338,8 @@ impl<'a> BqSpeedupStorage<'a> { index: index_relation, distance_fn: default_distance, quantizer: Self::load_quantizer(index_relation, meta_page, stats), - heap_rel: Some(heap_rel), - heap_attr: Some(heap_attr), + heap_rel: heap_rel, + heap_attr: heap_attr, qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), } } @@ -354,8 +354,8 @@ impl<'a> BqSpeedupStorage<'a> { distance_fn: default_distance, //OPT: get rid of clone quantizer: quantizer.clone(), - heap_rel: Some(heap_relation), - heap_attr: Some(get_attribute_number_from_index(index_relation)), + heap_rel: heap_relation, + heap_attr: get_attribute_number_from_index(index_relation), qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), } } @@ -440,12 +440,7 @@ impl<'a> BqSpeedupStorage<'a> { heap_pointer: HeapPointer, stats: &mut T, ) -> TableSlot { - TableSlot::new( - self.heap_rel.unwrap(), - heap_pointer, - self.heap_attr.unwrap(), - stats, - ) + TableSlot::new(self.heap_rel, heap_pointer, self.heap_attr, stats) } } diff --git a/timescale_vector/src/access_method/pq_storage.rs b/timescale_vector/src/access_method/pq_storage.rs index 974ceda7..5df4d700 100644 --- a/timescale_vector/src/access_method/pq_storage.rs +++ b/timescale_vector/src/access_method/pq_storage.rs @@ -23,43 +23,6 @@ use crate::util::{ use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance}; -/*pub struct PqNodeDistanceMeasure<'a> { - storage: &'a PqCompressionStorage<'a>, - table: PqDistanceTable, -} - -impl<'a> PqNodeDistanceMeasure<'a> { - pub unsafe fn with_index_pointer( - 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 table = storage - .quantizer - .get_distance_table_pq_query(node.pq_vector.as_slice()); - - Self { - storage: storage, - table: table, - } - } -} - -impl<'a> NodeDistanceMeasure for PqNodeDistanceMeasure<'a> { - unsafe fn get_distance( - &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.table.distance(node1.pq_vector.as_slice()) - } -}*/ - pub struct PqNodeDistanceMeasure<'a> { storage: &'a PqCompressionStorage<'a>, vector: Vec, @@ -100,8 +63,8 @@ pub struct PqCompressionStorage<'a> { pub index: &'a PgRelation, pub distance_fn: fn(&[f32], &[f32]) -> f32, quantizer: PqQuantizer, - heap_rel: Option<&'a PgRelation>, - heap_attr: Option, + heap_rel: &'a PgRelation, + heap_attr: pgrx::pg_sys::AttrNumber, } impl<'a> PqCompressionStorage<'a> { @@ -114,8 +77,8 @@ impl<'a> PqCompressionStorage<'a> { index: index, distance_fn: default_distance, quantizer: PqQuantizer::new(), - heap_rel: Some(heap_rel), - heap_attr: Some(heap_attr), + heap_rel: heap_rel, + heap_attr: heap_attr, } } @@ -138,8 +101,8 @@ impl<'a> PqCompressionStorage<'a> { index: index_relation, distance_fn: default_distance, quantizer: Self::load_quantizer(index_relation, meta_page, stats), - heap_rel: Some(heap_rel), - heap_attr: Some(heap_attr), + heap_rel: heap_rel, + heap_attr: heap_attr, } } @@ -153,8 +116,8 @@ impl<'a> PqCompressionStorage<'a> { distance_fn: default_distance, //OPT: get rid of clone quantizer: quantizer.clone(), - heap_rel: Some(heap_relation), - heap_attr: Some(get_attribute_number_from_index(heap_relation)), + heap_rel: heap_relation, + heap_attr: get_attribute_number_from_index(index_relation), } } @@ -206,12 +169,7 @@ impl<'a> PqCompressionStorage<'a> { heap_pointer: HeapPointer, stats: &mut T, ) -> TableSlot { - TableSlot::new( - self.heap_rel.unwrap(), - heap_pointer, - self.heap_attr.unwrap(), - stats, - ) + TableSlot::new(self.heap_rel, heap_pointer, self.heap_attr, stats) } } From 7cd2310a1e95283bafc1ee8fd57d242b7919db9a Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Wed, 27 Mar 2024 16:04:46 -0400 Subject: [PATCH 21/44] Make meta page backwards compatible from v1 --- timescale_vector/src/access_method/bq.rs | 4 +- timescale_vector/src/access_method/build.rs | 2 +- timescale_vector/src/access_method/graph.rs | 4 +- .../src/access_method/meta_page.rs | 248 ++++++++++++++---- .../src/access_method/plain_node.rs | 2 - .../src/access_method/pq_quantizer_storage.rs | 3 +- .../src/access_method/pq_storage.rs | 4 +- timescale_vector/src/access_method/scan.rs | 6 +- timescale_vector/src/access_method/vacuum.rs | 2 +- timescale_vector/src/util/mod.rs | 14 +- timescale_vector/src/util/page.rs | 96 ++++++- timescale_vector/src/util/tape.rs | 15 +- .../timescale_vector_derive/src/lib.rs | 30 ++- 13 files changed, 317 insertions(+), 113 deletions(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index 92c447b4..5877c351 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -370,10 +370,10 @@ impl<'a> BqSpeedupStorage<'a> { node.bq_vector.as_slice().to_vec() } - fn write_quantizer_metadata(&self, stats: &mut S) { + fn write_quantizer_metadata(&self, stats: &mut S) { if self.quantizer.use_mean { let index_pointer = unsafe { BqMeans::store(&self.index, &self.quantizer, stats) }; - super::meta_page::MetaPage::update_pq_pointer(&self.index, index_pointer); + super::meta_page::MetaPage::update_pq_pointer(&self.index, index_pointer, stats); } } diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index 68e21d32..d056f4f8 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -120,7 +120,7 @@ pub unsafe extern "C" fn aminsert( } let vec = vec.unwrap(); let heap_pointer = ItemPointer::with_item_pointer_data(*heap_tid); - let mut meta_page = MetaPage::read(&index_relation); + let mut meta_page = MetaPage::fetch(&index_relation); let mut storage = meta_page.get_storage_type(); let mut stats = InsertStats::new(); diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index 95a2dffd..b82a98ee 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -448,8 +448,8 @@ impl<'a> Graph<'a> { ) { if self.meta_page.get_init_ids().is_none() { //TODO probably better set off of centeroids - MetaPage::update_init_ids(index, vec![index_pointer]); - *self.meta_page = MetaPage::read(index); + MetaPage::update_init_ids(index, vec![index_pointer], stats); + *self.meta_page = MetaPage::fetch(index); self.neighbor_store.set_neighbors( storage, diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index ebc0d9ee..b842f817 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -1,23 +1,95 @@ -use pgrx::pg_sys::{BufferGetBlockNumber, Pointer}; +use pgrx::pg_sys::BufferGetBlockNumber; use pgrx::*; +use rkyv::{Archive, Deserialize, Serialize}; +use timescale_vector_derive::{Readable, Writeable}; use crate::access_method::options::TSVIndexOptions; use crate::util::page; use crate::util::*; +use super::stats::StatsNodeModify; use super::storage::StorageType; const TSV_MAGIC_NUMBER: u32 = 768756476; //Magic number, random -const TSV_VERSION: u32 = 1; +const TSV_VERSION: u32 = 2; const GRAPH_SLACK_FACTOR: f64 = 1.3_f64; -/// This is metadata about the entire index. -/// Stored as the first page in the index relation. + +const META_BLOCK_NUMBER: pg_sys::BlockNumber = 0; +const META_HEADER_OFFSET: pgrx::pg_sys::OffsetNumber = 1; +const META_OFFSET: pgrx::pg_sys::OffsetNumber = 2; +/// This is old metadata version for extension versions <=0.0.2. +/// Note it is NOT repr(C) #[derive(Clone)] -pub struct MetaPage { +pub struct MetaPageV1 { + /// random magic number for identifying the index + magic_number: u32, + /// version number for future-proofing + version: u32, + /// number of dimensions in the vector + num_dimensions: u32, + /// max number of outgoing edges a node in the graph can have (R in the papers) + num_neighbors: u32, + search_list_size: u32, + max_alpha: f64, + init_ids_block_number: pg_sys::BlockNumber, + init_ids_offset: pg_sys::OffsetNumber, + use_pq: bool, + pq_vector_length: usize, + pq_block_number: pg_sys::BlockNumber, + pq_block_offset: pg_sys::OffsetNumber, +} + +impl MetaPageV1 { + /// Returns the MetaPage from a page. + /// Should only be called from the very first page in a relation. + unsafe fn page_get_meta(page: pg_sys::Page, buffer: pg_sys::Buffer) -> *mut MetaPageV1 { + assert_eq!(BufferGetBlockNumber(buffer), 0); + let meta_page = ports::PageGetContents(page) as *mut MetaPageV1; + assert_eq!((*meta_page).magic_number, TSV_MAGIC_NUMBER); + assert_eq!((*meta_page).version, 1); + meta_page + } + + pub fn get_new_meta(&self) -> MetaPage { + MetaPage { + magic_number: TSV_MAGIC_NUMBER, + version: TSV_VERSION, + num_dimensions: self.num_dimensions, + num_neighbors: self.num_neighbors, + search_list_size: self.search_list_size, + max_alpha: self.max_alpha, + init_ids_block_number: self.init_ids_block_number, + init_ids_offset: self.init_ids_offset, + use_pq: self.use_pq, + pq_vector_length: self.pq_vector_length, + pq_block_number: self.pq_block_number, + pq_block_offset: self.pq_block_offset, + use_bq: false, + } + } +} + +/// This is metadata header. It contains just the magic number and version number. +/// Stored as the first page (offset 1) in the index relation. +/// The header is separate from the actual metadata to allow for future-proofing. +/// In particular, if the metadata format changes, we can still read the header to check the version. +#[derive(Clone, PartialEq, Archive, Deserialize, Serialize, Readable, Writeable)] +#[archive(check_bytes)] +pub struct MetaPageHeader { /// random magic number for identifying the index magic_number: u32, /// version number for future-proofing version: u32, +} + +/// This is metadata about the entire index. +/// Stored as the first page (offset 2) in the index relation. +#[derive(Clone, PartialEq, Archive, Deserialize, Serialize, Readable, Writeable)] +#[archive(check_bytes)] +pub struct MetaPage { + /// repeat the magic number and version from MetaPageHeader for sanity checks + magic_number: u32, + version: u32, /// number of dimensions in the vector num_dimensions: u32, /// max number of outgoing edges a node in the graph can have (R in the papers) @@ -96,21 +168,6 @@ impl MetaPage { Some(ptr) } - /// Returns the MetaPage from a page. - /// Should only be called from the very first page in a relation. - unsafe fn page_get_meta( - page: pg_sys::Page, - buffer: pg_sys::Buffer, - new: bool, - ) -> *mut MetaPage { - assert_eq!(BufferGetBlockNumber(buffer), 0); - let meta_page = ports::PageGetContents(page) as *mut MetaPage; - if !new { - assert_eq!((*meta_page).magic_number, TSV_MAGIC_NUMBER); - } - meta_page - } - /// Write out a new meta page. /// Has to be done as the first write to a new relation. pub unsafe fn create( @@ -118,62 +175,137 @@ impl MetaPage { num_dimensions: u32, opt: PgBox, ) -> MetaPage { + let meta = MetaPage { + magic_number: TSV_MAGIC_NUMBER, + version: TSV_VERSION, + num_dimensions, + num_neighbors: (*opt).num_neighbors, + search_list_size: (*opt).search_list_size, + max_alpha: (*opt).max_alpha, + init_ids_block_number: 0, + init_ids_offset: 0, + use_pq: (*opt).use_pq, + pq_vector_length: (*opt).pq_vector_length, + pq_block_number: 0, + pq_block_offset: 0, + use_bq: (*opt).use_bq, + }; let page = page::WritablePage::new(index, crate::util::page::PageType::Meta); - let meta = Self::page_get_meta(*page, *(*(page.get_buffer())), true); - (*meta).magic_number = TSV_MAGIC_NUMBER; - (*meta).version = TSV_VERSION; - (*meta).num_dimensions = num_dimensions; - (*meta).num_neighbors = (*opt).num_neighbors; - (*meta).search_list_size = (*opt).search_list_size; - (*meta).max_alpha = (*opt).max_alpha; - (*meta).use_pq = (*opt).use_pq; - (*meta).pq_vector_length = (*opt).pq_vector_length; - (*meta).pq_block_number = 0; - (*meta).pq_block_offset = 0; - (*meta).init_ids_block_number = 0; - (*meta).init_ids_offset = 0; - (*meta).use_bq = (*opt).use_bq; - let header = page.cast::(); - - let meta_end = (meta as Pointer).add(std::mem::size_of::()); - let page_start = (*page) as Pointer; - (*header).pd_lower = meta_end.offset_from(page_start) as _; - - let mp = (*meta).clone(); + meta.write_to_page(page); + meta + } + + unsafe fn write_to_page(&self, mut page: page::WritablePage) { + let header = MetaPageHeader { + magic_number: self.magic_number, + version: self.version, + }; + + assert!(header.magic_number == TSV_MAGIC_NUMBER); + assert!(header.version == TSV_VERSION); + + //serialize the header + let bytes = header.serialize_to_vec(); + let off = page.add_item(&bytes); + assert!(off == META_HEADER_OFFSET); + + //serialize the meta + let bytes = self.serialize_to_vec(); + let off = page.add_item(&bytes); + assert!(off == META_OFFSET); + page.commit(); - mp + } + + unsafe fn overwrite(index: &PgRelation, new_meta: &MetaPage) { + let mut page = page::WritablePage::modify(index, META_BLOCK_NUMBER); + page.reinit(crate::util::page::PageType::Meta); + new_meta.write_to_page(page); + + let page = page::ReadablePage::read(index, META_BLOCK_NUMBER); + let page_type = page.get_type(); + if page_type != crate::util::page::PageType::Meta { + pgrx::error!( + "Problem upgrading meta page: wrong page type: {:?}", + page_type + ); + } + let meta = Self::get_meta_from_page(page); + if meta != *new_meta { + pgrx::error!("Problem upgrading meta page: meta mismatch"); + } } /// Read the meta page for an index - pub fn read(index: &PgRelation) -> MetaPage { + pub fn fetch(index: &PgRelation) -> MetaPage { unsafe { - let page = page::ReadablePage::read(index, 0); - let meta = Self::page_get_meta(*page, *(*(page.get_buffer())), false); - (*meta).clone() + let page = page::ReadablePage::read(index, META_BLOCK_NUMBER); + let page_type = page.get_type(); + if page_type == crate::util::page::PageType::MetaV1 { + let old_meta = MetaPageV1::page_get_meta(*page, *(*(page.get_buffer()))); + let new_meta = (*old_meta).get_new_meta(); + + //release the page + std::mem::drop(page); + + Self::overwrite(index, &new_meta); + return new_meta; + } + Self::get_meta_from_page(page) } } + unsafe fn get_meta_from_page(page: page::ReadablePage) -> MetaPage { + //check the header. In the future, we can use this to check the version + let rb = page.get_item_unchecked(META_HEADER_OFFSET); + let meta = ReadableMetaPageHeader::with_readable_buffer(rb); + let archived = meta.get_archived_node(); + assert!(archived.magic_number == TSV_MAGIC_NUMBER); + assert!(archived.version == TSV_VERSION); + + let page = meta.get_owned_page(); + + //retrieve the MetaPage itself and deserialize it + let rb = page.get_item_unchecked(META_OFFSET); + let meta = ReadableMetaPage::with_readable_buffer(rb); + let archived = meta.get_archived_node(); + assert!(archived.magic_number == TSV_MAGIC_NUMBER); + assert!(archived.version == TSV_VERSION); + + archived.deserialize(&mut rkyv::Infallible).unwrap() + } + /// Change the init ids for an index. - pub fn update_init_ids(index: &PgRelation, init_ids: Vec) { + pub fn update_init_ids( + index: &PgRelation, + init_ids: Vec, + stats: &mut S, + ) { assert_eq!(init_ids.len(), 1); //change this if we support multiple let id = init_ids[0]; unsafe { - let page = page::WritablePage::modify(index, 0); - let meta = Self::page_get_meta(*page, *(*(page.get_buffer())), false); - (*meta).init_ids_block_number = id.block_number; - (*meta).init_ids_offset = id.offset; - page.commit() + let ip = ItemPointer::new(META_BLOCK_NUMBER, META_OFFSET); + let m = MetaPage::modify(index, ip, stats); + let mut archived = m.get_archived_node(); + archived.init_ids_block_number = id.block_number; + archived.init_ids_offset = id.offset; + m.commit() } } - pub fn update_pq_pointer(index: &PgRelation, pq_pointer: IndexPointer) { + pub fn update_pq_pointer( + index: &PgRelation, + pq_pointer: IndexPointer, + stats: &mut S, + ) { unsafe { - let page = page::WritablePage::modify(index, 0); - let meta = Self::page_get_meta(*page, *(*(page.get_buffer())), false); - (*meta).pq_block_number = pq_pointer.block_number; - (*meta).pq_block_offset = pq_pointer.offset; - page.commit() + let ip = ItemPointer::new(META_BLOCK_NUMBER, META_OFFSET); + let m = MetaPage::modify(index, ip, stats); + let mut archived = m.get_archived_node(); + archived.pq_block_number = pq_pointer.block_number; + archived.pq_block_offset = pq_pointer.offset; + m.commit(); } } } diff --git a/timescale_vector/src/access_method/plain_node.rs b/timescale_vector/src/access_method/plain_node.rs index 1e660eb3..37859193 100644 --- a/timescale_vector/src/access_method/plain_node.rs +++ b/timescale_vector/src/access_method/plain_node.rs @@ -8,9 +8,7 @@ use timescale_vector_derive::{Readable, Writeable}; use super::neighbor_with_distance::NeighborWithDistance; use super::pq_quantizer::PqVectorElement; -use super::stats::{StatsNodeModify, StatsNodeRead, StatsNodeWrite}; use super::storage::ArchivedData; -use crate::util::tape::Tape; use crate::util::{ArchivedItemPointer, HeapPointer, ItemPointer, ReadableBuffer, WritableBuffer}; use super::meta_page::MetaPage; diff --git a/timescale_vector/src/access_method/pq_quantizer_storage.rs b/timescale_vector/src/access_method/pq_quantizer_storage.rs index f76668a3..b00a6318 100644 --- a/timescale_vector/src/access_method/pq_quantizer_storage.rs +++ b/timescale_vector/src/access_method/pq_quantizer_storage.rs @@ -1,5 +1,4 @@ use std::mem::size_of; -use std::pin::Pin; use ndarray::Array3; use pgrx::pg_sys::BLCKSZ; @@ -13,7 +12,7 @@ use crate::util::page::PageType; use crate::util::tape::Tape; use crate::util::{IndexPointer, ItemPointer, ReadableBuffer, WritableBuffer}; -use super::stats::{StatsNodeModify, StatsNodeRead, StatsNodeWrite}; +use super::stats::{StatsNodeRead, StatsNodeWrite}; #[derive(Archive, Deserialize, Serialize, Readable, Writeable)] #[archive(check_bytes)] diff --git a/timescale_vector/src/access_method/pq_storage.rs b/timescale_vector/src/access_method/pq_storage.rs index 5df4d700..5bbd6918 100644 --- a/timescale_vector/src/access_method/pq_storage.rs +++ b/timescale_vector/src/access_method/pq_storage.rs @@ -121,10 +121,10 @@ impl<'a> PqCompressionStorage<'a> { } } - fn write_quantizer_metadata(&self, stats: &mut S) { + fn write_quantizer_metadata(&self, stats: &mut S) { let pq = self.quantizer.must_get_pq(); let index_pointer: IndexPointer = unsafe { write_pq(pq, &self.index, stats) }; - super::meta_page::MetaPage::update_pq_pointer(&self.index, index_pointer); + super::meta_page::MetaPage::update_pq_pointer(&self.index, index_pointer, stats); } fn visit_lsn_internal( diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index c6d3a126..3ace7d94 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -53,7 +53,7 @@ impl TSVScanState { query: PgVector, search_list_size: usize, ) { - let meta_page = MetaPage::read(&index); + let meta_page = MetaPage::fetch(&index); let storage = meta_page.get_storage_type(); let store_type = match storage { @@ -133,7 +133,7 @@ impl TSVResponseIterator { _meta_page: MetaPage, quantizer_stats: QuantizerStats, ) -> Self { - let mut meta_page = MetaPage::read(&index); + let mut meta_page = MetaPage::fetch(&index); let graph = Graph::new(GraphNeighborStore::Disk, &mut meta_page); let lsr = graph.greedy_search_streaming_init(query, search_list_size, storage); @@ -272,7 +272,7 @@ pub extern "C" fn amrescan( let mut scan: PgBox = unsafe { PgBox::from_pg(scan) }; let indexrel = unsafe { PgRelation::from_pg(scan.indexRelation) }; let heaprel = unsafe { PgRelation::from_pg(scan.heapRelation) }; - let meta_page = MetaPage::read(&indexrel); + let meta_page = MetaPage::fetch(&indexrel); let _storage = meta_page.get_storage_type(); if nkeys > 0 { diff --git a/timescale_vector/src/access_method/vacuum.rs b/timescale_vector/src/access_method/vacuum.rs index 20049aa9..cf9260d7 100644 --- a/timescale_vector/src/access_method/vacuum.rs +++ b/timescale_vector/src/access_method/vacuum.rs @@ -40,7 +40,7 @@ pub extern "C" fn ambulkdelete( ) }; - let meta_page = MetaPage::read(&index_relation); + let meta_page = MetaPage::fetch(&index_relation); let storage = meta_page.get_storage_type(); match storage { StorageType::BqSpeedup => { diff --git a/timescale_vector/src/util/mod.rs b/timescale_vector/src/util/mod.rs index 0f23fc38..dc72524c 100644 --- a/timescale_vector/src/util/mod.rs +++ b/timescale_vector/src/util/mod.rs @@ -36,6 +36,10 @@ impl<'a> ReadableBuffer<'a> { pub fn get_data_slice(&self) -> &[u8] { unsafe { std::slice::from_raw_parts(self.ptr, self.len) } } + + pub fn get_owned_page(self) -> ReadablePage<'a> { + self._page + } } pub struct WritableBuffer<'a> { @@ -84,15 +88,9 @@ impl ItemPointer { pub unsafe fn read_bytes(self, index: &PgRelation) -> ReadableBuffer { let page = ReadablePage::read(index, self.block_number); - let item_id = PageGetItemId(*page, self.offset); - let item = PageGetItem(*page, item_id) as *mut u8; - let len = (*item_id).lp_len(); - ReadableBuffer { - _page: page, - ptr: item, - len: len as _, - } + page.get_item_unchecked(self.offset) } + pub unsafe fn modify_bytes(self, index: &PgRelation) -> WritableBuffer { let page = WritablePage::modify(index, self.block_number); let item_id = PageGetItemId(*page, self.offset); diff --git a/timescale_vector/src/util/page.rs b/timescale_vector/src/util/page.rs index 823a0259..92284931 100644 --- a/timescale_vector/src/util/page.rs +++ b/timescale_vector/src/util/page.rs @@ -3,12 +3,16 @@ use pg_sys::Page; use pgrx::{ - pg_sys::{BlockNumber, BufferGetPage}, + pg_sys::{BlockNumber, BufferGetPage, OffsetNumber, BLCKSZ}, *, }; use std::ops::Deref; -use super::buffer::{LockedBufferExclusive, LockedBufferShare}; +use super::{ + buffer::{LockedBufferExclusive, LockedBufferShare}, + ports::{PageGetItem, PageGetItemId}, + ReadableBuffer, +}; pub struct WritablePage<'a> { buffer: LockedBufferExclusive<'a>, page: Page, @@ -20,25 +24,27 @@ pub const TSV_PAGE_ID: u16 = 0xAE24; /* magic number, generated randomly */ /// PageType identifies different types of pages in our index. /// The layout of any one type should be consistent -#[derive(Clone, Copy, PartialEq, Eq)] +#[derive(Clone, Copy, PartialEq, Eq, Debug)] pub enum PageType { - Meta = 0, + MetaV1 = 0, Node = 1, PqQuantizerDef = 2, PqQuantizerVector = 3, BqMeans = 4, BqNode = 5, + Meta = 6, } impl PageType { fn from_u8(value: u8) -> Self { match value { - 0 => PageType::Meta, + 0 => PageType::MetaV1, 1 => PageType::Node, 2 => PageType::PqQuantizerDef, 3 => PageType::PqQuantizerVector, 4 => PageType::BqMeans, 5 => PageType::BqNode, + 6 => PageType::Meta, _ => panic!("Unknown PageType number {}", value), } } @@ -96,18 +102,25 @@ impl<'a> WritablePage<'a> { let state = pg_sys::GenericXLogStart(index.as_ptr()); //TODO do we need a GENERIC_XLOG_FULL_IMAGE option? let page = pg_sys::GenericXLogRegisterBuffer(state, *buffer, 0); - pg_sys::PageInit( - page, - pg_sys::BLCKSZ as usize, - std::mem::size_of::(), - ); - *TsvPageOpaqueData::with_page(page) = TsvPageOpaqueData::new(page_type); - Self { + let mut new = Self { buffer: buffer, page: page, state: state, committed: false, - } + }; + new.reinit(page_type); + new + } + } + + pub fn reinit(&mut self, page_type: PageType) { + unsafe { + pg_sys::PageInit( + self.page, + pg_sys::BLCKSZ as usize, + std::mem::size_of::(), + ); + *TsvPageOpaqueData::with_page(self.page) = TsvPageOpaqueData::new(page_type); } } @@ -116,6 +129,28 @@ impl<'a> WritablePage<'a> { Self::modify_with_buffer(index, buffer) } + pub fn add_item(&mut self, data: &[u8]) -> OffsetNumber { + let size = data.len(); + assert!(self.get_free_space() >= size); + unsafe { self.add_item_unchecked(data) } + } + + pub unsafe fn add_item_unchecked(&mut self, data: &[u8]) -> OffsetNumber { + let size = data.len(); + assert!(size < BLCKSZ as usize); + + let offset_number = pg_sys::PageAddItemExtended( + self.page, + data.as_ptr() as _, + size, + pg_sys::InvalidOffsetNumber, + 0, + ); + + assert!(offset_number != pg_sys::InvalidOffsetNumber); + offset_number + } + /// get a writable page for cleanup(vacuum) operations. pub unsafe fn cleanup(index: &'a PgRelation, block: BlockNumber) -> Self { let buffer = LockedBufferExclusive::read_for_cleanup(index, block); @@ -159,6 +194,16 @@ impl<'a> WritablePage<'a> { PageType::from_u8((*opaque_data).page_type) } } + + pub fn set_types(&self, new: PageType) { + unsafe { + let opaque_data = + //safe to do because self.page was already verified during construction + TsvPageOpaqueData::with_page(self.page); + + (*opaque_data).page_type = new as u8; + } + } /// commit saves all the changes to the page. /// Note that this will consume the page and make it unusable after the call. pub fn commit(mut self) { @@ -204,9 +249,34 @@ impl<'a> ReadablePage<'a> { } } + pub fn get_type(&self) -> PageType { + unsafe { + let opaque_data = + //safe to do because self.page was already verified during construction + TsvPageOpaqueData::with_page(self.page); + + PageType::from_u8((*opaque_data).page_type) + } + } + pub fn get_buffer(&self) -> &LockedBufferShare { &self.buffer } + + // Safety: unsafe because no verification of the offset is done. + pub unsafe fn get_item_unchecked( + self, + offset: pgrx::pg_sys::OffsetNumber, + ) -> ReadableBuffer<'a> { + let item_id = PageGetItemId(self.page, offset); + let item = PageGetItem(self.page, item_id) as *mut u8; + let len = (*item_id).lp_len(); + ReadableBuffer { + _page: self, + ptr: item, + len: len as _, + } + } } impl<'a> Deref for ReadablePage<'a> { diff --git a/timescale_vector/src/util/tape.rs b/timescale_vector/src/util/tape.rs index 961dffda..f028f173 100644 --- a/timescale_vector/src/util/tape.rs +++ b/timescale_vector/src/util/tape.rs @@ -41,18 +41,11 @@ impl<'a> Tape<'a> { panic!("Not enough free space on new page"); } } - let offset_number = pg_sys::PageAddItemExtended( - *current_page, - data.as_ptr() as _, - size, - pg_sys::InvalidOffsetNumber, - 0, - ); - - assert!(offset_number != pg_sys::InvalidOffsetNumber); - let index_pointer = super::ItemPointer::with_page(¤t_page, offset_number); + let offset_number = current_page.add_item_unchecked(data); + + let item_pointer = super::ItemPointer::with_page(¤t_page, offset_number); current_page.commit(); - index_pointer + item_pointer } pub fn close(self) { diff --git a/timescale_vector/timescale_vector_derive/src/lib.rs b/timescale_vector/timescale_vector_derive/src/lib.rs index 2843b75a..b82485a9 100644 --- a/timescale_vector/timescale_vector_derive/src/lib.rs +++ b/timescale_vector/timescale_vector_derive/src/lib.rs @@ -27,19 +27,27 @@ fn impl_readable_macro(ast: &syn::DeriveInput) -> TokenStream { } impl<'a> #readable_name<'a> { + pub fn with_readable_buffer(rb: ReadableBuffer<'a>) -> Self { + Self { _rb: rb } + } + pub fn get_archived_node(&self) -> & #archived_name { // checking the code here is expensive during build, so skip it. // TODO: should we check the data during queries? //rkyv::check_archived_root::(self._rb.get_data_slice()).unwrap() unsafe { rkyv::archived_root::<#name>(self._rb.get_data_slice()) } } + + pub fn get_owned_page(self) -> crate::util::page::ReadablePage<'a> { + self._rb.get_owned_page() + } } impl #name { - pub unsafe fn read<'a, 'b, S: StatsNodeRead>(index: &'a PgRelation, index_pointer: ItemPointer, stats: &'b mut S) -> #readable_name<'a> { + 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> { let rb = index_pointer.read_bytes(index); stats.record_read(); - #readable_name { _rb: rb } + #readable_name::with_readable_buffer(rb) } } }; @@ -51,19 +59,20 @@ fn impl_writeable_macro(ast: &syn::DeriveInput) -> TokenStream { let writeable_name = format_ident!("Writable{}", name); let archived_name = format_ident!("Archived{}", name); let gen = quote! { + pub struct #writeable_name<'a> { wb: WritableBuffer<'a>, } impl #archived_name { - pub fn with_data(data: &mut [u8]) -> Pin<&mut #archived_name> { - let pinned_bytes = Pin::new(data); + pub fn with_data(data: &mut [u8]) -> std::pin::Pin<&mut #archived_name> { + let pinned_bytes = std::pin::Pin::new(data); unsafe { rkyv::archived_root_mut::<#name>(pinned_bytes) } } } impl<'a> #writeable_name<'a> { - pub fn get_archived_node(&self) -> Pin<&mut #archived_name> { + pub fn get_archived_node(&self) -> std::pin::Pin<&mut #archived_name> { #archived_name::with_data(self.wb.get_data_slice()) } @@ -73,18 +82,23 @@ fn impl_writeable_macro(ast: &syn::DeriveInput) -> TokenStream { } impl #name { - pub unsafe fn modify<'a, 'b, S: StatsNodeModify>(index: &'a PgRelation, index_pointer: ItemPointer, stats: &'b mut S) -> #writeable_name<'a> { + 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> { let wb = index_pointer.modify_bytes(index); stats.record_modify(); #writeable_name { wb: wb } } - pub fn write(&self, tape: &mut Tape, stats: &mut S) -> ItemPointer { + pub fn write(&self, tape: &mut crate::util::tape::Tape, stats: &mut S) -> ItemPointer { //TODO 256 probably too small - let bytes = rkyv::to_bytes::<_, 256>(self).unwrap(); + let bytes = self.serialize_to_vec(); stats.record_write(); unsafe { tape.write(&bytes) } } + + pub fn serialize_to_vec(&self) -> rkyv::util::AlignedVec { + //TODO 256 probably too small + rkyv::to_bytes::<_, 256>(self).unwrap() + } } }; gen.into() From 5089a9566adf53b291680df576b85be14a0f83d3 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Thu, 28 Mar 2024 21:11:12 -0400 Subject: [PATCH 22/44] Switch to versioned .so and support upgrades --- timescale_vector/Cargo.toml | 5 +- .../timescale_vector--0.0.2--0.0.3-dev.sql | 48 +++++ timescale_vector/src/access_method/bq.rs | 10 +- timescale_vector/src/access_method/build.rs | 8 +- .../src/access_method/meta_page.rs | 60 ++++-- timescale_vector/src/access_method/mod.rs | 42 +++- .../src/access_method/plain_storage.rs | 22 +- .../src/access_method/pq_storage.rs | 9 +- timescale_vector/src/access_method/scan.rs | 36 +++- .../src/access_method/upgrade_test.rs | 191 ++++++++++++++++++ timescale_vector/timescale_vector.control | 2 +- 11 files changed, 388 insertions(+), 45 deletions(-) create mode 100644 timescale_vector/sql/timescale_vector--0.0.2--0.0.3-dev.sql create mode 100644 timescale_vector/src/access_method/upgrade_test.rs diff --git a/timescale_vector/Cargo.toml b/timescale_vector/Cargo.toml index a2b9d622..4c724aaf 100644 --- a/timescale_vector/Cargo.toml +++ b/timescale_vector/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "timescale_vector" -version = "0.0.2" +version = "0.0.3-dev" edition = "2021" [lib] @@ -27,10 +27,13 @@ rand_core = "0.6" rand_xorshift = "0.3" rayon = "1" timescale_vector_derive = { path = "timescale_vector_derive" } +semver = "1.0.22" [dev-dependencies] pgrx-tests = "=0.11.2" +pgrx-pg-config = "=0.11.2" criterion = "0.5.1" +tempfile = "3.3.0" [profile.dev] panic = "unwind" diff --git a/timescale_vector/sql/timescale_vector--0.0.2--0.0.3-dev.sql b/timescale_vector/sql/timescale_vector--0.0.2--0.0.3-dev.sql new file mode 100644 index 00000000..bcafdb51 --- /dev/null +++ b/timescale_vector/sql/timescale_vector--0.0.2--0.0.3-dev.sql @@ -0,0 +1,48 @@ +/* +This file is auto generated by pgrx. + +The ordering of items is not stable, it is driven by a dependency graph. +*/ + +-- src/access_method/mod.rs:48 +-- timescale_vector::access_method::amhandler + + CREATE OR REPLACE FUNCTION tsv_amhandler(internal) RETURNS index_am_handler PARALLEL SAFE IMMUTABLE STRICT COST 0.0001 LANGUAGE c AS '$libdir/timescale_vector-0.0.3-dev', 'amhandler_wrapper'; + + DO $$ + DECLARE + c int; + BEGIN + SELECT count(*) + INTO c + FROM pg_catalog.pg_am a + WHERE a.amname = 'tsv'; + + IF c = 0 THEN + CREATE ACCESS METHOD tsv TYPE INDEX HANDLER tsv_amhandler; + END IF; + END; + $$; + + + + +-- src/access_method/mod.rs:91 + +DO $$ +DECLARE + c int; +BEGIN + SELECT count(*) + INTO c + FROM pg_catalog.pg_opclass c + WHERE c.opcname = 'vector_cosine_ops' + AND c.opcmethod = (SELECT oid FROM pg_catalog.pg_am am WHERE am.amname = 'tsv'); + + IF c = 0 THEN + CREATE OPERATOR CLASS vector_cosine_ops DEFAULT + FOR TYPE vector USING tsv AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops; + END IF; +END; +$$; \ No newline at end of file diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index 5877c351..a1fa6874 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -1,5 +1,5 @@ use super::{ - distance::{distance_cosine as default_distance, distance_xor_optimized}, + distance::distance_xor_optimized, graph::{ListSearchNeighbor, ListSearchResult}, graph_neighbor_store::GraphNeighborStore, pg_vector::PgVector, @@ -308,10 +308,11 @@ impl<'a> BqSpeedupStorage<'a> { index: &'a PgRelation, heap_rel: &'a PgRelation, heap_attr: pgrx::pg_sys::AttrNumber, + distance_fn: fn(&[f32], &[f32]) -> f32, ) -> BqSpeedupStorage<'a> { Self { index: index, - distance_fn: default_distance, + distance_fn: distance_fn, quantizer: BqQuantizer::new(), heap_rel: heap_rel, heap_attr: heap_attr, @@ -336,7 +337,7 @@ impl<'a> BqSpeedupStorage<'a> { ) -> BqSpeedupStorage<'a> { Self { index: index_relation, - distance_fn: default_distance, + distance_fn: meta_page.get_distance_function(), quantizer: Self::load_quantizer(index_relation, meta_page, stats), heap_rel: heap_rel, heap_attr: heap_attr, @@ -348,10 +349,11 @@ impl<'a> BqSpeedupStorage<'a> { index_relation: &'a PgRelation, heap_relation: &'a PgRelation, quantizer: &BqQuantizer, + distance_fn: fn(&[f32], &[f32]) -> f32, ) -> BqSpeedupStorage<'a> { Self { index: index_relation, - distance_fn: default_distance, + distance_fn: distance_fn, //OPT: get rid of clone quantizer: quantizer.clone(), heap_rel: heap_relation, diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index d056f4f8..788616cb 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -126,7 +126,8 @@ pub unsafe extern "C" fn aminsert( let mut stats = InsertStats::new(); match &mut storage { StorageType::Plain => { - let plain = PlainStorage::load_for_insert(&index_relation); + let plain = + PlainStorage::load_for_insert(&index_relation, meta_page.get_distance_function()); insert_storage( &plain, &index_relation, @@ -221,7 +222,8 @@ fn do_heap_scan<'a>( let mut write_stats = WriteStats::new(); match storage { StorageType::Plain => { - let mut plain = PlainStorage::new_for_build(index_relation); + let mut plain = + PlainStorage::new_for_build(index_relation, meta_page.get_distance_function()); plain.start_training(&meta_page); let page_type = PlainStorage::page_type(); let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); @@ -244,6 +246,7 @@ fn do_heap_scan<'a>( index_relation, heap_relation, get_attribute_number(index_info), + meta_page.get_distance_function(), ); pq.start_training(&meta_page); unsafe { @@ -278,6 +281,7 @@ fn do_heap_scan<'a>( index_relation, heap_relation, get_attribute_number(index_info), + meta_page.get_distance_function(), ); bq.start_training(&meta_page); unsafe { diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index b842f817..c878d3a5 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -1,12 +1,14 @@ use pgrx::pg_sys::BufferGetBlockNumber; use pgrx::*; use rkyv::{Archive, Deserialize, Serialize}; +use semver::Version; use timescale_vector_derive::{Readable, Writeable}; use crate::access_method::options::TSVIndexOptions; use crate::util::page; use crate::util::*; +use super::distance; use super::stats::StatsNodeModify; use super::storage::StorageType; @@ -54,6 +56,8 @@ impl MetaPageV1 { MetaPage { magic_number: TSV_MAGIC_NUMBER, version: TSV_VERSION, + extension_version_when_built: "0.0.2".to_string(), + distance_type: DistanceType::L2 as u16, num_dimensions: self.num_dimensions, num_neighbors: self.num_neighbors, search_list_size: self.search_list_size, @@ -82,6 +86,21 @@ pub struct MetaPageHeader { version: u32, } +pub enum DistanceType { + Cosine = 0, + L2 = 1, +} + +impl DistanceType { + fn from_u16(value: u16) -> Self { + match value { + 0 => DistanceType::Cosine, + 1 => DistanceType::L2, + _ => panic!("Unknown DistanceType number {}", value), + } + } +} + /// This is metadata about the entire index. /// Stored as the first page (offset 2) in the index relation. #[derive(Clone, PartialEq, Archive, Deserialize, Serialize, Readable, Writeable)] @@ -90,6 +109,8 @@ pub struct MetaPage { /// repeat the magic number and version from MetaPageHeader for sanity checks magic_number: u32, version: u32, + extension_version_when_built: String, + distance_type: u16, /// number of dimensions in the vector num_dimensions: u32, /// max number of outgoing edges a node in the graph can have (R in the papers) @@ -134,6 +155,13 @@ impl MetaPage { self.use_pq } + pub fn get_distance_function(&self) -> fn(&[f32], &[f32]) -> f32 { + match DistanceType::from_u16(self.distance_type) { + DistanceType::Cosine => distance::distance_cosine, + DistanceType::L2 => distance::distance_l2, + } + } + pub fn get_storage_type(&self) -> StorageType { if self.get_use_pq() { StorageType::PqCompression @@ -175,9 +203,13 @@ impl MetaPage { num_dimensions: u32, opt: PgBox, ) -> MetaPage { + let version = Version::parse(env!("CARGO_PKG_VERSION")).unwrap(); + let meta = MetaPage { magic_number: TSV_MAGIC_NUMBER, version: TSV_VERSION, + extension_version_when_built: version.to_string(), + distance_type: DistanceType::Cosine as u16, num_dimensions, num_neighbors: (*opt).num_neighbors, search_list_size: (*opt).search_list_size, @@ -284,14 +316,14 @@ impl MetaPage { assert_eq!(init_ids.len(), 1); //change this if we support multiple let id = init_ids[0]; + let mut meta = Self::fetch(index); + meta.init_ids_block_number = id.block_number; + meta.init_ids_offset = id.offset; + unsafe { - let ip = ItemPointer::new(META_BLOCK_NUMBER, META_OFFSET); - let m = MetaPage::modify(index, ip, stats); - let mut archived = m.get_archived_node(); - archived.init_ids_block_number = id.block_number; - archived.init_ids_offset = id.offset; - m.commit() - } + Self::overwrite(index, &meta); + stats.record_modify(); + }; } pub fn update_pq_pointer( @@ -299,13 +331,13 @@ impl MetaPage { pq_pointer: IndexPointer, stats: &mut S, ) { + let mut meta = Self::fetch(index); + meta.pq_block_number = pq_pointer.block_number; + meta.pq_block_offset = pq_pointer.offset; + unsafe { - let ip = ItemPointer::new(META_BLOCK_NUMBER, META_OFFSET); - let m = MetaPage::modify(index, ip, stats); - let mut archived = m.get_archived_node(); - archived.pq_block_number = pq_pointer.block_number; - archived.pq_block_offset = pq_pointer.offset; - m.commit(); - } + Self::overwrite(index, &meta); + stats.record_modify(); + }; } } diff --git a/timescale_vector/src/access_method/mod.rs b/timescale_vector/src/access_method/mod.rs index 1f9620ab..853a871f 100644 --- a/timescale_vector/src/access_method/mod.rs +++ b/timescale_vector/src/access_method/mod.rs @@ -15,6 +15,7 @@ mod scan; pub mod stats; mod storage; mod storage_common; +mod upgrade_test; mod vacuum; extern crate blas_src; @@ -28,8 +29,22 @@ mod pq_quantizer_storage; mod pq_storage; #[pg_extern(sql = " - CREATE OR REPLACE FUNCTION tsv_amhandler(internal) RETURNS index_am_handler PARALLEL SAFE IMMUTABLE STRICT COST 0.0001 LANGUAGE c AS 'MODULE_PATHNAME', '@FUNCTION_NAME@'; - CREATE ACCESS METHOD tsv TYPE INDEX HANDLER tsv_amhandler; + CREATE OR REPLACE FUNCTION tsv_amhandler(internal) RETURNS index_am_handler PARALLEL SAFE IMMUTABLE STRICT COST 0.0001 LANGUAGE c AS '@MODULE_PATHNAME@', '@FUNCTION_NAME@'; + + DO $$ + DECLARE + c int; + BEGIN + SELECT count(*) + INTO c + FROM pg_catalog.pg_am a + WHERE a.amname = 'tsv'; + + IF c = 0 THEN + CREATE ACCESS METHOD tsv TYPE INDEX HANDLER tsv_amhandler; + END IF; + END; + $$; ")] fn amhandler(_fcinfo: pg_sys::FunctionCallInfo) -> PgBox { let mut amroutine = @@ -73,12 +88,27 @@ fn amhandler(_fcinfo: pg_sys::FunctionCallInfo) -> PgBox amroutine.into_pg_boxed() } +// This SQL is made idempotent so that we can use the same script for the installation and the upgrade. extension_sql!( r#" -CREATE OPERATOR CLASS vector_cosine_ops DEFAULT -FOR TYPE vector USING tsv AS - OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops -; +DO $$ +DECLARE + c int; +BEGIN + SELECT count(*) + INTO c + FROM pg_catalog.pg_opclass c + WHERE c.opcname = 'vector_cosine_ops' + AND c.opcmethod = (SELECT oid FROM pg_catalog.pg_am am WHERE am.amname = 'tsv'); + + IF c = 0 THEN + CREATE OPERATOR CLASS vector_cosine_ops DEFAULT + FOR TYPE vector USING tsv AS + OPERATOR 1 <=> (vector, vector) FOR ORDER BY float_ops; + END IF; +END; +$$; + "#, name = "tsv_ops_operator" ); diff --git a/timescale_vector/src/access_method/plain_storage.rs b/timescale_vector/src/access_method/plain_storage.rs index b7a1ec56..423cc8e1 100644 --- a/timescale_vector/src/access_method/plain_storage.rs +++ b/timescale_vector/src/access_method/plain_storage.rs @@ -1,5 +1,4 @@ use super::{ - distance::distance_cosine as default_distance, graph::{ListSearchNeighbor, ListSearchResult}, graph_neighbor_store::GraphNeighborStore, pg_vector::PgVector, @@ -23,24 +22,33 @@ pub struct PlainStorage<'a> { } impl<'a> PlainStorage<'a> { - pub fn new_for_build(index: &'a PgRelation) -> PlainStorage<'a> { + pub fn new_for_build( + index: &'a PgRelation, + distance_fn: fn(&[f32], &[f32]) -> f32, + ) -> PlainStorage<'a> { Self { index: index, - distance_fn: default_distance, + distance_fn: distance_fn, } } - pub fn load_for_insert(index_relation: &'a PgRelation) -> PlainStorage<'a> { + pub fn load_for_insert( + index_relation: &'a PgRelation, + distance_fn: fn(&[f32], &[f32]) -> f32, + ) -> PlainStorage<'a> { Self { index: index_relation, - distance_fn: default_distance, + distance_fn: distance_fn, } } - pub fn load_for_search(index_relation: &'a PgRelation) -> PlainStorage<'a> { + pub fn load_for_search( + index_relation: &'a PgRelation, + distance_fn: fn(&[f32], &[f32]) -> f32, + ) -> PlainStorage<'a> { Self { index: index_relation, - distance_fn: default_distance, + distance_fn: distance_fn, } } } diff --git a/timescale_vector/src/access_method/pq_storage.rs b/timescale_vector/src/access_method/pq_storage.rs index 5bbd6918..d3fae389 100644 --- a/timescale_vector/src/access_method/pq_storage.rs +++ b/timescale_vector/src/access_method/pq_storage.rs @@ -1,5 +1,4 @@ use super::{ - distance::distance_cosine as default_distance, graph::{ListSearchNeighbor, ListSearchResult}, graph_neighbor_store::GraphNeighborStore, pg_vector::PgVector, @@ -72,10 +71,11 @@ impl<'a> PqCompressionStorage<'a> { index: &'a PgRelation, heap_rel: &'a PgRelation, heap_attr: pgrx::pg_sys::AttrNumber, + distance_fn: fn(&[f32], &[f32]) -> f32, ) -> PqCompressionStorage<'a> { Self { index: index, - distance_fn: default_distance, + distance_fn, quantizer: PqQuantizer::new(), heap_rel: heap_rel, heap_attr: heap_attr, @@ -99,7 +99,7 @@ impl<'a> PqCompressionStorage<'a> { ) -> PqCompressionStorage<'a> { Self { index: index_relation, - distance_fn: default_distance, + distance_fn: meta_page.get_distance_function(), quantizer: Self::load_quantizer(index_relation, meta_page, stats), heap_rel: heap_rel, heap_attr: heap_attr, @@ -110,10 +110,11 @@ impl<'a> PqCompressionStorage<'a> { index_relation: &'a PgRelation, heap_relation: &'a PgRelation, quantizer: &PqQuantizer, + distance_fn: fn(&[f32], &[f32]) -> f32, ) -> PqCompressionStorage<'a> { Self { index: index_relation, - distance_fn: default_distance, + distance_fn, //OPT: get rid of clone quantizer: quantizer.clone(), heap_rel: heap_relation, diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 3ace7d94..a3fcaf9b 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -37,12 +37,14 @@ enum StorageState { /* no lifetime usage here. */ struct TSVScanState { storage: *mut StorageState, + distance_fn: Option f32>, } impl TSVScanState { fn new() -> Self { Self { storage: std::ptr::null_mut(), + distance_fn: None, } } @@ -55,11 +57,12 @@ impl TSVScanState { ) { let meta_page = MetaPage::fetch(&index); let storage = meta_page.get_storage_type(); + let distance = meta_page.get_distance_function(); let store_type = match storage { StorageType::Plain => { let stats = QuantizerStats::new(); - let bq = PlainStorage::load_for_search(index); + let bq = PlainStorage::load_for_search(index, meta_page.get_distance_function()); let it = TSVResponseIterator::new(&bq, index, query, search_list_size, meta_page, stats); StorageState::Plain(it) @@ -67,7 +70,12 @@ impl TSVScanState { StorageType::PqCompression => { let mut stats = QuantizerStats::new(); let quantizer = PqQuantizer::load(index, &meta_page, &mut stats); - let pq = PqCompressionStorage::load_for_search(index, heap, &quantizer); + let pq = PqCompressionStorage::load_for_search( + index, + heap, + &quantizer, + meta_page.get_distance_function(), + ); let it = TSVResponseIterator::new(&pq, index, query, search_list_size, meta_page, stats); StorageState::PqCompression(quantizer, it) @@ -75,7 +83,12 @@ impl TSVScanState { StorageType::BqSpeedup => { let mut stats = QuantizerStats::new(); let quantizer = unsafe { BqMeans::load(index, &meta_page, &mut stats) }; - let bq = BqSpeedupStorage::load_for_search(index, heap, &quantizer); + let bq = BqSpeedupStorage::load_for_search( + index, + heap, + &quantizer, + meta_page.get_distance_function(), + ); let it = TSVResponseIterator::new(&bq, index, query, search_list_size, meta_page, stats); StorageState::BqSpeedup(quantizer, it) @@ -83,6 +96,7 @@ impl TSVScanState { }; self.storage = PgMemoryContexts::CurrentMemoryContext.leak_and_drop_on_delete(store_type); + self.distance_fn = Some(distance); } } @@ -318,17 +332,27 @@ pub extern "C" fn amgettuple( let mut storage = unsafe { state.storage.as_mut() }.expect("no storage in state"); match &mut storage { StorageState::BqSpeedup(quantizer, iter) => { - let bq = BqSpeedupStorage::load_for_search(&indexrel, &heaprel, quantizer); + let bq = BqSpeedupStorage::load_for_search( + &indexrel, + &heaprel, + quantizer, + state.distance_fn.unwrap(), + ); let next = iter.next_with_resort(&indexrel, &bq); get_tuple(next, scan) } StorageState::PqCompression(quantizer, iter) => { - let pq = PqCompressionStorage::load_for_search(&indexrel, &heaprel, quantizer); + let pq = PqCompressionStorage::load_for_search( + &indexrel, + &heaprel, + quantizer, + state.distance_fn.unwrap(), + ); let next = iter.next_with_resort(&indexrel, &pq); get_tuple(next, scan) } StorageState::Plain(iter) => { - let storage = PlainStorage::load_for_search(&indexrel); + let storage = PlainStorage::load_for_search(&indexrel, state.distance_fn.unwrap()); let next = iter.next(&indexrel, &storage); get_tuple(next, scan) } diff --git a/timescale_vector/src/access_method/upgrade_test.rs b/timescale_vector/src/access_method/upgrade_test.rs new file mode 100644 index 00000000..16ee3c54 --- /dev/null +++ b/timescale_vector/src/access_method/upgrade_test.rs @@ -0,0 +1,191 @@ +#[cfg(any(test, feature = "pg_test"))] +#[pgrx::pg_schema] +pub mod tests { + use pgrx::*; + use std::{fs, path::Path, process::Stdio}; + + fn copy_dir_all(src: impl AsRef, dst: impl AsRef) -> std::io::Result<()> { + fs::create_dir_all(&dst)?; + for entry in fs::read_dir(src)? { + let entry = entry?; + let ty = entry.file_type()?; + if ty.is_dir() { + copy_dir_all(entry.path(), dst.as_ref().join(entry.file_name()))?; + } else { + fs::copy(entry.path(), dst.as_ref().join(entry.file_name()))?; + } + } + Ok(()) + } + + #[test] + #[ignore] + ///This function is only a mock to bring up the test framewokr in test_delete_vacuum + fn test_upgrade() { + pgrx_tests::run_test( + "test_delete_mock_fn", + None, + crate::pg_test::postgresql_conf_options(), + ) + .unwrap(); + + let (mut client, _) = pgrx_tests::client().unwrap(); + + client + .execute( + &format!("DROP EXTENSION IF EXISTS timescale_vector CASCADE;"), + &[], + ) + .unwrap(); + + let current_file = file!(); + + // Convert the file path to an absolute path + let current_dir = std::env::current_dir().unwrap(); + let mut absolute_path = std::path::Path::new(¤t_dir).join(current_file); + absolute_path = absolute_path.ancestors().nth(4).unwrap().to_path_buf(); + + let temp_dir = tempfile::tempdir().unwrap(); + let temp_path = temp_dir.path(); + + copy_dir_all(absolute_path.clone(), temp_dir.path()).unwrap(); + + let pgrx = pgrx_pg_config::Pgrx::from_config().unwrap(); + let pg_version = pg_sys::get_pg_major_version_num(); + let pg_config = pgrx.get(&format!("pg{}", pg_version)).unwrap(); + + let version = "0.0.2"; + let res = std::process::Command::new("git") + .current_dir(temp_path) + .arg("checkout") + .arg("-f") + .arg(version) + .output() + .unwrap(); + assert!( + res.status.success(), + "failed: {:?} {:?} {:?}", + res, + absolute_path, + temp_dir.path() + ); + + // use latest pgrx + let res = std::process::Command::new("cargo") + .current_dir(temp_path.join("timescale_vector")) + .args(["rm", "pgrx"]) + .stdout(Stdio::inherit()) + .stderr(Stdio::piped()) + .output() + .unwrap(); + assert!(res.status.success(), "failed: {:?}", res); + + let res = std::process::Command::new("cargo") + .current_dir(temp_path.join("timescale_vector")) + .args(["rm", "--dev", "pgrx-tests"]) + .stdout(Stdio::inherit()) + .stderr(Stdio::piped()) + .output() + .unwrap(); + assert!(res.status.success(), "failed: {:?}", res); + + let res = std::process::Command::new("cargo") + .current_dir(temp_path.join("timescale_vector")) + .args(["add", "-F", &format!("pg{}", pg_version), "pgrx"]) + .stdout(Stdio::inherit()) + .stderr(Stdio::piped()) + .output() + .unwrap(); + assert!(res.status.success(), "failed: {:?}", res); + + //let contents = fs::read_to_string(temp_path.join("timescale_vector/Cargo.toml")).unwrap(); + //print!("cargo {}", contents); + + let res = std::process::Command::new("cargo") + .current_dir(temp_path.join("timescale_vector")) + .arg("pgrx") + .arg("install") + .arg("--test") + .arg("--pg-config") + .arg(pg_config.path().unwrap()) + .stdout(Stdio::inherit()) + .stderr(Stdio::piped()) + .output() + .unwrap(); + assert!(res.status.success(), "failed: {:?}", res); + + client + .execute( + &format!( + "CREATE EXTENSION timescale_vector VERSION '{}' CASCADE;", + version + ), + &[], + ) + .unwrap(); + + let suffix = (1..=253) + .map(|i| format!("{}", i)) + .collect::>() + .join(", "); + + client + .batch_execute(&format!( + "CREATE TABLE test(embedding vector(256)); + + select setseed(0.5); + -- generate 300 vectors + INSERT INTO test(embedding) + SELECT + * + FROM ( + SELECT + ('[ 0 , ' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding + FROM + generate_series(1, 255 * 300) i + GROUP BY + i % 300) g; + + INSERT INTO test(embedding) VALUES ('[1,2,3,{suffix}]'), ('[4,5,6,{suffix}]'), ('[7,8,10,{suffix}]'); + + CREATE INDEX idxtest + ON test + USING tsv(embedding); + " + )) + .unwrap(); + + client.execute("set enable_seqscan = 0;", &[]).unwrap(); + let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); + assert_eq!(cnt, 303, "count before upgrade"); + + //reinstall myself + let res = std::process::Command::new("cargo") + .arg("pgrx") + .arg("install") + .arg("--test") + .arg("--pg-config") + .arg(pg_config.path().unwrap()) + .stdout(Stdio::inherit()) + .stderr(Stdio::piped()) + .output() + .unwrap(); + assert!(res.status.success(), "failed: {:?}", res); + + //need to recreate the client to avoid double load of GUC. Look into this later. + let (mut client, _) = pgrx_tests::client().unwrap(); + client + .execute( + &format!( + "ALTER EXTENSION timescale_vector UPDATE TO '{}'", + env!("CARGO_PKG_VERSION") + ), + &[], + ) + .unwrap(); + + client.execute("set enable_seqscan = 0;", &[]).unwrap(); + let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); + assert_eq!(cnt, 303, "count after upgrade"); + } +} diff --git a/timescale_vector/timescale_vector.control b/timescale_vector/timescale_vector.control index 8e9c6c90..e6d31610 100644 --- a/timescale_vector/timescale_vector.control +++ b/timescale_vector/timescale_vector.control @@ -1,6 +1,6 @@ comment = 'timescale_vector: Advanced indexing for vector data' default_version = '@CARGO_VERSION@' -module_pathname = '$libdir/timescale_vector' +#module_pathname = '$libdir/timescale_vector' relocatable = false superuser = true requires = 'vector' From 6e205b5d6ff67a44a04d66e557181c763005d607 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Thu, 4 Apr 2024 12:10:10 -0400 Subject: [PATCH 23/44] Adjust index options --- timescale_vector/src/access_method/README.md | 20 ++ timescale_vector/src/access_method/bq.rs | 78 ++++++-- timescale_vector/src/access_method/build.rs | 13 +- .../src/access_method/meta_page.rs | 46 +++-- timescale_vector/src/access_method/options.rs | 177 +++++++++++++----- .../src/access_method/plain_storage.rs | 18 +- timescale_vector/src/access_method/storage.rs | 18 +- 7 files changed, 276 insertions(+), 94 deletions(-) diff --git a/timescale_vector/src/access_method/README.md b/timescale_vector/src/access_method/README.md index 839466d4..ed0e2cca 100644 --- a/timescale_vector/src/access_method/README.md +++ b/timescale_vector/src/access_method/README.md @@ -1,3 +1,23 @@ +# Graph + +The graph abstraction implements the 2 primary algorithms: +- greedy_search - which finds the index nodes closest to a given query +- prune_neighbors - which reduces the number of neighbors assigned to a particular nodes in order to fit within the num_neighbor limit. + +Graph also implements the insertion algoritm for when a node needs to be added to the graph. Insertion is mostly a combination of the 2 algorithms above, greedy_search and prune_neighbors. + +We support multiple storage layouts (described below). The logic in graph is works on an abstract storage object and is not concerned with the particular storage implementation. +Thus, any logic that needs to differ between storage implementations has to fall within the responsibility of the storage object and not be included in graph. + +## Greedy search + +Refer to the DiskANN paper for an overview. The greedy search algorithm works by traversing the graph to find the closest nodes to a given query. It does this by: +- starting with a set (right now implemented as just one) initial nodes (called init_ids). +- iteratively: +- - + + + # On Disk Layout Meta Page diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index a1fa6874..b3d06a89 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -13,7 +13,7 @@ use super::{ use std::{cell::RefCell, collections::HashMap, iter::once, marker::PhantomData, pin::Pin}; use pgrx::{ - pg_sys::{InvalidBlockNumber, InvalidOffsetNumber}, + pg_sys::{InvalidBlockNumber, InvalidOffsetNumber, BLCKSZ}, PgRelation, }; use rkyv::{vec::ArchivedVec, Archive, Archived, Deserialize, Serialize}; @@ -104,6 +104,10 @@ impl BqQuantizer { } } + fn quantized_size_bytes(num_dimensions: usize) -> usize { + Self::quantized_size(num_dimensions) * std::mem::size_of::() + } + fn quantize(&self, full_vector: &[f32]) -> Vec { assert!(!self.training); if self.use_mean { @@ -468,7 +472,7 @@ impl<'a> Storage for BqSpeedupStorage<'a> { ) -> ItemPointer { let bq_vector = self.quantizer.vector_for_new_node(meta_page, full_vector); - let node = BqNode::new(heap_pointer, &meta_page, bq_vector.as_slice()); + let node = BqNode::with_meta(heap_pointer, &meta_page, bq_vector.as_slice()); let index_pointer: IndexPointer = node.write(tape, stats); index_pointer @@ -652,19 +656,32 @@ pub struct BqNode { } impl BqNode { - pub fn new( + pub fn with_meta( heap_pointer: HeapPointer, meta_page: &MetaPage, bq_vector: &[BqVectorElement], ) -> Self { - let num_neighbors = meta_page.get_num_neighbors(); + Self::new( + heap_pointer, + meta_page.get_num_neighbors() as usize, + meta_page.get_num_dimensions() as usize, + bq_vector, + ) + } + + fn new( + heap_pointer: HeapPointer, + num_neighbors: usize, + num_dimensions: usize, + bq_vector: &[BqVectorElement], + ) -> Self { // always use vectors of num_neighbors in length because we never want the serialized size of a Node to change let neighbor_index_pointers: Vec<_> = (0..num_neighbors) .map(|_| ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber)) .collect(); let neighbor_vectors: Vec<_> = (0..num_neighbors) - .map(|_| vec![0; BqQuantizer::quantized_size(meta_page.get_num_dimensions() as _)]) + .map(|_| vec![0; BqQuantizer::quantized_size(num_dimensions as _)]) .collect(); Self { @@ -674,6 +691,43 @@ impl BqNode { neighbor_vectors: neighbor_vectors, } } + + fn test_size(num_neighbors: usize, num_dimensions: usize) -> usize { + let v: Vec = vec![0; BqQuantizer::quantized_size(num_dimensions)]; + let hp = HeapPointer::new(InvalidBlockNumber, InvalidOffsetNumber); + let n = Self::new(hp, num_neighbors, num_dimensions, &v); + n.serialize_to_vec().len() + } + + pub fn get_default_num_neighbors(num_dimensions: usize) -> usize { + //how many neighbors can fit on one page? That's what we choose. + + //we first overapproximate the number of neighbors and then double check by actually calculating the size of the BqNode. + + //blocksize - 100 bytes for the padding/header/etc. + let page_size = BLCKSZ as usize - 50; + //one quantized_vector takes this many bytes + let vec_size = BqQuantizer::quantized_size_bytes(num_dimensions as usize) + 1; + //start from the page size then subtract the heap_item_pointer and bq_vector elements of BqNode. + let starting = BLCKSZ as usize - std::mem::size_of::() - vec_size; + //one neigbors contribution to neighbor_index_pointers + neighbor_vectors in BqNode. + let one_neighbor = vec_size + std::mem::size_of::(); + + let mut num_neighbors_overapproximate: usize = starting / one_neighbor; + while num_neighbors_overapproximate > 0 { + let serialized_size = BqNode::test_size( + num_neighbors_overapproximate as usize, + num_dimensions as usize, + ); + if serialized_size <= page_size { + return num_neighbors_overapproximate; + } + num_neighbors_overapproximate -= 1; + } + pgrx::error!( + "Could not find a valid number of neighbors for the default value. Please specify one." + ); + } } impl ArchivedBqNode { @@ -767,9 +821,9 @@ mod tests { use pgrx::*; #[pg_test] - unsafe fn test_bq_storage_index_creation() -> spi::Result<()> { + unsafe fn test_bq_storage_index_creation_default_neighbors() -> spi::Result<()> { crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( - "num_neighbors=38, USE_BQ = TRUE", + "storage_layout = io_optimized", )?; Ok(()) } @@ -778,7 +832,7 @@ mod tests { unsafe fn test_bq_storage_index_creation_few_neighbors() -> spi::Result<()> { //a test with few neighbors tests the case that nodes share a page, which has caused deadlocks in the past. crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( - "num_neighbors=10, USE_BQ = TRUE", + "num_neighbors=10, storage_layout = io_optimized", )?; Ok(()) } @@ -786,28 +840,28 @@ mod tests { #[test] fn test_bq_storage_delete_vacuum_plain() { crate::access_method::vacuum::tests::test_delete_vacuum_plain_scaffold( - "num_neighbors = 10, use_bq = TRUE", + "num_neighbors = 10, storage_layout = io_optimized", ); } #[test] fn test_bq_storage_delete_vacuum_full() { crate::access_method::vacuum::tests::test_delete_vacuum_full_scaffold( - "num_neighbors = 38, use_bq = TRUE", + "num_neighbors = 38, storage_layout = io_optimized", ); } #[pg_test] unsafe fn test_bq_storage_empty_table_insert() -> spi::Result<()> { crate::access_method::build::tests::test_empty_table_insert_scaffold( - "num_neighbors=38, use_bq = TRUE", + "num_neighbors=38, storage_layout = io_optimized", ) } #[pg_test] unsafe fn test_bq_storage_insert_empty_insert() -> spi::Result<()> { crate::access_method::build::tests::test_insert_empty_insert_scaffold( - "num_neighbors=38, use_bq = TRUE", + "num_neighbors=38, storage_layout = io_optimized", ) } } diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index 788616cb..c4d5e004 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -70,26 +70,27 @@ pub extern "C" fn ambuild( let opt = TSVIndexOptions::from_relation(&index_relation); notice!( - "Starting index build. num_neighbors={} search_list_size={}, max_alpha={}, use_pq={}, pq_vector_length={}", - opt.num_neighbors, + "Starting index build. num_neighbors={} search_list_size={}, max_alpha={}, storage_layout={:?}, pq_vector_length={}", + opt.get_num_neighbors(), opt.search_list_size, opt.max_alpha, - opt.use_pq, + opt.get_storage_type(), opt.pq_vector_length ); let dimensions = index_relation.tuple_desc().get(0).unwrap().atttypmod; // PQ is only applicable to high dimension vectors. - if opt.use_pq { + //FIXME: uncomment/delete + /*if opt.get_storage_layout() == { if dimensions < opt.pq_vector_length as i32 { error!("use_pq can only be applied to vectors with greater than {} dimensions. {} dimensions provided", opt.pq_vector_length, dimensions) }; if dimensions % opt.pq_vector_length as i32 != 0 { error!("use_pq can only be applied to vectors where the number of dimensions {} is divisible by the pq_vector_length {} ", dimensions, opt.pq_vector_length) }; - } + }*/ assert!(dimensions > 0 && dimensions < 2000); - let meta_page = unsafe { MetaPage::create(&index_relation, dimensions as _, opt.clone()) }; + let meta_page = unsafe { MetaPage::create(&index_relation, dimensions as _, opt) }; let ntuples = do_heap_scan(index_info, &heap_relation, &index_relation, meta_page); diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index c878d3a5..a9ea36b9 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -8,7 +8,9 @@ use crate::access_method::options::TSVIndexOptions; use crate::util::page; use crate::util::*; +use super::bq::BqNode; use super::distance; +use super::options::NUM_NEIGHBORS_DEFAULT_SENTINEL; use super::stats::StatsNodeModify; use super::storage::StorageType; @@ -53,6 +55,10 @@ impl MetaPageV1 { } pub fn get_new_meta(&self) -> MetaPage { + if self.use_pq { + pgrx::error!("PQ is no longer supported. Please rebuild the TSV index."); + } + MetaPage { magic_number: TSV_MAGIC_NUMBER, version: TSV_VERSION, @@ -60,15 +66,14 @@ impl MetaPageV1 { distance_type: DistanceType::L2 as u16, num_dimensions: self.num_dimensions, num_neighbors: self.num_neighbors, + storage_type: StorageType::Plain as u8, search_list_size: self.search_list_size, max_alpha: self.max_alpha, init_ids_block_number: self.init_ids_block_number, init_ids_offset: self.init_ids_offset, - use_pq: self.use_pq, pq_vector_length: self.pq_vector_length, pq_block_number: self.pq_block_number, pq_block_offset: self.pq_block_offset, - use_bq: false, } } } @@ -110,20 +115,21 @@ pub struct MetaPage { magic_number: u32, version: u32, extension_version_when_built: String, + /// The value of the DistanceType enum distance_type: u16, /// number of dimensions in the vector num_dimensions: u32, + /// the value of the TSVStorageLayout enum + storage_type: u8, /// max number of outgoing edges a node in the graph can have (R in the papers) num_neighbors: u32, search_list_size: u32, max_alpha: f64, init_ids_block_number: pg_sys::BlockNumber, init_ids_offset: pg_sys::OffsetNumber, - use_pq: bool, pq_vector_length: usize, pq_block_number: pg_sys::BlockNumber, pq_block_offset: pg_sys::OffsetNumber, - use_bq: bool, } impl MetaPage { @@ -151,10 +157,6 @@ impl MetaPage { self.max_alpha } - fn get_use_pq(&self) -> bool { - self.use_pq - } - pub fn get_distance_function(&self) -> fn(&[f32], &[f32]) -> f32 { match DistanceType::from_u16(self.distance_type) { DistanceType::Cosine => distance::distance_cosine, @@ -163,13 +165,7 @@ impl MetaPage { } pub fn get_storage_type(&self) -> StorageType { - if self.get_use_pq() { - StorageType::PqCompression - } else if self.use_bq { - StorageType::BqSpeedup - } else { - StorageType::Plain - } + StorageType::from_u8(self.storage_type) } pub fn get_max_neighbors_during_build(&self) -> usize { @@ -186,7 +182,7 @@ impl MetaPage { } pub fn get_pq_pointer(&self) -> Option { - if (!self.use_pq && !self.use_bq) + if (self.storage_type != StorageType::BqSpeedup as u8) || (self.pq_block_number == 0 && self.pq_block_offset == 0) { return None; @@ -196,6 +192,19 @@ impl MetaPage { Some(ptr) } + fn calculate_num_neighbors(num_dimensions: u32, opt: &PgBox) -> u32 { + let num_neighbors = (*opt).get_num_neighbors(); + if num_neighbors == NUM_NEIGHBORS_DEFAULT_SENTINEL { + if (*opt).get_storage_type() == StorageType::Plain { + 50 + } else { + BqNode::get_default_num_neighbors(num_dimensions as usize) as u32 + } + } else { + num_neighbors as u32 + } + } + /// Write out a new meta page. /// Has to be done as the first write to a new relation. pub unsafe fn create( @@ -211,16 +220,15 @@ impl MetaPage { extension_version_when_built: version.to_string(), distance_type: DistanceType::Cosine as u16, num_dimensions, - num_neighbors: (*opt).num_neighbors, + storage_type: (*opt).get_storage_type() as u8, + num_neighbors: Self::calculate_num_neighbors(num_dimensions, &opt), search_list_size: (*opt).search_list_size, max_alpha: (*opt).max_alpha, init_ids_block_number: 0, init_ids_offset: 0, - use_pq: (*opt).use_pq, pq_vector_length: (*opt).pq_vector_length, pq_block_number: 0, pq_block_offset: 0, - use_bq: (*opt).use_bq, }; let page = page::WritablePage::new(index, crate::util::page::PageType::Meta); meta.write_to_page(page); diff --git a/timescale_vector/src/access_method/options.rs b/timescale_vector/src/access_method/options.rs index 45b7f8ab..b492ba25 100644 --- a/timescale_vector/src/access_method/options.rs +++ b/timescale_vector/src/access_method/options.rs @@ -1,34 +1,40 @@ use memoffset::*; -use pgrx::{pg_sys::AsPgCStr, prelude::*, set_varsize, PgRelation}; -use std::fmt::Debug; +use pgrx::{pg_sys::AsPgCStr, prelude::*, set_varsize, void_ptr, PgRelation}; +use std::{ffi::CStr, fmt::Debug}; -#[derive(Copy, Clone, Debug, PartialEq)] +use super::storage::StorageType; + +//DO NOT derive Clone for this struct. The storage layout string comes at the end and wouldn't be copied properly. +#[derive(Debug, PartialEq)] #[repr(C)] pub struct TSVIndexOptions { /* varlena header (do not touch directly!) */ #[allow(dead_code)] vl_len_: i32, - pub num_neighbors: u32, + pub storage_layout_offset: i32, + num_neighbors: i32, pub search_list_size: u32, pub max_alpha: f64, - pub use_pq: bool, - pub use_bq: bool, pub pq_vector_length: usize, } +pub const NUM_NEIGHBORS_DEFAULT_SENTINEL: i32 = -1; +const DEFAULT_MAX_ALPHA: f64 = 1.2; + impl TSVIndexOptions { + //note: this should only be used when building a new index. The options aren't really versioned. + //therefore, we should move all the options to the meta page when building the index (meta pages are properly versioned). pub fn from_relation(relation: &PgRelation) -> PgBox { if relation.rd_index.is_null() { panic!("'{}' is not a TSV index", relation.name()) } else if relation.rd_options.is_null() { // use defaults let mut ops = unsafe { PgBox::::alloc0() }; - ops.num_neighbors = 50; + ops.storage_layout_offset = 0; + ops.num_neighbors = NUM_NEIGHBORS_DEFAULT_SENTINEL; ops.search_list_size = 100; - ops.max_alpha = 1.0; - ops.use_pq = false; - ops.use_bq = false; + ops.max_alpha = DEFAULT_MAX_ALPHA; ops.pq_vector_length = 256; unsafe { set_varsize( @@ -41,11 +47,54 @@ impl TSVIndexOptions { unsafe { PgBox::from_pg(relation.rd_options as *mut TSVIndexOptions) } } } + + pub fn get_num_neighbors(&self) -> i32 { + if self.num_neighbors == NUM_NEIGHBORS_DEFAULT_SENTINEL { + //specify to use the default value here + //we can't derive the default at this point in the code because the default is based on the number of dimensions in the vector in the io_optimized case. + NUM_NEIGHBORS_DEFAULT_SENTINEL + } else { + if self.num_neighbors < 10 { + panic!("num_neighbors must be greater than 10, or -1 for default") + } + self.num_neighbors + } + } + + pub fn get_storage_type(&self) -> StorageType { + let s = self.get_str(self.storage_layout_offset, || "io_optimized".to_string()); + match s.as_str() { + "io_optimized" => StorageType::BqSpeedup, + "plain" => StorageType::Plain, + _ => panic!("invalid storage_layout: {}", s), + } + } + + fn get_str String>(&self, offset: i32, default: F) -> String { + if offset == 0 { + default() + } else { + let opts = self as *const _ as void_ptr as usize; + let value = + unsafe { CStr::from_ptr((opts + offset as usize) as *const std::os::raw::c_char) }; + + value.to_str().unwrap().to_owned() + } + } } -const NUM_REL_OPTS: usize = 6; +const NUM_REL_OPTS: usize = 5; static mut RELOPT_KIND_TSV: pg_sys::relopt_kind = 0; +// amoptions is a function that gets a datum of text[] data from pg_class.reloptions (which contains text in the format "key=value") and returns a bytea for the struct for the parsed options. +// this is used to fill the rd_options field in the index relation. +// except for during build the validate parameter should be false. +// any option that is no longer recognized that exists in the reloptions will simply be ignored when validate is false. +// therefore, it is safe to change the options struct and add/remove new options without breaking existing indexes. +// but note that the standard parsing way has no ability to put "migration" logic in here. So all new options will have to have defaults value when reading old indexes. +// we could do additional logic to fix this here, but instead we just move the option values to the meta page when building the index, and do versioning there. +// side note: this logic is not used in \d+ and similar psql commands to get description info. Those commands use the text array in pg_class.reloptions directly. +// so when displaying the info, they'll show the old options and their values as set when the index was created. #[allow(clippy::unneeded_field_pattern)] // b/c of offset_of!() #[pg_guard] pub unsafe extern "C" fn amoptions( @@ -54,6 +103,11 @@ pub unsafe extern "C" fn amoptions( ) -> *mut pg_sys::bytea { // TODO: how to make this const? we can't use offset_of!() macro in const definitions, apparently let tab: [pg_sys::relopt_parse_elt; NUM_REL_OPTS] = [ + pg_sys::relopt_parse_elt { + optname: "storage_layout".as_pg_cstr(), + opttype: pg_sys::relopt_type_RELOPT_TYPE_STRING, + offset: offset_of!(TSVIndexOptions, storage_layout_offset) as i32, + }, pg_sys::relopt_parse_elt { optname: "num_neighbors".as_pg_cstr(), opttype: pg_sys::relopt_type_RELOPT_TYPE_INT, @@ -69,16 +123,6 @@ pub unsafe extern "C" fn amoptions( opttype: pg_sys::relopt_type_RELOPT_TYPE_REAL, offset: offset_of!(TSVIndexOptions, max_alpha) as i32, }, - pg_sys::relopt_parse_elt { - optname: "use_pq".as_pg_cstr(), - opttype: pg_sys::relopt_type_RELOPT_TYPE_BOOL, - offset: offset_of!(TSVIndexOptions, use_pq) as i32, - }, - pg_sys::relopt_parse_elt { - optname: "use_bq".as_pg_cstr(), - opttype: pg_sys::relopt_type_RELOPT_TYPE_BOOL, - offset: offset_of!(TSVIndexOptions, use_bq) as i32, - }, pg_sys::relopt_parse_elt { optname: "pq_vector_length".as_pg_cstr(), opttype: pg_sys::relopt_type_RELOPT_TYPE_INT, @@ -110,15 +154,43 @@ unsafe fn build_relopts( rdopts as *mut pg_sys::bytea } +#[pg_guard] +extern "C" fn validate_storage_layout(value: *const std::os::raw::c_char) { + if value.is_null() { + // use a default value + return; + } + + let value = unsafe { CStr::from_ptr(value) } + .to_str() + .expect("failed to parse storage_layout value") + .to_lowercase(); + if value != "io_optimized" && value != "plain" { + panic!( + "invalid storage_layout. Must be one of 'io_optimized' or 'plain': {}", + value + ) + } +} + pub unsafe fn init() { RELOPT_KIND_TSV = pg_sys::add_reloption_kind(); + pg_sys::add_string_reloption( + RELOPT_KIND_TSV, + "storage_layout".as_pg_cstr(), + "Storage layout: either io_optimized or plain".as_pg_cstr(), + "io_optimized".as_pg_cstr(), + Some(validate_storage_layout), + pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, + ); + pg_sys::add_int_reloption( RELOPT_KIND_TSV, "num_neighbors".as_pg_cstr(), "Maximum number of neighbors in the graph".as_pg_cstr(), - 50, - 10, + NUM_NEIGHBORS_DEFAULT_SENTINEL, + -1, 1000, pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, ); @@ -137,25 +209,11 @@ pub unsafe fn init() { RELOPT_KIND_TSV, "max_alpha".as_pg_cstr(), "The maximum alpha used in pruning".as_pg_cstr(), - 1.0, + DEFAULT_MAX_ALPHA, 1.0, 5.0, pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, ); - pg_sys::add_bool_reloption( - RELOPT_KIND_TSV, - "use_pq".as_pg_cstr(), - "Enable product quantization".as_pg_cstr(), - false, - pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, - ); - pg_sys::add_bool_reloption( - RELOPT_KIND_TSV, - "use_bq".as_pg_cstr(), - "Enable binary quantization".as_pg_cstr(), - false, - pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, - ); pg_sys::add_int_reloption( RELOPT_KIND_TSV, "pq_vector_length".as_pg_cstr(), @@ -170,7 +228,10 @@ pub unsafe fn init() { #[cfg(any(test, feature = "pg_test"))] #[pgrx::pg_schema] mod tests { - use crate::access_method::options::TSVIndexOptions; + use crate::access_method::{ + options::{TSVIndexOptions, DEFAULT_MAX_ALPHA, NUM_NEIGHBORS_DEFAULT_SENTINEL}, + storage::StorageType, + }; use pgrx::*; #[pg_test] @@ -204,11 +265,10 @@ mod tests { Spi::get_one::("SELECT 'idxtest'::regclass::oid")?.expect("oid was null"); let indexrel = PgRelation::from_pg(pg_sys::RelationIdGetRelation(index_oid)); let options = TSVIndexOptions::from_relation(&indexrel); - assert_eq!(options.num_neighbors, 50); + assert_eq!(options.get_num_neighbors(), NUM_NEIGHBORS_DEFAULT_SENTINEL); assert_eq!(options.search_list_size, 100); - assert_eq!(options.max_alpha, 1.0); - assert_eq!(options.use_pq, false); - assert_eq!(options.use_bq, false); + assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); + assert_eq!(options.get_storage_type(), StorageType::BqSpeedup); assert_eq!(options.pq_vector_length, 256); Ok(()) } @@ -220,18 +280,39 @@ mod tests { CREATE INDEX idxtest ON test USING tsv(encoding) - WITH (use_bq = TRUE);", + WITH (storage_layout = io_optimized);", + ))?; + + let index_oid = + Spi::get_one::("SELECT 'idxtest'::regclass::oid")?.expect("oid was null"); + let indexrel = PgRelation::from_pg(pg_sys::RelationIdGetRelation(index_oid)); + let options = TSVIndexOptions::from_relation(&indexrel); + assert_eq!(options.get_num_neighbors(), NUM_NEIGHBORS_DEFAULT_SENTINEL); + assert_eq!(options.search_list_size, 100); + assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); + assert_eq!(options.get_storage_type(), StorageType::BqSpeedup); + assert_eq!(options.pq_vector_length, 256); + Ok(()) + } + + #[pg_test] + unsafe fn test_index_options_plain() -> spi::Result<()> { + Spi::run(&format!( + "CREATE TABLE test(encoding vector(3)); + CREATE INDEX idxtest + ON test + USING tsv(encoding) + WITH (storage_layout = plain);", ))?; let index_oid = Spi::get_one::("SELECT 'idxtest'::regclass::oid")?.expect("oid was null"); let indexrel = PgRelation::from_pg(pg_sys::RelationIdGetRelation(index_oid)); let options = TSVIndexOptions::from_relation(&indexrel); - assert_eq!(options.num_neighbors, 50); + assert_eq!(options.get_num_neighbors(), NUM_NEIGHBORS_DEFAULT_SENTINEL); assert_eq!(options.search_list_size, 100); - assert_eq!(options.max_alpha, 1.0); - assert_eq!(options.use_pq, false); - assert_eq!(options.use_bq, true); + assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); + assert_eq!(options.get_storage_type(), StorageType::Plain); assert_eq!(options.pq_vector_length, 256); Ok(()) } diff --git a/timescale_vector/src/access_method/plain_storage.rs b/timescale_vector/src/access_method/plain_storage.rs index 423cc8e1..5f052283 100644 --- a/timescale_vector/src/access_method/plain_storage.rs +++ b/timescale_vector/src/access_method/plain_storage.rs @@ -325,7 +325,7 @@ mod tests { #[pg_test] unsafe fn test_plain_storage_index_creation() -> spi::Result<()> { crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( - "num_neighbors=38", + "num_neighbors=38, storage_layout = plain", )?; Ok(()) } @@ -334,7 +334,7 @@ mod tests { unsafe fn test_plain_storage_index_creation_few_neighbors() -> spi::Result<()> { //a test with few neighbors tests the case that nodes share a page, which has caused deadlocks in the past. crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( - "num_neighbors=10", + "num_neighbors=10, storage_layout = plain", )?; Ok(()) } @@ -342,22 +342,28 @@ mod tests { #[test] fn test_plain_storage_delete_vacuum_plain() { crate::access_method::vacuum::tests::test_delete_vacuum_plain_scaffold( - "num_neighbors = 38", + "num_neighbors = 38, storage_layout = plain", ); } #[test] fn test_plain_storage_delete_vacuum_full() { - crate::access_method::vacuum::tests::test_delete_vacuum_full_scaffold("num_neighbors = 38"); + crate::access_method::vacuum::tests::test_delete_vacuum_full_scaffold( + "num_neighbors = 38, storage_layout = plain", + ); } #[pg_test] unsafe fn test_plain_storage_empty_table_insert() -> spi::Result<()> { - crate::access_method::build::tests::test_empty_table_insert_scaffold("num_neighbors=38") + crate::access_method::build::tests::test_empty_table_insert_scaffold( + "num_neighbors=38, storage_layout = plain", + ) } #[pg_test] unsafe fn test_plain_storage_insert_empty_insert() -> spi::Result<()> { - crate::access_method::build::tests::test_insert_empty_insert_scaffold("num_neighbors=38") + crate::access_method::build::tests::test_insert_empty_insert_scaffold( + "num_neighbors=38, storage_layout = plain", + ) } } diff --git a/timescale_vector/src/access_method/storage.rs b/timescale_vector/src/access_method/storage.rs index 0bfcbd2f..4692c4f1 100644 --- a/timescale_vector/src/access_method/storage.rs +++ b/timescale_vector/src/access_method/storage.rs @@ -123,8 +123,20 @@ pub trait Storage { fn get_distance_function(&self) -> fn(&[f32], &[f32]) -> f32; } +#[derive(PartialEq, Debug)] pub enum StorageType { - BqSpeedup, - PqCompression, - Plain, + Plain = 0, + BqSpeedup = 1, + PqCompression = 2, +} + +impl StorageType { + pub fn from_u8(value: u8) -> Self { + match value { + 0 => StorageType::Plain, + 1 => StorageType::BqSpeedup, + 2 => StorageType::PqCompression, + _ => panic!("Invalid storage type"), + } + } } From ca569b509ceae53ee38131cd2bf4ebfc4a15095e Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Thu, 4 Apr 2024 14:02:40 -0400 Subject: [PATCH 24/44] Remove PQ --- timescale_vector/src/access_method/bq.rs | 16 +- timescale_vector/src/access_method/build.rs | 86 +--- timescale_vector/src/access_method/graph.rs | 1 - .../src/access_method/meta_page.rs | 39 +- timescale_vector/src/access_method/mod.rs | 3 - timescale_vector/src/access_method/options.rs | 21 +- .../src/access_method/plain_node.rs | 17 +- .../src/access_method/pq_quantizer.rs | 316 ------------- .../src/access_method/pq_quantizer_storage.rs | 122 ----- .../src/access_method/pq_storage.rs | 418 ------------------ timescale_vector/src/access_method/scan.rs | 41 -- timescale_vector/src/access_method/storage.rs | 2 - timescale_vector/src/access_method/vacuum.rs | 14 +- 13 files changed, 32 insertions(+), 1064 deletions(-) delete mode 100644 timescale_vector/src/access_method/pq_quantizer.rs delete mode 100644 timescale_vector/src/access_method/pq_quantizer_storage.rs delete mode 100644 timescale_vector/src/access_method/pq_storage.rs diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index b3d06a89..a567a750 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -45,14 +45,14 @@ impl BqMeans { ) -> BqQuantizer { let mut quantizer = BqQuantizer::new(); if quantizer.use_mean { - if meta_page.get_pq_pointer().is_none() { + if meta_page.get_quantizer_metadata_pointer().is_none() { pgrx::error!("No BQ pointer found in meta page"); } - let pq_item_pointer = meta_page.get_pq_pointer().unwrap(); - let rpq = BqMeans::read(index, pq_item_pointer, stats); - let rpn = rpq.get_archived_node(); + let quantizer_item_pointer = meta_page.get_quantizer_metadata_pointer().unwrap(); + let bq = BqMeans::read(index, quantizer_item_pointer, stats); + let archived = bq.get_archived_node(); - quantizer.load(rpn.count, rpn.means.to_vec()); + quantizer.load(archived.count, archived.means.to_vec()); } quantizer } @@ -379,7 +379,11 @@ impl<'a> BqSpeedupStorage<'a> { fn write_quantizer_metadata(&self, stats: &mut S) { if self.quantizer.use_mean { let index_pointer = unsafe { BqMeans::store(&self.index, &self.quantizer, stats) }; - super::meta_page::MetaPage::update_pq_pointer(&self.index, index_pointer, stats); + super::meta_page::MetaPage::update_quantizer_metadata_pointer( + &self.index, + index_pointer, + stats, + ); } } diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index c4d5e004..fb5afb3a 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -18,12 +18,10 @@ use super::graph_neighbor_store::BuilderNeighborCache; use super::meta_page::MetaPage; use super::plain_storage::PlainStorage; -use super::pq_storage::PqCompressionStorage; use super::storage::{Storage, StorageType}; enum StorageBuildState<'a, 'b, 'c, 'd, 'e> { BqSpeedup(&'a mut BqSpeedupStorage<'b>, &'c mut BuildState<'d, 'e>), - PqCompression(&'a mut PqCompressionStorage<'b>, &'c mut BuildState<'d, 'e>), Plain(&'a mut PlainStorage<'b>, &'c mut BuildState<'d, 'e>), } @@ -70,25 +68,14 @@ pub extern "C" fn ambuild( let opt = TSVIndexOptions::from_relation(&index_relation); notice!( - "Starting index build. num_neighbors={} search_list_size={}, max_alpha={}, storage_layout={:?}, pq_vector_length={}", + "Starting index build. num_neighbors={} search_list_size={}, max_alpha={}, storage_layout={:?}", opt.get_num_neighbors(), opt.search_list_size, opt.max_alpha, opt.get_storage_type(), - opt.pq_vector_length ); let dimensions = index_relation.tuple_desc().get(0).unwrap().atttypmod; - // PQ is only applicable to high dimension vectors. - //FIXME: uncomment/delete - /*if opt.get_storage_layout() == { - if dimensions < opt.pq_vector_length as i32 { - error!("use_pq can only be applied to vectors with greater than {} dimensions. {} dimensions provided", opt.pq_vector_length, dimensions) - }; - if dimensions % opt.pq_vector_length as i32 != 0 { - error!("use_pq can only be applied to vectors where the number of dimensions {} is divisible by the pq_vector_length {} ", dimensions, opt.pq_vector_length) - }; - }*/ assert!(dimensions > 0 && dimensions < 2000); let meta_page = unsafe { MetaPage::create(&index_relation, dimensions as _, opt) }; @@ -138,23 +125,6 @@ pub unsafe extern "C" fn aminsert( &mut stats, ); } - StorageType::PqCompression => { - let pq = PqCompressionStorage::load_for_insert( - &heap_relation, - get_attribute_number(index_info), - &index_relation, - &meta_page, - &mut stats.quantizer_stats, - ); - insert_storage( - &pq, - &index_relation, - vec, - heap_pointer, - &mut meta_page, - &mut stats, - ); - } StorageType::BqSpeedup => { let bq = BqSpeedupStorage::load_for_insert( &heap_relation, @@ -242,41 +212,6 @@ fn do_heap_scan<'a>( finalize_index_build(&mut plain, &mut bs, write_stats) } - StorageType::PqCompression => { - let mut pq = PqCompressionStorage::new_for_build( - index_relation, - heap_relation, - get_attribute_number(index_info), - meta_page.get_distance_function(), - ); - pq.start_training(&meta_page); - unsafe { - pg_sys::IndexBuildHeapScan( - heap_relation.as_ptr(), - index_relation.as_ptr(), - index_info, - Some(build_callback_pq_train), - &mut pq, - ); - } - pq.finish_training(&mut write_stats); - - let page_type = PqCompressionStorage::page_type(); - let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); - let mut state = StorageBuildState::PqCompression(&mut pq, &mut bs); - - unsafe { - pg_sys::IndexBuildHeapScan( - heap_relation.as_ptr(), - index_relation.as_ptr(), - index_info, - Some(build_callback), - &mut state, - ); - } - - finalize_index_build(&mut pq, &mut bs, write_stats) - } StorageType::BqSpeedup => { let mut bq = BqSpeedupStorage::new_for_build( index_relation, @@ -397,22 +332,6 @@ unsafe extern "C" fn build_callback_bq_train( } } -#[pg_guard] -unsafe extern "C" fn build_callback_pq_train( - _index: pg_sys::Relation, - _ctid: pg_sys::ItemPointer, - values: *mut pg_sys::Datum, - isnull: *mut bool, - _tuple_is_alive: bool, - state: *mut std::os::raw::c_void, -) { - let vec = PgVector::from_pg_parts(values, isnull, 0); - if let Some(vec) = vec { - let pq = (state as *mut PqCompressionStorage).as_mut().unwrap(); - pq.add_sample(vec.to_slice()); - } -} - #[pg_guard] unsafe extern "C" fn build_callback( index: pg_sys::Relation, @@ -432,9 +351,6 @@ unsafe extern "C" fn build_callback( StorageBuildState::BqSpeedup(bq, state) => { build_callback_memory_wrapper(index_relation, heap_pointer, vec, state, *bq); } - StorageBuildState::PqCompression(pq, state) => { - build_callback_memory_wrapper(index_relation, heap_pointer, vec, state, *pq); - } StorageBuildState::Plain(plain, state) => { build_callback_memory_wrapper(index_relation, heap_pointer, vec, state, *plain); } diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index b82a98ee..55785369 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -385,7 +385,6 @@ impl<'a> Graph<'a> { continue; } - //todo handle the non-pq case let mut distance_between_candidate_and_existing_neighbor = unsafe { dist_state .get_distance(candidate_neighbor.get_index_pointer_to_neighbor(), stats) diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index a9ea36b9..f48a3115 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -38,7 +38,7 @@ pub struct MetaPageV1 { init_ids_block_number: pg_sys::BlockNumber, init_ids_offset: pg_sys::OffsetNumber, use_pq: bool, - pq_vector_length: usize, + _pq_vector_length: usize, pq_block_number: pg_sys::BlockNumber, pq_block_offset: pg_sys::OffsetNumber, } @@ -71,9 +71,8 @@ impl MetaPageV1 { max_alpha: self.max_alpha, init_ids_block_number: self.init_ids_block_number, init_ids_offset: self.init_ids_offset, - pq_vector_length: self.pq_vector_length, - pq_block_number: self.pq_block_number, - pq_block_offset: self.pq_block_offset, + quantizer_metadata_block_number: self.pq_block_number, + quantizer_metadata_block_offset: self.pq_block_offset, } } } @@ -127,9 +126,8 @@ pub struct MetaPage { max_alpha: f64, init_ids_block_number: pg_sys::BlockNumber, init_ids_offset: pg_sys::OffsetNumber, - pq_vector_length: usize, - pq_block_number: pg_sys::BlockNumber, - pq_block_offset: pg_sys::OffsetNumber, + quantizer_metadata_block_number: pg_sys::BlockNumber, + quantizer_metadata_block_offset: pg_sys::OffsetNumber, } impl MetaPage { @@ -145,10 +143,6 @@ impl MetaPage { self.num_neighbors } - pub fn get_pq_vector_length(&self) -> usize { - self.pq_vector_length - } - pub fn get_search_list_size_for_build(&self) -> u32 { self.search_list_size } @@ -181,14 +175,18 @@ impl MetaPage { Some(vec![ptr]) } - pub fn get_pq_pointer(&self) -> Option { + pub fn get_quantizer_metadata_pointer(&self) -> Option { if (self.storage_type != StorageType::BqSpeedup as u8) - || (self.pq_block_number == 0 && self.pq_block_offset == 0) + || (self.quantizer_metadata_block_number == 0 + && self.quantizer_metadata_block_offset == 0) { return None; } - let ptr = IndexPointer::new(self.pq_block_number, self.pq_block_offset); + let ptr = IndexPointer::new( + self.quantizer_metadata_block_number, + self.quantizer_metadata_block_offset, + ); Some(ptr) } @@ -226,9 +224,8 @@ impl MetaPage { max_alpha: (*opt).max_alpha, init_ids_block_number: 0, init_ids_offset: 0, - pq_vector_length: (*opt).pq_vector_length, - pq_block_number: 0, - pq_block_offset: 0, + quantizer_metadata_block_number: 0, + quantizer_metadata_block_offset: 0, }; let page = page::WritablePage::new(index, crate::util::page::PageType::Meta); meta.write_to_page(page); @@ -334,14 +331,14 @@ impl MetaPage { }; } - pub fn update_pq_pointer( + pub fn update_quantizer_metadata_pointer( index: &PgRelation, - pq_pointer: IndexPointer, + quantizer_pointer: IndexPointer, stats: &mut S, ) { let mut meta = Self::fetch(index); - meta.pq_block_number = pq_pointer.block_number; - meta.pq_block_offset = pq_pointer.offset; + meta.quantizer_metadata_block_number = quantizer_pointer.block_number; + meta.quantizer_metadata_block_offset = quantizer_pointer.offset; unsafe { Self::overwrite(index, &meta); diff --git a/timescale_vector/src/access_method/mod.rs b/timescale_vector/src/access_method/mod.rs index 853a871f..4a859f8f 100644 --- a/timescale_vector/src/access_method/mod.rs +++ b/timescale_vector/src/access_method/mod.rs @@ -24,9 +24,6 @@ mod bq; pub mod distance; #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] mod distance_x86; -mod pq_quantizer; -mod pq_quantizer_storage; -mod pq_storage; #[pg_extern(sql = " CREATE OR REPLACE FUNCTION tsv_amhandler(internal) RETURNS index_am_handler PARALLEL SAFE IMMUTABLE STRICT COST 0.0001 LANGUAGE c AS '@MODULE_PATHNAME@', '@FUNCTION_NAME@'; diff --git a/timescale_vector/src/access_method/options.rs b/timescale_vector/src/access_method/options.rs index b492ba25..3b637c90 100644 --- a/timescale_vector/src/access_method/options.rs +++ b/timescale_vector/src/access_method/options.rs @@ -16,7 +16,6 @@ pub struct TSVIndexOptions { num_neighbors: i32, pub search_list_size: u32, pub max_alpha: f64, - pub pq_vector_length: usize, } pub const NUM_NEIGHBORS_DEFAULT_SENTINEL: i32 = -1; @@ -35,7 +34,6 @@ impl TSVIndexOptions { ops.num_neighbors = NUM_NEIGHBORS_DEFAULT_SENTINEL; ops.search_list_size = 100; ops.max_alpha = DEFAULT_MAX_ALPHA; - ops.pq_vector_length = 256; unsafe { set_varsize( ops.as_ptr().cast(), @@ -83,7 +81,7 @@ impl TSVIndexOptions { } } -const NUM_REL_OPTS: usize = 5; +const NUM_REL_OPTS: usize = 4; static mut RELOPT_KIND_TSV: pg_sys::relopt_kind = 0; // amoptions is a function that gets a datum of text[] data from pg_class.reloptions (which contains text in the format "key=value") and returns a bytea for the struct for the parsed options. @@ -123,11 +121,6 @@ pub unsafe extern "C" fn amoptions( opttype: pg_sys::relopt_type_RELOPT_TYPE_REAL, offset: offset_of!(TSVIndexOptions, max_alpha) as i32, }, - pg_sys::relopt_parse_elt { - optname: "pq_vector_length".as_pg_cstr(), - opttype: pg_sys::relopt_type_RELOPT_TYPE_INT, - offset: offset_of!(TSVIndexOptions, pq_vector_length) as i32, - }, ]; build_relopts(reloptions, validate, tab) @@ -214,15 +207,6 @@ pub unsafe fn init() { 5.0, pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, ); - pg_sys::add_int_reloption( - RELOPT_KIND_TSV, - "pq_vector_length".as_pg_cstr(), - "Length of the quantized vector representation".as_pg_cstr(), - 256, - 8, - 256, - pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, - ); } #[cfg(any(test, feature = "pg_test"))] @@ -269,7 +253,6 @@ mod tests { assert_eq!(options.search_list_size, 100); assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); assert_eq!(options.get_storage_type(), StorageType::BqSpeedup); - assert_eq!(options.pq_vector_length, 256); Ok(()) } @@ -291,7 +274,6 @@ mod tests { assert_eq!(options.search_list_size, 100); assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); assert_eq!(options.get_storage_type(), StorageType::BqSpeedup); - assert_eq!(options.pq_vector_length, 256); Ok(()) } @@ -313,7 +295,6 @@ mod tests { assert_eq!(options.search_list_size, 100); assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); assert_eq!(options.get_storage_type(), StorageType::Plain); - assert_eq!(options.pq_vector_length, 256); Ok(()) } } diff --git a/timescale_vector/src/access_method/plain_node.rs b/timescale_vector/src/access_method/plain_node.rs index 37859193..b63096ea 100644 --- a/timescale_vector/src/access_method/plain_node.rs +++ b/timescale_vector/src/access_method/plain_node.rs @@ -3,11 +3,10 @@ use std::pin::Pin; use pgrx::pg_sys::{InvalidBlockNumber, InvalidOffsetNumber}; use pgrx::*; use rkyv::vec::ArchivedVec; -use rkyv::{Archive, Archived, Deserialize, Serialize}; +use rkyv::{Archive, Deserialize, Serialize}; use timescale_vector_derive::{Readable, Writeable}; use super::neighbor_with_distance::NeighborWithDistance; -use super::pq_quantizer::PqVectorElement; use super::storage::ArchivedData; use crate::util::{ArchivedItemPointer, HeapPointer, ItemPointer, ReadableBuffer, WritableBuffer}; @@ -50,16 +49,6 @@ impl Node { let pq_vector = Vec::with_capacity(0); Self::new_internal(vector, pq_vector, heap_item_pointer, meta_page) } - - pub fn new_for_pq( - heap_item_pointer: ItemPointer, - pq_vector: Vec, - meta_page: &MetaPage, - ) -> Self { - let vector = Vec::with_capacity(0); - - Self::new_internal(vector, pq_vector, heap_item_pointer, meta_page) - } } /// contains helpers for mutate-in-place. See struct_mutable_refs in test_alloc.rs in rkyv @@ -81,10 +70,6 @@ impl ArchivedNode { unsafe { self.map_unchecked_mut(|s| &mut s.neighbor_index_pointers) } } - pub fn pq_vectors(self: Pin<&mut Self>) -> Pin<&mut Archived>> { - unsafe { self.map_unchecked_mut(|s| &mut s.pq_vector) } - } - pub fn num_neighbors(&self) -> usize { self.neighbor_index_pointers .iter() diff --git a/timescale_vector/src/access_method/pq_quantizer.rs b/timescale_vector/src/access_method/pq_quantizer.rs deleted file mode 100644 index dfb44005..00000000 --- a/timescale_vector/src/access_method/pq_quantizer.rs +++ /dev/null @@ -1,316 +0,0 @@ -use ndarray::{Array1, Array2, Axis}; -use pgrx::{error, notice, PgRelation}; -use rand::Rng; -use reductive::pq::{Pq, QuantizeVector, TrainPq}; - -use crate::access_method::{ - distance::distance_l2_optimized_for_few_dimensions, pq_quantizer_storage::read_pq, -}; - -use super::{ - meta_page::MetaPage, - pg_vector::PgVector, - stats::{StatsDistanceComparison, StatsNodeRead}, -}; - -/// pq aka Product quantization (PQ) is one of the most widely used algorithms for memory-efficient approximated nearest neighbor search, -/// This module encapsulates a vanilla implementation of PQ that we use for the vector index. -/// More details: https://lear.inrialpes.fr/pubs/2011/JDS11/jegou_searching_with_quantization.pdf - -/// PQ_TRAINING_ITERATIONS is the number of times we train each independent Kmeans cluster. -/// 20 - 40 iterations is considered an industry best practice -/// https://github.com/matsui528/nanopq/blob/main/nanopq/pq.py#L60 -const PQ_TRAINING_ITERATIONS: usize = 20; - -/// NUM_SUBQUANTIZER_BITS is the number of code words used for quantization. We pin it to 8 so we can -/// use u8 to represent a subspace. -const NUM_SUBQUANTIZER_BITS: u32 = 8; - -/// NUM_TRAINING_ATTEMPTS is the number of times we'll attempt to train the quantizer. -const NUM_TRAINING_ATTEMPTS: usize = 1; - -/// NUM_TRAINING_SET_SIZE is the maximum number of vectors we want to consider for the quantizer training set. -/// We pick a value used by DiskANN implementations. -const NUM_TRAINING_SET_SIZE: usize = 256000; - -pub type PqVectorElement = u8; -/// PqTrainer is a utility that produces a product quantizer from training with sample vectors. - -#[derive(Clone)] -pub struct PqTrainer { - /// training_set contains the vectors we'll use to train PQ. - training_set: Vec>, - /// considered_samples is the number of samples we considered for the training set. - /// It is useful for reservoir sampling as we add samples. - considered_samples: usize, - /// num_subquantizers is the number of independent kmeans we want to partition the vectors into. - /// the more we have the more accurate the PQ, but the more space we use in memory. - num_subquantizers: usize, - // rng is the random number generator for reservoir sampling - //rng: ThreadRng, -} - -impl PqTrainer { - pub fn new(meta_page: &super::meta_page::MetaPage) -> PqTrainer { - PqTrainer { - training_set: Vec::with_capacity(NUM_TRAINING_SET_SIZE), - num_subquantizers: meta_page.get_pq_vector_length(), - considered_samples: 0, - } - } - - /// add_sample adds vectors to the training set via uniform reservoir sampling to keep the - /// number of vectors within a reasonable memory limit. - pub fn add_sample(&mut self, sample: &[f32]) { - if self.training_set.len() >= NUM_TRAINING_SET_SIZE { - // TODO: Cache this somehow. - let mut rng = rand::thread_rng(); - let index = rng.gen_range(0..self.considered_samples + 1); - if index < NUM_TRAINING_SET_SIZE { - self.training_set[index] = sample.to_vec(); - } - } else { - self.training_set.push(sample.to_vec()); - } - self.considered_samples += 1; - } - - pub fn train_pq(self) -> Pq { - notice!( - "Training Product Quantization with {} vectors", - self.training_set.len() - ); - if (self.training_set.len() as i32) < (2_i32.pow(NUM_SUBQUANTIZER_BITS)) { - error!("training set is too small, please run with use_pq as false.") - } - let training_set = self - .training_set - .iter() - .map(|x| x.to_vec()) - .flatten() - .collect(); - let shape = (self.training_set.len(), self.training_set[0].len()); - let instances = Array2::::from_shape_vec(shape, training_set).unwrap(); - Pq::train_pq( - self.num_subquantizers, - NUM_SUBQUANTIZER_BITS, - PQ_TRAINING_ITERATIONS, - NUM_TRAINING_ATTEMPTS, - instances, - ) - .unwrap() - } -} - -/// build_distance_table produces an Asymmetric Distance Table to quickly compute distances. -/// We compute the distance from every centroid and cache that so actual distance calculations -/// can be fast. -// TODO: This function could return a table that fits in SIMD registers. -fn build_distance_table( - pq: &Pq, - query: &[f32], - _distance_fn: fn(&[f32], &[f32]) -> f32, -) -> Vec { - let sq = pq.subquantizers(); - let num_centroids = pq.n_quantizer_centroids(); - let num_subquantizers = sq.len_of(Axis(0)); - let dt_size = num_subquantizers * num_centroids; - let mut distance_table = vec![0.0; dt_size]; - - let ds = query.len() / num_subquantizers; - let mut elements_for_assert = 0; - for (subquantizer_index, subquantizer) in sq.outer_iter().enumerate() { - let sl = &query[subquantizer_index * ds..(subquantizer_index + 1) * ds]; - for (centroid_index, c) in subquantizer.outer_iter().enumerate() { - /* always use l2 for pq measurements since centeroids use k-means (which uses euclidean/l2 distance) - * The quantization also uses euclidean distance too. In the future we can experiment with k-mediods - * using a different distance measure, but this may make little difference. */ - let dist = distance_l2_optimized_for_few_dimensions(sl, c.to_slice().unwrap()); - assert!(subquantizer_index < num_subquantizers); - assert!(centroid_index * num_subquantizers + subquantizer_index < dt_size); - distance_table[centroid_index * num_subquantizers + subquantizer_index] = dist; - elements_for_assert += 1; - } - } - assert_eq!(dt_size, elements_for_assert); - distance_table -} - -/* -It seems that the node-node comparisons don't benefit from a table. remove this later if still true. -The node comparisons are done in the prune step where we do too few comparisons to benefit from a table. -fn build_distance_table_pq_query(pq: &Pq, query: &[u8]) -> Vec { - let sq = pq.subquantizers(); - let num_centroids = pq.n_quantizer_centroids(); - let num_subquantizers = sq.len_of(Axis(0)); - let dt_size = num_subquantizers * num_centroids; - let mut distance_table = vec![0.0; dt_size]; - - let mut elements_for_assert = 0; - for (subquantizer_index, subquantizer) in sq.outer_iter().enumerate() { - let query_centroid_index = query[subquantizer_index] as usize; - let query_slice = subquantizer.index_axis(Axis(0), query_centroid_index); - for (centroid_index, c) in subquantizer.outer_iter().enumerate() { - /* always use l2 for pq measurements since centeroids use k-means (which uses euclidean/l2 distance) - * The quantization also uses euclidean distance too. In the future we can experiment with k-mediods - * using a different distance measure, but this may make little difference. */ - let dist = if query_centroid_index == centroid_index { - debug_assert!(query_slice.as_slice().unwrap() == c.as_slice().unwrap()); - 0.0 - } else { - distance_l2_optimized_for_few_dimensions( - query_slice.as_slice().unwrap(), - c.as_slice().unwrap(), - ) - }; - assert!(subquantizer_index < num_subquantizers); - assert!(centroid_index * num_subquantizers + subquantizer_index < dt_size); - distance_table[centroid_index * num_subquantizers + subquantizer_index] = dist; - elements_for_assert += 1; - } - } - assert_eq!(dt_size, elements_for_assert); - distance_table -} -*/ - -#[derive(Clone)] -pub struct PqQuantizer { - pq_trainer: Option, - pq: Option>, -} - -impl PqQuantizer { - pub fn new() -> PqQuantizer { - Self { - pq_trainer: None, - pq: None, - } - } - - pub fn load( - index_relation: &PgRelation, - meta_page: &super::meta_page::MetaPage, - stats: &mut S, - ) -> Self { - let pq_item_pointer = meta_page.get_pq_pointer().unwrap(); - let pq = unsafe { Some(read_pq(&index_relation, pq_item_pointer, stats)) }; - - Self { - pq_trainer: None, - pq: pq, - } - } - - pub fn must_get_pq(&self) -> &Pq { - self.pq.as_ref().unwrap() - } - - pub fn quantize(&self, full_vector: &[f32]) -> Vec { - assert!(self.pq.is_some()); - let pq = self.pq.as_ref().unwrap(); - //OPT is the copy really necessary? - let array_vec = Array1::from(full_vector.to_vec()); - pq.quantize_vector(array_vec).to_vec() - } - - pub fn start_training(&mut self, meta_page: &super::meta_page::MetaPage) { - self.pq_trainer = Some(PqTrainer::new(meta_page)); - } - - pub fn add_sample(&mut self, sample: &[f32]) { - self.pq_trainer.as_mut().unwrap().add_sample(sample); - } - - pub fn finish_training(&mut self) { - self.pq = Some(self.pq_trainer.take().unwrap().train_pq()); - } - - pub fn vector_for_new_node( - &self, - meta_page: &MetaPage, - full_vector: &[f32], - ) -> Vec { - assert!(self.pq_trainer.is_none() && self.pq.is_some()); - let pq_vec_len = meta_page.get_pq_vector_length(); - let res = self.quantize(full_vector); - assert!(res.len() == pq_vec_len); - res - } - - pub fn get_distance_table_full_query( - &self, - query: &[f32], - distance_fn: fn(&[f32], &[f32]) -> f32, - ) -> PqDistanceTable { - PqDistanceTable::new_for_full_query(&self.pq.as_ref().unwrap(), distance_fn, query) - } - - 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. -pub struct PqDistanceTable { - distance_table: Vec, -} - -impl PqDistanceTable { - pub fn new_for_full_query( - pq: &Pq, - distance_fn: fn(&[f32], &[f32]) -> f32, - query: &[f32], - ) -> PqDistanceTable { - PqDistanceTable { - distance_table: build_distance_table(pq, query, distance_fn), - } - } - - /// distance emits the sum of distances between each centroid in the quantized vector. - pub fn distance(&self, pq_vector: &[u8]) -> f32 { - let mut d = 0.0; - let num_subquantizers = pq_vector.len(); - // maybe we should unroll this loop? - for subquantizer_index in 0..num_subquantizers { - let centroid_index = pq_vector[subquantizer_index] as usize; - d += self.distance_table[centroid_index * num_subquantizers + subquantizer_index] - //d += self.distance_table[m][pq_vector[m] as usize]; - } - d - } -} - -pub enum PqSearchDistanceMeasure { - Pq(PqDistanceTable, PgVector), -} - -impl PqSearchDistanceMeasure { - pub fn calculate_pq_distance( - table: &PqDistanceTable, - pq_vector: &[PqVectorElement], - stats: &mut S, - ) -> f32 { - assert!(pq_vector.len() > 0); - let vec = pq_vector; - stats.record_quantized_distance_comparison(); - table.distance(vec) - } -} diff --git a/timescale_vector/src/access_method/pq_quantizer_storage.rs b/timescale_vector/src/access_method/pq_quantizer_storage.rs deleted file mode 100644 index b00a6318..00000000 --- a/timescale_vector/src/access_method/pq_quantizer_storage.rs +++ /dev/null @@ -1,122 +0,0 @@ -use std::mem::size_of; - -use ndarray::Array3; -use pgrx::pg_sys::BLCKSZ; -use pgrx::*; -use reductive::pq::Pq; - -use rkyv::{Archive, Deserialize, Serialize}; -use timescale_vector_derive::{Readable, Writeable}; - -use crate::util::page::PageType; -use crate::util::tape::Tape; -use crate::util::{IndexPointer, ItemPointer, ReadableBuffer, WritableBuffer}; - -use super::stats::{StatsNodeRead, StatsNodeWrite}; - -#[derive(Archive, Deserialize, Serialize, Readable, Writeable)] -#[archive(check_bytes)] -#[repr(C)] -pub struct PqQuantizerDef { - dim_0: usize, - dim_1: usize, - dim_2: usize, - vec_len: usize, - next_vector_pointer: ItemPointer, -} - -impl PqQuantizerDef { - pub fn new(dim_0: usize, dim_1: usize, dim_2: usize, vec_len: usize) -> PqQuantizerDef { - { - Self { - dim_0, - dim_1, - dim_2, - vec_len, - next_vector_pointer: ItemPointer { - block_number: 0, - offset: 0, - }, - } - } - } -} - -#[derive(Archive, Deserialize, Serialize, Readable, Writeable)] -#[archive(check_bytes)] -#[repr(C)] -pub struct PqQuantizerVector { - vec: Vec, - next_vector_pointer: ItemPointer, -} - -pub unsafe fn read_pq( - index: &PgRelation, - index_pointer: IndexPointer, - stats: &mut S, -) -> Pq { - //TODO: handle stats better - let rpq = PqQuantizerDef::read(index, index_pointer, stats); - stats.record_read(); - let rpn = rpq.get_archived_node(); - let size = rpn.dim_0 * rpn.dim_1 * rpn.dim_2; - let mut result: Vec = Vec::with_capacity(size as usize); - let mut next = rpn.next_vector_pointer.deserialize_item_pointer(); - loop { - if next.offset == 0 && next.block_number == 0 { - break; - } - let qvn = PqQuantizerVector::read(index, next, stats); - let vn = qvn.get_archived_node(); - result.extend(vn.vec.iter()); - next = vn.next_vector_pointer.deserialize_item_pointer(); - } - let sq = Array3::from_shape_vec( - (rpn.dim_0 as usize, rpn.dim_1 as usize, rpn.dim_2 as usize), - result, - ) - .unwrap(); - Pq::new(None, sq) -} - -pub unsafe fn write_pq( - pq: &Pq, - index: &PgRelation, - stats: &mut S, -) -> ItemPointer { - let vec = pq.subquantizers().to_slice_memory_order().unwrap().to_vec(); - let shape = pq.subquantizers().dim(); - let mut pq_node = PqQuantizerDef::new(shape.0, shape.1, shape.2, vec.len()); - - let mut pqt = Tape::new(index, PageType::PqQuantizerDef); - - // write out the large vector bits. - // we write "from the back" - let mut prev: IndexPointer = ItemPointer { - block_number: 0, - offset: 0, - }; - let mut prev_vec = vec; - - // get numbers that can fit in a page by subtracting the item pointer. - let block_fit = (BLCKSZ as usize / size_of::()) - size_of::() - 64; - let mut tape = Tape::new(index, PageType::PqQuantizerVector); - loop { - let l = prev_vec.len(); - if l == 0 { - pq_node.next_vector_pointer = prev; - return pq_node.write(&mut pqt, stats); - } - let lv = prev_vec; - let ni = if l > block_fit { l - block_fit } else { 0 }; - let (b, a) = lv.split_at(ni); - - let pqv_node = PqQuantizerVector { - vec: a.to_vec(), - next_vector_pointer: prev, - }; - let index_pointer: IndexPointer = pqv_node.write(&mut tape, stats); - prev = index_pointer; - prev_vec = b.to_vec(); - } -} diff --git a/timescale_vector/src/access_method/pq_storage.rs b/timescale_vector/src/access_method/pq_storage.rs deleted file mode 100644 index d3fae389..00000000 --- a/timescale_vector/src/access_method/pq_storage.rs +++ /dev/null @@ -1,418 +0,0 @@ -use super::{ - graph::{ListSearchNeighbor, ListSearchResult}, - graph_neighbor_store::GraphNeighborStore, - pg_vector::PgVector, - plain_node::{ArchivedNode, Node}, - plain_storage::PlainStorageLsnPrivateData, - pq_quantizer::{PqQuantizer, PqSearchDistanceMeasure, PqVectorElement}, - pq_quantizer_storage::write_pq, - stats::{ - GreedySearchStats, StatsDistanceComparison, StatsHeapNodeRead, StatsNodeModify, - StatsNodeRead, StatsNodeWrite, WriteStats, - }, - storage::{NodeDistanceMeasure, Storage}, - storage_common::get_attribute_number_from_index, -}; - -use pgrx::PgRelation; - -use crate::util::{ - page::PageType, table_slot::TableSlot, tape::Tape, HeapPointer, IndexPointer, ItemPointer, -}; - -use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance}; - -pub struct PqNodeDistanceMeasure<'a> { - storage: &'a PqCompressionStorage<'a>, - vector: Vec, -} - -impl<'a> PqNodeDistanceMeasure<'a> { - pub unsafe fn with_index_pointer( - 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( - &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> { - pub index: &'a PgRelation, - pub distance_fn: fn(&[f32], &[f32]) -> f32, - quantizer: PqQuantizer, - heap_rel: &'a PgRelation, - heap_attr: pgrx::pg_sys::AttrNumber, -} - -impl<'a> PqCompressionStorage<'a> { - pub fn new_for_build( - index: &'a PgRelation, - heap_rel: &'a PgRelation, - heap_attr: pgrx::pg_sys::AttrNumber, - distance_fn: fn(&[f32], &[f32]) -> f32, - ) -> PqCompressionStorage<'a> { - Self { - index: index, - distance_fn, - quantizer: PqQuantizer::new(), - heap_rel: heap_rel, - heap_attr: heap_attr, - } - } - - fn load_quantizer( - index_relation: &PgRelation, - meta_page: &super::meta_page::MetaPage, - stats: &mut S, - ) -> PqQuantizer { - PqQuantizer::load(&index_relation, meta_page, stats) - } - - pub fn load_for_insert( - heap_rel: &'a PgRelation, - heap_attr: pgrx::pg_sys::AttrNumber, - index_relation: &'a PgRelation, - meta_page: &super::meta_page::MetaPage, - stats: &mut S, - ) -> PqCompressionStorage<'a> { - Self { - index: index_relation, - distance_fn: meta_page.get_distance_function(), - quantizer: Self::load_quantizer(index_relation, meta_page, stats), - heap_rel: heap_rel, - heap_attr: heap_attr, - } - } - - pub fn load_for_search( - index_relation: &'a PgRelation, - heap_relation: &'a PgRelation, - quantizer: &PqQuantizer, - distance_fn: fn(&[f32], &[f32]) -> f32, - ) -> PqCompressionStorage<'a> { - Self { - index: index_relation, - distance_fn, - //OPT: get rid of clone - quantizer: quantizer.clone(), - heap_rel: heap_relation, - heap_attr: get_attribute_number_from_index(index_relation), - } - } - - fn write_quantizer_metadata(&self, stats: &mut S) { - let pq = self.quantizer.must_get_pq(); - let index_pointer: IndexPointer = unsafe { write_pq(pq, &self.index, stats) }; - super::meta_page::MetaPage::update_pq_pointer(&self.index, index_pointer, stats); - } - - fn visit_lsn_internal( - &self, - lsr: &mut ListSearchResult< - as Storage>::QueryDistanceMeasure, - as Storage>::LSNPrivateData, - >, - neighbors: &[ItemPointer], - gns: &GraphNeighborStore, - ) { - for &neighbor_index_pointer in neighbors.iter() { - if !lsr.prepare_insert(neighbor_index_pointer) { - continue; - } - - let rn_neighbor = - unsafe { Node::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() { - PqSearchDistanceMeasure::Pq(table, _) => { - PqSearchDistanceMeasure::calculate_pq_distance( - table, - node_neighbor.pq_vector.as_slice(), - &mut lsr.stats, - ) - } - }; - let lsn = ListSearchNeighbor::new( - neighbor_index_pointer, - distance, - PlainStorageLsnPrivateData::new(neighbor_index_pointer, node_neighbor, gns), - ); - - lsr.insert_neighbor(lsn); - } - } - - unsafe fn get_heap_table_slot_from_heap_pointer( - &self, - heap_pointer: HeapPointer, - stats: &mut T, - ) -> TableSlot { - TableSlot::new(self.heap_rel, heap_pointer, self.heap_attr, stats) - } -} - -impl<'a> Storage for PqCompressionStorage<'a> { - type QueryDistanceMeasure = PqSearchDistanceMeasure; - type NodeDistanceMeasure<'b> = PqNodeDistanceMeasure<'b> where Self: 'b; - type ArchivedType = ArchivedNode; - type LSNPrivateData = PlainStorageLsnPrivateData; //no data stored - - fn page_type() -> PageType { - PageType::Node - } - - fn create_node( - &self, - full_vector: &[f32], - heap_pointer: HeapPointer, - meta_page: &MetaPage, - tape: &mut Tape, - stats: &mut S, - ) -> ItemPointer { - let pq_vector = self.quantizer.vector_for_new_node(meta_page, full_vector); - let node = Node::new_for_pq(heap_pointer, pq_vector, meta_page); - let index_pointer: IndexPointer = node.write(tape, stats); - index_pointer - } - - fn start_training(&mut self, meta_page: &super::meta_page::MetaPage) { - self.quantizer.start_training(meta_page); - } - - fn add_sample(&mut self, sample: &[f32]) { - self.quantizer.add_sample(sample); - } - - fn finish_training(&mut self, stats: &mut WriteStats) { - self.quantizer.finish_training(); - self.write_quantizer_metadata(stats); - } - - fn finalize_node_at_end_of_build( - &mut self, - meta: &MetaPage, - index_pointer: IndexPointer, - neighbors: &Vec, - stats: &mut S, - ) { - let node = unsafe { Node::modify(self.index, index_pointer, stats) }; - let mut archived = node.get_archived_node(); - archived.as_mut().set_neighbors(neighbors, &meta); - - node.commit(); - } - - unsafe fn get_node_distance_measure<'b, S: StatsNodeRead>( - &'b self, - index_pointer: IndexPointer, - stats: &mut S, - ) -> PqNodeDistanceMeasure<'b> { - PqNodeDistanceMeasure::with_index_pointer(self, index_pointer, stats) - } - - fn get_query_distance_measure(&self, query: PgVector) -> PqSearchDistanceMeasure { - return PqSearchDistanceMeasure::Pq( - self.quantizer - .get_distance_table_full_query(query.to_slice(), self.distance_fn), - query, - ); - } - - fn get_full_distance_for_resort( - &self, - qdm: &Self::QueryDistanceMeasure, - _index_pointer: IndexPointer, - heap_pointer: HeapPointer, - stats: &mut S, - ) -> f32 { - let slot = unsafe { self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) }; - match qdm { - PqSearchDistanceMeasure::Pq(_, query) => self.get_distance_function()( - unsafe { slot.get_pg_vector().to_slice() }, - query.to_slice(), - ), - } - } - - //todo: same as Bq code? - fn get_neighbors_with_distances_from_disk( - &self, - neighbors_of: ItemPointer, - result: &mut Vec, - stats: &mut S, - ) { - let rn = unsafe { Node::read(self.index, neighbors_of, stats) }; - let dist_state = unsafe { self.get_node_distance_measure(neighbors_of, stats) }; - for n in rn.get_archived_node().iter_neighbors() { - let dist = unsafe { dist_state.get_distance(n, stats) }; - result.push(NeighborWithDistance::new(n, dist)) - } - } - - /* get_lsn and visit_lsn are different because the distance - comparisons for BQ get the vector from different places */ - fn create_lsn_for_init_id( - &self, - lsr: &mut ListSearchResult, - index_pointer: ItemPointer, - gns: &GraphNeighborStore, - ) -> ListSearchNeighbor { - if !lsr.prepare_insert(index_pointer) { - panic!("should not have had an init id already inserted"); - } - - let rn = unsafe { Node::read(self.index, index_pointer, &mut lsr.stats) }; - let node = rn.get_archived_node(); - - let distance = match lsr.sdm.as_ref().unwrap() { - PqSearchDistanceMeasure::Pq(table, _) => { - PqSearchDistanceMeasure::calculate_pq_distance( - table, - node.pq_vector.as_slice(), - &mut lsr.stats, - ) - } - }; - - ListSearchNeighbor::new( - index_pointer, - distance, - PlainStorageLsnPrivateData::new(index_pointer, node, gns), - ) - } - - fn visit_lsn( - &self, - lsr: &mut ListSearchResult, - lsn_idx: usize, - gns: &GraphNeighborStore, - ) { - let lsn = lsr.get_lsn_by_idx(lsn_idx); - //clone needed so we don't continue to borrow lsr - self.visit_lsn_internal(lsr, &lsn.get_private_data().neighbors.clone(), gns); - } - - fn return_lsn( - &self, - lsn: &ListSearchNeighbor, - _stats: &mut GreedySearchStats, - ) -> HeapPointer { - lsn.get_private_data().heap_pointer - } - - fn set_neighbors_on_disk( - &self, - meta: &MetaPage, - index_pointer: IndexPointer, - neighbors: &[NeighborWithDistance], - stats: &mut S, - ) { - let node = unsafe { Node::modify(self.index, index_pointer, stats) }; - let mut archived = node.get_archived_node(); - archived.as_mut().set_neighbors(neighbors, &meta); - node.commit(); - } - - fn get_distance_function(&self) -> fn(&[f32], &[f32]) -> f32 { - self.distance_fn - } -} - -#[cfg(any(test, feature = "pg_test"))] -#[pgrx::pg_schema] -mod tests { - use pgrx::*; - - #[pg_test] - unsafe fn test_pq_storage_index_creation_default() -> spi::Result<()> { - crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( - "num_neighbors=38, USE_PQ = TRUE", - )?; - Ok(()) - } - - #[pg_test] - unsafe fn test_pq_storage_index_creation_few_neighbors() -> spi::Result<()> { - //a test with few neighbors tests the case that nodes share a page, which has caused deadlocks in the past. - crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( - "num_neighbors=10, USE_PQ = TRUE", - )?; - Ok(()) - } - - #[test] - fn test_pq_storage_delete_vacuum_plain() { - crate::access_method::vacuum::tests::test_delete_vacuum_plain_scaffold( - "num_neighbors = 38, use_pq = TRUE", - ); - } - - #[test] - fn test_pq_storage_delete_vacuum_full() { - crate::access_method::vacuum::tests::test_delete_vacuum_full_scaffold( - "num_neighbors = 38, use_pq = TRUE", - ); - } - - /* can't run test_pq_storage_empty_table_insert because can't create pq index on pq table */ - - #[pg_test] - unsafe fn test_pq_storage_insert_empty_insert() -> spi::Result<()> { - let suffix = (1..=253) - .map(|i| format!("{}", i)) - .collect::>() - .join(", "); - - Spi::run(&format!( - "CREATE TABLE test(embedding vector(256)); - - INSERT INTO test (embedding) - SELECT - ('[' || i || ',2,3,{suffix}]')::vector - FROM generate_series(1, 300) i; - - CREATE INDEX idxtest - ON test - USING tsv(embedding) - WITH (num_neighbors = 10, use_pq = TRUE); - - DELETE FROM test; - - INSERT INTO test(embedding) VALUES ('[1,2,3,{suffix}]'), ('[14,15,16,{suffix}]'); - ", - ))?; - - let res: Option = Spi::get_one(&format!( - " set enable_seqscan = 0; - WITH cte as (select * from test order by embedding <=> '[0,0,0,{suffix}]') SELECT count(*) from cte;", - ))?; - assert_eq!(2, res.unwrap()); - - Spi::run(&format!("drop index idxtest;",))?; - - Ok(()) - } -} diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index a3fcaf9b..1e1baee1 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -14,8 +14,6 @@ use super::{ bq::{BqMeans, BqQuantizer, BqSearchDistanceMeasure, BqSpeedupStorageLsnPrivateData}, graph::{Graph, ListSearchResult}, plain_storage::{PlainDistanceMeasure, PlainStorage, PlainStorageLsnPrivateData}, - pq_quantizer::{PqQuantizer, PqSearchDistanceMeasure}, - pq_storage::PqCompressionStorage, stats::QuantizerStats, storage::{Storage, StorageType}, }; @@ -27,10 +25,6 @@ enum StorageState { BqQuantizer, TSVResponseIterator, ), - PqCompression( - PqQuantizer, - TSVResponseIterator, - ), Plain(TSVResponseIterator), } @@ -67,19 +61,6 @@ impl TSVScanState { TSVResponseIterator::new(&bq, index, query, search_list_size, meta_page, stats); StorageState::Plain(it) } - StorageType::PqCompression => { - let mut stats = QuantizerStats::new(); - let quantizer = PqQuantizer::load(index, &meta_page, &mut stats); - let pq = PqCompressionStorage::load_for_search( - index, - heap, - &quantizer, - meta_page.get_distance_function(), - ); - let it = - TSVResponseIterator::new(&pq, index, query, search_list_size, meta_page, stats); - StorageState::PqCompression(quantizer, it) - } StorageType::BqSpeedup => { let mut stats = QuantizerStats::new(); let quantizer = unsafe { BqMeans::load(index, &meta_page, &mut stats) }; @@ -304,17 +285,6 @@ pub extern "C" fn amrescan( let state = unsafe { (scan.opaque as *mut TSVScanState).as_mut() }.expect("no scandesc state"); state.initialize(&indexrel, &heaprel, query, search_list_size); - /*match &mut storage { - Storage::None => pgrx::error!("not implemented"), - Storage::PQ(_pq) => pgrx::error!("not implemented"), - Storage::BQ(_bq) => { - let state = - unsafe { (scan.opaque as *mut TSVScanState).as_mut() }.expect("no scandesc state"); - - let res = TSVResponseIterator::new(&indexrel, query, search_list_size); - state.iterator = PgMemoryContexts::CurrentMemoryContext.leak_and_drop_on_delete(res); - } - }*/ } #[pg_guard] @@ -341,16 +311,6 @@ pub extern "C" fn amgettuple( let next = iter.next_with_resort(&indexrel, &bq); get_tuple(next, scan) } - StorageState::PqCompression(quantizer, iter) => { - let pq = PqCompressionStorage::load_for_search( - &indexrel, - &heaprel, - quantizer, - state.distance_fn.unwrap(), - ); - let next = iter.next_with_resort(&indexrel, &pq); - get_tuple(next, scan) - } StorageState::Plain(iter) => { let storage = PlainStorage::load_for_search(&indexrel, state.distance_fn.unwrap()); let next = iter.next(&indexrel, &storage); @@ -389,7 +349,6 @@ pub extern "C" fn amendscan(scan: pg_sys::IndexScanDesc) { let mut storage = unsafe { state.storage.as_mut() }.expect("no storage in state"); match &mut storage { StorageState::BqSpeedup(_bq, iter) => end_scan::(iter), - StorageState::PqCompression(_pq, iter) => end_scan::(iter), StorageState::Plain(iter) => end_scan::(iter), } } diff --git a/timescale_vector/src/access_method/storage.rs b/timescale_vector/src/access_method/storage.rs index 4692c4f1..e8dd23f9 100644 --- a/timescale_vector/src/access_method/storage.rs +++ b/timescale_vector/src/access_method/storage.rs @@ -127,7 +127,6 @@ pub trait Storage { pub enum StorageType { Plain = 0, BqSpeedup = 1, - PqCompression = 2, } impl StorageType { @@ -135,7 +134,6 @@ impl StorageType { match value { 0 => StorageType::Plain, 1 => StorageType::BqSpeedup, - 2 => StorageType::PqCompression, _ => panic!("Invalid storage type"), } } diff --git a/timescale_vector/src/access_method/vacuum.rs b/timescale_vector/src/access_method/vacuum.rs index cf9260d7..cd3104e1 100644 --- a/timescale_vector/src/access_method/vacuum.rs +++ b/timescale_vector/src/access_method/vacuum.rs @@ -4,10 +4,7 @@ use pgrx::{ }; use crate::{ - access_method::{ - bq::BqSpeedupStorage, meta_page::MetaPage, plain_storage::PlainStorage, - pq_storage::PqCompressionStorage, - }, + access_method::{bq::BqSpeedupStorage, meta_page::MetaPage, plain_storage::PlainStorage}, util::{ page::WritablePage, ports::{PageGetItem, PageGetItemId, PageGetMaxOffsetNumber}, @@ -52,15 +49,6 @@ pub extern "C" fn ambulkdelete( callback_state, ); } - StorageType::PqCompression => { - bulk_delete_for_storage::( - &index_relation, - nblocks, - results, - callback, - callback_state, - ); - } StorageType::Plain => { bulk_delete_for_storage::( &index_relation, From b98d7413011621da00af72c9f47a153dcb33e8c2 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Thu, 4 Apr 2024 14:34:38 -0400 Subject: [PATCH 25/44] Meta page cleanup --- .../src/access_method/meta_page.rs | 46 +++++++------------ timescale_vector/src/util/mod.rs | 5 ++ 2 files changed, 21 insertions(+), 30 deletions(-) diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index f48a3115..bfeda410 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -1,4 +1,4 @@ -use pgrx::pg_sys::BufferGetBlockNumber; +use pgrx::pg_sys::{BufferGetBlockNumber, InvalidBlockNumber, InvalidOffsetNumber}; use pgrx::*; use rkyv::{Archive, Deserialize, Serialize}; use semver::Version; @@ -39,8 +39,8 @@ pub struct MetaPageV1 { init_ids_offset: pg_sys::OffsetNumber, use_pq: bool, _pq_vector_length: usize, - pq_block_number: pg_sys::BlockNumber, - pq_block_offset: pg_sys::OffsetNumber, + _pq_block_number: pg_sys::BlockNumber, + _pq_block_offset: pg_sys::OffsetNumber, } impl MetaPageV1 { @@ -69,10 +69,8 @@ impl MetaPageV1 { storage_type: StorageType::Plain as u8, search_list_size: self.search_list_size, max_alpha: self.max_alpha, - init_ids_block_number: self.init_ids_block_number, - init_ids_offset: self.init_ids_offset, - quantizer_metadata_block_number: self.pq_block_number, - quantizer_metadata_block_offset: self.pq_block_offset, + init_ids: ItemPointer::new(self.init_ids_block_number, self.init_ids_offset), + quantizer_metadata: ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber), } } } @@ -90,7 +88,7 @@ pub struct MetaPageHeader { version: u32, } -pub enum DistanceType { +enum DistanceType { Cosine = 0, L2 = 1, } @@ -124,10 +122,8 @@ pub struct MetaPage { num_neighbors: u32, search_list_size: u32, max_alpha: f64, - init_ids_block_number: pg_sys::BlockNumber, - init_ids_offset: pg_sys::OffsetNumber, - quantizer_metadata_block_number: pg_sys::BlockNumber, - quantizer_metadata_block_offset: pg_sys::OffsetNumber, + init_ids: ItemPointer, + quantizer_metadata: ItemPointer, } impl MetaPage { @@ -167,27 +163,21 @@ impl MetaPage { } pub fn get_init_ids(&self) -> Option> { - if self.init_ids_block_number == 0 && self.init_ids_offset == 0 { + if !self.init_ids.is_valid() { return None; } - let ptr = HeapPointer::new(self.init_ids_block_number, self.init_ids_offset); - Some(vec![ptr]) + Some(vec![self.init_ids]) } pub fn get_quantizer_metadata_pointer(&self) -> Option { if (self.storage_type != StorageType::BqSpeedup as u8) - || (self.quantizer_metadata_block_number == 0 - && self.quantizer_metadata_block_offset == 0) + || !self.quantizer_metadata.is_valid() { return None; } - let ptr = IndexPointer::new( - self.quantizer_metadata_block_number, - self.quantizer_metadata_block_offset, - ); - Some(ptr) + Some(self.quantizer_metadata) } fn calculate_num_neighbors(num_dimensions: u32, opt: &PgBox) -> u32 { @@ -222,10 +212,8 @@ impl MetaPage { num_neighbors: Self::calculate_num_neighbors(num_dimensions, &opt), search_list_size: (*opt).search_list_size, max_alpha: (*opt).max_alpha, - init_ids_block_number: 0, - init_ids_offset: 0, - quantizer_metadata_block_number: 0, - quantizer_metadata_block_offset: 0, + init_ids: ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber), + quantizer_metadata: ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber), }; let page = page::WritablePage::new(index, crate::util::page::PageType::Meta); meta.write_to_page(page); @@ -322,8 +310,7 @@ impl MetaPage { let id = init_ids[0]; let mut meta = Self::fetch(index); - meta.init_ids_block_number = id.block_number; - meta.init_ids_offset = id.offset; + meta.init_ids = id; unsafe { Self::overwrite(index, &meta); @@ -337,8 +324,7 @@ impl MetaPage { stats: &mut S, ) { let mut meta = Self::fetch(index); - meta.quantizer_metadata_block_number = quantizer_pointer.block_number; - meta.quantizer_metadata_block_offset = quantizer_pointer.offset; + meta.quantizer_metadata = quantizer_pointer; unsafe { Self::overwrite(index, &meta); diff --git a/timescale_vector/src/util/mod.rs b/timescale_vector/src/util/mod.rs index dc72524c..37c580c0 100644 --- a/timescale_vector/src/util/mod.rs +++ b/timescale_vector/src/util/mod.rs @@ -69,6 +69,11 @@ impl ItemPointer { } } + pub fn is_valid(&self) -> bool { + self.block_number != pgrx::pg_sys::InvalidBlockNumber + && self.offset != pgrx::pg_sys::InvalidOffsetNumber + } + pub unsafe fn with_page(page: &page::WritablePage, offset: pgrx::pg_sys::OffsetNumber) -> Self { Self { block_number: pgrx::pg_sys::BufferGetBlockNumber(**(page.get_buffer())), From 904d1928dbfbdea00b3051fa2e27f32361ae6411 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Thu, 4 Apr 2024 15:52:00 -0400 Subject: [PATCH 26/44] change resort->rescore --- timescale_vector/src/access_method/guc.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/timescale_vector/src/access_method/guc.rs b/timescale_vector/src/access_method/guc.rs index 6a83b1bb..3275d221 100644 --- a/timescale_vector/src/access_method/guc.rs +++ b/timescale_vector/src/access_method/guc.rs @@ -1,7 +1,7 @@ use pgrx::*; pub static TSV_QUERY_SEARCH_LIST_SIZE: GucSetting = GucSetting::::new(100); -pub static TSV_RESORT_SIZE: GucSetting = GucSetting::::new(10); +pub static TSV_RESORT_SIZE: GucSetting = GucSetting::::new(50); pub fn init() { GucRegistry::define_int_guc( @@ -16,9 +16,9 @@ pub fn init() { ); GucRegistry::define_int_guc( - "tsv.query_resort", - "The resort size used in queries", - "Resort size.", + "tsv.query_rescore", + "The number of elements rescored (0 to disable rescoring)", + "Rescoring takes the query_rescore number of elements that have the smallest approximate distance, rescores them with the exact distance, returning the closest ones with the exact distance.", &TSV_RESORT_SIZE, 1, 1000, From f7d72df863161884ce6290de13e60bb9a8e3e73d Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Fri, 5 Apr 2024 16:22:55 -0400 Subject: [PATCH 27/44] Support matryoshka embeddings with the num_dimensions option --- timescale_vector/src/access_method/bq.rs | 31 +++--- timescale_vector/src/access_method/build.rs | 77 ++++++++------- .../src/access_method/meta_page.rs | 18 +++- timescale_vector/src/access_method/options.rs | 50 +++++++++- .../src/access_method/pg_vector.rs | 95 ++++++++++++++++--- .../src/access_method/plain_storage.rs | 49 ++++++++-- timescale_vector/src/access_method/scan.rs | 39 ++++++-- timescale_vector/src/access_method/storage.rs | 1 + timescale_vector/src/util/table_slot.rs | 15 +-- 9 files changed, 284 insertions(+), 91 deletions(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index a567a750..3f40dfaf 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -137,7 +137,7 @@ impl BqQuantizer { self.training = true; if self.use_mean { self.count = 0; - self.mean = vec![0.0; meta_page.get_num_dimensions() as _]; + self.mean = vec![0.0; meta_page.get_num_dimensions_to_index() as _]; } } @@ -186,7 +186,6 @@ impl BqDistanceTable { } } - /// distance emits the sum of distances between each centroid in the quantized vector. pub fn distance(&self, bq_vector: &[BqVectorElement]) -> f32 { let count_ones = distance_xor_optimized(&self.quantized_vector, bq_vector); //dot product is LOWER the more xors that lead to 1 becaues that means a negative times a positive = negative component @@ -311,7 +310,6 @@ impl<'a> BqSpeedupStorage<'a> { pub fn new_for_build( index: &'a PgRelation, heap_rel: &'a PgRelation, - heap_attr: pgrx::pg_sys::AttrNumber, distance_fn: fn(&[f32], &[f32]) -> f32, ) -> BqSpeedupStorage<'a> { Self { @@ -319,7 +317,7 @@ impl<'a> BqSpeedupStorage<'a> { distance_fn: distance_fn, quantizer: BqQuantizer::new(), heap_rel: heap_rel, - heap_attr: heap_attr, + heap_attr: get_attribute_number_from_index(index), qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), } } @@ -334,7 +332,6 @@ impl<'a> BqSpeedupStorage<'a> { pub fn load_for_insert( heap_rel: &'a PgRelation, - heap_attr: pgrx::pg_sys::AttrNumber, index_relation: &'a PgRelation, meta_page: &super::meta_page::MetaPage, stats: &mut S, @@ -344,7 +341,7 @@ impl<'a> BqSpeedupStorage<'a> { distance_fn: meta_page.get_distance_function(), quantizer: Self::load_quantizer(index_relation, meta_page, stats), heap_rel: heap_rel, - heap_attr: heap_attr, + heap_attr: get_attribute_number_from_index(index_relation), qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), } } @@ -450,7 +447,7 @@ impl<'a> BqSpeedupStorage<'a> { heap_pointer: HeapPointer, stats: &mut T, ) -> TableSlot { - TableSlot::new(self.heap_rel, heap_pointer, self.heap_attr, stats) + TableSlot::new(self.heap_rel, heap_pointer, stats) } } @@ -529,7 +526,7 @@ impl<'a> Storage for BqSpeedupStorage<'a> { fn get_query_distance_measure(&self, query: PgVector) -> BqSearchDistanceMeasure { return BqSearchDistanceMeasure::Bq( self.quantizer - .get_distance_table(query.to_slice(), self.distance_fn), + .get_distance_table(query.to_index_slice(), self.distance_fn), query, ); } @@ -539,14 +536,16 @@ impl<'a> Storage for BqSpeedupStorage<'a> { qdm: &Self::QueryDistanceMeasure, _index_pointer: IndexPointer, heap_pointer: HeapPointer, + meta_page: &MetaPage, stats: &mut S, ) -> f32 { let slot = unsafe { self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) }; match qdm { - BqSearchDistanceMeasure::Bq(_, query) => self.get_distance_function()( - unsafe { slot.get_pg_vector().to_slice() }, - query.to_slice(), - ), + BqSearchDistanceMeasure::Bq(_, query) => { + let datum = unsafe { slot.get_attribute(self.heap_attr).unwrap() }; + let vec = unsafe { PgVector::from_datum(datum, meta_page, false, true) }; + self.get_distance_function()(vec.to_full_slice(), query.to_full_slice()) + } } } @@ -868,4 +867,12 @@ mod tests { "num_neighbors=38, storage_layout = io_optimized", ) } + + #[pg_test] + unsafe fn test_bq_storage_index_creation_num_dimensions() -> spi::Result<()> { + crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( + "storage_layout = io_optimized, num_dimensions=768", + )?; + Ok(()) + } } diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index fb5afb3a..b8152fe0 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -44,7 +44,6 @@ impl<'a, 'b> BuildState<'a, 'b> { ) -> Self { let tape = unsafe { Tape::new(index_relation, page_type) }; - //TODO: some ways to get rid of meta_page.clone? BuildState { memcxt: PgMemoryContexts::new("tsv build context"), ntuples: 0, @@ -97,25 +96,28 @@ pub unsafe extern "C" fn aminsert( heaprel: pg_sys::Relation, _check_unique: pg_sys::IndexUniqueCheck, _index_unchanged: bool, - index_info: *mut pg_sys::IndexInfo, + _index_info: *mut pg_sys::IndexInfo, ) -> bool { let index_relation = unsafe { PgRelation::from_pg(indexrel) }; let heap_relation = unsafe { PgRelation::from_pg(heaprel) }; - let vec = PgVector::from_pg_parts(values, isnull, 0); + let mut meta_page = MetaPage::fetch(&index_relation); + let vec = PgVector::from_pg_parts(values, isnull, 0, &meta_page, true, false); if let None = vec { //todo handle NULLs? return false; } let vec = vec.unwrap(); let heap_pointer = ItemPointer::with_item_pointer_data(*heap_tid); - let mut meta_page = MetaPage::fetch(&index_relation); let mut storage = meta_page.get_storage_type(); let mut stats = InsertStats::new(); match &mut storage { StorageType::Plain => { - let plain = - PlainStorage::load_for_insert(&index_relation, meta_page.get_distance_function()); + let plain = PlainStorage::load_for_insert( + &index_relation, + &heap_relation, + meta_page.get_distance_function(), + ); insert_storage( &plain, &index_relation, @@ -128,7 +130,6 @@ pub unsafe extern "C" fn aminsert( StorageType::BqSpeedup => { let bq = BqSpeedupStorage::load_for_insert( &heap_relation, - get_attribute_number(index_info), &index_relation, &meta_page, &mut stats.quantizer_stats, @@ -156,7 +157,7 @@ unsafe fn insert_storage( ) { let mut tape = Tape::new(&index_relation, S::page_type()); let index_pointer = storage.create_node( - vector.to_slice(), + vector.to_index_slice(), heap_pointer, &meta_page, &mut tape, @@ -172,11 +173,6 @@ pub extern "C" fn ambuildempty(_index_relation: pg_sys::Relation) { panic!("ambuildempty: not yet implemented") } -pub fn get_attribute_number(index_info: *mut pg_sys::IndexInfo) -> pg_sys::AttrNumber { - unsafe { assert!((*index_info).ii_NumIndexAttrs == 1) }; - unsafe { (*index_info).ii_IndexAttrNumbers[0] } -} - fn do_heap_scan<'a>( index_info: *mut pg_sys::IndexInfo, heap_relation: &'a PgRelation, @@ -193,8 +189,11 @@ fn do_heap_scan<'a>( let mut write_stats = WriteStats::new(); match storage { StorageType::Plain => { - let mut plain = - PlainStorage::new_for_build(index_relation, meta_page.get_distance_function()); + let mut plain = PlainStorage::new_for_build( + index_relation, + heap_relation, + meta_page.get_distance_function(), + ); plain.start_training(&meta_page); let page_type = PlainStorage::page_type(); let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); @@ -216,23 +215,27 @@ fn do_heap_scan<'a>( let mut bq = BqSpeedupStorage::new_for_build( index_relation, heap_relation, - get_attribute_number(index_info), meta_page.get_distance_function(), ); + + let page_type = BqSpeedupStorage::page_type(); + bq.start_training(&meta_page); + + let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); + let mut state = StorageBuildState::BqSpeedup(&mut bq, &mut bs); + unsafe { pg_sys::IndexBuildHeapScan( heap_relation.as_ptr(), index_relation.as_ptr(), index_info, Some(build_callback_bq_train), - &mut bq, + &mut state, ); } bq.finish_training(&mut write_stats); - let page_type = BqSpeedupStorage::page_type(); - let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); let mut state = StorageBuildState::BqSpeedup(&mut bq, &mut bs); unsafe { @@ -325,10 +328,17 @@ unsafe extern "C" fn build_callback_bq_train( _tuple_is_alive: bool, state: *mut std::os::raw::c_void, ) { - let vec = PgVector::from_pg_parts(values, isnull, 0); - if let Some(vec) = vec { - let bq = (state as *mut BqSpeedupStorage).as_mut().unwrap(); - bq.add_sample(vec.to_slice()); + let state = (state as *mut StorageBuildState).as_mut().unwrap(); + match state { + StorageBuildState::BqSpeedup(bq, state) => { + let vec = PgVector::from_pg_parts(values, isnull, 0, &state.meta_page, true, false); + if let Some(vec) = vec { + bq.add_sample(vec.to_index_slice()); + } + } + StorageBuildState::Plain(_, _) => { + panic!("Should not be training with plain storage"); + } } } @@ -342,16 +352,19 @@ unsafe extern "C" fn build_callback( state: *mut std::os::raw::c_void, ) { let index_relation = unsafe { PgRelation::from_pg(index) }; - let vec = PgVector::from_pg_parts(values, isnull, 0); - if let Some(vec) = vec { - let state = (state as *mut StorageBuildState).as_mut().unwrap(); - let heap_pointer = ItemPointer::with_item_pointer_data(*ctid); - - match state { - StorageBuildState::BqSpeedup(bq, state) => { + let state = (state as *mut StorageBuildState).as_mut().unwrap(); + match state { + StorageBuildState::BqSpeedup(bq, state) => { + let vec = PgVector::from_pg_parts(values, isnull, 0, &state.meta_page, true, false); + if let Some(vec) = vec { + let heap_pointer = ItemPointer::with_item_pointer_data(*ctid); build_callback_memory_wrapper(index_relation, heap_pointer, vec, state, *bq); } - StorageBuildState::Plain(plain, state) => { + } + StorageBuildState::Plain(plain, state) => { + let vec = PgVector::from_pg_parts(values, isnull, 0, &state.meta_page, true, false); + if let Some(vec) = vec { + let heap_pointer = ItemPointer::with_item_pointer_data(*ctid); build_callback_memory_wrapper(index_relation, heap_pointer, vec, state, *plain); } } @@ -399,7 +412,7 @@ fn build_callback_internal( } let index_pointer = storage.create_node( - vector.to_slice(), + vector.to_index_slice(), heap_pointer, &state.meta_page, &mut state.tape, diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index bfeda410..f02dac7d 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -10,7 +10,7 @@ use crate::util::*; use super::bq::BqNode; use super::distance; -use super::options::NUM_NEIGHBORS_DEFAULT_SENTINEL; +use super::options::{NUM_DIMENSIONS_DEFAULT_SENTINEL, NUM_NEIGHBORS_DEFAULT_SENTINEL}; use super::stats::StatsNodeModify; use super::storage::StorageType; @@ -65,6 +65,7 @@ impl MetaPageV1 { extension_version_when_built: "0.0.2".to_string(), distance_type: DistanceType::L2 as u16, num_dimensions: self.num_dimensions, + num_dimensions_to_index: self.num_dimensions, num_neighbors: self.num_neighbors, storage_type: StorageType::Plain as u8, search_list_size: self.search_list_size, @@ -114,8 +115,10 @@ pub struct MetaPage { extension_version_when_built: String, /// The value of the DistanceType enum distance_type: u16, - /// number of dimensions in the vector + /// number of total_dimensions in the vector num_dimensions: u32, + //number of dimensions in the vectors stored in the index + num_dimensions_to_index: u32, /// the value of the TSVStorageLayout enum storage_type: u8, /// max number of outgoing edges a node in the graph can have (R in the papers) @@ -133,6 +136,10 @@ impl MetaPage { self.num_dimensions } + pub fn get_num_dimensions_to_index(&self) -> u32 { + self.num_dimensions_to_index + } + /// Maximum number of neigbors per node. Given that we pre-allocate /// these many slots for each node, this cannot change after the graph is built. pub fn get_num_neighbors(&self) -> u32 { @@ -202,12 +209,19 @@ impl MetaPage { ) -> MetaPage { let version = Version::parse(env!("CARGO_PKG_VERSION")).unwrap(); + let num_dimensions_to_index = if (*opt).num_dimensions == NUM_DIMENSIONS_DEFAULT_SENTINEL { + num_dimensions + } else { + (*opt).num_dimensions + }; + let meta = MetaPage { magic_number: TSV_MAGIC_NUMBER, version: TSV_VERSION, extension_version_when_built: version.to_string(), distance_type: DistanceType::Cosine as u16, num_dimensions, + num_dimensions_to_index, storage_type: (*opt).get_storage_type() as u8, num_neighbors: Self::calculate_num_neighbors(num_dimensions, &opt), search_list_size: (*opt).search_list_size, diff --git a/timescale_vector/src/access_method/options.rs b/timescale_vector/src/access_method/options.rs index 3b637c90..3fd1f3a4 100644 --- a/timescale_vector/src/access_method/options.rs +++ b/timescale_vector/src/access_method/options.rs @@ -15,10 +15,12 @@ pub struct TSVIndexOptions { pub storage_layout_offset: i32, num_neighbors: i32, pub search_list_size: u32, + pub num_dimensions: u32, pub max_alpha: f64, } pub const NUM_NEIGHBORS_DEFAULT_SENTINEL: i32 = -1; +pub const NUM_DIMENSIONS_DEFAULT_SENTINEL: u32 = 0; const DEFAULT_MAX_ALPHA: f64 = 1.2; impl TSVIndexOptions { @@ -34,6 +36,7 @@ impl TSVIndexOptions { ops.num_neighbors = NUM_NEIGHBORS_DEFAULT_SENTINEL; ops.search_list_size = 100; ops.max_alpha = DEFAULT_MAX_ALPHA; + ops.num_dimensions = NUM_DIMENSIONS_DEFAULT_SENTINEL; unsafe { set_varsize( ops.as_ptr().cast(), @@ -81,7 +84,7 @@ impl TSVIndexOptions { } } -const NUM_REL_OPTS: usize = 4; +const NUM_REL_OPTS: usize = 5; static mut RELOPT_KIND_TSV: pg_sys::relopt_kind = 0; // amoptions is a function that gets a datum of text[] data from pg_class.reloptions (which contains text in the format "key=value") and returns a bytea for the struct for the parsed options. @@ -116,6 +119,11 @@ pub unsafe extern "C" fn amoptions( opttype: pg_sys::relopt_type_RELOPT_TYPE_INT, offset: offset_of!(TSVIndexOptions, search_list_size) as i32, }, + pg_sys::relopt_parse_elt { + optname: "num_dimensions".as_pg_cstr(), + opttype: pg_sys::relopt_type_RELOPT_TYPE_INT, + offset: offset_of!(TSVIndexOptions, num_dimensions) as i32, + }, pg_sys::relopt_parse_elt { optname: "max_alpha".as_pg_cstr(), opttype: pg_sys::relopt_type_RELOPT_TYPE_REAL, @@ -207,13 +215,26 @@ pub unsafe fn init() { 5.0, pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, ); + + pg_sys::add_int_reloption( + RELOPT_KIND_TSV, + "num_dimensions".as_pg_cstr(), + "The number of dimensions to index (0 to index all dimensions)".as_pg_cstr(), + 0, + 0, + 5000, + pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, + ); } #[cfg(any(test, feature = "pg_test"))] #[pgrx::pg_schema] mod tests { use crate::access_method::{ - options::{TSVIndexOptions, DEFAULT_MAX_ALPHA, NUM_NEIGHBORS_DEFAULT_SENTINEL}, + options::{ + TSVIndexOptions, DEFAULT_MAX_ALPHA, NUM_DIMENSIONS_DEFAULT_SENTINEL, + NUM_NEIGHBORS_DEFAULT_SENTINEL, + }, storage::StorageType, }; use pgrx::*; @@ -233,6 +254,7 @@ mod tests { let indexrel = PgRelation::from_pg(pg_sys::RelationIdGetRelation(index_oid)); let options = TSVIndexOptions::from_relation(&indexrel); assert_eq!(options.num_neighbors, 30); + assert_eq!(options.num_dimensions, NUM_DIMENSIONS_DEFAULT_SENTINEL); Ok(()) } @@ -252,6 +274,7 @@ mod tests { assert_eq!(options.get_num_neighbors(), NUM_NEIGHBORS_DEFAULT_SENTINEL); assert_eq!(options.search_list_size, 100); assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); + assert_eq!(options.num_dimensions, NUM_DIMENSIONS_DEFAULT_SENTINEL); assert_eq!(options.get_storage_type(), StorageType::BqSpeedup); Ok(()) } @@ -273,6 +296,7 @@ mod tests { assert_eq!(options.get_num_neighbors(), NUM_NEIGHBORS_DEFAULT_SENTINEL); assert_eq!(options.search_list_size, 100); assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); + assert_eq!(options.num_dimensions, NUM_DIMENSIONS_DEFAULT_SENTINEL); assert_eq!(options.get_storage_type(), StorageType::BqSpeedup); Ok(()) } @@ -297,4 +321,26 @@ mod tests { assert_eq!(options.get_storage_type(), StorageType::Plain); Ok(()) } + + #[pg_test] + unsafe fn test_index_options_custom() -> spi::Result<()> { + Spi::run(&format!( + "CREATE TABLE test(encoding vector(3)); + CREATE INDEX idxtest + ON test + USING tsv(encoding) + WITH (storage_layout = plain, num_neighbors=40, search_list_size=18, num_dimensions=20, max_alpha=1.4);", + ))?; + + let index_oid = + Spi::get_one::("SELECT 'idxtest'::regclass::oid")?.expect("oid was null"); + let indexrel = PgRelation::from_pg(pg_sys::RelationIdGetRelation(index_oid)); + let options = TSVIndexOptions::from_relation(&indexrel); + assert_eq!(options.get_num_neighbors(), 40); + assert_eq!(options.search_list_size, 18); + assert_eq!(options.max_alpha, 1.4); + assert_eq!(options.get_storage_type(), StorageType::Plain); + assert_eq!(options.num_dimensions, 20); + Ok(()) + } } diff --git a/timescale_vector/src/access_method/pg_vector.rs b/timescale_vector/src/access_method/pg_vector.rs index 2ab4fe8d..66abba0f 100644 --- a/timescale_vector/src/access_method/pg_vector.rs +++ b/timescale_vector/src/access_method/pg_vector.rs @@ -1,6 +1,6 @@ use pgrx::*; -use super::distance::preprocess_cosine; +use super::{distance::preprocess_cosine, meta_page}; //Ported from pg_vector code #[repr(C)] @@ -21,15 +21,26 @@ impl PgVectorInternal { } pub struct PgVector { - inner: *mut PgVectorInternal, - need_pfree: bool, + index_distance: Option<*mut PgVectorInternal>, + index_distance_needs_pfree: bool, + full_distance: Option<*mut PgVectorInternal>, + full_distance_needs_pfree: bool, } impl Drop for PgVector { fn drop(&mut self) { - if self.need_pfree { + if self.index_distance_needs_pfree { unsafe { - pg_sys::pfree(self.inner.cast()); + if self.index_distance.is_some() { + pg_sys::pfree(self.index_distance.unwrap().cast()); + } + } + } + if self.full_distance_needs_pfree { + unsafe { + if self.full_distance.is_some() { + pg_sys::pfree(self.full_distance.unwrap().cast()); + } } } } @@ -40,17 +51,29 @@ impl PgVector { datum_parts: *mut pg_sys::Datum, isnull_parts: *mut bool, index: usize, + meta_page: &meta_page::MetaPage, + index_distance: bool, + full_distance: bool, ) -> Option { let isnulls = std::slice::from_raw_parts(isnull_parts, index + 1); if isnulls[index] { return None; } let datums = std::slice::from_raw_parts(datum_parts, index + 1); - Some(Self::from_datum(datums[index])) + Some(Self::from_datum( + datums[index], + meta_page, + index_distance, + full_distance, + )) } - pub unsafe fn from_datum(datum: pg_sys::Datum) -> PgVector { - //FIXME: we are using a copy here to avoid lifetime issues and because in some cases we have to + unsafe fn create_inner( + datum: pg_sys::Datum, + meta_page: &meta_page::MetaPage, + is_index_distance: bool, + ) -> *mut PgVectorInternal { + //TODO: we are using a copy here to avoid lifetime issues and because in some cases we have to //modify the datum in preprocess_cosine. We should find a way to avoid the copy if the vector is //normalized and preprocess_cosine is a noop; let detoasted = pg_sys::pg_detoast_datum_copy(datum.cast_mut_ptr()); @@ -58,19 +81,67 @@ impl PgVector { detoasted.cast::(), datum.cast_mut_ptr::(), ); + + /* if is_copy every changes, need to change needs_pfree */ + assert!(is_copy, "Datum should be a copy"); let casted = detoasted.cast::(); + if is_index_distance + && meta_page.get_num_dimensions() != meta_page.get_num_dimensions_to_index() + { + assert!((*casted).dim > meta_page.get_num_dimensions_to_index() as _); + (*casted).dim = meta_page.get_num_dimensions_to_index() as _; + } + let dim = (*casted).dim; let raw_slice = unsafe { (*casted).x.as_mut_slice(dim as _) }; + preprocess_cosine(raw_slice); + casted + } + + pub unsafe fn from_datum( + datum: pg_sys::Datum, + meta_page: &meta_page::MetaPage, + index_distance: bool, + full_distance: bool, + ) -> PgVector { + if meta_page.get_num_dimensions() == meta_page.get_num_dimensions_to_index() { + /* optimization if the num dimensions are the same */ + let inner = Self::create_inner(datum, meta_page, true); + return PgVector { + index_distance: Some(inner), + index_distance_needs_pfree: true, + full_distance: Some(inner), + full_distance_needs_pfree: false, + }; + } + + let idx = if index_distance { + Some(Self::create_inner(datum, meta_page, true)) + } else { + None + }; + + let full = if full_distance { + Some(Self::create_inner(datum, meta_page, false)) + } else { + None + }; PgVector { - inner: casted, - need_pfree: is_copy, + index_distance: idx, + index_distance_needs_pfree: true, + full_distance: full, + full_distance_needs_pfree: true, } } - pub fn to_slice(&self) -> &[f32] { - unsafe { (*self.inner).to_slice() } + pub fn to_index_slice(&self) -> &[f32] { + unsafe { (*self.index_distance.unwrap()).to_slice() } + } + + pub fn to_full_slice(&self) -> &[f32] { + unsafe { (*self.full_distance.unwrap()).to_slice() } } } diff --git a/timescale_vector/src/access_method/plain_storage.rs b/timescale_vector/src/access_method/plain_storage.rs index 5f052283..70eef9a1 100644 --- a/timescale_vector/src/access_method/plain_storage.rs +++ b/timescale_vector/src/access_method/plain_storage.rs @@ -8,47 +8,61 @@ use super::{ StatsNodeRead, StatsNodeWrite, WriteStats, }, storage::{ArchivedData, NodeDistanceMeasure, Storage}, + storage_common::get_attribute_number_from_index, }; use pgrx::PgRelation; -use crate::util::{page::PageType, tape::Tape, HeapPointer, IndexPointer, ItemPointer}; +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: fn(&[f32], &[f32]) -> f32, + heap_rel: &'a PgRelation, + heap_attr: pgrx::pg_sys::AttrNumber, } impl<'a> PlainStorage<'a> { pub fn new_for_build( index: &'a PgRelation, + heap_rel: &'a PgRelation, distance_fn: fn(&[f32], &[f32]) -> f32, ) -> PlainStorage<'a> { Self { index: index, distance_fn: distance_fn, + heap_rel: heap_rel, + heap_attr: get_attribute_number_from_index(index), } } pub fn load_for_insert( index_relation: &'a PgRelation, + heap_rel: &'a PgRelation, distance_fn: fn(&[f32], &[f32]) -> f32, ) -> PlainStorage<'a> { Self { index: index_relation, distance_fn: distance_fn, + heap_rel: heap_rel, + heap_attr: get_attribute_number_from_index(&index_relation), } } pub fn load_for_search( index_relation: &'a PgRelation, + heap_rel: &'a PgRelation, distance_fn: fn(&[f32], &[f32]) -> f32, ) -> PlainStorage<'a> { Self { index: index_relation, distance_fn: distance_fn, + heap_rel: heap_rel, + heap_attr: get_attribute_number_from_index(&index_relation), } } } @@ -199,12 +213,23 @@ impl<'a> Storage for PlainStorage<'a> { } fn get_full_distance_for_resort( &self, - _qdm: &Self::QueryDistanceMeasure, + qdm: &Self::QueryDistanceMeasure, _index_pointer: IndexPointer, - _heap_pointer: HeapPointer, - _stats: &mut S, + heap_pointer: HeapPointer, + meta_page: &MetaPage, + stats: &mut S, ) -> f32 { - pgrx::error!("Plain node should never be resorted"); + /* Plain storage only needs to resort when the index is using less dimensions than the underlying data. */ + assert!(meta_page.get_num_dimensions() > meta_page.get_num_dimensions_to_index()); + + let slot = unsafe { TableSlot::new(self.heap_rel, heap_pointer, stats) }; + match qdm { + PlainDistanceMeasure::Full(query) => { + let datum = unsafe { slot.get_attribute(self.heap_attr).unwrap() }; + let vec = unsafe { PgVector::from_datum(datum, meta_page, false, true) }; + self.get_distance_function()(vec.to_full_slice(), query.to_full_slice()) + } + } } fn get_neighbors_with_distances_from_disk( &self, @@ -240,7 +265,7 @@ impl<'a> Storage for PlainStorage<'a> { let distance = match lsr.sdm.as_ref().unwrap() { PlainDistanceMeasure::Full(query) => PlainDistanceMeasure::calculate_distance( self.distance_fn, - query.to_slice(), + query.to_index_slice(), node.vector.as_slice(), &mut lsr.stats, ), @@ -275,7 +300,7 @@ impl<'a> Storage for PlainStorage<'a> { let distance = match lsr.sdm.as_ref().unwrap() { PlainDistanceMeasure::Full(query) => PlainDistanceMeasure::calculate_distance( self.distance_fn, - query.to_slice(), + query.to_index_slice(), node_neighbor.vector.as_slice(), &mut lsr.stats, ), @@ -323,7 +348,7 @@ mod tests { use pgrx::*; #[pg_test] - unsafe fn test_plain_storage_index_creation() -> spi::Result<()> { + unsafe fn test_plain_storage_index_creation_many_neighbors() -> spi::Result<()> { crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( "num_neighbors=38, storage_layout = plain", )?; @@ -366,4 +391,12 @@ mod tests { "num_neighbors=38, storage_layout = plain", ) } + + #[pg_test] + unsafe fn test_plain_storage_num_dimensions() -> spi::Result<()> { + crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( + "num_neighbors=38, storage_layout = plain, num_dimensions=768", + )?; + Ok(()) + } } diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 1e1baee1..8e6d6aef 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -32,13 +32,15 @@ enum StorageState { struct TSVScanState { storage: *mut StorageState, distance_fn: Option f32>, + meta_page: MetaPage, } impl TSVScanState { - fn new() -> Self { + fn new(meta_page: MetaPage) -> Self { Self { storage: std::ptr::null_mut(), distance_fn: None, + meta_page: meta_page, } } @@ -56,7 +58,8 @@ impl TSVScanState { let store_type = match storage { StorageType::Plain => { let stats = QuantizerStats::new(); - let bq = PlainStorage::load_for_search(index, meta_page.get_distance_function()); + let bq = + PlainStorage::load_for_search(index, heap, meta_page.get_distance_function()); let it = TSVResponseIterator::new(&bq, index, query, search_list_size, meta_page, stats); StorageState::Plain(it) @@ -125,6 +128,7 @@ impl TSVResponseIterator { index: &PgRelation, query: PgVector, search_list_size: usize, + //FIXME? _meta_page: MetaPage, quantizer_stats: QuantizerStats, ) -> Self { @@ -202,6 +206,7 @@ impl TSVResponseIterator { self.lsr.sdm.as_ref().unwrap(), index_pointer, heap_pointer, + &self.meta_page, &mut self.lsr.stats, ); @@ -242,8 +247,10 @@ pub extern "C" fn ambeginscan( norderbys, )) }; + let indexrel = unsafe { PgRelation::from_pg(index_relation) }; + let meta_page = MetaPage::fetch(&indexrel); - let state: TSVScanState = TSVScanState::new(); + let state: TSVScanState = TSVScanState::new(meta_page); scandesc.opaque = PgMemoryContexts::CurrentMemoryContext.leak_and_drop_on_delete(state) as void_mut_ptr; @@ -267,8 +274,6 @@ pub extern "C" fn amrescan( let mut scan: PgBox = unsafe { PgBox::from_pg(scan) }; let indexrel = unsafe { PgRelation::from_pg(scan.indexRelation) }; let heaprel = unsafe { PgRelation::from_pg(scan.heapRelation) }; - let meta_page = MetaPage::fetch(&indexrel); - let _storage = meta_page.get_storage_type(); if nkeys > 0 { scan.xs_recheck = true; @@ -277,13 +282,19 @@ pub extern "C" fn amrescan( let orderby_keys = unsafe { std::slice::from_raw_parts(orderbys as *const pg_sys::ScanKeyData, norderbys as _) }; - let query = unsafe { PgVector::from_datum(orderby_keys[0].sk_argument) }; - //TODO need to set search_list_size correctly - //TODO right now doesn't handle more than LIMIT 100; let search_list_size = super::guc::TSV_QUERY_SEARCH_LIST_SIZE.get() as usize; let state = unsafe { (scan.opaque as *mut TSVScanState).as_mut() }.expect("no scandesc state"); + + let query = unsafe { + PgVector::from_datum( + orderby_keys[0].sk_argument, + &state.meta_page, + true, /* needed for search */ + true, /* needed for resort */ + ) + }; state.initialize(&indexrel, &heaprel, query, search_list_size); } @@ -312,8 +323,16 @@ pub extern "C" fn amgettuple( get_tuple(next, scan) } StorageState::Plain(iter) => { - let storage = PlainStorage::load_for_search(&indexrel, state.distance_fn.unwrap()); - let next = iter.next(&indexrel, &storage); + let storage = + PlainStorage::load_for_search(&indexrel, &heaprel, state.distance_fn.unwrap()); + let next = if state.meta_page.get_num_dimensions() + == state.meta_page.get_num_dimensions_to_index() + { + /* no need to resort */ + iter.next(&indexrel, &storage) + } else { + iter.next_with_resort(&indexrel, &storage) + }; get_tuple(next, scan) } } diff --git a/timescale_vector/src/access_method/storage.rs b/timescale_vector/src/access_method/storage.rs index e8dd23f9..bedad284 100644 --- a/timescale_vector/src/access_method/storage.rs +++ b/timescale_vector/src/access_method/storage.rs @@ -77,6 +77,7 @@ pub trait Storage { query: &Self::QueryDistanceMeasure, index_pointer: IndexPointer, heap_pointer: HeapPointer, + meta_page: &MetaPage, stats: &mut S, ) -> f32; diff --git a/timescale_vector/src/util/table_slot.rs b/timescale_vector/src/util/table_slot.rs index dc27d8bb..56774e5a 100644 --- a/timescale_vector/src/util/table_slot.rs +++ b/timescale_vector/src/util/table_slot.rs @@ -1,21 +1,18 @@ use pgrx::pg_sys::{Datum, TupleTableSlot}; use pgrx::{pg_sys, PgBox, PgRelation}; -use crate::access_method::pg_vector::PgVector; use crate::access_method::stats::StatsHeapNodeRead; use crate::util::ports::slot_getattr; use crate::util::HeapPointer; pub struct TableSlot { slot: PgBox, - attribute_number: pg_sys::AttrNumber, } impl TableSlot { pub unsafe fn new( heap_rel: &PgRelation, heap_pointer: HeapPointer, - attribute_number: pg_sys::AttrNumber, stats: &mut S, ) -> Self { let slot = PgBox::from_pg(pg_sys::table_slot_create( @@ -37,20 +34,12 @@ impl TableSlot { ); stats.record_heap_read(); - Self { - slot, - attribute_number, - } + Self { slot } } - unsafe fn get_attribute(&self, attribute_number: pg_sys::AttrNumber) -> Option { + pub unsafe fn get_attribute(&self, attribute_number: pg_sys::AttrNumber) -> Option { slot_getattr(&self.slot, attribute_number) } - - pub unsafe fn get_pg_vector(&self) -> PgVector { - let vector = PgVector::from_datum(self.get_attribute(self.attribute_number).unwrap()); - vector - } } impl Drop for TableSlot { From 3ff4dccf3d688aaf6ca352a5c0b70b74e2b66c58 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 9 Apr 2024 14:47:05 -0400 Subject: [PATCH 28/44] Add update test --- timescale_vector/src/access_method/bq.rs | 9 ++ timescale_vector/src/access_method/build.rs | 101 ++++++++++++++++++ .../src/access_method/plain_storage.rs | 9 ++ 3 files changed, 119 insertions(+) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index 3f40dfaf..62bc2853 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -875,4 +875,13 @@ mod tests { )?; Ok(()) } + + #[pg_test] + unsafe fn test_bq_storage_index_updates() -> spi::Result<()> { + crate::access_method::build::tests::test_index_updates( + "storage_layout = io_optimized, num_neighbors=10", + 300, + )?; + Ok(()) + } } diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index b8152fe0..d4cc88d3 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -703,4 +703,105 @@ pub mod tests { Ok(()) } + + #[cfg(any(test, feature = "pg_test"))] + pub unsafe fn test_index_updates(index_options: &str, expected_cnt: i64) -> spi::Result<()> { + Spi::run(&format!( + "CREATE TABLE test_data ( + id int, + embedding vector (1536) + ); + + select setseed(0.5); + -- generate 300 vectors + INSERT INTO test_data (id, embedding) + SELECT + * + FROM ( + SELECT + i % {expected_cnt}, + ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding + FROM + generate_series(1, 1536 * {expected_cnt}) i + GROUP BY + i % {expected_cnt}) g; + + CREATE INDEX idx_tsv_bq ON test_data USING tsv (embedding) WITH ({index_options}); + + + SET enable_seqscan = 0; + -- perform index scans on the vectors + SELECT + * + FROM + test_data + ORDER BY + embedding <=> ( + SELECT + ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding + FROM generate_series(1, 1536));"))?; + + let test_vec: Option> = Spi::get_one(&format!( + "SELECT('{{' || array_to_string(array_agg(1.0), ',', '0') || '}}')::real[] AS embedding + FROM generate_series(1, 1536)" + ))?; + + let cnt: Option = Spi::get_one_with_args( + &format!( + " + SET enable_seqscan = 0; + SET enable_indexscan = 1; + SET tsv.query_search_list_size = 2; + WITH cte as (select * from test_data order by embedding <=> $1::vector) SELECT count(*) from cte; + ", + ), + vec![( + pgrx::PgOid::Custom(pgrx::pg_sys::FLOAT4ARRAYOID), + test_vec.clone().into_datum(), + )], + )?; + + assert!(cnt.unwrap() == expected_cnt, "initial count"); + + Spi::run(&format!( + " + + --CREATE INDEX idx_id ON test_data(id); + + WITH CTE as ( + SELECT + i % {expected_cnt} as id, + ('[' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding + FROM + generate_series(1, 1536 * {expected_cnt}) i + GROUP BY + i % {expected_cnt} + ) + UPDATE test_data SET embedding = cte.embedding + FROM cte + WHERE test_data.id = cte.id; + + --DROP INDEX idx_id; + ", + ))?; + + let cnt: Option = Spi::get_one_with_args( + &format!( + " + SET enable_seqscan = 0; + SET enable_indexscan = 1; + SET tsv.query_search_list_size = 2; + WITH cte as (select * from test_data order by embedding <=> $1::vector) SELECT count(*) from cte; + ", + ), + vec![( + pgrx::PgOid::Custom(pgrx::pg_sys::FLOAT4ARRAYOID), + test_vec.clone().into_datum(), + )], + )?; + + assert!(cnt.unwrap() == expected_cnt, "after update count"); + + Ok(()) + } } diff --git a/timescale_vector/src/access_method/plain_storage.rs b/timescale_vector/src/access_method/plain_storage.rs index 70eef9a1..1c4c8f73 100644 --- a/timescale_vector/src/access_method/plain_storage.rs +++ b/timescale_vector/src/access_method/plain_storage.rs @@ -399,4 +399,13 @@ mod tests { )?; Ok(()) } + + #[pg_test] + unsafe fn test_plain_storage_index_updates() -> spi::Result<()> { + crate::access_method::build::tests::test_index_updates( + "storage_layout = plain, num_neighbors=30", + 50, + )?; + Ok(()) + } } From 24fecf1de85270c153ffa996d088b837b9461eb0 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Mon, 22 Apr 2024 11:04:10 -0400 Subject: [PATCH 29/44] Bug fix: fix locking of buffer in the resort case --- timescale_vector/src/access_method/scan.rs | 48 +++++++++++----------- 1 file changed, 25 insertions(+), 23 deletions(-) diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 8e6d6aef..7a714e23 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -33,6 +33,7 @@ struct TSVScanState { storage: *mut StorageState, distance_fn: Option f32>, meta_page: MetaPage, + last_buffer: Option, } impl TSVScanState { @@ -41,6 +42,7 @@ impl TSVScanState { storage: std::ptr::null_mut(), distance_fn: None, meta_page: meta_page, + last_buffer: None, } } @@ -115,8 +117,6 @@ impl Ord for ResortData { struct TSVResponseIterator { lsr: ListSearchResult, search_list_size: usize, - current: usize, - last_buffer: Option, meta_page: MetaPage, quantizer_stats: QuantizerStats, resort_buffer: BinaryHeap, @@ -141,8 +141,6 @@ impl TSVResponseIterator { Self { search_list_size, lsr, - current: 0, - last_buffer: None, meta_page, quantizer_stats, resort_buffer: BinaryHeap::with_capacity(resort_size), @@ -153,7 +151,6 @@ impl TSVResponseIterator { impl TSVResponseIterator { fn next>( &mut self, - index: &PgRelation, storage: &S, ) -> Option<(HeapPointer, IndexPointer)> { let graph = Graph::new(GraphNeighborStore::Disk, &mut self.meta_page); @@ -166,16 +163,6 @@ impl TSVResponseIterator { match item { Some((heap_pointer, index_pointer)) => { - /* - * An index scan must maintain a pin on the index page holding the - * item last returned by amgettuple - * - * https://www.postgresql.org/docs/current/index-locking.html - */ - self.last_buffer = - Some(PinnedBufferShare::read(index, index_pointer.block_number)); - - self.current = self.current + 1; if heap_pointer.offset == InvalidOffsetNumber { /* deleted tuple */ continue; @@ -183,7 +170,6 @@ impl TSVResponseIterator { return Some((heap_pointer, index_pointer)); } None => { - self.last_buffer = None; return None; } } @@ -196,11 +182,11 @@ impl TSVResponseIterator { storage: &S, ) -> Option<(HeapPointer, IndexPointer)> { if self.resort_buffer.capacity() == 0 { - return self.next(index, storage); + return self.next(storage); } while self.resort_buffer.len() < self.resort_buffer.capacity() { - match self.next(index, storage) { + match self.next(storage) { Some((heap_pointer, index_pointer)) => { let distance = storage.get_full_distance_for_resort( self.lsr.sdm.as_ref().unwrap(), @@ -320,7 +306,7 @@ pub extern "C" fn amgettuple( state.distance_fn.unwrap(), ); let next = iter.next_with_resort(&indexrel, &bq); - get_tuple(next, scan) + get_tuple(state, next, scan) } StorageState::Plain(iter) => { let storage = @@ -329,27 +315,43 @@ pub extern "C" fn amgettuple( == state.meta_page.get_num_dimensions_to_index() { /* no need to resort */ - iter.next(&indexrel, &storage) + iter.next(&storage) } else { iter.next_with_resort(&indexrel, &storage) }; - get_tuple(next, scan) + get_tuple(state, next, scan) } } } fn get_tuple( + state: &mut TSVScanState, next: Option<(HeapPointer, IndexPointer)>, mut scan: PgBox, ) -> bool { scan.xs_recheckorderby = false; match next { - Some((heap_pointer, _)) => { + Some((heap_pointer, index_pointer)) => { let tid_to_set = &mut scan.xs_heaptid; heap_pointer.to_item_pointer_data(tid_to_set); + + /* + * An index scan must maintain a pin on the index page holding the + * item last returned by amgettuple + * + * https://www.postgresql.org/docs/current/index-locking.html + */ + let indexrel = unsafe { PgRelation::from_pg(scan.indexRelation) }; + state.last_buffer = Some(PinnedBufferShare::read( + &indexrel, + index_pointer.block_number, + )); true } - None => false, + None => { + state.last_buffer = None; + false + } } } From 3ae4b36752d230364af63dd8f01343c6e3980e83 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Fri, 19 Apr 2024 12:32:32 -0400 Subject: [PATCH 30/44] Optimize prune by using the cache in BqNodeDistanceMeasure --- timescale_vector/src/access_method/bq.rs | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index 62bc2853..6a4d067a 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -214,7 +214,7 @@ impl BqSearchDistanceMeasure { } pub struct BqNodeDistanceMeasure<'a> { - readable_node: ReadableBqNode<'a>, + vec: Vec, storage: &'a BqSpeedupStorage<'a>, } @@ -224,9 +224,9 @@ impl<'a> BqNodeDistanceMeasure<'a> { index_pointer: IndexPointer, stats: &mut T, ) -> Self { - let rn = unsafe { BqNode::read(storage.index, index_pointer, stats) }; + let cache = &mut storage.qv_cache.borrow_mut(); Self { - readable_node: rn, + vec: cache.get(index_pointer, storage, stats).to_vec(), storage: storage, } } @@ -238,16 +238,9 @@ impl<'a> NodeDistanceMeasure for BqNodeDistanceMeasure<'a> { index_pointer: IndexPointer, stats: &mut T, ) -> f32 { - //OPT: should I get and memoize the vector from self.readable_node in with_index_pointer above? - let rn1 = BqNode::read(self.storage.index, index_pointer, stats); - let rn2 = &self.readable_node; - let node1 = rn1.get_archived_node(); - let node2 = rn2.get_archived_node(); - assert!(node1.bq_vector.len() > 0); - assert!(node1.bq_vector.len() == node2.bq_vector.len()); - let vec1 = node1.bq_vector.as_slice(); - let vec2 = node2.bq_vector.as_slice(); - distance_xor_optimized(vec1, vec2) as f32 + let cache = &mut self.storage.qv_cache.borrow_mut(); + let vec1 = cache.get(index_pointer, self.storage, stats); + distance_xor_optimized(vec1, self.vec.as_slice()) as f32 } } From 926f1d50ff4a995b8e3bb56b70d33c478fcb8bfc Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Fri, 19 Apr 2024 13:32:24 -0400 Subject: [PATCH 31/44] Optimize prune by getting rid of another read --- timescale_vector/src/access_method/bq.rs | 69 ++++++++++++++---------- 1 file changed, 41 insertions(+), 28 deletions(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index 6a4d067a..32f830f5 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -386,26 +386,20 @@ impl<'a> BqSpeedupStorage<'a> { lsn_index_pointer: IndexPointer, gns: &GraphNeighborStore, ) { - //Opt shouldn't need to read the node in the builder graph case. - let rn_visiting = unsafe { BqNode::read(self.index, lsn_index_pointer, &mut lsr.stats) }; - let node_visiting = rn_visiting.get_archived_node(); - - let neighbors = match gns { - GraphNeighborStore::Disk => node_visiting.get_index_pointer_to_neighbors(), - GraphNeighborStore::Builder(b) => b.get_neighbors(lsn_index_pointer), - }; - - for (i, &neighbor_index_pointer) in neighbors.iter().enumerate() { - if !lsr.prepare_insert(neighbor_index_pointer) { - continue; - } + match gns { + GraphNeighborStore::Disk => { + let rn_visiting = + unsafe { BqNode::read(self.index, lsn_index_pointer, &mut lsr.stats) }; + let node_visiting = rn_visiting.get_archived_node(); + let neighbors = node_visiting.get_index_pointer_to_neighbors(); + + for (i, &neighbor_index_pointer) in neighbors.iter().enumerate() { + if !lsr.prepare_insert(neighbor_index_pointer) { + continue; + } - let distance = match lsr.sdm.as_ref().unwrap() { - BqSearchDistanceMeasure::Bq(table, _) => { - /* Note: there is no additional node reads here. We get all of our info from node_visiting - * This is what gives us a speedup in BQ Speedup */ - match gns { - GraphNeighborStore::Disk => { + let distance = match lsr.sdm.as_ref().unwrap() { + BqSearchDistanceMeasure::Bq(table, _) => { let bq_vector = node_visiting.neighbor_vectors[i].as_slice(); BqSearchDistanceMeasure::calculate_bq_distance( table, @@ -413,7 +407,24 @@ impl<'a> BqSpeedupStorage<'a> { &mut lsr.stats, ) } - GraphNeighborStore::Builder(_) => { + }; + let lsn = ListSearchNeighbor::new( + neighbor_index_pointer, + distance, + PhantomData::, + ); + + lsr.insert_neighbor(lsn); + } + } + GraphNeighborStore::Builder(b) => { + let neighbors = b.get_neighbors(lsn_index_pointer); + for &neighbor_index_pointer in neighbors.iter() { + if !lsr.prepare_insert(neighbor_index_pointer) { + continue; + } + let distance = match lsr.sdm.as_ref().unwrap() { + BqSearchDistanceMeasure::Bq(table, _) => { let mut cache = self.qv_cache.borrow_mut(); let bq_vector = cache.get(neighbor_index_pointer, self, &mut lsr.stats); let dist = BqSearchDistanceMeasure::calculate_bq_distance( @@ -423,15 +434,17 @@ impl<'a> BqSpeedupStorage<'a> { ); dist } - } - //let bq_vector = node_visiting.neighbor_vectors[i].as_slice(); - //BqSearchDistanceMeasure::calculate_bq_distance(table, bq_vector, &mut lsr.stats) - } - }; - let lsn = - ListSearchNeighbor::new(neighbor_index_pointer, distance, PhantomData::); + }; - lsr.insert_neighbor(lsn); + let lsn = ListSearchNeighbor::new( + neighbor_index_pointer, + distance, + PhantomData::, + ); + + lsr.insert_neighbor(lsn); + } + } } } From f61d20b6274bbaa05ea470ae8e4e28ce7b8575c9 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Fri, 19 Apr 2024 14:21:52 -0400 Subject: [PATCH 32/44] Optimize write in finalize_node_at_end_of_build to be ordered --- .../src/access_method/graph_neighbor_store.rs | 8 +++++--- timescale_vector/src/util/mod.rs | 14 ++++++++++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/timescale_vector/src/access_method/graph_neighbor_store.rs b/timescale_vector/src/access_method/graph_neighbor_store.rs index e0e4e596..50933a71 100644 --- a/timescale_vector/src/access_method/graph_neighbor_store.rs +++ b/timescale_vector/src/access_method/graph_neighbor_store.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::BTreeMap; use crate::util::{IndexPointer, ItemPointer}; @@ -15,13 +15,15 @@ use super::storage::Storage; /// the neighbors to the right pages. pub struct BuilderNeighborCache { //maps node's pointer to the representation on disk - neighbor_map: HashMap>, + //use a btree to provide ordering on the item pointers in iter(). + //this ensures the write in finalize_node_at_end_of_build() is ordered, not random. + neighbor_map: BTreeMap>, } impl BuilderNeighborCache { pub fn new() -> Self { Self { - neighbor_map: HashMap::with_capacity(200), + neighbor_map: BTreeMap::new(), } } pub fn iter(&self) -> impl Iterator)> { diff --git a/timescale_vector/src/util/mod.rs b/timescale_vector/src/util/mod.rs index 37c580c0..35d98e96 100644 --- a/timescale_vector/src/util/mod.rs +++ b/timescale_vector/src/util/mod.rs @@ -20,6 +20,20 @@ pub struct ItemPointer { pub offset: pgrx::pg_sys::OffsetNumber, } +impl PartialOrd for ItemPointer { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for ItemPointer { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.block_number + .cmp(&other.block_number) + .then_with(|| self.offset.cmp(&other.offset)) + } +} + impl ArchivedItemPointer { pub fn deserialize_item_pointer(&self) -> ItemPointer { self.deserialize(&mut rkyv::Infallible).unwrap() From 89c19e917a121bf0cdffa81c3184a3863be5f21b Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 23 Apr 2024 11:30:45 -0400 Subject: [PATCH 33/44] cleanup resort to use explicit parameter, not capacity --- timescale_vector/src/access_method/scan.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 7a714e23..1d33d39a 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -119,6 +119,7 @@ struct TSVResponseIterator { search_list_size: usize, meta_page: MetaPage, quantizer_stats: QuantizerStats, + resort_size: usize, resort_buffer: BinaryHeap, } @@ -143,6 +144,7 @@ impl TSVResponseIterator { lsr, meta_page, quantizer_stats, + resort_size, resort_buffer: BinaryHeap::with_capacity(resort_size), } } @@ -185,7 +187,7 @@ impl TSVResponseIterator { return self.next(storage); } - while self.resort_buffer.len() < self.resort_buffer.capacity() { + while self.resort_buffer.len() < self.resort_size { match self.next(storage) { Some((heap_pointer, index_pointer)) => { let distance = storage.get_full_distance_for_resort( From ea74738b765d724728b155c66980962b839b9d00 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 23 Apr 2024 11:53:25 -0400 Subject: [PATCH 34/44] Make rescore parameter work off if the std dev of the distances instead of on the number of items --- timescale_vector/src/access_method/scan.rs | 70 +++++++++++++++++++++- 1 file changed, 69 insertions(+), 1 deletion(-) diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 1d33d39a..ca4b569d 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -114,6 +114,54 @@ impl Ord for ResortData { } } +struct StreamingStats { + resort_size: usize, + count: i32, + mean: f32, + m2: f32, + max_distance: f32, +} + +impl StreamingStats { + fn new(resort_size: usize) -> Self { + Self { + resort_size, + count: 0, + mean: 0.0, + m2: 0.0, + max_distance: 0.0, + } + } + + fn update_base_stats(&mut self, distance: f32) { + if distance == 0.0 { + return; + } + self.count += 1; + let delta = distance - self.mean; + self.mean += delta / self.count as f32; + let delta2 = distance - self.mean; + self.m2 += delta * delta2; + } + + fn variance(&self) -> f32 { + if self.count < 2 { + return 0.0; + } + self.m2 / (self.count - 1) as f32 + } + + fn mean(&self) -> f32 { + self.mean + } + + fn update(&mut self, distance: f32, diff: f32) { + //base stats only on first resort_size elements + self.update_base_stats(diff); + self.max_distance = self.max_distance.max(distance); + } +} + struct TSVResponseIterator { lsr: ListSearchResult, search_list_size: usize, @@ -121,6 +169,7 @@ struct TSVResponseIterator { quantizer_stats: QuantizerStats, resort_size: usize, resort_buffer: BinaryHeap, + streaming_stats: StreamingStats, } impl TSVResponseIterator { @@ -146,6 +195,7 @@ impl TSVResponseIterator { quantizer_stats, resort_size, resort_buffer: BinaryHeap::with_capacity(resort_size), + streaming_stats: StreamingStats::new(resort_size), } } } @@ -187,7 +237,11 @@ impl TSVResponseIterator { return self.next(storage); } - while self.resort_buffer.len() < self.resort_size { + while self.resort_buffer.len() < 2 + || self.streaming_stats.count < 2 + || (self.streaming_stats.max_distance - self.resort_buffer.peek().unwrap().distance) + < self.streaming_stats.variance().sqrt() * (self.resort_size as f32 / 100.0) + { match self.next(storage) { Some((heap_pointer, index_pointer)) => { let distance = storage.get_full_distance_for_resort( @@ -198,6 +252,11 @@ impl TSVResponseIterator { &mut self.lsr.stats, ); + if self.resort_buffer.len() > 1 { + self.streaming_stats + .update(distance, distance - self.streaming_stats.max_distance); + } + self.resort_buffer.push(ResortData { heap_pointer, index_pointer, @@ -210,6 +269,15 @@ impl TSVResponseIterator { } } + /*error!( + "Resort buffer size: {}, mean: {}, variance: {}, max_distance: {}: diff: {}", + self.resort_buffer.len(), + self.streaming_stats.mean(), + self.streaming_stats.variance().sqrt(), + self.streaming_stats.max_distance, + self.streaming_stats.max_distance - self.resort_buffer.peek().unwrap().distance + );*/ + match self.resort_buffer.pop() { Some(rd) => Some((rd.heap_pointer, rd.index_pointer)), None => None, From 5cdcd43072e49997f3ab1813fe937950baccdd26 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 23 Apr 2024 21:03:43 -0400 Subject: [PATCH 35/44] Bug fix for num_dimensions --- timescale_vector/src/access_method/bq.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index 32f830f5..843d8c28 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -673,7 +673,7 @@ impl BqNode { Self::new( heap_pointer, meta_page.get_num_neighbors() as usize, - meta_page.get_num_dimensions() as usize, + meta_page.get_num_dimensions_to_index() as usize, bq_vector, ) } From cdf07199bc45e1ea8cd8eaf79d2f185754583805 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Wed, 24 Apr 2024 10:58:03 -0400 Subject: [PATCH 36/44] Implement bq_compressed --- timescale_vector/src/access_method/bq.rs | 315 +++++++++++------- timescale_vector/src/access_method/build.rs | 10 +- .../src/access_method/meta_page.rs | 30 +- timescale_vector/src/access_method/options.rs | 25 +- timescale_vector/src/access_method/scan.rs | 27 +- timescale_vector/src/access_method/storage.rs | 15 + timescale_vector/src/access_method/vacuum.rs | 2 +- 7 files changed, 259 insertions(+), 165 deletions(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index 843d8c28..c3eb87b5 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -16,7 +16,7 @@ use pgrx::{ pg_sys::{InvalidBlockNumber, InvalidOffsetNumber, BLCKSZ}, PgRelation, }; -use rkyv::{vec::ArchivedVec, Archive, Archived, Deserialize, Serialize}; +use rkyv::{vec::ArchivedVec, Archive, Deserialize, Serialize}; use crate::util::{ page::PageType, table_slot::TableSlot, tape::Tape, ArchivedItemPointer, HeapPointer, @@ -164,52 +164,62 @@ impl BqQuantizer { ) -> Vec { self.quantize(&full_vector) } - - fn get_distance_table( - &self, - query: &[f32], - _distance_fn: fn(&[f32], &[f32]) -> f32, - ) -> BqDistanceTable { - BqDistanceTable::new(self.quantize(query)) - } } -/// DistanceCalculator encapsulates the code to generate distances between a BQ vector and a query. -pub struct BqDistanceTable { +pub struct BqSearchDistanceMeasure { quantized_vector: Vec, + query: PgVector, + num_dimensions_for_neighbors: usize, } -impl BqDistanceTable { - pub fn new(query: Vec) -> BqDistanceTable { - BqDistanceTable { - quantized_vector: query, +impl BqSearchDistanceMeasure { + pub fn new( + quantized_vector: Vec, + query: PgVector, + num_dimensions_for_neighbors: usize, + ) -> BqSearchDistanceMeasure { + BqSearchDistanceMeasure { + quantized_vector, + query, + num_dimensions_for_neighbors, } } - pub fn distance(&self, bq_vector: &[BqVectorElement]) -> f32 { - let count_ones = distance_xor_optimized(&self.quantized_vector, bq_vector); - //dot product is LOWER the more xors that lead to 1 becaues that means a negative times a positive = negative component - //but the distance is 1 - dot product, so the more count_ones the higher the distance. - // one other check for distance(a,a), xor=0, count_ones=0, distance=0 - count_ones as f32 - } -} - -//FIXME: cleanup make this into a struct -pub enum BqSearchDistanceMeasure { - Bq(BqDistanceTable, PgVector), -} - -impl BqSearchDistanceMeasure { pub fn calculate_bq_distance( - table: &BqDistanceTable, + &self, bq_vector: &[BqVectorElement], + gns: &GraphNeighborStore, stats: &mut S, ) -> f32 { assert!(bq_vector.len() > 0); - let vec = bq_vector; stats.record_quantized_distance_comparison(); - table.distance(vec) + let (a, b) = match gns { + GraphNeighborStore::Disk => { + if self.num_dimensions_for_neighbors > 0 { + let quantized_dimensions = + BqQuantizer::quantized_size(self.num_dimensions_for_neighbors); + debug_assert!(self.quantized_vector.len() >= quantized_dimensions); + debug_assert!(bq_vector.len() >= quantized_dimensions); + ( + &self.quantized_vector.as_slice()[..quantized_dimensions], + &bq_vector[..quantized_dimensions], + ) + } else { + debug_assert!(self.quantized_vector.len() == bq_vector.len()); + (self.quantized_vector.as_slice(), bq_vector) + } + } + GraphNeighborStore::Builder(_b) => { + debug_assert!(self.quantized_vector.len() == bq_vector.len()); + (self.quantized_vector.as_slice(), bq_vector) + } + }; + + let count_ones = distance_xor_optimized(a, b); + //dot product is LOWER the more xors that lead to 1 becaues that means a negative times a positive = negative component + //but the distance is 1 - dot product, so the more count_ones the higher the distance. + // one other check for distance(a,a), xor=0, count_ones=0, distance=0 + count_ones as f32 } } @@ -297,21 +307,23 @@ pub struct BqSpeedupStorage<'a> { heap_rel: &'a PgRelation, heap_attr: pgrx::pg_sys::AttrNumber, qv_cache: RefCell, + num_dimensions_for_neighbors: usize, } impl<'a> BqSpeedupStorage<'a> { pub fn new_for_build( index: &'a PgRelation, heap_rel: &'a PgRelation, - distance_fn: fn(&[f32], &[f32]) -> f32, + meta_page: &super::meta_page::MetaPage, ) -> BqSpeedupStorage<'a> { Self { index: index, - distance_fn: distance_fn, + distance_fn: meta_page.get_distance_function(), quantizer: BqQuantizer::new(), heap_rel: heap_rel, heap_attr: get_attribute_number_from_index(index), qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), + num_dimensions_for_neighbors: meta_page.get_num_dimensions_for_neighbors() as usize, } } @@ -336,6 +348,7 @@ impl<'a> BqSpeedupStorage<'a> { heap_rel: heap_rel, heap_attr: get_attribute_number_from_index(index_relation), qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), + num_dimensions_for_neighbors: meta_page.get_num_dimensions_for_neighbors() as usize, } } @@ -343,16 +356,17 @@ impl<'a> BqSpeedupStorage<'a> { index_relation: &'a PgRelation, heap_relation: &'a PgRelation, quantizer: &BqQuantizer, - distance_fn: fn(&[f32], &[f32]) -> f32, + meta_page: &super::meta_page::MetaPage, ) -> BqSpeedupStorage<'a> { Self { index: index_relation, - distance_fn: distance_fn, + distance_fn: meta_page.get_distance_function(), //OPT: get rid of clone quantizer: quantizer.clone(), heap_rel: heap_relation, heap_attr: get_attribute_number_from_index(index_relation), qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), + num_dimensions_for_neighbors: meta_page.get_num_dimensions_for_neighbors() as usize, } } @@ -391,6 +405,7 @@ impl<'a> BqSpeedupStorage<'a> { let rn_visiting = unsafe { BqNode::read(self.index, lsn_index_pointer, &mut lsr.stats) }; let node_visiting = rn_visiting.get_archived_node(); + //OPT: get neighbors from private data just like plain storage in the self.num_dimensions_for_neighbors == 0 case let neighbors = node_visiting.get_index_pointer_to_neighbors(); for (i, &neighbor_index_pointer) in neighbors.iter().enumerate() { @@ -398,16 +413,26 @@ impl<'a> BqSpeedupStorage<'a> { continue; } - let distance = match lsr.sdm.as_ref().unwrap() { - BqSearchDistanceMeasure::Bq(table, _) => { - let bq_vector = node_visiting.neighbor_vectors[i].as_slice(); - BqSearchDistanceMeasure::calculate_bq_distance( - table, - bq_vector, - &mut lsr.stats, - ) - } + let distance = if self.num_dimensions_for_neighbors > 0 { + let bq_vector = node_visiting.neighbor_vectors[i].as_slice(); + lsr.sdm.as_ref().unwrap().calculate_bq_distance( + bq_vector, + gns, + &mut lsr.stats, + ) + } else { + let rn_neighbor = unsafe { + BqNode::read(self.index, neighbor_index_pointer, &mut lsr.stats) + }; + let node_neighbor = rn_neighbor.get_archived_node(); + let bq_vector = node_neighbor.bq_vector.as_slice(); + lsr.sdm.as_ref().unwrap().calculate_bq_distance( + bq_vector, + gns, + &mut lsr.stats, + ) }; + let lsn = ListSearchNeighbor::new( neighbor_index_pointer, distance, @@ -423,18 +448,13 @@ impl<'a> BqSpeedupStorage<'a> { if !lsr.prepare_insert(neighbor_index_pointer) { continue; } - let distance = match lsr.sdm.as_ref().unwrap() { - BqSearchDistanceMeasure::Bq(table, _) => { - let mut cache = self.qv_cache.borrow_mut(); - let bq_vector = cache.get(neighbor_index_pointer, self, &mut lsr.stats); - let dist = BqSearchDistanceMeasure::calculate_bq_distance( - table, - bq_vector, - &mut lsr.stats, - ); - dist - } - }; + let mut cache = self.qv_cache.borrow_mut(); + let bq_vector = cache.get(neighbor_index_pointer, self, &mut lsr.stats); + let distance = lsr.sdm.as_ref().unwrap().calculate_bq_distance( + bq_vector, + gns, + &mut lsr.stats, + ); let lsn = ListSearchNeighbor::new( neighbor_index_pointer, @@ -530,10 +550,10 @@ impl<'a> Storage for BqSpeedupStorage<'a> { } fn get_query_distance_measure(&self, query: PgVector) -> BqSearchDistanceMeasure { - return BqSearchDistanceMeasure::Bq( - self.quantizer - .get_distance_table(query.to_index_slice(), self.distance_fn), + return BqSearchDistanceMeasure::new( + self.quantizer.quantize(query.to_index_slice()), query, + self.num_dimensions_for_neighbors, ); } @@ -546,13 +566,10 @@ impl<'a> Storage for BqSpeedupStorage<'a> { stats: &mut S, ) -> f32 { let slot = unsafe { self.get_heap_table_slot_from_heap_pointer(heap_pointer, stats) }; - match qdm { - BqSearchDistanceMeasure::Bq(_, query) => { - let datum = unsafe { slot.get_attribute(self.heap_attr).unwrap() }; - let vec = unsafe { PgVector::from_datum(datum, meta_page, false, true) }; - self.get_distance_function()(vec.to_full_slice(), query.to_full_slice()) - } - } + + let datum = unsafe { slot.get_attribute(self.heap_attr).unwrap() }; + let vec = unsafe { PgVector::from_datum(datum, meta_page, false, true) }; + self.get_distance_function()(vec.to_full_slice(), qdm.query.to_full_slice()) } fn get_neighbors_with_distances_from_disk( @@ -565,12 +582,11 @@ impl<'a> Storage for BqSpeedupStorage<'a> { let archived = rn.get_archived_node(); let q = archived.bq_vector.as_slice(); - for (i, n) in rn.get_archived_node().iter_neighbors().enumerate() { - //let dist = unsafe { dist_state.get_distance(n, stats) }; - assert!(i < archived.neighbor_vectors.len()); - let neighbor_q = archived.neighbor_vectors[i].as_slice(); + for n in rn.get_archived_node().iter_neighbors() { + //OPT: we can optimize this if num_dimensions_for_neighbors == num_dimensions_to_index + let rn1 = unsafe { BqNode::read(self.index, n, stats) }; stats.record_quantized_distance_comparison(); - let dist = distance_xor_optimized(q, neighbor_q); + let dist = distance_xor_optimized(q, rn1.get_archived_node().bq_vector.as_slice()); result.push(NeighborWithDistance::new(n, dist as f32)) } } @@ -581,7 +597,7 @@ impl<'a> Storage for BqSpeedupStorage<'a> { &self, lsr: &mut ListSearchResult, index_pointer: ItemPointer, - _gns: &GraphNeighborStore, + gns: &GraphNeighborStore, ) -> ListSearchNeighbor { if !lsr.prepare_insert(index_pointer) { panic!("should not have had an init id already inserted"); @@ -590,15 +606,11 @@ impl<'a> Storage for BqSpeedupStorage<'a> { let rn = unsafe { BqNode::read(self.index, index_pointer, &mut lsr.stats) }; let node = rn.get_archived_node(); - let distance = match lsr.sdm.as_ref().unwrap() { - BqSearchDistanceMeasure::Bq(table, _) => { - BqSearchDistanceMeasure::calculate_bq_distance( - table, - node.bq_vector.as_slice(), - &mut lsr.stats, - ) - } - }; + let distance = lsr.sdm.as_ref().unwrap().calculate_bq_distance( + node.bq_vector.as_slice(), + gns, + &mut lsr.stats, + ); ListSearchNeighbor::new(index_pointer, distance, PhantomData::) } @@ -674,6 +686,7 @@ impl BqNode { heap_pointer, meta_page.get_num_neighbors() as usize, meta_page.get_num_dimensions_to_index() as usize, + meta_page.get_num_dimensions_for_neighbors() as usize, bq_vector, ) } @@ -681,7 +694,8 @@ impl BqNode { fn new( heap_pointer: HeapPointer, num_neighbors: usize, - num_dimensions: usize, + _num_dimensions: usize, + num_dimensions_for_neighbors: usize, bq_vector: &[BqVectorElement], ) -> Self { // always use vectors of num_neighbors in length because we never want the serialized size of a Node to change @@ -689,9 +703,13 @@ impl BqNode { .map(|_| ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber)) .collect(); - let neighbor_vectors: Vec<_> = (0..num_neighbors) - .map(|_| vec![0; BqQuantizer::quantized_size(num_dimensions as _)]) - .collect(); + let neighbor_vectors: Vec<_> = if num_dimensions_for_neighbors > 0 { + (0..num_neighbors) + .map(|_| vec![0; BqQuantizer::quantized_size(num_dimensions_for_neighbors as _)]) + .collect() + } else { + vec![] + }; Self { heap_item_pointer: heap_pointer, @@ -701,14 +719,27 @@ impl BqNode { } } - fn test_size(num_neighbors: usize, num_dimensions: usize) -> usize { + fn test_size( + num_neighbors: usize, + num_dimensions: usize, + num_dimensions_for_neighbors: usize, + ) -> usize { let v: Vec = vec![0; BqQuantizer::quantized_size(num_dimensions)]; let hp = HeapPointer::new(InvalidBlockNumber, InvalidOffsetNumber); - let n = Self::new(hp, num_neighbors, num_dimensions, &v); + let n = Self::new( + hp, + num_neighbors, + num_dimensions, + num_dimensions_for_neighbors, + &v, + ); n.serialize_to_vec().len() } - pub fn get_default_num_neighbors(num_dimensions: usize) -> usize { + pub fn get_default_num_neighbors( + num_dimensions: usize, + num_dimensions_for_neighbors: usize, + ) -> usize { //how many neighbors can fit on one page? That's what we choose. //we first overapproximate the number of neighbors and then double check by actually calculating the size of the BqNode. @@ -727,6 +758,7 @@ impl BqNode { let serialized_size = BqNode::test_size( num_neighbors_overapproximate as usize, num_dimensions as usize, + num_dimensions_for_neighbors as usize, ); if serialized_size <= page_size { return num_neighbors_overapproximate; @@ -740,20 +772,14 @@ impl BqNode { } impl ArchivedBqNode { - pub fn neighbor_index_pointer( - self: Pin<&mut Self>, - ) -> Pin<&mut ArchivedVec> { + fn neighbor_index_pointer(self: Pin<&mut Self>) -> Pin<&mut ArchivedVec> { unsafe { self.map_unchecked_mut(|s| &mut s.neighbor_index_pointers) } } - pub fn neighbor_vector(self: Pin<&mut Self>) -> Pin<&mut ArchivedVec>> { + fn neighbor_vector(self: Pin<&mut Self>) -> Pin<&mut ArchivedVec>> { unsafe { self.map_unchecked_mut(|s| &mut s.neighbor_vectors) } } - pub fn bq_vector(self: Pin<&mut Self>) -> Pin<&mut Archived>> { - unsafe { self.map_unchecked_mut(|s| &mut s.bq_vector) } - } - fn set_neighbors( mut self: Pin<&mut Self>, neighbors: &[NeighborWithDistance], @@ -767,12 +793,15 @@ impl ArchivedBqNode { a_index_pointer.block_number = ip.block_number; a_index_pointer.offset = ip.offset; - let quantized = cache.must_get(ip); - - let mut neighbor_vector = self.as_mut().neighbor_vector().index_pin(i); - for (index_in_q_vec, val) in quantized.iter().enumerate() { - let mut x = neighbor_vector.as_mut().index_pin(index_in_q_vec); - *x = *val; + if meta_page.get_num_dimensions_for_neighbors() > 0 { + let quantized = &cache.must_get(ip)[..BqQuantizer::quantized_size( + meta_page.get_num_dimensions_for_neighbors() as _, + )]; + let mut neighbor_vector = self.as_mut().neighbor_vector().index_pin(i); + for (index_in_q_vec, val) in quantized.iter().enumerate() { + let mut x = neighbor_vector.as_mut().index_pin(index_in_q_vec); + *x = *val; + } } } //set the marker that the list ended @@ -830,7 +859,7 @@ mod tests { use pgrx::*; #[pg_test] - unsafe fn test_bq_storage_index_creation_default_neighbors() -> spi::Result<()> { + unsafe fn test_bq_speedup_storage_index_creation_default_neighbors() -> spi::Result<()> { crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( "storage_layout = io_optimized", )?; @@ -838,7 +867,7 @@ mod tests { } #[pg_test] - unsafe fn test_bq_storage_index_creation_few_neighbors() -> spi::Result<()> { + unsafe fn test_bq_speedup_storage_index_creation_few_neighbors() -> spi::Result<()> { //a test with few neighbors tests the case that nodes share a page, which has caused deadlocks in the past. crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( "num_neighbors=10, storage_layout = io_optimized", @@ -847,35 +876,35 @@ mod tests { } #[test] - fn test_bq_storage_delete_vacuum_plain() { + fn test_bq_speedup_storage_delete_vacuum_plain() { crate::access_method::vacuum::tests::test_delete_vacuum_plain_scaffold( "num_neighbors = 10, storage_layout = io_optimized", ); } #[test] - fn test_bq_storage_delete_vacuum_full() { + fn test_bq_speedup_storage_delete_vacuum_full() { crate::access_method::vacuum::tests::test_delete_vacuum_full_scaffold( "num_neighbors = 38, storage_layout = io_optimized", ); } #[pg_test] - unsafe fn test_bq_storage_empty_table_insert() -> spi::Result<()> { + unsafe fn test_bq_speedup_storage_empty_table_insert() -> spi::Result<()> { crate::access_method::build::tests::test_empty_table_insert_scaffold( "num_neighbors=38, storage_layout = io_optimized", ) } #[pg_test] - unsafe fn test_bq_storage_insert_empty_insert() -> spi::Result<()> { + unsafe fn test_bq_speedup_storage_insert_empty_insert() -> spi::Result<()> { crate::access_method::build::tests::test_insert_empty_insert_scaffold( "num_neighbors=38, storage_layout = io_optimized", ) } #[pg_test] - unsafe fn test_bq_storage_index_creation_num_dimensions() -> spi::Result<()> { + unsafe fn test_bq_speedup_storage_index_creation_num_dimensions() -> spi::Result<()> { crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( "storage_layout = io_optimized, num_dimensions=768", )?; @@ -883,11 +912,73 @@ mod tests { } #[pg_test] - unsafe fn test_bq_storage_index_updates() -> spi::Result<()> { + unsafe fn test_bq_speedup_storage_index_updates() -> spi::Result<()> { crate::access_method::build::tests::test_index_updates( "storage_layout = io_optimized, num_neighbors=10", 300, )?; Ok(()) } + + #[pg_test] + unsafe fn test_bq_speedup_compressed_index_creation_default_neighbors() -> spi::Result<()> { + crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( + "storage_layout = memory_optimized", + )?; + Ok(()) + } + + #[pg_test] + unsafe fn test_bq_compressed_storage_index_creation_few_neighbors() -> spi::Result<()> { + //a test with few neighbors tests the case that nodes share a page, which has caused deadlocks in the past. + crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( + "num_neighbors=10, storage_layout = memory_optimized", + )?; + Ok(()) + } + + #[test] + fn test_bq_compressed_storage_delete_vacuum_plain() { + crate::access_method::vacuum::tests::test_delete_vacuum_plain_scaffold( + "num_neighbors = 10, storage_layout = memory_optimized", + ); + } + + #[test] + fn test_bq_compressed_storage_delete_vacuum_full() { + crate::access_method::vacuum::tests::test_delete_vacuum_full_scaffold( + "num_neighbors = 38, storage_layout = memory_optimized", + ); + } + + #[pg_test] + unsafe fn test_bq_compressed_storage_empty_table_insert() -> spi::Result<()> { + crate::access_method::build::tests::test_empty_table_insert_scaffold( + "num_neighbors=38, storage_layout = memory_optimized", + ) + } + + #[pg_test] + unsafe fn test_bq_compressed_storage_insert_empty_insert() -> spi::Result<()> { + crate::access_method::build::tests::test_insert_empty_insert_scaffold( + "num_neighbors=38, storage_layout = memory_optimized", + ) + } + + #[pg_test] + unsafe fn test_bq_compressed_storage_index_creation_num_dimensions() -> spi::Result<()> { + crate::access_method::build::tests::test_index_creation_and_accuracy_scaffold( + "storage_layout = memory_optimized, num_dimensions=768", + )?; + Ok(()) + } + + #[pg_test] + unsafe fn test_bq_compressed_storage_index_updates() -> spi::Result<()> { + crate::access_method::build::tests::test_index_updates( + "storage_layout = memory_optimized, num_neighbors=10", + 300, + )?; + Ok(()) + } } diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index d4cc88d3..5e24ac35 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -127,7 +127,7 @@ pub unsafe extern "C" fn aminsert( &mut stats, ); } - StorageType::BqSpeedup => { + StorageType::BqSpeedup | StorageType::BqCompression => { let bq = BqSpeedupStorage::load_for_insert( &heap_relation, &index_relation, @@ -211,12 +211,8 @@ fn do_heap_scan<'a>( finalize_index_build(&mut plain, &mut bs, write_stats) } - StorageType::BqSpeedup => { - let mut bq = BqSpeedupStorage::new_for_build( - index_relation, - heap_relation, - meta_page.get_distance_function(), - ); + StorageType::BqSpeedup | StorageType::BqCompression => { + let mut bq = BqSpeedupStorage::new_for_build(index_relation, heap_relation, &meta_page); let page_type = BqSpeedupStorage::page_type(); diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index f02dac7d..9aeb0066 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -140,6 +140,16 @@ impl MetaPage { self.num_dimensions_to_index } + pub fn get_num_dimensions_for_neighbors(&self) -> u32 { + match StorageType::from_u8(self.storage_type) { + StorageType::Plain => { + error!("get_num_dimensions_for_neighbors should not be called for Plain storage") + } + StorageType::BqSpeedup => self.num_dimensions_to_index, + StorageType::BqCompression => 0, + } + } + /// Maximum number of neigbors per node. Given that we pre-allocate /// these many slots for each node, this cannot change after the graph is built. pub fn get_num_neighbors(&self) -> u32 { @@ -178,22 +188,26 @@ impl MetaPage { } pub fn get_quantizer_metadata_pointer(&self) -> Option { - if (self.storage_type != StorageType::BqSpeedup as u8) - || !self.quantizer_metadata.is_valid() - { + if !self.quantizer_metadata.is_valid() { return None; } - Some(self.quantizer_metadata) + match self.get_storage_type() { + StorageType::Plain => None, + StorageType::BqSpeedup | StorageType::BqCompression => Some(self.quantizer_metadata), + } } fn calculate_num_neighbors(num_dimensions: u32, opt: &PgBox) -> u32 { let num_neighbors = (*opt).get_num_neighbors(); if num_neighbors == NUM_NEIGHBORS_DEFAULT_SENTINEL { - if (*opt).get_storage_type() == StorageType::Plain { - 50 - } else { - BqNode::get_default_num_neighbors(num_dimensions as usize) as u32 + match (*opt).get_storage_type() { + StorageType::Plain => 50, + StorageType::BqSpeedup => BqNode::get_default_num_neighbors( + num_dimensions as usize, + num_dimensions as usize, + ) as u32, + StorageType::BqCompression => 50, } } else { num_neighbors as u32 diff --git a/timescale_vector/src/access_method/options.rs b/timescale_vector/src/access_method/options.rs index 3fd1f3a4..f66a56f0 100644 --- a/timescale_vector/src/access_method/options.rs +++ b/timescale_vector/src/access_method/options.rs @@ -63,12 +63,11 @@ impl TSVIndexOptions { } pub fn get_storage_type(&self) -> StorageType { - let s = self.get_str(self.storage_layout_offset, || "io_optimized".to_string()); - match s.as_str() { - "io_optimized" => StorageType::BqSpeedup, - "plain" => StorageType::Plain, - _ => panic!("invalid storage_layout: {}", s), - } + let s = self.get_str(self.storage_layout_offset, || { + super::storage::DEFAULT_STORAGE_TYPE_STR.to_owned() + }); + + StorageType::from_str(s.as_str()) } fn get_str String>(&self, offset: i32, default: F) -> String { @@ -164,14 +163,8 @@ extern "C" fn validate_storage_layout(value: *const std::os::raw::c_char) { let value = unsafe { CStr::from_ptr(value) } .to_str() - .expect("failed to parse storage_layout value") - .to_lowercase(); - if value != "io_optimized" && value != "plain" { - panic!( - "invalid storage_layout. Must be one of 'io_optimized' or 'plain': {}", - value - ) - } + .expect("failed to parse storage_layout value"); + _ = StorageType::from_str(value); } pub unsafe fn init() { @@ -180,8 +173,8 @@ pub unsafe fn init() { pg_sys::add_string_reloption( RELOPT_KIND_TSV, "storage_layout".as_pg_cstr(), - "Storage layout: either io_optimized or plain".as_pg_cstr(), - "io_optimized".as_pg_cstr(), + "Storage layout: either memory_optimized, io_optimized, or plain".as_pg_cstr(), + super::storage::DEFAULT_STORAGE_TYPE_STR.as_pg_cstr(), Some(validate_storage_layout), pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, ); diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index ca4b569d..7b1322c0 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -66,15 +66,10 @@ impl TSVScanState { TSVResponseIterator::new(&bq, index, query, search_list_size, meta_page, stats); StorageState::Plain(it) } - StorageType::BqSpeedup => { + StorageType::BqSpeedup | StorageType::BqCompression => { let mut stats = QuantizerStats::new(); let quantizer = unsafe { BqMeans::load(index, &meta_page, &mut stats) }; - let bq = BqSpeedupStorage::load_for_search( - index, - heap, - &quantizer, - meta_page.get_distance_function(), - ); + let bq = BqSpeedupStorage::load_for_search(index, heap, &quantizer, &meta_page); let it = TSVResponseIterator::new(&bq, index, query, search_list_size, meta_page, stats); StorageState::BqSpeedup(quantizer, it) @@ -115,7 +110,6 @@ impl Ord for ResortData { } struct StreamingStats { - resort_size: usize, count: i32, mean: f32, m2: f32, @@ -123,9 +117,8 @@ struct StreamingStats { } impl StreamingStats { - fn new(resort_size: usize) -> Self { + fn new(_resort_size: usize) -> Self { Self { - resort_size, count: 0, mean: 0.0, m2: 0.0, @@ -151,10 +144,6 @@ impl StreamingStats { self.m2 / (self.count - 1) as f32 } - fn mean(&self) -> f32 { - self.mean - } - fn update(&mut self, distance: f32, diff: f32) { //base stats only on first resort_size elements self.update_base_stats(diff); @@ -230,7 +219,7 @@ impl TSVResponseIterator { fn next_with_resort>( &mut self, - index: &PgRelation, + _index: &PgRelation, storage: &S, ) -> Option<(HeapPointer, IndexPointer)> { if self.resort_buffer.capacity() == 0 { @@ -369,12 +358,8 @@ pub extern "C" fn amgettuple( let mut storage = unsafe { state.storage.as_mut() }.expect("no storage in state"); match &mut storage { StorageState::BqSpeedup(quantizer, iter) => { - let bq = BqSpeedupStorage::load_for_search( - &indexrel, - &heaprel, - quantizer, - state.distance_fn.unwrap(), - ); + let bq = + BqSpeedupStorage::load_for_search(&indexrel, &heaprel, quantizer, &state.meta_page); let next = iter.next_with_resort(&indexrel, &bq); get_tuple(state, next, scan) } diff --git a/timescale_vector/src/access_method/storage.rs b/timescale_vector/src/access_method/storage.rs index bedad284..30287e31 100644 --- a/timescale_vector/src/access_method/storage.rs +++ b/timescale_vector/src/access_method/storage.rs @@ -128,14 +128,29 @@ pub trait Storage { pub enum StorageType { Plain = 0, BqSpeedup = 1, + BqCompression = 2, } +pub const DEFAULT_STORAGE_TYPE_STR: &str = "memory_optimized"; + impl StorageType { pub fn from_u8(value: u8) -> Self { match value { 0 => StorageType::Plain, 1 => StorageType::BqSpeedup, + 2 => StorageType::BqCompression, _ => panic!("Invalid storage type"), } } + + pub fn from_str(value: &str) -> Self { + match value.to_lowercase().as_str() { + "plain" => StorageType::Plain, + "bq_speedup" | "io_optimized" => StorageType::BqSpeedup, + "bq_compression" | "memory_optimized" => StorageType::BqCompression, + _ => panic!( + "Invalid storage type. Must be one of 'plain', 'bq_speedup', 'bq_compression'" + ), + } + } } diff --git a/timescale_vector/src/access_method/vacuum.rs b/timescale_vector/src/access_method/vacuum.rs index cd3104e1..2c46524a 100644 --- a/timescale_vector/src/access_method/vacuum.rs +++ b/timescale_vector/src/access_method/vacuum.rs @@ -40,7 +40,7 @@ pub extern "C" fn ambulkdelete( let meta_page = MetaPage::fetch(&index_relation); let storage = meta_page.get_storage_type(); match storage { - StorageType::BqSpeedup => { + StorageType::BqSpeedup | StorageType::BqCompression => { bulk_delete_for_storage::( &index_relation, nblocks, From 7a4d5fd1f9fb9c371207a8fd17886d7780085d52 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Wed, 24 Apr 2024 17:10:05 -0400 Subject: [PATCH 37/44] Better progress tracking and statistics --- timescale_vector/src/access_method/build.rs | 34 +++++++++++++++++++ timescale_vector/src/access_method/graph.rs | 4 ++- timescale_vector/src/access_method/mod.rs | 2 ++ timescale_vector/src/access_method/options.rs | 2 +- timescale_vector/src/access_method/scan.rs | 23 ++++++++++--- timescale_vector/src/access_method/stats.rs | 31 +++++++++++++++++ timescale_vector/src/util/ports.rs | 3 ++ 7 files changed, 93 insertions(+), 6 deletions(-) diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index 5e24ac35..82a7421e 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -1,5 +1,6 @@ use std::time::Instant; +use pgrx::pg_sys::{pgstat_progress_update_param, AsPgCStr}; use pgrx::*; use crate::access_method::graph::Graph; @@ -12,6 +13,8 @@ use crate::util::page::PageType; use crate::util::tape::Tape; use crate::util::*; +use self::ports::PROGRESS_CREATE_IDX_SUBPHASE; + use super::bq::BqSpeedupStorage; use super::graph_neighbor_store::BuilderNeighborCache; @@ -216,6 +219,10 @@ fn do_heap_scan<'a>( let page_type = BqSpeedupStorage::page_type(); + unsafe { + pgstat_progress_update_param(PROGRESS_CREATE_IDX_SUBPHASE, BUILD_PHASE_TRAINING); + } + bq.start_training(&meta_page); let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); @@ -232,6 +239,13 @@ fn do_heap_scan<'a>( } bq.finish_training(&mut write_stats); + unsafe { + pgstat_progress_update_param( + PROGRESS_CREATE_IDX_SUBPHASE, + BUILD_PHASE_BUILDING_GRAPH, + ); + } + let mut state = StorageBuildState::BqSpeedup(&mut bq, &mut bs); unsafe { @@ -244,6 +258,12 @@ fn do_heap_scan<'a>( ); } + unsafe { + pgstat_progress_update_param( + PROGRESS_CREATE_IDX_SUBPHASE, + BUILD_PHASE_FINALIZING_GRAPH, + ); + } finalize_index_build(&mut bq, &mut bs, write_stats) } } @@ -420,6 +440,20 @@ fn build_callback_internal( .insert(&index, index_pointer, vector, storage, &mut state.stats); } +const BUILD_PHASE_TRAINING: i64 = 0; +const BUILD_PHASE_BUILDING_GRAPH: i64 = 1; +const BUILD_PHASE_FINALIZING_GRAPH: i64 = 2; + +#[pg_guard] +pub unsafe extern "C" fn ambuildphasename(phasenum: i64) -> *mut ffi::c_char { + match phasenum { + BUILD_PHASE_TRAINING => "training quantizer".as_pg_cstr(), + BUILD_PHASE_BUILDING_GRAPH => "building graph".as_pg_cstr(), + BUILD_PHASE_FINALIZING_GRAPH => "finalizing graph".as_pg_cstr(), + _ => error!("Unknown phase number {}", phasenum), + } +} + #[cfg(any(test, feature = "pg_test"))] #[pgrx::pg_schema] pub mod tests { diff --git a/timescale_vector/src/access_method/graph.rs b/timescale_vector/src/access_method/graph.rs index 55785369..b939a969 100644 --- a/timescale_vector/src/access_method/graph.rs +++ b/timescale_vector/src/access_method/graph.rs @@ -11,7 +11,7 @@ use crate::util::{HeapPointer, IndexPointer, ItemPointer}; use super::graph_neighbor_store::GraphNeighborStore; use super::pg_vector::PgVector; -use super::stats::{GreedySearchStats, InsertStats, PruneNeighborStats}; +use super::stats::{GreedySearchStats, InsertStats, PruneNeighborStats, StatsNodeVisit}; use super::storage::Storage; use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance}; @@ -108,6 +108,7 @@ impl ListSearchResult { /// Internal function pub fn insert_neighbor(&mut self, n: ListSearchNeighbor) { + self.stats.record_candidate(); self.candidates.push(Reverse(n)); } @@ -318,6 +319,7 @@ impl<'a> Graph<'a> { )); } } + lsr.stats.record_visit(); storage.visit_lsn(lsr, list_search_entry_idx, &self.neighbor_store); } } diff --git a/timescale_vector/src/access_method/mod.rs b/timescale_vector/src/access_method/mod.rs index 4a859f8f..7ffd4a26 100644 --- a/timescale_vector/src/access_method/mod.rs +++ b/timescale_vector/src/access_method/mod.rs @@ -82,6 +82,8 @@ fn amhandler(_fcinfo: pg_sys::FunctionCallInfo) -> PgBox amroutine.amgetbitmap = None; amroutine.amendscan = Some(scan::amendscan); + amroutine.ambuildphasename = Some(build::ambuildphasename); + amroutine.into_pg_boxed() } diff --git a/timescale_vector/src/access_method/options.rs b/timescale_vector/src/access_method/options.rs index f66a56f0..2bd5a0a2 100644 --- a/timescale_vector/src/access_method/options.rs +++ b/timescale_vector/src/access_method/options.rs @@ -268,7 +268,7 @@ mod tests { assert_eq!(options.search_list_size, 100); assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); assert_eq!(options.num_dimensions, NUM_DIMENSIONS_DEFAULT_SENTINEL); - assert_eq!(options.get_storage_type(), StorageType::BqSpeedup); + assert_eq!(options.get_storage_type(), StorageType::BqCompression); Ok(()) } diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 7b1322c0..1bbde192 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -159,6 +159,9 @@ struct TSVResponseIterator { resort_size: usize, resort_buffer: BinaryHeap, streaming_stats: StreamingStats, + next_calls: i32, + next_calls_with_resort: i32, + full_distance_comparisons: i32, } impl TSVResponseIterator { @@ -185,6 +188,9 @@ impl TSVResponseIterator { resort_size, resort_buffer: BinaryHeap::with_capacity(resort_size), streaming_stats: StreamingStats::new(resort_size), + next_calls: 0, + next_calls_with_resort: 0, + full_distance_comparisons: 0, } } } @@ -194,6 +200,7 @@ impl TSVResponseIterator { &mut self, storage: &S, ) -> Option<(HeapPointer, IndexPointer)> { + self.next_calls += 1; let graph = Graph::new(GraphNeighborStore::Disk, &mut self.meta_page); /* Iterate until we find a non-deleted tuple */ @@ -222,6 +229,7 @@ impl TSVResponseIterator { _index: &PgRelation, storage: &S, ) -> Option<(HeapPointer, IndexPointer)> { + self.next_calls_with_resort += 1; if self.resort_buffer.capacity() == 0 { return self.next(storage); } @@ -233,6 +241,7 @@ impl TSVResponseIterator { { match self.next(storage) { Some((heap_pointer, index_pointer)) => { + self.full_distance_comparisons += 1; let distance = storage.get_full_distance_for_resort( self.lsr.sdm.as_ref().unwrap(), index_pointer, @@ -433,13 +442,19 @@ pub extern "C" fn amendscan(scan: pg_sys::IndexScanDesc) { fn end_scan( iter: &mut TSVResponseIterator, ) { + debug_assert!(iter.quantizer_stats.node_reads == 1); + debug_assert!(iter.quantizer_stats.node_writes == 0); + debug1!( - "Query stats - node reads:{}, calls: {}, total distance comparisons: {}, quantized distance comparisons: {}, quantizer r/w: {}/{}", + "Query stats - reads_index={} reads_heap={} d_total={} d_quantized={} d_full={} next={} resort={} visits={} candidate={}", iter.lsr.stats.get_node_reads(), - iter.lsr.stats.get_calls(), + iter.lsr.stats.get_node_heap_reads(), iter.lsr.stats.get_total_distance_comparisons(), iter.lsr.stats.get_quantized_distance_comparisons(), - iter.quantizer_stats.node_reads, - iter.quantizer_stats.node_writes, + iter.full_distance_comparisons, + iter.next_calls, + iter.next_calls_with_resort, + iter.lsr.stats.get_visited_nodes(), + iter.lsr.stats.get_candidate_nodes(), ); } diff --git a/timescale_vector/src/access_method/stats.rs b/timescale_vector/src/access_method/stats.rs index ab2c6ceb..f6e5bd2a 100644 --- a/timescale_vector/src/access_method/stats.rs +++ b/timescale_vector/src/access_method/stats.rs @@ -21,6 +21,11 @@ pub trait StatsDistanceComparison { fn record_quantized_distance_comparison(&mut self); } +pub trait StatsNodeVisit { + fn record_visit(&mut self); + fn record_candidate(&mut self); +} + #[derive(Debug)] pub struct PruneNeighborStats { pub calls: usize, @@ -73,6 +78,8 @@ pub struct GreedySearchStats { node_reads: usize, node_heap_reads: usize, quantized_distance_comparisons: usize, + visited_nodes: usize, + candidate_nodes: usize, } impl GreedySearchStats { @@ -83,6 +90,8 @@ impl GreedySearchStats { node_reads: 0, node_heap_reads: 0, quantized_distance_comparisons: 0, + visited_nodes: 0, + candidate_nodes: 0, } } @@ -102,6 +111,10 @@ impl GreedySearchStats { self.node_reads } + pub fn get_node_heap_reads(&self) -> usize { + self.node_heap_reads + } + pub fn get_total_distance_comparisons(&self) -> usize { self.full_distance_comparisons + self.quantized_distance_comparisons } @@ -110,6 +123,14 @@ impl GreedySearchStats { self.quantized_distance_comparisons } + pub fn get_visited_nodes(&self) -> usize { + self.visited_nodes + } + + pub fn get_candidate_nodes(&self) -> usize { + self.candidate_nodes + } + pub fn get_full_distance_comparisons(&self) -> usize { self.full_distance_comparisons } @@ -141,6 +162,16 @@ impl StatsDistanceComparison for GreedySearchStats { } } +impl StatsNodeVisit for GreedySearchStats { + fn record_visit(&mut self) { + self.visited_nodes += 1; + } + + fn record_candidate(&mut self) { + self.candidate_nodes += 1; + } +} + #[derive(Debug)] pub struct QuantizerStats { pub node_reads: usize, diff --git a/timescale_vector/src/util/ports.rs b/timescale_vector/src/util/ports.rs index 6c8f32c4..f2bbcf13 100644 --- a/timescale_vector/src/util/ports.rs +++ b/timescale_vector/src/util/ports.rs @@ -2,6 +2,8 @@ //! Following pgrx conventions, we keep function names as close to Postgres as possible. //! Thus, we don't follow rust naming conventions. +use std::os::raw::c_int; + use memoffset::*; use pgrx::pg_sys::{Datum, ItemId, OffsetNumber, Pointer, TupleTableSlot}; use pgrx::{pg_sys, PgBox}; @@ -35,6 +37,7 @@ pub unsafe fn PageValidateSpecialPointer(page: pgrx::pg_sys::Page) { #[allow(non_upper_case_globals)] const SizeOfPageHeaderData: usize = offset_of!(pgrx::pg_sys::PageHeaderData, pd_linp); +pub const PROGRESS_CREATE_IDX_SUBPHASE: c_int = 10; #[allow(non_snake_case)] pub unsafe fn PageGetContents(page: pgrx::pg_sys::Page) -> *mut std::os::raw::c_char { From 35f0a79b3a6df2dde63d03c29cb42f9b23330503 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Sat, 27 Apr 2024 13:36:25 -0400 Subject: [PATCH 38/44] Allow using multiple bits in bq compression We found that when using 3 bits instead of one, with dimensions=768 then we get a better accuracy with the same number of shared buffers read. This isn't true for higher dimensions (e.g. 1538) --- timescale_vector/src/access_method/bq.rs | 164 ++++++++++++++---- .../src/access_method/distance.rs | 19 ++ .../src/access_method/meta_page.rs | 49 +++++- timescale_vector/src/access_method/options.rs | 20 ++- timescale_vector/src/access_method/vacuum.rs | 56 +++--- timescale_vector/src/util/page.rs | 9 +- 6 files changed, 248 insertions(+), 69 deletions(-) diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/bq.rs index c3eb87b5..520715b8 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/bq.rs @@ -35,6 +35,7 @@ const BITS_STORE_TYPE_SIZE: usize = 64; pub struct BqMeans { count: u64, means: Vec, + m2: Vec, } impl BqMeans { @@ -43,7 +44,7 @@ impl BqMeans { meta_page: &super::meta_page::MetaPage, stats: &mut S, ) -> BqQuantizer { - let mut quantizer = BqQuantizer::new(); + let mut quantizer = BqQuantizer::new(meta_page); if quantizer.use_mean { if meta_page.get_quantizer_metadata_pointer().is_none() { pgrx::error!("No BQ pointer found in meta page"); @@ -52,7 +53,11 @@ impl BqMeans { let bq = BqMeans::read(index, quantizer_item_pointer, stats); let archived = bq.get_archived_node(); - quantizer.load(archived.count, archived.means.to_vec()); + quantizer.load( + archived.count, + archived.means.to_vec(), + archived.m2.to_vec(), + ); } quantizer } @@ -66,6 +71,7 @@ impl BqMeans { let node = BqMeans { count: quantizer.count, means: quantizer.mean.to_vec(), + m2: quantizer.m2.to_vec(), }; let ptr = node.write(&mut tape, stats); tape.close(); @@ -79,49 +85,88 @@ pub struct BqQuantizer { training: bool, pub count: u64, pub mean: Vec, + pub m2: Vec, + pub num_bits_per_dimension: u8, } impl BqQuantizer { - fn new() -> BqQuantizer { + fn new(meta_page: &super::meta_page::MetaPage) -> BqQuantizer { Self { use_mean: true, training: false, count: 0, mean: vec![], + m2: vec![], + num_bits_per_dimension: meta_page.get_bq_num_bits_per_dimension(), } } - fn load(&mut self, count: u64, mean: Vec) { + fn load(&mut self, count: u64, mean: Vec, m2: Vec) { self.count = count; self.mean = mean; + self.m2 = m2 + } + + fn quantized_size(&self, full_vector_size: usize) -> usize { + Self::quantized_size_internal(full_vector_size, self.num_bits_per_dimension) } - fn quantized_size(full_vector_size: usize) -> usize { - if full_vector_size % BITS_STORE_TYPE_SIZE == 0 { - full_vector_size / BITS_STORE_TYPE_SIZE + fn quantized_size_internal(full_vector_size: usize, num_bits_per_dimension: u8) -> usize { + let num_bits = full_vector_size * num_bits_per_dimension as usize; + + if num_bits % BITS_STORE_TYPE_SIZE == 0 { + num_bits / BITS_STORE_TYPE_SIZE } else { - (full_vector_size / BITS_STORE_TYPE_SIZE) + 1 + (num_bits / BITS_STORE_TYPE_SIZE) + 1 } } - fn quantized_size_bytes(num_dimensions: usize) -> usize { - Self::quantized_size(num_dimensions) * std::mem::size_of::() + fn quantized_size_bytes(num_dimensions: usize, num_bits_per_dimension: u8) -> usize { + Self::quantized_size_internal(num_dimensions, num_bits_per_dimension) + * std::mem::size_of::() } fn quantize(&self, full_vector: &[f32]) -> Vec { assert!(!self.training); if self.use_mean { - let mut res_vector = vec![0; Self::quantized_size(full_vector.len())]; + let mut res_vector = vec![0; self.quantized_size(full_vector.len())]; - for (i, &v) in full_vector.iter().enumerate() { - if v > self.mean[i] { - res_vector[i / BITS_STORE_TYPE_SIZE] |= 1 << (i % BITS_STORE_TYPE_SIZE); + if self.num_bits_per_dimension == 1 { + for (i, &v) in full_vector.iter().enumerate() { + if v > self.mean[i] { + res_vector[i / BITS_STORE_TYPE_SIZE] |= 1 << (i % BITS_STORE_TYPE_SIZE); + } + } + } else { + for (i, &v) in full_vector.iter().enumerate() { + let mean = self.mean[i]; + let variance = self.m2[i] / self.count as f32; + let std_dev = variance.sqrt(); + let ranges = self.num_bits_per_dimension + 1; + + let v_z_score = (v - mean) / std_dev; + let index = (v_z_score + 2.0) / (4.0 / ranges as f32); //we consider z scores between -2 and 2 and divide them into {ranges} ranges + + let bit_position = i * self.num_bits_per_dimension as usize; + if index < 1.0 { + //all zeros + } else { + let count_ones = + (index.floor() as usize).min(self.num_bits_per_dimension as usize); + //fill in count_ones bits from the left + // ex count_ones=1: 100 + // ex count_ones=2: 110 + // ex count_ones=3: 111 + for j in 0..count_ones { + res_vector[(bit_position + j) / BITS_STORE_TYPE_SIZE] |= + 1 << ((bit_position + j) % BITS_STORE_TYPE_SIZE); + } + } } } - res_vector } else { - let mut res_vector = vec![0; Self::quantized_size(full_vector.len())]; + let mut res_vector = vec![0; self.quantized_size(full_vector.len())]; for (i, &v) in full_vector.iter().enumerate() { if v > 0.0 { @@ -138,6 +183,9 @@ impl BqQuantizer { if self.use_mean { self.count = 0; self.mean = vec![0.0; meta_page.get_num_dimensions_to_index() as _]; + if self.num_bits_per_dimension > 1 { + self.m2 = vec![0.0; meta_page.get_num_dimensions_to_index() as _]; + } } } @@ -146,10 +194,33 @@ impl BqQuantizer { self.count += 1; assert!(self.mean.len() == sample.len()); - self.mean - .iter_mut() - .zip(sample.iter()) - .for_each(|(m, s)| *m += (s - *m) / self.count as f32); + if self.num_bits_per_dimension > 1 { + assert!(self.m2.len() == sample.len()); + let delta: Vec<_> = self + .mean + .iter() + .zip(sample.iter()) + .map(|(m, s)| s - *m) + .collect(); + + self.mean + .iter_mut() + .zip(sample.iter()) + .for_each(|(m, s)| *m += (s - *m) / self.count as f32); + + let delta2 = self.mean.iter().zip(sample.iter()).map(|(m, s)| s - *m); + + self.m2 + .iter_mut() + .zip(delta.iter()) + .zip(delta2) + .for_each(|((m2, d), d2)| *m2 += d * d2); + } else { + self.mean + .iter_mut() + .zip(sample.iter()) + .for_each(|(m, s)| *m += (s - *m) / self.count as f32); + } } } @@ -170,18 +241,20 @@ pub struct BqSearchDistanceMeasure { quantized_vector: Vec, query: PgVector, num_dimensions_for_neighbors: usize, + quantized_dimensions: usize, } impl BqSearchDistanceMeasure { pub fn new( - quantized_vector: Vec, + quantizer: &BqQuantizer, query: PgVector, num_dimensions_for_neighbors: usize, ) -> BqSearchDistanceMeasure { BqSearchDistanceMeasure { - quantized_vector, + quantized_vector: quantizer.quantize(query.to_index_slice()), query, num_dimensions_for_neighbors, + quantized_dimensions: quantizer.quantized_size(num_dimensions_for_neighbors), } } @@ -196,13 +269,11 @@ impl BqSearchDistanceMeasure { let (a, b) = match gns { GraphNeighborStore::Disk => { if self.num_dimensions_for_neighbors > 0 { - let quantized_dimensions = - BqQuantizer::quantized_size(self.num_dimensions_for_neighbors); - debug_assert!(self.quantized_vector.len() >= quantized_dimensions); - debug_assert!(bq_vector.len() >= quantized_dimensions); + debug_assert!(self.quantized_vector.len() >= self.quantized_dimensions); + debug_assert!(bq_vector.len() >= self.quantized_dimensions); ( - &self.quantized_vector.as_slice()[..quantized_dimensions], - &bq_vector[..quantized_dimensions], + &self.quantized_vector.as_slice()[..self.quantized_dimensions], + &bq_vector[..self.quantized_dimensions], ) } else { debug_assert!(self.quantized_vector.len() == bq_vector.len()); @@ -319,7 +390,7 @@ impl<'a> BqSpeedupStorage<'a> { Self { index: index, distance_fn: meta_page.get_distance_function(), - quantizer: BqQuantizer::new(), + quantizer: BqQuantizer::new(meta_page), heap_rel: heap_rel, heap_attr: get_attribute_number_from_index(index), qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), @@ -499,7 +570,12 @@ impl<'a> Storage for BqSpeedupStorage<'a> { ) -> ItemPointer { let bq_vector = self.quantizer.vector_for_new_node(meta_page, full_vector); - let node = BqNode::with_meta(heap_pointer, &meta_page, bq_vector.as_slice()); + let node = BqNode::with_meta( + &self.quantizer, + heap_pointer, + &meta_page, + bq_vector.as_slice(), + ); let index_pointer: IndexPointer = node.write(tape, stats); index_pointer @@ -551,7 +627,7 @@ impl<'a> Storage for BqSpeedupStorage<'a> { fn get_query_distance_measure(&self, query: PgVector) -> BqSearchDistanceMeasure { return BqSearchDistanceMeasure::new( - self.quantizer.quantize(query.to_index_slice()), + &self.quantizer, query, self.num_dimensions_for_neighbors, ); @@ -678,6 +754,7 @@ pub struct BqNode { impl BqNode { pub fn with_meta( + quantizer: &BqQuantizer, heap_pointer: HeapPointer, meta_page: &MetaPage, bq_vector: &[BqVectorElement], @@ -687,6 +764,7 @@ impl BqNode { meta_page.get_num_neighbors() as usize, meta_page.get_num_dimensions_to_index() as usize, meta_page.get_num_dimensions_for_neighbors() as usize, + quantizer.num_bits_per_dimension, bq_vector, ) } @@ -696,6 +774,7 @@ impl BqNode { num_neighbors: usize, _num_dimensions: usize, num_dimensions_for_neighbors: usize, + num_bits_per_dimension: u8, bq_vector: &[BqVectorElement], ) -> Self { // always use vectors of num_neighbors in length because we never want the serialized size of a Node to change @@ -705,7 +784,15 @@ impl BqNode { let neighbor_vectors: Vec<_> = if num_dimensions_for_neighbors > 0 { (0..num_neighbors) - .map(|_| vec![0; BqQuantizer::quantized_size(num_dimensions_for_neighbors as _)]) + .map(|_| { + vec![ + 0; + BqQuantizer::quantized_size_internal( + num_dimensions_for_neighbors as _, + num_bits_per_dimension + ) + ] + }) .collect() } else { vec![] @@ -723,14 +810,17 @@ impl BqNode { num_neighbors: usize, num_dimensions: usize, num_dimensions_for_neighbors: usize, + num_bits_per_dimension: u8, ) -> usize { - let v: Vec = vec![0; BqQuantizer::quantized_size(num_dimensions)]; + let v: Vec = + vec![0; BqQuantizer::quantized_size_internal(num_dimensions, num_bits_per_dimension)]; let hp = HeapPointer::new(InvalidBlockNumber, InvalidOffsetNumber); let n = Self::new( hp, num_neighbors, num_dimensions, num_dimensions_for_neighbors, + num_bits_per_dimension, &v, ); n.serialize_to_vec().len() @@ -739,6 +829,7 @@ impl BqNode { pub fn get_default_num_neighbors( num_dimensions: usize, num_dimensions_for_neighbors: usize, + num_bits_per_dimension: u8, ) -> usize { //how many neighbors can fit on one page? That's what we choose. @@ -747,7 +838,8 @@ impl BqNode { //blocksize - 100 bytes for the padding/header/etc. let page_size = BLCKSZ as usize - 50; //one quantized_vector takes this many bytes - let vec_size = BqQuantizer::quantized_size_bytes(num_dimensions as usize) + 1; + let vec_size = + BqQuantizer::quantized_size_bytes(num_dimensions as usize, num_bits_per_dimension) + 1; //start from the page size then subtract the heap_item_pointer and bq_vector elements of BqNode. let starting = BLCKSZ as usize - std::mem::size_of::() - vec_size; //one neigbors contribution to neighbor_index_pointers + neighbor_vectors in BqNode. @@ -759,6 +851,7 @@ impl BqNode { num_neighbors_overapproximate as usize, num_dimensions as usize, num_dimensions_for_neighbors as usize, + num_bits_per_dimension, ); if serialized_size <= page_size { return num_neighbors_overapproximate; @@ -794,8 +887,9 @@ impl ArchivedBqNode { a_index_pointer.offset = ip.offset; if meta_page.get_num_dimensions_for_neighbors() > 0 { - let quantized = &cache.must_get(ip)[..BqQuantizer::quantized_size( + let quantized = &cache.must_get(ip)[..BqQuantizer::quantized_size_internal( meta_page.get_num_dimensions_for_neighbors() as _, + meta_page.get_bq_num_bits_per_dimension(), )]; let mut neighbor_vector = self.as_mut().neighbor_vector().index_pin(i); for (index_in_q_vec, val) in quantized.iter().enumerate() { diff --git a/timescale_vector/src/access_method/distance.rs b/timescale_vector/src/access_method/distance.rs index ad552f2f..eb13b821 100644 --- a/timescale_vector/src/access_method/distance.rs +++ b/timescale_vector/src/access_method/distance.rs @@ -182,6 +182,25 @@ pub fn distance_xor_optimized(a: &[u64], b: &[u64]) -> usize { 28 => xor_arm!(a, b, 28), 29 => xor_arm!(a, b, 29), 30 => xor_arm!(a, b, 30), + 31 => xor_arm!(a, b, 31), + 32 => xor_arm!(a, b, 32), + 33 => xor_arm!(a, b, 33), + 34 => xor_arm!(a, b, 34), + 35 => xor_arm!(a, b, 35), + 36 => xor_arm!(a, b, 36), + 37 => xor_arm!(a, b, 37), + 38 => xor_arm!(a, b, 38), + 39 => xor_arm!(a, b, 39), + 40 => xor_arm!(a, b, 40), + 41 => xor_arm!(a, b, 41), + 42 => xor_arm!(a, b, 42), + 43 => xor_arm!(a, b, 43), + 44 => xor_arm!(a, b, 44), + 45 => xor_arm!(a, b, 45), + 46 => xor_arm!(a, b, 46), + 47 => xor_arm!(a, b, 47), + 48 => xor_arm!(a, b, 48), + 49 => xor_arm!(a, b, 49), _ => a .iter() .zip(b.iter()) diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index 9aeb0066..e53c1778 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -10,7 +10,10 @@ use crate::util::*; use super::bq::BqNode; use super::distance; -use super::options::{NUM_DIMENSIONS_DEFAULT_SENTINEL, NUM_NEIGHBORS_DEFAULT_SENTINEL}; +use super::options::{ + BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, NUM_DIMENSIONS_DEFAULT_SENTINEL, + NUM_NEIGHBORS_DEFAULT_SENTINEL, +}; use super::stats::StatsNodeModify; use super::storage::StorageType; @@ -66,6 +69,7 @@ impl MetaPageV1 { distance_type: DistanceType::L2 as u16, num_dimensions: self.num_dimensions, num_dimensions_to_index: self.num_dimensions, + bq_num_bits_per_dimension: 1, num_neighbors: self.num_neighbors, storage_type: StorageType::Plain as u8, search_list_size: self.search_list_size, @@ -119,6 +123,7 @@ pub struct MetaPage { num_dimensions: u32, //number of dimensions in the vectors stored in the index num_dimensions_to_index: u32, + bq_num_bits_per_dimension: u8, /// the value of the TSVStorageLayout enum storage_type: u8, /// max number of outgoing edges a node in the graph can have (R in the papers) @@ -140,6 +145,10 @@ impl MetaPage { self.num_dimensions_to_index } + pub fn get_bq_num_bits_per_dimension(&self) -> u8 { + self.bq_num_bits_per_dimension + } + pub fn get_num_dimensions_for_neighbors(&self) -> u32 { match StorageType::from_u8(self.storage_type) { StorageType::Plain => { @@ -198,7 +207,11 @@ impl MetaPage { } } - fn calculate_num_neighbors(num_dimensions: u32, opt: &PgBox) -> u32 { + fn calculate_num_neighbors( + num_dimensions: u32, + num_bits_per_dimension: u8, + opt: &PgBox, + ) -> u32 { let num_neighbors = (*opt).get_num_neighbors(); if num_neighbors == NUM_NEIGHBORS_DEFAULT_SENTINEL { match (*opt).get_storage_type() { @@ -206,6 +219,7 @@ impl MetaPage { StorageType::BqSpeedup => BqNode::get_default_num_neighbors( num_dimensions as usize, num_dimensions as usize, + num_bits_per_dimension, ) as u32, StorageType::BqCompression => 50, } @@ -229,6 +243,30 @@ impl MetaPage { (*opt).num_dimensions }; + let bq_num_bits_per_dimension = + if (*opt).bq_num_bits_per_dimension == BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL { + if (*opt).get_storage_type() == StorageType::BqCompression + && num_dimensions_to_index < 900 + { + 3 + } else { + 1 + } + } else { + (*opt).bq_num_bits_per_dimension as u8 + }; + + if bq_num_bits_per_dimension > 1 && num_dimensions_to_index > 930 { + //limited by BqMeans fitting on a page + pgrx::error!("BQ with more than 1 bit per dimension is not supported for more than 900 dimensions"); + } + if bq_num_bits_per_dimension > 1 && (*opt).get_storage_type() != StorageType::BqCompression + { + pgrx::error!( + "BQ with more than 1 bit per dimension is only supported with the memory_optimized storage layout" + ); + } + let meta = MetaPage { magic_number: TSV_MAGIC_NUMBER, version: TSV_VERSION, @@ -237,7 +275,12 @@ impl MetaPage { num_dimensions, num_dimensions_to_index, storage_type: (*opt).get_storage_type() as u8, - num_neighbors: Self::calculate_num_neighbors(num_dimensions, &opt), + num_neighbors: Self::calculate_num_neighbors( + num_dimensions, + bq_num_bits_per_dimension, + &opt, + ), + bq_num_bits_per_dimension, search_list_size: (*opt).search_list_size, max_alpha: (*opt).max_alpha, init_ids: ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber), diff --git a/timescale_vector/src/access_method/options.rs b/timescale_vector/src/access_method/options.rs index 2bd5a0a2..fa3ccf8f 100644 --- a/timescale_vector/src/access_method/options.rs +++ b/timescale_vector/src/access_method/options.rs @@ -17,10 +17,12 @@ pub struct TSVIndexOptions { pub search_list_size: u32, pub num_dimensions: u32, pub max_alpha: f64, + pub bq_num_bits_per_dimension: u32, } pub const NUM_NEIGHBORS_DEFAULT_SENTINEL: i32 = -1; pub const NUM_DIMENSIONS_DEFAULT_SENTINEL: u32 = 0; +pub const BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL: u32 = 0; const DEFAULT_MAX_ALPHA: f64 = 1.2; impl TSVIndexOptions { @@ -37,6 +39,7 @@ impl TSVIndexOptions { ops.search_list_size = 100; ops.max_alpha = DEFAULT_MAX_ALPHA; ops.num_dimensions = NUM_DIMENSIONS_DEFAULT_SENTINEL; + ops.bq_num_bits_per_dimension = BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL; unsafe { set_varsize( ops.as_ptr().cast(), @@ -83,7 +86,7 @@ impl TSVIndexOptions { } } -const NUM_REL_OPTS: usize = 5; +const NUM_REL_OPTS: usize = 6; static mut RELOPT_KIND_TSV: pg_sys::relopt_kind = 0; // amoptions is a function that gets a datum of text[] data from pg_class.reloptions (which contains text in the format "key=value") and returns a bytea for the struct for the parsed options. @@ -123,6 +126,11 @@ pub unsafe extern "C" fn amoptions( opttype: pg_sys::relopt_type_RELOPT_TYPE_INT, offset: offset_of!(TSVIndexOptions, num_dimensions) as i32, }, + pg_sys::relopt_parse_elt { + optname: "num_bits_per_dimension".as_pg_cstr(), + opttype: pg_sys::relopt_type_RELOPT_TYPE_INT, + offset: offset_of!(TSVIndexOptions, bq_num_bits_per_dimension) as i32, + }, pg_sys::relopt_parse_elt { optname: "max_alpha".as_pg_cstr(), opttype: pg_sys::relopt_type_RELOPT_TYPE_REAL, @@ -218,6 +226,16 @@ pub unsafe fn init() { 5000, pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, ); + + pg_sys::add_int_reloption( + RELOPT_KIND_TSV, + "num_bits_per_dimension".as_pg_cstr(), + "The number of bits to use per dimension for compressed storage".as_pg_cstr(), + BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL as _, + 0, + 32, + pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, + ); } #[cfg(any(test, feature = "pg_test"))] diff --git a/timescale_vector/src/access_method/vacuum.rs b/timescale_vector/src/access_method/vacuum.rs index 2c46524a..50c8e04c 100644 --- a/timescale_vector/src/access_method/vacuum.rs +++ b/timescale_vector/src/access_method/vacuum.rs @@ -146,6 +146,8 @@ pub mod tests { #[cfg(test)] pub fn test_delete_vacuum_plain_scaffold(index_options: &str) { //do not run this test in parallel. (pgrx tests run in a txn rolled back after each test, but we do not have that luxury here). + + use rand::Rng; let _lock = VAC_PLAIN_MUTEX.lock().unwrap(); //we need to run vacuum in this test which cannot be run from SPI. @@ -160,16 +162,11 @@ pub mod tests { ) .unwrap(); - let suffix = (1..=253) - .map(|i| format!("{}", i)) - .collect::>() - .join(", "); - let (mut client, _) = pgrx_tests::client().unwrap(); client .batch_execute(&format!( - "CREATE TABLE test_vac(embedding vector(256)); + "CREATE TABLE test_vac(id INT GENERATED ALWAYS AS IDENTITY, embedding vector(256)); select setseed(0.5); -- generate 300 vectors @@ -178,15 +175,13 @@ pub mod tests { * FROM ( SELECT - ('[ 0 , ' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding + ('[ ' || array_to_string(array_agg(random()), ',', '0') || ']')::vector AS embedding FROM - generate_series(1, 255 * 300) i + generate_series(1, 256 * 303) i GROUP BY - i % 300) g; + i % 303) g; - INSERT INTO test_vac(embedding) VALUES ('[1,2,3,{suffix}]'), ('[4,5,6,{suffix}]'), ('[7,8,10,{suffix}]'); - CREATE INDEX idxtest_vac ON test_vac USING tsv(embedding) @@ -195,16 +190,29 @@ pub mod tests { )) .unwrap(); + let test_vec: Option> = client + .query_one( + &format!( + "SELECT('{{' || array_to_string(array_agg(1.0), ',', '0') || '}}')::real[] AS embedding + FROM generate_series(1, 256)" + ), + &[], + ) + .unwrap() + .get(0); + let test_vec = test_vec + .unwrap() + .into_iter() + .map(|x| Some(x)) + .collect::>(); + client.execute("set enable_seqscan = 0;", &[]).unwrap(); - let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); + let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> $1::float4[]::vector) SELECT count(*) from cte;"), &[&test_vec]).unwrap().get(0); assert_eq!(cnt, 303, "initial count"); client - .execute( - &format!("DELETE FROM test_vac WHERE embedding = '[1,2,3,{suffix}]';"), - &[], - ) + .execute(&format!("DELETE FROM test_vac WHERE id = 301;"), &[]) .unwrap(); client.close().unwrap(); @@ -213,29 +221,31 @@ pub mod tests { client.execute("VACUUM test_vac", &[]).unwrap(); + let mut rng = rand::thread_rng(); + let rand_vec = (1..=256) + .map(|_i| format!("{}", rng.gen::())) + .collect::>() + .join(", "); //inserts into the previous 1,2,3 spot that was deleted client .execute( - &format!("INSERT INTO test_vac(embedding) VALUES ('[10,12,13,{suffix}]');"), + &format!("INSERT INTO test_vac(embedding) VALUES ('[{rand_vec}]');"), &[], ) .unwrap(); client.execute("set enable_seqscan = 0;", &[]).unwrap(); - let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); + let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> $1::float4[]::vector) SELECT count(*) from cte;"), &[&test_vec]).unwrap().get(0); //if the old index is still used the count is 304 assert_eq!(cnt, 303, "count after vacuum"); //do another delete for same items (noop) client - .execute( - &format!("DELETE FROM test_vac WHERE embedding = '[1,2,3,{suffix}]';"), - &[], - ) + .execute(&format!("DELETE FROM test_vac WHERE id=301;"), &[]) .unwrap(); client.execute("set enable_seqscan = 0;", &[]).unwrap(); - let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> '[1,1,1,{suffix}]') SELECT count(*) from cte;"), &[]).unwrap().get(0); + let cnt: i64 = client.query_one(&format!("WITH cte as (select * from test_vac order by embedding <=> $1::float4[]::vector) SELECT count(*) from cte;"), &[&test_vec]).unwrap().get(0); //if the old index is still used the count is 304 assert_eq!(cnt, 303, "count after delete"); diff --git a/timescale_vector/src/util/page.rs b/timescale_vector/src/util/page.rs index 92284931..918eb30c 100644 --- a/timescale_vector/src/util/page.rs +++ b/timescale_vector/src/util/page.rs @@ -250,13 +250,8 @@ impl<'a> ReadablePage<'a> { } pub fn get_type(&self) -> PageType { - unsafe { - let opaque_data = - //safe to do because self.page was already verified during construction - TsvPageOpaqueData::with_page(self.page); - - PageType::from_u8((*opaque_data).page_type) - } + let opaque_data = TsvPageOpaqueData::read_from_page(&self.page); + PageType::from_u8((*opaque_data).page_type) } pub fn get_buffer(&self) -> &LockedBufferShare { From b5f5b3414bc2c6c2aee9c91ab6815846f1de85ba Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Fri, 3 May 2024 09:24:37 -0700 Subject: [PATCH 39/44] change default num_bits to 2 for dim=768 --- timescale_vector/src/access_method/meta_page.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index e53c1778..c1fbfcd1 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -248,7 +248,7 @@ impl MetaPage { if (*opt).get_storage_type() == StorageType::BqCompression && num_dimensions_to_index < 900 { - 3 + 2 } else { 1 } From edb94fe20e2665b99ead7c6d239745f079dc3017 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 7 May 2024 11:54:12 -0700 Subject: [PATCH 40/44] cleanup tests + nits --- README.md | 6 +++ timescale_vector/src/access_method/options.rs | 39 ++++++++++++++++++- .../src/access_method/upgrade_test.rs | 2 +- timescale_vector/src/access_method/vacuum.rs | 10 +++-- timescale_vector/src/util/table_slot.rs | 4 +- 5 files changed, 53 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 33524b31..b1b49e5b 100644 --- a/README.md +++ b/README.md @@ -52,6 +52,12 @@ You can run tests against a postgres version pg16 using ```shell cargo pgrx test ${postgres_version} ``` + +To run all tests run: +```shell +cargo test -- --ignored && cargo pgrx test ${postgres_version} +``` + 🐯 About Timescale TimescaleDB is a distributed time-series database built on PostgreSQL that scales to over 10 million of metrics per second, supports native compression, handles high cardinality, and offers native time-series capabilities, such as data retention policies, continuous aggregate views, downsampling, data gap-filling and interpolation. diff --git a/timescale_vector/src/access_method/options.rs b/timescale_vector/src/access_method/options.rs index fa3ccf8f..a0eee2de 100644 --- a/timescale_vector/src/access_method/options.rs +++ b/timescale_vector/src/access_method/options.rs @@ -243,8 +243,8 @@ pub unsafe fn init() { mod tests { use crate::access_method::{ options::{ - TSVIndexOptions, DEFAULT_MAX_ALPHA, NUM_DIMENSIONS_DEFAULT_SENTINEL, - NUM_NEIGHBORS_DEFAULT_SENTINEL, + TSVIndexOptions, BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, DEFAULT_MAX_ALPHA, + NUM_DIMENSIONS_DEFAULT_SENTINEL, NUM_NEIGHBORS_DEFAULT_SENTINEL, }, storage::StorageType, }; @@ -266,6 +266,10 @@ mod tests { let options = TSVIndexOptions::from_relation(&indexrel); assert_eq!(options.num_neighbors, 30); assert_eq!(options.num_dimensions, NUM_DIMENSIONS_DEFAULT_SENTINEL); + assert_eq!( + options.bq_num_bits_per_dimension, + BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, + ); Ok(()) } @@ -287,6 +291,10 @@ mod tests { assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); assert_eq!(options.num_dimensions, NUM_DIMENSIONS_DEFAULT_SENTINEL); assert_eq!(options.get_storage_type(), StorageType::BqCompression); + assert_eq!( + options.bq_num_bits_per_dimension, + BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, + ); Ok(()) } @@ -352,6 +360,33 @@ mod tests { assert_eq!(options.max_alpha, 1.4); assert_eq!(options.get_storage_type(), StorageType::Plain); assert_eq!(options.num_dimensions, 20); + assert_eq!( + options.bq_num_bits_per_dimension, + BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL + ); + Ok(()) + } + + #[pg_test] + unsafe fn test_index_options_custom_mem_optimized() -> spi::Result<()> { + Spi::run(&format!( + "CREATE TABLE test(encoding vector(3)); + CREATE INDEX idxtest + ON test + USING tsv(encoding) + WITH (storage_layout = memory_optimized, num_neighbors=40, search_list_size=18, num_dimensions=20, max_alpha=1.4, num_bits_per_dimension=5);", + ))?; + + let index_oid = + Spi::get_one::("SELECT 'idxtest'::regclass::oid")?.expect("oid was null"); + let indexrel = PgRelation::from_pg(pg_sys::RelationIdGetRelation(index_oid)); + let options = TSVIndexOptions::from_relation(&indexrel); + assert_eq!(options.get_num_neighbors(), 40); + assert_eq!(options.search_list_size, 18); + assert_eq!(options.max_alpha, 1.4); + assert_eq!(options.get_storage_type(), StorageType::BqCompression); + assert_eq!(options.num_dimensions, 20); + assert_eq!(options.bq_num_bits_per_dimension, 5); Ok(()) } } diff --git a/timescale_vector/src/access_method/upgrade_test.rs b/timescale_vector/src/access_method/upgrade_test.rs index 16ee3c54..0f408b0b 100644 --- a/timescale_vector/src/access_method/upgrade_test.rs +++ b/timescale_vector/src/access_method/upgrade_test.rs @@ -1,4 +1,4 @@ -#[cfg(any(test, feature = "pg_test"))] +#[cfg(test)] #[pgrx::pg_schema] pub mod tests { use pgrx::*; diff --git a/timescale_vector/src/access_method/vacuum.rs b/timescale_vector/src/access_method/vacuum.rs index 50c8e04c..3d16b29b 100644 --- a/timescale_vector/src/access_method/vacuum.rs +++ b/timescale_vector/src/access_method/vacuum.rs @@ -137,11 +137,11 @@ pub extern "C" fn amvacuumcleanup( #[cfg(any(test, feature = "pg_test"))] #[pgrx::pg_schema] pub mod tests { - use once_cell::sync::Lazy; use pgrx::*; - use std::sync::Mutex; - static VAC_PLAIN_MUTEX: Lazy> = Lazy::new(Mutex::default); + #[cfg(test)] + static VAC_PLAIN_MUTEX: once_cell::sync::Lazy> = + once_cell::sync::Lazy::new(std::sync::Mutex::default); #[cfg(test)] pub fn test_delete_vacuum_plain_scaffold(index_options: &str) { @@ -253,7 +253,9 @@ pub mod tests { client.execute("DROP TABLE test_vac", &[]).unwrap(); } - static VAC_FULL_MUTEX: Lazy> = Lazy::new(Mutex::default); + #[cfg(test)] + static VAC_FULL_MUTEX: once_cell::sync::Lazy> = + once_cell::sync::Lazy::new(std::sync::Mutex::default); #[cfg(test)] pub fn test_delete_vacuum_full_scaffold(index_options: &str) { diff --git a/timescale_vector/src/util/table_slot.rs b/timescale_vector/src/util/table_slot.rs index 56774e5a..092f5ac0 100644 --- a/timescale_vector/src/util/table_slot.rs +++ b/timescale_vector/src/util/table_slot.rs @@ -1,3 +1,5 @@ +use std::ptr::addr_of_mut; + use pgrx::pg_sys::{Datum, TupleTableSlot}; use pgrx::{pg_sys, PgBox, PgRelation}; @@ -29,7 +31,7 @@ impl TableSlot { fetch_row_version( heap_rel.as_ptr(), &mut ctid, - &mut pg_sys::SnapshotAnyData, + addr_of_mut!(pg_sys::SnapshotAnyData), slot.as_ptr(), ); stats.record_heap_read(); From 4334c6e0449470d7b99dd78d9e7ca4b424e77bd8 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Mon, 13 May 2024 15:06:54 -0400 Subject: [PATCH 41/44] Adjust debug level down during index builds --- timescale_vector/src/access_method/build.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index 82a7421e..fbc1ae4f 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -306,14 +306,14 @@ fn finalize_index_build( } } - info!("write done"); + debug1!("write done"); assert_eq!(write_stats.num_nodes, state.ntuples); let writing_took = Instant::now() .duration_since(write_stats.started) .as_secs_f64(); if write_stats.num_nodes > 0 { - info!( + debug1!( "Writing took {}s or {}s/tuple. Avg neighbors: {}", writing_took, writing_took / write_stats.num_nodes as f64, @@ -321,7 +321,7 @@ fn finalize_index_build( ); } if write_stats.prune_stats.calls > 0 { - info!( + debug1!( "When pruned for cleanup: avg neighbors before/after {}/{} of {} prunes", write_stats.prune_stats.num_neighbors_before_prune / write_stats.prune_stats.calls, write_stats.prune_stats.num_neighbors_after_prune / write_stats.prune_stats.calls, @@ -416,7 +416,7 @@ fn build_callback_internal( state.ntuples = state.ntuples + 1; if state.ntuples % 1000 == 0 { - info!( + debug1!( "Processed {} tuples in {}s which is {}s/tuple. Dist/tuple: Prune: {} search: {}. Stats: {:?}", state.ntuples, Instant::now().duration_since(state.started).as_secs_f64(), From d6f485e449884de0725ff2d5a952b693f0f29b8c Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Wed, 15 May 2024 17:35:37 -0400 Subject: [PATCH 42/44] cleanup --- README.md | 17 ++++++--- timescale_vector/src/access_method/README.md | 38 -------------------- 2 files changed, 12 insertions(+), 43 deletions(-) delete mode 100644 timescale_vector/src/access_method/README.md diff --git a/README.md b/README.md index b1b49e5b..56995ba1 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,14 @@ Timescale Vector -Say something chat gpt. +A vector index for speeding up ANN search in `pgvector`. 🔧 Tools Setup -Building the extension requires valid rust (we build and test on 1.65), rustfmt, and clang installs, along with the postgres headers for whichever version of postgres you are running, and pgx. We recommend installing rust using the official instructions: + +Building the extension requires valid rust, rustfmt, and clang installs, along with the postgres headers for whichever version of postgres you are running, and pgx. We recommend installing rust using the official instructions: ```shell curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh ``` -and build tools, the postgres headers, in the preferred manner for your system. You may also need to install OpenSSl. For Ubuntu you can follow the postgres install instructions then run +and build tools, the postgres headers, in the preferred manner for your system. You may also need to install OpenSSL. For Ubuntu you can follow the postgres install instructions then run ```shell sudo apt-get install make gcc pkg-config clang postgresql-server-dev-16 libssl-dev @@ -28,10 +29,11 @@ cargo pgrx init --pg16 pg_config Installing from source is also available on macOS and requires the same set of prerequisites and set up commands listed above. 💾 Building and Installing the extension + Download or clone this repository, and switch to the extension subdirectory, e.g. ```shell git clone https://github.com/timescale/timescale-vector && \ -cd timescale-vector/extension +cd timescale-vector/timescale_vector ``` Then run @@ -41,9 +43,13 @@ cargo pgrx install --release To initialize the extension after installation, enter the following into psql: +```sql CREATE EXTENSION timescale_vector; +``` + ✏️ Get Involved -The Timescale Vecotr project is still in the initial planning stage as we decide our priorities and what to implement first. As such, now is a great time to help shape the project's direction! Have a look at the list of features we're thinking of working on and feel free to comment on the features, expand the list, or hop on the Discussions forum for more in-depth discussions. + +The Timescale Vector project is still in it's early stage as we decide our priorities and what to implement. As such, now is a great time to help shape the project's direction! Have a look at the list of features we're thinking of working on and feel free to comment on the features, expand the list, or hop on the Discussions forum for more in-depth discussions. 🔨 Testing See above for prerequisites and installation instructions. @@ -59,6 +65,7 @@ cargo test -- --ignored && cargo pgrx test ${postgres_version} ``` 🐯 About Timescale + TimescaleDB is a distributed time-series database built on PostgreSQL that scales to over 10 million of metrics per second, supports native compression, handles high cardinality, and offers native time-series capabilities, such as data retention policies, continuous aggregate views, downsampling, data gap-filling and interpolation. TimescaleDB also supports full SQL, a variety of data types (numerics, text, arrays, JSON, booleans), and ACID semantics. Operationally mature capabilities include high availability, streaming backups, upgrades over time, roles and permissions, and security. diff --git a/timescale_vector/src/access_method/README.md b/timescale_vector/src/access_method/README.md deleted file mode 100644 index ed0e2cca..00000000 --- a/timescale_vector/src/access_method/README.md +++ /dev/null @@ -1,38 +0,0 @@ -# Graph - -The graph abstraction implements the 2 primary algorithms: -- greedy_search - which finds the index nodes closest to a given query -- prune_neighbors - which reduces the number of neighbors assigned to a particular nodes in order to fit within the num_neighbor limit. - -Graph also implements the insertion algoritm for when a node needs to be added to the graph. Insertion is mostly a combination of the 2 algorithms above, greedy_search and prune_neighbors. - -We support multiple storage layouts (described below). The logic in graph is works on an abstract storage object and is not concerned with the particular storage implementation. -Thus, any logic that needs to differ between storage implementations has to fall within the responsibility of the storage object and not be included in graph. - -## Greedy search - -Refer to the DiskANN paper for an overview. The greedy search algorithm works by traversing the graph to find the closest nodes to a given query. It does this by: -- starting with a set (right now implemented as just one) initial nodes (called init_ids). -- iteratively: -- - - - - -# On Disk Layout - -Meta Page -- basic metadata -- future proof -- page 0 -- start_node tid - -Graph pages --- start node first -- foreach node --- vector for node --- array of tids of neighbors --- array of distances? - -- in "special area" --- bitmap of deletes? - From 5479b325d2610e37adfc710d30b4ec86573e2654 Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Tue, 21 May 2024 19:44:39 +0200 Subject: [PATCH 43/44] Remove docker infra for now --- .github/workflows/docker.yaml | 109 ---------------------------------- Dockerfile | 105 -------------------------------- 2 files changed, 214 deletions(-) delete mode 100644 .github/workflows/docker.yaml delete mode 100644 Dockerfile diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml deleted file mode 100644 index 0bcba5de..00000000 --- a/.github/workflows/docker.yaml +++ /dev/null @@ -1,109 +0,0 @@ -name: docker -on: - pull_request: - push: - branches: - - main - tags: - - "*" - -jobs: - docker: - runs-on: ubuntu-latest - outputs: - branch_name: ${{steps.metadata.outputs.branch_name}} - image_branch_name: ${{steps.metadata.outputs.image_branch_name}} - strategy: - fail-fast: false - matrix: - cloud_image_tag_prefix: - - pg16 - pgversion: - - 16 - tsversion: - - 2.13 - steps: - - name: Checkout Timescale Vector - uses: actions/checkout@v3 - - - name: Login to GitHub Container Registry - uses: docker/login-action@v2 - with: - registry: ghcr.io - username: ${{ github.actor }} - password: ${{ secrets.GITHUB_TOKEN }} - - - name: Configure AWS credentials - uses: aws-actions/configure-aws-credentials@v1-node16 - with: - aws-access-key-id: ${{ secrets.ORG_AWS_ACCESS_KEY_ID }} - aws-secret-access-key: ${{ secrets.ORG_AWS_SECRET_ACCESS_KEY }} - aws-region: us-east-1 - - - name: Login to Amazon ECR - uses: aws-actions/amazon-ecr-login@v1 - - - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v2 - - - name: Gather metadata - id: metadata - run: | - tsmajor=$(echo ${{ matrix.tsversion }} | cut -d. -f1) - tsmajmin=$(echo ${{ matrix.tsversion }} | cut -d. -f1,2) - commit_sha=$(git rev-parse --short "${{ github.event.pull_request.head.sha || github.sha }}") - branch_name=$(echo ${{github.head_ref || github.ref_name}}) - image_branch_name=$(echo ${branch_name} | sed 's#/#-#') - base_cloud_image_tag=$(aws ecr describe-images --repository-name 'timescaledb-cloud' --region us-east-1 --query 'imageDetails[?imageTags[?starts_with(@,`${{ matrix.cloud_image_tag_prefix }}`) && contains(@, `ts${{ matrix.tsversion }}`) && contains(@, `amd64`)]].imageTags' --output text | sort -V | tail -n1) - echo "tsmajor=${tsmajor}" >> ${GITHUB_OUTPUT} - echo "tsmajmin=${tsmajmin}" >> ${GITHUB_OUTPUT} - echo "branch_name=${branch_name}" >> ${GITHUB_OUTPUT} - echo "image_branch_name=${image_branch_name}" >> ${GITHUB_OUTPUT} - echo "commit_sha=${commit_sha}" >> ${GITHUB_OUTPUT} - echo "base_cloud_image_tag=${base_cloud_image_tag}" >> ${GITHUB_OUTPUT} - - - name: Build and push - uses: docker/build-push-action@v3 - env: - DOCKER_PUSH_REQUIRED: ${{ github.event_name != 'pull_request' || github.event.pull_request.head.repo.owner.login == 'timescale' }} # Don't run docker push when this is a PR from a fork - with: - build-args: | - PG_VERSION=${{ matrix.pgversion }} - TIMESCALEDB_VERSION_MAJMIN=${{ steps.metadata.outputs.tsmajmin }} - BASE_IMAGE=142548018081.dkr.ecr.us-east-1.amazonaws.com/timescaledb-cloud:${{ steps.metadata.outputs.base_cloud_image_tag }} - secrets: | - "GITHUB_TOKEN=${{ secrets.GITHUB_TOKEN }}" - context: . - push: ${{ env.DOCKER_PUSH_REQUIRED == 'true' }} - load: ${{ env.DOCKER_PUSH_REQUIRED != 'true' }} - tags: | - ghcr.io/timescale/dev_timescale_vector:${{steps.metadata.outputs.image_branch_name}}-ts${{matrix.tsversion}}-pg${{matrix.pgversion}} - labels: | - org.opencontainers.image.source= - org.opencontainers.image.revision= - org.opencontainers.image.created= - cache-from: type=gha,scope=${{matrix.pgversion}}-${{matrix.tsversion}} - cache-to: type=gha,mode=max,scope=${{matrix.pgversion}}-${{matrix.tsversion}} - - - name: Publish images to ECR - uses: akhilerm/tag-push-action@v2.0.0 - with: - src: ghcr.io/timescale/dev_timescale_vector:${{steps.metadata.outputs.image_branch_name}}-ts${{matrix.tsversion}}-pg${{matrix.pgversion}} - dst: 142548018081.dkr.ecr.us-east-1.amazonaws.com/timescale-vector:${{steps.metadata.outputs.image_branch_name}}-${{steps.metadata.outputs.commit_sha}}-ts${{matrix.tsversion}}-pg${{matrix.pgversion}} - - - # This allows us to set a single job which must pass in GitHub's branch protection rules, - # otherwise we have to keep updating them as we add or remove postgres versions etc. - docker-result: - name: docker result - if: always() - needs: - - docker - runs-on: ubuntu-latest - steps: - - name: Mark the job as a success - if: needs.docker.result == 'success' - run: exit 0 - - name: Mark the job as a failure - if: needs.docker.result != 'success' - run: exit 1 diff --git a/Dockerfile b/Dockerfile deleted file mode 100644 index fe389128..00000000 --- a/Dockerfile +++ /dev/null @@ -1,105 +0,0 @@ -# syntax=docker/dockerfile:1.3-labs -ARG PG_VERSION=16 -ARG TIMESCALEDB_VERSION_MAJMIN=2.13 -ARG PGRX_VERSION=0.11.1 -ARG BASE_IMAGE=timescale/timescaledb-ha:pg${PG_VERSION}-ts${TIMESCALEDB_VERSION_MAJMIN}-all - -FROM timescale/timescaledb-ha:pg${PG_VERSION}-ts${TIMESCALEDB_VERSION_MAJMIN}-all AS ha-build-tools -ARG PG_VERSION -ARG PGRX_VERSION - -ENV DEBIAN_FRONTEND=noninteractive -USER root - -RUN apt-get update -RUN apt-get install -y \ - clang \ - gcc \ - pkg-config \ - wget \ - lsb-release \ - libssl-dev \ - curl \ - gnupg2 \ - binutils \ - devscripts \ - equivs \ - git \ - libkrb5-dev \ - libopenblas-dev \ - libopenblas-base \ - libperl-dev \ - make \ - cmake - -RUN wget -O - https://www.postgresql.org/media/keys/ACCC4CF8.asc | apt-key add - -RUN for t in deb deb-src; do \ - echo "$t [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/postgresql.keyring] http://apt.postgresql.org/pub/repos/apt/ $(lsb_release -s -c)-pgdg main" >> /etc/apt/sources.list.d/pgdg.list; \ - done - -RUN apt-get update && apt-get install -y \ - postgresql-${PG_VERSION} \ - postgresql-server-dev-${PG_VERSION} - -USER postgres -WORKDIR /build - -ENV HOME=/build \ - PATH=/build/.cargo/bin:$PATH \ - CARGO_HOME=/build/.cargo \ - RUSTUP_HOME=/build/.rustup - -RUN chown postgres:postgres /build - -# if you need bleeding edge timescaledb -# RUN cd /build && git clone https://github.com/timescale/timescaledb.git /build/timescaledb \ -# && cd /build/timescaledb && rm -fr build \ -# && git checkout ${TS_VERSION} \ -# && ./bootstrap -DCMAKE_BUILD_TYPE=RelWithDebInfo -DREGRESS_CHECKS=OFF -DTAP_CHECKS=OFF -DGENERATE_DOWNGRADE_SCRIPT=OFF -DWARNINGS_AS_ERRORS=OFF \ -# && cd build && make install \ -# && cd ~ - -RUN curl https://sh.rustup.rs -sSf | bash -s -- -y --profile=minimal -c rustfmt -ENV PATH="${CARGO_HOME}/bin:${PATH}" - -RUN set -ex \ - && mkdir /build/timescale-vector \ - && mkdir /build/timescale-vector/scripts \ - && mkdir /build/timescale-vector/timescale_vector - -## Install pgrx taking into account selected rust toolchain version. -## Making this a separate step to improve layer caching -#COPY --chown=postgres:postgres timescale_vector/rust-toolchain.toml /build/timescale-vector/timescale_vector/rust-toolchain.toml -COPY --chown=postgres:postgres scripts /build/timescale-vector/scripts -USER postgres -WORKDIR /build/timescale-vector/timescale_vector -RUN set -ex \ - && rm -rf "${CARGO_HOME}/registry" "${CARGO_HOME}/git" \ - && chown postgres:postgres -R "${CARGO_HOME}" \ - && cargo install cargo-pgrx --version ${PGRX_VERSION} --config net.git-fetch-with-cli=true - -## Copy and build Vector itself -USER postgres -COPY --chown=postgres:postgres timescale_vector /build/timescale-vector/timescale_vector -COPY --chown=postgres:postgres Makefile /build/timescale-vector/Makefile - -WORKDIR /build/timescale-vector -RUN PG_CONFIG="/usr/lib/postgresql/${PG_VERSION}/bin/pg_config" make package - -## COPY over the new files to the image. Done as a seperate stage so we don't -## ship the build tools. Fixed pg16 image is intentional. The image ships with -## PG 12, 13, 14, 15 and 16 binaries. The PATH environment variable below is used -## to specify the runtime version. -FROM ${BASE_IMAGE} -ARG PG_VERSION - -## Copy old versions and/or bleeding edge timescaledb if any were installed -COPY --from=ha-build-tools --chown=root:postgres /usr/share/postgresql /usr/share/postgresql -COPY --from=ha-build-tools --chown=root:postgres /usr/lib/postgresql /usr/lib/postgresql - -## Copy freshly build current Vector version -COPY --from=ha-build-tools --chown=root:postgres /build/timescale-vector/timescale_vector/target/release/timescale_vector-pg${PG_VERSION}/usr/lib/postgresql /usr/lib/postgresql -COPY --from=ha-build-tools --chown=root:postgres /build/timescale-vector/timescale_vector/target/release/timescale_vector-pg${PG_VERSION}/usr/share/postgresql /usr/share/postgresql -ENV PATH="/usr/lib/postgresql/${PG_VERSION}/bin:${PATH}" - -USER postgres From 54fe7fc1979c5d8ccc75793265221a8f1b57baad Mon Sep 17 00:00:00 2001 From: Matvey Arye Date: Wed, 22 May 2024 15:16:35 +0200 Subject: [PATCH 44/44] Rename BQ->SBQ --- timescale_vector/Cargo.toml | 6 +- timescale_vector/src/access_method/build.rs | 23 +-- .../src/access_method/meta_page.rs | 28 +-- timescale_vector/src/access_method/mod.rs | 2 +- timescale_vector/src/access_method/options.rs | 22 +-- .../src/access_method/plain_storage.rs | 2 +- .../src/access_method/{bq.rs => sbq.rs} | 184 +++++++++--------- timescale_vector/src/access_method/scan.rs | 32 +-- timescale_vector/src/access_method/storage.rs | 12 +- timescale_vector/src/access_method/vacuum.rs | 6 +- timescale_vector/src/util/page.rs | 8 +- 11 files changed, 165 insertions(+), 160 deletions(-) rename timescale_vector/src/access_method/{bq.rs => sbq.rs} (87%) diff --git a/timescale_vector/Cargo.toml b/timescale_vector/Cargo.toml index 4c724aaf..b4b202e5 100644 --- a/timescale_vector/Cargo.toml +++ b/timescale_vector/Cargo.toml @@ -14,7 +14,7 @@ pg_test = [] [dependencies] memoffset = "0.9.0" -pgrx = "=0.11.2" +pgrx = "=0.11.4" rkyv = { version="0.7.42", features=["validation"]} simdeez = {version = "1.0.8"} reductive = { version = "0.9.0"} @@ -30,8 +30,8 @@ timescale_vector_derive = { path = "timescale_vector_derive" } semver = "1.0.22" [dev-dependencies] -pgrx-tests = "=0.11.2" -pgrx-pg-config = "=0.11.2" +pgrx-tests = "=0.11.4" +pgrx-pg-config = "=0.11.4" criterion = "0.5.1" tempfile = "3.3.0" diff --git a/timescale_vector/src/access_method/build.rs b/timescale_vector/src/access_method/build.rs index fbc1ae4f..7893346e 100644 --- a/timescale_vector/src/access_method/build.rs +++ b/timescale_vector/src/access_method/build.rs @@ -15,8 +15,8 @@ use crate::util::*; use self::ports::PROGRESS_CREATE_IDX_SUBPHASE; -use super::bq::BqSpeedupStorage; use super::graph_neighbor_store::BuilderNeighborCache; +use super::sbq::SbqSpeedupStorage; use super::meta_page::MetaPage; @@ -24,7 +24,7 @@ use super::plain_storage::PlainStorage; use super::storage::{Storage, StorageType}; enum StorageBuildState<'a, 'b, 'c, 'd, 'e> { - BqSpeedup(&'a mut BqSpeedupStorage<'b>, &'c mut BuildState<'d, 'e>), + SbqSpeedup(&'a mut SbqSpeedupStorage<'b>, &'c mut BuildState<'d, 'e>), Plain(&'a mut PlainStorage<'b>, &'c mut BuildState<'d, 'e>), } @@ -130,8 +130,8 @@ pub unsafe extern "C" fn aminsert( &mut stats, ); } - StorageType::BqSpeedup | StorageType::BqCompression => { - let bq = BqSpeedupStorage::load_for_insert( + StorageType::SbqSpeedup | StorageType::SbqCompression => { + let bq = SbqSpeedupStorage::load_for_insert( &heap_relation, &index_relation, &meta_page, @@ -214,10 +214,11 @@ fn do_heap_scan<'a>( finalize_index_build(&mut plain, &mut bs, write_stats) } - StorageType::BqSpeedup | StorageType::BqCompression => { - let mut bq = BqSpeedupStorage::new_for_build(index_relation, heap_relation, &meta_page); + StorageType::SbqSpeedup | StorageType::SbqCompression => { + let mut bq = + SbqSpeedupStorage::new_for_build(index_relation, heap_relation, &meta_page); - let page_type = BqSpeedupStorage::page_type(); + let page_type = SbqSpeedupStorage::page_type(); unsafe { pgstat_progress_update_param(PROGRESS_CREATE_IDX_SUBPHASE, BUILD_PHASE_TRAINING); @@ -226,7 +227,7 @@ fn do_heap_scan<'a>( bq.start_training(&meta_page); let mut bs = BuildState::new(index_relation, meta_page, graph, page_type); - let mut state = StorageBuildState::BqSpeedup(&mut bq, &mut bs); + let mut state = StorageBuildState::SbqSpeedup(&mut bq, &mut bs); unsafe { pg_sys::IndexBuildHeapScan( @@ -246,7 +247,7 @@ fn do_heap_scan<'a>( ); } - let mut state = StorageBuildState::BqSpeedup(&mut bq, &mut bs); + let mut state = StorageBuildState::SbqSpeedup(&mut bq, &mut bs); unsafe { pg_sys::IndexBuildHeapScan( @@ -346,7 +347,7 @@ unsafe extern "C" fn build_callback_bq_train( ) { let state = (state as *mut StorageBuildState).as_mut().unwrap(); match state { - StorageBuildState::BqSpeedup(bq, state) => { + StorageBuildState::SbqSpeedup(bq, state) => { let vec = PgVector::from_pg_parts(values, isnull, 0, &state.meta_page, true, false); if let Some(vec) = vec { bq.add_sample(vec.to_index_slice()); @@ -370,7 +371,7 @@ unsafe extern "C" fn build_callback( let index_relation = unsafe { PgRelation::from_pg(index) }; let state = (state as *mut StorageBuildState).as_mut().unwrap(); match state { - StorageBuildState::BqSpeedup(bq, state) => { + StorageBuildState::SbqSpeedup(bq, state) => { let vec = PgVector::from_pg_parts(values, isnull, 0, &state.meta_page, true, false); if let Some(vec) = vec { let heap_pointer = ItemPointer::with_item_pointer_data(*ctid); diff --git a/timescale_vector/src/access_method/meta_page.rs b/timescale_vector/src/access_method/meta_page.rs index c1fbfcd1..1f03ef72 100644 --- a/timescale_vector/src/access_method/meta_page.rs +++ b/timescale_vector/src/access_method/meta_page.rs @@ -8,12 +8,12 @@ use crate::access_method::options::TSVIndexOptions; use crate::util::page; use crate::util::*; -use super::bq::BqNode; use super::distance; use super::options::{ - BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, NUM_DIMENSIONS_DEFAULT_SENTINEL, - NUM_NEIGHBORS_DEFAULT_SENTINEL, + NUM_DIMENSIONS_DEFAULT_SENTINEL, NUM_NEIGHBORS_DEFAULT_SENTINEL, + SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, }; +use super::sbq::SbqNode; use super::stats::StatsNodeModify; use super::storage::StorageType; @@ -154,8 +154,8 @@ impl MetaPage { StorageType::Plain => { error!("get_num_dimensions_for_neighbors should not be called for Plain storage") } - StorageType::BqSpeedup => self.num_dimensions_to_index, - StorageType::BqCompression => 0, + StorageType::SbqSpeedup => self.num_dimensions_to_index, + StorageType::SbqCompression => 0, } } @@ -203,7 +203,7 @@ impl MetaPage { match self.get_storage_type() { StorageType::Plain => None, - StorageType::BqSpeedup | StorageType::BqCompression => Some(self.quantizer_metadata), + StorageType::SbqSpeedup | StorageType::SbqCompression => Some(self.quantizer_metadata), } } @@ -216,12 +216,12 @@ impl MetaPage { if num_neighbors == NUM_NEIGHBORS_DEFAULT_SENTINEL { match (*opt).get_storage_type() { StorageType::Plain => 50, - StorageType::BqSpeedup => BqNode::get_default_num_neighbors( + StorageType::SbqSpeedup => SbqNode::get_default_num_neighbors( num_dimensions as usize, num_dimensions as usize, num_bits_per_dimension, ) as u32, - StorageType::BqCompression => 50, + StorageType::SbqCompression => 50, } } else { num_neighbors as u32 @@ -244,8 +244,8 @@ impl MetaPage { }; let bq_num_bits_per_dimension = - if (*opt).bq_num_bits_per_dimension == BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL { - if (*opt).get_storage_type() == StorageType::BqCompression + if (*opt).bq_num_bits_per_dimension == SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL { + if (*opt).get_storage_type() == StorageType::SbqCompression && num_dimensions_to_index < 900 { 2 @@ -257,13 +257,13 @@ impl MetaPage { }; if bq_num_bits_per_dimension > 1 && num_dimensions_to_index > 930 { - //limited by BqMeans fitting on a page - pgrx::error!("BQ with more than 1 bit per dimension is not supported for more than 900 dimensions"); + //limited by SbqMeans fitting on a page + pgrx::error!("SBQ with more than 1 bit per dimension is not supported for more than 900 dimensions"); } - if bq_num_bits_per_dimension > 1 && (*opt).get_storage_type() != StorageType::BqCompression + if bq_num_bits_per_dimension > 1 && (*opt).get_storage_type() != StorageType::SbqCompression { pgrx::error!( - "BQ with more than 1 bit per dimension is only supported with the memory_optimized storage layout" + "SBQ with more than 1 bit per dimension is only supported with the memory_optimized storage layout" ); } diff --git a/timescale_vector/src/access_method/mod.rs b/timescale_vector/src/access_method/mod.rs index 7ffd4a26..8c23b77c 100644 --- a/timescale_vector/src/access_method/mod.rs +++ b/timescale_vector/src/access_method/mod.rs @@ -20,10 +20,10 @@ mod vacuum; extern crate blas_src; -mod bq; pub mod distance; #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] mod distance_x86; +mod sbq; #[pg_extern(sql = " CREATE OR REPLACE FUNCTION tsv_amhandler(internal) RETURNS index_am_handler PARALLEL SAFE IMMUTABLE STRICT COST 0.0001 LANGUAGE c AS '@MODULE_PATHNAME@', '@FUNCTION_NAME@'; diff --git a/timescale_vector/src/access_method/options.rs b/timescale_vector/src/access_method/options.rs index a0eee2de..892df0d7 100644 --- a/timescale_vector/src/access_method/options.rs +++ b/timescale_vector/src/access_method/options.rs @@ -22,7 +22,7 @@ pub struct TSVIndexOptions { pub const NUM_NEIGHBORS_DEFAULT_SENTINEL: i32 = -1; pub const NUM_DIMENSIONS_DEFAULT_SENTINEL: u32 = 0; -pub const BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL: u32 = 0; +pub const SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL: u32 = 0; const DEFAULT_MAX_ALPHA: f64 = 1.2; impl TSVIndexOptions { @@ -39,7 +39,7 @@ impl TSVIndexOptions { ops.search_list_size = 100; ops.max_alpha = DEFAULT_MAX_ALPHA; ops.num_dimensions = NUM_DIMENSIONS_DEFAULT_SENTINEL; - ops.bq_num_bits_per_dimension = BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL; + ops.bq_num_bits_per_dimension = SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL; unsafe { set_varsize( ops.as_ptr().cast(), @@ -231,7 +231,7 @@ pub unsafe fn init() { RELOPT_KIND_TSV, "num_bits_per_dimension".as_pg_cstr(), "The number of bits to use per dimension for compressed storage".as_pg_cstr(), - BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL as _, + SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL as _, 0, 32, pg_sys::AccessExclusiveLock as pg_sys::LOCKMODE, @@ -243,8 +243,8 @@ pub unsafe fn init() { mod tests { use crate::access_method::{ options::{ - TSVIndexOptions, BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, DEFAULT_MAX_ALPHA, - NUM_DIMENSIONS_DEFAULT_SENTINEL, NUM_NEIGHBORS_DEFAULT_SENTINEL, + TSVIndexOptions, DEFAULT_MAX_ALPHA, NUM_DIMENSIONS_DEFAULT_SENTINEL, + NUM_NEIGHBORS_DEFAULT_SENTINEL, SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, }, storage::StorageType, }; @@ -268,7 +268,7 @@ mod tests { assert_eq!(options.num_dimensions, NUM_DIMENSIONS_DEFAULT_SENTINEL); assert_eq!( options.bq_num_bits_per_dimension, - BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, + SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, ); Ok(()) } @@ -290,10 +290,10 @@ mod tests { assert_eq!(options.search_list_size, 100); assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); assert_eq!(options.num_dimensions, NUM_DIMENSIONS_DEFAULT_SENTINEL); - assert_eq!(options.get_storage_type(), StorageType::BqCompression); + assert_eq!(options.get_storage_type(), StorageType::SbqCompression); assert_eq!( options.bq_num_bits_per_dimension, - BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, + SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL, ); Ok(()) } @@ -316,7 +316,7 @@ mod tests { assert_eq!(options.search_list_size, 100); assert_eq!(options.max_alpha, DEFAULT_MAX_ALPHA); assert_eq!(options.num_dimensions, NUM_DIMENSIONS_DEFAULT_SENTINEL); - assert_eq!(options.get_storage_type(), StorageType::BqSpeedup); + assert_eq!(options.get_storage_type(), StorageType::SbqSpeedup); Ok(()) } @@ -362,7 +362,7 @@ mod tests { assert_eq!(options.num_dimensions, 20); assert_eq!( options.bq_num_bits_per_dimension, - BQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL + SBQ_NUM_BITS_PER_DIMENSION_DEFAULT_SENTINEL ); Ok(()) } @@ -384,7 +384,7 @@ mod tests { assert_eq!(options.get_num_neighbors(), 40); assert_eq!(options.search_list_size, 18); assert_eq!(options.max_alpha, 1.4); - assert_eq!(options.get_storage_type(), StorageType::BqCompression); + assert_eq!(options.get_storage_type(), StorageType::SbqCompression); assert_eq!(options.num_dimensions, 20); assert_eq!(options.bq_num_bits_per_dimension, 5); Ok(()) diff --git a/timescale_vector/src/access_method/plain_storage.rs b/timescale_vector/src/access_method/plain_storage.rs index 1c4c8f73..d5872173 100644 --- a/timescale_vector/src/access_method/plain_storage.rs +++ b/timescale_vector/src/access_method/plain_storage.rs @@ -248,7 +248,7 @@ impl<'a> Storage for PlainStorage<'a> { } /* get_lsn and visit_lsn are different because the distance - comparisons for BQ get the vector from different places */ + comparisons for SBQ get the vector from different places */ fn create_lsn_for_init_id( &self, lsr: &mut ListSearchResult, diff --git a/timescale_vector/src/access_method/bq.rs b/timescale_vector/src/access_method/sbq.rs similarity index 87% rename from timescale_vector/src/access_method/bq.rs rename to timescale_vector/src/access_method/sbq.rs index 520715b8..1fd4cd26 100644 --- a/timescale_vector/src/access_method/bq.rs +++ b/timescale_vector/src/access_method/sbq.rs @@ -26,31 +26,31 @@ use crate::util::{ use super::{meta_page::MetaPage, neighbor_with_distance::NeighborWithDistance}; use crate::util::WritableBuffer; -type BqVectorElement = u64; +type SbqVectorElement = u64; const BITS_STORE_TYPE_SIZE: usize = 64; #[derive(Archive, Deserialize, Serialize, Readable, Writeable)] #[archive(check_bytes)] #[repr(C)] -pub struct BqMeans { +pub struct SbqMeans { count: u64, means: Vec, m2: Vec, } -impl BqMeans { +impl SbqMeans { pub unsafe fn load( index: &PgRelation, meta_page: &super::meta_page::MetaPage, stats: &mut S, - ) -> BqQuantizer { - let mut quantizer = BqQuantizer::new(meta_page); + ) -> SbqQuantizer { + let mut quantizer = SbqQuantizer::new(meta_page); if quantizer.use_mean { if meta_page.get_quantizer_metadata_pointer().is_none() { - pgrx::error!("No BQ pointer found in meta page"); + pgrx::error!("No SBQ pointer found in meta page"); } let quantizer_item_pointer = meta_page.get_quantizer_metadata_pointer().unwrap(); - let bq = BqMeans::read(index, quantizer_item_pointer, stats); + let bq = SbqMeans::read(index, quantizer_item_pointer, stats); let archived = bq.get_archived_node(); quantizer.load( @@ -64,11 +64,11 @@ impl BqMeans { pub unsafe fn store( index: &PgRelation, - quantizer: &BqQuantizer, + quantizer: &SbqQuantizer, stats: &mut S, ) -> ItemPointer { - let mut tape = Tape::new(index, PageType::BqMeans); - let node = BqMeans { + let mut tape = Tape::new(index, PageType::SbqMeans); + let node = SbqMeans { count: quantizer.count, means: quantizer.mean.to_vec(), m2: quantizer.m2.to_vec(), @@ -80,7 +80,7 @@ impl BqMeans { } #[derive(Clone)] -pub struct BqQuantizer { +pub struct SbqQuantizer { pub use_mean: bool, training: bool, pub count: u64, @@ -89,8 +89,8 @@ pub struct BqQuantizer { pub num_bits_per_dimension: u8, } -impl BqQuantizer { - fn new(meta_page: &super::meta_page::MetaPage) -> BqQuantizer { +impl SbqQuantizer { + fn new(meta_page: &super::meta_page::MetaPage) -> SbqQuantizer { Self { use_mean: true, training: false, @@ -123,10 +123,10 @@ impl BqQuantizer { fn quantized_size_bytes(num_dimensions: usize, num_bits_per_dimension: u8) -> usize { Self::quantized_size_internal(num_dimensions, num_bits_per_dimension) - * std::mem::size_of::() + * std::mem::size_of::() } - fn quantize(&self, full_vector: &[f32]) -> Vec { + fn quantize(&self, full_vector: &[f32]) -> Vec { assert!(!self.training); if self.use_mean { let mut res_vector = vec![0; self.quantized_size(full_vector.len())]; @@ -232,25 +232,25 @@ impl BqQuantizer { &self, _meta_page: &super::meta_page::MetaPage, full_vector: &[f32], - ) -> Vec { + ) -> Vec { self.quantize(&full_vector) } } -pub struct BqSearchDistanceMeasure { - quantized_vector: Vec, +pub struct SbqSearchDistanceMeasure { + quantized_vector: Vec, query: PgVector, num_dimensions_for_neighbors: usize, quantized_dimensions: usize, } -impl BqSearchDistanceMeasure { +impl SbqSearchDistanceMeasure { pub fn new( - quantizer: &BqQuantizer, + quantizer: &SbqQuantizer, query: PgVector, num_dimensions_for_neighbors: usize, - ) -> BqSearchDistanceMeasure { - BqSearchDistanceMeasure { + ) -> SbqSearchDistanceMeasure { + SbqSearchDistanceMeasure { quantized_vector: quantizer.quantize(query.to_index_slice()), query, num_dimensions_for_neighbors, @@ -260,7 +260,7 @@ impl BqSearchDistanceMeasure { pub fn calculate_bq_distance( &self, - bq_vector: &[BqVectorElement], + bq_vector: &[SbqVectorElement], gns: &GraphNeighborStore, stats: &mut S, ) -> f32 { @@ -294,14 +294,14 @@ impl BqSearchDistanceMeasure { } } -pub struct BqNodeDistanceMeasure<'a> { - vec: Vec, - storage: &'a BqSpeedupStorage<'a>, +pub struct SbqNodeDistanceMeasure<'a> { + vec: Vec, + storage: &'a SbqSpeedupStorage<'a>, } -impl<'a> BqNodeDistanceMeasure<'a> { +impl<'a> SbqNodeDistanceMeasure<'a> { pub unsafe fn with_index_pointer( - storage: &'a BqSpeedupStorage<'a>, + storage: &'a SbqSpeedupStorage<'a>, index_pointer: IndexPointer, stats: &mut T, ) -> Self { @@ -313,7 +313,7 @@ impl<'a> BqNodeDistanceMeasure<'a> { } } -impl<'a> NodeDistanceMeasure for BqNodeDistanceMeasure<'a> { +impl<'a> NodeDistanceMeasure for SbqNodeDistanceMeasure<'a> { unsafe fn get_distance( &self, index_pointer: IndexPointer, @@ -326,7 +326,7 @@ impl<'a> NodeDistanceMeasure for BqNodeDistanceMeasure<'a> { } struct QuantizedVectorCache { - quantized_vector_map: HashMap>, + quantized_vector_map: HashMap>, } /* should be a LRU cache for quantized vector. For now cheat and never evict @@ -342,9 +342,9 @@ impl QuantizedVectorCache { fn get( &mut self, index_pointer: IndexPointer, - storage: &BqSpeedupStorage, + storage: &SbqSpeedupStorage, stats: &mut S, - ) -> &[BqVectorElement] { + ) -> &[SbqVectorElement] { self.quantized_vector_map .entry(index_pointer) .or_insert_with(|| { @@ -352,7 +352,7 @@ impl QuantizedVectorCache { }) } - fn must_get(&self, index_pointer: IndexPointer) -> &[BqVectorElement] { + fn must_get(&self, index_pointer: IndexPointer) -> &[SbqVectorElement] { self.quantized_vector_map.get(&index_pointer).unwrap() } @@ -362,7 +362,7 @@ impl QuantizedVectorCache { fn preload, S: StatsNodeRead>( &mut self, index_pointers: I, - storage: &BqSpeedupStorage, + storage: &SbqSpeedupStorage, stats: &mut S, ) { for index_pointer in index_pointers { @@ -371,26 +371,26 @@ impl QuantizedVectorCache { } } -pub struct BqSpeedupStorage<'a> { +pub struct SbqSpeedupStorage<'a> { pub index: &'a PgRelation, pub distance_fn: fn(&[f32], &[f32]) -> f32, - quantizer: BqQuantizer, + quantizer: SbqQuantizer, heap_rel: &'a PgRelation, heap_attr: pgrx::pg_sys::AttrNumber, qv_cache: RefCell, num_dimensions_for_neighbors: usize, } -impl<'a> BqSpeedupStorage<'a> { +impl<'a> SbqSpeedupStorage<'a> { pub fn new_for_build( index: &'a PgRelation, heap_rel: &'a PgRelation, meta_page: &super::meta_page::MetaPage, - ) -> BqSpeedupStorage<'a> { + ) -> SbqSpeedupStorage<'a> { Self { index: index, distance_fn: meta_page.get_distance_function(), - quantizer: BqQuantizer::new(meta_page), + quantizer: SbqQuantizer::new(meta_page), heap_rel: heap_rel, heap_attr: get_attribute_number_from_index(index), qv_cache: RefCell::new(QuantizedVectorCache::new(1000)), @@ -402,8 +402,8 @@ impl<'a> BqSpeedupStorage<'a> { index_relation: &PgRelation, meta_page: &super::meta_page::MetaPage, stats: &mut S, - ) -> BqQuantizer { - unsafe { BqMeans::load(&index_relation, meta_page, stats) } + ) -> SbqQuantizer { + unsafe { SbqMeans::load(&index_relation, meta_page, stats) } } pub fn load_for_insert( @@ -411,7 +411,7 @@ impl<'a> BqSpeedupStorage<'a> { index_relation: &'a PgRelation, meta_page: &super::meta_page::MetaPage, stats: &mut S, - ) -> BqSpeedupStorage<'a> { + ) -> SbqSpeedupStorage<'a> { Self { index: index_relation, distance_fn: meta_page.get_distance_function(), @@ -426,9 +426,9 @@ impl<'a> BqSpeedupStorage<'a> { pub fn load_for_search( index_relation: &'a PgRelation, heap_relation: &'a PgRelation, - quantizer: &BqQuantizer, + quantizer: &SbqQuantizer, meta_page: &super::meta_page::MetaPage, - ) -> BqSpeedupStorage<'a> { + ) -> SbqSpeedupStorage<'a> { Self { index: index_relation, distance_fn: meta_page.get_distance_function(), @@ -445,15 +445,15 @@ impl<'a> BqSpeedupStorage<'a> { &self, index_pointer: IndexPointer, stats: &mut S, - ) -> Vec { - let rn = unsafe { BqNode::read(self.index, index_pointer, stats) }; + ) -> Vec { + let rn = unsafe { SbqNode::read(self.index, index_pointer, stats) }; let node = rn.get_archived_node(); node.bq_vector.as_slice().to_vec() } fn write_quantizer_metadata(&self, stats: &mut S) { if self.quantizer.use_mean { - let index_pointer = unsafe { BqMeans::store(&self.index, &self.quantizer, stats) }; + let index_pointer = unsafe { SbqMeans::store(&self.index, &self.quantizer, stats) }; super::meta_page::MetaPage::update_quantizer_metadata_pointer( &self.index, index_pointer, @@ -465,8 +465,8 @@ impl<'a> BqSpeedupStorage<'a> { fn visit_lsn_internal( &self, lsr: &mut ListSearchResult< - as Storage>::QueryDistanceMeasure, - as Storage>::LSNPrivateData, + as Storage>::QueryDistanceMeasure, + as Storage>::LSNPrivateData, >, lsn_index_pointer: IndexPointer, gns: &GraphNeighborStore, @@ -474,7 +474,7 @@ impl<'a> BqSpeedupStorage<'a> { match gns { GraphNeighborStore::Disk => { let rn_visiting = - unsafe { BqNode::read(self.index, lsn_index_pointer, &mut lsr.stats) }; + unsafe { SbqNode::read(self.index, lsn_index_pointer, &mut lsr.stats) }; let node_visiting = rn_visiting.get_archived_node(); //OPT: get neighbors from private data just like plain storage in the self.num_dimensions_for_neighbors == 0 case let neighbors = node_visiting.get_index_pointer_to_neighbors(); @@ -493,7 +493,7 @@ impl<'a> BqSpeedupStorage<'a> { ) } else { let rn_neighbor = unsafe { - BqNode::read(self.index, neighbor_index_pointer, &mut lsr.stats) + SbqNode::read(self.index, neighbor_index_pointer, &mut lsr.stats) }; let node_neighbor = rn_neighbor.get_archived_node(); let bq_vector = node_neighbor.bq_vector.as_slice(); @@ -548,16 +548,16 @@ impl<'a> BqSpeedupStorage<'a> { } } -pub type BqSpeedupStorageLsnPrivateData = PhantomData; //no data stored +pub type SbqSpeedupStorageLsnPrivateData = PhantomData; //no data stored -impl<'a> Storage for BqSpeedupStorage<'a> { - type QueryDistanceMeasure = BqSearchDistanceMeasure; - type NodeDistanceMeasure<'b> = BqNodeDistanceMeasure<'b> where Self: 'b; - type ArchivedType = ArchivedBqNode; - type LSNPrivateData = BqSpeedupStorageLsnPrivateData; //no data stored +impl<'a> Storage for SbqSpeedupStorage<'a> { + type QueryDistanceMeasure = SbqSearchDistanceMeasure; + type NodeDistanceMeasure<'b> = SbqNodeDistanceMeasure<'b> where Self: 'b; + type ArchivedType = ArchivedSbqNode; + type LSNPrivateData = SbqSpeedupStorageLsnPrivateData; //no data stored fn page_type() -> PageType { - PageType::BqNode + PageType::SbqNode } fn create_node( @@ -570,7 +570,7 @@ impl<'a> Storage for BqSpeedupStorage<'a> { ) -> ItemPointer { let bq_vector = self.quantizer.vector_for_new_node(meta_page, full_vector); - let node = BqNode::with_meta( + let node = SbqNode::with_meta( &self.quantizer, heap_pointer, &meta_page, @@ -603,14 +603,14 @@ impl<'a> Storage for BqSpeedupStorage<'a> { ) { let mut cache = self.qv_cache.borrow_mut(); /* It's important to preload cache with all the items since you can run into deadlocks - if you try to fetch a quantized vector while holding the BqNode::modify lock */ + if you try to fetch a quantized vector while holding the SbqNode::modify lock */ let iter = neighbors .iter() .map(|n| n.get_index_pointer_to_neighbor()) .chain(once(index_pointer)); cache.preload(iter, self, stats); - let node = unsafe { BqNode::modify(self.index, index_pointer, stats) }; + let node = unsafe { SbqNode::modify(self.index, index_pointer, stats) }; let mut archived = node.get_archived_node(); archived.as_mut().set_neighbors(neighbors, &meta, &cache); @@ -621,12 +621,12 @@ impl<'a> Storage for BqSpeedupStorage<'a> { &'b self, index_pointer: IndexPointer, stats: &mut S, - ) -> BqNodeDistanceMeasure<'b> { - BqNodeDistanceMeasure::with_index_pointer(self, index_pointer, stats) + ) -> SbqNodeDistanceMeasure<'b> { + SbqNodeDistanceMeasure::with_index_pointer(self, index_pointer, stats) } - fn get_query_distance_measure(&self, query: PgVector) -> BqSearchDistanceMeasure { - return BqSearchDistanceMeasure::new( + fn get_query_distance_measure(&self, query: PgVector) -> SbqSearchDistanceMeasure { + return SbqSearchDistanceMeasure::new( &self.quantizer, query, self.num_dimensions_for_neighbors, @@ -654,13 +654,13 @@ impl<'a> Storage for BqSpeedupStorage<'a> { result: &mut Vec, stats: &mut S, ) { - let rn = unsafe { BqNode::read(self.index, neighbors_of, stats) }; + let rn = unsafe { SbqNode::read(self.index, neighbors_of, stats) }; let archived = rn.get_archived_node(); let q = archived.bq_vector.as_slice(); for n in rn.get_archived_node().iter_neighbors() { //OPT: we can optimize this if num_dimensions_for_neighbors == num_dimensions_to_index - let rn1 = unsafe { BqNode::read(self.index, n, stats) }; + let rn1 = unsafe { SbqNode::read(self.index, n, stats) }; stats.record_quantized_distance_comparison(); let dist = distance_xor_optimized(q, rn1.get_archived_node().bq_vector.as_slice()); result.push(NeighborWithDistance::new(n, dist as f32)) @@ -668,7 +668,7 @@ impl<'a> Storage for BqSpeedupStorage<'a> { } /* get_lsn and visit_lsn are different because the distance - comparisons for BQ get the vector from different places */ + comparisons for SBQ get the vector from different places */ fn create_lsn_for_init_id( &self, lsr: &mut ListSearchResult, @@ -679,7 +679,7 @@ impl<'a> Storage for BqSpeedupStorage<'a> { panic!("should not have had an init id already inserted"); } - let rn = unsafe { BqNode::read(self.index, index_pointer, &mut lsr.stats) }; + let rn = unsafe { SbqNode::read(self.index, index_pointer, &mut lsr.stats) }; let node = rn.get_archived_node(); let distance = lsr.sdm.as_ref().unwrap().calculate_bq_distance( @@ -707,7 +707,7 @@ impl<'a> Storage for BqSpeedupStorage<'a> { stats: &mut GreedySearchStats, ) -> HeapPointer { let lsn_index_pointer = lsn.index_pointer; - let rn = unsafe { BqNode::read(self.index, lsn_index_pointer, stats) }; + let rn = unsafe { SbqNode::read(self.index, lsn_index_pointer, stats) }; let node = rn.get_archived_node(); let heap_pointer = node.heap_item_pointer.deserialize_item_pointer(); heap_pointer @@ -723,14 +723,14 @@ impl<'a> Storage for BqSpeedupStorage<'a> { let mut cache = QuantizedVectorCache::new(neighbors.len() + 1); /* It's important to preload cache with all the items since you can run into deadlocks - if you try to fetch a quantized vector while holding the BqNode::modify lock */ + if you try to fetch a quantized vector while holding the SbqNode::modify lock */ let iter = neighbors .iter() .map(|n| n.get_index_pointer_to_neighbor()) .chain(once(index_pointer)); cache.preload(iter, self, stats); - let node = unsafe { BqNode::modify(self.index, index_pointer, stats) }; + let node = unsafe { SbqNode::modify(self.index, index_pointer, stats) }; let mut archived = node.get_archived_node(); archived.as_mut().set_neighbors(neighbors, &meta, &cache); node.commit(); @@ -745,19 +745,19 @@ use timescale_vector_derive::{Readable, Writeable}; #[derive(Archive, Deserialize, Serialize, Readable, Writeable)] #[archive(check_bytes)] -pub struct BqNode { +pub struct SbqNode { pub heap_item_pointer: HeapPointer, - pub bq_vector: Vec, //don't use BqVectorElement because we don't want to change the size in on-disk format by accident + pub bq_vector: Vec, //don't use SbqVectorElement because we don't want to change the size in on-disk format by accident neighbor_index_pointers: Vec, - neighbor_vectors: Vec>, //don't use BqVectorElement because we don't want to change the size in on-disk format by accident + neighbor_vectors: Vec>, //don't use SbqVectorElement because we don't want to change the size in on-disk format by accident } -impl BqNode { +impl SbqNode { pub fn with_meta( - quantizer: &BqQuantizer, + quantizer: &SbqQuantizer, heap_pointer: HeapPointer, meta_page: &MetaPage, - bq_vector: &[BqVectorElement], + bq_vector: &[SbqVectorElement], ) -> Self { Self::new( heap_pointer, @@ -775,7 +775,7 @@ impl BqNode { _num_dimensions: usize, num_dimensions_for_neighbors: usize, num_bits_per_dimension: u8, - bq_vector: &[BqVectorElement], + bq_vector: &[SbqVectorElement], ) -> Self { // always use vectors of num_neighbors in length because we never want the serialized size of a Node to change let neighbor_index_pointers: Vec<_> = (0..num_neighbors) @@ -787,7 +787,7 @@ impl BqNode { .map(|_| { vec![ 0; - BqQuantizer::quantized_size_internal( + SbqQuantizer::quantized_size_internal( num_dimensions_for_neighbors as _, num_bits_per_dimension ) @@ -812,8 +812,8 @@ impl BqNode { num_dimensions_for_neighbors: usize, num_bits_per_dimension: u8, ) -> usize { - let v: Vec = - vec![0; BqQuantizer::quantized_size_internal(num_dimensions, num_bits_per_dimension)]; + let v: Vec = + vec![0; SbqQuantizer::quantized_size_internal(num_dimensions, num_bits_per_dimension)]; let hp = HeapPointer::new(InvalidBlockNumber, InvalidOffsetNumber); let n = Self::new( hp, @@ -833,21 +833,21 @@ impl BqNode { ) -> usize { //how many neighbors can fit on one page? That's what we choose. - //we first overapproximate the number of neighbors and then double check by actually calculating the size of the BqNode. + //we first overapproximate the number of neighbors and then double check by actually calculating the size of the SbqNode. //blocksize - 100 bytes for the padding/header/etc. let page_size = BLCKSZ as usize - 50; //one quantized_vector takes this many bytes let vec_size = - BqQuantizer::quantized_size_bytes(num_dimensions as usize, num_bits_per_dimension) + 1; - //start from the page size then subtract the heap_item_pointer and bq_vector elements of BqNode. + SbqQuantizer::quantized_size_bytes(num_dimensions as usize, num_bits_per_dimension) + 1; + //start from the page size then subtract the heap_item_pointer and bq_vector elements of SbqNode. let starting = BLCKSZ as usize - std::mem::size_of::() - vec_size; - //one neigbors contribution to neighbor_index_pointers + neighbor_vectors in BqNode. + //one neigbors contribution to neighbor_index_pointers + neighbor_vectors in SbqNode. let one_neighbor = vec_size + std::mem::size_of::(); let mut num_neighbors_overapproximate: usize = starting / one_neighbor; while num_neighbors_overapproximate > 0 { - let serialized_size = BqNode::test_size( + let serialized_size = SbqNode::test_size( num_neighbors_overapproximate as usize, num_dimensions as usize, num_dimensions_for_neighbors as usize, @@ -864,7 +864,7 @@ impl BqNode { } } -impl ArchivedBqNode { +impl ArchivedSbqNode { fn neighbor_index_pointer(self: Pin<&mut Self>) -> Pin<&mut ArchivedVec> { unsafe { self.map_unchecked_mut(|s| &mut s.neighbor_index_pointers) } } @@ -887,7 +887,7 @@ impl ArchivedBqNode { a_index_pointer.offset = ip.offset; if meta_page.get_num_dimensions_for_neighbors() > 0 { - let quantized = &cache.must_get(ip)[..BqQuantizer::quantized_size_internal( + let quantized = &cache.must_get(ip)[..SbqQuantizer::quantized_size_internal( meta_page.get_num_dimensions_for_neighbors() as _, meta_page.get_bq_num_bits_per_dimension(), )]; @@ -922,9 +922,9 @@ impl ArchivedBqNode { } } -impl ArchivedData for ArchivedBqNode { - fn with_data(data: &mut [u8]) -> Pin<&mut ArchivedBqNode> { - ArchivedBqNode::with_data(data) +impl ArchivedData for ArchivedSbqNode { + fn with_data(data: &mut [u8]) -> Pin<&mut ArchivedSbqNode> { + ArchivedSbqNode::with_data(data) } fn get_index_pointer_to_neighbors(&self) -> Vec { diff --git a/timescale_vector/src/access_method/scan.rs b/timescale_vector/src/access_method/scan.rs index 1bbde192..9f316b58 100644 --- a/timescale_vector/src/access_method/scan.rs +++ b/timescale_vector/src/access_method/scan.rs @@ -4,16 +4,16 @@ use pgrx::{pg_sys::InvalidOffsetNumber, *}; use crate::{ access_method::{ - bq::BqSpeedupStorage, graph_neighbor_store::GraphNeighborStore, meta_page::MetaPage, - pg_vector::PgVector, + graph_neighbor_store::GraphNeighborStore, meta_page::MetaPage, pg_vector::PgVector, + sbq::SbqSpeedupStorage, }, util::{buffer::PinnedBufferShare, HeapPointer, IndexPointer}, }; use super::{ - bq::{BqMeans, BqQuantizer, BqSearchDistanceMeasure, BqSpeedupStorageLsnPrivateData}, graph::{Graph, ListSearchResult}, plain_storage::{PlainDistanceMeasure, PlainStorage, PlainStorageLsnPrivateData}, + sbq::{SbqMeans, SbqQuantizer, SbqSearchDistanceMeasure, SbqSpeedupStorageLsnPrivateData}, stats::QuantizerStats, storage::{Storage, StorageType}, }; @@ -21,9 +21,9 @@ use super::{ /* Be very careful not to transfer PgRelations in the state, as they can change between calls. That means we shouldn't be using lifetimes here. Everything should be owned */ enum StorageState { - BqSpeedup( - BqQuantizer, - TSVResponseIterator, + SbqSpeedup( + SbqQuantizer, + TSVResponseIterator, ), Plain(TSVResponseIterator), } @@ -66,13 +66,13 @@ impl TSVScanState { TSVResponseIterator::new(&bq, index, query, search_list_size, meta_page, stats); StorageState::Plain(it) } - StorageType::BqSpeedup | StorageType::BqCompression => { + StorageType::SbqSpeedup | StorageType::SbqCompression => { let mut stats = QuantizerStats::new(); - let quantizer = unsafe { BqMeans::load(index, &meta_page, &mut stats) }; - let bq = BqSpeedupStorage::load_for_search(index, heap, &quantizer, &meta_page); + let quantizer = unsafe { SbqMeans::load(index, &meta_page, &mut stats) }; + let bq = SbqSpeedupStorage::load_for_search(index, heap, &quantizer, &meta_page); let it = TSVResponseIterator::new(&bq, index, query, search_list_size, meta_page, stats); - StorageState::BqSpeedup(quantizer, it) + StorageState::SbqSpeedup(quantizer, it) } }; @@ -366,9 +366,13 @@ pub extern "C" fn amgettuple( let mut storage = unsafe { state.storage.as_mut() }.expect("no storage in state"); match &mut storage { - StorageState::BqSpeedup(quantizer, iter) => { - let bq = - BqSpeedupStorage::load_for_search(&indexrel, &heaprel, quantizer, &state.meta_page); + StorageState::SbqSpeedup(quantizer, iter) => { + let bq = SbqSpeedupStorage::load_for_search( + &indexrel, + &heaprel, + quantizer, + &state.meta_page, + ); let next = iter.next_with_resort(&indexrel, &bq); get_tuple(state, next, scan) } @@ -433,7 +437,7 @@ pub extern "C" fn amendscan(scan: pg_sys::IndexScanDesc) { let mut storage = unsafe { state.storage.as_mut() }.expect("no storage in state"); match &mut storage { - StorageState::BqSpeedup(_bq, iter) => end_scan::(iter), + StorageState::SbqSpeedup(_bq, iter) => end_scan::(iter), StorageState::Plain(iter) => end_scan::(iter), } } diff --git a/timescale_vector/src/access_method/storage.rs b/timescale_vector/src/access_method/storage.rs index 30287e31..4cb91b95 100644 --- a/timescale_vector/src/access_method/storage.rs +++ b/timescale_vector/src/access_method/storage.rs @@ -127,8 +127,8 @@ pub trait Storage { #[derive(PartialEq, Debug)] pub enum StorageType { Plain = 0, - BqSpeedup = 1, - BqCompression = 2, + SbqSpeedup = 1, + SbqCompression = 2, } pub const DEFAULT_STORAGE_TYPE_STR: &str = "memory_optimized"; @@ -137,8 +137,8 @@ impl StorageType { pub fn from_u8(value: u8) -> Self { match value { 0 => StorageType::Plain, - 1 => StorageType::BqSpeedup, - 2 => StorageType::BqCompression, + 1 => StorageType::SbqSpeedup, + 2 => StorageType::SbqCompression, _ => panic!("Invalid storage type"), } } @@ -146,8 +146,8 @@ impl StorageType { pub fn from_str(value: &str) -> Self { match value.to_lowercase().as_str() { "plain" => StorageType::Plain, - "bq_speedup" | "io_optimized" => StorageType::BqSpeedup, - "bq_compression" | "memory_optimized" => StorageType::BqCompression, + "bq_speedup" | "io_optimized" => StorageType::SbqSpeedup, + "bq_compression" | "memory_optimized" => StorageType::SbqCompression, _ => panic!( "Invalid storage type. Must be one of 'plain', 'bq_speedup', 'bq_compression'" ), diff --git a/timescale_vector/src/access_method/vacuum.rs b/timescale_vector/src/access_method/vacuum.rs index 3d16b29b..11945b5e 100644 --- a/timescale_vector/src/access_method/vacuum.rs +++ b/timescale_vector/src/access_method/vacuum.rs @@ -4,7 +4,7 @@ use pgrx::{ }; use crate::{ - access_method::{bq::BqSpeedupStorage, meta_page::MetaPage, plain_storage::PlainStorage}, + access_method::{meta_page::MetaPage, plain_storage::PlainStorage, sbq::SbqSpeedupStorage}, util::{ page::WritablePage, ports::{PageGetItem, PageGetItemId, PageGetMaxOffsetNumber}, @@ -40,8 +40,8 @@ pub extern "C" fn ambulkdelete( let meta_page = MetaPage::fetch(&index_relation); let storage = meta_page.get_storage_type(); match storage { - StorageType::BqSpeedup | StorageType::BqCompression => { - bulk_delete_for_storage::( + StorageType::SbqSpeedup | StorageType::SbqCompression => { + bulk_delete_for_storage::( &index_relation, nblocks, results, diff --git a/timescale_vector/src/util/page.rs b/timescale_vector/src/util/page.rs index 918eb30c..82339852 100644 --- a/timescale_vector/src/util/page.rs +++ b/timescale_vector/src/util/page.rs @@ -30,8 +30,8 @@ pub enum PageType { Node = 1, PqQuantizerDef = 2, PqQuantizerVector = 3, - BqMeans = 4, - BqNode = 5, + SbqMeans = 4, + SbqNode = 5, Meta = 6, } @@ -42,8 +42,8 @@ impl PageType { 1 => PageType::Node, 2 => PageType::PqQuantizerDef, 3 => PageType::PqQuantizerVector, - 4 => PageType::BqMeans, - 5 => PageType::BqNode, + 4 => PageType::SbqMeans, + 5 => PageType::SbqNode, 6 => PageType::Meta, _ => panic!("Unknown PageType number {}", value), }