Skip to content

Commit

Permalink
Add new aggregate for collecting distinct values of a field into a Ro…
Browse files Browse the repository at this point in the history
…aringBitmap object.

The response is a byte array from which you can construct a bitmap, using

new ImmutableRoaringBitmap(ByteBuffer.wrap(bitmapBytes))

Fix for solr8
  • Loading branch information
timatbw authored and richardbadman committed Feb 11, 2021
1 parent ed4a196 commit 69ccd13
Show file tree
Hide file tree
Showing 3 changed files with 141 additions and 0 deletions.
2 changes: 2 additions & 0 deletions solr/core/ivy.xml
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@
<dependency org="io.opentracing" name="opentracing-util" rev="${/io.opentracing/opentracing-util}" conf="compile"/>
<dependency org="io.opentracing" name="opentracing-mock" rev="${/io.opentracing/opentracing-mock}" conf="test"/>

<dependency org="org.roaringbitmap" name="RoaringBitmap" rev="0.8.6" conf="compile"/>
<dependency org="org.roaringbitmap" name="shims" rev="0.8.6" conf="compile"/>
<dependency org="commons-codec" name="commons-codec" rev="${/commons-codec/commons-codec}" conf="compile"/>
<dependency org="commons-io" name="commons-io" rev="${/commons-io/commons-io}" conf="compile"/>
<dependency org="org.apache.commons" name="commons-exec" rev="${/org.apache.commons/commons-exec}" conf="compile"/>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@
import org.apache.solr.schema.TextField;
import org.apache.solr.search.facet.AggValueSource;
import org.apache.solr.search.facet.AvgAgg;
import org.apache.solr.search.facet.BitmapCollectorAgg;
import org.apache.solr.search.facet.CountAgg;
import org.apache.solr.search.facet.CountValsAgg;
import org.apache.solr.search.facet.HLLAgg;
Expand Down Expand Up @@ -1098,6 +1099,8 @@ public ValueSource parse(FunctionQParser fp) throws SyntaxError {

addParser("agg_termfreqfreq", new TermFrequencyOfFrequenciesAgg.Parser());

addParser("agg_bitmapcollector", new BitmapCollectorAgg.Parser());

addParser("childfield", new ChildFieldValueSourceParser());
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,136 @@
package org.apache.solr.search.facet;

import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.Arrays;
import java.util.function.IntFunction;

import org.apache.lucene.queries.function.ValueSource;
import org.apache.solr.common.util.SimpleOrderedMap;
import org.apache.solr.search.FunctionQParser;
import org.apache.solr.search.SyntaxError;
import org.apache.solr.search.ValueSourceParser;
import org.apache.solr.search.facet.SlotAcc.FuncSlotAcc;
import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
import org.roaringbitmap.buffer.MutableRoaringBitmap;


public class BitmapCollectorAgg extends SimpleAggValueSource {

private static final String KEY = "bitmap";

public static class Parser extends ValueSourceParser {
@Override
public ValueSource parse(FunctionQParser fp) throws SyntaxError {
return new BitmapCollectorAgg(fp.parseValueSource());
}
}

public BitmapCollectorAgg(ValueSource vs) {
super("bitmapcollector", vs);
}

@Override
public SlotAcc createSlotAcc(FacetContext fcontext, int numDocs, int numSlots) {
return new Acc(getArg(), fcontext, numSlots);
}

@Override
public FacetMerger createFacetMerger(Object prototype) {
return new Merger();
}

@Override
public String description() {
return "bitmapcollector";
}


private class Acc extends FuncSlotAcc {
MutableRoaringBitmap[] result;

Acc(ValueSource vs, FacetContext fcontext, int numSlots) {
super(vs, fcontext, numSlots);
this.result = new MutableRoaringBitmap[numSlots];
}

@Override
public void collect(int doc, int slot, IntFunction<SlotContext> slotContext) throws IOException {
if (result[slot] == null) {
result[slot] = new MutableRoaringBitmap();
}
result[slot].add(values.intVal(doc));
}

@Override
public int compare(int slotA, int slotB) {
return slotA - slotB;
}

@Override
public Object getValue(int slotNum) {
byte[] serialised;
if (result[slotNum] != null) {
result[slotNum].runOptimize();
serialised = bitmapToBytes(result[slotNum]);
} else {
serialised = new byte[0];
}
SimpleOrderedMap map = new SimpleOrderedMap();
map.add(KEY, serialised);
return map;
}

@Override
public void reset() {
Arrays.fill(result, null);
}

@Override
public void resize(Resizer resizer) {
result = resizer.resize(result, null);
}
}

public class Merger extends FacetMerger {

private MutableRoaringBitmap combined = new MutableRoaringBitmap();

@Override
public void merge(Object facetResult, Context mcontext) {
if (facetResult instanceof SimpleOrderedMap) {
byte[] bitmapBytes = (byte[])((SimpleOrderedMap)facetResult).get(KEY);
if (bitmapBytes.length != 0) {
combined.or(new ImmutableRoaringBitmap(ByteBuffer.wrap(bitmapBytes)));
}
}
}

@Override
public void finish(Context mcontext) {
// never called
}

@Override
public Object getMergedResult() {
combined.runOptimize();
SimpleOrderedMap map = new SimpleOrderedMap();
map.add(KEY, bitmapToBytes(combined));
return map;
}
}

private static byte[] bitmapToBytes(MutableRoaringBitmap bitmap) {
ByteArrayOutputStream bos = new ByteArrayOutputStream();
DataOutputStream dos = new DataOutputStream(bos);
try {
bitmap.serialize(dos);
dos.close();
return bos.toByteArray();
} catch (IOException ioe) {
throw new RuntimeException("Failed to serialise RoaringBitmap to bytes", ioe);
}
}
}

0 comments on commit 69ccd13

Please sign in to comment.