Fix bitmap aggs for non-distributed mode

BrandwatchLtd · Nov 3, 2021 · 4249609 · 4249609
1 parent d3bac6b
commit 4249609
Show file tree

Hide file tree

Showing 8 changed files with 113 additions and 36 deletions.
diff --git a/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencyCounter.java b/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencyCounter.java
@@ -2,6 +2,7 @@
 
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -253,6 +254,22 @@ public void normalize() {
     });
   }
 
+  public Map<Integer, Integer> toFrequencyOfFrequencies() {
+    Map<Integer, Integer> map = new LinkedHashMap<>();
+
+    int[] lowFrequencies = decode();
+    for (int i = 0; i < lowFrequencies.length; i++) {
+      int value = lowFrequencies[i];
+      if (value > 0) {
+        map.put(i, value);
+      }
+    }
+
+    overflow.forEach((value, freq) -> map.merge(freq, 1, Integer::sum));
+
+    return map;
+  }
+
   public int[] decode() {
     int endIndex = 0;
     while (endIndex < bitmaps.length && bitmaps[endIndex] != null) {

diff --git a/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencyCounter64.java b/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencyCounter64.java
@@ -2,6 +2,7 @@
 
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -254,6 +255,22 @@ public void normalize() {
     });
   }
 
+  public Map<Integer, Integer> toFrequencyOfFrequencies() {
+    Map<Integer, Integer> map = new LinkedHashMap<>();
+
+    int[] lowFrequencies = decode();
+    for (int i = 0; i < lowFrequencies.length; i++) {
+      int value = lowFrequencies[i];
+      if (value > 0) {
+        map.put(i, value);
+      }
+    }
+
+    overflow.forEach((value, freq) -> map.merge(freq, 1, Integer::sum));
+
+    return map;
+  }
+
   public int[] decode() {
     int endIndex = 0;
     while (endIndex < bitmaps.length && bitmaps[endIndex] != null) {

diff --git a/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencyOfFrequenciesSlotAcc.java b/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencyOfFrequenciesSlotAcc.java
@@ -0,0 +1,21 @@
+package org.apache.solr.search.facet;
+
+import java.util.Collections;
+
+import org.apache.lucene.queries.function.ValueSource;
+
+public class BitmapFrequencyOfFrequenciesSlotAcc extends BitmapFrequencySlotAcc {
+  public BitmapFrequencyOfFrequenciesSlotAcc(ValueSource values, FacetContext fcontext, int numSlots, int maxFrequency) {
+    super(values, fcontext, numSlots, maxFrequency);
+  }
+
+  @Override
+  public Object getFinalValue(BitmapFrequencyCounter result) {
+    if (result != null) {
+      result.normalize();
+      return result.toFrequencyOfFrequencies();
+    } else {
+      return Collections.emptyMap();
+    }
+  }
+}
diff --git a/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencyOfFrequenciesSlotAcc64.java b/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencyOfFrequenciesSlotAcc64.java
@@ -0,0 +1,24 @@
+package org.apache.solr.search.facet;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.function.IntFunction;
+
+import org.apache.lucene.queries.function.ValueSource;
+
+public class BitmapFrequencyOfFrequenciesSlotAcc64 extends BitmapFrequencySlotAcc64 {
+  public BitmapFrequencyOfFrequenciesSlotAcc64(ValueSource values, FacetContext fcontext, int numSlots, int maxFrequency) {
+    super(values, fcontext, numSlots, maxFrequency);
+  }
+
+  @Override
+  public Object getFinalValue(BitmapFrequencyCounter64 result) {
+    if (result != null) {
+      result.normalize();
+      return result.toFrequencyOfFrequencies();
+    } else {
+      return Collections.emptyMap();
+    }
+  }
+}
diff --git a/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencySlotAcc.java b/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencySlotAcc.java
@@ -33,13 +33,26 @@ public int compare(int slotA, int slotB) {
 
   @Override
   public Object getValue(int slotNum) {
-    if (result[slotNum] != null) {
-      return result[slotNum].serialize();
+    if (fcontext.isShard()) {
+      return getShardValue(result[slotNum]);
+    } else {
+      return getFinalValue(result[slotNum]);
+    }
+  }
+
+  public Object getShardValue(BitmapFrequencyCounter result) {
+    if (result != null) {
+      return result.serialize();
     } else {
       return Collections.emptyList();
     }
   }
 
+  public Object getFinalValue(BitmapFrequencyCounter result) {
+    result.normalize();
+    return getShardValue(result);
+  }
+
   @Override
   public void reset() {
     Arrays.fill(result, null);

diff --git a/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencySlotAcc64.java b/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencySlotAcc64.java
@@ -33,13 +33,26 @@ public int compare(int slotA, int slotB) {
 
   @Override
   public Object getValue(int slotNum) {
-    if (result[slotNum] != null) {
-      return result[slotNum].serialize();
+    if (fcontext.isShard()) {
+      return getShardValue(result[slotNum]);
+    } else {
+      return getFinalValue(result[slotNum]);
+    }
+  }
+
+  public Object getShardValue(BitmapFrequencyCounter64 result) {
+    if (result != null) {
+      return result.serialize();
     } else {
       return Collections.emptyList();
     }
   }
 
+  public Object getFinalValue(BitmapFrequencyCounter64 result) {
+    result.normalize();
+    return getShardValue(result);
+  }
+
   @Override
   public void reset() {
     Arrays.fill(result, null);

diff --git a/solr/core/src/java/org/apache/solr/search/facet/FrequencyOfFrequenciesAgg.java b/solr/core/src/java/org/apache/solr/search/facet/FrequencyOfFrequenciesAgg.java
@@ -28,7 +28,7 @@ public FrequencyOfFrequenciesAgg(ValueSource vs, Integer size) {
 
   @Override
   public SlotAcc createSlotAcc(FacetContext fcontext, int numDocs, int numSlots) {
-    return new BitmapFrequencySlotAcc(getArg(), fcontext, numSlots, size);
+    return new BitmapFrequencyOfFrequenciesSlotAcc(getArg(), fcontext, numSlots, size);
   }
 
   @Override
@@ -76,22 +76,8 @@ public void finish(Context mcontext) {
 
     @Override
     public Object getMergedResult() {
-      Map<Integer, Integer> map = new LinkedHashMap<>();
-
       result.normalize();
-
-      int[] lowFrequencies = result.decode();
-      for (int i = 0; i < lowFrequencies.length; i++) {
-        int value = lowFrequencies[i];
-        if (value > 0) {
-          map.put(i, value);
-        }
-      }
-
-      result.getOverflow()
-        .forEach((value, freq) -> map.merge(freq, 1, Integer::sum));
-
-      return map;
+      return result.toFrequencyOfFrequencies();
     }
   }
 }
diff --git a/solr/core/src/java/org/apache/solr/search/facet/FrequencyOfFrequenciesAgg64.java b/solr/core/src/java/org/apache/solr/search/facet/FrequencyOfFrequenciesAgg64.java
@@ -28,7 +28,7 @@ public FrequencyOfFrequenciesAgg64(ValueSource vs, Integer size) {
 
   @Override
   public SlotAcc createSlotAcc(FacetContext fcontext, int numDocs, int numSlots) {
-    return new BitmapFrequencySlotAcc64(getArg(), fcontext, numSlots, size);
+    return new BitmapFrequencyOfFrequenciesSlotAcc64(getArg(), fcontext, numSlots, size);
   }
 
   @Override
@@ -76,22 +76,8 @@ public void finish(Context mcontext) {
 
     @Override
     public Object getMergedResult() {
-      Map<Integer, Integer> map = new LinkedHashMap<>();
-
       result.normalize();
-
-      int[] lowFrequencies = result.decode();
-      for (int i = 0; i < lowFrequencies.length; i++) {
-        int value = lowFrequencies[i];
-        if (value > 0) {
-          map.put(i, value);
-        }
-      }
-
-      result.getOverflow()
-        .forEach((value, freq) -> map.merge(freq, 1, Integer::sum));
-
-      return map;
+      return result.toFrequencyOfFrequencies();
     }
   }
 }