diff --git a/src/common/query/src/native_histogram.rs b/src/common/query/src/native_histogram.rs index b10c78104a5..ca77326ed97 100644 --- a/src/common/query/src/native_histogram.rs +++ b/src/common/query/src/native_histogram.rs @@ -109,3 +109,110 @@ pub fn native_histogram_value_type() -> &'static ConcreteDataType { pub fn is_native_histogram_value_schema(name: &str, data_type: &ConcreteDataType) -> bool { name == NATIVE_HISTOGRAM_FIELD && data_type == native_histogram_value_type() } + +// --------------------------------------------------------------------------- +// Stable Parquet field ids for native-histogram sub-fields. +// +// External readers resolve nested struct fields by `PARQUET:field_id`, so each +// sub-field (and list element) needs a stable positive id. The struct schema +// is fixed (always the same 18 fields). Each histogram column owns a block of +// ids (offset from a reserved base by the column's id), so several histogram +// columns in one table get disjoint sub-field ids. The reserved base is +// disjoint from user column ids and mito2 internal ids (`1 << 30`). +// +// The id is computed with checked arithmetic: the reserved base plus +// `column_id * stride` cannot always fit in a positive `i32` (a `ColumnId` is +// `u32`), so derivation returns `None` once the representable range is +// exceeded. Callers must handle `None` explicitly — the SST parquet writer +// surfaces it as an error rather than wrapping, panicking, or silently +// dropping the field id. +// --------------------------------------------------------------------------- + +/// Reserved base for native-histogram struct sub-field ids. +pub const NATIVE_HISTOGRAM_SUBFIELD_ID_BASE: i32 = 0x5000_0000; + +/// Number of ids reserved per histogram column (18 sub-fields + headroom for +/// list element ids), so multiple histogram columns get disjoint ids. +pub const NATIVE_HISTOGRAM_SUBFIELD_ID_STRIDE: i32 = 64; + +/// Offset of list element ids within a column's id block. +pub const NATIVE_HISTOGRAM_LIST_ELEMENT_OFFSET: i32 = 32; + +/// Returns the stable field id for a native-histogram struct sub-field, +/// namespaced by its parent `column_id`, or `None` if `name` is not a known +/// sub-field or the derived id overflows a positive `i32`. +pub fn native_histogram_subfield_id(column_id: i32, name: &str) -> Option { + let idx = subfield_index(name)?; + NATIVE_HISTOGRAM_SUBFIELD_ID_BASE + .checked_add(column_id.checked_mul(NATIVE_HISTOGRAM_SUBFIELD_ID_STRIDE)?) + .and_then(|v| v.checked_add(idx)) +} + +/// Returns the stable list `element-id` for a list-typed native-histogram +/// sub-field, namespaced by its parent `column_id`, or `None` if `name` is not +/// a known sub-field or the derived id overflows a positive `i32`. +pub fn native_histogram_list_element_id(column_id: i32, name: &str) -> Option { + let idx = subfield_index(name)?; + NATIVE_HISTOGRAM_SUBFIELD_ID_BASE + .checked_add(column_id.checked_mul(NATIVE_HISTOGRAM_SUBFIELD_ID_STRIDE)?) + .and_then(|v| v.checked_add(NATIVE_HISTOGRAM_LIST_ELEMENT_OFFSET)) + .and_then(|v| v.checked_add(idx)) +} + +fn subfield_index(name: &str) -> Option { + NATIVE_HISTOGRAM_FIELD_NAMES + .iter() + .position(|n| *n == name) + .map(|i| i as i32) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn subfield_ids_are_namespaced_and_disjoint() { + // SCHEMA=0, ZERO_THRESHOLD=1, SUM=2, ..., CUSTOM_VALUES=5. + let sum_idx = 2; + let custom_values_idx = 5; + + // Ids are offset from the reserved base by the parent column id and + // the sub-field index. + assert_eq!( + native_histogram_subfield_id(1, SUM_FIELD), + Some(NATIVE_HISTOGRAM_SUBFIELD_ID_BASE + NATIVE_HISTOGRAM_SUBFIELD_ID_STRIDE + sum_idx,) + ); + // List element ids additionally carry the list offset. + assert_eq!( + native_histogram_list_element_id(1, CUSTOM_VALUES_FIELD), + Some( + NATIVE_HISTOGRAM_SUBFIELD_ID_BASE + + NATIVE_HISTOGRAM_SUBFIELD_ID_STRIDE + + NATIVE_HISTOGRAM_LIST_ELEMENT_OFFSET + + custom_values_idx, + ) + ); + // Different parent columns get disjoint ids for the same sub-field. + assert_ne!( + native_histogram_subfield_id(1, SUM_FIELD), + native_histogram_subfield_id(7, SUM_FIELD) + ); + // Unknown sub-field name -> None. + assert_eq!(native_histogram_subfield_id(1, "not_a_field"), None); + } + + #[test] + fn subfield_ids_overflow_returns_none() { + // `column_id` is u32-sized, but the derived id must fit in a positive + // i32. At column_id = 12_582_912, BASE + column_id*64 == i32::MAX + 1, + // which previously overflowed (debug panic / release wrap). Checked + // arithmetic must yield None instead of wrapping or panicking. + assert_eq!(native_histogram_subfield_id(12_582_912, SUM_FIELD), None); + assert_eq!( + native_histogram_list_element_id(12_582_912, CUSTOM_VALUES_FIELD), + None + ); + // One below that boundary is still representable. + assert!(native_histogram_subfield_id(12_582_911, SUM_FIELD).is_some()); + } +} diff --git a/src/datatypes/src/extension.rs b/src/datatypes/src/extension.rs index 83776cdcc17..a54478b07e1 100644 --- a/src/datatypes/src/extension.rs +++ b/src/datatypes/src/extension.rs @@ -12,4 +12,5 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub mod histogram; pub mod json; diff --git a/src/datatypes/src/extension/histogram.rs b/src/datatypes/src/extension/histogram.rs new file mode 100644 index 00000000000..79ad5b3350e --- /dev/null +++ b/src/datatypes/src/extension/histogram.rs @@ -0,0 +1,98 @@ +// 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. + +//! Arrow extension type for native-histogram struct columns. + +use arrow_schema::extension::ExtensionType; +use arrow_schema::{ArrowError, DataType, FieldRef}; + +/// Arrow extension type identifying a native-histogram struct column. +/// +/// Applied to the struct field at parquet-write time so that readers can +/// identify native-histogram columns by extension (`greptime.histogram`) rather +/// than relying on the field name. +#[derive(Debug, Clone, Default)] +pub struct HistogramExtensionType; + +impl ExtensionType for HistogramExtensionType { + const NAME: &'static str = "greptime.histogram"; + type Metadata = (); + + fn metadata(&self) -> &Self::Metadata { + &() + } + + fn serialize_metadata(&self) -> Option { + None + } + + fn deserialize_metadata(_metadata: Option<&str>) -> Result { + Ok(()) + } + + fn supports_data_type(&self, data_type: &DataType) -> Result<(), ArrowError> { + match data_type { + DataType::Struct(_) => Ok(()), + dt => Err(ArrowError::SchemaError(format!( + "Unexpected data type {dt}" + ))), + } + } + + fn try_new(data_type: &DataType, _metadata: Self::Metadata) -> Result { + let ext = Self; + ext.supports_data_type(data_type)?; + Ok(ext) + } +} + +/// Check if this field is a native-histogram extension type. +pub fn is_histogram_extension_type(field: &FieldRef) -> bool { + field.extension_type_name() == Some(HistogramExtensionType::NAME) +} + +#[cfg(test)] +mod tests { + use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; + use arrow_schema::{DataType, Field, Fields}; + + use super::*; + + #[test] + fn test_extension_name_and_detection() { + assert_eq!(HistogramExtensionType::NAME, "greptime.histogram"); + + // A plain struct field is not a histogram extension type. + let empty: Fields = Vec::::new().into(); + let plain = std::sync::Arc::new(Field::new("s", DataType::Struct(empty), true)); + assert!(!is_histogram_extension_type(&plain)); + + // Tagging the field with the extension makes it detectable. + let mut tagged = (*plain).clone(); + tagged.metadata_mut().insert( + EXTENSION_TYPE_NAME_KEY.to_string(), + HistogramExtensionType::NAME.to_string(), + ); + let tagged = std::sync::Arc::new(tagged); + assert!(is_histogram_extension_type(&tagged)); + } + + #[test] + fn test_supports_struct_only() { + let ext = HistogramExtensionType; + let empty: Fields = Vec::::new().into(); + assert!(ext.supports_data_type(&DataType::Struct(empty)).is_ok()); + assert!(ext.supports_data_type(&DataType::Int32).is_err()); + } +} diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index 219f3761952..5e93f011fa2 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -192,6 +192,28 @@ pub enum Error { location: Location, }, + #[snafu(display( + "Cannot assign a stable field id to native histogram sub-field '{}' of column id {} (unknown sub-field name or derived id overflows i32)", + field_name, + column_id + ))] + InvalidNativeHistogramSubfield { + column_id: i32, + field_name: String, + #[snafu(implicit)] + location: Location, + }, + + #[snafu(display( + "Native histogram column '{}' has no usable PARQUET:field_id to namespace its sub-field ids (missing, malformed, or exceeds i32::MAX)", + field_name + ))] + InvalidNativeHistogramFieldId { + field_name: String, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Region {} not found", region_id))] RegionNotFound { region_id: RegionId, @@ -1437,7 +1459,9 @@ impl ErrorExt for Error { | PuffinPurgeStager { source, .. } => source.status_code(), CleanDir { .. } => StatusCode::Unexpected, InvalidConfig { .. } => StatusCode::InvalidArguments, - StaleLogEntry { .. } => StatusCode::Unexpected, + StaleLogEntry { .. } + | InvalidNativeHistogramSubfield { .. } + | InvalidNativeHistogramFieldId { .. } => StatusCode::Unexpected, External { source, .. } => source.status_code(), diff --git a/src/mito2/src/sst.rs b/src/mito2/src/sst.rs index 20c9383f1f5..170f25a2b45 100644 --- a/src/mito2/src/sst.rs +++ b/src/mito2/src/sst.rs @@ -19,11 +19,17 @@ 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_schema, 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}; @@ -33,6 +39,7 @@ 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; @@ -64,7 +71,10 @@ pub enum FormatType { /// 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 an Arrow field. +/// 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() @@ -72,6 +82,125 @@ pub fn with_field_id(mut field: Field, column_id: u32) -> Field { 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 native-histogram column name and struct type +/// (`is_native_histogram_value_schema`); 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_schema( + field.name(), + &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::().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 = 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 { + // 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> = 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; @@ -557,4 +686,368 @@ mod tests { 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() { + use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; + + let schema = Arc::new(Schema::new(vec![ + Field::new( + "greptime_timestamp", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + histogram_field(NATIVE_HISTOGRAM_FIELD, 1), + ])); + + let wrapped = maybe_wrap_schema(&schema).unwrap(); + let hist = wrapped + .field_with_name(NATIVE_HISTOGRAM_FIELD) + .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_FIELD, native_histogram_subfield_id, + }; + + let schema = Arc::new(Schema::new(vec![ + histogram_field(NATIVE_HISTOGRAM_FIELD, 1), + histogram_field(NATIVE_HISTOGRAM_FIELD, 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_requires_canonical_name() { + // Detection requires the canonical column name: a histogram-typed field + // named differently is left untouched. + use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; + use common_query::native_histogram::native_histogram_value_type; + use datatypes::data_type::DataType; + + let hist_arrow = native_histogram_value_type().as_arrow_type(); + let schema = Arc::new(Schema::new(vec![Field::new( + "custom_histogram", + hist_arrow, + true, + )])); + + let wrapped = maybe_wrap_schema(&schema).unwrap(); + let hist = wrapped.field_with_name("custom_histogram").unwrap(); + assert!( + hist.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(), + "a histogram-typed field without the canonical name must not be stamped" + ); + } + + #[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 = 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. + use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; + + let schema = Arc::new(Schema::new(vec![ + Field::new( + "greptime_timestamp", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + histogram_field(NATIVE_HISTOGRAM_FIELD, 3), + ])); + + let on_disk = parquet_footer_arrow_schema(&schema); + let hist = on_disk + .field_with_name(NATIVE_HISTOGRAM_FIELD) + .expect("histogram field present"); + assert_histogram_stamped(hist, 3); + } + + #[test] + fn test_parquet_roundtrip_noncanonical_struct_untouched() { + // A histogram-shaped struct without the canonical column name is left + // untouched on disk: no extension, no nested field ids. + use arrow_schema::extension::EXTENSION_TYPE_NAME_KEY; + use common_query::native_histogram::native_histogram_value_type; + use datatypes::data_type::DataType; + + let hist_arrow = native_histogram_value_type().as_arrow_type(); + let schema = Arc::new(Schema::new(vec![ + Field::new( + "ts", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + Field::new("custom_histogram", hist_arrow, true), + ])); + + let on_disk = parquet_footer_arrow_schema(&schema); + let hist = on_disk.field_with_name("custom_histogram").unwrap(); + assert!( + hist.metadata().get(EXTENSION_TYPE_NAME_KEY).is_none(), + "a histogram-typed field without the canonical name must not be stamped on disk" + ); + if let ArrowDataType::Struct(children) = hist.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 on disk", + child.name() + ); + } + } else { + panic!("expected a struct, got {:?}", hist.data_type()); + } + } + + #[test] + fn test_maybe_wrap_schema_overflows_return_error() { + use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; + + // 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(NATIVE_HISTOGRAM_FIELD, 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_FIELD, native_histogram_value_type}; + use datatypes::data_type::DataType; + + let field = Field::new( + NATIVE_HISTOGRAM_FIELD, + 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. + use common_query::native_histogram::NATIVE_HISTOGRAM_FIELD; + + let schema = Arc::new(Schema::new(vec![ + Field::new( + "greptime_timestamp", + ArrowDataType::Timestamp(TimeUnit::Millisecond, None), + false, + ), + histogram_field(NATIVE_HISTOGRAM_FIELD, u32::MAX), + ])); + let err = maybe_wrap_schema(&schema).unwrap_err(); + assert!( + matches!( + err, + crate::error::Error::InvalidNativeHistogramFieldId { .. } + ), + "expected InvalidNativeHistogramFieldId, got {:?}", + err + ); + } } diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index e97965d8e70..5a6f894730d 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -60,6 +60,7 @@ use crate::sst::parquet::format::PrimaryKeyWriteFormat; use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, WriteOptions}; use crate::sst::{ DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, + maybe_wrap_schema, }; /// Converts a flat RecordBatch for writing to parquet. @@ -457,7 +458,7 @@ where self.bytes_written.clone(), ); let arrow_writer = - AsyncArrowWriter::try_new(writer, schema.clone(), Some(writer_props)) + AsyncArrowWriter::try_new(writer, maybe_wrap_schema(schema)?, Some(writer_props)) .context(WriteParquetSnafu)?; self.writer = Some(arrow_writer);