duckdb 1.0.1-dev21.0 → 1.0.1-dev27.0

package/src/duckdb/src/storage/table/struct_column_data.cpp

@@ -36,38 +36,17 @@ void StructColumnData::SetStart(idx_t new_start) {
 	validity.SetStart(new_start);
 }
 
-bool StructColumnData::CheckZonemap(ColumnScanState &state, TableFilter &filter) {
-	if (!state.segment_checked) {
-		if (!state.current) {
-			return true;
-		}
-		state.segment_checked = true;
-
-		FilterPropagateResult prune_result;
-		{
-			lock_guard<mutex> l(stats_lock);
-			prune_result = filter.CheckStatistics(state.current->stats.statistics);
-			if (prune_result != FilterPropagateResult::FILTER_ALWAYS_FALSE) {
-				return true;
-			}
-		}
-		lock_guard<mutex> l(update_lock);
-		if (updates) {
-			auto update_stats = updates->GetStatistics();
-			prune_result = filter.CheckStatistics(*update_stats);
-			return prune_result != FilterPropagateResult::FILTER_ALWAYS_FALSE;
-		} else {
-			return false;
-		}
-	} else {
-		return true;
-	}
-}
-
 idx_t StructColumnData::GetMaxEntry() {
 	return sub_columns[0]->GetMaxEntry();
 }
 
+void StructColumnData::InitializePrefetch(PrefetchState &prefetch_state, ColumnScanState &scan_state, idx_t rows) {
+	validity.InitializePrefetch(prefetch_state, scan_state.child_states[0], rows);
+	for (idx_t i = 0; i < sub_columns.size(); i++) {
+		sub_columns[i]->InitializePrefetch(prefetch_state, scan_state.child_states[i + 1], rows);
+	}
+}
+
 void StructColumnData::InitializeScan(ColumnScanState &state) {
 	D_ASSERT(state.child_states.size() == sub_columns.size() + 1);
 	state.row_index = 0;
@@ -274,13 +253,13 @@ public:
 		return std::move(global_stats);
 	}
 
-	void WriteDataPointers(RowGroupWriter &writer, Serializer &serializer) override {
-		serializer.WriteObject(101, "validity",
-		                       [&](Serializer &serializer) { validity_state->WriteDataPointers(writer, serializer); });
-		serializer.WriteList(102, "sub_columns", child_states.size(), [&](Serializer::List &list, idx_t i) {
-			auto &state = child_states[i];
-			list.WriteObject([&](Serializer &serializer) { state->WriteDataPointers(writer, serializer); });
-		});
+	PersistentColumnData ToPersistentData() override {
+		PersistentColumnData data(PhysicalType::STRUCT);
+		data.child_columns.push_back(validity_state->ToPersistentData());
+		for (auto &child_state : child_states) {
+			data.child_columns.push_back(child_state->ToPersistentData());
+		}
+		return data;
 	}
 };
 
@@ -299,15 +278,33 @@ unique_ptr<ColumnCheckpointState> StructColumnData::Checkpoint(RowGroup &row_gro
 	return std::move(checkpoint_state);
 }
 
-void StructColumnData::DeserializeColumn(Deserializer &deserializer, BaseStatistics &target_stats) {
-	deserializer.ReadObject(
-	    101, "validity", [&](Deserializer &deserializer) { validity.DeserializeColumn(deserializer, target_stats); });
+bool StructColumnData::IsPersistent() {
+	if (!validity.IsPersistent()) {
+		return false;
+	}
+	for (auto &child_col : sub_columns) {
+		if (!child_col->IsPersistent()) {
+			return false;
+		}
+	}
+	return true;
+}
 
-	deserializer.ReadList(102, "sub_columns", [&](Deserializer::List &list, idx_t i) {
-		auto &child_stats = StructStats::GetChildStats(target_stats, i);
-		list.ReadObject([&](Deserializer &item) { sub_columns[i]->DeserializeColumn(item, child_stats); });
-	});
+PersistentColumnData StructColumnData::Serialize() {
+	PersistentColumnData persistent_data(PhysicalType::STRUCT);
+	persistent_data.child_columns.push_back(validity.Serialize());
+	for (auto &sub_column : sub_columns) {
+		persistent_data.child_columns.push_back(sub_column->Serialize());
+	}
+	return persistent_data;
+}
 
+void StructColumnData::InitializeColumn(PersistentColumnData &column_data, BaseStatistics &target_stats) {
+	validity.InitializeColumn(column_data.child_columns[0], target_stats);
+	for (idx_t c_idx = 0; c_idx < sub_columns.size(); c_idx++) {
+		auto &child_stats = StructStats::GetChildStats(target_stats, c_idx);
+		sub_columns[c_idx]->InitializeColumn(column_data.child_columns[c_idx + 1], child_stats);
+	}
 	this->count = validity.count.load();
 }
 

package/src/duckdb/src/storage/table/table_statistics.cpp

@@ -142,7 +142,8 @@ void TableStatistics::Deserialize(Deserializer &deserializer, ColumnList &column
 
 		deserializer.Unset<LogicalType>();
 	});
