Reservoir sampling Algorithm A-ExpJ for C++17

Question

Implementation of the Reservoir sampling algorithm A-ExpJ that allows sampling K random elements from a stream of elements according to their weights when we don't know the size of the stream in advance and don't have the memory to store all the elements of the stream.

Description of the algorithm: https://en.wikipedia.org/wiki/Reservoir_sampling#Algorithm_A-ExpJ

I tried to make it as efficient time-wise as I could:

• one allocation for all the data
• not do anything if the object is not going to be added
• the ability to construct objects in-place
• not move objects while adding new ones

also tried to make it safe and support as many types as I could:

• should work with custom arithmetic types as weights (bigint?), but not tested
• should work with non-copyable and/or non-movable types

template<typename T, typename WeightType = float, typename URBG = std::mt19937, typename RandType = float>
class ReservoirSamplerWeighted
{
public:
    ReservoirSamplerWeighted(size_t samplesCount, URBG&& rand = std::mt19937{std::random_device{}()})
        : mSamplesCount(samplesCount)
        , mRand(std::forward<URBG>(rand))
    {
    }

    ~ReservoirSamplerWeighted()
    {
        if (mData)
        {
            for (size_t i = 0; i < mAllocatedElementsCount; ++i)
            {
                mElements[i].~T();
            }
            for (size_t i = 0; i < mSamplesCount; ++i)
            {
                mQueuePrios[i].~RandType();
            }
            std::free(mData);
        }
    }

    ReservoirSamplerWeighted(const ReservoirSamplerWeighted& other)
        : mSamplesCount(other.mSamplesCount)
        , mWeightJumpOver(other.mWeightJumpOver)
        , mRand(other.mRand)
        , mUniformDist(other.mUniformDist)
        , mAllocatedElementsCount(other.mAllocatedElementsCount)
    {
        if (other.mData)
        {
            allocateData();
            for (size_t i = 0; i < mSamplesCount; ++i)
            {
                new (mQueuePrios + i) RandType(other.mQueuePrios[i]);
            }

            std::memcpy(mQueueIndexes, other.mQueueIndexes, sizeof(size_t)*mAllocatedElementsCount);

            for (size_t i = 0; i < mAllocatedElementsCount; ++i)
            {
                new (mElements + i) T(other.mElements[i]);
            }
        }
    }

    ReservoirSamplerWeighted(ReservoirSamplerWeighted&& other)
        : mSamplesCount(other.mSamplesCount)
        , mWeightJumpOver(other.mWeightJumpOver)
        , mRand(other.mRand)
        , mUniformDist(other.mUniformDist)
        , mAllocatedElementsCount(other.mAllocatedElementsCount)
        , mData(other.mData)
        , mQueuePrios(other.mQueuePrios)
        , mQueueIndexes(other.mQueueIndexes)
        , mElements(other.mElements)
    {
        other.mWeightJumpOver = {};
        other.mAllocatedElementsCount = 0;
        other.mData = nullptr;
        other.mQueuePrios = nullptr;
        other.mQueueIndexes = nullptr;
        other.mElements = nullptr;
    }

    ReservoirSamplerWeighted& operator=(const ReservoirSamplerWeighted&) = delete;
    ReservoirSamplerWeighted& operator=(ReservoirSamplerWeighted&&) = delete;

    // if creation of an object to store is expensive you can check this before calling addElement
    // and provide a "dummy" object in case this returns false, because it will be ignored in that case
    bool willNextBeConsidered(WeightType weight) const
    {
        return (mWeightJumpOver - weight) <= 0;
    }

    template<typename E, typename = std::enable_if_t<std::is_move_constructible_v<std::decay_t<E>> && std::is_move_assignable_v<std::decay_t<E>> && std::is_same_v<std::decay_t<E>, T>>>
    void addElement(WeightType weight, E&& element)
    {
        emplaceElement(weight, std::move(element));
    }

    void addElement(WeightType weight, const T& element)
    {
        emplaceElement(weight, element);
    }

