From eb49cea6282fd7dcbe3960d9f477ff4fc83e4240 Mon Sep 17 00:00:00 2001
From: Hongze Zhang <mailtozhz@126.com>
Date: Wed, 4 Dec 2024 18:09:52 -0800
Subject: [PATCH] test: A test case to guard against output order of
 LocalPartiton(kGather) in serial execution mode (#11726)

Summary:
As title. The test case guards against the output data order of `LocalPartiton(kGather)` to match the input data order when serial execution is on.

Pull Request resolved: https://github.com/facebookincubator/velox/pull/11726

Reviewed By: gggrace14

Differential Revision: D66784276

Pulled By: xiaoxmeng

fbshipit-source-id: dddc5410476b2173f887bcba850dd906a0b56257
---
 velox/exec/tests/LocalPartitionTest.cpp | 34 +++++++++++++++++++++++++
 1 file changed, 34 insertions(+)

diff --git a/velox/exec/tests/LocalPartitionTest.cpp b/velox/exec/tests/LocalPartitionTest.cpp
index 97ff74d71c5a..1a8dc480c45d 100644
--- a/velox/exec/tests/LocalPartitionTest.cpp
+++ b/velox/exec/tests/LocalPartitionTest.cpp
@@ -17,6 +17,7 @@
 #include "velox/exec/tests/utils/AssertQueryBuilder.h"
 #include "velox/exec/tests/utils/HiveConnectorTestBase.h"
 #include "velox/exec/tests/utils/PlanBuilder.h"
+#include "velox/functions/prestosql/window/WindowFunctionsRegistration.h"
 
 using namespace facebook::velox;
 using namespace facebook::velox::exec;
@@ -26,6 +27,7 @@ class LocalPartitionTest : public HiveConnectorTestBase {
  protected:
   void SetUp() override {
     HiveConnectorTestBase::SetUp();
+    window::prestosql::registerAllWindowFunctions();
   }
 
   template <typename T>
@@ -149,6 +151,38 @@ TEST_F(LocalPartitionTest, gather) {
   verifyExchangeSourceOperatorStats(task, 300, 3, 1);
 }
 
+TEST_F(LocalPartitionTest, gatherPreserveInputOrderWithSerialExecutionMode) {
+  const std::vector<RowVectorPtr> vectors = {
+      makeRowVector({makeFlatVector<int64_t>({10, 20})}),
+      makeRowVector({makeFlatVector<int64_t>({30, 40})}),
+      makeRowVector({makeFlatVector<int64_t>({50, 60})}),
+      makeRowVector({makeFlatVector<int64_t>({70, 80})}),
+      makeRowVector({makeFlatVector<int64_t>({90, 100})}),
+      makeRowVector({makeFlatVector<int64_t>({110, 120})})};
+
+  auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
+
+  auto valuesNode = [&](const std::vector<int>& indices) {
+    std::vector<RowVectorPtr> values;
+    for (const auto& index : indices) {
+      values.push_back(vectors[index]);
+    }
+    return PlanBuilder(planNodeIdGenerator).values(values).planNode();
+  };
+
+  auto op =
+      PlanBuilder(planNodeIdGenerator)
+          .localPartition(
+              {}, {valuesNode({0, 1, 2}), valuesNode({3}), valuesNode({4, 5})})
+          .window({"row_number() over () as r"})
+          .planNode();
+
+  AssertQueryBuilder(op, duckDbQueryRunner_)
+      .serialExecution(true)
+      .assertResults(
+          "VALUES (10, 1), (20, 2), (30, 3), (40, 4), (50, 5), (60, 6), (70, 7), (80, 8), (90, 9), (100, 10), (110, 11), (120, 12)");
+}
+
 TEST_F(LocalPartitionTest, partition) {
   std::vector<RowVectorPtr> vectors = {
       makeRowVector({makeFlatSequence<int32_t>(0, 100)}),