CountMinSketch serialization

Author: liancheng

common/sketch/src/main/java/org/apache/spark/util/sketch/CountMinSketch.java

@@ -17,6 +17,8 @@
 
 package org.apache.spark.util.sketch;
 
+import java.io.DataInputStream;
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 
@@ -54,6 +56,25 @@
  * This implementation is largely based on the {@code CountMinSketch} class from stream-lib.
  */
 abstract public class CountMinSketch {
+  /**
+   * Version number of the serialized binary format.
+   */
+  public enum Version {
+    V1(1);
+
+    private final int versionNumber;
+
+    Version(int versionNumber) {
+      this.versionNumber = versionNumber;
+    }
+
+    public int getVersionNumber() {
+      return versionNumber;
+    }
+  }
+
+  public abstract Version version();
+
   /**
    * Returns the relative error (or {@code eps}) of this {@link CountMinSketch}.
    */
@@ -99,19 +120,44 @@ abstract public class CountMinSketch {
    *
    * Note that only Count-Min sketches with the same {@code depth}, {@code width}, and random seed
    * can be merged.
+   *
+   * @exception CountMinSketchMergeException if the {@code other} {@link CountMinSketch} has
+   *            incompatible depth, width, relative-error, confidence, or random seed.
    */
-  public abstract CountMinSketch mergeInPlace(CountMinSketch other);
+  public abstract CountMinSketch mergeInPlace(CountMinSketch other)
+      throws CountMinSketchMergeException;
 
   /**
    * Writes out this {@link CountMinSketch} to an output stream in binary format.
    */
-  public abstract void writeTo(OutputStream out);
+  public abstract void writeTo(OutputStream out) throws IOException;
 
   /**
    * Reads in a {@link CountMinSketch} from an input stream.
    */
-  public static CountMinSketch readFrom(InputStream in) {
-    throw new UnsupportedOperationException("Not implemented yet");
+  public static CountMinSketch readFrom(InputStream in) throws IOException {
+    DataInputStream dis = new DataInputStream(in);
+
+    // Ignores version number
+    dis.readInt();
+
+    long totalCount = dis.readLong();
+    int depth = dis.readInt();
+    int width = dis.readInt();
+
+    long hashA[] = new long[depth];
+    for (int i = 0; i < depth; ++i) {
+      hashA[i] = dis.readLong();
+    }
+
+    long table[][] = new long[depth][width];
+    for (int i = 0; i < depth; ++i) {
+      for (int j = 0; j < width; ++j) {
+        table[i][j] = dis.readLong();
+      }
+    }
+
+    return new CountMinSketchImpl(depth, width, totalCount, hashA, table);
   }
 
   /**

common/sketch/src/main/java/org/apache/spark/util/sketch/CountMinSketchImpl.java

@@ -17,6 +17,8 @@
 
 package org.apache.spark.util.sketch;
 
+import java.io.DataOutputStream;
+import java.io.IOException;
 import java.io.OutputStream;
 import java.io.UnsupportedEncodingException;
 import java.util.Arrays;
@@ -51,6 +53,53 @@ public CountMinSketchImpl(double eps, double confidence, int seed) {
     initTablesWith(depth, width, seed);
   }
 
+  public CountMinSketchImpl(int depth, int width, long totalCount, long hashA[], long table[][]) {
+    this.depth = depth;
+    this.width = width;
+    this.eps = 2.0 / width;
+    this.confidence = 1 - 1 / Math.pow(2, depth);
+    this.hashA = hashA;
+    this.table = table;
+    this.totalCount = totalCount;
+  }
+
+  @Override
+  public boolean equals(Object other) {
+    if (other == this) {
+      return true;
+    }
+
+    if (!(other instanceof CountMinSketchImpl)) {
+      return false;
+    }
+
+    CountMinSketchImpl that = (CountMinSketchImpl) other;
+
+    if (this.depth == that.depth &&
+        this.width == that.width &&
+        this.totalCount == that.totalCount) {
+      for (int i = 0; i < depth; ++i) {
+        if (this.hashA[i] != that.hashA[i]) {
+          return false;
+        }
+
+        for (int j = 0; j < width; ++j) {
+          if (this.table[i][j] != that.table[i][j]) {
+            return false;
+          }
+        }
+      }
+      return true;
+    } else {
+      return false;
+    }
+  }
+
+  @Override
+  public Version version() {
+    return Version.V1;
+  }
+
   private void initTablesWith(int depth, int width, int seed) {
     this.table = new long[depth][width];
     this.hashA = new long[depth];
@@ -221,27 +270,27 @@ private long estimateCountForStringItem(String item) {
   }
 
   @Override
-  public CountMinSketch mergeInPlace(CountMinSketch other) {
+  public CountMinSketch mergeInPlace(CountMinSketch other) throws CountMinSketchMergeException {
     if (other == null) {
-      throw new CMSMergeException("Cannot merge null estimator");
+      throw new CountMinSketchMergeException("Cannot merge null estimator");
     }
 
     if (!(other instanceof CountMinSketchImpl)) {
-      throw new CMSMergeException("Cannot merge estimator of class " + other.getClass().getName());
+      throw new CountMinSketchMergeException("Cannot merge estimator of class " + other.getClass().getName());
     }
 
     CountMinSketchImpl that = (CountMinSketchImpl) other;
 
     if (this.depth != that.depth) {
-      throw new CMSMergeException("Cannot merge estimators of different depth");
+      throw new CountMinSketchMergeException("Cannot merge estimators of different depth");
     }
 
     if (this.width != that.width) {
-      throw new CMSMergeException("Cannot merge estimators of different width");
+      throw new CountMinSketchMergeException("Cannot merge estimators of different width");
     }
 
     if (!Arrays.equals(this.hashA, that.hashA)) {
-      throw new CMSMergeException("Cannot merge estimators of different seed");
+      throw new CountMinSketchMergeException("Cannot merge estimators of different seed");
     }
 
     for (int i = 0; i < this.table.length; ++i) {
@@ -256,13 +305,23 @@ public CountMinSketch mergeInPlace(CountMinSketch other) {
   }
 
   @Override
-  public void writeTo(OutputStream out) {
-    throw new UnsupportedOperationException("Not implemented yet");
-  }
+  public void writeTo(OutputStream out) throws IOException {
+    DataOutputStream dos = new DataOutputStream(out);
+
+    dos.writeInt(version().getVersionNumber());
+
+    dos.writeLong(this.totalCount);
+    dos.writeInt(this.depth);
+    dos.writeInt(this.width);
 
-  protected static class CMSMergeException extends RuntimeException {
-    public CMSMergeException(String message) {
-      super(message);
+    for (int i = 0; i < this.depth; ++i) {
+      dos.writeLong(this.hashA[i]);
+    }
+
+    for (int i = 0; i < this.depth; ++i) {
+      for (int j = 0; j < this.width; ++j) {
+        dos.writeLong(table[i][j]);
+      }
     }
   }
 }

common/sketch/src/main/java/org/apache/spark/util/sketch/CountMinSketchMergeException.java

@@ -0,0 +1,24 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.util.sketch;
+
+public class CountMinSketchMergeException extends Exception {
+  public CountMinSketchMergeException(String message) {
+    super(message);
+  }
+}

common/sketch/src/test/scala/org/apache/spark/util/sketch/CountMinSketchSuite.scala

@@ -17,6 +17,8 @@
 
 package org.apache.spark.util.sketch
 
+import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
+
 import scala.reflect.ClassTag
 import scala.util.Random
 
@@ -29,6 +31,16 @@ class CountMinSketchSuite extends FunSuite { // scalastyle:ignore funsuite
 
   private val seed = 42
 
+  private def checkSerDe(sketch: CountMinSketch): Unit = {
+    val out = new ByteArrayOutputStream()
+    sketch.writeTo(out)
+
+    val in = new ByteArrayInputStream(out.toByteArray)
+    val deserialized = CountMinSketch.readFrom(in)
+
+    assert(sketch === deserialized)
+  }
+
   def testAccuracy[T: ClassTag](typeName: String)(itemGenerator: Random => T): Unit = {
     test(s"accuracy - $typeName") {
       val r = new Random()
@@ -45,7 +57,10 @@ class CountMinSketchSuite extends FunSuite { // scalastyle:ignore funsuite
       }
 
       val sketch = CountMinSketch.create(epsOfTotalCount, confidence, seed)
+      checkSerDe(sketch)
+
       sampledItemIndices.foreach(i => sketch.add(allItems(i)))
+      checkSerDe(sketch)
 
       val probCorrect = {
         val numErrors = allItems.map { item =>
@@ -75,11 +90,16 @@ class CountMinSketchSuite extends FunSuite { // scalastyle:ignore funsuite
 
       val sketches = perSketchItems.map { items =>
         val sketch = CountMinSketch.create(epsOfTotalCount, confidence, seed)
+        checkSerDe(sketch)
+
         items.foreach(sketch.add)
+        checkSerDe(sketch)
+
         sketch
       }
 
       val mergedSketch = sketches.reduce(_ mergeInPlace _)
+      checkSerDe(mergedSketch)
 
       val expectedSketch = {
         val sketch = CountMinSketch.create(epsOfTotalCount, confidence, seed)
@@ -109,4 +129,22 @@ class CountMinSketchSuite extends FunSuite { // scalastyle:ignore funsuite
   testItemType[Long]("Long") { _.nextLong() }
 
   testItemType[String]("String") { r => r.nextString(r.nextInt(20)) }
+
+  test("incompatible merge") {
+    intercept[CountMinSketchMergeException] {
+      CountMinSketch.create(10, 10, 1).mergeInPlace(null)
+    }
+
+    intercept[CountMinSketchMergeException] {
+      val sketch1 = CountMinSketch.create(10, 20, 1)
+      val sketch2 = CountMinSketch.create(10, 20, 2)
+      sketch1.mergeInPlace(sketch2)
+    }
+
+    intercept[CountMinSketchMergeException] {
+      val sketch1 = CountMinSketch.create(10, 10, 1)
+      val sketch2 = CountMinSketch.create(10, 20, 2)
+      sketch1.mergeInPlace(sketch2)
+    }
+  }
 }