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

src/tdigest.c

@@ -48,55 +48,161 @@
     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);
+// Optional key-comparison counter, used only by the complexity regression test.
+// Zero-cost unless TD_INSTRUMENT_SORT is defined at build time.
+#ifdef TD_INSTRUMENT_SORT
+unsigned long long td_sort_comparisons = 0;
+#define TD_SORT_CMP() (++td_sort_comparisons)
+#else
+#define TD_SORT_CMP() ((void)0)
+#endif
+
+#define TD_INSORT_THRESHOLD 16
+
+// Insertion sort of the inclusive range [lo, hi] (used for small ranges).
+static void td_insertion_sort(double *means, long long *weights, int lo, int hi) {
+    for (int i = lo + 1; i <= hi; i++) {
+        const double m = means[i];
+        const long long w = weights[i];
+        int j = i - 1;
+        while (j >= lo && (TD_SORT_CMP(), means[j] > m)) {
+            means[j + 1] = means[j];
+            weights[j + 1] = weights[j];
+            j--;
+        }
+        means[j + 1] = m;
+        weights[j + 1] = w;
+    }
+}
+
+// Max-heap sift-down over the range starting at `lo` (heap size `n`, root heap
+// index `i`), keyed on means with weights moved in lock-step.
+static void td_sift_down(double *means, long long *weights, int lo, int i, int n) {
+    for (;;) {
+        int child = 2 * i + 1;
+        if (child >= n) {
+            break;
+        }
+        if (child + 1 < n && (TD_SORT_CMP(), means[lo + child] < means[lo + child + 1])) {
+            child++;
         }
+        if (!(TD_SORT_CMP(), means[lo + i] < means[lo + child])) {
+            break;
+        }
+        swap(means, lo + i, lo + child);
+        swap_l(weights, lo + i, lo + child);
+        i = child;
+    }
+}
+
+// Heapsort of the inclusive range [lo, hi]: the O(n log n) worst-case fallback.
+static void td_heap_sort(double *means, long long *weights, int lo, int hi) {
+    const int n = hi - lo + 1;
+    for (int i = n / 2 - 1; i >= 0; i--) {
+        td_sift_down(means, weights, lo, i, n);
+    }
+    for (int end = n - 1; end > 0; end--) {
+        swap(means, lo, lo + end);
+        swap_l(weights, lo, lo + end);
+        td_sift_down(means, weights, lo, 0, end);
     }
-    swap(means, i + 1, end);
-    swap_l(weights, i + 1, end);
-    return i + 1;
+}
+
+// Move the median of means[lo], means[mid], means[hi] into `mid` (weights follow)
+// and return it, so an already-sorted / organ-pipe input does not pick a bad pivot.
+static double td_median3(double *means, long long *weights, int lo, int mid, int hi) {
+    if ((TD_SORT_CMP(), means[mid] < means[lo])) {
+        swap(means, lo, mid);
+        swap_l(weights, lo, mid);
+    }
+    if ((TD_SORT_CMP(), means[hi] < means[lo])) {
+        swap(means, lo, hi);
+        swap_l(weights, lo, hi);
+    }
+    if ((TD_SORT_CMP(), means[hi] < means[mid])) {
+        swap(means, mid, hi);
+        swap_l(weights, mid, hi);
+    }
+    return means[mid];
 }
 
 /**
- * Standard quick sort except that sorting rearranges parallel arrays
+ * Introsort over two parallel arrays keyed on `means` (weights permuted in
+ * lock-step). Guaranteed O(n log n) time and O(log n) stack:
+ *   - 3-way (Dutch-flag) partitioning collapses runs of equal keys in a single
+ *     pass -- t-digest's common duplicate-heavy input;
+ *   - a median-of-three pivot avoids the trivial already-sorted worst case;
+ *   - a recursion-depth limit falls back to heapsort, so an adversarial
+ *     distinct-value permutation crafted to defeat the pivot (which a
+ *     fixed-position pivot cannot resist) cannot force quadratic time;
+ *   - small ranges finish with insertion sort;
+ *   - recursing only the smaller side bounds the stack to O(log n).
  *
- * @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
+ * Note: this is NOT output-identical to a naive single-pivot sort for runs of
+ * equal keys -- the order within an equal-key run differs, which can change how
+ * weighted duplicate centroids are subsequently merged. Results stay within
+ * t-digest's accuracy guarantee; exact centroid-for-centroid reproduction of the
+ * old sorter is not preserved.
  */
