From 64ac8a8eb8849d13c5c594d9093d23353aca0a70 Mon Sep 17 00:00:00 2001 From: Dennis Zhuang Date: Sat, 17 Oct 2026 23:04:22 +0800 Subject: [PATCH] fix(promql): expand replacements like Go and merge label sets field by field label_replace expands its template with Go's regexp.Expand rules, where $01 or a reference of more than nine digits names a group. Series that share a label set merge field by field, colliding only when one field has two values at one point. Signed-off-by: Dennis Zhuang --- src/promql/src/grid_series.rs | 42 ++++++++++++++++--------- src/promql/src/grid_series/tests.rs | 30 ++++++++++++++++++ src/promql/src/map_label.rs | 48 ++++++++++++++++++++++++++++- src/promql/src/map_label/tests.rs | 28 +++++++++++++++++ 4 files changed, 133 insertions(+), 15 deletions(-) diff --git a/src/promql/src/grid_series.rs b/src/promql/src/grid_series.rs index e5e9bd53532..4028fe8a056 100644 --- a/src/promql/src/grid_series.rs +++ b/src/promql/src/grid_series.rs @@ -162,7 +162,7 @@ fn label_set(metric: &MetricName) -> LabelSet { } /// Merges series that share a label set into one, as Prometheus does after a function that -/// rewrites labels or drops `__name__`; two of them present at one point fail. +/// rewrites labels or drops `__name__`; two of them holding one field at one point fail. pub fn merge_same_labelset(schema: &Arc, series: Vec) -> Result> { let mut groups: HashMap> = HashMap::with_capacity(series.len()); let mut collided = false; @@ -186,6 +186,7 @@ pub fn merge_same_labelset(schema: &Arc, series: Vec) -> R let mut builder = GridSeriesBuilder::try_new(schema.clone())?; let mut rows = Vec::new(); let mut values = vec![None; schema.fields.len()]; + let mut merged = vec![None; schema.fields.len()]; for (labels, members) in groups { if let [member] = members[..] { output.extend(series[member].take()); @@ -208,19 +209,32 @@ pub fn merge_same_labelset(schema: &Arc, series: Vec) -> R } } rows.sort_unstable_by_key(|&(point, _, _)| point); - if rows.windows(2).any(|pair| pair[0].0 == pair[1].0) { - let labels: Vec<_> = labels - .iter() - .map(|(name, value)| format!("{name}=\"{value}\"")) - .collect(); - return Err(DataFusionError::Execution(format!( - "vector cannot contain metrics with the same labelset {{{}}}", - labels.join(", ") - ))); - } - for &(point, member, row) in &rows { - read_row(series[member].as_ref().unwrap(), row, &mut values)?; - builder.push(point, &values)?; + let mut start = 0; + while start < rows.len() { + let point = rows[start].0; + let end = start + rows[start..].partition_point(|&(p, _, _)| p == point); + // Rows of one point merge field by field; only a field present twice collides. + merged.iter_mut().for_each(|value| *value = None); + for &(_, member, row) in &rows[start..end] { + read_row(series[member].as_ref().unwrap(), row, &mut values)?; + for (slot, value) in merged.iter_mut().zip(&values) { + if value.is_some() { + if slot.is_some() { + let labels: Vec<_> = labels + .iter() + .map(|(name, value)| format!("{name}=\"{value}\"")) + .collect(); + return Err(DataFusionError::Execution(format!( + "vector cannot contain metrics with the same labelset {{{}}}", + labels.join(", ") + ))); + } + *slot = *value; + } + } + } + builder.push(point, &merged)?; + start = end; } let metric = series[members[0]].as_ref().unwrap().metric.clone(); output.extend(builder.finish(metric)?); diff --git a/src/promql/src/grid_series/tests.rs b/src/promql/src/grid_series/tests.rs index aaa98e9312d..4193beb09ab 100644 --- a/src/promql/src/grid_series/tests.rs +++ b/src/promql/src/grid_series/tests.rs @@ -112,3 +112,33 @@ fn merge_treats_empty_label_values_as_absent() { Points::from([(0, 1f64.to_bits()), (1, 2f64.to_bits())]) ); } + +/// Series of one label set merge field by field: different fields at one point combine, and +/// only one field present in both collides. +#[test] +fn merge_combines_disjoint_fields_at_one_point() { + let schema = grid_schema(2, 1); + let series = |row: [Option; 2]| { + let mut builder = GridSeriesBuilder::try_new(schema.clone()).unwrap(); + builder.push(0, &row).unwrap(); + builder.finish(metric(0)).unwrap().unwrap() + }; + let merged = merge_same_labelset( + &schema, + vec![series([Some(1.0), None]), series([None, Some(2.0)])], + ) + .unwrap(); + assert_eq!(merged.len(), 1); + assert_series_eq( + Some(&merged[0]), + &vec![vec![Some(1.0), Some(2.0)]], + &metric(0), + ); + assert!( + merge_same_labelset( + &schema, + vec![series([Some(1.0), None]), series([Some(2.0), None])] + ) + .is_err() + ); +} diff --git a/src/promql/src/map_label.rs b/src/promql/src/map_label.rs index eea21700630..1df90c290d9 100644 --- a/src/promql/src/map_label.rs +++ b/src/promql/src/map_label.rs @@ -53,12 +53,58 @@ pub fn label_replace( continue; }; value.clear(); - captures.expand(replacement, &mut value); + expand(&compiled, &captures, replacement, &mut value); set_label(&mut series.metric, dst, &value)?; } merge_same_labelset(schema, series) } +/// Go `regexp.Expand`, whose template syntax differs from `regex::Captures::expand`: a numeric +/// reference with a leading zero, like `$01`, or of more than nine digits is a group name. +fn expand(regex: &Regex, captures: ®ex::Captures<'_>, template: &str, out: &mut String) { + let mut rest = template; + while let Some(dollar) = rest.find('$') { + out.push_str(&rest[..dollar]); + rest = &rest[dollar + 1..]; + if let Some(after) = rest.strip_prefix('$') { + out.push('$'); + rest = after; + continue; + } + let (braced, body) = match rest.strip_prefix('{') { + Some(body) => (true, body), + None => (false, rest), + }; + let len = body + .find(|c: char| !c.is_alphanumeric() && c != '_') + .unwrap_or(body.len()); + let name = &body[..len]; + let closed = !braced || body[len..].starts_with('}'); + if name.is_empty() || !closed { + // Malformed; Go keeps the `$` as text. + out.push('$'); + continue; + } + rest = &body[len + usize::from(braced)..]; + let number = (name.bytes().all(|b| b.is_ascii_digit()) + && !(name.len() > 1 && name.starts_with('0'))) + .then(|| name.parse::().ok()) + .flatten() + .filter(|_| name.len() <= 9); + let group = match number { + Some(number) => captures.get(number), + None => regex + .capture_names() + .position(|group| group == Some(name)) + .and_then(|index| captures.get(index)), + }; + if let Some(group) = group { + out.push_str(group.as_str()); + } + } + out.push_str(rest); +} + /// `label_join(v, dst, separator, src...)`: `dst` is set to the values of `src` joined by /// `separator`, an absent label joining as the empty string. pub fn label_join( diff --git a/src/promql/src/map_label/tests.rs b/src/promql/src/map_label/tests.rs index 5d779eb2c99..c24f955b40b 100644 --- a/src/promql/src/map_label/tests.rs +++ b/src/promql/src/map_label/tests.rs @@ -84,3 +84,31 @@ fn label_functions_match_prometheus() { "{matched} matched, {merged} merged, {failed} failed" ); } + +/// Replacement templates expand as Go's `regexp.Expand` does; expectations come from Go. +#[test] +fn replacements_expand_like_go() { + let regex = regex::Regex::new("^(?s:(?P[a-z]+)-([0-9]+))$").unwrap(); + let captures = regex.captures("abc-42").unwrap(); + for (template, expected) in [ + ("$01", ""), + ("${01}", ""), + ("$1", "abc"), + ("${1}x", "abcx"), + ("$1x", ""), + ("$$", "$"), + ("$", "$"), + ("${1", "${1"), + ("$word", "abc"), + ("${word}!", "abc!"), + ("$2$1", "42abc"), + ("$-", "$-"), + ("$0", "abc-42"), + ("$9", ""), + ("$1234567890", ""), + ] { + let mut out = String::new(); + super::expand(®ex, &captures, template, &mut out); + assert_eq!(out, expected, "{template}"); + } +}