Support longs in BitArray (backport of #61867) (#61871)

We frequently use `long`s with `BitArray` in aggs and right now we have
to assert that the `long` fits in an `int`. This adds support for `long`
to `BitArray` so we don't need those assertions.
This commit is contained in:
Nik Everett 2020-09-02 17:24:31 -04:00 committed by GitHub
parent e0eafec897
commit c19f67ce30
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 36 additions and 36 deletions

View File

@ -190,7 +190,7 @@ public abstract class ParentJoinAggregator extends BucketsAggregator implements
private final BitArray ordsBits;
public DenseCollectionStrategy(long maxOrd, BigArrays bigArrays) {
ordsBits = new BitArray((int) maxOrd, context.bigArrays());
ordsBits = new BitArray(maxOrd, context.bigArrays());
}
@Override

View File

@ -36,7 +36,7 @@ public final class BitArray implements Releasable {
* Create the {@linkplain BitArray}.
* @param initialSize the initial size of underlying storage.
*/
public BitArray(int initialSize, BigArrays bigArrays) {
public BitArray(long initialSize, BigArrays bigArrays) {
this.bigArrays = bigArrays;
this.bits = bigArrays.newLongArray(initialSize, true);
}
@ -44,8 +44,8 @@ public final class BitArray implements Releasable {
/**
* Set the {@code index}th bit.
*/
public void set(int index) {
int wordNum = wordNum(index);
public void set(long index) {
long wordNum = wordNum(index);
bits = bigArrays.grow(bits, wordNum + 1);
bits.set(wordNum, bits.get(wordNum) | bitmask(index));
}
@ -53,8 +53,8 @@ public final class BitArray implements Releasable {
/**
* Clear the {@code index}th bit.
*/
public void clear(int index) {
int wordNum = wordNum(index);
public void clear(long index) {
long wordNum = wordNum(index);
if (wordNum >= bits.size()) {
/*
* No need to resize the array just to clear the bit because we'll
@ -68,8 +68,8 @@ public final class BitArray implements Releasable {
/**
* Is the {@code index}th bit set?
*/
public boolean get(int index) {
int wordNum = wordNum(index);
public boolean get(long index) {
long wordNum = wordNum(index);
if (wordNum >= bits.size()) {
/*
* If the word is bigger than the array then it could *never* have
@ -81,11 +81,11 @@ public final class BitArray implements Releasable {
return (bits.get(wordNum) & bitmask) != 0;
}
private static int wordNum(int index) {
private static long wordNum(long index) {
return index >> 6;
}
private static long bitmask(int index) {
private static long bitmask(long index) {
return 1L << index;
}

View File

@ -209,7 +209,7 @@ public abstract class BucketedSort implements Releasable {
* Is this bucket a min heap {@code true} or in gathering mode {@code false}?
*/
private boolean inHeapMode(long bucket) {
return heapMode.get((int) bucket);
return heapMode.get(bucket);
}
/**
@ -419,11 +419,7 @@ public abstract class BucketedSort implements Releasable {
setIndexToDocValue(index);
loader().loadFromDoc(index, doc);
if (next == 0) {
if (bucket > Integer.MAX_VALUE) {
throw new UnsupportedOperationException("Bucketed sort doesn't support more than [" + Integer.MAX_VALUE + "] buckets");
// BitArray needs int keys and this'd take a ton of memory to use that many buckets. So we just don't.
}
heapMode.set((int) bucket);
heapMode.set(bucket);
heapify(rootIndex);
} else {
setNextGatherOffset(rootIndex, next - 1);

View File

@ -31,6 +31,8 @@ import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
import static org.junit.Assume.assumeThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
@ -60,6 +62,18 @@ public class BitArrayTests extends ESTestCase {
}
}
public void testVeryLarge() {
assumeThat(Runtime.getRuntime().maxMemory(), greaterThanOrEqualTo(ByteSizeUnit.MB.toBytes(512)));
try (BitArray bitArray = new BitArray(1, BigArrays.NON_RECYCLING_INSTANCE)) {
long index = randomLongBetween(Integer.MAX_VALUE, (long) (Integer.MAX_VALUE * 1.5));
assertFalse(bitArray.get(index));
bitArray.set(index);
assertTrue(bitArray.get(index));
bitArray.clear(index);
assertFalse(bitArray.get(index));
}
}
public void testTooBigIsNotSet() {
try (BitArray bitArray = new BitArray(1, BigArrays.NON_RECYCLING_INSTANCE)) {
for (int i = 0; i < 1000; i++) {

View File

@ -441,36 +441,33 @@ class TopMetricsAggregator extends NumericMetricsAggregator.MultiValue {
}
void markMissing(long index) {
int i = asInt(index);
if (tracker == null) {
tracker = new BitArray(i, bigArrays);
tracker = new BitArray(index, bigArrays);
}
tracker.set(i);
tracker.set(index);
}
void markNotMissing(long index) {
if (tracker == null) {
return;
}
tracker.clear(asInt(index));
tracker.clear(index);
}
void swap(long lhs, long rhs) {
if (tracker == null) {
return;
}
int l = asInt(lhs);
int r = asInt(rhs);
boolean backup = tracker.get(l);
if (tracker.get(r)) {
tracker.set(l);
boolean backup = tracker.get(lhs);
if (tracker.get(rhs)) {
tracker.set(lhs);
} else {
tracker.clear(l);
tracker.clear(lhs);
}
if (backup) {
tracker.set(r);
tracker.set(rhs);
} else {
tracker.clear(r);
tracker.clear(rhs);
}
}
@ -478,7 +475,7 @@ class TopMetricsAggregator extends NumericMetricsAggregator.MultiValue {
if (tracker == null) {
return false;
}
return tracker.get(asInt(index));
return tracker.get(index);
}
@Override
@ -487,12 +484,5 @@ class TopMetricsAggregator extends NumericMetricsAggregator.MultiValue {
tracker.close();
}
}
private int asInt(long index) {
if (index > Integer.MAX_VALUE) {
throw new IllegalArgumentException("top_metrics can't track more than " + Integer.MAX_VALUE + " values.");
}
return (int) index;
}
}
}