FED-32: Port `tryOptimizeSubselectionWithFragments` #5195

duckki · 2024-05-17T01:34:14Z

Implemented try_optimize_with_fragments function.

Added a new file optimize.rs under query_plan.
Refactored some SelectionSet methods and moved them to SelectionMap for reuse.

Checklist

Complete the checklist (and note appropriate exceptions) before the PR is marked ready-for-review.

Exceptions

Unfortunately, testing is very cumbersome without the whole optimize implementation (FED-191).
I'll implement optimize and test this together.

Notes

It may be appropriate to bring upcoming changes to the attention of other (impacted) groups. Please endeavour to do this before seeking PR approval. The mechanism for doing this will vary considerably, so use your judgement as to how and when to do this. ↩
Configuration is an important part of many changes. Where applicable please try to document configuration examples. ↩
Tick whichever testing boxes are applicable. If you are adding Manual Tests, please document the manual testing (extensively) in the Exceptions. ↩

- partially implemented `try_optimize_with_fragments`

…SelectionSet

github-actions · 2024-05-17T01:34:28Z

@duckki, please consider creating a changeset entry in /.changesets/. These instructions describe the process and tooling.

router-perf · 2024-05-17T01:34:45Z

lrlna

Wow this is a big one!!

First of all, your Rust has gotten so good. I am gobsmacked and super impressed at how quickly you got to writing this level of Rust. Incredible!

Second of all, I think we should keep the optimised functionality in operation.rs for a little while longer, and have a conversation about doing a refactor after we've finished everything with in the rewrite. It'll be easier to reason about the changes holistically once everything is in.

I am also really not certain working on the SelectionMap directly is a good idea. I played around a little bit with your changes to see if it's possible to work off a SelectionSet, and I think it should be. We just have to reason through the minus and the intersection functionality a little bit. Nothing wrong with doing a SelectionMap::new(), but I am more concerned about moving functions that are not reliant on being updated upon modifications.

There are also a few smaller comments and suggestions along the way.

lrlna · 2024-05-21T09:19:21Z

apollo-federation/src/query_plan/operation.rs

+            let key = TYPENAME_KEY.get_or_init(|| SelectionKey::Field {
+                response_name: TYPENAME_FIELD,
+                directives: Arc::new(Default::default()),
+            });
+
+            self.contains_key(key)
+        }
+
+        pub(crate) fn fields_in_set(&self) -> Vec<CollectedFieldInSet> {
+            let mut fields = Vec::new();
+
+            for (_key, selection) in self.iter() {
+                match selection {
+                    Selection::Field(field) => fields.push(CollectedFieldInSet {
+                        path: Vec::new(),
+                        field: field.clone(),
+                    }),
+                    Selection::FragmentSpread(_fragment) => {
+                        todo!()
+                    }
+                    Selection::InlineFragment(inline_fragment) => {
+                        let condition = inline_fragment
+                            .inline_fragment
+                            .data()
+                            .type_condition_position
+                            .as_ref();
+                        let header = match condition {
+                            Some(cond) => vec![FetchDataPathElement::TypenameEquals(
+                                cond.type_name().clone().into(),
+                            )],
+                            None => vec![],
+                        };
+                        for CollectedFieldInSet { path, field } in inline_fragment
+                            .selection_set
+                            .selections
+                            .fields_in_set()
+                            .into_iter()
+                        {
+                            let mut new_path = header.clone();
+                            new_path.extend(path);
+                            fields.push(CollectedFieldInSet {
+                                path: new_path,
+                                field,
+                            })
+                        }
+                    }
+                }
+            }
+            fields
+        }
+
+        pub(crate) fn containment(&self, other: &Self, options: ContainmentOptions) -> Containment {
+            if other.len() > self.len() {
+                // If `other` has more selections but we're ignoring missing __typename, then in the case where
+                // `other` has a __typename but `self` does not, then we need the length of `other` to be at
+                // least 2 more than other of `self` to be able to conclude there is no contains.
+                if !options.ignore_missing_typename
+                    || other.len() > self.len() + 1
+                    || self.has_top_level_typename_field()
+                    || !other.has_top_level_typename_field()
+                {
+                    return Containment::NotContained;
+                }
+            }
+
+            let mut is_equal = true;
+            let mut did_ignore_typename = false;
+
+            for (key, other_selection) in other.iter() {
+                if key.is_typename_field() && options.ignore_missing_typename {
+                    if !self.has_top_level_typename_field() {
+                        did_ignore_typename = true;
+                    }
+                    continue;
+                }
+
+                let Some(self_selection) = self.get(key) else {
+                    return Containment::NotContained;
+                };
+
+                match self_selection.containment(other_selection, options) {
+                    Containment::NotContained => return Containment::NotContained,
+                    Containment::StrictlyContained if is_equal => is_equal = false,
+                    Containment::StrictlyContained | Containment::Equal => {}
+                }
+            }
+
+            let expected_len = if did_ignore_typename {
+                self.len() + 1
+            } else {
+                self.len()
+            };
+
+            if is_equal && other.len() == expected_len {
+                Containment::Equal
+            } else {
+                Containment::StrictlyContained
+            }
+        }
    }


Do we need these methods to kept up to date when a modification happens? The SelectionMap is an "element" type, which keeps the keys updated. The SelectionSet is the right place to have these methods if that's not the case. Given optimisation happens pretty much at the very end of the process, I'd wager we want to keep these out of SelectionMap.

lrlna · 2024-05-21T09:21:33Z