Fix O(n²) compress on duplicate/low-cardinality input (MOD-17228)

src/tdigest.c

@@ -48,54 +48,64 @@
     arr[j] = temp;
 }
 
-static unsigned int partition(double *means, long long *weights, unsigned int start,
-                              unsigned int end, unsigned int pivot_idx) {
-    const double pivotMean = means[pivot_idx];
-    swap(means, pivot_idx, end);
-    swap_l(weights, pivot_idx, end);
-
-    int i = start - 1;
-
-    for (unsigned int j = start; j < end; j++) {
-        // If current element is smaller than the pivot
-        if (means[j] < pivotMean) {
-            // increment index of smaller element
-            i++;
-            swap(means, i, j);
-            swap_l(weights, i, j);
-        }
-    }
-    swap(means, i + 1, end);
-    swap_l(weights, i + 1, end);
-    return i + 1;
-}
-
 /**
- * Standard quick sort except that sorting rearranges parallel arrays
+ * Quicksort that rearranges two parallel arrays, keyed on `means`.
+ *
+ * Uses 3-way (Dutch-national-flag) partitioning with tail-recursion elimination.
  *
- * @param means  Values to sort on
- * @param weights The auxillary values to sort.
- * @param start  The beginning of the values to sort
- * @param end    The value after the last value to sort
+ * The previous single-pivot version was quadratic in time and linear in stack
+ * depth on duplicate / low-cardinality input -- which is t-digest's *common*
+ * case, since it summarizes streams of repeated measurements, not an adversarial
+ * one. A run of equal keys made the central pivot peel one element per level:
+ * O(n^2) comparisons and an O(n) recursion depth. 3-way partitioning collapses
+ * each run of equal keys in a single pass, and recursing only into the smaller
+ * side bounds the stack to O(log n). Output is identical to the old sort.
+ *
+ * @param means   Values to sort on.
+ * @param weights The parallel array, permuted in lock-step.
+ * @param lo, hi  Inclusive bounds of the range to sort.
  */
