Add new aggregate for collecting distinct values of a field into a Ro…

…aringBitmap object.

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

new ImmutableRoaringBitmap(ByteBuffer.wrap(bitmapBytes))

solr/core/ivy.xml

@@ -30,6 +30,8 @@
   </configurations>
 
   <dependencies>
+    <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"/>

solr/core/src/java/org/apache/solr/search/ValueSourceParser.java

@@ -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.HLLAgg;
 import org.apache.solr.search.facet.MinMaxAgg;
@@ -1056,6 +1057,8 @@ public ValueSource parse(FunctionQParser fp) throws SyntaxError {
 
     addParser("agg_topdocs", new TopDocsAgg.Parser());
 
+    addParser("agg_bitmapcollector", new BitmapCollectorAgg.Parser());
+
     addParser("childfield", new ChildFieldValueSourceParser());
   }
 

solr/core/src/java/org/apache/solr/search/facet/BitmapCollectorAgg.java

@@ -0,0 +1,135 @@
+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.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);
+    }
+  }
+}