duckdb 0.8.2-dev11.0 → 0.8.2-dev1182.0

package/src/duckdb/src/common/local_file_system.cpp

@@ -700,7 +700,7 @@ void LocalFileSystem::CreateDirectory(const string &directory) {
 	}
 	auto unicode_path = WindowsUtil::UTF8ToUnicode(directory.c_str());
 	if (directory.empty() || !CreateDirectoryW(unicode_path.c_str(), NULL) || !DirectoryExists(directory)) {
-		throw IOException("Could not create directory!");
+		throw IOException("Could not create directory: \'%s\'", directory.c_str());
 	}
 }
 

package/src/duckdb/src/common/multi_file_reader.cpp

@@ -6,6 +6,7 @@
 #include "duckdb/common/exception.hpp"
 #include "duckdb/function/function_set.hpp"
 #include "duckdb/common/hive_partitioning.hpp"
+#include "duckdb/common/types.hpp"
 
 namespace duckdb {
 
@@ -13,6 +14,8 @@ void MultiFileReader::AddParameters(TableFunction &table_function) {
 	table_function.named_parameters["filename"] = LogicalType::BOOLEAN;
 	table_function.named_parameters["hive_partitioning"] = LogicalType::BOOLEAN;
 	table_function.named_parameters["union_by_name"] = LogicalType::BOOLEAN;
+	table_function.named_parameters["hive_types"] = LogicalType::ANY;
+	table_function.named_parameters["hive_types_autocast"] = LogicalType::BOOLEAN;
 }
 
 vector<string> MultiFileReader::GetFileList(ClientContext &context, const Value &input, const string &name,
@@ -49,7 +52,8 @@ vector<string> MultiFileReader::GetFileList(ClientContext &context, const Value
 	return files;
 }
 
-bool MultiFileReader::ParseOption(const string &key, const Value &val, MultiFileReaderOptions &options) {
+bool MultiFileReader::ParseOption(const string &key, const Value &val, MultiFileReaderOptions &options,
+                                  ClientContext &context) {
 	auto loption = StringUtil::Lower(key);
 	if (loption == "filename") {
 		options.filename = BooleanValue::Get(val);
@@ -58,6 +62,28 @@ bool MultiFileReader::ParseOption(const string &key, const Value &val, MultiFile
 		options.auto_detect_hive_partitioning = false;
 	} else if (loption == "union_by_name") {
 		options.union_by_name = BooleanValue::Get(val);
+	} else if (loption == "hive_types_autocast" || loption == "hive_type_autocast") {
+		options.hive_types_autocast = BooleanValue::Get(val);
+	} else if (loption == "hive_types" || loption == "hive_type") {
+		if (val.type().id() != LogicalTypeId::STRUCT) {
+			throw InvalidInputException(
+			    "'hive_types' only accepts a STRUCT('name':VARCHAR, ...), but '%s' was provided",
+			    val.type().ToString());
+		}
+		// verify that that all the children of the struct value are VARCHAR
+		auto &children = StructValue::GetChildren(val);
+		for (idx_t i = 0; i < children.size(); i++) {
+			const Value &child = children[i];
+			if (child.type().id() != LogicalType::VARCHAR) {
+				throw InvalidInputException("hive_types: '%s' must be a VARCHAR, instead: '%s' was provided",
+				                            StructType::GetChildName(val.type(), i), child.type().ToString());
+			}
+			// for every child of the struct, get the logical type
+			LogicalType transformed_type = TransformStringToLogicalType(child.ToString(), context);
+			const string &name = StructType::GetChildName(val.type(), i);
+			options.hive_types_schema[name] = transformed_type;
+		}
+		D_ASSERT(!options.hive_types_schema.empty());
 	} else {
 		return false;
 	}
@@ -111,28 +137,28 @@ MultiFileReaderBindData MultiFileReader::BindOptions(MultiFileReaderOptions &opt
 			auto file_partitions = HivePartitioning::Parse(f);
 			for (auto &part_info : partitions) {
 				if (file_partitions.find(part_info.first) == file_partitions.end()) {
+					string error = "Hive partition mismatch between file \"%s\" and \"%s\": key \"%s\" not found";
 					if (options.auto_detect_hive_partitioning == true) {
-						throw BinderException(
-						    "Hive partitioning was enabled automatically, but an error was encountered: Hive partition "
-						    "mismatch between file \"%s\" and \"%s\": key \"%s\" not found\n\nTo switch off hive "
-						    "partition, set: HIVE_PARTITIONING=0",
-						    files[0], f, part_info.first);
+						throw InternalException(error + "(hive partitioning was autodetected)", files[0], f,
+						                        part_info.first);
 					}
-					throw BinderException(
-					    "Hive partition mismatch between file \"%s\" and \"%s\": key \"%s\" not found", files[0], f,
-					    part_info.first);
+					throw BinderException(error.c_str(), files[0], f, part_info.first);
 				}
 			}
 			if (partitions.size() != file_partitions.size()) {
+				string error_msg = "Hive partition mismatch between file \"%s\" and \"%s\"";
 				if (options.auto_detect_hive_partitioning == true) {
-					throw BinderException("Hive partitioning was enabled automatically, but an error was encountered: "
-					                      "Hive partition mismatch between file \"%s\" and \"%s\"\n\nTo switch off "
-					                      "hive partition, set: HIVE_PARTITIONING=0",
-					                      files[0], f);
+					throw InternalException(error_msg + "(hive partitioning was autodetected)", files[0], f);
 				}
-				throw BinderException("Hive partition mismatch between file \"%s\" and \"%s\"", files[0], f);
+				throw BinderException(error_msg.c_str(), files[0], f);
 			}
 		}
+
+		if (!options.hive_types_schema.empty()) {
+			// verify that all hive_types are existing partitions
+			options.VerifyHiveTypesArePartitions(partitions);
+		}
+
 		for (auto &part : partitions) {
 			idx_t hive_partitioning_index = DConstants::INVALID_INDEX;
 			auto lookup = std::find(names.begin(), names.end(), part.first);
@@ -140,11 +166,11 @@ MultiFileReaderBindData MultiFileReader::BindOptions(MultiFileReaderOptions &opt
 				// hive partitioning column also exists in file - override
 				auto idx = lookup - names.begin();
 				hive_partitioning_index = idx;
-				return_types[idx] = LogicalType::VARCHAR;
+				return_types[idx] = options.GetHiveLogicalType(part.first);
 			} else {
 				// hive partitioning column does not exist in file - add a new column containing the key
 				hive_partitioning_index = names.size();
-				return_types.emplace_back(LogicalType::VARCHAR);
+				return_types.emplace_back(options.GetHiveLogicalType(part.first));
 				names.emplace_back(part.first);
 			}
 			bind_data.hive_partitioning_indexes.emplace_back(part.first, hive_partitioning_index);
@@ -156,7 +182,9 @@ MultiFileReaderBindData MultiFileReader::BindOptions(MultiFileReaderOptions &opt
 void MultiFileReader::FinalizeBind(const MultiFileReaderOptions &file_options, const MultiFileReaderBindData &options,
                                    const string &filename, const vector<string> &local_names,
                                    const vector<LogicalType> &global_types, const vector<string> &global_names,
-                                   const vector<column_t> &global_column_ids, MultiFileReaderData &reader_data) {
+                                   const vector<column_t> &global_column_ids, MultiFileReaderData &reader_data,
+                                   ClientContext &context) {
+
 	// create a map of name -> column index
 	case_insensitive_map_t<idx_t> name_map;
 	if (file_options.union_by_name) {
@@ -183,7 +211,8 @@ void MultiFileReader::FinalizeBind(const MultiFileReaderOptions &file_options, c
 			bool found_partition = false;
 			for (auto &entry : options.hive_partitioning_indexes) {
 				if (column_id == entry.index) {
-					reader_data.constant_map.emplace_back(i, Value(partitions[entry.value]));
+					Value value = file_options.GetHivePartitionValue(partitions[entry.value], entry.value, context);
+					reader_data.constant_map.emplace_back(i, value);
 					found_partition = true;
 					break;
 				}
@@ -313,7 +342,16 @@ void MultiFileReaderOptions::Serialize(Serializer &serializer) const {
 	FieldWriter writer(serializer);
 	writer.WriteField<bool>(filename);
 	writer.WriteField<bool>(hive_partitioning);
+	writer.WriteField<bool>(auto_detect_hive_partitioning);
 	writer.WriteField<bool>(union_by_name);
+	writer.WriteField<bool>(hive_types_autocast);
+	// serialize hive_types_schema
+	const uint32_t schema_size = hive_types_schema.size();
+	writer.WriteField<uint32_t>(schema_size);
+	for (auto &hive_type : hive_types_schema) {
+		writer.WriteString(hive_type.first);
+		writer.WriteString(hive_type.second.ToString());
+	}
 	writer.Finalize();
 }
 
@@ -322,7 +360,16 @@ MultiFileReaderOptions MultiFileReaderOptions::Deserialize(Deserializer &source)
 	FieldReader reader(source);
 	result.filename = reader.ReadRequired<bool>();
 	result.hive_partitioning = reader.ReadRequired<bool>();
+	result.auto_detect_hive_partitioning = reader.ReadRequired<bool>();
 	result.union_by_name = reader.ReadRequired<bool>();
+	result.hive_types_autocast = reader.ReadRequired<bool>();
+	// deserialize hive_types_schema
+	const uint32_t schema_size = reader.ReadRequired<uint32_t>();
+	for (idx_t i = 0; i < schema_size; i++) {
+		const string name = reader.ReadRequired<string>();
+		const LogicalType type = TransformStringToLogicalType(reader.ReadRequired<string>());
+		result.hive_types_schema[name] = type;
+	}
 	reader.Finalize();
 	return result;
 }
@@ -365,7 +412,9 @@ HivePartitioningIndex HivePartitioningIndex::Deserialize(Deserializer &source) {
 void MultiFileReaderOptions::AddBatchInfo(BindInfo &bind_info) const {
 	bind_info.InsertOption("filename", Value::BOOLEAN(filename));
 	bind_info.InsertOption("hive_partitioning", Value::BOOLEAN(hive_partitioning));
+	bind_info.InsertOption("auto_detect_hive_partitioning", Value::BOOLEAN(auto_detect_hive_partitioning));
 	bind_info.InsertOption("union_by_name", Value::BOOLEAN(union_by_name));
+	bind_info.InsertOption("hive_types_autocast", Value::BOOLEAN(hive_types_autocast));
 }
 
 void UnionByName::CombineUnionTypes(const vector<string> &col_names, const vector<LogicalType> &sql_types,
@@ -390,4 +439,118 @@ void UnionByName::CombineUnionTypes(const vector<string> &col_names, const vecto
 	}
 }
 
+bool MultiFileReaderOptions::AutoDetectHivePartitioningInternal(const vector<string> &files) {
+	std::unordered_set<string> partitions;
+
+	auto splits_first_file = StringUtil::Split(files.front(), FileSystem::PathSeparator());
+	if (splits_first_file.size() < 2) {
+		return false;
+	}
+	for (auto it = splits_first_file.begin(); it != splits_first_file.end(); it++) {
+		auto partition = StringUtil::Split(*it, "=");
+		if (partition.size() == 2) {
+			partitions.insert(partition.front());
+		}
+	}
+	if (partitions.empty()) {
+		return false;
+	}
+	for (auto &file : files) {
+		auto splits = StringUtil::Split(file, FileSystem::PathSeparator());
+		if (splits.size() != splits_first_file.size()) {
+			return false;
+		}
+		for (auto it = splits.begin(); it != std::prev(splits.end()); it++) {
+			auto part = StringUtil::Split(*it, "=");
+			if (part.size() != 2) {
+				continue;
+			}
+			if (partitions.find(part.front()) == partitions.end()) {
+				return false;
+			}
+		}
+	}
+	return true;
+}
+void MultiFileReaderOptions::AutoDetectHiveTypesInternal(const string &file, ClientContext &context) {
+	std::map<string, string> partitions;
+	auto splits = StringUtil::Split(file, FileSystem::PathSeparator());
+	if (splits.size() < 2) {
+		return;
+	}
+	for (auto it = splits.begin(); it != std::prev(splits.end()); it++) {
+		auto part = StringUtil::Split(*it, "=");
+		if (part.size() == 2) {
+			partitions[part.front()] = part.back();
+		}
+	}
+	if (partitions.empty()) {
+		return;
+	}
+
+	const LogicalType candidates[] = {LogicalType::DATE, LogicalType::TIMESTAMP, LogicalType::BIGINT};
+	for (auto &part : partitions) {
+		const string &name = part.first;
+		if (hive_types_schema.find(name) != hive_types_schema.end()) {
+			continue;
+		}
+		Value value(part.second);
+		for (auto &candidate : candidates) {
+			const bool success = value.TryCastAs(context, candidate);
+			if (success) {
+				hive_types_schema[name] = candidate;
+				break;
+			}
+		}
+	}
+}
+void MultiFileReaderOptions::AutoDetectHivePartitioning(const vector<string> &files, ClientContext &context) {
+	D_ASSERT(!files.empty());
+	const bool hp_explicitly_disabled = !auto_detect_hive_partitioning && !hive_partitioning;
+	const bool ht_enabled = !hive_types_schema.empty();
+	if (hp_explicitly_disabled && ht_enabled) {
+		throw InvalidInputException("cannot disable hive_partitioning when hive_types is enabled");
+	}
+	if (ht_enabled && auto_detect_hive_partitioning && !hive_partitioning) {
+		// hive_types flag implies hive_partitioning
+		hive_partitioning = true;
+		auto_detect_hive_partitioning = false;
+	}
+	if (auto_detect_hive_partitioning) {
+		hive_partitioning = AutoDetectHivePartitioningInternal(files);
+	}
+	if (hive_partitioning && hive_types_autocast) {
+		AutoDetectHiveTypesInternal(files.front(), context);
+	}
+}
+void MultiFileReaderOptions::VerifyHiveTypesArePartitions(const std::map<string, string> &partitions) const {
+	for (auto &hive_type : hive_types_schema) {
+		if (partitions.find(hive_type.first) == partitions.end()) {
+			throw InvalidInputException("Unknown hive_type: \"%s\" does not appear to be a partition", hive_type.first);
+		}
+	}
+}
+LogicalType MultiFileReaderOptions::GetHiveLogicalType(const string &hive_partition_column) const {
+	if (!hive_types_schema.empty()) {
+		auto it = hive_types_schema.find(hive_partition_column);
+		if (it != hive_types_schema.end()) {
+			return it->second;
+		}
+	}
+	return LogicalType::VARCHAR;
+}
+Value MultiFileReaderOptions::GetHivePartitionValue(const string &base, const string &entry,
+                                                    ClientContext &context) const {
+	Value value(base);
+	auto it = hive_types_schema.find(entry);
+	if (it == hive_types_schema.end()) {
+		return value;
+	}
+	if (!value.TryCastAs(context, it->second)) {
+		throw InvalidInputException("Unable to cast '%s' (from hive partition column '%s') to: '%s'", value.ToString(),
+		                            StringUtil::Upper(it->first), it->second.ToString());
+	}
+	return value;
+}
+
 } // namespace duckdb

package/src/duckdb/src/common/radix_partitioning.cpp

@@ -8,9 +8,26 @@
 
 namespace duckdb {
 
+//! Templated radix partitioning constants, can be templated to the number of radix bits
+template <idx_t radix_bits>
+struct RadixPartitioningConstants {
+public:
+	//! Bitmask of the upper bits starting at the 5th byte
+	static constexpr const idx_t NUM_PARTITIONS = RadixPartitioning::NumberOfPartitions(radix_bits);
+	static constexpr const idx_t SHIFT = RadixPartitioning::Shift(radix_bits);
+	static constexpr const hash_t MASK = RadixPartitioning::Mask(radix_bits);
+
+public:
+	//! Apply bitmask and right shift to get a number between 0 and NUM_PARTITIONS
+	static inline hash_t ApplyMask(hash_t hash) {
+		D_ASSERT((hash & MASK) >> SHIFT < NUM_PARTITIONS);
+		return (hash & MASK) >> SHIFT;
+	}
+};
+
 template <class OP, class RETURN_TYPE, typename... ARGS>
-RETURN_TYPE RadixBitsSwitch(idx_t radix_bits, ARGS &&... args) {
-	D_ASSERT(radix_bits <= sizeof(hash_t) * 8);
+RETURN_TYPE RadixBitsSwitch(idx_t radix_bits, ARGS &&...args) {
+	D_ASSERT(radix_bits <= RadixPartitioning::MAX_RADIX_BITS);
 	switch (radix_bits) {
 	case 1:
 		return OP::template Operation<1>(std::forward<ARGS>(args)...);
@@ -20,7 +37,7 @@ RETURN_TYPE RadixBitsSwitch(idx_t radix_bits, ARGS &&... args) {
 		return OP::template Operation<3>(std::forward<ARGS>(args)...);
 	case 4:
 		return OP::template Operation<4>(std::forward<ARGS>(args)...);
-	case 5:
+	case 5: // LCOV_EXCL_START
 		return OP::template Operation<5>(std::forward<ARGS>(args)...);
 	case 6:
 		return OP::template Operation<6>(std::forward<ARGS>(args)...);
@@ -32,9 +49,14 @@ RETURN_TYPE RadixBitsSwitch(idx_t radix_bits, ARGS &&... args) {
 		return OP::template Operation<9>(std::forward<ARGS>(args)...);
 	case 10:
 		return OP::template Operation<10>(std::forward<ARGS>(args)...);
+	case 11:
+		return OP::template Operation<10>(std::forward<ARGS>(args)...);
+	case 12:
+		return OP::template Operation<10>(std::forward<ARGS>(args)...);
 	default:
-		throw InternalException("TODO");
-	}
+		throw InternalException(
+		    "radix_bits higher than RadixPartitioning::MAX_RADIX_BITS encountered in RadixBitsSwitch");
+	} // LCOV_EXCL_STOP
 }
 
 template <idx_t radix_bits>
@@ -69,10 +91,6 @@ struct HashsToBinsFunctor {
 	}
 };
 
-void RadixPartitioning::HashesToBins(Vector &hashes, idx_t radix_bits, Vector &bins, idx_t count) {
-	return RadixBitsSwitch<HashsToBinsFunctor, void>(radix_bits, hashes, bins, count);
-}
-
 //===--------------------------------------------------------------------===//
 // Row Data Partitioning
 //===--------------------------------------------------------------------===//

package/src/duckdb/src/common/row_operations/row_external.cpp

@@ -99,7 +99,7 @@ void RowOperations::UnswizzleHeapPointer(const RowLayout &layout, const data_ptr
 
 static inline void VerifyUnswizzledString(const RowLayout &layout, const idx_t &col_idx, const data_ptr_t &row_ptr) {
 #ifdef DEBUG
-	if (layout.GetTypes()[col_idx] == LogicalTypeId::BLOB) {
+	if (layout.GetTypes()[col_idx].id() != LogicalTypeId::VARCHAR) {
 		return;
 	}
 	idx_t entry_idx;

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

@@ -273,16 +273,13 @@ void MergeSorter::ComputeMerge(const idx_t &count, bool left_smaller[]) {
 			break;
 		}
 		// Pin the radix sorting data
-		if (!l_done) {
-			left->PinRadix(l.block_idx);
-			l_radix_ptr = left->RadixPtr();
-		}
-		if (!r_done) {
-			right->PinRadix(r.block_idx);
-			r_radix_ptr = right->RadixPtr();
-		}
-		const idx_t &l_count = !l_done ? l_sorted_block.radix_sorting_data[l.block_idx]->count : 0;
-		const idx_t &r_count = !r_done ? r_sorted_block.radix_sorting_data[r.block_idx]->count : 0;
+		left->PinRadix(l.block_idx);
+		l_radix_ptr = left->RadixPtr();
+		right->PinRadix(r.block_idx);
+		r_radix_ptr = right->RadixPtr();
+
+		const idx_t l_count = l_sorted_block.radix_sorting_data[l.block_idx]->count;
+		const idx_t r_count = r_sorted_block.radix_sorting_data[r.block_idx]->count;
 		// Compute the merge
 		if (sort_layout.all_constant) {
 			// All sorting columns are constant size
@@ -298,12 +295,8 @@ void MergeSorter::ComputeMerge(const idx_t &count, bool left_smaller[]) {
 			}
 		} else {
 			// Pin the blob data
-			if (!l_done) {
-				left->PinData(*l_sorted_block.blob_sorting_data);
-			}
-			if (!r_done) {
-				right->PinData(*r_sorted_block.blob_sorting_data);
-			}
+			left->PinData(*l_sorted_block.blob_sorting_data);
+			right->PinData(*r_sorted_block.blob_sorting_data);
 			// Merge with variable size sorting columns
 			for (; compared < count && l.entry_idx < l_count && r.entry_idx < r_count; compared++) {
 				left_smaller[compared] =

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

@@ -87,7 +87,7 @@ PartitionGlobalSinkState::PartitionGlobalSinkState(ClientContext &context,
                                                    const vector<unique_ptr<BaseStatistics>> &partition_stats,
                                                    idx_t estimated_cardinality)
     : context(context), buffer_manager(BufferManager::GetBufferManager(context)), allocator(Allocator::Get(context)),
-      payload_types(payload_types), memory_per_thread(0), count(0) {
+      fixed_bits(0), payload_types(payload_types), memory_per_thread(0), count(0) {
 
 	GenerateOrderings(partitions, orders, partition_bys, order_bys, partition_stats);
 
@@ -102,9 +102,19 @@ PartitionGlobalSinkState::PartitionGlobalSinkState(ClientContext &context,
 	}
 }
 
+void PartitionGlobalSinkState::SyncPartitioning(const PartitionGlobalSinkState &other) {
+	fixed_bits = other.grouping_data ? other.grouping_data->GetRadixBits() : 0;
+
+	const auto old_bits = grouping_data ? grouping_data->GetRadixBits() : 0;
+	if (fixed_bits != old_bits) {
+		const auto hash_col_idx = payload_types.size();
+		grouping_data = make_uniq<RadixPartitionedColumnData>(context, grouping_types, fixed_bits, hash_col_idx);
+	}
+}
+
 void PartitionGlobalSinkState::ResizeGroupingData(idx_t cardinality) {
 	//	Have we started to combine? Then just live with it.
-	if (grouping_data && !grouping_data->GetPartitions().empty()) {
+	if (fixed_bits || (grouping_data && !grouping_data->GetPartitions().empty())) {
 		return;
 	}
 	//	Is the average partition size too large?
@@ -186,9 +196,7 @@ void PartitionGlobalSinkState::CombineLocalPartition(GroupingPartition &local_pa
 	grouping_data->Combine(*local_partition);
 }
 
-void PartitionGlobalSinkState::BuildSortState(ColumnDataCollection &group_data, PartitionGlobalHashGroup &hash_group) {
-	auto &global_sort = *hash_group.global_sort;
-
+void PartitionGlobalSinkState::BuildSortState(ColumnDataCollection &group_data, GlobalSortState &global_sort) const {
 	//	 Set up the sort expression computation.
 	vector<LogicalType> sort_types;
 	ExpressionExecutor executor(context);
@@ -234,6 +242,10 @@ void PartitionGlobalSinkState::BuildSortState(ColumnDataCollection &group_data,
 	}
 
 	global_sort.AddLocalState(local_sort);
+}
+
+void PartitionGlobalSinkState::BuildSortState(ColumnDataCollection &group_data, PartitionGlobalHashGroup &hash_group) {
+	BuildSortState(group_data, *hash_group.global_sort);
 
 	hash_group.count += group_data.Count();
 }
@@ -482,18 +494,29 @@ public:
 	TaskExecutionResult ExecuteTask(TaskExecutionMode mode) override;
 
 private:
+	struct ExecutorCallback : public PartitionGlobalMergeStates::Callback {
+		explicit ExecutorCallback(Executor &executor) : executor(executor) {
+		}
+
+		bool HasError() const override {
+			return executor.HasError();
+		}
+
+		Executor &executor;
+	};
+
 	shared_ptr<Event> event;
 	PartitionLocalMergeState local_state;
 	PartitionGlobalMergeStates &hash_groups;
 };
 
-TaskExecutionResult PartitionMergeTask::ExecuteTask(TaskExecutionMode mode) {
+bool PartitionGlobalMergeStates::ExecuteTask(PartitionLocalMergeState &local_state, Callback &callback) {
 	// Loop until all hash groups are done
 	size_t sorted = 0;
-	while (sorted < hash_groups.states.size()) {
+	while (sorted < states.size()) {
 		// First check if there is an unfinished task for this thread
-		if (executor.HasError()) {
-			return TaskExecutionResult::TASK_ERROR;
+		if (callback.HasError()) {
+			return false;
 		}
 		if (!local_state.TaskFinished()) {
 			local_state.ExecuteTask();
@@ -501,8 +524,8 @@ TaskExecutionResult PartitionMergeTask::ExecuteTask(TaskExecutionMode mode) {
 		}
 
 		// Thread is done with its assigned task, try to fetch new work
-		for (auto group = sorted; group < hash_groups.states.size(); ++group) {
-			auto &global_state = hash_groups.states[group];
+		for (auto group = sorted; group < states.size(); ++group) {
+			auto &global_state = states[group];
 			if (global_state->IsSorted()) {
 				// This hash group is done
 				// Update the high water mark of densely completed groups
@@ -543,6 +566,16 @@ TaskExecutionResult PartitionMergeTask::ExecuteTask(TaskExecutionMode mode) {
 		}
 	}
 
+	return true;
+}
+
+TaskExecutionResult PartitionMergeTask::ExecuteTask(TaskExecutionMode mode) {
+	ExecutorCallback callback(executor);
+
+	if (!hash_groups.ExecuteTask(local_state, callback)) {
+		return TaskExecutionResult::TASK_ERROR;
+	}
+
 	event->FinishTask();
 	return TaskExecutionResult::TASK_FINISHED;
 }

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

@@ -1,11 +1,14 @@
 #include "duckdb/common/types/batched_data_collection.hpp"
+
+#include "duckdb/common/optional_ptr.hpp"
 #include "duckdb/common/printer.hpp"
 #include "duckdb/storage/buffer_manager.hpp"
-#include "duckdb/common/optional_ptr.hpp"
 
 namespace duckdb {
 
-BatchedDataCollection::BatchedDataCollection(vector<LogicalType> types_p) : types(std::move(types_p)) {
+BatchedDataCollection::BatchedDataCollection(ClientContext &context_p, vector<LogicalType> types_p,
+                                             bool buffer_managed_p)
+    : context(context_p), types(std::move(types_p)), buffer_managed(buffer_managed_p) {
 }
 
 void BatchedDataCollection::Append(DataChunk &input, idx_t batch_index) {
@@ -20,6 +23,8 @@ void BatchedDataCollection::Append(DataChunk &input, idx_t batch_index) {
 		unique_ptr<ColumnDataCollection> new_collection;
 		if (last_collection.collection) {
 			new_collection = make_uniq<ColumnDataCollection>(*last_collection.collection);
+		} else if (buffer_managed) {
+			new_collection = make_uniq<ColumnDataCollection>(BufferManager::GetBufferManager(context), types);
 		} else {
 			new_collection = make_uniq<ColumnDataCollection>(Allocator::DefaultAllocator(), types);
 		}

package/src/duckdb/src/common/types/column/column_data_allocator.cpp

@@ -1,8 +1,8 @@
 #include "duckdb/common/types/column/column_data_allocator.hpp"
 
 #include "duckdb/common/types/column/column_data_collection_segment.hpp"
-#include "duckdb/storage/buffer_manager.hpp"
 #include "duckdb/storage/buffer/block_handle.hpp"
+#include "duckdb/storage/buffer_manager.hpp"
 
 namespace duckdb {
 
@@ -19,6 +19,7 @@ ColumnDataAllocator::ColumnDataAllocator(ClientContext &context, ColumnDataAlloc
     : type(allocator_type) {
 	switch (type) {
 	case ColumnDataAllocatorType::BUFFER_MANAGER_ALLOCATOR:
+	case ColumnDataAllocatorType::HYBRID:
 		alloc.buffer_manager = &BufferManager::GetBufferManager(context);
 		break;
 	case ColumnDataAllocatorType::IN_MEMORY_ALLOCATOR:
@@ -33,6 +34,7 @@ ColumnDataAllocator::ColumnDataAllocator(ColumnDataAllocator &other) {
 	type = other.GetType();
 	switch (type) {
 	case ColumnDataAllocatorType::BUFFER_MANAGER_ALLOCATOR:
+	case ColumnDataAllocatorType::HYBRID:
 		alloc.allocator = other.alloc.allocator;
 		break;
 	case ColumnDataAllocatorType::IN_MEMORY_ALLOCATOR:
@@ -44,7 +46,7 @@ ColumnDataAllocator::ColumnDataAllocator(ColumnDataAllocator &other) {
 }
 
 BufferHandle ColumnDataAllocator::Pin(uint32_t block_id) {
-	D_ASSERT(type == ColumnDataAllocatorType::BUFFER_MANAGER_ALLOCATOR);
+	D_ASSERT(type == ColumnDataAllocatorType::BUFFER_MANAGER_ALLOCATOR || type == ColumnDataAllocatorType::HYBRID);
 	shared_ptr<BlockHandle> handle;
 	if (shared) {
 		// we only need to grab the lock when accessing the vector, because vector access is not thread-safe:
@@ -58,7 +60,7 @@ BufferHandle ColumnDataAllocator::Pin(uint32_t block_id) {
 }
 
 BufferHandle ColumnDataAllocator::AllocateBlock(idx_t size) {
-	D_ASSERT(type == ColumnDataAllocatorType::BUFFER_MANAGER_ALLOCATOR);
+	D_ASSERT(type == ColumnDataAllocatorType::BUFFER_MANAGER_ALLOCATOR || type == ColumnDataAllocatorType::HYBRID);
 	auto block_size = MaxValue<idx_t>(size, Storage::BLOCK_SIZE);
 	BlockMetaData data;
 	data.size = 0;
@@ -136,6 +138,7 @@ void ColumnDataAllocator::AllocateData(idx_t size, uint32_t &block_id, uint32_t
                                        ChunkManagementState *chunk_state) {
 	switch (type) {
 	case ColumnDataAllocatorType::BUFFER_MANAGER_ALLOCATOR:
+	case ColumnDataAllocatorType::HYBRID:
 		if (shared) {
 			lock_guard<mutex> guard(lock);
 			AllocateBuffer(size, block_id, offset, chunk_state);
@@ -174,8 +177,8 @@ data_ptr_t ColumnDataAllocator::GetDataPointer(ChunkManagementState &state, uint
 	return state.handles[block_id].Ptr() + offset;
 }
 
-void ColumnDataAllocator::UnswizzlePointers(ChunkManagementState &state, Vector &result, uint16_t v_offset,
-                                            uint16_t count, uint32_t block_id, uint32_t offset) {
+void ColumnDataAllocator::UnswizzlePointers(ChunkManagementState &state, Vector &result, idx_t v_offset, uint16_t count,
+                                            uint32_t block_id, uint32_t offset) {
 	D_ASSERT(result.GetType().InternalType() == PhysicalType::VARCHAR);
 	lock_guard<mutex> guard(lock);
 
@@ -225,7 +228,7 @@ Allocator &ColumnDataAllocator::GetAllocator() {
 }
 
 void ColumnDataAllocator::InitializeChunkState(ChunkManagementState &state, ChunkMetaData &chunk) {
-	if (type != ColumnDataAllocatorType::BUFFER_MANAGER_ALLOCATOR) {
+	if (type != ColumnDataAllocatorType::BUFFER_MANAGER_ALLOCATOR && type != ColumnDataAllocatorType::HYBRID) {
 		// nothing to pin
 		return;
 	}