feat(qp): introduce new optimization to batch siblings using multi-type steps

Author: dotansimha

lib/query-planner/src/ast/merge_path.rs

@@ -179,6 +179,17 @@ impl MergePath {
         }
         self.common_prefix_len(other) == other.len()
     }
+
+    pub fn without_type_castings(&self) -> Self {
+        let new_segments = self
+            .inner
+            .iter()
+            .filter(|segment| !matches!(segment, Segment::Cast(_, _)))
+            .cloned()
+            .collect::<Vec<_>>();
+
+        Self::new(new_segments)
+    }
 }
 
 impl Display for MergePath {

lib/query-planner/src/planner/fetch/fetch_step_data.rs

@@ -1,19 +1,16 @@
 use std::{collections::BTreeSet, fmt::Display};
 
 use bitflags::bitflags;
-use petgraph::{graph::NodeIndex, visit::EdgeRef};
-use tracing::trace;
+use petgraph::graph::NodeIndex;
 
 use crate::{
     ast::{
         merge_path::{Condition, MergePath},
         operation::VariableDefinition,
         safe_merge::AliasesRecords,
-        selection_set::find_arguments_conflicts,
     },
     planner::{
         fetch::{
-            fetch_graph::FetchGraph,
             selections::FetchStepSelections,
             state::{MultiTypeFetchStep, SingleTypeFetchStep},
         },
@@ -121,76 +118,9 @@ impl<State> FetchStepData<State> {
     }
 }
 
-impl FetchStepData<MultiTypeFetchStep> {
-    pub fn can_merge(
-        &self,
-        self_index: NodeIndex,
-        other_index: NodeIndex,
-        other: &Self,
-        fetch_graph: &FetchGraph<MultiTypeFetchStep>,
-    ) -> bool {
-        if self_index == other_index {
-            return false;
-        }
-
-        if self.service_name != other.service_name {
-            return false;
-        }
-
-        // We allow to merge root with entity calls by adding an inline fragment with the @include/@skip
-        if self.is_entity_call() && other.is_entity_call() && self.condition != other.condition {
-            return false;
-        }
-
-        // If both are entities, their response_paths should match,
-        // as we can't merge entity calls resolving different entities
-        if matches!(self.kind, FetchStepKind::Entity) && self.kind == other.kind {
-            if !self.response_path.eq(&other.response_path) {
-                return false;
-            }
-        } else {
-            // otherwise we can merge
-            if !other.response_path.starts_with(&self.response_path) {
-                return false;
-            }
-        }
-
-        let has_input_conflicts =
-            FetchStepSelections::iter_matching_types(&self.input, &other.input, |_, s1, s2| {
-                find_arguments_conflicts(s1, s2)
-            })
-            .iter()
-            .any(|(_, conflicts)| !conflicts.is_empty());
-
-        if has_input_conflicts {
-            trace!(
-                "preventing merge of [{}]+[{}] due to input conflicts",
-                self_index.index(),
-                other_index.index(),
-            );
-
-            return false;
-        }
-
-        // if the `other` FetchStep has a single parent and it's `this` FetchStep
-        if fetch_graph.parents_of(other_index).count() == 1
-            && fetch_graph
-                .parents_of(other_index)
-                .all(|edge| edge.source() == self_index)
-        {
-            return true;
-        }
-
-        // if they do not share parents, they can't be merged
-        if !fetch_graph.parents_of(self_index).all(|self_edge| {
-            fetch_graph
-                .parents_of(other_index)
-                .any(|other_edge| other_edge.source() == self_edge.source())
-        }) {
-            return false;
-        }
-
-        true
+impl<State> FetchStepData<State> {
+    pub fn is_fetching_multiple_types(&self) -> bool {
+        self.input.is_fetching_multiple_types() || self.output.is_fetching_multiple_types()
     }
 }
 

lib/query-planner/src/planner/fetch/optimize/batch_multi_type.rs

@@ -0,0 +1,129 @@
+use std::collections::{HashMap, VecDeque};
+
+use petgraph::{graph::NodeIndex, Direction};
+use tracing::{instrument, trace};
+
+use crate::planner::fetch::{
+    error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepData,
+    optimize::utils::perform_fetch_step_merge, state::MultiTypeFetchStep,
+};
+
+impl FetchGraph<MultiTypeFetchStep> {
+    #[instrument(level = "trace", skip_all)]
+    pub(crate) fn batch_multi_type(&mut self) -> Result<(), FetchGraphError> {
+        let root_index = self
+            .root_index
+            .ok_or(FetchGraphError::NonSingleRootStep(0))?;
+        // Breadth-First Search (BFS) starting from the root node.
+        let mut queue = VecDeque::from([root_index]);
+
+        while let Some(parent_index) = queue.pop_front() {
+            let mut merges_to_perform = Vec::<(NodeIndex, NodeIndex)>::new();
+            let mut node_indexes: HashMap<NodeIndex, NodeIndex> = HashMap::new();
+            let siblings_indices = self
+                .graph
+                .neighbors_directed(parent_index, Direction::Outgoing)
+                .collect::<Vec<NodeIndex>>();
+
+            for (i, sibling_index) in siblings_indices.iter().enumerate() {
+                queue.push_back(*sibling_index);
+                let current = self.get_step_data(*sibling_index)?;
+
+                for other_sibling_index in siblings_indices.iter().skip(i + 1) {
+                    trace!(
+                        "checking if [{}] and [{}] can be batched",
+                        sibling_index.index(),
+                        other_sibling_index.index()
+                    );
+
+                    let other_sibling = self.get_step_data(*other_sibling_index)?;
+
+                    if current.can_be_batched_with(other_sibling) {
+                        trace!(
+                            "Found multi-type batching optimization: [{}] <- [{}]",
+                            sibling_index.index(),
+                            other_sibling_index.index()
+                        );
+                        // Register their original indexes in the map.
+                        node_indexes.insert(*sibling_index, *sibling_index);
+                        node_indexes.insert(*other_sibling_index, *other_sibling_index);
+
+                        merges_to_perform.push((*sibling_index, *other_sibling_index));
+                    }
+                }
+            }
+
+            for (child_index, other_child_index) in merges_to_perform {
+                // Get the latest indexes for the nodes, accounting for previous merges.
+                let child_index_latest = node_indexes
+                    .get(&child_index)
+                    .ok_or(FetchGraphError::IndexMappingLost)?;
+                let other_child_index_latest = node_indexes
+                    .get(&other_child_index)
+                    .ok_or(FetchGraphError::IndexMappingLost)?;
+
+                if child_index_latest == other_child_index_latest {
+                    continue;
+                }
+
+                if self.is_ancestor_or_descendant(*child_index_latest, *other_child_index_latest) {
+                    continue;
+                }
+
+                let (me, other) =
+                    self.get_pair_of_steps_mut(*child_index_latest, *other_child_index_latest)?;
+
+                // We "declare" the known type for the step, so later merging will be possible into that type instead of failing with an error.
+                for (input_type_name, _) in other.input.iter_selections() {
+                    me.input.declare_known_type(input_type_name);
+                }
+
+                // We "declare" the known type for the step, so later merging will be possible into that type instead of failing with an error.
+                for (output_type_name, _) in other.output.iter_selections() {
+                    me.output.declare_known_type(output_type_name);
+                }
+
+                perform_fetch_step_merge(
+                    *child_index_latest,
+                    *other_child_index_latest,
+                    self,
+                    true,
+                )?;
+
+                // Because `other_child` was merged into `child`,
+                // then everything that was pointing to `other_child`
+                // has to point to the `child`.
+                node_indexes.insert(*other_child_index_latest, *child_index_latest);
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl FetchStepData<MultiTypeFetchStep> {
+    pub fn can_be_batched_with(&self, other: &Self) -> bool {
+        if self.kind != other.kind {
+            return false;
+        }
+
+        if self.service_name != other.service_name {
+            return false;
+        }
+
+        if !self.is_entity_call() || !other.is_entity_call() {
+            return false;
+        }
+
+        if self.response_path.without_type_castings() != other.response_path.without_type_castings()
+        {
+            return false;
+        }
+
+        if self.has_arguments_conflicts_with(other) {
+            return false;
+        }
+
+        true
+    }
+}

lib/query-planner/src/planner/fetch/optimize/merge_children_with_parents.rs

@@ -63,7 +63,7 @@ impl FetchGraph<MultiTypeFetchStep> {
                     .get(&child_index)
                     .ok_or(FetchGraphError::IndexMappingLost)?;
 
-                perform_fetch_step_merge(*parent_index_latest, *child_index_latest, self)?;
+                perform_fetch_step_merge(*parent_index_latest, *child_index_latest, self, false)?;
 
                 // Because `child` was merged into `parent`,
                 // then everything that was pointing to `child`

lib/query-planner/src/planner/fetch/optimize/merge_leafs.rs

@@ -1,13 +1,9 @@
 use petgraph::graph::NodeIndex;
 use tracing::{instrument, trace};
 
-use crate::{
-    ast::selection_set::find_arguments_conflicts,
-    planner::fetch::{
-        error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepData,
-        optimize::utils::perform_fetch_step_merge, selections::FetchStepSelections,
-        state::MultiTypeFetchStep,
-    },
+use crate::planner::fetch::{
+    error::FetchGraphError, fetch_graph::FetchGraph, fetch_step_data::FetchStepData,
+    optimize::utils::perform_fetch_step_merge, state::MultiTypeFetchStep,
 };
 
 impl FetchStepData<MultiTypeFetchStep> {
@@ -57,19 +53,7 @@ impl FetchStepData<MultiTypeFetchStep> {
             return false;
         }
 
-        let input_conflicts = FetchStepSelections::<MultiTypeFetchStep>::iter_matching_types(
-            &self.input,
-            &other.input,
-            |_, self_selections, other_selections| {
-                find_arguments_conflicts(self_selections, other_selections)
-            },
-        );
-
-        let has_conflicts = input_conflicts
-            .iter()
-            .any(|(_, conflicts)| !conflicts.is_empty());
-
-        if has_conflicts {
+        if self.has_arguments_conflicts_with(other) {
             return false;
         }
 
@@ -85,7 +69,7 @@ impl FetchGraph<MultiTypeFetchStep> {
     /// meaning the overall depth (amount of parallel layers) is not increased.
     pub(crate) fn merge_leafs(&mut self) -> Result<(), FetchGraphError> {
         while let Some((target_idx, leaf_idx)) = self.find_merge_candidate()? {
-            perform_fetch_step_merge(target_idx, leaf_idx, self)?;
+            perform_fetch_step_merge(target_idx, leaf_idx, self, false)?;
         }
 
         Ok(())

lib/query-planner/src/planner/fetch/optimize/merge_siblings.rs

@@ -92,7 +92,12 @@ impl FetchGraph<MultiTypeFetchStep> {
                     .get(&other_child_index)
                     .ok_or(FetchGraphError::IndexMappingLost)?;
 
-                perform_fetch_step_merge(*child_index_latest, *other_child_index_latest, self)?;
+                perform_fetch_step_merge(
+                    *child_index_latest,
+                    *other_child_index_latest,
+                    self,
+                    false,
+                )?;
 
                 // Because `other_child` was merged into `child`,
                 // then everything that was pointing to `other_child`

lib/query-planner/src/planner/fetch/optimize/mod.rs

@@ -1,4 +1,5 @@
 mod apply_internal_aliases_patching;
+mod batch_multi_type;
 mod deduplicate_and_prune_fetch_steps;
 mod merge_children_with_parents;
 mod merge_leafs;
@@ -29,6 +30,7 @@ impl FetchGraph<MultiTypeFetchStep> {
             self.merge_siblings()?;
             self.merge_leafs()?;
             self.deduplicate_and_prune_fetch_steps()?;
+            self.batch_multi_type()?;
 
             let node_count_after = self.graph.node_count();
             let edge_count_after = self.graph.edge_count();