@@ -8,9 +8,12 @@ use sketch_db_common::enums::CleanupPolicy;
88
99use crate :: config:: input:: ControllerConfig ;
1010use crate :: error:: ControllerError ;
11- use crate :: planner:: single_query:: { IntermediateAggConfig , SingleQueryProcessor } ;
11+ use crate :: planner:: single_query:: { BinaryArm , IntermediateAggConfig , SingleQueryProcessor } ;
1212use crate :: RuntimeOptions ;
1313
14+ /// `(query_string, Vec<(identifying_key, cleanup_param)>)` pairs produced by binary leaf decomposition.
15+ type LeafEntries = Vec < ( String , Vec < ( String , Option < u64 > ) > ) > ;
16+
1417/// Run the full planning pipeline and produce YAML outputs
1518pub fn generate_plan (
1619 controller_config : & ControllerConfig ,
@@ -92,6 +95,14 @@ pub fn generate_plan(
9295 }
9396 Err ( e) => return Err ( e) ,
9497 }
98+ } else if let Some ( arm_entries) =
99+ collect_binary_leaf_entries ( & processor, & mut dedup_map) ?
100+ {
101+ // Binary arithmetic: register each leaf arm in dedup_map and query_keys_map
102+ for ( arm_query, keys_for_arm) in arm_entries {
103+ // Use `entry` so a standalone query that duplicates an arm wins
104+ query_keys_map. entry ( arm_query) . or_insert ( keys_for_arm) ;
105+ }
95106 }
96107 }
97108 }
@@ -123,6 +134,62 @@ pub fn generate_plan(
123134 } )
124135}
125136
137+ /// Recursively collect (arm_query_string, Vec<(dedup_key, cleanup_param)>) pairs
138+ /// from a binary arithmetic expression, registering new configs in `dedup_map`.
139+ ///
140+ /// Returns `Some(Vec<...>)` when every leaf arm is acceleratable.
141+ /// Returns `None` if any arm is unsupported (caller should skip the query).
142+ /// Returns `Err` only on internal planner errors.
143+ fn collect_binary_leaf_entries (
144+ processor : & SingleQueryProcessor ,
145+ dedup_map : & mut IndexMap < String , IntermediateAggConfig > ,
146+ ) -> Result < Option < LeafEntries > , ControllerError > {
147+ let arms = match processor. get_binary_arm_queries ( ) {
148+ Some ( arms) => arms,
149+ None => return Ok ( None ) , // not a binary expression
150+ } ;
151+
152+ let mut all_entries: LeafEntries = Vec :: new ( ) ;
153+
154+ for arm in [ arms. 0 , arms. 1 ] {
155+ match arm {
156+ BinaryArm :: Scalar ( _) => {
157+ // Scalar literals need no aggregation config — skip silently.
158+ }
159+ BinaryArm :: Query ( arm_query) => {
160+ let arm_processor = processor. make_arm_processor ( arm_query. clone ( ) ) ;
161+
162+ if arm_processor. is_supported ( ) {
163+ // Leaf arm: gather its streaming aggregation configs.
164+ let ( configs, cleanup_param) =
165+ arm_processor. get_streaming_aggregation_configs ( ) ?;
166+ let mut keys_for_arm = Vec :: new ( ) ;
167+ for config in configs {
168+ let key = config. identifying_key ( ) ;
169+ keys_for_arm. push ( ( key. clone ( ) , cleanup_param) ) ;
170+ dedup_map. entry ( key) . or_insert ( config) ;
171+ }
172+ all_entries. push ( ( arm_query, keys_for_arm) ) ;
173+ } else {
174+ // The arm might itself be a binary expression — recurse.
175+ match collect_binary_leaf_entries ( & arm_processor, dedup_map) ? {
176+ Some ( sub_entries) => {
177+ all_entries. extend ( sub_entries) ;
178+ }
179+ None => {
180+ // Arm is neither a supported leaf nor a binary expression.
181+ // This entire query cannot be accelerated.
182+ return Ok ( None ) ;
183+ }
184+ }
185+ }
186+ }
187+ }
188+ }
189+
190+ Ok ( Some ( all_entries) )
191+ }
192+
126193pub fn parse_cleanup_policy ( s : & str ) -> Result < CleanupPolicy , ControllerError > {
127194 match s {
128195 "circular_buffer" => Ok ( CleanupPolicy :: CircularBuffer ) ,
0 commit comments