diff --git a/velox/experimental/stateful/CMakeLists.txt b/velox/experimental/stateful/CMakeLists.txt index 16a0752e9a2..69f3c34b24c 100644 --- a/velox/experimental/stateful/CMakeLists.txt +++ b/velox/experimental/stateful/CMakeLists.txt @@ -37,7 +37,8 @@ velox_add_library( GroupWindowAggregator.cpp WatermarkGenerator.cpp WatermarkSource.cpp - WatermarkAssigner.cpp) + WatermarkAssigner.cpp + StreamRecordTimestampInserter.cpp) velox_link_libraries( velox_stateful_exec diff --git a/velox/experimental/stateful/StatefulPlanNode.cpp b/velox/experimental/stateful/StatefulPlanNode.cpp index 2fe4fb8806d..c88aeefd683 100644 --- a/velox/experimental/stateful/StatefulPlanNode.cpp +++ b/velox/experimental/stateful/StatefulPlanNode.cpp @@ -67,6 +67,9 @@ void StatefulPlanNode::registerSerDe() { auto& registry = DeserializationWithContextRegistryForSharedPtr(); registry.Register("WatermarkAssignerNode", WatermarkAssignerNode::create); + registry.Register( + "StreamRecordTimestampInserterNode", + StreamRecordTimestampInserterNode::create); registry.Register( "WatermarkPushDownSpec", WatermarkPushDownSpec::deserialize); registry.Register( @@ -121,6 +124,36 @@ core::PlanNodePtr WatermarkAssignerNode::create( planNodeId, project, idleTimeout, rowtimeFieldIndex, watermarkInterval); } +const std::vector& +StreamRecordTimestampInserterNode::sources() const { + return kEmptySources; +} + +void StreamRecordTimestampInserterNode::addDetails( + std::stringstream& stream) const { + stream << project_->toString(); +} + +folly::dynamic StreamRecordTimestampInserterNode::serialize() const { + auto obj = PlanNode::serialize(); + obj["project"] = project_->serialize(); + obj["rowtimeFieldIndex"] = rowtimeFieldIndex_; + return obj; +} + +// static +core::PlanNodePtr StreamRecordTimestampInserterNode::create( + const folly::dynamic& obj, + void* context) { + auto planNodeId = obj["id"].asString(); + auto project = + ISerializable::deserialize(obj["project"], context); + int rowtimeFieldIndex = obj["rowtimeFieldIndex"].asInt(); + + return std::make_shared( + planNodeId, project, rowtimeFieldIndex); +} + void StreamJoinNode::addDetails(std::stringstream& stream) const { stream << "leftPartFuncSpec: " << leftPartFuncSpec_->toString(); stream << ", rightPartFuncSpec: " << rightPartFuncSpec_->toString(); diff --git a/velox/experimental/stateful/StatefulPlanNode.h b/velox/experimental/stateful/StatefulPlanNode.h index ceb4e5fc8e9..09124834fec 100644 --- a/velox/experimental/stateful/StatefulPlanNode.h +++ b/velox/experimental/stateful/StatefulPlanNode.h @@ -115,6 +115,49 @@ class WatermarkAssignerNode : public core::PlanNode { const int64_t watermarkInterval_; }; +/// Plan node for StreamRecordTimestampInserter. Wraps a project that selects +/// the rowtime column (single-column output) and records the rowtime field +/// index in the original input row. Used to annotate StreamRecord with the +/// rowtime without row-column conversion. +class StreamRecordTimestampInserterNode : public core::PlanNode { + public: + StreamRecordTimestampInserterNode( + const core::PlanNodeId& id, + std::shared_ptr project, + int rowtimeFieldIndex) + : PlanNode(id), + project_(std::move(project)), + rowtimeFieldIndex_(rowtimeFieldIndex) {} + + const RowTypePtr& outputType() const override { + return project_->outputType(); + } + + const std::vector& sources() const override; + + std::string_view name() const override { + return "StreamRecordTimestampInserter"; + } + + folly::dynamic serialize() const override; + + static core::PlanNodePtr create(const folly::dynamic& obj, void* context); + + const std::shared_ptr& project() const { + return project_; + } + + int rowtimeFieldIndex() const { + return rowtimeFieldIndex_; + } + + private: + void addDetails(std::stringstream& stream) const override; + + const std::shared_ptr project_; + const int rowtimeFieldIndex_; +}; + class WatermarkPushDownSpec : public ISerializable { public: WatermarkPushDownSpec( diff --git a/velox/experimental/stateful/StatefulPlanner.cpp b/velox/experimental/stateful/StatefulPlanner.cpp index afe44aa636b..e3c7a8621df 100644 --- a/velox/experimental/stateful/StatefulPlanner.cpp +++ b/velox/experimental/stateful/StatefulPlanner.cpp @@ -48,6 +48,7 @@ #include "velox/experimental/stateful/StreamJoin.h" #include "velox/experimental/stateful/StreamKeyedOperator.h" #include "velox/experimental/stateful/StreamPartition.h" +#include "velox/experimental/stateful/StreamRecordTimestampInserter.h" #include "velox/experimental/stateful/WatermarkAssigner.h" #include "velox/experimental/stateful/WatermarkGenerator.h" #include "velox/experimental/stateful/WatermarkSource.h" @@ -84,6 +85,10 @@ StatefulOperatorPtr StatefulPlanner::transformStatefulOperators( if (std::dynamic_pointer_cast( statefulNode->node()) != nullptr) { result = transformWatermarkAssignerOperator(*statefulNode); + } else if ( + std::dynamic_pointer_cast( + statefulNode->node()) != nullptr) { + result = transformStreamRecordTimestampInserterOperator(*statefulNode); } else if ( std::dynamic_pointer_cast( statefulNode->node()) != nullptr) { @@ -153,6 +158,23 @@ StatefulOperatorPtr StatefulPlanner::transformWatermarkAssignerOperator( watermarkAssignerNode->watermarkInterval()); } +StatefulOperatorPtr +StatefulPlanner::transformStreamRecordTimestampInserterOperator( + const StatefulPlanNode& planNode) { + std::vector targets = + transformStatefulOperators(planNode.targets()); + + auto inserterNode = + std::dynamic_pointer_cast( + planNode.node()); + + auto op = std::make_unique( + nextOperatorId(), ctx_, nullptr, inserterNode->project()); + + return std::make_unique( + std::move(op), std::move(targets), inserterNode->rowtimeFieldIndex()); +} + StatefulOperatorPtr StatefulPlanner::transformStreamPartitionOperator( const StatefulPlanNode& planNode) { std::vector targets = diff --git a/velox/experimental/stateful/StatefulPlanner.h b/velox/experimental/stateful/StatefulPlanner.h index afebee68545..e75a3ac6cf4 100644 --- a/velox/experimental/stateful/StatefulPlanner.h +++ b/velox/experimental/stateful/StatefulPlanner.h @@ -50,6 +50,8 @@ class StatefulPlanner { const core::PlanNodePtr& planNode); StatefulOperatorPtr transformWatermarkAssignerOperator( const StatefulPlanNode& planNode); + StatefulOperatorPtr transformStreamRecordTimestampInserterOperator( + const StatefulPlanNode& planNode); StatefulOperatorPtr transformStreamPartitionOperator( const StatefulPlanNode& planNode); StatefulOperatorPtr transformStreamJoinOperator( diff --git a/velox/experimental/stateful/StreamRecordTimestampInserter.cpp b/velox/experimental/stateful/StreamRecordTimestampInserter.cpp new file mode 100644 index 00000000000..9f58a86a084 --- /dev/null +++ b/velox/experimental/stateful/StreamRecordTimestampInserter.cpp @@ -0,0 +1,80 @@ +/* + * 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/StreamRecordTimestampInserter.h" + +#include "velox/experimental/stateful/WatermarkGenerator.h" +#include "velox/type/Timestamp.h" +#include "velox/vector/SimpleVector.h" + +namespace facebook::velox::stateful { + +StreamRecordTimestampInserter::StreamRecordTimestampInserter( + std::unique_ptr op, + std::vector> targets, + int rowtimeFieldIndex) + : StatefulOperator(std::move(op), std::move(targets)), + rowtimeFieldIndex_(rowtimeFieldIndex) {} + +void StreamRecordTimestampInserter::addInput(StreamElementPtr input) { + auto record = std::static_pointer_cast(input); + input_ = record->record(); + op()->addInput(input_); +} + +void StreamRecordTimestampInserter::advance() { + if (!input_) { + return; + } + watermark::validateRowtimeNoNulls(input_, rowtimeFieldIndex_); + // The wrapped FilterProject projects the rowtime column into a single-column + // RowVector; drain it to read timestamps. + RowVectorPtr timestampVector = + watermark::getTimestampVector(op().get(), input_); + + VELOX_CHECK_EQ( + timestampVector->childrenSize(), + 1, + "Expected single-column timestamp vector from wrapped FilterProject"); + auto child = timestampVector->childAt(0); + auto tsVector = child->as>(); + VELOX_CHECK_NOT_NULL( + tsVector, "Rowtime column is not a SimpleVector"); + const vector_size_t timestampSize = timestampVector->size(); + if (timestampSize == 0) { + pushOutput( + std::make_shared(getPlanNodeId(), std::move(input_))); + input_.reset(); + return; + } + int64_t maxTs = std::numeric_limits::min(); + for (vector_size_t i = 0; i < timestampSize; ++i) { + const int64_t millis = tsVector->valueAt(i).toMillis(); + if (millis > maxTs) { + maxTs = millis; + } + } + + pushOutput(std::make_shared( + getPlanNodeId(), std::move(input_), maxTs)); + input_.reset(); +} + +void StreamRecordTimestampInserter::close() { + StatefulOperator::close(); + input_.reset(); +} + +} // namespace facebook::velox::stateful diff --git a/velox/experimental/stateful/StreamRecordTimestampInserter.h b/velox/experimental/stateful/StreamRecordTimestampInserter.h new file mode 100644 index 00000000000..0dd400bb1ac --- /dev/null +++ b/velox/experimental/stateful/StreamRecordTimestampInserter.h @@ -0,0 +1,52 @@ +/* + * 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. + */ +#pragma once + +#include + +#include "velox/experimental/stateful/StatefulOperator.h" +#include "velox/experimental/stateful/StreamElement.h" + +namespace facebook::velox::stateful { + +/// Native counterpart of Flink's +/// org.apache.flink.table.runtime.operators.sink.StreamRecordTimestampInserter. +/// Reads the rowtime column of each input batch, takes the max timestamp, and +/// emits the original RowVector wrapped in a StreamRecord carrying that +/// timestamp. +class StreamRecordTimestampInserter : public StatefulOperator { + public: + StreamRecordTimestampInserter( + std::unique_ptr op, + std::vector> targets, + int rowtimeFieldIndex); + + void addInput(StreamElementPtr input) override; + + void advance() override; + + std::string name() const override { + return "StreamRecordTimestampInserter"; + } + + void close() override; + + private: + RowVectorPtr input_; + const int rowtimeFieldIndex_; +}; + +} // namespace facebook::velox::stateful diff --git a/velox/experimental/stateful/tests/CMakeLists.txt b/velox/experimental/stateful/tests/CMakeLists.txt index 5631d2f60e7..1e826475088 100644 --- a/velox/experimental/stateful/tests/CMakeLists.txt +++ b/velox/experimental/stateful/tests/CMakeLists.txt @@ -43,3 +43,25 @@ target_link_libraries( velox_common_base GTest::gtest GTest::gtest_main) + +add_executable(velox_stateful_stream_record_timestamp_inserter_test + StreamRecordTimestampInserterTest.cpp) + +add_test(velox_stateful_stream_record_timestamp_inserter_test + velox_stateful_stream_record_timestamp_inserter_test) + +target_link_libraries( + velox_stateful_stream_record_timestamp_inserter_test + velox_stateful_exec + velox_stateful_agg + velox_stateful_join + velox_stateful_rank + velox_stateful_state + velox_stateful_window + velox_stateful_udf + velox_vector_test_lib + velox_exec_test_lib + velox_exec + velox_common_base + GTest::gtest + GTest::gtest_main) diff --git a/velox/experimental/stateful/tests/StreamRecordTimestampInserterTest.cpp b/velox/experimental/stateful/tests/StreamRecordTimestampInserterTest.cpp new file mode 100644 index 00000000000..3733a205b84 --- /dev/null +++ b/velox/experimental/stateful/tests/StreamRecordTimestampInserterTest.cpp @@ -0,0 +1,294 @@ +/* + * 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/StreamRecordTimestampInserter.h" + +#include +#include + +#include "velox/common/base/tests/GTestUtils.h" +#include "velox/core/PlanFragment.h" +#include "velox/exec/Driver.h" +#include "velox/exec/Task.h" +#include "velox/exec/tests/utils/OperatorTestBase.h" +#include "velox/experimental/stateful/StatefulOperator.h" +#include "velox/type/Timestamp.h" + +namespace facebook::velox::stateful::test { +namespace { + +class TimestampPassthroughOperator : public exec::Operator { + public: + explicit TimestampPassthroughOperator(exec::DriverCtx* driverCtx) + : Operator( + driverCtx, + ROW({"rowtime"}, {TIMESTAMP()}), + 0, + "timestamp_passthrough", + "TimestampPassthrough") {} + + bool needsInput() const override { + return !input_; + } + + void addInput(RowVectorPtr input) override { + input_ = std::move(input); + } + + RowVectorPtr getOutput() override { + VELOX_CHECK_NOT_NULL(input_); + auto output = std::make_shared( + input_->pool(), + outputType_, + nullptr, + input_->size(), + std::vector{input_->childAt(0)}); + input_.reset(); + return output; + } + + exec::BlockingReason isBlocked(ContinueFuture*) override { + return exec::BlockingReason::kNotBlocked; + } + + bool isFinished() override { + return false; + } + + private: + RowVectorPtr input_; +}; + +class NoOutputOperator : public exec::Operator { + public: + explicit NoOutputOperator(exec::DriverCtx* driverCtx) + : Operator( + driverCtx, + ROW({"rowtime"}, {TIMESTAMP()}), + 0, + "spy_target", + "SpyTarget") {} + + bool needsInput() const override { + return true; + } + + void addInput(RowVectorPtr) override {} + + RowVectorPtr getOutput() override { + return nullptr; + } + + exec::BlockingReason isBlocked(ContinueFuture*) override { + return exec::BlockingReason::kNotBlocked; + } + + bool isFinished() override { + return false; + } +}; + +class SpyStatefulOperator : public StatefulOperator { + public: + explicit SpyStatefulOperator(exec::DriverCtx* driverCtx) + : StatefulOperator(std::make_unique(driverCtx), {}) {} + + void addInput(StreamElementPtr input) override { + if (input->isRecord()) { + auto record = std::static_pointer_cast(input); + recordSizes_.push_back(record->record()->size()); + timestamps_.push_back(record->timestamp()); + hasTimestamps_.push_back(record->hasTimestamp()); + } + } + + void advance() override {} + + const std::vector& recordSizes() const { + return recordSizes_; + } + + const std::vector& timestamps() const { + return timestamps_; + } + + const std::vector& hasTimestamps() const { + return hasTimestamps_; + } + + private: + std::vector recordSizes_; + std::vector timestamps_; + std::vector hasTimestamps_; +}; + +class StreamRecordTimestampInserterTest : public exec::test::OperatorTestBase { + protected: + void SetUp() override { + OperatorTestBase::SetUp(); + + core::PlanFragment planFragment; + planFragment.planNode = std::make_shared( + core::PlanNodeId{"values"}, std::vector{tsBatch({0})}); + executor_ = std::make_shared(1); + task_ = exec::Task::create( + "StreamRecordTimestampInserterTest_task", + std::move(planFragment), + 0, + core::QueryCtx::create(executor_.get()), + exec::Task::ExecutionMode::kParallel); + driver_ = exec::Driver::testingCreate(); + driverCtx_ = std::make_unique(task_, 0, 0, 0, 0); + driverCtx_->driver = driver_.get(); + } + + void TearDown() override { + driverCtx_.reset(); + driver_.reset(); + task_.reset(); + executor_.reset(); + OperatorTestBase::TearDown(); + } + + RowVectorPtr tsBatch(const std::vector& millis) { + std::vector ts; + ts.reserve(millis.size()); + for (auto m : millis) { + ts.push_back(Timestamp::fromMillis(m)); + } + return makeRowVector({makeFlatVector(ts)}); + } + + RowVectorPtr nullableTsBatch( + const std::vector>& millis) { + std::vector> ts; + ts.reserve(millis.size()); + for (auto m : millis) { + ts.push_back( + m.has_value() ? std::make_optional(Timestamp::fromMillis(*m)) + : std::nullopt); + } + return makeRowVector({makeNullableFlatVector(ts, TIMESTAMP())}); + } + + std::unique_ptr makeInserter( + std::unique_ptr spy, + int rowtimeFieldIndex) { + std::vector targets; + targets.push_back(std::move(spy)); + return std::make_unique( + std::make_unique(driverCtx_.get()), + std::move(targets), + rowtimeFieldIndex); + } + + std::shared_ptr executor_; + std::shared_ptr task_; + std::shared_ptr driver_; + std::unique_ptr driverCtx_; +}; + +TEST_F(StreamRecordTimestampInserterTest, emitsBatchMaxTimestamp) { + auto spy = std::make_unique(driverCtx_.get()); + auto* spyPtr = spy.get(); + auto inserter = makeInserter(std::move(spy), /*rowtimeFieldIndex=*/0); + + inserter->addInput( + std::make_shared("input", tsBatch({100, 200, 50, 150}))); + inserter->advance(); + + ASSERT_EQ(1, spyPtr->recordSizes().size()); + EXPECT_EQ(4, spyPtr->recordSizes()[0]); + ASSERT_EQ(1, spyPtr->timestamps().size()); + EXPECT_EQ(200, spyPtr->timestamps()[0]); + ASSERT_EQ(1, spyPtr->hasTimestamps().size()); + EXPECT_TRUE(spyPtr->hasTimestamps()[0]); +} + +TEST_F(StreamRecordTimestampInserterTest, emitsNoOutputWithoutInput) { + auto spy = std::make_unique(driverCtx_.get()); + auto* spyPtr = spy.get(); + auto inserter = makeInserter(std::move(spy), /*rowtimeFieldIndex=*/0); + + inserter->advance(); + + EXPECT_TRUE(spyPtr->recordSizes().empty()); + EXPECT_TRUE(spyPtr->timestamps().empty()); +} + +TEST_F(StreamRecordTimestampInserterTest, rejectsNullRowtime) { + auto spy = std::make_unique(driverCtx_.get()); + auto inserter = makeInserter(std::move(spy), /*rowtimeFieldIndex=*/0); + + inserter->addInput(std::make_shared( + "input", nullableTsBatch({100, std::nullopt, 2000}))); + VELOX_ASSERT_THROW( + inserter->advance(), "RowTime field should not have nulls"); +} + +TEST_F(StreamRecordTimestampInserterTest, handlesSingleRowBatch) { + auto spy = std::make_unique(driverCtx_.get()); + auto* spyPtr = spy.get(); + auto inserter = makeInserter(std::move(spy), /*rowtimeFieldIndex=*/0); + + inserter->addInput(std::make_shared("input", tsBatch({1234}))); + inserter->advance(); + + ASSERT_EQ(1, spyPtr->recordSizes().size()); + EXPECT_EQ(1, spyPtr->recordSizes()[0]); + EXPECT_EQ(1234, spyPtr->timestamps()[0]); +} + +TEST_F(StreamRecordTimestampInserterTest, handlesMultipleBatchesIndependently) { + auto spy = std::make_unique(driverCtx_.get()); + auto* spyPtr = spy.get(); + auto inserter = makeInserter(std::move(spy), /*rowtimeFieldIndex=*/0); + + inserter->addInput( + std::make_shared("input", tsBatch({100, 300}))); + inserter->advance(); + inserter->addInput( + std::make_shared("input", tsBatch({50, 80}))); + inserter->advance(); + + ASSERT_EQ(2, spyPtr->recordSizes().size()); + EXPECT_EQ(300, spyPtr->timestamps()[0]); + EXPECT_EQ(80, spyPtr->timestamps()[1]); +} + +TEST_F(StreamRecordTimestampInserterTest, emitsEmptyBatchWithoutTimestamp) { + auto spy = std::make_unique(driverCtx_.get()); + auto* spyPtr = spy.get(); + auto inserter = makeInserter(std::move(spy), /*rowtimeFieldIndex=*/0); + + inserter->addInput(std::make_shared("input", tsBatch({}))); + inserter->advance(); + + ASSERT_EQ(1, spyPtr->recordSizes().size()); + EXPECT_EQ(0, spyPtr->recordSizes()[0]); + ASSERT_EQ(1, spyPtr->hasTimestamps().size()); + EXPECT_FALSE(spyPtr->hasTimestamps()[0]); +} + +} // namespace +} // namespace facebook::velox::stateful::test + +int main(int argc, char* argv[]) { + testing::InitGoogleTest(&argc, argv); + folly::Init init(&argc, &argv, false); + gflags::ParseCommandLineFlags(&argc, &argv, true); + return RUN_ALL_TESTS(); +}