diff --git a/velox/experimental/stateful/CombinedWatermarkStatus.cpp b/velox/experimental/stateful/CombinedWatermarkStatus.cpp index ee41cabb437..7e302cec937 100644 --- a/velox/experimental/stateful/CombinedWatermarkStatus.cpp +++ b/velox/experimental/stateful/CombinedWatermarkStatus.cpp @@ -18,15 +18,19 @@ 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(); } @@ -34,6 +38,10 @@ int64_t CombinedWatermarkStatus::getCombinedWatermark() { return combinedWatermark_; } +bool CombinedWatermarkStatus::isIdle() const { + return idle_; +} + bool CombinedWatermarkStatus::updateCombinedWatermark() { int64_t minimumOverAll = INT64_MAX; bool allIdle = true; diff --git a/velox/experimental/stateful/CombinedWatermarkStatus.h b/velox/experimental/stateful/CombinedWatermarkStatus.h index d4c4d52a10a..d0ec188ca3e 100644 --- a/velox/experimental/stateful/CombinedWatermarkStatus.h +++ b/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(); diff --git a/velox/experimental/stateful/StatefulOperator.cpp b/velox/experimental/stateful/StatefulOperator.cpp index 999d31b4907..5a2ebbf6764 100644 --- a/velox/experimental/stateful/StatefulOperator.cpp +++ b/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( + operator_->operatorCtx()->driverCtx()->task); + task->addOutput(std::make_shared(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( diff --git a/velox/experimental/stateful/StatefulOperator.h b/velox/experimental/stateful/StatefulOperator.h index 8b7ae6c2153..2ae7e3eb102 100644 --- a/velox/experimental/stateful/StatefulOperator.h +++ b/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; } diff --git a/velox/experimental/stateful/StatefulTask.cpp b/velox/experimental/stateful/StatefulTask.cpp index 5829390fb3e..d3feb9598d1 100644 --- a/velox/experimental/stateful/StatefulTask.cpp +++ b/velox/experimental/stateful/StatefulTask.cpp @@ -185,12 +185,20 @@ std::shared_ptr 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( diff --git a/velox/experimental/stateful/StatefulTask.h b/velox/experimental/stateful/StatefulTask.h index 53bbf7c7cf2..96b591efd75 100644 --- a/velox/experimental/stateful/StatefulTask.h +++ b/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 params, const std::vector& checkpointRecords = {}); diff --git a/velox/experimental/stateful/StreamElement.h b/velox/experimental/stateful/StreamElement.h index bd16935e24a..ea04fdc49e0 100644 --- a/velox/experimental/stateful/StreamElement.h +++ b/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; } diff --git a/velox/experimental/stateful/tests/CMakeLists.txt b/velox/experimental/stateful/tests/CMakeLists.txt index 46683adff34..5631d2f60e7 100644 --- a/velox/experimental/stateful/tests/CMakeLists.txt +++ b/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 diff --git a/velox/experimental/stateful/tests/CombinedWatermarkStatusTest.cpp b/velox/experimental/stateful/tests/CombinedWatermarkStatusTest.cpp new file mode 100644 index 00000000000..027aacba3d1 --- /dev/null +++ b/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 + +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