feat(query): answer planned range queries one evaluation per step

README.md

@@ -53,8 +53,16 @@ sample minus `--allowed-lateness-ms`), and later evaluations reuse it; a pane
 read before it closes is computed again by later evaluations. Kept panes are
 evicted like stored ones, measured from the newest evaluation. Anything else is forwarded
 to Prometheus (`x-asap-execution: forwarded`); a planned query that could not
-be answered also says why in `x-asap-forward-reason`. Range queries, SQL and
-ClickHouse are not supported yet.
+be answered also says why in `x-asap-forward-reason`.
+
+A range query (`/api/v1/query_range`) whose text is one of the plan's queries
+is answered as one evaluation per step, `start + k * step` through `end`,
+and returns a Prometheus matrix. Each query-time raw range is read from
+Prometheus once for all steps and sliced per step; stored and kept panes are
+read per step, so a kept pane built at one step is reused by the later ones.
+If any step cannot be answered, for example because a step is off the pane
+grid or reads an evicted pane, the whole range query is forwarded with the
+reason. SQL and ClickHouse are not supported yet.
 
 ## Running
 

data_plane/Cargo.toml

@@ -22,6 +22,7 @@ prost = "0.13"
 snap = "1.1"
 futures = "0.3"
 form_urlencoded = "1"
+chrono = { version = "0.4", default-features = false, features = ["alloc"] }
 
 [dev-dependencies]
 tower = { version = "0.5", features = ["util"] }

data_plane/src/engine.rs

@@ -32,6 +32,17 @@ pub enum Command {
         raw: BTreeMap<u64, Vec<RawSample>>,
         reply: Reply<Vec<ResultSample>>,
     },
+    RawRangesForSteps {
+        query: String,
+        steps: Vec<i64>,
+        reply: Reply<(usize, Vec<RawRange>)>,
+    },
+    EvaluateSteps {
+        root: usize,
+        steps: Vec<i64>,
+        raw: BTreeMap<u64, Vec<RawSample>>,
+        reply: Reply<Vec<Vec<ResultSample>>>,
+    },
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -104,6 +115,37 @@ impl EngineHandle {
         .await
     }
 
+    /// The plan root answering `query` at every one of `steps` and the raw
+    /// samples they need; see [`QueryEngine::raw_ranges_for_steps`].
+    pub async fn raw_ranges_for_steps(
+        &self,
+        query: String,
+        steps: Vec<i64>,
+    ) -> Result<(usize, Vec<RawRange>), Unanswered> {
+        self.ask(|reply| Command::RawRangesForSteps {
+            query,
+            steps,
+            reply,
+        })
+        .await
+    }
+
+    /// Runs a root at every one of `steps`; see [`QueryEngine::evaluate_steps`].
+    pub async fn evaluate_steps(
+        &self,
+        root: usize,
+        steps: Vec<i64>,
+        raw: BTreeMap<u64, Vec<RawSample>>,
+    ) -> Result<Vec<Vec<ResultSample>>, Unanswered> {
+        self.ask(|reply| Command::EvaluateSteps {
+            root,
+            steps,
+            raw,
+            reply,
+        })
+        .await
+    }
+
     async fn ask<T>(&self, command: impl FnOnce(Reply<T>) -> Command) -> Result<T, Unanswered> {
         let busy = || Unanswered::Execution("the engine is busy".into());
         let (reply, answer) = oneshot::channel();
@@ -178,6 +220,32 @@ impl Engine {
                         &raw,
                     ));
                 }
+                Command::RawRangesForSteps {
+                    query,
+                    steps,
+                    reply,
+                } => {
+                    let _ = reply.send(self.query.raw_ranges_for_steps(
+                        &query,
+                        &steps,
+                        &self.precompute.store,
+                    ));
+                }
+                Command::EvaluateSteps {
+                    root,
+                    steps,
+                    raw,
+                    reply,
+                } => {
+                    // One watermark for every step of the request.
+                    let _ = reply.send(self.query.evaluate_steps(
+                        root,
+                        &steps,
+                        &self.precompute.store,
+                        self.precompute.watermark_ms(),
+                        &raw,
+                    ));
+                }
             }
         }
     }

data_plane/src/lib.rs

@@ -3,8 +3,8 @@
 //!
 //! The installed plan is ASAPPlanner's selected `PhysicalASAPDAG`. The
 //! engine thread builds its ingestion-time outputs pane by pane and answers
-//! the plan's instant queries from them; every other query, and any planned
-//! query it cannot answer, is forwarded to Prometheus.
+//! the plan's instant and range queries from them; every other query, and
+//! any planned query it cannot answer, is forwarded to Prometheus.
 
 pub mod engine;
 pub mod precompute;
@@ -26,7 +26,7 @@ use axum::Router;
 use control_plane::InstalledPlan;
 use engine::EngineHandle;
 use precompute::RawSample;
