Skip to content

Commit fd08ef7

Browse files
karenyrxkarenyrx
committed
[GRPC] Add support for missing proto fields in FunctionScore and Highlight
Signed-off-by: Karen X <karenxyr@gmail.com>
1 parent b7f013f commit fd08ef7

File tree

13 files changed

+537
-18
lines changed

13 files changed

+537
-18
lines changed

‎CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
7272
- Bump opensearch-protobufs dependency to 0.24.0 and update transport-grpc module compatibility ([#20059](https://github.com/opensearch-project/OpenSearch/pull/20059))
7373

7474
- Refactor the ShardStats, WarmerStats and IndexingPressureStats class to use the Builder pattern instead of constructors ([#19966](https://github.com/opensearch-project/OpenSearch/pull/19966))
75+
- Add support for missing proto fields in GRPC FunctionScore and Highlight ([#20169](https://github.com/opensearch-project/OpenSearch/pull/20169))
76+
7577

7678
### Fixed
7779
- Fix Allocation and Rebalance Constraints of WeightFunction are incorrectly reset ([#19012](https://github.com/opensearch-project/OpenSearch/pull/19012))

‎modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/HighlightBuilderProtoUtils.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,9 @@ static HighlightBuilder fromProto(Highlight highlightProto, QueryBuilderProtoCon
152152
highlightBuilder.encoder(ProtobufEnumUtils.convertToString(highlightProto.getEncoder()));
153153
}
154154

155+
// TODO: Support useExplicitFieldOrder
156+
// A spec fix is required and a corresponding protobuf version upgrade.
157+
155158
if (highlightProto.getFieldsCount() > 0) {
156159
for (java.util.Map.Entry<String, org.opensearch.protobufs.HighlightField> entry : highlightProto.getFieldsMap().entrySet()) {
157160
String fieldName = entry.getKey();
@@ -252,6 +255,22 @@ static HighlightBuilder fromProto(Highlight highlightProto, QueryBuilderProtoCon
252255
fieldBuilder.forceSource(fieldProto.getForceSource());
253256
}
254257

258+
if (fieldProto.hasOrder() && fieldProto.getOrder() == HighlighterOrder.HIGHLIGHTER_ORDER_SCORE) {
259+
fieldBuilder.order(HighlightBuilder.Order.SCORE);
260+
}
261+
262+
if (fieldProto.hasPhraseLimit()) {
263+
fieldBuilder.phraseLimit(fieldProto.getPhraseLimit());
264+
}
265+
266+
if (fieldProto.hasRequireFieldMatch()) {
267+
fieldBuilder.requireFieldMatch(fieldProto.getRequireFieldMatch());
268+
}
269+
270+
if (fieldProto.hasTagsSchema() && fieldProto.getTagsSchema() != HighlighterTagsSchema.HIGHLIGHTER_TAGS_SCHEMA_UNSPECIFIED) {
271+
applyTagsSchemaToField(fieldBuilder, fieldProto.getTagsSchema());
272+
}
273+
255274
highlightBuilder.field(fieldBuilder);
256275
}
257276
}
@@ -274,4 +293,28 @@ private static HighlightBuilder.BoundaryScannerType parseBoundaryScanner(Boundar
274293
throw new IllegalArgumentException("Unknown BoundaryScanner value: " + boundaryScanner);
275294
}
276295
}
296+
297+
/**
298+
* Apply tags schema to a highlight field by setting pre/post tags.
299+
* This mirrors the behavior of {@link HighlightBuilder#tagsSchema(String)} but for field-level settings.
300+
* <p>
301+
* The {@code tagsSchema} method on {@link HighlightBuilder} sets built-in pre and post tags based on a schema name.
302+
* Since {@link HighlightBuilder.Field} doesn't have a {@code tagsSchema} method, this helper applies the same
303+
* logic by directly setting the pre/post tags on the field.
304+
*
305+
* @param fieldBuilder the highlight field builder to apply the tags schema to
306+
* @param tagsSchema the tags schema to apply (e.g., STYLED)
307+
* @see HighlightBuilder#tagsSchema(String)
308+
*/
309+
private static void applyTagsSchemaToField(HighlightBuilder.Field fieldBuilder, HighlighterTagsSchema tagsSchema) {
310+
switch (tagsSchema) {
311+
// TODO add HIGHLIGHTER_TAGS_SCHEMA_DEFAULT once new protobufs are published after spec fix
312+
case HIGHLIGHTER_TAGS_SCHEMA_STYLED:
313+
fieldBuilder.preTags(HighlightBuilder.DEFAULT_STYLED_PRE_TAG);
314+
fieldBuilder.postTags(HighlightBuilder.DEFAULT_STYLED_POST_TAGS);
315+
break;
316+
default:
317+
throw new IllegalArgumentException("Unknown tags schema: " + tagsSchema);
318+
}
319+
}
277320
}

‎modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/query/functionscore/DecayFunctionProtoUtils.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
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.protobufs.MultiValueMode;
11+
12+
/**
13+
* Common utility class for decay function Protocol Buffer conversions.
14+
* Contains shared methods used by {@link ExpDecayFunctionProtoUtils},
15+
* {@link GaussDecayFunctionProtoUtils}, and {@link LinearDecayFunctionProtoUtils}.
16+
*/
17+
class DecayFunctionProtoUtils {
18+
19+
private DecayFunctionProtoUtils() {
20+
// Utility class, no instances
21+
}
22+
23+
/**
24+
* Converts a Protocol Buffer MultiValueMode enum to OpenSearch MultiValueMode.
25+
*
26+
* @param multiValueMode the Protocol Buffer MultiValueMode enum value
27+
* @return the corresponding OpenSearch MultiValueMode
28+
* @throws IllegalArgumentException if the multiValueMode is unsupported
29+
*/
30+
static org.opensearch.search.MultiValueMode parseMultiValueMode(MultiValueMode multiValueMode) {
31+
return switch (multiValueMode) {
32+
case MULTI_VALUE_MODE_AVG -> org.opensearch.search.MultiValueMode.AVG;
33+
case MULTI_VALUE_MODE_MAX -> org.opensearch.search.MultiValueMode.MAX;
34+
case MULTI_VALUE_MODE_MIN -> org.opensearch.search.MultiValueMode.MIN;
35+
default -> throw new IllegalArgumentException("Unsupported multi value mode: " + multiValueMode);
36+
};
37+
}
38+
}

‎modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/query/functionscore/ExpDecayFunctionProtoUtils.java‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.opensearch.protobufs.DecayFunction;
1515
import org.opensearch.protobufs.DecayPlacement;
1616
import org.opensearch.protobufs.GeoDecayPlacement;
17+
import org.opensearch.protobufs.MultiValueMode;
1718
import org.opensearch.protobufs.NumericDecayPlacement;
1819
import org.opensearch.transport.grpc.proto.request.common.GeoPointProtoUtils;
1920

@@ -49,21 +50,29 @@ static ScoreFunctionBuilder<?> fromProto(DecayFunction decayFunction) {
4950
String fieldName = entry.getKey();
5051
DecayPlacement decayPlacement = entry.getValue();
5152

53+
ExponentialDecayFunctionBuilder builder;
5254
if (decayPlacement.hasNumericDecayPlacement()) {
53-
return parseNumericExpDecay(fieldName, decayPlacement.getNumericDecayPlacement());
55+
builder = parseNumericExpDecay(fieldName, decayPlacement.getNumericDecayPlacement());
5456
} else if (decayPlacement.hasGeoDecayPlacement()) {
55-
return parseGeoExpDecay(fieldName, decayPlacement.getGeoDecayPlacement());
57+
builder = parseGeoExpDecay(fieldName, decayPlacement.getGeoDecayPlacement());
5658
} else if (decayPlacement.hasDateDecayPlacement()) {
57-
return parseDateExpDecay(fieldName, decayPlacement.getDateDecayPlacement());
59+
builder = parseDateExpDecay(fieldName, decayPlacement.getDateDecayPlacement());
5860
} else {
5961
throw new IllegalArgumentException("Unsupported decay placement type");
6062
}
63+
64+
// Set multi_value_mode if present
65+
if (decayFunction.hasMultiValueMode() && decayFunction.getMultiValueMode() != MultiValueMode.MULTI_VALUE_MODE_UNSPECIFIED) {
66+
builder.setMultiValueMode(DecayFunctionProtoUtils.parseMultiValueMode(decayFunction.getMultiValueMode()));
67+
}
68+
69+
return builder;
6170
}
6271

6372
/**
6473
* Parses a numeric decay placement for exponential decay.
6574
*/
66-
private static ScoreFunctionBuilder<?> parseNumericExpDecay(String fieldName, NumericDecayPlacement numericPlacement) {
75+
private static ExponentialDecayFunctionBuilder parseNumericExpDecay(String fieldName, NumericDecayPlacement numericPlacement) {
6776
ExponentialDecayFunctionBuilder builder;
6877
if (numericPlacement.hasDecay()) {
6978
builder = new ExponentialDecayFunctionBuilder(
@@ -88,7 +97,7 @@ private static ScoreFunctionBuilder<?> parseNumericExpDecay(String fieldName, Nu
8897
/**
8998
* Parses a geo decay placement for exponential decay.
9099
*/
91-
private static ScoreFunctionBuilder<?> parseGeoExpDecay(String fieldName, GeoDecayPlacement geoPlacement) {
100+
private static ExponentialDecayFunctionBuilder parseGeoExpDecay(String fieldName, GeoDecayPlacement geoPlacement) {
92101
GeoPoint geoPoint = GeoPointProtoUtils.parseGeoPoint(geoPlacement.getOrigin());
93102

94103
ExponentialDecayFunctionBuilder builder;
@@ -115,7 +124,7 @@ private static ScoreFunctionBuilder<?> parseGeoExpDecay(String fieldName, GeoDec
115124
/**
116125
* Parses a date decay placement for exponential decay.
117126
*/
118-
private static ScoreFunctionBuilder<?> parseDateExpDecay(String fieldName, DateDecayPlacement datePlacement) {
127+
private static ExponentialDecayFunctionBuilder parseDateExpDecay(String fieldName, DateDecayPlacement datePlacement) {
119128
Object origin = datePlacement.hasOrigin() ? datePlacement.getOrigin() : null;
120129

121130
ExponentialDecayFunctionBuilder builder;

‎modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/query/functionscore/GaussDecayFunctionProtoUtils.java‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.opensearch.protobufs.DecayFunction;
1515
import org.opensearch.protobufs.DecayPlacement;
1616
import org.opensearch.protobufs.GeoDecayPlacement;
17+
import org.opensearch.protobufs.MultiValueMode;
1718
import org.opensearch.protobufs.NumericDecayPlacement;
1819
import org.opensearch.transport.grpc.proto.request.common.GeoPointProtoUtils;
1920

@@ -49,21 +50,29 @@ static ScoreFunctionBuilder<?> fromProto(DecayFunction decayFunction) {
4950
String fieldName = entry.getKey();
5051
DecayPlacement decayPlacement = entry.getValue();
5152

53+
GaussDecayFunctionBuilder builder;
5254
if (decayPlacement.hasNumericDecayPlacement()) {
53-
return parseNumericGaussDecay(fieldName, decayPlacement.getNumericDecayPlacement());
55+
builder = parseNumericGaussDecay(fieldName, decayPlacement.getNumericDecayPlacement());
5456
} else if (decayPlacement.hasGeoDecayPlacement()) {
55-
return parseGeoGaussDecay(fieldName, decayPlacement.getGeoDecayPlacement());
57+
builder = parseGeoGaussDecay(fieldName, decayPlacement.getGeoDecayPlacement());
5658
} else if (decayPlacement.hasDateDecayPlacement()) {
57-
return parseDateGaussDecay(fieldName, decayPlacement.getDateDecayPlacement());
59+
builder = parseDateGaussDecay(fieldName, decayPlacement.getDateDecayPlacement());
5860
} else {
5961
throw new IllegalArgumentException("Unsupported decay placement type");
6062
}
63+
64+
// Set multi_value_mode if present
65+
if (decayFunction.hasMultiValueMode() && decayFunction.getMultiValueMode() != MultiValueMode.MULTI_VALUE_MODE_UNSPECIFIED) {
66+
builder.setMultiValueMode(DecayFunctionProtoUtils.parseMultiValueMode(decayFunction.getMultiValueMode()));
67+
}
68+
69+
return builder;
6170
}
6271

6372
/**
6473
* Parses a numeric decay placement for Gaussian decay.
6574
*/
66-
private static ScoreFunctionBuilder<?> parseNumericGaussDecay(String fieldName, NumericDecayPlacement numericPlacement) {
75+
private static GaussDecayFunctionBuilder parseNumericGaussDecay(String fieldName, NumericDecayPlacement numericPlacement) {
6776
GaussDecayFunctionBuilder builder;
6877
if (numericPlacement.hasDecay()) {
6978
builder = new GaussDecayFunctionBuilder(
@@ -88,7 +97,7 @@ private static ScoreFunctionBuilder<?> parseNumericGaussDecay(String fieldName,
8897
/**
8998
* Parses a geo decay placement for Gaussian decay.
9099
*/
91-
private static ScoreFunctionBuilder<?> parseGeoGaussDecay(String fieldName, GeoDecayPlacement geoPlacement) {
100+
private static GaussDecayFunctionBuilder parseGeoGaussDecay(String fieldName, GeoDecayPlacement geoPlacement) {
92101
GeoPoint geoPoint = GeoPointProtoUtils.parseGeoPoint(geoPlacement.getOrigin());
93102

94103
GaussDecayFunctionBuilder builder;
@@ -115,7 +124,7 @@ private static ScoreFunctionBuilder<?> parseGeoGaussDecay(String fieldName, GeoD
115124
/**
116125
* Parses a date decay placement for Gaussian decay.
117126
*/
118-
private static ScoreFunctionBuilder<?> parseDateGaussDecay(String fieldName, DateDecayPlacement datePlacement) {
127+
private static GaussDecayFunctionBuilder parseDateGaussDecay(String fieldName, DateDecayPlacement datePlacement) {
119128
Object origin = datePlacement.hasOrigin() ? datePlacement.getOrigin() : null;
120129

121130
GaussDecayFunctionBuilder builder;

‎modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/query/functionscore/LinearDecayFunctionProtoUtils.java‎

Lines changed: 15 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.opensearch.protobufs.DecayFunction;
1515
import org.opensearch.protobufs.DecayPlacement;
1616
import org.opensearch.protobufs.GeoDecayPlacement;
17+
import org.opensearch.protobufs.MultiValueMode;
1718
import org.opensearch.protobufs.NumericDecayPlacement;
1819
import org.opensearch.transport.grpc.proto.request.common.GeoPointProtoUtils;
1920

@@ -49,15 +50,23 @@ static ScoreFunctionBuilder<?> fromProto(DecayFunction decayFunction) {
4950
String fieldName = entry.getKey();
5051
DecayPlacement decayPlacement = entry.getValue();
5152

53+
LinearDecayFunctionBuilder builder;
5254
if (decayPlacement.hasNumericDecayPlacement()) {
53-
return parseNumericLinearDecay(fieldName, decayPlacement.getNumericDecayPlacement());
55+
builder = parseNumericLinearDecay(fieldName, decayPlacement.getNumericDecayPlacement());
5456
} else if (decayPlacement.hasGeoDecayPlacement()) {
55-
return parseGeoLinearDecay(fieldName, decayPlacement.getGeoDecayPlacement());
57+
builder = parseGeoLinearDecay(fieldName, decayPlacement.getGeoDecayPlacement());
5658
} else if (decayPlacement.hasDateDecayPlacement()) {
57-
return parseDateLinearDecay(fieldName, decayPlacement.getDateDecayPlacement());
59+
builder = parseDateLinearDecay(fieldName, decayPlacement.getDateDecayPlacement());
5860
} else {
5961
throw new IllegalArgumentException("Unsupported decay placement type");
6062
}
63+
64+
// Set multi_value_mode if present
65+
if (decayFunction.hasMultiValueMode() && decayFunction.getMultiValueMode() != MultiValueMode.MULTI_VALUE_MODE_UNSPECIFIED) {
66+
builder.setMultiValueMode(DecayFunctionProtoUtils.parseMultiValueMode(decayFunction.getMultiValueMode()));
67+
}
68+
69+
return builder;
6170
}
6271

6372
/**
@@ -69,7 +78,7 @@ static ScoreFunctionBuilder<?> fromProto(DecayFunction decayFunction) {
6978
* @param numericPlacement the protobuf numeric decay placement containing origin, scale, offset, and decay
7079
* @return the corresponding OpenSearch LinearDecayFunctionBuilder
7180
*/
72-
private static ScoreFunctionBuilder<?> parseNumericLinearDecay(String fieldName, NumericDecayPlacement numericPlacement) {
81+
private static LinearDecayFunctionBuilder parseNumericLinearDecay(String fieldName, NumericDecayPlacement numericPlacement) {
7382
LinearDecayFunctionBuilder builder;
7483
if (numericPlacement.hasDecay()) {
7584
builder = new LinearDecayFunctionBuilder(
@@ -100,7 +109,7 @@ private static ScoreFunctionBuilder<?> parseNumericLinearDecay(String fieldName,
100109
* @param geoPlacement the protobuf geo decay placement containing origin (lat/lon), scale, offset, and decay
101110
* @return the corresponding OpenSearch LinearDecayFunctionBuilder
102111
*/
103-
private static ScoreFunctionBuilder<?> parseGeoLinearDecay(String fieldName, GeoDecayPlacement geoPlacement) {
112+
private static LinearDecayFunctionBuilder parseGeoLinearDecay(String fieldName, GeoDecayPlacement geoPlacement) {
104113
GeoPoint geoPoint = GeoPointProtoUtils.parseGeoPoint(geoPlacement.getOrigin());
105114

106115
LinearDecayFunctionBuilder builder;
@@ -133,7 +142,7 @@ private static ScoreFunctionBuilder<?> parseGeoLinearDecay(String fieldName, Geo
133142
* @param datePlacement the protobuf date decay placement containing origin (date), scale, offset, and decay
134143
* @return the corresponding OpenSearch LinearDecayFunctionBuilder
135144
*/
136-
private static ScoreFunctionBuilder<?> parseDateLinearDecay(String fieldName, DateDecayPlacement datePlacement) {
145+
private static LinearDecayFunctionBuilder parseDateLinearDecay(String fieldName, DateDecayPlacement datePlacement) {
137146
Object origin = datePlacement.hasOrigin() ? datePlacement.getOrigin() : null;
138147

139148
LinearDecayFunctionBuilder builder;

‎modules/transport-grpc/src/main/java/org/opensearch/transport/grpc/proto/request/search/query/functionscore/RandomScoreFunctionProtoUtils.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ static ScoreFunctionBuilder<?> fromProto(RandomScoreFunction randomScore) {
4848
org.opensearch.protobufs.RandomScoreFunctionSeed seed = randomScore.getSeed();
4949
if (seed.hasInt32()) {
5050
builder.seed(seed.getInt32());
51+
} else if (seed.hasInt64()) {
52+
builder.seed(seed.getInt64());
5153
} else if (seed.hasString()) {
5254
builder.seed(seed.getString());
5355
}

0 commit comments

Comments
 (0)