Enhancements to the aggregation-based eager load mechanism

lib/mongoid/association/eager.rb

@@ -4,6 +4,15 @@ module Mongoid
   module Association
     # Base class for eager load preload functions.
     class Eager
+      # Build a preloader for the given arguments and run it.
+      #
+      # @param (see #initialize)
+      #
+      # @return [ Array ] The list of documents given.
+      def self.run(associations, docs, use_lookup = false, pipeline = [])
+        new(associations, docs, use_lookup, pipeline).run
+      end
+
       # Instantiate the eager load class.
       #
       # @example Create the new belongs to eager load preloader.

lib/mongoid/association/eager_load/discriminated_inclusion.rb

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+require 'mongoid/association/eager_load/inclusion'
+
+module Mongoid
+  module Association
+    module EagerLoad
+      # Loads an inclusion that more than one subclass defines under the same name
+      # but pointing at different targets. A single $lookup can't serve them: they
+      # would all write to the same field and overwrite one another. So each
+      # subclass's inclusion is contributed into its own temporary field, carrying
+      # its own nested children, and a $set then routes every document to the field
+      # for its own type, by the discriminator.
+      #
+      # For Machine.eager_load(:widgets), where Lathe#widgets => Cog and
+      # Press#widgets => Belt, it emits:
+      #
+      #   { '$lookup' => { 'from' => 'cogs',  ..., 'as' => '__eager_load_widgets_Lathe' } },
+      #   { '$lookup' => { 'from' => 'belts', ..., 'as' => '__eager_load_widgets_Press' } },
+      #   { '$set' => {
+      #     'widgets' => { '$switch' => { 'branches' => [   # route each document to its
+      #       { 'case' => { '$eq' => [ '$_type', 'Lathe' ] }, 'then' => '$__eager_load_widgets_Lathe' },  # own type's matches
+      #       { 'case' => { '$eq' => [ '$_type', 'Press' ] }, 'then' => '$__eager_load_widgets_Press' }
+      #     ], 'default' => [] } }
+      #   } },
+      #   { '$unset' => [ '__eager_load_widgets_Lathe', '__eager_load_widgets_Press' ] }
+      #
+      # @api private
+      class DiscriminatedInclusion < Inclusion
+        def initialize(nodes)
+          super()
+          @nodes = nodes
+        end
+
+        # Append each subclass's lookup (into its own temporary field), the routing
+        # $set, and the cleanup $unset.
+        #
+        # @param [ Array<Hash> ] destination The pipeline the stages are appended to.
+        def contribute(destination, _chain)
+          fields = @nodes.map { |node| [ node, contribute_into_temporary(destination, node) ] }
+          destination << route_by_type(fields)
+          destination << { '$unset' => fields.map { |_node, field| field } }
+        end
+
+        private
+
+        # Let the node build its own $lookup (with its nested children) and redirect
+        # it to write into a temporary field instead of the shared association name.
+        def contribute_into_temporary(destination, node)
+          field = temporary_field(node)
+          captured = []
+          node.contribute(captured, [])
+          captured.first['$lookup']['as'] = field
+          destination.concat(captured)
+          field
+        end
+
+        def temporary_field(node)
+          "__eager_load_#{node.association.name}_#{owner(node).discriminator_value}"
+        end
+
+        # The $set that fills the association on each document from the temporary
+        # field matching its own type.
+        def route_by_type(fields)
+          branches = fields.map do |node, field|
+            {
+              'case' => { '$eq' => [ "$#{discriminator_key}", owner(node).discriminator_value ] },
+              'then' => "$#{field}"
+            }
+          end
+          { '$set' => { name => { '$switch' => { 'branches' => branches, 'default' => [] } } } }
+        end
+
+        def owner(node)
+          node.association.inverse_class
+        end
+
+        def name
+          @nodes.first.association.name.to_s
+        end
+
+        def discriminator_key
+          owner(@nodes.first).discriminator_key
+        end
+      end
+    end
+  end
+end

lib/mongoid/association/eager_load/embedded_distributor.rb

