Add inner hits support to semantic query

docs/changelog/111834.yaml

@@ -0,0 +1,5 @@
+pr: 111834
+summary: Add inner hits support to semantic query
+area: Search
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -188,6 +188,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion ESQL_SINGLE_VALUE_QUERY_SOURCE = def(8_718_00_0);
     public static final TransportVersion ESQL_ORIGINAL_INDICES = def(8_719_00_0);
     public static final TransportVersion ML_INFERENCE_EIS_INTEGRATION_ADDED = def(8_720_00_0);
+    public static final TransportVersion SEMANTIC_QUERY_INNER_HITS = def(8_721_00_0);
     /*
      * STOP! READ THIS FIRST! No, really,
      *        ____ _____ ___  ____  _        ____  _____    _    ____    _____ _   _ ___ ____    _____ ___ ____  ____ _____ _

server/src/main/java/org/elasticsearch/index/query/InnerHitBuilder.java

@@ -49,9 +49,9 @@ public final class InnerHitBuilder implements Writeable, ToXContentObject {
     public static final ParseField COLLAPSE_FIELD = new ParseField("collapse");
     public static final ParseField FIELD_FIELD = new ParseField("field");
 
+    public static final int DEFAULT_FROM = 0;
+    public static final int DEFAULT_SIZE = 3;
     private static final boolean DEFAULT_IGNORE_UNAMPPED = false;
-    private static final int DEFAULT_FROM = 0;
-    private static final int DEFAULT_SIZE = 3;
     private static final boolean DEFAULT_VERSION = false;
     private static final boolean DEFAULT_SEQ_NO_AND_PRIMARY_TERM = false;
     private static final boolean DEFAULT_EXPLAIN = false;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/mapper/SemanticTextFieldMapper.java

@@ -38,6 +38,7 @@
 import org.elasticsearch.index.mapper.ValueFetcher;
 import org.elasticsearch.index.mapper.vectors.DenseVectorFieldMapper;
 import org.elasticsearch.index.mapper.vectors.SparseVectorFieldMapper;
+import org.elasticsearch.index.query.InnerHitBuilder;
 import org.elasticsearch.index.query.MatchNoneQueryBuilder;
 import org.elasticsearch.index.query.NestedQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
@@ -52,6 +53,7 @@
 import org.elasticsearch.xcontent.XContentParserConfiguration;
 import org.elasticsearch.xpack.core.ml.inference.results.MlTextEmbeddingResults;
 import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
+import org.elasticsearch.xpack.inference.queries.InnerChunkBuilder;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -383,7 +385,12 @@ public IndexFieldData.Builder fielddataBuilder(FieldDataContext fieldDataContext
             throw new IllegalArgumentException("[semantic_text] fields do not support sorting, scripting or aggregating");
         }
 
-        public QueryBuilder semanticQuery(InferenceResults inferenceResults, float boost, String queryName) {
+        public QueryBuilder semanticQuery(
+            InferenceResults inferenceResults,
+            float boost,
+            String queryName,
+            InnerChunkBuilder innerChunkBuilder
+        ) {
             String nestedFieldPath = getChunksFieldName(name());
             String inferenceResultsFieldName = getEmbeddingsFieldName(name());
             QueryBuilder childQueryBuilder;
@@ -459,7 +466,10 @@ public QueryBuilder semanticQuery(InferenceResults inferenceResults, float boost
                 };
             }
 
-            return new NestedQueryBuilder(nestedFieldPath, childQueryBuilder, ScoreMode.Max).boost(boost).queryName(queryName);
+            InnerHitBuilder innerHitBuilder = innerChunkBuilder != null ? innerChunkBuilder.toInnerHitBuilder() : null;
+            return new NestedQueryBuilder(nestedFieldPath, childQueryBuilder, ScoreMode.Max).boost(boost)
+                .queryName(queryName)
+                .innerHit(innerHitBuilder);
         }
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/InnerChunkBuilder.java

@@ -0,0 +1,124 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.inference.queries;
+
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.common.io.stream.Writeable;
+import org.elasticsearch.index.query.InnerHitBuilder;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.xcontent.ObjectParser;
+import org.elasticsearch.xcontent.ToXContentObject;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentParser;
+
+import java.io.IOException;
+import java.util.Objects;
+
+import static org.elasticsearch.index.query.InnerHitBuilder.DEFAULT_FROM;
+import static org.elasticsearch.index.query.InnerHitBuilder.DEFAULT_SIZE;
+
+public class InnerChunkBuilder implements Writeable, ToXContentObject {
+    private static final ObjectParser<InnerChunkBuilder, Void> PARSER = new ObjectParser<>("inner_chunks", InnerChunkBuilder::new);
+
+    static {
+        PARSER.declareInt(InnerChunkBuilder::setFrom, SearchSourceBuilder.FROM_FIELD);
+        PARSER.declareInt(InnerChunkBuilder::setSize, SearchSourceBuilder.SIZE_FIELD);
+    }
+
+    private String name;
+    private int from = DEFAULT_FROM;
+    private int size = DEFAULT_SIZE;
+
+    public InnerChunkBuilder() {
+        this.name = null;
+    }
+
+    public InnerChunkBuilder(StreamInput in) throws IOException {
+        name = in.readString();
+        from = in.readVInt();
+        size = in.readVInt();
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeString(name);
+        out.writeVInt(from);
+        out.writeVInt(size);
+    }
+
+    public String getName() {
+        return name;
+    }
+
+    public void setName(String name) {
+        this.name = name;
+    }
+
+    public int getFrom() {
+        return from;
+    }
+
+    public InnerChunkBuilder setFrom(int from) {
+        this.from = from;
+        return this;
+    }
+
+    public int getSize() {
+        return size;
+    }
+
+    public InnerChunkBuilder setSize(int size) {
+        this.size = size;
+        return this;
+    }
+
+    public InnerHitBuilder toInnerHitBuilder() {
+        if (name == null) {
+            throw new IllegalStateException("name must have a value");
+        }
+        return new InnerHitBuilder(name).setFrom(from).setSize(size);
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+        // Don't include name in XContent because it is hard-coded
+        builder.startObject();
+        if (from != DEFAULT_FROM) {
+            builder.field(SearchSourceBuilder.FROM_FIELD.getPreferredName(), from);
+        }
+        if (size != DEFAULT_SIZE) {
+            builder.field(SearchSourceBuilder.SIZE_FIELD.getPreferredName(), size);
+        }
+        builder.endObject();
+        return builder;
+    }
+
+    public static InnerChunkBuilder fromXContent(XContentParser parser) throws IOException {
+        return PARSER.parse(parser, new InnerChunkBuilder(), null);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        InnerChunkBuilder that = (InnerChunkBuilder) o;
+        return from == that.from && size == that.size && Objects.equals(name, that.name);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(name, from, size);
+    }
+
+    @Override
+    public String toString() {
+        return Strings.toString(this, true, true);
+    }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilder.java

@@ -16,6 +16,7 @@
 import org.elasticsearch.cluster.metadata.InferenceFieldMetadata;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
 import org.elasticsearch.index.query.MatchNoneQueryBuilder;
@@ -44,7 +45,9 @@
 import java.util.Map;
 import java.util.Objects;
 
+import static org.elasticsearch.TransportVersions.SEMANTIC_QUERY_INNER_HITS;
 import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg;
+import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg;
 import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN;
 import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin;
 
@@ -53,26 +56,33 @@ public class SemanticQueryBuilder extends AbstractQueryBuilder<SemanticQueryBuil
 
     private static final ParseField FIELD_FIELD = new ParseField("field");
     private static final ParseField QUERY_FIELD = new ParseField("query");
+    private static final ParseField INNER_CHUNKS_FIELD = new ParseField("inner_chunks");
 
     private static final ConstructingObjectParser<SemanticQueryBuilder, Void> PARSER = new ConstructingObjectParser<>(
         NAME,
         false,
-        args -> new SemanticQueryBuilder((String) args[0], (String) args[1])
+        args -> new SemanticQueryBuilder((String) args[0], (String) args[1], (InnerChunkBuilder) args[2])
     );
 
     static {
         PARSER.declareString(constructorArg(), FIELD_FIELD);
         PARSER.declareString(constructorArg(), QUERY_FIELD);
+        PARSER.declareObject(optionalConstructorArg(), (p, c) -> InnerChunkBuilder.fromXContent(p), INNER_CHUNKS_FIELD);
         declareStandardFields(PARSER);
     }
 
     private final String fieldName;
     private final String query;
+    private final InnerChunkBuilder innerChunkBuilder;
     private final SetOnce<InferenceServiceResults> inferenceResultsSupplier;
     private final InferenceResults inferenceResults;
     private final boolean noInferenceResults;
 
     public SemanticQueryBuilder(String fieldName, String query) {
+        this(fieldName, query, null);
+    }
+
+    public SemanticQueryBuilder(String fieldName, String query, @Nullable InnerChunkBuilder innerChunkBuilder) {
         if (fieldName == null) {
             throw new IllegalArgumentException("[" + NAME + "] requires a " + FIELD_FIELD.getPreferredName() + " value");
         }
@@ -81,15 +91,25 @@ public SemanticQueryBuilder(String fieldName, String query) {
         }
         this.fieldName = fieldName;
         this.query = query;
+        this.innerChunkBuilder = innerChunkBuilder;
         this.inferenceResults = null;
         this.inferenceResultsSupplier = null;
         this.noInferenceResults = false;
+
+        if (this.innerChunkBuilder != null) {
+            this.innerChunkBuilder.setName(fieldName);
+        }
     }
 
     public SemanticQueryBuilder(StreamInput in) throws IOException {
         super(in);
         this.fieldName = in.readString();
         this.query = in.readString();
+        if (in.getTransportVersion().onOrAfter(SEMANTIC_QUERY_INNER_HITS)) {
+            this.innerChunkBuilder = in.readOptionalWriteable(InnerChunkBuilder::new);
+        } else {
+            this.innerChunkBuilder = null;
+        }
         this.inferenceResults = in.readOptionalNamedWriteable(InferenceResults.class);
         this.noInferenceResults = in.readBoolean();
         this.inferenceResultsSupplier = null;
@@ -102,6 +122,9 @@ protected void doWriteTo(StreamOutput out) throws IOException {
         }
         out.writeString(fieldName);
         out.writeString(query);
+        if (out.getTransportVersion().onOrAfter(SEMANTIC_QUERY_INNER_HITS)) {
+            out.writeOptionalWriteable(innerChunkBuilder);
+        }
         out.writeOptionalNamedWriteable(inferenceResults);
         out.writeBoolean(noInferenceResults);
     }
@@ -114,13 +137,18 @@ private SemanticQueryBuilder(
     ) {
         this.fieldName = other.fieldName;
         this.query = other.query;
+        this.innerChunkBuilder = other.innerChunkBuilder;
         this.boost = other.boost;
         this.queryName = other.queryName;
         this.inferenceResultsSupplier = inferenceResultsSupplier;
         this.inferenceResults = inferenceResults;
         this.noInferenceResults = noInferenceResults;
     }
 
+    public InnerChunkBuilder innerChunk() {
+        return innerChunkBuilder;
+    }
+
     @Override
     public String getWriteableName() {
         return NAME;
@@ -140,6 +168,9 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep
         builder.startObject(NAME);
         builder.field(FIELD_FIELD.getPreferredName(), fieldName);
         builder.field(QUERY_FIELD.getPreferredName(), query);
+        if (innerChunkBuilder != null) {
+            builder.field(INNER_CHUNKS_FIELD.getPreferredName(), innerChunkBuilder);
+        }
         boostAndQueryNameToXContent(builder);
         builder.endObject();
     }
@@ -166,7 +197,7 @@ private QueryBuilder doRewriteBuildSemanticQuery(SearchExecutionContext searchEx
                 );
             }
 
-            return semanticTextFieldType.semanticQuery(inferenceResults, boost(), queryName());
+            return semanticTextFieldType.semanticQuery(inferenceResults, boost(), queryName(), innerChunkBuilder);
         } else {
             throw new IllegalArgumentException(
                 "Field [" + fieldName + "] of type [" + fieldType.typeName() + "] does not support " + NAME + " queries"
@@ -300,11 +331,12 @@ private static String getInferenceIdForForField(Collection<IndexMetadata> indexM
     protected boolean doEquals(SemanticQueryBuilder other) {
         return Objects.equals(fieldName, other.fieldName)
             && Objects.equals(query, other.query)
+            && Objects.equals(innerChunkBuilder, other.innerChunkBuilder)
             && Objects.equals(inferenceResults, other.inferenceResults);
     }
 
     @Override
     protected int doHashCode() {
-        return Objects.hash(fieldName, query, inferenceResults);
+        return Objects.hash(fieldName, query, innerChunkBuilder, inferenceResults);
     }
 }