mito2/

sst.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Sorted strings tables.

use std::collections::HashMap;
use std::sync::Arc;

use api::v1::SemanticType;
use arrow_schema::DataType;
use arrow_schema::extension::{EXTENSION_TYPE_NAME_KEY, ExtensionType};
use common_base::readable_size::ReadableSize;
use common_query::native_histogram::{
    is_native_histogram_value_type, native_histogram_list_element_id, native_histogram_subfield_id,
};
use datatypes::arrow::datatypes::{
    DataType as ArrowDataType, Field, FieldRef, Fields, Schema, SchemaRef,
};
use datatypes::arrow::record_batch::RecordBatch;
use datatypes::extension::histogram::HistogramExtensionType;
use datatypes::prelude::ConcreteDataType;
use datatypes::timestamp::timestamp_array_to_primitive;
use serde::{Deserialize, Serialize};
use store_api::codec::PrimaryKeyEncoding;
use store_api::metadata::RegionMetadata;
use store_api::storage::consts::{
    OP_TYPE_COLUMN_NAME, PRIMARY_KEY_COLUMN_NAME, SEQUENCE_COLUMN_NAME,
};

use crate::error::{InvalidNativeHistogramFieldIdSnafu, InvalidNativeHistogramSubfieldSnafu};
use crate::sst::parquet::flat_format::time_index_column_index;

pub mod file;
pub mod file_purger;
pub mod file_ref;
pub mod index;
pub mod location;
pub mod parquet;
pub(crate) mod primary_key;
pub mod range_index;
pub(crate) mod version;

/// Default write buffer size, it should be greater than the default minimum upload part of S3 (5mb).
pub const DEFAULT_WRITE_BUFFER_SIZE: ReadableSize = ReadableSize::mb(8);

/// Default number of concurrent write, it only works on object store backend(e.g., S3).
pub const DEFAULT_WRITE_CONCURRENCY: usize = 8;

/// Format type of the SST file.
#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, strum::EnumString)]
#[serde(rename_all = "snake_case")]
#[strum(serialize_all = "snake_case")]
pub enum FormatType {
    /// Parquet with primary key encoded.
    #[default]
    PrimaryKey,
    /// Flat Parquet format.
    Flat,
}

/// Iceberg-compatible column field ID key stored in Parquet column metadata.
pub const PARQUET_FIELD_ID_KEY: &str = "PARQUET:field_id";

/// Adds `PARQUET:field_id` metadata to a top-level Arrow field.
///
/// Native-histogram sub-field ids are stamped separately at parquet-write
/// time by [`stamp_native_histogram_subfield_ids`], not here.
pub fn with_field_id(mut field: Field, column_id: u32) -> Field {
    field
        .metadata_mut()
        .insert(PARQUET_FIELD_ID_KEY.to_string(), column_id.to_string());
    field
}

/// Stamps the `greptime.histogram` extension and reserved `PARQUET:field_id`s
/// onto a native-histogram struct field (and its sub-fields / list element
/// fields), so external readers can identify it by extension and resolve
/// nested fields by id.
///
/// Detection is by the exact native-histogram struct type
/// (`is_native_histogram_value_type`); other struct columns are left untouched.
/// mito2 reads SST columns by schema position, never by field metadata, so this
/// only affects external readers.
///
/// Returns an error if the parent column's `PARQUET:field_id` is missing,
/// malformed, or exceeds `i32::MAX`, or if a sub-field id cannot be derived
/// — because the sub-field name is not a known native-histogram field, or
/// the derived id overflows a positive `i32` (an absurdly large parent
/// `column_id`); see [`native_histogram_subfield_id`].
fn stamp_native_histogram_subfield_ids(field: &mut Field) -> crate::error::Result<()> {
    if !is_native_histogram_value_type(&ConcreteDataType::from_arrow_type(field.data_type())) {
        return Ok(());
    }
    // Namespace sub-field ids by the parent column's field id (its
    // `PARQUET:field_id`, stamped earlier by `with_field_id`) so several
    // histogram columns in one table get disjoint ids. Fail loudly if the id
    // is absent, malformed, or too large to fit a positive `i32`.
    let column_id = field
        .metadata()
        .get(PARQUET_FIELD_ID_KEY)
        .and_then(|s| s.parse::<i32>().ok())
        .ok_or_else(|| {
            InvalidNativeHistogramFieldIdSnafu {
                field_name: field.name().clone(),
            }
            .build()
        })?;
    // Tag the field with the greptime.histogram extension.
    field.metadata_mut().insert(
        EXTENSION_TYPE_NAME_KEY.to_string(),
        HistogramExtensionType::NAME.to_string(),
    );
    let ArrowDataType::Struct(children) = field.data_type() else {
        return Ok(());
    };
    let new_children: crate::error::Result<Fields> = children
        .iter()
        .map(|child| {
            let mut c = (**child).clone();
            // `None` here means either the sub-field name is not a known
            // native-histogram field, or the derived id overflowed i32.
            // Surface it as an error rather than silently leaving the field
            // without an id.
            let id = native_histogram_subfield_id(column_id, c.name()).ok_or_else(|| {
                InvalidNativeHistogramSubfieldSnafu {
                    column_id,
                    field_name: c.name().clone(),
                }
                .build()
            })?;
            // Stamp the sub-field's own id.
            c.metadata_mut()
                .insert(PARQUET_FIELD_ID_KEY.to_string(), id.to_string());
            // If the sub-field is a list, stamp its element field's id.
            if let ArrowDataType::List(elem) = c.data_type() {
                let elem_id =
                    native_histogram_list_element_id(column_id, c.name()).ok_or_else(|| {
                        InvalidNativeHistogramSubfieldSnafu {
                            column_id,
                            field_name: c.name().clone(),
                        }
                        .build()
                    })?;
                let mut new_elem = (**elem).clone();
                new_elem
                    .metadata_mut()
                    .insert(PARQUET_FIELD_ID_KEY.to_string(), elem_id.to_string());
                c.set_data_type(ArrowDataType::List(Arc::new(new_elem)));
            }
            Ok(Arc::new(c))
        })
        .collect();
    field.set_data_type(ArrowDataType::Struct(new_children?));
    Ok(())
}