@@ -0,0 +1,151 @@
+# frozen_string_literal: true
+
+module Mongoid
+  module Association
+    # Objects that carry out the aggregation-based eager load triggered by
+    # Criteria#eager_load.
+    module EagerLoad
+      # Distributes the results of a $lookup onto the embedded documents they
+      # belong to.
+      #
+      # A $lookup overwrites the single field it writes to, and it can't distribute
+      # its matches across the elements of an embedded array. So when the reference
+      # being eager-loaded lives on an embedded document, the matches are first
+      # collected in a temporary top-level field and then distributed down the
+      # embedded path onto each embedded document, merging into it (so the rest of
+      # the document is kept) and correlating by key. The temporary field is then
+      # dropped.
+      #
+      # For Computer.eager_load(port: :device) (Port belongs_to :device) it emits:
+      #
+      #   { '$lookup' => {                     # devices can't be written into
+      #     'from' => 'devices',               # the embedded port, so they are
+      #     'localField' => 'port.device_id',  # collected in a temp top-level
+      #     'foreignField' => '_id',           # field instead
+      #     'as' => '__eager_load_port_device'
+      #   } },
+      #   { '$set' => {
+      #     'port' => { '$mergeObjects' => [   # merge a 'device' key onto the port
+      #       '$port',
+      #       { 'device' => { '$filter' => { ... } } }   # this port's matches
+      #     ] }
+      #   } },
+      #   { '$unset' => '__eager_load_port_device' }  # drop the temp field
+      #
+      # @api private
+      class EmbeddedDistributor
+        # @param [ Mongoid::Association::Relatable ] association The referenced
+        #   inclusion being eager-loaded from within an embedded document.
+        # @param [ Array<Mongoid::Association::Relatable> ] chain The embedded
+        #   ancestors, from the root document inward, down to the association's owner.
+        # @param [ Hash ] lookup_stage The $lookup stage built for the association.
+        #
+        # @return [ EmbeddedDistributor ] The distributor.
+        class << self
+          def for(association:, chain:, lookup_stage:)
+            lookup = lookup_stage['$lookup']
+            new(association, chain, lookup_stage, lookup['localField'], lookup['foreignField'])
+          end
+
+          private :new
+        end
+
+        def initialize(association, chain, lookup_stage, local_field, foreign_field)
+          @association = association
+          @chain = chain
+          @lookup_stage = lookup_stage
+          @local_field = local_field
+          @foreign_field = foreign_field
+        end
+
+        # The stages that run the $lookup into a temporary field and then
+        # distribute its matches onto the embedded documents along the path.
+        #
+        # @return [ Array<Hash> ] The stages to append to the pipeline.
+        def stages
+          redirect_lookup_to_temporary_field
+          [
+            @lookup_stage,
+            { '$set' => {
+              root => distributed_value(@chain, "$#{root}")
+            } },
+            { '$unset' => temporary_field }
+          ]
+        end
+
+        private
+
+        # The $lookup runs at the top level, so it reads the local field by its
+        # full embedded path and writes the matches into the temporary field.
+        def redirect_lookup_to_temporary_field
+          lookup = @lookup_stage['$lookup']
+          lookup['localField'] = "#{path}.#{@local_field}"
+          lookup['as'] = temporary_field
+        end
+
+        def path
+          @chain.map(&:store_as).join('.')
+        end
+
+        def root
+          @chain.first.store_as
+        end
+
+        def temporary_field
+          "__eager_load_#{path.tr('.', '_')}_#{@association.name}"
+        end
+
+        # An embedded collection (embeds_many) is rebuilt with $map so each element
+        # keeps its own matches instead of collapsing onto the first; a single
+        # embedded document (embeds_one) receives its matches in place.
+        def distributed_value(chain, node)
+          head, *rest = chain
+          many = head.many?
+          element = many ? "$$#{head.store_as}" : node
+          child =
+            if rest.empty?
+              { @association.name.to_s => correlated_matches(element) }
+            else
+              segment = rest.first.store_as
+              { segment => distributed_value(rest, "#{element}.#{segment}") }
+            end
+          merged = { '$mergeObjects' => [ element, child ] }
+          return merge_into_present(node, merged) unless many
+
+          { '$map' => {
+            'input' => node,
+            'as' => head.store_as,
+            'in' => merged
+          } }
+        end
+
+        # Merge the matches into a single embedded document only when it exists, so
+        # an absent embeds_one stays absent instead of being synthesized from its
+        # matches alone.
+        def merge_into_present(node, merged)
+          { '$cond' => {
+            'if' => { '$ifNull' => [ node, false ] },
+            'then' => merged,
+            'else' => node
+          } }
+        end
+
+        # The matches that belong to a single embedded element.
+        def correlated_matches(element)
+          { '$filter' => {
+            'input' => "$#{temporary_field}",
+            'as' => 'match',
+            'cond' => { match_operator => [ "$$match.#{@foreign_field}", "#{element}.#{@local_field}" ] }
+          } }
+        end
+
+        # A has_and_belongs_to_many holds an array of foreign keys, so a match
+        # belongs when its key is among them ($in); every other association points
+        # at a single key ($eq).
+        def match_operator
+          @association.many_to_many? ? '$in' : '$eq'
+        end
+      end
+    end
+  end
+end