Java Code Examples for org.roaringbitmap.buffer.MutableRoaringBitmap#add()

private MutableRoaringBitmap getPinotDocIds(IndexSearcher indexSearcher, MutableRoaringBitmap luceneDocIds) {
  IntIterator luceneDocIDIterator = luceneDocIds.getIntIterator();
  MutableRoaringBitmap actualDocIDs = new MutableRoaringBitmap();
  try {
    while (luceneDocIDIterator.hasNext()) {
      int luceneDocId = luceneDocIDIterator.next();
      Document document = indexSearcher.doc(luceneDocId);
      int pinotDocId = Integer.valueOf(document.get(LuceneTextIndexCreator.LUCENE_INDEX_DOC_ID_COLUMN_NAME));
      actualDocIDs.add(pinotDocId);
    }
  } catch (Exception e) {
    LOGGER.error("Failure while retrieving document from index for column {}, exception {}", _column, e.getMessage());
    throw new RuntimeException(e);
  }
  return actualDocIDs;
}
 

@Test
public void testAndDocIdIterator() {
  // AND result: [2, 7, 13, 15, 16, 20]
  int[] docIds1 = new int[]{0, 1, 2, 3, 5, 7, 10, 12, 13, 15, 16, 18, 20};
  int[] docIds2 = new int[]{1, 2, 4, 5, 6, 7, 9, 11, 12, 13, 15, 16, 17, 19, 20};
  int[] docIds3 = new int[]{0, 2, 3, 4, 7, 8, 10, 11, 13, 15, 16, 19, 20};

  MutableRoaringBitmap bitmap1 = new MutableRoaringBitmap();
  bitmap1.add(docIds1);
  MutableRoaringBitmap bitmap2 = new MutableRoaringBitmap();
  bitmap2.add(docIds2);
  MutableRoaringBitmap bitmap3 = new MutableRoaringBitmap();
  bitmap3.add(docIds3);
  AndDocIdIterator andDocIdIterator = new AndDocIdIterator(
      new BlockDocIdIterator[]{new RangelessBitmapDocIdIterator(bitmap1), new RangelessBitmapDocIdIterator(
          bitmap2), new RangelessBitmapDocIdIterator(bitmap3)});

  assertEquals(andDocIdIterator.next(), 2);
  assertEquals(andDocIdIterator.next(), 7);
  assertEquals(andDocIdIterator.advance(10), 13);
  assertEquals(andDocIdIterator.advance(16), 16);
  assertEquals(andDocIdIterator.next(), 20);
  assertEquals(andDocIdIterator.next(), Constants.EOF);
}
 

@Test
public void testBitmapDocIdIterator() {
  int[] docIds = new int[]{1, 2, 4, 5, 6, 8, 12, 15, 16, 18, 20, 21};
  MutableRoaringBitmap bitmap = new MutableRoaringBitmap();
  bitmap.add(docIds);
  int numDocs = 25;
  BitmapDocIdIterator docIdIterator = new BitmapDocIdIterator(bitmap, numDocs);
  assertEquals(docIdIterator.advance(2), 2);
  assertEquals(docIdIterator.advance(3), 4);
  assertEquals(docIdIterator.next(), 5);
  assertEquals(docIdIterator.advance(6), 6);
  assertEquals(docIdIterator.next(), 8);
  assertEquals(docIdIterator.advance(13), 15);
  assertEquals(docIdIterator.advance(19), 20);
  assertEquals(docIdIterator.next(), 21);
  assertEquals(docIdIterator.next(), Constants.EOF);
}
 

@Setup
public void setup() {
  bitmap1 = new MutableRoaringBitmap();
  bitmap2 = new MutableRoaringBitmap();
  int k = 1 << 16;
  int i = 0;
  for (; i < 10000; ++i) {
    bitmap1.add(i * k);
  }
  for (; i < 10050; ++i) {
    bitmap2.add(i * k);
    bitmap1.add(i * k + 13);
  }
  for (; i < 20000; ++i) {
    bitmap2.add(i * k);
  }
  bitmap1.add(i * k);
}
 

@Setup
public void setup() {
  bitmap1 = new MutableRoaringBitmap();
  bitmap2 = new MutableRoaringBitmap();
  int k = 1 << 16;
  int i = 0;
  for (; i < 10000; ++i) {
    bitmap1.add(i * k);
  }
  for (; i < 10050; ++i) {
    bitmap2.add(i * k);
    bitmap1.add(i * k + 13);
  }
  for (; i < 20000; ++i) {
    bitmap2.add(i * k);
  }
  bitmap1.add(i * k);
}
 

