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, 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    #[strum(to_string = "LastRow {{ after_merge: {after_merge} }}")]
86    LastRow {
87        /// Whether selection runs after cross-source merge and deduplication.
88        after_merge: bool,
89    },
90}
91
92/// A hint on how to distribute time-series data on the scan output.
93#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Display)]
94pub enum TimeSeriesDistribution {
95    /// Data are distributed by time window first. The scanner will
96    /// return all data within one time window before moving to the next one.
97    TimeWindowed,
98    /// Data are organized by time-series first. The scanner will return
99    /// all data for one time-series before moving to the next one.
100    PerSeries,
101}
102
103#[derive(Default, Clone, Debug, PartialEq)]
104pub struct ScanRequest {
105    /// Optional projection information for the scan. `None` reads all root
106    /// columns.
107    pub projection: Option<Vec<usize>>,
108    /// Filters pushed down
109    pub filters: Vec<Expr>,
110    /// Expected output ordering. This is only a hint and isn't guaranteed.
111    pub output_ordering: Option<Vec<OrderOption>>,
112    /// limit can be used to reduce the amount scanned
113    /// from the datasource as a performance optimization.
114    /// If set, it contains the amount of rows needed by the caller,
115    /// The data source should return *at least* this number of rows if available.
116    pub limit: Option<usize>,
117    /// Optional hint to select rows from time-series.
118    pub series_row_selector: Option<TimeSeriesRowSelector>,
119    /// Optional constraint on the sequence number of the rows to read.
120    /// If set, only rows with a sequence number **lesser or equal** to this value
121    /// will be returned.
122    /// This is the effective memtable upper bound used by the scan, whether provided
123    /// explicitly or bound on scan open.
124    pub memtable_max_sequence: Option<SequenceNumber>,
125    /// Optional constraint on the minimal sequence number in the memtable.
126    /// If set, only the memtables that contain sequences **greater than** this value will be scanned
127    pub memtable_min_sequence: Option<SequenceNumber>,
128    /// Optional constraint on the minimal sequence number in the SST files.
129    /// If set, only the SST files that contain sequences greater than this value will be scanned.
130    pub sst_min_sequence: Option<SequenceNumber>,
131    /// Whether to skip all SST files.
132    /// This is stronger than `sst_min_sequence` and also skips SST files without sequence metadata.
133    pub skip_sst_files: bool,
134    /// Whether to bind the effective snapshot upper bound when opening the scan.
135    pub snapshot_on_scan: bool,
136    /// Explicit intent to read an exact row-level sequence delta `(min, max]`
137    /// across memtables and all SST files (Flow's `sequence_range` incremental
138    /// mode). The engine performs exact row-level filtering only when the region
139    /// preserves per-row sequences and every participating SST file is trusted;
140    /// otherwise it returns a structured stale/unsupported error so the caller
141    /// falls back instead of silently approximating.
142    ///
143    /// Historical `memtable_only` reads must never set this flag.
144    pub exact_sequence_range: bool,
145    /// Optional hint for the distribution of time-series data.
146    pub distribution: Option<TimeSeriesDistribution>,
147    /// Optional hint for KNN vector search. When set, the scan should use
148    /// vector index to find the k nearest neighbors.
149    pub vector_search: Option<VectorSearchRequest>,
150    /// Optional hint from query-driven JSON type concretization.
151    pub json_type_hint: HashMap<String, JsonNativeType>,
152    /// Whether Mito should keep string primary-key columns dictionary encoded in its output.
153    pub preserve_pk_dictionary_encoding: bool,
154}
155
156impl Display for ScanRequest {
157    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
158        enum Delimiter {
159            None,
160            Init,
161        }
162
163        impl Delimiter {
164            fn as_str(&mut self) -> &str {
165                match self {
166                    Delimiter::None => {
167                        *self = Delimiter::Init;
168                        ""
169                    }
170                    Delimiter::Init => ", ",
171                }
172            }
173        }
174
175        let mut delimiter = Delimiter::None;
176
177        write!(f, "ScanRequest {{ ")?;
178        if let Some(projection) = &self.projection {
179            write!(f, "{}projection: {:?}", delimiter.as_str(), projection)?;
180        }
181        if !self.filters.is_empty() {
182            write!(
183                f,
184                "{}filters: [{}]",
185                delimiter.as_str(),
186                self.filters
187                    .iter()
188                    .map(|f| f.to_string())
189                    .collect::<Vec<_>>()
190                    .join(", ")
191            )?;
192        }
193        if let Some(output_ordering) = &self.output_ordering {
194            write!(
195                f,
196                "{}output_ordering: {:?}",
197                delimiter.as_str(),
198                output_ordering
199            )?;
200        }
201        if let Some(limit) = &self.limit {
202            write!(f, "{}limit: {}", delimiter.as_str(), limit)?;
203        }
204        if let Some(series_row_selector) = &self.series_row_selector {
205            write!(
206                f,
207                "{}series_row_selector: {}",
208                delimiter.as_str(),
209                series_row_selector
210            )?;
211        }
212        if let Some(sequence) = &self.memtable_max_sequence {
213            write!(f, "{}sequence: {}", delimiter.as_str(), sequence)?;
214        }
215        if let Some(sst_min_sequence) = &self.sst_min_sequence {
216            write!(
217                f,
218                "{}sst_min_sequence: {}",
219                delimiter.as_str(),
220                sst_min_sequence
221            )?;
222        }
223        if self.skip_sst_files {
224            write!(
225                f,
226                "{}skip_sst_files: {}",
227                delimiter.as_str(),
228                self.skip_sst_files
229            )?;
230        }
231        if self.snapshot_on_scan {
232            write!(
233                f,
234                "{}snapshot_on_scan: {}",
235                delimiter.as_str(),
236                self.snapshot_on_scan
237            )?;
238        }
239        if self.exact_sequence_range {
240            write!(f, "{}exact_sequence_range: true", delimiter.as_str())?;
241        }
242        if self.preserve_pk_dictionary_encoding {
243            write!(
244                f,
245                "{}preserve_pk_dictionary_encoding: true",
246                delimiter.as_str()
247            )?;
248        }
249        if let Some(distribution) = &self.distribution {
250            write!(f, "{}distribution: {}", delimiter.as_str(), distribution)?;
251        }
252        if let Some(vector_search) = &self.vector_search {
253            write!(
254                f,
255                "{}vector_search: column_id={}, k={}, metric={}",
256                delimiter.as_str(),
257                vector_search.column_id,
258                vector_search.k,
259                vector_search.metric
260            )?;
261        }
262        if !self.json_type_hint.is_empty() {
263            write!(
264                f,
265                "{}json_type_hint: {}",
266                delimiter.as_str(),
267                self.json_type_hint
268                    .iter()
269                    .map(|(column, json_type)| format!("({column}: {json_type})"))
270                    .join(", ")
271            )?;
272        }
273        write!(f, " }}")
274    }
275}
276
277#[cfg(test)]
278mod tests {
279    use datafusion_expr::{Operator, binary_expr, col, lit};
280
281    use super::*;
282
283    #[test]
284    fn test_display_scan_request() {
285        let request = ScanRequest {
286            ..Default::default()
287        };
288        assert_eq!(request.to_string(), "ScanRequest {  }");
289
290        let projection = Some(vec![1, 2]);
291        let request = ScanRequest {
292            projection,
293            filters: vec![
294                binary_expr(col("i"), Operator::Gt, lit(1)),
295                binary_expr(col("s"), Operator::Eq, lit("x")),
296            ],
297            limit: Some(10),
298            ..Default::default()
299        };
300        assert_eq!(
301            request.to_string(),
302            r#"ScanRequest { projection: [1, 2], filters: [i > Int32(1), s = Utf8("x")], limit: 10 }"#
303        );
304
305        let request = ScanRequest {
306            filters: vec![
307                binary_expr(col("i"), Operator::Gt, lit(1)),
308                binary_expr(col("s"), Operator::Eq, lit("x")),
309            ],
310            limit: Some(10),
311            ..Default::default()
312        };
313        assert_eq!(
314            request.to_string(),
315            r#"ScanRequest { filters: [i > Int32(1), s = Utf8("x")], limit: 10 }"#
316        );
317
318        let projection = Some(vec![1, 2]);
319        let request = ScanRequest {
320            projection,
321            limit: Some(10),
322            ..Default::default()
323        };
324        assert_eq!(
325            request.to_string(),
326            "ScanRequest { projection: [1, 2], limit: 10 }"
327        );
328
329        let request = ScanRequest {
330            series_row_selector: Some(TimeSeriesRowSelector::LastRow { after_merge: true }),
331            snapshot_on_scan: true,
332            exact_sequence_range: true,
333            ..Default::default()
334        };
335        assert_eq!(
336            request.to_string(),
337            "ScanRequest { series_row_selector: LastRow { after_merge: true }, snapshot_on_scan: true, exact_sequence_range: true }"
338        );
339
340        assert_eq!(
341            TimeSeriesRowSelector::LastRow { after_merge: false }.to_string(),
342            "LastRow { after_merge: false }"
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
