Skip to content

Commit

Permalink
Spotless
Browse files Browse the repository at this point in the history
Signed-off-by: Harsha Vamsi Kalluri <[email protected]>
  • Loading branch information
harshavamsi committed May 22, 2024
1 parent 95236d6 commit 76b4abe
Show file tree
Hide file tree
Showing 4 changed files with 87 additions and 87 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -82,8 +82,8 @@
import java.util.function.LongSupplier;
import java.util.function.Supplier;

import static org.apache.lucene.document.LongPoint.pack;
import static org.opensearch.common.time.DateUtils.toLong;
import static org.apache.lucene.document.LongPoint.pack;

/**
* A {@link FieldMapper} for dates.
Expand Down Expand Up @@ -467,16 +467,24 @@ public Query rangeQuery(
}
DateMathParser parser = forcedDateParser == null ? dateMathParser : forcedDateParser;
return dateRangeQuery(lowerTerm, upperTerm, includeLower, includeUpper, timeZone, parser, context, resolution, (l, u) -> {
Query query = new ApproximateableQuery(new PointRangeQuery(name(), pack(new long[]{l}).bytes, pack(new long[]{u}).bytes, new long[]{l}.length) {
protected String toString(int dimension, byte[] value) {
return Long.toString(LongPoint.decodeDimension(value, 0));
Query query = new ApproximateableQuery(
new PointRangeQuery(name(), pack(new long[] { l }).bytes, pack(new long[] { u }).bytes, new long[] { l }.length) {
protected String toString(int dimension, byte[] value) {
return Long.toString(LongPoint.decodeDimension(value, 0));
}
},
new ApproximatePointRangeQuery(
name(),
pack(new long[] { l }).bytes,
pack(new long[] { u }).bytes,
new long[] { l }.length
) {
@Override
protected String toString(int dimension, byte[] value) {
return Long.toString(LongPoint.decodeDimension(value, 0));
}
}
}, new ApproximatePointRangeQuery(name(), pack(new long[]{l}).bytes, pack(new long[]{u}).bytes, new long[]{l}.length) {
@Override
protected String toString(int dimension, byte[] value) {
return Long.toString(LongPoint.decodeDimension(value, 0));
}
});
);
if (hasDocValues()) {
Query dvQuery = SortedNumericDocValuesField.newSlowRangeQuery(name(), l, u);
query = new IndexOrDocValuesQuery(query, dvQuery);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upp
this(field, lowerPoint, upperPoint, numDims, 10_000);
}

protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upperPoint, int numDims, int size){
protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upperPoint, int numDims, int size) {
checkArgs(field, lowerPoint, upperPoint);
this.field = field;
if (numDims <= 0) {
Expand All @@ -62,10 +62,8 @@ protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upp
}
if (lowerPoint.length != upperPoint.length) {
throw new IllegalArgumentException(
"lowerPoint has length="
+ lowerPoint.length
+ " but upperPoint has different length="
+ upperPoint.length);
"lowerPoint has length=" + lowerPoint.length + " but upperPoint has different length=" + upperPoint.length
);
}
this.numDims = numDims;
this.bytesPerDim = lowerPoint.length / numDims;
Expand All @@ -74,11 +72,11 @@ protected ApproximatePointRangeQuery(String field, byte[] lowerPoint, byte[] upp
this.upperPoint = upperPoint;
}

public int getSize(){
public int getSize() {
return this.size;
}

public void setSize(int size){
public void setSize(int size) {
this.size = size;
}

Expand Down Expand Up @@ -126,9 +124,8 @@ private PointValues.Relation relate(byte[] minPackedValue, byte[] maxPackedValue
return PointValues.Relation.CELL_OUTSIDE_QUERY;
}

crosses |=
comparator.compare(minPackedValue, offset, lowerPoint, offset) < 0
|| comparator.compare(maxPackedValue, offset, upperPoint, offset) > 0;
crosses |= comparator.compare(minPackedValue, offset, lowerPoint, offset) < 0
|| comparator.compare(maxPackedValue, offset, upperPoint, offset) > 0;
}

if (crosses) {
Expand Down Expand Up @@ -256,7 +253,8 @@ private boolean checkValidPointValues(PointValues values) throws IOException {
+ "\" was indexed with numIndexDimensions="
+ values.getNumIndexDimensions()
+ " but this query has numDims="
+ numDims);
+ numDims
);
}
if (bytesPerDim != values.getBytesPerDimension()) {
throw new IllegalArgumentException(
Expand All @@ -265,7 +263,8 @@ private boolean checkValidPointValues(PointValues values) throws IOException {
+ "\" was indexed with bytesPerDim="
+ values.getBytesPerDimension()
+ " but this query has bytesPerDim="
+ bytesPerDim);
+ bytesPerDim
);
}
return true;
}
Expand All @@ -275,7 +274,8 @@ private void intersect(PointValues.PointTree pointTree, PointValues.IntersectVis
assert pointTree.moveToParent() == false;
}

private long intersect(PointValues.IntersectVisitor visitor, PointValues.PointTree pointTree, int count, long docCount) throws IOException {
private long intersect(PointValues.IntersectVisitor visitor, PointValues.PointTree pointTree, int count, long docCount)
throws IOException {
PointValues.Relation r = visitor.compare(pointTree.getMinPackedValue(), pointTree.getMaxPackedValue());
if (docCount >= count) {
return 0;
Expand Down Expand Up @@ -303,8 +303,7 @@ private long intersect(PointValues.IntersectVisitor visitor, PointValues.PointTr
// Leaf node; scan and filter all points in this block:
if (docCount <= count) {
pointTree.visitDocValues(visitor);
}
else break;
} else break;
}
break;
default:
Expand Down Expand Up @@ -363,8 +362,7 @@ public ScorerSupplier scorerSupplier(LeafReaderContext context) throws IOExcepti
return new ScorerSupplier() {
@Override
public Scorer get(long leadCost) {
return new ConstantScoreScorer(
weight, score(), scoreMode, DocIdSetIterator.all(reader.maxDoc()));
return new ConstantScoreScorer(weight, score(), scoreMode, DocIdSetIterator.all(reader.maxDoc()));
}

@Override
Expand All @@ -389,7 +387,7 @@ && cost() > reader.maxDoc() / 2) {
// by computing the set of documents that do NOT match the range
final FixedBitSet result = new FixedBitSet(reader.maxDoc());
result.set(0, reader.maxDoc());
long[] cost = new long[]{reader.maxDoc()};
long[] cost = new long[] { reader.maxDoc() };
intersect(values.getPointTree(), getInverseIntersectVisitor(result, cost), size);
final DocIdSetIterator iterator = new BitSetIterator(result, cost[0]);
return new ConstantScoreScorer(weight, score(), scoreMode, iterator);
Expand Down Expand Up @@ -432,8 +430,7 @@ public int count(LeafReaderContext context) throws IOException {
}

if (reader.hasDeletions() == false) {
if (relate(values.getMinPackedValue(), values.getMaxPackedValue())
== PointValues.Relation.CELL_INSIDE_QUERY) {
if (relate(values.getMinPackedValue(), values.getMaxPackedValue()) == PointValues.Relation.CELL_INSIDE_QUERY) {
return values.getDocCount();
}
// only 1D: we have the guarantee that it will actually run fast since there are at most 2
Expand Down Expand Up @@ -463,40 +460,36 @@ public int count(LeafReaderContext context) throws IOException {
private long pointCount(
PointValues.PointTree pointTree,
BiFunction<byte[], byte[], PointValues.Relation> nodeComparator,
Predicate<byte[]> leafComparator)
throws IOException {
final long[] matchingNodeCount = {0};
Predicate<byte[]> leafComparator
) throws IOException {
final long[] matchingNodeCount = { 0 };
// create a custom IntersectVisitor that records the number of leafNodes that matched
final PointValues.IntersectVisitor visitor =
new PointValues.IntersectVisitor() {
@Override
public void visit(int docID) {
// this branch should be unreachable
throw new UnsupportedOperationException(
"This IntersectVisitor does not perform any actions on a "
+ "docID="
+ docID
+ " node being visited");
}
final PointValues.IntersectVisitor visitor = new PointValues.IntersectVisitor() {
@Override
public void visit(int docID) {
// this branch should be unreachable
throw new UnsupportedOperationException(
"This IntersectVisitor does not perform any actions on a " + "docID=" + docID + " node being visited"
);
}

@Override
public void visit(int docID, byte[] packedValue) {
if (leafComparator.test(packedValue)) {
matchingNodeCount[0]++;
}
@Override
public void visit(int docID, byte[] packedValue) {
if (leafComparator.test(packedValue)) {
matchingNodeCount[0]++;
}
}

@Override
public PointValues.Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
return nodeComparator.apply(minPackedValue, maxPackedValue);
}
};
@Override
public PointValues.Relation compare(byte[] minPackedValue, byte[] maxPackedValue) {
return nodeComparator.apply(minPackedValue, maxPackedValue);
}
};
pointCount(visitor, pointTree, matchingNodeCount);
return matchingNodeCount[0];
}

private void pointCount(
PointValues.IntersectVisitor visitor, PointValues.PointTree pointTree, long[] matchingNodeCount)
private void pointCount(PointValues.IntersectVisitor visitor, PointValues.PointTree pointTree, long[] matchingNodeCount)
throws IOException {
PointValues.Relation r = visitor.compare(pointTree.getMinPackedValue(), pointTree.getMaxPackedValue());
switch (r) {
Expand All @@ -509,10 +502,10 @@ private void pointCount(
matchingNodeCount[0] += pointTree.size();
return;
case CELL_CROSSES_QUERY:
/*
The cell crosses the shape boundary, or the cell fully contains the query, so we fall
through and do full counting.
*/
/*
The cell crosses the shape boundary, or the cell fully contains the query, so we fall
through and do full counting.
*/
if (pointTree.moveToChild()) {
do {
pointCount(visitor, pointTree, matchingNodeCount);
Expand Down Expand Up @@ -597,13 +590,9 @@ public final String toString(String field) {
int startOffset = bytesPerDim * i;

sb.append('[');
sb.append(
toString(
i, ArrayUtil.copyOfSubArray(lowerPoint, startOffset, startOffset + bytesPerDim)));
sb.append(toString(i, ArrayUtil.copyOfSubArray(lowerPoint, startOffset, startOffset + bytesPerDim)));
sb.append(" TO ");
sb.append(
toString(
i, ArrayUtil.copyOfSubArray(upperPoint, startOffset, startOffset + bytesPerDim)));
sb.append(toString(i, ArrayUtil.copyOfSubArray(upperPoint, startOffset, startOffset + bytesPerDim)));
sb.append(']');
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -32,21 +32,26 @@ public final class ApproximateableQuery extends Query {
private Weight originalQueryWeight, approximationQueryWeight;

public ApproximateableQuery(Query originalQuery, Query approximationQuery) {
this(originalQuery, approximationQuery, null, null);
this(originalQuery, approximationQuery, null, null);
}

public ApproximateableQuery(Query originalQuery, Query approximationQuery, Weight originalQueryWeight, Weight approximationQueryWeight){
public ApproximateableQuery(
Query originalQuery,
Query approximationQuery,
Weight originalQueryWeight,
Weight approximationQueryWeight
) {
this.originalQuery = originalQuery;
this.approximationQuery = approximationQuery;
this.originalQueryWeight = originalQueryWeight;
this.approximationQueryWeight = approximationQueryWeight;
}

public void setOriginalQueryWeight(Weight originalQueryWeight){
public void setOriginalQueryWeight(Weight originalQueryWeight) {
this.originalQueryWeight = originalQueryWeight;
}

public void setApproximationQueryWeight(Weight approximationQueryWeight){
public void setApproximationQueryWeight(Weight approximationQueryWeight) {
this.approximationQueryWeight = approximationQueryWeight;
}

Expand Down Expand Up @@ -95,7 +100,8 @@ public ScorerSupplier scorerSupplier(LeafReaderContext leafReaderContext) throws
return new ScorerSupplier() {
@Override
public Scorer get(long l) throws IOException {
// TODO: we need to figure out how to compute the cost of running two different queries, by default return the original query's scoreSupplier
// TODO: we need to figure out how to compute the cost of running two different queries, by default return the
// original query's scoreSupplier
return originalQueryScoreSupplier.get(l);
}

Expand All @@ -122,7 +128,6 @@ public boolean isCacheable(LeafReaderContext leafReaderContext) {
};
}


@Override
public String toString(String s) {
return "ApproximateableQuery(originalQuery="
Expand All @@ -141,7 +146,7 @@ public void visit(QueryVisitor queryVisitor) {

@Override
public boolean equals(Object o) {
if(!sameClassAs(o)){
if (!sameClassAs(o)) {
return false;
}
ApproximateableQuery that = (ApproximateableQuery) o;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@

import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.lucene.document.LongPoint;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexReader;
import org.apache.lucene.index.LeafReaderContext;
Expand All @@ -57,7 +56,6 @@
import org.apache.lucene.search.ScoreDoc;
import org.apache.lucene.search.ScoreMode;
import org.apache.lucene.search.Scorer;
import org.apache.lucene.search.ScorerSupplier;
import org.apache.lucene.search.TermStatistics;
import org.apache.lucene.search.TopFieldDocs;
import org.apache.lucene.search.TotalHits;
Expand All @@ -71,8 +69,6 @@
import org.opensearch.common.annotation.PublicApi;
import org.opensearch.common.lease.Releasable;
import org.opensearch.common.lucene.search.TopDocsAndMaxScore;
import org.opensearch.index.mapper.DateFieldMapper;
import org.opensearch.index.query.RangeQueryBuilder;
import org.opensearch.search.DocValueFormat;
import org.opensearch.search.SearchService;
import org.opensearch.search.approximate.ApproximatePointRangeQuery;
Expand Down Expand Up @@ -329,10 +325,9 @@ private void searchLeaf(LeafReaderContext ctx, Weight weight, Collector collecto
// catch early terminated exception and rethrow?
Bits liveDocs = ctx.reader().getLiveDocs();
BitSet liveDocsBitSet = getSparseBitSetOrNull(liveDocs);
if(isApproximateableRangeQuery()){
if (isApproximateableRangeQuery()) {
ApproximateableQuery query = ((ApproximateableQuery) ((IndexOrDocValuesQuery) searchContext.query()).getIndexQuery());
if (searchContext.size() > 10_000)
((ApproximatePointRangeQuery) query.getApproximationQuery()).setSize(searchContext.size());
if (searchContext.size() > 10_000) ((ApproximatePointRangeQuery) query.getApproximationQuery()).setSize(searchContext.size());
weight = query.getApproximationQueryWeight();
}
if (liveDocsBitSet == null) {
Expand Down Expand Up @@ -426,18 +421,21 @@ private static BitSet getSparseBitSetOrNull(Bits liveDocs) {

}

private boolean isApproximateableRangeQuery(){
boolean isTopLevelRangeQuery = searchContext.query() instanceof IndexOrDocValuesQuery &&
((ApproximateableQuery) ((IndexOrDocValuesQuery) searchContext.query()).getIndexQuery()).getOriginalQuery() instanceof PointRangeQuery;
private boolean isApproximateableRangeQuery() {
boolean isTopLevelRangeQuery = searchContext.query() instanceof IndexOrDocValuesQuery
&& ((ApproximateableQuery) ((IndexOrDocValuesQuery) searchContext.query()).getIndexQuery())
.getOriginalQuery() instanceof PointRangeQuery;

boolean hasSort = false;

if (searchContext.request() != null && searchContext.request().source() != null ) {
if (searchContext.request() != null && searchContext.request().source() != null) {
FieldSortBuilder primarySortField = FieldSortBuilder.getPrimaryFieldSortOrNull(searchContext.request().source());
if (primarySortField != null && primarySortField.missing() == null && Objects.equals(searchContext.trackTotalHitsUpTo(), SearchContext.TRACK_TOTAL_HITS_DISABLED)) {
hasSort = true;
}
if (primarySortField != null
&& primarySortField.missing() == null
&& Objects.equals(searchContext.trackTotalHitsUpTo(), SearchContext.TRACK_TOTAL_HITS_DISABLED)) {
hasSort = true;
}
}

return isTopLevelRangeQuery && !hasSort;
}
Expand Down

0 comments on commit 76b4abe

Please sign in to comment.