mito_codec/row_converter/
sparse.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
15mod checked;
16
17use std::collections::{HashMap, HashSet};
18use std::sync::Arc;
19
20use bytes::BufMut;
21use common_recordbatch::filter::SimpleFilterEvaluator;
22use datatypes::prelude::ConcreteDataType;
23use datatypes::value::{Value, ValueRef};
24use memcomparable::{Deserializer, Serializer};
25use serde::{Deserialize, Serialize};
26use snafu::{ResultExt, ensure};
27use store_api::codec::PrimaryKeyEncoding;
28use store_api::metadata::RegionMetadataRef;
29use store_api::storage::ColumnId;
30use store_api::storage::consts::ReservedColumnId;
31
32use crate::error::{
33    DeserializeFieldSnafu, InvalidSparsePrimaryKeySnafu, Result, SerializeFieldSnafu,
34    UnsupportedOperationSnafu,
35};
36use crate::key_values::KeyValue;
37use crate::primary_key_filter::SparsePrimaryKeyFilter;
38use crate::row_converter::dense::SortField;
39pub use crate::row_converter::sparse::checked::SparsePrimaryKeyView;
40use crate::row_converter::{CompositeValues, PrimaryKeyCodec, PrimaryKeyFilter};
41
42/// A codec for sparse key of metrics.
43///
44/// ## Encoding format
45/// Each primary key is encoded as a sequence of `(column_id, value)` pairs:
46/// - The first two fields are always the reserved `table_id` (uint32) and `tsid` (uint64).
47/// - User-defined labels follow, sorted by **column name** in lexicographical order.
48/// - Null values are omitted (not encoded).
49///
50/// The `column_id` is encoded as a 4-byte big-endian integer, and the value is encoded
51/// using memcomparable serialization.
52///
53/// `decode_leftmost` always decodes the first value from the encoded bytes (i.e., the
54/// `table_id` field).
55///
56/// ## Requirements
57/// It requires the input primary key columns are sorted by the column name in lexicographical order.
58/// It encodes the column id of the physical region.
59#[derive(Clone, Debug)]
60pub struct SparsePrimaryKeyCodec {
61    inner: Arc<SparsePrimaryKeyCodecInner>,
62}
63
64#[derive(Debug)]
65struct SparsePrimaryKeyCodecInner {
66    // Internal fields
67    table_id_field: SortField,
68    // Internal fields
69    tsid_field: SortField,
70    // User defined label field
71    label_field: SortField,
72    // Columns in primary key
73    //
74    // None means all unknown columns is primary key(`Self::label_field`).
75    columns: Option<HashSet<ColumnId>>,
76}
77
78/// Sparse values representation.
79///
80/// Callers must not insert a column id that is already present; otherwise
81/// the existing entry will shadow the newly inserted value on lookup.
82#[derive(Debug, Clone, PartialEq, Eq, Default)]
83pub struct SparseValues {
84    values: Vec<(ColumnId, Value)>,
85}
86
87impl SparseValues {
88    /// Creates an empty [`SparseValues`].
89    pub fn new() -> Self {
90        Self { values: Vec::new() }
91    }
92
93    /// Creates an empty [`SparseValues`] with space reserved for `cap` entries.
94    pub fn with_capacity(cap: usize) -> Self {
95        Self {
96            values: Vec::with_capacity(cap),
97        }
98    }
99
100    /// Returns the value of the given column, or [`Value::Null`] if the column is not present.
101    pub fn get_or_null(&self, column_id: ColumnId) -> &Value {
102        for (id, value) in &self.values {
103            if *id == column_id {
104                return value;
105            }
106        }
107        &Value::Null
108    }
109
110    /// Returns the value of the given column, or [`None`] if the column is not present.
111    pub fn get(&self, column_id: &ColumnId) -> Option<&Value> {
112        for (id, value) in &self.values {
113            if id == column_id {
114                return Some(value);
115            }
116        }
117        None
118    }
119
120    /// Appends a new `(column_id, value)` pair.
121    ///
122    /// Append-only: the caller must ensure `column_id` is not already present.
123    pub fn insert(&mut self, column_id: ColumnId, value: Value) {
124        self.values.push((column_id, value));
125    }
126
127    /// Returns an iterator over all stored column id/value pairs.
128    pub fn iter(&self) -> impl Iterator<Item = (&ColumnId, &Value)> {
129        self.values.iter().map(|(id, value)| (id, value))
130    }
131}
132
133/// The column id of the tsid.
134pub const RESERVED_COLUMN_ID_TSID: ColumnId = ReservedColumnId::tsid();
135/// The column id of the table id.
136pub const RESERVED_COLUMN_ID_TABLE_ID: ColumnId = ReservedColumnId::table_id();
137/// The size of the column id in the encoded sparse row.
138pub const COLUMN_ID_ENCODE_SIZE: usize = 4;
139
140// Fixed byte offsets for reserved columns in the sparse encoding.
141// Layout: [table_id_col_id: 4B][marker: 1B][table_id: 4B][tsid_col_id: 4B][marker: 1B][tsid: 8B]
142/// Byte offset to the table_id value (after its 4-byte column id).
143const TABLE_ID_VALUE_OFFSET: usize = COLUMN_ID_ENCODE_SIZE;
144/// Byte offset to the tsid value (after 9-byte table_id entry + 4-byte tsid column id).
145const TSID_VALUE_OFFSET: usize = COLUMN_ID_ENCODE_SIZE + 5 + COLUMN_ID_ENCODE_SIZE;
146/// Byte offset where tag columns start (after 9-byte table_id + 13-byte tsid entries).
147const TAGS_START_OFFSET: usize = COLUMN_ID_ENCODE_SIZE + 5 + COLUMN_ID_ENCODE_SIZE + 9;
148
149/// Inline capacity for the small-vec fast path of [`SparseOffsetsCache`].
150///
151/// Most sparse primary keys carry only a handful of tags, so a linear scan
152/// over a short `Vec` beats a `HashMap` lookup. Tags beyond this capacity
153/// spill into the overflow `HashMap`.
154const SPARSE_OFFSETS_INLINE_CAP: usize = 32;
155
156/// A lazily populated cache of tag column offsets inside a sparse primary key.
157#[derive(Debug, Clone)]
158pub struct SparseOffsetsCache {
159    /// Small-vec fast path. Reserves [`SPARSE_OFFSETS_INLINE_CAP`] slots on
160    /// the first insert.
161    inline: Vec<(ColumnId, usize)>,
162    /// Overflow for columns beyond the inline capacity. Lazily allocated.
163    overflow: HashMap<ColumnId, usize>,
164    /// Next byte position in the pk to resume parsing from.
165    cursor: usize,
166    /// True once the decoder has walked past the last tag column (or stopped
167    /// on an unknown column id); no further offsets can be discovered.
168    finished: bool,
169}
170
171impl Default for SparseOffsetsCache {
172    fn default() -> Self {
173        Self::new()
174    }
175}
176
177impl SparseOffsetsCache {
178    pub fn new() -> Self {
179        Self {
180            inline: Vec::new(),
181            overflow: HashMap::new(),
182            cursor: TAGS_START_OFFSET,
183            finished: false,
184        }
185    }
186
187    pub fn clear(&mut self) {
188        self.inline.clear();
189        self.overflow.clear();
190        self.cursor = TAGS_START_OFFSET;
191        self.finished = false;
192    }
193
194    /// Returns the cached offset for `column_id`, if any.
195    fn get(&self, column_id: ColumnId) -> Option<usize> {
196        for entry in &self.inline {
197            if entry.0 == column_id {
198                return Some(entry.1);
199            }
200        }
201        self.overflow.get(&column_id).copied()
202    }
203
204    /// Records a new `(column_id, offset)` entry.
205    fn insert(&mut self, column_id: ColumnId, offset: usize) {
206        if self.inline.len() < SPARSE_OFFSETS_INLINE_CAP {
207            if self.inline.capacity() == 0 {
208                self.inline.reserve_exact(SPARSE_OFFSETS_INLINE_CAP);
209            }
210            self.inline.push((column_id, offset));
211        } else {
212            self.overflow.insert(column_id, offset);
213        }
214    }
215
216    #[cfg(test)]
217    fn contains(&self, column_id: ColumnId) -> bool {
218        self.get(column_id).is_some()
219    }
220}
221
222impl SparsePrimaryKeyCodec {
223    /// Creates a new [`SparsePrimaryKeyCodec`] instance.
224    pub fn from_columns(columns_ids: impl Iterator<Item = ColumnId>) -> Self {
225        let columns = columns_ids.collect();
226        Self {
227            inner: Arc::new(SparsePrimaryKeyCodecInner {
228                table_id_field: SortField::new(ConcreteDataType::uint32_datatype()),
229                tsid_field: SortField::new(ConcreteDataType::uint64_datatype()),
230                label_field: SortField::new(ConcreteDataType::string_datatype()),
231                columns: Some(columns),
232            }),
233        }
234    }
235
236    /// Creates a new [`SparsePrimaryKeyCodec`] instance.
237    pub fn new(region_metadata: &RegionMetadataRef) -> Self {
238        Self::from_columns(region_metadata.primary_key_columns().map(|c| c.column_id))
239    }
240
241    /// Returns a new [`SparsePrimaryKeyCodec`] instance.
242    ///
243    /// It treats all unknown columns as primary key(label field).
244    pub fn schemaless() -> Self {
245        Self {
246            inner: Arc::new(SparsePrimaryKeyCodecInner {
247                table_id_field: SortField::new(ConcreteDataType::uint32_datatype()),
248                tsid_field: SortField::new(ConcreteDataType::uint64_datatype()),
249                label_field: SortField::new(ConcreteDataType::string_datatype()),
250                columns: None,
251            }),
252        }
253    }
254
255    /// Creates a new [`SparsePrimaryKeyCodec`] instance with additional label `fields`.
256    pub fn with_fields(fields: Vec<(ColumnId, SortField)>) -> Self {
257        Self {
258            inner: Arc::new(SparsePrimaryKeyCodecInner {
259                columns: Some(fields.iter().map(|f| f.0).collect()),
260                table_id_field: SortField::new(ConcreteDataType::uint32_datatype()),
261                tsid_field: SortField::new(ConcreteDataType::uint64_datatype()),
262                label_field: SortField::new(ConcreteDataType::string_datatype()),
263            }),
264        }
265    }
266
267    /// Returns the field of the given column id.
268    fn get_field(&self, column_id: ColumnId) -> Option<&SortField> {
269        // if the `columns` is not specified, all unknown columns is primary key(label field).
270        if let Some(columns) = &self.inner.columns
271            && !columns.contains(&column_id)
272        {
273            return None;
274        }
275
276        match column_id {
277            RESERVED_COLUMN_ID_TABLE_ID => Some(&self.inner.table_id_field),
278            RESERVED_COLUMN_ID_TSID => Some(&self.inner.tsid_field),
279            _ => Some(&self.inner.label_field),
280        }
281    }
282
283    /// Encodes the given bytes into a [`SparseValues`].
284    pub fn encode_to_vec<'a, I>(&self, row: I, buffer: &mut Vec<u8>) -> Result<()>
285    where
286        I: Iterator<Item = (ColumnId, ValueRef<'a>)>,
287    {
288        let mut serializer = Serializer::new(buffer);
289        for (column_id, value) in row {
290            if value.is_null() {
291                continue;
292            }
293
294            if let Some(field) = self.get_field(column_id) {
295                column_id
296                    .serialize(&mut serializer)
297                    .context(SerializeFieldSnafu)?;
298                field.serialize(&mut serializer, &value)?;
299            } else {
300                // TODO(weny): handle the error.
301                common_telemetry::warn!("Column {} is not in primary key, skipping", column_id);
302            }
303        }
304        Ok(())
305    }
306
307    pub fn encode_raw_tag_value<'a, I>(&self, row: I, buffer: &mut Vec<u8>) -> Result<()>
308    where
309        I: Iterator<Item = (ColumnId, &'a [u8])>,
310    {
311        for (tag_column_id, tag_value) in row {
312            let value_len = tag_value.len();
313            buffer.reserve(6 + value_len.div_ceil(8) * 9);
314            buffer.put_u32(tag_column_id);
315            buffer.put_u8(1);
316            buffer.put_u8(!tag_value.is_empty() as u8);
317
318            // Manual implementation of memcomparable::ser::Serializer::serialize_bytes
319            // to avoid byte-by-byte put.
320            let mut len = 0;
321            let num_chucks = value_len / 8;
322            let remainder = value_len % 8;
323
324            for idx in 0..num_chucks {
325                buffer.extend_from_slice(&tag_value[idx * 8..idx * 8 + 8]);
326                len += 8;
327                // append an extra byte that signals the number of significant bytes in this chunk
328                // 1-8: many bytes were significant and this group is the last group
329                // 9: all 8 bytes were significant and there is more data to come
330                let extra = if len == value_len { 8 } else { 9 };
331                buffer.put_u8(extra);
332            }
333
334            if remainder != 0 {
335                buffer.extend_from_slice(&tag_value[len..value_len]);
336                buffer.put_bytes(0, 8 - remainder);
337                buffer.put_u8(remainder as u8);
338            }
339        }
340        Ok(())
341    }
342
343    /// Encodes the given bytes into a [`SparseValues`].
344    pub fn encode_internal(&self, table_id: u32, tsid: u64, buffer: &mut Vec<u8>) -> Result<()> {
345        buffer.reserve_exact(22);
346        buffer.put_u32(RESERVED_COLUMN_ID_TABLE_ID);
347        buffer.put_u8(1);
348        buffer.put_u32(table_id);
349        buffer.put_u32(RESERVED_COLUMN_ID_TSID);
350        buffer.put_u8(1);
351        buffer.put_u64(tsid);
352        Ok(())
353    }
354
355    /// Decodes the reserved `(table_id, tsid)` prefix from a sparse primary key.
356    pub fn decode_ids(&self, bytes: &[u8]) -> Result<(u32, u64)> {
357        // Two column IDs, two non-null markers, a u32 table ID, and a u64 TSID.
358        const INTERNAL_PREFIX_LEN: usize = 4 + 1 + 4 + 4 + 1 + 8;
359        ensure!(
360            bytes.len() >= INTERNAL_PREFIX_LEN,
361            InvalidSparsePrimaryKeySnafu {
362                reason: format!(
363                    "internal prefix requires at least {INTERNAL_PREFIX_LEN} bytes, got {}",
364                    bytes.len()
365                ),
366            }
367        );
368
369        let mut deserializer = Deserializer::new(bytes);
370
371        let table_id_column = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?;
372        ensure!(
373            table_id_column == RESERVED_COLUMN_ID_TABLE_ID,
374            InvalidSparsePrimaryKeySnafu {
375                reason: format!(
376                    "expected table id column {}, got {}",
377                    RESERVED_COLUMN_ID_TABLE_ID, table_id_column
378                ),
379            }
380        );
381        let table_id = self.inner.table_id_field.deserialize(&mut deserializer)?;
382
383        let tsid_column = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?;
384        ensure!(
385            tsid_column == RESERVED_COLUMN_ID_TSID,
386            InvalidSparsePrimaryKeySnafu {
387                reason: format!(
388                    "expected tsid column {}, got {}",
389                    RESERVED_COLUMN_ID_TSID, tsid_column
390                ),
391            }
392        );
393        let tsid = self.inner.tsid_field.deserialize(&mut deserializer)?;
394
395        match (table_id, tsid) {
396            (Value::UInt32(table_id), Value::UInt64(tsid)) => Ok((table_id, tsid)),
397            (table_id, tsid) => InvalidSparsePrimaryKeySnafu {
398                reason: format!(
399                    "expected UInt32 table id and UInt64 tsid, got {table_id:?} and {tsid:?}"
400                ),
401            }
402            .fail(),
403        }
404    }
405
406    /// Decodes the given bytes into a [`SparseValues`].
407    fn decode_sparse(&self, bytes: &[u8]) -> Result<SparseValues> {
408        let mut deserializer = Deserializer::new(bytes);
409        let mut values = SparseValues::with_capacity(16);
410
411        let column_id = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?;
412        let value = self.inner.table_id_field.deserialize(&mut deserializer)?;
413        values.insert(column_id, value);
414
415        let column_id = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?;
416        let value = self.inner.tsid_field.deserialize(&mut deserializer)?;
417        values.insert(column_id, value);
418        while deserializer.has_remaining() {
419            let column_id = u32::deserialize(&mut deserializer).context(DeserializeFieldSnafu)?;
420            let value = self.inner.label_field.deserialize(&mut deserializer)?;
421            values.insert(column_id, value);
422        }
423
424        Ok(values)
425    }
426
427    /// Decodes the given bytes into a [`Value`].
428    fn decode_leftmost(&self, bytes: &[u8]) -> Result<Option<Value>> {
429        let mut deserializer = Deserializer::new(bytes);
430        // Skip the column id.
431        deserializer.advance(COLUMN_ID_ENCODE_SIZE);
432        let value = self.inner.table_id_field.deserialize(&mut deserializer)?;
433        Ok(Some(value))
434    }
435
436    /// Returns the offset of the given column id in the given primary key.
437    ///
438    /// The pk must start with the table_id + tsid prefix written by
439    /// `encode_internal`.
440    ///
441    /// # Panics
442    ///
443    /// Panics if `pk` is not a well-formed sparse primary key produced by
444    /// this codec (e.g. truncated or otherwise malformed bytes).
445    pub fn has_column(
446        &self,
447        pk: &[u8],
448        cache: &mut SparseOffsetsCache,
449        column_id: ColumnId,
450    ) -> Option<usize> {
451        // Decoding is lazy: on each call we only advance the cache's cursor as
452        // far as needed to answer the query. A column that has already been
453        // seen returns immediately; a column we haven't reached yet causes the
454        // parser to resume from `cache.cursor` and stop as soon as the column
455        // is located. Once the cursor walks off the end (or hits an unknown
456        // column id) the cache is marked finished, so subsequent misses are
457        // O(1).
458        // table_id and tsid are at fixed offsets.
459        match column_id {
460            RESERVED_COLUMN_ID_TABLE_ID => return Some(TABLE_ID_VALUE_OFFSET),
461            RESERVED_COLUMN_ID_TSID => return Some(TSID_VALUE_OFFSET),
462            _ => {}
463        }
464
465        if let Some(offset) = cache.get(column_id) {
466            return Some(offset);
467        }
468        if cache.finished {
469            return None;
470        }
471
472        let mut deserializer = Deserializer::new(pk);
473        deserializer.advance(cache.cursor);
474        let mut offset = cache.cursor;
475        while deserializer.has_remaining() {
476            let col = u32::deserialize(&mut deserializer).unwrap();
477            offset += COLUMN_ID_ENCODE_SIZE;
478            let value_offset = offset;
479            cache.insert(col, value_offset);
480            let Some(field) = self.get_field(col) else {
481                cache.finished = true;
482                cache.cursor = offset;
483                return None;
484            };
485
486            let skip = field.skip_deserialize(pk, &mut deserializer).unwrap();
487            offset += skip;
488            cache.cursor = offset;
489            if col == column_id {
490                return Some(value_offset);
491            }
492        }
493
494        cache.finished = true;
495        None
496    }
497
498    /// Decode value at `offset` in `pk`.
499    pub fn decode_value_at(&self, pk: &[u8], offset: usize, column_id: ColumnId) -> Result<Value> {
500        let mut deserializer = Deserializer::new(pk);
501        deserializer.advance(offset);
502        // Safety: checked by `has_column`
503        let field = self.get_field(column_id).unwrap();
504        field.deserialize(&mut deserializer)
505    }
506
507    /// Returns the encoded bytes of the given `column_id` in `pk`.
508    ///
509    /// Returns `Ok(None)` if the `column_id` is missing in `pk`.
510    pub fn encoded_value_for_column<'a>(
511        &self,
512        pk: &'a [u8],
513        cache: &mut SparseOffsetsCache,
514        column_id: ColumnId,
515    ) -> Result<Option<&'a [u8]>> {
516        let Some(offset) = self.has_column(pk, cache, column_id) else {
517            return Ok(None);
518        };
519
520        let Some(field) = self.get_field(column_id) else {
521            return Ok(None);
522        };
523
524        let mut deserializer = Deserializer::new(pk);
525        deserializer.advance(offset);
526        let len = field.skip_deserialize(pk, &mut deserializer)?;
527        Ok(Some(&pk[offset..offset + len]))
528    }
529}
530
531impl PrimaryKeyCodec for SparsePrimaryKeyCodec {
532    fn encode_key_value(&self, _key_value: &KeyValue, _buffer: &mut Vec<u8>) -> Result<()> {
533        UnsupportedOperationSnafu {
534            err_msg: "The encode_key_value method is not supported in SparsePrimaryKeyCodec.",
535        }
536        .fail()
537    }
538
539    fn encode_values(&self, values: &[(ColumnId, Value)], buffer: &mut Vec<u8>) -> Result<()> {
540        self.encode_to_vec(values.iter().map(|v| (v.0, v.1.as_value_ref())), buffer)
541    }
542
543    fn encode_value_refs(
544        &self,
545        values: &[(ColumnId, ValueRef)],
546        buffer: &mut Vec<u8>,
547    ) -> Result<()> {
548        self.encode_to_vec(values.iter().map(|v| (v.0, v.1.clone())), buffer)
549    }
550
551    fn estimated_size(&self) -> Option<usize> {
552        None
553    }
554
555    fn num_fields(&self) -> Option<usize> {
556        None
557    }
558
559    fn encoding(&self) -> PrimaryKeyEncoding {
560        PrimaryKeyEncoding::Sparse
561    }
562
563    fn primary_key_filter(
564        &self,
565        metadata: &RegionMetadataRef,
566        filters: Arc<Vec<SimpleFilterEvaluator>>,
567    ) -> Box<dyn PrimaryKeyFilter> {
568        Box::new(SparsePrimaryKeyFilter::new(
569            metadata.clone(),
570            filters,
571            self.clone(),
572        ))
573    }
574
575    fn decode(&self, bytes: &[u8]) -> Result<CompositeValues> {
576        Ok(CompositeValues::Sparse(self.decode_sparse(bytes)?))
577    }
578
579    fn decode_leftmost(&self, bytes: &[u8]) -> Result<Option<Value>> {
580        self.decode_leftmost(bytes)
581    }
582}
583
584/// Field with column id.
585pub struct FieldWithId {
586    pub field: SortField,
587    pub column_id: ColumnId,
588}
589
590/// A special encoder for memtable.
591pub struct SparseEncoder {
592    fields: Vec<FieldWithId>,
593}
594
595impl SparseEncoder {
596    pub fn new(fields: Vec<FieldWithId>) -> Self {
597        Self { fields }
598    }
599
600    pub fn encode_to_vec<'a, I>(&self, row: I, buffer: &mut Vec<u8>) -> Result<()>
601    where
602        I: Iterator<Item = ValueRef<'a>>,
603    {
604        let mut serializer = Serializer::new(buffer);
605        for (value, field) in row.zip(self.fields.iter()) {
606            if !value.is_null() {
607                field
608                    .column_id
609                    .serialize(&mut serializer)
610                    .context(SerializeFieldSnafu)?;
611                field.field.serialize(&mut serializer, &value)?;
612            }
613        }
614        Ok(())
615    }
616}
617
618#[cfg(test)]
619mod tests {
620    use std::sync::Arc;
621
622    use api::v1::SemanticType;
623    use common_query::prelude::{greptime_timestamp, greptime_value};
624    use common_time::Timestamp;
625    use common_time::timestamp::TimeUnit;
626    use datatypes::schema::ColumnSchema;
627    use datatypes::value::{OrderedFloat, Value};
628    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
629    use store_api::metric_engine_consts::{
630        DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME,
631    };
632    use store_api::storage::{ColumnId, RegionId};
633
634    use super::*;
635
636    fn test_region_metadata() -> RegionMetadataRef {
637        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1));
638        builder
639            .push_column_metadata(ColumnMetadata {
640                column_schema: ColumnSchema::new(
641                    DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
642                    ConcreteDataType::uint32_datatype(),
643                    false,
644                ),
645                semantic_type: SemanticType::Tag,
646                column_id: ReservedColumnId::table_id(),
647            })
648            .push_column_metadata(ColumnMetadata {
649                column_schema: ColumnSchema::new(
650                    DATA_SCHEMA_TSID_COLUMN_NAME,
651                    ConcreteDataType::uint64_datatype(),
652                    false,
653                ),
654                semantic_type: SemanticType::Tag,
655                column_id: ReservedColumnId::tsid(),
656            })
657            .push_column_metadata(ColumnMetadata {
658                column_schema: ColumnSchema::new("pod", ConcreteDataType::string_datatype(), true),
659                semantic_type: SemanticType::Tag,
660                column_id: 1,
661            })
662            .push_column_metadata(ColumnMetadata {
663                column_schema: ColumnSchema::new(
664                    "namespace",
665                    ConcreteDataType::string_datatype(),
666                    true,
667                ),
668                semantic_type: SemanticType::Tag,
669                column_id: 2,
670            })
671            .push_column_metadata(ColumnMetadata {
672                column_schema: ColumnSchema::new(
673                    "container",
674                    ConcreteDataType::string_datatype(),
675                    true,
676                ),
677                semantic_type: SemanticType::Tag,
678                column_id: 3,
679            })
680            .push_column_metadata(ColumnMetadata {
681                column_schema: ColumnSchema::new(
682                    "pod_name",
683                    ConcreteDataType::string_datatype(),
684                    true,
685                ),
686                semantic_type: SemanticType::Tag,
687                column_id: 4,
688            })
689            .push_column_metadata(ColumnMetadata {
690                column_schema: ColumnSchema::new(
691                    "pod_ip",
692                    ConcreteDataType::string_datatype(),
693                    true,
694                ),
695                semantic_type: SemanticType::Tag,
696                column_id: 5,
697            })
698            .push_column_metadata(ColumnMetadata {
699                column_schema: ColumnSchema::new(
700                    greptime_value(),
701                    ConcreteDataType::float64_datatype(),
702                    false,
703                ),
704                semantic_type: SemanticType::Field,
705                column_id: 6,
706            })
707            .push_column_metadata(ColumnMetadata {
708                column_schema: ColumnSchema::new(
709                    greptime_timestamp(),
710                    ConcreteDataType::timestamp_nanosecond_datatype(),
711                    false,
712                ),
713                semantic_type: SemanticType::Timestamp,
714                column_id: 7,
715            })
716            .primary_key(vec![
717                ReservedColumnId::table_id(),
718                ReservedColumnId::tsid(),
719                1,
720                2,
721                3,
722                4,
723                5,
724            ]);
725        let metadata = builder.build().unwrap();
726        Arc::new(metadata)
727    }
728
729    #[test]
730    fn test_sparse_value_new_and_get_or_null() {
731        let mut sparse_value = SparseValues::new();
732        sparse_value.insert(1, Value::Int32(42));
733
734        assert_eq!(sparse_value.get_or_null(1), &Value::Int32(42));
735        assert_eq!(sparse_value.get_or_null(2), &Value::Null);
736    }
737
738    #[test]
739    fn test_sparse_value_insert() {
740        let mut sparse_value = SparseValues::new();
741        sparse_value.insert(1, Value::Int32(42));
742
743        assert_eq!(sparse_value.get_or_null(1), &Value::Int32(42));
744    }
745
746    fn test_row() -> Vec<(ColumnId, ValueRef<'static>)> {
747        vec![
748            (RESERVED_COLUMN_ID_TABLE_ID, ValueRef::UInt32(42)),
749            (
750                RESERVED_COLUMN_ID_TSID,
751                ValueRef::UInt64(123843349035232323),
752            ),
753            // label: pod
754            (1, ValueRef::String("greptime-frontend-6989d9899-22222")),
755            // label: namespace
756            (2, ValueRef::String("greptime-cluster")),
757            // label: container
758            (3, ValueRef::String("greptime-frontend-6989d9899-22222")),
759            // label: pod_name
760            (4, ValueRef::String("greptime-frontend-6989d9899-22222")),
761            // label: pod_ip
762            (5, ValueRef::String("10.10.10.10")),
763            // field: greptime_value
764            (6, ValueRef::Float64(OrderedFloat(1.0))),
765            // field: greptime_timestamp
766            (
767                7,
768                ValueRef::Timestamp(Timestamp::new(1618876800000000000, TimeUnit::Nanosecond)),
769            ),
770        ]
771    }
772
773    #[test]
774    fn test_encode_by_short_cuts() {
775        let region_metadata = test_region_metadata();
776        let codec = SparsePrimaryKeyCodec::new(&region_metadata);
777        let mut buffer = Vec::new();
778        let internal_columns = [
779            (RESERVED_COLUMN_ID_TABLE_ID, ValueRef::UInt32(1024)),
780            (RESERVED_COLUMN_ID_TSID, ValueRef::UInt64(42)),
781        ];
782        let tags = [
783            (1, "greptime-frontend-6989d9899-22222"),
784            (2, "greptime-cluster"),
785            (3, "greptime-frontend-6989d9899-22222"),
786            (4, "greptime-frontend-6989d9899-22222"),
787            (5, "10.10.10.10"),
788        ];
789        codec
790            .encode_to_vec(internal_columns.into_iter(), &mut buffer)
791            .unwrap();
792        codec
793            .encode_to_vec(
794                tags.iter()
795                    .map(|(col_id, tag_value)| (*col_id, ValueRef::String(tag_value))),
796                &mut buffer,
797            )
798            .unwrap();
799
800        let mut buffer_by_raw_encoding = Vec::new();
801        codec
802            .encode_internal(1024, 42, &mut buffer_by_raw_encoding)
803            .unwrap();
804        let tags: Vec<_> = tags
805            .into_iter()
806            .map(|(col_id, tag_value)| (col_id, tag_value.as_bytes()))
807            .collect();
808        codec
809            .encode_raw_tag_value(
810                tags.iter().map(|(c, b)| (*c, *b)),
811                &mut buffer_by_raw_encoding,
812            )
813            .unwrap();
814        assert_eq!(buffer, buffer_by_raw_encoding);
815    }
816
817    #[test]
818    fn raw_encoding_matches_serde_at_chunk_boundaries() {
819        let codec = SparsePrimaryKeyCodec::schemaless();
820        let mut actual = Vec::new();
821        for len in [0, 7, 8, 9, 15, 16, 17, 1024] {
822            let label = "x".repeat(len);
823            let mut expected = Vec::new();
824            codec.encode_internal(42, 7, &mut expected).unwrap();
825            codec
826                .encode_to_vec([(1, ValueRef::String(&label))].into_iter(), &mut expected)
827                .unwrap();
828            actual.clear();
829            codec.encode_internal(42, 7, &mut actual).unwrap();
830            codec
831                .encode_raw_tag_value([(1, label.as_bytes())].into_iter(), &mut actual)
832                .unwrap();
833            assert_eq!(actual, expected, "label length {len}");
834        }
835    }
836
837    #[test]
838    fn test_encode_to_vec() {
839        let region_metadata = test_region_metadata();
840        let codec = SparsePrimaryKeyCodec::new(&region_metadata);
841        let mut buffer = Vec::new();
842
843        let row = test_row();
844        codec.encode_to_vec(row.into_iter(), &mut buffer).unwrap();
845        assert!(!buffer.is_empty());
846        let sparse_value = codec.decode_sparse(&buffer).unwrap();
847        assert_eq!(
848            sparse_value.get_or_null(RESERVED_COLUMN_ID_TABLE_ID),
849            &Value::UInt32(42)
850        );
851        assert_eq!(
852            sparse_value.get_or_null(1),
853            &Value::String("greptime-frontend-6989d9899-22222".into())
854        );
855        assert_eq!(
856            sparse_value.get_or_null(2),
857            &Value::String("greptime-cluster".into())
858        );
859        assert_eq!(
860            sparse_value.get_or_null(3),
861            &Value::String("greptime-frontend-6989d9899-22222".into())
862        );
863        assert_eq!(
864            sparse_value.get_or_null(4),
865            &Value::String("greptime-frontend-6989d9899-22222".into())
866        );
867        assert_eq!(
868            sparse_value.get_or_null(5),
869            &Value::String("10.10.10.10".into())
870        );
871    }
872
873    #[test]
874    fn test_decode_leftmost() {
875        let region_metadata = test_region_metadata();
876        let codec = SparsePrimaryKeyCodec::new(&region_metadata);
877        let mut buffer = Vec::new();
878        let row = test_row();
879        codec.encode_to_vec(row.into_iter(), &mut buffer).unwrap();
880        assert!(!buffer.is_empty());
881        let result = codec.decode_leftmost(&buffer).unwrap().unwrap();
882        assert_eq!(result, Value::UInt32(42));
883    }
884
885    #[test]
886    fn test_decode_ids() {
887        let region_metadata = test_region_metadata();
888        let codec = SparsePrimaryKeyCodec::new(&region_metadata);
889        let mut buffer = Vec::new();
890        codec.encode_internal(42, 100, &mut buffer).unwrap();
891
892        assert_eq!((42, 100), codec.decode_ids(&buffer).unwrap());
893
894        let mut invalid = buffer.clone();
895        invalid[0..4].copy_from_slice(&1_u32.to_be_bytes());
896        assert!(codec.decode_ids(&invalid).is_err());
897        assert!(codec.decode_ids(&buffer[..buffer.len() - 1]).is_err());
898    }
899
900    #[test]
901    fn test_has_column() {
902        let region_metadata = test_region_metadata();
903        let codec = SparsePrimaryKeyCodec::new(&region_metadata);
904        let mut buffer = Vec::new();
905        let row = test_row();
906        codec.encode_to_vec(row.into_iter(), &mut buffer).unwrap();
907        assert!(!buffer.is_empty());
908
909        let mut offsets_map = SparseOffsetsCache::new();
910        for column_id in [
911            RESERVED_COLUMN_ID_TABLE_ID,
912            RESERVED_COLUMN_ID_TSID,
913            1,
914            2,
915            3,
916            4,
917            5,
918        ] {
919            let offset = codec.has_column(&buffer, &mut offsets_map, column_id);
920            assert!(offset.is_some());
921        }
922
923        let offset = codec.has_column(&buffer, &mut offsets_map, 6);
924        assert!(offset.is_none());
925    }
926
927    #[test]
928    fn test_has_column_lazy_resume() {
929        let region_metadata = test_region_metadata();
930        let codec = SparsePrimaryKeyCodec::new(&region_metadata);
931        let mut buffer = Vec::new();
932        codec
933            .encode_to_vec(test_row().into_iter(), &mut buffer)
934            .unwrap();
935
936        let mut cache = SparseOffsetsCache::new();
937        // Look up an early column: only a prefix of tags is decoded.
938        assert!(codec.has_column(&buffer, &mut cache, 1).is_some());
939        assert!(!cache.finished);
940        assert!(cache.contains(1));
941        assert!(!cache.contains(5));
942
943        // A later column resumes from the cursor.
944        assert!(codec.has_column(&buffer, &mut cache, 5).is_some());
945        assert!(cache.contains(5));
946
947        // An earlier column that was already cached still resolves.
948        assert!(codec.has_column(&buffer, &mut cache, 2).is_some());
949
950        // A non-existent column walks off the end and marks the cache finished.
951        assert!(codec.has_column(&buffer, &mut cache, 999).is_none());
952        assert!(cache.finished);
953        // Further misses are O(1).
954        assert!(codec.has_column(&buffer, &mut cache, 998).is_none());
955    }
956
957    #[test]
958    fn test_decode_value_at() {
959        let region_metadata = test_region_metadata();
960        let codec = SparsePrimaryKeyCodec::new(&region_metadata);
961        let mut buffer = Vec::new();
962        let row = test_row();
963        codec.encode_to_vec(row.into_iter(), &mut buffer).unwrap();
964        assert!(!buffer.is_empty());
965
966        let row = test_row();
967        let mut offsets_map = SparseOffsetsCache::new();
968        for column_id in [
969            RESERVED_COLUMN_ID_TABLE_ID,
970            RESERVED_COLUMN_ID_TSID,
971            1,
972            2,
973            3,
974            4,
975            5,
976        ] {
977            let offset = codec
978                .has_column(&buffer, &mut offsets_map, column_id)
979                .unwrap();
980            let value = codec.decode_value_at(&buffer, offset, column_id).unwrap();
981            let expected_value = row
982                .iter()
983                .find(|(id, _)| *id == column_id)
984                .unwrap()
985                .1
986                .clone();
987            assert_eq!(value.as_value_ref(), expected_value);
988        }
989    }
990
991    #[test]
992    fn test_encoded_value_for_column() {
993        let region_metadata = test_region_metadata();
994        let codec = SparsePrimaryKeyCodec::new(&region_metadata);
995        let mut buffer = Vec::new();
996        let row = test_row();
997        codec
998            .encode_to_vec(row.clone().into_iter(), &mut buffer)
999            .unwrap();
1000        assert!(!buffer.is_empty());
1001
1002        let mut offsets_map = SparseOffsetsCache::new();
1003        for column_id in [
1004            RESERVED_COLUMN_ID_TABLE_ID,
1005            RESERVED_COLUMN_ID_TSID,
1006            1,
1007            2,
1008            3,
1009            4,
1010            5,
1011        ] {
1012            let encoded_value = codec
1013                .encoded_value_for_column(&buffer, &mut offsets_map, column_id)
1014                .unwrap()
1015                .unwrap();
1016            let expected_value = row
1017                .iter()
1018                .find(|(id, _)| *id == column_id)
1019                .unwrap()
1020                .1
1021                .clone();
1022            let data_type = match column_id {
1023                RESERVED_COLUMN_ID_TABLE_ID => ConcreteDataType::uint32_datatype(),
1024                RESERVED_COLUMN_ID_TSID => ConcreteDataType::uint64_datatype(),
1025                _ => ConcreteDataType::string_datatype(),
1026            };
1027            let field = SortField::new(data_type);
1028            let mut expected_encoded = Vec::new();
1029            let mut serializer = Serializer::new(&mut expected_encoded);
1030            field.serialize(&mut serializer, &expected_value).unwrap();
1031            assert_eq!(encoded_value, expected_encoded.as_slice());
1032        }
1033
1034        for column_id in [6_u32, 7_u32, 999_u32] {
1035            let encoded_value = codec
1036                .encoded_value_for_column(&buffer, &mut offsets_map, column_id)
1037                .unwrap();
1038            assert!(encoded_value.is_none());
1039        }
1040    }
1041}
mito_codec/row_converter/sparse.rs

mito_codec/row_converter/
sparse.rs
