Narrow queried indices based on _typename filters.

When a query on an abstract type filters on `__typename`, we can skip
any index that contains none of the requested concrete types. For
example, querying `DistributionChannel` with `_typename: PhysicalStore`
only needs to hit `physical_stores`, not `distribution_channels`.

`TypenameIndexPicker` implements this using `FilterValueSetExtractor`
for full set-algebra support across `not`, `any_of`, and `all_of`
filter combinators.

Generated with Claude Code

Author: marcdaniels-toast

elasticgraph-graphql/lib/elastic_graph/graphql/datastore_query.rb

@@ -28,7 +28,8 @@ class GraphQL
     # with minimal effort.
     class DatastoreQuery < Support::MemoizableData.define(
       :total_document_count_needed, :aggregations, :logger, :filter_interpreter, :routing_picker,
-      :index_expression_builder, :default_page_size, :search_index_definitions, :max_page_size,
+      :index_expression_builder, :typename_index_picker, :default_page_size, :search_index_definitions,
+      :index_definitions_by_possible_typename, :max_page_size,
       :client_filters, :internal_filters, :sort, :document_pagination,
       :requested_fields, :request_all_fields, :requested_highlights, :request_all_highlights,
       :individual_docs_needed, :size_multiplier, :monotonic_clock_deadline, :schema_element_names
@@ -39,6 +40,7 @@ class DatastoreQuery < Support::MemoizableData.define(
       require "elastic_graph/graphql/datastore_query/index_expression_builder"
       require "elastic_graph/graphql/datastore_query/paginator"
       require "elastic_graph/graphql/datastore_query/routing_picker"
+      require "elastic_graph/graphql/datastore_query/typename_index_picker"
 
       # Performs a list of queries by building a hash of datastore msearch header/body tuples (keyed
       # by query), yielding them to the caller, and then post-processing the results. The caller is
@@ -300,6 +302,14 @@ def ignored_values_for_routing
         @ignored_values_for_routing ||= search_index_definitions.flat_map { |i| i.ignored_values_for_routing.to_a }.to_set
       end
 
+      def search_index_definitions
+        # Narrows the stored index definitions based on a __typename filter. `super` calls the
+        # generated Data accessor to get the value passed to `new_query` as `search_index_definitions:`
+        # at construction time, before narrowing is applied.
+        @narrowed_search_index_definitions ||=
+          typename_index_picker.pick_index_definitions(client_filters.to_a, index_definitions_by_possible_typename, super)
+      end
+
       def to_datastore_body
         @to_datastore_body ||= aggregations_datastore_body
           .merge(document_paginator.to_datastore_body)
@@ -360,8 +370,16 @@ def index_expression_builder
           )
         end
 
+        def typename_index_picker
+          @typename_index_picker ||= TypenameIndexPicker.new(
+            filter_node_interpreter: filter_node_interpreter,
+            schema_names: runtime_metadata.schema_element_names
+          )
+        end
+
         def new_query(
           search_index_definitions:,
+          index_definitions_by_possible_typename: {},
           client_filters: [],
           internal_filters: [],
           sort: [],
@@ -388,9 +406,11 @@ def new_query(
           DatastoreQuery.new(
             routing_picker: routing_picker,
             index_expression_builder: index_expression_builder,
+            typename_index_picker: typename_index_picker,
             logger: logger,
             schema_element_names: runtime_metadata.schema_element_names,
             search_index_definitions: search_index_definitions,
+            index_definitions_by_possible_typename: index_definitions_by_possible_typename,
             client_filters: client_filters.to_set,
             internal_filters: internal_filters.to_set,
             sort: sort,

elasticgraph-graphql/lib/elastic_graph/graphql/datastore_query/typename_index_picker.rb

@@ -0,0 +1,160 @@
+# Copyright 2024 - 2026 Block, Inc.
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+#
+# frozen_string_literal: true
+
+require "elastic_graph/graphql/filtering/filter_value_set_extractor"
+
+module ElasticGraph
+  class GraphQL
+    class DatastoreQuery
+      # Responsible for narrowing the set of searched indices based on a `__typename` filter.
+      # When a query filters on `__typename`, we can skip any index that contains none of the
+      # requested concrete types, avoiding unnecessary datastore queries.
+      #
+      # Mirrors the structure of `RoutingPicker`, which performs a similar narrowing for shards.
+      class TypenameIndexPicker
+        def initialize(filter_node_interpreter:, schema_names:)
+          @filter_value_set_extractor = Filtering::FilterValueSetExtractor.new(
+            filter_node_interpreter,
+            schema_names,
+            TypenameValueSet::ALL,
+            TypenameValueSet::EMPTY
+          ) do |operator, filter_value|
+            if operator == :equal_to_any_of
+              TypenameValueSet.of(filter_value.compact)
+            end
+          end
+        end
+
+        # Given a list of `filter_hashes`, a map of `{ type_name => [index_definition] }`, and the
+        # full set of `search_index_definitions`, returns the narrowed set of index definitions that
+        # need to be searched. Returns `search_index_definitions` unchanged if there is no `__typename`
+        # filter or if no type mapping is available.
+        def pick_index_definitions(filter_hashes, index_definitions_by_possible_typename, search_index_definitions)
+          # An empty map means the caller did not opt into typename-based narrowing (e.g. a direct
+          # `new_query` call that omits `index_definitions_by_possible_typename`)
+          return search_index_definitions if index_definitions_by_possible_typename.empty?
+
+          typename_set = @filter_value_set_extractor.extract_filter_value_set(filter_hashes, ["__typename"])
+          return search_index_definitions if typename_set.nil?
+
+          index_definitions_by_possible_typename.flat_map do |type_name, index_defs|
+            typename_set.include?(type_name) ? index_defs : []
+          end
+        end
+      end
+
+      class TypenameValueSet < Data.define(:type, :type_names)
+        def self.of(type_names)
+          new(:inclusive, type_names.to_set)
+        end
+
+        def self.of_all_except(type_names)
+          new(:exclusive, type_names.to_set)
+        end
+
+        ALL = of_all_except([])
+        EMPTY = of([])
+
+        def include?(type_name)
+          inclusive? ? type_names.include?(type_name) : !type_names.include?(type_name)
+        end
+
+        def intersection(other)
+          if inclusive? && other.inclusive?
+            # Since both sets are inclusive, we can just delegate to `Set#intersection` here.
+            TypenameValueSet.of(type_names.intersection(other.type_names))
+          elsif exclusive? && other.exclusive?
+            # Since both sets are exclusive, we need to return an exclusive set of the union of the
+            # excluded values. For example:
+            #
+            #   s1 = TypenameValueSet.of_all_except(["Widget"])   # everything except Widget
+            #   s2 = TypenameValueSet.of_all_except(["Address"])  # everything except Address
+            #
+            #   s3 = s1.intersection(s2)
+            #
+            # Here s3 would be everything except Widget or Address (the same as `of_all_except(["Widget", "Address"])`)
+            TypenameValueSet.of_all_except(type_names.union(other.type_names))
+          else
+            # Since one set is inclusive and one is exclusive, we need an inclusive set of
+            # `included_names - excluded_names`. For example:
+            #
+            #   s1 = TypenameValueSet.of(["Widget", "Address"])   # Widget, Address
+            #   s2 = TypenameValueSet.of_all_except(["Address"])  # everything except Address
+            #
+            #   s3 = s1.intersection(s2)
+            #
+            # Here s3 would be just Widget.
+            included_names = inclusive? ? type_names : other.type_names
+            excluded_names = inclusive? ? other.type_names : type_names
+            TypenameValueSet.of(included_names - excluded_names)
+          end
+        end
+
+        def union(other)
+          if inclusive? && other.inclusive?
+            # Since both sets are inclusive, we can just delegate to `Set#union` here.
+            TypenameValueSet.of(type_names.union(other.type_names))
+          elsif exclusive? && other.exclusive?
+            # Since both sets are exclusive, only names excluded from *both* are excluded from the union.
+            # For example:
+            #
+            #   s1 = TypenameValueSet.of_all_except(["Widget", "Address"])    # everything except Widget, Address
+            #   s2 = TypenameValueSet.of_all_except(["Address", "Manufacturer"]) # everything except Address, Manufacturer
+            #
+            #   s3 = s1.union(s2)
+            #
+            # Widget is in s2 and Manufacturer is in s1, so both appear in the union.
+            # Only Address is excluded from both, so s3 is everything except Address
+            # (the same as `of_all_except(["Address"])`, i.e. the intersection of the exclusion sets).
+            TypenameValueSet.of_all_except(type_names.intersection(other.type_names))
+          elsif inclusive?
+            # s1 is inclusive, s2 is exclusive: return exclusive set of `excluded - included`.
+            # For example:
+            #
+            #   s1 = TypenameValueSet.of(["Widget"])                             # Widget
+            #   s2 = TypenameValueSet.of_all_except(["Widget", "Address"])       # everything except Widget, Address
+            #
+            #   s3 = s1.union(s2)
+            #
+            # Widget is already in s1, so it appears in the union.
+            # Address is excluded from s2 and not in s1, so it remains excluded.
+            # Here s3 is everything except Address (the same as `of_all_except(["Address"])`)
+            TypenameValueSet.of_all_except(other.type_names - type_names)
+          else
+            # s1 is exclusive, s2 is inclusive: symmetric to the above.
+            TypenameValueSet.of_all_except(type_names - other.type_names)
+          end
+        end
+
+        def negate
+          with(type: INVERTED_TYPES.fetch(type))
+        end
+
+        INVERTED_TYPES = {inclusive: :exclusive, exclusive: :inclusive}
+
+        protected
+
+        def inclusive?
+          type == :inclusive
+        end
+
+        def exclusive?
+          type == :exclusive
+        end
+      end
+
+      # `TypenameIndexPicker` exists only for use by `DatastoreQuery` and is effectively private.
+      private_constant :TypenameIndexPicker
+      # `TypenameValueSet` exists only for use here and is effectively private.
+      private_constant :TypenameValueSet
+
+      # Steep is complaining that it can't find some `DatastoreQuery` methods but they are not in this file...
+      # @dynamic shard_routing_values, effective_size, merge_with, search_index_expression, with, to_datastore_msearch_header_and_body
+    end
+  end
+end

elasticgraph-graphql/lib/elastic_graph/graphql/resolvers/query_adapter.rb

@@ -65,6 +65,7 @@ def build_new_query_from(field, args, lookahead, context, monotonic_clock_deadli
 
           initial_query = @datastore_query_builder.new_query(
             search_index_definitions: unwrapped_type.search_index_definitions,
+            index_definitions_by_possible_typename: unwrapped_type.index_definitions_by_possible_typename,
             monotonic_clock_deadline: monotonic_clock_deadline
           )
 

elasticgraph-graphql/lib/elastic_graph/graphql/schema/type.rb

@@ -134,6 +134,21 @@ def subtypes
             .to_set
         end
 
+        # Returns a map of each possible concrete typename to its index definitions.
+        # For abstract types, maps each concrete subtype; for concrete types, maps this type itself.
+        def index_definitions_by_possible_typename
+          @index_definitions_by_possible_typename ||=
+            if (st = source_type)
+              # Derived types (e.g. indexed aggregations) share their source document type's indices.
+              # Delegate so that _typename filters match concrete document typenames, not the derived type name.
+              st.index_definitions_by_possible_typename
+            elsif abstract?
+              subtypes.to_h { |subtype| [subtype.name, subtype.search_index_definitions] }
+            else
+              {name => search_index_definitions}
+            end
+        end
+
         # For derived types (e.g. indexed aggregations), returns the underlying source document type.
         # Returns `nil` for non-derived types.
         def source_type

elasticgraph-graphql/sig/elastic_graph/graphql/datastore_query/typename_index_picker.rbs

@@ -0,0 +1,55 @@
+module ElasticGraph
+  class GraphQL
+    class DatastoreQuery
+      class TypenameIndexPicker
+        def initialize: (
+          filter_node_interpreter: Filtering::FilterNodeInterpreter,
+          schema_names: SchemaArtifacts::RuntimeMetadata::SchemaElementNames
+        ) -> void
+
+        def pick_index_definitions: (
+          ::Array[::Hash[::String, untyped]],
+          ::Hash[::String, ::Array[DatastoreCore::_IndexDefinition]],
+          ::Array[DatastoreCore::_IndexDefinition]
+        ) -> ::Array[DatastoreCore::_IndexDefinition]
+
+        private
+
+        @filter_value_set_extractor: Filtering::FilterValueSetExtractor[TypenameValueSet]
+      end
+
+      type typenameValueSetType = :inclusive | :exclusive
+
+      class TypenameValueSetSupertype
+        attr_reader type: typenameValueSetType
+        attr_reader type_names: ::Set[::String]
+
+        def initialize: (typenameValueSetType, ::Set[::String]) -> void
+        def self.with: (
+          type: typenameValueSetType,
+          type_names: ::Set[::String]
+        ) -> TypenameValueSet
+
+        def with: (
+          ?type: typenameValueSetType,
+          ?type_names: ::Set[::String]
+        ) -> TypenameValueSet
+      end
+
+      class TypenameValueSet < TypenameValueSetSupertype
+        include Support::_NegatableSet[TypenameValueSet]
+        def self.of: (::Enumerable[::String]) -> TypenameValueSet
+        def self.of_all_except: (::Enumerable[::String]) -> TypenameValueSet
+
+        ALL: TypenameValueSet
+        EMPTY: TypenameValueSet
+        INVERTED_TYPES: ::Hash[typenameValueSetType, typenameValueSetType]
+
+        def include?: (::String) -> bool
+
+        def inclusive?: () -> bool
+        def exclusive?: () -> bool
+      end
+    end
+  end
+end

elasticgraph-graphql/sig/elastic_graph/graphql/schema/type.rbs

@@ -11,6 +11,7 @@ module ElasticGraph
         def search_index_definitions: () -> ::Array[DatastoreCore::_IndexDefinition]
         def source_type: () -> Type?
         def subtypes: () -> ::Set[Type]
+        def index_definitions_by_possible_typename: () -> ::Hash[::String, ::Array[DatastoreCore::_IndexDefinition]]
         def shares_index_with_non_subtypes?: () -> bool
         def unwrap_fully: () -> Type
         def field_named: (::String) -> Field

elasticgraph-graphql/spec/unit/elastic_graph/graphql/datastore_query/merge_spec.rb

@@ -21,7 +21,7 @@ class GraphQL
         # These are derived from app state and don't vary in two different queries for the same app,
         # so we don't have to deal with merging them.
         app_level_attributes = %i[
-          logger filter_interpreter routing_picker index_expression_builder
+          logger filter_interpreter routing_picker index_expression_builder typename_index_picker
           default_page_size max_page_size schema_element_names
         ]
 
@@ -50,6 +50,16 @@ class GraphQL
         }.to raise_error ArgumentError, a_string_including("search_index_definitions")
       end
 
+      it "does not allow `index_definitions_by_possible_typename` to be overridden", covers: :index_definitions_by_possible_typename do
+        widgets_def = graphql.datastore_core.index_definitions_by_name.fetch("widgets")
+
+        query = new_query(index_definitions_by_possible_typename: {"Widget" => [widgets_def]})
+
+        expect {
+          query.merge_with(index_definitions_by_possible_typename: {"Widget" => [widgets_def]})
+        }.to raise_error ArgumentError, a_string_including("index_definitions_by_possible_typename")
+      end
+
       %i[client_filters internal_filters].each do |filter_attr|
         describe ":#{filter_attr}", covers: filter_attr do
           it "can merge `equal_to_any_of` conditions from two separate queries that are on separate fields" do