mito_codec/
row_converter.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
15pub mod dense;
16pub mod sparse;
17
18use std::fmt::Debug;
19use std::sync::Arc;
20
21use common_recordbatch::filter::SimpleFilterEvaluator;
22use datatypes::value::{Value, ValueRef};
23pub use dense::{DensePrimaryKeyCodec, SortField};
24pub use sparse::{COLUMN_ID_ENCODE_SIZE, SparseOffsetsCache, SparsePrimaryKeyCodec, SparseValues};
25use store_api::codec::PrimaryKeyEncoding;
26use store_api::metadata::{RegionMetadata, RegionMetadataRef};
27use store_api::storage::ColumnId;
28
29use crate::error::Result;
30use crate::key_values::KeyValue;
31
32/// Row value encoder/decoder.
33pub trait PrimaryKeyCodecExt {
34    /// Encodes rows to bytes.
35    /// # Note
36    /// Ensure the length of row iterator matches the length of fields.
37    fn encode<'a, I>(&self, row: I) -> Result<Vec<u8>>
38    where
39        I: Iterator<Item = ValueRef<'a>>,
40    {
41        let mut buffer = Vec::new();
42        self.encode_to_vec(row, &mut buffer)?;
43        Ok(buffer)
44    }
45
46    /// Encodes rows to specific vec.
47    /// # Note
48    /// Ensure the length of row iterator matches the length of fields.
49    fn encode_to_vec<'a, I>(&self, row: I, buffer: &mut Vec<u8>) -> Result<()>
50    where
51        I: Iterator<Item = ValueRef<'a>>;
52}
53
54pub trait PrimaryKeyFilter: Send + Sync {
55    /// Returns true if the primary key matches the filter.
56    fn matches(&mut self, pk: &[u8]) -> Result<bool>;
57}
58
59/// Composite values decoded from primary key bytes.
60#[derive(Debug, Clone, PartialEq, Eq)]
61pub enum CompositeValues {
62    Dense(Vec<(ColumnId, Value)>),
63    Sparse(SparseValues),
64}
65
66impl CompositeValues {
67    /// Extends the composite values with the given values.
68    ///
69    /// Append-only: `values` must not contain a column id already present in
70    /// the composite; otherwise the existing entry would shadow the new one on
71    /// `SparseValues` lookup.
72    pub fn extend(&mut self, values: &[(ColumnId, Value)]) {
73        match self {
74            CompositeValues::Dense(dense_values) => {
75                for (column_id, value) in values {
76                    dense_values.push((*column_id, value.clone()));
77                }
78            }
79            CompositeValues::Sparse(sprase_value) => {
80                for (column_id, value) in values {
81                    sprase_value.insert(*column_id, value.clone());
82                }
83            }
84        }
85    }
86}
87
88#[cfg(any(test, feature = "testing"))]
89impl CompositeValues {
90    pub fn into_sparse(self) -> SparseValues {
91        match self {
92            CompositeValues::Sparse(v) => v,
93            _ => panic!("CompositeValues is not sparse"),
94        }
95    }
96
97    pub fn into_dense(self) -> Vec<Value> {
98        match self {
99            CompositeValues::Dense(v) => v.into_iter().map(|(_, v)| v).collect(),
100            _ => panic!("CompositeValues is not dense"),
101        }
102    }
103}
104
105pub trait PrimaryKeyCodec: Send + Sync + Debug {
106    /// Encodes a key value to bytes.
107    fn encode_key_value(&self, key_value: &KeyValue, buffer: &mut Vec<u8>) -> Result<()>;
108
109    /// Encodes values to bytes.
110    fn encode_values(&self, values: &[(ColumnId, Value)], buffer: &mut Vec<u8>) -> Result<()>;
111
112    /// Encodes values to bytes.
113    fn encode_value_refs(
114        &self,
115        values: &[(ColumnId, ValueRef)],
116        buffer: &mut Vec<u8>,
117    ) -> Result<()>;
118
119    /// Returns the number of fields in the primary key.
120    fn num_fields(&self) -> Option<usize>;
121
122    /// Returns a primary key filter factory.
123    fn primary_key_filter(
124        &self,
125        metadata: &RegionMetadataRef,
126        filters: Arc<Vec<SimpleFilterEvaluator>>,
127    ) -> Box<dyn PrimaryKeyFilter>;
128
129    /// Returns the estimated size of the primary key.
130    fn estimated_size(&self) -> Option<usize> {
131        None
132    }
133
134    /// Returns the encoding type of the primary key.
135    fn encoding(&self) -> PrimaryKeyEncoding;
136
137    /// Returns the dense codec for schema-aware positional access, if applicable.
138    fn as_dense(&self) -> Option<&DensePrimaryKeyCodec> {
139        None
140    }
141
142    /// Decodes the primary key from the given bytes.
143    ///
144    /// Returns a [`CompositeValues`] that follows the primary key ordering.
145    fn decode(&self, bytes: &[u8]) -> Result<CompositeValues>;
146
147    /// Decode the leftmost value from bytes.
148    fn decode_leftmost(&self, bytes: &[u8]) -> Result<Option<Value>>;
149}
150
151/// Builds a primary key codec from region metadata.
152pub fn build_primary_key_codec(region_metadata: &RegionMetadata) -> Arc<dyn PrimaryKeyCodec> {
153    let fields = region_metadata.primary_key_columns().map(|col| {
154        (
155            col.column_id,
156            SortField::new(col.column_schema.data_type.clone()),
157        )
158    });
159    build_primary_key_codec_with_fields(region_metadata.primary_key_encoding, fields)
160}
161
162/// Builds a primary key codec from region metadata.
163pub fn build_primary_key_codec_with_fields(
164    encoding: PrimaryKeyEncoding,
165    fields: impl Iterator<Item = (ColumnId, SortField)>,
166) -> Arc<dyn PrimaryKeyCodec> {
167    match encoding {
168        PrimaryKeyEncoding::Dense => Arc::new(DensePrimaryKeyCodec::with_fields(fields.collect())),
169        PrimaryKeyEncoding::Sparse => {
170            Arc::new(SparsePrimaryKeyCodec::with_fields(fields.collect()))
171        }
172    }
173}
174
175/// Finds the checked boundary of an Option<String>, shared by Dense and Sparse.
176/// This validates framing, not UTF-8; consumers decoding strings validate UTF-8.
177pub(crate) fn encoded_string_len(bytes: &[u8]) -> memcomparable::Result<usize> {
178    match bytes.first().copied().ok_or(memcomparable::Error::Eof)? {
179        0 => return Ok(1),
180        1 => {}
181        marker => return Err(memcomparable::Error::InvalidTagEncoding(marker as usize)),
182    }
183    match bytes.get(1).copied().ok_or(memcomparable::Error::Eof)? {
184        0 => return Ok(2),
185        1 => {}
186        marker => return Err(memcomparable::Error::InvalidBytesEncoding(marker)),
187    }
188    let mut end = 2;
189    loop {
190        let chunk = bytes.get(end..end + 9).ok_or(memcomparable::Error::Eof)?;
191        end += 9;
192        match chunk[8] {
193            1..=8 => return Ok(end),
194            9 => {}
195            marker => return Err(memcomparable::Error::InvalidBytesEncoding(marker)),
196        }
197    }
198}
mito_codec/row_converter.rs

mito_codec/
row_converter.rs