@Override
public MutableRoaringBitmap applyAnd(ImmutableRoaringBitmap docIds) {
  MutableRoaringBitmap result = new MutableRoaringBitmap();
  IntIterator docIdIterator = docIds.getIntIterator();
  int nextDocId;
  while (docIdIterator.hasNext() && (nextDocId = docIdIterator.next()) < _numDocs) {
    int length = _valueSet.getIntValues(nextDocId, _dictIdBuffer);
    _numEntriesScanned += length;
    if (_predicateEvaluator.applyMV(_dictIdBuffer, length)) {
      result.add(nextDocId);
    }
  }
  return result;
}
 

@Override
public MutableRoaringBitmap applyAnd(ImmutableRoaringBitmap docIds) {
  MutableRoaringBitmap result = new MutableRoaringBitmap();
  IntIterator docIdIterator = docIds.getIntIterator();
  int nextDocId;
  while (docIdIterator.hasNext() && (nextDocId = docIdIterator.next()) < _numDocs) {
    _numEntriesScanned++;
    if (_valueMatcher.doesValueMatch(nextDocId)) {
      result.add(nextDocId);
    }
  }
  return result;
}
 

private void addToMetricIndex(final String key, final Integer id) {
  if (metricIndex.containsKey(key)) {
    MutableRoaringBitmap currentMap = lookupMetricIndex(key).toMutableRoaringBitmap();
    currentMap.add(id);
    metricIndex.put(key, RoaringBitMapUtils.toByteBuffer(currentMap));
  } else {
    MutableRoaringBitmap newMap = MutableRoaringBitmap.bitmapOf(id);
    metricIndex.put(key, RoaringBitMapUtils.toByteBuffer(newMap));
  }
}
 

@Benchmark
public MutableRoaringBitmap createBitmapUnordered_mutable() {

  MutableRoaringBitmap r = new MutableRoaringBitmap();

  for (int k = 65536 - 1; k >= 0; k--) {
    r.add(k * 32);
  }
  return r;
}
 

@Benchmark
public MutableRoaringBitmap createBitmapOrdered_mutable() {

  MutableRoaringBitmap r = new MutableRoaringBitmap();

  for (int k = 0; k < 65536; k++) {
    r.add(k * 32);
  }
  return r;
}
 

@Test
public void testOrDocIdIterator() {
  // OR result: [0, 1, 2, 4, 5, 6, 8, 10, 13, 15, 16, 17, 18, 19, 20]
  int[] docIds1 = new int[]{1, 4, 6, 10, 15, 17, 18, 20};
  int[] docIds2 = new int[]{0, 1, 5, 8, 15, 18};
  int[] docIds3 = new int[]{1, 2, 6, 13, 16, 19};

  MutableRoaringBitmap bitmap1 = new MutableRoaringBitmap();
  bitmap1.add(docIds1);
  MutableRoaringBitmap bitmap2 = new MutableRoaringBitmap();
  bitmap2.add(docIds2);
  MutableRoaringBitmap bitmap3 = new MutableRoaringBitmap();
  bitmap3.add(docIds3);
  OrDocIdIterator andDocIdIterator = new OrDocIdIterator(
      new BlockDocIdIterator[]{new RangelessBitmapDocIdIterator(bitmap1), new RangelessBitmapDocIdIterator(
          bitmap2), new RangelessBitmapDocIdIterator(bitmap3)});

  assertEquals(andDocIdIterator.advance(1), 1);
  assertEquals(andDocIdIterator.next(), 2);
  assertEquals(andDocIdIterator.next(), 4);
  assertEquals(andDocIdIterator.advance(7), 8);
  assertEquals(andDocIdIterator.advance(13), 13);
  assertEquals(andDocIdIterator.next(), 15);
  assertEquals(andDocIdIterator.advance(18), 18);
  assertEquals(andDocIdIterator.next(), 19);
  assertEquals(andDocIdIterator.advance(21), Constants.EOF);
}
 

@Setup
public void setup() {
  bitmap1 = new MutableRoaringBitmap();
  bitmap2 = new MutableRoaringBitmap();
  int k = 1 << 16;
  for (int i = 0; i < 10000; ++i) {
    bitmap1.add(2 * i * k);
    bitmap2.add(2 * i * k + 1);
  }
}
 

@Setup
public void setup() {
  bitmap1 = new MutableRoaringBitmap();
  bitmap2 = new MutableRoaringBitmap();
  int k = 1 << 16;
  for (int i = 0; i < 10000; ++i) {
    bitmap1.add(i * k);
    bitmap2.add(i * k);
  }
}
 

