metric_engine/
batch_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::hash::Hasher;
16use std::sync::Arc;
17
18use datatypes::arrow::array::{Array, ArrayRef, BinaryBuilder, UInt64Array};
19use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
20use datatypes::arrow::record_batch::RecordBatch;
21use datatypes::arrow_array::{is_string_null_at, string_array_value_at_index};
22use fxhash::FxHasher;
23use mito_codec::row_converter::SparsePrimaryKeyCodec;
24use snafu::ResultExt;
25use store_api::storage::ColumnId;
26use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
27
28use crate::error::{EncodePrimaryKeySnafu, Result, UnexpectedRequestSnafu};
29
30/// Info about a tag column for TSID computation and sparse primary key encoding.
31#[allow(dead_code)]
32pub struct TagColumnInfo {
33    /// Column name (used for label-name hash).
34    pub name: String,
35    /// Column index in the RecordBatch.
36    pub index: usize,
37    /// Column ID in the physical region.
38    pub column_id: ColumnId,
39}
40
41/// Computes the TSID for each row in a [RecordBatch].
42///
43/// The TSID is a stable hash of the set of labels (tags) present in each row.
44/// It accounts for both the names and values of all non-null tag columns.
45///
46/// # Logic
47/// - If a row has no nulls across all `sorted_tag_columns`, it uses a precomputed hash of all label names.
48/// - If a row has nulls, it dynamically computes a hash of the names of labels that are present (non-null).
49/// - In both cases, it then hashes the values of those present labels in the order specified by `sorted_tag_columns`.
50pub fn compute_tsid_array(
51    batch: &RecordBatch,
52    sorted_tag_columns: &[TagColumnInfo],
53    tag_arrays: &[&ArrayRef],
54) -> UInt64Array {
55    let num_rows = batch.num_rows();
56
57    let label_name_hash = {
58        let mut hasher = FxHasher::default();
59        for tag_col in sorted_tag_columns {
60            hasher.write(tag_col.name.as_bytes());
61            hasher.write_u8(0xff);
62        }
63        hasher.finish()
64    };
65
66    let mut tsid_values = Vec::with_capacity(num_rows);
67    for row in 0..num_rows {
68        let has_null = tag_arrays.iter().any(|arr| is_string_null_at(arr, row));
69
70        let tsid = if !has_null {
71            let mut hasher = FxHasher::default();
72            hasher.write_u64(label_name_hash);
73            for arr in tag_arrays {
74                if let Some(value) = string_array_value_at_index(arr, row) {
75                    hasher.write(value.as_bytes());
76                    hasher.write_u8(0xff);
77                }
78            }
79            hasher.finish()
80        } else {
81            let mut name_hasher = FxHasher::default();
82            for (tc, arr) in sorted_tag_columns.iter().zip(tag_arrays.iter()) {
83                if !is_string_null_at(arr, row) {
84                    name_hasher.write(tc.name.as_bytes());
85                    name_hasher.write_u8(0xff);
86                }
87            }
88            let row_label_hash = name_hasher.finish();
89
90            let mut val_hasher = FxHasher::default();
91            val_hasher.write_u64(row_label_hash);
92            for arr in tag_arrays {
93                if let Some(value) = string_array_value_at_index(arr, row) {
94                    val_hasher.write(value.as_bytes());
95                    val_hasher.write_u8(0xff);
96                }
97            }
98            val_hasher.finish()
99        };
100
101        tsid_values.push(tsid);
102    }
103
104    UInt64Array::from(tsid_values)
105}
106
107fn build_tag_arrays<'a>(
108    batch: &'a RecordBatch,
109    sorted_tag_columns: &[TagColumnInfo],
110) -> Result<Vec<&'a ArrayRef>> {
111    sorted_tag_columns
112        .iter()
113        .map(|tc| {
114            let array = batch.column(tc.index);
115            match array.data_type() {
116                DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Ok(array),
117                DataType::Dictionary(key_type, value_type)
118                    if key_type.is_integer() && value_type.is_string() =>
119                {
120                    Ok(array)
121                }
122                data_type => UnexpectedRequestSnafu {
123                    reason: format!(
124                        "Tag column '{}' must be a string, given: {data_type}",
125                        tc.name
126                    ),
127                }
128                .fail(),
129            }
130        })
131        .collect()
132}
133
134/// Modifies a [RecordBatch] to include a sparse primary key column.
135///
136/// This function transforms the input `batch` into a new `RecordBatch` where the first column
137/// is the generated primary key (named [PRIMARY_KEY_COLUMN_NAME]), followed by columns
138/// indicated by `extra_column_indices`.
139///
140/// The primary key uses a "sparse" encoding, which compactly represents the row's identity
141/// by only including non-null tag values. The encoding, handled by [SparsePrimaryKeyCodec],
142/// consists of:
143/// 1. The `table_id`.
144/// 2. A `tsid` (Time Series ID), which is a hash of the present tags.
145/// 3. The actual non-null tag values paired with their `column_id`.
146///
147/// # Parameters
148/// - `batch`: The source [RecordBatch].
149/// - `table_id`: The ID of the table.
150/// - `sorted_tag_columns`: Metadata for tag columns, used for both TSID computation and PK encoding.
151/// - `extra_column_indices`: Indices of columns from the original batch to keep in the output
152///   (typically the timestamp and value fields).
153pub fn modify_batch_sparse(
154    batch: RecordBatch,
155    table_id: u32,
156    sorted_tag_columns: &[TagColumnInfo],
157    extra_column_indices: &[usize],
158) -> Result<RecordBatch> {
159    let num_rows = batch.num_rows();
160    let codec = SparsePrimaryKeyCodec::schemaless();
161    let tag_arrays = build_tag_arrays(&batch, sorted_tag_columns)?;
162    let tsid_array = compute_tsid_array(&batch, sorted_tag_columns, &tag_arrays);
163
164    let mut pk_builder = BinaryBuilder::with_capacity(num_rows, 0);
165    let mut buffer = Vec::new();
166    for row in 0..num_rows {
167        buffer.clear();
168        codec
169            .encode_internal(table_id, tsid_array.value(row), &mut buffer)
170            .context(EncodePrimaryKeySnafu)?;
171
172        let tags = sorted_tag_columns
173            .iter()
174            .zip(tag_arrays.iter())
175            .filter_map(|(tc, arr)| {
176                string_array_value_at_index(arr, row).map(|value| (tc.column_id, value.as_bytes()))
177            });
178        codec
179            .encode_raw_tag_value(tags, &mut buffer)
180            .context(EncodePrimaryKeySnafu)?;
181
182        pk_builder.append_value(&buffer);
183    }
184
185    let pk_array = pk_builder.finish();
186
187    let mut fields = vec![Arc::new(Field::new(
188        PRIMARY_KEY_COLUMN_NAME,
189        DataType::Binary,
190        false,
191    ))];
192    let mut columns: Vec<Arc<dyn Array>> = vec![Arc::new(pk_array)];
193
194    for &idx in extra_column_indices {
195        fields.push(batch.schema().fields()[idx].clone());
196        columns.push(batch.column(idx).clone());
197    }
198
199    let new_schema = Arc::new(ArrowSchema::new(fields));
200    RecordBatch::try_new(new_schema, columns).map_err(|e| {
201        UnexpectedRequestSnafu {
202            reason: format!("Failed to build modified sparse RecordBatch: {e}"),
203        }
204        .build()
205    })
206}
207
208#[cfg(test)]
209mod tests {
210    use std::collections::HashMap;
211    use std::sync::Arc;
212
213    use api::v1::value::ValueData;
214    use api::v1::{ColumnDataType, ColumnSchema, Row, Rows, SemanticType, Value};
215    use datatypes::arrow::array::{BinaryArray, Int64Array, StringArray};
216    use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema};
217    use datatypes::arrow::record_batch::RecordBatch;
218    use store_api::codec::PrimaryKeyEncoding;
219    use store_api::metadata::ColumnMetadata;
220    use store_api::storage::consts::PRIMARY_KEY_COLUMN_NAME;
221
222    use super::*;
223    use crate::row_modifier::{RowModifier, RowsIter, TableIdInput};
224
225    fn build_sparse_test_batch() -> RecordBatch {
226        let schema = Arc::new(ArrowSchema::new(vec![
227            Field::new("greptime_timestamp", DataType::Int64, false),
228            Field::new("greptime_value", DataType::Float64, true),
229            Field::new("namespace", DataType::Utf8, true),
230            Field::new("host", DataType::Utf8, true),
231        ]));
232        RecordBatch::try_new(
233            schema,
234            vec![
235                Arc::new(Int64Array::from(vec![1000])),
236                Arc::new(datatypes::arrow::array::Float64Array::from(vec![42.0])),
237                Arc::new(StringArray::from(vec!["greptimedb"])),
238                Arc::new(StringArray::from(vec!["127.0.0.1"])),
239            ],
240        )
241        .unwrap()
242    }
243
244    fn sparse_tag_columns() -> Vec<TagColumnInfo> {
245        vec![
246            TagColumnInfo {
247                name: "host".to_string(),
248                index: 3,
249                column_id: 3,
250            },
251            TagColumnInfo {
252                name: "namespace".to_string(),
253                index: 2,
254                column_id: 2,
255            },
256        ]
257    }
258
259    #[test]
260    fn test_compute_tsid_basic() {
261        let schema = Arc::new(ArrowSchema::new(vec![
262            Field::new("namespace", DataType::Utf8, true),
263            Field::new("host", DataType::Utf8, true),
264        ]));
265        let batch = RecordBatch::try_new(
266            schema,
267            vec![
268                Arc::new(StringArray::from(vec!["greptimedb"])),
269                Arc::new(StringArray::from(vec!["127.0.0.1"])),
270            ],
271        )
272        .unwrap();
273
274        let tag_columns: Vec<TagColumnInfo> = vec![
275            TagColumnInfo {
276                name: "host".to_string(),
277                index: 1,
278                column_id: 2,
279            },
280            TagColumnInfo {
281                name: "namespace".to_string(),
282                index: 0,
283                column_id: 1,
284            },
285        ];
286        let tag_arrays = build_tag_arrays(&batch, &tag_columns).unwrap();
287        let tsid_array = compute_tsid_array(&batch, &tag_columns, &tag_arrays);
288
289        assert_eq!(tsid_array.value(0), 2721566936019240841);
290    }
291
292    #[test]
293    fn test_compute_tsid_with_nulls() {
294        let schema = Arc::new(ArrowSchema::new(vec![
295            Field::new("a", DataType::Utf8, true),
296            Field::new("b", DataType::Utf8, true),
297        ]));
298        let batch_no_null = RecordBatch::try_new(
299            schema.clone(),
300            vec![
301                Arc::new(StringArray::from(vec!["A"])),
302                Arc::new(StringArray::from(vec!["B"])),
303            ],
304        )
305        .unwrap();
306        let tag_cols_2: Vec<TagColumnInfo> = vec![
307            TagColumnInfo {
308                name: "a".to_string(),
309                index: 0,
310                column_id: 1,
311            },
312            TagColumnInfo {
313                name: "b".to_string(),
314                index: 1,
315                column_id: 2,
316            },
317        ];
318        let tag_arrays_2 = build_tag_arrays(&batch_no_null, &tag_cols_2).unwrap();
319        let tsid_no_null = compute_tsid_array(&batch_no_null, &tag_cols_2, &tag_arrays_2);
320
321        let schema3 = Arc::new(ArrowSchema::new(vec![
322            Field::new("a", DataType::Utf8, true),
323            Field::new("b", DataType::Utf8, true),
324            Field::new("c", DataType::Utf8, true),
325        ]));
326        let batch_with_null = RecordBatch::try_new(
327            schema3,
328            vec![
329                Arc::new(StringArray::from(vec!["A"])),
330                Arc::new(StringArray::from(vec!["B"])),
331                Arc::new(StringArray::from(vec![None as Option<&str>])),
332            ],
333        )
334        .unwrap();
335        let tag_cols_3: Vec<TagColumnInfo> = vec![
336            TagColumnInfo {
337                name: "a".to_string(),
338                index: 0,
339                column_id: 1,
340            },
341            TagColumnInfo {
342                name: "b".to_string(),
343                index: 1,
344                column_id: 2,
345            },
346            TagColumnInfo {
347                name: "c".to_string(),
348                index: 2,
349                column_id: 3,
350            },
351        ];
352        let tag_arrays_3 = build_tag_arrays(&batch_with_null, &tag_cols_3).unwrap();
353        let tsid_with_null = compute_tsid_array(&batch_with_null, &tag_cols_3, &tag_arrays_3);
354
355        assert_eq!(tsid_no_null.value(0), tsid_with_null.value(0));
356    }
357
358    #[test]
359    fn test_modify_batch_sparse() {
360        let batch = build_sparse_test_batch();
361        let tag_columns = sparse_tag_columns();
362        let non_tag_indices = vec![0, 1];
363        let table_id: u32 = 1025;
364
365        let modified =
366            modify_batch_sparse(batch, table_id, &tag_columns, &non_tag_indices).unwrap();
367
368        assert_eq!(modified.num_columns(), 3);
369        assert_eq!(modified.schema().field(0).name(), PRIMARY_KEY_COLUMN_NAME);
370        assert_eq!(modified.schema().field(1).name(), "greptime_timestamp");
371        assert_eq!(modified.schema().field(2).name(), "greptime_value");
372    }
373
374    #[test]
375    fn test_modify_batch_sparse_matches_row_modifier() {
376        let batch = build_sparse_test_batch();
377        let tag_columns = sparse_tag_columns();
378        let non_tag_indices = vec![0, 1];
379        let table_id: u32 = 1025;
380        let modified =
381            modify_batch_sparse(batch, table_id, &tag_columns, &non_tag_indices).unwrap();
382
383        let make_info = |name: &str, column_id: ColumnId| ColumnMetadata {
384            column_schema: datatypes::schema::ColumnSchema::new(
385                name.to_string(),
386                datatypes::prelude::ConcreteDataType::string_datatype(),
387                false,
388            ),
389            semantic_type: SemanticType::Tag,
390            column_id,
391        };
392        let name_to_column_id: HashMap<String, ColumnMetadata> = [
393            (
394                "greptime_timestamp".to_string(),
395                make_info("greptime_timestamp", 0),
396            ),
397            ("greptime_value".to_string(), make_info("greptime_value", 1)),
398            ("namespace".to_string(), make_info("namespace", 2)),
399            ("host".to_string(), make_info("host", 3)),
400        ]
401        .into_iter()
402        .collect();
403
404        let rows = Rows {
405            schema: vec![
406                ColumnSchema {
407                    column_name: "greptime_timestamp".to_string(),
408                    datatype: ColumnDataType::TimestampMillisecond as i32,
409                    semantic_type: SemanticType::Timestamp as i32,
410                    ..Default::default()
411                },
412                ColumnSchema {
413                    column_name: "greptime_value".to_string(),
414                    datatype: ColumnDataType::Float64 as i32,
415                    semantic_type: SemanticType::Field as i32,
416                    ..Default::default()
417                },
418                ColumnSchema {
419                    column_name: "namespace".to_string(),
420                    datatype: ColumnDataType::String as i32,
421                    semantic_type: SemanticType::Tag as i32,
422                    ..Default::default()
423                },
424                ColumnSchema {
425                    column_name: "host".to_string(),
426                    datatype: ColumnDataType::String as i32,
427                    semantic_type: SemanticType::Tag as i32,
428                    ..Default::default()
429                },
430            ],
431            rows: vec![Row {
432                values: vec![
433                    Value {
434                        value_data: Some(ValueData::TimestampMillisecondValue(1000)),
435                    },
436                    Value {
437                        value_data: Some(ValueData::F64Value(42.0)),
438                    },
439                    Value {
440                        value_data: Some(ValueData::StringValue("greptimedb".to_string())),
441                    },
442                    Value {
443                        value_data: Some(ValueData::StringValue("127.0.0.1".to_string())),
444                    },
445                ],
446            }],
447        };
448
449        let row_iter = RowsIter::new(rows, &name_to_column_id);
450        let rows = RowModifier::default()
451            .modify_rows(
452                row_iter,
453                TableIdInput::Single(table_id),
454                PrimaryKeyEncoding::Sparse,
455            )
456            .unwrap();
457        let ValueData::BinaryValue(expected_pk) =
458            rows.rows[0].values[0].value_data.clone().unwrap()
459        else {
460            panic!("expected binary primary key");
461        };
462
463        let actual_array = modified
464            .column(0)
465            .as_any()
466            .downcast_ref::<BinaryArray>()
467            .unwrap();
468        assert_eq!(actual_array.value(0), expected_pk.as_slice());
469    }
470
471    #[test]
472    fn label_replace_with_utf8view_labels_does_not_panic() {
473        // Tag (label) columns may arrive as any string representation (`Utf8`,
474        // `LargeUtf8`, `Utf8View`, or dictionary-encoded). `modify_batch_sparse`
475        // must not assume the tag columns are plain `StringArray`s.
476        let tag_arrays: Vec<ArrayRef> = vec![
477            Arc::new(StringArray::from(vec!["greptimedb"])),
478            Arc::new(datatypes::arrow::array::LargeStringArray::from(vec![
479                "greptimedb",
480            ])),
481            Arc::new(datatypes::arrow::array::StringViewArray::from(vec![
482                "greptimedb",
483            ])),
484            Arc::new(datatypes::arrow::array::DictionaryArray::<
485                datatypes::arrow::datatypes::UInt32Type,
486            >::new(
487                datatypes::arrow::array::UInt32Array::from(vec![0]),
488                Arc::new(StringArray::from(vec!["greptimedb"])),
489            )),
490        ];
491        let tag_columns = vec![TagColumnInfo {
492            name: "namespace".to_string(),
493            index: 2,
494            column_id: 2,
495        }];
496        let non_tag_indices = vec![0, 1];
497
498        let primary_keys = tag_arrays
499            .into_iter()
500            .map(|tag_array| {
501                let schema = Arc::new(ArrowSchema::new(vec![
502                    Field::new("greptime_timestamp", DataType::Int64, false),
503                    Field::new("greptime_value", DataType::Float64, true),
504                    Field::new("namespace", tag_array.data_type().clone(), true),
505                ]));
506                let batch = RecordBatch::try_new(
507                    schema,
508                    vec![
509                        Arc::new(Int64Array::from(vec![1000])),
510                        Arc::new(datatypes::arrow::array::Float64Array::from(vec![42.0])),
511                        tag_array,
512                    ],
513                )
514                .unwrap();
515                let modified =
516                    modify_batch_sparse(batch, 1025, &tag_columns, &non_tag_indices).unwrap();
517                modified
518                    .column(0)
519                    .as_any()
520                    .downcast_ref::<BinaryArray>()
521                    .unwrap()
522                    .value(0)
523                    .to_vec()
524            })
525            .collect::<Vec<_>>();
526
527        assert!(primary_keys.windows(2).all(|keys| keys[0] == keys[1]));
528    }
529
530    #[test]
531    fn test_compute_tsid_utf8view_with_nulls_matches_utf8() {
532        // Regression test for Utf8View tag/label columns: the whole
533        // `build_tag_arrays` -> `compute_tsid_array` -> sparse primary key path must handle
534        // `Utf8View` columns exactly like `Utf8`, including rows with null tags, instead of
535        // panicking or producing different TSIDs.
536        let namespace_values: Vec<Option<&str>> =
537            vec![Some("ns-a"), Some("ns-b"), None, Some("ns-c")];
538        let host_values: Vec<Option<&str>> = vec![Some("host-1"), None, Some("host-2"), None];
539
540        let build_batch = |data_type: DataType| {
541            let (namespace, host): (ArrayRef, ArrayRef) = match &data_type {
542                DataType::Utf8 => (
543                    Arc::new(StringArray::from(namespace_values.clone())),
544                    Arc::new(StringArray::from(host_values.clone())),
545                ),
546                DataType::Utf8View => (
547                    Arc::new(datatypes::arrow::array::StringViewArray::from(
548                        namespace_values.clone(),
549                    )),
550                    Arc::new(datatypes::arrow::array::StringViewArray::from(
551                        host_values.clone(),
552                    )),
553                ),
554                _ => unreachable!(),
555            };
556            let schema = Arc::new(ArrowSchema::new(vec![
557                Field::new("greptime_timestamp", DataType::Int64, false),
558                Field::new("greptime_value", DataType::Float64, true),
559                Field::new("namespace", data_type.clone(), true),
560                Field::new("host", data_type, true),
561            ]));
562            RecordBatch::try_new(
563                schema,
564                vec![
565                    Arc::new(Int64Array::from(vec![1000, 1001, 1002, 1003])),
566                    Arc::new(datatypes::arrow::array::Float64Array::from(vec![
567                        1.0, 2.0, 3.0, 4.0,
568                    ])),
569                    namespace,
570                    host,
571                ],
572            )
573            .unwrap()
574        };
575
576        let tag_columns = vec![
577            TagColumnInfo {
578                name: "host".to_string(),
579                index: 3,
580                column_id: 3,
581            },
582            TagColumnInfo {
583                name: "namespace".to_string(),
584                index: 2,
585                column_id: 2,
586            },
587        ];
588
589        let utf8_batch = build_batch(DataType::Utf8);
590        let utf8view_batch = build_batch(DataType::Utf8View);
591
592        // Both representations must be accepted by `build_tag_arrays`.
593        let utf8_tag_arrays = build_tag_arrays(&utf8_batch, &tag_columns).unwrap();
594        let utf8view_tag_arrays = build_tag_arrays(&utf8view_batch, &tag_columns).unwrap();
595
596        // TSIDs must be identical across representations, including rows with null tags.
597        let utf8_tsids = compute_tsid_array(&utf8_batch, &tag_columns, &utf8_tag_arrays);
598        let utf8view_tsids =
599            compute_tsid_array(&utf8view_batch, &tag_columns, &utf8view_tag_arrays);
600        assert_eq!(utf8_tsids, utf8view_tsids);
601
602        // The full sparse primary key (TSID + tag values) must also be identical.
603        let non_tag_indices = vec![0, 1];
604        let modified_utf8 =
605            modify_batch_sparse(utf8_batch, 1025, &tag_columns, &non_tag_indices).unwrap();
606        let modified_utf8view =
607            modify_batch_sparse(utf8view_batch, 1025, &tag_columns, &non_tag_indices).unwrap();
608        let utf8_pks = modified_utf8
609            .column(0)
610            .as_any()
611            .downcast_ref::<BinaryArray>()
612            .unwrap();
613        let utf8view_pks = modified_utf8view
614            .column(0)
615            .as_any()
616            .downcast_ref::<BinaryArray>()
617            .unwrap();
618        for row in 0..utf8_pks.len() {
619            assert_eq!(utf8_pks.value(row), utf8view_pks.value(row));
620        }
621    }
622}
metric_engine/batch_modifier.rs

metric_engine/
batch_modifier.rs
