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

servers/batcher/logical_table/
batch_convert.rs