/// Returns a copy of `schema` with native-histogram sub-field ids stamped,
/// for the parquet writer.
///
/// This is called on the schema handed to `AsyncArrowWriter`, not in
/// [`with_field_id`], because the SST arrow schema is also the memtable's
/// in-memory schema, whose `Struct` equality (`PartialEq`) is
/// metadata-sensitive — stamping there would break writes. The parquet writer
/// compares types with `DataType::equals_datatype`, which ignores field
/// metadata, so a stamped schema accepts an unstamped batch.
pub fn maybe_wrap_schema(schema: &SchemaRef) -> crate::error::Result<SchemaRef> {
    // Fast path: only a struct column can be a native histogram; if there are
    // none, skip the rebuild.
    if !schema
        .fields()
        .iter()
        .any(|f| matches!(f.data_type(), ArrowDataType::Struct(_)))
    {
        return Ok(schema.clone());
    }
    let new_fields: crate::error::Result<Vec<FieldRef>> = schema
        .fields()
        .iter()
        .map(|f| {
            let mut field = (**f).clone();
            stamp_native_histogram_subfield_ids(&mut field)?;
            Ok(Arc::new(field))
        })
        .collect();
    Ok(Arc::new(Schema::new_with_metadata(
        Fields::from(new_fields?),
        schema.metadata().clone(),
    )))
}

/// Parquet field ID base for internal columns (__primary_key, __sequence, __op_type).
/// Uses bit 30 to distinguish from user column IDs and fit in positive i32 range.
pub(crate) const INTERNAL_PARQUET_FIELD_ID_BASE: u32 = 1 << 30;

/// Parquet field ID for the __primary_key column.
pub(crate) const PRIMARY_KEY_PARQUET_FIELD_ID: u32 = INTERNAL_PARQUET_FIELD_ID_BASE;
/// Parquet field ID for the __sequence column.
pub(crate) const SEQUENCE_PARQUET_FIELD_ID: u32 = INTERNAL_PARQUET_FIELD_ID_BASE + 1;
/// Parquet field ID for the __op_type column.
pub(crate) const OP_TYPE_PARQUET_FIELD_ID: u32 = INTERNAL_PARQUET_FIELD_ID_BASE + 2;

/// Gets the arrow schema to store in parquet.
pub fn to_sst_arrow_schema(metadata: &RegionMetadata) -> SchemaRef {
    let fields = Fields::from_iter(
        metadata
            .schema
            .arrow_schema()
            .fields()
            .iter()
            .zip(&metadata.column_metadatas)
            .filter_map(|(field, column_meta)| {
                if column_meta.semantic_type == SemanticType::Field {
                    Some(Arc::new(with_field_id(
                        (**field).clone(),
                        column_meta.column_id,
                    )))
                } else {
                    // We have fixed positions for tags (primary key) and time index.
                    None
                }
            })
            .chain([Arc::new(with_field_id(
                (*metadata.time_index_field()).clone(),
                metadata.time_index_column().column_id,
            ))])
            .chain(internal_fields()),
    );

    Arc::new(Schema::new(fields))
}

/// Options of flat schema.
pub struct FlatSchemaOptions {
    /// Whether to store primary key columns additionally instead of an encoded column.
    pub raw_pk_columns: bool,
    /// Whether to use dictionary encoding for string primary key columns
    /// when storing primary key columns.
    /// Only takes effect when `raw_pk_columns` is true.
    pub string_pk_use_dict: bool,
    /// The column's concretized JSON types, to be set into Arrow schema.
    /// Otherwise it's empty struct in the Arrow schema.
    pub concretized_json_types: HashMap<String, DataType>,
}

impl Default for FlatSchemaOptions {
    fn default() -> Self {
        Self {
            raw_pk_columns: true,
            string_pk_use_dict: true,
            concretized_json_types: HashMap::new(),
        }
    }
}

