YQ-4785 fixed checkpoints for map join (#26591)

Author: GrigoriyPA

ydb/core/kqp/executer_actor/kqp_tasks_graph.cpp

@@ -1015,6 +1015,8 @@ void TKqpTasksGraph::FillChannelDesc(NDqProto::TChannel& channelDesc, const TCha
     channelDesc.SetSrcTaskId(channel.SrcTask);
     channelDesc.SetDstTaskId(channel.DstTask);
     channelDesc.SetEnableSpilling(enableSpilling);
+    channelDesc.SetCheckpointingMode(channel.CheckpointingMode);
+    channelDesc.SetWatermarksMode(channel.WatermarksMode);
 
     const auto& resultChannelProxies = GetMeta().ResultChannelProxies;
 

ydb/core/kqp/ut/federated_query/datastreams/datastreams_ut.cpp

@@ -249,16 +249,16 @@ class TStreamingTestFixture : public NUnitTest::TBaseFixture {
         }
     }
 
-    void ReadTopicMessage(const TString& topicName, const TString& expectedMessage, TDuration disposition = TDuration::Seconds(100)) {
-        ReadTopicMessages(topicName, {expectedMessage}, disposition);
+    void ReadTopicMessage(const TString& topicName, const TString& expectedMessage, TInstant disposition = TInstant::Now() - TDuration::Seconds(100), bool sort = false) {
+        ReadTopicMessages(topicName, {expectedMessage}, disposition, sort);
     }
 
-    void ReadTopicMessages(const TString& topicName, const TVector<TString>& expectedMessages, TDuration disposition = TDuration::Seconds(100)) {
+    void ReadTopicMessages(const TString& topicName, const TVector<TString>& expectedMessages, TInstant disposition = TInstant::Now() - TDuration::Seconds(100), bool sort = false) {
         NTopic::TReadSessionSettings readSettings;
         readSettings
             .WithoutConsumer()
             .AppendTopics(
-                NTopic::TTopicReadSettings(topicName).ReadFromTimestamp(TInstant::Now() - disposition)
+                NTopic::TTopicReadSettings(topicName).ReadFromTimestamp(disposition)
                     .AppendPartitionIds(0)
             );
 
@@ -287,12 +287,19 @@ class TStreamingTestFixture : public NUnitTest::TBaseFixture {
                 }
             }
 
-            UNIT_ASSERT_GE(expectedMessages.size(), received.size());
+            UNIT_ASSERT_C(expectedMessages.size() >= received.size(), TStringBuilder()
+                << "expected #" << expectedMessages.size() << " messages ("
+                << JoinSeq(", ", expectedMessages) << "), got #" << received.size() << " messages ("
+                << JoinSeq(", ", received) << ")");
 
             error = TStringBuilder() << "got new event, received #" << received.size() << " / " << expectedMessages.size() << " messages";
             return false;
         });
 
+        if (sort) {
+            std::sort(received.begin(), received.end());
+        }
+
         UNIT_ASSERT_VALUES_EQUAL(received.size(), expectedMessages.size());
         for (size_t i = 0; i < received.size(); ++i) {
             UNIT_ASSERT_VALUES_EQUAL(received[i], expectedMessages[i]);
@@ -2142,6 +2149,134 @@ Y_UNIT_TEST_SUITE(KqpStreamingQueriesDdl) {
             UNIT_ASSERT_STRING_CONTAINS(*ast, "DqCnStreamLookup");
         });
     }
