mito_codec/
index.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
15//! Index codec utilities.
16
17use std::collections::HashMap;
18use std::sync::Arc;
19
20use datatypes::value::ValueRef;
21use memcomparable::Serializer;
22use snafu::{OptionExt, ResultExt, ensure};
23use store_api::codec::PrimaryKeyEncoding;
24use store_api::metadata::ColumnMetadata;
25use store_api::storage::ColumnId;
26
27use crate::error::{FieldTypeMismatchSnafu, IndexEncodeNullSnafu, Result};
28use crate::row_converter::sparse::{
29    RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID, SparsePrimaryKeyView,
30};
31use crate::row_converter::{PrimaryKeyCodec, SortField, build_primary_key_codec_with_fields};
32
33/// Encodes index values according to their data types for sorting and storage use.
34pub struct IndexValueCodec;
35
36impl IndexValueCodec {
37    /// Returns the index bytes of a value, or `None` for NULL.
38    /// Strings borrow their original UTF-8 bytes. For other non-null values, this
39    /// clears and reuses `buffer` rather than appending as [`Self::encode_nonnull_value`] does.
40    pub fn encode_value<'a>(
41        value: ValueRef<'a>,
42        field: &SortField,
43        buffer: &'a mut Vec<u8>,
44    ) -> Result<Option<&'a [u8]>> {
45        if value.is_null() {
46            return Ok(None);
47        }
48        if field.encode_data_type().is_string() {
49            return Ok(value
50                .try_into_string()
51                .context(FieldTypeMismatchSnafu)?
52                .map(str::as_bytes));
53        }
54        buffer.clear();
55        Self::encode_nonnull_value(value, field, buffer)?;
56        Ok(Some(buffer))
57    }
58
59    /// Extracts one sparse PK column in index format, without constructing a Value.
60    /// Numeric reserved fields borrow the PK; strings are unchunked into the reusable buffer.
61    /// Missing and null labels return None, whereas an empty string returns Some(&[]).
62    pub fn encode_sparse_value<'a, 'pk: 'a>(
63        pk: &mut SparsePrimaryKeyView<'pk, '_>,
64        column_id: ColumnId,
65        buffer: &'a mut Vec<u8>,
66    ) -> Result<Option<&'a [u8]>> {
67        if matches!(
68            column_id,
69            RESERVED_COLUMN_ID_TABLE_ID | RESERVED_COLUMN_ID_TSID
70        ) {
71            return pk.encoded_value(column_id);
72        }
73        pk.label(column_id, buffer)
74            .map(|value| value.map(str::as_bytes))
75    }
76
77    /// Serializes a non-null `ValueRef` using the data type defined in `SortField` and writes
78    /// the result into a buffer.
79    ///
80    /// For `String` data types, we don't serialize it via memcomparable, but directly write the
81    /// bytes into the buffer, since we have to keep the original string for searching with regex.
82    ///
83    /// # Arguments
84    /// * `value` - The value to be encoded.
85    /// * `field` - Contains data type to guide serialization.
86    /// * `buffer` - Destination buffer for the serialized value.
87    pub fn encode_nonnull_value(
88        value: ValueRef,
89        field: &SortField,
90        buffer: &mut Vec<u8>,
91    ) -> Result<()> {
92        ensure!(!value.is_null(), IndexEncodeNullSnafu);
93
94        if field.encode_data_type().is_string() {
95            let value = value
96                .try_into_string()
97                .context(FieldTypeMismatchSnafu)?
98                .context(IndexEncodeNullSnafu)?;
99            buffer.extend_from_slice(value.as_bytes());
100            Ok(())
101        } else {
102            buffer.reserve(field.estimated_size());
103            let mut serializer = Serializer::new(buffer);
104            field.serialize(&mut serializer, &value)
105        }
106    }
107}
108
109pub struct PkColInfo {
110    pub idx: usize,
111    pub field: SortField,
112}
113
114impl PkColInfo {
115    pub fn new(idx: usize, field: SortField) -> Self {
116        Self { idx, field }
117    }
118}
119
120/// Decodes primary key values into their corresponding column ids, data types and values.
121pub struct IndexValuesCodec {
122    /// Column ids -> column info mapping.
123    columns_mapping: HashMap<ColumnId, PkColInfo>,
124    /// The decoder for the primary key.
125    decoder: Arc<dyn PrimaryKeyCodec>,
126}
127
128impl IndexValuesCodec {
129    /// Creates a new `IndexValuesCodec` from a list of `ColumnMetadata` of tag columns.
130    pub fn from_tag_columns<'a>(
131        primary_key_encoding: PrimaryKeyEncoding,
132        tag_columns: impl Iterator<Item = &'a ColumnMetadata>,
133    ) -> Self {
134        let (columns_mapping, fields): (HashMap<ColumnId, PkColInfo>, Vec<(ColumnId, SortField)>) =
135            tag_columns
136                .enumerate()
137                .map(|(idx, column)| {
138                    let col_id = column.column_id;
139                    let field = SortField::new(column.column_schema.data_type.clone());
140                    let pk_col_info = PkColInfo::new(idx, field.clone());
141                    ((col_id, pk_col_info), (col_id, field))
142                })
143                .unzip();
144
145        let decoder = build_primary_key_codec_with_fields(primary_key_encoding, fields.into_iter());
146
147        Self {
148            columns_mapping,
149            decoder,
150        }
151    }
152
153    pub fn pk_col_info(&self, column_id: ColumnId) -> Option<&PkColInfo> {
154        self.columns_mapping.get(&column_id)
155    }
156
157    pub fn decoder(&self) -> &dyn PrimaryKeyCodec {
158        self.decoder.as_ref()
159    }
160}
161
162#[cfg(test)]
163mod tests {
164    use datatypes::data_type::ConcreteDataType;
165    use datatypes::schema::ColumnSchema;
166    use datatypes::value::Value;
167    use store_api::metadata::ColumnMetadata;
168
169    use super::*;
170    use crate::error::Error;
171    use crate::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField};
172
173    #[test]
174    fn borrowed_values_preserve_index_encoding_with_reused_buffer() {
175        let mut buffer = vec![0xff; 64];
176        for value in [
177            Value::from("中文\0abcdefgh"),
178            Value::from(""),
179            Value::Int64(-42),
180            Value::UInt64(u64::MAX),
181            Value::Boolean(true),
182            Value::Binary(vec![0, 1, 255].into()),
183            Value::Null,
184            Value::from("after-null"),
185        ] {
186            let field = SortField::new(value.data_type());
187            let mut expected = Vec::new();
188            if !value.is_null() {
189                IndexValueCodec::encode_nonnull_value(value.as_value_ref(), &field, &mut expected)
190                    .unwrap();
191            }
192            let encoded =
193                IndexValueCodec::encode_value(value.as_value_ref(), &field, &mut buffer).unwrap();
194            assert_eq!(encoded, (!value.is_null()).then_some(expected.as_slice()));
195        }
196        for (value, field) in [
197            (ValueRef::UInt64(1), ConcreteDataType::string_datatype()),
198            (ValueRef::String("x"), ConcreteDataType::uint64_datatype()),
199        ] {
200            assert!(matches!(
201                IndexValueCodec::encode_value(value, &SortField::new(field), &mut buffer),
202                Err(Error::FieldTypeMismatch { .. })
203            ));
204        }
205    }
206
207    #[test]
208    fn test_encode_value_basic() {
209        let value = ValueRef::from("hello");
210        let field = SortField::new(ConcreteDataType::string_datatype());
211
212        let mut buffer = Vec::new();
213        IndexValueCodec::encode_nonnull_value(value, &field, &mut buffer).unwrap();
214        assert!(!buffer.is_empty());
215    }
216
217    #[test]
218    fn test_encode_value_type_mismatch() {
219        let value = ValueRef::from("hello");
220        let field = SortField::new(ConcreteDataType::int64_datatype());
221
222        let mut buffer = Vec::new();
223        let res = IndexValueCodec::encode_nonnull_value(value, &field, &mut buffer);
224        assert!(matches!(res, Err(Error::FieldTypeMismatch { .. })));
225    }
226
227    #[test]
228    fn test_encode_null_value() {
229        let value = ValueRef::Null;
230        let field = SortField::new(ConcreteDataType::string_datatype());
231
232        let mut buffer = Vec::new();
233        let res = IndexValueCodec::encode_nonnull_value(value, &field, &mut buffer);
234        assert!(matches!(res, Err(Error::IndexEncodeNull { .. })));
235    }
236
237    #[test]
238    fn test_decode_primary_key_basic() {
239        let tag_columns = [
240            ColumnMetadata {
241                column_schema: ColumnSchema::new("tag0", ConcreteDataType::string_datatype(), true),
242                semantic_type: api::v1::SemanticType::Tag,
243                column_id: 1,
244            },
245            ColumnMetadata {
246                column_schema: ColumnSchema::new("tag1", ConcreteDataType::int64_datatype(), false),
247                semantic_type: api::v1::SemanticType::Tag,
248                column_id: 2,
249            },
250        ];
251
252        let primary_key = DensePrimaryKeyCodec::with_fields(vec![
253            (0, SortField::new(ConcreteDataType::string_datatype())),
254            (1, SortField::new(ConcreteDataType::int64_datatype())),
255        ])
256        .encode([ValueRef::Null, ValueRef::Int64(10)].into_iter())
257        .unwrap();
258
259        let codec =
260            IndexValuesCodec::from_tag_columns(PrimaryKeyEncoding::Dense, tag_columns.iter());
261        let values = codec.decoder().decode(&primary_key).unwrap().into_dense();
262
263        assert_eq!(values.len(), 2);
264        assert_eq!(values[0], Value::Null);
265        assert_eq!(values[1], Value::Int64(10));
266    }
267}
mito_codec/index.rs

mito_codec/
index.rs
