Return virtual timestamp for scan query (#14001) (#18152)

dcherednik · web-flow · commit 92895832bc39 · 2025-05-09T17:24:39.000+02:00
Conflicts:
ydb/core/kqp/executer_actor/kqp_executer_impl.h
ydb/public/api/protos/ydb_table.proto
ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/table/table.h
ydb/public/sdk/cpp/src/client/table/impl/readers.cpp
diff --git a/ydb/core/grpc_services/query/rpc_execute_query.cpp b/ydb/core/grpc_services/query/rpc_execute_query.cpp
@@ -327,6 +327,13 @@ class TExecuteQueryRPC : public TActorBootstrapped<TExecuteQueryRPC> {
         response.set_result_set_index(ev->Get()->Record.GetQueryResultIndex());
         response.mutable_result_set()->Swap(ev->Get()->Record.MutableResultSet());
 
+        if (ev->Get()->Record.HasVirtualTimestamp()) {
+            auto snap = response.mutable_snapshot_timestamp();
+            auto& ts = ev->Get()->Record.GetVirtualTimestamp();
+            snap->set_plan_step(ts.GetStep());
+            snap->set_tx_id(ts.GetTxId());
+        }
+
         TString out;
         Y_PROTOBUF_SUPPRESS_NODISCARD response.SerializeToString(&out);
 
diff --git a/ydb/core/grpc_services/rpc_stream_execute_scan_query.cpp b/ydb/core/grpc_services/rpc_stream_execute_scan_query.cpp
@@ -340,27 +340,40 @@ class TStreamExecuteScanQueryRPC : public TActorBootstrapped<TStreamExecuteScanQ
             ExecuterActorId_ = ev->Sender;
         }
 
+        auto& evRecord = ev->Get()->Record;
+
         Ydb::Table::ExecuteScanQueryPartialResponse response;
-        response.set_status(StatusIds::SUCCESS);
-        response.mutable_result()->mutable_result_set()->Swap(ev->Get()->Record.MutableResultSet());
+
+        {
+            response.set_status(StatusIds::SUCCESS);
+            auto result = response.mutable_result();
+            result->mutable_result_set()->Swap(evRecord.MutableResultSet());
+
+            if (evRecord.HasVirtualTimestamp()) {
+                auto snap = result->mutable_snapshot();
+                auto ts = evRecord.GetVirtualTimestamp();
+                snap->set_plan_step(ts.GetStep());
+                snap->set_tx_id(ts.GetTxId());
+            }
+        }
 
         TString out;
         Y_PROTOBUF_SUPPRESS_NODISCARD response.SerializeToString(&out);
 
         FlowControl_.PushResponse(out.size());
         const i64 freeSpaceBytes = FlowControl_.FreeSpaceBytes();
-        LastSeqNo_ = ev->Get()->Record.GetSeqNo();
+        LastSeqNo_ = evRecord.GetSeqNo();
         AckedFreeSpaceBytes_ = freeSpaceBytes;
 
         Request_->SendSerializedResult(std::move(out), StatusIds::SUCCESS);
 
         LOG_DEBUG_S(ctx, NKikimrServices::RPC_REQUEST, this->SelfId() << " Send stream data ack"
-            << ", seqNo: " << ev->Get()->Record.GetSeqNo()
+            << ", seqNo: " << evRecord.GetSeqNo()
             << ", freeSpace: " << freeSpaceBytes
             << ", to: " << ev->Sender
             << ", queue: " << FlowControl_.QueueSize());
 
-        auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>(ev->Get()->Record.GetSeqNo(), ev->Get()->Record.GetChannelId());
+        auto resp = MakeHolder<NKqp::TEvKqpExecuter::TEvStreamDataAck>(evRecord.GetSeqNo(), evRecord.GetChannelId());
         resp->Record.SetFreeSpace(freeSpaceBytes);
 
         ctx.Send(ev->Sender, resp.Release());
diff --git a/ydb/core/kqp/executer_actor/kqp_executer_impl.h b/ydb/core/kqp/executer_actor/kqp_executer_impl.h
@@ -304,6 +304,13 @@ class TKqpExecuterBase : public TActorBootstrapped<TDerived> {
                 streamEv->Record.SetChannelId(channel.Id);
                 streamEv->Record.MutableResultSet()->Swap(&resultSet);
 
+                const auto& snap = GetSnapshot();
+                if (snap.IsValid()) {
+                    auto vt = streamEv->Record.MutableVirtualTimestamp();
+                    vt->SetStep(snap.Step);
+                    vt->SetTxId(snap.TxId);
+                }
+
                 LOG_D("Send TEvStreamData to " << Target << ", seqNo: " << streamEv->Record.GetSeqNo()
                     << ", nRows: " << streamEv->Record.GetResultSet().rows().size());
 
diff --git a/ydb/core/kqp/ut/scan/kqp_scan_ut.cpp b/ydb/core/kqp/ut/scan/kqp_scan_ut.cpp
@@ -530,17 +530,52 @@ Y_UNIT_TEST_SUITE(KqpScan) {
                 .Build()
             .Build();
 
-        auto it = db.StreamExecuteScanQuery(R"(
-            DECLARE $key AS Uint64;
 
-            SELECT * FROM `/Root/EightShard` WHERE Key = $key;
-        )", params).GetValueSync();
+        {
+            auto it = db.StreamExecuteScanQuery(R"(
+                DECLARE $key AS Uint64;
 
-        UNIT_ASSERT(it.IsSuccess());
+                SELECT * FROM `/Root/EightShard` WHERE Key = $key;
+            )", params).GetValueSync();
 
-        CompareYson(R"([
-            [[1];[202u];["Value2"]]
-        ])", StreamResultToYson(it));
+            UNIT_ASSERT(it.IsSuccess());
+
+                   CompareYson(R"([
+                [[1];[202u];["Value2"]]
+            ])", StreamResultToYson(it));
+        }
+
+        {
+            auto it = db.StreamExecuteScanQuery(R"(
+                DECLARE $key AS Uint64;
+
+                SELECT * FROM `/Root/EightShard` WHERE Key = $key;
+            )", params).GetValueSync();
+
+            UNIT_ASSERT(it.IsSuccess());
+            auto part = it.ReadNext().GetValueSync();
+            UNIT_ASSERT(part.IsSuccess());
+
+
+            UNIT_ASSERT(part.HasVirtualTimestamp());
+            UNIT_ASSERT(part.GetVirtualTimestamp().GetStep() != 0);
+            UNIT_ASSERT(part.GetVirtualTimestamp().GetTxId() != 0);
+        }
+
+        {
+            auto it = db.StreamExecuteScanQuery(R"(
+                SELECT * FROM `/Root/EightShard` WHERE Key = 9876554123;
+            )", params).GetValueSync();
+
+            UNIT_ASSERT(it.IsSuccess());
+            auto part = it.ReadNext().GetValueSync();
+            UNIT_ASSERT(part.IsSuccess());
+
+
+            UNIT_ASSERT(part.HasVirtualTimestamp());
+            UNIT_ASSERT(part.GetVirtualTimestamp().GetStep() != 0);
+            UNIT_ASSERT(part.GetVirtualTimestamp().GetTxId() != 0);
+        }
     }
 
     Y_UNIT_TEST(AggregateByColumn) {
diff --git a/ydb/core/protos/kqp.proto b/ydb/core/protos/kqp.proto
@@ -498,6 +498,7 @@ message TEvExecuterStreamData {
     optional uint32 QueryResultIndex = 3;
     optional uint32 ChannelId = 4;
     optional NActorsProto.TActorId ChannelActorId = 5;
+    optional TKqpSnapshot VirtualTimestamp = 6;
 };
 
 message TEvExecuterStreamDataAck {
diff --git a/ydb/public/api/protos/ydb_query.proto b/ydb/public/api/protos/ydb_query.proto
@@ -8,6 +8,7 @@ option java_outer_classname = "QueryProtos";
 import "google/protobuf/duration.proto";
 
 import "ydb/public/api/protos/annotations/validation.proto";
+import "ydb/public/api/protos/ydb_common.proto";
 import "ydb/public/api/protos/ydb_issue_message.proto";
 import "ydb/public/api/protos/ydb_operation.proto";
 import "ydb/public/api/protos/ydb_query_stats.proto";
@@ -191,6 +192,7 @@ message ExecuteQueryResponsePart {
     Ydb.TableStats.QueryStats exec_stats = 5;
 
     TransactionMeta tx_meta = 6;
+    VirtualTimestamp snapshot_timestamp = 7;
 }
 
 message ExecuteScriptRequest {
diff --git a/ydb/public/api/protos/ydb_table.proto b/ydb/public/api/protos/ydb_table.proto
@@ -1223,4 +1223,6 @@ message ExecuteScanQueryPartialResult {
     // works only in mode: MODE_EXPLAIN,
     // collects additional diagnostics about query compilation, including query plan and scheme
     string query_full_diagnostics = 7;
+    // Optional snapshot that corresponds to the returned data
+    VirtualTimestamp snapshot = 8;
 }
diff --git a/ydb/public/sdk/cpp/client/ydb_table/impl/readers.cpp b/ydb/public/sdk/cpp/client/ydb_table/impl/readers.cpp
@@ -90,9 +90,16 @@ TAsyncScanQueryPart TScanQueryPartIterator::TReaderImpl::ReadNext(std::shared_pt
 
             diagnostics = self->Response_.result().query_full_diagnostics();
 
+            std::optional<TVirtualTimestamp> vt;
+
+            if (self->Response_.result().has_snapshot()) {
+                const auto& snap = self->Response_.result().snapshot();
+                vt = TVirtualTimestamp(snap.plan_step(), snap.tx_id());
+            }
+
             if (self->Response_.result().has_result_set()) {
                 promise.SetValue({std::move(status),
-                    TResultSet(std::move(*self->Response_.mutable_result()->mutable_result_set())), queryStats, diagnostics});
+                    TResultSet(std::move(*self->Response_.mutable_result()->mutable_result_set())), queryStats, diagnostics, std::move(vt)});
             } else {
                 promise.SetValue({std::move(status), queryStats, diagnostics});
             }
diff --git a/ydb/public/sdk/cpp/client/ydb_table/table.h b/ydb/public/sdk/cpp/client/ydb_table/table.h
@@ -1927,6 +1927,8 @@ class TReadTableSnapshot {
     ui64 TxId_;
 };
 
+using TVirtualTimestamp = TReadTableSnapshot;
+
 template<typename TPart>
 class TSimpleStreamPart : public TStreamPartStatus {
 public:
@@ -1980,6 +1982,10 @@ class TScanQueryPart : public TStreamPartStatus {
     const TString& GetDiagnostics() const { return *Diagnostics_; }
     TString&& ExtractDiagnostics() { return std::move(*Diagnostics_); }
 
+    bool HasVirtualTimestamp() const { return Vt_.has_value(); }
+    const TVirtualTimestamp& GetVirtualTimestamp() const { return *Vt_; }
+    TVirtualTimestamp&& ExtractVirtualTimestamp() { return std::move(*Vt_); }
+
     TScanQueryPart(TStatus&& status)
         : TStreamPartStatus(std::move(status))
     {}
@@ -1990,17 +1996,20 @@ class TScanQueryPart : public TStreamPartStatus {
         , Diagnostics_(diagnostics)
     {}
 
-    TScanQueryPart(TStatus&& status, TResultSet&& resultSet, const TMaybe<TQueryStats>& queryStats, const TMaybe<TString>& diagnostics)
+    TScanQueryPart(TStatus&& status, TResultSet&& resultSet, const TMaybe<TQueryStats>& queryStats,
+        const TMaybe<TString>& diagnostics, std::optional<TVirtualTimestamp>&& vt)
         : TStreamPartStatus(std::move(status))
         , ResultSet_(std::move(resultSet))
         , QueryStats_(queryStats)
         , Diagnostics_(diagnostics)
+        , Vt_(std::move(vt))
     {}
 
 private:
     TMaybe<TResultSet> ResultSet_;
     TMaybe<TQueryStats> QueryStats_;
     TMaybe<TString> Diagnostics_;
+    std::optional<TVirtualTimestamp> Vt_;
 };
 
 using TAsyncScanQueryPart = NThreading::TFuture<TScanQueryPart>;

Original file line number	Diff line number	Diff line change
`@@ -1223,4 +1223,6 @@ message ExecuteScanQueryPartialResult {`
`1223`	`1223`	`// works only in mode: MODE_EXPLAIN,`
`1224`	`1224`	`// collects additional diagnostics about query compilation, including query plan and scheme`
`1225`	`1225`	`string query_full_diagnostics = 7;`
	`1226`	`+ // Optional snapshot that corresponds to the returned data`
	`1227`	`+ VirtualTimestamp snapshot = 8;`
`1226`	`1228`	`}`