diff --git a/ddprof-lib/src/main/cpp/reservoirSampler.h b/ddprof-lib/src/main/cpp/reservoirSampler.h index 8ffb6ad3f1..d087c05a13 100644 --- a/ddprof-lib/src/main/cpp/reservoirSampler.h +++ b/ddprof-lib/src/main/cpp/reservoirSampler.h @@ -1,5 +1,5 @@ /* - * Copyright 2024 Datadog + * Copyright 2024, 2026 Datadog * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -18,7 +18,6 @@ #define RESERVOIR_SAMPLER_H #include "xorshift.h" -#include #include #include @@ -69,16 +68,35 @@ class ReservoirSampler { for (int i = 0; i < _size && i < (int)input.size(); i++) { _reservoir.push_back(input[i]); } + // Algorithm L (Li, 1994). `target` is the 0-based index of the next + // input item to place, so the first candidate is input[_size] plus a + // skip, and every later one is one past the item just placed plus a + // skip. Without that +1 a zero skip places the same item again, + // duplicating it in the sample. + const int n = (int)input.size(); double weight = exp(log(nextUniform()) / _size); - int target = _size + (int) (log(nextUniform()) / log(1 - weight)); - assert(target >= 0); - while (target < (int)input.size()) { + int target = advance(_size, n, weight); + while (target < n) { _reservoir[xorshift::boundedIndex(xorshift::next(_rng), (u32)_size)] = input[target]; weight *= exp(log(nextUniform()) / _size); - target += (int) (log(nextUniform()) / log(1 - weight)); + target = advance(target + 1, n, weight); } return _reservoir; } + +private: + // Returns `from` plus a geometric skip, or `n` when the skip runs past the + // input. The skip is compared as a double before the int conversion: as + // `weight` shrinks it can exceed INT_MAX, and once `1 - weight` rounds to + // 1.0 the quotient is -inf (log(u) < 0 over log(1.0) == +0.0). Either way + // the stream is exhausted, and converting such a value would be undefined. + int advance(int from, int n, double weight) { + double skip = floor(log(nextUniform()) / log(1 - weight)); + if (!(skip >= 0.0 && skip < (double)(n - from))) { + return n; + } + return from + (int)skip; + } }; #endif //RESERVOIR_SAMPLER_H \ No newline at end of file diff --git a/ddprof-lib/src/test/cpp/xorshift_ut.cpp b/ddprof-lib/src/test/cpp/xorshift_ut.cpp index e1aa598650..e6222defd1 100644 --- a/ddprof-lib/src/test/cpp/xorshift_ut.cpp +++ b/ddprof-lib/src/test/cpp/xorshift_ut.cpp @@ -407,3 +407,55 @@ TEST(ReservoirSamplerTest, ReuseAcrossCallsDoesNotAccumulate) { ASSERT_EQ(4u, out.size()) << "call " << call; } } + +// Algorithm L steps one past the item it just placed before adding the skip. +// Without that step a zero skip places the same input item a second time, so a +// sample of distinct threads can name one thread twice and miss another. +TEST(ReservoirSamplerTest, SampleOfDistinctInputHasNoDuplicates) { + const int cases[][2] = {{4, 1}, {10, 3}, {17, 16}, {100, 16}, {2048, 16}, {16, 16}}; + for (const auto& c : cases) { + const int n = c[0]; + const int k = c[1]; + std::vector input; + for (int i = 0; i < n; i++) { + input.push_back(i); + } + for (int run = 0; run < 2000; run++) { + ReservoirSampler sampler(k, (u64)run); + std::vector& out = sampler.sample(input); + ASSERT_EQ((size_t)k, out.size()) << "n=" << n << " k=" << k; + std::set distinct(out.begin(), out.end()); + ASSERT_EQ((size_t)k, distinct.size()) + << "duplicate selection for n=" << n << " k=" << k << " run=" << run; + } + } +} + +// Each input item must land in the sample with probability k/n. The bound is +// +/-5% of the expectation, which is over ten standard deviations wide for +// these draw counts; the missing step put item k at roughly 1.25x (n=4, k=1) +// and 1.8x (n=10, k=3) its fair share. +TEST(ReservoirSamplerTest, InclusionFrequencyIsUniform) { + const int cases[][2] = {{4, 1}, {10, 3}, {40, 16}}; + const int draws = 100000; + for (const auto& c : cases) { + const int n = c[0]; + const int k = c[1]; + std::vector input; + for (int i = 0; i < n; i++) { + input.push_back(i); + } + std::vector hits(n, 0); + for (int run = 0; run < draws; run++) { + ReservoirSampler sampler(k, (u64)run); + for (int value : sampler.sample(input)) { + hits[value]++; + } + } + const double expected = (double)draws * k / n; + for (int i = 0; i < n; i++) { + EXPECT_NEAR((double)hits[i], expected, expected * 0.05) + << "item " << i << " for n=" << n << " k=" << k; + } + } +}