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

servers/batcher/logical_table/
batch_convert.rs
