From 8612db55bbc545a0b8bc0d7549591b75dc509a0e Mon Sep 17 00:00:00 2001 From: Patrick Ribbsaeter Date: Wed, 19 Aug 2026 22:06:40 +0200 Subject: [PATCH] fix(exec): Emit TopNRowNumber in-memory partitions in ascending rank order (#18494) Differential Revision: D116638042 Pull Request resolved: https://github.com/facebookincubator/velox/pull/18529 --- velox/core/PlanNode.h | 3 +- velox/exec/TopNRowNumber.cpp | 101 +++++++++++++------------ velox/exec/TopNRowNumber.h | 33 ++++---- velox/exec/tests/TopNRowNumberTest.cpp | 77 +++++++++++++++++++ 4 files changed, 145 insertions(+), 69 deletions(-) diff --git a/velox/core/PlanNode.h b/velox/core/PlanNode.h index f73254d..1bec4a2 100644 --- a/velox/core/PlanNode.h +++ b/velox/core/PlanNode.h @@ -5994,7 +5994,8 @@ using MarkSortedNodePtr = std::shared_ptr; /// Optimized version of a WindowNode for a single row_number, rank or /// dense_rank function with a limit over sorted partitions. The output of this /// node contains all input columns followed by an optional -/// 'rowNumberColumnName' BIGINT column. +/// 'rowNumberColumnName' BIGINT column, with rows within each partition emitted +/// in ascending order of sorting keys (matching WindowNode). /// TODO: This node will be renamed to TopNRank or TopNRowNode once all the /// support for handling rank and dense_rank is committed to Velox. class TopNRowNumberNode : public PlanNode { diff --git a/velox/exec/TopNRowNumber.cpp b/velox/exec/TopNRowNumber.cpp index b26b851..53f48e9 100644 --- a/velox/exec/TopNRowNumber.cpp +++ b/velox/exec/TopNRowNumber.cpp @@ -544,15 +544,24 @@ vector_size_t TopNRowNumber::fixTopRank(TopRows& partition) { } TopNRowNumber::TopRows* TopNRowNumber::nextPartition() { - auto setNextRankAndPeer = [&](TopRows& partition) { - nextRank_ = fixTopRank(partition); + auto preparePartition = [&](TopRows& partition) { + fixTopRank(partition); + const auto numRows = partition.rows.size(); + currentPartitionRows_.resize(numRows); + for (int32_t i = static_cast(numRows) - 1; i >= 0; --i) { + currentPartitionRows_[i] = partition.rows.top(); + partition.rows.pop(); + } + currentPartitionRowOffset_ = 0; + nextRank_ = 1; numPeers_ = 1; + previousRow_ = nullptr; }; if (!table_) { if (!outputPartitionNumber_) { outputPartitionNumber_ = 0; - setNextRankAndPeer(*singlePartition_); + preparePartition(*singlePartition_); return singlePartition_.get(); } return nullptr; @@ -578,56 +587,51 @@ TopNRowNumber::TopRows* TopNRowNumber::nextPartition() { } auto partition = &partitionAt(partitions_[outputPartitionNumber_.value()]); - setNextRankAndPeer(*partition); + preparePartition(*partition); return partition; } -template -void TopNRowNumber::computeNextRankInMemory( - TopRows& partition, - vector_size_t outputIndex) { - if constexpr (TRank == core::TopNRowNumberNode::RankFunction::kRowNumber) { - nextRank_ -= 1; - return; - } - - // This is the logic for rank() and dense_rank(). - // If the next row is a peer of the current one, then the rank remains the - // same. - if (comparator_.compare(outputRows_[outputIndex], partition.rows.top()) == - 0) { - return; - } - - // The new row is not a peer of the current one. So dense_rank drops the - // rank by 1, but rank drops by the number of peers of the new top - // row (new rank) in TopRows queue. - if constexpr (TRank == core::TopNRowNumberNode::RankFunction::kDenseRank) { - nextRank_ -= 1; - } else { - nextRank_ -= partition.numTopRankRows(); - } -} - template void TopNRowNumber::appendPartitionRows( - TopRows& partition, vector_size_t numRows, vector_size_t outputOffset, FlatVector* rankValues) { - // The partition.rows priority queue pops rows in order of reverse - // ranks. Output rows based on nextRank_ and update it with each row. - for (auto i = 0; i < numRows; ++i) { - auto index = outputOffset + i; + // Output rows in ascending order of ranks (lowest sorting key / rank 1 + // first). + for (vector_size_t i = 0; i < numRows; ++i) { + const auto rowIndex = currentPartitionRowOffset_ + i; + const auto outputIndex = outputOffset + i; + char* row = currentPartitionRows_[rowIndex]; + outputRows_[outputIndex] = row; + if (rankValues) { - rankValues->set(index, nextRank_); - } - outputRows_[index] = partition.rows.top(); - partition.rows.pop(); - if (!partition.rows.empty()) { - computeNextRankInMemory(partition, index); + if (rowIndex == 0) { + nextRank_ = 1; + numPeers_ = 1; + } else { + if constexpr ( + TRank == core::TopNRowNumberNode::RankFunction::kRowNumber) { + nextRank_ = rowIndex + 1; + } else if constexpr ( + TRank == core::TopNRowNumberNode::RankFunction::kDenseRank) { + if (comparator_.compare(previousRow_, row) != 0) { + nextRank_ += 1; + } + } else { + // kRank + if (comparator_.compare(previousRow_, row) == 0) { + numPeers_ += 1; + } else { + nextRank_ += numPeers_; + numPeers_ = 1; + } + } + } + rankValues->set(outputIndex, nextRank_); + previousRow_ = row; } } + currentPartitionRowOffset_ += numRows; } RowVectorPtr TopNRowNumber::getOutput() { @@ -698,13 +702,14 @@ RowVectorPtr TopNRowNumber::getOutputFromMemory() { } } + const auto numPartitionRowsLeft = + currentPartitionRows_.size() - currentPartitionRowOffset_; const auto numOutputRowsLeft = outputBatchSize_ - offset; - if (outputPartition_->rows.size() > numOutputRowsLeft) { + if (numPartitionRowsLeft > numOutputRowsLeft) { // Output as many rows as possible. RANK_FUNCTION_DISPATCH( appendPartitionRows, rankFunction_, - *outputPartition_, numOutputRowsLeft, offset, rowNumbers); @@ -712,16 +717,14 @@ RowVectorPtr TopNRowNumber::getOutputFromMemory() { break; } - // Add all partition rows. - const auto numPartitionRows = outputPartition_->rows.size(); + // Add all remaining partition rows. RANK_FUNCTION_DISPATCH( appendPartitionRows, rankFunction_, - *outputPartition_, - numPartitionRows, + numPartitionRowsLeft, offset, rowNumbers); - offset += numPartitionRows; + offset += numPartitionRowsLeft; outputPartition_ = nullptr; } diff --git a/velox/exec/TopNRowNumber.h b/velox/exec/TopNRowNumber.h index 12daa35..300bea6 100644 --- a/velox/exec/TopNRowNumber.h +++ b/velox/exec/TopNRowNumber.h @@ -61,17 +61,12 @@ class TopNRowNumberSpiller; /// reclamation was triggered during processing. /// /// If the rows are in memory, then the operator iterates over each partition -/// in the HashTable, and starts outputting rows from the partition. The -/// TopRows structure maintains the rows in descending order of their ranks -/// (greatest rank at the top of the priority queue). So when outputting, -/// the operator first fixes the top rank of the partition using fixTopRank() -/// and then computes the ranks of each row using computeNextRankInMemory(). -/// The logic of the next rank differs based on the ranking function. +/// in the HashTable, and outputs rows in ascending order of their ranks +/// (lowest sorting key / rank 1 first, matching the Window operator). /// /// If the rows are in the spill, then the spiller iterates over each spilled -/// partition in order of the ranks. For each row from the spill, the next -/// rank is computed using computeNextRankInSpill() function. The logic of -/// the next rank differs based on the ranking function. +/// partition in ascending order of the ranks. For each row from the spill, the +/// next rank is computed using computeNextRankInSpill() function. /// Note : The spill could have > limit rows for a partition as each spill /// resets the TopRows for the partition. So stop outputting rows after /// reaching the limit for each partition. @@ -215,18 +210,10 @@ class TopNRowNumber : public Operator { // output the partition. vector_size_t fixTopRank(TopRows& partition); - // Computes the rank for the next row to be output - // (all output rows in memory). - template - void computeNextRankInMemory(TopRows& partition, vector_size_t rowIndex); - - // Appends numRows of the current partition to the output. Note: The rows are - // popped in reverse order of the rank. - // NOTE: This function erases the yielded output rows from the partition - // and the next call starts with the remaining rows. + // Appends numRows of the current partition to the output in ascending + // order of ranks. template void appendPartitionRows( - TopRows& partition, vector_size_t numRows, vector_size_t outputOffset, FlatVector* rowNumbers); @@ -384,6 +371,14 @@ class TopNRowNumber : public Operator { // partition is output across multiple output blocks. TopNRowNumber::TopRows* outputPartition_{nullptr}; + // Rows of current partition ordered in ascending order of ranks. + std::vector currentPartitionRows_; + // Offset of the next row to output from currentPartitionRows_. + vector_size_t currentPartitionRowOffset_{0}; + // Previous row output in current partition (used for rank/dense_rank peer + // comparison). + char* previousRow_{nullptr}; + // The below variables are used when outputting from the spiller. // Spiller for contents of the 'data_'. std::unique_ptr spiller_; diff --git a/velox/exec/tests/TopNRowNumberTest.cpp b/velox/exec/tests/TopNRowNumberTest.cpp index 7f79ab8..f5c11ba 100644 --- a/velox/exec/tests/TopNRowNumberTest.cpp +++ b/velox/exec/tests/TopNRowNumberTest.cpp @@ -20,6 +20,7 @@ #include "velox/exec/tests/utils/AssertQueryBuilder.h" #include "velox/exec/tests/utils/OperatorTestBase.h" #include "velox/exec/tests/utils/PlanBuilder.h" +#include "velox/functions/prestosql/window/WindowFunctionsRegistration.h" using namespace facebook::velox::exec::test; @@ -36,6 +37,8 @@ class TopNRowNumberTest : public OperatorTestBase { void SetUp() override { exec::test::OperatorTestBase::SetUp(); filesystems::registerLocalFileSystem(); + velox::window::prestosql::registerAllWindowFunctions(); + common::testutil::TestValue::enable(); } const std::string functionName_; @@ -701,6 +704,80 @@ DEBUG_ONLY_TEST_P(MultiTopNRowNumberTest, oomInGroupProbe) { AssertQueryBuilder(plan).copyResults(pool_.get()), errorMessage); } +TEST_P(MultiTopNRowNumberTest, windowVsTopNRowNumber) { + auto data = makeRowVector({ + makeFlatVector({5, 3, 1, 4, 2}), + }); + + auto topNPlan = PlanBuilder() + .values({data}) + .topNRank(functionName_, {}, {"c0"}, 5, true) + .planNode(); + auto windowPlan = + PlanBuilder() + .values({data}) + .window({fmt::format("{}() over (order by c0) as rn", functionName_)}) + .planNode(); + + auto topNResult = + AssertQueryBuilder(topNPlan).maxDrivers(1).copyResults(pool_.get()); + auto windowResult = + AssertQueryBuilder(windowPlan).maxDrivers(1).copyResults(pool_.get()); + + assertEqualResults({windowResult}, {topNResult}); +} + +DEBUG_ONLY_TEST_P(MultiTopNRowNumberTest, inMemoryVsSpilled) { + const vector_size_t size = 10'000; + auto data = split( + makeRowVector( + {"s", "p"}, + { + makeFlatVector( + size, [](auto row) { return (size - row) * 10; }), + makeFlatVector( + size, [](auto row) { return row % 5000; }), + }), + 10); + + core::PlanNodeId topNId; + auto plan = PlanBuilder() + .values(data) + .topNRank(functionName_, {"p"}, {"s"}, 1000, true) + .capturePlanNodeId(topNId) + .planNode(); + + auto inMemoryResult = + AssertQueryBuilder(plan).maxDrivers(1).copyResults(pool_.get()); + + auto spillDirectory = TempDirectoryPath::create(); + std::atomic_int inputCount{0}; + SCOPED_TESTVALUE_SET( + "facebook::velox::exec::Driver::runInternal::addInput", + std::function(([&](exec::Operator* op) { + if (op->operatorCtx()->operatorType() != "TopNRowNumber") { + return; + } + if (++inputCount == 3) { + testingRunArbitration(op->pool()); + } + }))); + + std::shared_ptr task; + auto spilledResult = + AssertQueryBuilder(plan) + .maxDrivers(1) + .spillDirectory(spillDirectory->getPath()) + .config(core::QueryConfig::kSpillEnabled, "true") + .config(core::QueryConfig::kTopNRowNumberSpillEnabled, "true") + .copyResults(pool_.get(), task); + + auto planStats = exec::toPlanStats(task->taskStats()); + ASSERT_GT(planStats.at(topNId).spilledRows, 0); + + assertEqualResults({inMemoryResult}, {spilledResult}); +} + VELOX_INSTANTIATE_TEST_SUITE_P( TopNRowNumberTest, MultiTopNRowNumberTest, -- 2.55.0.windows.3