    template<typename... Args>
    void emplaceElement(WeightType weight, Args&&... arguments)
    {
        if (mData == nullptr)
        {
            prepareData();
        }

        if (weight > WeightType(0.0))
        {
            if (mAllocatedElementsCount < mSamplesCount)
            {
                const RandType r = std::pow(mUniformDist(mRand), (static_cast<RandType>(1.0) / weight));
                insertSorted(r, std::forward<Args>(arguments)...);
                if (mAllocatedElementsCount == mSamplesCount)
                {
                    mWeightJumpOver = log(mUniformDist(mRand)) / log(mQueuePrios[0]);
                }
            }
            else
            {
                mWeightJumpOver -= weight;
                if (mWeightJumpOver <= 0)
                {
                    const RandType t = std::pow(mQueuePrios[0], weight);
                    const RandType r = std::pow(std::uniform_real_distribution<RandType>(t, static_cast<RandType>(1.0))(mRand), static_cast<RandType>(1.0) / weight);

                    insertSortedRemoveFirst(r, std::forward<Args>(arguments)...);

                    mWeightJumpOver = log(mUniformDist(mRand)) / log(mQueuePrios[0]);
                }
            }
        }
    }

    const std::pair<const T*, size_t> getResult() const { return std::make_pair(mElements, mAllocatedElementsCount); }

private:
    void prepareData()
    {
        allocateData();
        for (size_t i = 0; i < mSamplesCount; ++i)
        {
            new (mQueuePrios + i) RandType();
        }
    }

    void allocateData()
    {
        assert(mData == nullptr);
        constexpr size_t alignment = std::max(std::max(std::alignment_of_v<RandType>, std::alignment_of_v<size_t>), std::alignment_of_v<T>);
        const size_t keysExtent = (sizeof(RandType)*mSamplesCount) % std::alignment_of_v<size_t>;
        const size_t indexesAlignmentGap = keysExtent > 0 ? (std::alignment_of_v<size_t> - keysExtent) : 0;
        const size_t indexesOffset = sizeof(RandType)*mSamplesCount + indexesAlignmentGap;
        const size_t indexesExtent = (indexesOffset + sizeof(size_t)*mSamplesCount) % std::alignment_of_v<T>;
        const size_t elementsAlignmentGap = indexesExtent > 0 ? (std::alignment_of_v<T> - indexesExtent) : 0;
        const size_t elementsOffset = indexesOffset + sizeof(size_t)*mSamplesCount + elementsAlignmentGap;
        const size_t alignedSize = elementsOffset + sizeof(T)*mSamplesCount;

        mData = std::aligned_alloc(alignment, alignedSize);
        mQueuePrios = reinterpret_cast<RandType*>(mData);
        mQueueIndexes = reinterpret_cast<size_t*>(static_cast<char*>(mData) + indexesOffset);
        mElements = reinterpret_cast<T*>(static_cast<char*>(mData) + elementsOffset);
    }

    template<typename... Args>
    void insertSorted(RandType r, Args&&... arguments)
    {
        RandType* it = std::upper_bound(mQueuePrios, mQueuePrios + mAllocatedElementsCount, r);
        const size_t firstMovedIdx = std::distance(mQueuePrios, it);

        std::move_backward(it, mQueuePrios + mAllocatedElementsCount, mQueuePrios + mAllocatedElementsCount + 1);
        std::move_backward(mQueueIndexes + firstMovedIdx, mQueueIndexes + mAllocatedElementsCount, mQueueIndexes + mAllocatedElementsCount + 1);

        mQueuePrios[firstMovedIdx] = r;
        mQueueIndexes[firstMovedIdx] = mAllocatedElementsCount;
        new (mElements + mAllocatedElementsCount) T(std::forward<Args>(arguments)...);
        ++mAllocatedElementsCount;
    }

    template<typename... Args>
    void insertSortedRemoveFirst(RandType r, Args&&... arguments)
    {
        RandType* it = std::upper_bound(mQueuePrios + 1, mQueuePrios + mSamplesCount, r);
        const size_t firstNotMovedIdx = std::distance(mQueuePrios, it);
        const size_t oldElementIdx = mQueueIndexes[0];

        std::move(mQueuePrios + 1, it, mQueuePrios);
        std::move(mQueueIndexes + 1, mQueueIndexes + firstNotMovedIdx, mQueueIndexes);

        mQueuePrios[firstNotMovedIdx - 1] = r;
        mQueueIndexes[firstNotMovedIdx - 1] = oldElementIdx;
        if constexpr (std::is_assignable_v<T, T> && sizeof...(Args) == 1 && std::is_same_v<std::decay_t<std::tuple_element_t<0, std::tuple<Args...>>>, T>)
        {
            mElements[oldElementIdx] = std::forward<Args...>(arguments...);
        }
        else if constexpr (std::is_move_assignable_v<T>)
        {
            mElements[oldElementIdx] = T(std::forward<Args>(arguments)...);
        }
        else
        {
            // support for non-moveable types
            mElements[oldElementIdx].~T();
            new (mElements + (oldElementIdx)) T(std::forward<Args>(arguments)...);
        }
    }

private:
    const size_t mSamplesCount;
    WeightType mWeightJumpOver {};
    URBG mRand;
    std::uniform_real_distribution<RandType> mUniformDist{static_cast<RandType>(0.0), static_cast<RandType>(1.0)};
    size_t mAllocatedElementsCount = 0;
    void* mData = nullptr;
    RandType* mQueuePrios = nullptr;
    size_t* mQueueIndexes = nullptr;
    T* mElements = nullptr;
};

