rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-08-18 12:08:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(metric-engine): handle Utf8View tag/label columns without panicking (#8772)

Browse Source 
* fix(metric-engine): handle Utf8View tag/label columns without panicking

label_replace (planned as DataFusion regexp_replace) coerces to Utf8View,
so label columns materialize as StringViewArray; build_tag_arrays'
StringArray downcast then panicked ('tag column must be utf8') — e.g. for
OTLP/json2 ingest. TSID computation, sparse-PK encoding and tag
extraction now accept generic ArrayRef tag columns (Utf8/LargeUtf8/
Utf8View/Dictionary) via string_array_value_at_index, and build_tag_arrays
errors instead of panicking on non-string columns. The mito2 time-series
memtable string-field paths are hardened the same way.

Adds label_replace_with_utf8view_labels_does_not_panic (issue #8732).

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* refactor(metric-engine): add is_string_null_at helper for tag null checks

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix(datatypes): use is_none_or to satisfy clippy unnecessary-map-or

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

* fix: reject oversized string batches before memtable append

Distinguish a full active string builder from a batch that cannot fit an
empty Arrow string builder at all. Scan every string field so a later
intrinsically oversized field cannot be skipped after an earlier field
requests a freeze. Return InvalidBatch instead of reaching Arrow's offset
overflow panic.

Also cover Utf8View tags with nulls through the metric-engine tag, TSID,
and sparse-primary-key path.

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>

---------

Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com>
This commit is contained in:
discord9
2026-08-12 07:49:46 +00:00
committed by GitHub
parent 3b3a9032e0
commit 9241e96fa2
4 changed files with 447 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/datatypes/src/arrow_array.rs
+37
View File
@@ -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::<i32>();
!array.is_valid(i)
}
DataType::LargeUtf8 => {
let array = array.as_string::<i64>();
!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
src/metric-engine/src/batch_modifier.rs
+187 -20
View File
@@ -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<Vec<&'a ArrayRef>> {
sorted_tag_columns
.iter()
.map(|tc| {
batch
.column(tc.index)
.as_any()
.downcast_ref::<StringArray>()
.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<RecordBatch> {
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<ArrayRef> = 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::<BinaryArray>()
.unwrap()
.value(0)
.to_vec()
})
.collect::<Vec<_>>();
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<Option<&str>> =
vec![Some("ns-a"), Some("ns-b"), None, Some("ns-c")];
let host_values: Vec<Option<&str>> = 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::<BinaryArray>()
.unwrap();
let utf8view_pks = modified_utf8view
.column(0)
.as_any()
.downcast_ref::<BinaryArray>()
.unwrap();
for row in 0..utf8_pks.len() {
assert_eq!(utf8_pks.value(row), utf8view_pks.value(row));
}
}
}
src/mito2/src/memtable/time_series.rs
+171 -31
View File
@@ -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<Item = usize>) -> Option<i32> {
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<i32>, limit: i32) -> Result<bool> {
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<FieldBuilder>],
field_types: &[ConcreteDataType],
limit: i32,
) -> Result<bool> {
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::<StringArray>()
{
i32::try_from(string_array.value_data().len()).ok()
} else {
let string_vector = field_src
.as_any()
.downcast_ref::<StringVector>()
.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<bool> {
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::<StringArray>()
.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::<StringArray>() {
builder.append_array(string_array);
} else {
let string_vector = field_src
.as_any()
.downcast_ref::<StringArray>()
.downcast_ref::<StringVector>()
.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()
);
}
}
src/servers/src/http/result/prometheus_resp.rs
+52
View File
@@ -479,6 +479,7 @@ fn merge_annotations(target: &mut Option<Vec<String>>, source: Option<Vec<String
mod tests {
use std::sync::Arc;
use arrow::array::StringViewArray;
use common_query::native_histogram::{
CUSTOM_BUCKETS_SCHEMA, CounterResetHint, NativeHistogram, Span, build_histogram_array,
native_histogram_value_type,
@@ -664,6 +665,57 @@ mod tests {
assert_eq!(histogram.sum, "3");
}
#[test]
fn label_replace_with_utf8view_labels_does_not_panic() {
// A PromQL `label_replace` query produces its new label through DataFusion's
// `regexp_replace`, whose output materializes as a `Utf8View` array even when
// the source label is a plain `Utf8`. Serializing such labels must not assume
// the column is a `StringArray`.
let schema = Arc::new(Schema::new(vec![
ColumnSchema::new(
"timestamp",
ConcreteDataType::timestamp_millisecond_datatype(),
false,
),
ColumnSchema::new("host", ConcreteDataType::string_datatype(), false),
ColumnSchema::new("host_copy", ConcreteDataType::utf8_view_datatype(), false),
ColumnSchema::new("value", ConcreteDataType::float64_datatype(), true),
]));
let batch = RecordBatch::new(
schema.clone(),
vec![
Arc::new(TimestampMillisecondVector::from_values([1_000])) as _,
Arc::new(StringVector::from(vec![Some("server-01")])) as _,
Arc::new(StringVector::from(StringViewArray::from(vec![Some(
"server-01",
)]))) as _,
Arc::new(Float64Vector::from(vec![Some(1.0)])) as _,
],
)
.unwrap();
let batches = RecordBatches::try_new(schema, vec![batch]).unwrap();
let response = PrometheusJsonResponse::record_batches_to_data(
batches,
Some("label_replace_repro".to_string()),
ValueType::Vector,
)
.unwrap();
let PrometheusResponse::PromData(PromData {
result: PromQueryResult::Vector(series),
..
}) = response
else {
panic!("expected vector response");
};
assert_eq!(series.len(), 1);
assert_eq!(series[0].metric["__name__"], "label_replace_repro");
assert_eq!(series[0].metric["host"], "server-01");
assert_eq!(series[0].metric["host_copy"], "server-01");
assert_eq!(series[0].value, Some((1.0, "1".to_string())));
}
#[test]
fn matrix_response_preserves_ordinary_histogram_nan_and_filters_stale_marker() {
let schema = Arc::new(Schema::new(vec![