impl FlatSchemaOptions {
    /// Creates a options according to the primary key encoding.
    pub fn from_encoding(encoding: PrimaryKeyEncoding) -> Self {
        if encoding == PrimaryKeyEncoding::Dense {
            Self::default()
        } else {
            Self {
                raw_pk_columns: false,
                string_pk_use_dict: false,
                concretized_json_types: HashMap::new(),
            }
        }
    }
}

/// Gets the arrow schema to store in parquet.
///
/// The schema is:
/// ```text
/// primary key columns, field columns, time index, __primary_key, __sequence, __op_type
/// ```
///
/// # Panics
/// Panics if the metadata is invalid.
pub fn to_flat_sst_arrow_schema(
    metadata: &RegionMetadata,
    options: &FlatSchemaOptions,
) -> SchemaRef {
    let num_fields = flat_sst_arrow_schema_column_num(metadata, options);
    let mut fields = Vec::with_capacity(num_fields);
    let schema = metadata.schema.arrow_schema();
    if options.raw_pk_columns {
        for pk_id in &metadata.primary_key {
            let pk_index = metadata.column_index_by_id(*pk_id).unwrap();
            let column_id = metadata.column_metadatas[pk_index].column_id;
            if options.string_pk_use_dict {
                let old_field = &schema.fields[pk_index];
                let new_field = tag_maybe_to_dictionary_field(
                    &metadata.column_metadatas[pk_index].column_schema.data_type,
                    old_field,
                );
                let new_field = concretize_json_type(new_field, options);
                fields.push(Arc::new(with_field_id((*new_field).clone(), column_id)));
            }
        }
    }
    let remaining_fields = schema
        .fields()
        .iter()
        .zip(&metadata.column_metadatas)
        .filter_map(|(field, column_meta)| {
            if column_meta.semantic_type == SemanticType::Field {
                let field = concretize_json_type(field.clone(), options);
                Some(Arc::new(with_field_id(
                    Arc::unwrap_or_clone(field),
                    column_meta.column_id,
                )))
            } else {
                None
            }
        })
        .chain([Arc::new(with_field_id(
            (*metadata.time_index_field()).clone(),
            metadata.time_index_column().column_id,
        ))])
        .chain(internal_fields());
    for field in remaining_fields {
        fields.push(field);
    }

    Arc::new(Schema::new(fields))
}

fn concretize_json_type(field: Arc<Field>, options: &FlatSchemaOptions) -> Arc<Field> {
    if let Some(data_type) = options.concretized_json_types.get(field.name()) {
        let mut field = Arc::unwrap_or_clone(field);
        field.set_data_type(data_type.clone());
        Arc::new(field)
    } else {
        field
    }
}

/// Returns the number of columns in the flat format.
pub fn flat_sst_arrow_schema_column_num(
    metadata: &RegionMetadata,
    options: &FlatSchemaOptions,
) -> usize {
    if options.raw_pk_columns {
        metadata.column_metadatas.len() + 3
    } else {
        metadata.column_metadatas.len() + 3 - metadata.primary_key.len()
    }
}

/// Helper function to create a dictionary field from a field.
fn to_dictionary_field(field: &Field) -> Field {
    let mut new_field = Field::new_dictionary(
        field.name(),
        datatypes::arrow::datatypes::DataType::UInt32,
        field.data_type().clone(),
        field.is_nullable(),
    );

    // retain field_id metadata
    if let Some(field_id) = field.metadata().get(PARQUET_FIELD_ID_KEY) {
        new_field
            .metadata_mut()
            .insert(PARQUET_FIELD_ID_KEY.to_string(), field_id.clone());
    }

    new_field
}

/// Helper function to create a dictionary field from a field if it is a string column.
pub(crate) fn tag_maybe_to_dictionary_field(
    data_type: &ConcreteDataType,
    field: &Arc<Field>,
) -> Arc<Field> {
    if data_type.is_string() {
        Arc::new(to_dictionary_field(field))
    } else {
        field.clone()
    }
}

/// Fields for internal columns.
pub(crate) fn internal_fields() -> [FieldRef; 3] {
    // Internal columns are always not null.
    [
        Arc::new(with_field_id(
            Field::new_dictionary(
                PRIMARY_KEY_COLUMN_NAME,
                ArrowDataType::UInt32,
                ArrowDataType::Binary,
                false,
            ),
            PRIMARY_KEY_PARQUET_FIELD_ID,
        )),
        Arc::new(with_field_id(
            Field::new(SEQUENCE_COLUMN_NAME, ArrowDataType::UInt64, false),
            SEQUENCE_PARQUET_FIELD_ID,
        )),
        Arc::new(with_field_id(
            Field::new(OP_TYPE_COLUMN_NAME, ArrowDataType::UInt8, false),
            OP_TYPE_PARQUET_FIELD_ID,
        )),
    ]
}

