Skip to main content

lance_index/
types.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::sync::Arc;
5
6use crate::IndexType;
7use lance_table::format::IndexMetadata;
8use roaring::RoaringBitmap;
9use uuid::Uuid;
10
11/// A single physical segment of a logical index.
12///
13/// Each segment is stored independently and will become one manifest entry when committed.
14/// The logical index identity (name / target column / dataset version) is provided separately
15/// by the commit API.
16#[derive(Debug, Clone, PartialEq)]
17pub struct IndexSegment {
18    /// Unique ID of the physical segment.
19    uuid: Uuid,
20    /// The fragments covered by this segment.
21    fragment_bitmap: RoaringBitmap,
22    /// Metadata specific to the index type.
23    index_details: Arc<prost_types::Any>,
24    /// The on-disk index version for this segment.
25    index_version: i32,
26}
27
28impl IndexSegment {
29    /// Create a fully described segment with the given UUID, fragment coverage, and index
30    /// metadata.
31    pub fn new<I>(
32        uuid: Uuid,
33        fragment_bitmap: I,
34        index_details: Arc<prost_types::Any>,
35        index_version: i32,
36    ) -> Self
37    where
38        I: IntoIterator<Item = u32>,
39    {
40        Self {
41            uuid,
42            fragment_bitmap: fragment_bitmap.into_iter().collect(),
43            index_details,
44            index_version,
45        }
46    }
47
48    /// Return the UUID of this segment.
49    pub fn uuid(&self) -> Uuid {
50        self.uuid
51    }
52
53    /// Return the fragment coverage of this segment.
54    pub fn fragment_bitmap(&self) -> &RoaringBitmap {
55        &self.fragment_bitmap
56    }
57
58    /// Return the serialized index details for this segment.
59    pub fn index_details(&self) -> &Arc<prost_types::Any> {
60        &self.index_details
61    }
62
63    /// Return the on-disk index version for this segment.
64    pub fn index_version(&self) -> i32 {
65        self.index_version
66    }
67
68    /// Consume the segment and return its component parts.
69    pub fn into_parts(self) -> (Uuid, RoaringBitmap, Arc<prost_types::Any>, i32) {
70        (
71            self.uuid,
72            self.fragment_bitmap,
73            self.index_details,
74            self.index_version,
75        )
76    }
77}
78
79/// A plan for building one physical segment from one or more existing
80/// vector index segments.
81#[derive(Debug, Clone, PartialEq)]
82pub struct IndexSegmentPlan {
83    segment: IndexSegment,
84    segments: Vec<IndexMetadata>,
85    estimated_bytes: u64,
86    requested_index_type: Option<IndexType>,
87}
88
89impl IndexSegmentPlan {
90    /// Create a plan for one built segment.
91    pub fn new(
92        segment: IndexSegment,
93        segments: Vec<IndexMetadata>,
94        estimated_bytes: u64,
95        requested_index_type: Option<IndexType>,
96    ) -> Self {
97        Self {
98            segment,
99            segments,
100            estimated_bytes,
101            requested_index_type,
102        }
103    }
104
105    /// Return the segment metadata that should be committed after this plan is built.
106    pub fn segment(&self) -> &IndexSegment {
107        &self.segment
108    }
109
110    /// Return the input segment metadata that should be combined into the segment.
111    pub fn segments(&self) -> &[IndexMetadata] {
112        &self.segments
113    }
114
115    /// Return the estimated number of bytes covered by this plan.
116    pub fn estimated_bytes(&self) -> u64 {
117        self.estimated_bytes
118    }
119
120    /// Return the requested logical index type, if one was supplied to the planner.
121    pub fn requested_index_type(&self) -> Option<IndexType> {
122        self.requested_index_type
123    }
124}