From a107761850de3d16302a304df04881b1e888109a Mon Sep 17 00:00:00 2001 From: Xinli shang Date: Mon, 7 Sep 2026 16:20:41 -0700 Subject: [PATCH] feat: expose ReplacePartitions on Table and Transaction Adds NewReplacePartitions() to Table and Transaction, with the StaticTable override returning NotSupported. Routes the existing ReplacePartitions tests through Table::NewReplacePartitions() and adds a transaction-scoped test. Part of #637, PR2 of 2 for #775. --- src/iceberg/table.cc | 11 +++++ src/iceberg/table.h | 6 +++ src/iceberg/test/replace_partitions_test.cc | 52 +++++++++++++++++++-- src/iceberg/test/table_test.cc | 1 + src/iceberg/transaction.cc | 8 ++++ src/iceberg/transaction.h | 4 ++ 6 files changed, 78 insertions(+), 4 deletions(-) diff --git a/src/iceberg/table.cc b/src/iceberg/table.cc index f9e7026f9..c512c11df 100644 --- a/src/iceberg/table.cc +++ b/src/iceberg/table.cc @@ -36,6 +36,7 @@ #include "iceberg/update/fast_append.h" #include "iceberg/update/merge_append.h" #include "iceberg/update/overwrite_files.h" +#include "iceberg/update/replace_partitions.h" #include "iceberg/update/rewrite_files.h" #include "iceberg/update/row_delta.h" #include "iceberg/update/set_snapshot.h" @@ -259,6 +260,12 @@ Result> Table::NewRewriteFiles() { return RewriteFiles::Make(name().name, std::move(ctx)); } +Result> Table::NewReplacePartitions() { + ICEBERG_ASSIGN_OR_RAISE( + auto ctx, TransactionContext::Make(shared_from_this(), TransactionKind::kUpdate)); + return ReplacePartitions::Make(name().name, std::move(ctx)); +} + Result> Table::NewUpdateStatistics() { ICEBERG_ASSIGN_OR_RAISE( auto ctx, TransactionContext::Make(shared_from_this(), TransactionKind::kUpdate)); @@ -379,6 +386,10 @@ Result> StaticTable::NewRewriteFiles() { return NotSupported("Cannot create a rewrite files for a static table"); } +Result> StaticTable::NewReplacePartitions() { + return NotSupported("Cannot replace partitions for a static table"); +} + Result> StaticTable::NewSnapshotManager() { return NotSupported("Cannot create a snapshot manager for a static table"); } diff --git a/src/iceberg/table.h b/src/iceberg/table.h index eb028c738..4c0470cb9 100644 --- a/src/iceberg/table.h +++ b/src/iceberg/table.h @@ -205,6 +205,10 @@ class ICEBERG_EXPORT Table : public std::enable_shared_from_this { /// changes. virtual Result> NewRewriteFiles(); + /// \brief Create a new ReplacePartitions to dynamically overwrite partitions and commit + /// the changes. + virtual Result> NewReplacePartitions(); + /// \brief Create a new SnapshotManager to manage snapshots and snapshot references. virtual Result> NewSnapshotManager(); @@ -287,6 +291,8 @@ class ICEBERG_EXPORT StaticTable : public Table { Result> NewRewriteFiles() override; + Result> NewReplacePartitions() override; + Result> NewSnapshotManager() override; private: diff --git a/src/iceberg/test/replace_partitions_test.cc b/src/iceberg/test/replace_partitions_test.cc index 86f4f20a2..159f2fd98 100644 --- a/src/iceberg/test/replace_partitions_test.cc +++ b/src/iceberg/test/replace_partitions_test.cc @@ -216,10 +216,8 @@ class ReplacePartitionsTest : public UpdateTestBase { return paths; } - Result> NewReplace() { - ICEBERG_ASSIGN_OR_RAISE(auto ctx, - TransactionContext::Make(table_, TransactionKind::kUpdate)); - return ReplacePartitions::Make(TableName(), std::move(ctx)); + Result> NewReplace() { + return table_->NewReplacePartitions(); } int64_t CommitFastAppend(const std::shared_ptr& file) { @@ -252,6 +250,31 @@ TEST_F(ReplacePartitionsTest, OperationIsOverwrite) { EXPECT_EQ(op->operation(), DataOperation::kOverwrite); } +// A replace created from a transaction commits with the rest of the transaction. +TEST_F(ReplacePartitionsTest, TxnNewReplacePartitions) { + CommitFastAppend(file_a_); + CommitFastAppend(file_b_); + + ICEBERG_UNWRAP_OR_FAIL(auto txn, Transaction::Make(table_, TransactionKind::kUpdate)); + ICEBERG_UNWRAP_OR_FAIL(auto op, txn->NewReplacePartitions()); + ASSERT_NE(op, nullptr); + + auto replacement = MakeDataFile("/data/file_a_new.parquet", /*partition_x=*/1L); + op->AddFile(replacement); + EXPECT_THAT(op->Commit(), IsOk()); + EXPECT_THAT(txn->Commit(), IsOk()); + + EXPECT_THAT(table_->Refresh(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto snapshot, table_->current_snapshot()); + EXPECT_EQ(snapshot->summary.at(SnapshotSummaryFields::kOperation), + DataOperation::kOverwrite); + EXPECT_EQ(snapshot->summary.at(SnapshotSummaryFields::kReplacePartitions), "true"); + + ICEBERG_UNWRAP_OR_FAIL(auto paths, LiveDataFilePaths()); + EXPECT_THAT( + paths, ::testing::UnorderedElementsAre(replacement->file_path, file_b_->file_path)); +} + // Replacing a partition drops its existing file and records the summary flag. TEST_F(ReplacePartitionsTest, PartitionedReplaceCommit) { CommitFastAppend(file_a_); @@ -289,6 +312,27 @@ TEST_F(ReplacePartitionsTest, ReplaceLeavesOtherPartitions) { file_b_->file_path, table_location_ + "/data/file_a_new.parquet")); } +// Several files may be staged for one partition; the partition is dropped once +// and all staged files become its new contents. +TEST_F(ReplacePartitionsTest, ReplacePartitionWithMultipleFiles) { + CommitFastAppend(file_a_); + + ICEBERG_UNWRAP_OR_FAIL(auto op, NewReplace()); + auto first = MakeDataFile("/data/file_a_new_1.parquet", /*partition_x=*/1L); + auto second = MakeDataFile("/data/file_a_new_2.parquet", /*partition_x=*/1L); + op->AddFile(first); + op->AddFile(second); + EXPECT_THAT(op->Commit(), IsOk()); + + EXPECT_THAT(table_->Refresh(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto snapshot, table_->current_snapshot()); + EXPECT_EQ(snapshot->summary.at(SnapshotSummaryFields::kAddedDataFiles), "2"); + EXPECT_EQ(snapshot->summary.at(SnapshotSummaryFields::kDeletedDataFiles), "1"); + + ICEBERG_UNWRAP_OR_FAIL(auto live, LiveDataFilePaths()); + EXPECT_THAT(live, ::testing::UnorderedElementsAre(first->file_path, second->file_path)); +} + // An unpartitioned spec triggers a table-wide replace of every existing file. TEST_F(ReplacePartitionsTest, UnpartitionedReplacesWholeTable) { CommitFastAppend(file_a_); diff --git a/src/iceberg/test/table_test.cc b/src/iceberg/test/table_test.cc index 94a00e879..a85dc3014 100644 --- a/src/iceberg/test/table_test.cc +++ b/src/iceberg/test/table_test.cc @@ -168,6 +168,7 @@ TEST(StaticTableTest, NewMutatingOperationsAreNotSupported) { EXPECT_THAT(table->NewRowDelta(), IsError(ErrorKind::kNotSupported)); EXPECT_THAT(table->NewOverwrite(), IsError(ErrorKind::kNotSupported)); EXPECT_THAT(table->NewRewriteFiles(), IsError(ErrorKind::kNotSupported)); + EXPECT_THAT(table->NewReplacePartitions(), IsError(ErrorKind::kNotSupported)); EXPECT_THAT(table->NewSnapshotManager(), IsError(ErrorKind::kNotSupported)); } diff --git a/src/iceberg/transaction.cc b/src/iceberg/transaction.cc index 80d39c8a8..bacd880de 100644 --- a/src/iceberg/transaction.cc +++ b/src/iceberg/transaction.cc @@ -38,6 +38,7 @@ #include "iceberg/update/merge_append.h" #include "iceberg/update/overwrite_files.h" #include "iceberg/update/pending_update.h" +#include "iceberg/update/replace_partitions.h" #include "iceberg/update/rewrite_files.h" #include "iceberg/update/row_delta.h" #include "iceberg/update/set_snapshot.h" @@ -534,6 +535,13 @@ Result> Transaction::NewRewriteFiles() { return rewrite_files; } +Result> Transaction::NewReplacePartitions() { + ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr replace_partitions, + ReplacePartitions::Make(ctx_->table->name().name, ctx_)); + ICEBERG_RETURN_UNEXPECTED(AddUpdate(replace_partitions)); + return replace_partitions; +} + Result> Transaction::NewUpdateStatistics() { ICEBERG_ASSIGN_OR_RAISE(std::shared_ptr update_statistics, UpdateStatistics::Make(ctx_)); diff --git a/src/iceberg/transaction.h b/src/iceberg/transaction.h index 007b1057e..3ee0372ad 100644 --- a/src/iceberg/transaction.h +++ b/src/iceberg/transaction.h @@ -128,6 +128,10 @@ class ICEBERG_EXPORT Transaction : public std::enable_shared_from_this> NewRewriteFiles(); + /// \brief Create a new ReplacePartitions to dynamically overwrite partitions and commit + /// the changes. + Result> NewReplacePartitions(); + /// \brief Create a new SnapshotManager to manage snapshots. Result> NewSnapshotManager();