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

package/src/duckdb/src/common/sort/partition_state.cpp

@@ -11,7 +11,7 @@ namespace duckdb {
 
 PartitionGlobalHashGroup::PartitionGlobalHashGroup(BufferManager &buffer_manager, const Orders &partitions,
                                                    const Orders &orders, const Types &payload_types, bool external)
-    : count(0) {
+    : count(0), batch_base(0) {
 
 	RowLayout payload_layout;
 	payload_layout.Initialize(payload_types);
@@ -191,52 +191,45 @@ void PartitionGlobalSinkState::CombineLocalPartition(GroupingPartition &local_pa
 	grouping_data->Combine(*local_partition);
 }
 
-void PartitionGlobalSinkState::BuildSortState(TupleDataCollection &group_data, GlobalSortState &global_sort) const {
+PartitionLocalMergeState::PartitionLocalMergeState(PartitionGlobalSinkState &gstate)
+    : merge_state(nullptr), stage(PartitionSortStage::INIT), finished(true), executor(gstate.context) {
+
 	//	 Set up the sort expression computation.
 	vector<LogicalType> sort_types;
-	ExpressionExecutor executor(context);
-	for (auto &order : orders) {
+	for (auto &order : gstate.orders) {
 		auto &oexpr = order.expression;
 		sort_types.emplace_back(oexpr->return_type);
 		executor.AddExpression(*oexpr);
 	}
-	DataChunk sort_chunk;
-	sort_chunk.Initialize(allocator, sort_types);
+	sort_chunk.Initialize(gstate.allocator, sort_types);
+	payload_chunk.Initialize(gstate.allocator, gstate.payload_types);
+}
 
+void PartitionLocalMergeState::Scan() {
+	auto &group_data = *merge_state->group_data;
+	auto &hash_group = *merge_state->hash_group;
+	auto &chunk_state = merge_state->chunk_state;
 	// Copy the data from the group into the sort code.
+	auto &global_sort = *hash_group.global_sort;
 	LocalSortState local_sort;
 	local_sort.Initialize(global_sort, global_sort.buffer_manager);
 
-	//	Strip hash column
-	DataChunk payload_chunk;
-	payload_chunk.Initialize(allocator, payload_types);
-
-	vector<column_t> column_ids;
-	column_ids.reserve(payload_types.size());
-	for (column_t i = 0; i < payload_types.size(); ++i) {
-		column_ids.emplace_back(i);
-	}
-	TupleDataScanState chunk_state;
-	group_data.InitializeScan(chunk_state, column_ids);
-	while (group_data.Scan(chunk_state, payload_chunk)) {
+	TupleDataScanState local_scan;
+	group_data.InitializeScan(local_scan, merge_state->column_ids);
+	while (group_data.Scan(chunk_state, local_scan, payload_chunk)) {
 		sort_chunk.Reset();
 		executor.Execute(payload_chunk, sort_chunk);
 
 		local_sort.SinkChunk(sort_chunk, payload_chunk);
-		if (local_sort.SizeInBytes() > memory_per_thread) {
+		if (local_sort.SizeInBytes() > merge_state->memory_per_thread) {
 			local_sort.Sort(global_sort, true);
 		}
+		hash_group.count += payload_chunk.size();
 	}
 
 	global_sort.AddLocalState(local_sort);
 }
 
-void PartitionGlobalSinkState::BuildSortState(TupleDataCollection &group_data, PartitionGlobalHashGroup &hash_group) {
-	BuildSortState(group_data, *hash_group.global_sort);
-
-	hash_group.count += group_data.Count();
-}
-
 //	Per-thread sink state
 PartitionLocalSinkState::PartitionLocalSinkState(ClientContext &context, PartitionGlobalSinkState &gstate_p)
     : gstate(gstate_p), allocator(Allocator::Get(context)), executor(context) {
@@ -349,10 +342,11 @@ void PartitionLocalSinkState::Combine() {
 	gstate.CombineLocalPartition(local_partition, local_append);
 }
 
-PartitionGlobalMergeState::PartitionGlobalMergeState(PartitionGlobalSinkState &sink, GroupDataPtr group_data,
+PartitionGlobalMergeState::PartitionGlobalMergeState(PartitionGlobalSinkState &sink, GroupDataPtr group_data_p,
                                                      hash_t hash_bin)
-    : sink(sink), group_data(std::move(group_data)), stage(PartitionSortStage::INIT), total_tasks(0), tasks_assigned(0),
-      tasks_completed(0) {
+    : sink(sink), group_data(std::move(group_data_p)), memory_per_thread(sink.memory_per_thread),
+      num_threads(TaskScheduler::GetScheduler(sink.context).NumberOfThreads()), stage(PartitionSortStage::INIT),
+      total_tasks(0), tasks_assigned(0), tasks_completed(0) {
 
 	const auto group_idx = sink.hash_groups.size();
 	auto new_group = make_uniq<PartitionGlobalHashGroup>(sink.buffer_manager, sink.partitions, sink.orders,
@@ -363,13 +357,18 @@ PartitionGlobalMergeState::PartitionGlobalMergeState(PartitionGlobalSinkState &s
 	global_sort = sink.hash_groups[group_idx]->global_sort.get();
 
 	sink.bin_groups[hash_bin] = group_idx;
+
+	column_ids.reserve(sink.payload_types.size());
+	for (column_t i = 0; i < sink.payload_types.size(); ++i) {
+		column_ids.emplace_back(i);
+	}
+	group_data->InitializeScan(chunk_state, column_ids);
 }
 
 void PartitionLocalMergeState::Prepare() {
-	auto &global_sort = *merge_state->global_sort;
-	merge_state->sink.BuildSortState(*merge_state->group_data, *merge_state->hash_group);
 	merge_state->group_data.reset();
 
+	auto &global_sort = *merge_state->global_sort;
 	global_sort.PrepareMergePhase();
 }
 
@@ -381,6 +380,9 @@ void PartitionLocalMergeState::Merge() {
 
 void PartitionLocalMergeState::ExecuteTask() {
 	switch (stage) {
+	case PartitionSortStage::SCAN:
+		Scan();
+		break;
 	case PartitionSortStage::PREPARE:
 		Prepare();
 		break;
@@ -427,6 +429,11 @@ bool PartitionGlobalMergeState::TryPrepareNextStage() {
 
 	switch (stage) {
 	case PartitionSortStage::INIT:
+		total_tasks = num_threads;
+		stage = PartitionSortStage::SCAN;
+		return true;
+
+	case PartitionSortStage::SCAN:
 		total_tasks = 1;
 		stage = PartitionSortStage::PREPARE;
 		return true;
@@ -474,8 +481,9 @@ PartitionGlobalMergeStates::PartitionGlobalMergeStates(PartitionGlobalSinkState
 
 class PartitionMergeTask : public ExecutorTask {
 public:
-	PartitionMergeTask(shared_ptr<Event> event_p, ClientContext &context_p, PartitionGlobalMergeStates &hash_groups_p)
-	    : ExecutorTask(context_p), event(std::move(event_p)), hash_groups(hash_groups_p) {
+	PartitionMergeTask(shared_ptr<Event> event_p, ClientContext &context_p, PartitionGlobalMergeStates &hash_groups_p,
+	                   PartitionGlobalSinkState &gstate)
+	    : ExecutorTask(context_p), event(std::move(event_p)), local_state(gstate), hash_groups(hash_groups_p) {
 	}
 
 	TaskExecutionResult ExecuteTask(TaskExecutionMode mode) override;
@@ -576,7 +584,7 @@ void PartitionMergeEvent::Schedule() {
 
 	vector<shared_ptr<Task>> merge_tasks;
 	for (idx_t tnum = 0; tnum < num_threads; tnum++) {
-		merge_tasks.emplace_back(make_uniq<PartitionMergeTask>(shared_from_this(), context, merge_states));
+		merge_tasks.emplace_back(make_uniq<PartitionMergeTask>(shared_from_this(), context, merge_states, gstate));
 	}
 	SetTasks(std::move(merge_tasks));
 }

package/src/duckdb/src/common/types/data_chunk.cpp

@@ -1,21 +1,15 @@
 #include "duckdb/common/types/data_chunk.hpp"
 
 #include "duckdb/common/array.hpp"
-#include "duckdb/common/arrow/arrow.hpp"
 #include "duckdb/common/exception.hpp"
 #include "duckdb/common/helper.hpp"
 #include "duckdb/common/printer.hpp"
 #include "duckdb/common/serializer.hpp"
-#include "duckdb/common/to_string.hpp"
-#include "duckdb/common/types/arrow_aux_data.hpp"
-#include "duckdb/common/types/date.hpp"
+#include "duckdb/common/serializer/format_serializer.hpp"
+#include "duckdb/common/serializer/format_deserializer.hpp"
 #include "duckdb/common/types/interval.hpp"
-#include "duckdb/common/types/null_value.hpp"
 #include "duckdb/common/types/sel_cache.hpp"
-#include "duckdb/common/types/timestamp.hpp"
-#include "duckdb/common/types/uuid.hpp"
 #include "duckdb/common/types/vector_cache.hpp"
-#include "duckdb/common/unordered_map.hpp"
 #include "duckdb/common/vector.hpp"
 #include "duckdb/common/vector_operations/vector_operations.hpp"
 #include "duckdb/execution/execution_context.hpp"
@@ -268,6 +262,48 @@ void DataChunk::Deserialize(Deserializer &source) {
 	Verify();
 }
 
+void DataChunk::FormatSerialize(FormatSerializer &serializer) const {
+	// write the count
+	auto row_count = size();
+	serializer.WriteProperty<sel_t>(100, "rows", row_count);
+	auto column_count = ColumnCount();
+
+	// Write the types
+	serializer.WriteList(101, "types", column_count,
+	                     [&](FormatSerializer::List &list, idx_t i) { list.WriteElement(data[i].GetType()); });
+
+	// Write the data
+	serializer.WriteList(102, "columns", column_count, [&](FormatSerializer::List &list, idx_t i) {
+		list.WriteObject([&](FormatSerializer &object) {
+			// Reference the vector to avoid potentially mutating it during serialization
+			Vector serialized_vector(data[i].GetType());
+			serialized_vector.Reference(data[i]);
+			serialized_vector.FormatSerialize(object, row_count);
+		});
+	});
+}
+
+void DataChunk::FormatDeserialize(FormatDeserializer &deserializer) {
+	// read the count
+	auto row_count = deserializer.ReadProperty<sel_t>(100, "rows");
+
+	// Read the types
+	vector<LogicalType> types;
+	deserializer.ReadList(101, "types", [&](FormatDeserializer::List &list, idx_t i) {
+		auto type = list.ReadElement<LogicalType>();
+		types.push_back(type);
+	});
+	Initialize(Allocator::DefaultAllocator(), types);
+
+	// now load the column data
+	SetCardinality(row_count);
+
+	// Read the data
+	deserializer.ReadList(102, "columns", [&](FormatDeserializer::List &list, idx_t i) {
+		list.ReadObject([&](FormatDeserializer &object) { data[i].FormatDeserialize(object, row_count); });
+	});
+}
+
 void DataChunk::Slice(const SelectionVector &sel_vector, idx_t count_p) {
 	this->count = count_p;
 	SelCache merge_cache;

package/src/duckdb/src/common/types/hyperloglog.cpp

@@ -2,6 +2,9 @@
 
 #include "duckdb/common/exception.hpp"
 #include "duckdb/common/field_writer.hpp"
+#include "duckdb/common/serializer/format_serializer.hpp"
+#include "duckdb/common/serializer/format_deserializer.hpp"
+
 #include "hyperloglog.hpp"
 
 namespace duckdb {
@@ -106,6 +109,24 @@ unique_ptr<HyperLogLog> HyperLogLog::Deserialize(FieldReader &reader) {
 	return result;
 }
 
+void HyperLogLog::FormatSerialize(FormatSerializer &serializer) const {
+	serializer.WriteProperty(100, "type", HLLStorageType::UNCOMPRESSED);
+	serializer.WriteProperty(101, "data", GetPtr(), GetSize());
+}
+
+unique_ptr<HyperLogLog> HyperLogLog::FormatDeserialize(FormatDeserializer &deserializer) {
+	auto result = make_uniq<HyperLogLog>();
+	auto storage_type = deserializer.ReadProperty<HLLStorageType>(100, "type");
+	switch (storage_type) {
+	case HLLStorageType::UNCOMPRESSED:
+		deserializer.ReadProperty(101, "data", result->GetPtr(), GetSize());
+		break;
+	default:
+		throw SerializationException("Unknown HyperLogLog storage type!");
+	}
+	return result;
+}
+
 //===--------------------------------------------------------------------===//
 // Vectorized HLL implementation
 //===--------------------------------------------------------------------===//

package/src/duckdb/src/common/types/interval.cpp

@@ -11,6 +11,9 @@
 #include "duckdb/common/operator/subtract.hpp"
 #include "duckdb/common/string_util.hpp"
 
+#include "duckdb/common/serializer/format_serializer.hpp"
+#include "duckdb/common/serializer/format_deserializer.hpp"
+
 namespace duckdb {
 
 bool Interval::FromString(const string &str, interval_t &result) {