Leaf reader is for primitive arrays and primitive children of nested arrays
| 452 | |
| 453 | // Leaf reader is for primitive arrays and primitive children of nested arrays |
| 454 | class LeafReader : public ColumnReaderImpl { |
| 455 | public: |
| 456 | LeafReader(std::shared_ptr<ReaderContext> ctx, std::shared_ptr<Field> field, |
| 457 | std::unique_ptr<FileColumnIterator> input, |
| 458 | ::parquet::internal::LevelInfo leaf_info) |
| 459 | : ctx_(std::move(ctx)), |
| 460 | field_(std::move(field)), |
| 461 | input_(std::move(input)), |
| 462 | descr_(input_->descr()) { |
| 463 | const auto type_id = field_->type()->id(); |
| 464 | // If binary-like, RecordReader is able to read directly as the concrete type |
| 465 | // so as to avoid offset limitations. |
| 466 | std::shared_ptr<DataType> type_for_reading = |
| 467 | (::arrow::is_base_binary_like(type_id) || ::arrow::is_binary_view_like(type_id)) |
| 468 | ? field_->type() |
| 469 | : nullptr; |
| 470 | record_reader_ = RecordReader::Make( |
| 471 | descr_, leaf_info, ctx_->pool, |
| 472 | /*read_dictionary=*/field_->type()->id() == ::arrow::Type::DICTIONARY, |
| 473 | /*read_dense_for_nullable=*/false, /*arrow_type=*/type_for_reading); |
| 474 | NextRowGroup(); |
| 475 | } |
| 476 | |
| 477 | Status GetDefLevels(const int16_t** data, int64_t* length) final { |
| 478 | *data = record_reader_->def_levels(); |
| 479 | *length = record_reader_->levels_position(); |
| 480 | return Status::OK(); |
| 481 | } |
| 482 | |
| 483 | Status GetRepLevels(const int16_t** data, int64_t* length) final { |
| 484 | *data = record_reader_->rep_levels(); |
| 485 | *length = record_reader_->levels_position(); |
| 486 | return Status::OK(); |
| 487 | } |
| 488 | |
| 489 | bool IsOrHasRepeatedChild() const final { return false; } |
| 490 | |
| 491 | Status LoadBatch(int64_t records_to_read) final { |
| 492 | BEGIN_PARQUET_CATCH_EXCEPTIONS |
| 493 | out_ = nullptr; |
| 494 | record_reader_->Reset(); |
| 495 | // Pre-allocation gives much better performance for flat columns |
| 496 | record_reader_->Reserve(records_to_read); |
| 497 | const bool should_load_statistics = ctx_->reader_properties->should_load_statistics(); |
| 498 | int64_t num_target_row_groups = 0; |
| 499 | while (records_to_read > 0) { |
| 500 | if (!record_reader_->HasMoreData()) { |
| 501 | break; |
| 502 | } |
| 503 | int64_t records_read = record_reader_->ReadRecords(records_to_read); |
| 504 | records_to_read -= records_read; |
| 505 | if (records_read == 0) { |
| 506 | NextRowGroup(); |
| 507 | } else { |
| 508 | num_target_row_groups++; |
| 509 | // We can't mix multiple row groups when we load statistics |
| 510 | // because statistics are associated with a row group. If we |
| 511 | // want to mix multiple row groups and keep valid statistics, |