Skip to content

Commit

Permalink
[BUG] Flaky search.aggregation/230_composite_unsigned tests due to ne…
Browse files Browse the repository at this point in the history
…sted aggregations order (#7548)

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>
(cherry picked from commit af78ef9)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
  • Loading branch information
github-actions[bot] committed May 16, 2023
1 parent 1dbdf76 commit 7511ece
Show file tree
Hide file tree
Showing 2 changed files with 151 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,7 @@ int hashCodeCurrent() {
}
}

int compareValues(long v1, long v2) {
private int compareValues(long v1, long v2) {
return Long.compare(v1, v2) * reverseMul;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,21 +14,35 @@
import org.apache.lucene.search.PointRangeQuery;
import org.apache.lucene.search.Query;
import org.opensearch.common.CheckedFunction;
import org.opensearch.common.Numbers;
import org.opensearch.common.lease.Releasables;
import org.opensearch.common.util.BigArrays;
import org.opensearch.common.util.BitArray;
import org.opensearch.common.util.LongArray;
import org.opensearch.index.mapper.MappedFieldType;
import org.opensearch.index.mapper.NumberFieldMapper;
import org.opensearch.search.DocValueFormat;
import org.opensearch.search.aggregations.LeafBucketCollector;
import org.opensearch.search.aggregations.bucket.missing.MissingOrder;

import java.io.IOException;
import java.util.function.LongUnaryOperator;
import java.math.BigInteger;
import java.util.Objects;

/**
* A {@link SingleDimensionValuesSource} for unsigned longs.
*
* @opensearch.internal
*/
public class UnsignedLongValuesSource extends LongValuesSource {
public class UnsignedLongValuesSource extends SingleDimensionValuesSource<BigInteger> {
private final BigArrays bigArrays;
private final CheckedFunction<LeafReaderContext, SortedNumericDocValues, IOException> docValuesFunc;

private BitArray bits;
private LongArray values;
private long currentValue;
private boolean missingCurrentValue;

public UnsignedLongValuesSource(
BigArrays bigArrays,
MappedFieldType fieldType,
Expand All @@ -39,18 +53,148 @@ public UnsignedLongValuesSource(
int size,
int reverseMul
) {
super(bigArrays, fieldType, docValuesFunc, LongUnaryOperator.identity(), format, missingBucket, missingOrder, size, reverseMul);
super(bigArrays, format, fieldType, missingBucket, missingOrder, size, reverseMul);
this.bigArrays = bigArrays;
this.docValuesFunc = docValuesFunc;
this.bits = missingBucket ? new BitArray(Math.min(size, 100), bigArrays) : null;
this.values = bigArrays.newLongArray(Math.min(size, 100), false);
}

@Override
void copyCurrent(int slot) {
values = bigArrays.grow(values, slot + 1);
if (missingBucket && missingCurrentValue) {
bits.clear(slot);
} else {
assert missingCurrentValue == false;
if (missingBucket) {
bits.set(slot);
}
values.set(slot, currentValue);
}
}

@Override
int compare(int from, int to) {
if (missingBucket) {
int result = missingOrder.compare(() -> bits.get(from) == false, () -> bits.get(to) == false, reverseMul);
if (MissingOrder.unknownOrder(result) == false) {
return result;
}
}
return compareValues(values.get(from), values.get(to));
}

@Override
int compareCurrent(int slot) {
if (missingBucket) {
int result = missingOrder.compare(() -> missingCurrentValue, () -> bits.get(slot) == false, reverseMul);
if (MissingOrder.unknownOrder(result) == false) {
return result;
}
}
return compareValues(currentValue, values.get(slot));
}

@Override
int compareCurrentWithAfter() {
if (missingBucket) {
int result = missingOrder.compare(() -> missingCurrentValue, () -> Objects.isNull(afterValue), reverseMul);
if (MissingOrder.unknownOrder(result) == false) {
return result;
}
}
return compareValues(currentValue, afterValue.longValue());
}

@Override
int hashCode(int slot) {
if (missingBucket && bits.get(slot) == false) {
return 0;
} else {
return Long.hashCode(values.get(slot));
}
}

@Override
int hashCodeCurrent() {
if (missingCurrentValue) {
return 0;
} else {
return Long.hashCode(currentValue);
}
}

@Override
protected void setAfter(Comparable value) {
if (missingBucket && value == null) {
afterValue = null;
} else {
// parse the value from a string in case it is a date or a formatted unsigned long.
afterValue = format.parseUnsignedLong(value.toString(), false, () -> {
throw new IllegalArgumentException("now() is not supported in [after] key");
});
}
}

@Override
int compareValues(long v1, long v2) {
BigInteger toComparable(int slot) {
if (missingBucket && bits.get(slot) == false) {
return null;
}
return Numbers.toUnsignedBigInteger(values.get(slot));
}

@Override
LeafBucketCollector getLeafCollector(LeafReaderContext context, LeafBucketCollector next) throws IOException {
final SortedNumericDocValues dvs = docValuesFunc.apply(context);
return new LeafBucketCollector() {
@Override
public void collect(int doc, long bucket) throws IOException {
if (dvs.advanceExact(doc)) {
int num = dvs.docValueCount();
for (int i = 0; i < num; i++) {
currentValue = dvs.nextValue();
missingCurrentValue = false;
next.collect(doc, bucket);
}
} else if (missingBucket) {
missingCurrentValue = true;
next.collect(doc, bucket);
}
}
};
}

@Override
LeafBucketCollector getLeafCollector(Comparable value, LeafReaderContext context, LeafBucketCollector next) {
// We still accept long here (not BigInteger)
if (value.getClass() != Long.class) {
throw new IllegalArgumentException("Expected Long, got " + value.getClass());
}
currentValue = (Long) value;
return new LeafBucketCollector() {
@Override
public void collect(int doc, long bucket) throws IOException {
next.collect(doc, bucket);
}
};
}

@Override
public void close() {
Releasables.close(values, bits);
}

private int compareValues(long v1, long v2) {
return Long.compareUnsigned(v1, v2) * reverseMul;
}

@Override
SortedDocsProducer createSortedDocsProducerOrNull(IndexReader reader, Query query) {
query = extractQuery(query);
if (checkIfSortedDocsIsApplicable(reader, fieldType) == false || checkMatchAllOrRangeQuery(query, fieldType.name()) == false) {
query = LongValuesSource.extractQuery(query);
if (checkIfSortedDocsIsApplicable(reader, fieldType) == false
|| LongValuesSource.checkMatchAllOrRangeQuery(query, fieldType.name()) == false) {
return null;
}
final byte[] lowerPoint;
Expand Down

0 comments on commit 7511ece

Please sign in to comment.