/// Returns a copy of `schema` with the `__primary_key` field replaced by a plain `Binary` field.
pub(crate) fn override_pk_field_to_binary(schema: &SchemaRef) -> SchemaRef {
    let new_fields = schema
        .fields()
        .iter()
        .map(|field| {
            if field.name() == PRIMARY_KEY_COLUMN_NAME {
                let mut new_field = Field::new(
                    PRIMARY_KEY_COLUMN_NAME,
                    ArrowDataType::Binary,
                    field.is_nullable(),
                );
                // Preserve the field_id metadata so parquet readers that require
                // all columns to carry a field_id don't fail.
                if let Some(field_id) = field.metadata().get(PARQUET_FIELD_ID_KEY) {
                    new_field
                        .metadata_mut()
                        .insert(PARQUET_FIELD_ID_KEY.to_string(), field_id.clone());
                }
                Arc::new(new_field)
            } else {
                field.clone()
            }
        })
        .collect::<Vec<_>>();
    Arc::new(Schema::new(new_fields))
}

/// Gets the estimated number of series from record batches.
///
/// This struct tracks the last timestamp value to detect series boundaries
/// by observing when timestamps decrease (indicating a new series).
#[derive(Default)]
pub(crate) struct SeriesEstimator {
    /// The last timestamp value seen
    last_timestamp: Option<i64>,
    /// The estimated number of series
    series_count: u64,
}

impl SeriesEstimator {
    /// Updates the estimator with a new record batch in flat format.
    ///
    /// This method examines the time index column to detect series boundaries.
    pub(crate) fn update_flat(&mut self, record_batch: &RecordBatch) {
        let batch_rows = record_batch.num_rows();
        if batch_rows == 0 {
            return;
        }

        let time_index_pos = time_index_column_index(record_batch.num_columns());
        let timestamps = record_batch.column(time_index_pos);
        let Some((ts_values, _unit)) = timestamp_array_to_primitive(timestamps) else {
            return;
        };
        let values = ts_values.values();

        // Checks if there's a boundary between the last batch and this batch
        if let Some(last_ts) = self.last_timestamp {
            if values[0] <= last_ts {
                self.series_count += 1;
            }
        } else {
            // First batch, counts as first series
            self.series_count = 1;
        }

        // Counts series boundaries within this batch.
        for i in 0..batch_rows - 1 {
            // We assumes the same timestamp as a new series, which is different from
            // how we split batches.
            if values[i] >= values[i + 1] {
                self.series_count += 1;
            }
        }

        // Updates the last timestamp
        self.last_timestamp = Some(values[batch_rows - 1]);
    }

