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 @@ -149,6 +151,38 @@ TEST_F(LocalPartitionTest, gather) { verifyExchangeSourceOperatorStats(task, 300, 3, 1); } +TEST_F(LocalPartitionTest, gatherPreserveInputOrderWithSerialExecutionMode) { + const std::vector vectors = { + makeRowVector({makeFlatVector({10, 20})}), + makeRowVector({makeFlatVector({30, 40})}), + makeRowVector({makeFlatVector({50, 60})}), + makeRowVector({makeFlatVector({70, 80})}), + makeRowVector({makeFlatVector({90, 100})}), + makeRowVector({makeFlatVector({110, 120})})}; + + auto planNodeIdGenerator = std::make_shared(); + + auto valuesNode = [&](const std::vector& indices) { + std::vector 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 vectors = { makeRowVector({makeFlatSequence(0, 100)}),