| 418 | } |
| 419 | |
| 420 | Status TopNNode::GetNextPartitioned( |
| 421 | RuntimeState* state, RowBatch* batch, bool* eos) { |
| 422 | DCHECK(is_partitioned()); |
| 423 | *eos = false; |
| 424 | while (!batch->AtCapacity()) { |
| 425 | RETURN_IF_CANCELLED(state); |
| 426 | RETURN_IF_ERROR(QueryMaintenance(state)); |
| 427 | if (sort_out_batch_pos_ >= sort_out_batch_->num_rows()) { |
| 428 | // Output rows will reference tuples from sorter output batches - make sure memory |
| 429 | // is transferred correctly. |
| 430 | sort_out_batch_->TransferResourceOwnership(batch); |
| 431 | sort_out_batch_->Reset(); |
| 432 | sort_out_batch_pos_ = 0; |
| 433 | if (batch->AtCapacity()) break; |
| 434 | bool sorter_eos = false; |
| 435 | RETURN_IF_ERROR(sorter_->GetNext(sort_out_batch_.get(), &sorter_eos)); |
| 436 | if (sorter_eos && sort_out_batch_->num_rows() == 0) { |
| 437 | sort_out_batch_->TransferResourceOwnership(batch); |
| 438 | *eos = true; |
| 439 | break; |
| 440 | } |
| 441 | } |
| 442 | // Copy rows within the partition limits from 'sort_out_batch_' to 'batch'. |
| 443 | // NOTE: this loop could be codegen'd, but is unlikely to be the bottleneck for |
| 444 | // most partitioned top-N queries. |
| 445 | while (sort_out_batch_pos_ < sort_out_batch_->num_rows()) { |
| 446 | TupleRow* curr_row = sort_out_batch_->GetRow(sort_out_batch_pos_); |
| 447 | ++sort_out_batch_pos_; |
| 448 | // If 'num_rows_returned_from_partition_' > 0, then 'prev_row' is the previous row |
| 449 | // returned from the current partition. |
| 450 | TupleRow* prev_row = reinterpret_cast<TupleRow*>(&tmp_tuple_); |
| 451 | bool add_row = false; |
| 452 | if (num_rows_returned_from_partition_ > 0 |
| 453 | && partition_cmp_->Compare(curr_row, prev_row) == 0) { |
| 454 | // Return rows up to the limit plus any ties that match the last returned row. |
| 455 | if (num_rows_returned_from_partition_ < per_partition_limit() |
| 456 | || (include_ties() && |
| 457 | intra_partition_order_cmp_->Compare(curr_row, prev_row) == 0)) { |
| 458 | add_row = true; |
| 459 | ++num_rows_returned_from_partition_; |
| 460 | } |
| 461 | } else { |
| 462 | // New partition. |
| 463 | DCHECK_GT(per_partition_limit(), 0); |
| 464 | COUNTER_ADD(num_partitions_counter_, 1); |
| 465 | add_row = true; |
| 466 | num_rows_returned_from_partition_ = 1; |
| 467 | } |
| 468 | if (add_row) { |
| 469 | Tuple* out_tuple = curr_row->GetTuple(0); |
| 470 | tmp_tuple_ = out_tuple; |
| 471 | TupleRow* out_row = batch->GetRow(batch->AddRow()); |
| 472 | out_row->SetTuple(0, out_tuple); |
| 473 | batch->CommitLastRow(); |
| 474 | IncrementNumRowsReturned(1); |
| 475 | if (batch->AtCapacity()) break; |
| 476 | } |
| 477 | } |
nothing calls this directly
no test coverage detected