metric_engine/
row_modifier.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::{BTreeMap, HashMap};
16use std::hash::Hasher;
17
18use api::v1::value::ValueData;
19use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, SemanticType, Value};
20use datatypes::value::ValueRef;
21use fxhash::FxHasher;
22use mito_codec::row_converter::SparsePrimaryKeyCodec;
23use smallvec::SmallVec;
24use snafu::ResultExt;
25use store_api::codec::PrimaryKeyEncoding;
26use store_api::metadata::ColumnMetadata;
27use store_api::metric_engine_consts::{
28    DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME,
29};
30use store_api::storage::consts::{PRIMARY_KEY_COLUMN_NAME, ReservedColumnId};
31use store_api::storage::{ColumnId, TableId};
32
33use crate::error::{EncodePrimaryKeySnafu, Result, TableIdCountMismatchSnafu};
34
35/// A row modifier modifies [`Rows`].
36///
37/// - For [`PrimaryKeyEncoding::Sparse`] encoding,
38///   it replaces the primary key columns with the encoded primary key column(`__primary_key`).
39///
40/// - For [`PrimaryKeyEncoding::Dense`] encoding,
41///   it adds two columns(`__table_id`, `__tsid`) to the row.
42pub struct RowModifier {
43    codec: SparsePrimaryKeyCodec,
44}
45
46/// Table id input for row modification.
47#[derive(Clone, Copy)]
48pub(crate) enum TableIdInput<'a> {
49    Single(TableId),
50    Batch(&'a [TableId]),
51}
52
53impl<'a> TableIdInput<'a> {
54    fn table_id_for_row(&self, row_idx: usize) -> TableId {
55        match self {
56            TableIdInput::Single(table_id) => *table_id,
57            TableIdInput::Batch(table_ids) => table_ids[row_idx],
58        }
59    }
60}
61
62impl Default for RowModifier {
63    fn default() -> Self {
64        Self {
65            codec: SparsePrimaryKeyCodec::schemaless(),
66        }
67    }
68}
69
70impl RowModifier {
71    /// Modify rows with the given primary key encoding and table ids.
72    pub(crate) fn modify_rows(
73        &self,
74        iter: RowsIter,
75        table_ids: TableIdInput<'_>,
76        encoding: PrimaryKeyEncoding,
77    ) -> Result<Rows> {
78        let row_count = iter.rows.rows.len();
79        Self::validate_table_id_count(table_ids, row_count)?;
80        match encoding {
81            PrimaryKeyEncoding::Sparse => self.modify_rows_sparse(iter, table_ids),
82            PrimaryKeyEncoding::Dense => self.modify_rows_dense(iter, table_ids),
83        }
84    }
85
86    /// Modifies rows with sparse primary key encoding.
87    /// It replaces the primary key columns with the encoded primary key column(`__primary_key`).
88    fn modify_rows_sparse(&self, mut iter: RowsIter, table_ids: TableIdInput<'_>) -> Result<Rows> {
89        let num_column = iter.rows.schema.len();
90        let num_primary_key_column = iter.index.num_primary_key_column;
91        // num_output_column = remaining columns(fields columns + timestamp column) + 1 (encoded primary key column)
92        let num_output_column = num_column - num_primary_key_column + 1;
93
94        let mut buffer = vec![];
95
96        for (row_index, mut row_iter) in iter.iter_mut().enumerate() {
97            let table_id = table_ids.table_id_for_row(row_index);
98            let (table_id_value, tsid) = Self::fill_internal_columns(table_id, &row_iter);
99            let mut values = Vec::with_capacity(num_output_column);
100            buffer.clear();
101            let internal_columns = [
102                (
103                    ReservedColumnId::table_id(),
104                    api::helper::pb_value_to_value_ref(&table_id_value, None),
105                ),
106                (
107                    ReservedColumnId::tsid(),
108                    api::helper::pb_value_to_value_ref(&tsid, None),
109                ),
110            ];
111            self.codec
112                .encode_to_vec(internal_columns.into_iter(), &mut buffer)
113                .context(EncodePrimaryKeySnafu)?;
114            self.codec
115                .encode_to_vec(row_iter.primary_keys(), &mut buffer)
116                .context(EncodePrimaryKeySnafu)?;
117
118            values.push(ValueData::BinaryValue(buffer.clone()).into());
119            values.extend(row_iter.remaining());
120            // Replace the row with the encoded row
121            *row_iter.row = Row { values };
122        }
123
124        // Update the schema
125        let mut schema = Vec::with_capacity(num_output_column);
126        schema.push(ColumnSchema {
127            column_name: PRIMARY_KEY_COLUMN_NAME.to_string(),
128            datatype: ColumnDataType::Binary as i32,
129            semantic_type: SemanticType::Tag as _,
130            datatype_extension: None,
131            options: None,
132        });
133        schema.extend(iter.remaining_columns());
134        iter.rows.schema = schema;
135
136        Ok(iter.rows)
137    }
138
139    /// Modifies rows with dense primary key encoding.
140    /// It adds two columns(`__table_id`, `__tsid`) to the row.
141    fn modify_rows_dense(&self, mut iter: RowsIter, table_ids: TableIdInput<'_>) -> Result<Rows> {
142        // add table_name column
143        iter.rows.schema.push(ColumnSchema {
144            column_name: DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
145            datatype: ColumnDataType::Uint32 as i32,
146            semantic_type: SemanticType::Tag as _,
147            datatype_extension: None,
148            options: None,
149        });
150        // add tsid column
151        iter.rows.schema.push(ColumnSchema {
152            column_name: DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
153            datatype: ColumnDataType::Uint64 as i32,
154            semantic_type: SemanticType::Tag as _,
155            datatype_extension: None,
156            options: None,
157        });
158        for (row_index, row_iter) in iter.iter_mut().enumerate() {
159            let table_id = table_ids.table_id_for_row(row_index);
160            let (table_id_value, tsid) = Self::fill_internal_columns(table_id, &row_iter);
161            row_iter.row.values.push(table_id_value);
162            row_iter.row.values.push(tsid);
163        }
164
165        Ok(iter.rows)
166    }
167
168    fn validate_table_id_count(table_ids: TableIdInput<'_>, row_count: usize) -> Result<()> {
169        if let TableIdInput::Batch(table_ids) = table_ids
170            && table_ids.len() != row_count
171        {
172            return TableIdCountMismatchSnafu {
173                expected: row_count,
174                actual: table_ids.len(),
175            }
176            .fail();
177        }
178        Ok(())
179    }
180
181    /// Fills internal columns of a row with table name and a hash of tag values.
182    pub fn fill_internal_columns(table_id: TableId, iter: &RowIter<'_>) -> (Value, Value) {
183        let ts_id = if !iter.has_null_labels() {
184            // No null labels in row, we can safely reuse the precomputed label name hash.
185            let mut ts_id_gen = TsidGenerator::new(iter.index.label_name_hash);
186            for (_, value) in iter.primary_keys_with_name() {
187                // The type is checked before. So only null is ignored.
188                if let Some(ValueData::StringValue(string)) = &value.value_data {
189                    ts_id_gen.write_str(string);
190                } else {
191                    unreachable!(
192                        "Should not contain null or non-string value: {:?}, table id: {}",
193                        value, table_id
194                    );
195                }
196            }
197            ts_id_gen.finish()
198        } else {
199            TsidGenerator::from_sorted_tags(iter.primary_keys_with_name().filter_map(
200                |(name, value)| match &value.value_data {
201                    Some(ValueData::StringValue(value)) => Some((name.as_str(), value.as_str())),
202                    _ => None,
203                },
204            ))
205        };
206
207        (
208            ValueData::U32Value(table_id).into(),
209            ValueData::U64Value(ts_id).into(),
210        )
211    }
212}
213
214/// Tsid generator.
215#[derive(Default)]
216pub struct TsidGenerator {
217    hasher: FxHasher,
218}
219
220impl TsidGenerator {
221    /// Computes the metric TSID from non-null tags sorted by name, retaining empty values.
222    pub fn from_sorted_tags<'a>(tags: impl Iterator<Item = (&'a str, &'a str)> + Clone) -> u64 {
223        let mut names = Self::default();
224        for (name, _) in tags.clone() {
225            names.write_str(name);
226        }
227        let mut values = Self::new(names.finish());
228        for (_, value) in tags {
229            values.write_str(value);
230        }
231        values.finish()
232    }
233
234    pub fn new(label_name_hash: u64) -> Self {
235        let mut hasher = FxHasher::default();
236        hasher.write_u64(label_name_hash);
237        Self { hasher }
238    }
239
240    /// Writes a label pair to the generator.
241    pub fn write_str(&mut self, value: &str) {
242        self.hasher.write(value.as_bytes());
243        self.hasher.write_u8(0xff);
244    }
245
246    /// Generates a new TSID.
247    pub fn finish(&mut self) -> u64 {
248        self.hasher.finish()
249    }
250}
251
252/// Index of a value.
253#[derive(Debug, Clone, Copy)]
254struct ValueIndex {
255    column_id: ColumnId,
256    index: usize,
257}
258
259/// Index of a row.
260struct IterIndex {
261    indices: Vec<ValueIndex>,
262    num_primary_key_column: usize,
263    /// Precomputed hash for label names.
264    label_name_hash: u64,
265}
266
267impl IterIndex {
268    fn new(
269        row_schema: &[ColumnSchema],
270        physical_columns: &HashMap<String, ColumnMetadata>,
271    ) -> Self {
272        let mut reserved_indices = SmallVec::<[ValueIndex; 2]>::new();
273        // Uses BTreeMap to keep the primary key column name order (lexicographical)
274        let mut primary_key_indices = BTreeMap::new();
275        let mut field_indices = SmallVec::<[ValueIndex; 1]>::new();
276        let mut ts_index = None;
277        for (idx, col) in row_schema.iter().enumerate() {
278            match col.semantic_type() {
279                SemanticType::Tag => match col.column_name.as_str() {
280                    DATA_SCHEMA_TABLE_ID_COLUMN_NAME => {
281                        reserved_indices.push(ValueIndex {
282                            column_id: ReservedColumnId::table_id(),
283                            index: idx,
284                        });
285                    }
286                    DATA_SCHEMA_TSID_COLUMN_NAME => {
287                        reserved_indices.push(ValueIndex {
288                            column_id: ReservedColumnId::tsid(),
289                            index: idx,
290                        });
291                    }
292                    _ => {
293                        // Inserts primary key column name follower the column name order (lexicographical)
294                        primary_key_indices.insert(
295                            col.column_name.as_str(),
296                            ValueIndex {
297                                column_id: physical_columns
298                                    .get(&col.column_name)
299                                    .unwrap()
300                                    .column_id,
301                                index: idx,
302                            },
303                        );
304                    }
305                },
306                SemanticType::Field => {
307                    field_indices.push(ValueIndex {
308                        column_id: physical_columns.get(&col.column_name).unwrap().column_id,
309                        index: idx,
310                    });
311                }
312                SemanticType::Timestamp => {
313                    ts_index = Some(ValueIndex {
314                        column_id: physical_columns.get(&col.column_name).unwrap().column_id,
315                        index: idx,
316                    });
317                }
318            }
319        }
320        let num_primary_key_column = primary_key_indices.len() + reserved_indices.len();
321        let mut indices = Vec::with_capacity(num_primary_key_column + 2);
322        indices.extend(reserved_indices);
323        let mut label_name_hasher = TsidGenerator::default();
324        for (pk_name, pk_index) in primary_key_indices {
325            // primary_key_indices already sorted.
326            label_name_hasher.write_str(pk_name);
327            indices.push(pk_index);
328        }
329        let label_name_hash = label_name_hasher.finish();
330
331        indices.extend(ts_index);
332        indices.extend(field_indices);
333        IterIndex {
334            indices,
335            num_primary_key_column,
336            label_name_hash,
337        }
338    }
339}
340
341/// Iterator of rows.
342pub struct RowsIter {
343    rows: Rows,
344    index: IterIndex,
345}
346
347impl RowsIter {
348    pub fn new(rows: Rows, physical_columns: &HashMap<String, ColumnMetadata>) -> Self {
349        let index: IterIndex = IterIndex::new(&rows.schema, physical_columns);
350        Self { rows, index }
351    }
352
353    /// Returns the iterator of rows.
354    pub fn iter_mut(&mut self) -> impl Iterator<Item = RowIter<'_>> {
355        self.rows.rows.iter_mut().map(|row| RowIter {
356            row,
357            index: &self.index,
358            schema: &self.rows.schema,
359        })
360    }
361
362    /// Returns the remaining columns.
363    fn remaining_columns(&mut self) -> impl Iterator<Item = ColumnSchema> + '_ {
364        self.index.indices[self.index.num_primary_key_column..]
365            .iter()
366            .map(|idx| std::mem::take(&mut self.rows.schema[idx.index]))
367    }
368}
369
370/// Iterator of a row.
371pub struct RowIter<'a> {
372    row: &'a mut Row,
373    index: &'a IterIndex,
374    schema: &'a Vec<ColumnSchema>,
375}
376
377impl RowIter<'_> {
378    /// Returns the primary keys with their names.
379    fn primary_keys_with_name(&self) -> impl Iterator<Item = (&String, &Value)> + Clone {
380        self.index.indices[..self.index.num_primary_key_column]
381            .iter()
382            .map(|idx| {
383                (
384                    &self.schema[idx.index].column_name,
385                    &self.row.values[idx.index],
386                )
387            })
388    }
389
390    /// Returns true if any label in current row is null.
391    fn has_null_labels(&self) -> bool {
392        self.index.indices[..self.index.num_primary_key_column]
393            .iter()
394            .any(|idx| self.row.values[idx.index].value_data.is_none())
395    }
396
397    /// Returns the primary keys.
398    pub fn primary_keys(&self) -> impl Iterator<Item = (ColumnId, ValueRef<'_>)> {
399        self.index.indices[..self.index.num_primary_key_column]
400            .iter()
401            .map(|idx| {
402                (
403                    idx.column_id,
404                    api::helper::pb_value_to_value_ref(
405                        &self.row.values[idx.index],
406                        self.schema[idx.index].datatype_extension.as_ref(),
407                    ),
408                )
409            })
410    }
411
412    /// Returns the remaining columns.
413    fn remaining(&mut self) -> impl Iterator<Item = Value> + '_ {
414        self.index.indices[self.index.num_primary_key_column..]
415            .iter()
416            .map(|idx| std::mem::take(&mut self.row.values[idx.index]))
417    }
418
419    /// Returns value at given offset.
420    /// # Panics
421    /// Panics if offset out-of-bound
422    pub fn value_at(&self, idx: usize) -> &Value {
423        &self.row.values[idx]
424    }
425}
426
427#[cfg(test)]
428mod tests {
429    use std::collections::HashMap;
430
431    use api::v1::{Row, Rows};
432    use store_api::codec::PrimaryKeyEncoding;
433
434    use super::*;
435    use crate::error::Error;
436
437    fn test_schema() -> Vec<ColumnSchema> {
438        vec![
439            ColumnSchema {
440                column_name: "namespace".to_string(),
441                datatype: ColumnDataType::String as i32,
442                semantic_type: SemanticType::Tag as _,
443                datatype_extension: None,
444                options: None,
445            },
446            ColumnSchema {
447                column_name: "host".to_string(),
448                datatype: ColumnDataType::String as i32,
449                semantic_type: SemanticType::Tag as _,
450                datatype_extension: None,
451                options: None,
452            },
453        ]
454    }
455
456    fn test_row(v1: &str, v2: &str) -> Row {
457        Row {
458            values: vec![
459                ValueData::StringValue(v1.to_string()).into(),
460                ValueData::StringValue(v2.to_string()).into(),
461            ],
462        }
463    }
464
465    fn make_info(name: &str, column_id: ColumnId) -> ColumnMetadata {
466        ColumnMetadata {
467            column_schema: datatypes::schema::ColumnSchema::new(
468                name.to_string(),
469                datatypes::prelude::ConcreteDataType::string_datatype(),
470                false,
471            ),
472            semantic_type: SemanticType::Tag,
473            column_id,
474        }
475    }
476
477    fn test_name_to_column_id() -> HashMap<String, ColumnMetadata> {
478        HashMap::from([
479            ("namespace".to_string(), make_info("namespace", 1)),
480            ("host".to_string(), make_info("host", 2)),
481        ])
482    }
483
484    #[test]
485    fn test_encode_sparse() {
486        let name_to_column_id = test_name_to_column_id();
487        let encoder = RowModifier::default();
488        let table_id = 1025;
489        let schema = test_schema();
490        let row = test_row("greptimedb", "127.0.0.1");
491        let rows = Rows {
492            schema,
493            rows: vec![row],
494        };
495        let rows_iter = RowsIter::new(rows, &name_to_column_id);
496        let result = encoder
497            .modify_rows(
498                rows_iter,
499                TableIdInput::Single(table_id),
500                PrimaryKeyEncoding::Sparse,
501            )
502            .unwrap();
503        assert_eq!(result.rows[0].values.len(), 1);
504        let encoded_primary_key = vec![
505            128, 0, 0, 4, 1, 0, 0, 4, 1, 128, 0, 0, 3, 1, 37, 196, 242, 181, 117, 224, 7, 137, 0,
506            0, 0, 2, 1, 1, 49, 50, 55, 46, 48, 46, 48, 46, 9, 49, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,
507            1, 1, 1, 103, 114, 101, 112, 116, 105, 109, 101, 9, 100, 98, 0, 0, 0, 0, 0, 0, 2,
508        ];
509        assert_eq!(
510            result.rows[0].values[0],
511            ValueData::BinaryValue(encoded_primary_key).into()
512        );
513        assert_eq!(result.schema, expected_sparse_schema());
514    }
515
516    fn expected_sparse_schema() -> Vec<ColumnSchema> {
517        vec![ColumnSchema {
518            column_name: PRIMARY_KEY_COLUMN_NAME.to_string(),
519            datatype: ColumnDataType::Binary as i32,
520            semantic_type: SemanticType::Tag as _,
521            datatype_extension: None,
522            options: None,
523        }]
524    }
525
526    fn expected_dense_schema() -> Vec<ColumnSchema> {
527        vec![
528            ColumnSchema {
529                column_name: "namespace".to_string(),
530                datatype: ColumnDataType::String as i32,
531                semantic_type: SemanticType::Tag as _,
532                datatype_extension: None,
533                options: None,
534            },
535            ColumnSchema {
536                column_name: "host".to_string(),
537                datatype: ColumnDataType::String as i32,
538                semantic_type: SemanticType::Tag as _,
539                datatype_extension: None,
540                options: None,
541            },
542            ColumnSchema {
543                column_name: DATA_SCHEMA_TABLE_ID_COLUMN_NAME.to_string(),
544                datatype: ColumnDataType::Uint32 as i32,
545                semantic_type: SemanticType::Tag as _,
546                datatype_extension: None,
547                options: None,
548            },
549            ColumnSchema {
550                column_name: DATA_SCHEMA_TSID_COLUMN_NAME.to_string(),
551                datatype: ColumnDataType::Uint64 as i32,
552                semantic_type: SemanticType::Tag as _,
553                datatype_extension: None,
554                options: None,
555            },
556        ]
557    }
558
559    #[test]
560    fn test_encode_dense() {
561        let name_to_column_id = test_name_to_column_id();
562        let encoder = RowModifier::default();
563        let table_id = 1025;
564        let schema = test_schema();
565        let row = test_row("greptimedb", "127.0.0.1");
566        let rows = Rows {
567            schema,
568            rows: vec![row],
569        };
570        let rows_iter = RowsIter::new(rows, &name_to_column_id);
571        let result = encoder
572            .modify_rows(
573                rows_iter,
574                TableIdInput::Single(table_id),
575                PrimaryKeyEncoding::Dense,
576            )
577            .unwrap();
578        assert_eq!(
579            result.rows[0].values[0],
580            ValueData::StringValue("greptimedb".to_string()).into()
581        );
582        assert_eq!(
583            result.rows[0].values[1],
584            ValueData::StringValue("127.0.0.1".to_string()).into()
585        );
586        assert_eq!(result.rows[0].values[2], ValueData::U32Value(1025).into());
587        assert_eq!(
588            result.rows[0].values[3],
589            ValueData::U64Value(2721566936019240841).into()
590        );
591        assert_eq!(result.schema, expected_dense_schema());
592    }
593
594    #[test]
595    fn test_table_id_count_mismatch() {
596        let name_to_column_id = test_name_to_column_id();
597        let encoder = RowModifier::default();
598        let schema = test_schema();
599        let rows = Rows {
600            schema,
601            rows: vec![test_row("a", "b"), test_row("c", "d")],
602        };
603        let rows_iter = RowsIter::new(rows, &name_to_column_id);
604        let table_ids = [1025];
605        let err = encoder
606            .modify_rows(
607                rows_iter,
608                TableIdInput::Batch(&table_ids),
609                PrimaryKeyEncoding::Dense,
610            )
611            .unwrap_err();
612        assert!(matches!(
613            err,
614            Error::TableIdCountMismatch {
615                expected: 2,
616                actual: 1,
617                ..
618            }
619        ));
620    }
621
622    #[test]
623    fn test_fill_internal_columns() {
624        let name_to_column_id = test_name_to_column_id();
625        let table_id = 1025;
626        let schema = test_schema();
627        let row = test_row("greptimedb", "127.0.0.1");
628        let rows = Rows {
629            schema,
630            rows: vec![row],
631        };
632        let mut rows_iter = RowsIter::new(rows, &name_to_column_id);
633        let row_iter = rows_iter.iter_mut().next().unwrap();
634        let (encoded_table_id, tsid) = RowModifier::fill_internal_columns(table_id, &row_iter);
635        assert_eq!(encoded_table_id, ValueData::U32Value(1025).into());
636        assert_eq!(tsid, ValueData::U64Value(2721566936019240841).into());
637
638        // Change the column order
639        let schema = vec![
640            ColumnSchema {
641                column_name: "host".to_string(),
642                datatype: ColumnDataType::String as i32,
643                semantic_type: SemanticType::Tag as _,
644                datatype_extension: None,
645                options: None,
646            },
647            ColumnSchema {
648                column_name: "namespace".to_string(),
649                datatype: ColumnDataType::String as i32,
650                semantic_type: SemanticType::Tag as _,
651                datatype_extension: None,
652                options: None,
653            },
654        ];
655        let row = test_row("127.0.0.1", "greptimedb");
656        let rows = Rows {
657            schema,
658            rows: vec![row],
659        };
660        let mut rows_iter = RowsIter::new(rows, &name_to_column_id);
661        let row_iter = rows_iter.iter_mut().next().unwrap();
662        let (encoded_table_id, tsid) = RowModifier::fill_internal_columns(table_id, &row_iter);
663        assert_eq!(encoded_table_id, ValueData::U32Value(1025).into());
664        assert_eq!(tsid, ValueData::U64Value(2721566936019240841).into());
665    }
666
667    /// Helper function to create a schema with multiple label columns
668    fn create_multi_label_schema(labels: &[&str]) -> Vec<ColumnSchema> {
669        labels
670            .iter()
671            .map(|name| ColumnSchema {
672                column_name: name.to_string(),
673                datatype: ColumnDataType::String as i32,
674                semantic_type: SemanticType::Tag as _,
675                datatype_extension: None,
676                options: None,
677            })
678            .collect()
679    }
680
681    /// Helper function to create a name_to_column_id map
682    fn create_name_to_column_id(labels: &[&str]) -> HashMap<String, ColumnMetadata> {
683        labels
684            .iter()
685            .enumerate()
686            .map(|(idx, name)| (name.to_string(), make_info(name, idx as ColumnId + 1)))
687            .collect()
688    }
689
690    /// Helper function to create a row with string values
691    fn create_row_with_values(values: &[&str]) -> Row {
692        Row {
693            values: values
694                .iter()
695                .map(|v| ValueData::StringValue(v.to_string()).into())
696                .collect(),
697        }
698    }
699
700    /// Helper function to create a row with some null values
701    fn create_row_with_nulls(values: &[Option<&str>]) -> Row {
702        Row {
703            values: values
704                .iter()
705                .map(|v| {
706                    v.map(|s| ValueData::StringValue(s.to_string()).into())
707                        .unwrap_or(Value { value_data: None })
708                })
709                .collect(),
710        }
711    }
712
713    /// Helper function to extract TSID from a row
714    fn extract_tsid(
715        schema: Vec<ColumnSchema>,
716        row: Row,
717        name_to_column_id: &HashMap<String, ColumnMetadata>,
718        table_id: TableId,
719    ) -> u64 {
720        let rows = Rows {
721            schema,
722            rows: vec![row],
723        };
724        let mut rows_iter = RowsIter::new(rows, name_to_column_id);
725        let row_iter = rows_iter.iter_mut().next().unwrap();
726        let (_, tsid_value) = RowModifier::fill_internal_columns(table_id, &row_iter);
727        match tsid_value.value_data {
728            Some(ValueData::U64Value(tsid)) => tsid,
729            _ => panic!("Expected U64Value for TSID"),
730        }
731    }
732
733    #[test]
734    fn test_tsid_same_for_different_label_orders() {
735        // Test that rows with the same label name-value pairs but in different orders
736        // produce the same TSID
737        let table_id = 1025;
738
739        // Schema 1: a, b, c
740        let schema1 = create_multi_label_schema(&["a", "b", "c"]);
741        let name_to_column_id1 = create_name_to_column_id(&["a", "b", "c"]);
742        let row1 = create_row_with_values(&["A", "B", "C"]);
743        let tsid1 = extract_tsid(schema1, row1, &name_to_column_id1, table_id);
744
745        // Schema 2: b, a, c (different order)
746        let schema2 = create_multi_label_schema(&["b", "a", "c"]);
747        let name_to_column_id2 = create_name_to_column_id(&["a", "b", "c"]);
748        let row2 = create_row_with_values(&["B", "A", "C"]);
749        let tsid2 = extract_tsid(schema2, row2, &name_to_column_id2, table_id);
750
751        // Schema 3: c, b, a (another different order)
752        let schema3 = create_multi_label_schema(&["c", "b", "a"]);
753        let name_to_column_id3 = create_name_to_column_id(&["a", "b", "c"]);
754        let row3 = create_row_with_values(&["C", "B", "A"]);
755        let tsid3 = extract_tsid(schema3, row3, &name_to_column_id3, table_id);
756
757        // All should have the same TSID since label names are sorted lexicographically
758        // and we're using the same label name-value pairs
759        assert_eq!(
760            tsid1, tsid2,
761            "TSID should be same for different column orders"
762        );
763        assert_eq!(
764            tsid2, tsid3,
765            "TSID should be same for different column orders"
766        );
767    }
768
769    #[test]
770    fn test_tsid_same_with_null_labels() {
771        // Test that rows that differ only by null label values produce the same TSID
772        let table_id = 1025;
773
774        // Row 1: a=A, b=B (no nulls, fast path)
775        let schema1 = create_multi_label_schema(&["a", "b"]);
776        let name_to_column_id1 = create_name_to_column_id(&["a", "b"]);
777        let row1 = create_row_with_values(&["A", "B"]);
778        let tsid1 = extract_tsid(schema1, row1, &name_to_column_id1, table_id);
779
780        // Row 2: a=A, b=B, c=null (has null, slow path)
781        let schema2 = create_multi_label_schema(&["a", "b", "c"]);
782        let name_to_column_id2 = create_name_to_column_id(&["a", "b", "c"]);
783        let row2 = create_row_with_nulls(&[Some("A"), Some("B"), None]);
784        let tsid2 = extract_tsid(schema2, row2, &name_to_column_id2, table_id);
785
786        // Both should have the same TSID since null labels are ignored
787        assert_eq!(
788            tsid1, tsid2,
789            "TSID should be same when only difference is null label values"
790        );
791    }
792
793    #[test]
794    fn test_sorted_tags_match_metric_rows() {
795        let names = ["b", "a", "absent"];
796        for values in [
797            [Some("B"), Some(""), Some("C")],
798            [Some("B"), Some(""), None],
799            [None, Some("A"), None],
800            [None, None, None],
801        ] {
802            let tags = names
803                .into_iter()
804                .zip(values)
805                .filter_map(|(name, value)| value.map(|value| (name, value)))
806                .collect::<BTreeMap<_, _>>();
807            let tsid =
808                TsidGenerator::from_sorted_tags(tags.iter().map(|(name, value)| (*name, *value)));
809            assert_eq!(
810                tsid,
811                extract_tsid(
812                    create_multi_label_schema(&names),
813                    create_row_with_nulls(&values),
814                    &create_name_to_column_id(&names),
815                    42
816                )
817            );
818        }
819    }
820
821    #[test]
822    fn test_tsid_same_with_multiple_null_labels() {
823        // Test with multiple null labels
824        let table_id = 1025;
825
826        // Row 1: a=A, b=B (no nulls)
827        let schema1 = create_multi_label_schema(&["a", "b"]);
828        let name_to_column_id1 = create_name_to_column_id(&["a", "b"]);
829        let row1 = create_row_with_values(&["A", "B"]);
830        let tsid1 = extract_tsid(schema1, row1, &name_to_column_id1, table_id);
831
832        // Row 2: a=A, b=B, c=null, d=null (multiple nulls)
833        let schema2 = create_multi_label_schema(&["a", "b", "c", "d"]);
834        let name_to_column_id2 = create_name_to_column_id(&["a", "b", "c", "d"]);
835        let row2 = create_row_with_nulls(&[Some("A"), Some("B"), None, None]);
836        let tsid2 = extract_tsid(schema2, row2, &name_to_column_id2, table_id);
837
838        assert_eq!(
839            tsid1, tsid2,
840            "TSID should be same when only difference is multiple null label values"
841        );
842    }
843
844    #[test]
845    fn test_tsid_different_with_different_non_null_values() {
846        // Test that rows with different non-null values produce different TSIDs
847        let table_id = 1025;
848
849        // Row 1: a=A, b=B
850        let schema1 = create_multi_label_schema(&["a", "b"]);
851        let name_to_column_id1 = create_name_to_column_id(&["a", "b"]);
852        let row1 = create_row_with_values(&["A", "B"]);
853        let tsid1 = extract_tsid(schema1, row1, &name_to_column_id1, table_id);
854
855        // Row 2: a=A, b=C (different value for b)
856        let schema2 = create_multi_label_schema(&["a", "b"]);
857        let name_to_column_id2 = create_name_to_column_id(&["a", "b"]);
858        let row2 = create_row_with_values(&["A", "C"]);
859        let tsid2 = extract_tsid(schema2, row2, &name_to_column_id2, table_id);
860
861        assert_ne!(
862            tsid1, tsid2,
863            "TSID should be different when label values differ"
864        );
865    }
866
867    #[test]
868    fn test_tsid_fast_path_vs_slow_path_consistency() {
869        // Test that fast path (no nulls) and slow path (with nulls) produce
870        // the same TSID for the same non-null label values
871        let table_id = 1025;
872
873        // Fast path: a=A, b=B (no nulls)
874        let schema_fast = create_multi_label_schema(&["a", "b"]);
875        let name_to_column_id_fast = create_name_to_column_id(&["a", "b"]);
876        let row_fast = create_row_with_values(&["A", "B"]);
877        let tsid_fast = extract_tsid(schema_fast, row_fast, &name_to_column_id_fast, table_id);
878
879        // Slow path: a=A, b=B, c=null (has null, triggers slow path)
880        let schema_slow = create_multi_label_schema(&["a", "b", "c"]);
881        let name_to_column_id_slow = create_name_to_column_id(&["a", "b", "c"]);
882        let row_slow = create_row_with_nulls(&[Some("A"), Some("B"), None]);
883        let tsid_slow = extract_tsid(schema_slow, row_slow, &name_to_column_id_slow, table_id);
884
885        assert_eq!(
886            tsid_fast, tsid_slow,
887            "Fast path and slow path should produce same TSID for same non-null values"
888        );
889    }
890
891    #[test]
892    fn test_tsid_with_null_in_middle() {
893        // Test with null in the middle of labels
894        let table_id = 1025;
895
896        // Row 1: a=A, b=B, c=C
897        let schema1 = create_multi_label_schema(&["a", "b", "c"]);
898        let name_to_column_id1 = create_name_to_column_id(&["a", "b", "c"]);
899        let row1 = create_row_with_values(&["A", "B", "C"]);
900        let tsid1 = extract_tsid(schema1, row1, &name_to_column_id1, table_id);
901
902        // Row 2: a=A, b=null, c=C (null in middle)
903        let schema2 = create_multi_label_schema(&["a", "b", "c"]);
904        let name_to_column_id2 = create_name_to_column_id(&["a", "b", "c"]);
905        let row2 = create_row_with_nulls(&[Some("A"), None, Some("C")]);
906        let tsid2 = extract_tsid(schema2, row2, &name_to_column_id2, table_id);
907
908        // Should be different because b is null in row2 but B in row1
909        // Actually wait, let me reconsider - if b is null, it should be ignored
910        // So row2 should be equivalent to a=A, c=C
911        // But row1 is a=A, b=B, c=C, so they should be different
912        assert_ne!(
913            tsid1, tsid2,
914            "TSID should be different when a non-null value becomes null"
915        );
916
917        // Row 3: a=A, c=C (no b at all, equivalent to row2)
918        let schema3 = create_multi_label_schema(&["a", "c"]);
919        let name_to_column_id3 = create_name_to_column_id(&["a", "c"]);
920        let row3 = create_row_with_values(&["A", "C"]);
921        let tsid3 = extract_tsid(schema3, row3, &name_to_column_id3, table_id);
922
923        // Row2 (a=A, b=null, c=C) should be same as row3 (a=A, c=C)
924        assert_eq!(
925            tsid2, tsid3,
926            "TSID should be same when null label is ignored"
927        );
928    }
929
930    #[test]
931    fn test_tsid_all_null_labels() {
932        // Test with all labels being null
933        let table_id = 1025;
934
935        // Row with all nulls
936        let schema = create_multi_label_schema(&["a", "b", "c"]);
937        let name_to_column_id = create_name_to_column_id(&["a", "b", "c"]);
938        let row = create_row_with_nulls(&[None, None, None]);
939        let tsid = extract_tsid(schema.clone(), row, &name_to_column_id, table_id);
940
941        // Should still produce a TSID (based on label names only when all values are null)
942        // This tests that the slow path handles the case where all values are null
943        // The TSID will be based on the label name hash only
944        // Test that it's consistent - same schema with all nulls should produce same TSID
945        let row2 = create_row_with_nulls(&[None, None, None]);
946        let tsid2 = extract_tsid(schema, row2, &name_to_column_id, table_id);
947        assert_eq!(
948            tsid, tsid2,
949            "TSID should be consistent when all label values are null"
950        );
951    }
952}
metric_engine/row_modifier.rs

metric_engine/
row_modifier.rs
