Skip to content

Commit

Permalink
Avoid creating a new page for every insert
Browse files Browse the repository at this point in the history
  • Loading branch information
syvb committed Nov 4, 2024
1 parent 8596d56 commit ae9c87c
Show file tree
Hide file tree
Showing 4 changed files with 35 additions and 4 deletions.
5 changes: 3 additions & 2 deletions pgvectorscale/src/access_method/build.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use std::time::Instant;

use pgrx::pg_sys::{pgstat_progress_update_param, AsPgCStr};
use pgrx::pg_sys::{pgstat_progress_update_param, AsPgCStr, BlockNumber};
use pgrx::*;

use crate::access_method::graph::Graph;
Expand Down Expand Up @@ -183,14 +183,15 @@ unsafe fn insert_storage<S: Storage>(
meta_page: &mut MetaPage,
stats: &mut InsertStats,
) {
let mut tape = Tape::new(&index_relation, S::page_type());
let mut tape = Tape::from_meta_page(index_relation, S::page_type(), meta_page);
let index_pointer = storage.create_node(
vector.to_index_slice(),
heap_pointer,
&meta_page,
&mut tape,
stats,
);
tape.save_block_number(meta_page);

let mut graph = Graph::new(GraphNeighborStore::Disk, meta_page);
graph.insert(&index_relation, index_pointer, vector, storage, stats)
Expand Down
12 changes: 11 additions & 1 deletion pgvectorscale/src/access_method/meta_page.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use pgrx::pg_sys::{BufferGetBlockNumber, InvalidBlockNumber, InvalidOffsetNumber};
use pgrx::pg_sys::{BlockNumber, BufferGetBlockNumber, InvalidBlockNumber, InvalidOffsetNumber};
use pgrx::*;
use pgvectorscale_derive::{Readable, Writeable};
use rkyv::{Archive, Deserialize, Serialize};
Expand Down Expand Up @@ -76,6 +76,7 @@ impl MetaPageV1 {
max_alpha: self.max_alpha,
init_ids: ItemPointer::new(self.init_ids_block_number, self.init_ids_offset),
quantizer_metadata: ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber),
tape_block_number: BlockNumber::MAX,
}
}
}
Expand Down Expand Up @@ -132,6 +133,7 @@ pub struct MetaPage {
max_alpha: f64,
init_ids: ItemPointer,
quantizer_metadata: ItemPointer,
tape_block_number: BlockNumber,
}

impl MetaPage {
Expand Down Expand Up @@ -207,6 +209,13 @@ impl MetaPage {
}
}

pub fn get_tape_block_number(&self) -> BlockNumber {
self.tape_block_number
}
pub fn set_tape_block_number(&mut self, block_number: BlockNumber) {
self.tape_block_number = block_number;
}

fn calculate_num_neighbors(
num_dimensions: u32,
num_bits_per_dimension: u8,
Expand Down Expand Up @@ -285,6 +294,7 @@ impl MetaPage {
max_alpha: (*opt).max_alpha,
init_ids: ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber),
quantizer_metadata: ItemPointer::new(InvalidBlockNumber, InvalidOffsetNumber),
tape_block_number: BlockNumber::MAX,
};
let page = page::WritablePage::new(index, crate::util::page::PageType::Meta);
meta.write_to_page(page);
Expand Down
2 changes: 1 addition & 1 deletion pgvectorscale/src/access_method/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ mod debugging;
mod graph;
mod graph_neighbor_store;
pub mod guc;
mod meta_page;
pub mod meta_page;
mod neighbor_with_distance;
pub mod options;
pub mod pg_vector;
Expand Down
20 changes: 20 additions & 0 deletions pgvectorscale/src/util/tape.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
//! Tape provides a simple infinite-tape-writing abstraction over postgres pages.

use super::page::{PageType, WritablePage};
use crate::access_method::meta_page::MetaPage;
use pgrx::{
pg_sys::{BlockNumber, BLCKSZ},
*,
Expand All @@ -24,6 +25,25 @@ impl<'a> Tape<'a> {
}
}

pub unsafe fn from_meta_page(
index: &'a PgRelation,
page_type: PageType,
meta_page: &MetaPage,
) -> Self {
match meta_page.get_tape_block_number() {
BlockNumber::MAX => Self::new(index, page_type),
current => Self {
index,
page_type,
current,
},
}
}

pub unsafe fn save_block_number(&self, meta_page: &mut MetaPage) {
meta_page.set_tape_block_number(self.current);
}

pub unsafe fn write(&mut self, data: &[u8]) -> super::ItemPointer {
let size = data.len();
assert!(size < BLCKSZ as usize);
Expand Down

0 comments on commit ae9c87c

Please sign in to comment.