| 504 | } |
| 505 | |
| 506 | Result<std::shared_ptr<Table>> AsyncScanner::TakeRows(const Array& indices) { |
| 507 | if (indices.null_count() != 0) { |
| 508 | return Status::NotImplemented("null take indices"); |
| 509 | } |
| 510 | |
| 511 | compute::ExecContext ctx(scan_options_->pool); |
| 512 | |
| 513 | const Array* original_indices; |
| 514 | // If we have to cast, this is the backing reference |
| 515 | std::shared_ptr<Array> original_indices_ptr; |
| 516 | if (indices.type_id() != Type::INT64) { |
| 517 | ARROW_ASSIGN_OR_RAISE( |
| 518 | original_indices_ptr, |
| 519 | compute::Cast(indices, int64(), compute::CastOptions::Safe(), &ctx)); |
| 520 | original_indices = original_indices_ptr.get(); |
| 521 | } else { |
| 522 | original_indices = &indices; |
| 523 | } |
| 524 | |
| 525 | std::shared_ptr<Array> unsort_indices; |
| 526 | { |
| 527 | ARROW_ASSIGN_OR_RAISE( |
| 528 | auto sort_indices, |
| 529 | compute::SortIndices(*original_indices, compute::SortOrder::Ascending, &ctx)); |
| 530 | ARROW_ASSIGN_OR_RAISE(original_indices_ptr, |
| 531 | compute::Take(*original_indices, *sort_indices, |
| 532 | compute::TakeOptions::Defaults(), &ctx)); |
| 533 | original_indices = original_indices_ptr.get(); |
| 534 | ARROW_ASSIGN_OR_RAISE( |
| 535 | unsort_indices, |
| 536 | compute::SortIndices(*sort_indices, compute::SortOrder::Ascending, &ctx)); |
| 537 | } |
| 538 | |
| 539 | RecordBatchVector out_batches; |
| 540 | |
| 541 | auto raw_indices = static_cast<const Int64Array&>(*original_indices).raw_values(); |
| 542 | int64_t offset = 0, row_begin = 0; |
| 543 | |
| 544 | ARROW_ASSIGN_OR_RAISE(auto batch_it, ScanBatches()); |
| 545 | while (true) { |
| 546 | ARROW_ASSIGN_OR_RAISE(auto batch, batch_it.Next()); |
| 547 | if (IsIterationEnd(batch)) break; |
| 548 | if (offset == original_indices->length()) break; |
| 549 | DCHECK_LT(offset, original_indices->length()); |
| 550 | |
| 551 | int64_t length = 0; |
| 552 | while (offset + length < original_indices->length()) { |
| 553 | auto rel_index = raw_indices[offset + length] - row_begin; |
| 554 | if (rel_index >= batch.record_batch->num_rows()) break; |
| 555 | ++length; |
| 556 | } |
| 557 | DCHECK_LE(offset + length, original_indices->length()); |
| 558 | if (length == 0) { |
| 559 | row_begin += batch.record_batch->num_rows(); |
| 560 | continue; |
| 561 | } |
| 562 | |
| 563 | Datum rel_indices = original_indices->Slice(offset, length); |