-use prometheus::InstantQuery;
+use prometheus::{InstantQuery, RangeQuery};
 use query::Unanswered;
 
 /// Prometheus' staleness marker: a NaN that ends a series, not a sample.
@@ -125,24 +125,57 @@ async fn query(
     headers: HeaderMap,
     body: Bytes,
 ) -> Response {
-    let Some(instant) = InstantQuery::parse(query.0.as_deref(), &body) else {
-        return forward(state, "query", method, query, headers, body).await;
+    let answered = match InstantQuery::parse(query.0.as_deref(), &body) {
+        Some(instant) => Some(answer(&state, instant).await),
+        None => None,
     };
-    let reason = match answer(&state, instant).await {
-        Ok(response) => {
+    respond(answered, state, "query", method, query, headers, body).await
+}
+
+/// Answers a planned range query from the plan, one evaluation per step;
+/// forwards anything else, and the whole range query if any step cannot
+/// be answered.
+async fn query_range(
+    state: State<Arc<DataPlane>>,
+    method: Method,
+    query: RawQuery,
+    headers: HeaderMap,
+    body: Bytes,
+) -> Response {
+    let answered = match RangeQuery::parse(query.0.as_deref(), &body) {
+        Some(range) => Some(answer_range(&state, range).await),
+        None => None,
+    };
+    respond(answered, state, "query_range", method, query, headers, body).await
+}
+
+/// The plan's answer with `x-asap-execution: plan`, or else the request
+/// forwarded to `endpoint`, with the reason if the query is planned.
+/// `answered` is `None` for a request the data plane does not parse.
+async fn respond(
+    answered: Option<Result<serde_json::Value, Unanswered>>,
+    state: State<Arc<DataPlane>>,
+    endpoint: &str,
+    method: Method,
+    query: RawQuery,
+    headers: HeaderMap,
+    body: Bytes,
+) -> Response {
+    let reason = match answered {
+        Some(Ok(response)) => {
             let mut response = axum::Json(response).into_response();
             response
                 .headers_mut()
                 .insert(EXECUTION_HEADER, HeaderValue::from_static("plan"));
             return response;
         }
-        Err(Unanswered::NotPlanned(_)) => None,
-        Err(reason) => {
-            tracing::debug!(%reason, "forwarding a planned query");
+        None | Some(Err(Unanswered::NotPlanned(_))) => None,
+        Some(Err(reason)) => {
+            tracing::debug!(%reason, endpoint, "forwarding a planned query");
             HeaderValue::from_str(&reason.to_string()).ok()
         }
     };
-    let mut response = forward(state, "query", method, query, headers, body).await;
+    let mut response = forward(state, endpoint, method, query, headers, body).await;
     if let Some(reason) = reason {
         response.headers_mut().insert(FORWARD_REASON_HEADER, reason);
     }
@@ -156,25 +189,42 @@ async fn answer(state: &DataPlane, query: InstantQuery) -> Result<serde_json::Va
             .map_or(0, |d| d.as_millis() as i64)
     });
     let (root, ranges) = state.engine.raw_ranges(query.query, time_ms).await?;
+    let raw = fetch(state, &ranges).await?;
+    let samples = state.engine.evaluate(root, time_ms, raw).await?;
+    Ok(prometheus::vector_response(time_ms, &samples))
+}
+
+/// Answers every step of `query`, fetching each raw range once for all.
+async fn answer_range(
+    state: &DataPlane,
+    query: RangeQuery,
+) -> Result<serde_json::Value, Unanswered> {
+    let steps = query.steps();
+    let (root, ranges) = state
+        .engine
+        .raw_ranges_for_steps(query.query, steps.clone())
+        .await?;
+    let raw = fetch(state, &ranges).await?;
+    let results = state
+        .engine
+        .evaluate_steps(root, steps.clone(), raw)
+        .await?;
+    Ok(prometheus::matrix_response(&steps, &results))
+}
+
+/// The samples of each of `ranges` from Prometheus, by node.
+async fn fetch(
+    state: &DataPlane,
+    ranges: &[query::RawRange],
+) -> Result<BTreeMap<u64, Vec<RawSample>>, Unanswered> {
     let mut raw = BTreeMap::new();
-    for range in &ranges {
+    for range in ranges {
         let samples = prometheus::fetch(&state.client, &state.prometheus_url, range)
             .await
             .map_err(|e| Unanswered::Execution(format!("fetching raw samples: {e}")))?;
         raw.insert(range.node, samples);
     }
-    let samples = state.engine.evaluate(root, time_ms, raw).await?;
-    Ok(prometheus::vector_response(time_ms, &samples))
-}
-
-async fn query_range(
-    state: State<Arc<DataPlane>>,
-    method: Method,
-    query: RawQuery,
-    headers: HeaderMap,
-    body: Bytes,
-) -> Response {
-    forward(state, "query_range", method, query, headers, body).await
+    Ok(raw)
 }
 
 /// Sends the request to Prometheus unchanged and returns its response.