| 239 | } |
| 240 | |
| 241 | void PartitionedOutput::addInput(RowVectorPtr input) { |
| 242 | initializeInput(std::move(input)); |
| 243 | |
| 244 | initializeDestinations(); |
| 245 | |
| 246 | initializeSizeBuffers(); |
| 247 | |
| 248 | estimateRowSizes(); |
| 249 | |
| 250 | for (auto& destination : destinations_) { |
| 251 | destination->beginBatch(); |
| 252 | } |
| 253 | |
| 254 | auto numInput = input_->size(); |
| 255 | if (numDestinations_ == 1) { |
| 256 | destinations_[0]->addRows(IndexRange{0, numInput}); |
| 257 | } else { |
| 258 | auto singlePartition = partitionFunction_->partition(*input_, partitions_); |
| 259 | if (replicateNullsAndAny_) { |
| 260 | collectNullRows(); |
| 261 | |
| 262 | vector_size_t start = 0; |
| 263 | if (!replicatedAny_) { |
| 264 | for (auto& destination : destinations_) { |
| 265 | destination->addRow(0); |
| 266 | } |
| 267 | replicatedAny_ = true; |
| 268 | // Make sure not to replicate first row twice. |
| 269 | start = 1; |
| 270 | } |
| 271 | for (auto i = start; i < numInput; ++i) { |
| 272 | if (nullRows_.isValid(i)) { |
| 273 | for (auto& destination : destinations_) { |
| 274 | destination->addRow(i); |
| 275 | } |
| 276 | } else { |
| 277 | if (singlePartition.has_value()) { |
| 278 | destinations_[singlePartition.value()]->addRow(i); |
| 279 | } else { |
| 280 | destinations_[partitions_[i]]->addRow(i); |
| 281 | } |
| 282 | } |
| 283 | } |
| 284 | } else { |
| 285 | if (singlePartition.has_value()) { |
| 286 | destinations_[singlePartition.value()]->addRows( |
| 287 | IndexRange{0, numInput}); |
| 288 | } else { |
| 289 | for (vector_size_t i = 0; i < numInput; ++i) { |
| 290 | destinations_[partitions_[i]]->addRow(i); |
| 291 | } |
| 292 | } |
| 293 | } |
| 294 | } |
| 295 | } |
| 296 | |
| 297 | void PartitionedOutput::collectNullRows() { |
| 298 | auto size = input_->size(); |