duckdb 1.1.4-dev2.0 → 1.1.4-dev9.0

package/src/duckdb/src/execution/sample/reservoir_sample.cpp

@@ -0,0 +1,930 @@
+#include "duckdb/execution/reservoir_sample.hpp"
+#include "duckdb/common/types/data_chunk.hpp"
+#include "duckdb/common/vector_operations/vector_operations.hpp"
+#include <unordered_set>
+
+namespace duckdb {
+
+std::pair<double, idx_t> BlockingSample::PopFromWeightQueue() {
+	D_ASSERT(base_reservoir_sample && !base_reservoir_sample->reservoir_weights.empty());
+	auto ret = base_reservoir_sample->reservoir_weights.top();
+	base_reservoir_sample->reservoir_weights.pop();
+
+	base_reservoir_sample->UpdateMinWeightThreshold();
+	D_ASSERT(base_reservoir_sample->min_weight_threshold > 0);
+	return ret;
+}
+
+double BlockingSample::GetMinWeightThreshold() {
+	return base_reservoir_sample->min_weight_threshold;
+}
+
+idx_t BlockingSample::GetPriorityQueueSize() {
+	return base_reservoir_sample->reservoir_weights.size();
+}
+
+void BlockingSample::Destroy() {
+	destroyed = true;
+}
+
+void ReservoirChunk::Serialize(Serializer &serializer) const {
+	chunk.Serialize(serializer);
+}
+
+unique_ptr<ReservoirChunk> ReservoirChunk::Deserialize(Deserializer &deserializer) {
+	auto result = make_uniq<ReservoirChunk>();
+	result->chunk.Deserialize(deserializer);
+	return result;
+}
+
+unique_ptr<ReservoirChunk> ReservoirChunk::Copy() const {
+	auto copy = make_uniq<ReservoirChunk>();
+	copy->chunk.Initialize(Allocator::DefaultAllocator(), chunk.GetTypes());
+
+	chunk.Copy(copy->chunk);
+	return copy;
+}
+
+ReservoirSample::ReservoirSample(idx_t sample_count, unique_ptr<ReservoirChunk> reservoir_chunk)
+    : ReservoirSample(Allocator::DefaultAllocator(), sample_count, 1) {
+	if (reservoir_chunk) {
+		this->reservoir_chunk = std::move(reservoir_chunk);
+		sel_size = this->reservoir_chunk->chunk.size();
+		sel = SelectionVector(0, sel_size);
+		ExpandSerializedSample();
+	}
+	stats_sample = true;
+}
+
+ReservoirSample::ReservoirSample(Allocator &allocator, idx_t sample_count, int64_t seed)
+    : BlockingSample(seed), sample_count(sample_count), allocator(allocator) {
+	base_reservoir_sample = make_uniq<BaseReservoirSampling>(seed);
+	type = SampleType::RESERVOIR_SAMPLE;
+	reservoir_chunk = nullptr;
+	stats_sample = false;
+	sel = SelectionVector(sample_count);
+	sel_size = 0;
+}
+
+idx_t ReservoirSample::GetSampleCount() {
+	return sample_count;
+}
+
+idx_t ReservoirSample::NumSamplesCollected() const {
+	if (!reservoir_chunk) {
+		return 0;
+	}
+	return reservoir_chunk->chunk.size();
+}
+
+SamplingState ReservoirSample::GetSamplingState() const {
+	if (base_reservoir_sample->reservoir_weights.empty()) {
+		return SamplingState::RANDOM;
+	}
+	return SamplingState::RESERVOIR;
+}
+
+idx_t ReservoirSample::GetActiveSampleCount() const {
+	switch (GetSamplingState()) {
+	case SamplingState::RANDOM:
+		return sel_size;
+	case SamplingState::RESERVOIR:
+		return base_reservoir_sample->reservoir_weights.size();
+	default:
+		throw InternalException("Sampling State is INVALID");
+	}
+}
+
+idx_t ReservoirSample::GetTuplesSeen() const {
+	return base_reservoir_sample->num_entries_seen_total;
+}
+
+DataChunk &ReservoirSample::Chunk() {
+	D_ASSERT(reservoir_chunk);
+	return reservoir_chunk->chunk;
+}
+
+unique_ptr<DataChunk> ReservoirSample::GetChunk() {
+	if (destroyed || !reservoir_chunk || Chunk().size() == 0) {
+		return nullptr;
+	}
+	// cannot destory internal samples.
+	auto ret = make_uniq<DataChunk>();
+
+	SelectionVector ret_sel(STANDARD_VECTOR_SIZE);
+	idx_t collected_samples = GetActiveSampleCount();
+
+	if (collected_samples == 0) {
+		return nullptr;
+	}
+
+	idx_t samples_remaining;
+	idx_t return_chunk_size;
+	if (collected_samples > STANDARD_VECTOR_SIZE) {
+		samples_remaining = collected_samples - STANDARD_VECTOR_SIZE;
+		return_chunk_size = STANDARD_VECTOR_SIZE;
+	} else {
+		samples_remaining = 0;
+		return_chunk_size = collected_samples;
+	}
+
+	for (idx_t i = samples_remaining; i < collected_samples; i++) {
+		// pop samples and reduce size of selection vector.
+		if (GetSamplingState() == SamplingState::RESERVOIR) {
+			auto top = PopFromWeightQueue();
+			ret_sel.set_index(i - samples_remaining, sel.get_index(top.second));
+		} else {
+			ret_sel.set_index(i - samples_remaining, sel.get_index(i));
+		}
+		sel_size -= 1;
+	}
+
+	auto reservoir_types = Chunk().GetTypes();
+
+	ret->Initialize(allocator, reservoir_types, STANDARD_VECTOR_SIZE);
+	ret->Slice(Chunk(), ret_sel, return_chunk_size);
+	ret->SetCardinality(return_chunk_size);
+	return ret;
+}
+
+unique_ptr<ReservoirChunk> ReservoirSample::CreateNewSampleChunk(vector<LogicalType> &types, idx_t size) const {
+	auto new_sample_chunk = make_uniq<ReservoirChunk>();
+	new_sample_chunk->chunk.Initialize(Allocator::DefaultAllocator(), types, size);
+
+	// set the NULL columns correctly
+	for (idx_t col_idx = 0; col_idx < types.size(); col_idx++) {
+		if (!ValidSampleType(types[col_idx]) && stats_sample) {
+			new_sample_chunk->chunk.data[col_idx].SetVectorType(VectorType::CONSTANT_VECTOR);
+			ConstantVector::SetNull(new_sample_chunk->chunk.data[col_idx], true);
+		}
+	}
+	return new_sample_chunk;
+}
+
+void ReservoirSample::Vacuum() {
+	Verify();
+	if (NumSamplesCollected() <= FIXED_SAMPLE_SIZE || !reservoir_chunk || destroyed) {
+		// sample is destroyed or too small to shrink
+		return;
+	}
+
+	auto ret = Copy();
+	auto ret_reservoir = duckdb::unique_ptr_cast<BlockingSample, ReservoirSample>(std::move(ret));
+	reservoir_chunk = std::move(ret_reservoir->reservoir_chunk);
+	sel = std::move(ret_reservoir->sel);
+	sel_size = ret_reservoir->sel_size;
+
+	Verify();
+	// We should only have one sample chunk now.
+	D_ASSERT(Chunk().size() > 0 && Chunk().size() <= sample_count);
+}
+
+unique_ptr<BlockingSample> ReservoirSample::Copy() const {
+
+	auto ret = make_uniq<ReservoirSample>(sample_count);
+	ret->stats_sample = stats_sample;
+
+	ret->base_reservoir_sample = base_reservoir_sample->Copy();
+	ret->destroyed = destroyed;
+
+	if (!reservoir_chunk || destroyed) {
+		return unique_ptr_cast<ReservoirSample, BlockingSample>(std::move(ret));
+	}
+
+	D_ASSERT(reservoir_chunk);
+
+	// create a new sample chunk to store new samples
+	auto types = reservoir_chunk->chunk.GetTypes();
+	// how many values should be copied
+	idx_t values_to_copy = MinValue<idx_t>(GetActiveSampleCount(), sample_count);
+
+	auto new_sample_chunk = CreateNewSampleChunk(types, GetReservoirChunkCapacity());
+
+	SelectionVector sel_copy(sel);
+
+	ret->reservoir_chunk = std::move(new_sample_chunk);
+	ret->UpdateSampleAppend(ret->reservoir_chunk->chunk, reservoir_chunk->chunk, sel_copy, values_to_copy);
+	ret->sel = SelectionVector(values_to_copy);
+	for (idx_t i = 0; i < values_to_copy; i++) {
+		ret->sel.set_index(i, i);
+	}
+	ret->sel_size = sel_size;
+	D_ASSERT(ret->reservoir_chunk->chunk.size() <= sample_count);
+	ret->Verify();
+	return unique_ptr_cast<ReservoirSample, BlockingSample>(std::move(ret));
+}
+
+void ReservoirSample::ConvertToReservoirSample() {
+	D_ASSERT(sel_size <= sample_count);
+	base_reservoir_sample->FillWeights(sel, sel_size);
+}
+
+vector<uint32_t> ReservoirSample::GetRandomizedVector(uint32_t range, uint32_t size) const {
+	vector<uint32_t> ret;
+	ret.reserve(range);
+	for (uint32_t i = 0; i < range; i++) {
+		ret.push_back(i);
+	}
+	if (size == FIXED_SAMPLE_SIZE) {
+		std::shuffle(ret.begin(), ret.end(), base_reservoir_sample->random);
+		return ret;
+	}
+	for (uint32_t i = 0; i < size; i++) {
+		uint32_t random_shuffle = base_reservoir_sample->random.NextRandomInteger32(i, range);
+		if (random_shuffle == i) {
+			// leave the value where it is
+			continue;
+		}
+		uint32_t tmp = ret[random_shuffle];
+		// basically replacing the tuple that was at index actual_sample_indexes[random_shuffle]
+		ret[random_shuffle] = ret[i];
+		ret[i] = tmp;
+	}
+	return ret;
+}
+
+void ReservoirSample::SimpleMerge(ReservoirSample &other) {
+	D_ASSERT(GetPriorityQueueSize() == 0);
+	D_ASSERT(other.GetPriorityQueueSize() == 0);
+	D_ASSERT(GetSamplingState() == SamplingState::RANDOM);
+	D_ASSERT(other.GetSamplingState() == SamplingState::RANDOM);
+
+	if (other.GetActiveSampleCount() == 0 && other.GetTuplesSeen() == 0) {
+		return;
+	}
+
+	if (GetActiveSampleCount() == 0 && GetTuplesSeen() == 0) {
+		sel = SelectionVector(other.sel);
+		sel_size = other.sel_size;
+		base_reservoir_sample->num_entries_seen_total = other.GetTuplesSeen();
+		return;
+	}
+
+	idx_t total_seen = GetTuplesSeen() + other.GetTuplesSeen();
+
+	auto weight_tuples_this = static_cast<double>(GetTuplesSeen()) / static_cast<double>(total_seen);
+	auto weight_tuples_other = static_cast<double>(other.GetTuplesSeen()) / static_cast<double>(total_seen);
+
+	// If weights don't add up to 1, most likely a simple merge occured and no new samples were added.
+	// if that is the case, add the missing weight to the lower weighted sample to adjust.
+	// this is to avoid cases where if you have a 20k row table and add another 20k rows row by row
+	// then eventually the missing weights will add up, and get you a more even distribution
+	if (weight_tuples_this + weight_tuples_other < 1) {
+		weight_tuples_other += 1 - (weight_tuples_other + weight_tuples_this);
+	}
+
+	idx_t keep_from_this = 0;
+	idx_t keep_from_other = 0;
+	D_ASSERT(stats_sample);
+	D_ASSERT(sample_count == FIXED_SAMPLE_SIZE);
+	D_ASSERT(sample_count == other.sample_count);
+	auto sample_count_double = static_cast<double>(sample_count);
+
+	if (weight_tuples_this > weight_tuples_other) {
+		keep_from_this = MinValue<idx_t>(static_cast<idx_t>(round(sample_count_double * weight_tuples_this)),
+		                                 GetActiveSampleCount());
+		keep_from_other = MinValue<idx_t>(sample_count - keep_from_this, other.GetActiveSampleCount());
+	} else {
+		keep_from_other = MinValue<idx_t>(static_cast<idx_t>(round(sample_count_double * weight_tuples_other)),
+		                                  other.GetActiveSampleCount());
+		keep_from_this = MinValue<idx_t>(sample_count - keep_from_other, GetActiveSampleCount());
+	}
+
+	D_ASSERT(keep_from_this <= GetActiveSampleCount());
+	D_ASSERT(keep_from_other <= other.GetActiveSampleCount());
+	D_ASSERT(keep_from_other + keep_from_this <= FIXED_SAMPLE_SIZE);
+	idx_t size_after_merge = MinValue<idx_t>(keep_from_other + keep_from_this, FIXED_SAMPLE_SIZE);
+
+	// Check if appending the other samples to this will go over the sample chunk size
+	if (reservoir_chunk->chunk.size() + keep_from_other > GetReservoirChunkCapacity()) {
+		Vacuum();
+	}
+
+	D_ASSERT(size_after_merge <= other.GetActiveSampleCount() + GetActiveSampleCount());
+	SelectionVector chunk_sel(keep_from_other);
+	auto offset = reservoir_chunk->chunk.size();
+	for (idx_t i = keep_from_this; i < size_after_merge; i++) {
+		if (i >= GetActiveSampleCount()) {
+			sel.set_index(GetActiveSampleCount(), offset);
+			sel_size += 1;
+		} else {
+			sel.set_index(i, offset);
+		}
+		chunk_sel.set_index(i - keep_from_this, other.sel.get_index(i - keep_from_this));
+		offset += 1;
+	}
+
+	D_ASSERT(GetActiveSampleCount() == size_after_merge);
+
+	// Copy the rows that make it to the sample from other and put them into this.
+	UpdateSampleAppend(reservoir_chunk->chunk, other.reservoir_chunk->chunk, chunk_sel, keep_from_other);
+	base_reservoir_sample->num_entries_seen_total += other.GetTuplesSeen();
+
+	// if THIS has too many samples now, we conver it to a slower sample.
+	if (GetTuplesSeen() >= FIXED_SAMPLE_SIZE * FAST_TO_SLOW_THRESHOLD) {
+		ConvertToReservoirSample();
+	}
+	Verify();
+}
+
+void ReservoirSample::WeightedMerge(ReservoirSample &other_sample) {
+	D_ASSERT(GetSamplingState() == SamplingState::RESERVOIR);
+	D_ASSERT(other_sample.GetSamplingState() == SamplingState::RESERVOIR);
+
+	// Find out how many samples we want to keep.
+	idx_t total_samples = GetActiveSampleCount() + other_sample.GetActiveSampleCount();
+	idx_t total_samples_seen =
+	    base_reservoir_sample->num_entries_seen_total + other_sample.base_reservoir_sample->num_entries_seen_total;
+	idx_t num_samples_to_keep = MinValue<idx_t>(total_samples, MinValue<idx_t>(sample_count, total_samples_seen));
+
+	D_ASSERT(GetActiveSampleCount() <= num_samples_to_keep);
+	D_ASSERT(total_samples <= FIXED_SAMPLE_SIZE * 2);
+
+	// pop from base base_reservoir weights until there are num_samples_to_keep left.
+	vector<idx_t> this_indexes_to_replace;
+	for (idx_t i = num_samples_to_keep; i < total_samples; i++) {
+		auto min_weight_this = base_reservoir_sample->min_weight_threshold;
+		auto min_weight_other = other_sample.base_reservoir_sample->min_weight_threshold;
+		// min weight threshol is always positive
+		if (min_weight_this > min_weight_other) {
+			// pop from other
+			other_sample.base_reservoir_sample->reservoir_weights.pop();
+			other_sample.base_reservoir_sample->UpdateMinWeightThreshold();
+		} else {
+			auto top_this = PopFromWeightQueue();
+			this_indexes_to_replace.push_back(top_this.second);
+			base_reservoir_sample->UpdateMinWeightThreshold();
+		}
+	}
+
+	D_ASSERT(other_sample.GetPriorityQueueSize() + GetPriorityQueueSize() <= FIXED_SAMPLE_SIZE);
+	D_ASSERT(other_sample.GetPriorityQueueSize() + GetPriorityQueueSize() == num_samples_to_keep);
+	D_ASSERT(other_sample.reservoir_chunk->chunk.GetTypes() == reservoir_chunk->chunk.GetTypes());
+
+	// Prepare a selection vector to copy data from the other sample chunk to this sample chunk
+	SelectionVector sel_other(other_sample.GetPriorityQueueSize());
+	D_ASSERT(GetPriorityQueueSize() <= num_samples_to_keep);
+	D_ASSERT(other_sample.GetPriorityQueueSize() >= this_indexes_to_replace.size());
+	idx_t chunk_offset = 0;
+
+	// Now push weights from other.base_reservoir_sample to this
+	// Depending on how many sample values "this" has, we either need to add to the selection vector
+	// Or replace values in "this'" selection vector
+	idx_t i = 0;
+	while (other_sample.GetPriorityQueueSize() > 0) {
+		auto other_top = other_sample.PopFromWeightQueue();
+		idx_t index_for_new_pair = chunk_offset + reservoir_chunk->chunk.size();
+
+		// update the sel used to copy values from other to this
+		sel_other.set_index(chunk_offset, other_top.second);
+		if (i < this_indexes_to_replace.size()) {
+			auto replacement_index = this_indexes_to_replace[i];
+			sel.set_index(replacement_index, index_for_new_pair);
+			other_top.second = replacement_index;
+		} else {
+			sel.set_index(sel_size, index_for_new_pair);
+			other_top.second = sel_size;
+			sel_size += 1;
+		}
+
+		// make sure that the sample indexes are (this.sample_chunk.size() + chunk_offfset)
+		base_reservoir_sample->reservoir_weights.push(other_top);
+		chunk_offset += 1;
+		i += 1;
+	}
+
+	D_ASSERT(GetPriorityQueueSize() == num_samples_to_keep);
+
+	base_reservoir_sample->UpdateMinWeightThreshold();
+	D_ASSERT(base_reservoir_sample->min_weight_threshold > 0);
+	base_reservoir_sample->num_entries_seen_total = GetTuplesSeen() + other_sample.GetTuplesSeen();
+
+	UpdateSampleAppend(reservoir_chunk->chunk, other_sample.reservoir_chunk->chunk, sel_other, chunk_offset);
+	if (reservoir_chunk->chunk.size() > FIXED_SAMPLE_SIZE * (FIXED_SAMPLE_SIZE_MULTIPLIER - 3)) {
+		Vacuum();
+	}
+
+	Verify();
+}
+
+void ReservoirSample::Merge(unique_ptr<BlockingSample> other) {
+	if (destroyed || other->destroyed) {
+		Destroy();
+		return;
+	}
+
+	D_ASSERT(other->type == SampleType::RESERVOIR_SAMPLE);
+	auto &other_sample = other->Cast<ReservoirSample>();
+
+	// if the other sample has not collected anything yet return
+	if (!other_sample.reservoir_chunk || other_sample.reservoir_chunk->chunk.size() == 0) {
+		return;
+	}
+
+	// this has not collected samples, take over the other
+	if (!reservoir_chunk || reservoir_chunk->chunk.size() == 0) {
+		base_reservoir_sample = std::move(other->base_reservoir_sample);
+		reservoir_chunk = std::move(other_sample.reservoir_chunk);
+		sel = SelectionVector(other_sample.sel);
+		sel_size = other_sample.sel_size;
+		Verify();
+		return;
+	}
+	//! Both samples are still in "fast sampling" method
+	if (GetSamplingState() == SamplingState::RANDOM && other_sample.GetSamplingState() == SamplingState::RANDOM) {
+		SimpleMerge(other_sample);
+		return;
+	}
+
+	// One or none of the samples are in "Fast Sampling" method.
+	// When this is the case, switch both to slow sampling
+	ConvertToReservoirSample();
+	other_sample.ConvertToReservoirSample();
+	WeightedMerge(other_sample);
+}
+
+void ReservoirSample::ShuffleSel(SelectionVector &sel, idx_t range, idx_t size) const {
+	auto randomized = GetRandomizedVector(static_cast<uint32_t>(range), static_cast<uint32_t>(size));
+	SelectionVector original_sel(range);
+	for (idx_t i = 0; i < range; i++) {
+		original_sel.set_index(i, sel.get_index(i));
+	}
+	for (idx_t i = 0; i < size; i++) {
+		sel.set_index(i, original_sel.get_index(randomized[i]));
+	}
+}
+
+void ReservoirSample::NormalizeWeights() {
+	vector<std::pair<double, idx_t>> tmp_weights;
+	while (!base_reservoir_sample->reservoir_weights.empty()) {
+		auto top = base_reservoir_sample->reservoir_weights.top();
+		tmp_weights.push_back(std::move(top));
+		base_reservoir_sample->reservoir_weights.pop();
+	}
+	std::sort(tmp_weights.begin(), tmp_weights.end(),
+	          [&](std::pair<double, idx_t> a, std::pair<double, idx_t> b) { return a.second < b.second; });
+	for (idx_t i = 0; i < tmp_weights.size(); i++) {
+		base_reservoir_sample->reservoir_weights.emplace(tmp_weights.at(i).first, i);
+	}
+	base_reservoir_sample->SetNextEntry();
+}
+
+void ReservoirSample::EvictOverBudgetSamples() {
+	Verify();
+	if (!reservoir_chunk || destroyed) {
+		return;
+	}
+
+	// since this is for serialization, we really need to make sure keep a
+	// minimum of 1% of the rows or 2048 rows
+	idx_t num_samples_to_keep =
+	    MinValue<idx_t>(FIXED_SAMPLE_SIZE, static_cast<idx_t>(SAVE_PERCENTAGE * static_cast<double>(GetTuplesSeen())));
+
+	if (num_samples_to_keep <= 0) {
+		reservoir_chunk->chunk.SetCardinality(0);
+		return;
+	}
+
+	if (num_samples_to_keep == sample_count) {
+		return;
+	}
+
+	// if we over sampled, make sure we only keep the highest percentage samples
+	std::unordered_set<idx_t> selections_to_delete;
+
+	while (num_samples_to_keep < GetPriorityQueueSize()) {
+		auto top = PopFromWeightQueue();
+		D_ASSERT(top.second < sel_size);
+		selections_to_delete.emplace(top.second);
+	}
+
+	// set up reservoir chunk for the reservoir sample
+	D_ASSERT(reservoir_chunk->chunk.size() <= sample_count);
+	// create a new sample chunk to store new samples
+	auto types = reservoir_chunk->chunk.GetTypes();
+	D_ASSERT(num_samples_to_keep <= sample_count);
+	D_ASSERT(stats_sample);
+	D_ASSERT(sample_count == FIXED_SAMPLE_SIZE);
+	auto new_reservoir_chunk = CreateNewSampleChunk(types, sample_count);
+
+	// The current selection vector can potentially have 2048 valid mappings.
+	// If we need to save a sample with less rows than that, we need to do the following
+	// 1. Create a new selection vector that doesn't point to the rows we are evicting
+	SelectionVector new_sel(num_samples_to_keep);
+	idx_t offset = 0;
+	for (idx_t i = 0; i < num_samples_to_keep + selections_to_delete.size(); i++) {
+		if (selections_to_delete.find(i) == selections_to_delete.end()) {
+			D_ASSERT(i - offset < num_samples_to_keep);
+			new_sel.set_index(i - offset, sel.get_index(i));
+		} else {
+			offset += 1;
+		}
+	}
+	// 2. Update row_ids in our weights so that they don't store rows ids to
+	//    indexes in the selection vector that have been evicted.
+	if (!selections_to_delete.empty()) {
+		NormalizeWeights();
+	}
+
+	D_ASSERT(reservoir_chunk->chunk.GetTypes() == new_reservoir_chunk->chunk.GetTypes());
+
+	UpdateSampleAppend(new_reservoir_chunk->chunk, reservoir_chunk->chunk, new_sel, num_samples_to_keep);
+	// set the cardinality
+	new_reservoir_chunk->chunk.SetCardinality(num_samples_to_keep);
+	reservoir_chunk = std::move(new_reservoir_chunk);
+	sel_size = num_samples_to_keep;
+	base_reservoir_sample->UpdateMinWeightThreshold();
+}
+
+void ReservoirSample::ExpandSerializedSample() {
+	if (!reservoir_chunk) {
+		return;
+	}
+
+	auto types = reservoir_chunk->chunk.GetTypes();
+	auto new_res_chunk = CreateNewSampleChunk(types, GetReservoirChunkCapacity());
+	auto copy_count = reservoir_chunk->chunk.size();
+	SelectionVector tmp_sel = SelectionVector(0, copy_count);
+	UpdateSampleAppend(new_res_chunk->chunk, reservoir_chunk->chunk, tmp_sel, copy_count);
+	new_res_chunk->chunk.SetCardinality(copy_count);
+	std::swap(reservoir_chunk, new_res_chunk);
+}
+
+idx_t ReservoirSample::GetReservoirChunkCapacity() const {
+	return sample_count + (FIXED_SAMPLE_SIZE_MULTIPLIER * FIXED_SAMPLE_SIZE);
+}
+
+idx_t ReservoirSample::FillReservoir(DataChunk &chunk) {
+
+	idx_t ingested_count = 0;
+	if (!reservoir_chunk) {
+		if (chunk.size() > FIXED_SAMPLE_SIZE) {
+			throw InternalException("Creating sample with DataChunk that is larger than the fixed sample size");
+		}
+		auto types = chunk.GetTypes();
+		// create a new sample chunk to store new samples
+		reservoir_chunk = CreateNewSampleChunk(types, GetReservoirChunkCapacity());
+	}
+
+	idx_t actual_sample_index_start = GetActiveSampleCount();
+	D_ASSERT(reservoir_chunk->chunk.ColumnCount() == chunk.ColumnCount());
+
+	if (reservoir_chunk->chunk.size() < sample_count) {
+		ingested_count = MinValue<idx_t>(sample_count - reservoir_chunk->chunk.size(), chunk.size());
+		auto random_other_sel =
+		    GetRandomizedVector(static_cast<uint32_t>(ingested_count), static_cast<uint32_t>(ingested_count));
+		SelectionVector sel_for_input_chunk(ingested_count);
+		for (idx_t i = 0; i < ingested_count; i++) {
+			sel.set_index(actual_sample_index_start + i, actual_sample_index_start + i);
+			sel_for_input_chunk.set_index(i, random_other_sel[i]);
+		}
+		UpdateSampleAppend(reservoir_chunk->chunk, chunk, sel_for_input_chunk, ingested_count);
+		sel_size += ingested_count;
+	}
+	D_ASSERT(GetActiveSampleCount() <= sample_count);
+	D_ASSERT(GetActiveSampleCount() >= ingested_count);
+	// always return how many tuples were ingested
+	return ingested_count;
+}
+
+void ReservoirSample::Destroy() {
+	destroyed = true;
+}
+
+SelectionVectorHelper ReservoirSample::GetReplacementIndexes(idx_t sample_chunk_offset,
+                                                             idx_t theoretical_chunk_length) {
+	if (GetSamplingState() == SamplingState::RANDOM) {
+		return GetReplacementIndexesFast(sample_chunk_offset, theoretical_chunk_length);
+	}
+	return GetReplacementIndexesSlow(sample_chunk_offset, theoretical_chunk_length);
+}
+
+SelectionVectorHelper ReservoirSample::GetReplacementIndexesFast(idx_t sample_chunk_offset, idx_t chunk_length) {
+
+	// how much weight to the other tuples have compared to the ones in this chunk?
+	auto weight_tuples_other = static_cast<double>(chunk_length) / static_cast<double>(GetTuplesSeen() + chunk_length);
+	auto num_to_pop = static_cast<uint32_t>(round(weight_tuples_other * static_cast<double>(sample_count)));
+	D_ASSERT(num_to_pop <= sample_count);
+	D_ASSERT(num_to_pop <= sel_size);
+	SelectionVectorHelper ret;
+
+	if (num_to_pop == 0) {
+		ret.sel = SelectionVector(num_to_pop);
+		ret.size = 0;
+		return ret;
+	}
+	std::unordered_map<idx_t, idx_t> replacement_indexes;
+	SelectionVector chunk_sel(num_to_pop);
+
+	auto random_indexes_chunk = GetRandomizedVector(static_cast<uint32_t>(chunk_length), num_to_pop);
+	auto random_sel_indexes = GetRandomizedVector(static_cast<uint32_t>(sel_size), num_to_pop);
+	for (idx_t i = 0; i < num_to_pop; i++) {
+		// update the selection vector for the reservoir sample
+		chunk_sel.set_index(i, random_indexes_chunk[i]);
+		// sel is not guaratneed to be random, so we update the indexes according to our
+		// random sel indexes.
+		sel.set_index(random_sel_indexes[i], sample_chunk_offset + i);
+	}
+
+	D_ASSERT(sel_size == sample_count);
+
+	ret.sel = SelectionVector(chunk_sel);
+	ret.size = num_to_pop;
+	return ret;
+}
+
+SelectionVectorHelper ReservoirSample::GetReplacementIndexesSlow(const idx_t sample_chunk_offset,
+                                                                 const idx_t chunk_length) {
+	idx_t remaining = chunk_length;
+	std::unordered_map<idx_t, idx_t> ret_map;
+	idx_t sample_chunk_index = 0;
+
+	idx_t base_offset = 0;
+
+	while (true) {
+		idx_t offset =
+		    base_reservoir_sample->next_index_to_sample - base_reservoir_sample->num_entries_to_skip_b4_next_sample;
+		if (offset >= remaining) {
+			// not in this chunk! increment current count and go to the next chunk
+			base_reservoir_sample->num_entries_to_skip_b4_next_sample += remaining;
+			break;
+		}
+		// in this chunk! replace the element
+		// ret[index_in_new_chunk] = index_in_sample_chunk (the sample chunk offset will be applied later)
+		// D_ASSERT(sample_chunk_index == ret.size());
+		ret_map[base_offset + offset] = sample_chunk_index;
+		double r2 = base_reservoir_sample->random.NextRandom32(base_reservoir_sample->min_weight_threshold, 1);
+		// replace element in our max_heap
+		// first get the top most pair
+		const auto top = PopFromWeightQueue();
+		const auto index = top.second;
+		const auto index_in_sample_chunk = sample_chunk_offset + sample_chunk_index;
+		sel.set_index(index, index_in_sample_chunk);
+		base_reservoir_sample->ReplaceElementWithIndex(index, r2, false);
+
+		sample_chunk_index += 1;
+		// shift the chunk forward
+		remaining -= offset;
+		base_offset += offset;
+	}
+
+	// create selection vector to return
+	SelectionVector ret_sel(ret_map.size());
+	D_ASSERT(sel_size == sample_count);
+	for (auto &kv : ret_map) {
+		ret_sel.set_index(kv.second, kv.first);
+	}
+	SelectionVectorHelper ret;
+	ret.sel = SelectionVector(ret_sel);
+	ret.size = static_cast<uint32_t>(ret_map.size());
+	return ret;
+}
+
+void ReservoirSample::Finalize() {
+}
+
+bool ReservoirSample::ValidSampleType(const LogicalType &type) {
+	return type.IsNumeric();
+}
+
+void ReservoirSample::UpdateSampleAppend(DataChunk &this_, DataChunk &other, SelectionVector &other_sel,
+                                         idx_t append_count) const {
+	idx_t new_size = this_.size() + append_count;
+	if (other.size() == 0) {
+		return;
+	}
+	D_ASSERT(this_.GetTypes() == other.GetTypes());
+
+	// UpdateSampleAppend(this_, other, other_sel, append_count);
+	D_ASSERT(this_.GetTypes() == other.GetTypes());
+	auto types = reservoir_chunk->chunk.GetTypes();
+
+	for (idx_t i = 0; i < reservoir_chunk->chunk.ColumnCount(); i++) {
+		auto col_type = types[i];
+		if (ValidSampleType(col_type) || !stats_sample) {
+			D_ASSERT(this_.data[i].GetVectorType() == VectorType::FLAT_VECTOR);
+			VectorOperations::Copy(other.data[i], this_.data[i], other_sel, append_count, 0, this_.size());
+		}
+	}
+	this_.SetCardinality(new_size);
+}
+
+void ReservoirSample::AddToReservoir(DataChunk &chunk) {
+	if (destroyed || chunk.size() == 0) {
+		return;
+	}
+
+	idx_t tuples_consumed = FillReservoir(chunk);
+	base_reservoir_sample->num_entries_seen_total += tuples_consumed;
+	D_ASSERT(sample_count == 0 || reservoir_chunk->chunk.size() >= 1);
+
+	if (tuples_consumed == chunk.size()) {
+		return;
+	}
+
+	// the chunk filled the first FIXED_SAMPLE_SIZE chunk but still has tuples remaining
+	// slice the chunk and call AddToReservoir again.
+	if (tuples_consumed != chunk.size() && tuples_consumed != 0) {
+		// Fill reservoir consumed some of the chunk to reach FIXED_SAMPLE_SIZE
+		// now we need to
+		// So we slice it and call AddToReservoir
+		auto slice = make_uniq<DataChunk>();
+		auto samples_remaining = chunk.size() - tuples_consumed;
+		auto types = chunk.GetTypes();
+		SelectionVector input_sel(samples_remaining);
+		for (idx_t i = 0; i < samples_remaining; i++) {
+			input_sel.set_index(i, tuples_consumed + i);
+		}
+		slice->Initialize(Allocator::DefaultAllocator(), types, samples_remaining);
+		slice->Slice(chunk, input_sel, samples_remaining);
+		slice->SetCardinality(samples_remaining);
+		AddToReservoir(*slice);
+		return;
+	}
+
+	// at this point we should have collected at least sample count samples
+	D_ASSERT(GetActiveSampleCount() >= sample_count);
+
+	auto chunk_sel = GetReplacementIndexes(reservoir_chunk->chunk.size(), chunk.size());
+
+	if (chunk_sel.size == 0) {
+		// not adding any samples
+		return;
+	}
+	idx_t size = chunk_sel.size;
+	D_ASSERT(size <= chunk.size());
+
+	UpdateSampleAppend(reservoir_chunk->chunk, chunk, chunk_sel.sel, size);
+
+	base_reservoir_sample->num_entries_seen_total += chunk.size();
+	D_ASSERT(base_reservoir_sample->reservoir_weights.size() == 0 ||
+	         base_reservoir_sample->reservoir_weights.size() == sample_count);
+
+	Verify();
+
+	// if we are over the threshold, we ned to swith to slow sampling.
+	if (GetSamplingState() == SamplingState::RANDOM && GetTuplesSeen() >= FIXED_SAMPLE_SIZE * FAST_TO_SLOW_THRESHOLD) {
+		ConvertToReservoirSample();
+	}
+	if (reservoir_chunk->chunk.size() >= (GetReservoirChunkCapacity() - (static_cast<idx_t>(FIXED_SAMPLE_SIZE) * 3))) {
+		Vacuum();
+	}
+}
+
+void ReservoirSample::Verify() {
+#ifdef DEBUG
+	if (destroyed) {
+		return;
+	}
+	if (GetPriorityQueueSize() == 0) {
+		D_ASSERT(GetActiveSampleCount() <= sample_count);
+		D_ASSERT(GetTuplesSeen() >= GetActiveSampleCount());
+		return;
+	}
+	if (NumSamplesCollected() > sample_count) {
+		D_ASSERT(GetPriorityQueueSize() == sample_count);
+	} else if (NumSamplesCollected() <= sample_count && GetPriorityQueueSize() > 0) {
+		// it's possible to collect more samples than your priority queue size.
+		// see sample_converts_to_reservoir_sample.test
+		D_ASSERT(NumSamplesCollected() >= GetPriorityQueueSize());
+	}
+	auto base_reservoir_copy = base_reservoir_sample->Copy();
+	std::unordered_map<idx_t, idx_t> index_count;
+	while (!base_reservoir_copy->reservoir_weights.empty()) {
+		auto &pair = base_reservoir_copy->reservoir_weights.top();
+		if (index_count.find(pair.second) == index_count.end()) {
+			index_count[pair.second] = 1;
+			base_reservoir_copy->reservoir_weights.pop();
+		} else {
+			index_count[pair.second] += 1;
+			base_reservoir_copy->reservoir_weights.pop();
+			throw InternalException("Duplicate selection index in reservoir weights");
+		}
+	}
+	// TODO: Verify the Sel as well. No duplicate indices.
+
+	if (reservoir_chunk) {
+		reservoir_chunk->chunk.Verify();
+	}
+#endif
+}
+
+ReservoirSamplePercentage::ReservoirSamplePercentage(double percentage, int64_t seed, idx_t reservoir_sample_size)
+    : BlockingSample(seed), allocator(Allocator::DefaultAllocator()), sample_percentage(percentage / 100.0),
+      reservoir_sample_size(reservoir_sample_size), current_count(0), is_finalized(false) {
+	current_sample = make_uniq<ReservoirSample>(allocator, reservoir_sample_size, base_reservoir_sample->random());
+	type = SampleType::RESERVOIR_PERCENTAGE_SAMPLE;
+}
+
+ReservoirSamplePercentage::ReservoirSamplePercentage(Allocator &allocator, double percentage, int64_t seed)
+    : BlockingSample(seed), allocator(allocator), sample_percentage(percentage / 100.0), current_count(0),
+      is_finalized(false) {
+	reservoir_sample_size = (idx_t)(sample_percentage * RESERVOIR_THRESHOLD);
+	current_sample = make_uniq<ReservoirSample>(allocator, reservoir_sample_size, base_reservoir_sample->random());
+	type = SampleType::RESERVOIR_PERCENTAGE_SAMPLE;
+}
+
+ReservoirSamplePercentage::ReservoirSamplePercentage(double percentage, int64_t seed)
+    : ReservoirSamplePercentage(Allocator::DefaultAllocator(), percentage, seed) {
+}
+
+void ReservoirSamplePercentage::AddToReservoir(DataChunk &input) {
+	base_reservoir_sample->num_entries_seen_total += input.size();
+	if (current_count + input.size() > RESERVOIR_THRESHOLD) {
+		// we don't have enough space in our current reservoir
+		// first check what we still need to append to the current sample
+		idx_t append_to_current_sample_count = RESERVOIR_THRESHOLD - current_count;
+		idx_t append_to_next_sample = input.size() - append_to_current_sample_count;
+		if (append_to_current_sample_count > 0) {
+			// we have elements remaining, first add them to the current sample
+			if (append_to_next_sample > 0) {
+				// we need to also add to the next sample
+				DataChunk new_chunk;
+				new_chunk.InitializeEmpty(input.GetTypes());
+				new_chunk.Slice(input, *FlatVector::IncrementalSelectionVector(), append_to_current_sample_count);
+				new_chunk.Flatten();
+				current_sample->AddToReservoir(new_chunk);
+			} else {
+				input.Flatten();
+				input.SetCardinality(append_to_current_sample_count);
+				current_sample->AddToReservoir(input);
+			}
+		}
+		if (append_to_next_sample > 0) {
+			// slice the input for the remainder
+			SelectionVector sel(append_to_next_sample);
+			for (idx_t i = append_to_current_sample_count; i < append_to_next_sample + append_to_current_sample_count;
+			     i++) {
+				sel.set_index(i - append_to_current_sample_count, i);
+			}
+			input.Slice(sel, append_to_next_sample);
+		}
+		// now our first sample is filled: append it to the set of finished samples
+		finished_samples.push_back(std::move(current_sample));
+
+		// allocate a new sample, and potentially add the remainder of the current input to that sample
+		current_sample = make_uniq<ReservoirSample>(allocator, reservoir_sample_size, base_reservoir_sample->random());
+		if (append_to_next_sample > 0) {
+			current_sample->AddToReservoir(input);
+		}
+		current_count = append_to_next_sample;
+	} else {
+		// we can just append to the current sample
+		current_count += input.size();
+		current_sample->AddToReservoir(input);
+	}
+}
+
+unique_ptr<DataChunk> ReservoirSamplePercentage::GetChunk() {
+	// reservoir sample percentage should never stay
+	if (!is_finalized) {
+		Finalize();
+	}
+	while (!finished_samples.empty()) {
+		auto &front = finished_samples.front();
+		auto chunk = front->GetChunk();
+		if (chunk && chunk->size() > 0) {
+			return chunk;
+		}
+		// move to the next sample
+		finished_samples.erase(finished_samples.begin());
+	}
+	return nullptr;
+}
+
+unique_ptr<BlockingSample> ReservoirSamplePercentage::Copy() const {
+	throw InternalException("Cannot call Copy on ReservoirSample Percentage");
+}
+
+void ReservoirSamplePercentage::Finalize() {
+	// need to finalize the current sample, if any
+	// we are finializing, so we are starting to return chunks. Our last chunk has
+	// sample_percentage * RESERVOIR_THRESHOLD entries that hold samples.
+	// if our current count is less than the sample_percentage * RESERVOIR_THRESHOLD
+	// then we have sampled too much for the current_sample and we need to redo the sample
+	// otherwise we can just push the current sample back
+	// Imagine sampling 70% of 100 rows (so 70 rows). We allocate sample_percentage * RESERVOIR_THRESHOLD
+	// -----------------------------------------
+	auto sampled_more_than_required =
+	    static_cast<double>(current_count) > sample_percentage * RESERVOIR_THRESHOLD || finished_samples.empty();
+	if (current_count > 0 && sampled_more_than_required) {
+		// create a new sample
+		auto new_sample_size = static_cast<idx_t>(round(sample_percentage * static_cast<double>(current_count)));
+		auto new_sample = make_uniq<ReservoirSample>(allocator, new_sample_size, base_reservoir_sample->random());
+		while (true) {
+			auto chunk = current_sample->GetChunk();
+			if (!chunk || chunk->size() == 0) {
+				break;
+			}
+			new_sample->AddToReservoir(*chunk);
+		}
+		finished_samples.push_back(std::move(new_sample));
+	} else {
+		finished_samples.push_back(std::move(current_sample));
+	}
+	// when finalizing, current_sample is null. All samples are now in finished samples.
+	current_sample = nullptr;
+	is_finalized = true;
+}
+
+} // namespace duckdb