@Setup
public void setup() {
  bitmap1 = new MutableRoaringBitmap();
  bitmap2 = new MutableRoaringBitmap();
  int k = 1 << 16;
  for (int i = 0; i < 10000; ++i) {
    bitmap1.add(2 * i * k);
    bitmap2.add(2 * i * k + 1);
  }
}
 

@Setup
public void setup() {
  bitmap1 = new MutableRoaringBitmap();
  bitmap2 = new MutableRoaringBitmap();
  int k = 1 << 16;
  for (int i = 0; i < 10000; ++i) {
    bitmap1.add(i * k);
    bitmap2.add(i * k);
  }
}
 

@BeforeClass
public static void prepareRandomRanges() throws Exception
{
  System.setProperty("jub.customkey", String.format("%06.5f", DENSITY));
  reset();

  final BitSet expectedUnion = new BitSet();
  for (int i = 0; i < SIZE; ++i) {
    ConciseSet c = new ConciseSet();
    MutableRoaringBitmap r = new MutableRoaringBitmap();
    {
      int k = 0;
      boolean fill = true;
      while (k < LENGTH) {
        int runLength = (int) (LENGTH * DENSITY) + rand.nextInt((int) (LENGTH * DENSITY));
        for (int j = k; fill && j < LENGTH && j < k + runLength; ++j) {
          c.add(j);
          r.add(j);
          expectedUnion.set(j);
        }
        k += runLength;
        fill = !fill;
      }
    }
    minIntersection = MIN_INTERSECT;
    for (int k = LENGTH / 2; k < LENGTH / 2 + minIntersection; ++k) {
      c.add(k);
      r.add(k);
      expectedUnion.set(k);
    }
    concise[i] = ImmutableConciseSet.newImmutableFromMutable(c);
    offheapConcise[i] = makeOffheapConcise(concise[i]);
    roaring[i] = r;
    immutableRoaring[i] = makeImmutableRoaring(r);
    offheapRoaring[i] = makeOffheapRoaring(r);
    genericConcise[i] = new WrappedImmutableConciseBitmap(offheapConcise[i]);
    genericRoaring[i] = new WrappedImmutableRoaringBitmap(offheapRoaring[i]);
  }
  unionCount = expectedUnion.cardinality();
  printSizeStats(DENSITY, "Random Alternating Bitmap");
}
 

@Test
public void testMutableRunSerializationBasicDeserialization() throws java.io.IOException {
  final int[] data = takeSortedAndDistinct(new Random(07734), 100000);
  RoaringBitmap bitmap_a = RoaringBitmap.bitmapOf(data);
  RoaringBitmap bitmap_ar = RoaringBitmap.bitmapOf(data);

  MutableRoaringBitmap bitmap_am = MutableRoaringBitmap.bitmapOf(data);
  MutableRoaringBitmap bitmap_amr = MutableRoaringBitmap.bitmapOf(data);

  for (int k = 100000; k < 200000; ++k) {
    bitmap_a.add(3 * k); // bitmap density and too many little runs
    bitmap_ar.add(3 * k);
    bitmap_am.add(3 * k);
    bitmap_amr.add(3 * k);
  }

  for (int k = 700000; k < 800000; ++k) { // will choose a runcontainer on this
    bitmap_a.add(k); // bitmap density and too many little runs
    bitmap_ar.add(k);
    bitmap_am.add(k);
    bitmap_amr.add(k);    
  }

  bitmap_ar.runOptimize();
  bitmap_amr.runOptimize();
  assertEquals(bitmap_a, bitmap_ar);
  assertEquals(bitmap_am, bitmap_amr);
  assertEquals(bitmap_am.serializedSizeInBytes(), bitmap_a.serializedSizeInBytes());
  assertEquals(bitmap_amr.serializedSizeInBytes(), bitmap_ar.serializedSizeInBytes());

  ByteBuffer outbuf = ByteBuffer.allocate(2*(bitmap_a.serializedSizeInBytes() + bitmap_ar.serializedSizeInBytes()));
  DataOutputStream out = new DataOutputStream(new ByteBufferBackedOutputStream(outbuf));
  try {
    bitmap_a.serialize(out);
    bitmap_ar.serialize(out);
    bitmap_am.serialize(out);
    bitmap_amr.serialize(out);
  } catch (Exception e) {
    e.printStackTrace();
  }
  outbuf.flip();

  RoaringBitmap bitmap_c1 = new RoaringBitmap();
  RoaringBitmap bitmap_c2 = new RoaringBitmap();
  RoaringBitmap bitmap_c3 = new RoaringBitmap();
  RoaringBitmap bitmap_c4 = new RoaringBitmap();

  DataInputStream in = new DataInputStream(new ByteBufferBackedInputStream(outbuf));
  bitmap_c1.deserialize(in);
  bitmap_c2.deserialize(in);
  bitmap_c3.deserialize(in);
  bitmap_c4.deserialize(in);

  assertEquals(bitmap_a, bitmap_c1);
  assertEquals(bitmap_a, bitmap_c2);
  assertEquals(bitmap_a, bitmap_c3);
  assertEquals(bitmap_a, bitmap_c4);
  assertEquals(bitmap_ar, bitmap_c1);
  assertEquals(bitmap_ar, bitmap_c2);
  assertEquals(bitmap_ar, bitmap_c3);
  assertEquals(bitmap_ar, bitmap_c4);

}
 

