diff --git a/src/datatypes/src/arrow_array.rs b/src/datatypes/src/arrow_array.rs index e5ef7add82..848cc7c91f 100644 --- a/src/datatypes/src/arrow_array.rs +++ b/src/datatypes/src/arrow_array.rs @@ -168,6 +168,43 @@ pub fn string_array_value_at_index(array: &ArrayRef, i: usize) -> Option<&str> { } } +/// Check whether the string value at index `i` is null for string or +/// dictionary-encoded string arrays. +/// +/// Returns `true` when the value is null or the array type is not a string +/// type, which corresponds to [`string_array_value_at_index`] returning `None`. +/// +/// # Panics +/// +/// If index `i` is out of bounds. +pub fn is_string_null_at(array: &ArrayRef, i: usize) -> bool { + match array.data_type() { + DataType::Utf8 => { + let array = array.as_string::(); + !array.is_valid(i) + } + DataType::LargeUtf8 => { + let array = array.as_string::(); + !array.is_valid(i) + } + DataType::Utf8View => { + let array = array.as_string_view(); + !array.is_valid(i) + } + DataType::Dictionary(key_type, value_type) + if key_type.is_integer() && value_type.is_string() => + { + downcast_dictionary_array! { + array => array + .key(i) + .is_none_or(|key| is_string_null_at(array.values(), key)), + _ => true, + } + } + _ => true, + } +} + /// Get the string value at index `i` for `Utf8`, `LargeUtf8`, or `Utf8View` arrays. /// /// Note: This method does not check for nulls and the value is arbitrary diff --git a/src/metric-engine/src/batch_modifier.rs b/src/metric-engine/src/batch_modifier.rs index 2162215d1b..575dffff1a 100644 --- a/src/metric-engine/src/batch_modifier.rs +++ b/src/metric-engine/src/batch_modifier.rs @@ -15,9 +15,10 @@ use std::hash::Hasher; use std::sync::Arc; -use datatypes::arrow::array::{Array, BinaryBuilder, StringArray, UInt64Array}; +use datatypes::arrow::array::{Array, ArrayRef, BinaryBuilder, UInt64Array}; use datatypes::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; use datatypes::arrow::record_batch::RecordBatch; +use datatypes::arrow_array::{is_string_null_at, string_array_value_at_index}; use fxhash::FxHasher; use mito_codec::row_converter::SparsePrimaryKeyCodec; use snafu::ResultExt; @@ -49,7 +50,7 @@ pub struct TagColumnInfo { pub fn compute_tsid_array( batch: &RecordBatch, sorted_tag_columns: &[TagColumnInfo], - tag_arrays: &[&StringArray], + tag_arrays: &[&ArrayRef], ) -> UInt64Array { let num_rows = batch.num_rows(); @@ -64,20 +65,22 @@ pub fn compute_tsid_array( let mut tsid_values = Vec::with_capacity(num_rows); for row in 0..num_rows { - let has_null = tag_arrays.iter().any(|arr| arr.is_null(row)); + let has_null = tag_arrays.iter().any(|arr| is_string_null_at(arr, row)); let tsid = if !has_null { let mut hasher = FxHasher::default(); hasher.write_u64(label_name_hash); for arr in tag_arrays { - hasher.write(arr.value(row).as_bytes()); - hasher.write_u8(0xff); + if let Some(value) = string_array_value_at_index(arr, row) { + hasher.write(value.as_bytes()); + hasher.write_u8(0xff); + } } hasher.finish() } else { let mut name_hasher = FxHasher::default(); for (tc, arr) in sorted_tag_columns.iter().zip(tag_arrays.iter()) { - if !arr.is_null(row) { + if !is_string_null_at(arr, row) { name_hasher.write(tc.name.as_bytes()); name_hasher.write_u8(0xff); } @@ -87,8 +90,8 @@ pub fn compute_tsid_array( let mut val_hasher = FxHasher::default(); val_hasher.write_u64(row_label_hash); for arr in tag_arrays { - if !arr.is_null(row) { - val_hasher.write(arr.value(row).as_bytes()); + if let Some(value) = string_array_value_at_index(arr, row) { + val_hasher.write(value.as_bytes()); val_hasher.write_u8(0xff); } } @@ -104,15 +107,26 @@ pub fn compute_tsid_array( fn build_tag_arrays<'a>( batch: &'a RecordBatch, sorted_tag_columns: &[TagColumnInfo], -) -> Vec<&'a StringArray> { +) -> Result> { sorted_tag_columns .iter() .map(|tc| { - batch - .column(tc.index) - .as_any() - .downcast_ref::() - .expect("tag column must be utf8") + let array = batch.column(tc.index); + match array.data_type() { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => Ok(array), + DataType::Dictionary(key_type, value_type) + if key_type.is_integer() && value_type.is_string() => + { + Ok(array) + } + data_type => UnexpectedRequestSnafu { + reason: format!( + "Tag column '{}' must be a string, given: {data_type}", + tc.name + ), + } + .fail(), + } }) .collect() } @@ -144,7 +158,7 @@ pub fn modify_batch_sparse( ) -> Result { let num_rows = batch.num_rows(); let codec = SparsePrimaryKeyCodec::schemaless(); - let tag_arrays: Vec<&StringArray> = build_tag_arrays(&batch, sorted_tag_columns); + let tag_arrays = build_tag_arrays(&batch, sorted_tag_columns)?; let tsid_array = compute_tsid_array(&batch, sorted_tag_columns, &tag_arrays); let mut pk_builder = BinaryBuilder::with_capacity(num_rows, 0); @@ -158,8 +172,9 @@ pub fn modify_batch_sparse( let tags = sorted_tag_columns .iter() .zip(tag_arrays.iter()) - .filter(|(_, arr)| !arr.is_null(row)) - .map(|(tc, arr)| (tc.column_id, arr.value(row).as_bytes())); + .filter_map(|(tc, arr)| { + string_array_value_at_index(arr, row).map(|value| (tc.column_id, value.as_bytes())) + }); codec .encode_raw_tag_value(tags, &mut buffer) .context(EncodePrimaryKeySnafu)?; @@ -268,7 +283,7 @@ mod tests { column_id: 1, }, ]; - let tag_arrays = build_tag_arrays(&batch, &tag_columns); + let tag_arrays = build_tag_arrays(&batch, &tag_columns).unwrap(); let tsid_array = compute_tsid_array(&batch, &tag_columns, &tag_arrays); assert_eq!(tsid_array.value(0), 2721566936019240841); @@ -300,7 +315,7 @@ mod tests { column_id: 2, }, ]; - let tag_arrays_2 = build_tag_arrays(&batch_no_null, &tag_cols_2); + let tag_arrays_2 = build_tag_arrays(&batch_no_null, &tag_cols_2).unwrap(); let tsid_no_null = compute_tsid_array(&batch_no_null, &tag_cols_2, &tag_arrays_2); let schema3 = Arc::new(ArrowSchema::new(vec![ @@ -334,7 +349,7 @@ mod tests { column_id: 3, }, ]; - let tag_arrays_3 = build_tag_arrays(&batch_with_null, &tag_cols_3); + let tag_arrays_3 = build_tag_arrays(&batch_with_null, &tag_cols_3).unwrap(); let tsid_with_null = compute_tsid_array(&batch_with_null, &tag_cols_3, &tag_arrays_3); assert_eq!(tsid_no_null.value(0), tsid_with_null.value(0)); @@ -452,4 +467,156 @@ mod tests { .unwrap(); assert_eq!(actual_array.value(0), expected_pk.as_slice()); } + + #[test] + fn label_replace_with_utf8view_labels_does_not_panic() { + // Tag (label) columns may arrive as any string representation (`Utf8`, + // `LargeUtf8`, `Utf8View`, or dictionary-encoded). `modify_batch_sparse` + // must not assume the tag columns are plain `StringArray`s. + let tag_arrays: Vec = vec![ + Arc::new(StringArray::from(vec!["greptimedb"])), + Arc::new(datatypes::arrow::array::LargeStringArray::from(vec![ + "greptimedb", + ])), + Arc::new(datatypes::arrow::array::StringViewArray::from(vec![ + "greptimedb", + ])), + Arc::new(datatypes::arrow::array::DictionaryArray::< + datatypes::arrow::datatypes::UInt32Type, + >::new( + datatypes::arrow::array::UInt32Array::from(vec![0]), + Arc::new(StringArray::from(vec!["greptimedb"])), + )), + ]; + let tag_columns = vec![TagColumnInfo { + name: "namespace".to_string(), + index: 2, + column_id: 2, + }]; + let non_tag_indices = vec![0, 1]; + + let primary_keys = tag_arrays + .into_iter() + .map(|tag_array| { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("greptime_timestamp", DataType::Int64, false), + Field::new("greptime_value", DataType::Float64, true), + Field::new("namespace", tag_array.data_type().clone(), true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1000])), + Arc::new(datatypes::arrow::array::Float64Array::from(vec![42.0])), + tag_array, + ], + ) + .unwrap(); + let modified = + modify_batch_sparse(batch, 1025, &tag_columns, &non_tag_indices).unwrap(); + modified + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + .to_vec() + }) + .collect::>(); + + assert!(primary_keys.windows(2).all(|keys| keys[0] == keys[1])); + } + + #[test] + fn test_compute_tsid_utf8view_with_nulls_matches_utf8() { + // Regression test for Utf8View tag/label columns: the whole + // `build_tag_arrays` -> `compute_tsid_array` -> sparse primary key path must handle + // `Utf8View` columns exactly like `Utf8`, including rows with null tags, instead of + // panicking or producing different TSIDs. + let namespace_values: Vec> = + vec![Some("ns-a"), Some("ns-b"), None, Some("ns-c")]; + let host_values: Vec> = vec![Some("host-1"), None, Some("host-2"), None]; + + let build_batch = |data_type: DataType| { + let (namespace, host): (ArrayRef, ArrayRef) = match &data_type { + DataType::Utf8 => ( + Arc::new(StringArray::from(namespace_values.clone())), + Arc::new(StringArray::from(host_values.clone())), + ), + DataType::Utf8View => ( + Arc::new(datatypes::arrow::array::StringViewArray::from( + namespace_values.clone(), + )), + Arc::new(datatypes::arrow::array::StringViewArray::from( + host_values.clone(), + )), + ), + _ => unreachable!(), + }; + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("greptime_timestamp", DataType::Int64, false), + Field::new("greptime_value", DataType::Float64, true), + Field::new("namespace", data_type.clone(), true), + Field::new("host", data_type, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int64Array::from(vec![1000, 1001, 1002, 1003])), + Arc::new(datatypes::arrow::array::Float64Array::from(vec![ + 1.0, 2.0, 3.0, 4.0, + ])), + namespace, + host, + ], + ) + .unwrap() + }; + + let tag_columns = vec![ + TagColumnInfo { + name: "host".to_string(), + index: 3, + column_id: 3, + }, + TagColumnInfo { + name: "namespace".to_string(), + index: 2, + column_id: 2, + }, + ]; + + let utf8_batch = build_batch(DataType::Utf8); + let utf8view_batch = build_batch(DataType::Utf8View); + + // Both representations must be accepted by `build_tag_arrays`. + let utf8_tag_arrays = build_tag_arrays(&utf8_batch, &tag_columns).unwrap(); + let utf8view_tag_arrays = build_tag_arrays(&utf8view_batch, &tag_columns).unwrap(); + + // TSIDs must be identical across representations, including rows with null tags. + let utf8_tsids = compute_tsid_array(&utf8_batch, &tag_columns, &utf8_tag_arrays); + let utf8view_tsids = + compute_tsid_array(&utf8view_batch, &tag_columns, &utf8view_tag_arrays); + assert_eq!(utf8_tsids, utf8view_tsids); + + // The full sparse primary key (TSID + tag values) must also be identical. + let non_tag_indices = vec![0, 1]; + let modified_utf8 = + modify_batch_sparse(utf8_batch, 1025, &tag_columns, &non_tag_indices).unwrap(); + let modified_utf8view = + modify_batch_sparse(utf8view_batch, 1025, &tag_columns, &non_tag_indices).unwrap(); + let utf8_pks = modified_utf8 + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let utf8view_pks = modified_utf8view + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for row in 0..utf8_pks.len() { + assert_eq!(utf8_pks.value(row), utf8view_pks.value(row)); + } + } } diff --git a/src/mito2/src/memtable/time_series.rs b/src/mito2/src/memtable/time_series.rs index f5b772b076..7e89d569ee 100644 --- a/src/mito2/src/memtable/time_series.rs +++ b/src/mito2/src/memtable/time_series.rs @@ -32,8 +32,8 @@ use datatypes::prelude::{ScalarVector, Vector, VectorRef}; use datatypes::types::TimestampType; use datatypes::value::{Value, ValueRef}; use datatypes::vectors::{ - Helper, TimestampMicrosecondVector, TimestampMillisecondVector, TimestampNanosecondVector, - TimestampSecondVector, UInt8Vector, UInt64Vector, + Helper, StringVector, TimestampMicrosecondVector, TimestampMillisecondVector, + TimestampNanosecondVector, TimestampSecondVector, UInt8Vector, UInt64Vector, }; use mito_codec::key_values::KeyValue; use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt}; @@ -72,6 +72,93 @@ const INITIAL_BUILDER_CAPACITY: usize = 4; /// Vector builder capacity. const BUILDER_CAPACITY: usize = 512; +fn checked_string_values_size(mut lengths: impl Iterator) -> Option { + lengths.try_fold(0_i32, |size, len| { + size.checked_add(i32::try_from(len).ok()?) + }) +} + +/// Checks whether a string field builder that currently holds `current_len` bytes can +/// additionally accommodate `space_needed` bytes of string data, given that the offset of a +/// single Arrow string array is limited to `limit` (`i32::MAX` in production). +/// +/// Returns: +/// - `Ok(true)` if `current_len + space_needed <= limit`; +/// - `Ok(false)` if `space_needed <= limit` but `current_len + space_needed > limit`, i.e. +/// the current builder is too full but an empty builder could accommodate the batch (the +/// caller may freeze the current builder and replay the batch onto an empty one); +/// - `Err(InvalidBatchSnafu)` if the batch's string data alone (`space_needed`, or the fact +/// that its total bytes cannot even be represented as `i32`) exceeds `limit`, i.e. the +/// batch can never be accommodated by any builder. +fn check_string_capacity(current_len: i32, space_needed: Option, limit: i32) -> Result { + let Some(space_needed) = space_needed.filter(|&space| space <= limit) else { + return error::InvalidBatchSnafu { + reason: format!( + "String data of the batch exceeds the Arrow string array offset limit ({limit}) and cannot be stored in a single column" + ), + } + .fail(); + }; + Ok(current_len + .checked_add(space_needed) + .is_some_and(|v| v <= limit)) +} + +/// Scans all string fields of a batch and checks whether they can be accommodated. +/// +/// Every string field is checked (no early return on `Ok(false)`) so that an intrinsic +/// oversize in *any* field surfaces as an error even when an earlier field only overflows +/// the current builder. +/// +/// Returns: +/// - `Ok(true)` if every string field fits into the current builders; +/// - `Ok(false)` if no field is intrinsically oversized but at least one field only fits +/// into an empty builder (the caller may freeze the current builder and replay the batch +/// onto an empty one); +/// - `Err(InvalidBatchSnafu)` if any single field's string data alone (`space_needed`, or +/// the fact that its total bytes cannot even be represented as `i32`) exceeds `limit`, +/// i.e. that field can never be accommodated by any builder. +fn scan_string_capacity( + fields: &[VectorRef], + field_builders: &[Option], + field_types: &[ConcreteDataType], + limit: i32, +) -> Result { + let mut can_fit_current = true; + for ((field_src, field_dest), field_type) in fields + .iter() + .zip(field_builders.iter()) + .zip(field_types.iter()) + { + if !matches!(field_type, ConcreteDataType::String(_)) { + continue; + } + let current_size = match field_dest { + Some(FieldBuilder::String(builder)) => builder.next_offset(), + None => 0, + Some(FieldBuilder::Other(_)) => unreachable!(), + }; + let array = field_src.to_arrow_array(); + let space_needed = if let Some(string_array) = array.as_any().downcast_ref::() + { + i32::try_from(string_array.value_data().len()).ok() + } else { + let string_vector = field_src + .as_any() + .downcast_ref::() + .with_context(|| error::InvalidBatchSnafu { + reason: format!( + "Field type mismatch, expecting String, given: {}", + field_src.data_type() + ), + })?; + checked_string_values_size(string_vector.iter_data().flatten().map(str::len)) + }; + can_fit_current &= check_string_capacity(current_size, space_needed, limit)?; + } + Ok(can_fit_current) +} + /// Builder to build [TimeSeriesMemtable]. #[derive(Debug, Default)] pub struct TimeSeriesMemtableBuilder { @@ -920,32 +1007,17 @@ impl ValueBuilder { } /// Checks if current value builder have sufficient space to accommodate `fields`. - /// Returns false if there is no space to accommodate fields due to offset overflow. + /// + /// Returns `Ok(false)` if the current builder lacks the remaining space to accommodate + /// the fields due to offset overflow, but an empty builder would be able to accommodate + /// the batch (the caller may freeze the current builder and replay the batch onto an + /// empty one). + /// + /// Returns `Err(InvalidBatchSnafu)` if the string data of a single batch itself exceeds + /// the Arrow string array offset limit and thus can never be accommodated, not even by an + /// empty builder. pub(crate) fn can_accommodate(&self, fields: &[VectorRef]) -> Result { - for (field_src, field_dest) in fields.iter().zip(self.fields.iter()) { - let Some(builder) = field_dest else { - continue; - }; - let FieldBuilder::String(builder) = builder else { - continue; - }; - let array = field_src.to_arrow_array(); - let string_array = array - .as_any() - .downcast_ref::() - .with_context(|| error::InvalidBatchSnafu { - reason: format!( - "Field type mismatch, expecting String, given: {}", - field_src.data_type() - ), - })?; - let space_needed = string_array.value_data().len() as i32; - // offset may overflow - if builder.next_offset().checked_add(space_needed).is_none() { - return Ok(false); - } - } - Ok(true) + scan_string_capacity(fields, &self.fields, &self.field_types, i32::MAX) } pub(crate) fn extend( @@ -1017,17 +1089,26 @@ impl ValueBuilder { match builder { FieldBuilder::String(builder) => { let array = field_src.to_arrow_array(); - let string_array = - array + if let Some(string_array) = array.as_any().downcast_ref::() { + builder.append_array(string_array); + } else { + let string_vector = field_src .as_any() - .downcast_ref::() + .downcast_ref::() .with_context(|| error::InvalidBatchSnafu { reason: format!( "Field type mismatch, expecting String, given: {}", field_src.data_type() ), })?; - builder.append_array(string_array); + for value in string_vector.iter_data() { + if let Some(value) = value { + builder.append(value); + } else { + builder.append_null(); + } + } + } } FieldBuilder::Other(builder) => { let len = field_src.len(); @@ -2296,4 +2377,63 @@ mod tests { assert!(iter.next().is_none()); } + + #[test] + fn test_can_accommodate_string_offset_overflow() { + // A batch whose string data alone exceeds the Arrow string array offset limit can + // never be accommodated, not even by an empty builder. `can_accommodate` must return + // a structured `InvalidBatch` error instead of `Ok(false)`: the latter would make + // `Series::extend` freeze the current builder and replay the same oversized batch + // onto an empty builder, which then panics on offset overflow in `StringBuilder`. + assert!(check_string_capacity(0, Some(101), 100).is_err()); + assert!(check_string_capacity(0, None, i32::MAX).is_err()); + + // A batch that fits in an empty builder but not in the current one reports + // `Ok(false)` so the caller can freeze and replay onto an empty builder. + assert!(!check_string_capacity(95, Some(10), 100).unwrap()); + assert!(!check_string_capacity(91, Some(10), 100).unwrap()); + + // The current builder can accommodate the batch when the combined size fits, + // including the case where it exactly reaches the limit. + assert!(check_string_capacity(0, Some(100), 100).unwrap()); + assert!(check_string_capacity(90, Some(10), 100).unwrap()); + } + + #[test] + fn test_can_accommodate_checks_all_string_fields() { + // Two string fields: the first only overflows the *current* builder (an empty + // builder would fit it), the second is intrinsically oversized. The scan must not + // early-return `Ok(false)` on the first field: it has to check every string field so + // the intrinsic oversize of the second field still surfaces as `InvalidBatch`. + let limit = 10; + // 8 bytes already in the current builder: 8 + 4 > 10, but 4 <= 10, so this field + // alone fits an empty builder. + let mut current = StringBuilder::with_capacity(1, 8); + current.append("12345678"); + let field_builders = [Some(FieldBuilder::String(current)), None]; + let field_types = [ + ConcreteDataType::string_datatype(), + ConcreteDataType::string_datatype(), + ]; + + let first = Arc::new(StringVector::from(StringArray::from(vec!["abcd"]))) as VectorRef; + // 15 bytes > limit: intrinsically oversized, can never fit any builder. + let second = Arc::new(StringVector::from(StringArray::from(vec![ + "123456789012345", + ]))) as VectorRef; + + let err = scan_string_capacity(&[first, second], &field_builders, &field_types, limit) + .unwrap_err(); + assert!( + matches!(err, error::Error::InvalidBatch { .. }), + "expected InvalidBatch, got {err:?}" + ); + + // Sanity: without the oversized field the result is just `Ok(false)` (current full, + // empty builder fits), which lets the caller freeze and replay onto an empty builder. + let first_only = Arc::new(StringVector::from(StringArray::from(vec!["abcd"]))) as VectorRef; + assert!( + !scan_string_capacity(&[first_only], &field_builders, &field_types, limit).unwrap() + ); + } } diff --git a/src/servers/src/http/result/prometheus_resp.rs b/src/servers/src/http/result/prometheus_resp.rs index da32103b79..e79628fac0 100644 --- a/src/servers/src/http/result/prometheus_resp.rs +++ b/src/servers/src/http/result/prometheus_resp.rs @@ -479,6 +479,7 @@ fn merge_annotations(target: &mut Option>, source: Option