Remove chunk collection from reservoir sampler

.github/config/uncovered_files.csv

@@ -72,7 +72,6 @@ common/types/batched_data_collection.cpp	11
 common/types/bit.cpp	9
 common/types/blob.cpp	3
 common/types/cast_helpers.cpp	2
-common/types/chunk_collection.cpp	94
 common/types/column/column_data_allocator.cpp	13
 common/types/column/column_data_collection.cpp	55
 common/types/column/partitioned_column_data.cpp	8
@@ -385,7 +384,6 @@ include/duckdb/common/sort/sorted_block.hpp	1
 include/duckdb/common/string_util.hpp	9
 include/duckdb/common/type_util.hpp	2
 include/duckdb/common/types.hpp	5
-include/duckdb/common/types/chunk_collection.hpp	6
 include/duckdb/common/types/column/column_data_allocator.hpp	3
 include/duckdb/common/types/column/partitioned_column_data.hpp	2
 include/duckdb/common/types/datetime.hpp	13

extension/parquet/column_reader.cpp

@@ -20,7 +20,6 @@
 #ifndef DUCKDB_AMALGAMATION
 #include "duckdb/common/types/bit.hpp"
 #include "duckdb/common/types/blob.hpp"
-#include "duckdb/common/types/chunk_collection.hpp"
 #endif
 
 namespace duckdb {

extension/parquet/column_writer.cpp

@@ -13,7 +13,6 @@
 #include "duckdb/common/serializer/memory_stream.hpp"
 #include "duckdb/common/serializer/write_stream.hpp"
 #include "duckdb/common/string_map_set.hpp"
-#include "duckdb/common/types/chunk_collection.hpp"
 #include "duckdb/common/types/date.hpp"
 #include "duckdb/common/types/hugeint.hpp"
 #include "duckdb/common/types/uhugeint.hpp"

extension/parquet/include/column_reader.hpp

@@ -19,7 +19,6 @@
 #ifndef DUCKDB_AMALGAMATION
 
 #include "duckdb/common/operator/cast_operators.hpp"
-#include "duckdb/common/types/chunk_collection.hpp"
 #include "duckdb/common/types/string_type.hpp"
 #include "duckdb/common/types/vector.hpp"
 #include "duckdb/common/types/vector_cache.hpp"

extension/parquet/parquet_extension.cpp

@@ -25,7 +25,6 @@
 #include "duckdb/common/multi_file_reader.hpp"
 #include "duckdb/common/serializer/deserializer.hpp"
 #include "duckdb/common/serializer/serializer.hpp"
-#include "duckdb/common/types/chunk_collection.hpp"
 #include "duckdb/function/copy_function.hpp"
 #include "duckdb/function/pragma_function.hpp"
 #include "duckdb/function/table_function.hpp"

src/common/symbols.cpp

@@ -5,7 +5,6 @@
 
 #include "duckdb/catalog/catalog.hpp"
 #include "duckdb/catalog/catalog_entry/list.hpp"
-#include "duckdb/common/types/chunk_collection.hpp"
 #include "duckdb/common/types/column/column_data_allocator.hpp"
 #include "duckdb/common/types/column/column_data_collection.hpp"
 #include "duckdb/common/vector.hpp"

src/common/types/CMakeLists.txt

@@ -13,7 +13,6 @@ add_library_unity(
   bit.cpp
   blob.cpp
   cast_helpers.cpp
-  chunk_collection.cpp
   conflict_manager.cpp
   conflict_info.cpp
   data_chunk.cpp

src/common/types/vector_buffer.cpp

@@ -1,7 +1,6 @@
 #include "duckdb/common/types/vector_buffer.hpp"
 
 #include "duckdb/common/assert.hpp"
-#include "duckdb/common/types/chunk_collection.hpp"
 #include "duckdb/common/types/vector.hpp"
 #include "duckdb/common/vector_operations/vector_operations.hpp"
 #include "duckdb/storage/buffer/buffer_handle.hpp"

src/common/vector_operations/vector_copy.cpp

@@ -6,7 +6,6 @@
 
 #include "duckdb/common/exception.hpp"
 #include "duckdb/common/types/null_value.hpp"
-#include "duckdb/common/types/chunk_collection.hpp"
 #include "duckdb/common/uhugeint.hpp"
 #include "duckdb/storage/segment/uncompressed.hpp"
 #include "duckdb/common/vector_operations/vector_operations.hpp"

src/core_functions/aggregate/holistic/quantile.cpp

@@ -6,7 +6,7 @@
 #include "duckdb/common/operator/cast_operators.hpp"
 #include "duckdb/common/operator/abs.hpp"
 #include "duckdb/common/operator/multiply.hpp"
-#include "duckdb/common/types/chunk_collection.hpp"
+
 #include "duckdb/common/types/timestamp.hpp"
 #include "duckdb/common/queue.hpp"
 #include "duckdb/common/serializer/serializer.hpp"

src/core_functions/aggregate/holistic/reservoir_quantile.cpp

@@ -32,7 +32,7 @@ struct ReservoirQuantileState {
 	}
 
 	void ReplaceElement(T &input) {
-		v[r_samp->min_entry] = input;
+		v[r_samp->min_weighted_entry_index] = input;
 		r_samp->ReplaceElement();
 	}
 
@@ -41,8 +41,8 @@ struct ReservoirQuantileState {
 			v[pos++] = element;
 			r_samp->InitializeReservoir(pos, len);
 		} else {
-			D_ASSERT(r_samp->next_index >= r_samp->current_count);
-			if (r_samp->next_index == r_samp->current_count) {
+			D_ASSERT(r_samp->next_index_to_sample >= r_samp->num_entries_to_skip_b4_next_sample);
+			if (r_samp->next_index_to_sample == r_samp->num_entries_to_skip_b4_next_sample) {
 				ReplaceElement(element);
 			}
 		}

src/core_functions/scalar/list/list_sort.cpp

@@ -1,6 +1,6 @@
 #include "duckdb/core_functions/scalar/list_functions.hpp"
 #include "duckdb/common/enum_util.hpp"
-#include "duckdb/common/types/chunk_collection.hpp"
+
 #include "duckdb/planner/expression/bound_function_expression.hpp"
 #include "duckdb/planner/expression/bound_cast_expression.hpp"
 #include "duckdb/execution/expression_executor.hpp"

src/execution/operator/aggregate/physical_window.cpp

@@ -8,7 +8,7 @@
 #include "duckdb/common/radix_partitioning.hpp"
 #include "duckdb/common/row_operations/row_operations.hpp"
 #include "duckdb/common/sort/partition_state.hpp"
-#include "duckdb/common/types/chunk_collection.hpp"
+
 #include "duckdb/common/types/column/column_data_consumer.hpp"
 #include "duckdb/common/types/row/row_data_collection_scanner.hpp"
 #include "duckdb/common/uhugeint.hpp"

src/execution/operator/helper/physical_materialized_collector.cpp

@@ -1,5 +1,5 @@
 #include "duckdb/execution/operator/helper/physical_materialized_collector.hpp"
-#include "duckdb/common/types/chunk_collection.hpp"
+
 #include "duckdb/main/materialized_query_result.hpp"
 #include "duckdb/main/client_context.hpp"
 

src/execution/operator/helper/physical_reservoir_sample.cpp

@@ -6,6 +6,7 @@ namespace duckdb {
 //===--------------------------------------------------------------------===//
 // Sink
 //===--------------------------------------------------------------------===//
+
 class SampleGlobalSinkState : public GlobalSinkState {
 public:
 	explicit SampleGlobalSinkState(Allocator &allocator, SampleOptions &options) {
@@ -24,7 +25,8 @@ class SampleGlobalSinkState : public GlobalSinkState {
 		}
 	}
 
-	//! The lock for updating the global aggregate state
+	//! The lock for updating the global aggoregate state
+	//! Also used to update the global sample when percentages are used
 	mutex lock;
 	//! The reservoir sample
 	unique_ptr<BlockingSample> sample;
@@ -36,24 +38,47 @@ unique_ptr<GlobalSinkState> PhysicalReservoirSample::GetGlobalSinkState(ClientCo
 
 SinkResultType PhysicalReservoirSample::Sink(ExecutionContext &context, DataChunk &chunk,
                                              OperatorSinkInput &input) const {
-	auto &gstate = input.global_state.Cast<SampleGlobalSinkState>();
-	if (!gstate.sample) {
-		return SinkResultType::FINISHED;
+	auto &global_state = input.global_state.Cast<SampleGlobalSinkState>();
+	// Percentage only has a global sample.
+	lock_guard<mutex> glock(global_state.lock);
+	if (!global_state.sample) {
+		// always gather full thread percentage
+		auto &allocator = Allocator::Get(context.client);
+		if (options->is_percentage) {
+			double percentage = options->sample_size.GetValue<double>();
+			if (percentage == 0) {
+				return SinkResultType::FINISHED;
+			}
+			global_state.sample = make_uniq<ReservoirSamplePercentage>(allocator, percentage, options->seed);
+		} else {
+			idx_t num_samples = options->sample_size.GetValue<idx_t>();
+			if (num_samples == 0) {
+				return SinkResultType::FINISHED;
+			}
+			global_state.sample = make_uniq<ReservoirSample>(allocator, num_samples, options->seed);
+		}
 	}
-	// we implement reservoir sampling without replacement and exponential jumps here
-	// the algorithm is adopted from the paper Weighted random sampling with a reservoir by Pavlos S. Efraimidis et al.
-	// note that the original algorithm is about weighted sampling; this is a simplified approach for uniform sampling
-	lock_guard<mutex> glock(gstate.lock);
-	gstate.sample->AddToReservoir(chunk);
+	global_state.sample->AddToReservoir(chunk);
 	return SinkResultType::NEED_MORE_INPUT;
 }
 
+SinkCombineResultType PhysicalReservoirSample::Combine(ExecutionContext &context,
+                                                       OperatorSinkCombineInput &input) const {
+	return SinkCombineResultType::FINISHED;
+}
+
+SinkFinalizeType PhysicalReservoirSample::Finalize(Pipeline &pipeline, Event &event, ClientContext &context,
+                                                   OperatorSinkFinalizeInput &input) const {
+	return SinkFinalizeType::READY;
+}
+
 //===--------------------------------------------------------------------===//
 // Source
 //===--------------------------------------------------------------------===//
 SourceResultType PhysicalReservoirSample::GetData(ExecutionContext &context, DataChunk &chunk,
                                                   OperatorSourceInput &input) const {
 	auto &sink = this->sink_state->Cast<SampleGlobalSinkState>();
+	lock_guard<mutex> glock(sink.lock);
 	if (!sink.sample) {
 		return SourceResultType::FINISHED;
 	}