duckdb 0.8.2-dev33.0 → 0.8.2-dev3334.0

package/src/duckdb/src/execution/operator/scan/physical_column_data_scan.cpp

@@ -64,6 +64,9 @@ void PhysicalColumnDataScan::BuildPipelines(Pipeline &current, MetaPipeline &met
 		state.SetPipelineSource(current, delim_join.distinct->Cast<PhysicalOperator>());
 		return;
 	}
+	case PhysicalOperatorType::CTE_SCAN: {
+		break;
+	}
 	case PhysicalOperatorType::RECURSIVE_CTE_SCAN:
 		if (!meta_pipeline.HasRecursiveCTE()) {
 			throw InternalException("Recursive CTE scan found without recursive CTE node");
@@ -76,4 +79,20 @@ void PhysicalColumnDataScan::BuildPipelines(Pipeline &current, MetaPipeline &met
 	state.SetPipelineSource(current, *this);
 }
 
+string PhysicalColumnDataScan::ParamsToString() const {
+	string result = "";
+	switch (type) {
+	case PhysicalOperatorType::CTE_SCAN:
+	case PhysicalOperatorType::RECURSIVE_CTE_SCAN: {
+		result += "\n[INFOSEPARATOR]\n";
+		result += StringUtil::Format("idx: %llu", cte_index);
+		break;
+	}
+	default:
+		break;
+	}
+
+	return result;
+}
+
 } // namespace duckdb

package/src/duckdb/src/execution/operator/scan/physical_table_scan.cpp

@@ -9,24 +9,15 @@
 
 namespace duckdb {
 
-PhysicalTableScan::PhysicalTableScan(vector<LogicalType> types, TableFunction function_p,
-                                     unique_ptr<FunctionData> bind_data_p, vector<column_t> column_ids_p,
-                                     vector<string> names_p, unique_ptr<TableFilterSet> table_filters_p,
-                                     idx_t estimated_cardinality)
-    : PhysicalOperator(PhysicalOperatorType::TABLE_SCAN, std::move(types), estimated_cardinality),
-      function(std::move(function_p)), bind_data(std::move(bind_data_p)), column_ids(std::move(column_ids_p)),
-      names(std::move(names_p)), table_filters(std::move(table_filters_p)) {
-}
-
 PhysicalTableScan::PhysicalTableScan(vector<LogicalType> types, TableFunction function_p,
                                      unique_ptr<FunctionData> bind_data_p, vector<LogicalType> returned_types_p,
                                      vector<column_t> column_ids_p, vector<idx_t> projection_ids_p,
                                      vector<string> names_p, unique_ptr<TableFilterSet> table_filters_p,
-                                     idx_t estimated_cardinality)
+                                     idx_t estimated_cardinality, ExtraOperatorInfo extra_info)
     : PhysicalOperator(PhysicalOperatorType::TABLE_SCAN, std::move(types), estimated_cardinality),
       function(std::move(function_p)), bind_data(std::move(bind_data_p)), returned_types(std::move(returned_types_p)),
       column_ids(std::move(column_ids_p)), projection_ids(std::move(projection_ids_p)), names(std::move(names_p)),
-      table_filters(std::move(table_filters_p)) {
+      table_filters(std::move(table_filters_p)), extra_info(extra_info) {
 }
 
 class TableScanGlobalSourceState : public GlobalSourceState {
@@ -149,8 +140,12 @@ string PhysicalTableScan::ParamsToString() const {
 			}
 		}
 	}
+	if (!extra_info.file_filters.empty()) {
+		result += "\n[INFOSEPARATOR]\n";
+		result += "File Filters: " + extra_info.file_filters;
+	}
 	result += "\n[INFOSEPARATOR]\n";
-	result += StringUtil::Format("EC: %llu", estimated_props->GetCardinality<idx_t>());
+	result += StringUtil::Format("EC: %llu", estimated_cardinality);
 	return result;
 }
 

package/src/duckdb/src/execution/operator/schema/physical_attach.cpp

