| 585 | |
| 586 | template <typename IndexType> |
| 587 | class ListReader : public ColumnReaderImpl { |
| 588 | public: |
| 589 | ListReader(std::shared_ptr<ReaderContext> ctx, std::shared_ptr<Field> field, |
| 590 | ::parquet::internal::LevelInfo level_info, |
| 591 | std::unique_ptr<ColumnReaderImpl> child_reader) |
| 592 | : ctx_(std::move(ctx)), |
| 593 | field_(std::move(field)), |
| 594 | level_info_(level_info), |
| 595 | item_reader_(std::move(child_reader)) {} |
| 596 | |
| 597 | Status GetDefLevels(const int16_t** data, int64_t* length) override { |
| 598 | return item_reader_->GetDefLevels(data, length); |
| 599 | } |
| 600 | |
| 601 | Status GetRepLevels(const int16_t** data, int64_t* length) override { |
| 602 | return item_reader_->GetRepLevels(data, length); |
| 603 | } |
| 604 | |
| 605 | bool IsOrHasRepeatedChild() const final { return true; } |
| 606 | |
| 607 | Status LoadBatch(int64_t number_of_records) final { |
| 608 | return item_reader_->LoadBatch(number_of_records); |
| 609 | } |
| 610 | |
| 611 | virtual ::arrow::Result<std::shared_ptr<ChunkedArray>> AssembleArray( |
| 612 | std::shared_ptr<ArrayData> data) { |
| 613 | if (field_->type()->id() == ::arrow::Type::MAP) { |
| 614 | // Error out if data is not map-compliant instead of aborting in MakeArray below |
| 615 | RETURN_NOT_OK(::arrow::MapArray::ValidateChildData(data->child_data)); |
| 616 | } |
| 617 | std::shared_ptr<Array> result = ::arrow::MakeArray(data); |
| 618 | return std::make_shared<ChunkedArray>(result); |
| 619 | } |
| 620 | |
| 621 | Status BuildArray(int64_t length_upper_bound, |
| 622 | std::shared_ptr<ChunkedArray>* out) override { |
| 623 | const int16_t* def_levels; |
| 624 | const int16_t* rep_levels; |
| 625 | int64_t num_levels; |
| 626 | RETURN_NOT_OK(item_reader_->GetDefLevels(&def_levels, &num_levels)); |
| 627 | RETURN_NOT_OK(item_reader_->GetRepLevels(&rep_levels, &num_levels)); |
| 628 | |
| 629 | std::shared_ptr<ResizableBuffer> validity_buffer; |
| 630 | ::parquet::internal::ValidityBitmapInputOutput validity_io; |
| 631 | validity_io.values_read_upper_bound = length_upper_bound; |
| 632 | if (field_->nullable()) { |
| 633 | ARROW_ASSIGN_OR_RAISE(validity_buffer, |
| 634 | AllocateResizableBuffer( |
| 635 | bit_util::BytesForBits(length_upper_bound), ctx_->pool)); |
| 636 | validity_io.valid_bits = validity_buffer->mutable_data(); |
| 637 | } |
| 638 | ARROW_ASSIGN_OR_RAISE( |
| 639 | std::shared_ptr<ResizableBuffer> offsets_buffer, |
| 640 | AllocateResizableBuffer( |
| 641 | sizeof(IndexType) * std::max(int64_t{1}, length_upper_bound + 1), |
| 642 | ctx_->pool)); |
| 643 | // Ensure zero initialization in case we have reached a zero length list (and |
| 644 | // because first entry is always zero). |