frontend/instance/
exemplar.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::{BTreeMap, BTreeSet};
16use std::sync::Arc;
17
18use auth::{PermissionTableTarget, PermissionTableTargets};
19use common_base::cancellation::CancellableFuture;
20use common_catalog::consts::METRICS_EXEMPLARS_TABLE;
21use common_error::ext::BoxedError;
22use common_query::OutputData;
23use common_query::prelude::{
24    EXEMPLAR_LABELS_COLUMN, EXEMPLAR_METRIC_COLUMN, greptime_timestamp, greptime_value,
25};
26use common_query::prometheus::format_prometheus_float;
27use common_recordbatch::RecordBatch;
28use common_time::{Timestamp, timestamp::TimeUnit};
29use datafusion::common::{Column, ScalarValue};
30use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation};
31use datafusion_expr::{BinaryExpr, Expr, Operator, lit};
32use datatypes::arrow::datatypes::DataType;
33use datatypes::data_type::ConcreteDataType;
34use datatypes::value::Value;
35use datatypes::vectors::json::array::JsonArray;
36use datatypes::vectors::{Helper, VectorRef};
37use futures::StreamExt;
38use promql_parser::label::{METRIC_NAME, MatchOp, Matcher};
39use servers::error::{self, Result};
40use servers::exemplar::query::{Observation, Query, Selector, Series};
41use servers::interceptor::{PromQueryInterceptor, PromQueryInterceptorRef};
42use servers::prometheus_handler::PrometheusHandler;
43use session::context::QueryContextRef;
44use snafu::{OptionExt, ResultExt, ensure};
45use store_api::metric_engine_consts::{
46    DATA_SCHEMA_TABLE_ID_COLUMN_NAME, DATA_SCHEMA_TSID_COLUMN_NAME,
47};
48use table::metadata::TableInfo;
49
50use crate::instance::Instance;
51
52type Labels = BTreeMap<String, String>;
53// Retain NaN payloads and signed zero as distinct observations, using their stored bits.
54type Observations = BTreeMap<Labels, BTreeMap<(i64, u64, Labels), Observation>>;
55
56impl Instance {
57    pub(super) async fn query_exemplars_inner(
58        &self,
59        query: Query,
60        ctx: QueryContextRef,
61    ) -> Result<Vec<Series>> {
62        self.check_prom_query_privilege(&ctx)?;
63        let interceptor = self.plugins.get::<PromQueryInterceptorRef<error::Error>>();
64        let mut groups = BTreeMap::<String, Vec<&Selector>>::new();
65        for selector in &query.selectors {
66            groups
67                .entry(selector.schema.clone())
68                .or_default()
69                .push(selector);
70        }
71        // Authorizes every branch before any exemplar scan, including absent destinations.
72        for (schema, selectors) in &groups {
73            let local = ctx.fork();
74            local.set_current_schema(schema);
75            let local = Arc::new(local);
76            let mut targets = vec![PermissionTableTarget::new(
77                ctx.current_catalog(),
78                schema,
79                METRICS_EXEMPLARS_TABLE,
80            )];
81            for selector in selectors {
82                let sources = selector
83                    .matchers
84                    .iter()
85                    .filter(|matcher| {
86                        matcher.name == METRIC_NAME
87                            && matches!(matcher.op, MatchOp::Equal)
88                            && !matcher.value.is_empty()
89                    })
90                    .map(|matcher| &matcher.value)
91                    .collect::<BTreeSet<_>>();
92                if sources.is_empty() {
93                    self.check_query_target_permission(PermissionTableTargets::Unresolved, &local)
94                        .await?;
95                }
96                targets.extend(sources.into_iter().map(|source| {
97                    PermissionTableTarget::new(ctx.current_catalog(), schema, source)
98                }));
99            }
100            self.check_query_target_permission(PermissionTableTargets::resolved(targets), &local)
101                .await?;
102        }
103        let ticket = self.process_manager.register_query(
104            ctx.current_catalog().to_string(),
105            groups.keys().cloned().collect(),
106            format!("query_exemplars({})", query.expression),
107            ctx.conn_info().to_string(),
108            Some(ctx.process_id()),
109            None,
110        );
111        let cancellation = ticket.cancellation_handle.clone();
112        let future = async {
113            let pool = self
114                .query_engine
115                .engine_state()
116                .session_state()
117                .runtime_env()
118                .memory_pool
119                .clone();
120            let mut reservation =
121                MemoryConsumer::new("Prometheus exemplar response").register(&pool);
122            let mut observations = Observations::new();
123            let mut skipped = 0;
124            for (schema, selectors) in groups {
125                let local = ctx.fork();
126                local.set_current_schema(&schema);
127                let local = Arc::new(local);
128                let Some(table) = self
129                    .catalog_manager
130                    .table(
131                        ctx.current_catalog(),
132                        &schema,
133                        METRICS_EXEMPLARS_TABLE,
134                        Some(&local),
135                    )
136                    .await
137                    .map_err(BoxedError::new)
138                    .context(error::ExecuteQuerySnafu)?
139                else {
140                    continue;
141                };
142                let info = table.table_info();
143                operator::insert::validate_exemplar_table(&info)
144                    .map_err(BoxedError::new)
145                    .context(error::ExecuteQuerySnafu)?;
146                let filter = filter(&info, &selectors, query.bounds)?;
147                let frame = self
148                    .query_engine
149                    .read_table(table)
150                    .map_err(BoxedError::new)
151                    .context(error::ExecuteQuerySnafu)?;
152                let plan = frame
153                    .filter(filter)
154                    .context(error::DataFusionSnafu)?
155                    .into_unoptimized_plan();
156                interceptor.pre_execute_exemplars(&query, info.table_ref(), local.clone())?;
157                let output = self
158                    .statement_executor
159                    .exec_plan(plan, local)
160                    .await
161                    .map_err(BoxedError::new)
162                    .context(error::ExecuteQuerySnafu)?;
163                match output.data {
164                    OutputData::Stream(mut stream) => {
165                        while let Some(batch) = stream.next().await {
166                            skipped += collect(
167                                batch.context(error::CollectRecordbatchSnafu)?,
168                                &info,
169                                &mut observations,
170                                &mut reservation,
171                            )?;
172                            tokio::task::yield_now().await;
173                        }
174                    }
175                    OutputData::RecordBatches(batches) => {
176                        for batch in batches.take() {
177                            skipped += collect(batch, &info, &mut observations, &mut reservation)?;
178                            tokio::task::yield_now().await;
179                        }
180                    }
181                    _ => {
182                        return error::UnexpectedResultSnafu {
183                            reason: "exemplar scan did not return rows",
184                        }
185                        .fail();
186                    }
187                }
188            }
189            if skipped > 0 {
190                let warning = format!("Skipped {skipped} malformed stored exemplars");
191                ctx.set_warning(match ctx.warning() {
192                    Some(existing) => format!("{existing}; {warning}"),
193                    None => warning,
194                });
195            }
196            Ok(observations
197                .into_iter()
198                .map(|(series_labels, observations)| Series {
199                    series_labels,
200                    exemplars: observations.into_values().collect(),
201                })
202                .collect())
203        };
204        let result = CancellableFuture::new(future, cancellation)
205            .await
206            .map_err(|_| error::CancelledSnafu.build())?;
207        drop(ticket);
208        result
209    }
210}
211
212fn column(name: &str) -> Expr {
213    Expr::Column(Column::from_name(name))
214}
215
216fn filter(info: &TableInfo, selectors: &[&Selector], bounds: (i64, i64)) -> Result<Expr> {
217    let schema = &info.meta.schema;
218    let tags = info
219        .meta
220        .primary_key_indices
221        .iter()
222        .map(|index| schema.column_schemas()[*index].name.as_str())
223        .filter(|name| {
224            ![
225                DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
226                DATA_SCHEMA_TSID_COLUMN_NAME,
227            ]
228            .contains(name)
229        })
230        .collect::<BTreeSet<_>>();
231    let labels = selectors
232        .iter()
233        .map(|selector| {
234            selector
235                .matchers
236                .iter()
237                .map(|matcher| matcher_filter(matcher, &tags))
238                .fold(lit(true), Expr::and)
239        })
240        .fold(lit(false), Expr::or);
241    let unit = schema
242        .timestamp_column()
243        .and_then(|column| column.data_type.as_timestamp())
244        .map(|t| t.unit())
245        .context(error::UnexpectedResultSnafu {
246            reason: "exemplar time index is missing",
247        })?;
248    let start = Timestamp::new_nanosecond(bounds.0)
249        .convert_to_ceil(unit)
250        .context(error::InvalidQuerySnafu {
251            reason: "exemplar start exceeds timestamp range",
252        })?;
253    let end = Timestamp::new_nanosecond(bounds.1)
254        .convert_to(unit)
255        .context(error::InvalidQuerySnafu {
256            reason: "exemplar end exceeds timestamp range",
257        })?;
258    let literal = |timestamp: Timestamp| {
259        Expr::Literal(
260            match unit {
261                TimeUnit::Second => ScalarValue::TimestampSecond(Some(timestamp.value()), None),
262                TimeUnit::Millisecond => {
263                    ScalarValue::TimestampMillisecond(Some(timestamp.value()), None)
264                }
265                TimeUnit::Microsecond => {
266                    ScalarValue::TimestampMicrosecond(Some(timestamp.value()), None)
267                }
268                TimeUnit::Nanosecond => {
269                    ScalarValue::TimestampNanosecond(Some(timestamp.value()), None)
270                }
271            },
272            None,
273        )
274    };
275    Ok(labels
276        .and(column(greptime_timestamp()).gt_eq(literal(start)))
277        .and(column(greptime_timestamp()).lt_eq(literal(end))))
278}
279
280fn matcher_filter(matcher: &Matcher, tags: &BTreeSet<&str>) -> Expr {
281    let name = if matcher.name == METRIC_NAME {
282        EXEMPLAR_METRIC_COLUMN
283    } else {
284        matcher.name.as_str()
285    };
286    // Storage fields and the physical metric column are not Prometheus series labels.
287    if (matcher.name == EXEMPLAR_METRIC_COLUMN || !tags.contains(name))
288        && matcher.name != METRIC_NAME
289    {
290        return lit(matcher.is_match(""));
291    }
292    let col = column(name);
293    let condition = match &matcher.op {
294        MatchOp::Equal => col.clone().eq(lit(matcher.value.clone())),
295        MatchOp::NotEqual => col.clone().not_eq(lit(matcher.value.clone())),
296        MatchOp::Re(regex) | MatchOp::NotRe(regex) => Expr::BinaryExpr(BinaryExpr {
297            left: Box::new(col.clone()),
298            op: if matches!(matcher.op, MatchOp::Re(_)) {
299                Operator::RegexMatch
300            } else {
301                Operator::RegexNotMatch
302            },
303            right: Box::new(lit(regex.as_str().to_string())),
304        }),
305    };
306    if matcher.is_match("") {
307        col.is_null().or(condition)
308    } else {
309        condition
310    }
311}
312
313fn collect(
314    batch: RecordBatch,
315    info: &TableInfo,
316    observations: &mut Observations,
317    reservation: &mut MemoryReservation,
318) -> Result<usize> {
319    let memory = batch.buffer_memory_size();
320    reservation
321        .try_grow(memory)
322        .context(error::DataFusionSnafu)?;
323    let schema = batch.df_record_batch().schema();
324    let index = |name: &str| schema.index_of(name).context(error::ArrowSnafu);
325    let metric = Helper::try_into_vector(batch.column(index(EXEMPLAR_METRIC_COLUMN)?))
326        .map_err(BoxedError::new)
327        .context(error::ExecuteQuerySnafu)?;
328    let time = Helper::try_into_vector(batch.column(index(greptime_timestamp())?))
329        .map_err(BoxedError::new)
330        .context(error::ExecuteQuerySnafu)?;
331    let value = Helper::try_into_vector(batch.column(index(greptime_value())?))
332        .map_err(BoxedError::new)
333        .context(error::ExecuteQuerySnafu)?;
334    let value_type = |vector: &VectorRef| {
335        let mut data_type = vector.data_type();
336        // Storage may dictionary-encode tags; Vector::get returns their decoded values.
337        while let ConcreteDataType::Dictionary(dictionary) = data_type {
338            data_type = dictionary.value_type().clone();
339        }
340        data_type
341    };
342    ensure!(
343        value_type(&metric).is_string()
344            && value_type(&time).is_timestamp()
345            && value_type(&value) == ConcreteDataType::float64_datatype(),
346        error::UnexpectedResultSnafu {
347            reason: "exemplar scan returned incompatible column types",
348        }
349    );
350    let label_index = index(EXEMPLAR_LABELS_COLUMN)?;
351    let labels = JsonArray::from(batch.column(label_index))
352        .project_to_v2(schema.field(label_index), &DataType::Binary)
353        .map_err(BoxedError::new)
354        .context(error::ExecuteQuerySnafu)?;
355    let mut tags = Vec::<(&str, VectorRef)>::new();
356    for tag in &info.meta.primary_key_indices {
357        let name = info.meta.schema.column_schemas()[*tag].name.as_str();
358        if ![
359            EXEMPLAR_METRIC_COLUMN,
360            DATA_SCHEMA_TABLE_ID_COLUMN_NAME,
361            DATA_SCHEMA_TSID_COLUMN_NAME,
362        ]
363        .contains(&name)
364        {
365            let values = Helper::try_into_vector(batch.column(index(name)?))
366                .map_err(BoxedError::new)
367                .context(error::ExecuteQuerySnafu)?;
368            ensure!(
369                value_type(&values).is_string(),
370                error::UnexpectedResultSnafu {
371                    reason: "exemplar scan returned a non-string tag",
372                }
373            );
374            tags.push((name, values));
375        }
376    }
377    let mut skipped = 0;
378    'rows: for row in 0..batch.num_rows() {
379        let Value::String(metric) = metric.get(row) else {
380            skipped += 1;
381            continue;
382        };
383        let Value::Timestamp(timestamp) = time.get(row) else {
384            skipped += 1;
385            continue;
386        };
387        let Some(timestamp) = timestamp.convert_to(TimeUnit::Nanosecond) else {
388            skipped += 1;
389            continue;
390        };
391        let timestamp = timestamp.value();
392        let Value::Float64(value) = value.get(row) else {
393            skipped += 1;
394            continue;
395        };
396        let mut series = BTreeMap::from([(METRIC_NAME.to_string(), metric.into_string())]);
397        for (name, values) in &tags {
398            if let Value::String(value) = values.get(row) {
399                series.insert((*name).to_string(), value.into_string());
400            }
401        }
402        let labels = JsonArray::from(&labels)
403            .try_get_value(row)
404            .map_err(BoxedError::new)
405            .context(error::ExecuteQuerySnafu)?;
406        let serde_json::Value::Object(labels) = labels else {
407            skipped += 1;
408            continue;
409        };
410        let mut string_labels = Labels::new();
411        for (name, value) in labels {
412            let serde_json::Value::String(value) = value else {
413                skipped += 1;
414                continue 'rows;
415            };
416            string_labels.insert(name, value);
417        }
418        let labels = string_labels;
419        let key = (timestamp, value.0.to_bits(), labels.clone());
420        if observations
421            .get(&series)
422            .is_some_and(|items| items.contains_key(&key))
423        {
424            continue;
425        }
426        let bytes = 256
427            + series
428                .iter()
429                .chain(&labels)
430                .map(|(name, value)| name.len() + value.len() + 128)
431                .sum::<usize>()
432                * 2;
433        reservation
434            .try_grow(bytes)
435            .context(error::DataFusionSnafu)?;
436        observations.entry(series).or_default().insert(
437            key,
438            Observation {
439                labels,
440                value: format_prometheus_float(value.0),
441                timestamp: timestamp.div_euclid(1_000_000) as f64 / 1000.0,
442            },
443        );
444    }
445    reservation.shrink(memory);
446    Ok(skipped)
447}
448
449#[cfg(test)]
450mod tests {
451    use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool};
452    use datatypes::prelude::ConcreteDataType;
453    use datatypes::schema::{ColumnSchema, Schema};
454    use datatypes::vectors::Float64Vector;
455    use table::metadata::{TableInfoBuilder, TableMeta};
456
457    use super::*;
458
459    #[test]
460    fn test_exemplar_scan_batch_respects_query_memory_budget() {
461        let schema = Arc::new(Schema::new(vec![ColumnSchema::new(
462            greptime_value(),
463            ConcreteDataType::float64_datatype(),
464            false,
465        )]));
466        let batch = RecordBatch::new(
467            schema,
468            vec![Arc::new(Float64Vector::from_vec(vec![1.0; 128])) as VectorRef],
469        )
470        .unwrap();
471        let info = TableInfoBuilder::new(METRICS_EXEMPLARS_TABLE, TableMeta::empty())
472            .build()
473            .unwrap();
474        let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(1));
475        let mut reservation = MemoryConsumer::new("exemplar test").register(&pool);
476        let mut observations = Observations::new();
477        let error = collect(batch, &info, &mut observations, &mut reservation).unwrap_err();
478        assert!(matches!(
479            error,
480            error::Error::DataFusion {
481                error: datafusion::error::DataFusionError::ResourcesExhausted(_),
482                ..
483            }
484        ));
485        assert!(observations.is_empty());
486        assert_eq!(pool.reserved(), 0);
487    }
488}
frontend/instance/exemplar.rs

frontend/instance/
exemplar.rs
