Merge pull request #11177 from Mytherin/fixrevertappendinternal

Fix RevertAppendInternal
duckdb · Mar 15, 2024 · badf7a2 · badf7a2
2 parents 01a9015 + dca7948
commit badf7a2
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 8 deletions.
diff --git a/src/include/duckdb/storage/table/segment_tree.hpp b/src/include/duckdb/storage/table/segment_tree.hpp
@@ -69,6 +69,9 @@ class SegmentTree {
 	}
 	idx_t GetSegmentCount() {
 		auto l = Lock();
+		return GetSegmentCount(l);
+	}
+	idx_t GetSegmentCount(SegmentLock &l) {
 		return nodes.size();
 	}
 	//! Gets a pointer to the nth segment. Negative numbers start from the back.
@@ -185,6 +188,9 @@ class SegmentTree {
 	//! Get the segment index of the column segment for the given row
 	idx_t GetSegmentIndex(SegmentLock &l, idx_t row_number) {
 		idx_t segment_index;
+		D_ASSERT(!nodes.empty());
+		D_ASSERT(row_number >= nodes[0].row_start);
+		D_ASSERT(row_number < nodes.back().row_start + nodes.back().node->count);
 		if (TryGetSegmentIndex(l, row_number, segment_index)) {
 			return segment_index;
 		}
@@ -207,9 +213,6 @@ class SegmentTree {
 		if (nodes.empty()) {
 			return false;
 		}
-		D_ASSERT(!nodes.empty());
-		D_ASSERT(row_number >= nodes[0].row_start);
-		D_ASSERT(row_number < nodes.back().row_start + nodes.back().node->count);
 		idx_t lower = 0;
 		idx_t upper = nodes.size() - 1;
 		// binary search to find the node

diff --git a/src/storage/table/row_group_collection.cpp b/src/storage/table/row_group_collection.cpp
@@ -417,14 +417,20 @@ void RowGroupCollection::CommitAppend(transaction_t commit_id, idx_t row_start,
 }
 
 void RowGroupCollection::RevertAppendInternal(idx_t start_row) {
-	if (total_rows <= start_row) {
-		return;
-	}
 	total_rows = start_row;
 
 	auto l = row_groups->Lock();
-	// find the segment index that the current row belongs to
-	idx_t segment_index = row_groups->GetSegmentIndex(l, start_row);
+	idx_t segment_count = row_groups->GetSegmentCount(l);
+	if (segment_count == 0) {
+		// we have no segments to revert
+		return;
+	}
+	idx_t segment_index;
+	// find the segment index that the start row belongs to
+	if (!row_groups->TryGetSegmentIndex(l, start_row, segment_index)) {
+		// revert from the last segment
+		segment_index = segment_count - 1;
+	}
 	auto &segment = *row_groups->GetSegmentByIndex(l, segment_index);
 
 	// remove any segments AFTER this segment: they should be deleted entirely

diff --git a/test/sql/transactions/test_index_rollback_flushed_data.test b/test/sql/transactions/test_index_rollback_flushed_data.test
@@ -0,0 +1,44 @@
+# name: test/sql/transactions/test_index_rollback_flushed_data.test
+# description: Test that we revert the global storage correctly after a constraint violation
+# group: [transactions]
+
+require skip_reload
+
+statement ok
+PRAGMA enable_verification
+
+statement ok con1
+CREATE TABLE integers(i INTEGER UNIQUE);
+
+statement ok con1
+BEGIN TRANSACTION;
+
+statement ok con2
+BEGIN TRANSACTION;
+
+statement ok con1
+INSERT INTO integers VALUES (-10);
+
+statement ok con2
+INSERT INTO integers SELECT range FROM range(2, 4097, 1);
+
+# constraint violation
+statement ok con2
+INSERT INTO integers VALUES (-10);
+
+# con1 commits first
+statement ok con1
+COMMIT;
+
+# con2 fails to commit because of the conflict
+statement error con2
+COMMIT;
+----
+
+statement ok
+INSERT INTO integers SELECT i FROM range(2, 4097, 1) t1(i)
+
+query I
+SELECT MAX(i) FROM integers
+----
+4096