소스 검색

Add scaled float to the scripting fields API (#82275)

This adds the mapped type scaled float the scripting fields API using double as the base type.
Jack Conradson 3 년 전
부모
커밋
e5efade6aa

+ 5 - 0
docs/changelog/82275.yaml

@@ -0,0 +1,5 @@
+pr: 82275
+summary: Add scaled float to the scripting fields API
+area: Infra/Scripting
+type: enhancement
+issues: []

+ 5 - 0
modules/lang-painless/src/main/resources/org/elasticsearch/painless/org.elasticsearch.script.fields.txt

@@ -56,6 +56,11 @@ class org.elasticsearch.script.field.FloatDocValuesField @dynamic_type {
   float get(int, float)
 }
 
+class org.elasticsearch.script.field.ScaledFloatDocValuesField @dynamic_type {
+  double get(double)
+  double get(int, double)
+}
+
 # defaults are cast to byte, taking an int facilitates resolution with constants without casting
 class org.elasticsearch.script.field.ByteDocValuesField @dynamic_type {
   byte get(int)

+ 40 - 0
modules/lang-painless/src/yamlRestTest/resources/rest-api-spec/test/painless/50_script_doc_values.yml

@@ -90,6 +90,7 @@ setup:
               byte: [16, 32, 64, 8, 4]
               double: [3.141592653588, 2.141592653587]
               float: [1.123, 2.234]
+              scaled_float: [-3.5, 2.5]
 
 
     - do:
@@ -1340,6 +1341,45 @@ setup:
                             source: "doc['scaled_float'].get(0)"
     - match: { hits.hits.0.fields.field.0: 3.14 }
 
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            sort: [ { rank: asc } ]
+            script_fields:
+              field:
+                script:
+                  source: "field('scaled_float').get(0.0)"
+    - match: { hits.hits.0.fields.field.0: 3.14 }
+    - match: { hits.hits.1.fields.field.0: 0.0 }
+    - match: { hits.hits.2.fields.field.0: -3.5 }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            sort: [ { rank: asc } ]
+            script_fields:
+              field:
+                script:
+                  source: "/* avoid stash */ $('scaled_float', 0.0)"
+    - match: { hits.hits.0.fields.field.0: 3.14 }
+    - match: { hits.hits.1.fields.field.0: 0.0 }
+    - match: { hits.hits.2.fields.field.0: -3.5 }
+
+    - do:
+        search:
+          rest_total_hits_as_int: true
+          body:
+            sort: [ { rank: asc } ]
+            script_fields:
+              field:
+                script:
+                  source: "field('scaled_float').get(1, 0.0)"
+    - match: { hits.hits.0.fields.field.0: 0.0 }
+    - match: { hits.hits.1.fields.field.0: 0.0 }
+    - match: { hits.hits.2.fields.field.0: 2.5 }
+
     - do:
         search:
             rest_total_hits_as_int: true

+ 2 - 8
modules/mapper-extras/src/main/java/org/elasticsearch/index/mapper/extras/ScaledFloatFieldMapper.java

@@ -25,8 +25,6 @@ import org.elasticsearch.index.fielddata.IndexFieldData;
 import org.elasticsearch.index.fielddata.IndexNumericFieldData;
 import org.elasticsearch.index.fielddata.LeafNumericFieldData;
 import org.elasticsearch.index.fielddata.NumericDoubleValues;
-import org.elasticsearch.index.fielddata.ScriptDocValues.Doubles;
-import org.elasticsearch.index.fielddata.ScriptDocValues.DoublesSupplier;
 import org.elasticsearch.index.fielddata.SortedBinaryDocValues;
 import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
 import org.elasticsearch.index.fielddata.plain.SortedNumericIndexFieldData;
@@ -40,8 +38,8 @@ import org.elasticsearch.index.mapper.TextSearchInfo;
 import org.elasticsearch.index.mapper.TimeSeriesParams;
 import org.elasticsearch.index.mapper.ValueFetcher;
 import org.elasticsearch.index.query.SearchExecutionContext;
-import org.elasticsearch.script.field.DelegateDocValuesField;
 import org.elasticsearch.script.field.DocValuesField;
+import org.elasticsearch.script.field.ScaledFloatDocValuesField;
 import org.elasticsearch.script.field.ToScriptField;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.aggregations.support.ValuesSourceType;
@@ -267,11 +265,7 @@ public class ScaledFloatFieldMapper extends FieldMapper {
                     IndexNumericFieldData.NumericType.LONG,
                     (dv, n) -> { throw new UnsupportedOperationException(); }
                 ).build(cache, breakerService);
-                return new ScaledFloatIndexFieldData(
-                    scaledValues,
-                    scalingFactor,
-                    (dv, n) -> new DelegateDocValuesField(new Doubles(new DoublesSupplier(dv)), n)
-                );
+                return new ScaledFloatIndexFieldData(scaledValues, scalingFactor, ScaledFloatDocValuesField::new);
             };
         }
 

