Skip to content

Commit

Permalink
Support optional function score variant
Browse files Browse the repository at this point in the history
This commit supports function score functions with an
optional function variant. The function is optional as it
is perfectly valid to have a function with only a weight e.g.
a function score query with multiple functions and a score mode
of "first", where the last weight function acts as some default value.

Closes opensearch-project#820
  • Loading branch information
russcam committed Feb 26, 2024
1 parent 0c88a6a commit 39fa124
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 14 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -106,20 +106,28 @@ public final Object _get() {
@Nullable
private final Double weight;

public FunctionScore(FunctionScoreVariant value) {

this._kind = ApiTypeHelper.requireNonNull(value._functionScoreKind(), this, "<variant kind>");
this._value = ApiTypeHelper.requireNonNull(value, this, "<variant value>");
public FunctionScore(@Nullable FunctionScoreVariant value) {
if (value != null) {
this._kind = ApiTypeHelper.requireNonNull(value._functionScoreKind(), this, "<variant kind>");
this._value = ApiTypeHelper.requireNonNull(value, this, "<variant value>");
} else {
this._kind = null;
this._value = null;
}

this.filter = null;
this.weight = null;

}

private FunctionScore(Builder builder) {

this._kind = ApiTypeHelper.requireNonNull(builder._kind, builder, "<variant kind>");
this._value = ApiTypeHelper.requireNonNull(builder._value, builder, "<variant value>");
if (builder._value != null) {
this._kind = ApiTypeHelper.requireNonNull(builder._kind, builder, "<variant kind>");
this._value = ApiTypeHelper.requireNonNull(builder._value, builder, "<variant value>");
} else {
this._kind = null;
this._value = null;
}

this.filter = builder.filter;
this.weight = builder.weight;
Expand Down Expand Up @@ -266,9 +274,11 @@ public void serialize(JsonGenerator generator, JsonpMapper mapper) {

}

generator.writeKey(_kind.jsonValue());
if (_value instanceof JsonpSerializable) {
((JsonpSerializable) _value).serialize(generator, mapper);
if (_value != null) {
generator.writeKey(_kind.jsonValue());
if (_value instanceof JsonpSerializable) {
((JsonpSerializable) _value).serialize(generator, mapper);
}
}

generator.writeEnd();
Expand All @@ -279,8 +289,10 @@ public Builder toBuilder() {
return new Builder()._kind(_kind)._value(_value).filter(filter).weight(weight);
}

public static class Builder extends ObjectBuilderBase {
public static class Builder extends ObjectBuilderBase implements ObjectBuilder<FunctionScore> {
@Nullable
private Kind _kind;
@Nullable
private Object _value;

@Nullable
Expand All @@ -289,12 +301,12 @@ public static class Builder extends ObjectBuilderBase {
@Nullable
private Double weight;

protected final Builder _kind(Kind v) {
protected final Builder _kind(@Nullable Kind v) {
this._kind = v;
return this;
}

protected final Builder _value(Object v) {
protected final Builder _value(@Nullable Object v) {
this._value = v;
return this;
}
Expand Down Expand Up @@ -384,7 +396,7 @@ public ContainerBuilder scriptScore(Function<ScriptScoreFunction.Builder, Object
return this.scriptScore(fn.apply(new ScriptScoreFunction.Builder()).build());
}

protected FunctionScore build() {
public FunctionScore build() {
_checkSingleUse();
return new FunctionScore(this);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,4 +11,37 @@ public void toBuilder() {

assertEquals(toJson(copied), toJson(origin));
}

@Test
public void canSupportWeightFunction() {
Query functionScoreQuery = FunctionScoreQuery.of(fs -> fs
.functions(f -> f.weight(5d))
)._toQuery();

String json = "{\"function_score\":{\"functions\":[{\"weight\":5.0}]}}";

assertEquals(json, toJson(functionScoreQuery));

Query roundtripQuery = checkJsonRoundtrip(functionScoreQuery, json);

assertNull(roundtripQuery.functionScore().functions().get(0)._kind());
assertEquals(5.0, roundtripQuery.functionScore().functions().get(0).weight(), 0.001);
}

@Test
public void canSupportFunctionVariant() {
Query functionScoreQuery = FunctionScoreQuery.of(fs -> fs
.functions(f -> f.weight(3d).linear(l -> l.field("field").placement(p -> p.decay(8.0))))
)._toQuery();

String json = "{\"function_score\":{\"functions\":[{\"weight\":3.0,\"linear\":{\"field\":{\"decay\":8.0}}}]}}";

assertEquals(json, toJson(functionScoreQuery));

Query roundtripQuery = checkJsonRoundtrip(functionScoreQuery, json);

assertEquals(FunctionScore.Kind.Linear, roundtripQuery.functionScore().functions().get(0)._kind());
assertEquals(3.0, roundtripQuery.functionScore().functions().get(0).weight(), 0.001);
assertEquals(8.0, roundtripQuery.functionScore().functions().get(0).linear().placement().decay(), 0.001);
}
}

0 comments on commit 39fa124

Please sign in to comment.