-	table_sample = deserializer.ReadPropertyWithDefault<unique_ptr<BlockingSample>>(101, "table_sample", nullptr);
+	table_sample =
+	    deserializer.ReadPropertyWithExplicitDefault<unique_ptr<BlockingSample>>(101, "table_sample", nullptr);
 }
 
 unique_ptr<TableStatisticsLock> TableStatistics::GetLock() {

package/src/duckdb/src/storage/table/update_segment.cpp

@@ -1,12 +1,12 @@
 #include "duckdb/storage/table/update_segment.hpp"
 
+#include "duckdb/common/exception/transaction_exception.hpp"
+#include "duckdb/common/printer.hpp"
 #include "duckdb/storage/statistics/distinct_statistics.hpp"
-
 #include "duckdb/storage/table/column_data.hpp"
 #include "duckdb/transaction/duck_transaction.hpp"
 #include "duckdb/transaction/update_info.hpp"
-#include "duckdb/common/printer.hpp"
-#include "duckdb/common/exception/transaction_exception.hpp"
+
 #include <algorithm>
 
 namespace duckdb {
@@ -928,7 +928,7 @@ idx_t TemplatedUpdateNumericStatistics(UpdateSegment *segment, SegmentStatistics
 
 	if (mask.AllValid()) {
 		for (idx_t i = 0; i < count; i++) {
-			NumericStats::Update<T>(stats.statistics, update_data[i]);
+			stats.statistics.UpdateNumericStats<T>(update_data[i]);
 		}
 		sel.Initialize(nullptr);
 		return count;
@@ -938,7 +938,7 @@ idx_t TemplatedUpdateNumericStatistics(UpdateSegment *segment, SegmentStatistics
 		for (idx_t i = 0; i < count; i++) {
 			if (mask.RowIsValid(i)) {
 				sel.set_index(not_null_count++, i);
-				NumericStats::Update<T>(stats.statistics, update_data[i]);
+				stats.statistics.UpdateNumericStats<T>(update_data[i]);
 			}
 		}
 		return not_null_count;
@@ -1062,7 +1062,8 @@ static idx_t SortSelectionVector(SelectionVector &sel, idx_t count, row_t *ids)
 
 UpdateInfo *CreateEmptyUpdateInfo(TransactionData transaction, idx_t type_size, idx_t count,
                                   unsafe_unique_array<char> &data) {
-	data = make_unsafe_uniq_array<char>(sizeof(UpdateInfo) + (sizeof(sel_t) + type_size) * STANDARD_VECTOR_SIZE);
+	data = make_unsafe_uniq_array_uninitialized<char>(sizeof(UpdateInfo) +
+	                                                  (sizeof(sel_t) + type_size) * STANDARD_VECTOR_SIZE);
 	auto update_info = reinterpret_cast<UpdateInfo *>(data.get());
 	update_info->max = STANDARD_VECTOR_SIZE;
 	update_info->tuples = reinterpret_cast<sel_t *>((data_ptr_cast(update_info)) + sizeof(UpdateInfo));
@@ -1163,8 +1164,8 @@ void UpdateSegment::Update(TransactionData transaction, idx_t column_index, Vect
 		auto result = make_uniq<UpdateNodeData>();
 
 		result->info = make_uniq<UpdateInfo>();
-		result->tuples = make_unsafe_uniq_array<sel_t>(STANDARD_VECTOR_SIZE);
-		result->tuple_data = make_unsafe_uniq_array<data_t>(STANDARD_VECTOR_SIZE * type_size);
+		result->tuples = make_unsafe_uniq_array_uninitialized<sel_t>(STANDARD_VECTOR_SIZE);
+		result->tuple_data = make_unsafe_uniq_array_uninitialized<data_t>(STANDARD_VECTOR_SIZE * type_size);
 		result->info->tuples = result->tuples.get();
 		result->info->tuple_data = result->tuple_data.get();
 		result->info->version_number = TRANSACTION_ID_START - 1;

package/src/duckdb/src/storage/table/validity_column_data.cpp

@@ -9,8 +9,8 @@ ValidityColumnData::ValidityColumnData(BlockManager &block_manager, DataTableInf
     : ColumnData(block_manager, info, column_index, start_row, LogicalType(LogicalTypeId::VALIDITY), &parent) {
 }
 
-bool ValidityColumnData::CheckZonemap(ColumnScanState &state, TableFilter &filter) {
-	return true;
+FilterPropagateResult ValidityColumnData::CheckZonemap(ColumnScanState &state, TableFilter &filter) {
+	return FilterPropagateResult::NO_PRUNING_POSSIBLE;
 }
 
 void ValidityColumnData::AppendData(BaseStatistics &stats, ColumnAppendState &state, UnifiedVectorFormat &vdata,

package/src/duckdb/src/storage/table_index_list.cpp

@@ -167,21 +167,22 @@ vector<column_t> TableIndexList::GetRequiredColumns() {
 	return result;
 }
 
-vector<IndexStorageInfo> TableIndexList::GetStorageInfos() {
+vector<IndexStorageInfo> TableIndexList::GetStorageInfos(const case_insensitive_map_t<Value> &options) {
 
 	vector<IndexStorageInfo> index_storage_infos;
 	for (auto &index : indexes) {
 		if (index->IsBound()) {
-			auto index_storage_info = index->Cast<BoundIndex>().GetStorageInfo(false);
-			D_ASSERT(index_storage_info.IsValid() && !index_storage_info.name.empty());
-			index_storage_infos.push_back(index_storage_info);
-		} else {
-			// TODO: Will/should this ever happen?
-			auto index_storage_info = index->Cast<UnboundIndex>().GetStorageInfo();
+			auto index_storage_info = index->Cast<BoundIndex>().GetStorageInfo(options, false);
 			D_ASSERT(index_storage_info.IsValid() && !index_storage_info.name.empty());
 			index_storage_infos.push_back(index_storage_info);
+			continue;
 		}
+
+		auto index_storage_info = index->Cast<UnboundIndex>().GetStorageInfo();
+		D_ASSERT(index_storage_info.IsValid() && !index_storage_info.name.empty());
+		index_storage_infos.push_back(index_storage_info);
 	}
+
 	return index_storage_infos;
 }
 

package/src/duckdb/src/storage/temporary_file_manager.cpp

@@ -58,7 +58,7 @@ bool BlockIndexManager::HasFreeBlocks() {
 }
 
 void BlockIndexManager::SetMaxIndex(idx_t new_index) {
-	static constexpr idx_t TEMPFILE_BLOCK_SIZE = Storage::BLOCK_ALLOC_SIZE;
+	static constexpr idx_t TEMP_FILE_BLOCK_SIZE = DEFAULT_BLOCK_ALLOC_SIZE;
 	if (!manager) {
 		max_index = new_index;
 	} else {
@@ -66,11 +66,11 @@ void BlockIndexManager::SetMaxIndex(idx_t new_index) {
 		if (new_index < old) {
 			max_index = new_index;
 			auto difference = old - new_index;
-			auto size_on_disk = difference * TEMPFILE_BLOCK_SIZE;
+			auto size_on_disk = difference * TEMP_FILE_BLOCK_SIZE;
 			manager->DecreaseSizeOnDisk(size_on_disk);
 		} else if (new_index > old) {
 			auto difference = new_index - old;
-			auto size_on_disk = difference * TEMPFILE_BLOCK_SIZE;
+			auto size_on_disk = difference * TEMP_FILE_BLOCK_SIZE;
 			manager->IncreaseSizeOnDisk(size_on_disk);
 			// Increase can throw, so this is only updated after it was succesfully updated
 			max_index = new_index;
@@ -118,15 +118,16 @@ TemporaryFileIndex TemporaryFileHandle::TryGetBlockIndex() {
 }
 
 void TemporaryFileHandle::WriteTemporaryFile(FileBuffer &buffer, TemporaryFileIndex index) {
-	D_ASSERT(buffer.size == Storage::BLOCK_SIZE);
+	// We group DEFAULT_BLOCK_ALLOC_SIZE blocks into the same file.
+	D_ASSERT(buffer.size == BufferManager::GetBufferManager(db).GetBlockSize());
 	buffer.Write(*handle, GetPositionInFile(index.block_index));
 }
 
 unique_ptr<FileBuffer> TemporaryFileHandle::ReadTemporaryBuffer(idx_t block_index,
                                                                 unique_ptr<FileBuffer> reusable_buffer) {
-	return StandardBufferManager::ReadTemporaryBufferInternal(BufferManager::GetBufferManager(db), *handle,
-	                                                          GetPositionInFile(block_index), Storage::BLOCK_SIZE,
-	                                                          std::move(reusable_buffer));
+	return StandardBufferManager::ReadTemporaryBufferInternal(
+	    BufferManager::GetBufferManager(db), *handle, GetPositionInFile(block_index),
+	    BufferManager::GetBufferManager(db).GetBlockSize(), std::move(reusable_buffer));
 }
 
 void TemporaryFileHandle::EraseBlockIndex(block_id_t block_index) {
@@ -180,7 +181,7 @@ void TemporaryFileHandle::RemoveTempBlockIndex(TemporaryFileLock &, idx_t index)
 }
 
 idx_t TemporaryFileHandle::GetPositionInFile(idx_t index) {
-	return index * Storage::BLOCK_ALLOC_SIZE;
+	return index * BufferManager::GetBufferManager(db).GetBlockAllocSize();
 }
 
 //===--------------------------------------------------------------------===//
@@ -277,7 +278,8 @@ TemporaryFileManager::TemporaryManagerLock::TemporaryManagerLock(mutex &mutex) :
 }
 
 void TemporaryFileManager::WriteTemporaryBuffer(block_id_t block_id, FileBuffer &buffer) {
-	D_ASSERT(buffer.size == Storage::BLOCK_SIZE);
+	// We group DEFAULT_BLOCK_ALLOC_SIZE blocks into the same file.
+	D_ASSERT(buffer.size == BufferManager::GetBufferManager(db).GetBlockSize());
 	TemporaryFileIndex index;
 	TemporaryFileHandle *handle = nullptr;
 

package/src/duckdb/src/storage/temporary_memory_manager.cpp

@@ -4,24 +4,38 @@
 #include "duckdb/parallel/task_scheduler.hpp"
 #include "duckdb/storage/buffer_manager.hpp"
 
+#include <cmath>
+
 namespace duckdb {
 
 TemporaryMemoryState::TemporaryMemoryState(TemporaryMemoryManager &temporary_memory_manager_p,
                                            idx_t minimum_reservation_p)
     : temporary_memory_manager(temporary_memory_manager_p), remaining_size(0),
-      minimum_reservation(minimum_reservation_p), reservation(0) {
+      minimum_reservation(minimum_reservation_p), reservation(0), materialization_penalty(1) {
 }
 
 TemporaryMemoryState::~TemporaryMemoryState() {
 	temporary_memory_manager.Unregister(*this);
 }
 
-void TemporaryMemoryState::SetRemainingSize(ClientContext &context, idx_t new_remaining_size) {
+void TemporaryMemoryState::SetRemainingSize(idx_t new_remaining_size) {
+	auto guard = temporary_memory_manager.Lock();
+	temporary_memory_manager.SetRemainingSize(*this, new_remaining_size);
+}
+
+void TemporaryMemoryState::SetRemainingSizeAndUpdateReservation(ClientContext &context, idx_t new_remaining_size) {
+	D_ASSERT(new_remaining_size != 0); // Use SetZero instead
 	auto guard = temporary_memory_manager.Lock();
 	temporary_memory_manager.SetRemainingSize(*this, new_remaining_size);
 	temporary_memory_manager.UpdateState(context, *this);
 }
 
+void TemporaryMemoryState::SetZero() {
+	auto guard = temporary_memory_manager.Lock();
+	temporary_memory_manager.SetRemainingSize(*this, 0);
+	temporary_memory_manager.SetReservation(*this, 0);
+}
+
 idx_t TemporaryMemoryState::GetRemainingSize() const {
 	return remaining_size;
 }
@@ -30,10 +44,28 @@ void TemporaryMemoryState::SetMinimumReservation(idx_t new_minimum_reservation)
 	minimum_reservation = new_minimum_reservation;
 }
 
+idx_t TemporaryMemoryState::GetMinimumReservation() const {
+	return minimum_reservation;
+}
+
+void TemporaryMemoryState::UpdateReservation(ClientContext &context) {
+	auto guard = temporary_memory_manager.Lock();
+	temporary_memory_manager.UpdateState(context, *this);
+}
+
 idx_t TemporaryMemoryState::GetReservation() const {
 	return reservation;
 }
 
+void TemporaryMemoryState::SetMaterializationPenalty(idx_t new_materialization_penalty) {
+	auto guard = temporary_memory_manager.Lock();
+	materialization_penalty = new_materialization_penalty;
+}
+
+idx_t TemporaryMemoryState::GetMaterializationPenalty() const {
+	return materialization_penalty;
+}
+
 TemporaryMemoryManager::TemporaryMemoryManager() : reservation(0), remaining_size(0) {
 }
 
@@ -41,11 +73,22 @@ unique_lock<mutex> TemporaryMemoryManager::Lock() {
 	return unique_lock<mutex>(lock);
 }
 
+void TemporaryMemoryManager::Unregister(TemporaryMemoryState &temporary_memory_state) {
+	auto guard = Lock();
+
+	SetReservation(temporary_memory_state, 0);
+	SetRemainingSize(temporary_memory_state, 0);
+	active_states.erase(temporary_memory_state);
+
+	Verify();
+}
+
 void TemporaryMemoryManager::UpdateConfiguration(ClientContext &context) {
 	auto &buffer_manager = BufferManager::GetBufferManager(context);
 	auto &task_scheduler = TaskScheduler::GetScheduler(context);
 
-	memory_limit = NumericCast<idx_t>(MAXIMUM_MEMORY_LIMIT_RATIO * static_cast<double>(buffer_manager.GetMaxMemory()));
+	memory_limit =
+	    LossyNumericCast<idx_t>(MAXIMUM_MEMORY_LIMIT_RATIO * static_cast<double>(buffer_manager.GetMaxMemory()));
 	has_temporary_directory = buffer_manager.HasTemporaryDirectory();
 	num_threads = NumericCast<idx_t>(task_scheduler.NumberOfThreads());
 	query_max_memory = buffer_manager.GetQueryMaxMemory();
@@ -62,8 +105,8 @@ unique_ptr<TemporaryMemoryState> TemporaryMemoryManager::Register(ClientContext
 	auto minimum_reservation = MinValue(num_threads * MINIMUM_RESERVATION_PER_STATE_PER_THREAD,
 	                                    memory_limit / MINIMUM_RESERVATION_MEMORY_LIMIT_DIVISOR);
 	auto result = unique_ptr<TemporaryMemoryState>(new TemporaryMemoryState(*this, minimum_reservation));
-	SetRemainingSize(*result, result->minimum_reservation);
-	SetReservation(*result, result->minimum_reservation);
+	SetRemainingSize(*result, result->GetMinimumReservation());
+	SetReservation(*result, result->GetMinimumReservation());
 	active_states.insert(*result);
 
 	Verify();
@@ -73,72 +116,217 @@ unique_ptr<TemporaryMemoryState> TemporaryMemoryManager::Register(ClientContext
 void TemporaryMemoryManager::UpdateState(ClientContext &context, TemporaryMemoryState &temporary_memory_state) {
 	UpdateConfiguration(context);
 
-	if (context.config.force_external) {
+	// The lower bound for the reservation of this state is either the minimum reservation or the remaining size
+	const auto lower_bound =
+	    MinValue(temporary_memory_state.GetMinimumReservation(), temporary_memory_state.GetRemainingSize());
+
+	if (temporary_memory_state.GetRemainingSize() == 0) {
+		// Sometimes set to 0 to denote end of state (before actually deleting the state)
+		SetReservation(temporary_memory_state, 0);
+	} else if (context.config.force_external) {
 		// We're forcing external processing. Give it the minimum
-		SetReservation(temporary_memory_state, temporary_memory_state.minimum_reservation);
+		SetReservation(temporary_memory_state, lower_bound);
 	} else if (!has_temporary_directory) {
 		// We cannot offload, so we cannot limit memory usage. Set reservation equal to the remaining size
-		SetReservation(temporary_memory_state, temporary_memory_state.remaining_size);
-	} else if (reservation - temporary_memory_state.reservation >= memory_limit) {
+		SetReservation(temporary_memory_state, temporary_memory_state.GetRemainingSize());
+	} else if (reservation - temporary_memory_state.GetReservation() + lower_bound >= memory_limit) {
 		// We overshot. Set reservation equal to the minimum
-		SetReservation(temporary_memory_state, temporary_memory_state.minimum_reservation);
+		SetReservation(temporary_memory_state, lower_bound);
 	} else {
-		// The lower bound for the reservation of this state is its minimum reservation
-		auto &lower_bound = temporary_memory_state.minimum_reservation;
-
 		// The upper bound for the reservation of this state is the minimum of:
 		// 1. Remaining size of the state
 		// 2. The max memory per query
 		// 3. MAXIMUM_FREE_MEMORY_RATIO * free memory
-		auto upper_bound = MinValue<idx_t>(temporary_memory_state.remaining_size, query_max_memory);
-		auto free_memory = memory_limit - (reservation - temporary_memory_state.reservation);
-		upper_bound = MinValue<idx_t>(upper_bound, NumericCast<idx_t>(MAXIMUM_FREE_MEMORY_RATIO * free_memory));
-
-		if (remaining_size > memory_limit) {
-			// We're processing more data than fits in memory, so we must further limit memory usage.
-			// The upper bound for the reservation of this state is now also the minimum of:
-			// 3. The ratio of the remaining size of this state and the total remaining size * memory limit
-			auto ratio_of_remaining = double(temporary_memory_state.remaining_size) / double(remaining_size);
-			upper_bound = MinValue<idx_t>(upper_bound, NumericCast<idx_t>(ratio_of_remaining * memory_limit));
+		auto upper_bound = MinValue(temporary_memory_state.GetRemainingSize(), query_max_memory);
+		const auto free_memory = memory_limit - (reservation - temporary_memory_state.GetReservation());
+		upper_bound = MinValue(upper_bound,
+		                       LossyNumericCast<idx_t>(MAXIMUM_FREE_MEMORY_RATIO * static_cast<double>(free_memory)));
+		upper_bound = MinValue(upper_bound, free_memory);
+
+		idx_t new_reservation;
+		if (lower_bound >= upper_bound) {
+			new_reservation = lower_bound;
+		} else {
+			new_reservation = remaining_size > memory_limit ? ComputeReservation(temporary_memory_state) : upper_bound;
 		}
 
-		SetReservation(temporary_memory_state, MaxValue<idx_t>(lower_bound, upper_bound));
+		SetReservation(temporary_memory_state, new_reservation);
 	}
 
 	Verify();
 }
 
 void TemporaryMemoryManager::SetRemainingSize(TemporaryMemoryState &temporary_memory_state, idx_t new_remaining_size) {
-	D_ASSERT(this->remaining_size >= temporary_memory_state.remaining_size);
-	this->remaining_size -= temporary_memory_state.remaining_size;
+	D_ASSERT(this->remaining_size >= temporary_memory_state.GetRemainingSize());
+	this->remaining_size -= temporary_memory_state.GetRemainingSize();
 	temporary_memory_state.remaining_size = new_remaining_size;
-	this->remaining_size += temporary_memory_state.remaining_size;
+	this->remaining_size += temporary_memory_state.GetRemainingSize();
 }
 
 void TemporaryMemoryManager::SetReservation(TemporaryMemoryState &temporary_memory_state, idx_t new_reservation) {
-	D_ASSERT(this->reservation >= temporary_memory_state.reservation);
-	this->reservation -= temporary_memory_state.reservation;
+	D_ASSERT(this->reservation >= temporary_memory_state.GetReservation());
+	this->reservation -= temporary_memory_state.GetReservation();
 	temporary_memory_state.reservation = new_reservation;
-	this->reservation += temporary_memory_state.reservation;
+	this->reservation += temporary_memory_state.GetReservation();
 }
 
-void TemporaryMemoryManager::Unregister(TemporaryMemoryState &temporary_memory_state) {
-	auto guard = Lock();
+//! Compute initial reservation for use in ComputeReservation
+static idx_t ComputeInitialReservation(const TemporaryMemoryState &temporary_memory_state) {
+	// Maximum of minimum reservation and the current reservation
+	auto result = MaxValue(temporary_memory_state.GetMinimumReservation(), temporary_memory_state.GetReservation());
+	// Bounded by the remaining size
+	result = MinValue(result, temporary_memory_state.GetRemainingSize());
+	// At least 1
+	return MaxValue<idx_t>(result, 1);
+}
 
-	SetReservation(temporary_memory_state, 0);
-	SetRemainingSize(temporary_memory_state, 0);
-	active_states.erase(temporary_memory_state);
+static void ComputeDerivatives(const vector<reference<const TemporaryMemoryState>> &states, const vector<idx_t> &res,
+                               vector<double> &der, const idx_t n) {
+	// Cost function takes "throughput" (reservation / size) of each operator as its principal input
+	double prod_siz = 1;
+	double prod_res = 1;
+	double mat_cost = 0;
+	for (idx_t i = 0; i < n; i++) {
+		auto &state = states[i].get();
+		const auto resd = static_cast<double>(res[i]);
+		const auto sizd = static_cast<double>(MaxValue<idx_t>(state.GetRemainingSize(), 1));
+		const auto pend = static_cast<double>(state.GetMaterializationPenalty());
+		prod_res *= resd;
+		prod_siz *= sizd;
+		mat_cost += pend * (1 - resd / sizd); // Materialization cost: sum of (1 - throughput)
+	}
+	const double nd = static_cast<double>(n);                    // n as double for convenience
+	const double tp_mult = 1 - pow(prod_res / prod_siz, 1 / nd); // Throughput multiplier: 1 - geomean throughputs
 
-	Verify();
+	// Cost function: materialization cost * (1 - throughput multiplier), but we don't actually need to compute it
+	// here. We need to compute the derivative with respect to every reservation, stored in "der"
+	// Just use https://www.derivative-calculator.net with this (n = 3) to see what's going on
+	// (3 - (a_1/s_1)-(a_2/s_2)-(a_3/s_3))*(1-((a_1/s_1)*(a_2/s_2)*(a_3/s_3))^(1/3))
+	const double intermediate = -(pow(prod_res, 1 / nd) * mat_cost) / (nd * pow(prod_siz, 1 / nd));
+	for (idx_t i = 0; i < n; i++) {
+		auto &state = states[i].get();
+		const auto resd = static_cast<double>(res[i]);
+		const auto sizd = static_cast<double>(MaxValue<idx_t>(state.GetRemainingSize(), 1));
+		const auto pend = static_cast<double>(state.GetMaterializationPenalty());
+		der[i] = intermediate / resd - pend * tp_mult / sizd;
+	}
+}
+
+idx_t TemporaryMemoryManager::ComputeReservation(const TemporaryMemoryState &temporary_memory_state) const {
+	static constexpr idx_t OPTIMIZATION_ITERATIONS_MULTIPLIER = 5;
+
+	// Use vectors for ease
+	optional_idx state_index;
+	vector<reference<const TemporaryMemoryState>> states;
+	vector<idx_t> res;
+	vector<double> der;
+	idx_t sum_of_initial_res = 0;
+
+	idx_t i = 0;
+	for (auto &state : active_states) {
+		const auto initial_reservation = ComputeInitialReservation(state);
+		sum_of_initial_res += initial_reservation;
+		if (RefersToSameObject(state.get(), temporary_memory_state)) {
+			state_index = i;
+		} else if (initial_reservation >= state.get().GetRemainingSize()) {
+			continue;
+		}
+		states.emplace_back(state);
+		res.push_back(initial_reservation);
+		der.push_back(NumericLimits<double>::Maximum());
+		i++;
+	}
+	const idx_t n = i;
+
+	if (sum_of_initial_res >= memory_limit) {
+		return res[state_index.GetIndex()];
+	}
+	const auto free_memory = memory_limit - sum_of_initial_res;
+
+	// Distribute memory in OPTIMIZATION_ITERATIONS
+	idx_t remaining_memory = free_memory;
+	const idx_t optimization_iterations = OPTIMIZATION_ITERATIONS_MULTIPLIER * n;
+	for (idx_t opt_idx = 0; opt_idx < optimization_iterations; opt_idx++) {
+		D_ASSERT(remaining_memory != 0);
+		ComputeDerivatives(states, res, der, n);
+
+		// Find the index of the state with the lowest derivative
+		idx_t min_idx = 0;
+		double min_der = NumericLimits<double>::Maximum();
+		for (i = 0; i < n; i++) {
+			auto &state = states[i].get();
+			if (res[i] >= state.GetRemainingSize()) {
+				continue; // We can't increase the reservation of "maxed" states, so we skip these
+			}
+			if (der[i] < min_der) {
+				min_idx = i;
+				min_der = der[i];
+			}
+		}
+		auto &min_state = states[min_idx].get();
+
+		// This is how much memory we will distribute in this round
+		const auto iter_memory = ExactNumericCast<idx_t>(
+		    std::ceil(static_cast<double>(remaining_memory) / static_cast<double>(optimization_iterations - opt_idx)));
+
+		// Compute how much we can add
+		const auto state_room = min_state.GetRemainingSize() - res[min_idx];
+		const auto delta = MinValue(iter_memory, state_room);
+		D_ASSERT(delta <= remaining_memory);
+
+		// Update counts
+		res[min_idx] += delta;
+		remaining_memory -= delta;
+	}
+	D_ASSERT(remaining_memory == 0);
+
+	// We computed how the memory should be assigned to the states,
+	// but we did not yet take into account the upper bound of MAXIMUM_FREE_MEMORY_RATIO * free_memory.
+	// We don't simply bound this state's reservation because order matters,
+	// and this function is called in arbitrary order for all active states.
+	// So, we order the states by derivative (as if they got the max reservation),
+	// and assign the states with the lowest derivate first.
+	// This way, order does not matter _as much_, and the "best" states will _likely_ get more memory
+	vector<idx_t> idxs(n, 0);
+	vector<idx_t> max_res(n, 0);
+	for (i = 0; i < n; i++) {
+		idxs[i] = i;
+		max_res[i] = states[i].get().GetRemainingSize() - 1;
+	}
+	ComputeDerivatives(states, max_res, der, n);
+	std::sort(idxs.begin(), idxs.end(), [&](const idx_t &lhs, const idx_t &rhs) { return der[lhs] < der[rhs]; });
+
+	// Loop through sorted indices until we encounter the state index
+	remaining_memory = free_memory;
+	for (const auto idx : idxs) {
+		auto &state = states[idx].get();
+		D_ASSERT(res[idx] <= state.GetRemainingSize());
+		const auto initial_state_reservation = ComputeInitialReservation(state);
+		const auto upper_bound = LossyNumericCast<idx_t>(
+		    MAXIMUM_FREE_MEMORY_RATIO * static_cast<double>(initial_state_reservation + remaining_memory));
+		auto state_reservation = MinValue(res[idx], upper_bound);
+		// But make sure it's never less than the initial reservation
+		state_reservation = MaxValue(state_reservation, initial_state_reservation);
+		// If this is the current state, we can just return
+		if (idx == state_index.GetIndex()) {
+			return state_reservation;
+		}
+		// Decrement the remaining memory
+		const auto delta = state_reservation - initial_state_reservation;
+		remaining_memory -= delta;
+	}
+
+	// We cannot have gotten through the loop above without finding the current state
+	throw InternalException("Did not find state_index in ComputeOptimalReservation");
 }
 
 void TemporaryMemoryManager::Verify() const {
 #ifdef DEBUG
 	idx_t total_reservation = 0;
 	idx_t total_remaining_size = 0;
-	for (auto &temporary_memory_state : active_states) {
-		total_reservation += temporary_memory_state.get().reservation;
-		total_remaining_size += temporary_memory_state.get().remaining_size;
+	for (auto &active_state : active_states) {
+		total_reservation += active_state.get().GetReservation();
+		total_remaining_size += active_state.get().GetRemainingSize();
 	}
 	D_ASSERT(total_reservation == this->reservation);
 	D_ASSERT(total_remaining_size == this->remaining_size);