duckdb · Mytherin · Nov 8, 2022 · Sep 29, 2022 · Sep 29, 2022 · Sep 30, 2022
diff --git a/src/common/radix_partitioning.cpp b/src/common/radix_partitioning.cpp
@@ -1,10 +1,12 @@
 #include "duckdb/common/radix_partitioning.hpp"
 
 #include "duckdb/common/row_operations/row_operations.hpp"
+#include "duckdb/common/types/partitioned_column_data.hpp"
 #include "duckdb/common/types/row_data_collection.hpp"
 #include "duckdb/common/types/row_layout.hpp"
 #include "duckdb/common/types/vector.hpp"
 #include "duckdb/common/vector_operations/binary_executor.hpp"
+#include "duckdb/common/vector_operations/unary_executor.hpp"
 
 namespace duckdb {
 
@@ -95,6 +97,32 @@ RETURN_TYPE DoubleRadixBitsSwitch1(idx_t radix_bits_1, idx_t radix_bits_2, ARGS
 	}
 }
 
+template <idx_t radix_bits>
+struct RadixLessThan {
+	static inline bool Operation(hash_t hash, hash_t cutoff) {
+		using CONSTANTS = RadixPartitioningConstants<radix_bits>;
+		return CONSTANTS::ApplyMask(hash) < cutoff;
+	}
+};
+
+struct SelectFunctor {
+	template <idx_t radix_bits>
+	static idx_t Operation(Vector &hashes, const SelectionVector *sel, idx_t count, idx_t cutoff,
+	                       SelectionVector *true_sel, SelectionVector *false_sel) {
+		Vector cutoff_vector(Value::HASH(cutoff));
+		return BinaryExecutor::Select<hash_t, hash_t, RadixLessThan<radix_bits>>(hashes, cutoff_vector, sel, count,
+		                                                                         true_sel, false_sel);
+	}
+};
+
+idx_t RadixPartitioning::Select(Vector &hashes, const SelectionVector *sel, idx_t count, idx_t radix_bits, idx_t cutoff,
+                                SelectionVector *true_sel, SelectionVector *false_sel) {
+	return RadixBitsSwitch<SelectFunctor, idx_t>(radix_bits, hashes, sel, count, cutoff, true_sel, false_sel);
+}
+
+//===--------------------------------------------------------------------===//
+// Row Data Partitioning
+//===--------------------------------------------------------------------===//
 template <idx_t radix_bits>
 static void InitPartitions(BufferManager &buffer_manager, vector<unique_ptr<RowDataCollection>> &partition_collections,
                            RowDataBlock *partition_blocks[], vector<BufferHandle> &partition_handles,
@@ -157,7 +185,10 @@ struct PartitionFunctor {
 
 		auto &data_blocks = block_collection.blocks;
 		auto &heap_blocks = string_heap.blocks;
-		for (idx_t block_idx = 0; block_idx < data_blocks.size(); block_idx++) {
+		for (idx_t block_idx_plus_one = data_blocks.size(); block_idx_plus_one > 0; block_idx_plus_one--) {
+			// We loop through blocks in reverse to save some of that PRECIOUS I/O
+			idx_t block_idx = block_idx_plus_one - 1;
+
 			RowDataBlock *data_block;
 			BufferHandle data_handle;
 			data_ptr_t data_ptr;
@@ -275,7 +306,7 @@ struct PartitionFunctor {
 	static inline void FlushTempBuf(data_ptr_t &data_ptr, const idx_t &row_width, uint32_t &block_count,
 	                                const data_ptr_t &tmp_buf, uint32_t &pos, const idx_t count) {
 		pos -= count;
-		memcpy(data_ptr, tmp_buf + pos * row_width, count * row_width);
+		FastMemcpy(data_ptr, tmp_buf + pos * row_width, count * row_width);
 		data_ptr += count * row_width;
 		block_count += count;
 	}
@@ -369,35 +400,69 @@ struct PartitionFunctor {
 	}
 };
 
-void RadixPartitioning::Partition(BufferManager &buffer_manager, const RowLayout &layout, const idx_t hash_offset,
-                                  RowDataCollection &block_collection, RowDataCollection &string_heap,
-                                  vector<unique_ptr<RowDataCollection>> &partition_block_collections,
-                                  vector<unique_ptr<RowDataCollection>> &partition_string_heaps, idx_t radix_bits) {
+void RadixPartitioning::PartitionRowData(BufferManager &buffer_manager, const RowLayout &layout,
+                                         const idx_t hash_offset, RowDataCollection &block_collection,
+                                         RowDataCollection &string_heap,
+                                         vector<unique_ptr<RowDataCollection>> &partition_block_collections,
+                                         vector<unique_ptr<RowDataCollection>> &partition_string_heaps,
+                                         idx_t radix_bits) {
 	return RadixBitsSwitch<PartitionFunctor, void>(radix_bits, buffer_manager, layout, hash_offset, block_collection,
 	                                               string_heap, partition_block_collections, partition_string_heaps);
 }
 
-template <idx_t radix_bits>
-struct RadixLessThan {
-	static inline bool Operation(hash_t hash, hash_t cutoff) {
-		using CONSTANTS = RadixPartitioningConstants<radix_bits>;
-		return CONSTANTS::ApplyMask(hash) < cutoff;
+//===--------------------------------------------------------------------===//
+// Column Data Partitioning
+//===--------------------------------------------------------------------===//
+RadixPartitionedColumnData::RadixPartitionedColumnData(ClientContext &context_p, vector<LogicalType> types_p,
+                                                       idx_t radix_bits_p, idx_t hash_col_idx_p)
+    : PartitionedColumnData(PartitionedColumnDataType::RADIX, context_p, move(types_p)), radix_bits(radix_bits_p),
+      hash_col_idx(hash_col_idx_p) {
+	D_ASSERT(hash_col_idx < types.size());
+	const auto num_partitions = RadixPartitioning::NumberOfPartitions(radix_bits);
+	allocators->allocators.reserve(num_partitions);
+	for (idx_t i = 0; i < num_partitions; i++) {
+		CreateAllocator();
 	}
-};
+	D_ASSERT(allocators->allocators.size() == num_partitions);
+}
 
-struct SelectFunctor {
+RadixPartitionedColumnData::RadixPartitionedColumnData(const RadixPartitionedColumnData &other)
+    : PartitionedColumnData(other), radix_bits(other.radix_bits), hash_col_idx(other.hash_col_idx) {
+	for (idx_t i = 0; i < RadixPartitioning::NumberOfPartitions(radix_bits); i++) {
+		partitions.emplace_back(CreatePartitionCollection(i));
+	}
+}
+
+RadixPartitionedColumnData::~RadixPartitionedColumnData() {
+}
+
+void RadixPartitionedColumnData::InitializeAppendStateInternal(PartitionedColumnDataAppendState &state) const {
+	const auto num_partitions = RadixPartitioning::NumberOfPartitions(radix_bits);
+	state.partition_buffers.reserve(num_partitions);
+	state.partition_append_states.reserve(num_partitions);
+	for (idx_t i = 0; i < num_partitions; i++) {
+		// TODO only initialize the append if partition idx > ...
+		state.partition_append_states.emplace_back(make_unique<ColumnDataAppendState>());
+		partitions[i]->InitializeAppend(*state.partition_append_states[i]);
+		state.partition_buffers.emplace_back(CreatePartitionBuffer());
+	}
+}
+
+struct ComputePartitionIndicesFunctor {
 	template <idx_t radix_bits>
-	static idx_t Operation(Vector &hashes, const SelectionVector *sel, idx_t count, idx_t cutoff,
-	                       SelectionVector *true_sel, SelectionVector *false_sel) {
-		Vector cutoff_vector(Value::HASH(cutoff));
-		return BinaryExecutor::Select<hash_t, hash_t, RadixLessThan<radix_bits>>(hashes, cutoff_vector, sel, count,
-		                                                                         true_sel, false_sel);
+	static void Operation(Vector &hashes, Vector &partition_indices, idx_t count) {
+		UnaryExecutor::Execute<hash_t, hash_t>(hashes, partition_indices, count, [&](hash_t hash) {
+			using CONSTANTS = RadixPartitioningConstants<radix_bits>;
+			return CONSTANTS::ApplyMask(hash);
+		});
 	}
 };
 
-idx_t RadixPartitioning::Select(Vector &hashes, const SelectionVector *sel, idx_t count, idx_t radix_bits, idx_t cutoff,
-                                SelectionVector *true_sel, SelectionVector *false_sel) {
-	return RadixBitsSwitch<SelectFunctor, idx_t>(radix_bits, hashes, sel, count, cutoff, true_sel, false_sel);
+void RadixPartitionedColumnData::ComputePartitionIndices(PartitionedColumnDataAppendState &state, DataChunk &input) {
+	D_ASSERT(partitions.size() == RadixPartitioning::NumberOfPartitions(radix_bits));
+	D_ASSERT(state.partition_buffers.size() == RadixPartitioning::NumberOfPartitions(radix_bits));
+	RadixBitsSwitch<ComputePartitionIndicesFunctor, void>(radix_bits, input.data[hash_col_idx], state.partition_indices,
+	                                                      input.size());
 }
 
 } // namespace duckdb
diff --git a/src/common/symbols.cpp b/src/common/symbols.cpp
@@ -18,6 +18,7 @@
 #include "duckdb/optimizer/join_order_optimizer.hpp"
 #include "duckdb/optimizer/rule.hpp"
 #include "duckdb/parallel/pipeline.hpp"
+#include "duckdb/parallel/meta_pipeline.hpp"
 #include "duckdb/parser/constraint.hpp"
 #include "duckdb/parser/constraints/list.hpp"
 #include "duckdb/parser/expression/list.hpp"
@@ -35,6 +36,8 @@
 #include "duckdb/storage/data_table.hpp"
 #include "duckdb/storage/write_ahead_log.hpp"
 #include "duckdb/transaction/transaction.hpp"
+#include "duckdb/common/types/column_data_collection.hpp"
+#include "duckdb/common/types/column_data_allocator.hpp"
 
 using namespace duckdb;
 
@@ -79,11 +82,16 @@ template class std::unique_ptr<CrossProductRef>;
 template class std::unique_ptr<JoinRef>;
 template class std::unique_ptr<SubqueryRef>;
    template class std::unique_ptr<TableFunctionRef>;
+template class std::shared_ptr<Event>;
 template class std::unique_ptr<Pipeline>;
 template class std::shared_ptr<Pipeline>;
 template class std::weak_ptr<Pipeline>;
+template class std::shared_ptr<MetaPipeline>;
 template class std::unique_ptr<RowGroup>;
 template class std::shared_ptr<RowGroupCollection>;
+template class std::unique_ptr<ColumnDataCollection>;
+template class std::shared_ptr<ColumnDataAllocator>;
+template class std::unique_ptr<PartitionedColumnData>;
 template class std::shared_ptr<PreparedStatementData>;
 template class std::unique_ptr<VacuumInfo>;
 
@@ -127,6 +135,7 @@ template class std::unique_ptr<Vector[]>;
 template class std::unique_ptr<DataChunk>;
 template class std::unique_ptr<JoinHashTable>;
 template class std::unique_ptr<JoinHashTable::ScanStructure>;
+template class std::unique_ptr<JoinHashTable::ProbeSpill>;
 template class std::unique_ptr<data_ptr_t[]>;
 template class std::unique_ptr<Rule>;
 template class std::unique_ptr<LogicalFilter>;
@@ -167,11 +176,16 @@ template class std::vector<PhysicalType>;
 template class std::vector<Value>;
 template class std::vector<int>;
 INSTANTIATE_VECTOR(std::vector<std::unique_ptr<Rule>>)
+INSTANTIATE_VECTOR(std::vector<std::shared_ptr<Event>>)
 INSTANTIATE_VECTOR(std::vector<std::unique_ptr<Pipeline>>)
 INSTANTIATE_VECTOR(std::vector<std::shared_ptr<Pipeline>>)
+INSTANTIATE_VECTOR(std::vector<std::weak_ptr<Pipeline>>)
+INSTANTIATE_VECTOR(std::vector<std::shared_ptr<MetaPipeline>>)
 template class std::vector<std::vector<Expression *>>;
 template class std::vector<LogicalType>;
 INSTANTIATE_VECTOR(std::vector<std::unique_ptr<JoinHashTable>>)
+INSTANTIATE_VECTOR(std::vector<std::unique_ptr<ColumnDataCollection>>)
+INSTANTIATE_VECTOR(std::vector<std::shared_ptr<ColumnDataAllocator>>)
 
 #if !defined(__clang__)
 template struct std::atomic<uint64_t>;

diff --git a/src/common/types/CMakeLists.txt b/src/common/types/CMakeLists.txt
@@ -13,6 +13,7 @@ add_library_unity(
   column_data_allocator.cpp
   column_data_collection.cpp
   column_data_collection_segment.cpp
+  column_data_consumer.cpp
   data_chunk.cpp
   date.cpp
   decimal.cpp
@@ -21,6 +22,7 @@ add_library_unity(
   uuid.cpp
   hyperloglog.cpp
   interval.cpp
+  partitioned_column_data.cpp
   row_data_collection.cpp
   row_data_collection_scanner.cpp
   row_layout.cpp

diff --git a/src/common/types/column_data_allocator.cpp b/src/common/types/column_data_allocator.cpp
@@ -1,6 +1,7 @@
 #include "duckdb/common/types/column_data_allocator.hpp"
-#include "duckdb/storage/buffer_manager.hpp"
+
 #include "duckdb/common/types/column_data_collection_segment.hpp"
+#include "duckdb/storage/buffer_manager.hpp"
 
 namespace duckdb {
 
@@ -29,7 +30,15 @@ ColumnDataAllocator::ColumnDataAllocator(ClientContext &context, ColumnDataAlloc
 
 BufferHandle ColumnDataAllocator::Pin(uint32_t block_id) {
 	D_ASSERT(type == ColumnDataAllocatorType::BUFFER_MANAGER_ALLOCATOR);
-	return alloc.buffer_manager->Pin(blocks[block_id].handle);
+	shared_ptr<BlockHandle> *block_handle;
+	if (shared) {
+		// need to grab handle from the vector within a lock else threadsan will complain
+		lock_guard<mutex> guard(lock);
+		block_handle = &blocks[block_id].handle;
+	} else {
+		block_handle = &blocks[block_id].handle;
+	}
+	return alloc.buffer_manager->Pin(*block_handle);
 }
 
 void ColumnDataAllocator::AllocateBlock() {
@@ -106,9 +115,15 @@ void ColumnDataAllocator::AllocateData(idx_t size, uint32_t &block_id, uint32_t
                                        ChunkManagementState *chunk_state) {
 	switch (type) {
 	case ColumnDataAllocatorType::BUFFER_MANAGER_ALLOCATOR:
-		AllocateBuffer(size, block_id, offset, chunk_state);
+		if (shared) {
+			lock_guard<mutex> guard(lock);
+			AllocateBuffer(size, block_id, offset, chunk_state);
+		} else {
+			AllocateBuffer(size, block_id, offset, chunk_state);
+		}
 		break;
 	case ColumnDataAllocatorType::IN_MEMORY_ALLOCATOR:
+		D_ASSERT(!shared);
 		AllocateMemory(size, block_id, offset, chunk_state);
 		break;
 	default:
@@ -138,6 +153,10 @@ data_ptr_t ColumnDataAllocator::GetDataPointer(ChunkManagementState &state, uint
 	return state.handles[block_id].Ptr() + offset;
 }
 
+void ColumnDataAllocator::DeleteBlock(uint32_t block_id) {
+	blocks[block_id].handle->SetCanDestroy(true);
+}
+
 Allocator &ColumnDataAllocator::GetAllocator() {
 	return type == ColumnDataAllocatorType::IN_MEMORY_ALLOCATOR ? *alloc.allocator
 	                                                            : alloc.buffer_manager->GetBufferAllocator();

diff --git a/src/common/types/column_data_collection.cpp b/src/common/types/column_data_collection.cpp
@@ -812,4 +812,8 @@ bool ColumnDataCollection::ResultEquals(const ColumnDataCollection &left, const
 	return true;
 }
 
+const vector<unique_ptr<ColumnDataCollectionSegment>> &ColumnDataCollection::GetSegments() const {
+	return segments;
+}
+
 } // namespace duckdb
diff --git a/src/common/types/column_data_collection_segment.cpp b/src/common/types/column_data_collection_segment.cpp
@@ -214,7 +214,6 @@ void ColumnDataCollectionSegment::FetchChunk(idx_t chunk_idx, DataChunk &result)
 void ColumnDataCollectionSegment::FetchChunk(idx_t chunk_idx, DataChunk &result, const vector<column_t> &column_ids) {
 	D_ASSERT(chunk_idx < chunk_data.size());
 	ChunkManagementState state;
-	InitializeChunkState(chunk_idx, state);
 	state.properties = ColumnDataScanProperties::DISALLOW_ZERO_COPY;
 	ReadChunk(chunk_idx, state, result, column_ids);
 }