Skip to content

Commit

Permalink
HACK
Browse files Browse the repository at this point in the history
  • Loading branch information
finnegancarroll committed Sep 23, 2024
1 parent 881c919 commit a4ead60
Show file tree
Hide file tree
Showing 5 changed files with 85 additions and 34 deletions.
3 changes: 1 addition & 2 deletions server/build.gradle
Original file line number Diff line number Diff line change
Expand Up @@ -145,8 +145,7 @@ tasks.withType(JavaCompile).configureEach {
}

compileJava {
options.compilerArgs += ['-processor', ['org.apache.logging.log4j.core.config.plugins.processor.PluginProcessor',
'org.opensearch.common.annotation.processor.ApiAnnotationProcessor'].join(',')]
options.compilerArgs += ['-processor', ['org.apache.logging.log4j.core.config.plugins.processor.PluginProcessor'].join(',')]
}

tasks.named("internalClusterTest").configure {
Expand Down
75 changes: 54 additions & 21 deletions server/src/main/java/org/opensearch/search/SearchHits.java
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,9 @@
import org.opensearch.core.xcontent.ToXContentFragment;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;
import org.opensearch.proto.search.SearchHitsProtoDef;
import org.opensearch.rest.action.search.RestSearchAction;
import org.opensearch.transport.protobuf.SearchHitProtobuf;

import java.io.IOException;
import java.util.ArrayList;
Expand All @@ -54,6 +56,8 @@
import java.util.Objects;

import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
import static org.opensearch.transport.protobuf.ProtoSerDeHelpers.sortFieldToProto;
import static org.opensearch.transport.protobuf.SearchHitProtobuf.sortValueToProto;

/**
* Encapsulates the results of a search operation
Expand Down Expand Up @@ -226,34 +230,63 @@ public static final class Fields {
public static final String MAX_SCORE = "max_score";
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject(Fields.HITS);
boolean totalHitAsInt = params.paramAsBoolean(RestSearchAction.TOTAL_HITS_AS_INT_PARAM, false);
if (totalHitAsInt) {
long total = totalHits == null ? -1 : totalHits.value;
builder.field(Fields.TOTAL, total);
} else if (totalHits != null) {
builder.startObject(Fields.TOTAL);
builder.field("value", totalHits.value);
builder.field("relation", totalHits.relation == Relation.EQUAL_TO ? "eq" : "gte");
builder.endObject();

////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////

public SearchHitsProtoDef.SearchHitsProto toProto() {
SearchHitsProtoDef.SearchHitsProto.Builder builder = SearchHitsProtoDef.SearchHitsProto.newBuilder().setMaxScore(maxScore);

for (SearchHit hit : hits) {
builder.addHits(new SearchHitProtobuf(hit).toProto());
}
if (Float.isNaN(maxScore)) {
builder.nullField(Fields.MAX_SCORE);
} else {
builder.field(Fields.MAX_SCORE, maxScore);

if (collapseField != null) {
builder.setCollapseField(collapseField);
}
builder.field(Fields.HITS);
builder.startArray();
for (SearchHit hit : hits) {
hit.toXContent(builder, params);

if (totalHits != null) {
SearchHitsProtoDef.TotalHitsProto.Builder totHitsBuilder = SearchHitsProtoDef.TotalHitsProto.newBuilder()
.setRelation(totalHits.relation.ordinal())
.setValue(totalHits.value);
builder.setTotalHits(totHitsBuilder);
}

if (sortFields != null) {
for (SortField field : sortFields) {
builder.addSortFields(sortFieldToProto(field));
}
}

if (collapseValues != null) {
for (Object col : collapseValues) {
builder.addCollapseValues(sortValueToProto(col));
}
}
builder.endArray();

return builder.build();
}


@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
System.out.println("SearchHits.toXContent");

builder.startObject(Fields.HITS);

builder.field("protobuf", toProto().toByteString());

builder.endObject();
return builder;
}


////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////


public static SearchHits fromXContent(XContentParser parser) throws IOException {
if (parser.currentToken() != XContentParser.Token.START_OBJECT) {
parser.nextToken();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,17 @@
import org.opensearch.common.annotation.PublicApi;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.proto.search.fetch.FetchSearchResultProtoDef;
import org.opensearch.search.SearchHit;
import org.opensearch.search.SearchHits;
import org.opensearch.search.SearchPhaseResult;
import org.opensearch.search.SearchShardTarget;
import org.opensearch.search.internal.ShardSearchContextId;
import org.opensearch.search.query.QuerySearchResult;

import org.opensearch.transport.protobuf.FetchSearchResultProtobuf;
import org.opensearch.transport.protobuf.SearchHitsProtobuf;

import java.io.IOException;

/**
Expand All @@ -58,12 +62,33 @@ public class FetchSearchResult extends SearchPhaseResult {

public FetchSearchResult() {}

////////////////////////////////////////////////
////////////////////////////////////////////////
////////////////////////////////////////////////

public FetchSearchResult(StreamInput in) throws IOException {
super(in);
contextId = new ShardSearchContextId(in);
hits = new SearchHits(in);
System.out.println("FetchSearchResult In - PROTOBUF");
FetchSearchResultProtoDef.FetchSearchResultProto proto = FetchSearchResultProtoDef.FetchSearchResultProto.parseFrom(in);
hits = new SearchHitsProtobuf(proto.getHits());
}

FetchSearchResultProtoDef.FetchSearchResultProto toProto() {
FetchSearchResultProtoDef.FetchSearchResultProto.Builder builder = FetchSearchResultProtoDef.FetchSearchResultProto.newBuilder()
.setHits(new SearchHitsProtobuf(hits).toProto())
.setCounter(this.counter);
return builder.build();
}

@Override
public void writeTo(StreamOutput out) throws IOException {
System.out.println("FetchSearchResult Out - PROTOBUF");
toProto().writeTo(out);
}

////////////////////////////////////////////////
////////////////////////////////////////////////
////////////////////////////////////////////////

public FetchSearchResult(ShardSearchContextId id, SearchShardTarget shardTarget) {
this.contextId = id;
setSearchShardTarget(shardTarget);
Expand Down Expand Up @@ -103,10 +128,4 @@ public FetchSearchResult initCounter() {
public int counterGetAndIncrement() {
return counter++;
}

@Override
public void writeTo(StreamOutput out) throws IOException {
contextId.writeTo(out);
hits.writeTo(out);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ public void fromProtobufStream(StreamInput in) throws IOException {
fromProto(proto);
}

SearchHitProto toProto() {
public SearchHitProto toProto() {
SearchHitProto.Builder builder = SearchHitProto.newBuilder()
.setScore(score)
.setVersion(version)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ public void fromProtobufStream(StreamInput in) throws IOException {
fromProto(proto);
}

SearchHitsProto toProto() {
public SearchHitsProto toProto() {
SearchHitsProto.Builder builder = SearchHitsProto.newBuilder().setMaxScore(maxScore);

for (SearchHit hit : hits) {
Expand Down

0 comments on commit a4ead60

Please sign in to comment.