Merge branch 'main' into esql_ts_agg_detect

Author: nik9000

docs/reference/elasticsearch/rest-apis/retrievers/linear-retriever.md

@@ -91,3 +91,37 @@ The `linear` retriever supports the following normalizers:
     score = (score - min) / (max - min)
     ```
 * `l2_norm`: Normalizes scores using the L2 norm of the score values {applies_to}`stack: ga 9.1`
+
+## Example
+
+This example of a hybrid search weights KNN results five times more heavily than BM25 results in the final ranking.
+
+```console
+GET my_index/_search
+{
+  "retriever": {
+    "linear": {
+      "retrievers": [
+        {
+          "retriever": {
+            "knn": {
+              ...
+            }
+          },
+          "weight": 5 # KNN query weighted 5x
+        },
+        {
+          "retriever": {
+            "standard": {
+              ...
+            }
+          },
+          "weight": 1.5 # BM25 query weighted 1.5x
+        }
+      ]
+    }
+  }
+}
+```
+
+See also [this hybrid search example](retrievers-examples.md#retrievers-examples-linear-retriever).

test/framework/src/main/java/org/elasticsearch/datageneration/matchers/source/SourceTransforms.java

@@ -18,7 +18,7 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
-class SourceTransforms {
+public class SourceTransforms {
     /**
      * This preprocessing step makes it easier to match the document using a unified structure.
      * It performs following modifications:

test/framework/src/main/java/org/elasticsearch/datageneration/queries/LeafQueryGenerator.java

@@ -0,0 +1,109 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.datageneration.queries;
+
+import org.elasticsearch.datageneration.FieldType;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.QueryBuilders;
+import org.elasticsearch.test.ESTestCase;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+
+public interface LeafQueryGenerator {
+
+    List<QueryBuilder> generate(Map<String, Object> fieldMapping, String path, Object value);
+
+    /**
+     * Build a query for a specific type. If the field is nested, this query will need to be wrapped in nested queries.
+     * @param type the type to build a query for
+     * @return a generator that can build queries for this type
+     */
+    static LeafQueryGenerator buildForType(String type) {
+        LeafQueryGenerator noQueries = (Map<String, Object> fieldMapping, String path, Object value) -> List.of();
+
+        FieldType fieldType = FieldType.tryParse(type);
+        if (fieldType == null) {
+            return noQueries;
+        }
+
+        return switch (fieldType) {
+            case KEYWORD -> new KeywordQueryGenerator();
+            case TEXT -> new TextQueryGenerator();
+            case WILDCARD -> new WildcardQueryGenerator();
+            default -> noQueries;
+        };
+    }
+
+    class KeywordQueryGenerator implements LeafQueryGenerator {
+        public List<QueryBuilder> generate(Map<String, Object> fieldMapping, String path, Object value) {
+            if (fieldMapping != null) {
+                boolean isIndexed = (Boolean) fieldMapping.getOrDefault("index", true);
+                boolean hasDocValues = (Boolean) fieldMapping.getOrDefault("doc_values", true);
+                if (isIndexed == false && hasDocValues == false) {
+                    return List.of();
+                }
+            }
+            return List.of(QueryBuilders.termQuery(path, value));
+        }
+    }
+
+    class WildcardQueryGenerator implements LeafQueryGenerator {
+        public List<QueryBuilder> generate(Map<String, Object> fieldMapping, String path, Object value) {
+            // Queries with emojis can currently fail due to https://github.com/elastic/elasticsearch/issues/132144
+            if (containsHighSurrogates((String) value)) {
+                return List.of();
+            }
+            return List.of(QueryBuilders.termQuery(path, value), QueryBuilders.wildcardQuery(path, value + "*"));
+        }
+    }
+
+    class TextQueryGenerator implements LeafQueryGenerator {
+        public List<QueryBuilder> generate(Map<String, Object> fieldMapping, String path, Object value) {
+            if (fieldMapping != null) {
+                boolean isIndexed = (Boolean) fieldMapping.getOrDefault("index", true);
+                if (isIndexed == false) {
+                    return List.of();
+                }
+            }
+
+            var results = new ArrayList<QueryBuilder>();
+            results.add(QueryBuilders.matchQuery(path, value));
+            var phraseQuery = buildPhraseQuery(path, (String) value);
+            if (phraseQuery != null) {
+                results.add(phraseQuery);
+            }
+            return results;
+        }
+
+        private static QueryBuilder buildPhraseQuery(String path, String value) {
+            var tokens = Arrays.asList(value.split("[^a-zA-Z0-9]"));
+            if (tokens.isEmpty()) {
+                return null;
+            }
+
+            int low = ESTestCase.randomIntBetween(0, tokens.size() - 1);
+            int hi = ESTestCase.randomIntBetween(low + 1, tokens.size());
+            var phrase = String.join(" ", tokens.subList(low, hi));
+            return QueryBuilders.matchPhraseQuery(path, phrase);
+        }
+    }
+
+    static boolean containsHighSurrogates(String s) {
+        for (int i = 0; i < s.length(); i++) {
+            if (Character.isHighSurrogate(s.charAt(i))) {
+                return true;
+            }
+        }
+        return false;
+    }
+}

test/framework/src/main/java/org/elasticsearch/datageneration/queries/QueryGenerator.java

@@ -0,0 +1,98 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.datageneration.queries;
+
+import org.apache.lucene.search.join.ScoreMode;
+import org.elasticsearch.datageneration.Mapping;
+import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.QueryBuilders;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+
+public class QueryGenerator {
+
+    private final Mapping mapping;
+
+    public QueryGenerator(Mapping mapping) {
+        this.mapping = mapping;
+    }
+
+    public List<QueryBuilder> generateQueries(String type, String path, Object value) {
+        // This query generator cannot handle fields with periods in the name.
+        if (path.equals("host.name")) {
+            return List.of();
+        }
+        // Can handle dynamically mapped fields, but not runtime fields
+        if (isRuntimeField(path)) {
+            return List.of();
+        }
+        var leafQueryGenerator = LeafQueryGenerator.buildForType(type);
+        var fieldMapping = mapping.lookup().get(path);
+        var leafQueries = leafQueryGenerator.generate(fieldMapping, path, value);
+        return leafQueries.stream().map(q -> wrapInNestedQuery(path, q)).toList();
+    }
+
+    private QueryBuilder wrapInNestedQuery(String path, QueryBuilder leafQuery) {
+        String[] parts = path.split("\\.");
+        List<String> nestedPaths = getNestedPathPrefixes(parts);
+        QueryBuilder query = leafQuery;
+        for (String nestedPath : nestedPaths.reversed()) {
+            query = QueryBuilders.nestedQuery(nestedPath, query, ScoreMode.Max);
+        }
+        return query;
+    }
+
+    @SuppressWarnings("unchecked")
+    private List<String> getNestedPathPrefixes(String[] path) {
+        Map<String, Object> mapping = this.mapping.raw();
+        mapping = (Map<String, Object>) mapping.get("_doc");
+        mapping = (Map<String, Object>) mapping.get("properties");
+
+        var result = new ArrayList<String>();
+        for (int i = 0; i < path.length - 1; i++) {
+            var field = path[i];
+            mapping = (Map<String, Object>) mapping.get(field);
+
+            // dynamic field
+            if (mapping == null) {
+                break;
+            }
+
+            boolean nested = "nested".equals(mapping.get("type"));
+            if (nested) {
+                result.add(String.join(".", Arrays.copyOfRange(path, 0, i + 1)));
+            }
+            mapping = (Map<String, Object>) mapping.get("properties");
+        }
+        return result;
+    }
+
+    @SuppressWarnings("unchecked")
+    private boolean isRuntimeField(String path) {
+        String[] parts = path.split("\\.");
+        var topLevelMapping = (Map<String, Object>) mapping.raw().get("_doc");
+        boolean inRuntimeContext = "runtime".equals(topLevelMapping.get("dynamic"));
+        for (int i = 0; i < parts.length - 1; i++) {
+            var pathToHere = String.join(".", Arrays.copyOfRange(parts, 0, i + 1));
+            Map<String, Object> fieldMapping = mapping.lookup().get(pathToHere);
+            if (fieldMapping == null) {
+                break;
+            }
+            if (fieldMapping.containsKey("dynamic")) {
+                // lower down dynamic definitions override higher up behavior
+                inRuntimeContext = "runtime".equals(fieldMapping.get("dynamic"));
+            }
+        }
+        return inRuntimeContext;
+    }
+}

x-pack/plugin/esql/compute/test/src/main/java/org/elasticsearch/compute/test/CannedSourceOperator.java

@@ -84,14 +84,20 @@ public static List<Page> deepCopyOf(BlockFactory blockFactory, List<Page> pages)
         try {
             for (Page p : pages) {
                 Block[] blocks = new Block[p.getBlockCount()];
-                for (int b = 0; b < blocks.length; b++) {
-                    Block orig = p.getBlock(b);
-                    try (Block.Builder builder = orig.elementType().newBlockBuilder(p.getPositionCount(), blockFactory)) {
-                        builder.copyFrom(orig, 0, p.getPositionCount());
-                        blocks[b] = builder.build();
+                try {
+                    for (int b = 0; b < blocks.length; b++) {
+                        Block orig = p.getBlock(b);
+                        try (Block.Builder builder = orig.elementType().newBlockBuilder(p.getPositionCount(), blockFactory)) {
+                            builder.copyFrom(orig, 0, p.getPositionCount());
+                            blocks[b] = builder.build();
+                        }
                     }
+                    out.add(new Page(blocks));
+                } catch (Exception e) {
+                    // Something went wrong, release the blocks.
+                    Releasables.closeExpectNoException(blocks);
+                    throw e;
                 }
-                out.add(new Page(blocks));
             }
         } finally {
             if (pages.size() != out.size()) {

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -12,7 +12,7 @@
 import org.apache.lucene.sandbox.document.HalfFloatPoint;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.ExceptionsHelper;
-import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.cluster.RemoteException;
 import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
 import org.elasticsearch.cluster.project.ProjectResolver;
@@ -76,7 +76,7 @@
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.NotEquals;
 import org.elasticsearch.xpack.esql.index.EsIndex;
 import org.elasticsearch.xpack.esql.inference.InferenceResolution;
-import org.elasticsearch.xpack.esql.inference.InferenceRunner;
+import org.elasticsearch.xpack.esql.inference.InferenceService;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.parser.QueryParam;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
@@ -161,8 +161,6 @@
 import static org.hamcrest.Matchers.instanceOf;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
-import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 
 public final class EsqlTestUtils {
@@ -422,20 +420,9 @@ public static LogicalOptimizerContext unboundLogicalOptimizerContext() {
         mock(ProjectResolver.class),
         mock(IndexNameExpressionResolver.class),
         null,
-        mockInferenceRunner()
+        new InferenceService(mock(Client.class))
     );
 
-    @SuppressWarnings("unchecked")
-    private static InferenceRunner mockInferenceRunner() {
-        InferenceRunner inferenceRunner = mock(InferenceRunner.class);
-        doAnswer(i -> {
-            i.getArgument(1, ActionListener.class).onResponse(emptyInferenceResolution());
-            return null;
-        }).when(inferenceRunner).resolveInferenceIds(any(), any());
-
-        return inferenceRunner;
-    }
-
     private EsqlTestUtils() {}
 
     public static Configuration configuration(QueryPragmas pragmas, String query) {