rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-11-14 14:56:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(promql): run date functions only through map_math and refine the review points

Browse Source 
Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
This commit is contained in:
Dennis Zhuang
2026-10-22 03:14:53 +08:00
parent e06f5b3f58
commit 7217bf43fb
7 changed files with 47 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/promql/src/map.rs
+3 -3
View File
@@ -83,14 +83,14 @@ pub enum MathOp {
min: f64,
max: f64,
},
/// A date function of the value read as Unix seconds, such as `hour(v)`.
Date(DateFunction),
ClampMin {
min: f64,
},
ClampMax {
max: f64,
},
/// A date function of the value read as Unix seconds, such as `hour(v)`.
Date(DateFunction),
}
/// Float arithmetic between a vector and a scalar, as in Prometheus.
@@ -297,7 +297,6 @@ impl MathOp {
}
}
Self::Clamp { min, max } => go_max(min, go_min(max, v)),
Self::Date(function) => function.of_seconds(v),
Self::ClampMin { min } => go_max(min, v),
Self::ClampMax { max } => go_min(max, v),
Self::Abs => v.abs(),
@@ -305,6 +304,7 @@ impl MathOp {
Self::Floor => v.floor(),
Self::Sqrt => v.sqrt(),
Self::Round { nearest } => Round::round_to_nearest(v, nearest),
Self::Date(function) => function.of_seconds(v),
}
}
src/promql/src/map_time.rs
+4 -88
View File
@@ -15,17 +15,13 @@
//! `time()` and the date functions over evaluated series, following Prometheus `funcTime` and
//! `dateWrapper`.
use std::sync::Arc;
use datafusion::arrow::array::{Array, Float64Array};
use datafusion::arrow::array::Float64Array;
use datafusion::error::{DataFusionError, Result};
use datatypes::series::MetricName;
use datatypes::series::samples::{
DenseSamplesRef, SampleRef, Samples, Series, SeriesFields, SeriesSchema,
};
use datatypes::series::samples::{Samples, Series, SeriesFields, SeriesSchema};
use datatypes::series::time::{GridRange, TimeAxis};
use crate::grid_series::{GridSeriesBuilder, METADATA_LABELS, external};
use crate::grid_series::external;
#[cfg(test)]
mod tests;
@@ -86,86 +82,6 @@ pub fn date_of_eval_time(function: DateFunction, schema: &SeriesSchema) -> Resul
eval_time_series(schema, |ms| function.apply(ms / 1000))
}
/// A date function over `series`: each float value is read as Unix seconds and replaced by the
/// date component, and the metric name is dropped. Missing values stay missing; histogram
/// samples are dropped, so the result is `None` when no float remains.
///
/// A missing value is a hole like any other, but a hole only skips a position: the result is one
/// series over the runs of consecutive positions that hold a value.
///
/// Series that share a label set once their names are dropped must then go through
/// [`merge_same_labelset`](crate::grid_series::merge_same_labelset).
pub fn map_date(
function: DateFunction,
schema: &Arc<SeriesSchema>,
series: &Series,
) -> Result<Option<Series>> {
if series.time.grid().is_none() {
return Err(DataFusionError::Execution(
"date functions require an evaluated series".into(),
));
}
let mut metric = series.metric.clone();
metric.remove_all(&METADATA_LABELS);
let map = |seconds: f64| function.of_seconds(seconds);
// The fast path maps every float in place, so it needs each field to be entirely present: it
// keeps the input axis as it is. A hole means the mapped rows are fewer than the axis rows,
// so it falls through to the builder, which rebuilds the axis from the rows that survive.
let all_floats = series.fields.iter().all(|samples| {
samples.segments().iter().all(|segment| {
matches!(segment.dense(), Some(DenseSamplesRef::Float(floats)) if floats.null_count() == 0)
})
});
if all_floats {
let fields = series
.fields
.iter()
.map(|samples| {
let segments = samples
.segments()
.iter()
.map(|segment| {
let Some(DenseSamplesRef::Float(floats)) = segment.dense() else {
unreachable!("checked above");
};
Samples::from_float(floats.unary(map))
})
.collect();
Samples::try_from_chunks(segments).map_err(external)
})
.collect::<Result<Vec<_>>>()?;
let fields = SeriesFields::try_from_vec(fields).map_err(external)?;
let series =
Series::try_new(schema, metric, series.time.clone(), fields).map_err(external)?;
return Ok(Some(series));
}
// Holes in any field drop rows, so the builder names the positions that survive.
let mut builder = GridSeriesBuilder::try_new(schema.clone())?;
let mut values = vec![None; series.fields.len()];
// The axis names the positions to evaluate and a cursor per field reads them in order, so an
// axis of several runs costs its rows and not a scan per row.
let mut cursors = series
.fields
.iter()
.map(|samples| samples.cursor())
.collect::<Vec<_>>();
for (position, row) in series.time.positions() {
for (value, samples) in values.iter_mut().zip(&mut cursors) {
*value = match samples
.at(row)
.expect("the axis and the fields are aligned")
{
SampleRef::Float(seconds) => Some(map(seconds)),
SampleRef::Missing | SampleRef::Histogram(_) => None,
};
}
builder.push(position, &values)?;
}
builder.finish(metric)
}
fn eval_time_series(schema: &SeriesSchema, value: impl Fn(i64) -> f64) -> Result<Series> {
let grid = schema.eval_grid.clone().ok_or_else(|| {
DataFusionError::Execution("time functions require an evaluation grid".into())
@@ -200,7 +116,7 @@ const UNIX_TO_ABSOLUTE: i64 = (1969 * 365 + 1969 / 4 - 1969 / 100 + 1969 / 400)
impl DateFunction {
/// The PromQL name of the function.
pub fn name(self) -> &'static str {
pub(crate) fn name(self) -> &'static str {
match self {
Self::DayOfMonth => "day_of_month",
Self::DayOfWeek => "day_of_week",
src/promql/src/map_time/tests.rs
+27 -22
View File
@@ -24,12 +24,13 @@ use datatypes::series::samples::{
};
use datatypes::series::time::{EvalGrid, GridRange, GridSpan, GridSpec, TimeAxis};
use super::{DateFunction, QueryContextFunction, date_of_eval_time, map_date, time};
use super::{DateFunction, QueryContextFunction, date_of_eval_time, time};
use crate::grid_series::merge_same_labelset;
use crate::grid_series::reference::{
Labels, Points, canonical_bits, golden_cases, labels_of, metric_of, points_by_labels,
points_of, run_ranges, series_of, sparse_series,
};
use crate::map::{MathOp, map_math};
fn schema(grid: Arc<EvalGrid>) -> Arc<SeriesSchema> {
Arc::new(SeriesSchema {
@@ -71,7 +72,8 @@ fn time_functions_match_prometheus() {
.iter()
.filter_map(|(labels, points)| series_of(&schema, labels, points))
{
actual.extend(map_date(function(name), &schema, &series).unwrap());
actual
.extend(map_math(&schema, &series, MathOp::Date(function(name))).unwrap());
}
actual
}
@@ -107,7 +109,7 @@ fn time_functions_match_prometheus() {
/// Histogram samples are dropped and the series trimmed to the remaining floats; chunked floats
/// are mapped chunk by chunk.
#[test]
fn map_date_drops_histograms_and_maps_chunks() {
fn date_functions_drop_histograms_and_map_chunks() {
let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1000, 4).unwrap()).unwrap());
let schema = schema(grid.clone());
let range = GridRange::try_new(grid, 1, 3).unwrap();
@@ -155,7 +157,7 @@ fn map_date_drops_histograms_and_maps_chunks() {
Samples::try_from_histogram(UInt32Array::from(vec![0]), histograms.finish().unwrap())
.unwrap();
let mixed = Samples::try_from_chunks(vec![histogram, floats(&[3600.0, 7200.0])]).unwrap();
let hours = map_date(DateFunction::Hour, &schema, &series(mixed))
let hours = map_math(&schema, &series(mixed), MathOp::Date(DateFunction::Hour))
.unwrap()
.expect("a mixed series with one float");
assert_eq!(
@@ -168,9 +170,13 @@ fn map_date_drops_histograms_and_maps_chunks() {
);
let chunked = Samples::try_from_chunks(vec![floats(&[60.0]), floats(&[120.0, 180.0])]).unwrap();
let minutes = map_date(DateFunction::Minute, &schema, &series(chunked))
.unwrap()
.expect("a chunked float series");
let minutes = map_math(
&schema,
&series(chunked),
MathOp::Date(DateFunction::Minute),
)
.unwrap()
.expect("a chunked float series");
assert_eq!(
points_of(&minutes),
BTreeMap::from([
@@ -181,13 +187,13 @@ fn map_date_drops_histograms_and_maps_chunks() {
);
}
/// `map_date` reads a series of several runs by the positions its axis carries.
/// A date function reads a series of several runs by the positions its axis carries.
///
/// A sparse input cannot be mapped through the in-place fast path when a field holds a hole, so
/// the piece-cutting path runs: the result must report the right positions, not an offset from
/// the first one.
#[test]
fn map_date_reads_a_sparse_axis_by_its_positions() {
fn date_functions_read_a_sparse_axis_by_its_positions() {
let grid = Arc::new(EvalGrid::try_new(GridSpec::try_new(0, 1000, 8).unwrap()).unwrap());
let schema = schema(grid.clone());
// Positions 1, 2 and 6 hold seconds; the axis skips 3, 4 and 5.
@@ -205,7 +211,7 @@ fn map_date_reads_a_sparse_axis_by_its_positions() {
assert_eq!(input.time.first_position(), Some(1));
assert_eq!(input.time.last_position(), Some(6));
let output = map_date(DateFunction::Minute, &schema, &input)
let output = map_math(&schema, &input, MathOp::Date(DateFunction::Minute))
.unwrap()
.expect("three rows");
// No field holds a hole, so the fast path keeps the axis as it is: one series, three rows.
@@ -227,8 +233,8 @@ fn map_date_reads_a_sparse_axis_by_its_positions() {
// `minute()` of Unix seconds 1000, 2000 and 6000.
assert_eq!(read, vec![(1, 16.0), (2, 33.0), (6, 40.0)]);
// A hole in the field keeps the fast path from applying, so the axis is rebuilt from the rows
// that survive: still one series, now over the runs of consecutive valued positions.
// A hole in a float field stays a missing row on the input's axis, as with every map, and the
// positions around it are mapped.
let mut holed = Points::new();
for (position, seconds) in [(2, 120), (3, 180), (4, 240), (6, 360), (7, 420)] {
holed.insert(position, f64::to_bits(seconds as f64));
@@ -239,10 +245,10 @@ fn map_date_reads_a_sparse_axis_by_its_positions() {
&holed,
)
.unwrap();
let output = map_date(DateFunction::Minute, &schema, &input)
let output = map_math(&schema, &input, MathOp::Date(DateFunction::Minute))
.unwrap()
.expect("five rows");
assert_eq!(run_ranges(&output), vec![(2, 3), (6, 2)]);
assert_eq!(run_ranges(&output), run_ranges(&input));
assert_eq!(
points_of(&output),
BTreeMap::from([
@@ -254,10 +260,10 @@ fn map_date_reads_a_sparse_axis_by_its_positions() {
])
);
// A sparse axis *and* a hole in the field, with the samples split over two chunks: the fast
// path cannot apply, so the piece-cutting path must walk the axis' positions and rebuild the
// axis from the rows that survive. Positions 0, 1, 3, 4 and 5 hold rows and position 3 is the
// hole, so a walk that counted rows from the first position would map the wrong seconds.
// A sparse axis *and* a hole in the field, with the samples split over two chunks: each chunk
// is mapped on the input's axis, the hole stays a missing row, and the values land on the
// positions the axis names. Positions 0, 1, 3, 4 and 5 hold rows and position 3 is the hole,
// so a walk that counted rows from the first position would map the wrong seconds.
let time =
TimeAxis::from_grid_spans(grid, vec![GridSpan::new(0, 2), GridSpan::new(3, 3)]).unwrap();
let chunked = Samples::try_from_chunks(vec![
@@ -273,11 +279,10 @@ fn map_date_reads_a_sparse_axis_by_its_positions() {
)
.unwrap();
assert_eq!(input.time.len(), 5);
let output = map_date(DateFunction::Minute, &schema, &input)
let output = map_math(&schema, &input, MathOp::Date(DateFunction::Minute))
.unwrap()
.expect("four rows");
assert_eq!(output.time.len(), 4);
assert_eq!(run_ranges(&output), vec![(0, 2), (4, 2)]);
.expect("five rows");
assert_eq!(run_ranges(&output), run_ranges(&input));
// `minute()` of Unix seconds 60, 120, 240 and 300.
assert_eq!(
points_of(&output),
src/promql/src/native/execute.rs
+5 -4
View File
@@ -3406,9 +3406,10 @@ impl NativeExecutor {
/// Maps one chunk with `ops`, keeping its axis and its closing.
///
/// Mapping is point-wise over valued positions, so a chunk whose values survive keeps exactly the
/// rows it came with: a position no field fills is not a row of any chunk, and mapping does not
/// invent one. `None` means the operation emptied every position, which yields no chunk at all
/// Mapping is point-wise over valued positions: a float chunk keeps the rows it came with, and a
/// chunk holding native histograms loses the rows only histograms filled, as math skips them; a
/// position no field fills is not a row of any chunk, and mapping does not invent one. `None` means
/// the operation emptied every position, which yields no chunk at all
/// unless this chunk closed its series: a closing chunk is still handed over, holding no position,
/// because a consumer tells "nothing more here" from "no more input" by that chunk alone.
fn map_chunk(
@@ -3421,7 +3422,7 @@ fn map_chunk(
series = match map_math(schema, &series, op)? {
Some(mapped) => mapped,
// `clamp` with `max < min` is the empty vector, and math skips native histograms, so a
// series of histograms only has no value left: every position it held was decided and
// series without a float value has none left: every position it held was decided and
// holds no value now, so there is no row to hand over.
None => {
return if chunk.is_last {
src/promql/src/native/execute/tests.rs
-34
View File
@@ -6417,37 +6417,3 @@ fn an_all_missing_match_group_fits_its_pool() {
assert_eq!(pool.reserved(), 0);
assert_eq!(execution.state().reserved_bytes(), 0);
}
/// A date function runs as a map over evaluated input: each value is read as Unix seconds and
/// replaced by its component, and the metric name is dropped.
#[test]
fn a_date_map_reads_values_as_unix_seconds() {
let grid = grid(0, 60_000, 3);
let schema = evaluated_schema(Arc::clone(&grid));
let tree = NativeTree::try_new(map(
evaluated_input(Arc::clone(&grid)),
vec![MathOp::Date(crate::map_time::DateFunction::Hour)],
))
.unwrap();
let mut execution = executor(&tree, pool(), &ExecutionPlanMetricsSet::new());
// 05:00:30 and 23:59:59 on 1970-01-01, then 1970-01-02 00:00:00.
execution
.push_evaluated(evaluated_chunk_of(
&schema,
named_metric("a", "x"),
0,
vec![Some(18_030.0), Some(86_399.0), Some(86_400.0)],
true,
))
.unwrap();
let chunks = run(&mut execution);
assert_eq!(
decided(&chunks),
[(0, Some(5.0)), (1, Some(23.0)), (2, Some(0.0))]
);
assert!(
chunks
.iter()
.all(|chunk| chunk.series.metric.name().is_none())
);
}
src/query/src/promql/native_compile.rs
+7 -1
View File
@@ -574,7 +574,7 @@ fn math_signature(name: &str) -> Option<MathSignature> {
"rad" => Some(MathSignature::Unary(MathOp::Rad)),
"sgn" => Some(MathSignature::Unary(MathOp::Sgn)),
// A date function takes the vector its values are read from; without one it reads the
// evaluation time, which has no input to map and is refused here.
// evaluation time, which has no input to map and is refused by `math_op`.
"day_of_month" => Some(MathSignature::Unary(MathOp::Date(DateFunction::DayOfMonth))),
"day_of_week" => Some(MathSignature::Unary(MathOp::Date(DateFunction::DayOfWeek))),
"day_of_year" => Some(MathSignature::Unary(MathOp::Date(DateFunction::DayOfYear))),
@@ -622,6 +622,12 @@ fn math_op<'a>(
// The vector is the function's first argument, so a call that wrote no argument at all is
// refused by the same arity check as one that wrote too few constants.
let Some((vector, parameters)) = arguments.split_first() else {
if let MathSignature::Unary(MathOp::Date(_)) = signature {
return Err(not_implemented(&format!(
"a call to `{name}` without an argument, which reads the evaluation time and has \
no native kernel"
)));
}
return Err(arity_error(name, signature, 0));
};
let given = parameters.len() + 1;
src/query/src/promql/native_compile/tests.rs
+1 -5
View File
@@ -429,11 +429,7 @@ fn a_call_with_the_wrong_number_of_arguments_names_the_count_it_takes() {
),
// A date function written without an argument is valid PromQL that reads the evaluation
// time; there is no input to map, so it is left to the ordinary path.
(
eval_stmt("hour()"),
"hour",
"with 0 argument(s), where it takes its vector",
),
(eval_stmt("hour()"), "hour", "reads the evaluation time"),
];
for (stmt, name, expected) in cases {