Sort t-digest centroids by value, then by weight (#739)

src/main/java/org/apache/datasketches/tdigest/Sort.java

@@ -32,11 +32,15 @@ public final class Sort {
   public Sort() { }
 
   /**
-   * Stable sort two arrays.
-   * The first array is sorted while the second array is permuted the same way.
+   * Sort two arrays so that the result is deterministic.
+   * The elements are ordered by key (ascending) and, if the keys are equal, by the paired value
+   * (ascending). The first array is sorted while the second array is permuted the same way.
+   *
+   * <p>Because the pairs (key, value) are totally ordered, the output depends only on the
+   * multiset of pairs and not on the initial order of the elements nor on the choice of pivots.
    *
    * @param keys array to be sorted
-   * @param values array to be permuted the same way
+   * @param values array to be permuted the same way, also used to break ties between equal keys
    * @param n number of elements to sort from the beginning of the arrays
    */
   public static void stableSort(final double[] keys, final long[] values, final int n) {
@@ -50,64 +54,36 @@ private static void stableLimitedQuickSort(final double[] keys, final long[] val
     while (end - start > limit) {
 
       final int pivotIndex = start + ThreadLocalRandom.current().nextInt(end - start);
-      final double pivotValue = keys[pivotIndex];
-
-      // move pivot to beginning of array
-      swap(keys, start, pivotIndex);
-      swap(values, start, pivotIndex);
+      final double pivotKey = keys[pivotIndex];
+      final long pivotValue = values[pivotIndex];
 
-      // use a three way partition because many duplicate values is an important case
-      int low = start + 1;   // low points to first value not known to be equal to pivotValue
-      int high = end;        // high points to first value > pivotValue
-      int i = low;           // i scans the array
+      // use a three way partition because many duplicate pairs is an important case
+      // invariant: pairs in [start, low) are less than the pivot
+      // invariant: pairs in [low, i) are equal to the pivot
+      // invariant: pairs in [high, end) are greater than the pivot
+      // pairs in [i, high) are not yet classified
+      int low = start;
+      int high = end;
+      int i = start;
       while (i < high) {
-        // invariant: values[k] == pivotValue for k in [0..low)
-        // invariant: values[k] < pivotValue for k in [low..i)
-        // invariant: values[k] > pivotValue for k in [high..end)
-        // in-loop:  i < high
-        // in-loop:  low < high
-        // in-loop:  i >= low
-        final double vi = keys[i];
-        if (vi == pivotValue && i == pivotIndex) {
-          if (low != i) {
-            swap(keys, low, i);
-            swap(values, low, i);
-          } else {
-            i++;
-          }
+        final int c = compare(keys[i], values[i], pivotKey, pivotValue);
+        if (c < 0) {
+          swap(keys, low, i);
+          swap(values, low, i);
           low++;
-        } else if (vi > pivotValue || (vi == pivotValue && i > pivotIndex)) {
+          i++;
+        } else if (c > 0) {
           high--;
           swap(keys, i, high);
           swap(values, i, high);
         } else {
           i++;
         }
       }
-      // assert i == high || low == high therefore, we are done with partition
-      // at this point, i == high, [start, low) == pivot,
-      // [low, high) < pivot and [high, end) > pivot
-      // we have to move the values equal to the pivot into the middle.
-      // To do this, we swap pivot values into the top end of the [low, high) range
-      // stopping when we run out of destinations or when we run out of values to copy
-      int from = start;
-      int to = high - 1;
-      for (i = 0; from < low && to >= low; i++) {
-        swap(keys, from, to);
-        swap(values, from++, to--);
-      }
-      if (from == low) {
-        // ran out of things to copy. This means that the last destination is the boundary
-        low = to + 1;
-      } else {
-        // ran out of places to copy to. This means that there are uncopied pivots and the
-        // boundary is at the beginning of those
-        low = from;
-      }
+      // at this point [low, high) holds pairs equal to the pivot (same key and same value),
+      // which are indistinguishable, so they don't need any further sorting.
 
       // now recurse, but arrange it to handle the longer limit by tail recursion
-      // we have to sort the pivot values because they may have different weights
-      // we can't do that, however until we know how much weight is in the left and right
       if (low - start < end - high) {
         // left side is smaller
         stableLimitedQuickSort(keys, values, start, low, limit);
@@ -132,7 +108,7 @@ private static void stableLimitedInsertionSort(final double[] keys, final long[]
       // values in [start, i) are ordered
       // scan backwards to find where to stick the current key
       for (int j = i; j >= m; j--) {
-        if (j == 0 || keys[j - 1] <= k) {
+        if (j == 0 || compare(keys[j - 1], values[j - 1], k, v) <= 0) {
           if (j < i) {
             System.arraycopy(keys, j, keys, j + 1, i - j);
             System.arraycopy(values, j, values, j + 1, i - j);
@@ -145,6 +121,16 @@ private static void stableLimitedInsertionSort(final double[] keys, final long[]
     }
   }
 
+  /**
+   * Compares two (key, value) pairs: first by key, then, if the keys are equal, by value.
+   */
+  private static int compare(final double key1, final long value1,
+      final double key2, final long value2) {
+    if (key1 < key2) { return -1; }
+    if (key1 > key2) { return 1; }
+    return Long.compare(value1, value2);
+  }
+
   private static void swap(final double[] values, final int i, final int j) {
     final double tmpValue = values[i];
     values[i] = values[j];

src/main/java/org/apache/datasketches/tdigest/TDigestDouble.java

@@ -42,6 +42,10 @@
  * and the following implementation:
  * https://github.com/tdunning/t-digest
  * This implementation is similar to MergingDigest in the above implementation
+ *
+ * <p>When centroids are merged, they are ordered by mean ascending and then, for equal means,
+ * by weight ascending. This tie-break rule makes the result of a merge independent of the
+ * order of the input and is part of the behavior of this sketch.</p>
  */
 public final class TDigestDouble {
 
@@ -684,13 +688,14 @@ private void merge(final double[] values, final long[] weights, final long weigh
     numCentroids_++;
     int current = 1;
     double weightSoFar = 0;
+    // the inputs don't change inside the loop, so compute the normalizer only once
+    final double normalizer = ScaleFunction.normalizer(k_ * 2, centroidsWeight_);
     while (current != num) {
       final double proposedWeight = centroidWeights_[numCentroids_ - 1] + weights[current];
       boolean addThis = false;
       if ((current != 1) && (current != (num - 1))) {
         final double q0 = weightSoFar / centroidsWeight_;
         final double q2 = (weightSoFar + proposedWeight) / centroidsWeight_;
-        final double normalizer = ScaleFunction.normalizer(k_ * 2, centroidsWeight_);
         addThis = proposedWeight <= (centroidsWeight_ * Math.min(ScaleFunction.max(q0, normalizer), ScaleFunction.max(q2, normalizer)));
       }
       if (addThis) { // merge into existing centroid
@@ -738,7 +743,9 @@ static double normalizer(final double compression, final double n) {
     }
 
     static double z(final double compression, final double n) {
-      return (4 * Math.log(n / compression)) + 24;
+      // StrictMath.log gives bit-identical results on every JVM and platform, while Math.log may
+      // differ by up to 1 ulp, which could change which centroids get combined in merge()
+      return (4 * StrictMath.log(n / compression)) + 24;
     }
   }
 

src/test/java/org/apache/datasketches/tdigest/SortTest.java

@@ -21,6 +21,9 @@
 
 import static org.testng.Assert.assertEquals;
 
+import java.util.Arrays;
+import java.util.Random;
+
 import org.apache.datasketches.common.Shuffle;
 import org.testng.annotations.Test;
 
@@ -62,4 +65,87 @@ public void large() {
     }
   }
 
+  @Test
+  public void equalKeysAreOrderedByWeight() {
+    // the example from the issue: (200,1) (100,3) (300,2) (100,1)
+    final double[] keys = {200, 100, 300, 100};
+    final long[] values = {1, 3, 2, 1};
+    Sort.stableSort(keys, values, keys.length);
+    assertEquals(keys, new double[] {100, 100, 200, 300});
+    assertEquals(values, new long[] {1, 3, 1, 2});
+  }
+
+  @Test
+  public void resultDoesNotDependOnInputOrder() {
+    // many duplicate keys with different weights, more elements than the insertion sort limit,
+    // so that both the quicksort and the insertion sort are exercised
+    final int n = 1000;
+    final Random rnd = new Random(1);
+    final double[] origKeys = new double[n];
+    final long[] origValues = new long[n];
+    for (int i = 0; i < n; i++) {
+      origKeys[i] = rnd.nextInt(20);
+      origValues[i] = 1 + rnd.nextInt(1000);
+    }
+    double[] expectedKeys = null;
+    long[] expectedValues = null;
+    for (int iteration = 0; iteration < 20; iteration++) {
+      final double[] keys = origKeys.clone();
+      final long[] values = origValues.clone();
+      // shuffle both arrays with the same permutation
+      for (int i = n - 1; i > 0; i--) {
+        final int j = rnd.nextInt(i + 1);
+        final double tmpKey = keys[i];
+        keys[i] = keys[j];
+        keys[j] = tmpKey;
+        final long tmpValue = values[i];
+        values[i] = values[j];
+        values[j] = tmpValue;
+      }
+      Sort.stableSort(keys, values, n);
+      if (expectedKeys == null) {
+        expectedKeys = keys;
+        expectedValues = values;
+      } else {
+        assertEquals(keys, expectedKeys);
+        assertEquals(values, expectedValues);
+      }
+    }
+  }
+
+  @Test
+  public void matchesReferenceSort() {
+    final Random rnd = new Random(2);
+    final int[] sizes = {0, 1, 2, 10, 63, 64, 65, 66, 200, 5000};
+    final int[] distinctKeys = {1, 3, 100, 100000};
+    for (final int n : sizes) {
+      for (final int d : distinctKeys) {
+        final double[] keys = new double[n];
+        final long[] values = new long[n];
+        final long[][] pairs = new long[n][];
+        for (int i = 0; i < n; i++) {
+          keys[i] = rnd.nextInt(d) - d / 2;
+          values[i] = 1 + rnd.nextInt(7);
+          pairs[i] = new long[] {(long) keys[i], values[i]};
+        }
+        // reference: order by key, then by weight
+        Arrays.sort(pairs, (x, y) -> x[0] != y[0] ? Long.compare(x[0], y[0]) : Long.compare(x[1], y[1]));
+        Sort.stableSort(keys, values, n);
+        for (int i = 0; i < n; i++) {
+          assertEquals((long) keys[i], pairs[i][0]);
+          assertEquals(values[i], pairs[i][1]);
+        }
+      }
+    }
+  }
+
+  @Test
+  public void sortsOnlyRequestedPrefix() {
+    final double[] keys = {5, 3, 4, 1, 99, 98};
+    final long[] values = {1, 1, 1, 1, 7, 7};
+    Sort.stableSort(keys, values, 4);
+    assertEquals(keys, new double[] {1, 3, 4, 5, 99, 98});
+    assertEquals(values, new long[] {1, 1, 1, 1, 7, 7});
+  }
+
 }