[Enhancement] Record generation version on lake PK index sstables (backport #76208) (#76343)

Co-authored-by: xiangguangyxg <110401425+xiangguangyxg@users.noreply.github.com>

Author: mergify[bot]

be/src/storage/lake/lake_persistent_index.cpp

@@ -14,6 +14,8 @@
 
 #include "storage/lake/lake_persistent_index.h"
 
+#include <unordered_map>
+
 #include "fs/fs_util.h"
 #include "fs/key_cache.h"
 #include "serde/column_array_serde.h"
@@ -238,6 +240,10 @@ Status LakePersistentIndex::ingest_sst(const FileMetaPB& sst_meta, const Persist
     sstable_pb.set_filesize(sst_meta.size());
     sstable_pb.set_shared_rssid(rssid);
     sstable_pb.set_shared_version(version);
+    // Record the generation version (PersistentIndexSstablePB.generation_version):
+    // the version at which this ingested sstable becomes visible in this tablet.
+    // Coexists with shared_version (read-time projection) and does not affect it.
+    sstable_pb.set_generation_version(version);
     sstable_pb.set_encryption_meta(sst_meta.encryption_meta());
     // Preserve the shared flag from sst_meta. During tablet split cross-publish,
     // newly-ingested SSTs may be shared with sibling split tablets; losing this flag
@@ -865,7 +871,32 @@ Status LakePersistentIndex::apply_opcompaction(const TabletMetadataPtr& metadata
     return Status::OK();
 }
 
-Status LakePersistentIndex::commit(MetaFileBuilder* builder) {
+void LakePersistentIndex::assign_generation_versions(PersistentIndexSstableMetaPB* new_meta,
+                                                     const PersistentIndexSstableMetaPB& prev_meta,
+                                                     int64_t generation_version) {
+    // Callers resolve |generation_version| to a real publish/reshard version first; a new file must
+    // never be stamped 0 ("unknown"). Benign if it ever slips through in release -- 0 is the
+    // conservative value consumers treat as retain -- so a debug assert suffices.
+    DCHECK_GT(generation_version, 0);
+    std::unordered_map<std::string, int64_t> prev_versions;
+    prev_versions.reserve(prev_meta.sstables_size());
+    for (const auto& s : prev_meta.sstables()) {
+        prev_versions.emplace(s.filename(), s.generation_version());
+    }
+    for (auto& s : *new_meta->mutable_sstables()) {
+        if (s.generation_version() != 0) {
+            continue;
+        }
+        auto it = prev_versions.find(s.filename());
+        if (it != prev_versions.end()) {
+            s.set_generation_version(it->second);
+        } else {
+            s.set_generation_version(generation_version);
+        }
+    }
+}
+
+Status LakePersistentIndex::commit(MetaFileBuilder* builder, int64_t generation_version) {
     if ((too_many_rebuild_files() || too_many_rebuild_rows()) && !_memtable->empty()) {
         // If we have too many files or rows need to be rebuilt,
         // we need to do flush to reduce index rebuild cost later.
@@ -888,6 +919,14 @@ Status LakePersistentIndex::commit(MetaFileBuilder* builder) {
             new_sstable_pb->CopyFrom(sstable_pb);
         }
     }
+    // generation_version == 0 means "use this publish's version"; a non-zero override is passed by
+    // the reshard flush (see UpdateManager::flush_pk_memtable). The builder still holds the
+    // previous persisted sstable_meta here (before finalize_sstable_meta below), which
+    // assign_generation_versions uses to carry existing versions forward.
+    if (generation_version == 0) {
+        generation_version = builder->tablet_meta()->version();
+    }
+    assign_generation_versions(&sstable_meta, builder->tablet_meta()->sstable_meta(), generation_version);
     builder->finalize_sstable_meta(sstable_meta);
     auto [file_cnt, row_cnt] = need_rebuild_counts(*builder->tablet_meta(), sstable_meta);
     _need_rebuild_file_cnt = file_cnt;

be/src/storage/lake/lake_persistent_index.h

@@ -121,7 +121,7 @@ class LakePersistentIndex : public PersistentIndex {
 
     Status apply_opcompaction(const TabletMetadataPtr& metadata, const TxnLogPB_OpCompaction& op_compaction);
 
-    Status commit(MetaFileBuilder* builder);
+    Status commit(MetaFileBuilder* builder, int64_t generation_version = 0);
 
     Status load_from_lake_tablet(TabletManager* tablet_mgr, const TabletMetadataPtr& metadata, int64_t base_version,
                                  const MetaFileBuilder* builder);
@@ -138,6 +138,14 @@ class LakePersistentIndex : public PersistentIndex {
     static void pick_sstables_for_merge(const PersistentIndexSstableMetaPB& sstable_meta,
                                         std::vector<PersistentIndexSstablePB>* sstables, bool* merge_base_level);
 
+    // Assign the generation version (PersistentIndexSstablePB.generation_version) to each
+    // sstable in |new_meta|: a sstable that already carries a non-zero value is left
+    // untouched; a file present in |prev_meta| keeps its recorded value (a legacy 0 stays 0,
+    // so lake vacuum retention treats it conservatively); a file absent from |prev_meta| is
+    // new this publish and is stamped with |generation_version|.
+    static void assign_generation_versions(PersistentIndexSstableMetaPB* new_meta,
+                                           const PersistentIndexSstableMetaPB& prev_meta, int64_t generation_version);
+
     // Check if this rowset need to rebuild, return `True` means need to rebuild this rowset.
     static bool needs_rowset_rebuild(const RowsetMetadataPB& rowset, uint32_t rebuild_rss_id);
 

be/src/storage/lake/lake_primary_index.cpp

@@ -227,7 +227,8 @@ Status LakePrimaryIndex::ingest_sst(const FileMetaPB& sst_meta, const Persistent
     }
 }
 
-Status LakePrimaryIndex::commit(const TabletMetadataPtr& metadata, MetaFileBuilder* builder) {
+Status LakePrimaryIndex::commit(const TabletMetadataPtr& metadata, MetaFileBuilder* builder,
+                                int64_t generation_version) {
     TRACE_COUNTER_SCOPE_LATENCY_US("primary_index_commit_latency_us");
     if (!_enable_persistent_index) {
         return Status::OK();
@@ -250,7 +251,7 @@ Status LakePrimaryIndex::commit(const TabletMetadataPtr& metadata, MetaFileBuild
     case PersistentIndexTypePB::CLOUD_NATIVE: {
         auto* lake_persistent_index = dynamic_cast<LakePersistentIndex*>(_persistent_index.get());
         if (lake_persistent_index != nullptr) {
-            return lake_persistent_index->commit(builder);
+            return lake_persistent_index->commit(builder, generation_version);
         } else {
             return Status::InternalError("Persistent index is not a LakePersistentIndex.");
         }

be/src/storage/lake/lake_primary_index.h

@@ -71,7 +71,7 @@ class LakePrimaryIndex : public PrimaryIndex {
 
     Status apply_opcompaction(const TabletMetadataPtr& metadata, const TxnLogPB_OpCompaction& op_compaction);
 
-    Status commit(const TabletMetadataPtr& metadata, MetaFileBuilder* builder);
+    Status commit(const TabletMetadataPtr& metadata, MetaFileBuilder* builder, int64_t generation_version = 0);
 
     // Force any in-memory memtables of the cloud-native persistent index to
     // be flushed into sstables on shared storage. A no-op for LOCAL /

be/src/storage/lake/tablet_merger.cpp

@@ -2512,6 +2512,8 @@ Status emit_legacy_shared_sstable_into_dest(TabletManager* tablet_manager, const
     RETURN_IF_ERROR(rebuild_legacy_shared_sstable(tablet_manager, new_metadata.id(), src_pb, ctx.metadata(),
                                                   new_metadata, *family_maps_ptr, &projected_pb));
     if (!projected_pb.filename().empty()) {
+        // Rebuilt = a brand-new physical file, first visible at the merge version.
+        projected_pb.set_generation_version(new_metadata.version());
         dest->Add()->Swap(&projected_pb);
     }
     return Status::OK();
@@ -2537,6 +2539,8 @@ Status emit_non_shared_legacy_sstable_into_dest(TabletManager* tablet_manager, c
         RETURN_IF_ERROR(rebuild_non_shared_legacy_sstable(tablet_manager, new_metadata.id(), src_pb, ctx,
                                                           ctx.metadata(), &rebuilt_pb));
         if (!rebuilt_pb.filename().empty()) {
+            // Rebuilt = a brand-new physical file, first visible at the merge version.
+            rebuilt_pb.set_generation_version(new_metadata.version());
             dest->Add()->Swap(&rebuilt_pb);
         }
         return Status::OK();
@@ -2588,7 +2592,8 @@ Status merge_sstables(TabletManager* tablet_manager, std::vector<TabletMergeCont
         // the case where an old tablet accumulated post-split DML that never
         // reached shared storage before merge; see the symmetric call in
         // split_tablet for the pre-split side of the invariant.
-        ASSIGN_OR_RETURN(auto flushed_metadata, update_manager->flush_pk_memtable(ctx.metadata()));
+        ASSIGN_OR_RETURN(auto flushed_metadata,
+                         update_manager->flush_pk_memtable(ctx.metadata(), new_metadata->version()));
         ctx.set_metadata(std::move(flushed_metadata));
         if (!ctx.metadata()->has_sstable_meta()) continue;
 

be/src/storage/lake/tablet_splitter.cpp

@@ -1536,7 +1536,7 @@ StatusOr<std::unordered_map<int64_t, MutableTabletMetadataPtr>> split_tablet(
     // the "reshard inputs must have full sstable coverage" invariant; merge
     // does the post-split half in merge_sstables.
     ASSIGN_OR_RETURN(TabletMetadataPtr old_tablet_metadata,
-                     tablet_manager->update_mgr()->flush_pk_memtable(tablet_metadata));
+                     tablet_manager->update_mgr()->flush_pk_memtable(tablet_metadata, new_version));
 
     // Dispatch on FE-supplied new_tablet_ranges. When set, FE has computed the
     // K-1 boundaries externally (external boundaries / external boundaries); BE computes

be/src/storage/lake/update_manager.cpp

@@ -213,7 +213,8 @@ void UpdateManager::unload_and_remove_primary_index(int64_t tablet_id) {
     }
 }
 
-StatusOr<TabletMetadataPtr> UpdateManager::flush_pk_memtable(const TabletMetadataPtr& metadata) {
+StatusOr<TabletMetadataPtr> UpdateManager::flush_pk_memtable(const TabletMetadataPtr& metadata,
+                                                             int64_t generation_version) {
     if (!is_primary_key(*metadata) || !metadata->enable_persistent_index() ||
         metadata->persistent_index_type() != PersistentIndexTypePB::CLOUD_NATIVE) {
         return metadata;
@@ -257,7 +258,9 @@ StatusOr<TabletMetadataPtr> UpdateManager::flush_pk_memtable(const TabletMetadat
     // mutable_metadata->sstable_meta(). builder.finalize() is NOT called
     // here — this flush does not produce its own metadata version; the
     // returned metadata is consumed by the surrounding reshard publish.
-    RETURN_IF_ERROR(index_entry->value().commit(metadata, &builder));
+    // Pass the reshard publish version as generation_version so freshly-flushed sstables are
+    // stamped with it (this flush runs at base_version); the metadata version is left untouched.
+    RETURN_IF_ERROR(index_entry->value().commit(metadata, &builder, generation_version));
 
     // Success: dismiss the failure cleanup and release the cache entry.
     // write_guard is released when it goes out of scope at function return.

be/src/storage/lake/update_manager.h

@@ -225,7 +225,7 @@ class UpdateManager {
     // exclusion between publish_resharding_tablet and publish_version, so
     // the cached _data_version cannot advance past metadata->version()
     // during this call.
-    StatusOr<TabletMetadataPtr> flush_pk_memtable(const TabletMetadataPtr& metadata);
+    StatusOr<TabletMetadataPtr> flush_pk_memtable(const TabletMetadataPtr& metadata, int64_t generation_version);
 
     StatusOr<IndexEntry*> rebuild_primary_index(const TabletMetadataPtr& metadata, MetaFileBuilder* builder,
                                                 int64_t base_version, int64_t new_version,

be/test/storage/lake/lake_persistent_index_test.cpp

@@ -1330,6 +1330,7 @@ TEST_F(LakePersistentIndexTest, test_ingest_sst_preserves_shared_flag_for_new_ss
         if (sst.filename() == sst_filename) {
             match_count++;
             EXPECT_TRUE(sst.shared()) << "sst_meta.shared=true must be preserved in committed sstable_meta";
+            EXPECT_EQ(1, sst.generation_version()) << "ingest_sst stamps the generation version (the ingest version)";
         }
     }
     EXPECT_EQ(match_count, 1);
@@ -1573,4 +1574,101 @@ TEST_F(LakePersistentIndexTest, test_load_dels_parallel_matches_single_pass) {
     }
 }
 
+TEST(AssignGenerationVersionsTest, stamps_new_preserves_existing_and_legacy) {
+    PersistentIndexSstableMetaPB prev;
+    {
+        auto* s = prev.add_sstables();
+        s->set_filename("existing_v3.sst");
+        s->set_generation_version(3);
+    }
+    {
+        auto* s = prev.add_sstables();
+        s->set_filename("legacy_zero.sst"); // version unset == 0
+    }
+
+    PersistentIndexSstableMetaPB cur;
+    {
+        auto* s = cur.add_sstables();
+        s->set_filename("existing_v3.sst"); // reappears, version dropped by copy
+    }
+    {
+        auto* s = cur.add_sstables();
+        s->set_filename("legacy_zero.sst"); // legacy, still unset
+    }
+    {
+        auto* s = cur.add_sstables();
+        s->set_filename("fresh_new.sst"); // brand new this publish
+    }
+    {
+        auto* s = cur.add_sstables();
+        s->set_filename("ingested.sst");
+        s->set_generation_version(9); // producer already stamped, absent from prev
+    }
+    {
+        auto* s = cur.add_sstables();
+        s->set_filename("existing_v3.sst");
+        s->set_generation_version(7); // same filename as prev@3 but already non-zero
+    }
+
+    LakePersistentIndex::assign_generation_versions(&cur, prev, /*generation_version=*/100);
+
+    ASSERT_EQ(5, cur.sstables_size());
+    EXPECT_EQ(3, cur.sstables(0).generation_version());   // existing, local 0 -> carried from prev (3)
+    EXPECT_EQ(0, cur.sstables(1).generation_version());   // legacy 0 -> STAYS 0 (present in prev); vacuum retains
+    EXPECT_EQ(100, cur.sstables(2).generation_version()); // new (absent from prev) -> new version
+    EXPECT_EQ(9, cur.sstables(3).generation_version());   // producer-stamped non-zero -> preserved (absent from prev)
+    EXPECT_EQ(7, cur.sstables(4).generation_version());   // non-zero -> NEVER overwritten, even though prev has @3
+}
+
+TEST_F(LakePersistentIndexTest, test_commit_stamps_version) {
+    const int N = 100;
+    const int64_t publish_version = 7;
+    auto tablet_id = _tablet_metadata->id();
+    auto index = std::make_unique<LakePersistentIndex>(_tablet_mgr.get(), tablet_id);
+    ASSERT_OK(index->init(_tablet_metadata));
+
+    using Key = uint64_t;
+    std::vector<Key> keys(N);
+    std::vector<Slice> key_slices;
+    std::vector<IndexValue> values;
+    key_slices.reserve(N);
+    values.reserve(N);
+    for (int j = 0; j < N; j++) {
+        keys[j] = j;
+        key_slices.emplace_back((uint8_t*)(&keys[j]), sizeof(Key));
+        values.emplace_back(j * 2);
+    }
+    index->prepare(EditVersion(publish_version, 0), 0);
+    std::vector<IndexValue> old_values(N);
+    ASSERT_OK(index->upsert(N, key_slices.data(), values.data(), old_values.data()));
+    ASSERT_OK(index->flush_memtable(true));
+    ASSERT_OK(index->sync_flush_all_memtables(10000000));
+
+    Tablet tablet(_tablet_mgr.get(), tablet_id);
+    auto meta = std::make_shared<TabletMetadata>();
+    meta->CopyFrom(*_tablet_metadata);
+    meta->set_version(publish_version);
+    MetaFileBuilder builder(tablet, meta);
+    ASSERT_OK(index->commit(&builder));
+
+    // A freshly-flushed sstable is stamped with the publish version.
+    ASSERT_GE(meta->sstable_meta().sstables_size(), 1);
+    for (const auto& s : meta->sstable_meta().sstables()) {
+        EXPECT_EQ(publish_version, s.generation_version()) << "freshly-flushed sstable must carry its publish version";
+    }
+
+    // Re-commit at a later version WITHOUT new data: existing sstables must keep
+    // their original version (no re-stamp), because they are present in
+    // the previous sstable_meta the builder carries forward.
+    auto meta2 = std::make_shared<TabletMetadata>();
+    meta2->CopyFrom(*meta); // carries the publish_version sstable_meta forward
+    meta2->set_version(publish_version + 10);
+    MetaFileBuilder builder2(tablet, meta2);
+    ASSERT_OK(index->commit(&builder2));
+    ASSERT_GE(meta2->sstable_meta().sstables_size(), 1);
+    for (const auto& s : meta2->sstable_meta().sstables()) {
+        EXPECT_EQ(publish_version, s.generation_version()) << "re-commit must not re-stamp an existing sstable";
+    }
+}
+
 } // namespace starrocks::lake