Add sample k collect for first level of vector index table (#7412)

Author: MBkkt

ydb/core/protos/index_builder.proto

@@ -117,3 +117,14 @@ enum EBuildStatus {
     BUILD_ERROR = 5;
     BAD_REQUEST = 6;
 }
+
+message TEvUploadSampleKResponse {
+    optional uint64 Id = 1;
+
+    optional Ydb.StatusIds.StatusCode UploadStatus = 2;
+    repeated Ydb.Issue.IssueMessage Issues = 3;
+
+    optional uint64 RowsDelta = 4;
+    optional uint64 BytesDelta = 5;
+}
+

ydb/core/protos/out/out.cpp

@@ -1,3 +1,5 @@
+#include <ydb/public/api/protos/ydb_table.pb.h>
+
 #include <ydb/core/protos/blobstorage.pb.h>
 #include <ydb/core/protos/blobstorage_vdisk_internal.pb.h>
 #include <ydb/core/protos/blobstorage_vdisk_config.pb.h>
@@ -251,3 +253,7 @@ Y_DECLARE_OUT_SPEC(, NKikimrStat::TEvAnalyzeStatusResponse_EStatus, stream, valu
 Y_DECLARE_OUT_SPEC(, NKikimrStat::TEvStatisticsResponse::EStatus, stream, value) {
     stream << NKikimrStat::TEvStatisticsResponse::EStatus_Name(value);
 }
+
+Y_DECLARE_OUT_SPEC(, Ydb::Table::IndexBuildState_State, stream, value) {
+    stream << IndexBuildState_State_Name(value);
+}

ydb/core/protos/tx_datashard.proto

@@ -1451,16 +1451,20 @@ message TEvSampleKRequest {
     optional NKikimrProto.TPathID PathId = 3;
 
     optional uint64 K = 4;
+    // For now we use tablet id as seed value, but it's not convinient for tests
     optional uint64 Seed = 5;
+    // We want to sample small instead of large probabilites
+    // because size of message will be smaller (varint)
+    optional uint64 MaxProbability = 6;
 
-    repeated string Columns = 6;
-    optional NKikimrTx.TKeyRange KeyRange = 7;
+    repeated string Columns = 7;
+    optional NKikimrTx.TKeyRange KeyRange = 8;
 
-    optional uint64 SnapshotTxId = 8;
-    optional uint64 SnapshotStep = 9;
+    optional uint64 SnapshotTxId = 9;
+    optional uint64 SnapshotStep = 10;
 
-    optional uint64 SeqNoGeneration = 10;
-    optional uint64 SeqNoRound = 11;
+    optional uint64 SeqNoGeneration = 11;
+    optional uint64 SeqNoRound = 12;
 }
 
 message TEvSampleKResponse {
@@ -1472,13 +1476,14 @@ message TEvSampleKResponse {
     optional NKikimrIndexBuilder.EBuildStatus Status = 4;
     repeated Ydb.Issue.IssueMessage Issues = 5;
 
-    optional NKikimrTx.TKeyRange RequestedKeyRange = 6;
+    optional uint64 RowsDelta = 6;
+    optional uint64 BytesDelta = 7;
 
-    optional uint64 RequestSeqNoGeneration = 7;
-    optional uint64 RequestSeqNoRound = 8;
+    optional uint64 RequestSeqNoGeneration = 8;
+    optional uint64 RequestSeqNoRound = 9;
 
-    repeated uint64 Probabilities = 9;
-    repeated bytes Rows = 10;
+    repeated uint64 Probabilities = 10;
+    repeated bytes Rows = 11;
 }
 
 message TEvCdcStreamScanRequest {

ydb/core/scheme/scheme_tablecell.cpp

@@ -122,9 +122,7 @@ namespace {
         if (!SerializeCellVecInit(cells, resultBuffer, resultCells))
             return;
 
-        size_t size = sizeof(ui16);
-        for (auto& cell : cells)
-            size += sizeof(TCellHeader) + cell.Size();
+        auto size = TSerializedCellVec::SerializedSize(cells);
 
         resultBuffer.ReserveAndResize(size);
         char* resultBufferData = resultBuffer.Detach();
@@ -254,6 +252,14 @@ TString TSerializedCellVec::Serialize(TConstArrayRef<TCell> cells) {
     return result;
 }
 
+size_t TSerializedCellVec::SerializedSize(TConstArrayRef<TCell> cells) {
+    size_t size = sizeof(ui16) + sizeof(TCellHeader) * cells.size();
+    for (auto& cell : cells) {
+        size += cell.Size();
+    }
+    return size;
+}
+
 bool TSerializedCellVec::DoTryParse(const TString& data) {
     return TryDeserializeCellVec(data, Buf, Cells);
 }

ydb/core/scheme/scheme_tablecell.h

@@ -567,6 +567,8 @@ class TSerializedCellVec {
 
     static TString Serialize(TConstArrayRef<TCell> cells);
 
+    static size_t SerializedSize(TConstArrayRef<TCell> cells);
+
     const TString &GetBuffer() const { return Buf; }
 
     TString ReleaseBuffer() {

ydb/core/tx/columnshard/defs.cpp

@@ -2,18 +2,6 @@
 
 namespace NKikimr::NColumnShard {
 
-namespace {
-    const constexpr ui64 MAX_BLOB_SIZE_LIMIT = 8 * 1024 * 1024;
-    ui64 MaxBlobSize = MAX_BLOB_SIZE_LIMIT;
-}
-
-ui64 TLimits::GetMaxBlobSize() {
-    return MaxBlobSize;
-}
-
-ui64 TLimits::GetBlobSizeLimit() {
-    return MAX_BLOB_SIZE_LIMIT;
-}
 
 void TLimits::SetMaxBlobSize(const ui64 value) {
     Y_ABORT_UNLESS(value <= MAX_BLOB_SIZE_LIMIT);

ydb/core/tx/columnshard/defs.h

@@ -15,14 +15,23 @@ namespace NKikimr::NColumnShard {
 using TLogThis = TCtorLogger<NKikimrServices::TX_COLUMNSHARD>;
 
 struct TLimits {
+private:
+    static constexpr ui64 MAX_BLOB_SIZE_LIMIT = 8 * 1024 * 1024;
+    static inline ui64 MaxBlobSize = MAX_BLOB_SIZE_LIMIT;
+
+public:
     static constexpr const ui32 MIN_SMALL_BLOBS_TO_INSERT = 200;
     static constexpr const ui32 MIN_BYTES_TO_INSERT = 4 * 1024 * 1024;
     static constexpr const ui64 MAX_BYTES_TO_INSERT = 16 * 1024 * 1024;
     static constexpr const ui32 MAX_TX_RECORDS = 100000;
     static constexpr const ui64 MAX_BLOBS_TO_DELETE = NOlap::TCompactionLimits::MAX_BLOBS_TO_DELETE;
 
-    static ui64 GetBlobSizeLimit();
-    static ui64 GetMaxBlobSize();
+    static constexpr ui64 GetBlobSizeLimit() {
+        return MAX_BLOB_SIZE_LIMIT;
+    }
+    static ui64 GetMaxBlobSize() {
+        return MaxBlobSize;
+    }
     static void SetMaxBlobSize(const ui64 value);
 
     class TMaxBlobSizeGuard: TNonCopyable {

ydb/core/tx/datashard/build_index.cpp

@@ -28,8 +28,8 @@ namespace NKikimr::NDataShard {
 #define LOG_E(stream) LOG_ERROR_S(*TlsActivationContext, NKikimrServices::TX_DATASHARD, stream)
 
 using TColumnsTypes = THashMap<TString, NScheme::TTypeInfo>;
-using TTypes = TVector<std::pair<TString, Ydb::Type>>;
-using TRows = TVector<std::pair<TSerializedCellVec, TString>>;
+using TTypes = NTxProxy::TUploadTypes;
+using TRows = NTxProxy::TUploadRows;
 
 static TColumnsTypes GetAllTypes(const TUserTable& tableInfo) {
     TColumnsTypes result;
@@ -119,42 +119,6 @@ bool BuildExtraColumns(TVector<TCell>& cells, const NKikimrIndexBuilder::TColumn
     return true;
 }
 
-struct TStatus {
-    Ydb::StatusIds::StatusCode StatusCode = Ydb::StatusIds::STATUS_CODE_UNSPECIFIED;
-    NYql::TIssues Issues;
-
-    bool IsNone() const {
-        return StatusCode == Ydb::StatusIds::STATUS_CODE_UNSPECIFIED;
-    }
-
-    bool IsSuccess() const {
-        return StatusCode == Ydb::StatusIds::SUCCESS;
-    }
-
-    bool IsRetriable() const {
-        return StatusCode == Ydb::StatusIds::UNAVAILABLE || StatusCode == Ydb::StatusIds::OVERLOADED;
-    }
-
-    TString ToString() const {
-        return TStringBuilder()
-               << "Status {"
-               << " Code: " << Ydb::StatusIds_StatusCode_Name(StatusCode)
-               << " Issues: " << Issues.ToString()
-               << " }";
-    }
-};
-
-struct TUploadLimits {
-    ui64 BatchRowsLimit = 500;
-    ui64 BatchBytesLimit = 1u << 23; // 8MB
-    ui32 MaxUploadRowsRetryCount = 50;
-    ui32 BackoffCeiling = 3;
-
-    TDuration GetTimeoutBackouff(ui32 retryNo) const {
-        return TDuration::Seconds(1u << Max(retryNo, BackoffCeiling));
-    }
-};
-
 class TBufferData: public IStatHolder, public TNonCopyable {
 public:
     TBufferData()
@@ -258,7 +222,7 @@ class TBuildScanUpload: public TActor<TBuildScanUpload<Activity>>, public NTable
     ui64 RetryCount = 0;
 
     TUploadMonStats Stats = TUploadMonStats("tablets", "build_index_upload");
-    TStatus UploadStatus;
+    TUploadStatus UploadStatus;
 
     TBuildScanUpload(ui64 buildIndexId,
                      const TString& target,

ydb/core/tx/datashard/datashard_ut_sample_k.cpp

@@ -62,6 +62,7 @@ Y_UNIT_TEST_SUITE (TTxDataShardSampleKScan) {
             rec.SetSnapshotTxId(snapshot.TxId);
             rec.SetSnapshotStep(snapshot.Step);
 
+            rec.SetMaxProbability(std::numeric_limits<uint64_t>::max());
             rec.SetSeed(1337);
             if (!rec.HasK()) {
                 rec.SetK(1);
@@ -104,6 +105,7 @@ Y_UNIT_TEST_SUITE (TTxDataShardSampleKScan) {
                 rec.SetSnapshotTxId(snapshot.TxId);
                 rec.SetSnapshotStep(snapshot.Step);
 
+                rec.SetMaxProbability(std::numeric_limits<uint64_t>::max());
                 rec.SetSeed(seed);
                 rec.SetK(k);
             };

ydb/core/tx/datashard/sample_k.cpp

@@ -36,8 +36,6 @@ class TSampleKScan final: public TActor<TSampleKScan>, public NTable::IScan {
     const TSerializedTableRange RequestedRange;
     const ui64 K;
 
-    IDriver* Driver = nullptr;
-
     struct TProbability {
         ui64 P = 0;
         ui64 I = 0;
@@ -46,11 +44,18 @@ class TSampleKScan final: public TActor<TSampleKScan>, public NTable::IScan {
         auto operator<=>(const TProbability&) const noexcept = default;
     };
 
-    TReallyFastRng32 Rng;
+    ui64 RowsCount = 0;
+    ui64 RowsBytes = 0;
+
+    // We are using binary heap, because we don't want to do batch processing here,
+    // serialization is more expensive than compare
     ui64 MaxProbability = 0;
+    TReallyFastRng32 Rng;
     std::vector<TProbability> MaxRows;
     std::vector<TString> DataRows;
 
+    IDriver* Driver = nullptr;
+
 public:
     static constexpr NKikimrServices::TActivity::EType ActorActivityType() {
         return NKikimrServices::TActivity::SAMPLE_K_SCAN_ACTOR;
@@ -61,6 +66,7 @@ class TSampleKScan final: public TActor<TSampleKScan>, public NTable::IScan {
                  const TSerializedTableRange& range,
                  ui64 k,
                  ui64 seed,
+                 ui64 maxProbability,
                  TProtoColumnsCRef columns,
                  const TUserTable& tableInfo)
         : TActor(&TThis::StateWork)
@@ -71,8 +77,10 @@ class TSampleKScan final: public TActor<TSampleKScan>, public NTable::IScan {
         , TableRange(tableInfo.Range)
         , RequestedRange(range)
         , K(k)
+        , MaxProbability(maxProbability)
         , Rng(seed)
     {
+        Y_ASSERT(MaxProbability != 0);
     }
 
     ~TSampleKScan() final = default;
@@ -109,19 +117,32 @@ class TSampleKScan final: public TActor<TSampleKScan>, public NTable::IScan {
 
     EScan Feed(TArrayRef<const TCell> key, const TRow& row) noexcept final {
         LOG_T("Feed key " << DebugPrintPoint(KeyTypes, key, *AppData()->TypeRegistry) << " " << Debug());
+        ++RowsCount;
+
+        const auto probability = GetProbability();
+        if (probability > MaxProbability) {
+            // TODO(mbkkt) it's not nice that we need to compute this, probably can be precomputed in TRow
+            RowsBytes += TSerializedCellVec::SerializedSize(*row);
+            return EScan::Feed;
+        }
+
+        auto serialized = TSerializedCellVec::Serialize(*row);
+        RowsBytes += serialized.size();
 
-        const auto probability = Rng.GenRand64();
         if (DataRows.size() < K) {
             MaxRows.push_back({probability, DataRows.size()});
-            DataRows.emplace_back(TSerializedCellVec::Serialize(*row));
+            DataRows.emplace_back(std::move(serialized));
             if (DataRows.size() == K) {
                 std::make_heap(MaxRows.begin(), MaxRows.end());
                 MaxProbability = MaxRows.front().P;
             }
-        } else if (probability < MaxProbability) {
-            ReplaceRow(row, probability);
+        } else {
+            ReplaceRow(std::move(serialized), probability);
         }
 
+        if (MaxProbability == 0) {
+            return EScan::Final;
+        }
         return EScan::Feed;
     }
 
@@ -167,21 +188,35 @@ class TSampleKScan final: public TActor<TSampleKScan>, public NTable::IScan {
         }
     }
 
-    void ReplaceRow(const TRow& row, ui64 p) {
+    void ReplaceRow(TString&& row, ui64 p) {
+        // TODO(mbkkt) use tournament tree to make less compare and swaps
         std::pop_heap(MaxRows.begin(), MaxRows.end());
-        DataRows[MaxRows.back().I] = TSerializedCellVec::Serialize(*row);
+        DataRows[MaxRows.back().I] = std::move(row);
         MaxRows.back().P = p;
         std::push_heap(MaxRows.begin(), MaxRows.end());
         MaxProbability = MaxRows.front().P;
     }
 
     void FillResponse() {
         std::sort(MaxRows.begin(), MaxRows.end());
+        auto& record = Response->Record;
         for (auto& [p, i] : MaxRows) {
-            Response->Record.AddProbabilities(p);
-            Response->Record.AddRows(std::move(DataRows[i]));
+            record.AddProbabilities(p);
+            record.AddRows(std::move(DataRows[i]));
+        }
+        record.SetRowsDelta(RowsCount);
+        record.SetBytesDelta(RowsBytes);
+        record.SetStatus(NKikimrIndexBuilder::EBuildStatus::DONE);
+    }
+
+    ui64 GetProbability() {
+        while (true) {
+            auto p = Rng.GenRand64();
+            // We exclude max ui64 from generated probabilities, so we can use this value as initial max
+            if (Y_LIKELY(p != std::numeric_limits<ui64>::max())) {
+                return p;
+            }
         }
-        Response->Record.SetStatus(NKikimrIndexBuilder::EBuildStatus::DONE);
     }
 };
 
@@ -317,6 +352,7 @@ void TDataShard::HandleSafe(TEvDataShard::TEvSampleKRequest::TPtr& ev, const TAc
                                       requestedRange,
                                       record.GetK(),
                                       record.GetSeed(),
+                                      record.GetMaxProbability(),
                                       record.GetColumns(),
                                       userTable),
                                   ev->Cookie,