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

package/src/duckdb/src/storage/compression/string_uncompressed.cpp

@@ -1,9 +1,9 @@
 #include "duckdb/storage/string_uncompressed.hpp"
 
 #include "duckdb/common/pair.hpp"
-#include "duckdb/storage/checkpoint/write_overflow_strings_to_disk.hpp"
-#include "duckdb/common/serializer/serializer.hpp"
 #include "duckdb/common/serializer/deserializer.hpp"
+#include "duckdb/common/serializer/serializer.hpp"
+#include "duckdb/storage/checkpoint/write_overflow_strings_to_disk.hpp"
 #include "duckdb/storage/table/column_data.hpp"
 
 namespace duckdb {
@@ -22,7 +22,8 @@ UncompressedStringSegmentState::~UncompressedStringSegmentState() {
 // Analyze
 //===--------------------------------------------------------------------===//
 struct StringAnalyzeState : public AnalyzeState {
-	StringAnalyzeState() : count(0), total_string_size(0), overflow_strings(0) {
+	explicit StringAnalyzeState(const CompressionInfo &info)
+	    : AnalyzeState(info), count(0), total_string_size(0), overflow_strings(0) {
 	}
 
 	idx_t count;
@@ -31,7 +32,8 @@ struct StringAnalyzeState : public AnalyzeState {
 };
 
 unique_ptr<AnalyzeState> UncompressedStringStorage::StringInitAnalyze(ColumnData &col_data, PhysicalType type) {
-	return make_uniq<StringAnalyzeState>();
+	CompressionInfo info(col_data.GetBlockManager().GetBlockSize());
+	return make_uniq<StringAnalyzeState>(info);
 }
 
 bool UncompressedStringStorage::StringAnalyze(AnalyzeState &state_p, Vector &input, idx_t count) {
@@ -46,7 +48,7 @@ bool UncompressedStringStorage::StringAnalyze(AnalyzeState &state_p, Vector &inp
 		if (vdata.validity.RowIsValid(idx)) {
 			auto string_size = data[idx].GetSize();
 			state.total_string_size += string_size;
-			if (string_size >= StringUncompressed::STRING_BLOCK_LIMIT) {
+			if (string_size >= StringUncompressed::GetStringBlockLimit(state.info.GetBlockSize())) {
 				state.overflow_strings++;
 			}
 		}
@@ -62,6 +64,19 @@ idx_t UncompressedStringStorage::StringFinalAnalyze(AnalyzeState &state_p) {
 //===--------------------------------------------------------------------===//
 // Scan
 //===--------------------------------------------------------------------===//
+void UncompressedStringInitPrefetch(ColumnSegment &segment, PrefetchState &prefetch_state) {
+	prefetch_state.AddBlock(segment.block);
+	auto segment_state = segment.GetSegmentState();
+	if (segment_state) {
+		auto &state = segment_state->Cast<UncompressedStringSegmentState>();
+		auto &block_manager = segment.GetBlockManager();
+		for (auto &block_id : state.on_disk_blocks) {
+			auto block_handle = state.GetHandle(block_manager, block_id);
+			prefetch_state.AddBlock(block_handle);
+		}
+	}
+}
+
 unique_ptr<SegmentScanState> UncompressedStringStorage::StringInitScan(ColumnSegment &segment) {
 	auto result = make_uniq<StringScanState>();
 	auto &buffer_manager = BufferManager::GetBufferManager(segment.db);
@@ -146,11 +161,10 @@ void UncompressedStringStorage::StringFetchRow(ColumnSegment &segment, ColumnFet
 struct SerializedStringSegmentState : public ColumnSegmentState {
 	SerializedStringSegmentState() {
 	}
-	explicit SerializedStringSegmentState(vector<block_id_t> blocks_p) : blocks(std::move(blocks_p)) {
+	explicit SerializedStringSegmentState(vector<block_id_t> blocks_p) {
+		blocks = std::move(blocks_p);
 	}
 
-	vector<block_id_t> blocks;
-
 	void Serialize(Serializer &serializer) const override {
 		serializer.WriteProperty(1, "overflow_blocks", blocks);
 	}
@@ -175,7 +189,7 @@ UncompressedStringStorage::StringInitSegment(ColumnSegment &segment, block_id_t
 	return std::move(result);
 }
 
-idx_t UncompressedStringStorage::FinalizeAppend(ColumnSegment &segment, SegmentStatistics &stats) {
+idx_t UncompressedStringStorage::FinalizeAppend(ColumnSegment &segment, SegmentStatistics &) {
 	auto &buffer_manager = BufferManager::GetBufferManager(segment.db);
 	auto handle = buffer_manager.Pin(segment.block);
 	auto dict = GetDictionary(segment, handle);
@@ -183,10 +197,13 @@ idx_t UncompressedStringStorage::FinalizeAppend(ColumnSegment &segment, SegmentS
 	// compute the total size required to store this segment
 	auto offset_size = DICTIONARY_HEADER_SIZE + segment.count * sizeof(int32_t);
 	auto total_size = offset_size + dict.size;
-	if (total_size >= COMPACTION_FLUSH_LIMIT) {
+
+	CompressionInfo info(segment.GetBlockManager().GetBlockSize());
+	if (total_size >= info.GetCompactionFlushLimit()) {
 		// the block is full enough, don't bother moving around the dictionary
 		return segment.SegmentSize();
 	}
+
 	// the block has space left: figure out how much space we can save
 	auto move_amount = segment.SegmentSize() - total_size;
 	// move the dictionary so it lines up exactly with the offsets
@@ -240,7 +257,7 @@ CompressionFunction StringUncompressed::GetFunction(PhysicalType data_type) {
 	                           UncompressedStringStorage::StringInitAppend, UncompressedStringStorage::StringAppend,
 	                           UncompressedStringStorage::FinalizeAppend, nullptr,
 	                           UncompressedStringStorage::SerializeState, UncompressedStringStorage::DeserializeState,
-	                           UncompressedStringStorage::CleanupState);
+	                           UncompressedStringStorage::CleanupState, UncompressedStringInitPrefetch);
 }
 
 //===--------------------------------------------------------------------===//
@@ -293,7 +310,7 @@ void UncompressedStringStorage::WriteStringMemory(ColumnSegment &segment, string
 	if (!state.head || state.head->offset + total_length >= state.head->size) {
 		// string does not fit, allocate space for it
 		// create a new string block
-		idx_t alloc_size = MaxValue<idx_t>(total_length, Storage::BLOCK_SIZE);
+		auto alloc_size = MaxValue<idx_t>(total_length, segment.GetBlockManager().GetBlockSize());
 		auto new_block = make_uniq<StringBlock>();
 		new_block->offset = 0;
 		new_block->size = alloc_size;
@@ -321,12 +338,13 @@ void UncompressedStringStorage::WriteStringMemory(ColumnSegment &segment, string
 
 string_t UncompressedStringStorage::ReadOverflowString(ColumnSegment &segment, Vector &result, block_id_t block,
                                                        int32_t offset) {
-	D_ASSERT(block != INVALID_BLOCK);
-	D_ASSERT(offset < int32_t(Storage::BLOCK_SIZE));
-
 	auto &block_manager = segment.GetBlockManager();
 	auto &buffer_manager = block_manager.buffer_manager;
 	auto &state = segment.GetSegmentState()->Cast<UncompressedStringSegmentState>();
+
+	D_ASSERT(block != INVALID_BLOCK);
+	D_ASSERT(offset < NumericCast<int32_t>(block_manager.GetBlockSize()));
+
 	if (block < MAXIMUM_BLOCK) {
 		// read the overflow string from disk
 		// pin the initial handle and read the length
@@ -339,7 +357,7 @@ string_t UncompressedStringStorage::ReadOverflowString(ColumnSegment &segment, V
 		offset += sizeof(uint32_t);
 
 		// allocate a buffer to store the string
-		auto alloc_size = MaxValue<idx_t>(Storage::BLOCK_SIZE, length);
+		auto alloc_size = MaxValue<idx_t>(block_manager.GetBlockSize(), length);
 		// allocate a buffer to store the compressed string
 		// TODO: profile this to check if we need to reuse buffer
 		auto target_handle = buffer_manager.Allocate(MemoryTag::OVERFLOW_STRINGS, alloc_size);
@@ -347,11 +365,11 @@ string_t UncompressedStringStorage::ReadOverflowString(ColumnSegment &segment, V
 
 		// now append the string to the single buffer
 		while (remaining > 0) {
-			idx_t to_write =
-			    MinValue<idx_t>(remaining, Storage::BLOCK_SIZE - sizeof(block_id_t) - UnsafeNumericCast<idx_t>(offset));
+			idx_t to_write = MinValue<idx_t>(remaining, block_manager.GetBlockSize() - sizeof(block_id_t) -
+			                                                UnsafeNumericCast<idx_t>(offset));
 			memcpy(target_ptr, handle.Ptr() + offset, to_write);
 			remaining -= to_write;
-			offset += to_write;
+			offset += UnsafeNumericCast<int32_t>(to_write);
 			target_ptr += to_write;
 			if (remaining > 0) {
 				// read the next block
@@ -365,16 +383,16 @@ string_t UncompressedStringStorage::ReadOverflowString(ColumnSegment &segment, V
 		auto final_buffer = target_handle.Ptr();
 		StringVector::AddHandle(result, std::move(target_handle));
 		return ReadString(final_buffer, 0, length);
-	} else {
-		// read the overflow string from memory
-		// first pin the handle, if it is not pinned yet
-		auto entry = state.overflow_blocks.find(block);
-		D_ASSERT(entry != state.overflow_blocks.end());
-		auto handle = buffer_manager.Pin(entry->second.get().block);
-		auto final_buffer = handle.Ptr();
-		StringVector::AddHandle(result, std::move(handle));
-		return ReadStringWithLength(final_buffer, offset);
 	}
+
+	// read the overflow string from memory
+	// first pin the handle, if it is not pinned yet
+	auto entry = state.overflow_blocks.find(block);
+	D_ASSERT(entry != state.overflow_blocks.end());
+	auto handle = buffer_manager.Pin(entry->second.get().block);
+	auto final_buffer = handle.Ptr();
+	StringVector::AddHandle(result, std::move(handle));
+	return ReadStringWithLength(final_buffer, offset);
 }
 
 string_t UncompressedStringStorage::ReadString(data_ptr_t target, int32_t offset, uint32_t string_length) {
@@ -402,44 +420,45 @@ void UncompressedStringStorage::ReadStringMarker(data_ptr_t target, block_id_t &
 	memcpy(&offset, target, sizeof(int32_t));
 }
 
-string_location_t UncompressedStringStorage::FetchStringLocation(StringDictionaryContainer dict, data_ptr_t baseptr,
-                                                                 int32_t dict_offset) {
-	D_ASSERT(dict_offset >= -1 * int32_t(Storage::BLOCK_SIZE) && dict_offset <= int32_t(Storage::BLOCK_SIZE));
-	if (dict_offset < 0) {
-		string_location_t result;
-		ReadStringMarker(baseptr + dict.end - idx_t(-1 * dict_offset), result.block_id, result.offset);
-		return result;
-	} else {
+string_location_t UncompressedStringStorage::FetchStringLocation(StringDictionaryContainer dict, data_ptr_t base_ptr,
+                                                                 int32_t dict_offset, const idx_t block_size) {
+	D_ASSERT(dict_offset + NumericCast<int32_t>(block_size) >= 0 && dict_offset <= NumericCast<int32_t>(block_size));
+	if (dict_offset >= 0) {
 		return string_location_t(INVALID_BLOCK, dict_offset);
 	}
+
+	string_location_t result;
+	ReadStringMarker(base_ptr + dict.end - NumericCast<idx_t>(-1 * dict_offset), result.block_id, result.offset);
+	return result;
 }
 
 string_t UncompressedStringStorage::FetchStringFromDict(ColumnSegment &segment, StringDictionaryContainer dict,
-                                                        Vector &result, data_ptr_t baseptr, int32_t dict_offset,
+                                                        Vector &result, data_ptr_t base_ptr, int32_t dict_offset,
                                                         uint32_t string_length) {
-	// fetch base data
-	D_ASSERT(dict_offset <= int32_t(Storage::BLOCK_SIZE));
-	string_location_t location = FetchStringLocation(dict, baseptr, dict_offset);
-	return FetchString(segment, dict, result, baseptr, location, string_length);
+	// Fetch the base data.
+	auto block_size = segment.GetBlockManager().GetBlockSize();
+	D_ASSERT(dict_offset <= NumericCast<int32_t>(block_size));
+	string_location_t location = FetchStringLocation(dict, base_ptr, dict_offset, block_size);
+	return FetchString(segment, dict, result, base_ptr, location, string_length);
 }
 
 string_t UncompressedStringStorage::FetchString(ColumnSegment &segment, StringDictionaryContainer dict, Vector &result,
-                                                data_ptr_t baseptr, string_location_t location,
+                                                data_ptr_t base_ptr, string_location_t location,
                                                 uint32_t string_length) {
 	if (location.block_id != INVALID_BLOCK) {
 		// big string marker: read from separate block
 		return ReadOverflowString(segment, result, location.block_id, location.offset);
-	} else {
-		if (location.offset == 0) {
-			return string_t(nullptr, 0);
-		}
-		// normal string: read string from this block
-		auto dict_end = baseptr + dict.end;
-		auto dict_pos = dict_end - location.offset;
-
-		auto str_ptr = char_ptr_cast(dict_pos);
-		return string_t(str_ptr, string_length);
 	}
+	if (location.offset == 0) {
+		return string_t(nullptr, 0);
+	}
+
+	// normal string: read string from this block
+	auto dict_end = base_ptr + dict.end;
+	auto dict_pos = dict_end - location.offset;
+
+	auto str_ptr = char_ptr_cast(dict_pos);
+	return string_t(str_ptr, string_length);
 }
 
 } // namespace duckdb

package/src/duckdb/src/storage/compression/uncompressed.cpp

@@ -1,5 +1,5 @@
-#include "duckdb/function/compression/compression.hpp"
 #include "duckdb/storage/segment/uncompressed.hpp"
+#include "duckdb/function/compression/compression.hpp"
 
 namespace duckdb {
 
@@ -30,7 +30,7 @@ CompressionFunction UncompressedFun::GetFunction(PhysicalType type) {
 	}
 }
 
-bool UncompressedFun::TypeIsSupported(PhysicalType type) {
+bool UncompressedFun::TypeIsSupported(const PhysicalType) {
 	return true;
 }
 

package/src/duckdb/src/storage/compression/validity_uncompressed.cpp

@@ -1,11 +1,11 @@
-#include "duckdb/storage/segment/uncompressed.hpp"
-#include "duckdb/storage/buffer_manager.hpp"
+#include "duckdb/common/types/null_value.hpp"
 #include "duckdb/common/types/vector.hpp"
+#include "duckdb/function/compression_function.hpp"
+#include "duckdb/storage/buffer_manager.hpp"
+#include "duckdb/storage/segment/uncompressed.hpp"
 #include "duckdb/storage/table/append_state.hpp"
-
-#include "duckdb/common/types/null_value.hpp"
+#include "duckdb/storage/table/column_data.hpp"
 #include "duckdb/storage/table/column_segment.hpp"
-#include "duckdb/function/compression_function.hpp"
 #include "duckdb/storage/table/scan_state.hpp"
 
 namespace duckdb {
@@ -177,14 +177,15 @@ const validity_t ValidityUncompressed::UPPER_MASKS[] = {0x0,
 // Analyze
 //===--------------------------------------------------------------------===//
 struct ValidityAnalyzeState : public AnalyzeState {
-	ValidityAnalyzeState() : count(0) {
+	explicit ValidityAnalyzeState(const CompressionInfo &info) : AnalyzeState(info), count(0) {
 	}
 
 	idx_t count;
 };
 
 unique_ptr<AnalyzeState> ValidityInitAnalyze(ColumnData &col_data, PhysicalType type) {
-	return make_uniq<ValidityAnalyzeState>();
+	CompressionInfo info(col_data.GetBlockManager().GetBlockSize());
+	return make_uniq<ValidityAnalyzeState>(info);
 }
 
 bool ValidityAnalyze(AnalyzeState &state_p, Vector &input, idx_t count) {

package/src/duckdb/src/storage/data_pointer.cpp

@@ -6,10 +6,32 @@
 
 namespace duckdb {
 
+DataPointer::DataPointer(BaseStatistics stats) : statistics(std::move(stats)) {
+}
+
+DataPointer::DataPointer(DataPointer &&other) noexcept : statistics(std::move(other.statistics)) {
+	std::swap(row_start, other.row_start);
+	std::swap(tuple_count, other.tuple_count);
+	std::swap(block_pointer, other.block_pointer);
+	std::swap(compression_type, other.compression_type);
+	std::swap(segment_state, other.segment_state);
+}
+
+DataPointer &DataPointer::operator=(DataPointer &&other) noexcept {
+	std::swap(row_start, other.row_start);
+	std::swap(tuple_count, other.tuple_count);
+	std::swap(block_pointer, other.block_pointer);
+	std::swap(compression_type, other.compression_type);
+	std::swap(statistics, other.statistics);
+	std::swap(segment_state, other.segment_state);
+	return *this;
+}
+
 unique_ptr<ColumnSegmentState> ColumnSegmentState::Deserialize(Deserializer &deserializer) {
 	auto compression_type = deserializer.Get<CompressionType>();
 	auto &db = deserializer.Get<DatabaseInstance &>();
 	auto &type = deserializer.Get<const LogicalType &>();
+
 	auto compression_function = DBConfig::GetConfig(db).GetCompressionFunction(compression_type, type.InternalType());
 	if (!compression_function || !compression_function->deserialize_state) {
 		throw SerializationException("Deserializing a ColumnSegmentState but could not find deserialize method");

package/src/duckdb/src/storage/data_table.cpp

@@ -225,13 +225,16 @@ TableIOManager &TableIOManager::Get(DataTable &table) {
 //===--------------------------------------------------------------------===//
 void DataTable::InitializeScan(TableScanState &state, const vector<column_t> &column_ids,
                                TableFilterSet *table_filters) {
-	state.checkpoint_lock = info->checkpoint_lock.GetSharedLock();
+	if (!state.checkpoint_lock) {
+		state.checkpoint_lock = make_shared_ptr<CheckpointLock>(info->checkpoint_lock.GetSharedLock());
+	}
 	state.Initialize(column_ids, table_filters);
 	row_groups->InitializeScan(state.table_state, column_ids, table_filters);
 }
 
 void DataTable::InitializeScan(DuckTransaction &transaction, TableScanState &state, const vector<column_t> &column_ids,
                                TableFilterSet *table_filters) {
+	state.checkpoint_lock = transaction.SharedLockTable(*info);
 	auto &local_storage = LocalStorage::Get(transaction);
 	InitializeScan(state, column_ids, table_filters);
 	local_storage.InitializeScan(*this, state.local_state, table_filters);
@@ -239,7 +242,9 @@ void DataTable::InitializeScan(DuckTransaction &transaction, TableScanState &sta
 
 void DataTable::InitializeScanWithOffset(TableScanState &state, const vector<column_t> &column_ids, idx_t start_row,
                                          idx_t end_row) {
-	state.checkpoint_lock = info->checkpoint_lock.GetSharedLock();
+	if (!state.checkpoint_lock) {
+		state.checkpoint_lock = make_shared_ptr<CheckpointLock>(info->checkpoint_lock.GetSharedLock());
+	}
 	state.Initialize(column_ids);
 	row_groups->InitializeScanWithOffset(state.table_state, column_ids, start_row, end_row);
 }
@@ -255,7 +260,8 @@ idx_t DataTable::MaxThreads(ClientContext &context) {
 
 void DataTable::InitializeParallelScan(ClientContext &context, ParallelTableScanState &state) {
 	auto &local_storage = LocalStorage::Get(context, db);
-	state.checkpoint_lock = info->checkpoint_lock.GetSharedLock();
+	auto &transaction = DuckTransaction::Get(context, db);
+	state.checkpoint_lock = transaction.SharedLockTable(*info);
 	row_groups->InitializeParallelScan(state.scan_state);
 
 	local_storage.InitializeParallelScan(*this, state.local_state);
@@ -326,6 +332,10 @@ void DataTable::VacuumIndexes() {
 	});
 }
 
+void DataTable::CleanupAppend(transaction_t lowest_transaction, idx_t start, idx_t count) {
+	row_groups->CleanupAppend(lowest_transaction, start, count);
+}
+
 bool DataTable::IndexNameIsUnique(const string &name) {
 	return info->indexes.NameIsUnique(name);
 }
@@ -515,7 +525,7 @@ void DataTable::VerifyForeignKeyConstraint(const BoundForeignKeyConstraint &bfk,
 	}
 
 	auto &table_entry_ptr =
-	    Catalog::GetEntry<TableCatalogEntry>(context, INVALID_CATALOG, bfk.info.schema, bfk.info.table);
+	    Catalog::GetEntry<TableCatalogEntry>(context, db.GetName(), bfk.info.schema, bfk.info.table);
 	// make the data chunk to check
 	vector<LogicalType> types;
 	for (auto &col : table_entry_ptr.GetColumns().Physical()) {
@@ -935,16 +945,34 @@ void DataTable::ScanTableSegment(idx_t row_start, idx_t count, const std::functi
 	}
 }
 
-void DataTable::MergeStorage(RowGroupCollection &data, TableIndexList &indexes) {
-	row_groups->MergeStorage(data);
+void DataTable::MergeStorage(RowGroupCollection &data, TableIndexList &,
+                             optional_ptr<StorageCommitState> commit_state) {
+	row_groups->MergeStorage(data, this, commit_state);
 	row_groups->Verify();
 }
 
-void DataTable::WriteToLog(WriteAheadLog &log, idx_t row_start, idx_t count) {
-	if (log.skip_writing) {
-		return;
-	}
+void DataTable::WriteToLog(WriteAheadLog &log, idx_t row_start, idx_t count,
+                           optional_ptr<StorageCommitState> commit_state) {
 	log.WriteSetTable(info->schema, info->table);
+	if (commit_state) {
+		idx_t optimistic_count = 0;
+		auto entry = commit_state->GetRowGroupData(*this, row_start, optimistic_count);
+		if (entry) {
+			D_ASSERT(optimistic_count > 0);
+			log.WriteRowGroupData(*entry);
+			if (optimistic_count > count) {
+				throw InternalException(
+				    "Optimistically written count cannot exceed actual count (got %llu, but expected count is %llu)",
+				    optimistic_count, count);
+			}
+			// write any remaining (non-optimistically written) rows to the WAL normally
+			row_start += optimistic_count;
+			count -= optimistic_count;
+			if (count == 0) {
+				return;
+			}
+		}
+	}
 	ScanTableSegment(row_start, count, [&](DataChunk &chunk) { log.WriteInsert(chunk); });
 }
 
@@ -1334,13 +1362,14 @@ void DataTable::Update(TableUpdateState &state, ClientContext &context, Vector &
 
 	// otherwise global storage
 	if (n_global_update > 0) {
+		auto &transaction = DuckTransaction::Get(context, db);
 		updates_slice.Slice(updates, sel_global_update, n_global_update);
 		updates_slice.Flatten();
 		row_ids_slice.Slice(row_ids, sel_global_update, n_global_update);
 		row_ids_slice.Flatten(n_global_update);
 
-		row_groups->Update(DuckTransaction::Get(context, db), FlatVector::GetData<row_t>(row_ids_slice), column_ids,
-		                   updates_slice);
+		transaction.UpdateCollection(row_groups);
+		row_groups->Update(transaction, FlatVector::GetData<row_t>(row_ids_slice), column_ids, updates_slice);
 	}
 }
 

package/src/duckdb/src/storage/local_storage.cpp

@@ -137,7 +137,7 @@ ErrorData LocalTableStorage::AppendToIndexes(DuckTransaction &transaction, RowGr
 		if (error.HasError()) {
 			return false;
 		}
-		start_row += chunk.size();
+		start_row += UnsafeNumericCast<row_t>(chunk.size());
 		return true;
 	});
 	return error;
@@ -180,7 +180,7 @@ void LocalTableStorage::AppendToIndexes(DuckTransaction &transaction, TableAppen
 				return false;
 			} // LCOV_EXCL_STOP
 
-			current_row += chunk.size();
+			current_row += UnsafeNumericCast<row_t>(chunk.size());
 			if (current_row >= append_state.current_row) {
 				// finished deleting all rows from the index: abort now
 				return false;
@@ -317,7 +317,7 @@ LocalStorage &LocalStorage::Get(ClientContext &context, Catalog &catalog) {
 void LocalStorage::InitializeScan(DataTable &table, CollectionScanState &state,
                                   optional_ptr<TableFilterSet> table_filters) {
 	auto storage = table_manager.GetStorage(table);
-	if (storage == nullptr) {
+	if (storage == nullptr || storage->row_groups->GetTotalRows() == 0) {
 		return;
 	}
 	storage->InitializeScan(state, table_filters);
@@ -385,7 +385,7 @@ void LocalStorage::LocalMerge(DataTable &table, RowGroupCollection &collection)
 			error.Throw();
 		}
 	}
-	storage.row_groups->MergeStorage(collection);
+	storage.row_groups->MergeStorage(collection, nullptr, nullptr);
 	storage.merged_storage = true;
 }
 
@@ -435,8 +435,14 @@ void LocalStorage::Update(DataTable &table, Vector &row_ids, const vector<Physic
 	storage->row_groups->Update(TransactionData(0, 0), ids, column_ids, updates);
 }
 
-void LocalStorage::Flush(DataTable &table, LocalTableStorage &storage) {
+void LocalStorage::Flush(DataTable &table, LocalTableStorage &storage, optional_ptr<StorageCommitState> commit_state) {
+	if (storage.is_dropped) {
+		return;
+	}
 	if (storage.row_groups->GetTotalRows() <= storage.deleted_rows) {
+		// all rows that we added were deleted
+		// rollback any partial blocks that are still outstanding
+		storage.Rollback();
 		return;
 	}
 	idx_t append_count = storage.row_groups->GetTotalRows() - storage.deleted_rows;
@@ -458,7 +464,7 @@ void LocalStorage::Flush(DataTable &table, LocalTableStorage &storage) {
 			storage.AppendToIndexes(transaction, append_state, append_count, false);
 		}
 		// finally move over the row groups
-		table.MergeStorage(*storage.row_groups, storage.indexes);
+		table.MergeStorage(*storage.row_groups, storage.indexes, commit_state);
 	} else {
 		// check if we have written data
 		// if we have, we cannot merge to disk after all
@@ -472,7 +478,7 @@ void LocalStorage::Flush(DataTable &table, LocalTableStorage &storage) {
 	table.VacuumIndexes();
 }
 
-void LocalStorage::Commit(LocalStorage::CommitState &commit_state, DuckTransaction &transaction) {
+void LocalStorage::Commit(optional_ptr<StorageCommitState> commit_state) {
 	// commit local storage
 	// iterate over all entries in the table storage map and commit them
 	// after this, the local storage is no longer required and can be cleared
@@ -480,7 +486,7 @@ void LocalStorage::Commit(LocalStorage::CommitState &commit_state, DuckTransacti
 	for (auto &entry : table_storage) {
 		auto table = entry.first;
 		auto storage = entry.second.get();
-		Flush(table, *storage);
+		Flush(table, *storage, commit_state);
 		entry.second.reset();
 	}
 }
@@ -508,6 +514,14 @@ idx_t LocalStorage::AddedRows(DataTable &table) {
 	return storage->row_groups->GetTotalRows() - storage->deleted_rows;
 }
 
+void LocalStorage::DropTable(DataTable &table) {
+	auto storage = table_manager.GetStorage(table);
+	if (!storage) {
+		return;
+	}
+	storage->is_dropped = true;
+}
+
 void LocalStorage::MoveStorage(DataTable &old_dt, DataTable &new_dt) {
 	// check if there are any pending appends for the old version of the table
 	auto new_storage = table_manager.MoveEntry(old_dt);

package/src/duckdb/src/storage/metadata/metadata_manager.cpp

@@ -26,7 +26,7 @@ MetadataHandle MetadataManager::AllocateHandle() {
 			break;
 		}
 	}
-	if (free_block == INVALID_BLOCK) {
+	if (free_block == INVALID_BLOCK || free_block > PeekNextBlockId()) {
 		free_block = AllocateNewBlock();
 	}
 	D_ASSERT(free_block != INVALID_BLOCK);
@@ -52,7 +52,7 @@ MetadataHandle MetadataManager::AllocateHandle() {
 
 MetadataHandle MetadataManager::Pin(MetadataPointer pointer) {
 	D_ASSERT(pointer.index < METADATA_BLOCK_COUNT);
-	auto &block = blocks[pointer.block_index];
+	auto &block = blocks[UnsafeNumericCast<int64_t>(pointer.block_index)];
 
 	MetadataHandle handle;
 	handle.pointer.block_index = pointer.block_index;
@@ -67,28 +67,27 @@ void MetadataManager::ConvertToTransient(MetadataBlock &block) {
 
 	// allocate a new transient block to replace it
 	shared_ptr<BlockHandle> new_block;
-	auto new_buffer = buffer_manager.Allocate(MemoryTag::METADATA, Storage::BLOCK_SIZE, false, &new_block);
+	auto new_buffer = buffer_manager.Allocate(MemoryTag::METADATA, block_manager.GetBlockSize(), false, &new_block);
 
 	// copy the data to the transient block
-	memcpy(new_buffer.Ptr(), old_buffer.Ptr(), Storage::BLOCK_SIZE);
-
+	memcpy(new_buffer.Ptr(), old_buffer.Ptr(), block_manager.GetBlockSize());
 	block.block = std::move(new_block);
 
 	// unregister the old block
-	block_manager.UnregisterBlock(block.block_id, false);
+	block_manager.UnregisterBlock(block.block_id);
 }
 
 block_id_t MetadataManager::AllocateNewBlock() {
 	auto new_block_id = GetNextBlockId();
 
 	MetadataBlock new_block;
-	auto handle = buffer_manager.Allocate(MemoryTag::METADATA, Storage::BLOCK_SIZE, false, &new_block.block);
+	auto handle = buffer_manager.Allocate(MemoryTag::METADATA, block_manager.GetBlockSize(), false, &new_block.block);
 	new_block.block_id = new_block_id;
 	for (idx_t i = 0; i < METADATA_BLOCK_COUNT; i++) {
 		new_block.free_blocks.push_back(NumericCast<uint8_t>(METADATA_BLOCK_COUNT - i - 1));
 	}
 	// zero-initialize the handle
-	memset(handle.Ptr(), 0, Storage::BLOCK_SIZE);
+	memset(handle.Ptr(), 0, block_manager.GetBlockSize());
 	AddBlock(std::move(new_block));
 	return new_block_id;
 }
@@ -147,22 +146,22 @@ MetadataPointer MetadataManager::RegisterDiskPointer(MetaBlockPointer pointer) {
 	return FromDiskPointer(pointer);
 }
 
-BlockPointer MetadataManager::ToBlockPointer(MetaBlockPointer meta_pointer) {
+BlockPointer MetadataManager::ToBlockPointer(MetaBlockPointer meta_pointer, const idx_t metadata_block_size) {
 	BlockPointer result;
 	result.block_id = meta_pointer.GetBlockId();
-	result.offset = meta_pointer.GetBlockIndex() * MetadataManager::METADATA_BLOCK_SIZE + meta_pointer.offset;
-	D_ASSERT(result.offset < MetadataManager::METADATA_BLOCK_SIZE * MetadataManager::METADATA_BLOCK_COUNT);
+	result.offset = meta_pointer.GetBlockIndex() * NumericCast<uint32_t>(metadata_block_size) + meta_pointer.offset;
+	D_ASSERT(result.offset < metadata_block_size * MetadataManager::METADATA_BLOCK_COUNT);
 	return result;
 }
 
-MetaBlockPointer MetadataManager::FromBlockPointer(BlockPointer block_pointer) {
+MetaBlockPointer MetadataManager::FromBlockPointer(BlockPointer block_pointer, const idx_t metadata_block_size) {
 	if (!block_pointer.IsValid()) {
 		return MetaBlockPointer();
 	}
-	idx_t index = block_pointer.offset / MetadataManager::METADATA_BLOCK_SIZE;
-	auto offset = block_pointer.offset % MetadataManager::METADATA_BLOCK_SIZE;
+	idx_t index = block_pointer.offset / metadata_block_size;
+	auto offset = block_pointer.offset % metadata_block_size;
 	D_ASSERT(index < MetadataManager::METADATA_BLOCK_COUNT);
-	D_ASSERT(offset < MetadataManager::METADATA_BLOCK_SIZE);
+	D_ASSERT(offset < metadata_block_size);
 	MetaBlockPointer result;
 	result.block_pointer = idx_t(block_pointer.block_id) | index << 56ULL;
 	result.offset = UnsafeNumericCast<uint32_t>(offset);
@@ -174,13 +173,14 @@ idx_t MetadataManager::BlockCount() {
 }
 
 void MetadataManager::Flush() {
-	const idx_t total_metadata_size = MetadataManager::METADATA_BLOCK_SIZE * MetadataManager::METADATA_BLOCK_COUNT;
+	const idx_t total_metadata_size = GetMetadataBlockSize() * MetadataManager::METADATA_BLOCK_COUNT;
+
 	// write the blocks of the metadata manager to disk
 	for (auto &kv : blocks) {
 		auto &block = kv.second;
 		auto handle = buffer_manager.Pin(block.block);
 		// there are a few bytes left-over at the end of the block, zero-initialize them
-		memset(handle.Ptr() + total_metadata_size, 0, Storage::BLOCK_SIZE - total_metadata_size);
+		memset(handle.Ptr() + total_metadata_size, 0, block_manager.GetBlockSize() - total_metadata_size);
 		D_ASSERT(kv.first == block.block_id);
 		if (block.block->BlockId() >= MAXIMUM_BLOCK) {
 			// temporary block - convert to persistent
@@ -220,6 +220,10 @@ void MetadataBlock::Write(WriteStream &sink) {
 	sink.Write<idx_t>(FreeBlocksToInteger());
 }
 
+idx_t MetadataManager::GetMetadataBlockSize() const {
+	return AlignValueFloor(block_manager.GetBlockSize() / METADATA_BLOCK_COUNT);
+}
+
 MetadataBlock MetadataBlock::Read(ReadStream &source) {
 	MetadataBlock result;
 	result.block_id = source.Read<block_id_t>();
@@ -315,6 +319,18 @@ vector<MetadataBlockInfo> MetadataManager::GetMetadataInfo() const {
 	return result;
 }
 
+vector<shared_ptr<BlockHandle>> MetadataManager::GetBlocks() const {
+	vector<shared_ptr<BlockHandle>> result;
+	for (auto &entry : blocks) {
+		result.push_back(entry.second.block);
+	}
+	return result;
+}
+
+block_id_t MetadataManager::PeekNextBlockId() {
+	return block_manager.PeekFreeBlockId();
+}
+
 block_id_t MetadataManager::GetNextBlockId() {
 	return block_manager.GetFreeBlockId();
 }