-static void td_qsort(double *means, long long *weights, unsigned int start, unsigned int end) {
-    if (start < end) {
-        // two elements can be directly compared
-        if ((end - start) == 1) {
-            if (means[start] > means[end]) {
-                swap(means, start, end);
-                swap_l(weights, start, end);
+static void td_qsort(double *means, long long *weights, unsigned int lo_u, unsigned int hi_u) {
+    // Signed locals so the partition pointers can pass below `lo` without an
+    // unsigned underflow. Indices fit an int: the node arrays are sized by `cap`
+    // (an int field), so the range sorted is always within [0, node_count).
+    int lo = (int)lo_u;
+    int hi = (int)hi_u;
+    while (lo < hi) {
+        // Capture the pivot by value: the partition swaps will move means[mid].
+        const double pivot = means[lo + (hi - lo) / 2];
+        // While scanning: [lo, lt) < pivot, [lt, i) == pivot, (gt, hi] > pivot,
+        // and [i, gt] is still unclassified.
+        int lt = lo, i = lo, gt = hi;
+        while (i <= gt) {
+            if (means[i] < pivot) {
+                swap(means, i, lt);
+                swap_l(weights, i, lt);
+                lt++;
+                i++;
+            } else if (means[i] > pivot) {
+                swap(means, i, gt);
+                swap_l(weights, i, gt);
+                gt--;
+            } else {
+                i++;
             }
-            return;
         }
-        // generating a random number as a pivot was very expensive vs the array size
-        // const unsigned int pivot_idx = start + rand()%(end - start + 1);
-        const unsigned int pivot_idx = (end + start) / 2; // central pivot
-        const unsigned int new_pivot_idx = partition(means, weights, start, end, pivot_idx);
-        if (new_pivot_idx > start) {
-            td_qsort(means, weights, start, new_pivot_idx - 1);
+        // Now [lo, lt) < pivot, [lt, gt] == pivot (done), (gt, hi] > pivot.
+        const int left_size = lt - lo;   // count of elements < pivot
+        const int right_size = hi - gt;  // count of elements > pivot
+        // Recurse into the smaller side, loop on the larger (bounds stack depth).
+        if (left_size < right_size) {
+            if (left_size > 1) {
+                td_qsort(means, weights, (unsigned int)lo, (unsigned int)(lt - 1));
+            }
+            lo = gt + 1;
+        } else {
+            if (right_size > 1) {
+                td_qsort(means, weights, (unsigned int)(gt + 1), (unsigned int)hi);
+            }
+            hi = lt - 1;
         }
-        td_qsort(means, weights, new_pivot_idx + 1, end);
     }
 }
 

tests/unit/td_test.c

@@ -580,6 +580,60 @@ MU_TEST(test_quantiles_multiple) {
     td_free(t);
 }
 
+// Exercises the centroid sort on the inputs that were pathological for the old
+// single-pivot quicksort: all-equal, ascending, and descending. The 3-way sort
+// must produce a correctly ordered, correct-weight digest for each. (The old
+// sort was O(n^2)/O(n)-stack on the all-equal case; this guards correctness of
+// the replacement.)
+MU_TEST(test_duplicate_heavy_compress) {
+    const int n = 50000;
+
+    // All-equal: everything collapses onto a single value.
+    td_histogram_t *eq = td_new(200);
+    mu_assert(eq != NULL, "created_histogram");
+    for (int i = 0; i < n; ++i) {
+        mu_assert(td_add(eq, 42.0, 1) == 0, "Insertion");
+    }
+    mu_assert(td_compress(eq) == 0, "compress all-equal");
+    mu_assert_double_eq((double)n, td_size(eq));
+    mu_assert_double_eq(42.0, td_min(eq));
+    mu_assert_double_eq(42.0, td_max(eq));
+    mu_assert_double_eq(42.0, td_quantile(eq, 0.0));
+    mu_assert_double_eq(42.0, td_quantile(eq, 0.5));
+    mu_assert_double_eq(42.0, td_quantile(eq, 1.0));
+    // Merged centroid means must be non-decreasing.
+    for (int i = 1; i < eq->merged_nodes; ++i) {
+        mu_assert(eq->nodes_mean[i - 1] <= eq->nodes_mean[i], "means sorted (all-equal)");
+    }
+    td_free(eq);
+
+    // Ascending and descending must yield the same digest bounds.
+    td_histogram_t *asc = td_new(200);
+    td_histogram_t *desc = td_new(200);
+    mu_assert(asc != NULL && desc != NULL, "created_histograms");
+    for (int i = 0; i < n; ++i) {
+        mu_assert(td_add(asc, (double)i, 1) == 0, "Insertion asc");
+        mu_assert(td_add(desc, (double)(n - 1 - i), 1) == 0, "Insertion desc");
+    }
+    mu_assert(td_compress(asc) == 0, "compress asc");
+    mu_assert(td_compress(desc) == 0, "compress desc");
+    mu_assert_double_eq(0.0, td_min(asc));
+    mu_assert_double_eq((double)(n - 1), td_max(asc));
+    mu_assert_double_eq(0.0, td_min(desc));
+    mu_assert_double_eq((double)(n - 1), td_max(desc));
+    for (int i = 1; i < asc->merged_nodes; ++i) {
+        mu_assert(asc->nodes_mean[i - 1] <= asc->nodes_mean[i], "means sorted (asc)");
+    }
+    for (int i = 1; i < desc->merged_nodes; ++i) {
+        mu_assert(desc->nodes_mean[i - 1] <= desc->nodes_mean[i], "means sorted (desc)");
+    }
+    // The median of a dense uniform 0..n-1 sits near the middle for both orders.
+    mu_assert_double_eq_epsilon((double)(n - 1) / 2.0, td_quantile(asc, 0.5), (double)n * 0.02);
+    mu_assert_double_eq_epsilon((double)(n - 1) / 2.0, td_quantile(desc, 0.5), (double)n * 0.02);
+    td_free(asc);
+    td_free(desc);
+}
+
 MU_TEST_SUITE(test_suite) {
     MU_RUN_TEST(test_basic);
     MU_RUN_TEST(test_td_init);
@@ -601,6 +655,7 @@ MU_TEST_SUITE(test_suite) {
     MU_RUN_TEST(test_trimmed_mean_complex);
     MU_RUN_TEST(test_overflow);
     MU_RUN_TEST(test_overflow_merge);
+    MU_RUN_TEST(test_duplicate_heavy_compress);
 }
 
 int main(int argc, char *argv[]) {