risingwave_storage/hummock/vector/
mod.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use risingwave_hummock_sdk::vector_index::VectorFileInfo;
16
17use crate::hummock::vector::file::{VectorBlock, VectorBlockBuilder, VectorBlockMeta};
18use crate::hummock::{HummockResult, SstableStoreRef, VectorBlockHolder};
19use crate::vector::VectorRef;
20use crate::vector::hnsw::VectorAccessor;
21
22pub(crate) mod file;
23pub(crate) mod writer;
24
25pub struct VectorBlockAccessor {
26    block: VectorBlockHolder,
27    idx: usize,
28}
29
30impl VectorAccessor for VectorBlockAccessor {
31    fn vec_ref(&self) -> VectorRef<'_> {
32        self.block.vec_ref(self.idx)
33    }
34
35    fn info(&self) -> &[u8] {
36        self.block.info(self.idx)
37    }
38}
39
40pub enum EnumVectorAccessor<'a> {
41    Builder(&'a VectorBlockBuilder, usize),
42    BlockRef(&'a VectorBlock, usize),
43    BlockHolder(VectorBlockAccessor),
44}
45
46impl VectorAccessor for EnumVectorAccessor<'_> {
47    fn vec_ref(&self) -> VectorRef<'_> {
48        match self {
49            EnumVectorAccessor::Builder(builder, offset) => builder.vec_ref(*offset),
50            EnumVectorAccessor::BlockRef(block, offset) => block.vec_ref(*offset),
51            EnumVectorAccessor::BlockHolder(accessor) => accessor.vec_ref(),
52        }
53    }
54
55    fn info(&self) -> &[u8] {
56        match self {
57            EnumVectorAccessor::Builder(builder, offset) => builder.info(*offset),
58            EnumVectorAccessor::BlockRef(block, offset) => block.info(*offset),
59            EnumVectorAccessor::BlockHolder(accessor) => accessor.info(),
60        }
61    }
62}
63
64pub async fn get_vector_block(
65    sstable_store: &SstableStoreRef,
66    files: &[VectorFileInfo],
67    idx: usize,
68) -> HummockResult<VectorBlockAccessor> {
69    let vector_file = search_vector_files(files, idx);
70    let meta = sstable_store.get_vector_file_meta(vector_file).await?;
71    let (block_meta, block_idx, offset) = search_blocks(&meta.block_metas, idx);
72    let block = sstable_store
73        .get_vector_block(vector_file, block_idx, block_meta)
74        .await?;
75    Ok(VectorBlockAccessor { block, idx: offset })
76}
77
78fn search_vector_files(files: &[VectorFileInfo], idx: usize) -> &VectorFileInfo {
79    let (file_idx, _) = search_vector(files, idx, |file| file.start_vector_id);
80    &files[file_idx]
81}
82
83fn search_blocks(blocks: &[VectorBlockMeta], idx: usize) -> (&VectorBlockMeta, usize, usize) {
84    let (block_idx, offset) = search_vector(blocks, idx, |block| block.start_vector_id);
85    (&blocks[block_idx], block_idx, offset)
86}
87
88/// return (`holder_idx`, `offset` inside holder)
89pub fn search_vector<T>(
90    vector_holders: &[T],
91    idx: usize,
92    get_start_vector_id: impl Fn(&T) -> usize,
93) -> (usize, usize) {
94    let holder_idx = vector_holders
95        .partition_point(|holder| get_start_vector_id(holder) <= idx)
96        .checked_sub(1)
97        .unwrap_or_else(|| {
98            panic!(
99                "idx {} too small for first vector id {}",
100                idx,
101                vector_holders.first().map_or(0, &get_start_vector_id)
102            )
103        });
104    let holder = &vector_holders[holder_idx];
105    let offset = idx - get_start_vector_id(holder);
106    (holder_idx, offset)
107}
risingwave_storage/hummock/vector/mod.rs

risingwave_storage/hummock/vector/
mod.rs
