22
33use crate :: engines:: simple_engine:: { RangeQueryExecutionContext , StoreQueryParams } ;
44use asap_types:: enums:: WindowType ;
5+ use promql_utilities:: data_model:: KeyByLabelNames ;
56use promql_utilities:: query_logics:: enums:: Statistic ;
7+ use tracing:: debug;
68
79#[ derive( Debug , Clone , Copy , PartialEq , Eq ) ]
810pub ( crate ) struct NodeId ( usize ) ;
@@ -38,10 +40,12 @@ pub(crate) enum QueryPlanNode {
3840 LimitTopK {
3941 input : NodeId ,
4042 k : String ,
43+ grouping_labels : KeyByLabelNames ,
4144 } ,
4245 Format {
4346 input : NodeId ,
4447 include_metric_name : bool ,
48+ metric : String ,
4549 } ,
4650}
4751
@@ -57,7 +61,44 @@ pub(crate) struct PlanOptions {
5761 pub format_output : bool ,
5862}
5963
64+ pub ( crate ) trait QueryPlanRuntime {
65+ type Output : Clone ;
66+ type Error : std:: fmt:: Display ;
67+
68+ fn execute_node (
69+ & self ,
70+ id : NodeId ,
71+ node : & QueryPlanNode ,
72+ inputs : & [ Self :: Output ] ,
73+ ) -> Result < Self :: Output , Self :: Error > ;
74+ }
75+
76+ #[ derive( Debug ) ]
77+ pub ( crate ) enum QueryPlanExecutionError < E > {
78+ InvalidPlan ( String ) ,
79+ Node { id : NodeId , source : E } ,
80+ }
81+
82+ impl < E : std:: fmt:: Display > std:: fmt:: Display for QueryPlanExecutionError < E > {
83+ fn fmt ( & self , formatter : & mut std:: fmt:: Formatter < ' _ > ) -> std:: fmt:: Result {
84+ match self {
85+ Self :: InvalidPlan ( error) => write ! ( formatter, "invalid query plan: {error}" ) ,
86+ Self :: Node { id, source } => {
87+ write ! ( formatter, "Query plan node n{} failed: {source}" , id. 0 )
88+ }
89+ }
90+ }
91+ }
92+
6093impl QueryPlan {
94+ #[ cfg( feature = "native_query_legacy_test_support" ) ]
95+ pub ( crate ) fn malformed_for_test ( ) -> Self {
96+ Self {
97+ nodes : Vec :: new ( ) ,
98+ root : NodeId ( 0 ) ,
99+ }
100+ }
101+
61102 pub ( crate ) fn compile_range (
62103 context : & RangeQueryExecutionContext ,
63104 options : PlanOptions ,
@@ -114,18 +155,90 @@ impl QueryPlan {
114155 . ok_or_else ( || "Topk query is missing required `k` parameter" . to_string ( ) ) ?;
115156 k. parse :: < usize > ( )
116157 . map_err ( |_| "Topk query has an invalid `k` parameter" . to_string ( ) ) ?;
117- root = Self :: push ( & mut nodes, QueryPlanNode :: LimitTopK { input : root, k } ) ;
158+ root = Self :: push (
159+ & mut nodes,
160+ QueryPlanNode :: LimitTopK {
161+ input : root,
162+ k,
163+ grouping_labels : context. base . grouping_labels . clone ( ) ,
164+ } ,
165+ ) ;
118166 }
119167 if options. format_output {
120168 root = Self :: push (
121169 & mut nodes,
122170 QueryPlanNode :: Format {
123171 input : root,
124- include_metric_name : context. base . metadata . keep_metric_name ,
172+ include_metric_name : context. base . metadata . statistic_to_compute
173+ == Statistic :: Topk
174+ && context. base . metadata . keep_metric_name ,
175+ metric : context. base . metric . clone ( ) ,
125176 } ,
126177 ) ;
127178 }
128- Ok ( Self { nodes, root } )
179+ let plan = Self { nodes, root } ;
180+ plan. validate ( ) ?;
181+ Ok ( plan)
182+ }
183+
184+ /// Rejects plans whose node dependencies cannot be executed safely.
185+ pub ( crate ) fn validate ( & self ) -> Result < ( ) , String > {
186+ if self . nodes . is_empty ( ) {
187+ return Err ( "Query plan has no nodes" . to_string ( ) ) ;
188+ }
189+ if self . root . 0 != self . nodes . len ( ) - 1 {
190+ return Err ( format ! (
191+ "Query plan root n{} does not include every node" ,
192+ self . root. 0
193+ ) ) ;
194+ }
195+ for ( index, node) in self . nodes . iter ( ) . enumerate ( ) {
196+ for input in node. inputs ( ) {
197+ if input. 0 >= index {
198+ return Err ( format ! (
199+ "Query plan node n{index} references unavailable input n{}" ,
200+ input. 0
201+ ) ) ;
202+ }
203+ }
204+ }
205+ Ok ( ( ) )
206+ }
207+
208+ pub ( crate ) fn execute < R : QueryPlanRuntime > (
209+ & self ,
210+ runtime : & R ,
211+ ) -> Result < R :: Output , QueryPlanExecutionError < R :: Error > > {
212+ self . validate ( )
213+ . map_err ( QueryPlanExecutionError :: InvalidPlan ) ?;
214+ let mut outputs: Vec < R :: Output > = Vec :: with_capacity ( self . nodes . len ( ) ) ;
215+ for ( index, node) in self . nodes . iter ( ) . enumerate ( ) {
216+ let inputs = node
217+ . inputs ( )
218+ . into_iter ( )
219+ . map ( |input| outputs[ input. 0 ] . clone ( ) )
220+ . collect :: < Vec < _ > > ( ) ;
221+ debug ! (
222+ node_id = index,
223+ node_kind = node. kind( ) ,
224+ input_count = inputs. len( ) ,
225+ "Executing native query plan node"
226+ ) ;
227+ let output = runtime
228+ . execute_node ( NodeId ( index) , node, & inputs)
229+ . map_err ( |source| QueryPlanExecutionError :: Node {
230+ id : NodeId ( index) ,
231+ source,
232+ } ) ?;
233+ debug ! (
234+ node_id = index,
235+ node_kind = node. kind( ) ,
236+ "Completed native query plan node"
237+ ) ;
238+ outputs. push ( output) ;
239+ }
240+ debug ! ( root_node_id = self . root. 0 , "Completed native query plan" ) ;
241+ Ok ( outputs[ self . root . 0 ] . clone ( ) )
129242 }
130243
131244 fn push ( nodes : & mut Vec < QueryPlanNode > , node : QueryPlanNode ) -> NodeId {
@@ -194,9 +307,11 @@ impl QueryPlan {
194307 kwargs. sort_unstable_by_key ( |( key, _) | * key) ;
195308 format ! ( "n{index} Estimate(n{}, {statistic}, {kwargs:?})" , input. 0 )
196309 } ,
197- QueryPlanNode :: LimitTopK { input, k } => format ! ( "n{index} LimitTopK(n{}, k={k})" , input. 0 ) ,
198- QueryPlanNode :: Format { input, include_metric_name } => format ! (
199- "n{index} Format(n{}, include_metric_name={include_metric_name})" , input. 0
310+ QueryPlanNode :: LimitTopK { input, k, .. } => {
311+ format ! ( "n{index} LimitTopK(n{}, k={k})" , input. 0 )
312+ }
313+ QueryPlanNode :: Format { input, include_metric_name, metric } => format ! (
314+ "n{index} Format(n{}, include_metric_name={include_metric_name}) metric={metric}" , input. 0
200315 ) ,
201316 } ;
202317 lines. push ( line) ;
@@ -206,13 +321,43 @@ impl QueryPlan {
206321 }
207322}
208323
324+ impl QueryPlanNode {
325+ fn kind ( & self ) -> & ' static str {
326+ match self {
327+ Self :: StoreRead { .. } => "StoreRead" ,
328+ Self :: ComposeWindows { .. } => "ComposeWindows" ,
329+ Self :: ResolveKeys { .. } => "ResolveKeys" ,
330+ Self :: Estimate { .. } => "Estimate" ,
331+ Self :: LimitTopK { .. } => "LimitTopK" ,
332+ Self :: Format { .. } => "Format" ,
333+ }
334+ }
335+
336+ fn inputs ( & self ) -> Vec < NodeId > {
337+ match self {
338+ Self :: StoreRead { .. } => Vec :: new ( ) ,
339+ Self :: ComposeWindows { input, .. }
340+ | Self :: Estimate { input, .. }
341+ | Self :: LimitTopK { input, .. }
342+ | Self :: Format { input, .. } => vec ! [ * input] ,
343+ Self :: ResolveKeys { values, keys } => {
344+ keys. iter ( ) . copied ( ) . fold ( vec ! [ * values] , |mut inputs, key| {
345+ inputs. push ( key) ;
346+ inputs
347+ } )
348+ }
349+ }
350+ }
351+ }
352+
209353#[ cfg( test) ]
210354mod tests {
211355 use super :: * ;
212356 use crate :: data_model:: AggregationIdInfo ;
213357 use crate :: engines:: simple_engine:: { QueryExecutionContext , QueryMetadata , StoreQueryPlan } ;
214358 use promql_utilities:: data_model:: KeyByLabelNames ;
215359 use promql_utilities:: query_logics:: enums:: AggregationType ;
360+ use std:: cell:: RefCell ;
216361 use std:: collections:: HashMap ;
217362
218363 fn context ( ) -> RangeQueryExecutionContext {
@@ -347,4 +492,66 @@ mod tests {
347492
348493 assert_eq ! ( error, "Topk query is missing required `k` parameter" ) ;
349494 }
495+
496+ #[ test]
497+ fn rejects_a_node_that_references_a_later_node ( ) {
498+ let plan = QueryPlan {
499+ nodes : vec ! [ QueryPlanNode :: Estimate {
500+ input: NodeId ( 1 ) ,
501+ statistic: Statistic :: Sum ,
502+ query_kwargs: HashMap :: new( ) ,
503+ } ] ,
504+ root : NodeId ( 0 ) ,
505+ } ;
506+
507+ assert_eq ! (
508+ plan. validate( ) . expect_err( "invalid plan must fail loudly" ) ,
509+ "Query plan node n0 references unavailable input n1"
510+ ) ;
511+ }
512+
513+ struct RecordingRuntime ( RefCell < Vec < usize > > ) ;
514+
515+ impl QueryPlanRuntime for RecordingRuntime {
516+ type Output = usize ;
517+ type Error = std:: convert:: Infallible ;
518+
519+ fn execute_node (
520+ & self ,
521+ id : NodeId ,
522+ _node : & QueryPlanNode ,
523+ inputs : & [ Self :: Output ] ,
524+ ) -> Result < Self :: Output , Self :: Error > {
525+ self . 0 . borrow_mut ( ) . push ( id. 0 ) ;
526+ Ok ( 1 + inputs. iter ( ) . sum :: < usize > ( ) )
527+ }
528+ }
529+
530+ #[ test]
531+ fn executes_nodes_once_in_dependency_order ( ) {
532+ let plan = QueryPlan {
533+ nodes : vec ! [
534+ QueryPlanNode :: StoreRead {
535+ query: StoreQueryParams {
536+ metric: "requests" . into( ) ,
537+ aggregation_id: 7 ,
538+ start_timestamp: 0 ,
539+ end_timestamp: 1 ,
540+ } ,
541+ strategy: StoreReadStrategy :: WindowGrid ,
542+ } ,
543+ QueryPlanNode :: ComposeWindows {
544+ input: NodeId ( 0 ) ,
545+ output_timestamps: vec![ 1 ] ,
546+ lookback_ms: 1 ,
547+ window_size_ms: 1 ,
548+ bucket_step_ms: 1 ,
549+ } ,
550+ ] ,
551+ root : NodeId ( 1 ) ,
552+ } ;
553+ let runtime = RecordingRuntime ( RefCell :: new ( Vec :: new ( ) ) ) ;
554+ assert_eq ! ( plan. execute( & runtime) . unwrap( ) , 2 ) ;
555+ assert_eq ! ( * runtime. 0 . borrow( ) , vec![ 0 , 1 ] ) ;
556+ }
350557}
0 commit comments