mito_codec/row_converter/
dense.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::sync::Arc;
16
17use bytes::Buf;
18use common_base::bytes::Bytes;
19use common_decimal::Decimal128;
20use common_recordbatch::filter::SimpleFilterEvaluator;
21use common_time::time::Time;
22use common_time::{Date, Duration, IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth};
23use datatypes::data_type::ConcreteDataType;
24use datatypes::prelude::Value;
25use datatypes::types::IntervalType;
26use datatypes::value::ValueRef;
27use memcomparable::{Deserializer, Serializer};
28use paste::paste;
29use serde::{Deserialize, Serialize};
30use snafu::{ResultExt, ensure};
31use store_api::codec::PrimaryKeyEncoding;
32use store_api::metadata::{RegionMetadata, RegionMetadataRef};
33use store_api::storage::ColumnId;
34
35use crate::error::{
36    self, FieldTypeMismatchSnafu, NotSupportedFieldSnafu, Result, SerializeFieldSnafu,
37};
38use crate::key_values::KeyValue;
39use crate::primary_key_filter::DensePrimaryKeyFilter;
40use crate::row_converter::{
41    CompositeValues, PrimaryKeyCodec, PrimaryKeyCodecExt, PrimaryKeyFilter, encoded_string_len,
42};
43
44/// Field to serialize and deserialize value in memcomparable format.
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct SortField {
47    data_type: ConcreteDataType,
48}
49
50impl SortField {
51    pub fn new(data_type: ConcreteDataType) -> Self {
52        Self { data_type }
53    }
54
55    /// Returns the data type of the field.
56    pub fn data_type(&self) -> &ConcreteDataType {
57        &self.data_type
58    }
59
60    /// Returns the physical data type to encode of the field.
61    ///
62    /// For example, a dictionary field will be encoded as its value type.
63    pub fn encode_data_type(&self) -> &ConcreteDataType {
64        match &self.data_type {
65            ConcreteDataType::Dictionary(dict_type) => dict_type.value_type(),
66            _ => &self.data_type,
67        }
68    }
69
70    pub fn estimated_size(&self) -> usize {
71        Self::estimated_size_by_type(self.encode_data_type())
72    }
73
74    fn estimated_size_by_type(data_type: &ConcreteDataType) -> usize {
75        match data_type {
76            ConcreteDataType::Boolean(_) => 2,
77            ConcreteDataType::Int8(_) | ConcreteDataType::UInt8(_) => 2,
78            ConcreteDataType::Int16(_) | ConcreteDataType::UInt16(_) => 3,
79            ConcreteDataType::Int32(_) | ConcreteDataType::UInt32(_) => 5,
80            ConcreteDataType::Int64(_) | ConcreteDataType::UInt64(_) => 9,
81            ConcreteDataType::Float32(_) => 5,
82            ConcreteDataType::Float64(_) => 9,
83            ConcreteDataType::Binary(_)
84            | ConcreteDataType::Json(_)
85            | ConcreteDataType::Vector(_) => 11,
86            ConcreteDataType::String(_) => 11, // a non-empty string takes at least 11 bytes.
87            ConcreteDataType::Date(_) => 5,
88            ConcreteDataType::Timestamp(_) => 10,
89            ConcreteDataType::Time(_) => 10,
90            ConcreteDataType::Duration(_) => 10,
91            ConcreteDataType::Interval(_) => 18,
92            ConcreteDataType::Decimal128(_) => 19,
93            ConcreteDataType::Null(_)
94            | ConcreteDataType::List(_)
95            | ConcreteDataType::Struct(_)
96            | ConcreteDataType::Dictionary(_) => 0,
97        }
98    }
99
100    /// Serialize a value to the serializer.
101    pub fn serialize(
102        &self,
103        serializer: &mut Serializer<&mut Vec<u8>>,
104        value: &ValueRef,
105    ) -> Result<()> {
106        Self::serialize_by_type(self.encode_data_type(), serializer, value)
107    }
108
109    fn serialize_by_type(
110        data_type: &ConcreteDataType,
111        serializer: &mut Serializer<&mut Vec<u8>>,
112        value: &ValueRef,
113    ) -> Result<()> {
114        macro_rules! cast_value_and_serialize {
115            (
116                $data_type: ident;
117                $serializer: ident;
118                $(
119                    $ty: ident, $f: ident
120                ),*
121            ) => {
122                match $data_type {
123                $(
124                    ConcreteDataType::$ty(_) => {
125                        paste!{
126                            value
127                            .[<try_into_ $f>]()
128                            .context(FieldTypeMismatchSnafu)?
129                            .serialize($serializer)
130                            .context(SerializeFieldSnafu)?;
131                        }
132                    }
133                )*
134                    ConcreteDataType::Timestamp(_) => {
135                        let timestamp = value.try_into_timestamp().context(FieldTypeMismatchSnafu)?;
136                        timestamp
137                            .map(|t|t.value())
138                            .serialize($serializer)
139                            .context(SerializeFieldSnafu)?;
140                    }
141                    ConcreteDataType::Interval(IntervalType::YearMonth(_)) => {
142                        let interval = value.try_into_interval_year_month().context(FieldTypeMismatchSnafu)?;
143                        interval.map(|i| i.to_i32())
144                            .serialize($serializer)
145                            .context(SerializeFieldSnafu)?;
146                    }
147                    ConcreteDataType::Interval(IntervalType::DayTime(_)) => {
148                        let interval = value.try_into_interval_day_time().context(FieldTypeMismatchSnafu)?;
149                        interval.map(|i| i.to_i64())
150                            .serialize($serializer)
151                            .context(SerializeFieldSnafu)?;
152                    }
153                    ConcreteDataType::Interval(IntervalType::MonthDayNano(_)) => {
154                        let interval = value.try_into_interval_month_day_nano().context(FieldTypeMismatchSnafu)?;
155                        interval.map(|i| i.to_i128())
156                            .serialize($serializer)
157                            .context(SerializeFieldSnafu)?;
158                    }
159                    ConcreteDataType::List(_) |
160                    ConcreteDataType::Struct(_) |
161                    ConcreteDataType::Dictionary(_) |
162                    ConcreteDataType::Null(_) => {
163                        return error::NotSupportedFieldSnafu {
164                            data_type: $data_type.clone()
165                        }.fail()
166                    }
167                }
168            };
169        }
170        cast_value_and_serialize!(data_type; serializer;
171            Boolean, boolean,
172            Binary, binary,
173            Int8, i8,
174            UInt8, u8,
175            Int16, i16,
176            UInt16, u16,
177            Int32, i32,
178            UInt32, u32,
179            Int64, i64,
180            UInt64, u64,
181            Float32, f32,
182            Float64, f64,
183            String, string,
184            Date, date,
185            Time, time,
186            Duration, duration,
187            Decimal128, decimal128,
188            Json, binary,
189            Vector, binary
190        );
191
192        Ok(())
193    }
194
195    /// Deserialize a value from the deserializer.
196    pub fn deserialize<B: Buf>(&self, deserializer: &mut Deserializer<B>) -> Result<Value> {
197        Self::deserialize_by_type(self.encode_data_type(), deserializer)
198    }
199
200    fn deserialize_by_type<B: Buf>(
201        data_type: &ConcreteDataType,
202        deserializer: &mut Deserializer<B>,
203    ) -> Result<Value> {
204        macro_rules! deserialize_and_build_value {
205            (
206                $data_type: ident;
207                $serializer: ident;
208                $(
209                    $ty: ident, $f: ident
210                ),*
211            ) => {
212
213                match $data_type {
214                    $(
215                        ConcreteDataType::$ty(_) => {
216                            Ok(Value::from(Option::<$f>::deserialize(deserializer).context(error::DeserializeFieldSnafu)?))
217                        }
218                    )*
219                    ConcreteDataType::Binary(_) | ConcreteDataType::Json(_) | ConcreteDataType::Vector(_) => Ok(Value::from(
220                        Option::<Vec<u8>>::deserialize(deserializer)
221                            .context(error::DeserializeFieldSnafu)?
222                            .map(Bytes::from),
223                    )),
224                    ConcreteDataType::Timestamp(ty) => {
225                        let timestamp = Option::<i64>::deserialize(deserializer)
226                            .context(error::DeserializeFieldSnafu)?
227                            .map(|t|ty.create_timestamp(t));
228                        Ok(Value::from(timestamp))
229                    }
230                    ConcreteDataType::Interval(IntervalType::YearMonth(_)) => {
231                        let interval = Option::<i32>::deserialize(deserializer)
232                            .context(error::DeserializeFieldSnafu)?
233                            .map(IntervalYearMonth::from_i32);
234                        Ok(Value::from(interval))
235                    }
236                    ConcreteDataType::Interval(IntervalType::DayTime(_)) => {
237                        let interval = Option::<i64>::deserialize(deserializer)
238                            .context(error::DeserializeFieldSnafu)?
239                            .map(IntervalDayTime::from_i64);
240                        Ok(Value::from(interval))
241                    }
242                    ConcreteDataType::Interval(IntervalType::MonthDayNano(_)) => {
243                        let interval = Option::<i128>::deserialize(deserializer)
244                            .context(error::DeserializeFieldSnafu)?
245                            .map(IntervalMonthDayNano::from_i128);
246                        Ok(Value::from(interval))
247                    }
248                    ConcreteDataType::List(l) => NotSupportedFieldSnafu {
249                        data_type: ConcreteDataType::List(l.clone()),
250                    }
251                    .fail(),
252                    ConcreteDataType::Struct(f) => NotSupportedFieldSnafu {
253                        data_type: ConcreteDataType::Struct(f.clone()),
254                    }
255                    .fail(),
256                    ConcreteDataType::Dictionary(d) => NotSupportedFieldSnafu {
257                        data_type: ConcreteDataType::Dictionary(d.clone()),
258                    }
259                    .fail(),
260                    ConcreteDataType::Null(n) => NotSupportedFieldSnafu {
261                        data_type: ConcreteDataType::Null(n.clone()),
262                    }
263                    .fail(),
264                }
265            };
266        }
267        deserialize_and_build_value!(data_type; deserializer;
268            Boolean, bool,
269            Int8, i8,
270            Int16, i16,
271            Int32, i32,
272            Int64, i64,
273            UInt8, u8,
274            UInt16, u16,
275            UInt32, u32,
276            UInt64, u64,
277            Float32, f32,
278            Float64, f64,
279            String, String,
280            Date, Date,
281            Time, Time,
282            Duration, Duration,
283            Decimal128, Decimal128
284        )
285    }
286
287    /// Skip deserializing this field, returns the length of it.
288    pub(crate) fn skip_deserialize(
289        &self,
290        bytes: &[u8],
291        deserializer: &mut Deserializer<&[u8]>,
292    ) -> Result<usize> {
293        let pos = deserializer.position();
294        let remaining = bytes
295            .get(pos..)
296            .ok_or(memcomparable::Error::Eof)
297            .context(error::DeserializeFieldSnafu)?;
298        let len = Self::encoded_len(self.encode_data_type(), remaining)?;
299        if len > 1 && self.encode_data_type().is_boolean() && remaining[1] > 1 {
300            return Err(memcomparable::Error::InvalidBoolEncoding(remaining[1]))
301                .context(error::DeserializeFieldSnafu);
302        }
303        deserializer.advance(len);
304        Ok(len)
305    }
306
307    /// Checks field boundaries before any unchecked reads in memcomparable.
308    fn encoded_len(data_type: &ConcreteDataType, bytes: &[u8]) -> Result<usize> {
309        let marker = bytes
310            .first()
311            .copied()
312            .ok_or(memcomparable::Error::Eof)
313            .context(error::DeserializeFieldSnafu)?;
314        match marker {
315            0 => return Ok(1),
316            1 => {}
317            value => {
318                return Err(memcomparable::Error::InvalidTagEncoding(value as usize))
319                    .context(error::DeserializeFieldSnafu);
320            }
321        }
322        let to_skip = match data_type {
323            ConcreteDataType::Boolean(_) => 2,
324            ConcreteDataType::Int8(_) | ConcreteDataType::UInt8(_) => 2,
325            ConcreteDataType::Int16(_) | ConcreteDataType::UInt16(_) => 3,
326            ConcreteDataType::Int32(_) | ConcreteDataType::UInt32(_) => 5,
327            ConcreteDataType::Int64(_) | ConcreteDataType::UInt64(_) => 9,
328            ConcreteDataType::Float32(_) => 5,
329            ConcreteDataType::Float64(_) => 9,
330            ConcreteDataType::Binary(_)
331            | ConcreteDataType::Json(_)
332            | ConcreteDataType::Vector(_) => {
333                // Binary is encoded as a sequence of bytes, not chunked strings.
334                return encoded_binary_len(bytes).context(error::DeserializeFieldSnafu);
335            }
336            ConcreteDataType::String(_) => {
337                return encoded_string_len(bytes).context(error::DeserializeFieldSnafu);
338            }
339            ConcreteDataType::Date(_) => 5,
340            ConcreteDataType::Timestamp(_) => 9, // We treat timestamp as Option<i64>
341            ConcreteDataType::Time(_) => 10,     // i64 and 1 byte time unit
342            ConcreteDataType::Duration(_) => 10,
343            ConcreteDataType::Interval(IntervalType::YearMonth(_)) => 5,
344            ConcreteDataType::Interval(IntervalType::DayTime(_)) => 9,
345            ConcreteDataType::Interval(IntervalType::MonthDayNano(_)) => 17,
346            ConcreteDataType::Decimal128(_) => 19,
347            ConcreteDataType::Null(_)
348            | ConcreteDataType::List(_)
349            | ConcreteDataType::Struct(_)
350            | ConcreteDataType::Dictionary(_) => {
351                return NotSupportedFieldSnafu {
352                    data_type: data_type.clone(),
353                }
354                .fail();
355            }
356        };
357        if bytes.len() < to_skip {
358            return Err(memcomparable::Error::Eof).context(error::DeserializeFieldSnafu);
359        }
360        Ok(to_skip)
361    }
362
363    /// Decodes primitive fields straight into a reference value for column
364    /// builders, avoiding owned Value construction and conversion. The caller has
365    /// already checked the encoded boundary.
366    fn deserialize_primitive_ref<B: Buf>(
367        data_type: &ConcreteDataType,
368        deserializer: &mut Deserializer<B>,
369    ) -> Option<Result<ValueRef<'static>>> {
370        macro_rules! decode_primitive {
371            ($($variant:ident, $native:ty),* $(,)?) => {
372                match data_type {
373                    $(ConcreteDataType::$variant(_) => Some(
374                        Option::<$native>::deserialize(deserializer)
375                            .map(ValueRef::from)
376                            .context(error::DeserializeFieldSnafu)
377                    ),)*
378                    _ => None,
379                }
380            };
381        }
382        decode_primitive!(
383            Boolean, bool, Int8, i8, Int16, i16, Int32, i32, Int64, i64, UInt8, u8, UInt16, u16,
384            UInt32, u32, UInt64, u64, Float32, f32, Float64, f64,
385        )
386    }
387}
388
389/// The Option marker has already been checked by SortField.
390fn encoded_binary_len(bytes: &[u8]) -> memcomparable::Result<usize> {
391    let mut current = 1;
392    loop {
393        match bytes.get(current).copied() {
394            Some(0) => return Ok(current + 1),
395            Some(1) if bytes.get(current + 1).is_some() => current += 2,
396            Some(1) | None => return Err(memcomparable::Error::Eof),
397            Some(marker) => return Err(memcomparable::Error::InvalidSeqEncoding(marker)),
398        }
399    }
400}
401
402impl PrimaryKeyCodecExt for DensePrimaryKeyCodec {
403    fn encode_to_vec<'a, I>(&self, row: I, buffer: &mut Vec<u8>) -> Result<()>
404    where
405        I: Iterator<Item = ValueRef<'a>>,
406    {
407        self.encode_dense(row, buffer)
408    }
409}
410
411/// A memory-comparable row [`Value`] encoder/decoder.
412#[derive(Clone, Debug)]
413pub struct DensePrimaryKeyCodec {
414    /// Primary key fields.
415    ordered_primary_key_columns: Arc<Vec<(ColumnId, SortField)>>,
416}
417
418impl DensePrimaryKeyCodec {
419    pub fn new(metadata: &RegionMetadata) -> Self {
420        let ordered_primary_key_columns = metadata
421            .primary_key_columns()
422            .map(|c| {
423                (
424                    c.column_id,
425                    SortField::new(c.column_schema.data_type.clone()),
426                )
427            })
428            .collect::<Vec<_>>();
429
430        Self::with_fields(ordered_primary_key_columns)
431    }
432
433    pub fn with_fields(fields: Vec<(ColumnId, SortField)>) -> Self {
434        Self {
435            ordered_primary_key_columns: Arc::new(fields),
436        }
437    }
438
439    fn encode_dense<'a, I>(&self, row: I, buffer: &mut Vec<u8>) -> Result<()>
440    where
441        I: Iterator<Item = ValueRef<'a>>,
442    {
443        let mut serializer = Serializer::new(buffer);
444        for (idx, value) in row.enumerate() {
445            self.field_at(idx).serialize(&mut serializer, &value)?;
446        }
447        Ok(())
448    }
449
450    /// Decode primary key values from bytes.
451    pub fn decode_dense(&self, bytes: &[u8]) -> Result<Vec<(ColumnId, Value)>> {
452        let mut deserializer = Deserializer::new(bytes);
453        let mut values = Vec::with_capacity(self.ordered_primary_key_columns.len());
454        for (column_id, field) in self.ordered_primary_key_columns.iter() {
455            let value = field.deserialize(&mut deserializer)?;
456            values.push((*column_id, value));
457        }
458        Ok(values)
459    }
460
461    /// Counts complete fields in a key encoded with a prefix of this codec's schema.
462    ///
463    /// Dense keys contain neither column ids nor types: callers must ensure that
464    /// existing fields have the same order and types. Only EOF between fields is
465    /// accepted; truncated fields and bytes beyond the schema return an error.
466    /// Field boundaries are checked in all builds; full value validation is only
467    /// enabled in debug builds and this crate's unit tests.
468    pub fn decode_prefix_len(&self, bytes: &[u8]) -> Result<usize> {
469        let mut deserializer = Deserializer::new(bytes);
470        for (index, (_, field)) in self.ordered_primary_key_columns.iter().enumerate() {
471            if !deserializer.has_remaining() {
472                return Ok(index);
473            }
474            let start = deserializer.position();
475            // Preserve the prefix error contract used by primary-key range mapping.
476            let len = SortField::encoded_len(field.encode_data_type(), &bytes[start..]).map_err(
477                |source| match source {
478                    error::Error::DeserializeField { .. } => error::InvalidDensePrimaryKeySnafu {
479                        reason: "truncated field or invalid encoding",
480                    }
481                    .build(),
482                    source => source,
483                },
484            )?;
485            // Production range mapping only needs boundaries, not allocated field values.
486            #[cfg(any(debug_assertions, test))]
487            field.deserialize(&mut Deserializer::new(&bytes[start..start + len]))?;
488            deserializer.advance(len);
489        }
490        snafu::ensure!(
491            !deserializer.has_remaining(),
492            error::InvalidDensePrimaryKeySnafu {
493                reason: "key contains bytes beyond the primary key schema",
494            }
495        );
496        Ok(self.num_fields())
497    }
498
499    /// Decode primary key values from bytes without column id.
500    pub fn decode_dense_without_column_id(&self, bytes: &[u8]) -> Result<Vec<Value>> {
501        let mut deserializer = Deserializer::new(bytes);
502        let mut values = Vec::with_capacity(self.ordered_primary_key_columns.len());
503        for (_, field) in self.ordered_primary_key_columns.iter() {
504            let value = field.deserialize(&mut deserializer)?;
505            values.push(value);
506        }
507        Ok(values)
508    }
509
510    /// Iterates over all source PK fields in schema order, checking each field's
511    /// encoded boundary before deserializing it. Stops after the first error.
512    ///
513    /// Unlike positional access, this requires no offsets or value cache. Callers
514    /// can consume each value immediately instead of retaining the entire key.
515    pub fn decode_dense_iter<'a>(
516        &'a self,
517        bytes: &'a [u8],
518    ) -> impl Iterator<Item = Result<(ColumnId, Value)>> + 'a {
519        self.ordered_primary_key_columns.iter().scan(
520            Some(Deserializer::new(bytes)),
521            move |state, (id, field)| {
522                let deserializer = state.as_mut()?;
523                let remaining = &bytes[deserializer.position()..];
524                let decoded = SortField::encoded_len(field.encode_data_type(), remaining)
525                    .and_then(|_| field.deserialize(deserializer).map(|value| (*id, value)));
526                if decoded.is_err() {
527                    *state = None;
528                }
529                Some(decoded)
530            },
531        )
532    }
533
534    /// Returns the column ids and sort fields in encoded primary-key order.
535    pub fn fields(&self) -> &[(ColumnId, SortField)] {
536        &self.ordered_primary_key_columns
537    }
538
539    /// Decodes all fields in source order into a consumer, checking boundaries.
540    /// Strings borrow the reusable buffer for the duration of each callback, so
541    /// materializing them into column builders needs no per-value owned string.
542    pub fn decode_dense_with(
543        &self,
544        bytes: &[u8],
545        value_buf: &mut Vec<u8>,
546        mut consume: impl FnMut(usize, ValueRef<'_>),
547    ) -> Result<()> {
548        let mut deserializer = Deserializer::new(bytes);
549        for (pos, (_, field)) in self.ordered_primary_key_columns.iter().enumerate() {
550            let data_type = field.encode_data_type();
551            let remaining = &bytes[deserializer.position()..];
552            SortField::encoded_len(data_type, remaining)?;
553            if data_type.is_string() && remaining[0] != 0 {
554                deserializer.advance(1);
555                deserializer
556                    .read_bytes_into(value_buf)
557                    .context(error::DeserializeFieldSnafu)?;
558                let value = std::str::from_utf8(value_buf).map_err(|err| {
559                    error::InvalidDensePrimaryKeySnafu {
560                        reason: format!("string is not valid UTF-8: {err}"),
561                    }
562                    .build()
563                })?;
564                consume(pos, ValueRef::String(value));
565            } else if let Some(value) =
566                SortField::deserialize_primitive_ref(data_type, &mut deserializer)
567            {
568                consume(pos, value?);
569            } else {
570                let value = field.deserialize(&mut deserializer)?;
571                consume(pos, value.as_value_ref());
572            }
573        }
574        Ok(())
575    }
576
577    /// Returns the field at `pos`.
578    ///
579    /// # Panics
580    /// Panics if `pos` is out of bounds.
581    fn field_at(&self, pos: usize) -> &SortField {
582        &self.ordered_primary_key_columns[pos].1
583    }
584
585    /// Advances `deserializer` to the start of value at `pos`.
586    ///
587    /// Returns the offset of the value at `pos` in `bytes`.
588    fn advance_to_value_at(
589        &self,
590        bytes: &[u8],
591        pos: usize,
592        offsets_buf: &mut Vec<usize>,
593        deserializer: &mut Deserializer<&[u8]>,
594    ) -> Result<usize> {
595        ensure!(
596            pos < self.num_fields(),
597            error::InvalidDensePrimaryKeySnafu {
598                reason: format!(
599                    "field position {pos} exceeds field count {}",
600                    self.num_fields()
601                ),
602            }
603        );
604        if offsets_buf.is_empty() {
605            // The schema bounds this cache. Reserve once even when the first
606            // requested field is near the end of a wide key.
607            offsets_buf.reserve_exact(self.num_fields());
608            offsets_buf.push(0);
609        }
610        // Start at the requested field if cached, otherwise resume discovery
611        // from the furthest known boundary.
612        let value_start = pos.min(offsets_buf.len() - 1);
613        let mut offset = offsets_buf[value_start];
614        ensure!(
615            offset <= bytes.len(),
616            error::InvalidDensePrimaryKeySnafu {
617                reason: "cached offset exceeds key length",
618            }
619        );
620        deserializer.advance(offset);
621        for i in value_start..pos {
622            offset += self.field_at(i).skip_deserialize(bytes, deserializer)?;
623            offsets_buf.push(offset);
624        }
625        Ok(offset)
626    }
627
628    /// Decode value at `pos` in `bytes`.
629    ///
630    /// The i-th element in the offsets buffer is the start of field i. The buffer
631    /// must only be reused for the same encoded key and codec.
632    #[inline]
633    pub fn decode_value_at(
634        &self,
635        bytes: &[u8],
636        pos: usize,
637        offsets_buf: &mut Vec<usize>,
638    ) -> Result<Value> {
639        let encoded = self.encoded_value_at(bytes, pos, offsets_buf)?;
640        self.field_at(pos)
641            .deserialize(&mut Deserializer::new(encoded))
642    }
643
644    /// Returns the encoded bytes at `pos` in `bytes`.
645    ///
646    /// The i-th element in the offsets buffer is the start of field i. The buffer
647    /// must only be reused for the same encoded key and codec. Validates the
648    /// framing up to this field, without validating or decoding later fields.
649    #[inline]
650    pub fn encoded_value_at<'a>(
651        &self,
652        bytes: &'a [u8],
653        pos: usize,
654        offsets_buf: &mut Vec<usize>,
655    ) -> Result<&'a [u8]> {
656        let mut deserializer = Deserializer::new(bytes);
657        let offset = self.advance_to_value_at(bytes, pos, offsets_buf, &mut deserializer)?;
658
659        let len = self
660            .field_at(pos)
661            .skip_deserialize(bytes, &mut deserializer)?;
662        // We already found the next field's start while validating this one.
663        // Reuse it instead of scanning this field again on the next lookup.
664        if offsets_buf.len() == pos + 1 && pos + 1 < self.num_fields() {
665            offsets_buf.push(offset + len);
666        }
667        Ok(&bytes[offset..offset + len])
668    }
669
670    pub fn estimated_size(&self) -> usize {
671        self.ordered_primary_key_columns
672            .iter()
673            .map(|(_, f)| f.estimated_size())
674            .sum()
675    }
676
677    pub fn num_fields(&self) -> usize {
678        self.ordered_primary_key_columns.len()
679    }
680}
681
682impl PrimaryKeyCodec for DensePrimaryKeyCodec {
683    fn encode_key_value(&self, key_value: &KeyValue, buffer: &mut Vec<u8>) -> Result<()> {
684        self.encode_dense(key_value.primary_keys(), buffer)
685    }
686
687    fn encode_values(&self, values: &[(ColumnId, Value)], buffer: &mut Vec<u8>) -> Result<()> {
688        self.encode_dense(values.iter().map(|(_, v)| v.as_value_ref()), buffer)
689    }
690
691    fn encode_value_refs(
692        &self,
693        values: &[(ColumnId, ValueRef)],
694        buffer: &mut Vec<u8>,
695    ) -> Result<()> {
696        let iter = values.iter().map(|(_, v)| v.clone());
697        self.encode_dense(iter, buffer)
698    }
699
700    fn estimated_size(&self) -> Option<usize> {
701        Some(self.estimated_size())
702    }
703
704    fn num_fields(&self) -> Option<usize> {
705        Some(self.num_fields())
706    }
707
708    fn encoding(&self) -> PrimaryKeyEncoding {
709        PrimaryKeyEncoding::Dense
710    }
711
712    fn as_dense(&self) -> Option<&DensePrimaryKeyCodec> {
713        Some(self)
714    }
715
716    fn primary_key_filter(
717        &self,
718        metadata: &RegionMetadataRef,
719        filters: Arc<Vec<SimpleFilterEvaluator>>,
720    ) -> Box<dyn PrimaryKeyFilter> {
721        Box::new(DensePrimaryKeyFilter::new(
722            metadata.clone(),
723            filters,
724            self.clone(),
725        ))
726    }
727
728    fn decode(&self, bytes: &[u8]) -> Result<CompositeValues> {
729        Ok(CompositeValues::Dense(self.decode_dense(bytes)?))
730    }
731
732    fn decode_leftmost(&self, bytes: &[u8]) -> Result<Option<Value>> {
733        // TODO(weny, yinwen): avoid decoding the whole primary key.
734        let mut values = self.decode_dense(bytes)?;
735        Ok(values.pop().map(|(_, v)| v))
736    }
737}
738
739#[cfg(test)]
740mod tests {
741    use common_base::bytes::StringBytes;
742    use common_time::{IntervalDayTime, IntervalMonthDayNano, IntervalYearMonth, Timestamp};
743    use datatypes::value::Value;
744
745    use super::*;
746
747    fn check_encode_and_decode(data_types: &[ConcreteDataType], row: Vec<Value>) {
748        let encoder = DensePrimaryKeyCodec::with_fields(
749            data_types
750                .iter()
751                .map(|t| (0, SortField::new(t.clone())))
752                .collect::<Vec<_>>(),
753        );
754
755        let value_ref = row.iter().map(|v| v.as_value_ref()).collect::<Vec<_>>();
756
757        let result = encoder.encode(value_ref.iter().cloned()).unwrap();
758        let boundaries: Vec<_> = (0..=row.len())
759            .map(|count| {
760                encoder
761                    .encode(value_ref[..count].iter().cloned())
762                    .unwrap()
763                    .len()
764            })
765            .collect();
766        for end in 0..=result.len() {
767            match boundaries.iter().position(|boundary| *boundary == end) {
768                Some(count) => {
769                    assert_eq!(count, encoder.decode_prefix_len(&result[..end]).unwrap())
770                }
771                None => assert!(
772                    encoder.decode_prefix_len(&result[..end]).is_err(),
773                    "truncated {data_types:?} at {end}"
774                ),
775            }
776        }
777        let decoded = encoder.decode(&result).unwrap().into_dense();
778        assert_eq!(decoded, row);
779        let mut value_buf = Vec::new();
780        let mut borrowed = Vec::new();
781        encoder
782            .decode_dense_with(&result, &mut value_buf, |_, value| {
783                borrowed.push(Value::from(value))
784            })
785            .unwrap();
786        assert_eq!(borrowed, row);
787        let sequential = encoder
788            .decode_dense_iter(&result)
789            .collect::<Result<Vec<_>>>()
790            .unwrap();
791        assert_eq!(
792            sequential
793                .iter()
794                .map(|(_, value)| value)
795                .collect::<Vec<_>>(),
796            row.iter().collect::<Vec<_>>()
797        );
798        for end in 0..result.len() {
799            assert!(
800                encoder
801                    .decode_dense_with(&result[..end], &mut value_buf, |_, _| {})
802                    .is_err()
803            );
804            assert!(
805                encoder
806                    .decode_dense_iter(&result[..end])
807                    .collect::<Result<Vec<_>>>()
808                    .is_err(),
809                "truncated at {end}"
810            );
811        }
812        // Every supported type must have the same boundaries in positional and
813        // full decoding. Exercise every truncation both at the requested field
814        // and while skipping it to reach a later field.
815        let mut offsets = Vec::new();
816        let mut end = 0;
817        for (pos, value) in row.iter().enumerate() {
818            let encoded = encoder
819                .encoded_value_at(&result, pos, &mut offsets)
820                .unwrap();
821            assert_eq!(
822                encoder
823                    .field_at(pos)
824                    .deserialize(&mut Deserializer::new(encoded))
825                    .unwrap(),
826                *value
827            );
828            let start = end;
829            end += encoded.len();
830            for len in start..end {
831                assert!(
832                    encoder
833                        .encoded_value_at(&result[..len], pos, &mut Vec::new())
834                        .is_err()
835                );
836                assert!(
837                    encoder
838                        .decode_value_at(&result[..len], pos, &mut Vec::new())
839                        .is_err()
840                );
841                if pos + 1 < row.len() {
842                    assert!(
843                        encoder
844                            .encoded_value_at(&result[..len], pos + 1, &mut Vec::new())
845                            .is_err()
846                    );
847                }
848            }
849        }
850        assert_eq!(end, result.len());
851        let mut decoded = Vec::new();
852        let mut offsets = Vec::new();
853        // Iter two times to test offsets buffer.
854        for _ in 0..2 {
855            decoded.clear();
856            for i in 0..data_types.len() {
857                let value = encoder.decode_value_at(&result, i, &mut offsets).unwrap();
858                decoded.push(value);
859            }
860            assert_eq!(data_types.len(), offsets.len(), "offsets: {offsets:?}");
861            assert_eq!(decoded, row);
862        }
863    }
864
865    #[test]
866    fn test_decode_prefix_len_accepts_only_complete_fields() {
867        let codec = DensePrimaryKeyCodec::with_fields(vec![
868            (0, SortField::new(ConcreteDataType::string_datatype())),
869            (1, SortField::new(ConcreteDataType::int64_datatype())),
870            (2, SortField::new(ConcreteDataType::binary_datatype())),
871            (3, SortField::new(ConcreteDataType::string_datatype())),
872        ]);
873        let values = [
874            Value::from("abcdefghijk"),
875            Value::Int64(42),
876            Value::Binary(vec![0, 255, 1].into()),
877            Value::Null,
878        ];
879        let mut boundaries = vec![0];
880        let mut encoded = Vec::new();
881        for count in 1..=values.len() {
882            encoded.clear();
883            codec
884                .encode_dense(
885                    values[..count].iter().map(Value::as_value_ref),
886                    &mut encoded,
887                )
888                .unwrap();
889            boundaries.push(encoded.len());
890            assert_eq!(count, codec.decode_prefix_len(&encoded).unwrap());
891        }
892        for end in 0..=encoded.len() {
893            match boundaries.iter().position(|boundary| *boundary == end) {
894                Some(count) => assert_eq!(count, codec.decode_prefix_len(&encoded[..end]).unwrap()),
895                None => assert!(
896                    codec.decode_prefix_len(&encoded[..end]).is_err(),
897                    "truncated at {end}"
898                ),
899            }
900        }
901        encoded.push(0);
902        assert!(codec.decode_prefix_len(&encoded).is_err());
903        assert!(codec.decode_prefix_len(&[2]).is_err());
904        let empty = DensePrimaryKeyCodec::with_fields(vec![]);
905        assert_eq!(0, empty.decode_prefix_len(&[]).unwrap());
906        assert!(empty.decode_prefix_len(&[0]).is_err());
907    }
908
909    #[test]
910    fn test_prefix_len_validates_values_in_unit_tests() {
911        let codec = DensePrimaryKeyCodec::with_fields(vec![(
912            0,
913            SortField::new(ConcreteDataType::boolean_datatype()),
914        )]);
915        // The field is complete, but 2 is not a boolean. Even release unit tests
916        // retain value validation through cfg(test).
917        assert!(codec.decode_prefix_len(&[1, 2]).is_err());
918    }
919
920    #[test]
921    fn test_memcmp() {
922        let encoder = DensePrimaryKeyCodec::with_fields(vec![
923            (0, SortField::new(ConcreteDataType::string_datatype())),
924            (1, SortField::new(ConcreteDataType::int64_datatype())),
925        ]);
926        let values = [Value::String("abcdefgh".into()), Value::Int64(128)];
927        let value_ref = values.iter().map(|v| v.as_value_ref()).collect::<Vec<_>>();
928        let result = encoder.encode(value_ref.iter().cloned()).unwrap();
929
930        let decoded = encoder.decode(&result).unwrap().into_dense();
931        assert_eq!(&values, &decoded as &[Value]);
932    }
933
934    #[test]
935    fn test_memcmp_timestamp() {
936        check_encode_and_decode(
937            &[
938                ConcreteDataType::timestamp_millisecond_datatype(),
939                ConcreteDataType::int64_datatype(),
940            ],
941            vec![
942                Value::Timestamp(Timestamp::new_millisecond(42)),
943                Value::Int64(43),
944            ],
945        );
946    }
947
948    #[test]
949    fn test_memcmp_duration() {
950        check_encode_and_decode(
951            &[
952                ConcreteDataType::duration_millisecond_datatype(),
953                ConcreteDataType::int64_datatype(),
954            ],
955            vec![
956                Value::Duration(Duration::new_millisecond(44)),
957                Value::Int64(45),
958            ],
959        )
960    }
961
962    #[test]
963    fn test_memcmp_binary() {
964        check_encode_and_decode(
965            &[
966                ConcreteDataType::binary_datatype(),
967                ConcreteDataType::int64_datatype(),
968            ],
969            vec![
970                Value::Binary(Bytes::from("hello".as_bytes())),
971                Value::Int64(43),
972            ],
973        );
974    }
975
976    #[test]
977    fn test_memcmp_string() {
978        check_encode_and_decode(
979            &[ConcreteDataType::string_datatype()],
980            vec![Value::String(StringBytes::from("hello"))],
981        );
982
983        check_encode_and_decode(&[ConcreteDataType::string_datatype()], vec![Value::Null]);
984
985        check_encode_and_decode(
986            &[ConcreteDataType::string_datatype()],
987            vec![Value::String("".into())],
988        );
989        check_encode_and_decode(
990            &[ConcreteDataType::string_datatype()],
991            vec![Value::String("world".into())],
992        );
993    }
994
995    #[test]
996    fn test_encode_null() {
997        check_encode_and_decode(
998            &[
999                ConcreteDataType::string_datatype(),
1000                ConcreteDataType::int32_datatype(),
1001            ],
1002            vec![Value::String(StringBytes::from("abcd")), Value::Null],
1003        )
1004    }
1005
1006    #[test]
1007    fn test_encoded_value_at() {
1008        let data_types = [
1009            ConcreteDataType::string_datatype(),
1010            ConcreteDataType::int32_datatype(),
1011            ConcreteDataType::string_datatype(),
1012        ];
1013        let encoder = DensePrimaryKeyCodec::with_fields(
1014            data_types
1015                .iter()
1016                .enumerate()
1017                .map(|(idx, t)| (idx as ColumnId, SortField::new(t.clone())))
1018                .collect::<Vec<_>>(),
1019        );
1020        let row = [Value::String("hello".into()), Value::Int32(42), Value::Null];
1021        let value_ref = row.iter().map(|v| v.as_value_ref()).collect::<Vec<_>>();
1022        let encoded_pk = encoder.encode(value_ref.iter().cloned()).unwrap();
1023
1024        let mut offsets = Vec::new();
1025        let mut combined = Vec::new();
1026        for (pos, value) in row.iter().enumerate() {
1027            let encoded_value = encoder
1028                .encoded_value_at(&encoded_pk, pos, &mut offsets)
1029                .unwrap();
1030            combined.extend_from_slice(encoded_value);
1031
1032            let field = SortField::new(data_types[pos].clone());
1033            let mut expected = Vec::new();
1034            let mut serializer = Serializer::new(&mut expected);
1035            field
1036                .serialize(&mut serializer, &value.as_value_ref())
1037                .unwrap();
1038            assert_eq!(encoded_value, expected.as_slice());
1039        }
1040        assert_eq!(combined, encoded_pk);
1041        assert_eq!(offsets.len(), row.len());
1042
1043        // Verify the offsets buffer can be reused for random access.
1044        for (pos, value) in row.iter().enumerate().rev() {
1045            let encoded_value = encoder
1046                .encoded_value_at(&encoded_pk, pos, &mut offsets)
1047                .unwrap();
1048
1049            let field = SortField::new(data_types[pos].clone());
1050            let mut expected = Vec::new();
1051            let mut serializer = Serializer::new(&mut expected);
1052            field
1053                .serialize(&mut serializer, &value.as_value_ref())
1054                .unwrap();
1055            assert_eq!(encoded_value, expected.as_slice());
1056        }
1057    }
1058
1059    #[test]
1060    fn test_positional_decode_invalid_encoding() {
1061        for (ty, bytes) in [
1062            (ConcreteDataType::int64_datatype(), vec![2]),
1063            (ConcreteDataType::boolean_datatype(), vec![1, 2]),
1064            (ConcreteDataType::binary_datatype(), vec![1, 2]),
1065            (ConcreteDataType::binary_datatype(), vec![1, 1, 42, 2]),
1066            (ConcreteDataType::string_datatype(), vec![1, 2]),
1067            (
1068                ConcreteDataType::string_datatype(),
1069                vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
1070            ),
1071            (
1072                ConcreteDataType::string_datatype(),
1073                vec![1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 10],
1074            ),
1075        ] {
1076            let codec = DensePrimaryKeyCodec::with_fields(vec![
1077                (0, SortField::new(ty)),
1078                (1, SortField::new(ConcreteDataType::int64_datatype())),
1079            ]);
1080            for pos in [0, 1, 2, usize::MAX] {
1081                assert!(
1082                    codec
1083                        .encoded_value_at(&bytes, pos, &mut Vec::new())
1084                        .is_err()
1085                );
1086                assert!(codec.decode_value_at(&bytes, pos, &mut Vec::new()).is_err());
1087            }
1088            assert!(
1089                codec
1090                    .decode_dense_iter(&bytes)
1091                    .collect::<Result<Vec<_>>>()
1092                    .is_err()
1093            );
1094            assert!(
1095                codec
1096                    .decode_dense_with(&bytes, &mut Vec::new(), |_, _| {})
1097                    .is_err()
1098            );
1099        }
1100        let codec = DensePrimaryKeyCodec::with_fields(vec![
1101            (0, SortField::new(ConcreteDataType::string_datatype())),
1102            (1, SortField::new(ConcreteDataType::int64_datatype())),
1103        ]);
1104        let bytes = codec
1105            .encode([ValueRef::String("abcdefghijk"), ValueRef::Int64(42)].into_iter())
1106            .unwrap();
1107        for pos in [2, usize::MAX] {
1108            assert!(
1109                codec
1110                    .encoded_value_at(&bytes, pos, &mut Vec::new())
1111                    .is_err()
1112            );
1113            assert!(codec.decode_value_at(&bytes, pos, &mut Vec::new()).is_err());
1114        }
1115        for offset in [bytes.len(), bytes.len() + 1, usize::MAX] {
1116            for pos in [0, 1] {
1117                assert!(
1118                    codec
1119                        .encoded_value_at(&bytes, pos, &mut vec![offset])
1120                        .is_err()
1121                );
1122                assert!(
1123                    codec
1124                        .decode_value_at(&bytes, pos, &mut vec![offset])
1125                        .is_err()
1126                );
1127            }
1128        }
1129        let mut invalid_utf8 = bytes;
1130        invalid_utf8[2] = 0xff;
1131        assert!(
1132            codec
1133                .decode_dense_with(&invalid_utf8, &mut Vec::new(), |_, _| {})
1134                .is_err()
1135        );
1136        assert!(
1137            codec
1138                .decode_dense_iter(&invalid_utf8)
1139                .collect::<Result<Vec<_>>>()
1140                .is_err()
1141        );
1142        assert!(
1143            codec
1144                .decode_value_at(&invalid_utf8, 0, &mut Vec::new())
1145                .is_err()
1146        );
1147    }
1148
1149    #[test]
1150    fn test_memcmp_dictionary() {
1151        // Test Dictionary<i32, string>
1152        check_encode_and_decode(
1153            &[ConcreteDataType::dictionary_datatype(
1154                ConcreteDataType::int32_datatype(),
1155                ConcreteDataType::string_datatype(),
1156            )],
1157            vec![Value::String("hello".into())],
1158        );
1159
1160        // Test Dictionary<i32, i64>
1161        check_encode_and_decode(
1162            &[ConcreteDataType::dictionary_datatype(
1163                ConcreteDataType::int32_datatype(),
1164                ConcreteDataType::int64_datatype(),
1165            )],
1166            vec![Value::Int64(42)],
1167        );
1168
1169        // Test Dictionary with null value
1170        check_encode_and_decode(
1171            &[ConcreteDataType::dictionary_datatype(
1172                ConcreteDataType::int32_datatype(),
1173                ConcreteDataType::string_datatype(),
1174            )],
1175            vec![Value::Null],
1176        );
1177
1178        // Test multiple Dictionary columns
1179        check_encode_and_decode(
1180            &[
1181                ConcreteDataType::dictionary_datatype(
1182                    ConcreteDataType::int32_datatype(),
1183                    ConcreteDataType::string_datatype(),
1184                ),
1185                ConcreteDataType::dictionary_datatype(
1186                    ConcreteDataType::int16_datatype(),
1187                    ConcreteDataType::int64_datatype(),
1188                ),
1189            ],
1190            vec![Value::String("world".into()), Value::Int64(123)],
1191        );
1192    }
1193
1194    #[test]
1195    fn test_encode_multiple_rows() {
1196        check_encode_and_decode(
1197            &[
1198                ConcreteDataType::string_datatype(),
1199                ConcreteDataType::int64_datatype(),
1200                ConcreteDataType::boolean_datatype(),
1201            ],
1202            vec![
1203                Value::String("hello".into()),
1204                Value::Int64(42),
1205                Value::Boolean(false),
1206            ],
1207        );
1208
1209        check_encode_and_decode(
1210            &[
1211                ConcreteDataType::string_datatype(),
1212                ConcreteDataType::int64_datatype(),
1213                ConcreteDataType::boolean_datatype(),
1214            ],
1215            vec![
1216                Value::String("world".into()),
1217                Value::Int64(43),
1218                Value::Boolean(true),
1219            ],
1220        );
1221
1222        check_encode_and_decode(
1223            &[
1224                ConcreteDataType::string_datatype(),
1225                ConcreteDataType::int64_datatype(),
1226                ConcreteDataType::boolean_datatype(),
1227            ],
1228            vec![Value::Null, Value::Int64(43), Value::Boolean(true)],
1229        );
1230
1231        // All types.
1232        check_encode_and_decode(
1233            &[
1234                ConcreteDataType::boolean_datatype(),
1235                ConcreteDataType::int8_datatype(),
1236                ConcreteDataType::uint8_datatype(),
1237                ConcreteDataType::int16_datatype(),
1238                ConcreteDataType::uint16_datatype(),
1239                ConcreteDataType::int32_datatype(),
1240                ConcreteDataType::uint32_datatype(),
1241                ConcreteDataType::int64_datatype(),
1242                ConcreteDataType::uint64_datatype(),
1243                ConcreteDataType::float32_datatype(),
1244                ConcreteDataType::float64_datatype(),
1245                ConcreteDataType::binary_datatype(),
1246                ConcreteDataType::string_datatype(),
1247                ConcreteDataType::date_datatype(),
1248                ConcreteDataType::timestamp_millisecond_datatype(),
1249                ConcreteDataType::time_millisecond_datatype(),
1250                ConcreteDataType::duration_millisecond_datatype(),
1251                ConcreteDataType::interval_year_month_datatype(),
1252                ConcreteDataType::interval_day_time_datatype(),
1253                ConcreteDataType::interval_month_day_nano_datatype(),
1254                ConcreteDataType::decimal128_default_datatype(),
1255                ConcreteDataType::vector_datatype(3),
1256                ConcreteDataType::dictionary_datatype(
1257                    ConcreteDataType::int32_datatype(),
1258                    ConcreteDataType::string_datatype(),
1259                ),
1260            ],
1261            vec![
1262                Value::Boolean(true),
1263                Value::Int8(8),
1264                Value::UInt8(8),
1265                Value::Int16(16),
1266                Value::UInt16(16),
1267                Value::Int32(32),
1268                Value::UInt32(32),
1269                Value::Int64(64),
1270                Value::UInt64(64),
1271                Value::Float32(1.0.into()),
1272                Value::Float64(1.0.into()),
1273                Value::Binary(b"hello"[..].into()),
1274                Value::String("world".into()),
1275                Value::Date(Date::new(10)),
1276                Value::Timestamp(Timestamp::new_millisecond(12)),
1277                Value::Time(Time::new_millisecond(13)),
1278                Value::Duration(Duration::new_millisecond(14)),
1279                Value::IntervalYearMonth(IntervalYearMonth::new(1)),
1280                Value::IntervalDayTime(IntervalDayTime::new(1, 15)),
1281                Value::IntervalMonthDayNano(IntervalMonthDayNano::new(1, 1, 15)),
1282                Value::Decimal128(Decimal128::from(16)),
1283                Value::Binary(Bytes::from(vec![0; 12])),
1284                Value::String("dict_value".into()),
1285            ],
1286        );
1287    }
1288}
mito_codec/row_converter/dense.rs

mito_codec/row_converter/
dense.rs
