com.couchbase.client.core.protostellar.search.ProtostellarCoreSearchOps Maven / Gradle / Ivy
Show all versions of core-io Show documentation
/*
* Copyright (c) 2023 Couchbase, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.couchbase.client.core.protostellar.search;
import com.couchbase.client.core.CoreProtostellar;
import com.couchbase.client.core.annotation.Stability;
import com.couchbase.client.core.api.kv.CoreAsyncResponse;
import com.couchbase.client.core.api.kv.CoreDurability;
import com.couchbase.client.core.api.manager.CoreBucketAndScope;
import com.couchbase.client.core.api.search.CoreSearchKeyset;
import com.couchbase.client.core.api.search.CoreSearchMetaData;
import com.couchbase.client.core.api.search.CoreSearchOps;
import com.couchbase.client.core.api.search.CoreSearchOptions;
import com.couchbase.client.core.api.search.CoreSearchQuery;
import com.couchbase.client.core.api.search.CoreSearchScanConsistency;
import com.couchbase.client.core.api.search.facet.CoreDateRange;
import com.couchbase.client.core.api.search.facet.CoreDateRangeFacet;
import com.couchbase.client.core.api.search.facet.CoreNumericRange;
import com.couchbase.client.core.api.search.facet.CoreNumericRangeFacet;
import com.couchbase.client.core.api.search.facet.CoreSearchFacet;
import com.couchbase.client.core.api.search.facet.CoreTermFacet;
import com.couchbase.client.core.api.search.result.CoreDateRangeSearchFacetResult;
import com.couchbase.client.core.api.search.result.CoreNumericRangeSearchFacetResult;
import com.couchbase.client.core.api.search.result.CoreReactiveSearchResult;
import com.couchbase.client.core.api.search.result.CoreSearchDateRange;
import com.couchbase.client.core.api.search.result.CoreSearchFacetResult;
import com.couchbase.client.core.api.search.result.CoreSearchMetrics;
import com.couchbase.client.core.api.search.result.CoreSearchNumericRange;
import com.couchbase.client.core.api.search.result.CoreSearchResult;
import com.couchbase.client.core.api.search.result.CoreSearchRow;
import com.couchbase.client.core.api.search.result.CoreSearchRowLocation;
import com.couchbase.client.core.api.search.result.CoreSearchRowLocations;
import com.couchbase.client.core.api.search.result.CoreSearchTermRange;
import com.couchbase.client.core.api.search.result.CoreTermSearchFacetResult;
import com.couchbase.client.core.cnc.RequestSpan;
import com.couchbase.client.core.cnc.TracingIdentifiers;
import com.couchbase.client.core.deps.com.google.protobuf.ByteString;
import com.couchbase.client.core.deps.com.google.protobuf.Timestamp;
import com.couchbase.client.core.error.CouchbaseException;
import com.couchbase.client.core.json.Mapper;
import com.couchbase.client.core.protostellar.CoreProtostellarAccessorsStreaming;
import com.couchbase.client.core.protostellar.CoreProtostellarErrorHandlingUtil;
import com.couchbase.client.core.protostellar.ProtostellarRequest;
import com.couchbase.client.core.retry.ProtostellarRequestBehaviour;
import com.couchbase.client.core.service.ServiceType;
import com.couchbase.client.protostellar.search.v1.DateRange;
import com.couchbase.client.protostellar.search.v1.DateRangeFacet;
import com.couchbase.client.protostellar.search.v1.Facet;
import com.couchbase.client.protostellar.search.v1.NumericRange;
import com.couchbase.client.protostellar.search.v1.NumericRangeFacet;
import com.couchbase.client.protostellar.search.v1.SearchQueryRequest;
import com.couchbase.client.protostellar.search.v1.SearchQueryResponse;
import com.couchbase.client.protostellar.search.v1.TermFacet;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.publisher.Sinks;
import reactor.util.annotation.NonNull;
import reactor.util.annotation.Nullable;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.time.Duration;
import java.time.Instant;
import java.util.ArrayList;
import java.util.List;
import java.util.ListIterator;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import static com.couchbase.client.core.protostellar.CoreProtostellarUtil.createSpan;
import static com.couchbase.client.core.protostellar.CoreProtostellarUtil.unsupportedInProtostellar;
import static com.couchbase.client.core.util.CbCollections.transform;
import static com.couchbase.client.core.util.CbCollections.transformValues;
import static com.couchbase.client.core.util.ProtostellarUtil.convert;
import static com.couchbase.client.core.util.Validators.notNull;
import static java.util.Collections.emptyMap;
import static java.util.Objects.requireNonNull;
@Stability.Internal
public class ProtostellarCoreSearchOps implements CoreSearchOps {
private final CoreProtostellar core;
public ProtostellarCoreSearchOps(CoreProtostellar core, @Nullable CoreBucketAndScope scope) {
this.core = requireNonNull(core);
// scope is silently ignored as it requires ING-381.
// throwing here would require creating a new CoreSearchOps object on every search operation.
}
@Override
public CoreAsyncResponse searchQueryAsync(String indexName,
CoreSearchQuery search,
CoreSearchOptions options) {
ProtostellarRequest request = request(core, indexName, search, options);
CoreAsyncResponse> responses = CoreProtostellarAccessorsStreaming.async(core,
request,
(endpoint, stream) -> endpoint.searchStub()
.withDeadline(request.deadline())
.searchQuery(request.request(), stream),
(error) -> CoreProtostellarErrorHandlingUtil.convertException(core, request, error)
);
return responses.map(results -> {
List rows = new ArrayList<>();
CoreSearchMetaData metaData = null;
Map facets = emptyMap();
for (SearchQueryResponse r : results) {
if (r.hasMetaData()) {
metaData = parseMetadata(r);
facets = parseFacets(r);
}
r.getHitsList()
.forEach(hit -> rows.add(parse(hit)));
}
return new CoreSearchResult(rows, facets, metaData);
});
}
private static CoreSearchMetaData parseMetadata(SearchQueryResponse response) {
SearchQueryResponse.MetaData md = response.getMetaData();
SearchQueryResponse.SearchMetrics metrics = md.getMetrics();
return new CoreSearchMetaData(
md.getErrorsMap(),
new CoreSearchMetrics(
convert(metrics.getExecutionTime()),
metrics.getTotalRows(),
metrics.getMaxScore(),
metrics.getSuccessPartitionCount(),
metrics.getErrorPartitionCount()
)
);
}
@Override
public Mono searchQueryReactive(String indexName,
CoreSearchQuery query,
CoreSearchOptions options) {
return Mono.defer(() -> {
try {
ProtostellarRequest request = request(core, indexName, query, options);
Sinks.One out = Sinks.one();
Flux responses = CoreProtostellarAccessorsStreaming.reactive(core,
request,
(endpoint, stream) -> endpoint.searchStub()
.withDeadline(request.deadline())
.searchQuery(request.request(), stream),
(error) -> CoreProtostellarErrorHandlingUtil.convertException(core, request, error)
);
Sinks.Many rows = Sinks.many().unicast().onBackpressureBuffer();
Sinks.One metaData = Sinks.one();
Sinks.One> facets = Sinks.one();
responses.publishOn(core.context().environment().scheduler())
.subscribe(response -> {
response.getHitsList().forEach(hit -> {
CoreSearchRow row = parse(hit);
rows.tryEmitNext(row).orThrow();
});
if (response.hasMetaData()) {
CoreSearchMetaData cmd = parseMetadata(response);
metaData.tryEmitValue(cmd).orThrow();
if (response.getFacetsCount() > 0) {
Map coreFacets = parseFacets(response);
facets.tryEmitValue(coreFacets).orThrow();
} else {
facets.tryEmitValue(emptyMap()).orThrow();
}
}
},
// Error has already been passed through CoreProtostellarErrorHandlingUtil in CoreProtostellarAccessorsStreaming
throwable -> rows.tryEmitError(throwable).orThrow(),
() -> rows.tryEmitComplete().orThrow());
return Mono.just(new CoreReactiveSearchResult(rows.asFlux(), facets.asMono(), metaData.asMono()));
} catch (Throwable err) {
// Any errors from initial option validation
return Mono.error(err);
}
});
}
private static Map parseFacets(SearchQueryResponse response) {
return transformValues(response.getFacetsMap(), ProtostellarCoreSearchOps::convertFacetResult);
}
private static CoreSearchFacetResult convertFacetResult(String facetName, SearchQueryResponse.FacetResult facet) {
if (facet.hasTermFacet()) {
SearchQueryResponse.TermFacetResult result = facet.getTermFacet();
return new CoreTermSearchFacetResult(
facetName,
result.getField(),
result.getTotal(),
result.getMissing(),
result.getOther(),
transform(result.getTermsList(), it -> new CoreSearchTermRange(
it.getName(),
it.getSize()
))
);
}
if (facet.hasNumericRangeFacet()) {
SearchQueryResponse.NumericRangeFacetResult result = facet.getNumericRangeFacet();
return new CoreNumericRangeSearchFacetResult(
facetName,
result.getField(),
result.getTotal(),
result.getMissing(),
result.getOther(),
transform(result.getNumericRangesList(), it -> new CoreSearchNumericRange(
it.getName(),
parseNumericRangeEndpoint(it.getMin()),
parseNumericRangeEndpoint(it.getMax()),
it.getSize()
))
);
}
if (facet.hasDateRangeFacet()) {
SearchQueryResponse.DateRangeFacetResult result = facet.getDateRangeFacet();
return new CoreDateRangeSearchFacetResult(
facetName,
result.getField(),
result.getTotal(),
result.getMissing(),
result.getOther(),
transform(result.getDateRangesList(), it -> new CoreSearchDateRange(
it.getName(),
it.hasStart() ? toInstant(it.getStart()) : null,
it.hasEnd() ? toInstant(it.getEnd()) : null,
it.getSize()
))
);
}
throw new RuntimeException("Unexpected facet result type: " + facet);
}
/**
* Workaround for a Protostellar API bug where these endpoints are represented
* as `long` instead of `Double`.
*
* TODO: remove this method once NumericRangeResult.min() and max() return Double.
* Might need to have caller add a hasMin() / hasMax() check as well.
*/
private static Double parseNumericRangeEndpoint(@Nullable Number n) {
return n == null ? null : n.doubleValue();
}
private static CoreSearchRow parse(SearchQueryResponse.SearchQueryRow row) {
return new CoreSearchRow(
row.getIndex(),
row.getId(),
row.getScore(),
row.getExplanation().toByteArray(),
parseLocations(row),
parseFragments(row),
parseFields(row),
() -> CoreSearchKeyset.EMPTY // pending ING-476
);
}
private static byte[] parseFields(SearchQueryResponse.SearchQueryRow hit) {
// Turn this map back into a JSON Object.
Map fields = hit.getFieldsMap();
if (fields.isEmpty()) {
return new byte[]{'{', '}'};
}
ByteArrayOutputStream os = new ByteArrayOutputStream();
os.write('{');
hit.getFieldsMap().forEach((key, value) -> {
try {
Mapper.writer().writeValue(os, key);
os.write(':');
os.write(value.toByteArray());
os.write(',');
} catch (IOException e) {
throw new RuntimeException(e);
}
});
byte[] result = os.toByteArray();
result[result.length - 1] = '}'; // overwrite trailing comma
return result;
}
private static Map> parseFragments(SearchQueryResponse.SearchQueryRow row) {
return transformValues(row.getFragmentsMap(), SearchQueryResponse.Fragment::getContentList);
}
private static Optional parseLocations(SearchQueryResponse.SearchQueryRow row) {
if (row.getLocationsCount() == 0) {
return Optional.empty();
}
List result = new ArrayList<>(row.getLocationsCount());
row.getLocationsList().forEach(loc -> result.add(parseOneLocation(loc)));
return Optional.of(CoreSearchRowLocations.from(result));
}
private static CoreSearchRowLocation parseOneLocation(SearchQueryResponse.Location loc) {
return new CoreSearchRowLocation(
loc.getField(),
loc.getTerm(),
loc.getPosition(),
loc.getStart(),
loc.getEnd(),
loc.getArrayPositionsCount() == 0 ? null : toPrimitiveLongArray(loc.getArrayPositionsList())
);
}
private static long[] toPrimitiveLongArray(List list) {
long[] result = new long[list.size()];
for (ListIterator i = list.listIterator(); i.hasNext(); ) {
result[i.nextIndex()] = i.next().longValue();
}
return result;
}
private static ProtostellarRequest request(CoreProtostellar core,
String indexName,
CoreSearchQuery query,
CoreSearchOptions opts) {
notNull(indexName, "IndexName");
notNull(query, "Query");
notNull(opts, "SearchOptions");
opts.validate();
Duration timeout = opts.commonOptions().timeout().orElse(core.context().environment().timeoutConfig().queryTimeout());
RequestSpan span = createSpan(core, TracingIdentifiers.SPAN_REQUEST_SEARCH, CoreDurability.NONE, opts.commonOptions().parentSpan().orElse(null));
SearchQueryRequest.Builder request = SearchQueryRequest.newBuilder()
.setIndexName(indexName)
.setQuery(query.asProtostellar());
if (opts.consistency() != null) {
if (opts.consistency() == CoreSearchScanConsistency.NOT_BOUNDED) {
request.setScanConsistency(SearchQueryRequest.ScanConsistency.SCAN_CONSISTENCY_NOT_BOUNDED);
}
}
if (opts.limit() != null) {
request.setLimit(opts.limit());
}
if (opts.skip() != null) {
request.setSkip(opts.skip());
}
if (opts.searchBefore() != null || opts.searchAfter() != null) {
throw unsupportedInProtostellar("keyset pagination with searchBefore/After");
}
if (opts.explain() != null) {
request.setIncludeExplanation(opts.explain());
}
if (opts.highlightStyle() != null) {
switch (opts.highlightStyle()) {
case HTML:
request.setHighlightStyle(SearchQueryRequest.HighlightStyle.HIGHLIGHT_STYLE_HTML);
break;
case ANSI:
request.setHighlightStyle(SearchQueryRequest.HighlightStyle.HIGHLIGHT_STYLE_ANSI);
break;
case SERVER_DEFAULT:
request.setHighlightStyle(SearchQueryRequest.HighlightStyle.HIGHLIGHT_STYLE_DEFAULT);
break;
}
}
if (!opts.highlightFields().isEmpty()) {
request.addAllHighlightFields(opts.highlightFields());
}
if (!opts.fields().isEmpty()) {
request.addAllFields(opts.fields());
}
opts.sort().forEach(sort -> request.addSort(sort.asProtostellar()));
if (opts.disableScoring() != null) {
request.setDisableScoring(opts.disableScoring());
}
if (!opts.collections().isEmpty()) {
request.addAllCollections(opts.collections());
}
if (opts.includeLocations() != null) {
request.setIncludeExplanation(opts.includeLocations());
}
opts.facets().forEach((name, facet) -> request.putFacets(name, convertFacet(facet)));
return new ProtostellarRequest<>(
request.build(),
core,
ServiceType.SEARCH,
TracingIdentifiers.SPAN_REQUEST_SEARCH,
span,
timeout,
false,
opts.commonOptions().retryStrategy().orElse(core.context().environment().retryStrategy()),
opts.commonOptions().clientContext(),
0L,
null
);
}
private static Facet convertFacet(CoreSearchFacet facet) {
if (facet instanceof CoreTermFacet) {
return convertTermFacet(facet);
}
if (facet instanceof CoreNumericRangeFacet) {
return convertNumericRangeFacet(facet);
}
if (facet instanceof CoreDateRangeFacet) {
return convertDateRangeFacet(facet);
}
throw new RuntimeException("Unexpected facet type: " + facet.getClass());
}
@NonNull
private static Facet convertDateRangeFacet(CoreSearchFacet facet) {
DateRangeFacet.Builder builder = DateRangeFacet.newBuilder()
.setField(facet.field());
Integer size = facet.size();
if (size != null) {
builder.setSize(size);
}
List coreRanges = ((CoreDateRangeFacet) facet).dateRanges();
coreRanges.forEach(it -> builder.addDateRanges(convertDateRange(it)));
return Facet.newBuilder()
.setDateRangeFacet(builder)
.build();
}
@NonNull
private static Facet convertNumericRangeFacet(CoreSearchFacet facet) {
NumericRangeFacet.Builder builder = NumericRangeFacet.newBuilder()
.setField(facet.field());
Integer size = facet.size();
if (size != null) {
builder.setSize(size);
}
List coreRanges = ((CoreNumericRangeFacet) facet).ranges();
coreRanges.forEach(it -> builder.addNumericRanges(convertNumericRange(it)));
return Facet.newBuilder()
.setNumericRangeFacet(builder)
.build();
}
@NonNull
private static Facet convertTermFacet(CoreSearchFacet facet) {
TermFacet.Builder builder = TermFacet.newBuilder()
.setField(facet.field());
Integer size = facet.size();
if (size != null) {
builder.setSize(size);
}
return Facet.newBuilder()
.setTermFacet(builder)
.build();
}
private static DateRange.Builder convertDateRange(CoreDateRange range) {
DateRange.Builder builder = DateRange.newBuilder()
.setName(range.name());
String start = range.start();
if (start != null) {
builder.setStart(start);
}
String end = range.end();
if (end != null) {
builder.setEnd(end);
}
return builder;
}
private static NumericRange.Builder convertNumericRange(CoreNumericRange range) {
NumericRange.Builder builder = NumericRange.newBuilder()
.setName(range.name());
Double min = range.min();
if (min != null) {
builder.setMin(min.floatValue());
}
Double max = range.max();
if (max != null) {
builder.setMax(max.floatValue());
}
return builder;
}
private static Instant toInstant(Timestamp ts) {
return Instant.ofEpochSecond(ts.getSeconds(), ts.getNanos());
}
}