Examples with test code: https://wandbox.org/permlink/gvRzEKxibM9sv5A1

Answer 1

Avoid manual memory management

A great deal of complexity comes from the fact that you allocate memory manually. I assume this was done to satisfy this requirement you had:

• should work with non-copyable and/or non-movable types

I cannot imagine a scenario where you would want to do reservoir sampling on non-copyable/movable types, but on the other hand it's always good to not make assumptions and to make your code work in as many situations as possible.

Instead of doing manual memory allocation, and then having to make your own destructors, copy and move constructors, I see three alternatives:

1. Use std::unique_ptr<T[]> to store the data. With C++20, you can use std::make_unique_for_overwrite() to allocate the memory, pre-C++20 you just have to manually new T[samplesCount] and store it in the std::unique_ptr. This uses default initialization, so for POD types it does what you want, but it might be less efficient for classes with default constructors. It also requires assignment operators to exist for T.

2. Use std::deque<T> to store the data. This is like a std::vector, but doesn't move elements when it is resized. You can emplace_back() new elements efficiently. One drawback is that it might still do multiple allocations, although since no moving of data in memory is involved it's much cheaper than std::vector.

3. Create your own class that acts like a fixed size std::vector but also allows non-copyable/movable types, and use that. This removes the complexity from class ReservoirSamplerWeighted.

There are lots of subtle issues that you need to think about when doing this yourself, like whether alignof(T) > sizeof(T) is possible, object lifetimes, and whether your pointers are valid.

But you do require T to be copyable

In your copy constructor, you call placement-new and pass a reference to other.mElements[i]; this invokes the copy constructor of T.

Make use of std::priority_queue

You already require that RandType is movable, and std::size_t certainly is, so I don't think you need to put the queue priorities and indexes in the same allocated memory as the data elements. And in that case, I strongly recommend you use std::priority_queue; it does exactly what you want, and it's more efficient since it is implemented as a heap instead of a sorted list: both inserting an element and popping the minimum element are \$O(\log N)\$ instead of \$O(N)\$.

To do this in a nice way, create a struct that holds an index and a priority, and has an operator<() overload that compares two objects based on their priority.

Returning the data

Your getResult() works correctly, but instead of returning a std::pair, returning a struct {T* data; std::size_t size;} would be nicer, as the members are explicitly named. If you can use C++20, then I would recommend returning a std::span instead. And of course if you use an STL container to store the data, then you can just return a reference to that container.

Something to think about is whether you want the caller to move the data out of your class. For example, they might want to pass it around while not wanting to keep the ReservoirSamplerWeighted object alive all the time. Maybe an extract() function would be nice, but of course that would only work if the data is held in a proper container object (but that can be any of the three alternatives I mentioned above).

Allowing custom arithmetic types

It would indeed be nice to support custom WeightTypes and RandTypes, but make sure your code handles it correctly. One issue is that you need "0" and "1" values. You write WeightType(0.0) and RandType(1.0) in your code. But would a bigint class have a constructor that takes a double argument? I think WeightType{} would be a safer way to construct a zero, but how to make a one? You could perhaps make a one_v template like C++20's mathematical constants, which can then be overloaded for T outside of the definition of your class ReservoirSamplerWeighted. But also consider that you use std::log() (except you forgot the std::) and std::pow() in your code, so it doesn't really doesn't make sense to use anything other than float or double.

You also might want to think about constraining the allowed types to ones that your class can work with. You can use SFINAE or concepts for that. You want to check for something like std::floating_point.