mirror of
https://github.com/openobserve/openobserve.git
synced 2026-11-05 00:09:35 +00:00
refactor(promql): extract shared fused evaluation and selector helpers (#14035)
This commit is contained in:
Huaijin
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,153 @@
|
||||
// Copyright 2026 OpenObserve Inc.
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
//! Aggregation dispatch over the fused and generic evaluators.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use config::meta::promql::value::*;
|
||||
use datafusion::error::{DataFusionError, Result};
|
||||
use promql_parser::parser::{Call, Expr as PromExpr, LabelModifier, token};
|
||||
|
||||
use super::Engine;
|
||||
use crate::{aggregations, functions, fused};
|
||||
|
||||
/// A recognized `agg(range_func(...))` expression, the shape the fused
|
||||
/// evaluator accepts.
|
||||
struct FusedAggShape<'a> {
|
||||
op: fused::FusedAggOp,
|
||||
func: Arc<dyn functions::RangeFunc>,
|
||||
range_arg: &'a PromExpr,
|
||||
}
|
||||
|
||||
impl Engine {
|
||||
pub(super) async fn aggregate_exprs(
|
||||
&mut self,
|
||||
op: &token::TokenType,
|
||||
expr: &PromExpr,
|
||||
param: &Option<Box<PromExpr>>,
|
||||
modifier: &Option<LabelModifier>,
|
||||
) -> Result<Value> {
|
||||
// fused shapes fold the range function into the aggregation; others stay generic
|
||||
if let Some(shape) = fused_agg_shape(op, expr) {
|
||||
let range_input = self.exec_expr(shape.range_arg).await?;
|
||||
return fused::fused_range_agg(
|
||||
modifier,
|
||||
range_input,
|
||||
shape.func.as_ref(),
|
||||
shape.op,
|
||||
&self.eval_ctx,
|
||||
);
|
||||
}
|
||||
|
||||
let input = self.exec_expr(expr).await?;
|
||||
|
||||
let eval_ctx = self.eval_ctx.clone();
|
||||
|
||||
Ok(match op.id() {
|
||||
token::T_SUM => aggregations::sum(modifier, input, &eval_ctx)?,
|
||||
token::T_AVG => aggregations::avg(modifier, input, &eval_ctx)?,
|
||||
token::T_COUNT => aggregations::count(modifier, input, &eval_ctx)?,
|
||||
token::T_MIN => aggregations::min(modifier, input, &eval_ctx)?,
|
||||
token::T_MAX => aggregations::max(modifier, input, &eval_ctx)?,
|
||||
token::T_GROUP => aggregations::group(modifier, input, &eval_ctx)?,
|
||||
token::T_STDDEV => aggregations::stddev(modifier, input, &eval_ctx)?,
|
||||
token::T_STDVAR => aggregations::stdvar(modifier, input, &eval_ctx)?,
|
||||
token::T_TOPK => {
|
||||
let param_expr = param.clone().unwrap();
|
||||
let k_value = self.exec_expr(¶m_expr).await?;
|
||||
let k = match k_value {
|
||||
Value::Float(f) => f as usize,
|
||||
_ => {
|
||||
return Err(DataFusionError::Plan(
|
||||
"[topk] param must be a number".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
aggregations::topk(k, modifier, input, &eval_ctx)?
|
||||
}
|
||||
token::T_BOTTOMK => {
|
||||
let param_expr = param.clone().unwrap();
|
||||
let k_value = self.exec_expr(¶m_expr).await?;
|
||||
let k = match k_value {
|
||||
Value::Float(f) => f as usize,
|
||||
_ => {
|
||||
return Err(DataFusionError::Plan(
|
||||
"[bottomk] param must be a number".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
aggregations::bottomk(k, modifier, input, &eval_ctx)?
|
||||
}
|
||||
token::T_COUNT_VALUES => {
|
||||
let param_expr = param.clone().unwrap();
|
||||
let label_name = self.exec_expr(¶m_expr).await?;
|
||||
let label_name_str = match label_name {
|
||||
Value::String(s) => s,
|
||||
_ => {
|
||||
return Err(DataFusionError::Plan(
|
||||
"[count_values] param must be a string".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
aggregations::count_values(&label_name_str, modifier, input, &eval_ctx)?
|
||||
}
|
||||
token::T_QUANTILE => {
|
||||
let param_expr = param.clone().unwrap();
|
||||
let qtile_value = self.exec_expr(¶m_expr).await?;
|
||||
let qtile = match qtile_value {
|
||||
Value::Float(f) => f,
|
||||
_ => {
|
||||
return Err(DataFusionError::Plan(
|
||||
"[quantile] param must be a number".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
aggregations::quantile(qtile, input, &eval_ctx)?
|
||||
}
|
||||
_ => {
|
||||
return Err(DataFusionError::NotImplemented(format!(
|
||||
"Unsupported Aggregate: {op:?}"
|
||||
)));
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn fused_agg_shape<'a>(op: &token::TokenType, expr: &'a PromExpr) -> Option<FusedAggShape<'a>> {
|
||||
if !config::get_config()
|
||||
.search
|
||||
.feature_metrics_fused_agg_enabled
|
||||
{
|
||||
return None;
|
||||
}
|
||||
let agg_op = fused::FusedAggOp::from_token(op.id())?;
|
||||
let PromExpr::Call(Call { func, args }) = expr else {
|
||||
return None;
|
||||
};
|
||||
if args.len() != 1 {
|
||||
return None;
|
||||
}
|
||||
let range_func = functions::fusable_range_func(func.name)?;
|
||||
let range_arg: &PromExpr = args
|
||||
.args
|
||||
.last()
|
||||
.expect("promql-parser validated the function argument");
|
||||
Some(FusedAggShape {
|
||||
op: agg_op,
|
||||
func: Arc::from(range_func),
|
||||
range_arg,
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,711 @@
|
||||
// Copyright 2026 OpenObserve Inc.
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
//! PromQL function-call dispatch and argument helpers. Touches only
|
||||
//! `eval_ctx` besides recursing through `exec_expr`.
|
||||
|
||||
use std::str::FromStr;
|
||||
|
||||
use config::meta::promql::value::*;
|
||||
use datafusion::error::{DataFusionError, Result};
|
||||
use hashbrown::HashSet;
|
||||
use promql_parser::parser::{Expr as PromExpr, Function, FunctionArgs};
|
||||
|
||||
use super::Engine;
|
||||
use crate::functions;
|
||||
|
||||
impl Engine {
|
||||
pub(super) async fn call_expr(
|
||||
&mut self,
|
||||
func: &Function,
|
||||
args: &FunctionArgs,
|
||||
) -> Result<Value> {
|
||||
use crate::functions::Func;
|
||||
|
||||
let func_name = Func::from_str(func.name).map_err(|_| {
|
||||
DataFusionError::NotImplemented(format!("Unsupported function: {}", func.name))
|
||||
})?;
|
||||
|
||||
// There are a few functions which need no arguments for e.g. time()
|
||||
let functions_without_args: HashSet<&str> = HashSet::from_iter(vec![
|
||||
"day_of_month",
|
||||
"day_of_week",
|
||||
"day_of_year",
|
||||
"days_in_month",
|
||||
"hour",
|
||||
"minute",
|
||||
"month",
|
||||
"time",
|
||||
"year",
|
||||
]);
|
||||
let input = match functions_without_args.contains(func.name) {
|
||||
true => match args.len() {
|
||||
0 => {
|
||||
// Found no arg to pass to, lets use a `matrix(time())` as the arg.
|
||||
// https://prometheus.io/docs/prometheus/latest/querying/functions/#functions
|
||||
let timestamps = self.eval_ctx.timestamps();
|
||||
let samples: Vec<Sample> = timestamps
|
||||
.iter()
|
||||
.map(|&ts| Sample::new(ts, ts as f64))
|
||||
.collect();
|
||||
let default_now_matrix = vec![RangeValue {
|
||||
labels: Labels::default(),
|
||||
samples,
|
||||
exemplars: None,
|
||||
time_window: None,
|
||||
}];
|
||||
Value::Matrix(default_now_matrix)
|
||||
}
|
||||
1 => self.call_expr_first_arg(args).await?,
|
||||
|
||||
_ => {
|
||||
return Err(DataFusionError::NotImplemented(
|
||||
"Invalid args passed to the function".into(),
|
||||
));
|
||||
}
|
||||
},
|
||||
false => {
|
||||
let last_arg = args
|
||||
.last()
|
||||
.expect("BUG: promql-parser should have validated function arguments");
|
||||
self.exec_expr(&last_arg).await?
|
||||
}
|
||||
};
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
let result = match func_name {
|
||||
Func::Abs => functions::abs(input)?,
|
||||
Func::Absent => functions::absent(input, &self.eval_ctx)?,
|
||||
Func::AbsentOverTime => functions::absent_over_time(input, &self.eval_ctx)?,
|
||||
Func::AvgOverTime => functions::avg_over_time(input, &self.eval_ctx)?,
|
||||
Func::Ceil => functions::ceil(input)?,
|
||||
Func::Changes => functions::changes(input, &self.eval_ctx)?,
|
||||
Func::Clamp => {
|
||||
let err =
|
||||
"Invalid args, expected \"clamp(v instant-vector, min scalar, max scalar)\"";
|
||||
self.ensure_three_args(args, err)?;
|
||||
|
||||
let input = self.call_expr_first_arg(args).await?;
|
||||
let min = self.call_expr_second_arg(args).await?;
|
||||
let max = self.call_expr_third_arg(args).await?;
|
||||
|
||||
let (min_f, max_f) = match (min, max) {
|
||||
(Value::Float(min), Value::Float(max)) => {
|
||||
if min > max {
|
||||
return Ok(Value::Matrix(vec![]));
|
||||
}
|
||||
(min, max)
|
||||
}
|
||||
_ => {
|
||||
return Err(DataFusionError::NotImplemented(err.into()));
|
||||
}
|
||||
};
|
||||
functions::clamp(input, min_f, max_f)?
|
||||
}
|
||||
Func::ClampMax => {
|
||||
let err = "Invalid args, expected \"clamp(v instant-vector, max scalar)\"";
|
||||
self.ensure_two_args(args, err)?;
|
||||
|
||||
let input = self.call_expr_first_arg(args).await?;
|
||||
let max = self.call_expr_second_arg(args).await?;
|
||||
let max_f = match max {
|
||||
Value::Float(max) => max,
|
||||
_ => {
|
||||
return Err(DataFusionError::NotImplemented(err.into()));
|
||||
}
|
||||
};
|
||||
functions::clamp(input, f64::MIN, max_f)?
|
||||
}
|
||||
Func::ClampMin => {
|
||||
let err = "Invalid args, expected \"clamp(v instant-vector, min scalar)\"";
|
||||
self.ensure_two_args(args, err)?;
|
||||
|
||||
let input = self.call_expr_first_arg(args).await?;
|
||||
let min = self.call_expr_second_arg(args).await?;
|
||||
let min_f = match min {
|
||||
Value::Float(min) => min,
|
||||
_ => {
|
||||
return Err(DataFusionError::NotImplemented(err.into()));
|
||||
}
|
||||
};
|
||||
functions::clamp(input, min_f, f64::MAX)?
|
||||
}
|
||||
Func::CountOverTime => functions::count_over_time(input, &self.eval_ctx)?,
|
||||
Func::DayOfMonth => functions::day_of_month(input)?,
|
||||
Func::DayOfWeek => functions::day_of_week(input)?,
|
||||
Func::DayOfYear => functions::day_of_year(input)?,
|
||||
Func::DaysInMonth => functions::days_in_month(input)?,
|
||||
Func::Delta => functions::delta(input, &self.eval_ctx)?,
|
||||
Func::Deriv => functions::deriv(input, &self.eval_ctx)?,
|
||||
Func::Exp => functions::exp(input)?,
|
||||
Func::Floor => functions::floor(input)?,
|
||||
Func::HistogramCount => {
|
||||
return Err(DataFusionError::NotImplemented(format!(
|
||||
"Unsupported Function: {func_name:?}"
|
||||
)));
|
||||
}
|
||||
Func::HistogramFraction => {
|
||||
return Err(DataFusionError::NotImplemented(format!(
|
||||
"Unsupported Function: {func_name:?}"
|
||||
)));
|
||||
}
|
||||
Func::HistogramQuantile => {
|
||||
let args = &args.args;
|
||||
if args.len() != 2 {
|
||||
return Err(DataFusionError::Plan(format!(
|
||||
"{}: expected 2 arguments, got {}",
|
||||
func.name,
|
||||
args.len()
|
||||
)));
|
||||
}
|
||||
let phi = {
|
||||
match *args[0] {
|
||||
PromExpr::NumberLiteral(ref num) => num.val,
|
||||
_ => {
|
||||
return Err(DataFusionError::Plan(format!(
|
||||
"{}: the first argument must be a number",
|
||||
func.name
|
||||
)));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Use range version if we have an eval context
|
||||
functions::histogram_quantile(phi, input, &self.eval_ctx)?
|
||||
}
|
||||
Func::HistogramSum => {
|
||||
return Err(DataFusionError::NotImplemented(format!(
|
||||
"Unsupported Function: {func_name:?}"
|
||||
)));
|
||||
}
|
||||
Func::HoltWinters => {
|
||||
let err =
|
||||
"Invalid args, expected \"holt_winters(v range-vector, sf scalar, tf scalar)\"";
|
||||
self.ensure_three_args(args, err)?;
|
||||
|
||||
let input = self.call_expr_first_arg(args).await?;
|
||||
let sf = self.call_expr_second_arg(args).await?;
|
||||
let tf = self.call_expr_third_arg(args).await?;
|
||||
|
||||
let scaling_factor = self.parse_f64_else_err(&sf, err)?;
|
||||
let trend_factor = self.parse_f64_else_err(&tf, err)?;
|
||||
|
||||
functions::holt_winters(input, scaling_factor, trend_factor, &self.eval_ctx)?
|
||||
}
|
||||
Func::Hour => functions::hour(input)?,
|
||||
Func::Idelta => functions::idelta(input, &self.eval_ctx)?,
|
||||
Func::Increase => functions::increase(input, &self.eval_ctx)?,
|
||||
Func::Irate => functions::irate(input, &self.eval_ctx)?,
|
||||
Func::LabelJoin => {
|
||||
let err = "Invalid args, expected \"label_join(v instant-vector, dst string, sep string, src_1 string, src_2 string, ...)\"";
|
||||
self.ensure_ge_three_args(args, err)?;
|
||||
|
||||
let input = self.call_expr_first_arg(args).await?;
|
||||
let dst_label = self.call_expr_second_arg(args).await?.get_string().ok_or(
|
||||
DataFusionError::NotImplemented("Invalid destination label found".into()),
|
||||
)?;
|
||||
let separator = self.call_expr_third_arg(args).await?.get_string().ok_or(
|
||||
DataFusionError::NotImplemented("Invalid separator label found".into()),
|
||||
)?;
|
||||
|
||||
let mut source_labels = vec![];
|
||||
for each_src in args.args[3..].iter() {
|
||||
if let Value::String(label) = self.exec_expr(each_src).await.unwrap() {
|
||||
source_labels.push(label);
|
||||
};
|
||||
}
|
||||
if source_labels.is_empty() {
|
||||
return Err(DataFusionError::NotImplemented(
|
||||
"source labels can not be empty or invalid".into(),
|
||||
));
|
||||
}
|
||||
functions::label_join(input, &dst_label, &separator, source_labels)?
|
||||
}
|
||||
Func::LabelReplace => {
|
||||
let err = "Invalid args, expected \"label_replace(v instant-vector, dst_label string, replacement string, src_label string, regex string)\"";
|
||||
|
||||
self.ensure_five_args(args, err)?;
|
||||
let input = self.call_expr_first_arg(args).await?;
|
||||
|
||||
let dst_label = self.call_expr_second_arg(args).await?.get_string().ok_or(
|
||||
DataFusionError::NotImplemented("Invalid destination label found".into()),
|
||||
)?;
|
||||
let replacement = self.call_expr_third_arg(args).await?.get_string().ok_or(
|
||||
DataFusionError::NotImplemented("Invalid replacement string found".into()),
|
||||
)?;
|
||||
|
||||
let src_label = self.call_expr_fourth_arg(args).await?.get_string().ok_or(
|
||||
DataFusionError::NotImplemented("Invalid source label string found".into()),
|
||||
)?;
|
||||
|
||||
let regex = self.call_expr_fifth_arg(args).await?.get_string().ok_or(
|
||||
DataFusionError::NotImplemented("Invalid regex string found".into()),
|
||||
)?;
|
||||
|
||||
functions::label_replace(input, &dst_label, &replacement, &src_label, ®ex)?
|
||||
}
|
||||
Func::LastOverTime => functions::last_over_time(input, &self.eval_ctx)?,
|
||||
Func::Ln => functions::ln(input)?,
|
||||
Func::Log10 => functions::log10(input)?,
|
||||
Func::Log2 => functions::log2(input)?,
|
||||
Func::MaxOverTime => functions::max_over_time(input, &self.eval_ctx)?,
|
||||
Func::MinOverTime => functions::min_over_time(input, &self.eval_ctx)?,
|
||||
Func::Minute => functions::minute(input)?,
|
||||
Func::Month => functions::month(input)?,
|
||||
Func::PredictLinear => {
|
||||
let err = "Invalid args, expected \"predict_linear(v range-vector, t scalar)\"";
|
||||
|
||||
self.ensure_two_args(args, err)?;
|
||||
let input = self.call_expr_first_arg(args).await?;
|
||||
|
||||
let prediction_steps = self.call_expr_second_arg(args).await?.get_float().ok_or(
|
||||
DataFusionError::NotImplemented(
|
||||
"Invalid prediction_steps, f64 expected".into(),
|
||||
),
|
||||
)?;
|
||||
functions::predict_linear(input, prediction_steps, &self.eval_ctx)?
|
||||
}
|
||||
Func::QuantileOverTime => {
|
||||
let err = "Invalid args, expected \"quantile_over_time(scalar, range-vector)\"";
|
||||
|
||||
self.ensure_two_args(args, err)?;
|
||||
let phi_quantile = match self.call_expr_first_arg(args).await {
|
||||
Ok(Value::Float(v)) => v,
|
||||
_ => {
|
||||
return Err(DataFusionError::Plan(
|
||||
"[quantile] param must be a NumberLiteral".into(),
|
||||
));
|
||||
}
|
||||
};
|
||||
let input = self.call_expr_second_arg(args).await?;
|
||||
functions::quantile_over_time(phi_quantile, input, &self.eval_ctx)?
|
||||
}
|
||||
Func::Rate => functions::rate(input, &self.eval_ctx)?,
|
||||
Func::Resets => functions::resets(input, &self.eval_ctx)?,
|
||||
Func::Round => functions::round(input)?,
|
||||
Func::Scalar => functions::scalar(input, &self.eval_ctx)?,
|
||||
Func::Sgn => functions::sgn(input)?,
|
||||
Func::Sort => {
|
||||
return Err(DataFusionError::NotImplemented(format!(
|
||||
"Unsupported Function: {func_name:?}"
|
||||
)));
|
||||
}
|
||||
Func::SortDesc => {
|
||||
return Err(DataFusionError::NotImplemented(format!(
|
||||
"Unsupported Function: {func_name:?}"
|
||||
)));
|
||||
}
|
||||
Func::Sqrt => functions::sqrt(input)?,
|
||||
Func::StddevOverTime => functions::stddev_over_time(input, &self.eval_ctx)?,
|
||||
Func::StdvarOverTime => functions::stdvar_over_time(input, &self.eval_ctx)?,
|
||||
Func::SumOverTime => functions::sum_over_time(input, &self.eval_ctx)?,
|
||||
// TODO: check this implementation
|
||||
Func::Time => Value::Float((self.eval_ctx.start / 1_000_000) as f64),
|
||||
Func::Timestamp => functions::timestamp(input)?,
|
||||
Func::Vector => functions::vector(input, &self.eval_ctx)?,
|
||||
Func::Year => functions::year(input)?,
|
||||
};
|
||||
log::info!(
|
||||
"[trace_id: {}] [PromQL Timing] call_expr({}) execution took: {:?}",
|
||||
self.trace_id,
|
||||
func.name,
|
||||
start.elapsed()
|
||||
);
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
async fn call_expr_first_arg(&mut self, args: &FunctionArgs) -> Result<Value> {
|
||||
self.exec_expr(args.args.first().expect("Missing arg 0"))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn call_expr_second_arg(&mut self, args: &FunctionArgs) -> Result<Value> {
|
||||
self.exec_expr(args.args.get(1).expect("Missing arg 1"))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn call_expr_third_arg(&mut self, args: &FunctionArgs) -> Result<Value> {
|
||||
self.exec_expr(args.args.get(2).expect("Missing arg 2"))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn call_expr_fourth_arg(&mut self, args: &FunctionArgs) -> Result<Value> {
|
||||
self.exec_expr(args.args.get(3).expect("Missing arg 3"))
|
||||
.await
|
||||
}
|
||||
|
||||
async fn call_expr_fifth_arg(&mut self, args: &FunctionArgs) -> Result<Value> {
|
||||
self.exec_expr(args.args.get(4).expect("Missing arg 4"))
|
||||
.await
|
||||
}
|
||||
|
||||
fn ensure_two_args(&self, args: &FunctionArgs, err: &str) -> Result<()> {
|
||||
if args.len() != 2 {
|
||||
return Err(DataFusionError::NotImplemented(err.into()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ensure_three_args(&self, args: &FunctionArgs, err: &str) -> Result<()> {
|
||||
if args.len() != 3 {
|
||||
return Err(DataFusionError::NotImplemented(err.into()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ensure_ge_three_args(&self, args: &FunctionArgs, err: &str) -> Result<()> {
|
||||
if args.len() < 3 {
|
||||
return Err(DataFusionError::NotImplemented(err.into()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn ensure_five_args(&self, args: &FunctionArgs, err: &str) -> Result<()> {
|
||||
if args.len() != 5 {
|
||||
return Err(DataFusionError::NotImplemented(err.into()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_f64_else_err<T: Into<String>>(&self, value: &Value, err: T) -> Result<f64> {
|
||||
match value {
|
||||
Value::Float(f) => Ok(*f),
|
||||
_ => Err(DataFusionError::NotImplemented(err.into())),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use promql_parser::parser::{FunctionArgs, NumberLiteral};
|
||||
|
||||
use super::*;
|
||||
use crate::{engine::tests::*, exec::PromqlContext};
|
||||
|
||||
#[test]
|
||||
fn test_ensure_two_args() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 1.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 2.0 })),
|
||||
],
|
||||
};
|
||||
let result = engine.ensure_two_args(&args, "test error");
|
||||
assert!(result.is_ok());
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![Box::new(PromExpr::NumberLiteral(NumberLiteral {
|
||||
val: 1.0,
|
||||
}))],
|
||||
};
|
||||
let result = engine.ensure_two_args(&args, "test error");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ensure_three_args() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 1.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 2.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 3.0 })),
|
||||
],
|
||||
};
|
||||
let result = engine.ensure_three_args(&args, "test error");
|
||||
assert!(result.is_ok());
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![Box::new(PromExpr::NumberLiteral(NumberLiteral {
|
||||
val: 1.0,
|
||||
}))],
|
||||
};
|
||||
let result = engine.ensure_three_args(&args, "test error");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ensure_ge_three_args() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 1.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 2.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 3.0 })),
|
||||
],
|
||||
};
|
||||
let result = engine.ensure_ge_three_args(&args, "test error");
|
||||
assert!(result.is_ok());
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 1.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 2.0 })),
|
||||
],
|
||||
};
|
||||
let result = engine.ensure_ge_three_args(&args, "test error");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ensure_five_args() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 1.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 2.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 3.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 4.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 5.0 })),
|
||||
],
|
||||
};
|
||||
let result = engine.ensure_five_args(&args, "test error");
|
||||
assert!(result.is_ok());
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![Box::new(PromExpr::NumberLiteral(NumberLiteral {
|
||||
val: 1.0,
|
||||
}))],
|
||||
};
|
||||
let result = engine.ensure_five_args(&args, "test error");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parse_f64_else_err() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let value = Value::Float(42.0);
|
||||
let result = engine.parse_f64_else_err(&value, "test error");
|
||||
assert!(result.is_ok());
|
||||
assert_eq!(result.unwrap(), 42.0);
|
||||
|
||||
let value = Value::String("not a float".to_string());
|
||||
let result = engine.parse_f64_else_err(&value, "test error");
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_call_expr_first_arg() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![Box::new(PromExpr::NumberLiteral(NumberLiteral {
|
||||
val: 42.0,
|
||||
}))],
|
||||
};
|
||||
let result = engine.call_expr_first_arg(&args).await;
|
||||
assert!(result.is_ok());
|
||||
|
||||
if let Ok(Value::Float(val)) = result {
|
||||
assert_eq!(val, 42.0);
|
||||
} else {
|
||||
panic!("Expected Value::Float");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_call_expr_second_arg() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 10.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 42.0 })),
|
||||
],
|
||||
};
|
||||
let result = engine.call_expr_second_arg(&args).await;
|
||||
assert!(result.is_ok());
|
||||
|
||||
if let Ok(Value::Float(val)) = result {
|
||||
assert_eq!(val, 42.0);
|
||||
} else {
|
||||
panic!("Expected Value::Float");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_call_expr_third_arg() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 10.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 20.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 42.0 })),
|
||||
],
|
||||
};
|
||||
let result = engine.call_expr_third_arg(&args).await;
|
||||
assert!(result.is_ok());
|
||||
|
||||
if let Ok(Value::Float(val)) = result {
|
||||
assert_eq!(val, 42.0);
|
||||
} else {
|
||||
panic!("Expected Value::Float");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_call_expr_fourth_arg() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 10.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 20.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 30.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 42.0 })),
|
||||
],
|
||||
};
|
||||
let result = engine.call_expr_fourth_arg(&args).await;
|
||||
assert!(result.is_ok());
|
||||
|
||||
if let Ok(Value::Float(val)) = result {
|
||||
assert_eq!(val, 42.0);
|
||||
} else {
|
||||
panic!("Expected Value::Float");
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_call_expr_fifth_arg() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let args = FunctionArgs {
|
||||
args: vec![
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 10.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 20.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 30.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 40.0 })),
|
||||
Box::new(PromExpr::NumberLiteral(NumberLiteral { val: 42.0 })),
|
||||
],
|
||||
};
|
||||
let result = engine.call_expr_fifth_arg(&args).await;
|
||||
assert!(result.is_ok());
|
||||
|
||||
if let Ok(Value::Float(val)) = result {
|
||||
assert_eq!(val, 42.0);
|
||||
} else {
|
||||
panic!("Expected Value::Float");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,650 @@
|
||||
// Copyright 2026 OpenObserve Inc.
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
//! Pre-scan of the query AST that decides which label columns to load.
|
||||
//! Writes `label_selector` and `disable_label_selector`.
|
||||
|
||||
use datafusion::error::{DataFusionError, Result};
|
||||
use promql_parser::parser::{
|
||||
AggregateExpr, BinModifier, BinaryExpr, Call, Expr as PromExpr, LabelModifier, ParenExpr,
|
||||
UnaryExpr, VectorMatchCardinality, token,
|
||||
};
|
||||
|
||||
use super::Engine;
|
||||
|
||||
impl Engine {
|
||||
/// Recursively filters the necessary columns needed for before executing the given PromQL
|
||||
/// expression.
|
||||
pub fn extract_columns_from_prom_expr(&mut self, prom_expr: &PromExpr) -> Result<()> {
|
||||
match prom_expr {
|
||||
PromExpr::Aggregate(AggregateExpr {
|
||||
op,
|
||||
expr,
|
||||
param,
|
||||
modifier,
|
||||
}) => {
|
||||
self.extract_columns_from_prom_expr(expr)?;
|
||||
if let Some(expr) = param {
|
||||
self.extract_columns_from_prom_expr(expr)?;
|
||||
}
|
||||
self.extract_columns_from_modifier(modifier, op);
|
||||
Ok(())
|
||||
}
|
||||
PromExpr::Unary(UnaryExpr { expr }) => self.extract_columns_from_prom_expr(expr),
|
||||
PromExpr::Binary(BinaryExpr {
|
||||
op,
|
||||
lhs,
|
||||
rhs,
|
||||
modifier,
|
||||
}) => {
|
||||
self.extract_columns_from_prom_expr(lhs)?;
|
||||
self.extract_columns_from_prom_expr(rhs)?;
|
||||
if let Some(BinModifier {
|
||||
card,
|
||||
matching,
|
||||
return_bool: _,
|
||||
}) = modifier
|
||||
{
|
||||
self.extract_columns_from_modifier(matching, op);
|
||||
// group_left or group_right -> no column selection
|
||||
match card {
|
||||
VectorMatchCardinality::ManyToOne(_)
|
||||
| VectorMatchCardinality::OneToMany(_) => {
|
||||
self.label_selector.clear();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
PromExpr::Paren(ParenExpr { expr }) => self.extract_columns_from_prom_expr(expr),
|
||||
PromExpr::Subquery(expr) => self.extract_columns_from_prom_expr(&expr.expr),
|
||||
PromExpr::Call(Call { func, args }) => {
|
||||
// `label_replace` / `label_join` create new labels that don't
|
||||
// exist in the source schema. Restricting column selection
|
||||
// based on the aggregation `by()` list would then drop both the
|
||||
// source labels these functions read from and leave the
|
||||
// newly-created label absent from the loaded data, so the
|
||||
// aggregation groups everything together. See issue #11321.
|
||||
if matches!(func.name, "label_replace" | "label_join") {
|
||||
self.disable_label_selector = true;
|
||||
}
|
||||
_ = args
|
||||
.args
|
||||
.iter()
|
||||
.map(|expr| self.extract_columns_from_prom_expr(expr))
|
||||
.collect::<Vec<_>>();
|
||||
Ok(())
|
||||
}
|
||||
PromExpr::Extension(expr) => Err(DataFusionError::NotImplemented(format!(
|
||||
"Unsupported Extension: {expr:?}",
|
||||
))),
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Help function to extract columns from [LabelModifier].
|
||||
/// Aggregation function topk & bottomk are special cases where
|
||||
/// modifier is applied to grouped result -> not columns filtered.
|
||||
/// For promql:
|
||||
/// sum(irate(zo_incoming_requests{namespace="ziox"}[5m])) by (exported_endpoint)
|
||||
/// we need to extract the columns `exported_endpoint` from the modifier. Because
|
||||
/// the result will be grouped by `exported_endpoint`, and don't consider other labesl.
|
||||
fn extract_columns_from_modifier(
|
||||
&mut self,
|
||||
modifier: &Option<LabelModifier>,
|
||||
op: &token::TokenType,
|
||||
) {
|
||||
if let Some(label_modifier) = modifier {
|
||||
match op.id() {
|
||||
// topk and bottomk query all columns when with modifiers
|
||||
token::T_TOPK | token::T_BOTTOMK => self.label_selector.clear(),
|
||||
_ => {
|
||||
if let (label_selector, LabelModifier::Include(labels)) =
|
||||
(&mut self.label_selector, label_modifier)
|
||||
{
|
||||
label_selector.extend(labels.labels.iter().cloned());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use promql_parser::parser::{
|
||||
AggregateExpr, BinModifier, BinaryExpr, Call, Extension, Function, FunctionArgs,
|
||||
NumberLiteral, ParenExpr, StringLiteral, SubqueryExpr, UnaryExpr, VectorMatchCardinality,
|
||||
value::ValueType,
|
||||
};
|
||||
|
||||
use super::*;
|
||||
use crate::{engine::tests::*, exec::PromqlContext};
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_number_literal() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
let expr = PromExpr::NumberLiteral(NumberLiteral { val: 42.0 });
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&expr);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_string_literal() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
let expr = PromExpr::StringLiteral(StringLiteral {
|
||||
val: "test".to_string(),
|
||||
});
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&expr);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_paren() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
let inner_expr = PromExpr::NumberLiteral(NumberLiteral { val: 42.0 });
|
||||
let expr = PromExpr::Paren(ParenExpr {
|
||||
expr: Box::new(inner_expr),
|
||||
});
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&expr);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_unary() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
let inner_expr = PromExpr::NumberLiteral(NumberLiteral { val: 42.0 });
|
||||
let expr = PromExpr::Unary(UnaryExpr {
|
||||
expr: Box::new(inner_expr),
|
||||
});
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&expr);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_binary() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let lhs = PromExpr::NumberLiteral(NumberLiteral { val: 42.0 });
|
||||
let rhs = PromExpr::NumberLiteral(NumberLiteral { val: 10.0 });
|
||||
let expr = PromExpr::Binary(BinaryExpr {
|
||||
lhs: Box::new(lhs),
|
||||
rhs: Box::new(rhs),
|
||||
op: create_test_token(),
|
||||
modifier: None,
|
||||
});
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&expr);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_binary_with_modifier() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let lhs = PromExpr::NumberLiteral(NumberLiteral { val: 42.0 });
|
||||
let rhs = PromExpr::NumberLiteral(NumberLiteral { val: 10.0 });
|
||||
let modifier = Some(BinModifier {
|
||||
card: VectorMatchCardinality::OneToOne,
|
||||
matching: Some(LabelModifier::Include(promql_parser::label::Labels {
|
||||
labels: vec!["env".to_string()],
|
||||
})),
|
||||
return_bool: false,
|
||||
});
|
||||
let expr = PromExpr::Binary(BinaryExpr {
|
||||
lhs: Box::new(lhs),
|
||||
rhs: Box::new(rhs),
|
||||
op: create_test_token(),
|
||||
modifier,
|
||||
});
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&expr);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_binary_with_many_to_one() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let lhs = PromExpr::NumberLiteral(NumberLiteral { val: 42.0 });
|
||||
let rhs = PromExpr::NumberLiteral(NumberLiteral { val: 10.0 });
|
||||
let modifier = Some(BinModifier {
|
||||
card: VectorMatchCardinality::ManyToOne(promql_parser::label::Labels {
|
||||
labels: vec!["env".to_string()],
|
||||
}),
|
||||
matching: Some(LabelModifier::Include(promql_parser::label::Labels {
|
||||
labels: vec!["env".to_string()],
|
||||
})),
|
||||
return_bool: false,
|
||||
});
|
||||
let expr = PromExpr::Binary(BinaryExpr {
|
||||
lhs: Box::new(lhs),
|
||||
rhs: Box::new(rhs),
|
||||
op: create_test_token(),
|
||||
modifier,
|
||||
});
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&expr);
|
||||
assert!(result.is_ok());
|
||||
// Should clear label_selector for ManyToOne
|
||||
assert!(engine.label_selector.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_aggregate() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let expr = PromExpr::NumberLiteral(NumberLiteral { val: 42.0 });
|
||||
let aggregate_expr = PromExpr::Aggregate(AggregateExpr {
|
||||
op: create_test_token(),
|
||||
expr: Box::new(expr),
|
||||
param: None,
|
||||
modifier: None,
|
||||
});
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&aggregate_expr);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_aggregate_with_param() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let expr = PromExpr::NumberLiteral(NumberLiteral { val: 42.0 });
|
||||
let param = PromExpr::NumberLiteral(NumberLiteral { val: 0.5 });
|
||||
let aggregate_expr = PromExpr::Aggregate(AggregateExpr {
|
||||
op: create_test_token(),
|
||||
expr: Box::new(expr),
|
||||
param: Some(Box::new(param)),
|
||||
modifier: None,
|
||||
});
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&aggregate_expr);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_subquery() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let expr = PromExpr::NumberLiteral(NumberLiteral { val: 42.0 });
|
||||
let subquery_expr = PromExpr::Subquery(SubqueryExpr {
|
||||
expr: Box::new(expr),
|
||||
range: Duration::from_secs(300),
|
||||
offset: None,
|
||||
step: None,
|
||||
at: None,
|
||||
});
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&subquery_expr);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_call() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let args = FunctionArgs { args: vec![] };
|
||||
let func = Function {
|
||||
name: "time",
|
||||
arg_types: vec![],
|
||||
variadic: false,
|
||||
return_type: ValueType::Scalar,
|
||||
};
|
||||
let expr = PromExpr::Call(Call { func, args });
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&expr);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_call_with_args() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let arg = PromExpr::NumberLiteral(NumberLiteral { val: 42.0 });
|
||||
let args = FunctionArgs {
|
||||
args: vec![Box::new(arg)],
|
||||
};
|
||||
let func = Function {
|
||||
name: "abs",
|
||||
arg_types: vec![],
|
||||
variadic: false,
|
||||
return_type: ValueType::Scalar,
|
||||
};
|
||||
let expr = PromExpr::Call(Call { func, args });
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&expr);
|
||||
assert!(result.is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_prom_expr_extension() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
// Create a mock extension expression
|
||||
let extension_expr = PromExpr::Extension(Extension {
|
||||
expr: Arc::new(TestExtension),
|
||||
});
|
||||
|
||||
let result = engine.extract_columns_from_prom_expr(&extension_expr);
|
||||
assert!(result.is_err());
|
||||
assert!(
|
||||
result
|
||||
.unwrap_err()
|
||||
.to_string()
|
||||
.contains("Unsupported Extension")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_modifier_none() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
engine.extract_columns_from_modifier(&None, &create_test_token());
|
||||
// Should not change label_selector
|
||||
assert!(engine.label_selector.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_modifier_topk() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let modifier = Some(LabelModifier::Include(promql_parser::label::Labels {
|
||||
labels: vec!["env".to_string()],
|
||||
}));
|
||||
|
||||
engine.extract_columns_from_modifier(&modifier, &create_test_token());
|
||||
// Should clear label_selector for topk
|
||||
assert!(!engine.label_selector.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_modifier_bottomk() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let modifier = Some(LabelModifier::Include(promql_parser::label::Labels {
|
||||
labels: vec!["env".to_string()],
|
||||
}));
|
||||
|
||||
engine.extract_columns_from_modifier(&modifier, &create_test_token());
|
||||
// Should clear label_selector for bottomk
|
||||
assert!(!engine.label_selector.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_modifier_include() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let modifier = Some(LabelModifier::Include(promql_parser::label::Labels {
|
||||
labels: vec!["env".to_string(), "service".to_string()],
|
||||
}));
|
||||
|
||||
engine.extract_columns_from_modifier(&modifier, &create_test_token());
|
||||
// Should add labels to label_selector
|
||||
assert!(!engine.label_selector.is_empty());
|
||||
assert!(engine.label_selector.contains("env"));
|
||||
assert!(engine.label_selector.contains("service"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_label_replace_disables_selector() {
|
||||
// Regression test for #11321: `count by (new) (label_replace(m, "new", ...))`
|
||||
// must not restrict loaded columns to `{"new"}`, since `new` is created
|
||||
// by `label_replace` and the source label it reads from would otherwise
|
||||
// not be loaded, so aggregation collapses all series into one group.
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let label_replace_call = PromExpr::Call(Call {
|
||||
func: Function {
|
||||
name: "label_replace",
|
||||
arg_types: vec![],
|
||||
variadic: false,
|
||||
return_type: ValueType::Vector,
|
||||
},
|
||||
args: FunctionArgs { args: vec![] },
|
||||
});
|
||||
let aggregate_expr = PromExpr::Aggregate(AggregateExpr {
|
||||
op: create_test_token(),
|
||||
expr: Box::new(label_replace_call),
|
||||
param: None,
|
||||
modifier: Some(LabelModifier::Include(promql_parser::label::Labels {
|
||||
labels: vec!["new".to_string()],
|
||||
})),
|
||||
});
|
||||
|
||||
engine
|
||||
.extract_columns_from_prom_expr(&aggregate_expr)
|
||||
.unwrap();
|
||||
assert!(
|
||||
engine.disable_label_selector,
|
||||
"label_replace must set disable_label_selector"
|
||||
);
|
||||
// `exec()` clears the selector when the flag is set; mimic that here.
|
||||
if engine.disable_label_selector {
|
||||
engine.label_selector.clear();
|
||||
}
|
||||
assert!(engine.label_selector.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_extract_columns_from_modifier_exclude() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let modifier = Some(LabelModifier::Exclude(promql_parser::label::Labels {
|
||||
labels: vec!["env".to_string()],
|
||||
}));
|
||||
|
||||
engine.extract_columns_from_modifier(&modifier, &create_test_token());
|
||||
// Should not change label_selector for exclude
|
||||
assert!(engine.label_selector.is_empty());
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,888 @@
|
||||
// Copyright 2026 OpenObserve Inc.
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
//! Vector/matrix selector evaluation and data loading. Reads `ctx`,
|
||||
//! `label_selector`, and `skip_labels`; writes `result_type`.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use config::meta::promql::{NAME_LABEL, value::*};
|
||||
use datafusion::error::{DataFusionError, Result};
|
||||
use futures::future::try_join_all;
|
||||
use hashbrown::HashMap;
|
||||
use infra::errors::{Error, ErrorCodes};
|
||||
use promql_parser::{
|
||||
label::{MatchOp, Matchers},
|
||||
parser::{Offset, VectorSelector},
|
||||
};
|
||||
use rayon::iter::{IntoParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
|
||||
|
||||
use super::Engine;
|
||||
use crate::{
|
||||
load_series::{LoadedMetrics, PartitionedMetrics, selector_load_data_from_datafusion},
|
||||
micros,
|
||||
};
|
||||
|
||||
impl Engine {
|
||||
/// Instant vector selector --- select a single sample at each evaluation
|
||||
/// timestamp.
|
||||
///
|
||||
/// See <https://promlabs.com/blog/2020/07/02/selecting-data-in-promql/#confusion-alert-instantrange-selectors-vs-instantrange-queries>
|
||||
pub(super) async fn eval_vector_selector(
|
||||
&mut self,
|
||||
selector: &VectorSelector,
|
||||
) -> Result<Vec<RangeValue>> {
|
||||
if self.result_type.is_none() {
|
||||
self.result_type = Some("vector".to_string());
|
||||
}
|
||||
|
||||
let mut selector = selector.clone();
|
||||
if selector.name.is_none() {
|
||||
let name = match selector.matchers.find_matchers(NAME_LABEL).first() {
|
||||
Some(mat) => mat.value.clone(),
|
||||
None => {
|
||||
return Err(DataFusionError::Plan(
|
||||
"VectorSelector: metric name is required".into(),
|
||||
));
|
||||
}
|
||||
};
|
||||
selector.name = Some(name);
|
||||
// the matcher is fully consumed by stream selection; leaving it in
|
||||
// would filter on the stored `__name__` column (which may keep the
|
||||
// pre-`format_stream_name` case), leak into partition pruning, and
|
||||
// make the selector's PromQL text unparseable on super-cluster peers
|
||||
selector
|
||||
.matchers
|
||||
.matchers
|
||||
.retain(|mat| mat.name != NAME_LABEL);
|
||||
}
|
||||
|
||||
let data = self.selector_load_data_owned(&selector, None).await?;
|
||||
|
||||
let metrics_cache = match data.get_range_values() {
|
||||
Some(v) => v,
|
||||
None => return Ok(vec![]),
|
||||
};
|
||||
|
||||
let mut offset_modifier = 0;
|
||||
if let Some(offset) = selector.offset {
|
||||
match offset {
|
||||
Offset::Pos(offset) => {
|
||||
offset_modifier = micros(offset);
|
||||
}
|
||||
Offset::Neg(offset) => {
|
||||
offset_modifier = -micros(offset);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
// Get all evaluation timestamps from the context
|
||||
let eval_timestamps = self.eval_ctx.timestamps();
|
||||
|
||||
// For each metric, select appropriate samples at each evaluation timestamp
|
||||
// TODO: make it parallel
|
||||
let mut result = Vec::with_capacity(metrics_cache.len());
|
||||
for metric in metrics_cache {
|
||||
let mut selected_samples = Vec::with_capacity(eval_timestamps.len());
|
||||
|
||||
for &eval_ts in &eval_timestamps {
|
||||
// Calculate lookback window for this evaluation timestamp
|
||||
let start = eval_ts - self.ctx.lookback_delta;
|
||||
|
||||
// Find the sample for this evaluation timestamp
|
||||
// Binary search for the last sample before or at eval_ts (considering offset)
|
||||
let end_index = metric
|
||||
.samples
|
||||
.partition_point(|v| v.timestamp + offset_modifier <= eval_ts);
|
||||
|
||||
let match_sample = if end_index > 0 {
|
||||
metric.samples.get(end_index - 1).and_then(|sample| {
|
||||
let adjusted_ts = sample.timestamp + offset_modifier;
|
||||
if adjusted_ts >= start && adjusted_ts <= eval_ts {
|
||||
Some(sample)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Add the matched sample (already validated to be within range)
|
||||
if let Some(sample) = match_sample {
|
||||
// Use eval_ts as the timestamp for the selected sample
|
||||
// See https://promlabs.com/blog/2020/06/18/the-anatomy-of-a-promql-query/#instant-queries
|
||||
selected_samples.push(Sample::new(eval_ts, sample.value));
|
||||
}
|
||||
}
|
||||
|
||||
// Only include metrics that have at least one sample
|
||||
if !selected_samples.is_empty() {
|
||||
result.push(RangeValue {
|
||||
labels: metric.labels,
|
||||
samples: selected_samples,
|
||||
exemplars: metric.exemplars,
|
||||
time_window: metric.time_window,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
/// Range vector selector --- select a whole time range at each evaluation
|
||||
/// timestamp.
|
||||
///
|
||||
/// See <https://promlabs.com/blog/2020/07/02/selecting-data-in-promql/#confusion-alert-instantrange-selectors-vs-instantrange-queries>
|
||||
///
|
||||
/// MatrixSelector is a special case of VectorSelector that returns a matrix
|
||||
/// of samples.
|
||||
pub(super) async fn eval_matrix_selector(
|
||||
&mut self,
|
||||
selector: &VectorSelector,
|
||||
range: Duration,
|
||||
) -> Result<Vec<RangeValue>> {
|
||||
if self.result_type.is_none() {
|
||||
self.result_type = Some("matrix".to_string());
|
||||
}
|
||||
|
||||
let mut selector = selector.clone();
|
||||
if selector.name.is_none() {
|
||||
let name = selector
|
||||
.matchers
|
||||
.find_matchers(NAME_LABEL)
|
||||
.first()
|
||||
.unwrap()
|
||||
.value
|
||||
.clone();
|
||||
|
||||
selector.name = Some(name);
|
||||
// see eval_vector_selector: the matcher is consumed by stream selection
|
||||
selector
|
||||
.matchers
|
||||
.matchers
|
||||
.retain(|mat| mat.name != NAME_LABEL);
|
||||
}
|
||||
|
||||
let data = self
|
||||
.selector_load_data_owned(&selector, Some(range))
|
||||
.await?;
|
||||
|
||||
let values = match data.get_range_values() {
|
||||
Some(v) => v,
|
||||
None => return Ok(vec![]),
|
||||
};
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
let mut values = values
|
||||
.into_par_iter()
|
||||
.map(|rv| RangeValue {
|
||||
labels: rv.labels,
|
||||
samples: rv.samples,
|
||||
exemplars: rv.exemplars,
|
||||
time_window: Some(TimeWindow::new(range)),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
log::info!(
|
||||
"[trace_id: {}] [PromQL Timing] eval_matrix_selector() processing took: {:?}",
|
||||
self.trace_id,
|
||||
start.elapsed()
|
||||
);
|
||||
|
||||
// apply offset to samples
|
||||
let offset_modifier = get_offset_modifier(selector.offset);
|
||||
if offset_modifier != 0 {
|
||||
values.par_iter_mut().for_each(|rv| {
|
||||
rv.samples
|
||||
.iter_mut()
|
||||
.for_each(|s| s.timestamp += offset_modifier);
|
||||
});
|
||||
}
|
||||
|
||||
Ok(values)
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "promql:engine:load_data_owned", skip_all)]
|
||||
async fn selector_load_data_owned(
|
||||
&mut self,
|
||||
selector: &VectorSelector,
|
||||
range: Option<Duration>,
|
||||
) -> Result<Value> {
|
||||
let mut metric_values = match self.selector_load_data_inner(selector, range).await {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
log::error!(
|
||||
"[trace_id: {}] [PromQL] Failed to load data for stream, error: {e:?}",
|
||||
self.trace_id
|
||||
);
|
||||
return Err(e);
|
||||
}
|
||||
};
|
||||
|
||||
// no data, return immediately
|
||||
if metric_values.is_empty() {
|
||||
return Ok(Value::None);
|
||||
}
|
||||
|
||||
let start = std::time::Instant::now();
|
||||
metric_values.par_iter_mut().for_each(|metric| {
|
||||
metric.samples.sort_unstable_by_key(|k| k.timestamp);
|
||||
if self.ctx.query_ctx.query_exemplars
|
||||
&& let Some(exemplars) = &mut metric.exemplars
|
||||
{
|
||||
exemplars.sort_by_key(|k| k.timestamp);
|
||||
}
|
||||
});
|
||||
let values = if metric_values.is_empty() {
|
||||
Value::None
|
||||
} else {
|
||||
Value::Matrix(metric_values)
|
||||
};
|
||||
log::info!(
|
||||
"[trace_id: {}] [PromQL] sort samples by timestamps took: {:?}",
|
||||
self.trace_id,
|
||||
start.elapsed()
|
||||
);
|
||||
Ok(values)
|
||||
}
|
||||
|
||||
#[tracing::instrument(name = "promql:engine:load_data", skip_all)]
|
||||
async fn selector_load_data_inner(
|
||||
&self,
|
||||
selector: &VectorSelector,
|
||||
range: Option<Duration>,
|
||||
) -> Result<Vec<RangeValue>> {
|
||||
let start_time = std::time::Instant::now();
|
||||
// https://promlabs.com/blog/2020/07/02/selecting-data-in-promql/#lookback-delta
|
||||
let offset_modifier = get_offset_modifier(selector.offset.clone());
|
||||
// Positive offset (e.g. `offset 10m`) looks into the past, so we shift
|
||||
// the data-load window backwards by `offset_modifier`.
|
||||
let start =
|
||||
self.ctx.start - range.map_or(self.ctx.lookback_delta, micros) - offset_modifier;
|
||||
let end = self.ctx.end - offset_modifier;
|
||||
|
||||
// 1. Group by metrics (sets of label name-value pairs)
|
||||
let table_name = selector.name.as_ref().unwrap();
|
||||
log::info!(
|
||||
"[trace_id: {}] [PromQL] loading data for stream: {table_name}, range: [{start},{end}), filter: {:?}",
|
||||
self.trace_id,
|
||||
selector.to_string(),
|
||||
);
|
||||
|
||||
let mut filters = equal_matcher_filters(&selector.matchers);
|
||||
|
||||
// check for super cluster
|
||||
let trace_id = self.ctx.query_ctx.trace_id.clone();
|
||||
#[cfg(feature = "enterprise")]
|
||||
let (super_tx, mut super_rx) = tokio::sync::mpsc::channel::<
|
||||
Result<(HashMap<u64, RangeValue>, config::meta::search::ScanStats)>,
|
||||
>(1);
|
||||
#[cfg(feature = "enterprise")]
|
||||
if self.ctx.query_ctx.is_super_cluster {
|
||||
let query_ctx = self.ctx.query_ctx.clone();
|
||||
let step = self.eval_ctx.step;
|
||||
let selector = selector.clone();
|
||||
let label_selector = self.label_selector.clone();
|
||||
let trace_id_clone = trace_id.clone();
|
||||
tokio::task::spawn(async move {
|
||||
let ret = o2_enterprise::enterprise::metrics::super_cluster::selector_load_data(
|
||||
query_ctx,
|
||||
selector,
|
||||
range,
|
||||
&label_selector,
|
||||
start,
|
||||
end,
|
||||
step,
|
||||
)
|
||||
.await;
|
||||
if let Err(e) = super_tx.send(ret).await {
|
||||
log::error!(
|
||||
"[trace_id: {trace_id_clone}] [PromQL] Failed to send super cluster result to channel, error: {e:?}",
|
||||
);
|
||||
}
|
||||
drop(super_tx);
|
||||
});
|
||||
} else {
|
||||
drop(super_tx);
|
||||
}
|
||||
|
||||
let ctxs = self
|
||||
.ctx
|
||||
.table_provider
|
||||
.create_context(
|
||||
&self.ctx.query_ctx.org_id,
|
||||
table_name,
|
||||
(start, end),
|
||||
selector.matchers.clone(),
|
||||
self.label_selector.clone(),
|
||||
&mut filters,
|
||||
)
|
||||
.await?;
|
||||
|
||||
// check if we need to load data from local cluster
|
||||
#[cfg(feature = "enterprise")]
|
||||
let ctxs = if self.ctx.query_ctx.is_super_cluster
|
||||
&& !o2_enterprise::enterprise::super_cluster::search::has_local_cluster(
|
||||
self.ctx.query_ctx.regions.clone(),
|
||||
self.ctx.query_ctx.clusters.clone(),
|
||||
)
|
||||
.await
|
||||
{
|
||||
vec![]
|
||||
} else {
|
||||
ctxs
|
||||
};
|
||||
|
||||
let mut label_selector = self.label_selector.clone();
|
||||
label_selector.extend(self.ctx.label_selector.iter().cloned());
|
||||
|
||||
// Calculate step and lookback for the optimization
|
||||
let start = self.eval_ctx.start - offset_modifier;
|
||||
let end = self.eval_ctx.end - offset_modifier;
|
||||
let step = self.eval_ctx.step;
|
||||
let lookback = range.map_or(self.ctx.lookback_delta, micros);
|
||||
|
||||
let skip_labels = self.skip_labels;
|
||||
let mut tasks = Vec::with_capacity(ctxs.len());
|
||||
let mut abort_handles = Vec::with_capacity(ctxs.len());
|
||||
for (ctx, schema, scan_stats, keep_filters) in ctxs {
|
||||
let query_ctx = self.ctx.query_ctx.clone();
|
||||
let mut selector = selector.clone();
|
||||
if !keep_filters {
|
||||
selector.matchers = Matchers::empty();
|
||||
};
|
||||
let label_selector = label_selector.clone();
|
||||
let task = tokio::spawn(async move {
|
||||
tokio::time::timeout(
|
||||
Duration::from_secs(query_ctx.timeout),
|
||||
selector_load_data_from_datafusion(
|
||||
query_ctx,
|
||||
ctx,
|
||||
schema,
|
||||
selector,
|
||||
label_selector,
|
||||
start,
|
||||
end,
|
||||
step,
|
||||
lookback,
|
||||
skip_labels,
|
||||
),
|
||||
)
|
||||
.await
|
||||
});
|
||||
abort_handles.push(task.abort_handle());
|
||||
tasks.push(task);
|
||||
// update stats
|
||||
let mut ctx_scan_stats = self.ctx.scan_stats.write().await;
|
||||
ctx_scan_stats.add(&scan_stats);
|
||||
}
|
||||
|
||||
let mut abort_receiver = self
|
||||
.ctx
|
||||
.table_provider
|
||||
.register_cancellation(&trace_id)
|
||||
.await?;
|
||||
|
||||
// run datafusion collect data task
|
||||
let timeout = self.ctx.query_ctx.timeout;
|
||||
let query_task = try_join_all(tasks);
|
||||
tokio::pin!(query_task);
|
||||
let task_results = tokio::select! {
|
||||
ret = &mut query_task => {
|
||||
match ret {
|
||||
Ok(ret) => {
|
||||
// Unwrap the nested Results: JoinHandle result -> timeout result -> actual result
|
||||
let mut unwrapped_results = Vec::new();
|
||||
for result in ret {
|
||||
match result {
|
||||
Ok(Ok(data)) => unwrapped_results.push(data),
|
||||
Ok(Err(_)) => {
|
||||
log::error!("[trace_id {trace_id}] [PromQL] grpc search load data task timeout");
|
||||
return Err(DataFusionError::Plan(
|
||||
Error::ErrorCode(ErrorCodes::SearchTimeout("[PromQL] grpc search load data task timeout".to_string())).to_string()
|
||||
));
|
||||
}
|
||||
Err(err) => {
|
||||
log::error!("[trace_id {trace_id}] [PromQL] grpc search execute error: {err}");
|
||||
return Err(DataFusionError::Plan(format!("task error: {err}")));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(unwrapped_results)
|
||||
},
|
||||
Err(err) => {
|
||||
log::error!("[trace_id {trace_id}] [PromQL] grpc search execute error: {err}");
|
||||
Err(Error::Message(err.to_string()))
|
||||
}
|
||||
}
|
||||
},
|
||||
_ = tokio::time::sleep(tokio::time::Duration::from_secs(timeout )) => {
|
||||
for handle in abort_handles {
|
||||
handle.abort();
|
||||
}
|
||||
log::error!("[trace_id {trace_id}] [PromQL] grpc search timeout");
|
||||
Err(Error::ErrorCode(ErrorCodes::SearchTimeout("[PromQL] grpc search timeout".to_string())))
|
||||
},
|
||||
_ = async {
|
||||
match abort_receiver.as_mut() {
|
||||
Some(receiver) => {
|
||||
let _ = receiver.await;
|
||||
}
|
||||
None => futures::future::pending::<()>().await,
|
||||
}
|
||||
} => {
|
||||
for handle in abort_handles {
|
||||
handle.abort();
|
||||
}
|
||||
log::info!("[trace_id {trace_id}] [PromQL] grpc search canceled");
|
||||
Err(Error::ErrorCode(ErrorCodes::SearchCancelQuery("[PromQL] grpc search canceled".to_string())))
|
||||
}
|
||||
};
|
||||
|
||||
let task_results =
|
||||
task_results.map_err(|e| DataFusionError::Plan(format!("task error: {e}")))?;
|
||||
|
||||
// check for super cluster
|
||||
#[cfg(feature = "enterprise")]
|
||||
let metrics = if self.ctx.query_ctx.is_super_cluster {
|
||||
let mut metrics = merge_loaded_metrics(task_results);
|
||||
let (metric, stats) = match super_rx.recv().await {
|
||||
Some(Ok(ret)) => ret,
|
||||
Some(Err(e)) => {
|
||||
log::error!(
|
||||
"[trace_id: {}] [PromQL] Super cluster result channel error: {e:?}",
|
||||
self.trace_id
|
||||
);
|
||||
return Err(e);
|
||||
}
|
||||
None => {
|
||||
log::error!(
|
||||
"[trace_id: {}] [PromQL] Super cluster result channel is closed",
|
||||
self.trace_id
|
||||
);
|
||||
return Err(DataFusionError::Plan(
|
||||
"super cluster result channel is closed".to_string(),
|
||||
));
|
||||
}
|
||||
};
|
||||
for (key, value) in metric {
|
||||
if let Some(metric) = metrics.get_mut(&key) {
|
||||
metric.extend(value);
|
||||
} else {
|
||||
metrics.insert(key, value);
|
||||
}
|
||||
}
|
||||
let mut ctx_scan_stats = self.ctx.scan_stats.write().await;
|
||||
ctx_scan_stats.add(&stats);
|
||||
|
||||
metrics.into_values().collect()
|
||||
} else {
|
||||
collect_loaded_metrics(task_results)
|
||||
};
|
||||
#[cfg(not(feature = "enterprise"))]
|
||||
let metrics = collect_loaded_metrics(task_results);
|
||||
|
||||
log::info!(
|
||||
"[trace_id: {}] load data done for stream: {}, took: {} ms",
|
||||
self.trace_id,
|
||||
table_name,
|
||||
start_time.elapsed().as_millis()
|
||||
);
|
||||
|
||||
Ok(metrics)
|
||||
}
|
||||
}
|
||||
|
||||
/// Discard the already-partitioned series hashes without rebuilding a global
|
||||
/// hash table. This is the common path when DataFusion creates one query
|
||||
/// context for the selected schema.
|
||||
fn flatten_partitioned_metrics(partitions: PartitionedMetrics) -> Vec<RangeValue> {
|
||||
let metrics_count = partitions.iter().map(HashMap::len).sum();
|
||||
let mut metrics = Vec::with_capacity(metrics_count);
|
||||
for partition in partitions {
|
||||
metrics.extend(partition.into_values());
|
||||
}
|
||||
metrics
|
||||
}
|
||||
|
||||
fn flatten_loaded_metrics(metrics: LoadedMetrics) -> Vec<RangeValue> {
|
||||
match metrics {
|
||||
LoadedMetrics::Partitioned(partitions) => flatten_partitioned_metrics(partitions),
|
||||
LoadedMetrics::Merged(metrics) => metrics.into_values().collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// A single DataFusion context is the common case and its series are already
|
||||
/// deduplicated, so flatten it; series from multiple contexts must be merged
|
||||
/// by series hash.
|
||||
fn collect_loaded_metrics(mut results: Vec<LoadedMetrics>) -> Vec<RangeValue> {
|
||||
if results.len() == 1 {
|
||||
flatten_loaded_metrics(results.pop().unwrap())
|
||||
} else {
|
||||
merge_loaded_metrics(results).into_values().collect()
|
||||
}
|
||||
}
|
||||
|
||||
fn equal_matcher_filters(matchers: &Matchers) -> Vec<(String, Vec<String>)> {
|
||||
matchers
|
||||
.matchers
|
||||
.iter()
|
||||
.filter_map(|mat| {
|
||||
if mat.op == MatchOp::Equal {
|
||||
Some((mat.name.to_string(), vec![mat.value.to_string()]))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Merge series that may occur in more than one DataFusion context. Contexts
|
||||
/// can have different partition counts, so this fallback deliberately merges
|
||||
/// by series hash instead of zipping partitions.
|
||||
fn merge_loaded_metrics(results: Vec<LoadedMetrics>) -> HashMap<u64, RangeValue> {
|
||||
// Multiple contexts can contain the same high-cardinality series. Grow
|
||||
// from the unique keys instead of reserving the sum of all context sizes.
|
||||
let mut metrics: HashMap<u64, RangeValue> = HashMap::default();
|
||||
for result in results {
|
||||
let maps = match result {
|
||||
LoadedMetrics::Partitioned(partitions) => partitions,
|
||||
LoadedMetrics::Merged(metrics) => vec![metrics],
|
||||
};
|
||||
for map in maps {
|
||||
for (hash, value) in map {
|
||||
match metrics.entry(hash) {
|
||||
hashbrown::hash_map::Entry::Occupied(mut entry) => {
|
||||
entry.get_mut().extend(value)
|
||||
}
|
||||
hashbrown::hash_map::Entry::Vacant(entry) => {
|
||||
entry.insert(value);
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
metrics
|
||||
}
|
||||
|
||||
fn get_offset_modifier(offset: Option<Offset>) -> i64 {
|
||||
if let Some(offset) = offset {
|
||||
match offset {
|
||||
Offset::Pos(offset) => micros(offset),
|
||||
Offset::Neg(offset) => -micros(offset),
|
||||
}
|
||||
} else {
|
||||
0
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use promql_parser::{
|
||||
label::{MatchOp, Matchers},
|
||||
parser::{Offset, VectorSelector},
|
||||
};
|
||||
|
||||
use super::*;
|
||||
use crate::{engine::tests::*, exec::PromqlContext};
|
||||
|
||||
#[test]
|
||||
fn test_flatten_partitioned_metrics_preserves_all_series() {
|
||||
let partitions = vec![
|
||||
HashMap::from([(11, range_value(100, 1.0))]),
|
||||
HashMap::from([(22, range_value(200, 2.0))]),
|
||||
];
|
||||
|
||||
let metrics = flatten_partitioned_metrics(partitions);
|
||||
assert_eq!(metrics.len(), 2);
|
||||
assert_eq!(
|
||||
metrics
|
||||
.iter()
|
||||
.map(|metric| metric.samples.len())
|
||||
.sum::<usize>(),
|
||||
2
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_merge_partitioned_metrics_extends_series_across_contexts() {
|
||||
// Deliberately use different partition counts: separate DataFusion
|
||||
// contexts are not required to have identical physical plans.
|
||||
let results = vec![
|
||||
LoadedMetrics::Partitioned(vec![HashMap::from([(11, range_value(100, 1.0))])]),
|
||||
LoadedMetrics::Partitioned(vec![
|
||||
HashMap::new(),
|
||||
HashMap::from([(11, range_value(200, 2.0)), (22, range_value(300, 3.0))]),
|
||||
]),
|
||||
];
|
||||
|
||||
let metrics = merge_loaded_metrics(results);
|
||||
assert_eq!(metrics.len(), 2);
|
||||
assert_eq!(metrics[&11].samples.len(), 2);
|
||||
assert_eq!(metrics[&22].samples.len(), 1);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_eval_vector_selector_strips_consumed_name_matcher() {
|
||||
let trace_id = "test_trace";
|
||||
let captured = Arc::new(std::sync::Mutex::new(None));
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, "test_org", 30),
|
||||
MatcherCapturingProvider {
|
||||
captured: captured.clone(),
|
||||
},
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
// `{__name__="test_metric", env="prod"}` form: the name matcher is lifted
|
||||
// into selector.name and must NOT survive as a column filter — the stored
|
||||
// `__name__` column can hold the pre-`format_stream_name` metric name
|
||||
let selector = VectorSelector {
|
||||
name: None,
|
||||
matchers: Matchers {
|
||||
matchers: vec![
|
||||
promql_parser::label::Matcher {
|
||||
name: NAME_LABEL.to_string(),
|
||||
op: MatchOp::Equal,
|
||||
value: "test_metric".to_string(),
|
||||
},
|
||||
promql_parser::label::Matcher {
|
||||
name: "env".to_string(),
|
||||
op: MatchOp::Equal,
|
||||
value: "prod".to_string(),
|
||||
},
|
||||
],
|
||||
or_matchers: vec![],
|
||||
},
|
||||
offset: None,
|
||||
at: None,
|
||||
};
|
||||
|
||||
engine.eval_vector_selector(&selector).await.unwrap();
|
||||
|
||||
let matchers = captured.lock().unwrap().take().unwrap();
|
||||
assert!(matchers.matchers.iter().all(|m| m.name != NAME_LABEL));
|
||||
assert_eq!(matchers.matchers.len(), 1);
|
||||
assert_eq!(matchers.matchers[0].name, "env");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_eval_vector_selector_basic() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let matchers = Matchers {
|
||||
matchers: vec![promql_parser::label::Matcher {
|
||||
name: "env".to_string(),
|
||||
op: MatchOp::Equal,
|
||||
value: "prod".to_string(),
|
||||
}],
|
||||
or_matchers: vec![],
|
||||
};
|
||||
|
||||
let selector = VectorSelector {
|
||||
name: Some("test_metric".to_string()),
|
||||
matchers,
|
||||
offset: None,
|
||||
at: None,
|
||||
};
|
||||
|
||||
let result = engine.eval_vector_selector(&selector).await;
|
||||
assert!(result.is_ok());
|
||||
let values = result.unwrap();
|
||||
assert_eq!(values.len(), 0); // Mock provider returns empty data
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_eval_vector_selector_with_offset() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let matchers = Matchers {
|
||||
matchers: vec![promql_parser::label::Matcher {
|
||||
name: "env".to_string(),
|
||||
op: MatchOp::Equal,
|
||||
value: "prod".to_string(),
|
||||
}],
|
||||
or_matchers: vec![],
|
||||
};
|
||||
|
||||
let selector = VectorSelector {
|
||||
name: Some("test_metric".to_string()),
|
||||
matchers,
|
||||
offset: Some(Offset::Pos(Duration::from_secs(60))),
|
||||
at: None,
|
||||
};
|
||||
|
||||
let result = engine.eval_vector_selector(&selector).await;
|
||||
assert!(result.is_ok());
|
||||
let values = result.unwrap();
|
||||
assert_eq!(values.len(), 0); // Mock provider returns empty data
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_eval_vector_selector_with_negative_offset() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let matchers = Matchers {
|
||||
matchers: vec![promql_parser::label::Matcher {
|
||||
name: "env".to_string(),
|
||||
op: MatchOp::Equal,
|
||||
value: "prod".to_string(),
|
||||
}],
|
||||
or_matchers: vec![],
|
||||
};
|
||||
|
||||
let selector = VectorSelector {
|
||||
name: Some("test_metric".to_string()),
|
||||
matchers,
|
||||
offset: Some(Offset::Neg(Duration::from_secs(60))),
|
||||
at: None,
|
||||
};
|
||||
|
||||
let result = engine.eval_vector_selector(&selector).await;
|
||||
assert!(result.is_ok());
|
||||
let values = result.unwrap();
|
||||
assert_eq!(values.len(), 0); // Mock provider returns empty data
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_eval_matrix_selector_basic() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let matchers = Matchers {
|
||||
matchers: vec![promql_parser::label::Matcher {
|
||||
name: "env".to_string(),
|
||||
op: MatchOp::Equal,
|
||||
value: "prod".to_string(),
|
||||
}],
|
||||
or_matchers: vec![],
|
||||
};
|
||||
|
||||
let selector = VectorSelector {
|
||||
name: Some("test_metric".to_string()),
|
||||
matchers,
|
||||
offset: None,
|
||||
at: None,
|
||||
};
|
||||
|
||||
let result = engine
|
||||
.eval_matrix_selector(&selector, Duration::from_secs(300))
|
||||
.await;
|
||||
assert!(result.is_ok());
|
||||
let values = result.unwrap();
|
||||
assert_eq!(values.len(), 0); // Mock provider returns empty data
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_eval_matrix_selector_with_offset() {
|
||||
let trace_id = "test_trace";
|
||||
let org_id = "test_org";
|
||||
let mut engine = Engine::new(
|
||||
trace_id,
|
||||
Arc::new(PromqlContext::new(
|
||||
create_test_query_ctx(trace_id, org_id, 30),
|
||||
SimpleMockProvider,
|
||||
vec![],
|
||||
)),
|
||||
create_test_eval_ctx(),
|
||||
);
|
||||
|
||||
let matchers = Matchers {
|
||||
matchers: vec![promql_parser::label::Matcher {
|
||||
name: "env".to_string(),
|
||||
op: MatchOp::Equal,
|
||||
value: "prod".to_string(),
|
||||
}],
|
||||
or_matchers: vec![],
|
||||
};
|
||||
|
||||
let selector = VectorSelector {
|
||||
name: Some("test_metric".to_string()),
|
||||
matchers,
|
||||
offset: Some(Offset::Pos(Duration::from_secs(120))),
|
||||
at: None,
|
||||
};
|
||||
|
||||
let result = engine
|
||||
.eval_matrix_selector(&selector, Duration::from_secs(300))
|
||||
.await;
|
||||
assert!(result.is_ok());
|
||||
let values = result.unwrap();
|
||||
assert_eq!(values.len(), 0); // Mock provider returns empty data
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_offset_modifier_none() {
|
||||
assert_eq!(get_offset_modifier(None), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_offset_modifier_positive() {
|
||||
let result = get_offset_modifier(Some(Offset::Pos(Duration::from_secs(60))));
|
||||
assert_eq!(result, 60_000_000); // 60s in micros
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_get_offset_modifier_negative() {
|
||||
let result = get_offset_modifier(Some(Offset::Neg(Duration::from_secs(30))));
|
||||
assert_eq!(result, -30_000_000); // -30s in micros
|
||||
}
|
||||
}
|
||||
@@ -13,9 +13,11 @@
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
use std::time::Duration;
|
||||
|
||||
use config::meta::promql::{
|
||||
NAME_LABEL,
|
||||
value::{CounterSeries, EvalContext, RangeValue, Value},
|
||||
value::{CounterSeries, EvalContext, ExtrapolationKind, RangeValue, Sample, Value},
|
||||
};
|
||||
use datafusion::error::{DataFusionError, Result};
|
||||
use promql_parser::parser::LabelModifier;
|
||||
@@ -94,37 +96,15 @@ pub(crate) fn fused_range_agg(
|
||||
.as_ref()
|
||||
.expect("range function input must have a time window")
|
||||
.range;
|
||||
let range_micros = micros(range);
|
||||
let mut start_index = 0;
|
||||
let mut end_index = 0;
|
||||
let counter = CounterSeries::try_new(
|
||||
fold_series(
|
||||
&mut acc,
|
||||
&metric.samples,
|
||||
range,
|
||||
func,
|
||||
counter_kind,
|
||||
eval_ctx,
|
||||
range_micros,
|
||||
×tamps,
|
||||
);
|
||||
|
||||
for (slot, &eval_ts) in timestamps.iter().enumerate() {
|
||||
let window_samples = advance_sample_window(
|
||||
&metric.samples,
|
||||
eval_ts - range_micros,
|
||||
eval_ts,
|
||||
&mut start_index,
|
||||
&mut end_index,
|
||||
);
|
||||
if window_samples.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let value = match &counter {
|
||||
Some(counter) => {
|
||||
counter.extrapolate(start_index, end_index, eval_ts, range)
|
||||
}
|
||||
None => func.exec(window_samples, eval_ts, &range),
|
||||
};
|
||||
if let Some(value) = value {
|
||||
acc.push(slot, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
acc
|
||||
};
|
||||
@@ -174,6 +154,43 @@ pub(crate) fn fused_range_agg(
|
||||
Ok(Value::Matrix(results))
|
||||
}
|
||||
|
||||
/// Evaluates `func` over one series' time-ordered samples and pushes each
|
||||
/// produced value into `acc` at its evaluation slot.
|
||||
pub(super) fn fold_series(
|
||||
acc: &mut FusedAccumulator,
|
||||
samples: &[Sample],
|
||||
range: Duration,
|
||||
func: &dyn RangeFunc,
|
||||
counter_kind: Option<ExtrapolationKind>,
|
||||
eval_ctx: &EvalContext,
|
||||
timestamps: &[i64],
|
||||
) {
|
||||
let range_micros = micros(range);
|
||||
let mut start_index = 0;
|
||||
let mut end_index = 0;
|
||||
let counter = CounterSeries::try_new(samples, counter_kind, eval_ctx, range_micros);
|
||||
|
||||
for (slot, &eval_ts) in timestamps.iter().enumerate() {
|
||||
let window_samples = advance_sample_window(
|
||||
samples,
|
||||
eval_ts - range_micros,
|
||||
eval_ts,
|
||||
&mut start_index,
|
||||
&mut end_index,
|
||||
);
|
||||
if window_samples.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let value = match &counter {
|
||||
Some(counter) => counter.extrapolate(start_index, end_index, eval_ts, range),
|
||||
None => func.exec(window_samples, eval_ts, &range),
|
||||
};
|
||||
if let Some(value) = value {
|
||||
acc.push(slot, value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
@@ -121,7 +121,21 @@ enum LabelColumn<'a> {
|
||||
Utf8View(&'a StringViewArray),
|
||||
}
|
||||
|
||||
impl LabelColumn<'_> {
|
||||
impl<'a> LabelColumn<'a> {
|
||||
fn try_from_array(column: &'a dyn Array) -> Option<Self> {
|
||||
match column.data_type() {
|
||||
DataType::Utf8 => column
|
||||
.as_any()
|
||||
.downcast_ref::<StringArray>()
|
||||
.map(Self::Utf8),
|
||||
DataType::Utf8View => column
|
||||
.as_any()
|
||||
.downcast_ref::<StringViewArray>()
|
||||
.map(Self::Utf8View),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn is_null(&self, row: usize) -> bool {
|
||||
match self {
|
||||
Self::Utf8(values) => values.is_null(row),
|
||||
@@ -218,18 +232,7 @@ pub(super) async fn load_labels(
|
||||
let column = columns.get(*index).ok_or_else(|| {
|
||||
DataFusionError::Execution(format!("label column {name} is missing"))
|
||||
})?;
|
||||
match column.data_type() {
|
||||
DataType::Utf8 => column
|
||||
.as_any()
|
||||
.downcast_ref::<StringArray>()
|
||||
.map(LabelColumn::Utf8),
|
||||
DataType::Utf8View => column
|
||||
.as_any()
|
||||
.downcast_ref::<StringViewArray>()
|
||||
.map(LabelColumn::Utf8View),
|
||||
_ => None,
|
||||
}
|
||||
.ok_or_else(|| {
|
||||
LabelColumn::try_from_array(column.as_ref()).ok_or_else(|| {
|
||||
DataFusionError::Execution(format!(
|
||||
"label column {name} is not Utf8 or Utf8View"
|
||||
))
|
||||
|
||||
@@ -90,6 +90,52 @@ fn with_hash_label(labels: Labels, hash: u64, include_hash_label: bool) -> Label
|
||||
const OPTIMIZATION_STEP_LOOKBACK_MULTIPLIER: i64 = 5;
|
||||
const OPTIMIZATION_MAX_STEPS: i64 = 30;
|
||||
|
||||
/// Restricts `df` to the rows the evaluation can observe: per-step lookback
|
||||
/// windows when the steps are sparse enough, the contiguous
|
||||
/// `[start - lookback, end]` range otherwise.
|
||||
pub(crate) fn apply_time_window(
|
||||
df: DataFrame,
|
||||
start: i64,
|
||||
end: i64,
|
||||
step: i64,
|
||||
lookback: i64,
|
||||
) -> Result<DataFrame> {
|
||||
// Optimization: When step > lookback, we don't need to load all data in
|
||||
// [start-lookback, end] Instead, we only need to load data windows around
|
||||
// each evaluation point
|
||||
let use_optimization = start != end
|
||||
&& step > 0
|
||||
&& step >= lookback * OPTIMIZATION_STEP_LOOKBACK_MULTIPLIER
|
||||
&& (((end - start) / step) + 1) < OPTIMIZATION_MAX_STEPS;
|
||||
if use_optimization {
|
||||
let num_steps = ((end - start) / step) + 1;
|
||||
let eval_timestamps: Vec<i64> = (0..num_steps).map(|i| start + (step * i)).collect();
|
||||
|
||||
let mut conditions: Vec<Expr> = Vec::new();
|
||||
for &eval_ts in &eval_timestamps {
|
||||
let window_start = eval_ts - lookback;
|
||||
let window_end = eval_ts;
|
||||
|
||||
conditions.push(
|
||||
col(TIMESTAMP_COL_NAME)
|
||||
.gt_eq(lit(window_start))
|
||||
.and(col(TIMESTAMP_COL_NAME).lt_eq(lit(window_end))),
|
||||
);
|
||||
}
|
||||
|
||||
let filters = disjunction(conditions).unwrap();
|
||||
df.filter(filters)
|
||||
} else {
|
||||
// Need to include lookback window before start for the first evaluation point
|
||||
let query_start = start - lookback;
|
||||
df.filter(
|
||||
col(TIMESTAMP_COL_NAME)
|
||||
.gt_eq(lit(query_start))
|
||||
.and(col(TIMESTAMP_COL_NAME).lt_eq(lit(end))),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub(super) async fn selector_load_data_from_datafusion(
|
||||
query_ctx: Arc<QueryContext>,
|
||||
@@ -107,43 +153,7 @@ pub(super) async fn selector_load_data_from_datafusion(
|
||||
let table_name = selector.name.as_ref().unwrap();
|
||||
|
||||
let mut df_group = match ctx.table(table_name).await {
|
||||
Ok(v) => {
|
||||
// Optimization: When step > lookback, we don't need to load all data in
|
||||
// [start-lookback, end] Instead, we only need to load data windows around
|
||||
// each evaluation point
|
||||
let use_optimization = start != end
|
||||
&& step > 0
|
||||
&& step >= lookback * OPTIMIZATION_STEP_LOOKBACK_MULTIPLIER
|
||||
&& (((end - start) / step) + 1) < OPTIMIZATION_MAX_STEPS;
|
||||
if use_optimization {
|
||||
let num_steps = ((end - start) / step) + 1;
|
||||
let eval_timestamps: Vec<i64> =
|
||||
(0..num_steps).map(|i| start + (step * i)).collect();
|
||||
|
||||
let mut conditions: Vec<Expr> = Vec::new();
|
||||
for &eval_ts in &eval_timestamps {
|
||||
let window_start = eval_ts - lookback;
|
||||
let window_end = eval_ts;
|
||||
|
||||
conditions.push(
|
||||
col(TIMESTAMP_COL_NAME)
|
||||
.gt_eq(lit(window_start))
|
||||
.and(col(TIMESTAMP_COL_NAME).lt_eq(lit(window_end))),
|
||||
);
|
||||
}
|
||||
|
||||
let filters = disjunction(conditions).unwrap();
|
||||
v.filter(filters)?
|
||||
} else {
|
||||
// Need to include lookback window before start for the first evaluation point
|
||||
let query_start = start - lookback;
|
||||
v.filter(
|
||||
col(TIMESTAMP_COL_NAME)
|
||||
.gt_eq(lit(query_start))
|
||||
.and(col(TIMESTAMP_COL_NAME).lt_eq(lit(end))),
|
||||
)?
|
||||
}
|
||||
}
|
||||
Ok(v) => apply_time_window(v, start, end, step, lookback)?,
|
||||
Err(_) => {
|
||||
return Ok(LoadedMetrics::Partitioned(Vec::new()));
|
||||
}
|
||||
|
||||
+24
-9
@@ -20,7 +20,7 @@ use config::{
|
||||
meta::promql::{BUCKET_LABEL, HASH_LABEL, NAME_LABEL, VALUE_LABEL},
|
||||
};
|
||||
use datafusion::{
|
||||
arrow::datatypes::{DataType, Schema},
|
||||
arrow::datatypes::{DataType, Field, Schema},
|
||||
common::ScalarValue,
|
||||
error::Result,
|
||||
functions::regex::regexp_like,
|
||||
@@ -28,7 +28,28 @@ use datafusion::{
|
||||
prelude::{DataFrame, Expr, col, lit},
|
||||
};
|
||||
use hashbrown::HashSet;
|
||||
use promql_parser::label::{MatchOp, Matchers};
|
||||
use promql_parser::label::{MatchOp, Matcher, Matchers};
|
||||
|
||||
/// The schema field a residual matcher filters on; `None` when
|
||||
/// `matcher_predicates` skips the matcher entirely.
|
||||
pub fn matcher_residual_field<'a>(schema: &'a Schema, matcher: &Matcher) -> Option<&'a Field> {
|
||||
// `__name__` is consumed by stream selection; the stored column may hold the
|
||||
// pre-`format_stream_name` metric name (e.g. mixed case), so filtering on it
|
||||
// would drop every row of a stream that was already selected by name.
|
||||
if matcher.name == TIMESTAMP_COL_NAME
|
||||
|| matcher.name == VALUE_LABEL
|
||||
|| matcher.name == NAME_LABEL
|
||||
{
|
||||
return None;
|
||||
}
|
||||
schema.field_with_name(&matcher.name).ok()
|
||||
}
|
||||
|
||||
/// Whether `matcher_predicates` evaluates this matcher on the given schema.
|
||||
/// Index pruning may claim exactness only for matchers covered under this rule.
|
||||
pub fn matcher_is_residual(schema: &Schema, matcher: &Matcher) -> bool {
|
||||
matcher_residual_field(schema, matcher).is_some()
|
||||
}
|
||||
|
||||
/// Build the DataFusion predicates used for PromQL label matchers.
|
||||
///
|
||||
@@ -37,13 +58,7 @@ use promql_parser::label::{MatchOp, Matchers};
|
||||
pub fn matcher_predicates(schema: &Schema, matchers: &Matchers) -> Vec<Expr> {
|
||||
let mut predicates = Vec::new();
|
||||
for mat in matchers.matchers.iter() {
|
||||
// `__name__` is consumed by stream selection; the stored column may hold the
|
||||
// pre-`format_stream_name` metric name (e.g. mixed case), so filtering on it
|
||||
// would drop every row of a stream that was already selected by name.
|
||||
if mat.name == TIMESTAMP_COL_NAME || mat.name == VALUE_LABEL || mat.name == NAME_LABEL {
|
||||
continue;
|
||||
}
|
||||
let Ok(field) = schema.field_with_name(&mat.name) else {
|
||||
let Some(field) = matcher_residual_field(schema, mat) else {
|
||||
continue;
|
||||
};
|
||||
let field_type = field.data_type().clone();
|
||||
|
||||
Reference in New Issue
Block a user