Count-Min: serialization version 2 (cross-language format)

src/main/java/org/apache/datasketches/count/CountMinSketch.java

@@ -21,6 +21,8 @@
 
 import static java.lang.foreign.ValueLayout.JAVA_BYTE;
 import static java.lang.foreign.ValueLayout.JAVA_LONG_UNALIGNED;
+import static org.apache.datasketches.hash.XxHash64.hash;
+import static org.apache.datasketches.hash.XxHash64.hashBytes;
 
 import java.lang.foreign.MemorySegment;
 import java.nio.charset.StandardCharsets;
@@ -42,8 +44,19 @@
  * providing approximate counts with configurable error bounds.
  *
  * Reference: http://dimacs.rutgers.edu/~graham/pubs/papers/cm-full.pdf
+ *
+ * <p>New sketches are written in serialization version 2, the cross-language format defined in
+ * <a href="https://github.com/apache/datasketches-tck/blob/main/docs/count-min-format.md">
+ * the DataSketches TCK</a>. Version 2 hashes items with XxHash64, so its images give the same results
+ * in every language. A sketch read from a version 1 image keeps the version 1 hashing, which used
+ * MurmurHash3 and seeds from {@link Random}, and is written back as version 1.
+ * Version 1 and version 2 sketches can't be merged.</p>
  */
 public class CountMinSketch {
+  static final byte SERIAL_VERSION_1 = 1;
+  static final byte SERIAL_VERSION_2 = 2;
+
+  private final byte serialVersion_;
   private final byte numHashes_;
   private final int numBuckets_;
   private final long seed_;
@@ -64,17 +77,52 @@ int mask() {
   }
 
   /**
-   * Creates a CountMin sketch with given number of hash functions and buckets,
-   * and a user-specified seed.
+   * Creates an empty CountMin sketch with the given number of hash functions and buckets,
+   * and the default seed, {@link Util#DEFAULT_UPDATE_SEED}.
+   * The sketch uses serialization version 2.
+   *
+   * <p>{@link #suggestNumHashes(double)} and {@link #suggestNumBuckets(double)} give suitable values
+   * for a desired confidence and relative error.</p>
+   *
+   * @param numHashes The number of hash functions (rows), from 1 to 127
+   * @param numBuckets The number of buckets per hash function, at least 3.
+   * The product of numHashes and numBuckets must be less than 2<sup>30</sup>.
+   * @throws SketchesArgumentException if a parameter is out of range
+   */
+  public CountMinSketch(final int numHashes, final int numBuckets) {
+    this(numHashes, numBuckets, Util.DEFAULT_UPDATE_SEED);
+  }
+
+  /**
+   * Creates an empty CountMin sketch with the given number of hash functions and buckets,
+   * and the given seed. The sketch uses serialization version 2.
+   *
+   * <p>{@link #suggestNumHashes(double)} and {@link #suggestNumBuckets(double)} give suitable values
+   * for a desired confidence and relative error. Sketches can be merged only if they use the same
+   * seed, and an image must be read back with the seed it was created with.</p>
    *
-   * @param numHashes The number of hash functions to apply to items
+   * @param numHashes The number of hash functions (rows), from 1 to 127
+   * @param numBuckets The number of buckets per hash function, at least 3.
+   * The product of numHashes and numBuckets must be less than 2<sup>30</sup>.
+   * @param seed The base hash seed
+   * @throws SketchesArgumentException if a parameter is out of range
+   */
+  public CountMinSketch(final int numHashes, final int numBuckets, final long seed) {
+    this(numHashes, numBuckets, seed, SERIAL_VERSION_2);
+  }
+
+  /**
+   * Creates a CountMin sketch that hashes items as the given serialization version does.
+   *
+   * @param numHashes The number of hash functions to apply to items, from 1 to 127
    * @param numBuckets Array size for each of the hashing function
    * @param seed The base hash seed
+   * @param serialVersion 1 for the legacy hashing, 2 for the cross-language hashing
    */
-  CountMinSketch(final byte numHashes, final int numBuckets, final long seed) {
+  private CountMinSketch(final int numHashes, final int numBuckets, final long seed, final byte serialVersion) {
     // Validate numHashes
-    if (numHashes <= 0) {
-      throw new SketchesArgumentException("Number of hash functions must be positive, got: " + numHashes);
+    if ((numHashes < 1) || (numHashes > Byte.MAX_VALUE)) {
+      throw new SketchesArgumentException("Number of hash functions must be from 1 to 127, got: " + numHashes);
     }
 
     // Validate numBuckets with clear mathematical justification
@@ -102,16 +150,24 @@ int mask() {
           + "Consider reducing numHashes or numBuckets.");
     }
 
-    numHashes_ = numHashes;
+    serialVersion_ = serialVersion;
+    numHashes_ = (byte) numHashes; //stored as one byte in the image
     numBuckets_ = numBuckets;
     seed_ = seed;
     hashSeeds_ = new long[numHashes];
     sketchArray_ = new long[(int) totalSize];
     totalWeight_ = 0;
 
-    final Random rand = new Random(seed);
-    for (int i = 0; i < numHashes; i++) {
-      hashSeeds_[i] = rand.nextLong();
+    if (serialVersion == SERIAL_VERSION_1) {
+      final Random rand = new Random(seed);
+      for (int i = 0; i < numHashes; i++) {
+        hashSeeds_[i] = rand.nextLong();
+      }
+    } else {
+      // rowSeed[i] = XXH64(LE64(i), seed): XxHash64.hash(long, seed) hashes the 8 little-endian bytes of i
+      for (int i = 0; i < numHashes; i++) {
+        hashSeeds_[i] = hash(i, seed);
+      }
     }
   }
 
@@ -128,13 +184,42 @@ private long[] getHashes(final byte[] item) {
     final long[] updateLocations = new long[numHashes_];
 
     for (int i = 0; i < numHashes_; i++) {
-      final long[] index = MurmurHash3.hash(item, hashSeeds_[i]);
-      updateLocations[i] = (i * (long)numBuckets_) + Math.floorMod(index[0], numBuckets_);
+      updateLocations[i] = (i * (long) numBuckets_) + bucket(item, hashSeeds_[i]);
     }
 
     return updateLocations;
   }
 
+  private long[] getHashes(final long item) {
+    if (serialVersion_ == SERIAL_VERSION_1) {
+      return getHashes(longToBytes(item));
+    }
+    final long[] updateLocations = new long[numHashes_];
+    for (int i = 0; i < numHashes_; i++) {
+      // XxHash64.hash(long, seed) hashes the 8 little-endian bytes of the item, as version 2 requires
+      final long h = hash(item, hashSeeds_[i]);
+      updateLocations[i] = (i * (long) numBuckets_) + ((h >>> 1) % numBuckets_);
+    }
+    return updateLocations;
+  }
+
+  private long bucket(final byte[] item, final long rowSeed) {
+    if (serialVersion_ == SERIAL_VERSION_1) {
+      return Math.floorMod(MurmurHash3.hash(item, rowSeed)[0], numBuckets_);
+    }
+    final long h = hashBytes(item, 0, item.length, rowSeed);
+    return (h >>> 1) % numBuckets_;
+  }
+
+  /**
+   * Returns the serialization version this sketch hashes and serializes with:
+   * 2 for new sketches, or 1 for a sketch read from a version 1 image.
+   * @return the serialization version
+   */
+  byte getSerialVersion() {
+    return serialVersion_;
+  }
+
   /**
    * Checks if the CountMinSketch has processed any items.
    * @return True if the sketch is empty, otherwise false.
@@ -214,7 +299,7 @@ public static int suggestNumBuckets(final double relativeError) {
    * @param weight The weight of the item.
    */
   public void update(final long item, final long weight) {
-    update(longToBytes(item), weight);
+    updateLocations(getHashes(item), weight);
   }
 
   /**
@@ -236,12 +321,14 @@ public void update(final String item, final long weight) {
    * @param weight The weight of the item.
    */
   public void update(final byte[] item, final long weight) {
-    if (item.length == 0) {
+    if ((item == null) || (item.length == 0)) {
       return;
     }
+    updateLocations(getHashes(item), weight);
+  }
 
+  private void updateLocations(final long[] hashLocations, final long weight) {
     totalWeight_ += weight > 0 ? weight : -weight;
-    final long[] hashLocations = getHashes(item);
     for (final long h : hashLocations) {
       sketchArray_[(int) h] += weight;
     }
@@ -253,7 +340,7 @@ public void update(final byte[] item, final long weight) {
    * @return Estimated frequency.
    */
   public long getEstimate(final long item) {
-    return getEstimate(longToBytes(item));
+    return estimateAt(getHashes(item));
   }
 
   /**
@@ -276,11 +363,13 @@ public long getEstimate(final String item) {
    * @return Estimated frequency.
    */
   public long getEstimate(final byte[] item) {
-    if (item.length == 0) {
+    if ((item == null) || (item.length == 0)) {
       return 0;
     }
+    return estimateAt(getHashes(item));
+  }
 
-    final long[] hashLocations = getHashes(item);
+  private long estimateAt(final long[] hashLocations) {
     long res = sketchArray_[(int) hashLocations[0]];
     // Start from index 1 to avoid processing first element twice
     for (int i = 1; i < hashLocations.length; i++) {
@@ -296,7 +385,7 @@ public long getEstimate(final byte[] item) {
    * @return Upper bound of estimated frequency.
    */
   public long getUpperBound(final long item) {
-    return getUpperBound(longToBytes(item));
+    return getEstimate(item) + (long) (getRelativeError() * getTotalWeight_());
   }
 
   /**
@@ -319,7 +408,7 @@ public long getUpperBound(final String item) {
    * @return Upper bound of estimated frequency.
    */
   public long getUpperBound(final byte[] item) {
-    if (item.length == 0) {
+    if ((item == null) || (item.length == 0)) {
       return 0;
     }
 
@@ -332,7 +421,7 @@ public long getUpperBound(final byte[] item) {
    * @return Lower bound of estimated frequency.
    */
   public long getLowerBound(final long item) {
-    return getLowerBound(longToBytes(item));
+    return getEstimate(item);
   }
 
   /**
@@ -373,6 +462,10 @@ public void merge(final CountMinSketch other) {
     if (!acceptableConfig) {
       throw new SketchesException("Incompatible sketch configuration.");
     }
+    if (serialVersion_ != other.serialVersion_) {
+      throw new SketchesException("Cannot merge sketches with different serialization versions: "
+          + serialVersion_ + " and " + other.serialVersion_);
+    }
 
     for (int i = 0; i < sketchArray_.length; i++) {
       sketchArray_[i] += other.sketchArray_[i];
@@ -404,8 +497,7 @@ public byte[] toByteArray() {
     // Long 0
     final int preambleLongs = Family.COUNTMIN.getMinPreLongs();
     posSeg.setByte((byte) preambleLongs);
-    final int serialVersion = 1;
-    posSeg.setByte((byte) serialVersion);
+    posSeg.setByte(serialVersion_);
     final int familyId = Family.COUNTMIN.getID();
     posSeg.setByte((byte) familyId);
     final int flagsByte = isEmpty() ? Flag.IS_EMPTY.mask() : 0;
@@ -457,10 +549,9 @@ public static CountMinSketch deserialize(final byte[] b, final long seed) {
       throw new SketchesArgumentException("Preamble longs mismatch: expected " + expectedPreambleLongs
           + ", actual " + preambleLongs);
     }
-    final int expectedSerialVersion = 1;
-    if (serialVersion != expectedSerialVersion) {
-      throw new SketchesArgumentException("Serial version mismatch: expected " + expectedSerialVersion
-          + ", actual " + serialVersion);
+    if ((serialVersion != SERIAL_VERSION_1) && (serialVersion != SERIAL_VERSION_2)) {
+      throw new SketchesArgumentException("Serial version mismatch: expected " + SERIAL_VERSION_1
+          + " or " + SERIAL_VERSION_2 + ", actual " + serialVersion);
     }
     final int expectedFamilyId = Family.COUNTMIN.getID();
     if (familyId != expectedFamilyId) {
@@ -483,7 +574,7 @@ public static CountMinSketch deserialize(final byte[] b, final long seed) {
       final long arrayBytes = (long) numHashes * numBuckets * Long.BYTES;
       Util.checkBounds(preambleBytes, Long.BYTES + Math.max(arrayBytes, 0), b.length);
     }
-    final CountMinSketch cms = new CountMinSketch(numHashes, numBuckets, seed);
+    final CountMinSketch cms = new CountMinSketch(numHashes, numBuckets, seed, serialVersion);
     if (empty) {
       return cms;
     }