From 1fe33eb166b46eb9b77721e273b93dc351f892a7 Mon Sep 17 00:00:00 2001
From: zhli1142015 <zhli@microsoft.com>
Date: Mon, 17 Nov 2025 16:00:31 +0800
Subject: [PATCH 1/2] [VL] Fix overflow of pageNumber in VeloxSortShuffleWriter

---
 cpp/velox/shuffle/VeloxSortShuffleWriter.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cpp/velox/shuffle/VeloxSortShuffleWriter.cc b/cpp/velox/shuffle/VeloxSortShuffleWriter.cc
index 3bcdacf3e9e..57166e7ab7e 100644
--- a/cpp/velox/shuffle/VeloxSortShuffleWriter.cc
+++ b/cpp/velox/shuffle/VeloxSortShuffleWriter.cc
@@ -31,6 +31,7 @@ constexpr uint32_t kMaskLower27Bits = (1 << 27) - 1;
 constexpr uint64_t kMaskLower40Bits = (1UL << 40) - 1;
 constexpr uint32_t kPartitionIdStartByteIndex = 5;
 constexpr uint32_t kPartitionIdEndByteIndex = 7;
+constexpr uint32_t kMaxPageNumber = (1 << 13) - 1; // 13-bit max = 8191
 
 uint64_t toCompactRowId(uint32_t partitionId, uint32_t pageNumber, uint32_t offsetInPage) {
   // |63 partitionId(24) |39 inputIndex(13) |26 rowIndex(27) |
@@ -216,7 +217,7 @@ void VeloxSortShuffleWriter::insertRows(
 }
 
 arrow::Status VeloxSortShuffleWriter::maybeSpill(uint32_t nextRows) {
-  if ((uint64_t)offset_ + nextRows > std::numeric_limits<uint32_t>::max()) {
+  if ((uint64_t)offset_ + nextRows > std::numeric_limits<uint32_t>::max() || pageNumber_ >= kMaxPageNumber) {
     RETURN_NOT_OK(evictAllPartitions());
   }
   return arrow::Status::OK();

From 40a9147c7982d0faf73cf53ab6dfd4ee34e4525a Mon Sep 17 00:00:00 2001
From: zhli1142015 <zhli@microsoft.com>
Date: Wed, 19 Nov 2025 18:56:25 +0800
Subject: [PATCH 2/2] address comments

---
 cpp/velox/shuffle/VeloxSortShuffleWriter.h | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/cpp/velox/shuffle/VeloxSortShuffleWriter.h b/cpp/velox/shuffle/VeloxSortShuffleWriter.h
index 6fc08434bda..9ab842718b1 100644
--- a/cpp/velox/shuffle/VeloxSortShuffleWriter.h
+++ b/cpp/velox/shuffle/VeloxSortShuffleWriter.h
@@ -106,7 +106,9 @@ class VeloxSortShuffleWriter final : public VeloxShuffleWriter {
   std::list<facebook::velox::BufferPtr> pages_;
   std::vector<char*> pageAddresses_;
   char* currentPage_;
+  // 13-bit: max 8192 pages
   uint32_t pageNumber_;
+  // 27-bit: max 128MB page size
   uint32_t pageCursor_;
   // For debug.
   uint32_t currenPageSize_;
@@ -116,7 +118,7 @@ class VeloxSortShuffleWriter final : public VeloxShuffleWriter {
 
   // Row ID -> Partition ID
   // subscript: The index of row in the current input RowVector
-  // value: Partition ID
+  // value: Partition ID (24-bit: max 16M partitions)
   // Updated for each input RowVector.
   std::vector<uint32_t> row2Partition_;