| 588 | } |
| 589 | |
| 590 | Result<std::shared_ptr<Table>> Table::CombineChunks(MemoryPool* pool) const { |
| 591 | const int ncolumns = num_columns(); |
| 592 | std::vector<std::shared_ptr<ChunkedArray>> compacted_columns(ncolumns); |
| 593 | for (int i = 0; i < ncolumns; ++i) { |
| 594 | const auto& col = column(i); |
| 595 | if (col->num_chunks() <= 1) { |
| 596 | compacted_columns[i] = col; |
| 597 | continue; |
| 598 | } |
| 599 | |
| 600 | if (is_binary_like(col->type()->id())) { |
| 601 | // ARROW-5744 Allow binary columns to be combined into multiple chunks to avoid |
| 602 | // buffer overflow |
| 603 | ArrayVector chunks; |
| 604 | int chunk_i = 0; |
| 605 | while (chunk_i < col->num_chunks()) { |
| 606 | ArrayVector safe_chunks; |
| 607 | int64_t data_length = 0; |
| 608 | for (; chunk_i < col->num_chunks(); ++chunk_i) { |
| 609 | const auto& chunk = col->chunk(chunk_i); |
| 610 | data_length += checked_cast<const BinaryArray&>(*chunk).total_values_length(); |
| 611 | if (data_length >= kBinaryMemoryLimit) { |
| 612 | break; |
| 613 | } |
| 614 | safe_chunks.push_back(chunk); |
| 615 | } |
| 616 | chunks.emplace_back(); |
| 617 | ARROW_ASSIGN_OR_RAISE(chunks.back(), Concatenate(safe_chunks, pool)); |
| 618 | } |
| 619 | compacted_columns[i] = std::make_shared<ChunkedArray>(std::move(chunks)); |
| 620 | } else { |
| 621 | ARROW_ASSIGN_OR_RAISE(auto compacted, Concatenate(col->chunks(), pool)); |
| 622 | compacted_columns[i] = std::make_shared<ChunkedArray>(compacted); |
| 623 | } |
| 624 | } |
| 625 | return Table::Make(schema(), std::move(compacted_columns), num_rows_); |
| 626 | } |
| 627 | |
| 628 | Result<std::shared_ptr<RecordBatch>> Table::CombineChunksToBatch(MemoryPool* pool) const { |
| 629 | ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Table> combined, CombineChunks(pool)); |