servers/batcher/logical_table/
batch_convert.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, HashSet};
16use std::sync::Arc;
17use std::time::{Duration, Instant};
18
19use arrow::compute::concat_batches;
20use arrow::datatypes::{DataType as ArrowDataType, Schema as ArrowSchema};
21use arrow::record_batch::RecordBatch;
22use common_query::prelude::{greptime_timestamp, greptime_value};
23use metric_engine::batch_modifier::{TagColumnInfo, modify_batch_sparse};
24use smallvec::SmallVec;
25use snafu::{OptionExt, ResultExt, ensure};
26use table::metadata::TableId;
27
28use crate::error;
29use crate::error::Result;
30use crate::metrics::PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED;
31
32/// An aligned logical record batch and its timestamp column index.
33#[derive(Debug, Clone)]
34pub struct RecordBatchWithTsIdx {
35    /// The aligned logical record batch.
36    pub(in crate::batcher::logical_table) batch: RecordBatch,
37    /// The timestamp column index in `batch`.
38    pub(in crate::batcher::logical_table) timestamp_index: usize,
39}
40
41impl RecordBatchWithTsIdx {
42    /// Creates a record batch with a validated timestamp column index.
43    pub fn try_new(batch: RecordBatch, timestamp_index: usize) -> Result<Self> {
44        let schema = batch.schema();
45        let timestamp_field = schema.fields().get(timestamp_index).with_context(|| {
46            error::InvalidPromRemoteRequestSnafu {
47                msg: format!(
48                    "Timestamp column index {} is out of bounds for record batch with {} columns",
49                    timestamp_index,
50                    batch.num_columns()
51                ),
52            }
53        })?;
54        ensure!(
55            matches!(timestamp_field.data_type(), ArrowDataType::Timestamp(_, _)),
56            error::InvalidPromRemoteRequestSnafu {
57                msg: format!(
58                    "Column at index {} is not a timestamp column: {:?}",
59                    timestamp_index,
60                    timestamp_field.data_type()
61                ),
62            }
63        );
64
65        Ok(Self {
66            batch,
67            timestamp_index,
68        })
69    }
70
71    #[cfg(test)]
72    pub(crate) fn into_parts(self) -> (RecordBatch, usize) {
73        (self.batch, self.timestamp_index)
74    }
75}
76
77#[derive(Debug, Clone)]
78pub struct TableBatch {
79    pub table_name: String,
80    pub table_id: TableId,
81    pub batches: Vec<RecordBatchWithTsIdx>,
82    pub row_count: usize,
83}
84
85/// Classifies columns in a logical-table batch for sparse primary-key conversion.
86///
87/// Returns:
88/// - `Vec<TagColumnInfo>`: all Utf8 tag columns sorted by tag name, used for
89///   TSID and sparse primary-key encoding.
90/// - `SmallVec<[usize; 3]>`: indices of columns copied into the physical batch
91///   after `__primary_key`, ordered as `[greptime_timestamp, greptime_value,
92///   partition_tag_columns...]`.
93pub(in crate::batcher::logical_table) fn columns_taxonomy(
94    batch_schema: &Arc<ArrowSchema>,
95    table_name: &str,
96    name_to_ids: &HashMap<String, u32>,
97    partition_columns: &HashSet<&str>,
98) -> Result<(Vec<TagColumnInfo>, SmallVec<[usize; 3]>)> {
99    let mut tag_columns = Vec::new();
100    let mut essential_column_indices =
101        SmallVec::<[usize; 3]>::with_capacity(2 + partition_columns.len());
102    // Placeholder for greptime_timestamp and greptime_value
103    essential_column_indices.push(0);
104    essential_column_indices.push(0);
105
106    let mut timestamp_index = None;
107    let mut value_index = None;
108
109    for (index, field) in batch_schema.fields().iter().enumerate() {
110        match field.data_type() {
111            ArrowDataType::Utf8 => {
112                let column_id = name_to_ids.get(field.name()).copied().with_context(|| {
113                    error::InvalidPromRemoteRequestSnafu {
114                        msg: format!(
115                            "Column '{}' from logical table '{}' not found in physical table column IDs",
116                            field.name(),
117                            table_name
118                        ),
119                    }
120                })?;
121                tag_columns.push(TagColumnInfo {
122                    name: field.name().clone(),
123                    index,
124                    column_id,
125                });
126
127                if partition_columns.contains(field.name().as_str()) {
128                    essential_column_indices.push(index);
129                }
130            }
131            ArrowDataType::Timestamp(_, _) => {
132                ensure!(
133                    timestamp_index.replace(index).is_none(),
134                    error::InvalidPromRemoteRequestSnafu {
135                        msg: format!(
136                            "Duplicated timestamp column in logical table '{}' batch schema",
137                            table_name
138                        ),
139                    }
140                );
141            }
142            ArrowDataType::Float64 => {
143                ensure!(
144                    value_index.replace(index).is_none(),
145                    error::InvalidPromRemoteRequestSnafu {
146                        msg: format!(
147                            "Duplicated value column in logical table '{}' batch schema",
148                            table_name
149                        ),
150                    }
151                );
152            }
153            datatype => {
154                return error::InvalidPromRemoteRequestSnafu {
155                    msg: format!(
156                        "Unexpected data type '{datatype:?}' in logical table '{}' batch schema",
157                        table_name
158                    ),
159                }
160                .fail();
161            }
162        }
163    }
164
165    let timestamp_index =
166        timestamp_index.with_context(|| error::InvalidPromRemoteRequestSnafu {
167            msg: format!(
168                "Missing essential column '{}' in logical table '{}' batch schema",
169                greptime_timestamp(),
170                table_name
171            ),
172        })?;
173    let value_index = value_index.with_context(|| error::InvalidPromRemoteRequestSnafu {
174        msg: format!(
175            "Missing essential column '{}' in logical table '{}' batch schema",
176            greptime_value(),
177            table_name
178        ),
179    })?;
180
181    tag_columns.sort_by(|a, b| a.name.cmp(&b.name));
182
183    essential_column_indices[0] = timestamp_index;
184    essential_column_indices[1] = value_index;
185
186    Ok((tag_columns, essential_column_indices))
187}
188
189/// Transforms logical table batches into physical format (sparse primary key encoding).
190///
191/// It identifies tag columns and essential columns (timestamp, value) for each logical batch
192/// and applies sparse primary key modification.
193pub(in crate::batcher::logical_table) fn transform_logical_batches_to_physical(
194    table_batches: &[TableBatch],
195    name_to_ids: &HashMap<String, u32>,
196    partition_columns_set: &HashSet<&str>,
197) -> Result<Vec<RecordBatch>> {
198    let mut modified_batches: Vec<RecordBatch> =
199        Vec::with_capacity(table_batches.iter().map(|b| b.batches.len()).sum());
200
201    let mut modify_elapsed = Duration::ZERO;
202    let mut columns_taxonomy_elapsed = Duration::ZERO;
203
204    for table_batch in table_batches {
205        let table_id = table_batch.table_id;
206
207        for batch in &table_batch.batches {
208            let batch = &batch.batch;
209            let batch_schema = batch.schema();
210            let start = Instant::now();
211            let (tag_columns, essential_col_indices) = columns_taxonomy(
212                &batch_schema,
213                &table_batch.table_name,
214                name_to_ids,
215                partition_columns_set,
216            )?;
217
218            columns_taxonomy_elapsed += start.elapsed();
219            if tag_columns.is_empty() && essential_col_indices.is_empty() {
220                continue;
221            }
222
223            let modified = {
224                let start = Instant::now();
225                // The schema of modified batch is: __primary_key, timestamp, value, other partition columns...
226                let batch = modify_batch_sparse(
227                    batch.clone(),
228                    table_id,
229                    &tag_columns,
230                    &essential_col_indices,
231                )?;
232                modify_elapsed += start.elapsed();
233                batch
234            };
235
236            modified_batches.push(modified);
237        }
238    }
239
240    PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
241        .with_label_values(&["flush_physical_modify_batch"])
242        .observe(modify_elapsed.as_secs_f64());
243    PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
244        .with_label_values(&["flush_physical_columns_taxonomy"])
245        .observe(columns_taxonomy_elapsed.as_secs_f64());
246
247    ensure!(
248        !modified_batches.is_empty(),
249        error::InternalSnafu {
250            err_msg: "No batches can be transformed during pending flush",
251        }
252    );
253    Ok(modified_batches)
254}
255
256/// Concatenates all modified batches into a single large batch.
257///
258/// All modified batches share the same physical schema.
259pub(in crate::batcher::logical_table) fn concat_modified_batches(
260    modified_batches: &[RecordBatch],
261) -> Result<RecordBatch> {
262    let combined_schema = modified_batches[0].schema();
263    let _timer = PENDING_ROWS_BATCH_FLUSH_STAGE_ELAPSED
264        .with_label_values(&["flush_physical_concat_all"])
265        .start_timer();
266    concat_batches(&combined_schema, modified_batches).context(error::ArrowSnafu)
267}
268
269#[cfg(test)]
270mod tests {
271    use std::collections::{HashMap, HashSet};
272    use std::sync::Arc;
273
274    use arrow::array::{Array, StringArray, TimestampMillisecondArray};
275    use arrow::datatypes::{DataType as ArrowDataType, Field, Schema as ArrowSchema};
276    use arrow::record_batch::RecordBatch;
277    use metric_engine::batch_modifier::modify_batch_sparse;
278    use smallvec::SmallVec;
279
280    use crate::batcher::logical_table::batch_convert::{
281        TableBatch, columns_taxonomy, transform_logical_batches_to_physical,
282    };
283    use crate::batcher::logical_table::test_util::mock_aligned_tag_batch;
284    use crate::error::Error;
285
286    #[test]
287    fn test_collect_tag_columns_and_non_tag_indices_keeps_partition_tag_column() {
288        let schema = Arc::new(ArrowSchema::new(vec![
289            Field::new(
290                "greptime_timestamp",
291                ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
292                false,
293            ),
294            Field::new("greptime_value", ArrowDataType::Float64, true),
295            Field::new("host", ArrowDataType::Utf8, true),
296            Field::new("region", ArrowDataType::Utf8, true),
297        ]));
298        let name_to_ids =
299            HashMap::from([("host".to_string(), 1_u32), ("region".to_string(), 2_u32)]);
300        let partition_columns = HashSet::from(["host"]);
301
302        let (tag_columns, non_tag_indices) =
303            columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns).unwrap();
304
305        assert_eq!(2, tag_columns.len());
306        assert_eq!(&[0, 1, 2], non_tag_indices.as_slice());
307    }
308
309    #[test]
310    fn test_collect_tag_columns_and_non_tag_indices_prioritizes_essential_columns() {
311        let schema = Arc::new(ArrowSchema::new(vec![
312            Field::new("host", ArrowDataType::Utf8, true),
313            Field::new("greptime_value", ArrowDataType::Float64, true),
314            Field::new(
315                "greptime_timestamp",
316                ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
317                false,
318            ),
319            Field::new("region", ArrowDataType::Utf8, true),
320        ]));
321        let name_to_ids =
322            HashMap::from([("host".to_string(), 1_u32), ("region".to_string(), 2_u32)]);
323        let partition_columns = HashSet::from(["host", "region"]);
324
325        let (_tag_columns, non_tag_indices): (_, SmallVec<[usize; 3]>) =
326            columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns).unwrap();
327
328        assert_eq!(&[2, 1, 0, 3], non_tag_indices.as_slice());
329    }
330
331    #[test]
332    fn test_collect_tag_columns_and_non_tag_indices_rejects_unexpected_data_type() {
333        let schema = Arc::new(ArrowSchema::new(vec![
334            Field::new(
335                "greptime_timestamp",
336                ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
337                false,
338            ),
339            Field::new("greptime_value", ArrowDataType::Float64, true),
340            Field::new("host", ArrowDataType::Utf8, true),
341            Field::new("invalid", ArrowDataType::Boolean, true),
342        ]));
343        let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]);
344        let partition_columns = HashSet::from(["host"]);
345
346        let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns);
347
348        assert!(matches!(
349            result,
350            Err(Error::InvalidPromRemoteRequest { .. })
351        ));
352    }
353
354    #[test]
355    fn test_collect_tag_columns_and_non_tag_indices_rejects_int64_timestamp_column() {
356        let schema = Arc::new(ArrowSchema::new(vec![
357            Field::new("greptime_timestamp", ArrowDataType::Int64, false),
358            Field::new("greptime_value", ArrowDataType::Float64, true),
359            Field::new("host", ArrowDataType::Utf8, true),
360        ]));
361        let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]);
362        let partition_columns = HashSet::from(["host"]);
363
364        let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns);
365
366        assert!(matches!(
367            result,
368            Err(Error::InvalidPromRemoteRequest { .. })
369        ));
370    }
371
372    #[test]
373    fn test_collect_tag_columns_and_non_tag_indices_rejects_duplicated_timestamp_column() {
374        let schema = Arc::new(ArrowSchema::new(vec![
375            Field::new(
376                "ts1",
377                ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
378                false,
379            ),
380            Field::new(
381                "ts2",
382                ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
383                false,
384            ),
385            Field::new("greptime_value", ArrowDataType::Float64, true),
386            Field::new("host", ArrowDataType::Utf8, true),
387        ]));
388        let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]);
389        let partition_columns = HashSet::from(["host"]);
390
391        let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns);
392
393        assert!(matches!(
394            result,
395            Err(Error::InvalidPromRemoteRequest { .. })
396        ));
397    }
398
399    #[test]
400    fn test_collect_tag_columns_and_non_tag_indices_rejects_duplicated_value_column() {
401        let schema = Arc::new(ArrowSchema::new(vec![
402            Field::new(
403                "greptime_timestamp",
404                ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
405                false,
406            ),
407            Field::new("value1", ArrowDataType::Float64, true),
408            Field::new("value2", ArrowDataType::Float64, true),
409            Field::new("host", ArrowDataType::Utf8, true),
410        ]));
411        let name_to_ids = HashMap::from([("host".to_string(), 1_u32)]);
412        let partition_columns = HashSet::from(["host"]);
413
414        let result = columns_taxonomy(&schema, "cpu", &name_to_ids, &partition_columns);
415
416        assert!(matches!(
417            result,
418            Err(Error::InvalidPromRemoteRequest { .. })
419        ));
420    }
421
422    #[test]
423    fn test_modify_batch_sparse_with_taxonomy_per_batch() {
424        use arrow::array::BinaryArray;
425
426        let schema1 = Arc::new(ArrowSchema::new(vec![
427            Field::new(
428                "greptime_timestamp",
429                ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
430                false,
431            ),
432            Field::new("greptime_value", ArrowDataType::Float64, true),
433            Field::new("tag1", ArrowDataType::Utf8, true),
434        ]));
435
436        let schema2 = Arc::new(ArrowSchema::new(vec![
437            Field::new(
438                "greptime_timestamp",
439                ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None),
440                false,
441            ),
442            Field::new("greptime_value", ArrowDataType::Float64, true),
443            Field::new("tag1", ArrowDataType::Utf8, true),
444            Field::new("tag2", ArrowDataType::Utf8, true),
445        ]));
446        let batch2 = RecordBatch::try_new(
447            schema2.clone(),
448            vec![
449                Arc::new(TimestampMillisecondArray::from(vec![2000])),
450                Arc::new(arrow::array::Float64Array::from(vec![2.0])),
451                Arc::new(StringArray::from(vec!["v1"])),
452                Arc::new(StringArray::from(vec!["v2"])),
453            ],
454        )
455        .unwrap();
456
457        let name_to_ids = HashMap::from([("tag1".to_string(), 1), ("tag2".to_string(), 2)]);
458        let partition_columns = HashSet::new();
459
460        // A batch that only has tag1, same values as batch2 for ts and val.
461        let batch3 = RecordBatch::try_new(
462            schema1.clone(),
463            vec![
464                Arc::new(TimestampMillisecondArray::from(vec![2000])),
465                Arc::new(arrow::array::Float64Array::from(vec![2.0])),
466                Arc::new(StringArray::from(vec!["v1"])),
467            ],
468        )
469        .unwrap();
470
471        // Simulate the new loop logic in flush_batch_physical:
472        // Resolve taxonomy FOR EACH BATCH.
473        let (tag_columns2, indices2) =
474            columns_taxonomy(&batch2.schema(), "table", &name_to_ids, &partition_columns).unwrap();
475        let modified2 = modify_batch_sparse(batch2, 123, &tag_columns2, &indices2).unwrap();
476
477        let (tag_columns3, indices3) =
478            columns_taxonomy(&batch3.schema(), "table", &name_to_ids, &partition_columns).unwrap();
479        let modified3 = modify_batch_sparse(batch3, 123, &tag_columns3, &indices3).unwrap();
480
481        let pk2 = modified2
482            .column(0)
483            .as_any()
484            .downcast_ref::<BinaryArray>()
485            .unwrap();
486        let pk3 = modified3
487            .column(0)
488            .as_any()
489            .downcast_ref::<BinaryArray>()
490            .unwrap();
491
492        // Now they SHOULD be different because tag2 is included in pk2 but not in pk3.
493        assert_ne!(
494            pk2.value(0),
495            pk3.value(0),
496            "PK should be different because batch2 has tag2!"
497        );
498    }
499
500    #[test]
501    fn test_transform_logical_batches_to_physical_success() {
502        let batch = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0);
503
504        let table_batches = vec![TableBatch {
505            table_name: "t1".to_string(),
506            table_id: 1,
507            batches: vec![batch],
508            row_count: 1,
509        }];
510
511        let name_to_ids = HashMap::from([("tag1".to_string(), 1)]);
512        let partition_columns = HashSet::new();
513        let modified =
514            transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns)
515                .unwrap();
516
517        assert_eq!(1, modified.len());
518        assert_eq!(3, modified[0].num_columns());
519        assert_eq!("__primary_key", modified[0].schema().field(0).name());
520        assert_eq!("greptime_timestamp", modified[0].schema().field(1).name());
521        assert_eq!("greptime_value", modified[0].schema().field(2).name());
522    }
523
524    #[test]
525    fn test_transform_logical_batches_to_physical_taxonomy_failure() {
526        let batch = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0);
527
528        let table_batches = vec![TableBatch {
529            table_name: "t1".to_string(),
530            table_id: 1,
531            batches: vec![batch],
532            row_count: 1,
533        }];
534
535        // tag1 is missing from name_to_ids, causing columns_taxonomy to fail.
536        let name_to_ids = HashMap::new();
537        let partition_columns = HashSet::new();
538        let err =
539            transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns)
540                .unwrap_err();
541
542        assert!(
543            err.to_string()
544                .contains("not found in physical table column IDs")
545        );
546    }
547
548    #[test]
549    fn test_transform_logical_batches_to_physical_multiple_batches() {
550        let batch1 = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0);
551        let batch2 = mock_aligned_tag_batch("tag2", "v2", 2000, 2.0);
552
553        let table_batches = vec![
554            TableBatch {
555                table_name: "t1".to_string(),
556                table_id: 1,
557                batches: vec![batch1],
558                row_count: 1,
559            },
560            TableBatch {
561                table_name: "t2".to_string(),
562                table_id: 2,
563                batches: vec![batch2],
564                row_count: 1,
565            },
566        ];
567
568        let name_to_ids = HashMap::from([("tag1".to_string(), 1), ("tag2".to_string(), 2)]);
569        let partition_columns = HashSet::new();
570        let modified =
571            transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns)
572                .unwrap();
573
574        assert_eq!(2, modified.len());
575    }
576
577    #[test]
578    fn test_transform_logical_batches_to_physical_mixed_success_failure() {
579        let batch1 = mock_aligned_tag_batch("tag1", "v1", 1000, 1.0);
580        let batch2 = mock_aligned_tag_batch("tag2", "v2", 2000, 2.0);
581
582        let table_batches = vec![
583            TableBatch {
584                table_name: "t1".to_string(),
585                table_id: 1,
586                batches: vec![batch1],
587                row_count: 1,
588            },
589            TableBatch {
590                table_name: "t2".to_string(),
591                table_id: 2,
592                batches: vec![batch2],
593                row_count: 1,
594            },
595        ];
596
597        // tag1 is missing from name_to_ids, causing batch1 to fail.
598        let name_to_ids = HashMap::from([("tag2".to_string(), 2)]);
599        let partition_columns = HashSet::new();
600        let err =
601            transform_logical_batches_to_physical(&table_batches, &name_to_ids, &partition_columns)
602                .unwrap_err();
603
604        assert!(err.to_string().contains("tag1"));
605    }
606}
servers/batcher/logical_table/batch_convert.rs

servers/batcher/logical_table/
batch_convert.rs
