routee_compass/plugin/input/default/grid_search/
plugin.rs

1use crate::app::search::SearchApp;
2use crate::plugin::input::input_field::InputField;
3use crate::plugin::input::input_plugin::InputPlugin;
4use crate::plugin::input::InputJsonExtensions;
5use crate::plugin::input::InputPluginError;
6use routee_compass_core::util::multiset::MultiSet;
7use std::sync::Arc;
8
9/// Builds an input plugin that duplicates queries if array-valued fields are present
10/// by stepping through each combination of value
11pub struct GridSearchPlugin {}
12
13const NAME: &str = "grid_search";
14
15impl InputPlugin for GridSearchPlugin {
16    fn name(&self) -> &str {
17        NAME
18    }
19
20    fn process(
21        &self,
22        input: &mut serde_json::Value,
23        _search_app: Arc<SearchApp>,
24    ) -> Result<(), InputPluginError> {
25        match process_grid_search(input)? {
26            None => Ok(()),
27            Some(grid_expansion) => {
28                let mut replacement = serde_json::json![grid_expansion];
29                std::mem::swap(&mut replacement, input);
30                Ok(())
31            }
32        }
33    }
34}
35
36fn process_grid_search(
37    input: &serde_json::Value,
38) -> Result<Option<Vec<serde_json::Value>>, InputPluginError> {
39    let grid_search_input = match input.get_grid_search() {
40        Some(gsi) => gsi,
41        None => return Ok(None),
42    };
43
44    // prevent recursion due to nested grid search keys
45    let recurses = serde_json::to_string(grid_search_input)
46        .map_err(|e| InputPluginError::JsonError { source: e })?
47        .contains("grid_search");
48    if recurses {
49        return Err(InputPluginError::InputPluginFailed(String::from(
50            "grid search section cannot contain the string 'grid_search'",
51        )));
52    }
53
54    let map = grid_search_input
55        .as_object()
56        .ok_or_else(|| InputPluginError::UnexpectedQueryStructure(format!("{input:?}")))?;
57    let mut keys: Vec<String> = vec![];
58    let mut multiset_input: Vec<Vec<serde_json::Value>> = vec![];
59    let mut multiset_indices: Vec<Vec<usize>> = vec![];
60    for (k, v) in map {
61        if let Some(v) = v.as_array() {
62            keys.push(k.to_string());
63            multiset_input.push(v.to_vec());
64            let indices = (0..v.len()).collect();
65            multiset_indices.push(indices);
66        }
67    }
68    // for each combination, copy the grid search values into a fresh
69    // copy of the source (minus the "grid_search" key)
70    // let remove_key = InputField::GridSearch.to_str();
71    let mut initial_map = input
72        .as_object()
73        .ok_or_else(|| InputPluginError::UnexpectedQueryStructure(format!("{input:?}")))?
74        .clone();
75    initial_map.remove(InputField::GridSearch.to_str());
76    let initial = serde_json::json!(initial_map);
77    let multiset = MultiSet::from(&multiset_indices);
78    let result: Vec<serde_json::Value> = multiset
79        .into_iter()
80        .map(|combination| {
81            let mut instance = initial.clone();
82            let it = keys.iter().zip(combination.iter()).enumerate();
83            for (set_idx, (key, val_idx)) in it {
84                let value = multiset_input[set_idx][*val_idx].clone();
85                match value {
86                    serde_json::Value::Object(o) => {
87                        for (k, v) in o.into_iter() {
88                            instance[k] = v.clone();
89                        }
90                    }
91                    _ => {
92                        instance[key] = multiset_input[set_idx][*val_idx].clone();
93                    }
94                }
95            }
96            instance
97        })
98        .collect();
99
100    Ok(Some(result))
101}
102
103#[cfg(test)]
104mod test {
105    use super::*;
106
107    use serde_json::json;
108
109    #[test]
110    fn test_grid_search_empty_parent_object() {
111        let input = serde_json::json!({
112            "grid_search": {
113                "bar": ["a", "b", "c"],
114                "foo": [1.2, 3.4]
115            }
116        });
117
118        let result = match process_grid_search(&input) {
119            Ok(Some(rows)) => rows,
120            Ok(None) => panic!("process_grid_search returned no expansions"),
121            Err(e) => panic!("{}", e),
122        };
123        let expected = vec![
124            json![{"bar":"a","foo":1.2}],
125            json![{"bar":"b","foo":1.2}],
126            json![{"bar":"c","foo":1.2}],
127            json![{"bar":"a","foo":3.4}],
128            json![{"bar":"b","foo":3.4}],
129            json![{"bar":"c","foo":3.4}],
130        ];
131        assert_eq!(result, expected)
132    }
133
134    #[test]
135    fn test_grid_search_persisted_parent_keys() {
136        let input = serde_json::json!({
137            "ignored_key": "ignored_value",
138            "grid_search": {
139                "bar": ["a", "b", "c"],
140                "foo": [1.2, 3.4]
141            }
142        });
143
144        let result = match process_grid_search(&input) {
145            Ok(Some(rows)) => rows,
146            Ok(None) => panic!("process_grid_search returned no expansions"),
147            Err(e) => panic!("{}", e),
148        };
149
150        let expected = vec![
151            json![{"bar":"a","foo":1.2,"ignored_key": "ignored_value"}],
152            json![{"bar":"b","foo":1.2,"ignored_key": "ignored_value"}],
153            json![{"bar":"c","foo":1.2,"ignored_key": "ignored_value"}],
154            json![{"bar":"a","foo":3.4,"ignored_key": "ignored_value"}],
155            json![{"bar":"b","foo":3.4,"ignored_key": "ignored_value"}],
156            json![{"bar":"c","foo":3.4,"ignored_key": "ignored_value"}],
157        ];
158
159        assert_eq!(result, expected)
160    }
161
162    #[test]
163    fn test_grid_search_using_objects() {
164        let input = serde_json::json!({
165            "ignored_key": "ignored_value",
166            "grid_search": {
167                "a": [1, 2],
168                "ignored_inner_key": [
169                    { "x": 0, "y": 0 },
170                    { "x": 1, "y": 1 }
171                ],
172            }
173        });
174
175        let result = match process_grid_search(&input) {
176            Ok(Some(rows)) => rows,
177            Ok(None) => panic!("process_grid_search returned no expansions"),
178            Err(e) => panic!("{}", e),
179        };
180
181        let expected = vec![
182            json![{"a":1,"ignored_key":"ignored_value","x":0,"y":0}],
183            json![{"a":2,"ignored_key":"ignored_value","x":0,"y":0}],
184            json![{"a":1,"ignored_key":"ignored_value","x":1,"y":1}],
185            json![ {"a":2,"ignored_key":"ignored_value","x":1,"y":1}],
186        ];
187
188        assert_eq!(result, expected)
189    }
190
191    #[test]
192    fn test_nested() {
193        let input = serde_json::json!({
194            "abc": 123,
195            "grid_search":{
196                "model_name": ["2016_TOYOTA_Camry_4cyl_2WD","2017_CHEVROLET_Bolt"],
197                "_ignore":[
198                    { "name":"d1", "weights": { "distance":1, "time":0, "energy_electric":0 } },
199                    { "name":"t1", "weights": { "distance":0, "time":1, "energy_electric":0 } },
200                    { "name":"e1", "weights": { "distance":0, "time":0, "energy_electric":1 } }
201                ]
202            }
203        });
204
205        let result = match process_grid_search(&input) {
206            Ok(Some(rows)) => rows,
207            Ok(None) => panic!("process_grid_search returned no expansions"),
208            Err(e) => panic!("{}", e),
209        };
210
211        let expected = vec![
212            json![{"abc":123,"model_name":"2016_TOYOTA_Camry_4cyl_2WD","name":"d1","weights":{"distance":1,"time":0,"energy_electric":0}}],
213            json![{"abc":123,"model_name":"2017_CHEVROLET_Bolt","name":"d1","weights":{"distance":1,"time":0,"energy_electric":0}}],
214            json![{"abc":123,"model_name":"2016_TOYOTA_Camry_4cyl_2WD","name":"t1","weights":{"distance":0,"time":1,"energy_electric":0}}],
215            json![{"abc":123,"model_name":"2017_CHEVROLET_Bolt","name":"t1","weights":{"distance":0,"time":1,"energy_electric":0}}],
216            json![{"abc":123,"model_name":"2016_TOYOTA_Camry_4cyl_2WD","name":"e1","weights":{"distance":0,"time":0,"energy_electric":1}}],
217            json![{"abc":123,"model_name":"2017_CHEVROLET_Bolt","name":"e1","weights":{"distance":0,"time":0,"energy_electric":1}}],
218        ];
219
220        assert_eq!(result, expected)
221    }
222
223    #[test]
224    pub fn test_handle_recursion() {
225        let input = serde_json::json!({
226            "abc": 123,
227            "grid_search":{
228                "grid_search": {
229                    "foo": [ "a", "b" ]
230                }
231            }
232        });
233
234        match process_grid_search(&input) {
235            Ok(Some(_)) => panic!("process_grid_search should return an error"),
236            Ok(None) => panic!("process_grid_search returned no error"),
237            Err(_) => {}
238        };
239    }
240}
routee_compass/plugin/input/default/grid_search/plugin.rs

routee_compass/plugin/input/default/grid_search/
plugin.rs