+
+    Y_UNIT_TEST_F(OffsetsAndStateRecoveryOnInternalRetry, TStreamingTestFixture) {
+        // Join with S3 used for introducing temporary failure and force retry on specific key
+
+        constexpr char sourceBucket[] = "test_streaming_query_recovery_on_internal_retry";
+        constexpr char objectContent[] = R"(
+{"fqdn": "host1.example.com", "payload": "P1"}
+{"fqdn": "host2.example.com"                              })";
+        constexpr char objectPath[] = "path/test_object.json";
+        CreateBucketWithObject(sourceBucket, objectPath, objectContent);
+
+        constexpr char inputTopicName[] = "internalRetryInputTopicName";
+        constexpr char outputTopicName[] = "internalRetryOutputTopicName";
+        CreateTopic(inputTopicName);
+        CreateTopic(outputTopicName);
+
+        constexpr char pqSourceName[] = "pqSourceName";
+        constexpr char s3SourceName[] = "s3Source";
+        CreatePqSource(pqSourceName);
+        CreateS3Source(sourceBucket, s3SourceName);
+
+        constexpr char queryName[] = "streamingQuery";
+        ExecQuery(fmt::format(R"(
+            CREATE STREAMING QUERY `{query_name}` AS
+            DO BEGIN
+                PRAGMA ydb.HashJoinMode = "map";
+                $s3_lookup = SELECT * FROM `{s3_source}`.`path/` WITH (
+                    FORMAT = "json_each_row",
+                    SCHEMA (
+                        fqdn String NOT NULL,
+                        payload String
+                    )
+                );
+
+                -- Test that offsets are recovered
+                $pq_source = SELECT * FROM `{pq_source}`.`{input_topic}` WITH (
+                    FORMAT = "json_each_row",
+                    SCHEMA (
+                        time String NOT NULL,
+                        event String,
+                        host String
+                    )
+                );
+
+                $joined = SELECT
+                    Unwrap(l.payload) AS payload, -- Test failure here
+                    p.*
+                FROM $pq_source AS p
+                LEFT JOIN $s3_lookup AS l
+                ON (l.fqdn = p.host);
+
+                -- Test that state also recovered
+                $grouped = SELECT
+                    event,
+                    CAST(SOME(time) AS String) AS time,
+                    SOME(payload) AS payload,
+                    CAST(COUNT(*) AS String) AS count
+                FROM $joined
+                GROUP BY
+                    HOP (CAST(time AS Timestamp), "PT1H", "PT1H", "PT0H"),
+                    event;
+
+                INSERT INTO `{pq_source}`.`{output_topic}`
+                SELECT Unwrap(event || "-" || time || "-" || payload || "-" || count) FROM $grouped
+            END DO;)",
+            "query_name"_a = queryName,
+            "pq_source"_a = pqSourceName,
+            "s3_source"_a = s3SourceName,
+            "input_topic"_a = inputTopicName,
+            "output_topic"_a = outputTopicName
+        ));
+
+        CheckScriptExecutionsCount(1, 1);
+        Sleep(TDuration::Seconds(1));
+
+        // Fill HOP state for key A
+        WriteTopicMessages(inputTopicName, {
+            R"({"time": "2025-08-24T00:00:00.000000Z", "event": "A", "host": "host1.example.com"})",
+            R"({"time": "2025-08-25T00:00:00.000000Z", "event": "A", "host": "host1.example.com"})",
+        });
+        ReadTopicMessage(outputTopicName, "A-2025-08-24T00:00:00.000000Z-P1-1");
+
+        Sleep(TDuration::Seconds(2));
+        auto readDisposition = TInstant::Now();
+
+        // Write failure message for key B
+        WriteTopicMessage(inputTopicName, R"({"time": "2025-08-24T00:00:00.000000Z", "event": "B", "host": "host2.example.com"})");
+
+        // Wait script execution retry
+        WaitFor(TDuration::Seconds(10), "wait retry", [&](TString& error) {
+            const auto& results = ExecQuery(R"(
+                SELECT MAX(lease_generation) AS generation FROM `.metadata/script_executions`;
+            )");
+            UNIT_ASSERT_VALUES_EQUAL(results.size(), 1);
+
+            std::optional<i64> generation;
+            CheckScriptResult(results[0], 1, 1, [&](TResultSetParser& result) {
+                generation = result.ColumnParser(0).GetOptionalInt64();
+            });
+
+            if (!generation || *generation < 2) {
+                error = TStringBuilder() << "generation is: " << (generation ? ToString(*generation) : "null");
+                return false;
+            }
+
+            return true;
+        });
+
+        // Resolve query failure
+        UploadObject(sourceBucket, objectPath, R"(
+{"fqdn": "host1.example.com", "payload": "P1"}
+{"fqdn": "host2.example.com", "payload": "P2"             })");
+        Sleep(TDuration::Seconds(2));
+
+        // Check that offset is restored
+        WriteTopicMessage(inputTopicName, R"({"time": "2025-08-25T00:00:00.000000Z", "event": "B", "host": "host2.example.com"})");
+        ReadTopicMessage(outputTopicName, "B-2025-08-24T00:00:00.000000Z-P2-1", readDisposition);
+
+        Sleep(TDuration::Seconds(1));
+        readDisposition = TInstant::Now();
+
+        // Check that HOP state is restored
+        WriteTopicMessage(inputTopicName, R"({"time": "2025-08-26T00:00:00.000000Z", "event": "A", "host": "host1.example.com"})");
+        ReadTopicMessages(outputTopicName, {
+            "A-2025-08-25T00:00:00.000000Z-P1-1",
+            "B-2025-08-25T00:00:00.000000Z-P2-1"
+        }, readDisposition, /* sort */ true);
+    }
 }
 
 } // namespace NKikimr::NKqp

