[FLINK]Refine: Event time window aggregate for nexmark q12

velox/core/QueryConfig.h

@@ -562,6 +562,13 @@ class QueryConfig {
   static constexpr const char* kFieldNamesInJsonCastEnabled =
       "field_names_in_json_cast_enabled";
 
+  static constexpr const char* kStatefulTaskParallelism =
+      "stateful_task_parallelism";
+
+  int32_t statefulTaskParallelism() const {
+    return get<int32_t>(kStatefulTaskParallelism, 0);
+  }
+
   bool selectiveNimbleReaderEnabled() const {
     return get<bool>(kSelectiveNimbleReaderEnabled, false);
   }

velox/exec/HashPartitionFunction.cpp

@@ -109,11 +109,11 @@ std::optional<uint32_t> HashPartitionFunction::partition(
   } else {
     if (localExchange_) {
       for (auto i = 0; i < size; ++i) {
-        partitions[i] = localExchangeHash(hashes_[i]) % numPartitions_;
+        partitions[i] = numPartitions_ > 0 ? localExchangeHash(hashes_[i]) % numPartitions_ : localExchangeHash(hashes_[i]);
       }
     } else {
       for (auto i = 0; i < size; ++i) {
-        partitions[i] = hashes_[i] % numPartitions_;
+        partitions[i] = numPartitions_ > 0 ? hashes_[i] % numPartitions_ : hashes_[i];
       }
     }
   }

velox/exec/HashPartitionFunction.h

@@ -59,7 +59,7 @@ class HashPartitionFunction : public core::PartitionFunction {
       const std::vector<VectorPtr>& constValues);
 
   const bool localExchange_;
-  const int numPartitions_;
+  int numPartitions_{0};
   const std::optional<HashBitRange> hashBitRange_ = std::nullopt;
   std::vector<std::unique_ptr<VectorHasher>> hashers_;
 

velox/experimental/stateful/KeySelector.cpp

@@ -15,76 +15,86 @@
  */
 #include "velox/experimental/stateful/KeySelector.h"
 #include "velox/experimental/stateful/window/WindowPartitionFunction.h"
+#include <folly/container/F14Map.h>
+#include <algorithm>
+#include <memory>
+#include <optional>
 
 namespace facebook::velox::stateful {
 
+namespace {
+
+void allocateIndexBuffers(
+    const folly::F14FastMap<int64_t, vector_size_t>& partitionCounts,
+    folly::F14FastMap<int64_t, BufferPtr>& keyToIndexBuffers,
+    folly::F14FastMap<int64_t, vector_size_t*>& keyToRawIndices,
+    memory::MemoryPool* pool) {
+  keyToIndexBuffers.reserve(partitionCounts.size());
+  keyToRawIndices.reserve(partitionCounts.size());
+  for (const auto& [key, num] : partitionCounts) {
+    keyToIndexBuffers[key] = allocateIndices(num, pool);
+    keyToRawIndices[key] = keyToIndexBuffers[key]->asMutable<vector_size_t>();
+  }
+}
+
+} // namespace
+
 KeySelector::KeySelector(
     std::unique_ptr<core::PartitionFunction> partitionFunction,
     memory::MemoryPool* pool,
     int numPartitions)
     : partitionFunction_(std::move(partitionFunction)),
       pool_(pool),
-      numPartitions_(numPartitions) {}
+      numPartitions_(numPartitions) {
+  windowPartitionFunction_ =
+    dynamic_cast<WindowPartitionFunction*>(partitionFunction_.get());
+}
 
-std::map<int64_t, RowVectorPtr> KeySelector::partition(
-    const RowVectorPtr& input) {
-  if (numPartitions_ == 1) {
-    return std::map<int64_t, RowVectorPtr>{{0, input}};
-  }
+std::map<int64_t, RowVectorPtr> KeySelector::partition(const RowVectorPtr& input) {
   prepareForInput(input);
-
-  // TODO: The partition function doesn't use max parallelism.
-  std::vector<int64_t> partitions(input->size());
+  const auto numInput = input->size();
+  std::vector<int64_t> partitions(numInput);
   std::optional<int64_t> res;
-  auto windowPartitionFunction = dynamic_cast<WindowPartitionFunction*>(partitionFunction_.get());
-  if (windowPartitionFunction) {
-    res = windowPartitionFunction->partition(*input, partitions);
+  if (windowPartitionFunction_) {
+    res = windowPartitionFunction_->partition(*input, partitions);
   } else {
-    std::vector<uint32_t> tmpPartitions(input->size());
+    std::vector<uint32_t> tmpPartitions(numInput);
     std::optional<uint32_t> tmpRes = partitionFunction_->partition(*input, tmpPartitions);
     if (tmpRes) {
       res = static_cast<int64_t>(*tmpRes);
     }
-    for (vector_size_t i = 0; i < tmpPartitions.size(); ++i) {
-      partitions[i] = static_cast<int64_t>(tmpPartitions[i]);
+    for (vector_size_t i = 0; i < numInput; ++i) {
+      partitions[i] =static_cast<int64_t>(tmpPartitions[i]);
     }
   }
   if (res) {
-    // TODO: this is a optimization, as the RowVector may have be partitioned in
-    // local aggregation, so need not to partition again in global agg, but need
-    // to verify whether the judge condition is enough.
     return std::map<int64_t, RowVectorPtr>{{*res, input}};
   }
-  const auto numInput = input->size();
-  std::map<int64_t, vector_size_t> numOfKeys;
-  for (auto i = 0; i < numInput; ++i) {
-    if (numOfKeys.count(partitions[i]) == 0) {
-      numOfKeys[partitions[i]] = 1;
-    } else {
-      numOfKeys[partitions[i]] = numOfKeys[partitions[i]] + 1;
-    }
+
+  folly::F14FastMap<int64_t, vector_size_t> partitionCounts;
+  partitionCounts.reserve(std::min<vector_size_t>(numInput, 4096));
+
+  for (vector_size_t i = 0; i < numInput; ++i) {
+    ++partitionCounts[partitions[i]];
   }
+  folly::F14FastMap<int64_t, BufferPtr> keyToIndexBuffers;
+  folly::F14FastMap<int64_t, vector_size_t*> keyToRawIndices;
+  allocateIndexBuffers(partitionCounts, keyToIndexBuffers, keyToRawIndices, pool_);
 
-  std::map<int64_t, BufferPtr> keyToIndexBuffers;
-  std::map<int64_t, vector_size_t*> keyToRawIndices;
-  allocateIndexBuffers(numOfKeys, keyToIndexBuffers, keyToRawIndices);
+  folly::F14FastMap<int64_t, vector_size_t> nextRowIndex;
+  nextRowIndex.reserve(partitionCounts.size());
 
-  numOfKeys.clear();
-  for (auto i = 0; i < numInput; ++i) {
-    auto partition = partitions[i];
-    int index = 0;
-    if (numOfKeys.count(partition)) {
-      index = numOfKeys[partition];
-    }
+  for (vector_size_t i = 0; i < numInput; ++i) {
+    const auto partition = partitions[i];
+    const auto index = nextRowIndex[partition];
     keyToRawIndices[partition][index] = i;
-    numOfKeys[partition] = index + 1;
+    nextRowIndex[partition] = index + 1;
   }
 
   std::map<int64_t, RowVectorPtr> results;
-  for (auto& [key, partitionSize] : numOfKeys) {
-    auto partitionData =
-        wrapChildren(input, partitionSize, keyToIndexBuffers[key]);
-    results[key] = partitionData;
+  for (const auto& [partitionKey, partitionSize] : partitionCounts) {
+    results.emplace(
+      partitionKey, wrapChildren(input, partitionSize, keyToIndexBuffers[partitionKey]));
   }
   return results;
 }
@@ -99,16 +109,6 @@ void KeySelector::prepareForInput(const RowVectorPtr& input) {
   }
 }
 
-void KeySelector::allocateIndexBuffers(
-    const std::map<int64_t, vector_size_t>& numOfKeys,
-    std::map<int64_t, BufferPtr>& keyToIndexBuffers,
-    std::map<int64_t, vector_size_t*>& keyToRawIndices) {
-  for (auto& [key, num] : numOfKeys) {
-    keyToIndexBuffers[key] = allocateIndices(num, pool_);
-    keyToRawIndices[key] = keyToIndexBuffers[key]->asMutable<vector_size_t>();
-  }
-}
-
 RowVectorPtr KeySelector::wrapChildren(
     const RowVectorPtr& input,
     vector_size_t size,

velox/experimental/stateful/KeySelector.h

@@ -17,38 +17,41 @@
 
 #include <climits>
 #include <map>
-#include "velox/experimental/stateful/StatefulPlanNode.h"
 #include "velox/vector/ComplexVector.h"
+#include "velox/core/PlanNode.h"
 
 namespace facebook::velox::stateful {
 
+class WindowPartitionFunction;
+
 /// This class is relevant to flink KeySelector.
 /// It can partition the RowVector according to the key fields.
 class KeySelector {
  public:
   KeySelector(
       std::unique_ptr<core::PartitionFunction> partitionFunction,
       memory::MemoryPool* pool,
-      int numPartitions = INT_MAX);
+      int numPartitions = 0);
 
   std::map<int64_t, RowVectorPtr> partition(const RowVectorPtr& input);
 
+  void setNumPartitions(int numPartitions) {
+    numPartitions_ = numPartitions;
+  }
+
  private:
   void prepareForInput(const RowVectorPtr& input);
 
-  void allocateIndexBuffers(
-      const std::map<int64_t, vector_size_t>& numOfKeys,
-      std::map<int64_t, BufferPtr>& keyToIndexBuffers,
-      std::map<int64_t, vector_size_t*>& keyToRawIndices);
-
   RowVectorPtr wrapChildren(
       const RowVectorPtr& input,
       vector_size_t size,
       const BufferPtr& indices);
 
   const std::unique_ptr<core::PartitionFunction> partitionFunction_;
   memory::MemoryPool* pool_;
-  const int numPartitions_;
+  int numPartitions_;
+  /// Non-owning; null if partitionFunction_ is not a WindowPartitionFunction.
+  WindowPartitionFunction* windowPartitionFunction_{nullptr};
 };
 
 } // namespace facebook::velox::stateful

velox/experimental/stateful/LocalWindowAggregator.cpp

@@ -16,7 +16,6 @@
 #include "velox/experimental/stateful/LocalWindowAggregator.h"
 #include "velox/experimental/stateful/window/SliceAssigner.h"
 #include <cstdint>
-#include "velox/experimental/stateful/window/TimeWindowUtil.h"
 
 namespace facebook::velox::stateful {
 
@@ -47,15 +46,13 @@ void LocalWindowAggregator::advance() {
   if (!input_) {
     return;
   }
-
+  // partition input by key
   std::map<int64_t, RowVectorPtr> keyToData = keySelector_->partition(input_);
   for (const auto& [key, data] : keyToData) {
-    std::map<int64_t, RowVectorPtr> sliceEndToData =
-        sliceAssigner_->assignSliceEnd(data);
-    for (const auto& [sliceEnd, data] : sliceEndToData) {
-      // TODO: addElement may have data output.
-      auto windowData = data;
-      windowBuffer_->addElement(key, sliceEnd, windowData);
+    // assign slice end to data
+    std::map<int64_t, RowVectorPtr> sliceEndToData = sliceAssigner_->assignSliceEnd(data);
+    for (auto& [sliceEnd, data] : sliceEndToData) {
+      windowBuffer_->addElement(key, sliceEnd, data);
     }
   }
   input_.reset();
@@ -68,29 +65,45 @@ void LocalWindowAggregator::processWatermarkInternal(int64_t timestamp) {
       // we only need to call advanceProgress() when current watermark may
       // trigger window
       auto windowKeyToData = windowBuffer_->advanceProgress(currentWatermark_);
+      auto* aggregator = op().get();
+      auto* aggPool = aggregator->pool();
+      int64_t windowTriggered = -1;
       for (const auto& [windowKey, datas] : windowKeyToData) {
-        if (datas.empty()) {
+        const int64_t window = windowKey.window();
+        if (datas.empty() || currentWatermark_ < window) {
+          continue;
+        }
+        RowVectorPtr mergedInput = (datas.size() == 1) ? datas.front() : TimeWindowUtil::mergeVectors(datas, aggPool);
+        if (!mergedInput) {
           continue;
         }
-        op()->addInput(TimeWindowUtil::mergeVectors(datas, op()->pool()));
-        RowVectorPtr output = op()->getOutput();
+        aggregator->addInput(mergedInput);
+        RowVectorPtr output = aggregator->getOutput();
         if (output) {
-          pushOutput(std::make_shared<StreamRecord>(
-              getPlanNodeId(),
-              addWindowEndToVector(std::move(output), windowKey.window())));
+          pushOutput(
+            std::make_shared<StreamRecord>(getPlanNodeId(),
+            windowKey.key(),
+            std::move(addWindowEndToVector(output, window))));
+          if (windowTriggered < window) {
+            windowTriggered = window;
+          }
         }
       }
+      if (windowTriggered >= 0) {
+        windowBuffer_->clear(windowTriggered);
+      }
       nextTriggerWatermark_ = TimeWindowUtil::getNextTriggerWatermark(
           currentWatermark_,
           windowInterval_,
           shiftTimeZone_,
           useDayLightSaving_);
     }
   }
+  pushOutput(std::make_shared<Watermark>(getPlanNodeId(), timestamp));
 }
 
 void LocalWindowAggregator::close() {
-  processWatermarkInternal(INT_MAX);
+  // processWatermarkInternal(INT_MAX);
   StatefulOperator::close();
   input_.reset();
   windowBuffer_->clear();
@@ -103,7 +116,6 @@ RowVectorPtr LocalWindowAggregator::addWindowEndToVector(
     int64_t sliceEnd) {
   auto newColumn = BaseVector::create(BIGINT(), vector->size(), vector->pool());
   auto windowEndCol = newColumn->as<FlatVector<int64_t>>();
-
   for (int i = 0; i < vector->size(); ++i) {
     windowEndCol->set(i, sliceEnd);
   }

velox/experimental/stateful/StatefulPlanNode.cpp

@@ -15,6 +15,7 @@
  */
 #include "velox/experimental/stateful/StatefulPlanNode.h"
 #include "velox/experimental/stateful/window/WindowPartitionFunction.h"
+#include "velox/exec/PartitionFunction.h"
 #include <cstdint>
 
 namespace facebook::velox::stateful {
@@ -62,7 +63,7 @@ core::PlanNodePtr StatefulPlanNode::create(
 }
 
 void StatefulPlanNode::registerSerDe() {
-  registerPartitionFunctionSerDe();
+  exec::registerPartitionFunctionSerDe();
 
   auto& registry = DeserializationWithContextRegistryForSharedPtr();
 

velox/experimental/stateful/StatefulPlanner.cpp

@@ -260,8 +260,13 @@ StatefulOperatorPtr StatefulPlanner::transformStreamWindowAggregationOperator(
   std::unique_ptr<KeySelector> keySelectorForSliceAssigner = std::make_unique<KeySelector>(
       windowAggNode->sliceAssignerSpec()->create(INT_MAX, true), op->pool());
   std::unique_ptr<SliceAssigner> sliceAssigner =
-      std::make_unique<SliceAssigner>(std::move(keySelectorForSliceAssigner), windowAggNode->size(), windowAggNode->step(),
-      windowAggNode->offset(),Window::getType(windowAggNode->windowType()),windowAggNode->rowtimeIndex());
+      std::make_unique<SliceAssigner>(
+          std::move(keySelectorForSliceAssigner),
+          windowAggNode->size(),
+          windowAggNode->step(),
+          windowAggNode->offset(),
+          Window::getType(windowAggNode->windowType()),
+          windowAggNode->rowtimeIndex());
   if (windowAggNode->isLocalAgg()) {
     return std::make_unique<LocalWindowAggregator>(
         std::move(op),
@@ -282,6 +287,7 @@ StatefulOperatorPtr StatefulPlanner::transformStreamWindowAggregationOperator(
         std::move(keySelector),
         std::move(sliceAssigner),
         windowAggNode->windowInterval(),
+        windowAggNode->size(),
         windowAggNode->useDayLightSaving(),
         windowAggNode->isEventTime(),
         windowAggNode->windowStartIndex(),

velox/experimental/stateful/StatefulPlanner.h

@@ -36,11 +36,12 @@ class StatefulPlanner {
 
  protected:
   StatefulPlanner(exec::DriverCtx* ctx, StateBackend* stateBackend)
-      : ctx_(ctx), stateBackend_(stateBackend) {}
+      : ctx_(ctx), stateBackend_(stateBackend), parallelism_(ctx->queryConfig().statefulTaskParallelism()) {}
 
  private:
   exec::DriverCtx* ctx_ = nullptr;
   StateBackend* stateBackend_ = nullptr;
+  int32_t parallelism_ = 0;
 
   StatefulOperatorPtr transformStatefulOperators(
       const core::PlanNodePtr& planNode);