feat: Support watermark status propagation

velox/experimental/stateful/CombinedWatermarkStatus.cpp

@@ -18,22 +18,30 @@
 
 namespace facebook::velox::stateful {
 
-bool CombinedWatermarkStatus::updateWatermark(int index, int64_t timestamp) {
+bool CombinedWatermarkStatus::updateWatermark(
+    int32_t index,
+    int64_t timestamp) {
   VELOX_CHECK(index < partialWatermarks_.size(), "Index out of range");
   auto& watermark = partialWatermarks_[index];
 
-  if (!watermark.setWatermark(timestamp)) {
-    return false;
-  }
+  watermark.setWatermark(timestamp);
+  return updateCombinedWatermark();
+}
 
-  // If the watermark is already set, we do not update it.
+bool CombinedWatermarkStatus::updateStatus(int32_t index, bool idle) {
+  VELOX_CHECK(index < partialWatermarks_.size(), "Index out of range");
+  partialWatermarks_[index].setIdle(idle);
   return updateCombinedWatermark();
 }
 
 int64_t CombinedWatermarkStatus::getCombinedWatermark() {
   return combinedWatermark_;
 }
 
+bool CombinedWatermarkStatus::isIdle() const {
+  return idle_;
+}
+
 bool CombinedWatermarkStatus::updateCombinedWatermark() {
   int64_t minimumOverAll = INT64_MAX;
   bool allIdle = true;

velox/experimental/stateful/CombinedWatermarkStatus.h

@@ -29,14 +29,12 @@ class CombinedWatermarkStatus {
   class PartialWatermark {
    public:
     bool setWatermark(int64_t watermark) {
-      if (watermark < watermark_) {
-        // If the new watermark is less than or equal to the current one, we do
-        // not update it.
+      idle_ = false;
+      if (watermark <= watermark_) {
         return false;
       }
 
       watermark_ = watermark;
-      idle_ = false;
       return true;
     }
 
@@ -57,15 +55,20 @@ class CombinedWatermarkStatus {
     bool idle_ = false;
   };
 
-  CombinedWatermarkStatus(int numWatermarks) {
+  CombinedWatermarkStatus(int32_t numWatermarks) {
     VELOX_CHECK(numWatermarks > 0, "numWatermarks must be greater than 0");
     partialWatermarks_.resize(numWatermarks);
   }
 
-  bool updateWatermark(int index, int64_t timestamp);
+  bool updateWatermark(int32_t index, int64_t timestamp);
+
+  // idle == true means WatermarkStatus.IDLE; false means ACTIVE.
+  bool updateStatus(int32_t index, bool idle);
 
   int64_t getCombinedWatermark();
 
+  bool isIdle() const;
+
  private:
   bool updateCombinedWatermark();
 

velox/experimental/stateful/StatefulOperator.cpp

@@ -158,12 +158,35 @@ void StatefulOperator::emitWatermark(int64_t timestamp) {
   }
 }
 
-void StatefulOperator::processWatermark(int64_t timestamp, int index) {
-  if (combinedWatermarkStatus_->updateWatermark(index, timestamp)) {
+void StatefulOperator::emitWatermarkStatus(bool idle) {
+  if (isSink()) {
+    return;
+  }
+
+  if (targets_.empty()) {
+    auto task = std::static_pointer_cast<StatefulTask>(
+        operator_->operatorCtx()->driverCtx()->task);
+    task->addOutput(std::make_shared<WatermarkStatus>(getPlanNodeId(), idle));
+    return;
+  }
+  for (auto& target : targets_) {
+    target->processWatermarkStatus(idle);
+  }
+}
+
+void StatefulOperator::processWatermark(int64_t timestamp, int32_t index) {
+  const bool wasIdle = combinedWatermarkStatus_->isIdle();
+  const bool watermarkUpdated =
+      combinedWatermarkStatus_->updateWatermark(index, timestamp);
+  const bool isIdle = combinedWatermarkStatus_->isIdle();
+  if (wasIdle != isIdle) {
+    emitWatermarkStatus(isIdle);
+  }
+  if (watermarkUpdated) {
     // If the watermark is updated, we need to advance the timer service.
     int64_t combinedWatermark =
         combinedWatermarkStatus_->getCombinedWatermark();
-    processWatermark(combinedWatermark);
+    emitWatermark(combinedWatermark);
   }
 }
 
@@ -172,6 +195,21 @@ void StatefulOperator::processWatermark(int64_t timestamp) {
   emitWatermark(timestamp);
 }
 
+void StatefulOperator::processWatermarkStatus(bool idle, int32_t index) {
+  const bool wasIdle = combinedWatermarkStatus_->isIdle();
+  if (combinedWatermarkStatus_->updateStatus(index, idle)) {
+    emitWatermark(combinedWatermarkStatus_->getCombinedWatermark());
+  }
+  const bool isIdle = combinedWatermarkStatus_->isIdle();
+  if (wasIdle != isIdle) {
+    emitWatermarkStatus(isIdle);
+  }
+}
+
+void StatefulOperator::processWatermarkStatus(bool idle) {
+  processWatermarkStatus(idle, 0);
+}
+
 void StatefulOperator::initializeStateBackend(StateBackend* stateBackend) {
   if (!stateHandler_) {
     stateHandler_ = std::make_shared<StreamOperatorStateHandler>(

velox/experimental/stateful/StatefulOperator.h

@@ -66,10 +66,14 @@ class StatefulOperator {
 
   virtual void finish();
 
-  virtual void processWatermark(int64_t timestamp, int index);
+  virtual void processWatermark(int64_t timestamp, int32_t index);
 
   virtual void processWatermark(int64_t timestamp);
 
+  virtual void processWatermarkStatus(bool idle, int32_t index);
+
+  virtual void processWatermarkStatus(bool idle);
+
   virtual void initializeState();
 
   void initializeStateBackend(StateBackend* stateBackend);
@@ -130,6 +134,7 @@ class StatefulOperator {
 
   void pushOutput(StreamElementPtr output);
   void emitWatermark(int64_t timestamp);
+  void emitWatermarkStatus(bool idle);
   void setSourceEmpty(bool sourceEmpty) {
     sourceEmpty_ = sourceEmpty;
   }

velox/experimental/stateful/StatefulTask.cpp

@@ -185,12 +185,20 @@ std::shared_ptr<NativeCallbackBridge> StatefulTask::nativeCallbackBridge()
   return nativeCallbackBridge_;
 }
 
-void StatefulTask::notifyWatermark(int64_t watermark, int index) {
+void StatefulTask::notifyWatermark(int64_t watermark, int32_t index) {
   operatorChain_->processWatermark(watermark, index);
 }
 
 void StatefulTask::notifyWatermark(int64_t watermark) {
-  operatorChain_->processWatermark(watermark);
+  operatorChain_->processWatermark(watermark, 0);
+}
+
+void StatefulTask::notifyWatermarkStatus(bool idle, int32_t index) {
+  operatorChain_->processWatermarkStatus(idle, index);
+}
+
+void StatefulTask::notifyWatermarkStatus(bool idle) {
+  operatorChain_->processWatermarkStatus(idle);
 }
 
 void StatefulTask::initializeState(

velox/experimental/stateful/StatefulTask.h

@@ -63,10 +63,15 @@ class StatefulTask : public exec::Task {
 
   StreamElementPtr next(ContinueFuture* future, int32_t& retCode);
 
-  void notifyWatermark(int64_t watermark, int index);
+  void notifyWatermark(int64_t watermark, int32_t index);
 
   void notifyWatermark(int64_t watermark);
 
+  // idle == true means WatermarkStatus.IDLE; false means ACTIVE.
+  void notifyWatermarkStatus(bool idle, int32_t index);
+
+  void notifyWatermarkStatus(bool idle);
+
   void initializeState(
       const std::shared_ptr<const KeyedStateBackendParameters> params,
       const std::vector<std::string>& checkpointRecords = {});

velox/experimental/stateful/StreamElement.h

@@ -24,9 +24,17 @@ class StreamElement {
  public:
   StreamElement(std::string nodeId) : nodeId_(std::move(nodeId)) {}
 
-  virtual bool isWatermark() = 0;
+  virtual bool isWatermark() {
+    return false;
+  }
+
+  virtual bool isRecord() {
+    return false;
+  }
 
-  virtual bool isRecord() = 0;
+  virtual bool isWatermarkStatus() {
+    return false;
+  }
 
   const std::string nodeId() const {
     return nodeId_;
@@ -52,12 +60,26 @@ class Watermark : public StreamElement {
     return true;
   }
 
-  bool isRecord() override {
-    return false;
+ private:
+  const int64_t timestamp_;
+};
+
+class WatermarkStatus : public StreamElement {
+ public:
+  // idle == true means WatermarkStatus.IDLE; false means ACTIVE.
+  WatermarkStatus(std::string nodeId, bool idle)
+      : StreamElement(nodeId), idle_(idle) {}
+
+  bool idle() const {
+    return idle_;
+  }
+
+  bool isWatermarkStatus() override {
+    return true;
   }
 
  private:
-  const int64_t timestamp_;
+  const bool idle_;
 };
 
 class StreamRecord : public StreamElement {
@@ -95,10 +117,6 @@ class StreamRecord : public StreamElement {
     return key_;
   }
 
-  bool isWatermark() override {
-    return false;
-  }
-
   bool isRecord() override {
     return true;
   }

velox/experimental/stateful/tests/CMakeLists.txt

@@ -15,9 +15,19 @@
 add_executable(velox_stateful_watermark_generator_test
                WatermarkGeneratorTest.cpp)
 
+add_executable(velox_stateful_combined_watermark_status_test
+               CombinedWatermarkStatusTest.cpp ../CombinedWatermarkStatus.cpp)
+
 add_test(velox_stateful_watermark_generator_test
          velox_stateful_watermark_generator_test)
 
+add_test(velox_stateful_combined_watermark_status_test
+         velox_stateful_combined_watermark_status_test)
+
+target_link_libraries(
+  velox_stateful_combined_watermark_status_test velox_common_base GTest::gtest
+  GTest::gtest_main)
+
 target_link_libraries(
   velox_stateful_watermark_generator_test
   velox_stateful_exec

velox/experimental/stateful/tests/CombinedWatermarkStatusTest.cpp

@@ -0,0 +1,107 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "velox/experimental/stateful/CombinedWatermarkStatus.h"
+
+#include <gtest/gtest.h>
+
+namespace facebook::velox::stateful::test {
+namespace {
+
+TEST(CombinedWatermarkStatusTest, activeInputsUseMinimumWatermark) {
+  CombinedWatermarkStatus status(2);
+
+  EXPECT_FALSE(status.isIdle());
+
+  EXPECT_FALSE(status.updateWatermark(0, 100));
+  EXPECT_FALSE(status.isIdle());
+  EXPECT_EQ(status.getCombinedWatermark(), INT64_MIN);
+
+  EXPECT_TRUE(status.updateWatermark(1, 90));
+  EXPECT_EQ(status.getCombinedWatermark(), 90);
+
+  EXPECT_TRUE(status.updateWatermark(1, 110));
+  EXPECT_EQ(status.getCombinedWatermark(), 100);
+
+  EXPECT_TRUE(status.updateWatermark(0, 120));
+  EXPECT_EQ(status.getCombinedWatermark(), 110);
+}
+
+TEST(CombinedWatermarkStatusTest, idleInputsAreExcludedFromMinimumWatermark) {
+  CombinedWatermarkStatus status(2);
+
+  EXPECT_FALSE(status.updateWatermark(0, 100));
+  EXPECT_TRUE(status.updateWatermark(1, 50));
+  EXPECT_EQ(status.getCombinedWatermark(), 50);
+
+  EXPECT_TRUE(status.updateStatus(1, true));
+  EXPECT_FALSE(status.isIdle());
+  EXPECT_EQ(status.getCombinedWatermark(), 100);
+
+  EXPECT_FALSE(status.updateStatus(0, true));
+  EXPECT_TRUE(status.isIdle());
+  EXPECT_EQ(status.getCombinedWatermark(), 100);
+}
+
+TEST(
+    CombinedWatermarkStatusTest,
+    watermarkReactivatesIdleInputWithoutRollback) {
+  CombinedWatermarkStatus status(1);
+
+  EXPECT_TRUE(status.updateWatermark(0, 100));
+  EXPECT_EQ(status.getCombinedWatermark(), 100);
+
+  EXPECT_FALSE(status.updateStatus(0, true));
+  EXPECT_TRUE(status.isIdle());
+
+  EXPECT_FALSE(status.updateWatermark(0, 100));
+  EXPECT_FALSE(status.isIdle());
+  EXPECT_EQ(status.getCombinedWatermark(), 100);
+
+  EXPECT_FALSE(status.updateStatus(0, true));
+  EXPECT_TRUE(status.isIdle());
+
+  EXPECT_FALSE(status.updateWatermark(0, 90));
+  EXPECT_FALSE(status.isIdle());
+  EXPECT_EQ(status.getCombinedWatermark(), 100);
+}
+
+TEST(CombinedWatermarkStatusTest, activeStatusRejoinsWithPreviousWatermark) {
+  CombinedWatermarkStatus status(2);
+
+  EXPECT_FALSE(status.updateWatermark(0, 100));
+  EXPECT_TRUE(status.updateWatermark(1, 50));
+  EXPECT_EQ(status.getCombinedWatermark(), 50);
+
+  EXPECT_TRUE(status.updateStatus(1, true));
+  EXPECT_EQ(status.getCombinedWatermark(), 100);
+
+  EXPECT_TRUE(status.updateWatermark(0, 120));
+  EXPECT_EQ(status.getCombinedWatermark(), 120);
+
+  EXPECT_FALSE(status.updateStatus(1, false));
+  EXPECT_FALSE(status.isIdle());
+  EXPECT_EQ(status.getCombinedWatermark(), 120);
+
+  EXPECT_FALSE(status.updateWatermark(1, 130));
+  EXPECT_EQ(status.getCombinedWatermark(), 120);
+
+  EXPECT_TRUE(status.updateWatermark(0, 140));
+  EXPECT_EQ(status.getCombinedWatermark(), 130);
+}
+
+} // namespace
+} // namespace facebook::velox::stateful::test