mito2/read/
compat.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
15//! Utilities to adapt readers with different schema.
16
17use std::collections::HashMap;
18use std::sync::Arc;
19
20use api::v1::SemanticType;
21use datatypes::arrow::array::{
22    Array, ArrayRef, BinaryArray, BinaryBuilder, DictionaryArray, UInt32Array, new_null_array,
23};
24use datatypes::arrow::compute::{TakeOptions, take};
25use datatypes::arrow::datatypes::{FieldRef, Schema, SchemaRef};
26use datatypes::arrow::record_batch::RecordBatch;
27use datatypes::data_type::ConcreteDataType;
28use datatypes::prelude::DataType;
29use datatypes::value::Value;
30use datatypes::vectors::{Helper, VectorRef};
31use mito_codec::row_converter::{
32    CompositeValues, PrimaryKeyCodec, SortField, build_primary_key_codec,
33    build_primary_key_codec_with_fields,
34};
35use snafu::{OptionExt, ResultExt, ensure};
36use store_api::codec::PrimaryKeyEncoding;
37use store_api::metadata::{RegionMetadata, RegionMetadataRef};
38use store_api::storage::ColumnId;
39
40use crate::error::{
41    CompatReaderSnafu, ComputeArrowSnafu, CreateDefaultSnafu, DecodeSnafu, EncodeSnafu,
42    NewRecordBatchSnafu, Result, UnexpectedSnafu, UnsupportedOperationSnafu,
43};
44use crate::read::flat_projection::{FlatProjectionMapper, flat_projected_columns};
45use crate::sst::parquet::flat_format::{FlatReadFormat, primary_key_column_index};
46use crate::sst::parquet::format::{INTERNAL_COLUMN_NUM, PrimaryKeyArray};
47use crate::sst::{internal_fields, tag_maybe_to_dictionary_field, with_field_id};
48
49/// Returns true if the columns in the `projection_mapper` and `read_format` have same data types
50/// and primary key encodings.
51pub(crate) fn has_same_columns_and_pk_encoding(
52    projection_mapper: &FlatProjectionMapper,
53    read_format: &FlatReadFormat,
54    compaction: bool,
55) -> bool {
56    let left = projection_mapper.metadata();
57    let right = read_format.metadata();
58    if left.primary_key_encoding != right.primary_key_encoding {
59        return false;
60    }
61
62    if left.column_metadatas.len() != right.column_metadatas.len() {
63        return false;
64    }
65
66    for (left_col, right_col) in left.column_metadatas.iter().zip(&right.column_metadatas) {
67        if left_col.column_id != right_col.column_id {
68            return false;
69        }
70        debug_assert_eq!(left_col.semantic_type, right_col.semantic_type);
71    }
72
73    &projection_mapper.input_arrow_schema(compaction) == read_format.arrow_schema()
74}
75
76/// A helper struct to adapt schema of the batch to an expected schema.
77pub(crate) struct FlatCompatBatch {
78    /// Indices to convert actual fields to expect fields.
79    index_or_defaults: Vec<IndexOrDefault>,
80    /// Expected arrow schema.
81    arrow_schema: SchemaRef,
82    /// Primary key adapter.
83    compat_pk: FlatCompatPrimaryKey,
84}
85
86impl FlatCompatBatch {
87    /// Creates a [FlatCompatBatch].
88    ///
89    /// - `mapper` is built from the metadata users expect to see.
90    /// - `read_format` is the [FlatReadFormat] of the input parquet.
91    /// - `compaction` indicates whether the reader is for compaction.
92    pub(crate) fn try_new(
93        mapper: &FlatProjectionMapper,
94        read_format: &FlatReadFormat,
95        compaction: bool,
96    ) -> Result<Option<Self>> {
97        let actual = read_format.metadata();
98        let actual_schema = flat_projected_columns(
99            actual,
100            read_format.format_projection(),
101            read_format.json_target_types(),
102        );
103
104        let expect_schema = mapper.batch_schema();
105        if expect_schema == actual_schema
106            && actual.primary_key == mapper.metadata().primary_key
107            && actual.primary_key_encoding == mapper.metadata().primary_key_encoding
108        {
109            // Although the SST has a different schema, but the schema after projection is the same
110            // as expected schema.
111            return Ok(None);
112        }
113
114        if actual.primary_key_encoding == PrimaryKeyEncoding::Sparse && compaction {
115            // Special handling for sparse encoding in compaction.
116            return FlatCompatBatch::try_new_compact_sparse(mapper, actual);
117        }
118
119        let (index_or_defaults, fields) =
120            Self::compute_index_and_fields(&actual_schema, expect_schema, mapper.metadata())?;
121
122        let compat_pk = FlatCompatPrimaryKey::new(mapper.metadata(), actual)?;
123
124        Ok(Some(Self {
125            index_or_defaults,
126            arrow_schema: Arc::new(Schema::new(fields)),
127            compat_pk,
128        }))
129    }
130
131    fn compute_index_and_fields(
132        actual_schema: &[(ColumnId, ConcreteDataType)],
133        expect_schema: &[(ColumnId, ConcreteDataType)],
134        expect_metadata: &RegionMetadata,
135    ) -> Result<(Vec<IndexOrDefault>, Vec<FieldRef>)> {
136        // Maps column id to the index and data type in the actual schema.
137        let actual_schema_index: HashMap<_, _> = actual_schema
138            .iter()
139            .enumerate()
140            .map(|(idx, (column_id, data_type))| (*column_id, (idx, data_type)))
141            .collect();
142
143        let mut index_or_defaults = Vec::with_capacity(expect_schema.len());
144        let mut fields = Vec::with_capacity(expect_schema.len());
145        for (column_id, expect_data_type) in expect_schema {
146            // Safety: expect_schema comes from the same mapper.
147            let column_index = expect_metadata.column_index_by_id(*column_id).unwrap();
148            let expect_column = &expect_metadata.column_metadatas[column_index];
149            let column_field = &expect_metadata.schema.arrow_schema().fields()[column_index];
150            // For tag columns, we need to create a dictionary field.
151            if expect_column.semantic_type == SemanticType::Tag {
152                let field = tag_maybe_to_dictionary_field(
153                    &expect_column.column_schema.data_type,
154                    column_field,
155                );
156                fields.push(Arc::new(with_field_id(
157                    (*field).clone(),
158                    expect_column.column_id,
159                )));
160            } else {
161                let field = with_field_id(
162                    Arc::unwrap_or_clone(column_field.clone()),
163                    expect_column.column_id,
164                )
165                .with_data_type(expect_data_type.as_arrow_type());
166                fields.push(Arc::new(field))
167            };
168
169            if let Some((index, actual_data_type)) = actual_schema_index.get(column_id) {
170                let mut cast_type = None;
171
172                // Same column different type.
173                if expect_data_type != *actual_data_type {
174                    ensure!(
175                        !expect_data_type.is_json2() && !actual_data_type.is_json2(),
176                        CompatReaderSnafu {
177                            region_id: expect_metadata.region_id,
178                            reason: format!(
179                                "JSON2 column '{}' must be aligned before FlatCompatBatch, actual: {}, expected: {}",
180                                expect_column.column_schema.name,
181                                actual_data_type,
182                                expect_data_type,
183                            ),
184                        }
185                    );
186                    cast_type = Some(expect_data_type.clone());
187                }
188                // Source has this column.
189                index_or_defaults.push(IndexOrDefault::Index {
190                    pos: *index,
191                    cast_type,
192                });
193            } else {
194                // Create a default vector with 1 element for that column.
195                let default_vector = if expect_data_type.is_json2() {
196                    // JSON2 columns do not support defaults, so missing columns use NULL.
197                    // `expect_data_type` may be a query-concretized JSON2 type,
198                    // while region metadata still carries its unconcretized type.
199                    // The default vector must match the target Arrow field above.
200                    Helper::try_into_vector(new_null_array(&expect_data_type.as_arrow_type(), 1))
201                        .context(CreateDefaultSnafu {
202                            region_id: expect_metadata.region_id,
203                            column: &expect_column.column_schema.name,
204                        })?
205                } else {
206                    expect_column
207                        .column_schema
208                        .create_default_vector(1)
209                        .context(CreateDefaultSnafu {
210                            region_id: expect_metadata.region_id,
211                            column: &expect_column.column_schema.name,
212                        })?
213                        .with_context(|| CompatReaderSnafu {
214                            region_id: expect_metadata.region_id,
215                            reason: format!(
216                                "column {} does not have a default value to read",
217                                expect_column.column_schema.name
218                            ),
219                        })?
220                };
221                index_or_defaults.push(IndexOrDefault::DefaultValue {
222                    default_vector,
223                    semantic_type: expect_column.semantic_type,
224                });
225            };
226        }
227        fields.extend_from_slice(&internal_fields());
228
229        Ok((index_or_defaults, fields))
230    }
231
232    fn try_new_compact_sparse(
233        mapper: &FlatProjectionMapper,
234        actual: &RegionMetadataRef,
235    ) -> Result<Option<Self>> {
236        // Currently, we don't support converting sparse encoding back to dense encoding in
237        // flat format.
238        ensure!(
239            mapper.metadata().primary_key_encoding == PrimaryKeyEncoding::Sparse,
240            UnsupportedOperationSnafu {
241                err_msg: "Flat format doesn't support converting sparse encoding back to dense encoding"
242            }
243        );
244
245        // For sparse encoding, we don't need to check the primary keys.
246        // Since this is for compaction, we always read all columns.
247        let actual_schema: Vec<_> = actual
248            .field_columns()
249            .chain([actual.time_index_column()])
250            .map(|col| (col.column_id, col.column_schema.data_type.clone()))
251            .collect();
252        let expect_schema: Vec<_> = mapper
253            .metadata()
254            .field_columns()
255            .chain([mapper.metadata().time_index_column()])
256            .map(|col| (col.column_id, col.column_schema.data_type.clone()))
257            .collect();
258
259        let (index_or_defaults, fields) =
260            Self::compute_index_and_fields(&actual_schema, &expect_schema, mapper.metadata())?;
261
262        let compat_pk = FlatCompatPrimaryKey::default();
263
264        Ok(Some(Self {
265            index_or_defaults,
266            arrow_schema: Arc::new(Schema::new(fields)),
267            compat_pk,
268        }))
269    }
270
271    /// Make columns of the `batch` compatible.
272    pub(crate) fn compat(&self, batch: RecordBatch) -> Result<RecordBatch> {
273        let len = batch.num_rows();
274        let columns = self
275            .index_or_defaults
276            .iter()
277            .map(|index_or_default| match index_or_default {
278                IndexOrDefault::Index { pos, cast_type } => {
279                    let old_column = batch.column(*pos);
280
281                    if let Some(ty) = cast_type {
282                        let casted =
283                            datatypes::arrow::compute::cast(old_column, &ty.as_arrow_type())
284                                .context(ComputeArrowSnafu)?;
285                        Ok(casted)
286                    } else {
287                        Ok(old_column.clone())
288                    }
289                }
290                IndexOrDefault::DefaultValue {
291                    default_vector,
292                    semantic_type,
293                } => repeat_vector(default_vector, len, *semantic_type == SemanticType::Tag),
294            })
295            .chain(
296                // Adds internal columns.
297                batch.columns()[batch.num_columns() - INTERNAL_COLUMN_NUM..]
298                    .iter()
299                    .map(|col| Ok(col.clone())),
300            )
301            .collect::<Result<Vec<_>>>()?;
302
303        let mut columns = columns;
304        let primary_key_index = primary_key_column_index(columns.len());
305        columns[primary_key_index] = self.compat_primary_key(&columns[primary_key_index])?;
306
307        RecordBatch::try_new(self.arrow_schema.clone(), columns).context(NewRecordBatchSnafu)
308    }
309
310    /// Makes an encoded primary-key array compatible with the expected metadata.
311    pub(crate) fn compat_primary_key(&self, primary_key: &ArrayRef) -> Result<ArrayRef> {
312        self.compat_pk.compat(primary_key)
313    }
314}
315
316/// Repeats the vector value `to_len` times.
317fn repeat_vector(vector: &VectorRef, to_len: usize, is_tag: bool) -> Result<ArrayRef> {
318    assert_eq!(1, vector.len());
319    let data_type = vector.data_type();
320    if is_tag && data_type.is_string() {
321        let values = vector.to_arrow_array();
322        if values.is_null(0) {
323            // Creates a dictionary array with `to_len` null keys.
324            let keys = UInt32Array::new_null(to_len);
325            Ok(Arc::new(DictionaryArray::new(keys, values.slice(0, 0))))
326        } else {
327            let keys = UInt32Array::from_value(0, to_len);
328            Ok(Arc::new(DictionaryArray::new(keys, values)))
329        }
330    } else {
331        let keys = UInt32Array::from_value(0, to_len);
332        take(
333            &vector.to_arrow_array(),
334            &keys,
335            Some(TakeOptions {
336                check_bounds: false,
337            }),
338        )
339        .context(ComputeArrowSnafu)
340    }
341}
342
343/// Returns true if the actual primary keys is the same as expected.
344fn is_primary_key_same(expect: &RegionMetadata, actual: &RegionMetadata) -> Result<bool> {
345    ensure!(
346        actual.primary_key.len() <= expect.primary_key.len(),
347        CompatReaderSnafu {
348            region_id: expect.region_id,
349            reason: format!(
350                "primary key has more columns {} than expect {}",
351                actual.primary_key.len(),
352                expect.primary_key.len()
353            ),
354        }
355    );
356    ensure!(
357        actual.primary_key == expect.primary_key[..actual.primary_key.len()],
358        CompatReaderSnafu {
359            region_id: expect.region_id,
360            reason: format!(
361                "primary key has different prefix, expect: {:?}, actual: {:?}",
362                expect.primary_key, actual.primary_key
363            ),
364        }
365    );
366
367    Ok(actual.primary_key.len() == expect.primary_key.len())
368}
369
370/// Index in source batch or a default value to fill a column.
371#[derive(Debug)]
372enum IndexOrDefault {
373    /// Index of the column in source batch.
374    Index {
375        pos: usize,
376        cast_type: Option<ConcreteDataType>,
377    },
378    /// Default value for the column.
379    DefaultValue {
380        /// Default value. The vector has only 1 element.
381        default_vector: VectorRef,
382        /// Semantic type of the column.
383        semantic_type: SemanticType,
384    },
385}
386
387/// Helper to rewrite primary key to another encoding for flat format.
388struct FlatRewritePrimaryKey {
389    /// New primary key encoder.
390    codec: Arc<dyn PrimaryKeyCodec>,
391    /// Metadata of the expected region.
392    metadata: RegionMetadataRef,
393    /// Original primary key codec.
394    /// If we need to rewrite the primary key.
395    old_codec: Arc<dyn PrimaryKeyCodec>,
396}
397
398impl FlatRewritePrimaryKey {
399    fn new(
400        expect: &RegionMetadataRef,
401        actual: &RegionMetadataRef,
402    ) -> Option<FlatRewritePrimaryKey> {
403        if expect.primary_key_encoding == actual.primary_key_encoding {
404            return None;
405        }
406        let codec = build_primary_key_codec(expect);
407        let old_codec = build_primary_key_codec(actual);
408
409        Some(FlatRewritePrimaryKey {
410            codec,
411            metadata: expect.clone(),
412            old_codec,
413        })
414    }
415
416    /// Rewrites the primary key of the `batch`.
417    /// It also appends the values to the primary key.
418    fn rewrite_key(
419        &self,
420        append_values: &[(ColumnId, Value)],
421        primary_key: &ArrayRef,
422    ) -> Result<ArrayRef> {
423        if let Some(old_pk_dict_array) = primary_key.as_any().downcast_ref::<PrimaryKeyArray>() {
424            let old_pk_values_array = old_pk_dict_array
425                .values()
426                .as_any()
427                .downcast_ref::<BinaryArray>()
428                .context(UnexpectedSnafu {
429                    reason: "Primary-key dictionary values are not binary",
430                })?;
431            let new_pk_values_array =
432                Arc::new(self.rewrite_values(append_values, old_pk_values_array)?);
433            return Ok(Arc::new(PrimaryKeyArray::new(
434                old_pk_dict_array.keys().clone(),
435                new_pk_values_array,
436            )));
437        }
438
439        let old_pk_values_array =
440            primary_key
441                .as_any()
442                .downcast_ref::<BinaryArray>()
443                .context(UnexpectedSnafu {
444                    reason: format!(
445                        "Primary-key column is neither binary nor dictionary, got {:?}",
446                        primary_key.data_type()
447                    ),
448                })?;
449        Ok(Arc::new(
450            self.rewrite_values(append_values, old_pk_values_array)?,
451        ))
452    }
453
454    fn rewrite_values(
455        &self,
456        append_values: &[(ColumnId, Value)],
457        old_pk_values_array: &BinaryArray,
458    ) -> Result<BinaryArray> {
459        let mut builder = BinaryBuilder::with_capacity(
460            old_pk_values_array.len(),
461            old_pk_values_array.value_data().len(),
462        );
463
464        // Binary buffer for the primary key.
465        let mut buffer = Vec::with_capacity(
466            old_pk_values_array.value_data().len() / old_pk_values_array.len().max(1),
467        );
468        let mut column_id_values = Vec::new();
469        // Iterates the binary array and rewrites the primary key.
470        for value in old_pk_values_array.iter() {
471            let Some(old_pk) = value else {
472                builder.append_null();
473                continue;
474            };
475            // Decodes the old primary key.
476            let mut pk_values = self.old_codec.decode(old_pk).context(DecodeSnafu)?;
477            pk_values.extend(append_values);
478
479            buffer.clear();
480            column_id_values.clear();
481            // Encodes the new primary key.
482            match pk_values {
483                CompositeValues::Dense(dense_values) => {
484                    self.codec
485                        .encode_values(dense_values.as_slice(), &mut buffer)
486                        .context(EncodeSnafu)?;
487                }
488                CompositeValues::Sparse(sparse_values) => {
489                    for id in &self.metadata.primary_key {
490                        let value = sparse_values.get_or_null(*id);
491                        column_id_values.push((*id, value.clone()));
492                    }
493                    self.codec
494                        .encode_values(&column_id_values, &mut buffer)
495                        .context(EncodeSnafu)?;
496                }
497            }
498            builder.append_value(&buffer);
499        }
500        Ok(builder.finish())
501    }
502}
503
504/// Helper to make primary key compatible for flat format.
505#[derive(Default)]
506struct FlatCompatPrimaryKey {
507    /// Primary key rewriter.
508    rewriter: Option<FlatRewritePrimaryKey>,
509    /// Converter to append values to primary keys.
510    converter: Option<Arc<dyn PrimaryKeyCodec>>,
511    /// Default values to append.
512    values: Vec<(ColumnId, Value)>,
513}
514
515impl FlatCompatPrimaryKey {
516    fn new(expect: &RegionMetadataRef, actual: &RegionMetadataRef) -> Result<Self> {
517        let rewriter = FlatRewritePrimaryKey::new(expect, actual);
518
519        if is_primary_key_same(expect, actual)? {
520            return Ok(Self {
521                rewriter,
522                converter: None,
523                values: Vec::new(),
524            });
525        }
526
527        // We need to append default values to the primary key.
528        let to_add = &expect.primary_key[actual.primary_key.len()..];
529        let mut values = Vec::with_capacity(to_add.len());
530        let mut fields = Vec::with_capacity(to_add.len());
531        for column_id in to_add {
532            // Safety: The id comes from expect region metadata.
533            let column = expect.column_by_id(*column_id).unwrap();
534            fields.push((
535                *column_id,
536                SortField::new(column.column_schema.data_type.clone()),
537            ));
538            let default_value = column
539                .column_schema
540                .create_default()
541                .context(CreateDefaultSnafu {
542                    region_id: expect.region_id,
543                    column: &column.column_schema.name,
544                })?
545                .with_context(|| CompatReaderSnafu {
546                    region_id: expect.region_id,
547                    reason: format!(
548                        "key column {} does not have a default value to read",
549                        column.column_schema.name
550                    ),
551                })?;
552            values.push((*column_id, default_value));
553        }
554        // is_primary_key_same() is false so we have different number of primary key columns.
555        debug_assert!(!fields.is_empty());
556
557        // Create converter to append values.
558        let converter = Some(build_primary_key_codec_with_fields(
559            expect.primary_key_encoding,
560            fields.into_iter(),
561        ));
562
563        Ok(Self {
564            rewriter,
565            converter,
566            values,
567        })
568    }
569
570    /// Makes an encoded primary-key array compatible.
571    fn compat(&self, primary_key: &ArrayRef) -> Result<ArrayRef> {
572        if let Some(rewriter) = &self.rewriter {
573            // If we have different encoding, rewrite the whole primary key.
574            return rewriter.rewrite_key(&self.values, primary_key);
575        }
576
577        self.append_key(primary_key)
578    }
579
580    /// Appends values to the primary key array.
581    fn append_key(&self, primary_key: &ArrayRef) -> Result<ArrayRef> {
582        let Some(converter) = &self.converter else {
583            return Ok(primary_key.clone());
584        };
585
586        if let Some(old_pk_dict_array) = primary_key.as_any().downcast_ref::<PrimaryKeyArray>() {
587            let old_pk_values_array = old_pk_dict_array
588                .values()
589                .as_any()
590                .downcast_ref::<BinaryArray>()
591                .context(UnexpectedSnafu {
592                    reason: "Primary-key dictionary values are not binary",
593                })?;
594            let new_pk_values_array =
595                Arc::new(self.append_values(old_pk_values_array, converter.as_ref())?);
596            return Ok(Arc::new(PrimaryKeyArray::new(
597                old_pk_dict_array.keys().clone(),
598                new_pk_values_array,
599            )));
600        }
601
602        let old_pk_values_array =
603            primary_key
604                .as_any()
605                .downcast_ref::<BinaryArray>()
606                .context(UnexpectedSnafu {
607                    reason: format!(
608                        "Primary-key column is neither binary nor dictionary, got {:?}",
609                        primary_key.data_type()
610                    ),
611                })?;
612        Ok(Arc::new(
613            self.append_values(old_pk_values_array, converter.as_ref())?,
614        ))
615    }
616
617    fn append_values(
618        &self,
619        old_pk_values_array: &BinaryArray,
620        converter: &dyn PrimaryKeyCodec,
621    ) -> Result<BinaryArray> {
622        let mut builder = BinaryBuilder::with_capacity(
623            old_pk_values_array.len(),
624            old_pk_values_array.value_data().len()
625                + converter.estimated_size().unwrap_or_default() * old_pk_values_array.len(),
626        );
627
628        // Binary buffer for the primary key.
629        let mut buffer = Vec::with_capacity(
630            old_pk_values_array.value_data().len() / old_pk_values_array.len().max(1)
631                + converter.estimated_size().unwrap_or_default(),
632        );
633
634        // Iterates the binary array and appends values to the primary key.
635        for value in old_pk_values_array.iter() {
636            let Some(old_pk) = value else {
637                builder.append_null();
638                continue;
639            };
640
641            buffer.clear();
642            buffer.extend_from_slice(old_pk);
643            converter
644                .encode_values(&self.values, &mut buffer)
645                .context(EncodeSnafu)?;
646
647            builder.append_value(&buffer);
648        }
649
650        Ok(builder.finish())
651    }
652}
653
654#[cfg(test)]
655mod tests {
656    use std::collections::BTreeMap;
657    use std::sync::Arc;
658
659    use api::v1::{OpType, SemanticType};
660    use datatypes::arrow::array::{
661        ArrayRef, BinaryArray, BinaryDictionaryBuilder, Int64Array, StringDictionaryBuilder,
662        TimestampMicrosecondArray, TimestampMillisecondArray, UInt8Array, UInt64Array,
663    };
664    use datatypes::arrow::datatypes::UInt32Type;
665    use datatypes::arrow::record_batch::RecordBatch;
666    use datatypes::prelude::ConcreteDataType;
667    use datatypes::schema::ColumnSchema;
668    use datatypes::types::json_type::JsonNativeType;
669    use datatypes::value::ValueRef;
670    use mito_codec::row_converter::{
671        DensePrimaryKeyCodec, PrimaryKeyCodecExt, SparsePrimaryKeyCodec,
672    };
673    use store_api::codec::PrimaryKeyEncoding;
674    use store_api::metadata::{ColumnMetadata, RegionMetadataBuilder};
675    use store_api::storage::RegionId;
676
677    use super::*;
678    use crate::read::flat_projection::FlatProjectionMapper;
679    use crate::read::read_columns::ReadColumns;
680    use crate::sst::parquet::flat_format::FlatReadFormat;
681    use crate::sst::{FlatSchemaOptions, to_flat_sst_arrow_schema};
682
683    /// Creates a new [RegionMetadata].
684    fn new_metadata(
685        semantic_types: &[(ColumnId, SemanticType, ConcreteDataType)],
686        primary_key: &[ColumnId],
687    ) -> RegionMetadata {
688        let mut builder = RegionMetadataBuilder::new(RegionId::new(1, 1));
689        for (id, semantic_type, data_type) in semantic_types {
690            let column_schema = match semantic_type {
691                SemanticType::Tag => {
692                    ColumnSchema::new(format!("tag_{id}"), data_type.clone(), true)
693                }
694                SemanticType::Field => {
695                    ColumnSchema::new(format!("field_{id}"), data_type.clone(), true)
696                }
697                SemanticType::Timestamp => ColumnSchema::new("ts", data_type.clone(), false),
698            };
699
700            builder.push_column_metadata(ColumnMetadata {
701                column_schema,
702                semantic_type: *semantic_type,
703                column_id: *id,
704            });
705        }
706        builder.primary_key(primary_key.to_vec());
707        builder.build().unwrap()
708    }
709
710    /// Encode primary key.
711    fn encode_key(keys: &[Option<&str>]) -> Vec<u8> {
712        let fields = (0..keys.len())
713            .map(|_| (0, SortField::new(ConcreteDataType::string_datatype())))
714            .collect();
715        let converter = DensePrimaryKeyCodec::with_fields(fields);
716        let row = keys.iter().map(|str_opt| match str_opt {
717            Some(v) => ValueRef::String(v),
718            None => ValueRef::Null,
719        });
720
721        converter.encode(row).unwrap()
722    }
723
724    /// Encode sparse primary key.
725    fn encode_sparse_key(keys: &[(ColumnId, Option<&str>)]) -> Vec<u8> {
726        let fields = (0..keys.len())
727            .map(|_| (1, SortField::new(ConcreteDataType::string_datatype())))
728            .collect();
729        let converter = SparsePrimaryKeyCodec::with_fields(fields);
730        let row = keys
731            .iter()
732            .map(|(id, str_opt)| match str_opt {
733                Some(v) => (*id, ValueRef::String(v)),
734                None => (*id, ValueRef::Null),
735            })
736            .collect::<Vec<_>>();
737        let mut buffer = vec![];
738        converter.encode_value_refs(&row, &mut buffer).unwrap();
739        buffer
740    }
741
742    /// Creates a primary key array for flat format testing.
743    fn build_flat_test_pk_array(primary_keys: &[&[u8]]) -> ArrayRef {
744        let mut builder = BinaryDictionaryBuilder::<UInt32Type>::new();
745        for &pk in primary_keys {
746            builder.append(pk).unwrap();
747        }
748        Arc::new(builder.finish())
749    }
750
751    #[test]
752    fn test_flat_compat_batch_with_missing_columns() {
753        let actual_metadata = Arc::new(new_metadata(
754            &[
755                (
756                    0,
757                    SemanticType::Timestamp,
758                    ConcreteDataType::timestamp_millisecond_datatype(),
759                ),
760                (1, SemanticType::Tag, ConcreteDataType::string_datatype()),
761                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
762            ],
763            &[1],
764        ));
765
766        let expected_metadata = Arc::new(new_metadata(
767            &[
768                (
769                    0,
770                    SemanticType::Timestamp,
771                    ConcreteDataType::timestamp_millisecond_datatype(),
772                ),
773                (1, SemanticType::Tag, ConcreteDataType::string_datatype()),
774                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
775                // Adds a new field.
776                (3, SemanticType::Field, ConcreteDataType::int64_datatype()),
777            ],
778            &[1],
779        ));
780
781        let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
782        let read_format = FlatReadFormat::new(
783            actual_metadata.clone(),
784            ReadColumns::new([0, 1, 2, 3]),
785            None,
786            "test",
787            false,
788        )
789        .unwrap();
790
791        let compat_batch = FlatCompatBatch::try_new(&mapper, &read_format, false)
792            .unwrap()
793            .unwrap();
794
795        let mut tag_builder = StringDictionaryBuilder::<UInt32Type>::new();
796        tag_builder.append_value("tag1");
797        tag_builder.append_value("tag1");
798        let tag_dict_array = Arc::new(tag_builder.finish());
799
800        let k1 = encode_key(&[Some("tag1")]);
801        let input_columns: Vec<ArrayRef> = vec![
802            tag_dict_array.clone(),
803            Arc::new(Int64Array::from(vec![100, 200])),
804            Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])),
805            build_flat_test_pk_array(&[&k1, &k1]),
806            Arc::new(UInt64Array::from_iter_values([1, 2])),
807            Arc::new(UInt8Array::from_iter_values([
808                OpType::Put as u8,
809                OpType::Put as u8,
810            ])),
811        ];
812        let input_schema =
813            to_flat_sst_arrow_schema(&actual_metadata, &FlatSchemaOptions::default());
814        let input_batch = RecordBatch::try_new(input_schema, input_columns).unwrap();
815
816        let result = compat_batch.compat(input_batch).unwrap();
817
818        let expected_schema =
819            to_flat_sst_arrow_schema(&expected_metadata, &FlatSchemaOptions::default());
820
821        let expected_columns: Vec<ArrayRef> = vec![
822            tag_dict_array.clone(),
823            Arc::new(Int64Array::from(vec![100, 200])),
824            Arc::new(Int64Array::from(vec![None::<i64>, None::<i64>])),
825            Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])),
826            build_flat_test_pk_array(&[&k1, &k1]),
827            Arc::new(UInt64Array::from_iter_values([1, 2])),
828            Arc::new(UInt8Array::from_iter_values([
829                OpType::Put as u8,
830                OpType::Put as u8,
831            ])),
832        ];
833        let expected_batch = RecordBatch::try_new(expected_schema, expected_columns).unwrap();
834
835        assert_eq!(expected_batch, result);
836    }
837
838    #[test]
839    fn test_flat_compat_batch_uses_projected_json2_type() -> Result<()> {
840        let json2 = ConcreteDataType::json2(JsonNativeType::object());
841        let actual_metadata = Arc::new(new_metadata(
842            &[
843                (
844                    0,
845                    SemanticType::Timestamp,
846                    ConcreteDataType::timestamp_millisecond_datatype(),
847                ),
848                (1, SemanticType::Field, json2.clone()),
849            ],
850            &[],
851        ));
852        let expected_metadata = Arc::new(new_metadata(
853            &[
854                (
855                    0,
856                    SemanticType::Timestamp,
857                    ConcreteDataType::timestamp_millisecond_datatype(),
858                ),
859                (1, SemanticType::Field, json2),
860                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
861            ],
862            &[],
863        ));
864        let read_columns = ReadColumns::new([0, 1, 2])
865            .with_json_target_types(BTreeMap::from([(1, JsonNativeType::Variant)]));
866        let mapper = FlatProjectionMapper::new_with_read_columns(
867            &expected_metadata,
868            vec![0, 1, 2],
869            read_columns.clone(),
870        )?;
871        let read_format = FlatReadFormat::new(actual_metadata, read_columns, None, "test", false)?;
872
873        let compat = FlatCompatBatch::try_new(&mapper, &read_format, false)?.unwrap();
874        let json_index = mapper
875            .batch_schema()
876            .iter()
877            .position(|(id, _)| *id == 1)
878            .unwrap();
879        assert!(matches!(
880            &compat.index_or_defaults[json_index],
881            IndexOrDefault::Index {
882                cast_type: None,
883                ..
884            }
885        ));
886        Ok(())
887    }
888
889    #[test]
890    fn test_flat_compat_batch_with_read_projection_superset() {
891        let actual_metadata = Arc::new(new_metadata(
892            &[
893                (
894                    0,
895                    SemanticType::Timestamp,
896                    ConcreteDataType::timestamp_millisecond_datatype(),
897                ),
898                (1, SemanticType::Tag, ConcreteDataType::string_datatype()),
899                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
900            ],
901            &[1],
902        ));
903
904        let expected_metadata = Arc::new(new_metadata(
905            &[
906                (
907                    0,
908                    SemanticType::Timestamp,
909                    ConcreteDataType::timestamp_millisecond_datatype(),
910                ),
911                (1, SemanticType::Tag, ConcreteDataType::string_datatype()),
912                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
913                // Adds a new field.
914                (3, SemanticType::Field, ConcreteDataType::int64_datatype()),
915            ],
916            &[1],
917        ));
918
919        let mapper = FlatProjectionMapper::new_with_read_columns(
920            &expected_metadata,
921            vec![1, 2],
922            ReadColumns::new([1, 2, 3]),
923        )
924        .unwrap();
925        let read_format = FlatReadFormat::new(
926            actual_metadata.clone(),
927            ReadColumns::new([1, 2, 3]),
928            None,
929            "test",
930            false,
931        )
932        .unwrap();
933
934        let compat_batch = FlatCompatBatch::try_new(&mapper, &read_format, false)
935            .unwrap()
936            .unwrap();
937
938        let mut tag_builder = StringDictionaryBuilder::<UInt32Type>::new();
939        tag_builder.append_value("tag1");
940        tag_builder.append_value("tag1");
941        let tag_dict_array = Arc::new(tag_builder.finish());
942
943        let k1 = encode_key(&[Some("tag1")]);
944        let input_columns: Vec<ArrayRef> = vec![
945            tag_dict_array.clone(),
946            Arc::new(Int64Array::from(vec![100, 200])),
947            Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])),
948            build_flat_test_pk_array(&[&k1, &k1]),
949            Arc::new(UInt64Array::from_iter_values([1, 2])),
950            Arc::new(UInt8Array::from_iter_values([
951                OpType::Put as u8,
952                OpType::Put as u8,
953            ])),
954        ];
955        let input_schema =
956            to_flat_sst_arrow_schema(&actual_metadata, &FlatSchemaOptions::default());
957        let input_batch = RecordBatch::try_new(input_schema, input_columns).unwrap();
958
959        let result = compat_batch.compat(input_batch).unwrap();
960
961        let expected_schema =
962            to_flat_sst_arrow_schema(&expected_metadata, &FlatSchemaOptions::default());
963        let expected_columns: Vec<ArrayRef> = vec![
964            tag_dict_array.clone(),
965            Arc::new(Int64Array::from(vec![100, 200])),
966            Arc::new(Int64Array::from(vec![None::<i64>, None::<i64>])),
967            Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])),
968            build_flat_test_pk_array(&[&k1, &k1]),
969            Arc::new(UInt64Array::from_iter_values([1, 2])),
970            Arc::new(UInt8Array::from_iter_values([
971                OpType::Put as u8,
972                OpType::Put as u8,
973            ])),
974        ];
975        let expected_batch = RecordBatch::try_new(expected_schema, expected_columns).unwrap();
976
977        assert_eq!(expected_batch, result);
978    }
979
980    #[test]
981    fn test_flat_compat_batch_with_different_pk_encoding() {
982        let mut actual_metadata = new_metadata(
983            &[
984                (
985                    0,
986                    SemanticType::Timestamp,
987                    ConcreteDataType::timestamp_millisecond_datatype(),
988                ),
989                (1, SemanticType::Tag, ConcreteDataType::string_datatype()),
990                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
991            ],
992            &[1],
993        );
994        actual_metadata.primary_key_encoding = PrimaryKeyEncoding::Dense;
995        let actual_metadata = Arc::new(actual_metadata);
996
997        let mut expected_metadata = new_metadata(
998            &[
999                (
1000                    0,
1001                    SemanticType::Timestamp,
1002                    ConcreteDataType::timestamp_millisecond_datatype(),
1003                ),
1004                (1, SemanticType::Tag, ConcreteDataType::string_datatype()),
1005                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
1006                (3, SemanticType::Tag, ConcreteDataType::string_datatype()),
1007            ],
1008            &[1, 3],
1009        );
1010        expected_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse;
1011        let expected_metadata = Arc::new(expected_metadata);
1012
1013        let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
1014        let read_format = FlatReadFormat::new(
1015            actual_metadata.clone(),
1016            ReadColumns::new([0, 1, 2, 3]),
1017            None,
1018            "test",
1019            false,
1020        )
1021        .unwrap();
1022
1023        let compat_batch = FlatCompatBatch::try_new(&mapper, &read_format, false)
1024            .unwrap()
1025            .unwrap();
1026
1027        // Tag array.
1028        let mut tag1_builder = StringDictionaryBuilder::<UInt32Type>::new();
1029        tag1_builder.append_value("tag1");
1030        tag1_builder.append_value("tag1");
1031        let tag1_dict_array = Arc::new(tag1_builder.finish());
1032
1033        let k1 = encode_key(&[Some("tag1")]);
1034        let input_columns: Vec<ArrayRef> = vec![
1035            tag1_dict_array.clone(),
1036            Arc::new(Int64Array::from(vec![100, 200])),
1037            Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])),
1038            build_flat_test_pk_array(&[&k1, &k1]),
1039            Arc::new(UInt64Array::from_iter_values([1, 2])),
1040            Arc::new(UInt8Array::from_iter_values([
1041                OpType::Put as u8,
1042                OpType::Put as u8,
1043            ])),
1044        ];
1045        let input_schema =
1046            to_flat_sst_arrow_schema(&actual_metadata, &FlatSchemaOptions::default());
1047        let input_batch = RecordBatch::try_new(input_schema, input_columns).unwrap();
1048
1049        let result = compat_batch.compat(input_batch).unwrap();
1050
1051        let sparse_k1 = encode_sparse_key(&[(1, Some("tag1")), (3, None)]);
1052        let mut null_tag_builder = StringDictionaryBuilder::<UInt32Type>::new();
1053        null_tag_builder.append_nulls(2);
1054        let null_tag_dict_array = Arc::new(null_tag_builder.finish());
1055        let expected_columns: Vec<ArrayRef> = vec![
1056            tag1_dict_array.clone(),
1057            null_tag_dict_array,
1058            Arc::new(Int64Array::from(vec![100, 200])),
1059            Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])),
1060            build_flat_test_pk_array(&[&sparse_k1, &sparse_k1]),
1061            Arc::new(UInt64Array::from_iter_values([1, 2])),
1062            Arc::new(UInt8Array::from_iter_values([
1063                OpType::Put as u8,
1064                OpType::Put as u8,
1065            ])),
1066        ];
1067        let output_schema =
1068            to_flat_sst_arrow_schema(&expected_metadata, &FlatSchemaOptions::default());
1069        let expected_batch = RecordBatch::try_new(output_schema, expected_columns).unwrap();
1070
1071        assert_eq!(expected_batch, result);
1072    }
1073
1074    #[test]
1075    fn test_compat_primary_key_with_different_encoding_only() {
1076        let mut actual_metadata = new_metadata(
1077            &[
1078                (
1079                    0,
1080                    SemanticType::Timestamp,
1081                    ConcreteDataType::timestamp_millisecond_datatype(),
1082                ),
1083                (1, SemanticType::Tag, ConcreteDataType::string_datatype()),
1084                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
1085            ],
1086            &[1],
1087        );
1088        actual_metadata.primary_key_encoding = PrimaryKeyEncoding::Dense;
1089        let actual_metadata = Arc::new(actual_metadata);
1090
1091        let mut expected_metadata = (*actual_metadata).clone();
1092        expected_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse;
1093        let expected_metadata = Arc::new(expected_metadata);
1094
1095        let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
1096        let read_format = FlatReadFormat::new(
1097            actual_metadata,
1098            ReadColumns::new([0, 1, 2]),
1099            None,
1100            "test",
1101            false,
1102        )
1103        .unwrap();
1104        let compat = FlatCompatBatch::try_new(&mapper, &read_format, false)
1105            .unwrap()
1106            .unwrap();
1107
1108        let dense_key = encode_key(&[Some("tag1")]);
1109        let sparse_key = encode_sparse_key(&[(1, Some("tag1"))]);
1110
1111        let dictionary_key = build_flat_test_pk_array(&[&dense_key, &dense_key]);
1112        let result = compat.compat_primary_key(&dictionary_key).unwrap();
1113        let result = result.as_any().downcast_ref::<PrimaryKeyArray>().unwrap();
1114        let values = result
1115            .values()
1116            .as_any()
1117            .downcast_ref::<BinaryArray>()
1118            .unwrap();
1119        assert_eq!(values.value(result.keys().value(0) as usize), sparse_key);
1120        assert_eq!(values.value(result.keys().value(1) as usize), sparse_key);
1121
1122        let binary_key: ArrayRef = Arc::new(BinaryArray::from(vec![
1123            Some(dense_key.as_slice()),
1124            Some(dense_key.as_slice()),
1125        ]));
1126        let result = compat.compat_primary_key(&binary_key).unwrap();
1127        let result = result.as_any().downcast_ref::<BinaryArray>().unwrap();
1128        assert_eq!(result.value(0), sparse_key);
1129        assert_eq!(result.value(1), sparse_key);
1130    }
1131
1132    /// A widened time index unit must be cast on compat (values rescaled to
1133    /// the expected unit), and identical units must skip compat entirely.
1134    #[test]
1135    fn test_flat_compat_batch_with_time_index_unit_change() {
1136        let actual_metadata = Arc::new(new_metadata(
1137            &[
1138                (
1139                    0,
1140                    SemanticType::Timestamp,
1141                    ConcreteDataType::timestamp_millisecond_datatype(),
1142                ),
1143                (1, SemanticType::Tag, ConcreteDataType::string_datatype()),
1144                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
1145            ],
1146            &[1],
1147        ));
1148
1149        let mut expected_metadata = (*actual_metadata).clone();
1150        for column in expected_metadata.column_metadatas.iter_mut() {
1151            if column.semantic_type == SemanticType::Timestamp {
1152                column.column_schema.data_type = ConcreteDataType::timestamp_microsecond_datatype();
1153            }
1154        }
1155        // Keep the schema consistent with the column metadata.
1156        let expected_metadata = Arc::new(
1157            RegionMetadataBuilder::from_existing(expected_metadata)
1158                .build()
1159                .unwrap(),
1160        );
1161
1162        // Same units: nothing to compat.
1163        let same_unit_metadata = Arc::new((*actual_metadata).clone());
1164        let mapper = FlatProjectionMapper::all(&same_unit_metadata).unwrap();
1165        let read_format = FlatReadFormat::new(
1166            actual_metadata.clone(),
1167            ReadColumns::new([0, 1, 2]),
1168            None,
1169            "test",
1170            false,
1171        )
1172        .unwrap();
1173        assert!(
1174            FlatCompatBatch::try_new(&mapper, &read_format, false)
1175                .unwrap()
1176                .is_none()
1177        );
1178
1179        // Widened unit: the time index column is rescaled, others pass through.
1180        let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
1181        let read_format = FlatReadFormat::new(
1182            actual_metadata.clone(),
1183            ReadColumns::new([0, 1, 2]),
1184            None,
1185            "test",
1186            false,
1187        )
1188        .unwrap();
1189        let compat_batch = FlatCompatBatch::try_new(&mapper, &read_format, false)
1190            .unwrap()
1191            .unwrap();
1192
1193        let mut tag_builder = StringDictionaryBuilder::<UInt32Type>::new();
1194        tag_builder.append_value("tag1");
1195        tag_builder.append_value("tag1");
1196        let tag_dict_array = Arc::new(tag_builder.finish());
1197
1198        let k1 = encode_key(&[Some("tag1")]);
1199        let input_columns: Vec<ArrayRef> = vec![
1200            tag_dict_array.clone(),
1201            Arc::new(Int64Array::from(vec![100, 200])),
1202            Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])),
1203            build_flat_test_pk_array(&[&k1, &k1]),
1204            Arc::new(UInt64Array::from_iter_values([1, 2])),
1205            Arc::new(UInt8Array::from_iter_values([
1206                OpType::Put as u8,
1207                OpType::Put as u8,
1208            ])),
1209        ];
1210        let input_schema =
1211            to_flat_sst_arrow_schema(&actual_metadata, &FlatSchemaOptions::default());
1212        let input_batch = RecordBatch::try_new(input_schema, input_columns).unwrap();
1213
1214        let result = compat_batch.compat(input_batch).unwrap();
1215
1216        let expected_columns: Vec<ArrayRef> = vec![
1217            tag_dict_array.clone(),
1218            Arc::new(Int64Array::from(vec![100, 200])),
1219            // Rescaled, not reinterpreted: 1000ms -> 1_000_000us.
1220            Arc::new(TimestampMicrosecondArray::from_iter_values([
1221                1_000_000, 2_000_000,
1222            ])),
1223            build_flat_test_pk_array(&[&k1, &k1]),
1224            Arc::new(UInt64Array::from_iter_values([1, 2])),
1225            Arc::new(UInt8Array::from_iter_values([
1226                OpType::Put as u8,
1227                OpType::Put as u8,
1228            ])),
1229        ];
1230        let expected_schema =
1231            to_flat_sst_arrow_schema(&expected_metadata, &FlatSchemaOptions::default());
1232        let expected_batch = RecordBatch::try_new(expected_schema, expected_columns).unwrap();
1233
1234        assert_eq!(expected_batch, result);
1235    }
1236
1237    /// A widened time index unit combined with an added column (the realistic
1238    /// multi-alter sequence) must rescale and default-fill in one pass.
1239    #[test]
1240    fn test_flat_compat_batch_time_index_unit_change_with_added_column() {
1241        let actual_metadata = Arc::new(new_metadata(
1242            &[
1243                (
1244                    0,
1245                    SemanticType::Timestamp,
1246                    ConcreteDataType::timestamp_millisecond_datatype(),
1247                ),
1248                (1, SemanticType::Tag, ConcreteDataType::string_datatype()),
1249                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
1250            ],
1251            &[1],
1252        ));
1253
1254        let mut expected_metadata = new_metadata(
1255            &[
1256                (
1257                    0,
1258                    SemanticType::Timestamp,
1259                    ConcreteDataType::timestamp_microsecond_datatype(),
1260                ),
1261                (1, SemanticType::Tag, ConcreteDataType::string_datatype()),
1262                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
1263                (3, SemanticType::Field, ConcreteDataType::int64_datatype()),
1264            ],
1265            &[1],
1266        );
1267        expected_metadata.primary_key_encoding = PrimaryKeyEncoding::Dense;
1268        let expected_metadata = Arc::new(expected_metadata);
1269
1270        let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
1271        let read_format = FlatReadFormat::new(
1272            actual_metadata.clone(),
1273            ReadColumns::new([0, 1, 2, 3]),
1274            None,
1275            "test",
1276            false,
1277        )
1278        .unwrap();
1279
1280        let compat_batch = FlatCompatBatch::try_new(&mapper, &read_format, false)
1281            .unwrap()
1282            .unwrap();
1283
1284        let mut tag_builder = StringDictionaryBuilder::<UInt32Type>::new();
1285        tag_builder.append_value("tag1");
1286        tag_builder.append_value("tag1");
1287        let tag_dict_array = Arc::new(tag_builder.finish());
1288
1289        let k1 = encode_key(&[Some("tag1")]);
1290        let input_columns: Vec<ArrayRef> = vec![
1291            tag_dict_array.clone(),
1292            Arc::new(Int64Array::from(vec![100, 200])),
1293            Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])),
1294            build_flat_test_pk_array(&[&k1, &k1]),
1295            Arc::new(UInt64Array::from_iter_values([1, 2])),
1296            Arc::new(UInt8Array::from_iter_values([
1297                OpType::Put as u8,
1298                OpType::Put as u8,
1299            ])),
1300        ];
1301        let input_schema =
1302            to_flat_sst_arrow_schema(&actual_metadata, &FlatSchemaOptions::default());
1303        let input_batch = RecordBatch::try_new(input_schema, input_columns).unwrap();
1304
1305        let result = compat_batch.compat(input_batch).unwrap();
1306
1307        let expected_columns: Vec<ArrayRef> = vec![
1308            tag_dict_array.clone(),
1309            Arc::new(Int64Array::from(vec![100, 200])),
1310            Arc::new(Int64Array::from(vec![None::<i64>, None::<i64>])),
1311            Arc::new(TimestampMicrosecondArray::from_iter_values([
1312                1_000_000, 2_000_000,
1313            ])),
1314            build_flat_test_pk_array(&[&k1, &k1]),
1315            Arc::new(UInt64Array::from_iter_values([1, 2])),
1316            Arc::new(UInt8Array::from_iter_values([
1317                OpType::Put as u8,
1318                OpType::Put as u8,
1319            ])),
1320        ];
1321        let expected_schema =
1322            to_flat_sst_arrow_schema(&expected_metadata, &FlatSchemaOptions::default());
1323        let expected_batch = RecordBatch::try_new(expected_schema, expected_columns).unwrap();
1324
1325        assert_eq!(expected_batch, result);
1326    }
1327
1328    /// The sparse compaction compat path must also rescale a widened time
1329    /// index unit (metric-engine regions compact with sparse encoding).
1330    #[test]
1331    fn test_flat_compat_batch_compact_sparse_time_index_unit_change() {
1332        let mut actual_metadata = new_metadata(
1333            &[
1334                (
1335                    0,
1336                    SemanticType::Timestamp,
1337                    ConcreteDataType::timestamp_millisecond_datatype(),
1338                ),
1339                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
1340            ],
1341            &[],
1342        );
1343        actual_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse;
1344        let actual_metadata = Arc::new(actual_metadata);
1345
1346        let mut expected_metadata = (*actual_metadata).clone();
1347        for column in expected_metadata.column_metadatas.iter_mut() {
1348            if column.semantic_type == SemanticType::Timestamp {
1349                column.column_schema.data_type = ConcreteDataType::timestamp_microsecond_datatype();
1350            }
1351        }
1352        let expected_metadata = Arc::new(
1353            RegionMetadataBuilder::from_existing(expected_metadata)
1354                .build()
1355                .unwrap(),
1356        );
1357
1358        let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
1359        let read_format = FlatReadFormat::new(
1360            actual_metadata.clone(),
1361            ReadColumns::new([0, 2]),
1362            None,
1363            "test",
1364            true,
1365        )
1366        .unwrap();
1367
1368        let compat_batch = FlatCompatBatch::try_new(&mapper, &read_format, true)
1369            .unwrap()
1370            .unwrap();
1371
1372        let sparse_k1 = encode_sparse_key(&[]);
1373        let input_columns: Vec<ArrayRef> = vec![
1374            Arc::new(Int64Array::from(vec![100, 200])),
1375            Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])),
1376            build_flat_test_pk_array(&[&sparse_k1, &sparse_k1]),
1377            Arc::new(UInt64Array::from_iter_values([1, 2])),
1378            Arc::new(UInt8Array::from_iter_values([
1379                OpType::Put as u8,
1380                OpType::Put as u8,
1381            ])),
1382        ];
1383        let input_schema =
1384            to_flat_sst_arrow_schema(&actual_metadata, &FlatSchemaOptions::default());
1385        let input_batch = RecordBatch::try_new(input_schema, input_columns).unwrap();
1386
1387        let result = compat_batch.compat(input_batch).unwrap();
1388
1389        let expected_columns: Vec<ArrayRef> = vec![
1390            Arc::new(Int64Array::from(vec![100, 200])),
1391            Arc::new(TimestampMicrosecondArray::from_iter_values([
1392                1_000_000, 2_000_000,
1393            ])),
1394            build_flat_test_pk_array(&[&sparse_k1, &sparse_k1]),
1395            Arc::new(UInt64Array::from_iter_values([1, 2])),
1396            Arc::new(UInt8Array::from_iter_values([
1397                OpType::Put as u8,
1398                OpType::Put as u8,
1399            ])),
1400        ];
1401        let output_schema =
1402            to_flat_sst_arrow_schema(&expected_metadata, &FlatSchemaOptions::default());
1403        let expected_batch = RecordBatch::try_new(output_schema, expected_columns).unwrap();
1404
1405        assert_eq!(expected_batch, result);
1406    }
1407
1408    #[test]
1409    fn test_flat_compat_batch_compact_sparse() {
1410        let mut actual_metadata = new_metadata(
1411            &[
1412                (
1413                    0,
1414                    SemanticType::Timestamp,
1415                    ConcreteDataType::timestamp_millisecond_datatype(),
1416                ),
1417                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
1418            ],
1419            &[],
1420        );
1421        actual_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse;
1422        let actual_metadata = Arc::new(actual_metadata);
1423
1424        let mut expected_metadata = new_metadata(
1425            &[
1426                (
1427                    0,
1428                    SemanticType::Timestamp,
1429                    ConcreteDataType::timestamp_millisecond_datatype(),
1430                ),
1431                (2, SemanticType::Field, ConcreteDataType::int64_datatype()),
1432                (3, SemanticType::Field, ConcreteDataType::int64_datatype()),
1433            ],
1434            &[],
1435        );
1436        expected_metadata.primary_key_encoding = PrimaryKeyEncoding::Sparse;
1437        let expected_metadata = Arc::new(expected_metadata);
1438
1439        let mapper = FlatProjectionMapper::all(&expected_metadata).unwrap();
1440        let read_format = FlatReadFormat::new(
1441            actual_metadata.clone(),
1442            ReadColumns::new([0, 2, 3]),
1443            None,
1444            "test",
1445            true,
1446        )
1447        .unwrap();
1448
1449        let compat_batch = FlatCompatBatch::try_new(&mapper, &read_format, true)
1450            .unwrap()
1451            .unwrap();
1452
1453        let sparse_k1 = encode_sparse_key(&[]);
1454        let input_columns: Vec<ArrayRef> = vec![
1455            Arc::new(Int64Array::from(vec![100, 200])),
1456            Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])),
1457            build_flat_test_pk_array(&[&sparse_k1, &sparse_k1]),
1458            Arc::new(UInt64Array::from_iter_values([1, 2])),
1459            Arc::new(UInt8Array::from_iter_values([
1460                OpType::Put as u8,
1461                OpType::Put as u8,
1462            ])),
1463        ];
1464        let input_schema =
1465            to_flat_sst_arrow_schema(&actual_metadata, &FlatSchemaOptions::default());
1466        let input_batch = RecordBatch::try_new(input_schema, input_columns).unwrap();
1467
1468        let result = compat_batch.compat(input_batch).unwrap();
1469
1470        let expected_columns: Vec<ArrayRef> = vec![
1471            Arc::new(Int64Array::from(vec![100, 200])),
1472            Arc::new(Int64Array::from(vec![None::<i64>, None::<i64>])),
1473            Arc::new(TimestampMillisecondArray::from_iter_values([1000, 2000])),
1474            build_flat_test_pk_array(&[&sparse_k1, &sparse_k1]),
1475            Arc::new(UInt64Array::from_iter_values([1, 2])),
1476            Arc::new(UInt8Array::from_iter_values([
1477                OpType::Put as u8,
1478                OpType::Put as u8,
1479            ])),
1480        ];
1481        let output_schema =
1482            to_flat_sst_arrow_schema(&expected_metadata, &FlatSchemaOptions::default());
1483        let expected_batch = RecordBatch::try_new(output_schema, expected_columns).unwrap();
1484
1485        assert_eq!(expected_batch, result);
1486    }
1487}
mito2/read/compat.rs

mito2/read/
compat.rs
