duckdb 0.8.2-dev3458.0 → 0.8.2-dev3949.0

package/src/duckdb/src/common/types/row/partitioned_tuple_data.cpp

@@ -8,7 +8,7 @@ namespace duckdb {
 
 PartitionedTupleData::PartitionedTupleData(PartitionedTupleDataType type_p, BufferManager &buffer_manager_p,
                                            const TupleDataLayout &layout_p)
-    : type(type_p), buffer_manager(buffer_manager_p), layout(layout_p.Copy()),
+    : type(type_p), buffer_manager(buffer_manager_p), layout(layout_p.Copy()), count(0), data_size(0),
       allocators(make_shared<PartitionTupleDataAllocators>()) {
 }
 
@@ -19,6 +19,10 @@ PartitionedTupleData::PartitionedTupleData(const PartitionedTupleData &other)
 PartitionedTupleData::~PartitionedTupleData() {
 }
 
+const TupleDataLayout &PartitionedTupleData::GetLayout() const {
+	return layout;
+}
+
 PartitionedTupleDataType PartitionedTupleData::GetType() const {
 	return type;
 }
@@ -26,6 +30,7 @@ PartitionedTupleDataType PartitionedTupleData::GetType() const {
 void PartitionedTupleData::InitializeAppendState(PartitionedTupleDataAppendState &state,
                                                  TupleDataPinProperties properties) const {
 	state.partition_sel.Initialize();
+	state.reverse_partition_sel.Initialize();
 
 	vector<column_t> column_ids;
 	column_ids.reserve(layout.ColumnCount());
@@ -36,97 +41,178 @@ void PartitionedTupleData::InitializeAppendState(PartitionedTupleDataAppendState
 	InitializeAppendStateInternal(state, properties);
 }
 
-void PartitionedTupleData::Append(PartitionedTupleDataAppendState &state, DataChunk &input) {
+void PartitionedTupleData::Append(PartitionedTupleDataAppendState &state, DataChunk &input,
+                                  const SelectionVector &append_sel, const idx_t append_count) {
+	TupleDataCollection::ToUnifiedFormat(state.chunk_state, input);
+	AppendUnified(state, input, append_sel, append_count);
+}
+
+bool PartitionedTupleData::UseFixedSizeMap() const {
+	return MaxPartitionIndex() < PartitionedTupleDataAppendState::MAP_THRESHOLD;
+}
+
+void PartitionedTupleData::AppendUnified(PartitionedTupleDataAppendState &state, DataChunk &input,
+                                         const SelectionVector &append_sel, const idx_t append_count) {
+	const idx_t actual_append_count = append_count == DConstants::INVALID_INDEX ? input.size() : append_count;
+
 	// Compute partition indices and store them in state.partition_indices
 	ComputePartitionIndices(state, input);
 
 	// Build the selection vector for the partitions
-	BuildPartitionSel(state, input.size());
+	BuildPartitionSel(state, append_sel, actual_append_count);
 
 	// Early out: check if everything belongs to a single partition
-	const auto &partition_entries = state.partition_entries;
-	if (partition_entries.size() == 1) {
-		const auto &partition_index = partition_entries.begin()->first;
-		auto &partition = *partitions[partition_index];
-		auto &partition_pin_state = *state.partition_pin_states[partition_index];
-		partition.Append(partition_pin_state, state.chunk_state, input);
-		return;
+	optional_idx partition_index;
+	if (UseFixedSizeMap()) {
+		if (state.fixed_partition_entries.size() == 1) {
+			partition_index = state.fixed_partition_entries.begin().GetKey();
+		}
+	} else {
+		if (state.partition_entries.size() == 1) {
+			partition_index = state.partition_entries.begin()->first;
+		}
 	}
+	if (partition_index.IsValid()) {
+		auto &partition = *partitions[partition_index.GetIndex()];
+		auto &partition_pin_state = *state.partition_pin_states[partition_index.GetIndex()];
 
-	TupleDataCollection::ToUnifiedFormat(state.chunk_state, input);
+		const auto size_before = partition.SizeInBytes();
+		partition.AppendUnified(partition_pin_state, state.chunk_state, input, append_sel, actual_append_count);
+		data_size += partition.SizeInBytes() - size_before;
+	} else {
+		// Compute the heap sizes for the whole chunk
+		if (!layout.AllConstant()) {
+			TupleDataCollection::ComputeHeapSizes(state.chunk_state, input, state.partition_sel, actual_append_count);
+		}
 
-	// Compute the heap sizes for the whole chunk
-	if (!layout.AllConstant()) {
-		TupleDataCollection::ComputeHeapSizes(state.chunk_state, input, state.partition_sel, input.size());
-	}
+		// Build the buffer space
+		BuildBufferSpace(state);
 
-	// Build the buffer space
-	BuildBufferSpace(state);
+		// Now scatter everything in one go
+		partitions[0]->Scatter(state.chunk_state, input, state.partition_sel, actual_append_count);
+	}
 
-	// Now scatter everything in one go
-	partitions[0]->Scatter(state.chunk_state, input, state.partition_sel, input.size());
+	count += actual_append_count;
+	Verify();
 }
 
-void PartitionedTupleData::Append(PartitionedTupleDataAppendState &state, TupleDataChunkState &input, idx_t count) {
+void PartitionedTupleData::Append(PartitionedTupleDataAppendState &state, TupleDataChunkState &input,
+                                  const idx_t append_count) {
 	// Compute partition indices and store them in state.partition_indices
-	ComputePartitionIndices(input.row_locations, count, state.partition_indices);
+	ComputePartitionIndices(input.row_locations, append_count, state.partition_indices);
 
 	// Build the selection vector for the partitions
-	BuildPartitionSel(state, count);
+	BuildPartitionSel(state, *FlatVector::IncrementalSelectionVector(), append_count);
 
 	// Early out: check if everything belongs to a single partition
-	auto &partition_entries = state.partition_entries;
-	if (partition_entries.size() == 1) {
-		const auto &partition_index = partition_entries.begin()->first;
-		auto &partition = *partitions[partition_index];
-		auto &partition_pin_state = *state.partition_pin_states[partition_index];
+	optional_idx partition_index;
+	if (UseFixedSizeMap()) {
+		if (state.fixed_partition_entries.size() == 1) {
+			partition_index = state.fixed_partition_entries.begin().GetKey();
+		}
+	} else {
+		if (state.partition_entries.size() == 1) {
+			partition_index = state.partition_entries.begin()->first;
+		}
+	}
+
+	if (partition_index.IsValid()) {
+		auto &partition = *partitions[partition_index.GetIndex()];
+		auto &partition_pin_state = *state.partition_pin_states[partition_index.GetIndex()];
 
 		state.chunk_state.heap_sizes.Reference(input.heap_sizes);
-		partition.Build(partition_pin_state, state.chunk_state, 0, count);
-		partition.CopyRows(state.chunk_state, input, *FlatVector::IncrementalSelectionVector(), count);
-		return;
+
+		const auto size_before = partition.SizeInBytes();
+		partition.Build(partition_pin_state, state.chunk_state, 0, append_count);
+		data_size += partition.SizeInBytes() - size_before;
+
+		partition.CopyRows(state.chunk_state, input, *FlatVector::IncrementalSelectionVector(), append_count);
+	} else {
+		// Build the buffer space
+		state.chunk_state.heap_sizes.Slice(input.heap_sizes, state.partition_sel, append_count);
+		state.chunk_state.heap_sizes.Flatten(append_count);
+		BuildBufferSpace(state);
+
+		// Copy the rows
+		partitions[0]->CopyRows(state.chunk_state, input, state.partition_sel, append_count);
+	}
+
+	count += append_count;
+	Verify();
+}
+
+// LCOV_EXCL_START
+template <class MAP_TYPE>
+struct UnorderedMapGetter {
+	static inline const typename MAP_TYPE::key_type &GetKey(typename MAP_TYPE::iterator &iterator) {
+		return iterator->first;
 	}
 
-	// Build the buffer space
-	state.chunk_state.heap_sizes.Slice(input.heap_sizes, state.partition_sel, count);
-	state.chunk_state.heap_sizes.Flatten(count);
-	BuildBufferSpace(state);
+	static inline const typename MAP_TYPE::key_type &GetKey(const typename MAP_TYPE::const_iterator &iterator) {
+		return iterator->first;
+	}
+
+	static inline typename MAP_TYPE::mapped_type &GetValue(typename MAP_TYPE::iterator &iterator) {
+		return iterator->second;
+	}
+
+	static inline const typename MAP_TYPE::mapped_type &GetValue(const typename MAP_TYPE::const_iterator &iterator) {
+		return iterator->second;
+	}
+};
+
+template <class T>
+struct FixedSizeMapGetter {
+	static inline const idx_t &GetKey(fixed_size_map_iterator_t<T> &iterator) {
+		return iterator.GetKey();
+	}
 
-	// Copy the rows
-	partitions[0]->CopyRows(state.chunk_state, input, state.partition_sel, count);
+	static inline const idx_t &GetKey(const const_fixed_size_map_iterator_t<T> &iterator) {
+		return iterator.GetKey();
+	}
+
+	static inline T &GetValue(fixed_size_map_iterator_t<T> &iterator) {
+		return iterator.GetValue();
+	}
+
+	static inline const T &GetValue(const const_fixed_size_map_iterator_t<T> &iterator) {
+		return iterator.GetValue();
+	}
+};
+// LCOV_EXCL_STOP
+
+void PartitionedTupleData::BuildPartitionSel(PartitionedTupleDataAppendState &state, const SelectionVector &append_sel,
+                                             const idx_t append_count) {
+	if (UseFixedSizeMap()) {
+		BuildPartitionSel<fixed_size_map_t<list_entry_t>, FixedSizeMapGetter<list_entry_t>>(
+		    state, state.fixed_partition_entries, append_sel, append_count);
+	} else {
+		BuildPartitionSel<perfect_map_t<list_entry_t>, UnorderedMapGetter<perfect_map_t<list_entry_t>>>(
+		    state, state.partition_entries, append_sel, append_count);
+	}
 }
 
-void PartitionedTupleData::BuildPartitionSel(PartitionedTupleDataAppendState &state, idx_t count) {
+template <class MAP_TYPE, class GETTER>
+void PartitionedTupleData::BuildPartitionSel(PartitionedTupleDataAppendState &state, MAP_TYPE &partition_entries,
+                                             const SelectionVector &append_sel, const idx_t append_count) {
 	const auto partition_indices = FlatVector::GetData<idx_t>(state.partition_indices);
-	auto &partition_entries = state.partition_entries;
-	auto &partition_entries_arr = state.partition_entries_arr;
 	partition_entries.clear();
 
-	const auto max_partition_index = MaxPartitionIndex();
-	const auto use_arr = max_partition_index < PartitionedTupleDataAppendState::MAP_THRESHOLD;
-
 	switch (state.partition_indices.GetVectorType()) {
 	case VectorType::FLAT_VECTOR:
-		if (use_arr) {
-			std::fill_n(partition_entries_arr, max_partition_index + 1, list_entry_t(0, 0));
-			for (idx_t i = 0; i < count; i++) {
-				const auto &partition_index = partition_indices[i];
-				partition_entries_arr[partition_index].length++;
-			}
-		} else {
-			for (idx_t i = 0; i < count; i++) {
-				const auto &partition_index = partition_indices[i];
-				auto partition_entry = partition_entries.find(partition_index);
-				if (partition_entry == partition_entries.end()) {
-					partition_entries.emplace(partition_index, list_entry_t(0, 1));
-				} else {
-					partition_entry->second.length++;
-				}
+		for (idx_t i = 0; i < append_count; i++) {
+			const auto index = append_sel.get_index(i);
+			const auto &partition_index = partition_indices[index];
+			auto partition_entry = partition_entries.find(partition_index);
+			if (partition_entry == partition_entries.end()) {
+				partition_entries[partition_index] = list_entry_t(0, 1);
+			} else {
+				GETTER::GetValue(partition_entry).length++;
 			}
 		}
 		break;
 	case VectorType::CONSTANT_VECTOR:
-		partition_entries[partition_indices[0]] = list_entry_t(0, count);
+		partition_entries[partition_indices[0]] = list_entry_t(0, append_count);
 		break;
 	default:
 		throw InternalException("Unexpected VectorType in PartitionedTupleData::Append");
@@ -134,64 +220,62 @@ void PartitionedTupleData::BuildPartitionSel(PartitionedTupleDataAppendState &st
 
 	// Early out: check if everything belongs to a single partition
 	if (partition_entries.size() == 1) {
+		// This needs to be initialized, even if we go the short path here
+		for (idx_t i = 0; i < append_count; i++) {
+			const auto index = append_sel.get_index(i);
+			state.reverse_partition_sel[index] = i;
+		}
 		return;
 	}
 
 	// Compute offsets from the counts
 	idx_t offset = 0;
-	if (use_arr) {
-		for (idx_t partition_index = 0; partition_index <= max_partition_index; partition_index++) {
-			auto &partition_entry = partition_entries_arr[partition_index];
-			partition_entry.offset = offset;
-			offset += partition_entry.length;
-		}
-	} else {
-		for (auto &pc : partition_entries) {
-			auto &partition_entry = pc.second;
-			partition_entry.offset = offset;
-			offset += partition_entry.length;
-		}
+	for (auto it = partition_entries.begin(); it != partition_entries.end(); ++it) {
+		auto &partition_entry = GETTER::GetValue(it);
+		partition_entry.offset = offset;
+		offset += partition_entry.length;
 	}
 
 	// Now initialize a single selection vector that acts as a selection vector for every partition
-	auto &all_partitions_sel = state.partition_sel;
-	if (use_arr) {
-		for (idx_t i = 0; i < count; i++) {
-			const auto &partition_index = partition_indices[i];
-			auto &partition_offset = partition_entries_arr[partition_index].offset;
-			all_partitions_sel[partition_offset++] = i;
-		}
-		// Now just add it to the map anyway so the rest of the functionality is shared
-		for (idx_t partition_index = 0; partition_index <= max_partition_index; partition_index++) {
-			const auto &partition_entry = partition_entries_arr[partition_index];
-			if (partition_entry.length != 0) {
-				partition_entries.emplace(partition_index, partition_entry);
-			}
-		}
-	} else {
-		for (idx_t i = 0; i < count; i++) {
-			const auto &partition_index = partition_indices[i];
-			auto &partition_offset = partition_entries[partition_index].offset;
-			all_partitions_sel[partition_offset++] = i;
-		}
+	auto &partition_sel = state.partition_sel;
+	auto &reverse_partition_sel = state.reverse_partition_sel;
+	for (idx_t i = 0; i < append_count; i++) {
+		const auto index = append_sel.get_index(i);
+		const auto &partition_index = partition_indices[index];
+		auto &partition_offset = partition_entries[partition_index].offset;
+		reverse_partition_sel[index] = partition_offset;
+		partition_sel[partition_offset++] = index;
 	}
 }
 
 void PartitionedTupleData::BuildBufferSpace(PartitionedTupleDataAppendState &state) {
-	for (auto &pc : state.partition_entries) {
-		const auto &partition_index = pc.first;
+	if (UseFixedSizeMap()) {
+		BuildBufferSpace<fixed_size_map_t<list_entry_t>, FixedSizeMapGetter<list_entry_t>>(
+		    state, state.fixed_partition_entries);
+	} else {
+		BuildBufferSpace<perfect_map_t<list_entry_t>, UnorderedMapGetter<perfect_map_t<list_entry_t>>>(
+		    state, state.partition_entries);
+	}
+}
+
+template <class MAP_TYPE, class GETTER>
+void PartitionedTupleData::BuildBufferSpace(PartitionedTupleDataAppendState &state, const MAP_TYPE &partition_entries) {
+	for (auto it = partition_entries.begin(); it != partition_entries.end(); ++it) {
+		const auto &partition_index = GETTER::GetKey(it);
 
 		// Partition, pin state for this partition index
 		auto &partition = *partitions[partition_index];
 		auto &partition_pin_state = *state.partition_pin_states[partition_index];
 
 		// Length and offset for this partition
-		const auto &partition_entry = pc.second;
+		const auto &partition_entry = GETTER::GetValue(it);
 		const auto &partition_length = partition_entry.length;
 		const auto partition_offset = partition_entry.offset - partition_length;
 
 		// Build out the buffer space for this partition
+		const auto size_before = partition.SizeInBytes();
 		partition.Build(partition_pin_state, state.chunk_state, partition_offset, partition_length);
+		data_size += partition.SizeInBytes() - size_before;
 	}
 }
 
@@ -210,7 +294,6 @@ void PartitionedTupleData::Combine(PartitionedTupleData &other) {
 
 	// Now combine the state's partitions into this
 	lock_guard<mutex> guard(lock);
-
 	if (partitions.empty()) {
 		// This is the first merge, we just copy them over
 		partitions = std::move(other.partitions);
@@ -221,40 +304,28 @@ void PartitionedTupleData::Combine(PartitionedTupleData &other) {
 			partitions[i]->Combine(*other.partitions[i]);
 		}
 	}
+	this->count += other.count;
+	this->data_size += other.data_size;
+	Verify();
 }
 
-void PartitionedTupleData::Partition(TupleDataCollection &source, TupleDataPinProperties properties) {
-	if (source.Count() == 0) {
-		return;
-	}
-#ifdef DEBUG
-	const auto count_before = source.Count();
-#endif
-
-	PartitionedTupleDataAppendState append_state;
-	InitializeAppendState(append_state, properties);
-
-	TupleDataChunkIterator iterator(source, TupleDataPinProperties::DESTROY_AFTER_DONE, true);
-	auto &chunk_state = iterator.GetChunkState();
-	do {
-		Append(append_state, chunk_state, iterator.GetCurrentChunkCount());
-	} while (iterator.Next());
-
-	FlushAppendState(append_state);
-	source.Reset();
-
-#ifdef DEBUG
-	idx_t count_after = 0;
-	for (const auto &partition : partitions) {
-		count_after += partition->Count();
+void PartitionedTupleData::Reset() {
+	for (auto &partition : partitions) {
+		partition->Reset();
 	}
-	D_ASSERT(count_before == count_after);
-#endif
+	this->count = 0;
+	this->data_size = 0;
+	Verify();
 }
 
 void PartitionedTupleData::Repartition(PartitionedTupleData &new_partitioned_data) {
 	D_ASSERT(layout.GetTypes() == new_partitioned_data.layout.GetTypes());
 
+	if (partitions.size() == new_partitioned_data.partitions.size()) {
+		new_partitioned_data.Combine(*this);
+		return;
+	}
+
 	PartitionedTupleDataAppendState append_state;
 	new_partitioned_data.InitializeAppendState(append_state);
 
@@ -279,20 +350,42 @@ void PartitionedTupleData::Repartition(PartitionedTupleData &new_partitioned_dat
 		}
 		partitions[actual_partition_idx]->Reset();
 	}
-
 	new_partitioned_data.FlushAppendState(append_state);
+
+	count = 0;
+	data_size = 0;
+
+	Verify();
+}
+
+void PartitionedTupleData::Unpin() {
+	for (auto &partition : partitions) {
+		partition->Unpin();
+	}
 }
 
 vector<unique_ptr<TupleDataCollection>> &PartitionedTupleData::GetPartitions() {
 	return partitions;
 }
 
-idx_t PartitionedTupleData::Count() const {
-	idx_t total_count = 0;
-	for (auto &partition : partitions) {
-		total_count += partition->Count();
+unique_ptr<TupleDataCollection> PartitionedTupleData::GetUnpartitioned() {
+	auto data_collection = std::move(partitions[0]);
+	partitions[0] = make_uniq<TupleDataCollection>(buffer_manager, layout);
+
+	for (idx_t i = 1; i < partitions.size(); i++) {
+		data_collection->Combine(*partitions[i]);
 	}
-	return total_count;
+	count = 0;
+	data_size = 0;
+
+	data_collection->Verify();
+	Verify();
+
+	return data_collection;
+}
+
+idx_t PartitionedTupleData::Count() const {
+	return count;
 }
 
 idx_t PartitionedTupleData::SizeInBytes() const {
@@ -303,6 +396,39 @@ idx_t PartitionedTupleData::SizeInBytes() const {
 	return total_size;
 }
 
+idx_t PartitionedTupleData::PartitionCount() const {
+	return partitions.size();
+}
+
+void PartitionedTupleData::Verify() const {
+#ifdef DEBUG
+	idx_t total_count = 0;
+	idx_t total_size = 0;
+	for (auto &partition : partitions) {
+		partition->Verify();
+		total_count += partition->Count();
+		total_size += partition->SizeInBytes();
+	}
+	D_ASSERT(total_count == this->count);
+	D_ASSERT(total_size == this->data_size);
+#endif
+}
+
+// LCOV_EXCL_START
+string PartitionedTupleData::ToString() {
+	string result =
+	    StringUtil::Format("PartitionedTupleData - [%llu Partitions, %llu Rows]\n", partitions.size(), Count());
+	for (idx_t partition_idx = 0; partition_idx < partitions.size(); partition_idx++) {
+		result += StringUtil::Format("Partition %llu: ", partition_idx) + partitions[partition_idx]->ToString();
+	}
+	return result;
+}
+
+void PartitionedTupleData::Print() {
+	Printer::Print(ToString());
+}
+// LCOV_EXCL_STOP
+
 void PartitionedTupleData::CreateAllocator() {
 	allocators->allocators.emplace_back(make_shared<TupleDataAllocator>(buffer_manager, layout));
 }

package/src/duckdb/src/common/types/row/row_layout.cpp

@@ -12,11 +12,10 @@
 
 namespace duckdb {
 
-RowLayout::RowLayout()
-    : flag_width(0), data_width(0), aggr_width(0), row_width(0), all_constant(true), heap_pointer_offset(0) {
+RowLayout::RowLayout() : flag_width(0), data_width(0), row_width(0), all_constant(true), heap_pointer_offset(0) {
 }
 
-void RowLayout::Initialize(vector<LogicalType> types_p, Aggregates aggregates_p, bool align) {
+void RowLayout::Initialize(vector<LogicalType> types_p, bool align) {
 	offsets.clear();
 	types = std::move(types_p);
 
@@ -31,7 +30,7 @@ void RowLayout::Initialize(vector<LogicalType> types_p, Aggregates aggregates_p,
 
 	// This enables pointer swizzling for out-of-core computation.
 	if (!all_constant) {
-		// When unswizzled the pointer lives here.
+		// When unswizzled, the pointer lives here.
 		// When swizzled, the pointer is replaced by an offset.
 		heap_pointer_offset = row_width;
 		// The 8 byte pointer will be replaced with an 8 byte idx_t when swizzled.
@@ -52,39 +51,12 @@ void RowLayout::Initialize(vector<LogicalType> types_p, Aggregates aggregates_p,
 		}
 	}
 
-	// Alignment padding for aggregates
-#ifndef DUCKDB_ALLOW_UNDEFINED
-	if (align) {
-		row_width = AlignValue(row_width);
-	}
-#endif
 	data_width = row_width - flag_width;
 
-	// Aggregate fields.
-	aggregates = std::move(aggregates_p);
-	for (auto &aggregate : aggregates) {
-		offsets.push_back(row_width);
-		row_width += aggregate.payload_size;
-#ifndef DUCKDB_ALLOW_UNDEFINED
-		D_ASSERT(aggregate.payload_size == AlignValue(aggregate.payload_size));
-#endif
-	}
-	aggr_width = row_width - data_width - flag_width;
-
 	// Alignment padding for the next row
-#ifndef DUCKDB_ALLOW_UNDEFINED
 	if (align) {
 		row_width = AlignValue(row_width);
 	}
-#endif
-}
-
-void RowLayout::Initialize(vector<LogicalType> types_p, bool align) {
-	Initialize(std::move(types_p), Aggregates(), align);
-}
-
-void RowLayout::Initialize(Aggregates aggregates_p, bool align) {
-	Initialize(vector<LogicalType>(), std::move(aggregates_p), align);
 }
 
 } // namespace duckdb