-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_introsort(double *means, long long *weights, int lo, int hi, int depth_limit) {
+    while (hi - lo > TD_INSORT_THRESHOLD) {
+        if (depth_limit == 0) {
+            td_heap_sort(means, weights, lo, hi);
             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);
+        depth_limit--;
+        const int mid = lo + (hi - lo) / 2;
+        const double pivot = td_median3(means, weights, lo, mid, hi);
+        // While scanning: [lo, lt) < pivot, [lt, i) == pivot, (gt, hi] > pivot.
+        int lt = lo, i = lo, gt = hi;
+        while (i <= gt) {
+            const double v = means[i];
+            if ((TD_SORT_CMP(), v < pivot)) {
+                swap(means, i, lt);
+                swap_l(weights, i, lt);
+                lt++;
+                i++;
+            } else if ((TD_SORT_CMP(), v > pivot)) {
+                swap(means, i, gt);
+                swap_l(weights, i, gt);
+                gt--;
+            } else {
+                i++;
+            }
         }
-        td_qsort(means, weights, new_pivot_idx + 1, end);
+        const int left_size = lt - lo;  // count of elements < pivot
+        const int right_size = hi - gt; // count of elements > pivot
+        // Recurse the smaller side, loop on the larger (bounds the stack).
+        if (left_size < right_size) {
+            if (left_size > 1) {
+                td_introsort(means, weights, lo, lt - 1, depth_limit);
+            }
+            lo = gt + 1;
+        } else {
+            if (right_size > 1) {
+                td_introsort(means, weights, gt + 1, hi, depth_limit);
+            }
+            hi = lt - 1;
+        }
+    }
+    td_insertion_sort(means, weights, lo, hi);
+}
+
+static void td_qsort(double *means, long long *weights, unsigned int lo_u, unsigned int hi_u) {
+    // Indices fit an int: the node arrays are sized by `cap` (an int field), so
+    // the sorted range is always within [0, node_count).
+    const int lo = (int)lo_u;
+    const int hi = (int)hi_u;
+    if (lo >= hi) {
+        return;
+    }
+    // Depth limit = 2*floor(log2(n)); exceeding it hands the range to heapsort,
+    // which is what guarantees the O(n log n) worst case.
+    int depth_limit = 0;
+    for (int t = hi - lo + 1; t > 1; t >>= 1) {
+        depth_limit += 2;
     }
+    td_introsort(means, weights, lo, hi, depth_limit);
 }
 
 static inline size_t cap_from_compression(double compression) {

tests/CMakeLists.txt

@@ -16,9 +16,20 @@ endif()
 if (BUILD_TESTS)
     set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -std=c99")
     add_executable(td_test unit/td_test.c unit/minunit.h)
+    # minunit.h uses clock_gettime(), hidden by the top-level -std=c99 unless
+    # _POSIX_C_SOURCE is defined on the command line before any header.
+    target_compile_definitions(td_test PRIVATE _POSIX_C_SOURCE=200809L)
     target_link_libraries(td_test tdigest m)
     enable_testing()
     add_test(td_test td_test)
+
+    # Complexity regression: compiles the library with TD_INSTRUMENT_SORT to count
+    # key comparisons and asserts one large compress stays O(n log n) (MOD-17228).
+    add_executable(td_sort_complexity_test unit/td_sort_complexity_test.c)
+    target_include_directories(td_sort_complexity_test PRIVATE ${CMAKE_SOURCE_DIR}/src)
+    target_compile_definitions(td_sort_complexity_test PRIVATE TD_INSTRUMENT_SORT)
+    target_link_libraries(td_sort_complexity_test m)
+    add_test(td_sort_complexity_test td_sort_complexity_test)
 endif()
 
 

tests/unit/td_sort_complexity_test.c

@@ -0,0 +1,103 @@
+/*
+ * Complexity regression for the centroid sort (MOD-17228).
+ *
+ * A correctness-only test cannot distinguish an O(n log n) sort from an O(n^2)
+ * one, and the ordinary unit test auto-compresses in small batches so it never
+ * exercises one large partition. This test compiles the library with
+ * TD_INSTRUMENT_SORT (which exposes a key-comparison counter), forces a SINGLE
+ * compress over a large adversarial permutation, and asserts the comparison
+ * count stays within an O(n log n) bound.
+ *
+ * The input is a DISTINCT-value "midpoint killer": values are assigned so that the
+ * midpoint of every recursion range is the minimum of that range. A deterministic
+ * midpoint-pivot sort therefore peels one element per level -> ~n^2/2 comparisons
+ * (this is exactly the distinct-permutation weakness of a fixed-position pivot).
+ * Introsort's median-of-three pivot plus heapsort fallback keeps it O(n log n).
+ */
+#define TD_INSTRUMENT_SORT
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "tdigest.c" /* brings in the static sort + td_sort_comparisons */
+
+/* Assign strictly increasing values to midpoints in recursion order, so the
+ * midpoint of any subrange is that subrange's minimum -> worst case for a
+ * midpoint-pivot quicksort, with all values distinct. */
+static void fill_midpoint_killer(double *a, int lo, int hi, int *next) {
+    if (lo > hi) {
+        return;
+    }
+    const int mid = lo + (hi - lo) / 2;
+    a[mid] = (double)(*next)++;
+    fill_midpoint_killer(a, lo, mid - 1, next);
+    fill_midpoint_killer(a, mid + 1, hi, next);
+}
+
+int main(void) {
+    const int n = 100000;
+
+    double *killer = (double *)malloc((size_t)n * sizeof(double));
+    if (killer == NULL) {
+        fprintf(stderr, "allocation failed\n");
+        return 1;
+    }
+    int next = 0;
+    fill_midpoint_killer(killer, 0, n - 1, &next);
+
+    /* Compression large enough that cap = 6*compression+10 > n, so all n points
+     * buffer and are sorted by a single td_compress() -- one big partition. */
+    td_histogram_t *h = td_new(n);
+    if (h == NULL) {
+        fprintf(stderr, "allocation failed\n");
+        free(killer);
+        return 1;
+    }
+    for (int i = 0; i < n; ++i) {
+        if (td_add(h, killer[i], 1) != 0) {
+            fprintf(stderr, "td_add failed at %d\n", i);
+            free(killer);
+            td_free(h);
+            return 1;
+        }
+    }
+    free(killer);
+    if (h->unmerged_nodes != n) {
+        fprintf(stderr, "expected one big compress, but auto-compress ran (unmerged=%d)\n",
+                h->unmerged_nodes);
+        td_free(h);
+        return 1;
+    }
+
+    td_sort_comparisons = 0;
+    if (td_compress(h) != 0) {
+        fprintf(stderr, "compress failed\n");
+        td_free(h);
+        return 1;
+    }
+
+    const double bound = 50.0 * (double)n * log2((double)n); /* generous O(n log n) */
+    const double quadratic = 0.5 * (double)n * (double)n;    /* what the old sort would need */
+    printf("n=%d  one-compress key-comparisons=%llu  O(n log n) bound=%.0f  (quadratic ~%.0f)\n", n,
+           td_sort_comparisons, bound, quadratic);
+
+    int rc = 0;
+    if ((double)td_sort_comparisons > bound) {
+        fprintf(stderr,
+                "FAIL: comparison count exceeds the O(n log n) bound -- sort is quadratic\n");
+        rc = 1;
+    }
+    /* sanity: the digest is actually sorted/usable after the compress */
+    for (int i = 1; i < h->merged_nodes; ++i) {
+        if (h->nodes_mean[i - 1] > h->nodes_mean[i]) {
+            fprintf(stderr, "FAIL: centroids not sorted at %d\n", i);
+            rc = 1;
+            break;
+        }
+    }
+    td_free(h);
+    if (rc == 0) {
+        printf("OK\n");
+    }
+    return rc;
+}