store_api/storage/
requests.rs

1// Copyright 2023 Greptime Team
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 std::collections::HashMap;
16use std::fmt::{Display, Formatter};
17
18use common_error::ext::BoxedError;
19use common_recordbatch::OrderOption;
20use datafusion_expr::expr::Expr;
21// Re-export vector types from datatypes to avoid duplication
22pub use datatypes::schema::{VectorDistanceMetric, VectorIndexEngineType};
23use datatypes::types::json_type::JsonNativeType;
24use itertools::Itertools;
25use strum::Display;
26
27use crate::storage::{ColumnId, ProjectionInput, SequenceNumber};
28
29/// A hint for KNN vector search.
30#[derive(Debug, Clone, PartialEq)]
31pub struct VectorSearchRequest {
32    /// Column ID of the vector column to search.
33    pub column_id: ColumnId,
34    /// The query vector to search for.
35    pub query_vector: Vec<f32>,
36    /// Number of nearest neighbors to return.
37    pub k: usize,
38    /// Distance metric to use (matches the index metric).
39    pub metric: VectorDistanceMetric,
40}
41
42/// Search results from vector index.
43#[derive(Debug, Clone, PartialEq)]
44pub struct VectorSearchMatches {
45    /// Keys (row offsets in the index).
46    pub keys: Vec<u64>,
47    /// Distances from the query vector.
48    pub distances: Vec<f32>,
49}
50
51/// Trait for vector index engines (HNSW implementations).
52///
53/// This trait defines the interface for pluggable vector index engines.
54/// Implementations (e.g., UsearchEngine) are provided by storage engines like mito2.
55pub trait VectorIndexEngine: Send + Sync {
56    /// Adds a vector with the given key.
57    fn add(&mut self, key: u64, vector: &[f32]) -> Result<(), BoxedError>;
58
59    /// Searches for k nearest neighbors.
60    fn search(&self, query: &[f32], k: usize) -> Result<VectorSearchMatches, BoxedError>;
61
62    /// Returns the serialized length.
63    fn serialized_length(&self) -> usize;
64
65    /// Serializes the index to a buffer.
66    fn save_to_buffer(&self, buffer: &mut [u8]) -> Result<(), BoxedError>;
67
68    /// Reserves capacity for vectors.
69    fn reserve(&mut self, capacity: usize) -> Result<(), BoxedError>;
70
71    /// Returns current size (number of vectors).
72    fn size(&self) -> usize;
73
74    /// Returns current capacity.
75    fn capacity(&self) -> usize;
76
77    /// Returns memory usage in bytes.
78    fn memory_usage(&self) -> usize;
79}
80
81/// A hint on how to select rows from a time-series.
82#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Display)]
83pub enum TimeSeriesRowSelector {
84    /// Only keep the last row of each time-series.
85    LastRow,
86}
87
88/// A hint on how to distribute time-series data on the scan output.
89#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Display)]
90pub enum TimeSeriesDistribution {
91    /// Data are distributed by time window first. The scanner will
92    /// return all data within one time window before moving to the next one.
93    TimeWindowed,
94    /// Data are organized by time-series first. The scanner will return
95    /// all data for one time-series before moving to the next one.
96    PerSeries,
97}
98
99#[derive(Default, Clone, Debug, PartialEq)]
100pub struct ScanRequest {
101    /// Optional projection information for the scan. `None` reads all root
102    /// columns.
103    pub projection_input: Option<ProjectionInput>,
104    /// Filters pushed down
105    pub filters: Vec<Expr>,
106    /// Expected output ordering. This is only a hint and isn't guaranteed.
107    pub output_ordering: Option<Vec<OrderOption>>,
108    /// limit can be used to reduce the amount scanned
109    /// from the datasource as a performance optimization.
110    /// If set, it contains the amount of rows needed by the caller,
111    /// The data source should return *at least* this number of rows if available.
112    pub limit: Option<usize>,
113    /// Optional hint to select rows from time-series.
114    pub series_row_selector: Option<TimeSeriesRowSelector>,
115    /// Optional constraint on the sequence number of the rows to read.
116    /// If set, only rows with a sequence number **lesser or equal** to this value
117    /// will be returned.
118    /// This is the effective memtable upper bound used by the scan, whether provided
119    /// explicitly or bound on scan open.
120    pub memtable_max_sequence: Option<SequenceNumber>,
121    /// Optional constraint on the minimal sequence number in the memtable.
122    /// If set, only the memtables that contain sequences **greater than** this value will be scanned
123    pub memtable_min_sequence: Option<SequenceNumber>,
124    /// Optional constraint on the minimal sequence number in the SST files.
125    /// If set, only the SST files that contain sequences greater than this value will be scanned.
126    pub sst_min_sequence: Option<SequenceNumber>,
127    /// Whether to skip all SST files.
128    /// This is stronger than `sst_min_sequence` and also skips SST files without sequence metadata.
129    pub skip_sst_files: bool,
130    /// Whether to bind the effective snapshot upper bound when opening the scan.
131    pub snapshot_on_scan: bool,
132    /// Optional hint for the distribution of time-series data.
133    pub distribution: Option<TimeSeriesDistribution>,
134    /// Optional hint for KNN vector search. When set, the scan should use
135    /// vector index to find the k nearest neighbors.
136    pub vector_search: Option<VectorSearchRequest>,
137    /// Optional hint from query-driven JSON type concretization.
138    pub json_type_hint: HashMap<String, JsonNativeType>,
139    /// Whether Mito should keep string primary-key columns dictionary encoded in its output.
140    pub preserve_pk_dictionary_encoding: bool,
141}
142
143impl ScanRequest {
144    /// Returns the top-level projected column indices.
145    pub fn projection_indices(&self) -> Option<&[usize]> {
146        self.projection_input
147            .as_ref()
148            .map(|projection_input| projection_input.projection.as_slice())
149    }
150}
151
152impl Display for ScanRequest {
153    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
154        enum Delimiter {
155            None,
156            Init,
157        }
158
159        impl Delimiter {
160            fn as_str(&mut self) -> &str {
161                match self {
162                    Delimiter::None => {
163                        *self = Delimiter::Init;
164                        ""
165                    }
166                    Delimiter::Init => ", ",
167                }
168            }
169        }
170
171        let mut delimiter = Delimiter::None;
172
173        write!(f, "ScanRequest {{ ")?;
174        if let Some(projection) = &self.projection_input {
175            write!(f, "{}projection: {:?}", delimiter.as_str(), projection)?;
176        }
177        if !self.filters.is_empty() {
178            write!(
179                f,
180                "{}filters: [{}]",
181                delimiter.as_str(),
182                self.filters
183                    .iter()
184                    .map(|f| f.to_string())
185                    .collect::<Vec<_>>()
186                    .join(", ")
187            )?;
188        }
189        if let Some(output_ordering) = &self.output_ordering {
190            write!(
191                f,
192                "{}output_ordering: {:?}",
193                delimiter.as_str(),
194                output_ordering
195            )?;
196        }
197        if let Some(limit) = &self.limit {
198            write!(f, "{}limit: {}", delimiter.as_str(), limit)?;
199        }
200        if let Some(series_row_selector) = &self.series_row_selector {
201            write!(
202                f,
203                "{}series_row_selector: {}",
204                delimiter.as_str(),
205                series_row_selector
206            )?;
207        }
208        if let Some(sequence) = &self.memtable_max_sequence {
209            write!(f, "{}sequence: {}", delimiter.as_str(), sequence)?;
210        }
211        if let Some(sst_min_sequence) = &self.sst_min_sequence {
212            write!(
213                f,
214                "{}sst_min_sequence: {}",
215                delimiter.as_str(),
216                sst_min_sequence
217            )?;
218        }
219        if self.skip_sst_files {
220            write!(
221                f,
222                "{}skip_sst_files: {}",
223                delimiter.as_str(),
224                self.skip_sst_files
225            )?;
226        }
227        if self.snapshot_on_scan {
228            write!(
229                f,
230                "{}snapshot_on_scan: {}",
231                delimiter.as_str(),
232                self.snapshot_on_scan
233            )?;
234        }
235        if self.preserve_pk_dictionary_encoding {
236            write!(
237                f,
238                "{}preserve_pk_dictionary_encoding: true",
239                delimiter.as_str()
240            )?;
241        }
242        if let Some(distribution) = &self.distribution {
243            write!(f, "{}distribution: {}", delimiter.as_str(), distribution)?;
244        }
245        if let Some(vector_search) = &self.vector_search {
246            write!(
247                f,
248                "{}vector_search: column_id={}, k={}, metric={}",
249                delimiter.as_str(),
250                vector_search.column_id,
251                vector_search.k,
252                vector_search.metric
253            )?;
254        }
255        if !self.json_type_hint.is_empty() {
256            write!(
257                f,
258                "{}json_type_hint: {}",
259                delimiter.as_str(),
260                self.json_type_hint
261                    .iter()
262                    .map(|(column, json_type)| format!("({column}: {json_type})"))
263                    .join(", ")
264            )?;
265        }
266        write!(f, " }}")
267    }
268}
269
270#[cfg(test)]
271mod tests {
272    use datafusion_expr::{Operator, binary_expr, col, lit};
273
274    use super::*;
275
276    #[test]
277    fn test_display_scan_request() {
278        let request = ScanRequest {
279            ..Default::default()
280        };
281        assert_eq!(request.to_string(), "ScanRequest {  }");
282
283        let projection_input = Some(vec![1, 2].into());
284        let request = ScanRequest {
285            projection_input,
286            filters: vec![
287                binary_expr(col("i"), Operator::Gt, lit(1)),
288                binary_expr(col("s"), Operator::Eq, lit("x")),
289            ],
290            limit: Some(10),
291            ..Default::default()
292        };
293        assert_eq!(
294            request.to_string(),
295            r#"ScanRequest { projection: ProjectionInput { projection: [1, 2], nested_paths: [] }, filters: [i > Int32(1), s = Utf8("x")], limit: 10 }"#
296        );
297
298        let request = ScanRequest {
299            filters: vec![
300                binary_expr(col("i"), Operator::Gt, lit(1)),
301                binary_expr(col("s"), Operator::Eq, lit("x")),
302            ],
303            limit: Some(10),
304            ..Default::default()
305        };
306        assert_eq!(
307            request.to_string(),
308            r#"ScanRequest { filters: [i > Int32(1), s = Utf8("x")], limit: 10 }"#
309        );
310
311        let projection_input = Some(vec![1, 2].into());
312        let request = ScanRequest {
313            projection_input,
314            limit: Some(10),
315            ..Default::default()
316        };
317        assert_eq!(
318            request.to_string(),
319            "ScanRequest { projection: ProjectionInput { projection: [1, 2], nested_paths: [] }, limit: 10 }"
320        );
321
322        let projection_input = Some(ProjectionInput::new(vec![1, 2]).with_nested_paths(vec![
323            vec!["j".to_string(), "a".to_string(), "b".to_string()],
324            vec!["s".to_string(), "x".to_string()],
325        ]));
326        let request = ScanRequest {
327            projection_input,
328            limit: Some(10),
329            ..Default::default()
330        };
331        assert_eq!(
332            request.to_string(),
333            r#"ScanRequest { projection: ProjectionInput { projection: [1, 2], nested_paths: [["j", "a", "b"], ["s", "x"]] }, limit: 10 }"#
334        );
335
336        let request = ScanRequest {
337            snapshot_on_scan: true,
338            ..Default::default()
339        };
340        assert_eq!(
341            request.to_string(),
342            "ScanRequest { snapshot_on_scan: true }"
343        );
344
345        let request = ScanRequest {
346            skip_sst_files: true,
347            ..Default::default()
348        };
349        assert_eq!(request.to_string(), "ScanRequest { skip_sst_files: true }");
350    }
351}
store_api/storage/requests.rs

store_api/storage/
requests.rs
