| 131 | } |
| 132 | |
| 133 | uint64_t SpillState::appendToPartition( |
| 134 | uint32_t partition, |
| 135 | const RowVectorPtr& rows) { |
| 136 | BOLT_CHECK( |
| 137 | isPartitionSpilled(partition), "Partition {} is not spilled", partition); |
| 138 | |
| 139 | BOLT_TEST_ADJUST( |
| 140 | "bytedance::bolt::exec::SpillState::appendToPartition", this); |
| 141 | |
| 142 | BOLT_CHECK_NOT_NULL( |
| 143 | ioConfig_.getSpillDirPathCb, "Spill directory callback not specified."); |
| 144 | const std::string& spillDir = ioConfig_.getSpillDirPathCb(); |
| 145 | BOLT_CHECK(!spillDir.empty(), "Spill directory does not exist"); |
| 146 | // Ensure that partition exist before writing. |
| 147 | if (partitionWriters_.at(partition) == nullptr) { |
| 148 | partitionWriters_[partition] = std::make_unique<SpillWriter>( |
| 149 | std::static_pointer_cast<const RowType>(rows->type()), |
| 150 | sortingKeys_, |
| 151 | fmt::format( |
| 152 | "{}/{}-spill-{}{}", |
| 153 | spillDir, |
| 154 | ioConfig_.fileNamePrefix, |
| 155 | partition, |
| 156 | (immediateFlush_ ? "-flags" : "")), |
| 157 | targetFileSize_, |
| 158 | ioConfig_, |
| 159 | pool_, |
| 160 | stats_, |
| 161 | maxBatchRows_, |
| 162 | std::nullopt); |
| 163 | } |
| 164 | |
| 165 | updateSpilledInputBytes(rows->estimateFlatSize()); |
| 166 | spilledRowCount_[partition] += rows->size(); |
| 167 | |
| 168 | IndexRange range{0, rows->size()}; |
| 169 | if (immediateFlush_) { |
| 170 | return partitionWriters_[partition]->writeAndFlush( |
| 171 | rows, folly::Range<IndexRange*>(&range, 1)); |
| 172 | } else { |
| 173 | return partitionWriters_[partition]->write( |
| 174 | rows, folly::Range<IndexRange*>(&range, 1)); |
| 175 | } |
| 176 | } |
| 177 | |
| 178 | uint64_t SpillState::appendToPartition( |
| 179 | uint32_t partition, |