[VL] Refine velox shuffle writer iteration using partitionUsed instea…

…d of all partitions (facebookincubator#3356)
marin-ma · Oct 9, 2023 · b05f5d8 · b05f5d8
1 parent 61577a5
commit b05f5d8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/cpp/velox/shuffle/VeloxShuffleWriter.cc b/cpp/velox/shuffle/VeloxShuffleWriter.cc
@@ -915,8 +915,8 @@ arrow::Status VeloxShuffleWriter::splitFixedWidthValueBuffer(const velox::RowVec
       auto column = rv.childAt(colIdx);
       if (vectorHasNull(column)) {
         auto& dstAddrs = partitionValidityAddrs_[col];
-        for (auto pid = 0; pid < numPartitions_; ++pid) {
-          if (partition2RowCount_[pid] > 0 && dstAddrs[pid] == nullptr) {
+        for (auto& pid : partitionUsed_) {
+          if (dstAddrs[pid] == nullptr) {
             // init bitmap if it's null, initialize the buffer as true
             auto newSize = std::max(partition2RowCount_[pid], (uint32_t)options_.buffer_size);
             GLUTEN_ASSIGN_OR_THROW(