public ConstantMVInvertedIndex(int numDocs) {
  _bitmap = new MutableRoaringBitmap();
  _bitmap.add(0, numDocs);
}
 

public ThreadSafeMutableRoaringBitmap(int firstDocId) {
  _mutableRoaringBitmap = new MutableRoaringBitmap();
  _mutableRoaringBitmap.add(firstDocId);
}
 

@Override
public BlockDocIdIterator iterator() {
  int numDocIdSets = _docIdSets.size();
  BlockDocIdIterator[] allDocIdIterators = new BlockDocIdIterator[numDocIdSets];
  List<SortedDocIdIterator> sortedDocIdIterators = new ArrayList<>();
  List<BitmapBasedDocIdIterator> bitmapBasedDocIdIterators = new ArrayList<>();
  List<BlockDocIdIterator> remainingDocIdIterators = new ArrayList<>();
  for (int i = 0; i < numDocIdSets; i++) {
    BlockDocIdIterator docIdIterator = _docIdSets.get(i).iterator();
    allDocIdIterators[i] = docIdIterator;
    if (docIdIterator instanceof SortedDocIdIterator) {
      sortedDocIdIterators.add((SortedDocIdIterator) docIdIterator);
    } else if (docIdIterator instanceof BitmapBasedDocIdIterator) {
      bitmapBasedDocIdIterators.add((BitmapBasedDocIdIterator) docIdIterator);
    } else {
      remainingDocIdIterators.add(docIdIterator);
    }
  }
  int numSortedDocIdIterators = sortedDocIdIterators.size();
  int numBitmapBasedDocIdIterators = bitmapBasedDocIdIterators.size();
  if (numSortedDocIdIterators + numBitmapBasedDocIdIterators > 1) {
    // When there are more than one index-base BlockDocIdIterator (SortedDocIdIterator or BitmapBasedDocIdIterator),
    // merge them and construct a BitmapDocIdIterator from the merged document ids. If there is no remaining
    // BlockDocIdIterator, directly return the merged BitmapDocIdIterator; otherwise, construct and return an
    // OrDocIdIterator with the merged BitmapDocIdIterator and the remaining BlockDocIdIterators.

    MutableRoaringBitmap docIds = new MutableRoaringBitmap();
    for (SortedDocIdIterator sortedDocIdIterator : sortedDocIdIterators) {
      for (Pairs.IntPair docIdRange : sortedDocIdIterator.getDocIdRanges()) {
        // NOTE: docIdRange has inclusive start and end.
        docIds.add(docIdRange.getLeft(), docIdRange.getRight() + 1L);
      }
    }
    for (BitmapBasedDocIdIterator bitmapBasedDocIdIterator : bitmapBasedDocIdIterators) {
      docIds.or(bitmapBasedDocIdIterator.getDocIds());
    }
    BitmapDocIdIterator bitmapDocIdIterator = new BitmapDocIdIterator(docIds, _numDocs);
    int numRemainingDocIdIterators = remainingDocIdIterators.size();
    if (numRemainingDocIdIterators == 0) {
      return bitmapDocIdIterator;
    } else {
      BlockDocIdIterator[] docIdIterators = new BlockDocIdIterator[numRemainingDocIdIterators + 1];
      docIdIterators[0] = bitmapDocIdIterator;
      for (int i = 0; i < numRemainingDocIdIterators; i++) {
        docIdIterators[i + 1] = remainingDocIdIterators.get(i);
      }
      return new OrDocIdIterator(docIdIterators);
    }
  } else {
    // Otherwise, construct and return an OrDocIdIterator with all BlockDocIdIterators.

    return new OrDocIdIterator(allDocIdIterators);
  }
}