@Override
public void doXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject(getName());
builder.field(fieldName);
XContentParser parser = XContentFactory.xContent(functionBytes).createParser(functionBytes);
builder.copyCurrentStructure(parser);
builder.field(DecayFunctionParser.MULTI_VALUE_MODE.getPreferredName(), multiValueMode.name());
builder.endObject();
}
@Override
@SuppressWarnings("rawtypes") // ValueSource uses a rawtype
public FunctionValues getValues(Map context, LeafReaderContext leaf) throws IOException {
AtomicNumericFieldData leafData = (AtomicNumericFieldData) fieldData.load(leaf);
NumericDoubleValues docValues = multiValueMode.select(leafData.getDoubleValues(), 0d);
return new DoubleDocValues(this) {
@Override
public double doubleVal(int doc) {
return docValues.get(doc);
}
};
}
@Override
public SortField parse(XContentParser parser, SearchContext context) throws Exception {
String script = null;
String scriptLang = null;
String type = null;
Map<String, Object> params = null;
boolean reverse = false;
MultiValueMode sortMode = null;
NestedInnerQueryParseSupport nestedHelper = null;
XContentParser.Token token;
String currentName = parser.currentName();
ScriptService.ScriptType scriptType = ScriptService.ScriptType.INLINE;
while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
if (token == XContentParser.Token.FIELD_NAME) {
currentName = parser.currentName();
} else if (token == XContentParser.Token.START_OBJECT) {
if ("params".equals(currentName)) {
params = parser.map();
} else if ("nested_filter".equals(currentName) || "nestedFilter".equals(currentName)) {
if (nestedHelper == null) {
nestedHelper = new NestedInnerQueryParseSupport(parser, context);
}
nestedHelper.filter();
}
} else if (token.isValue()) {
if ("reverse".equals(currentName)) {
reverse = parser.booleanValue();
} else if ("order".equals(currentName)) {
reverse = "desc".equals(parser.text());
} else if (ScriptService.SCRIPT_INLINE.match(currentName)) {
script = parser.text();
scriptType = ScriptService.ScriptType.INLINE;
} else if (ScriptService.SCRIPT_ID.match(currentName)) {
script = parser.text();
scriptType = ScriptService.ScriptType.INDEXED;
} else if (ScriptService.SCRIPT_FILE.match(currentName)) {
script = parser.text();
scriptType = ScriptService.ScriptType.FILE;
} else if (ScriptService.SCRIPT_LANG.match(currentName)) {
scriptLang = parser.text();
} else if ("type".equals(currentName)) {
type = parser.text();
} else if ("mode".equals(currentName)) {
sortMode = MultiValueMode.fromString(parser.text());
} else if ("nested_path".equals(currentName) || "nestedPath".equals(currentName)) {
if (nestedHelper == null) {
nestedHelper = new NestedInnerQueryParseSupport(parser, context);
}
nestedHelper.setPath(parser.text());
}
}
}
if (script == null) {
throw new SearchParseException(
context, "_script sorting requires setting the script to sort by");
}
if (type == null) {
throw new SearchParseException(
context, "_script sorting requires setting the type of the script");
}
final SearchScript searchScript =
context
.scriptService()
.search(
context.lookup(),
scriptLang,
script,
scriptType,
ScriptContext.Standard.SEARCH,
params);
if (STRING_SORT_TYPE.equals(type)
&& (sortMode == MultiValueMode.SUM || sortMode == MultiValueMode.AVG)) {
throw new SearchParseException(
context, "type [string] doesn't support mode [" + sortMode + "]");
}
if (sortMode == null) {
sortMode = reverse ? MultiValueMode.MAX : MultiValueMode.MIN;
}
// If nested_path is specified, then wrap the `fieldComparatorSource` in a
// `NestedFieldComparatorSource`
final Nested nested;
if (nestedHelper != null && nestedHelper.getPath() != null) {
FixedBitSetFilter rootDocumentsFilter =
context.fixedBitSetFilterCache().getFixedBitSetFilter(NonNestedDocsFilter.INSTANCE);
FixedBitSetFilter innerDocumentsFilter;
if (nestedHelper.filterFound()) {
innerDocumentsFilter =
context.fixedBitSetFilterCache().getFixedBitSetFilter(nestedHelper.getInnerFilter());
} else {
innerDocumentsFilter =
context
.fixedBitSetFilterCache()
.getFixedBitSetFilter(nestedHelper.getNestedObjectMapper().nestedTypeFilter());
}
nested = new Nested(rootDocumentsFilter, innerDocumentsFilter);
} else {
nested = null;
}
final IndexFieldData.XFieldComparatorSource fieldComparatorSource;
switch (type) {
case STRING_SORT_TYPE:
fieldComparatorSource =
new BytesRefFieldComparatorSource(null, null, sortMode, nested) {
@Override
protected SortedBinaryDocValues getValues(AtomicReaderContext context) {
searchScript.setNextReader(context);
final BinaryDocValues values =
new BinaryDocValues() {
final BytesRefBuilder spare = new BytesRefBuilder();
@Override
public BytesRef get(int docID) {
searchScript.setNextDocId(docID);
spare.copyChars(searchScript.run().toString());
return spare.get();
}
};
return FieldData.singleton(values, null);
}
@Override
protected void setScorer(Scorer scorer) {
searchScript.setScorer(scorer);
}
};
break;
case NUMBER_SORT_TYPE:
// TODO: should we rather sort missing values last?
fieldComparatorSource =
new DoubleValuesComparatorSource(null, Double.MAX_VALUE, sortMode, nested) {
@Override
protected SortedNumericDoubleValues getValues(AtomicReaderContext context) {
searchScript.setNextReader(context);
final NumericDoubleValues values =
new NumericDoubleValues() {
@Override
public double get(int docID) {
searchScript.setNextDocId(docID);
return searchScript.runAsDouble();
}
};
return FieldData.singleton(values, null);
}
@Override
protected void setScorer(Scorer scorer) {
searchScript.setScorer(scorer);
}
};
break;
default:
throw new SearchParseException(
context, "custom script sort type [" + type + "] not supported");
}
return new SortField("_script", fieldComparatorSource, reverse);
}
@Override
public int hashCode() {
int result = fieldData.hashCode();
result = 31 * result + multiValueMode.hashCode();
return result;
}
@Override
protected void doWriteTo(StreamOutput out) throws IOException {
out.writeString(fieldName);
out.writeBytesReference(functionBytes);
multiValueMode.writeTo(out);
}
@Override
protected DFB doReadFrom(StreamInput in) throws IOException {
DFB decayFunctionBuilder = createFunctionBuilder(in.readString(), in.readBytesReference());
decayFunctionBuilder.setMultiValueMode(MultiValueMode.readMultiValueModeFrom(in));
return decayFunctionBuilder;
}
@Override
public SortFieldAndFormat build(QueryShardContext context) throws IOException {
final SearchScript searchScript =
context.getSearchScript(script, ScriptContext.Standard.SEARCH, Collections.emptyMap());
MultiValueMode valueMode = null;
if (sortMode != null) {
valueMode = MultiValueMode.fromString(sortMode.toString());
}
boolean reverse = (order == SortOrder.DESC);
if (valueMode == null) {
valueMode = reverse ? MultiValueMode.MAX : MultiValueMode.MIN;
}
final Nested nested = resolveNested(context, nestedPath, nestedFilter);
final IndexFieldData.XFieldComparatorSource fieldComparatorSource;
switch (type) {
case STRING:
fieldComparatorSource =
new BytesRefFieldComparatorSource(null, null, valueMode, nested) {
LeafSearchScript leafScript;
@Override
protected SortedBinaryDocValues getValues(LeafReaderContext context)
throws IOException {
leafScript = searchScript.getLeafSearchScript(context);
final BinaryDocValues values =
new BinaryDocValues() {
final BytesRefBuilder spare = new BytesRefBuilder();
@Override
public BytesRef get(int docID) {
leafScript.setDocument(docID);
spare.copyChars(leafScript.run().toString());
return spare.get();
}
};
return FieldData.singleton(values, null);
}
@Override
protected void setScorer(Scorer scorer) {
leafScript.setScorer(scorer);
}
};
break;
case NUMBER:
fieldComparatorSource =
new DoubleValuesComparatorSource(null, Double.MAX_VALUE, valueMode, nested) {
LeafSearchScript leafScript;
@Override
protected SortedNumericDoubleValues getValues(LeafReaderContext context)
throws IOException {
leafScript = searchScript.getLeafSearchScript(context);
final NumericDoubleValues values =
new NumericDoubleValues() {
@Override
public double get(int docID) {
leafScript.setDocument(docID);
return leafScript.runAsDouble();
}
};
return FieldData.singleton(values, null);
}
@Override
protected void setScorer(Scorer scorer) {
leafScript.setScorer(scorer);
}
};
break;
default:
throw new QueryShardException(
context, "custom script sort type [" + type + "] not supported");
}
return new SortFieldAndFormat(
new SortField("_script", fieldComparatorSource, reverse), DocValueFormat.RAW);
}