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

mito_codec/row_converter/sparse/
checked.rs
