Skip to content

Commit 921dc23

Browse files
lucy66hwluckyxilu66karenyrx
authored
Implement GRPC FunctionScoreQuery (#19888)
* Implement GRPC FunctionScoreQuery Signed-off-by: xil <fridalu66@gmail.com> * update changelog Signed-off-by: xil <fridalu66@gmail.com> * spotlessApply Signed-off-by: xil <fridalu66@gmail.com> * remove unused change Signed-off-by: xil <xil@uber.com> * remove unused Signed-off-by: xil <xil@uber.com> * Address comment Signed-off-by: xil <fridalu66@gmail.com> * rebase Signed-off-by: xil <fridalu66@gmail.com> * revert Signed-off-by: xil <fridalu66@gmail.com> * update changelog Signed-off-by: xil <fridalu66@gmail.com> * rename converter to utils Signed-off-by: xil <fridalu66@gmail.com> * revert removed tests Signed-off-by: xil <fridalu66@gmail.com> * spotlessApply Signed-off-by: xil <fridalu66@gmail.com> --------- Signed-off-by: xil <fridalu66@gmail.com> Signed-off-by: xil <xil@uber.com> Signed-off-by: Karen X <karenxyr@gmail.com> Co-authored-by: xil <xil@uber.com> Co-authored-by: Karen X <karenxyr@gmail.com>
1 parent 332efe3 commit 921dc23

20 files changed

+2645
-0
lines changed

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
2424
- Implement GRPC ConstantScoreQuery, FuzzyQuery, MatchBoolPrefixQuery, MatchPhrasePrefix, PrefixQuery, MatchQuery ([#19854](https://github.com/opensearch-project/OpenSearch/pull/19854))
2525
- Add async periodic flush task support for pull-based ingestion ([#19878](https://github.com/opensearch-project/OpenSearch/pull/19878))
2626
- Add support for context aware segments ([#19098](https://github.com/opensearch-project/OpenSearch/pull/19098))
27+
- Implement GRPC FunctionScoreQuery ([#19888](https://github.com/opensearch-project/OpenSearch/pull/19888))
2728

2829
### Changed
2930
- Faster `terms` query creation for `keyword` field with index and docValues enabled ([#19350](https://github.com/opensearch-project/OpenSearch/pull/19350))

modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/query/QueryBuilderProtoConverterRegistryImpl.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import org.opensearch.common.inject.Singleton;
1414
import org.opensearch.index.query.QueryBuilder;
1515
import org.opensearch.protobufs.QueryContainer;
16+
import org.opensearch.transport.grpc.proto.request.search.query.functionscore.FunctionScoreQueryBuilderProtoConverter;
1617
import org.opensearch.transport.grpc.spi.QueryBuilderProtoConverter;
1718
import org.opensearch.transport.grpc.spi.QueryBuilderProtoConverterRegistry;
1819

@@ -67,6 +68,7 @@ protected void registerBuiltInConverters() {
6768
delegate.registerConverter(new MatchQueryBuilderProtoConverter());
6869
delegate.registerConverter(new MatchBoolPrefixQueryBuilderProtoConverter());
6970
delegate.registerConverter(new MatchPhrasePrefixQueryBuilderProtoConverter());
71+
delegate.registerConverter(new FunctionScoreQueryBuilderProtoConverter());
7072

7173
// Set the registry on all converters so they can access each other
7274
delegate.setRegistryOnAllConverters(this);
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,141 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
package org.opensearch.transport.grpc.proto.request.search.query.functionscore;
9+
10+
import org.opensearch.common.geo.GeoPoint;
11+
import org.opensearch.index.query.functionscore.ExponentialDecayFunctionBuilder;
12+
import org.opensearch.index.query.functionscore.ScoreFunctionBuilder;
13+
import org.opensearch.protobufs.DateDecayPlacement;
14+
import org.opensearch.protobufs.DecayFunction;
15+
import org.opensearch.protobufs.DecayPlacement;
16+
import org.opensearch.protobufs.GeoDecayPlacement;
17+
import org.opensearch.protobufs.NumericDecayPlacement;
18+
import org.opensearch.transport.grpc.proto.request.common.GeoPointProtoUtils;
19+
20+
import java.util.Map;
21+
22+
/**
23+
* Utility class for converting Protocol Buffer DecayFunction to OpenSearch ExponentialDecayFunctionBuilder.
24+
* This utility handles the transformation of Protocol Buffer DecayFunction objects
25+
* into OpenSearch ExponentialDecayFunctionBuilder instances.
26+
*/
27+
class ExpDecayFunctionProtoUtils {
28+
29+
private ExpDecayFunctionProtoUtils() {
30+
// Utility class, no instances
31+
}
32+
33+
/**
34+
* Converts a Protocol Buffer DecayFunction to an OpenSearch ScoreFunctionBuilder.
35+
* Similar to {@link org.opensearch.index.query.functionscore.DecayFunctionParser#fromXContent(org.opensearch.core.xcontent.XContentParser)},
36+
* this method parses the Protocol Buffer representation and creates a properly configured
37+
* ExponentialDecayFunctionBuilder with decay placement parameters (numeric, geo, or date).
38+
*
39+
* @param decayFunction the Protocol Buffer DecayFunction
40+
* @return the corresponding OpenSearch ScoreFunctionBuilder
41+
* @throws IllegalArgumentException if the decayFunction is null or doesn't contain placements
42+
*/
43+
static ScoreFunctionBuilder<?> fromProto(DecayFunction decayFunction) {
44+
if (decayFunction == null || decayFunction.getPlacementCount() == 0) {
45+
throw new IllegalArgumentException("DecayFunction must have at least one placement");
46+
}
47+
48+
Map.Entry<String, DecayPlacement> entry = decayFunction.getPlacementMap().entrySet().iterator().next();
49+
String fieldName = entry.getKey();
50+
DecayPlacement decayPlacement = entry.getValue();
51+
52+
if (decayPlacement.hasNumericDecayPlacement()) {
53+
return parseNumericExpDecay(fieldName, decayPlacement.getNumericDecayPlacement());
54+
} else if (decayPlacement.hasGeoDecayPlacement()) {
55+
return parseGeoExpDecay(fieldName, decayPlacement.getGeoDecayPlacement());
56+
} else if (decayPlacement.hasDateDecayPlacement()) {
57+
return parseDateExpDecay(fieldName, decayPlacement.getDateDecayPlacement());
58+
} else {
59+
throw new IllegalArgumentException("Unsupported decay placement type");
60+
}
61+
}
62+
63+
/**
64+
* Parses a numeric decay placement for exponential decay.
65+
*/
66+
private static ScoreFunctionBuilder<?> parseNumericExpDecay(String fieldName, NumericDecayPlacement numericPlacement) {
67+
ExponentialDecayFunctionBuilder builder;
68+
if (numericPlacement.hasDecay()) {
69+
builder = new ExponentialDecayFunctionBuilder(
70+
fieldName,
71+
numericPlacement.getOrigin(),
72+
numericPlacement.getScale(),
73+
numericPlacement.hasOffset() ? numericPlacement.getOffset() : null,
74+
numericPlacement.getDecay()
75+
);
76+
} else {
77+
builder = new ExponentialDecayFunctionBuilder(
78+
fieldName,
79+
numericPlacement.getOrigin(),
80+
numericPlacement.getScale(),
81+
numericPlacement.hasOffset() ? numericPlacement.getOffset() : null
82+
);
83+
}
84+
85+
return builder;
86+
}
87+
88+
/**
89+
* Parses a geo decay placement for exponential decay.
90+
*/
91+
private static ScoreFunctionBuilder<?> parseGeoExpDecay(String fieldName, GeoDecayPlacement geoPlacement) {
92+
GeoPoint geoPoint = GeoPointProtoUtils.parseGeoPoint(geoPlacement.getOrigin());
93+
94+
ExponentialDecayFunctionBuilder builder;
95+
if (geoPlacement.hasDecay()) {
96+
builder = new ExponentialDecayFunctionBuilder(
97+
fieldName,
98+
geoPoint,
99+
geoPlacement.getScale(),
100+
geoPlacement.hasOffset() ? geoPlacement.getOffset() : null,
101+
geoPlacement.getDecay()
102+
);
103+
} else {
104+
builder = new ExponentialDecayFunctionBuilder(
105+
fieldName,
106+
geoPoint,
107+
geoPlacement.getScale(),
108+
geoPlacement.hasOffset() ? geoPlacement.getOffset() : null
109+
);
110+
}
111+
112+
return builder;
113+
}
114+
115+
/**
116+
* Parses a date decay placement for exponential decay.
117+
*/
118+
private static ScoreFunctionBuilder<?> parseDateExpDecay(String fieldName, DateDecayPlacement datePlacement) {
119+
Object origin = datePlacement.hasOrigin() ? datePlacement.getOrigin() : null;
120+
121+
ExponentialDecayFunctionBuilder builder;
122+
if (datePlacement.hasDecay()) {
123+
builder = new ExponentialDecayFunctionBuilder(
124+
fieldName,
125+
origin,
126+
datePlacement.getScale(),
127+
datePlacement.hasOffset() ? datePlacement.getOffset() : null,
128+
datePlacement.getDecay()
129+
);
130+
} else {
131+
builder = new ExponentialDecayFunctionBuilder(
132+
fieldName,
133+
origin,
134+
datePlacement.getScale(),
135+
datePlacement.hasOffset() ? datePlacement.getOffset() : null
136+
);
137+
}
138+
139+
return builder;
140+
}
141+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
package org.opensearch.transport.grpc.proto.request.search.query.functionscore;
9+
10+
import org.opensearch.index.query.functionscore.FieldValueFactorFunctionBuilder;
11+
import org.opensearch.index.query.functionscore.ScoreFunctionBuilder;
12+
import org.opensearch.protobufs.FieldValueFactorModifier;
13+
import org.opensearch.protobufs.FieldValueFactorScoreFunction;
14+
15+
/**
16+
* Utility class for converting Protocol Buffer FieldValueFactorScoreFunction to OpenSearch objects.
17+
* This utility handles the transformation of Protocol Buffer FieldValueFactorScoreFunction objects
18+
* into OpenSearch FieldValueFactorFunctionBuilder instances.
19+
*/
20+
class FieldValueFactorFunctionProtoUtils {
21+
22+
private FieldValueFactorFunctionProtoUtils() {
23+
// Utility class, no instances
24+
}
25+
26+
/**
27+
* Converts a Protocol Buffer FieldValueFactorScoreFunction to an OpenSearch ScoreFunctionBuilder.
28+
* Similar to {@link FieldValueFactorFunctionBuilder#fromXContent(XContentParser)}, this method
29+
* parses the field, factor, missing value, and modifier parameters.
30+
*
31+
* @param fieldValueFactor the Protocol Buffer FieldValueFactorScoreFunction
32+
* @return the corresponding OpenSearch ScoreFunctionBuilder
33+
* @throws IllegalArgumentException if the fieldValueFactor is null
34+
*/
35+
static ScoreFunctionBuilder<?> fromProto(FieldValueFactorScoreFunction fieldValueFactor) {
36+
if (fieldValueFactor == null) {
37+
throw new IllegalArgumentException("FieldValueFactorScoreFunction cannot be null");
38+
}
39+
40+
FieldValueFactorFunctionBuilder builder = new FieldValueFactorFunctionBuilder(fieldValueFactor.getField());
41+
42+
if (fieldValueFactor.hasFactor()) {
43+
builder.factor(fieldValueFactor.getFactor());
44+
}
45+
46+
if (fieldValueFactor.hasMissing()) {
47+
builder.missing(fieldValueFactor.getMissing());
48+
}
49+
50+
if (fieldValueFactor.getModifier() != FieldValueFactorModifier.FIELD_VALUE_FACTOR_MODIFIER_NONE) {
51+
builder.modifier(parseFieldValueFactorModifier(fieldValueFactor.getModifier()));
52+
}
53+
54+
return builder;
55+
}
56+
57+
/**
58+
* Parses a protobuf FieldValueFactorModifier and returns the corresponding OpenSearch modifier.
59+
*
60+
* @param modifier the protobuf FieldValueFactorModifier
61+
* @return the corresponding OpenSearch FieldValueFactorFunction.Modifier
62+
*/
63+
private static org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier parseFieldValueFactorModifier(
64+
FieldValueFactorModifier modifier
65+
) {
66+
return switch (modifier) {
67+
case FIELD_VALUE_FACTOR_MODIFIER_NONE -> org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.NONE;
68+
case FIELD_VALUE_FACTOR_MODIFIER_LOG -> org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.LOG;
69+
case FIELD_VALUE_FACTOR_MODIFIER_LOG1P -> org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.LOG1P;
70+
case FIELD_VALUE_FACTOR_MODIFIER_LOG2P -> org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.LOG2P;
71+
case FIELD_VALUE_FACTOR_MODIFIER_LN -> org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.LN;
72+
case FIELD_VALUE_FACTOR_MODIFIER_LN1P -> org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.LN1P;
73+
case FIELD_VALUE_FACTOR_MODIFIER_LN2P -> org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.LN2P;
74+
case FIELD_VALUE_FACTOR_MODIFIER_SQUARE ->
75+
org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.SQUARE;
76+
case FIELD_VALUE_FACTOR_MODIFIER_SQRT -> org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.SQRT;
77+
case FIELD_VALUE_FACTOR_MODIFIER_RECIPROCAL ->
78+
org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.RECIPROCAL;
79+
case FIELD_VALUE_FACTOR_MODIFIER_UNSPECIFIED ->
80+
org.opensearch.common.lucene.search.function.FieldValueFactorFunction.Modifier.NONE;
81+
default -> throw new IllegalArgumentException("Unknown FieldValueFactorModifier: " + modifier);
82+
};
83+
}
84+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
package org.opensearch.transport.grpc.proto.request.search.query.functionscore;
9+
10+
import org.opensearch.index.query.QueryBuilder;
11+
import org.opensearch.protobufs.QueryContainer;
12+
import org.opensearch.transport.grpc.spi.QueryBuilderProtoConverter;
13+
import org.opensearch.transport.grpc.spi.QueryBuilderProtoConverterRegistry;
14+
15+
/**
16+
* Converter for FunctionScore queries.
17+
* This class implements the QueryBuilderProtoConverter interface to provide FunctionScore query support
18+
* for the gRPC transport module.
19+
*/
20+
public class FunctionScoreQueryBuilderProtoConverter implements QueryBuilderProtoConverter {
21+
22+
/**
23+
* Default constructor for FunctionScoreQueryBuilderProtoConverter.
24+
*/
25+
public FunctionScoreQueryBuilderProtoConverter() {}
26+
27+
private QueryBuilderProtoConverterRegistry registry;
28+
29+
@Override
30+
public void setRegistry(QueryBuilderProtoConverterRegistry registry) {
31+
this.registry = registry;
32+
}
33+
34+
@Override
35+
public QueryContainer.QueryContainerCase getHandledQueryCase() {
36+
return QueryContainer.QueryContainerCase.FUNCTION_SCORE;
37+
}
38+
39+
@Override
40+
public QueryBuilder fromProto(QueryContainer queryContainer) {
41+
if (queryContainer == null || !queryContainer.hasFunctionScore()) {
42+
throw new IllegalArgumentException("QueryContainer does not contain a FunctionScore query");
43+
}
44+
45+
return FunctionScoreQueryBuilderProtoUtils.fromProto(queryContainer.getFunctionScore(), registry);
46+
}
47+
}

0 commit comments

Comments
 (0)