feat: Add StreamRecordTimestampInserter operator

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

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<core::PlanNodePtr>&
+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<core::ProjectNode>(obj["project"], context);
+  int rowtimeFieldIndex = obj["rowtimeFieldIndex"].asInt();
+
+  return std::make_shared<const StreamRecordTimestampInserterNode>(
+      planNodeId, project, rowtimeFieldIndex);
+}
+
 void StreamJoinNode::addDetails(std::stringstream& stream) const {
   stream << "leftPartFuncSpec: " << leftPartFuncSpec_->toString();
   stream << ", rightPartFuncSpec: " << rightPartFuncSpec_->toString();

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<const core::ProjectNode> project,
+      int rowtimeFieldIndex)
+      : PlanNode(id),
+        project_(std::move(project)),
+        rowtimeFieldIndex_(rowtimeFieldIndex) {}
+
+  const RowTypePtr& outputType() const override {
+    return project_->outputType();
+  }
+
+  const std::vector<core::PlanNodePtr>& 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<const core::ProjectNode>& project() const {
+    return project_;
+  }
+
+  int rowtimeFieldIndex() const {
+    return rowtimeFieldIndex_;
+  }
+
+ private:
+  void addDetails(std::stringstream& stream) const override;
+
+  const std::shared_ptr<const core::ProjectNode> project_;
+  const int rowtimeFieldIndex_;
+};
+
 class WatermarkPushDownSpec : public ISerializable {
  public:
   WatermarkPushDownSpec(

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<const WatermarkAssignerNode>(
           statefulNode->node()) != nullptr) {
     result = transformWatermarkAssignerOperator(*statefulNode);
+  } else if (
+      std::dynamic_pointer_cast<const StreamRecordTimestampInserterNode>(
+          statefulNode->node()) != nullptr) {
+    result = transformStreamRecordTimestampInserterOperator(*statefulNode);
   } else if (
       std::dynamic_pointer_cast<const StreamPartitionNode>(
           statefulNode->node()) != nullptr) {
@@ -153,6 +158,23 @@ StatefulOperatorPtr StatefulPlanner::transformWatermarkAssignerOperator(
       watermarkAssignerNode->watermarkInterval());
 }
 
+StatefulOperatorPtr StatefulPlanner::
+    transformStreamRecordTimestampInserterOperator(
+        const StatefulPlanNode& planNode) {
+  std::vector<StatefulOperatorPtr> targets =
+      transformStatefulOperators(planNode.targets());
+
+  auto inserterNode =
+      std::dynamic_pointer_cast<const StreamRecordTimestampInserterNode>(
+          planNode.node());
+
+  auto op = std::make_unique<exec::FilterProject>(
+      nextOperatorId(), ctx_, nullptr, inserterNode->project());
+
+  return std::make_unique<StreamRecordTimestampInserter>(
+      std::move(op), std::move(targets), inserterNode->rowtimeFieldIndex());
+}
+
 StatefulOperatorPtr StatefulPlanner::transformStreamPartitionOperator(
     const StatefulPlanNode& planNode) {
   std::vector<StatefulOperatorPtr> targets =

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(

velox/experimental/stateful/StreamRecordTimestampInserter.cpp

@@ -0,0 +1,79 @@
+/*
+ * 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<exec::Operator> op,
+    std::vector<std::unique_ptr<StatefulOperator>> targets,
+    int rowtimeFieldIndex)
+    : StatefulOperator(std::move(op), std::move(targets)),
+      rowtimeFieldIndex_(rowtimeFieldIndex) {}
+
+void StreamRecordTimestampInserter::addInput(StreamElementPtr input) {
+  auto record = std::static_pointer_cast<StreamRecord>(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<SimpleVector<Timestamp>>();
+  VELOX_CHECK_NOT_NULL(tsVector, "Rowtime column is not a SimpleVector<Timestamp>");
+  const vector_size_t timestampSize = timestampVector->size();
+  if (timestampSize == 0) {
+    pushOutput(std::make_shared<StreamRecord>(
+        getPlanNodeId(), std::move(input_)));
+    input_.reset();
+    return;
+  }
+  int64_t maxTs = std::numeric_limits<int64_t>::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<StreamRecord>(
+      getPlanNodeId(), std::move(input_), maxTs));
+  input_.reset();
+}
+
+void StreamRecordTimestampInserter::close() {
+  StatefulOperator::close();
+  input_.reset();
+}
+
+} // namespace facebook::velox::stateful

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 <cstdint>
+
+#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<exec::Operator> op,
+      std::vector<std::unique_ptr<StatefulOperator>> 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

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)