Skip to content

Commit

Permalink
In progress unit testing for DerivedFieldScript
Browse files Browse the repository at this point in the history
  • Loading branch information
qreshi committed Mar 18, 2024
1 parent cc4a634 commit c5a62ca
Show file tree
Hide file tree
Showing 4 changed files with 51 additions and 61 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,6 @@ class org.opensearch.script.DerivedFieldScript @no_import {
}

static_import {
void emitLong(org.opensearch.script.DerivedFieldScript, long) bound_to org.opensearch.script.ScriptEmitValues$Long
void emitDouble(org.opensearch.script.DerivedFieldScript, double) bound_to org.opensearch.script.ScriptEmitValues$Double
void emitGeoPoint(org.opensearch.script.DerivedFieldScript, double, double) bound_to org.opensearch.script.ScriptEmitValues$GeoPoint
void emitBoolean(org.opensearch.script.DerivedFieldScript, boolean) bound_to org.opensearch.script.ScriptEmitValues$Boolean
void emitString(org.opensearch.script.DerivedFieldScript, String) bound_to org.opensearch.script.ScriptEmitValues$Strings
void emit(org.opensearch.script.DerivedFieldScript, Object) bound_to org.opensearch.script.ScriptEmitValues$EmitSingle
void emit(org.opensearch.script.DerivedFieldScript, double, double) bound_to org.opensearch.script.ScriptEmitValues$GeoPoint
}
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@

package org.opensearch.painless;

import org.apache.lucene.index.LeafReaderContext;
import org.apache.lucene.index.memory.MemoryIndex;
import org.opensearch.common.settings.Settings;
import org.opensearch.index.fielddata.IndexNumericFieldData;
import org.opensearch.index.fielddata.LeafNumericFieldData;
Expand All @@ -19,6 +21,7 @@
import org.opensearch.painless.spi.AllowlistLoader;
import org.opensearch.script.DerivedFieldScript;
import org.opensearch.script.ScriptContext;
import org.opensearch.search.lookup.LeafSearchLookup;
import org.opensearch.search.lookup.SearchLookup;

import java.io.IOException;
Expand All @@ -29,7 +32,9 @@

import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.when;

