Skip to content

Commit

Permalink
fix: page read counts exceed maximum locks (#23)
Browse files Browse the repository at this point in the history
close #22 

## bench results

trec-covid dataset
prev:  [pg_bm25.rs] Query: 15.8767s (31.49/s)
after: [pg_bm25.rs] Query: 15.7073s (31.83/s)

---------

Signed-off-by: Mingzhuo Yin <[email protected]>
  • Loading branch information
silver-ymz authored Dec 30, 2024
1 parent 98bf1f6 commit 52cbdfb
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 17 deletions.
21 changes: 15 additions & 6 deletions src/page/virtual.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use super::{
bm25_page_size, page_alloc_init_forknum, page_alloc_with_fsm, page_read, page_write, PageFlags,
PageReadGuard, PageWriteGuard,
PageWriteGuard,
};

const DIRECT_COUNT: usize = bm25_page_size() / 4;
Expand All @@ -9,14 +9,24 @@ const INDIRECT2_COUNT: usize = INDIRECT1_COUNT * DIRECT_COUNT;

pub struct VirtualPageReader {
relation: pgrx::pg_sys::Relation,
direct_inode: PageReadGuard,
direct_inode: Box<[u32]>,
indirect1_inode_blkno: u32,
}

impl VirtualPageReader {
pub fn new(relation: pgrx::pg_sys::Relation, blkno: u32) -> Self {
assert!(blkno != pgrx::pg_sys::InvalidBlockNumber);
let direct_inode_page = page_read(relation, blkno);
let data = direct_inode_page.data();
let mut direct_inode: Vec<u32> = Vec::with_capacity(data.len() / 4);
direct_inode.extend_from_slice(bytemuck::cast_slice(data));
let direct_inode = direct_inode.into_boxed_slice();
let indirect1_inode_blkno = direct_inode_page.opaque.next_blkno;

Self {
relation,
direct_inode: page_read(relation, blkno),
direct_inode,
indirect1_inode_blkno,
}
}

Expand All @@ -43,12 +53,11 @@ impl VirtualPageReader {
pub fn get_block_id(&self, virtual_id: u32) -> u32 {
let mut virtual_id = virtual_id as usize;
if virtual_id < DIRECT_COUNT {
let slice = &self.direct_inode.data()[virtual_id * 4..][..4];
return u32::from_le_bytes(slice.try_into().unwrap());
return self.direct_inode[virtual_id];
}

virtual_id -= DIRECT_COUNT;
let indirect1_inode = page_read(self.relation, self.direct_inode.opaque.next_blkno);
let indirect1_inode = page_read(self.relation, self.indirect1_inode_blkno);
if virtual_id < INDIRECT1_COUNT {
let indirect1_id = virtual_id / DIRECT_COUNT;
let indirect1_offset = virtual_id % DIRECT_COUNT;
Expand Down
16 changes: 5 additions & 11 deletions src/segment/posting/reader.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use std::{fmt::Debug, io::Read, mem::MaybeUninit};

use crate::{
page::{page_read, PageReadGuard, PageReader, VirtualPageReader},
page::{page_read, PageReader, VirtualPageReader},
segment::{field_norm::id_to_fieldnorm, posting::SkipBlockFlags},
utils::compress_block::BlockDecoder,
weight::Bm25Weight,
Expand Down Expand Up @@ -52,7 +52,6 @@ pub struct PostingReader<const WITH_FREQ: bool> {
block_data_reader: VirtualPageReader,
cur_page: pgrx::pg_sys::BlockNumber,
page_offset: usize,
page_inner: Option<PageReadGuard>,
cur_block: usize,
block_offset: usize,
remain_doc_cnt: u32,
Expand Down Expand Up @@ -108,7 +107,6 @@ impl<const WITH_FREQ: bool> PostingReader<WITH_FREQ> {
block_data_reader,
cur_page: 0,
page_offset: 0,
page_inner: None,
cur_block: 0,
block_offset: 0,
remain_doc_cnt: term_info.doc_count,
Expand Down Expand Up @@ -234,12 +232,10 @@ impl<const WITH_FREQ: bool> PostingReader<WITH_FREQ> {
self.skip_blocks[self.cur_block - 1].last_doc
};

let page = self.page_inner.get_or_insert_with(|| {
page_read(
self.index,
self.block_data_reader.get_block_id(self.cur_page),
)
});
let page = page_read(
self.index,
self.block_data_reader.get_block_id(self.cur_page),
);

if self.remain_doc_cnt < COMPRESSION_BLOCK_SIZE as u32 {
debug_assert!(skip.flag.contains(SkipBlockFlags::UNFULLED));
Expand Down Expand Up @@ -285,7 +281,6 @@ impl<const WITH_FREQ: bool> PostingReader<WITH_FREQ> {

if self.completed() {
self.page_offset = 0;
self.page_inner = None;
return;
}

Expand All @@ -295,7 +290,6 @@ impl<const WITH_FREQ: bool> PostingReader<WITH_FREQ> {
{
self.cur_page += 1;
self.page_offset = 0;
self.page_inner = None;
}
}
}

0 comments on commit 52cbdfb

Please sign in to comment.