+ 119 - 0
server/src/main/java/org/elasticsearch/script/field/ScaledFloatDocValuesField.java

@@ -0,0 +1,119 @@
+/*
+ * 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.field;
+
+import org.apache.lucene.util.ArrayUtil;
+import org.elasticsearch.index.fielddata.ScriptDocValues;
+import org.elasticsearch.index.fielddata.SortedNumericDoubleValues;
+
+import java.io.IOException;
+import java.util.NoSuchElementException;
+import java.util.PrimitiveIterator;
+
+public class ScaledFloatDocValuesField implements DocValuesField<Double>, ScriptDocValues.Supplier<Double> {
+
+    protected final SortedNumericDoubleValues input;
+    protected final String name;
+
+    protected double[] values = new double[0];
+    protected int count;
+
+    private ScriptDocValues.Doubles doubles = null;
+
+    public ScaledFloatDocValuesField(SortedNumericDoubleValues input, String name) {
+        this.input = input;
+        this.name = name;
+    }
+
+    @Override
+    public void setNextDocId(int docId) throws IOException {
+        if (input.advanceExact(docId)) {
+            resize(input.docValueCount());
+            for (int i = 0; i < count; i++) {
+                values[i] = input.nextValue();
+            }
+        } else {
+            resize(0);
+        }
+    }
+
+    protected void resize(int newSize) {
+        count = newSize;
+
+        assert count >= 0 : "size must be positive (got " + count + "): likely integer overflow?";
+        values = ArrayUtil.grow(values, count);
+    }
+
+    @Override
+    public ScriptDocValues<Double> getScriptDocValues() {
+        if (doubles == null) {
+            doubles = new ScriptDocValues.Doubles(this);
+        }
+
+        return doubles;
+    }
+
+    @Override
+    public Double getInternal(int index) {
+        return values[index];
+    }
+
+    @Override
+    public String getName() {
+        return name;
+    }
+
+    @Override
+    public boolean isEmpty() {
+        return count == 0;
+    }
+
+    @Override
+    public int size() {
+        return count;
+    }
+
+    public double get(double defaultValue) {
+        return get(0, defaultValue);
+    }
+
+    public double get(int index, double defaultValue) {
+        if (isEmpty() || index < 0 || index >= count) {
+            return defaultValue;
+        }
+
+        return values[index];
+    }
+
+    @Override
+    public PrimitiveIterator.OfDouble iterator() {
+        return new PrimitiveIterator.OfDouble() {
+            private int index = 0;
+
+            @Override
+            public boolean hasNext() {
+                return index < count;
+            }
+
+            @Override
+            public Double next() {
+                return nextDouble();
+            }
+
+            @Override
+            public double nextDouble() {
+                if (hasNext() == false) {
+                    throw new NoSuchElementException();
+                }
+
+                return values[index++];
+            }
+        };
+    }
+}

+ 22 - 0
server/src/test/java/org/elasticsearch/index/fielddata/FloatDocValuesFieldTests.java

@@ -10,6 +10,7 @@ package org.elasticsearch.index.fielddata;
 
 import org.elasticsearch.script.field.DoubleDocValuesField;
 import org.elasticsearch.script.field.FloatDocValuesField;
+import org.elasticsearch.script.field.ScaledFloatDocValuesField;
 import org.elasticsearch.test.ESTestCase;
 
 import java.io.IOException;
@@ -58,6 +59,27 @@ public class FloatDocValuesFieldTests extends ESTestCase {
         }
     }
 
+    public void testScaledFloatField() throws IOException {
+        double[][] values = generate(ESTestCase::randomDouble);
+        ScaledFloatDocValuesField scaledFloatField = new ScaledFloatDocValuesField(wrap(values), "test");
+        for (int round = 0; round < 10; round++) {
+            int d = between(0, values.length - 1);
+            scaledFloatField.setNextDocId(d);
+            if (values[d].length > 0) {
+                assertEquals(values[d][0], scaledFloatField.get(Double.MIN_VALUE), 0.0);
+                assertEquals(values[d][0], scaledFloatField.get(0, Double.MIN_VALUE), 0.0);
+            }
+            assertEquals(values[d].length, scaledFloatField.size());
+            for (int i = 0; i < values[d].length; i++) {
+                assertEquals(values[d][i], scaledFloatField.get(i, Double.MIN_VALUE), 0.0);
+            }
+            int i = 0;
+            for (double dbl : scaledFloatField) {
+                assertEquals(values[d][i++], dbl, 0.0);
+            }
+        }
+    }
+
     protected double[][] generate(DoubleSupplier supplier) {
         double[][] values = new double[between(3, 10)][];
         for (int d = 0; d < values.length; d++) {