From 0ceecfbf3b875b2a4cb4705ad032baa7ea648d5e Mon Sep 17 00:00:00 2001
From: Matthew Kavanagh <mattk@brandwatch.com>
Date: Thu, 6 Aug 2020 19:30:35 +0100
Subject: [PATCH] Fix BitmapFrequencyCounter serialization

---
 .../search/facet/BitmapFrequencyCounter.java  | 10 ++--
 .../facet/BitmapFrequencyCounterTest.java     | 57 +++++++++++++++++++
 2 files changed, 63 insertions(+), 4 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
index e1d245145d08..3f24ff3ae71b 100644
--- a/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencyCounter.java
+++ b/solr/core/src/java/org/apache/solr/search/facet/BitmapFrequencyCounter.java
@@ -1,6 +1,8 @@
 package org.apache.solr.search.facet;
 
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import org.apache.solr.common.util.SimpleOrderedMap;
@@ -84,7 +86,7 @@ public void add(int value) {
   public SimpleOrderedMap<Object> serialize() {
     SimpleOrderedMap<Object> serialized = new SimpleOrderedMap<>();
 
-    byte[][] serializedBitmaps = new byte[bitmaps.length][];
+    List<byte[]> serializedBitmaps = new ArrayList<>(bitmaps.length);
 
     int i = 0;
     while (i < bitmaps.length) {
@@ -94,7 +96,7 @@ public SimpleOrderedMap<Object> serialize() {
       }
 
       bitmap.runOptimize();
-      serializedBitmaps[i] = BitmapUtil.bitmapToBytes(bitmap);
+      serializedBitmaps.add(BitmapUtil.bitmapToBytes(bitmap));
 
       i++;
     }
@@ -119,10 +121,10 @@ public SimpleOrderedMap<Object> serialize() {
    * @param serialized The serialized data
    */
   public void deserialize(SimpleOrderedMap<Object> serialized) {
-    byte[][] serializedBitmaps = (byte[][]) serialized.get("bitmaps");
+    List<byte[]> serializedBitmaps = (List<byte[]>) serialized.get("bitmaps");
     if (serializedBitmaps != null) {
       for (int i = 0; i < bitmaps.length; i++) {
-        bitmaps[i] = BitmapUtil.bytesToBitmap(serializedBitmaps[i]);
+        bitmaps[i] = BitmapUtil.bytesToBitmap(serializedBitmaps.get(i));
       }
     }
 
diff --git a/solr/core/src/test/org/apache/solr/search/facet/BitmapFrequencyCounterTest.java b/solr/core/src/test/org/apache/solr/search/facet/BitmapFrequencyCounterTest.java
index a55bb6152b84..1196d7f7e60d 100644
--- a/solr/core/src/test/org/apache/solr/search/facet/BitmapFrequencyCounterTest.java
+++ b/solr/core/src/test/org/apache/solr/search/facet/BitmapFrequencyCounterTest.java
@@ -1,6 +1,13 @@
 package org.apache.solr.search.facet;
 
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.InputStream;
+
 import org.apache.lucene.util.LuceneTestCase;
+import org.apache.solr.common.util.JavaBinCodec;
+import org.apache.solr.common.util.SimpleOrderedMap;
 import org.junit.Test;
 
 public class BitmapFrequencyCounterTest extends LuceneTestCase {
@@ -293,4 +300,54 @@ public void givenSize4_whenMergingValues() {
     assertEquals(decoded[14], 0);
     assertEquals(decoded[15], 1);
   }
+
+  @Test
+  public void testSerialization() throws IOException {
+    BitmapFrequencyCounter x = new BitmapFrequencyCounter(2);
+
+    x.add(101);
+
+    x.add(102);
+    x.add(102);
+    x.add(202);
+    x.add(202);
+
+    x.add(103);
+    x.add(103);
+    x.add(103);
+    x.add(203);
+    x.add(203);
+    x.add(203);
+    x.add(303);
+    x.add(303);
+    x.add(303);
+
+    JavaBinCodec codec = new JavaBinCodec();
+
+    ByteArrayOutputStream out = new ByteArrayOutputStream();
+    codec.marshal(x.serialize(), out);
+
+    InputStream in = new ByteArrayInputStream(out.toByteArray());
+    BitmapFrequencyCounter y = new BitmapFrequencyCounter(2);
+    y.deserialize((SimpleOrderedMap<Object>) codec.unmarshal(in));
+
+    assertEquals(y.getBitmaps().length, 2);
+
+    assertTrue(y.getBitmaps()[0].contains(101));
+    assertFalse(y.getBitmaps()[1].contains(101));
+
+    assertFalse(y.getBitmaps()[0].contains(102));
+    assertTrue(y.getBitmaps()[1].contains(102));
+    assertFalse(y.getBitmaps()[0].contains(202));
+    assertTrue(y.getBitmaps()[1].contains(202));
+
+    assertTrue(y.getBitmaps()[0].contains(103));
+    assertTrue(y.getBitmaps()[1].contains(103));
+    assertTrue(y.getBitmaps()[0].contains(203));
+    assertTrue(y.getBitmaps()[1].contains(203));
+    assertTrue(y.getBitmaps()[0].contains(303));
+    assertTrue(y.getBitmaps()[1].contains(303));
+
+    assertTrue(y.getOverflow().isEmpty());
+  }
 }