| 278 | } |
| 279 | |
| 280 | RowVectorPtr RowNumber::getOutput() { |
| 281 | if (input_ == nullptr) { |
| 282 | return nullptr; |
| 283 | } |
| 284 | |
| 285 | if (!table_) { |
| 286 | // No partition keys. |
| 287 | return getOutputForSinglePartition(); |
| 288 | } |
| 289 | |
| 290 | const auto numInput = input_->size(); |
| 291 | |
| 292 | BufferPtr mapping; |
| 293 | vector_size_t* rawMapping; |
| 294 | vector_size_t index = 0; |
| 295 | if (limit_) { |
| 296 | mapping = allocateIndices(numInput, pool()); |
| 297 | rawMapping = mapping->asMutable<vector_size_t>(); |
| 298 | } |
| 299 | |
| 300 | // Compute row numbers if needed. |
| 301 | FlatVector<int64_t>* rowNumbers = nullptr; |
| 302 | if (generateRowNumber_) { |
| 303 | rowNumbers = &getOrCreateRowNumberVector(numInput); |
| 304 | } |
| 305 | |
| 306 | for (auto i = 0; i < numInput; ++i) { |
| 307 | auto* partition = lookup_->hits[i]; |
| 308 | const auto rowNumber = numRows(partition) + 1; |
| 309 | |
| 310 | if (limit_) { |
| 311 | if (rowNumber > limit_) { |
| 312 | // Exceeded the limit for this partition. Drop rows. |
| 313 | continue; |
| 314 | } |
| 315 | rawMapping[index++] = i; |
| 316 | } |
| 317 | |
| 318 | if (generateRowNumber_) { |
| 319 | rowNumbers->set(i, rowNumber); |
| 320 | } |
| 321 | setNumRows(partition, rowNumber); |
| 322 | } |
| 323 | |
| 324 | RowVectorPtr output; |
| 325 | if (limit_) { |
| 326 | if (index == 0) { |
| 327 | // Drop all rows. |
| 328 | output = nullptr; |
| 329 | } else { |
| 330 | output = fillOutput(index, mapping); |
| 331 | } |
| 332 | } else { |
| 333 | output = fillOutput(numInput, nullptr); |
| 334 | } |
| 335 | |
| 336 | if (spillInputReader_ != nullptr) { |
| 337 | if (spillInputReader_->nextBatch(input_)) { |