public class DerivedFieldScriptTests extends ScriptTestCase {
Expand All @@ -46,9 +51,6 @@ public void setUp() throws Exception {
allowlists.add(AllowlistLoader.loadFromResourceFiles(Allowlist.class, "org.opensearch.derived.txt"));
contexts.put(DerivedFieldScript.CONTEXT, allowlists);

// Mocking field values to be returned


SCRIPT_ENGINE = new PainlessScriptEngine(Settings.EMPTY, contexts);
}

Expand Down Expand Up @@ -99,13 +101,30 @@ public void testEmittingDoubleField() throws IOException {
when(fieldData.getFieldName()).thenReturn("test_double_field");
when(fieldData.load(any())).thenReturn(atomicFieldData);

SearchLookup lookup = new SearchLookup(mapperService, (ignored, searchLookup) -> fieldData);
SearchLookup lookup = spy(new SearchLookup(mapperService, (ignored, searchLookup) -> fieldData));

// We don't need a real index, just need to construct a LeafReaderContext which cannot be mocked
MemoryIndex index = new MemoryIndex();
LeafReaderContext leafReaderContext = index.createSearcher().getIndexReader().leaves().get(0);

LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
doReturn(leafLookup).when(lookup).getLeafSearchLookup(leafReaderContext);

// when(leafLookup.asMap()).thenReturn(Collections.emptyMap());

// SearchLookup lookup = mock(SearchLookup.class);
// LeafSearchLookup leafLookup = mock(LeafSearchLookup.class);
// when(lookup.getLeafSearchLookup(leafReaderContext)).thenReturn(leafLookup);
// SourceLookup sourceLookup = mock(SourceLookup.class);
// when(leafLookup.asMap()).thenReturn(Collections.singletonMap("_source", sourceLookup));
// when(sourceLookup.loadSourceIfNeeded()).thenReturn(Collections.singletonMap("test", 1));

// Execute the script
DerivedFieldScript script = compile("emitDouble(doc['test_double_field'].value)", lookup).newInstance(null);
DerivedFieldScript script = compile("emit(doc['test_double_field'].value)", lookup).newInstance(leafReaderContext);
script.setDocument(1);
script.execute();

List<Object> result = script.execute();
List<Object> result = script.getEmittedValues();
assertEquals(List.of(2.718), result);
}

Expand Down
15 changes: 11 additions & 4 deletions server/src/main/java/org/opensearch/script/DerivedFieldScript.java
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
import org.opensearch.search.lookup.SourceLookup;

import java.io.IOException;
import java.util.Collections;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
Expand Down Expand Up @@ -70,13 +70,13 @@ public DerivedFieldScript(Map<String, Object> params, SearchLookup lookup, LeafR
this.leafLookup = lookup.getLeafSearchLookup(leafContext);
parameters.putAll(leafLookup.asMap());
this.params = new DynamicMap(parameters, PARAMS_FUNCTIONS);
this.emittedValues = Collections.emptyList();
this.emittedValues = new ArrayList<>();
}

public DerivedFieldScript() {
this.params = null;
this.leafLookup = null;
this.emittedValues = Collections.emptyList();
this.emittedValues = new ArrayList<>();
}

/**
Expand All @@ -93,16 +93,23 @@ public Map<String, ScriptDocValues<?>> getDoc() {
return leafLookup.doc();
}

/**
* Return the emitted values from the script execution.
*/
public List<Object> getEmittedValues() { return emittedValues; }

/**
* Set the current document to run the script on next.
* Clears the emittedValues as well since they should be scoped per document.
*/
public void setDocument(int docid) {
this.emittedValues = new ArrayList<>();
leafLookup.setDocument(docid);
}

public void addEmittedValue(Object o) { emittedValues.add(o); }

public List<Object> execute() { return emittedValues; }
public void execute() {}

/**
* A factory to construct {@link DerivedFieldScript} instances.
Expand Down
59 changes: 13 additions & 46 deletions server/src/main/java/org/opensearch/script/ScriptEmitValues.java
Original file line number Diff line number Diff line change
Expand Up @@ -18,35 +18,29 @@
*/
public final class ScriptEmitValues {

// Emits a Long value (ex. from a long or date field type)
public static final class Long {
// Takes in a single value and emits it
// Could be a long, double, String, etc.
public static final class EmitSingle {

private final DerivedFieldScript derivedFieldScript;

public Long(DerivedFieldScript derivedFieldScript) {
public EmitSingle(DerivedFieldScript derivedFieldScript) {
this.derivedFieldScript = derivedFieldScript;
}

public void emitLong(long val) {
// TODO: Keeping this generic for the time being due to limitations with
// binding methods with the same name and arity.
// Ideally, we should have an emit signature per derived field type and try to scope
// that to the respective script execution so the other emits aren't allowed.
// One way to do this could be to create implementations of the DerivedFieldScript.LeafFactory
// per field type where they each define their own emit() method and then the engine that executes
// it can have custom compilation logic to perform class bindings on that emit implementation.
public void emit(Object val) {
derivedFieldScript.addEmittedValue(val);
}

}

// Emits a Double value
public static final class Double {

private final DerivedFieldScript derivedFieldScript;

public Double(DerivedFieldScript derivedFieldScript) {
this.derivedFieldScript = derivedFieldScript;
}

public void emitDouble(double val) {
derivedFieldScript.addEmittedValue(val);
}
}

// Emits a GeoPoint value
public static final class GeoPoint {

Expand All @@ -56,37 +50,10 @@ public GeoPoint(DerivedFieldScript derivedFieldScript) {
this.derivedFieldScript = derivedFieldScript;
}

public void emitGeoPoint(double lat, double lon) {
public void emit(double lat, double lon) {
derivedFieldScript.addEmittedValue(new Tuple<>(lat, lon));
}

}

// Emits a Boolean value
public static final class Boolean {

private final DerivedFieldScript derivedFieldScript;

public Boolean(DerivedFieldScript derivedFieldScript) {
this.derivedFieldScript = derivedFieldScript;
}

public void emitBoolean(boolean val) {
derivedFieldScript.addEmittedValue(val);
}
}

// Emits a String value
public static final class Strings {

private final DerivedFieldScript derivedFieldScript;

public Strings(DerivedFieldScript derivedFieldScript) {
this.derivedFieldScript = derivedFieldScript;
}

public void emitString(String val) {
derivedFieldScript.addEmittedValue(val);
}
}
}

0 comments on commit c5a62ca

Please sign in to comment.