ES|QL - Remove vectors from _source when applicable (#138013)

Author: carlosdelest

docs/changelog/138013.yaml

@@ -0,0 +1,5 @@
+pr: 138013
+summary: ES|QL - Remove vectors from `_source` when applicable
+area: "ES|QL"
+type: enhancement
+issues: []

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/DenseVectorFieldTypeIT.java

@@ -20,6 +20,7 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentFactory;
 import org.elasticsearch.xpack.esql.action.AbstractEsqlIntegTestCase;
+import org.hamcrest.Matchers;
 import org.junit.Before;
 
 import java.io.IOException;
@@ -40,6 +41,7 @@
 import static org.elasticsearch.index.mapper.SourceFieldMapper.Mode.SYNTHETIC;
 import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
 import static org.elasticsearch.xpack.esql.action.EsqlCapabilities.Cap.L2_NORM_VECTOR_SIMILARITY_FUNCTION;
+import static org.hamcrest.Matchers.hasKey;
 
 public class DenseVectorFieldTypeIT extends AbstractEsqlIntegTestCase {
 
@@ -68,7 +70,6 @@ private enum VectorSourceOptions {
     @ParametersFactory
     public static Iterable<Object[]> parameters() throws Exception {
         List<Object[]> params = new ArrayList<>();
-
         for (ElementType elementType : List.of(ElementType.BYTE, ElementType.FLOAT, ElementType.BIT)) {
             // Test all similarities
             for (DenseVectorFieldMapper.VectorSimilarity similarity : DenseVectorFieldMapper.VectorSimilarity.values()) {
@@ -156,7 +157,6 @@ public void testRetrieveDenseVectorFieldData() {
         try (var resp = run(query)) {
             List<List<Object>> valuesList = EsqlTestUtils.getValuesList(resp);
             assertEquals(valuesList.size(), indexedVectors.size());
-            // print all values for debugging
             valuesList.forEach(value -> {
                 assertEquals(2, value.size());
                 Integer id = (Integer) value.get(0);
@@ -180,6 +180,33 @@ public void testRetrieveDenseVectorFieldData() {
         }
     }
 
+    @SuppressWarnings("unchecked")
+    public void testDenseVectorsIncludedInSource() {
+        var query = """
+            FROM test METADATA _source
+            | KEEP _source
+            """;
+
+        try (var resp = run(query)) {
+            List<List<Object>> valuesList = EsqlTestUtils.getValuesList(resp);
+            assertEquals(valuesList.size(), indexedVectors.size());
+            valuesList.forEach(value -> {
+                assertEquals(1, value.size());
+                Map<String, Object> source = (Map<String, Object>) value.get(0);
+                assertThat(source, hasKey("id"));
+                assertNotNull(source.get("id"));
+                Integer id = Integer.valueOf(source.get("id").toString());
+                // Vectors should be in _source if they are included in the index settings, and the vector is not null
+                if (sourceOptions == VectorSourceOptions.INCLUDE_SOURCE_VECTORS && indexedVectors.get(id) != null) {
+                    assertThat(source, hasKey("vector"));
+                    assertNotNull(source.get("vector"));
+                } else {
+                    assertThat(source, Matchers.aMapWithSize(1));
+                }
+            });
+        }
+    }
+
     public void testNonIndexedDenseVectorField() throws IOException {
         createIndexWithDenseVector("no_dense_vectors", 64);
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/planner/EsPhysicalOperationProviders.java

@@ -61,6 +61,7 @@
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 import org.elasticsearch.search.fetch.StoredFieldsSpec;
+import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
 import org.elasticsearch.search.internal.AliasFilter;
 import org.elasticsearch.search.lookup.SearchLookup;
 import org.elasticsearch.search.lookup.SourceFilter;
@@ -97,6 +98,7 @@
 
 import static org.elasticsearch.common.lucene.search.Queries.newNonNestedFilter;
 import static org.elasticsearch.compute.lucene.LuceneSourceOperator.NO_LIMIT;
+import static org.elasticsearch.index.get.ShardGetService.maybeExcludeVectorFields;
 
 public class EsPhysicalOperationProviders extends AbstractPhysicalOperationProviders {
     private static final Logger logger = LogManager.getLogger(EsPhysicalOperationProviders.class);
@@ -464,6 +466,13 @@ public String shardIdentifier() {
         @Override
         public SourceLoader newSourceLoader(Set<String> sourcePaths) {
             var filter = sourcePaths != null ? new SourceFilter(sourcePaths.toArray(new String[0]), null) : null;
+            // Apply vector exclusion logic similar to ShardGetService
+            var fetchSourceContext = filter != null ? FetchSourceContext.of(true, null, filter.getIncludes(), filter.getExcludes()) : null;
+            var result = maybeExcludeVectorFields(ctx.getMappingLookup(), ctx.getIndexSettings(), fetchSourceContext, null);
+            var vectorFilter = result.v2();
+            if (vectorFilter != null) {
+                filter = vectorFilter;
+            }
             return ctx.newSourceLoader(filter, false);
         }