    /// Returns the estimated number of series.
    pub(crate) fn finish(&mut self) -> u64 {
        self.last_timestamp = None;
        let count = self.series_count;
        self.series_count = 0;

        count
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use ::parquet::arrow::AsyncArrowWriter;
    use ::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
    use ::parquet::basic::LogicalType;
    use ::parquet::variant::{VariantArray, VariantType, json_to_variant};
    use common_query::prelude::greptime_native_histogram;
    use datatypes::arrow::array::{
        ArrayRef, BinaryArray, DictionaryArray, Int64Array, StringArray, StructArray,
        TimestampMillisecondArray, UInt8Array, UInt32Array, UInt64Array,
    };
    use datatypes::arrow::datatypes::{DataType as ArrowDataType, Field, Schema, TimeUnit};
    use datatypes::arrow::record_batch::RecordBatch;
    use datatypes::extension::json::{Json2ExtensionType, Json2PhysicalLayout};
    use datatypes::vectors::json::array::JsonArray;
    use serde_json::json;

    use super::*;

    fn new_flat_record_batch(timestamps: &[i64]) -> RecordBatch {
        // Flat format has: [fields..., time_index, __primary_key, __sequence, __op_type]
        let num_cols = 4; // time_index + 3 internal columns
        let time_index_pos = time_index_column_index(num_cols);
        assert_eq!(time_index_pos, 0); // For 4 columns, time index should be at position 0

        let time_array = Arc::new(TimestampMillisecondArray::from(timestamps.to_vec()));
        let pk_array = Arc::new(DictionaryArray::new(
            UInt32Array::from(vec![0; timestamps.len()]),
            Arc::new(BinaryArray::from(vec![b"test".as_slice()])),
        ));
        let seq_array = Arc::new(UInt64Array::from(vec![1; timestamps.len()]));
        let op_array = Arc::new(UInt8Array::from(vec![1; timestamps.len()]));

        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "time",
                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            Field::new_dictionary(
                "__primary_key",
                ArrowDataType::UInt32,
                ArrowDataType::Binary,
                false,
            ),
            Field::new("__sequence", ArrowDataType::UInt64, false),
            Field::new("__op_type", ArrowDataType::UInt8, false),
        ]));

        RecordBatch::try_new(schema, vec![time_array, pk_array, seq_array, op_array]).unwrap()
    }

    #[test]
    fn test_series_estimator_flat_empty_batch() {
        let mut estimator = SeriesEstimator::default();
        let record_batch = new_flat_record_batch(&[]);
        estimator.update_flat(&record_batch);
        assert_eq!(0, estimator.finish());
    }

    #[test]
    fn test_series_estimator_flat_single_batch() {
        let mut estimator = SeriesEstimator::default();
        let record_batch = new_flat_record_batch(&[1, 2, 3]);
        estimator.update_flat(&record_batch);
        assert_eq!(1, estimator.finish());
    }

    #[test]
    fn test_series_estimator_flat_series_boundary_within_batch() {
        let mut estimator = SeriesEstimator::default();
        // Timestamps decrease from 3 to 2, indicating a series boundary
        let record_batch = new_flat_record_batch(&[1, 2, 3, 2, 4, 5]);
        estimator.update_flat(&record_batch);
        // Should detect boundary at position 3 (3 >= 2)
        assert_eq!(2, estimator.finish());
    }

    #[test]
    fn test_series_estimator_flat_multiple_boundaries_within_batch() {
        let mut estimator = SeriesEstimator::default();
        // Multiple series boundaries: 5>=4, 6>=3
        let record_batch = new_flat_record_batch(&[1, 2, 5, 4, 6, 3, 7]);
        estimator.update_flat(&record_batch);
        assert_eq!(3, estimator.finish());
    }

    #[test]
    fn test_series_estimator_flat_equal_timestamps() {
        let mut estimator = SeriesEstimator::default();
        // Equal timestamps are considered as new series
        let record_batch = new_flat_record_batch(&[1, 2, 2, 3, 3, 3, 4]);
        estimator.update_flat(&record_batch);
        // Boundaries at: 2>=2, 3>=3, 3>=3
        assert_eq!(4, estimator.finish());
    }

    #[test]
    fn test_series_estimator_flat_multiple_batches_continuation() {
        let mut estimator = SeriesEstimator::default();

        // First batch: timestamps 1, 2, 3
        let batch1 = new_flat_record_batch(&[1, 2, 3]);
        estimator.update_flat(&batch1);

        // Second batch: timestamps 4, 5, 6 (continuation)
        let batch2 = new_flat_record_batch(&[4, 5, 6]);
        estimator.update_flat(&batch2);

        assert_eq!(1, estimator.finish());
    }

    #[test]
    fn test_series_estimator_flat_multiple_batches_new_series() {
        let mut estimator = SeriesEstimator::default();

        // First batch: timestamps 1, 2, 3
        let batch1 = new_flat_record_batch(&[1, 2, 3]);
        estimator.update_flat(&batch1);

        // Second batch: timestamps 2, 3, 4 (goes back to 2, new series)
        let batch2 = new_flat_record_batch(&[2, 3, 4]);
        estimator.update_flat(&batch2);

        assert_eq!(2, estimator.finish());
    }

    #[test]
    fn test_series_estimator_flat_boundary_at_batch_edge_equal() {
        let mut estimator = SeriesEstimator::default();

        // First batch ending at 5
        let batch1 = new_flat_record_batch(&[1, 2, 5]);
        estimator.update_flat(&batch1);

        // Second batch starting at 5 (equal timestamp, new series)
        let batch2 = new_flat_record_batch(&[5, 6, 7]);
        estimator.update_flat(&batch2);

        assert_eq!(2, estimator.finish());
    }

    #[test]
    fn test_series_estimator_flat_mixed_batches() {
        let mut estimator = SeriesEstimator::default();

        // Batch 1: single series [10, 20, 30]
        let batch1 = new_flat_record_batch(&[10, 20, 30]);
        estimator.update_flat(&batch1);

        // Batch 2: starts new series [5, 15], boundary within batch [15, 10, 25]
        let batch2 = new_flat_record_batch(&[5, 15, 10, 25]);
        estimator.update_flat(&batch2);

        // Batch 3: continues from 25 to [30, 35]
        let batch3 = new_flat_record_batch(&[30, 35]);
        estimator.update_flat(&batch3);

        // Expected: 1 (batch1) + 1 (batch2 start) + 1 (within batch2) = 3
        assert_eq!(3, estimator.finish());
    }

    #[test]
    fn test_series_estimator_flat_descending_timestamps() {
        let mut estimator = SeriesEstimator::default();
        // Strictly descending timestamps - each pair creates a boundary
        let record_batch = new_flat_record_batch(&[10, 9, 8, 7, 6]);
        estimator.update_flat(&record_batch);
        // Boundaries: 10>=9, 9>=8, 8>=7, 7>=6 = 4 boundaries + 1 initial = 5 series
        assert_eq!(5, estimator.finish());
    }

    #[test]
    fn test_series_estimator_flat_finish_resets_state() {
        let mut estimator = SeriesEstimator::default();

        let batch1 = new_flat_record_batch(&[1, 2, 3]);
        estimator.update_flat(&batch1);

        assert_eq!(1, estimator.finish());

        // After finish, state should be reset
        let batch2 = new_flat_record_batch(&[4, 5, 6]);
        estimator.update_flat(&batch2);

        assert_eq!(1, estimator.finish());
    }

    /// Build a native-histogram struct field whose top-level `PARQUET:field_id`
    /// is `column_id` (as `with_field_id` does on the real write path).
    fn histogram_field(name: &str, column_id: u32) -> Field {
        use common_query::native_histogram::native_histogram_value_type;
        use datatypes::data_type::DataType;
        with_field_id(
            Field::new(name, native_histogram_value_type().as_arrow_type(), true),
            column_id,
        )
    }

    /// Asserts `field` is a stamped native-histogram struct: it carries the
    /// `greptime.histogram` extension and every sub-field (and list element)
    /// carries its reserved `PARQUET:field_id` namespaced by `column_id`.
    fn assert_histogram_stamped(field: &Field, column_id: i32) {
        use arrow_schema::extension::ExtensionType;
        use common_query::native_histogram::{
            native_histogram_list_element_id, native_histogram_subfield_id,
        };
        use datatypes::extension::histogram::HistogramExtensionType;

        assert_eq!(
            field
                .metadata()
                .get(arrow_schema::extension::EXTENSION_TYPE_NAME_KEY)
                .map(|s| s.as_str()),
            Some(HistogramExtensionType::NAME),
            "histogram field must carry the greptime.histogram extension"
        );
        let ArrowDataType::Struct(children) = field.data_type() else {
            panic!("expected a struct, got {:?}", field.data_type());
        };
        for child in children {
            let expected = native_histogram_subfield_id(column_id, child.name())
                .unwrap_or_else(|| panic!("no id for sub-field {}", child.name()));
            let got: i32 = child
                .metadata()
                .get(PARQUET_FIELD_ID_KEY)
                .unwrap_or_else(|| panic!("sub-field {} missing field id", child.name()))
                .parse()
                .unwrap();
            assert_eq!(got, expected, "sub-field {} id", child.name());
            if let ArrowDataType::List(elem) = child.data_type() {
                let elem_expected =
                    native_histogram_list_element_id(column_id, child.name()).unwrap();
                let elem_got: i32 = elem
                    .metadata()
                    .get(PARQUET_FIELD_ID_KEY)
                    .unwrap_or_else(|| panic!("list element of {} missing id", child.name()))
                    .parse()
                    .unwrap();
                assert_eq!(
                    elem_got,
                    elem_expected,
                    "list element id of {}",
                    child.name()
                );
            }
        }
    }

    #[test]
    fn test_maybe_wrap_schema_native_histogram() {
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "greptime_timestamp",
                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            histogram_field(greptime_native_histogram(), 1),
        ]));

        let wrapped = maybe_wrap_schema(&schema).unwrap();
        let hist = wrapped
            .field_with_name(greptime_native_histogram())
            .expect("histogram field present");
        // The struct has 18 sub-fields.
        let ArrowDataType::Struct(children) = hist.data_type() else {
            unreachable!()
        };
        assert_eq!(children.len(), 18);
        assert_histogram_stamped(hist, 1);
    }

    #[test]
    fn test_maybe_wrap_schema_multiple_histograms_disjoint_ids() {
        // Two histogram columns with distinct parent column ids get disjoint
        // sub-field ids (defensive: the metric engine yields at most one
        // histogram column, but the scheme must stay correct if more appear).
        use common_query::native_histogram::native_histogram_subfield_id;

        let schema = Arc::new(Schema::new(vec![
            histogram_field(greptime_native_histogram(), 1),
            histogram_field(greptime_native_histogram(), 7),
        ]));
        let wrapped = maybe_wrap_schema(&schema).unwrap();
        let h1 = &wrapped.fields()[0];
        let h2 = &wrapped.fields()[1];
        assert_histogram_stamped(h1, 1);
        assert_histogram_stamped(h2, 7);
        // The same sub-field name resolves to different ids across columns.
        assert_ne!(
            native_histogram_subfield_id(1, "sum"),
            native_histogram_subfield_id(7, "sum")
        );
    }

    #[test]
    fn test_maybe_wrap_schema_recognizes_histogram_by_type() {
        let schema = Arc::new(Schema::new(vec![histogram_field("custom_histogram", 5)]));

        let wrapped = maybe_wrap_schema(&schema).unwrap();
        let hist = wrapped.field_with_name("custom_histogram").unwrap();
        assert_histogram_stamped(hist, 5);
    }

    #[test]
    fn test_maybe_wrap_schema_plain_struct_not_stamped() {
        use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY;

        let plain = ArrowDataType::Struct(
            vec![
                Arc::new(Field::new("a", ArrowDataType::Int32, true)),
                Arc::new(Field::new("b", ArrowDataType::Utf8, true)),
            ]
            .into(),
        );
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "ts",
                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            Field::new("data", plain, true),
        ]));

        let wrapped = maybe_wrap_schema(&schema).unwrap();
        let data = wrapped.field_with_name("data").unwrap();
        assert!(
            data.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(),
            "non-histogram struct must not get the extension"
        );
        if let ArrowDataType::Struct(children) = data.data_type() {
            for child in children {
                assert!(
                    child.metadata().get(PARQUET_FIELD_ID_KEY).is_none(),
                    "non-histogram sub-field {} must not get a field id",
                    child.name()
                );
            }
        }
    }

    #[test]
    fn test_maybe_wrap_schema_no_struct_unchanged() {
        let schema: Arc<Schema> = Arc::new(Schema::new(vec![
            Field::new(
                "ts",
                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            Field::new("v", ArrowDataType::Float64, true),
        ]));
        let wrapped = maybe_wrap_schema(&schema).unwrap();
        assert!(
            Arc::ptr_eq(&wrapped, &schema),
            "a schema without any struct column must be returned unchanged"
        );
    }

    /// Writes `schema` through `maybe_wrap_schema` and a real parquet
    /// [`ArrowWriter`], then returns the arrow schema read back from the file
    /// footer. This proves the `greptime.histogram` extension and the nested
    /// `PARQUET:field_id`s actually land on disk, not just in memory.
    ///
    /// `maybe_wrap_schema` is exactly what the SST parquet writer hands to
    /// `AsyncArrowWriter` (see `writer.rs`); the sync [`ArrowWriter`] shares
    /// the same arrow-to-parquet schema conversion, so the footer it emits is
    /// the on-disk contract this change introduces. An empty batch suffices
    /// because the parquet footer always carries the schema.
    fn parquet_footer_arrow_schema(schema: &SchemaRef) -> SchemaRef {
        use ::parquet::arrow::ArrowWriter;
        use ::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
        use ::parquet::file::properties::WriterProperties;
        use bytes::Bytes;

        let wrapped = maybe_wrap_schema(schema).unwrap();
        let mut bytes = Vec::new();
        let props = WriterProperties::builder().build();
        let mut writer = ArrowWriter::try_new(&mut bytes, wrapped.clone(), Some(props)).unwrap();
        writer
            .write(&RecordBatch::new_empty(wrapped.clone()))
            .unwrap();
        writer.close().unwrap();

        ParquetRecordBatchReaderBuilder::try_new(Bytes::from(bytes))
            .unwrap()
            .schema()
            .clone()
    }

    #[test]
    fn test_maybe_wrap_schema_survives_parquet_roundtrip() {
        // On-disk contract: after writing through the parquet writer path, the
        // footer still carries the greptime.histogram extension and every
        // nested (sub-field + list-element) PARQUET:field_id.
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "greptime_timestamp",
                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            histogram_field(greptime_native_histogram(), 3),
        ]));

        let on_disk = parquet_footer_arrow_schema(&schema);
        let hist = on_disk
            .field_with_name(greptime_native_histogram())
            .expect("histogram field present");
        assert_histogram_stamped(hist, 3);
    }

    #[test]
    fn test_parquet_roundtrip_recognizes_histogram_by_type() {
        // The persisted type, rather than a process-local configured name,
        // identifies native histograms across upgrades and prefix changes.
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "ts",
                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            histogram_field("custom_histogram", 9),
        ]));

        let on_disk = parquet_footer_arrow_schema(&schema);
        let hist = on_disk.field_with_name("custom_histogram").unwrap();
        assert_histogram_stamped(hist, 9);
    }

    #[test]
    fn test_maybe_wrap_schema_overflows_return_error() {
        // A column id of 12_582_912 makes the derived sub-field id overflow
        // i32 (BASE + column_id*64 == i32::MAX + 1). The write path must
        // surface this as an error rather than silently dropping the field
        // id, wrapping, or panicking.
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "greptime_timestamp",
                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            histogram_field(greptime_native_histogram(), 12_582_912),
        ]));
        let err = maybe_wrap_schema(&schema).unwrap_err();
        assert!(
            matches!(
                err,
                crate::error::Error::InvalidNativeHistogramSubfield { .. }
            ),
            "expected InvalidNativeHistogramSubfield, got {:?}",
            err
        );
    }

    #[test]
    fn test_maybe_wrap_schema_missing_field_id_returns_error() {
        // The parent column's PARQUET:field_id namespaces every sub-field id.
        // If it is absent (e.g. a histogram struct handed to the writer
        // without the write path's stamping), the writer must fail loudly
        // rather than silently namespace under column 0, which would collide
        // with that column's nested ids.
        use common_query::native_histogram::native_histogram_value_type;
        use datatypes::data_type::DataType;

        let field = Field::new(
            greptime_native_histogram(),
            native_histogram_value_type().as_arrow_type(),
            true,
        );
        assert!(
            field.metadata().get(PARQUET_FIELD_ID_KEY).is_none(),
            "fixture must not carry a field id"
        );
        let schema = Arc::new(Schema::new(vec![field]));
        let err = maybe_wrap_schema(&schema).unwrap_err();
        assert!(
            matches!(
                err,
                crate::error::Error::InvalidNativeHistogramFieldId { .. }
            ),
            "expected InvalidNativeHistogramFieldId, got {:?}",
            err
        );
    }

    #[test]
    fn test_maybe_wrap_schema_field_id_above_i32_max_returns_error() {
        // `with_field_id` serializes the column id from a u32, so a valid id
        // above i32::MAX (e.g. u32::MAX) must not be silently parsed as a
        // failed i32 and collapsed onto column 0's nested ids. It must
        // surface a checked-conversion error instead.
        let schema = Arc::new(Schema::new(vec![
            Field::new(
                "greptime_timestamp",
                ArrowDataType::Timestamp(TimeUnit::Millisecond, None),
                false,
            ),
            histogram_field(greptime_native_histogram(), u32::MAX),
        ]));
        let err = maybe_wrap_schema(&schema).unwrap_err();
        assert!(
            matches!(
                err,
                crate::error::Error::InvalidNativeHistogramFieldId { .. }
            ),
            "expected InvalidNativeHistogramFieldId, got {:?}",
            err
        );
    }

    fn json2_v2_test_type() -> ArrowDataType {
        ArrowDataType::Struct(
            vec![
                Arc::new(Field::new("active", ArrowDataType::Boolean, true)),
                Arc::new(Field::new("hot", ArrowDataType::Int64, true)),
                Arc::new(Field::new("name", ArrowDataType::Utf8, true)),
            ]
            .into(),
        )
    }

    /// Validates the persisted-format foundation for the JSON2 v2 remainder.
    ///
    /// JSON2 will store `!__remainder__!` as a nested Variant child of its root
    /// Struct. Before enabling that layout in production, this test ensures the
    /// SST schema wrapper and Arrow writer preserve the Variant extension,
    /// encode the Parquet Variant logical type, and round-trip the values
    /// without changing the surrounding Struct.
    #[tokio::test]
    async fn test_nested_variant_survives_sst_writer_schema_roundtrip()
    -> Result<(), Box<dyn std::error::Error>> {
        let json: ArrayRef = Arc::new(StringArray::from(vec![
            Some(r#"{}"#),
            Some(r#"{"name":"Alice","active":true}"#),
            Some(r#"{"nested":{"count":42},"items":[1,"two",null]}"#),
            Some(r#"{"\u5b57\u6bb5":"\u503c"}"#),
            None,
        ]));
        let remainder = json_to_variant(&json)?;
        let remainder_field = remainder.field("!__remainder__!");
        let remainder_array = ArrayRef::from(remainder);
        let hot_field = Field::new("hot", ArrowDataType::Int64, true);
        let data_array = Arc::new(StructArray::new(
            vec![remainder_field.clone(), hot_field.clone()].into(),
            vec![
                remainder_array,
                Arc::new(Int64Array::from(vec![
                    Some(1),
                    Some(2),
                    Some(3),
                    Some(4),
                    None,
                ])),
            ],
            None,
        ));
        let data_field = Field::new(
            "data",
            ArrowDataType::Struct(vec![remainder_field, hot_field].into()),
            true,
        )
        .with_extension_type(Json2ExtensionType::default());
        let schema = Arc::new(Schema::new(vec![data_field]));
        let source = RecordBatch::try_new(schema.clone(), vec![data_array])?;

        let wrapped = maybe_wrap_schema(&schema)?;
        let mut buffer = Vec::new();
        let mut writer = AsyncArrowWriter::try_new(&mut buffer, wrapped, None)?;
        writer.write(&source).await?;
        writer.close().await?;

        let builder = ParquetRecordBatchReaderBuilder::try_new(bytes::Bytes::from(buffer))?;
        let parquet_remainder =
            &builder.parquet_schema().root_schema().get_fields()[0].get_fields()[0];
        assert_eq!(
            parquet_remainder.get_basic_info().logical_type_ref(),
            Some(&LogicalType::variant(None))
        );

        let ArrowDataType::Struct(children) = builder.schema().field_with_name("data")?.data_type()
        else {
            unreachable!();
        };
        assert!(children[0].has_valid_extension_type::<VariantType>());

        let mut reader = builder.build()?;
        let result = reader.next().unwrap()?;
        assert_eq!(source, result);
        let result_field = result.schema().field(0).clone();
        let result = result
            .column(0)
            .as_any()
            .downcast_ref::<StructArray>()
            .unwrap();
        VariantArray::try_new(result.column(0))?;
        let result: ArrayRef = Arc::new(result.clone());
        let result =
            JsonArray::from(&result).project_to_v2(&result_field, &json2_v2_test_type())?;
        assert_eq!(
            json!({"active": true, "hot": 2, "name": "Alice"}),
            JsonArray::from(&result).try_get_value(1)?
        );
        Ok(())
    }

    /// Ensures future readers retain compatibility with the first JSON2 v2 layout.
    #[test]
    fn test_read_json2_v2_fixture() -> Result<(), Box<dyn std::error::Error>> {
        let bytes = bytes::Bytes::from_static(include_bytes!("../test-data/json2-v2.parquet"));
        let builder = ParquetRecordBatchReaderBuilder::try_new(bytes)?;
        let field = builder.schema().field(0).clone();
        assert!(Json2PhysicalLayout::try_from_root(&field)?.is_version_2());

        let batch = builder.build()?.next().unwrap()?;
        let data = batch
            .column(0)
            .as_any()
            .downcast_ref::<StructArray>()
            .unwrap();
        VariantArray::try_new(data.column(0))?;
        let data: ArrayRef = Arc::new(data.clone());
        let data = JsonArray::from(&data).project_to_v2(&field, &json2_v2_test_type())?;
        assert_eq!(
            json!({"active": true, "hot": 2, "name": "Alice"}),
            JsonArray::from(&data).try_get_value(1)?
        );
        Ok(())
    }
}