mito_codec/
key_values.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::collections::HashMap;
16
17use api::v1::{ColumnSchema, Mutation, OpType, Row, Rows};
18use datatypes::prelude::ConcreteDataType;
19use datatypes::value::ValueRef;
20use memcomparable::Deserializer;
21use store_api::codec::{PrimaryKeyEncoding, infer_primary_key_encoding_from_hint};
22use store_api::metadata::RegionMetadata;
23use store_api::storage::SequenceNumber;
24
25use crate::row_converter::{COLUMN_ID_ENCODE_SIZE, SortField};
26
27/// Key value view of a mutation.
28#[derive(Debug)]
29pub struct KeyValues {
30    /// Mutation to read.
31    ///
32    /// This mutation must be a valid mutation and rows in the mutation
33    /// must not be `None`.
34    pub mutation: Mutation,
35    /// Key value read helper.
36    helper: SparseReadRowHelper,
37    /// Primary key encoding hint.
38    primary_key_encoding: PrimaryKeyEncoding,
39}
40
41impl KeyValues {
42    /// Creates [KeyValues] from specific `mutation`.
43    ///
44    /// Returns `None` if `rows` of the `mutation` is `None`.
45    pub fn new(metadata: &RegionMetadata, mutation: Mutation) -> Option<KeyValues> {
46        let rows = mutation.rows.as_ref()?;
47        let primary_key_encoding =
48            infer_primary_key_encoding_from_hint(mutation.write_hint.as_ref());
49        let helper = SparseReadRowHelper::new(metadata, rows, primary_key_encoding);
50
51        Some(KeyValues {
52            mutation,
53            helper,
54            primary_key_encoding,
55        })
56    }
57
58    /// Returns a key value iterator.
59    pub fn iter(&self) -> impl Iterator<Item = KeyValue<'_>> {
60        let rows = self.mutation.rows.as_ref().unwrap();
61        let schema = &rows.schema;
62        rows.rows.iter().enumerate().map(|(idx, row)| {
63            KeyValue {
64                row,
65                schema,
66                helper: &self.helper,
67                sequence: self.mutation.sequence + idx as u64, // Calculate sequence for each row.
68                // Safety: This is a valid mutation.
69                op_type: OpType::try_from(self.mutation.op_type).unwrap(),
70                primary_key_encoding: self.primary_key_encoding,
71            }
72        })
73    }
74
75    /// Returns number of rows.
76    pub fn num_rows(&self) -> usize {
77        // Safety: rows is not None.
78        self.mutation.rows.as_ref().unwrap().rows.len()
79    }
80
81    /// Returns if this container is empty
82    pub fn is_empty(&self) -> bool {
83        self.mutation.rows.is_none()
84    }
85
86    /// Return the max sequence in this container.
87    ///
88    /// When the mutation has no rows, the sequence is the same as the mutation sequence.
89    pub fn max_sequence(&self) -> SequenceNumber {
90        let mut sequence = self.mutation.sequence;
91        let num_rows = self.mutation.rows.as_ref().unwrap().rows.len() as u64;
92        sequence += num_rows;
93        if num_rows > 0 {
94            sequence -= 1;
95        }
96
97        sequence
98    }
99}
100
101/// Key value view of a row.
102///
103/// A key value view divides primary key columns and field (value) columns.
104/// Primary key columns have the same order as region's primary key. Field
105/// columns are ordered by their position in the region schema (The same order
106/// as users defined while creating the region).
107#[derive(Debug, Clone, Copy)]
108pub struct KeyValue<'a> {
109    row: &'a Row,
110    schema: &'a Vec<ColumnSchema>,
111    helper: &'a SparseReadRowHelper,
112    sequence: SequenceNumber,
113    op_type: OpType,
114    primary_key_encoding: PrimaryKeyEncoding,
115}
116
117impl KeyValue<'_> {
118    /// Returns primary key encoding.
119    pub fn primary_key_encoding(&self) -> PrimaryKeyEncoding {
120        self.primary_key_encoding
121    }
122
123    /// Returns the partition key.
124    pub fn partition_key(&self) -> u32 {
125        // TODO(yingwen): refactor this code
126        if self.primary_key_encoding == PrimaryKeyEncoding::Sparse {
127            let Some(primary_key) = self.primary_keys().next() else {
128                return 0;
129            };
130            let key = primary_key.try_into_binary().unwrap().unwrap();
131
132            let mut deserializer = Deserializer::new(key);
133            deserializer.advance(COLUMN_ID_ENCODE_SIZE);
134            let field = SortField::new(ConcreteDataType::uint32_datatype());
135            let table_id = field.deserialize(&mut deserializer).unwrap();
136            table_id.as_value_ref().try_into_u32().unwrap().unwrap()
137        } else {
138            let Some(value) = self.primary_keys().next() else {
139                return 0;
140            };
141
142            value.try_into_u32().unwrap().unwrap()
143        }
144    }
145
146    /// Get primary key columns.
147    pub fn primary_keys(&self) -> impl Iterator<Item = ValueRef<'_>> {
148        self.helper.indices[..self.helper.num_primary_key_column]
149            .iter()
150            .map(|idx| match idx {
151                Some(i) => api::helper::pb_value_to_value_ref(
152                    &self.row.values[*i],
153                    self.schema[*i].datatype_extension.as_ref(),
154                ),
155                None => ValueRef::Null,
156            })
157    }
158
159    /// Get field columns.
160    pub fn fields(&self) -> impl Iterator<Item = ValueRef<'_>> {
161        self.helper.indices[self.helper.num_primary_key_column + 1..]
162            .iter()
163            .map(|idx| match idx {
164                Some(i) => api::helper::pb_value_to_value_ref(
165                    &self.row.values[*i],
166                    self.schema[*i].datatype_extension.as_ref(),
167                ),
168                None => ValueRef::Null,
169            })
170    }
171
172    /// Get timestamp.
173    pub fn timestamp(&self) -> ValueRef<'_> {
174        // Timestamp is primitive, we clone it.
175        let index = self.helper.indices[self.helper.num_primary_key_column].unwrap();
176        api::helper::pb_value_to_value_ref(
177            &self.row.values[index],
178            self.schema[index].datatype_extension.as_ref(),
179        )
180    }
181
182    /// Get number of primary key columns.
183    pub fn num_primary_keys(&self) -> usize {
184        self.helper.num_primary_key_column
185    }
186
187    /// Get number of field columns.
188    pub fn num_fields(&self) -> usize {
189        self.helper.indices.len() - self.helper.num_primary_key_column - 1
190    }
191
192    /// Get sequence.
193    pub fn sequence(&self) -> SequenceNumber {
194        self.sequence
195    }
196
197    /// Get op type.
198    pub fn op_type(&self) -> OpType {
199        self.op_type
200    }
201}
202
203/// Helper to read rows in key, value order for sparse data.
204#[derive(Debug)]
205struct SparseReadRowHelper {
206    /// Key and value column indices.
207    ///
208    /// `indices[..num_primary_key_column]` are primary key columns, `indices[num_primary_key_column]`
209    /// is the timestamp column and remainings are field columns.
210    indices: Vec<Option<usize>>,
211    /// Number of primary key columns.
212    num_primary_key_column: usize,
213}
214
215impl SparseReadRowHelper {
216    /// Creates a [SparseReadRowHelper] for specific `rows`.
217    ///
218    /// # Panics
219    /// Time index column must exist.
220    fn new(
221        metadata: &RegionMetadata,
222        rows: &Rows,
223        primary_key_encoding: PrimaryKeyEncoding,
224    ) -> SparseReadRowHelper {
225        if primary_key_encoding == PrimaryKeyEncoding::Sparse {
226            // Optimized case: when schema has exactly 3 columns (primary key, timestamp, and one field),
227            // we can directly use their indices in order without building an explicit mapping.
228            // The column order is: encoded primary key, timestamp, and field.
229            if rows.schema.len() == 3 {
230                let indices = rows
231                    .schema
232                    .iter()
233                    .enumerate()
234                    .map(|(index, _)| Some(index))
235                    .collect();
236                return SparseReadRowHelper {
237                    indices,
238                    num_primary_key_column: 1,
239                };
240            };
241
242            let mut indices = Vec::with_capacity(rows.schema.len());
243            let name_to_index: HashMap<_, _> = rows
244                .schema
245                .iter()
246                .enumerate()
247                .map(|(index, col)| (&col.column_name, index))
248                .collect();
249            indices.extend(
250                rows.schema[0..2]
251                    .iter()
252                    .enumerate()
253                    .map(|(index, _)| Some(index)),
254            );
255            // Iterate columns and find field columns.
256            for column in metadata.field_columns() {
257                // Get index in request for each field column.
258                let index = name_to_index.get(&column.column_schema.name);
259                indices.push(index.copied());
260            }
261            return SparseReadRowHelper {
262                indices,
263                num_primary_key_column: 1,
264            };
265        }
266        // Build a name to index mapping for rows.
267        let name_to_index: HashMap<_, _> = rows
268            .schema
269            .iter()
270            .enumerate()
271            .map(|(index, col)| (&col.column_name, index))
272            .collect();
273        let mut indices = Vec::with_capacity(metadata.column_metadatas.len());
274
275        // Get primary key indices.
276        for pk_column_id in &metadata.primary_key {
277            // Safety: Id comes from primary key.
278            let column = metadata.column_by_id(*pk_column_id).unwrap();
279            let index = name_to_index.get(&column.column_schema.name);
280            indices.push(index.copied());
281        }
282        // Get timestamp index.
283        // Safety: time index must exist
284        let ts_index = name_to_index
285            .get(&metadata.time_index_column().column_schema.name)
286            .unwrap();
287        indices.push(Some(*ts_index));
288
289        // Iterate columns and find field columns.
290        for column in metadata.field_columns() {
291            // Get index in request for each field column.
292            let index = name_to_index.get(&column.column_schema.name);
293            indices.push(index.copied());
294        }
295
296        SparseReadRowHelper {
297            indices,
298            num_primary_key_column: metadata.primary_key.len(),
299        }
300    }
301}
302
303#[cfg(test)]
304mod tests {
305    use api::v1::{self, ColumnDataType, SemanticType};
306
307    use super::*;
308    use crate::test_util::{TestRegionMetadataBuilder, i64_value};
309
310    const TS_NAME: &str = "ts";
311    const START_SEQ: SequenceNumber = 100;
312
313    /// Creates a region: `ts, k0, k1, ..., v0, v1, ...`
314    fn new_region_metadata(num_tags: usize, num_fields: usize) -> RegionMetadata {
315        TestRegionMetadataBuilder::default()
316            .ts_name(TS_NAME)
317            .num_tags(num_tags)
318            .num_fields(num_fields)
319            .build()
320    }
321
322    /// Creates rows `[ 0, 1, ..., n ] x num_rows`
323    fn new_rows(column_names: &[&str], num_rows: usize) -> Rows {
324        let mut rows = Vec::with_capacity(num_rows);
325        for _ in 0..num_rows {
326            // For simplicity, we use i64 for timestamp millisecond type. This violates the column schema
327            // but it's acceptable for tests.
328            let values: Vec<_> = (0..column_names.len())
329                .map(|idx| i64_value(idx as i64))
330                .collect();
331            rows.push(Row { values });
332        }
333
334        let schema = column_names
335            .iter()
336            .map(|column_name| {
337                let datatype = if *column_name == TS_NAME {
338                    ColumnDataType::TimestampMillisecond as i32
339                } else {
340                    ColumnDataType::Int64 as i32
341                };
342                let semantic_type = if column_name.starts_with('k') {
343                    SemanticType::Tag as i32
344                } else if column_name.starts_with('v') {
345                    SemanticType::Field as i32
346                } else {
347                    SemanticType::Timestamp as i32
348                };
349                v1::ColumnSchema {
350                    column_name: column_name.to_string(),
351                    datatype,
352                    semantic_type,
353                    ..Default::default()
354                }
355            })
356            .collect();
357
358        Rows { rows, schema }
359    }
360
361    fn new_mutation(column_names: &[&str], num_rows: usize) -> Mutation {
362        let rows = new_rows(column_names, num_rows);
363        Mutation {
364            op_type: OpType::Put as i32,
365            sequence: START_SEQ,
366            rows: Some(rows),
367            write_hint: None,
368        }
369    }
370
371    fn check_key_values(
372        kvs: &KeyValues,
373        num_rows: usize,
374        keys: &[Option<i64>],
375        ts: i64,
376        values: &[Option<i64>],
377    ) {
378        assert_eq!(num_rows, kvs.num_rows());
379        let expect_ts = ValueRef::Int64(ts);
380        for (expect_seq, kv) in (START_SEQ..).zip(kvs.iter()) {
381            assert_eq!(expect_seq, kv.sequence());
382            assert_eq!(OpType::Put, kv.op_type);
383            assert_eq!(keys.len(), kv.num_primary_keys());
384            assert_eq!(values.len(), kv.num_fields());
385
386            assert_eq!(expect_ts, kv.timestamp());
387            let expect_keys: Vec<_> = keys.iter().map(|k| ValueRef::from(*k)).collect();
388            let actual_keys: Vec<_> = kv.primary_keys().collect();
389            assert_eq!(expect_keys, actual_keys);
390            let expect_values: Vec<_> = values.iter().map(|v| ValueRef::from(*v)).collect();
391            let actual_values: Vec<_> = kv.fields().collect();
392            assert_eq!(expect_values, actual_values);
393        }
394    }
395
396    #[test]
397    fn test_empty_key_values() {
398        let meta = new_region_metadata(1, 1);
399        let mutation = Mutation {
400            op_type: OpType::Put as i32,
401            sequence: 100,
402            rows: None,
403            write_hint: None,
404        };
405        let kvs = KeyValues::new(&meta, mutation);
406        assert!(kvs.is_none());
407    }
408
409    #[test]
410    fn test_ts_only() {
411        let meta = new_region_metadata(0, 0);
412        let mutation = new_mutation(&["ts"], 2);
413        let kvs = KeyValues::new(&meta, mutation).unwrap();
414        check_key_values(&kvs, 2, &[], 0, &[]);
415    }
416
417    #[test]
418    fn test_no_field() {
419        let meta = new_region_metadata(2, 0);
420        // The value of each row:
421        // k1=0, ts=1, k0=2,
422        let mutation = new_mutation(&["k1", "ts", "k0"], 3);
423        let kvs = KeyValues::new(&meta, mutation).unwrap();
424        // KeyValues
425        // keys: [k0=2, k1=0]
426        // ts: 1,
427        check_key_values(&kvs, 3, &[Some(2), Some(0)], 1, &[]);
428    }
429
430    #[test]
431    fn test_no_tag() {
432        let meta = new_region_metadata(0, 2);
433        // The value of each row:
434        // v1=0, v0=1, ts=2,
435        let mutation = new_mutation(&["v1", "v0", "ts"], 3);
436        let kvs = KeyValues::new(&meta, mutation).unwrap();
437        // KeyValues (note that v0 is in front of v1 in region schema)
438        // ts: 2,
439        // fields: [v0=1, v1=0]
440        check_key_values(&kvs, 3, &[], 2, &[Some(1), Some(0)]);
441    }
442
443    #[test]
444    fn test_tag_field() {
445        let meta = new_region_metadata(2, 2);
446        // The value of each row:
447        // k0=0, v0=1, ts=2, k1=3, v1=4,
448        let mutation = new_mutation(&["k0", "v0", "ts", "k1", "v1"], 3);
449        let kvs = KeyValues::new(&meta, mutation).unwrap();
450        // KeyValues
451        // keys: [k0=0, k1=3]
452        // ts: 2,
453        // fields: [v0=1, v1=4]
454        check_key_values(&kvs, 3, &[Some(0), Some(3)], 2, &[Some(1), Some(4)]);
455    }
456
457    #[test]
458    fn test_sparse_field() {
459        let meta = new_region_metadata(2, 2);
460        // The value of each row:
461        // k0=0, v0=1, ts=2, k1=3, (v1 will be null)
462        let mutation = new_mutation(&["k0", "v0", "ts", "k1"], 3);
463        let kvs = KeyValues::new(&meta, mutation).unwrap();
464        // KeyValues
465        // keys: [k0=0, k1=3]
466        // ts: 2,
467        // fields: [v0=1, v1=null]
468        check_key_values(&kvs, 3, &[Some(0), Some(3)], 2, &[Some(1), None]);
469    }
470
471    #[test]
472    fn test_sparse_tag_field() {
473        let meta = new_region_metadata(2, 2);
474        // The value of each row:
475        // k0 = 0, v0=1, ts=2, (k1, v1 will be null)
476        let mutation = new_mutation(&["k0", "v0", "ts"], 3);
477        let kvs = KeyValues::new(&meta, mutation).unwrap();
478        // KeyValues
479        // keys: [k0=0, k1=null]
480        // ts: 2,
481        // fields: [v0=1, v1=null]
482        check_key_values(&kvs, 3, &[Some(0), None], 2, &[Some(1), None]);
483    }
484}
mito_codec/key_values.rs

mito_codec/
key_values.rs
