From 8692f43311760f3114f12703ac4cd1492160b722 Mon Sep 17 00:00:00 2001 From: ggjh-159 Date: Thu, 16 Jul 2026 15:44:56 +0800 Subject: [PATCH] feat(stateful): carry per-row RowKind via $row_kind column --- velox/experimental/stateful/CMakeLists.txt | 3 +- .../stateful/GroupWindowAggregator.cpp | 1 + .../stateful/LocalWindowAggregator.cpp | 1 + velox/experimental/stateful/RowKind.h | 34 +++ .../stateful/StatefulOperator.cpp | 5 +- .../experimental/stateful/StatefulOperator.h | 15 +- .../experimental/stateful/StatefulPlanner.cpp | 5 + .../stateful/StatefulSourceOperator.h | 49 ++++ velox/experimental/stateful/StreamElement.cpp | 100 +++++++ velox/experimental/stateful/StreamElement.h | 68 ++++- .../stateful/StreamKeyedOperator.cpp | 1 + .../experimental/stateful/StreamPartition.cpp | 73 ++--- velox/experimental/stateful/StreamPartition.h | 12 +- .../stateful/WatermarkAssigner.cpp | 1 + .../stateful/WindowAggregator.cpp | 1 + .../stateful/tests/CMakeLists.txt | 16 ++ .../tests/StatefulSourceOperatorTest.cpp | 259 ++++++++++++++++++ 17 files changed, 600 insertions(+), 44 deletions(-) create mode 100644 velox/experimental/stateful/RowKind.h create mode 100644 velox/experimental/stateful/StatefulSourceOperator.h create mode 100644 velox/experimental/stateful/StreamElement.cpp create mode 100644 velox/experimental/stateful/tests/StatefulSourceOperatorTest.cpp diff --git a/velox/experimental/stateful/CMakeLists.txt b/velox/experimental/stateful/CMakeLists.txt index a99edde56ac..d038305fcc8 100644 --- a/velox/experimental/stateful/CMakeLists.txt +++ b/velox/experimental/stateful/CMakeLists.txt @@ -39,7 +39,8 @@ velox_add_library( WatermarkSource.cpp WatermarkAssigner.cpp StreamRecordTimestampInserter.cpp - WatermarkIdleTracker.cpp) + WatermarkIdleTracker.cpp + StreamElement.cpp) velox_link_libraries( velox_stateful_exec diff --git a/velox/experimental/stateful/GroupWindowAggregator.cpp b/velox/experimental/stateful/GroupWindowAggregator.cpp index 7f12dce35bf..ada3b1ffac1 100644 --- a/velox/experimental/stateful/GroupWindowAggregator.cpp +++ b/velox/experimental/stateful/GroupWindowAggregator.cpp @@ -75,6 +75,7 @@ void GroupWindowAggregator::initializeState() { void GroupWindowAggregator::addInput(StreamElementPtr input) { VELOX_CHECK(!input_, "Last input has not been processed"); auto record = std::static_pointer_cast(input); + // TODO: preserve rowKind across this operator. input_ = record->record(); } diff --git a/velox/experimental/stateful/LocalWindowAggregator.cpp b/velox/experimental/stateful/LocalWindowAggregator.cpp index 46e32de3f60..1e9d3d442cd 100644 --- a/velox/experimental/stateful/LocalWindowAggregator.cpp +++ b/velox/experimental/stateful/LocalWindowAggregator.cpp @@ -41,6 +41,7 @@ LocalWindowAggregator::LocalWindowAggregator( void LocalWindowAggregator::addInput(StreamElementPtr input) { VELOX_CHECK(!input_, "Last input has not been processed"); auto record = std::static_pointer_cast(input); + // TODO: preserve rowKind across this operator. input_ = record->record(); } diff --git a/velox/experimental/stateful/RowKind.h b/velox/experimental/stateful/RowKind.h new file mode 100644 index 00000000000..b5afd30cd04 --- /dev/null +++ b/velox/experimental/stateful/RowKind.h @@ -0,0 +1,34 @@ +/* + * 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 + +namespace facebook::velox::stateful { + +enum class RowKind : int8_t { + INSERT = 0, + UPDATE_BEFORE = 1, + UPDATE_AFTER = 2, + DELETE = 3, +}; + +// Name of the synthetic trailing column that carries RowKind across the JNI +// boundary inside a merged RowVector. +constexpr std::string_view kRowKindColumnName = "$row_kind"; + +} // namespace facebook::velox::stateful diff --git a/velox/experimental/stateful/StatefulOperator.cpp b/velox/experimental/stateful/StatefulOperator.cpp index 9b6ceba6d71..45676d8af2f 100644 --- a/velox/experimental/stateful/StatefulOperator.cpp +++ b/velox/experimental/stateful/StatefulOperator.cpp @@ -54,8 +54,9 @@ bool StatefulOperator::isFinished() { void StatefulOperator::addInput(StreamElementPtr input) { auto record = std::static_pointer_cast(input); - operator_->traceInput(record->record()); - operator_->addInput(record->record()); + const auto& rowVector = record->record(); + operator_->traceInput(rowVector); + operator_->addInput(rowVector); } bool StatefulOperator::sourceEmpty() { diff --git a/velox/experimental/stateful/StatefulOperator.h b/velox/experimental/stateful/StatefulOperator.h index a641e6d24a6..763db341a0e 100644 --- a/velox/experimental/stateful/StatefulOperator.h +++ b/velox/experimental/stateful/StatefulOperator.h @@ -47,7 +47,7 @@ class StatefulOperator { : keyedStateBackendParameters_(keyedStateBackendParameters), operator_(std::move(op)), targets_(std::move(targets)) { - sink = operator_->operatorType() == "TableWrite"; + sink_ = operator_->operatorType() == "TableWrite"; } virtual ~StatefulOperator() = default; @@ -56,8 +56,17 @@ class StatefulOperator { virtual bool isFinished(); + // Receives a StreamRecord and feeds value() to the underlying velox operator. + // The base implementation drops rowKind, which is correct for stateless + // passthrough. Operators with retract semantics (group aggregation, stream + // join, etc.) must override this and advance() to act on rowKind. virtual void addInput(StreamElementPtr input); + // Pulls the next output batch from the underlying velox operator and wraps it + // as a StreamRecord. The base implementation assumes appendOnly output (no + // $row_kind column), which matches stateless passthrough. Operators with + // retract semantics must override this together with addInput() to emit + // rowKind according to their changelog semantics. virtual void advance(); void advanceWithFuture(ContinueFuture* future); @@ -161,7 +170,7 @@ class StatefulOperator { private: bool isSink() { - return sink; + return sink_; } bool isSource() const { @@ -170,7 +179,7 @@ class StatefulOperator { std::unique_ptr operator_; std::vector> targets_; - bool sink; + bool sink_; bool sourceEmpty_ = true; std::unique_ptr combinedWatermarkStatus_; StreamOperatorStateHandlerPtr stateHandler_; diff --git a/velox/experimental/stateful/StatefulPlanner.cpp b/velox/experimental/stateful/StatefulPlanner.cpp index e3c7a8621df..84bcf919276 100644 --- a/velox/experimental/stateful/StatefulPlanner.cpp +++ b/velox/experimental/stateful/StatefulPlanner.cpp @@ -45,6 +45,7 @@ #include "velox/experimental/stateful/KeySelector.h" #include "velox/experimental/stateful/LocalWindowAggregator.h" #include "velox/experimental/stateful/StatefulPlanNode.h" +#include "velox/experimental/stateful/StatefulSourceOperator.h" #include "velox/experimental/stateful/StreamJoin.h" #include "velox/experimental/stateful/StreamKeyedOperator.h" #include "velox/experimental/stateful/StreamPartition.h" @@ -420,6 +421,10 @@ StatefulOperatorPtr StatefulPlanner::transformGenericOperator( std::move(op), std::move(targets), std::move(watermarkGenerator)); } std::unique_ptr op = transformOperator(planNode.node()); + if (std::dynamic_pointer_cast(planNode.node())) { + return std::make_unique( + std::move(op), std::move(targets)); + } return std::make_unique(std::move(op), std::move(targets)); } diff --git a/velox/experimental/stateful/StatefulSourceOperator.h b/velox/experimental/stateful/StatefulSourceOperator.h new file mode 100644 index 00000000000..ae0c5d8e136 --- /dev/null +++ b/velox/experimental/stateful/StatefulSourceOperator.h @@ -0,0 +1,49 @@ +/* + * 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 "velox/experimental/stateful/StatefulOperator.h" +#include "velox/experimental/stateful/StreamElement.h" + +namespace facebook::velox::stateful { + +// Wraps a source operator (TableScan, etc.) that has no upstream input. +// StreamRecord::create splits the source RowVector: if it carries a trailing +// $row_kind column, the result carries per-row RowKind; otherwise the result +// is appendOnly. +class StatefulSourceOperator : public StatefulOperator { + public: + using StatefulOperator::StatefulOperator; + + // Sources never receive input from upstream. Reject loudly to catch wiring + // bugs in the operator chain. + void addInput(StreamElementPtr /*input*/) final { + VELOX_FAIL("StatefulSourceOperator does not support addInput"); + } + + void advance() override { + setSourceEmpty(true); + auto intermediateResult = op()->getOutput(); + if (!intermediateResult) { + return; + } + setSourceEmpty(false); + pushOutput( + StreamRecord::create(getPlanNodeId(), std::move(intermediateResult))); + } +}; + +} // namespace facebook::velox::stateful diff --git a/velox/experimental/stateful/StreamElement.cpp b/velox/experimental/stateful/StreamElement.cpp new file mode 100644 index 00000000000..90c703569c0 --- /dev/null +++ b/velox/experimental/stateful/StreamElement.cpp @@ -0,0 +1,100 @@ +/* + * 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/StreamElement.h" + +#include "velox/common/base/Exceptions.h" +#include "velox/experimental/stateful/RowKind.h" +#include "velox/type/Type.h" +#include "velox/vector/ConstantVector.h" + +namespace facebook::velox::stateful { + +RowVectorPtr StreamRecord::toMergedRowVector(bool alwaysAppendRowKind) const { + if (appendOnly() && !alwaysAppendRowKind) { + return record_; + } + auto rowType = asRowType(record_->type()); + std::vector names; + std::vector types; + std::vector children; + auto n = rowType->size(); + names.reserve(n + 1); + types.reserve(n + 1); + children.reserve(n + 1); + for (size_t i = 0; i < n; ++i) { + names.emplace_back(rowType->nameOf(i)); + types.emplace_back(rowType->childAt(i)); + children.emplace_back(record_->childAt(i)); + } + names.emplace_back(kRowKindColumnName); + types.emplace_back(TINYINT()); + if (appendOnly()) { + children.emplace_back(std::make_shared>( + record_->pool(), + size(), + false /*isNull*/, + TINYINT(), + static_cast(RowKind::INSERT))); + } else { + children.emplace_back(rowKind_); + } + auto mergedType = ROW(std::move(names), std::move(types)); + return std::make_shared( + record_->pool(), mergedType, nullptr, size(), std::move(children)); +} + +std::shared_ptr StreamRecord::create( + std::string nodeId, + RowVectorPtr merged) { + VELOX_USER_CHECK_NOT_NULL(merged, "merged RowVector must not be null"); + auto rowType = asRowType(merged->type()); + if (rowType->size() == 0 || + rowType->nameOf(rowType->size() - 1) != kRowKindColumnName) { + return std::make_shared(std::move(nodeId), std::move(merged)); + } + auto lastIdx = rowType->size() - 1; + auto kindVector = merged->childAt(lastIdx); + auto simpleKind = std::dynamic_pointer_cast>(kindVector); + VELOX_USER_CHECK_NOT_NULL( + simpleKind, + "$row_kind column must be a SimpleVector, got encoding: {}", + kindVector->encoding()); + std::vector names; + std::vector types; + std::vector children; + names.reserve(lastIdx); + types.reserve(lastIdx); + children.reserve(lastIdx); + for (size_t i = 0; i < lastIdx; ++i) { + names.emplace_back(rowType->nameOf(i)); + types.emplace_back(rowType->childAt(i)); + children.emplace_back(merged->childAt(i)); + } + auto valueType = ROW(std::move(names), std::move(types)); + auto value = std::make_shared( + merged->pool(), valueType, nullptr, merged->size(), std::move(children)); + // Normalize a constant INSERT row kind back to appendOnly (rowKind=null) so + // downstream code can rely on appendOnly() without re-checking the encoding. + if (simpleKind->isConstantEncoding() && + simpleKind->as>()->valueAt(0) == + static_cast(RowKind::INSERT)) { + return std::make_shared(std::move(nodeId), std::move(value)); + } + return std::make_shared( + std::move(nodeId), std::move(value), std::move(simpleKind)); +} + +} // namespace facebook::velox::stateful diff --git a/velox/experimental/stateful/StreamElement.h b/velox/experimental/stateful/StreamElement.h index ea04fdc49e0..68eae249299 100644 --- a/velox/experimental/stateful/StreamElement.h +++ b/velox/experimental/stateful/StreamElement.h @@ -17,6 +17,7 @@ #include #include "velox/core/PlanNode.h" +#include "velox/vector/SimpleVector.h" namespace facebook::velox::stateful { @@ -82,25 +83,55 @@ class WatermarkStatus : public StreamElement { const bool idle_; }; +// StreamRecord carries a RowVector of user columns together with an optional +// per-row RowKind vector (TINYINT, byte ordinals matching Flink RowKind: +// 0=INSERT, 1=UPDATE_BEFORE, 2=UPDATE_AFTER, 3=DELETE). When rowKind is null +// the record is appendOnly (all INSERT) and toMergedRowVector() returns the +// user RowVector without a trailing $row_kind column. class StreamRecord : public StreamElement { public: - StreamRecord(std::string nodeId, RowVectorPtr record) + // Primary constructors: user RowVector + optional per-row RowKind vector. + // rowKind == nullptr means appendOnly. + StreamRecord( + std::string nodeId, + RowVectorPtr record, + SimpleVectorPtr rowKind = nullptr) : StreamElement(nodeId), record_(std::move(record)), + rowKind_(std::move(rowKind)), timestamp_(-1), hasTimestamp_(false), key_(-1) {} + StreamRecord( + std::string nodeId, + RowVectorPtr record, + SimpleVectorPtr rowKind, + int64_t timestamp) + : StreamElement(nodeId), + record_(std::move(record)), + rowKind_(std::move(rowKind)), + timestamp_(timestamp), + hasTimestamp_(true), + key_(-1) {} + + // Convenience constructor for appendOnly records carrying a timestamp. StreamRecord(std::string nodeId, RowVectorPtr record, int64_t timestamp) : StreamElement(nodeId), record_(std::move(record)), + rowKind_(nullptr), timestamp_(timestamp), hasTimestamp_(true), key_(-1) {} - StreamRecord(std::string nodeId, int key, RowVectorPtr record) + StreamRecord( + std::string nodeId, + int key, + RowVectorPtr record, + SimpleVectorPtr rowKind = nullptr) : StreamElement(nodeId), record_(std::move(record)), + rowKind_(std::move(rowKind)), timestamp_(-1), hasTimestamp_(false), key_(key) {} @@ -109,6 +140,19 @@ class StreamRecord : public StreamElement { return record_; } + // May be null when appendOnly. + const SimpleVectorPtr& rowKind() const { + return rowKind_; + } + + bool appendOnly() const { + return !rowKind_; + } + + vector_size_t size() const { + return record_->size(); + } + int64_t timestamp() const { return timestamp_; } @@ -125,10 +169,30 @@ class StreamRecord : public StreamElement { return hasTimestamp_; } + // Merges record_ + rowKind_ into a single RowVector. By default an appendOnly + // record returns record_ as-is (no trailing $row_kind column). When + // alwaysAppendRowKind is true, the trailing $row_kind column is always + // present: rowKind_ for changelog records, or a ConstantVector + // (INSERT) for appendOnly records. Used by callers (e.g. JNI boundary, + // StatefulCalcOperator) that need a RowVector schema matching an + // N+1-column output type regardless of appendOnly state. + RowVectorPtr toMergedRowVector(bool alwaysAppendRowKind = false) const; + + // Splits a merged RowVector (user columns + optional trailing $row_kind) + // back into a StreamRecord. If the trailing column is absent or is a + // ConstantVector(INSERT), the result is appendOnly (rowKind=null). + static std::shared_ptr create( + std::string nodeId, + RowVectorPtr merged); + private: const RowVectorPtr record_; + const SimpleVectorPtr rowKind_; const int64_t timestamp_; bool hasTimestamp_ = false; const int key_; }; + +using StreamRecordPtr = std::shared_ptr; + } // namespace facebook::velox::stateful diff --git a/velox/experimental/stateful/StreamKeyedOperator.cpp b/velox/experimental/stateful/StreamKeyedOperator.cpp index 2d032c96a9a..065ff36c337 100644 --- a/velox/experimental/stateful/StreamKeyedOperator.cpp +++ b/velox/experimental/stateful/StreamKeyedOperator.cpp @@ -42,6 +42,7 @@ bool StreamKeyedOperator::isFinished() { void StreamKeyedOperator::addInput(StreamElementPtr input) { VELOX_CHECK_NULL(input_); auto record = std::static_pointer_cast(input); + // TODO: preserve rowKind across this operator. input_ = record->record(); } diff --git a/velox/experimental/stateful/StreamPartition.cpp b/velox/experimental/stateful/StreamPartition.cpp index f9cbf563079..bf0c1bcaef2 100644 --- a/velox/experimental/stateful/StreamPartition.cpp +++ b/velox/experimental/stateful/StreamPartition.cpp @@ -37,23 +37,27 @@ bool StreamPartition::isFinished() { } void StreamPartition::addInput(StreamElementPtr input) { - VELOX_CHECK_NULL(input_); + VELOX_CHECK_NULL(inputRowVector_); + VELOX_CHECK_NULL(inputRowKind_); auto record = std::static_pointer_cast(input); - input_ = record->record(); + inputRowVector_ = record->record(); + inputRowKind_ = record->rowKind(); } void StreamPartition::advance() { - prepareForInput(input_); + prepareForInput(inputRowVector_); if (numPartitions_ == 1) { - pushToTask(std::make_shared(getPlanNodeId(), 0, input_)); - input_.reset(); + pushToTask(std::make_shared( + getPlanNodeId(), 0, inputRowVector_, inputRowKind_)); + inputRowVector_.reset(); + inputRowKind_.reset(); return; } // TODO: The partition function doesn't use max parallelism. - partitionFunction_->partition(*input_, partitions_); - const auto numInput = input_->size(); + partitionFunction_->partition(*inputRowVector_, partitions_); + const auto numInput = inputRowVector_->size(); std::vector maxIndex(numPartitions_, 0); for (auto i = 0; i < numInput; ++i) { ++maxIndex[partitions_[i]]; @@ -67,18 +71,19 @@ void StreamPartition::advance() { ++maxIndex[partition]; } - const int64_t totalSize = input_->retainedSize(); for (auto i = 0; i < numPartitions_; i++) { auto partitionSize = maxIndex[i]; if (partitionSize == 0) { // Do not enqueue empty partitions. continue; } - auto partitionData = wrapChildren(input_, partitionSize, indexBuffers_[i]); - pushToTask( - std::make_shared(getPlanNodeId(), i, partitionData)); + auto [value, rowKind] = wrapForPartition( + inputRowVector_, inputRowKind_, partitionSize, indexBuffers_[i]); + pushToTask(std::make_shared( + getPlanNodeId(), i, std::move(value), std::move(rowKind))); } - input_.reset(); + inputRowVector_.reset(); + inputRowKind_.reset(); } void StreamPartition::pushToTask(StreamElementPtr output) { @@ -87,8 +92,7 @@ void StreamPartition::pushToTask(StreamElementPtr output) { task->addOutput(std::move(output)); } -// These methods are copied from LocalPartition.cpp, maybe we can refactor -// them to reuse the code in LocalPartition.cpp. +// prepareForInput and allocateIndexBuffers are adapted from LocalPartition.cpp. void StreamPartition::prepareForInput(RowVectorPtr& input) { // Lazy vectors must be loaded or processed to ensure the late materialized in // order. @@ -116,32 +120,35 @@ void StreamPartition::allocateIndexBuffers( } } -RowVectorPtr StreamPartition::wrapChildren( - const RowVectorPtr& input, +std::pair> +StreamPartition::wrapForPartition( + const RowVectorPtr& value, + const SimpleVectorPtr& rowKind, vector_size_t size, const BufferPtr& indices) { - RowVectorPtr result = std::make_shared( + RowVectorPtr wrappedValue = std::make_shared( op()->pool(), - input->type(), + value->type(), nullptr, size, - std::vector(input->childrenSize())); - - for (auto i = 0; i < input->childrenSize(); ++i) { - auto& child = result->childAt(i); - if (child && child->encoding() == VectorEncoding::Simple::DICTIONARY && - child.use_count() == 1) { - child->BaseVector::resize(size); - child->setWrapInfo(indices); - child->setValueVector(input->childAt(i)); - } else { - child = BaseVector::wrapInDictionary( - nullptr, indices, size, input->childAt(i)); - } + std::vector(value->childrenSize())); + for (auto i = 0; i < value->childrenSize(); ++i) { + wrappedValue->childAt(i) = + BaseVector::wrapInDictionary(nullptr, indices, size, value->childAt(i)); + } + wrappedValue->updateContainsLazyNotLoaded(); + + if (rowKind == nullptr) { + return {std::move(wrappedValue), nullptr}; } - result->updateContainsLazyNotLoaded(); - return result; + // Re-apply the same indices to rowKind so each partition carries the original + // row kind for its rows. wrapInDictionary may return a DictionaryVector or, + // for constant input, a ConstantVector; both inherit from SimpleVector. + auto wrapped = BaseVector::wrapInDictionary(nullptr, indices, size, rowKind); + auto wrappedRowKind = + std::dynamic_pointer_cast>(wrapped); + return {std::move(wrappedValue), std::move(wrappedRowKind)}; } } // namespace facebook::velox::stateful diff --git a/velox/experimental/stateful/StreamPartition.h b/velox/experimental/stateful/StreamPartition.h index 04532689bf2..7ee385bbc65 100644 --- a/velox/experimental/stateful/StreamPartition.h +++ b/velox/experimental/stateful/StreamPartition.h @@ -45,15 +45,21 @@ class StreamPartition : public StatefulOperator { void allocateIndexBuffers(const std::vector& sizes); - RowVectorPtr wrapChildren( - const RowVectorPtr& input, + // Applies partition indices to a (value, rowKind) pair: wraps each child + // column of value and (when present) the rowKind vector with the same + // dictionary. appendOnly inputs (rowKind == nullptr) produce appendOnly + // outputs. + std::pair> wrapForPartition( + const RowVectorPtr& value, + const SimpleVectorPtr& rowKind, vector_size_t size, const BufferPtr& indices); const std::unique_ptr partitionFunction_; const int numPartitions_; - RowVectorPtr input_; + RowVectorPtr inputRowVector_; + SimpleVectorPtr inputRowKind_; /// Reusable memory for hash calculation. std::vector partitions_; diff --git a/velox/experimental/stateful/WatermarkAssigner.cpp b/velox/experimental/stateful/WatermarkAssigner.cpp index eb1560c9852..0b67463d3cd 100644 --- a/velox/experimental/stateful/WatermarkAssigner.cpp +++ b/velox/experimental/stateful/WatermarkAssigner.cpp @@ -37,6 +37,7 @@ WatermarkAssigner::~WatermarkAssigner() { void WatermarkAssigner::addInput(StreamElementPtr input) { auto record = std::static_pointer_cast(input); + // TODO: preserve rowKind across this operator. input_ = record->record(); op()->addInput(input_); diff --git a/velox/experimental/stateful/WindowAggregator.cpp b/velox/experimental/stateful/WindowAggregator.cpp index 31f2f893f8f..1be15f41e18 100644 --- a/velox/experimental/stateful/WindowAggregator.cpp +++ b/velox/experimental/stateful/WindowAggregator.cpp @@ -65,6 +65,7 @@ void WindowAggregator::initializeState() { void WindowAggregator::addInput(StreamElementPtr input) { VELOX_CHECK(!input_, "Last input has not been processed"); auto record = std::static_pointer_cast(input); + // TODO: preserve rowKind across this operator. input_ = record->record(); } diff --git a/velox/experimental/stateful/tests/CMakeLists.txt b/velox/experimental/stateful/tests/CMakeLists.txt index 02d9eacbe61..f4877171325 100644 --- a/velox/experimental/stateful/tests/CMakeLists.txt +++ b/velox/experimental/stateful/tests/CMakeLists.txt @@ -23,6 +23,9 @@ add_executable(velox_stateful_watermark_idle_tracker_test add_executable(velox_stateful_idle_timer_manager_test IdleTimerManagerTest.cpp) +add_executable(velox_stateful_source_operator_test + StatefulSourceOperatorTest.cpp) + add_test(velox_stateful_watermark_generator_test velox_stateful_watermark_generator_test) @@ -35,6 +38,9 @@ add_test(velox_stateful_watermark_idle_tracker_test add_test(velox_stateful_idle_timer_manager_test velox_stateful_idle_timer_manager_test) +add_test(velox_stateful_source_operator_test + velox_stateful_source_operator_test) + target_link_libraries( velox_stateful_combined_watermark_status_test velox_common_base GTest::gtest GTest::gtest_main) @@ -84,3 +90,13 @@ target_link_libraries( velox_common_base GTest::gtest GTest::gtest_main) + +target_link_libraries( + velox_stateful_source_operator_test + velox_stateful_exec + 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/StatefulSourceOperatorTest.cpp b/velox/experimental/stateful/tests/StatefulSourceOperatorTest.cpp new file mode 100644 index 00000000000..8e83dd76f0f --- /dev/null +++ b/velox/experimental/stateful/tests/StatefulSourceOperatorTest.cpp @@ -0,0 +1,259 @@ +/* + * 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/StatefulSourceOperator.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/Values.h" +#include "velox/exec/tests/utils/OperatorTestBase.h" +#include "velox/experimental/stateful/RowKind.h" +#include "velox/experimental/stateful/StatefulOperator.h" +#include "velox/experimental/stateful/StreamElement.h" + +namespace facebook::velox::stateful::test { +namespace { + +class PresetOutputOperator : public exec::Operator { + public: + PresetOutputOperator( + exec::DriverCtx* driverCtx, + std::vector outputs, + std::string planNodeId = "preset_output") + : Operator( + driverCtx, + outputs.empty() ? ROW({"c"}, {BIGINT()}) + : asRowType(outputs[0]->type()), + 0, + std::move(planNodeId), + "PresetOutput"), + outputs_(std::move(outputs)) {} + + bool needsInput() const override { + return false; + } + + void addInput(RowVectorPtr) override {} + + RowVectorPtr getOutput() override { + if (nextOutput_ >= outputs_.size()) { + return nullptr; + } + return outputs_[nextOutput_++]; + } + + exec::BlockingReason isBlocked(ContinueFuture*) override { + return exec::BlockingReason::kNotBlocked; + } + + bool isFinished() override { + return nextOutput_ >= outputs_.size(); + } + + private: + std::vector outputs_; + size_t nextOutput_{0}; +}; + +class NoOutputOperator : public exec::Operator { + public: + explicit NoOutputOperator(exec::DriverCtx* driverCtx) + : Operator( + driverCtx, + ROW({"c"}, {BIGINT()}), + 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 CaptureTarget : public StatefulOperator { + public: + explicit CaptureTarget(exec::DriverCtx* driverCtx) + : StatefulOperator(std::make_unique(driverCtx), {}) {} + + void addInput(StreamElementPtr input) override { + lastInput_ = std::move(input); + } + + void advance() override {} + + StreamRecord* lastRecord() const { + if (!lastInput_ || !lastInput_->isRecord()) { + return nullptr; + } + return std::static_pointer_cast(lastInput_).get(); + } + + private: + StreamElementPtr lastInput_; +}; + +class StatefulSourceOperatorTest : public exec::test::OperatorTestBase { + protected: + void SetUp() override { + OperatorTestBase::SetUp(); + + core::PlanFragment planFragment; + planFragment.planNode = std::make_shared( + core::PlanNodeId{"values"}, std::vector{plainBatch()}); + executor_ = std::make_shared(1); + task_ = exec::Task::create( + "StatefulSourceOperatorTest_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 plainBatch() { + return makeRowVector({makeFlatVector({1, 2, 3})}); + } + + RowVectorPtr mergedBatch() { + auto value = makeFlatVector({10, 20, 30, 40}); + auto rowKind = makeFlatVector({0, 1, 2, 3}); + return makeRowVector({"c", "$row_kind"}, {value, rowKind}); + } + + std::shared_ptr executor_; + std::shared_ptr task_; + std::shared_ptr driver_; + std::unique_ptr driverCtx_; +}; + +TEST_F(StatefulSourceOperatorTest, rejectsAddInput) { + StatefulSourceOperator sourceOp( + std::make_unique( + driverCtx_.get(), std::vector{plainBatch()}), + {}); + + VELOX_ASSERT_THROW( + sourceOp.addInput(std::make_shared("source", plainBatch())), + "StatefulSourceOperator does not support addInput"); +} + +TEST_F(StatefulSourceOperatorTest, advanceWrapsPlainRowVectorAsAppendOnly) { + auto capture = std::make_unique(driverCtx_.get()); + auto* capturePtr = capture.get(); + std::vector targets; + targets.push_back(std::move(capture)); + + StatefulSourceOperator sourceOp( + std::make_unique( + driverCtx_.get(), std::vector{plainBatch()}), + std::move(targets)); + + EXPECT_TRUE(sourceOp.sourceEmpty()); + sourceOp.advance(); + EXPECT_FALSE(sourceOp.sourceEmpty()); + + auto* record = capturePtr->lastRecord(); + ASSERT_NE(record, nullptr); + EXPECT_TRUE(record->appendOnly()); + EXPECT_EQ(record->rowKind(), nullptr); + EXPECT_EQ(record->size(), 3); +} + +TEST_F(StatefulSourceOperatorTest, advanceSplitsMergedRowVector) { + auto capture = std::make_unique(driverCtx_.get()); + auto* capturePtr = capture.get(); + std::vector targets; + targets.push_back(std::move(capture)); + + StatefulSourceOperator sourceOp( + std::make_unique( + driverCtx_.get(), std::vector{mergedBatch()}), + std::move(targets)); + + sourceOp.advance(); + + auto* record = capturePtr->lastRecord(); + ASSERT_NE(record, nullptr); + EXPECT_FALSE(record->appendOnly()); + ASSERT_NE(record->rowKind(), nullptr); + EXPECT_EQ( + record->rowKind()->valueAt(0), static_cast(RowKind::INSERT)); + EXPECT_EQ( + record->rowKind()->valueAt(1), + static_cast(RowKind::UPDATE_BEFORE)); + EXPECT_EQ( + record->rowKind()->valueAt(2), + static_cast(RowKind::UPDATE_AFTER)); + EXPECT_EQ( + record->rowKind()->valueAt(3), static_cast(RowKind::DELETE)); + EXPECT_EQ(record->record()->type()->size(), 1); +} + +TEST_F(StatefulSourceOperatorTest, advanceTracksSourceEmpty) { + auto capture = std::make_unique(driverCtx_.get()); + std::vector targets; + targets.push_back(std::move(capture)); + + StatefulSourceOperator sourceOp( + std::make_unique( + driverCtx_.get(), std::vector{plainBatch()}), + std::move(targets)); + + EXPECT_TRUE(sourceOp.sourceEmpty()); + sourceOp.advance(); + EXPECT_FALSE(sourceOp.sourceEmpty()); + sourceOp.advance(); + EXPECT_TRUE(sourceOp.sourceEmpty()); +} + +} // 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(); +}