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::vector::monitor::VectorStoreCacheStats;
19use crate::hummock::{HummockResult, SstableStoreRef, VectorBlockHolder};
20use crate::vector::VectorRef;
21use crate::vector::hnsw::VectorAccessor;
22
23pub(crate) mod file;
24pub(crate) mod monitor;
25pub(crate) mod writer;
26
27pub struct VectorBlockAccessor {
28    block: VectorBlockHolder,
29    idx: usize,
30}
31
32impl VectorAccessor for VectorBlockAccessor {
33    fn vec_ref(&self) -> VectorRef<'_> {
34        self.block.vec_ref(self.idx)
35    }
36
37    fn info(&self) -> &[u8] {
38        self.block.info(self.idx)
39    }
40}
41
42pub enum EnumVectorAccessor<'a> {
43    Builder(&'a VectorBlockBuilder, usize),
44    BlockRef(&'a VectorBlock, usize),
45    BlockHolder(VectorBlockAccessor),
46}
47
48impl VectorAccessor for EnumVectorAccessor<'_> {
49    fn vec_ref(&self) -> VectorRef<'_> {
50        match self {
51            EnumVectorAccessor::Builder(builder, offset) => builder.vec_ref(*offset),
52            EnumVectorAccessor::BlockRef(block, offset) => block.vec_ref(*offset),
53            EnumVectorAccessor::BlockHolder(accessor) => accessor.vec_ref(),
54        }
55    }
56
57    fn info(&self) -> &[u8] {
58        match self {
59            EnumVectorAccessor::Builder(builder, offset) => builder.info(*offset),
60            EnumVectorAccessor::BlockRef(block, offset) => block.info(*offset),
61            EnumVectorAccessor::BlockHolder(accessor) => accessor.info(),
62        }
63    }
64}
65
66pub async fn get_vector_block(
67    sstable_store: &SstableStoreRef,
68    files: &[VectorFileInfo],
69    idx: usize,
70    stats: &mut VectorStoreCacheStats,
71) -> HummockResult<VectorBlockAccessor> {
72    let vector_file = search_vector_files(files, idx);
73    let meta = sstable_store
74        .get_vector_file_meta(vector_file, stats)
75        .await?;
76    let (block_meta, block_idx, offset) = search_blocks(&meta.block_metas, idx);
77    let block = sstable_store
78        .get_vector_block(vector_file, block_idx, block_meta, stats)
79        .await?;
80    Ok(VectorBlockAccessor { block, idx: offset })
81}
82
83fn search_vector_files(files: &[VectorFileInfo], idx: usize) -> &VectorFileInfo {
84    let (file_idx, _) = search_vector(files, idx, |file| file.start_vector_id);
85    &files[file_idx]
86}
87
88fn search_blocks(blocks: &[VectorBlockMeta], idx: usize) -> (&VectorBlockMeta, usize, usize) {
89    let (block_idx, offset) = search_vector(blocks, idx, |block| block.start_vector_id);
90    (&blocks[block_idx], block_idx, offset)
91}
92
93/// return (`holder_idx`, `offset` inside holder)
94pub fn search_vector<T>(
95    vector_holders: &[T],
96    idx: usize,
97    get_start_vector_id: impl Fn(&T) -> usize,
98) -> (usize, usize) {
99    let holder_idx = vector_holders
100        .partition_point(|holder| get_start_vector_id(holder) <= idx)
101        .checked_sub(1)
102        .unwrap_or_else(|| {
103            panic!(
104                "idx {} too small for first vector id {}",
105                idx,
106                vector_holders.first().map_or(0, &get_start_vector_id)
107            )
108        });
109    let holder = &vector_holders[holder_idx];
110    let offset = idx - get_start_vector_id(holder);
111    (holder_idx, offset)
112}
risingwave_storage/hummock/vector/mod.rs

risingwave_storage/hummock/vector/
mod.rs