ydb/library/yql/dq/actors/compute/dq_compute_actor_checkpoints.cpp

@@ -586,13 +586,28 @@ NDqProto::ECheckpointingMode GetTaskCheckpointingMode(const TDqTaskSettings& tas
 }
 
 bool IsIngress(const TDqTaskSettings& task) {
-    // No inputs at all or the only inputs are sources.
-    for (const auto& input : task.GetInputs()) {
-        if (!input.HasSource()) {
-            return false;
+    // No inputs at all or there is no input channels with checkpoints.
+
+    const auto& inputs = task.GetInputs();
+    if (inputs.empty()) {
+        return true;
+    }
+
+    bool hasSource = false;
+    for (const auto& input : inputs) {
+        if (input.HasSource()) {
+            hasSource = true;
+            continue;
+        }
+
+        for (const auto& channel : input.GetChannels()) {
+            if (channel.GetCheckpointingMode() != NDqProto::CHECKPOINTING_MODE_DISABLED) {
+                return false;
+            }
         }
     }
-    return true;
+
+    return hasSource;
 }
 
 bool IsEgress(const TDqTaskSettings& task) {

ydb/library/yql/dq/actors/compute/dq_compute_actor_impl.h

@@ -779,7 +779,9 @@ class TDqComputeActorBase : public NActors::TActorBootstrapped<TDerived>
 
     void ResumeInputsByCheckpoint() override final {
         for (auto& [id, channelInfo] : InputChannelsMap) {
-            channelInfo.ResumeByCheckpoint();
+            if (channelInfo.PendingCheckpoint) {
+                channelInfo.ResumeByCheckpoint();
+            }
         }
     }
 

ydb/library/yql/dq/tasks/dq_tasks_graph.h

@@ -325,12 +325,27 @@ class TDqTasksGraph : private TMoveOnly {
     }
 
     bool IsIngress(const TTaskType& task) const {
+        // No inputs at all or there is no input channels with checkpoints.
+
+        if (!task.Inputs) {
+            return true;
+        }
+
+        bool hasSource = false;
         for (const auto& input : task.Inputs) {
-            if (!input.SourceType) {
-                return false;
+            if (input.SourceType) {
+                hasSource = true;
+                continue;
+            }
+
+            for (ui64 channelId : input.Channels) {
+                if (GetChannel(channelId).CheckpointingMode != NDqProto::CHECKPOINTING_MODE_DISABLED) {
+                    return false;
+                }
             }
         }
-        return true;
+
+        return hasSource;
     }
 
     static bool IsInfiniteSourceType(const TString& sourceType) {

ydb/library/yql/providers/dq/task_runner_actor/task_runner_actor.cpp

@@ -521,8 +521,8 @@ class TTaskRunnerActor
             ev->Get()->Index,
             std::move(batch),
             std::move(ev->Get()->Checkpoint),
-            ev->Get()->CheckpointSize,
             ev->Get()->Size,
+            ev->Get()->CheckpointSize,
             ev->Get()->Finished,
             ev->Get()->Changed);
     }