select_object optimizations (#3810)

Small optimizations when selecting entity types and keys for federated
queries

Co-authored-by: o0Ignition0o <[email protected]>

apollo-router/src/query_planner/fetch.rs

@@ -1,4 +1,3 @@
-use std::collections::HashMap;
 use std::fmt::Display;
 use std::sync::Arc;
 
@@ -98,13 +97,13 @@ pub(crate) struct FetchNode {
 
 pub(crate) struct Variables {
     pub(crate) variables: Object,
-    pub(crate) paths: HashMap<Path, usize>,
+    pub(crate) inverted_paths: Vec<Vec<Path>>,
 }
 
 impl Variables {
     #[instrument(skip_all, level = "debug", name = "make_variables")]
     #[allow(clippy::too_many_arguments)]
-    pub(super) async fn new(
+    pub(super) fn new(
         requires: &[Selection],
         variable_usages: &[String],
         data: &Value,
@@ -123,7 +122,7 @@ impl Variables {
                     .map(|(variable_key, value)| (variable_key.clone(), value.clone()))
             }));
 
-            let mut paths: HashMap<Path, usize> = HashMap::new();
+            let mut inverted_paths: Vec<Vec<Path>> = Vec::new();
             let mut values: IndexSet<Value> = IndexSet::new();
 
             data.select_values_and_paths(schema, current_dir, |path, value| {
@@ -132,11 +131,12 @@ impl Variables {
                         rewrites::apply_rewrites(schema, &mut value, input_rewrites);
                         match values.get_index_of(&value) {
                             Some(index) => {
-                                paths.insert(path.clone(), index);
+                                inverted_paths[index].push(path.clone());
                             }
                             None => {
-                                paths.insert(path.clone(), values.len());
+                                inverted_paths.push(vec![path.clone()]);
                                 values.insert(value);
+                                debug_assert!(inverted_paths.len() == values.len());
                             }
                         }
                     }
@@ -151,7 +151,10 @@ impl Variables {
 
             variables.insert("representations", representations);
 
-            Some(Variables { variables, paths })
+            Some(Variables {
+                variables,
+                inverted_paths,
+            })
         } else {
             // with nested operations (Query or Mutation has an operation returning a Query or Mutation),
             // when the first fetch fails, the query plan will still execute up until the second fetch,
@@ -176,7 +179,7 @@ impl Variables {
                             .map(|(variable_key, value)| (variable_key.clone(), value.clone()))
                     })
                     .collect::<Object>(),
-                paths: HashMap::new(),
+                inverted_paths: Vec::new(),
             })
         }
     }
@@ -198,7 +201,10 @@ impl FetchNode {
             ..
         } = self;
 
-        let Variables { variables, paths } = match Variables::new(
+        let Variables {
+            variables,
+            inverted_paths: paths,
+        } = match Variables::new(
             &self.requires,
             self.variable_usages.as_ref(),
             data,
@@ -207,9 +213,7 @@ impl FetchNode {
             parameters.supergraph_request,
             parameters.schema,
             &self.input_rewrites,
-        )
-        .await
-        {
+        ) {
             Some(variables) => variables,
             None => {
                 return Ok((Value::Object(Object::default()), Vec::new()));
@@ -304,15 +308,9 @@ impl FetchNode {
         &'a self,
         schema: &Schema,
         current_dir: &'a Path,
-        paths: HashMap<Path, usize>,
+        inverted_paths: Vec<Vec<Path>>,
         response: graphql::Response,
     ) -> (Value, Vec<Error>) {
-        // for each entity in the response, find out the path where it must be inserted
-        let mut inverted_paths: HashMap<usize, Vec<&Path>> = HashMap::new();
-        for (path, index) in paths.iter() {
-            (*inverted_paths.entry(*index).or_default()).push(path);
-        }
-
         if !self.requires.is_empty() {
             let entities_path = Path(vec![json_ext::PathElement::Key("_entities".to_string())]);
 
@@ -330,7 +328,7 @@ impl FetchNode {
                         match path.0.get(1) {
                             Some(json_ext::PathElement::Index(i)) => {
                                 for values_path in
-                                    inverted_paths.get(i).iter().flat_map(|v| v.iter())
+                                    inverted_paths.get(*i).iter().flat_map(|v| v.iter())
                                 {
                                     errors.push(Error {
                                         locations: error.locations.clone(),
@@ -370,7 +368,7 @@ impl FetchNode {
                         for (index, mut entity) in array.into_iter().enumerate() {
                             rewrites::apply_rewrites(schema, &mut entity, &self.output_rewrites);
 
-                            if let Some(paths) = inverted_paths.get(&index) {
+                            if let Some(paths) = inverted_paths.get(index) {
                                 if paths.len() > 1 {
                                     for path in &paths[1..] {
                                         let _ = value.insert(path, entity.clone());

apollo-router/src/query_planner/selection.rs

@@ -1,6 +1,6 @@
 use serde::Deserialize;
 use serde::Serialize;
-use serde_json_bytes::Entry;
+use serde_json_bytes::ByteString;
 
 use crate::error::FetchError;
 use crate::json_ext::Object;
@@ -53,16 +53,15 @@ pub(crate) fn select_object(
     selections: &[Selection],
     schema: &Schema,
 ) -> Result<Option<Value>, FetchError> {
-    let mut output = Object::new();
+    let mut output = Object::with_capacity(selections.len());
     for selection in selections {
         match selection {
             Selection::Field(field) => {
-                if let Some(value) = select_field(content, field, schema)? {
-                    match output.entry(field.name.to_owned()) {
-                        Entry::Occupied(mut existing) => existing.get_mut().deep_merge(value),
-                        Entry::Vacant(vacant) => {
-                            vacant.insert(value);
-                        }
+                if let Some((key, value)) = select_field(content, field, schema)? {
+                    if let Some(o) = output.get_mut(field.name.as_str()) {
+                        o.deep_merge(value);
+                    } else {
+                        output.insert(key.to_owned(), value);
                     }
                 }
             }
@@ -81,13 +80,16 @@ pub(crate) fn select_object(
     Ok(Some(Value::Object(output)))
 }
 
-fn select_field(
-    content: &Object,
+fn select_field<'a>(
+    content: &'a Object,
     field: &Field,
     schema: &Schema,
-) -> Result<Option<Value>, FetchError> {
-    let res = match (content.get(field.name.as_str()), &field.selections) {
-        (Some(v), _) => select_value(v, field, schema),
+) -> Result<Option<(&'a ByteString, Value)>, FetchError> {
+    let res = match (
+        content.get_key_value(field.name.as_str()),
+        &field.selections,
+    ) {
+        (Some((k, v)), _) => select_value(v, field, schema).map(|opt| opt.map(|v| (k, v))),
         (None, _) => Err(FetchError::ExecutionFieldNotFound {
             field: field.name.to_owned(),
         }),

apollo-router/src/query_planner/subscription.rs

@@ -209,9 +209,7 @@ impl SubscriptionNode {
             parameters.supergraph_request,
             parameters.schema,
             &self.input_rewrites,
-        )
-        .await
-        {
+        ) {
             Some(variables) => variables,
             None => {
                 return Ok(Vec::new());