Allow to sort search results by VERSION value returned by Painless sc…

…ript (#85990)


Version field types have specific semantics in terms of ordering, ie. SemVer.
This fix adds a new "version" type for script sorting (in addition to the existing "number" and "string" options)
to specify SemVer ordering semantics.

docs/changelog/85990.yaml

@@ -0,0 +1,7 @@
+pr: 85990
+summary: Allow to sort by script value using `SemVer` semantics
+area: Infra/Scripting
+type: bug
+issues:
+ - 85989
+ - 82287

modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.bytesref_sort.txt

@@ -0,0 +1,18 @@
+#
+# 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 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 or the Server
+# Side Public License, v 1.
+#
+
+# The whitelist for the fields api
+# The scripts must be whitelisted for painless to find the classes for the field API
+
+class org.elasticsearch.script.BytesRefProducer @no_import {
+}
+
+class org.elasticsearch.script.BytesRefSortScript @no_import {
+}
+class org.elasticsearch.script.BytesRefSortScript$Factory @no_import {
+}

server/src/main/java/org/elasticsearch/script/BytesRefProducer.java

@@ -0,0 +1,20 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.script;
+
+import org.apache.lucene.util.BytesRef;
+
+/**
+ * used by {@link org.elasticsearch.search.sort.ScriptSortBuilder} to refer to classes in x-pack
+ * (eg. org.elasticsearch.xpack.versionfield.Version) that need a custom FieldComparatorSource
+ */
+public interface BytesRefProducer {
+
+    BytesRef toBytesRef();
+}

server/src/main/java/org/elasticsearch/script/BytesRefSortScript.java

@@ -0,0 +1,38 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+package org.elasticsearch.script;
+
+import java.io.IOException;
+import java.util.Map;
+
+public abstract class BytesRefSortScript extends AbstractSortScript {
+
+    public static final String[] PARAMETERS = {};
+
+    public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>("bytesref_sort", Factory.class);
+
+    public BytesRefSortScript(Map<String, Object> params, DocReader docReader) {
+        super(params, docReader);
+    }
+
+    public abstract Object execute();
+
+    /**
+     * A factory to construct {@link BytesRefSortScript} instances.
+     */
+    public interface LeafFactory {
+        BytesRefSortScript newInstance(DocReader reader) throws IOException;
+    }
+
+    /**
+     * A factory to construct stateful {@link BytesRefSortScript} factories for a specific index.
+     */
+    public interface Factory extends ScriptFactory {
+        LeafFactory newFactory(Map<String, Object> params);
+    }
+}

server/src/main/java/org/elasticsearch/script/ScriptModule.java

@@ -48,6 +48,7 @@ public class ScriptModule {
                 ScoreScript.CONTEXT,
                 NumberSortScript.CONTEXT,
                 StringSortScript.CONTEXT,
+                BytesRefSortScript.CONTEXT,
                 TermsSetQueryScript.CONTEXT,
                 UpdateScript.CONTEXT,
                 BucketAggregationScript.CONTEXT,

server/src/main/java/org/elasticsearch/search/sort/ScriptSortBuilder.java

@@ -29,10 +29,13 @@
 import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
 import org.elasticsearch.index.fielddata.fieldcomparator.BytesRefFieldComparatorSource;
 import org.elasticsearch.index.fielddata.fieldcomparator.DoubleValuesComparatorSource;
+import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.query.QueryBuilder;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.QueryShardException;
 import org.elasticsearch.index.query.SearchExecutionContext;
+import org.elasticsearch.script.BytesRefProducer;
+import org.elasticsearch.script.BytesRefSortScript;
 import org.elasticsearch.script.DocValuesDocReader;
 import org.elasticsearch.script.NumberSortScript;
 import org.elasticsearch.script.Script;
@@ -72,14 +75,16 @@ public class ScriptSortBuilder extends SortBuilder<ScriptSortBuilder> {
 
     private NestedSortBuilder nestedSort;
 
+    private DocValueFormat scriptResultValueFormat = DocValueFormat.RAW;
+
     /**
      * Constructs a script sort builder with the given script.
      *
      * @param script
      *            The script to use.
      * @param type
-     *            The type of the script, can be either {@link ScriptSortType#STRING} or
-     *            {@link ScriptSortType#NUMBER}
+     *            The type of the script, can be {@link ScriptSortType#STRING},
+     *            {@link ScriptSortType#NUMBER} or {@link ScriptSortType#VERSION}
      */
     public ScriptSortBuilder(Script script, ScriptSortType type) {
         Objects.requireNonNull(script, "script cannot be null");
@@ -246,9 +251,18 @@ public static ScriptSortBuilder fromXContent(XContentParser parser, String eleme
 
     @Override
     public SortFieldAndFormat build(SearchExecutionContext context) throws IOException {
+        if ("version".equals(this.type.toString())) {
+            try {
+                // TODO there must be a better way to get the field type...
+                MappedFieldType scriptFieldType = context.buildAnonymousFieldType(this.type.toString());
+                scriptResultValueFormat = scriptFieldType.docValueFormat(null, null);
+            } catch (Exception e) {
+                // "version" type is not available, fall back to RAW and sort as a string
+            }
+        }
         return new SortFieldAndFormat(
             new SortField("_script", fieldComparatorSource(context), order == SortOrder.DESC),
-            DocValueFormat.RAW
+            scriptResultValueFormat == null ? DocValueFormat.RAW : scriptResultValueFormat
         );
     }
 
@@ -355,6 +369,64 @@ protected void setScorer(Scorable scorer) {
                     }
                 };
             }
+            case VERSION -> {
+                final BytesRefSortScript.Factory factory = context.compile(script, BytesRefSortScript.CONTEXT);
+                final BytesRefSortScript.LeafFactory searchScript = factory.newFactory(script.getParams());
+                return new BytesRefFieldComparatorSource(null, null, valueMode, nested) {
+                    BytesRefSortScript leafScript;
+
+                    @Override
+                    protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOException {
+                        leafScript = searchScript.newInstance(new DocValuesDocReader(searchLookup, context));
+                        final BinaryDocValues values = new AbstractBinaryDocValues() {
+
+                            @Override
+                            public boolean advanceExact(int doc) throws IOException {
+                                leafScript.setDocument(doc);
+                                return true;
+                            }
+
+                            @Override
+                            public BytesRef binaryValue() {
+                                Object result = leafScript.execute();
+                                if (result == null) {
+                                    return null;
+                                }
+                                if (result instanceof BytesRefProducer) {
+                                    return ((BytesRefProducer) result).toBytesRef();
+                                }
+
+                                if (scriptResultValueFormat == null) {
+                                    throw new IllegalArgumentException("Invalid sort type: version");
+                                }
+                                return scriptResultValueFormat.parseBytesRef(result);
+                            }
+                        };
+                        return FieldData.singleton(values);
+                    }
+
+                    @Override
+                    protected void setScorer(Scorable scorer) {
+                        leafScript.setScorer(scorer);
+                    }
+
+                    @Override
+                    public BucketedSort newBucketedSort(
+                        BigArrays bigArrays,
+                        SortOrder sortOrder,
+                        DocValueFormat format,
+                        int bucketSize,
+                        BucketedSort.ExtraData extra
+                    ) {
+                        throw new IllegalArgumentException(
+                            "error building sort for [_script]: "
+                                + "script sorting only supported on [numeric] scripts but was ["
+                                + type
+                                + "]"
+                        );
+                    }
+                };
+            }
             default -> throw new QueryShardException(context, "custom script sort type [" + type + "] not supported");
         }
     }
@@ -394,7 +466,9 @@ public enum ScriptSortType implements Writeable {
         /** script sort for a string value **/
         STRING,
         /** script sort for a numeric value **/
-        NUMBER;
+        NUMBER,
+        /** script sort for a Version field value **/
+        VERSION;
 
         @Override
         public void writeTo(final StreamOutput out) throws IOException {
@@ -413,6 +487,7 @@ public static ScriptSortType fromString(final String str) {
             return switch (str.toLowerCase(Locale.ROOT)) {
                 case ("string") -> ScriptSortType.STRING;
                 case ("number") -> ScriptSortType.NUMBER;
+                case ("version") -> ScriptSortType.VERSION;
                 default -> throw new IllegalArgumentException("Unknown ScriptSortType [" + str + "]");
             };
         }

server/src/test/java/org/elasticsearch/search/sort/ScriptSortBuilderTests.java

@@ -130,16 +130,21 @@ public void testScriptSortType() {
         // we rely on these ordinals in serialization, so changing them breaks bwc.
         assertEquals(0, ScriptSortType.STRING.ordinal());
         assertEquals(1, ScriptSortType.NUMBER.ordinal());
+        assertEquals(2, ScriptSortType.VERSION.ordinal());
 
         assertEquals("string", ScriptSortType.STRING.toString());
         assertEquals("number", ScriptSortType.NUMBER.toString());
+        assertEquals("version", ScriptSortType.VERSION.toString());
 
         assertEquals(ScriptSortType.STRING, ScriptSortType.fromString("string"));
         assertEquals(ScriptSortType.STRING, ScriptSortType.fromString("String"));
         assertEquals(ScriptSortType.STRING, ScriptSortType.fromString("STRING"));
         assertEquals(ScriptSortType.NUMBER, ScriptSortType.fromString("number"));
         assertEquals(ScriptSortType.NUMBER, ScriptSortType.fromString("Number"));
         assertEquals(ScriptSortType.NUMBER, ScriptSortType.fromString("NUMBER"));
+        assertEquals(ScriptSortType.VERSION, ScriptSortType.fromString("version"));
+        assertEquals(ScriptSortType.VERSION, ScriptSortType.fromString("Version"));
+        assertEquals(ScriptSortType.VERSION, ScriptSortType.fromString("VERSION"));
     }
 
     public void testScriptSortTypeNull() {
@@ -301,6 +306,10 @@ public void testBuildCorrectComparatorType() throws IOException {
         sortBuilder = new ScriptSortBuilder(mockScript(MOCK_SCRIPT_NAME), ScriptSortType.NUMBER);
         sortField = sortBuilder.build(createMockSearchExecutionContext()).field;
         assertThat(sortField.getComparatorSource(), instanceOf(DoubleValuesComparatorSource.class));
+
+        sortBuilder = new ScriptSortBuilder(mockScript(MOCK_SCRIPT_NAME), ScriptSortType.VERSION);
+        sortField = sortBuilder.build(createMockSearchExecutionContext()).field;
+        assertThat(sortField.getComparatorSource(), instanceOf(BytesRefFieldComparatorSource.class));
     }
 
     /**

test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java

@@ -145,6 +145,8 @@ public boolean needs_score() {
             return context.factoryClazz.cast(factory);
         } else if (context.instanceClazz.equals(StringSortScript.class)) {
             return context.factoryClazz.cast(new MockStringSortScriptFactory(script));
+        } else if (context.instanceClazz.equals(BytesRefSortScript.class)) {
+            return context.factoryClazz.cast(new MockBytesRefSortScriptFactory(script));
         } else if (context.instanceClazz.equals(IngestScript.class)) {
             IngestScript.Factory factory = vars -> new IngestScript(vars) {
                 @Override
@@ -804,4 +806,30 @@ public String execute() {
             };
         }
     }
+
+    class MockBytesRefSortScriptFactory implements BytesRefSortScript.Factory {
+        private final MockDeterministicScript script;
+
+        MockBytesRefSortScriptFactory(MockDeterministicScript script) {
+            this.script = script;
+        }
+
+        @Override
+        public boolean isResultDeterministic() {
+            return script.isResultDeterministic();
+        }
+
+        @Override
+        public BytesRefSortScript.LeafFactory newFactory(Map<String, Object> parameters) {
+            return docReader -> new BytesRefSortScript(parameters, docReader) {
+                @Override
+                public BytesRefProducer execute() {
+                    Map<String, Object> vars = new HashMap<>(parameters);
+                    vars.put("params", parameters);
+                    vars.put("doc", getDoc());
+                    return (BytesRefProducer) script.apply(vars);
+                }
+            };
+        }
+    }
 }

x-pack/plugin/mapper-version/src/main/java/org/elasticsearch/xpack/versionfield/Version.java

@@ -7,21 +7,28 @@
 
 package org.elasticsearch.xpack.versionfield;
 
-import org.elasticsearch.xcontent.ToXContent;
+import org.apache.lucene.util.BytesRef;
+import org.elasticsearch.script.BytesRefProducer;
+import org.elasticsearch.xcontent.ToXContentFragment;
 import org.elasticsearch.xcontent.XContentBuilder;
 
 import java.io.IOException;
 
 /**
  * Script value class.
- * TODO(stu): implement {@code Comparable<Version>} based on {@code VersionEncoder#prefixDigitGroupsWithLength(String, BytesRefBuilder)}
- * See: https://github.com/elastic/elasticsearch/issues/82287
  */
-public class Version implements ToXContent {
+public class Version implements ToXContentFragment, BytesRefProducer, Comparable<Version> {
     protected String version;
+    protected BytesRef bytes;
 
     public Version(String version) {
         this.version = version;
+        this.bytes = VersionEncoder.encodeVersion(version).bytesRef;
+    }
+
+    protected Version(BytesRef bytes) {
+        this.version = VersionEncoder.decodeVersion(bytes);
+        this.bytes = bytes;
     }
 
     @Override
@@ -35,7 +42,12 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
     }
 
     @Override
-    public boolean isFragment() {
-        return false;
+    public BytesRef toBytesRef() {
+        return bytes;
+    }
+
+    @Override
+    public int compareTo(Version o) {
+        return toBytesRef().compareTo(o.toBytesRef());
     }
 }

x-pack/plugin/mapper-version/src/main/java/org/elasticsearch/xpack/versionfield/VersionFieldDocValuesExtension.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.painless.spi.Whitelist;
 import org.elasticsearch.painless.spi.WhitelistLoader;
 import org.elasticsearch.script.AggregationScript;
+import org.elasticsearch.script.BytesRefSortScript;
 import org.elasticsearch.script.FieldScript;
 import org.elasticsearch.script.FilterScript;
 import org.elasticsearch.script.NumberSortScript;
@@ -41,6 +42,7 @@ public Map<ScriptContext<?>, List<Whitelist>> getContextWhitelists() {
         whitelist.put(FieldScript.CONTEXT, list);
         whitelist.put(NumberSortScript.CONTEXT, list);
         whitelist.put(StringSortScript.CONTEXT, list);
+        whitelist.put(BytesRefSortScript.CONTEXT, list);
         return whitelist;
     }
 }