@@ -66,7 +66,8 @@ SourceResultType PhysicalAttach::GetData(ExecutionContext &context, DataChunk &c
 	const auto &path = info->path;
 
 	if (name.empty()) {
-		name = AttachedDatabase::ExtractDatabaseName(path);
+		auto &fs = FileSystem::GetFileSystem(context.client);
+		name = AttachedDatabase::ExtractDatabaseName(path, fs);
 	}
 	auto &db_manager = DatabaseManager::Get(context.client);
 	auto existing_db = db_manager.GetDatabaseFromPath(context.client, path);

package/src/duckdb/src/execution/operator/schema/physical_create_art_index.cpp

@@ -0,0 +1,198 @@
+#include "duckdb/execution/operator/schema/physical_create_art_index.hpp"
+
+#include "duckdb/catalog/catalog_entry/duck_table_entry.hpp"
+#include "duckdb/catalog/catalog_entry/table_catalog_entry.hpp"
+#include "duckdb/catalog/catalog_entry/duck_index_entry.hpp"
+#include "duckdb/main/client_context.hpp"
+#include "duckdb/storage/index.hpp"
+#include "duckdb/storage/storage_manager.hpp"
+#include "duckdb/storage/table/append_state.hpp"
+#include "duckdb/main/database_manager.hpp"
+#include "duckdb/execution/index/art/art_key.hpp"
+#include "duckdb/execution/index/art/node.hpp"
+#include "duckdb/execution/index/art/leaf.hpp"
+
+namespace duckdb {
+
+PhysicalCreateARTIndex::PhysicalCreateARTIndex(LogicalOperator &op, TableCatalogEntry &table_p,
+                                               const vector<column_t> &column_ids, unique_ptr<CreateIndexInfo> info,
+                                               vector<unique_ptr<Expression>> unbound_expressions,
+                                               idx_t estimated_cardinality, const bool sorted)
+    : PhysicalOperator(PhysicalOperatorType::CREATE_INDEX, op.types, estimated_cardinality),
+      table(table_p.Cast<DuckTableEntry>()), info(std::move(info)), unbound_expressions(std::move(unbound_expressions)),
+      sorted(sorted) {
+	// convert virtual column ids to storage column ids
+	for (auto &column_id : column_ids) {
+		storage_ids.push_back(table.GetColumns().LogicalToPhysical(LogicalIndex(column_id)).index);
+	}
+}
+
+//===--------------------------------------------------------------------===//
+// Sink
+//===--------------------------------------------------------------------===//
+
+class CreateARTIndexGlobalSinkState : public GlobalSinkState {
+public:
+	//! Global index to be added to the table
+	unique_ptr<Index> global_index;
+};
+
+class CreateARTIndexLocalSinkState : public LocalSinkState {
+public:
+	explicit CreateARTIndexLocalSinkState(ClientContext &context) : arena_allocator(Allocator::Get(context)) {};
+
+	unique_ptr<Index> local_index;
+	ArenaAllocator arena_allocator;
+	vector<ARTKey> keys;
+	DataChunk key_chunk;
+	vector<column_t> key_column_ids;
+};
+
+unique_ptr<GlobalSinkState> PhysicalCreateARTIndex::GetGlobalSinkState(ClientContext &context) const {
+	auto state = make_uniq<CreateARTIndexGlobalSinkState>();
+
+	// create the global index
+	auto &storage = table.GetStorage();
+	state->global_index = make_uniq<ART>(storage_ids, TableIOManager::Get(storage), unbound_expressions,
+	                                     info->constraint_type, storage.db);
+
+	return (std::move(state));
+}
+
+unique_ptr<LocalSinkState> PhysicalCreateARTIndex::GetLocalSinkState(ExecutionContext &context) const {
+	auto state = make_uniq<CreateARTIndexLocalSinkState>(context.client);
+
+	// create the local index
+
+	auto &storage = table.GetStorage();
+	state->local_index = make_uniq<ART>(storage_ids, TableIOManager::Get(storage), unbound_expressions,
+	                                    info->constraint_type, storage.db);
+
+	state->keys = vector<ARTKey>(STANDARD_VECTOR_SIZE);
+	state->key_chunk.Initialize(Allocator::Get(context.client), state->local_index->logical_types);
+
+	for (idx_t i = 0; i < state->key_chunk.ColumnCount(); i++) {
+		state->key_column_ids.push_back(i);
+	}
+	return std::move(state);
+}
+
+SinkResultType PhysicalCreateARTIndex::SinkUnsorted(Vector &row_identifiers, OperatorSinkInput &input) const {
+
+	auto &l_state = input.local_state.Cast<CreateARTIndexLocalSinkState>();
+	auto count = l_state.key_chunk.size();
+
+	// get the corresponding row IDs
+	row_identifiers.Flatten(count);
+	auto row_ids = FlatVector::GetData<row_t>(row_identifiers);
+
+	// insert the row IDs
+	auto &art = l_state.local_index->Cast<ART>();
+	for (idx_t i = 0; i < count; i++) {
+		if (!art.Insert(*art.tree, l_state.keys[i], 0, row_ids[i])) {
+			throw ConstraintException("Data contains duplicates on indexed column(s)");
+		}
+	}
+
+	return SinkResultType::NEED_MORE_INPUT;
+}
+
+SinkResultType PhysicalCreateARTIndex::SinkSorted(Vector &row_identifiers, OperatorSinkInput &input) const {
+
+	auto &l_state = input.local_state.Cast<CreateARTIndexLocalSinkState>();
+	auto &storage = table.GetStorage();
+	auto &l_index = l_state.local_index;
+
+	// create an ART from the chunk
+	auto art = make_uniq<ART>(l_index->column_ids, l_index->table_io_manager, l_index->unbound_expressions,
+	                          l_index->constraint_type, storage.db, l_index->Cast<ART>().allocators);
+	if (!art->ConstructFromSorted(l_state.key_chunk.size(), l_state.keys, row_identifiers)) {
+		throw ConstraintException("Data contains duplicates on indexed column(s)");
+	}
+
+	// merge into the local ART
+	if (!l_index->MergeIndexes(*art)) {
+		throw ConstraintException("Data contains duplicates on indexed column(s)");
+	}
+
+	return SinkResultType::NEED_MORE_INPUT;
+}
+
+SinkResultType PhysicalCreateARTIndex::Sink(ExecutionContext &context, DataChunk &chunk,
+                                            OperatorSinkInput &input) const {
+
+	D_ASSERT(chunk.ColumnCount() >= 2);
+
+	// generate the keys for the given input
+	auto &l_state = input.local_state.Cast<CreateARTIndexLocalSinkState>();
+	l_state.key_chunk.ReferenceColumns(chunk, l_state.key_column_ids);
+	l_state.arena_allocator.Reset();
+	ART::GenerateKeys(l_state.arena_allocator, l_state.key_chunk, l_state.keys);
+
+	// insert the keys and their corresponding row IDs
+	auto &row_identifiers = chunk.data[chunk.ColumnCount() - 1];
+	if (sorted) {
+		return SinkSorted(row_identifiers, input);
+	}
+	return SinkUnsorted(row_identifiers, input);
+}
+
+SinkCombineResultType PhysicalCreateARTIndex::Combine(ExecutionContext &context,
+                                                      OperatorSinkCombineInput &input) const {
+
+	auto &gstate = input.global_state.Cast<CreateARTIndexGlobalSinkState>();
+	auto &lstate = input.local_state.Cast<CreateARTIndexLocalSinkState>();
+
+	// merge the local index into the global index
+	if (!gstate.global_index->MergeIndexes(*lstate.local_index)) {
+		throw ConstraintException("Data contains duplicates on indexed column(s)");
+	}
+
+	return SinkCombineResultType::FINISHED;
+}
+
+SinkFinalizeType PhysicalCreateARTIndex::Finalize(Pipeline &pipeline, Event &event, ClientContext &context,
+                                                  OperatorSinkFinalizeInput &input) const {
+
+	// here, we set the resulting global index as the newly created index of the table
+	auto &state = input.global_state.Cast<CreateARTIndexGlobalSinkState>();
+
+	// vacuum excess memory and verify
+	state.global_index->Vacuum();
+	D_ASSERT(!state.global_index->VerifyAndToString(true).empty());
+
+	auto &storage = table.GetStorage();
+	if (!storage.IsRoot()) {
+		throw TransactionException("Transaction conflict: cannot add an index to a table that has been altered!");
+	}
+
+	auto &schema = table.schema;
+	auto index_entry = schema.CreateIndex(context, *info, table).get();
+	if (!index_entry) {
+		D_ASSERT(info->on_conflict == OnCreateConflict::IGNORE_ON_CONFLICT);
+		// index already exists, but error ignored because of IF NOT EXISTS
+		return SinkFinalizeType::READY;
+	}
+	auto &index = index_entry->Cast<DuckIndexEntry>();
+
+	index.index = state.global_index.get();
+	index.info = storage.info;
+	for (auto &parsed_expr : info->parsed_expressions) {
+		index.parsed_expressions.push_back(parsed_expr->Copy());
+	}
+
+	// add index to storage
+	storage.info->indexes.AddIndex(std::move(state.global_index));
+	return SinkFinalizeType::READY;
+}
+
+//===--------------------------------------------------------------------===//
+// Source
+//===--------------------------------------------------------------------===//
+
+SourceResultType PhysicalCreateARTIndex::GetData(ExecutionContext &context, DataChunk &chunk,
+                                                 OperatorSourceInput &input) const {
+	return SourceResultType::FINISHED;
+}
+
+} // namespace duckdb

package/src/duckdb/src/execution/operator/schema/physical_create_type.cpp

@@ -74,15 +74,11 @@ SourceResultType PhysicalCreateType::GetData(ExecutionContext &context, DataChun
 	if (IsSink()) {
 		D_ASSERT(info->type == LogicalType::INVALID);
 		auto &g_sink_state = sink_state->Cast<CreateTypeGlobalState>();
-		info->type = LogicalType::ENUM(info->name, g_sink_state.result, g_sink_state.size);
+		info->type = LogicalType::ENUM(g_sink_state.result, g_sink_state.size);
 	}
 
 	auto &catalog = Catalog::GetCatalog(context.client, info->catalog);
-	auto catalog_entry = catalog.CreateType(context.client, *info);
-	D_ASSERT(catalog_entry->type == CatalogType::TYPE_ENTRY);
-	auto &catalog_type = catalog_entry->Cast<TypeCatalogEntry>();
-	EnumType::SetCatalog(info->type, &catalog_type);
-
+	catalog.CreateType(context.client, *info);
 	return SourceResultType::FINISHED;
 }
 

package/src/duckdb/src/execution/operator/set/physical_cte.cpp

@@ -0,0 +1,160 @@
+#include "duckdb/execution/operator/set/physical_cte.hpp"
+
+#include "duckdb/common/types/column/column_data_collection.hpp"
+#include "duckdb/common/vector_operations/vector_operations.hpp"
+#include "duckdb/execution/aggregate_hashtable.hpp"
+#include "duckdb/execution/executor.hpp"
+#include "duckdb/parallel/event.hpp"
+#include "duckdb/parallel/meta_pipeline.hpp"
+#include "duckdb/parallel/pipeline.hpp"
+#include "duckdb/parallel/task_scheduler.hpp"
+#include "duckdb/storage/buffer_manager.hpp"
+
+namespace duckdb {
+
+PhysicalCTE::PhysicalCTE(string ctename, idx_t table_index, vector<LogicalType> types, unique_ptr<PhysicalOperator> top,
+                         unique_ptr<PhysicalOperator> bottom, idx_t estimated_cardinality)
+    : PhysicalOperator(PhysicalOperatorType::CTE, std::move(types), estimated_cardinality), table_index(table_index),
+      ctename(std::move(ctename)) {
+	children.push_back(std::move(top));
+	children.push_back(std::move(bottom));
+}
+
+PhysicalCTE::~PhysicalCTE() {
+}
+
+//===--------------------------------------------------------------------===//
+// Sink
+//===--------------------------------------------------------------------===//
+class CTEState : public GlobalSinkState {
+public:
+	explicit CTEState(ClientContext &context, const PhysicalCTE &op)
+	    : intermediate_table(context, op.children[1]->GetTypes()) {
+	}
+	ColumnDataCollection intermediate_table;
+	ColumnDataScanState scan_state;
+	bool initialized = false;
+	bool finished_scan = false;
+};
+
+unique_ptr<GlobalSinkState> PhysicalCTE::GetGlobalSinkState(ClientContext &context) const {
+	working_table->Reset();
+	return make_uniq<CTEState>(context, *this);
+}
+
+SinkResultType PhysicalCTE::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const {
+	auto &gstate = input.global_state.Cast<CTEState>();
+	if (!gstate.finished_scan) {
+		working_table->Append(chunk);
+	} else {
+		gstate.intermediate_table.Append(chunk);
+	}
+	return SinkResultType::NEED_MORE_INPUT;
+}
+
+//===--------------------------------------------------------------------===//
+// Source
+//===--------------------------------------------------------------------===//
+SourceResultType PhysicalCTE::GetData(ExecutionContext &context, DataChunk &chunk, OperatorSourceInput &input) const {
+	auto &gstate = sink_state->Cast<CTEState>();
+	if (!gstate.initialized) {
+		gstate.intermediate_table.InitializeScan(gstate.scan_state);
+		gstate.finished_scan = false;
+		gstate.initialized = true;
+	}
+	if (!gstate.finished_scan) {
+		gstate.finished_scan = true;
+		ExecuteRecursivePipelines(context);
+	}
+
+	gstate.intermediate_table.Scan(gstate.scan_state, chunk);
+
+	return chunk.size() == 0 ? SourceResultType::FINISHED : SourceResultType::HAVE_MORE_OUTPUT;
+}
+
+void PhysicalCTE::ExecuteRecursivePipelines(ExecutionContext &context) const {
+	if (!recursive_meta_pipeline) {
+		throw InternalException("Missing meta pipeline for recursive CTE");
+	}
+
+	// get and reset pipelines
+	vector<shared_ptr<Pipeline>> pipelines;
+	recursive_meta_pipeline->GetPipelines(pipelines, true);
+	for (auto &pipeline : pipelines) {
+		auto sink = pipeline->GetSink();
+		if (sink.get() != this) {
+			sink->sink_state.reset();
+		}
+		for (auto &op_ref : pipeline->GetOperators()) {
+			auto &op = op_ref.get();
+			op.op_state.reset();
+		}
+		pipeline->ClearSource();
+	}
+
+	// get the MetaPipelines in the recursive_meta_pipeline and reschedule them
+	vector<shared_ptr<MetaPipeline>> meta_pipelines;
+	recursive_meta_pipeline->GetMetaPipelines(meta_pipelines, true, false);
+	auto &executor = recursive_meta_pipeline->GetExecutor();
+	vector<shared_ptr<Event>> events;
+	executor.ReschedulePipelines(meta_pipelines, events);
+
+	while (true) {
+		executor.WorkOnTasks();
+		if (executor.HasError()) {
+			executor.ThrowException();
+		}
+		bool finished = true;
+		for (auto &event : events) {
+			if (!event->IsFinished()) {
+				finished = false;
+				break;
+			}
+		}
+		if (finished) {
+			// all pipelines finished: done!
+			break;
+		}
+	}
+}
+
+//===--------------------------------------------------------------------===//
+// Pipeline Construction
+//===--------------------------------------------------------------------===//
+void PhysicalCTE::BuildPipelines(Pipeline &current, MetaPipeline &meta_pipeline) {
+	D_ASSERT(children.size() == 2);
+	op_state.reset();
+	sink_state.reset();
+	recursive_meta_pipeline.reset();
+
+	auto &state = meta_pipeline.GetState();
+	state.SetPipelineSource(current, *this);
+
+	auto &executor = meta_pipeline.GetExecutor();
+	executor.AddMaterializedCTE(*this);
+
+	auto &child_meta_pipeline = meta_pipeline.CreateChildMetaPipeline(current, *this);
+	child_meta_pipeline.Build(*children[0]);
+
+	// the RHS is the recursive pipeline
+	recursive_meta_pipeline = make_shared<MetaPipeline>(executor, state, this);
+	if (meta_pipeline.HasRecursiveCTE()) {
+		recursive_meta_pipeline->SetRecursiveCTE();
+	}
+	recursive_meta_pipeline->Build(*children[1]);
+}
+
+vector<const_reference<PhysicalOperator>> PhysicalCTE::GetSources() const {
+	return {*this};
+}
+
+string PhysicalCTE::ParamsToString() const {
+	string result = "";
+	result += "\n[INFOSEPARATOR]\n";
+	result += ctename;
+	result += "\n[INFOSEPARATOR]\n";
+	result += StringUtil::Format("idx: %llu", table_index);
+	return result;
+}
+
+} // namespace duckdb

package/src/duckdb/src/execution/operator/set/physical_recursive_cte.cpp

@@ -12,10 +12,11 @@
 
 namespace duckdb {
 
-PhysicalRecursiveCTE::PhysicalRecursiveCTE(vector<LogicalType> types, bool union_all, unique_ptr<PhysicalOperator> top,
-                                           unique_ptr<PhysicalOperator> bottom, idx_t estimated_cardinality)
+PhysicalRecursiveCTE::PhysicalRecursiveCTE(string ctename, idx_t table_index, vector<LogicalType> types, bool union_all,
+                                           unique_ptr<PhysicalOperator> top, unique_ptr<PhysicalOperator> bottom,
+                                           idx_t estimated_cardinality)
     : PhysicalOperator(PhysicalOperatorType::RECURSIVE_CTE, std::move(types), estimated_cardinality),
-      union_all(union_all) {
+      ctename(std::move(ctename)), table_index(table_index), union_all(union_all) {
 	children.push_back(std::move(top));
 	children.push_back(std::move(bottom));
 }
@@ -30,8 +31,8 @@ class RecursiveCTEState : public GlobalSinkState {
 public:
 	explicit RecursiveCTEState(ClientContext &context, const PhysicalRecursiveCTE &op)
 	    : intermediate_table(context, op.GetTypes()), new_groups(STANDARD_VECTOR_SIZE) {
-		ht = make_uniq<GroupedAggregateHashTable>(context, Allocator::Get(context), op.types, vector<LogicalType>(),
-		                                          vector<BoundAggregateExpression *>());
+		ht = make_uniq<GroupedAggregateHashTable>(context, BufferAllocator::Get(context), op.types,
+		                                          vector<LogicalType>(), vector<BoundAggregateExpression *>());
 	}
 
 	unique_ptr<GroupedAggregateHashTable> ht;
@@ -195,4 +196,13 @@ vector<const_reference<PhysicalOperator>> PhysicalRecursiveCTE::GetSources() con
 	return {*this};
 }
 
+string PhysicalRecursiveCTE::ParamsToString() const {
+	string result = "";
+	result += "\n[INFOSEPARATOR]\n";
+	result += ctename;
+	result += "\n[INFOSEPARATOR]\n";
+	result += StringUtil::Format("idx: %llu", table_index);
+	return result;
+}
+
 } // namespace duckdb

package/src/duckdb/src/execution/partitionable_hashtable.cpp

@@ -9,10 +9,10 @@ RadixPartitionInfo::RadixPartitionInfo(const idx_t n_partitions_upper_bound)
       radix_bits(RadixPartitioning::RadixBits(n_partitions)), radix_mask(RadixPartitioning::Mask(radix_bits)),
       radix_shift(RadixPartitioning::Shift(radix_bits)) {
 
+	D_ASSERT(radix_bits <= RadixPartitioning::MAX_RADIX_BITS);
 	D_ASSERT(n_partitions > 0);
-	D_ASSERT(n_partitions <= 256);
+	D_ASSERT(n_partitions == RadixPartitioning::NumberOfPartitions(radix_bits));
 	D_ASSERT(IsPowerOfTwo(n_partitions));
-	D_ASSERT(radix_bits <= 8);
 }
 
 PartitionableHashTable::PartitionableHashTable(ClientContext &context, Allocator &allocator,
@@ -47,11 +47,21 @@ HtEntryType PartitionableHashTable::GetHTEntrySize() {
 	return HtEntryType::HT_WIDTH_32;
 }
 
+bool OverMemoryLimit(ClientContext &context, const bool is_partitioned, const RadixPartitionInfo &partition_info,
+                     const GroupedAggregateHashTable &ht) {
+	const auto n_partitions = is_partitioned ? partition_info.n_partitions : 1;
+	const auto max_memory = BufferManager::GetBufferManager(context).GetMaxMemory();
+	const auto num_threads = TaskScheduler::GetScheduler(context).NumberOfThreads();
+	const auto memory_per_partition = 0.6 * max_memory / num_threads / n_partitions;
+	return ht.TotalSize() > memory_per_partition;
+}
+
 idx_t PartitionableHashTable::ListAddChunk(HashTableList &list, DataChunk &groups, Vector &group_hashes,
                                            DataChunk &payload, const unsafe_vector<idx_t> &filter) {
 	// If this is false, a single AddChunk would overflow the max capacity
 	D_ASSERT(list.empty() || groups.size() <= list.back()->MaxCapacity());
-	if (list.empty() || list.back()->Count() + groups.size() >= list.back()->MaxCapacity()) {
+	if (list.empty() || list.back()->Count() + groups.size() >= list.back()->MaxCapacity() ||
+	    OverMemoryLimit(context, is_partitioned, partition_info, *list.back())) {
 		idx_t new_capacity = GroupedAggregateHashTable::InitialCapacity();
 		if (!list.empty()) {
 			new_capacity = list.back()->Capacity();
@@ -70,7 +80,7 @@ idx_t PartitionableHashTable::AddChunk(DataChunk &groups, DataChunk &payload, bo
 
 	// we partition when we are asked to or when the unpartitioned ht runs out of space
 	if (!IsPartitioned() && do_partition) {
-		Partition();
+		Partition(false);
 	}
 
 	if (!IsPartitioned()) {
@@ -117,7 +127,7 @@ idx_t PartitionableHashTable::AddChunk(DataChunk &groups, DataChunk &payload, bo
 	return group_count;
 }
 
-void PartitionableHashTable::Partition() {
+void PartitionableHashTable::Partition(bool sink_done) {
 	D_ASSERT(!IsPartitioned());
 	D_ASSERT(radix_partitioned_hts.empty());
 	D_ASSERT(partition_info.n_partitions > 1);
@@ -130,7 +140,7 @@ void PartitionableHashTable::Partition() {
 			    context, allocator, group_types, payload_types, bindings, GetHTEntrySize()));
 			partition_hts[r] = radix_partitioned_hts[r].back().get();
 		}
-		unpartitioned_ht->Partition(partition_hts, partition_info.radix_bits);
+		unpartitioned_ht->Partition(partition_hts, partition_info.radix_bits, sink_done);
 		unpartitioned_ht.reset();
 	}
 	unpartitioned_hts.clear();
@@ -153,6 +163,22 @@ HashTableList PartitionableHashTable::GetUnpartitioned() {
 	return std::move(unpartitioned_hts);
 }
 
+idx_t PartitionableHashTable::GetPartitionCount(idx_t partition) const {
+	idx_t total_size = 0;
+	for (const auto &ht : radix_partitioned_hts[partition]) {
+		total_size += ht->Count();
+	}
+	return total_size;
+}
+
+idx_t PartitionableHashTable::GetPartitionSize(idx_t partition) const {
+	idx_t total_size = 0;
+	for (const auto &ht : radix_partitioned_hts[partition]) {
+		total_size += ht->DataSize();
+	}
+	return total_size;
+}
+
 void PartitionableHashTable::Finalize() {
 	if (IsPartitioned()) {
 		for (auto &ht_list : radix_partitioned_hts) {
@@ -169,4 +195,13 @@ void PartitionableHashTable::Finalize() {
 	}
 }
 
+void PartitionableHashTable::Append(GroupedAggregateHashTable &ht) {
+	if (unpartitioned_hts.empty()) {
+		unpartitioned_hts.push_back(make_uniq<GroupedAggregateHashTable>(context, allocator, group_types, payload_types,
+		                                                                 bindings, GetHTEntrySize(),
+		                                                                 GroupedAggregateHashTable::InitialCapacity()));
+	}
+	unpartitioned_hts.back()->Append(ht);
+}
+
 } // namespace duckdb