mito_codec/row_converter/sparse/
checked.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 bytes::Buf;
16use snafu::{OptionExt, ensure};
17use store_api::storage::ColumnId;
18
19use crate::error::{InvalidSparsePrimaryKeySnafu, Result};
20use crate::row_converter::encoded_string_len;
21use crate::row_converter::sparse::{
22    COLUMN_ID_ENCODE_SIZE, RESERVED_COLUMN_ID_TABLE_ID, RESERVED_COLUMN_ID_TSID,
23    SparseOffsetsCache, TABLE_ID_VALUE_OFFSET, TAGS_START_OFFSET, TSID_VALUE_OFFSET,
24    validate_prefix,
25};
26
27/// A borrowing view of one sparse primary key, with checked, lazy field lookup.
28///
29/// The scratch cache is reset on construction and retains its capacity between keys.
30/// Ordinary labels can be skipped without schema lookup: their encoding is always String.
31pub struct SparsePrimaryKeyView<'a, 'b> {
32    pk: &'a [u8],
33    cache: &'b mut SparseOffsetsCache,
34}
35
36impl<'a, 'b> SparsePrimaryKeyView<'a, 'b> {
37    /// Validates the fixed prefix and binds fresh offset discovery to this key.
38    pub fn new(pk: &'a [u8], cache: &'b mut SparseOffsetsCache) -> Result<Self> {
39        validate_prefix(pk)?;
40        cache.clear();
41        Ok(Self { pk, cache })
42    }
43
44    /// Returns the encoded value, including its null marker, or None for an absent label.
45    /// The result borrows the key, not the scratch cache.
46    pub fn encoded_value(&mut self, column_id: ColumnId) -> Result<Option<&'a [u8]>> {
47        match column_id {
48            RESERVED_COLUMN_ID_TABLE_ID => {
49                return Ok(Some(
50                    &self.pk[TABLE_ID_VALUE_OFFSET..TSID_VALUE_OFFSET - COLUMN_ID_ENCODE_SIZE],
51                ));
52            }
53            RESERVED_COLUMN_ID_TSID => {
54                return Ok(Some(&self.pk[TSID_VALUE_OFFSET..TAGS_START_OFFSET]));
55            }
56            _ => {}
57        }
58        if let Some(offset) = self.cache.get(column_id) {
59            return encoded_label(&self.pk[offset..]).map(Some);
60        }
61        while !self.cache.finished && self.cache.cursor < self.pk.len() {
62            let bytes = &self.pk[self.cache.cursor..];
63            let id = bytes.first_chunk::<COLUMN_ID_ENCODE_SIZE>().context(
64                InvalidSparsePrimaryKeySnafu {
65                    reason: "truncated label column id",
66                },
67            )?;
68            let id = u32::from_be_bytes(*id);
69            let offset = self.cache.cursor + COLUMN_ID_ENCODE_SIZE;
70            let value = encoded_label(&self.pk[offset..])?;
71            self.cache.insert(id, offset);
72            self.cache.cursor = offset + value.len();
73            if id == column_id {
74                return Ok(Some(value));
75            }
76        }
77        self.cache.finished = true;
78        Ok(None)
79    }
80
81    /// Extracts a UTF-8 label into a reusable buffer, validating it once.
82    /// Missing and null labels return `None` without modifying the buffer;
83    /// an empty label returns `Some("")`. Reserved numeric columns are rejected.
84    /// The returned string borrows `buffer`, whose contents are replaced for each label.
85    pub fn label<'buf>(
86        &mut self,
87        column_id: ColumnId,
88        buffer: &'buf mut Vec<u8>,
89    ) -> Result<Option<&'buf str>> {
90        ensure!(
91            !matches!(
92                column_id,
93                RESERVED_COLUMN_ID_TABLE_ID | RESERVED_COLUMN_ID_TSID
94            ),
95            InvalidSparsePrimaryKeySnafu {
96                reason: "reserved column is not a string label"
97            }
98        );
99        let Some(encoded) = self.encoded_value(column_id)? else {
100            return Ok(None);
101        };
102        if encoded[0] == 0 {
103            return Ok(None);
104        }
105
106        buffer.clear();
107        // encoded_value checked the markers and every chunk length. Reserve once
108        // before unchunking so the buffer does not grow for each 8-byte chunk.
109        buffer.reserve(encoded.len() - 2);
110        for chunk in encoded[2..].as_chunks::<9>().0 {
111            let len = usize::from(chunk[8]).min(8);
112            buffer.extend_from_slice(&chunk[..len]);
113        }
114        std::str::from_utf8(buffer).map(Some).map_err(|_| {
115            InvalidSparsePrimaryKeySnafu {
116                reason: "label is not valid UTF-8",
117            }
118            .build()
119        })
120    }
121
122    /// Returns the table id from the non-null prefix validated by [`Self::new`].
123    pub fn table_id(&self) -> u32 {
124        (&self.pk[TABLE_ID_VALUE_OFFSET + 1..]).get_u32()
125    }
126
127    /// Returns the series id from the non-null prefix validated by [`Self::new`].
128    pub fn tsid(&self) -> u64 {
129        (&self.pk[TSID_VALUE_OFFSET + 1..]).get_u64()
130    }
131}
132
133/// Finds the end of an Option<String> without allocating or reading its payload.
134/// Unlike memcomparable's skip_bytes, all advances are checked for truncated input.
135fn encoded_label(bytes: &[u8]) -> Result<&[u8]> {
136    let len = encoded_string_len(bytes).map_err(|error| {
137        InvalidSparsePrimaryKeySnafu {
138            reason: error.to_string(),
139        }
140        .build()
141    })?;
142    Ok(&bytes[..len])
143}
144
145#[cfg(test)]
146mod tests {
147    use datatypes::data_type::ConcreteDataType;
148    use datatypes::value::ValueRef;
149
150    use super::*;
151    use crate::index::IndexValueCodec;
152    use crate::row_converter::{PrimaryKeyCodec, SortField, SparsePrimaryKeyCodec};
153
154    #[test]
155    fn index_bytes_match_full_decode_across_chunks_and_lookup_orders() {
156        let codec = SparsePrimaryKeyCodec::schemaless();
157        let mut cache = SparseOffsetsCache::new();
158        let mut buffer = Vec::new();
159        // Reusing the cache across different keys must not reuse their field offsets.
160        for len in [0, 1, 7, 8, 9, 16, 24, 65] {
161            let label = "x".repeat(len);
162            let values = [
163                (RESERVED_COLUMN_ID_TABLE_ID, ValueRef::UInt32(42)),
164                (RESERVED_COLUMN_ID_TSID, ValueRef::UInt64(u64::MAX)),
165                (90, ValueRef::String(&label)),
166                (7, ValueRef::String("")),
167                (1, ValueRef::Null),
168                (3, ValueRef::String("last中文")),
169            ];
170            let mut pk = Vec::new();
171            codec.encode_value_refs(&values, &mut pk).unwrap();
172            // Exercise the offset cache's overflow storage as well as its inline entries.
173            codec
174                .encode_raw_tag_value((100..140).map(|id| (id, label.as_bytes())), &mut pk)
175                .unwrap();
176            // Writers normally omit null labels; existing encoded nulls must remain null too.
177            pk.extend_from_slice(&2_u32.to_be_bytes());
178            pk.push(0);
179            let decoded = codec.decode(&pk).unwrap().into_sparse();
180            for order in [
181                [
182                    RESERVED_COLUMN_ID_TABLE_ID,
183                    RESERVED_COLUMN_ID_TSID,
184                    90,
185                    7,
186                    1,
187                    3,
188                ],
189                [
190                    3,
191                    1,
192                    7,
193                    90,
194                    RESERVED_COLUMN_ID_TSID,
195                    RESERVED_COLUMN_ID_TABLE_ID,
196                ],
197            ] {
198                let mut view = SparsePrimaryKeyView::new(&pk, &mut cache).unwrap();
199                for id in order.into_iter().chain([139, 100, 132, 2]) {
200                    let expected = decoded.get(&id).filter(|v| !v.is_null()).map(|value| {
201                        let field = SortField::new(match id {
202                            RESERVED_COLUMN_ID_TABLE_ID => ConcreteDataType::uint32_datatype(),
203                            RESERVED_COLUMN_ID_TSID => ConcreteDataType::uint64_datatype(),
204                            _ => ConcreteDataType::string_datatype(),
205                        });
206                        let mut bytes = Vec::new();
207                        IndexValueCodec::encode_nonnull_value(
208                            value.as_value_ref(),
209                            &field,
210                            &mut bytes,
211                        )
212                        .unwrap();
213                        bytes
214                    });
215                    let actual =
216                        IndexValueCodec::encode_sparse_value(&mut view, id, &mut buffer).unwrap();
217                    assert_eq!(actual, expected.as_deref(), "len={len}, column={id}");
218                }
219                assert!(view.encoded_value(999).unwrap().is_none());
220            }
221        }
222    }
223
224    #[test]
225    fn reserved_ids_preserve_unsigned_values() {
226        let codec = SparsePrimaryKeyCodec::schemaless();
227        let mut cache = SparseOffsetsCache::new();
228        for (table_id, tsid) in [(0, 0), (42, u64::MAX), (u32::MAX, 42)] {
229            let mut pk = Vec::new();
230            codec.encode_internal(table_id, tsid, &mut pk).unwrap();
231            let view = SparsePrimaryKeyView::new(&pk, &mut cache).unwrap();
232            assert_eq!(view.table_id(), table_id);
233            assert_eq!(view.tsid(), tsid);
234        }
235    }
236
237    #[test]
238    fn malformed_sparse_values_return_errors() {
239        let codec = SparsePrimaryKeyCodec::schemaless();
240        let mut pk = Vec::new();
241        codec.encode_internal(1, 2, &mut pk).unwrap();
242        codec
243            .encode_raw_tag_value([(1, b"0123456789".as_slice())].into_iter(), &mut pk)
244            .unwrap();
245        let mut cache = SparseOffsetsCache::new();
246        for end in 0..pk.len() {
247            // A complete reserved prefix is a valid key with no labels.
248            if end == TAGS_START_OFFSET {
249                continue;
250            }
251            let result = SparsePrimaryKeyView::new(&pk[..end], &mut cache)
252                .and_then(|mut view| view.encoded_value(1));
253            assert!(result.is_err(), "truncation at {end}");
254        }
255        for (offset, value) in [(0, 0), (4, 0), (13, 2), (26, 2), (27, 2), (36, 0), (45, 10)] {
256            let mut invalid = pk.clone();
257            invalid[offset] = value;
258            assert!(
259                SparsePrimaryKeyView::new(&invalid, &mut cache)
260                    .and_then(|mut view| view.encoded_value(1))
261                    .is_err()
262            );
263        }
264        let mut invalid_utf8 = pk;
265        invalid_utf8[28] = 0xff;
266        let mut view = SparsePrimaryKeyView::new(&invalid_utf8, &mut cache).unwrap();
267        assert!(IndexValueCodec::encode_sparse_value(&mut view, 1, &mut Vec::new()).is_err());
268    }
269
270    #[test]
271    fn label_reuses_buffer_across_null_empty_invalid_and_valid_values() {
272        let codec = SparsePrimaryKeyCodec::schemaless();
273        let mut cache = SparseOffsetsCache::new();
274        let mut buffer = Vec::with_capacity(256);
275        let ptr = buffer.as_ptr();
276        for label in [
277            b"1234567\xe4\xb8\xad\0\xe6\x96\x87".as_slice(),
278            b"",
279            b"\xff",
280            b"after-error",
281        ] {
282            let mut pk = Vec::new();
283            codec.encode_internal(42, 7, &mut pk).unwrap();
284            codec
285                .encode_raw_tag_value([(1, label)].into_iter(), &mut pk)
286                .unwrap();
287            pk.extend_from_slice(&2_u32.to_be_bytes());
288            pk.push(0);
289            let mut view = SparsePrimaryKeyView::new(&pk, &mut cache).unwrap();
290            let actual = view.label(1, &mut buffer);
291            match std::str::from_utf8(label) {
292                Ok(expected) => assert_eq!(actual.unwrap(), Some(expected)),
293                Err(_) => assert!(actual.is_err()),
294            }
295            assert_eq!(view.label(2, &mut buffer).unwrap(), None);
296            assert_eq!(view.label(3, &mut buffer).unwrap(), None);
297            // Null/missing values do not overwrite the scratch buffer, matching index extraction.
298            assert_eq!(buffer, label);
299            assert_eq!(buffer.as_ptr(), ptr);
300            assert!(
301                view.label(RESERVED_COLUMN_ID_TABLE_ID, &mut buffer)
302                    .is_err()
303            );
304            assert!(view.label(RESERVED_COLUMN_ID_TSID, &mut buffer).is_err());
305        }
306    }
307}
mito_codec/row_converter/sparse/checked.rs

mito_codec/row_converter/sparse/
checked.rs
