src/paimon/format/parquet/file_reader_wrapper.cpp - paimon-cpp - Git at Google

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */

 #include "paimon/format/parquet/file_reader_wrapper.h"

 #include <algorithm>
 #include <cassert>
 #include <cstddef>

 #include "arrow/io/interfaces.h"
 #include "arrow/record_batch.h"
 #include "arrow/util/range.h"
 #include "fmt/format.h"
 #include "paimon/common/utils/arrow/arrow_utils.h"
 #include "paimon/format/parquet/column_index_filter.h"
 #include "paimon/format/parquet/page_filtered_row_group_reader.h"
 #include "paimon/format/parquet/parquet_format_defs.h"
 #include "paimon/macros.h"
 #include "parquet/arrow/reader.h"
 #include "parquet/arrow/schema.h"
 #include "parquet/file_reader.h"
 #include "parquet/metadata.h"
 #include "parquet/page_index.h"

 namespace paimon::parquet {

 namespace {

 // Merge overlapping or adjacent ReadRanges into a minimal set of non-overlapping ranges.
 // PreBufferRanges requires non-overlapping ranges, so this is necessary when combining
 // ranges from multiple sources (page-level ranges, column chunk ranges, etc.).
 std::vector<::arrow::io::ReadRange> MergeOverlappingRanges(
     std::vector<::arrow::io::ReadRange> ranges) {
     if (ranges.empty()) {
         return ranges;
     }

     // Sort by offset
     std::sort(ranges.begin(), ranges.end(),
               [](const ::arrow::io::ReadRange& a, const ::arrow::io::ReadRange& b) {
                   return a.offset < b.offset;
               });

     std::vector<::arrow::io::ReadRange> merged;
     merged.push_back(ranges[0]);

     for (size_t i = 1; i < ranges.size(); ++i) {
         auto& last = merged.back();
         const auto& curr = ranges[i];
         // Check if current range overlaps or is adjacent to the last merged range
         int64_t last_end = last.offset + last.length;
         if (curr.offset <= last_end) {
             // Merge: extend the last range if current extends beyond it
             int64_t curr_end = curr.offset + curr.length;
             if (curr_end > last_end) {
                 last.length = curr_end - last.offset;
             }
         } else {
             // No overlap, add as new range
             merged.push_back(curr);
         }
     }

     return merged;
 }

 }  // namespace

 Result<std::unique_ptr<FileReaderWrapper>> FileReaderWrapper::Create(
     std::unique_ptr<::parquet::arrow::FileReader>&& file_reader, int64_t batch_size,
     std::shared_ptr<::arrow::MemoryPool> pool) {
     try {
         if (file_reader == nullptr) {
             return Status::Invalid("file reader wrapper create failed. file reader is nullptr");
         }
         std::vector<std::pair<uint64_t, uint64_t>> all_row_group_ranges;
         auto meta_data = file_reader->parquet_reader()->metadata();
         // prepare [start_row_idx, end_row_idx) for all row groups
         uint64_t start_row_idx = 0;
         for (int32_t i = 0; i < meta_data->num_row_groups(); i++) {
             uint64_t end_row_idx = start_row_idx + meta_data->RowGroup(i)->num_rows();
             all_row_group_ranges.emplace_back(start_row_idx, end_row_idx);
             start_row_idx = end_row_idx;
         }
         uint64_t num_rows = file_reader->parquet_reader()->metadata()->num_rows();
         if (start_row_idx != num_rows) {
             assert(false);
             return Status::Invalid(fmt::format(
                 "unexpected error. row group ranges not match with num rows {}", num_rows));
         }
         std::vector<int32_t> columns_indices =
             arrow::internal::Iota(file_reader->parquet_reader()->metadata()->num_columns());
         auto file_reader_wrapper = std::unique_ptr<FileReaderWrapper>(new FileReaderWrapper(
             std::move(file_reader), all_row_group_ranges, num_rows, batch_size, pool));
         std::vector<TargetRowGroup> all_target_row_groups;
         for (int32_t i = 0; i < file_reader_wrapper->GetNumberOfRowGroups(); i++) {
             all_target_row_groups.emplace_back(/*rg_index=*/i, /*is_partially_matched=*/false,
                                                /*ranges=*/RowRanges());
         }
         PAIMON_RETURN_NOT_OK(
             file_reader_wrapper->PrepareForReadingLazy(all_target_row_groups, columns_indices));
         return file_reader_wrapper;
     }
     PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::Create")
 }

 FileReaderWrapper::~FileReaderWrapper() {
     WaitForPendingPreBuffer();
 }

 Result<std::shared_ptr<arrow::Schema>> FileReaderWrapper::GetSchema() const {
     try {
         std::shared_ptr<arrow::Schema> file_schema;
         PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_->GetSchema(&file_schema));
         return file_schema;
     }
     PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::GetSchema")
 }

 Status FileReaderWrapper::Close() {
     try {
         if (batch_reader_) {
             PAIMON_RETURN_NOT_OK_FROM_ARROW(batch_reader_->Close());
         }
         return Status::OK();
     }
     PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::Close")
 }

 FileReaderWrapper::FileReaderWrapper(
     std::unique_ptr<::parquet::arrow::FileReader>&& file_reader,
     const std::vector<std::pair<uint64_t, uint64_t>>& all_row_group_ranges, uint64_t num_rows,
     int64_t batch_size, std::shared_ptr<::arrow::MemoryPool> pool)
     : file_reader_(std::move(file_reader)),
       all_row_group_ranges_(all_row_group_ranges),
       pool_(std::move(pool)),
       batch_size_(batch_size),
       num_rows_(num_rows) {}

 void FileReaderWrapper::WaitForPendingPreBuffer() {
     if (!prebuffered_ranges_.empty() && file_reader_) {
         // Wait for all outstanding PreBuffer async reads to complete before destruction.
         // Without this, JindoSDK async pread callbacks may fire after the underlying
         // buffers and memory pool are freed, causing use-after-free crashes.
         auto status =
             file_reader_->parquet_reader()->WhenBufferedRanges(prebuffered_ranges_).status();
         (void)status;  // Best-effort; ignore errors during cleanup
         prebuffered_ranges_.clear();
     }
 }

 void FileReaderWrapper::AdvanceToNextRowGroup() {
     current_row_group_idx_++;
     // Skip row groups excluded by read range.
     while (current_row_group_idx_ < target_row_groups_.size() &&
            target_row_groups_[current_row_group_idx_].IsExcludedByReadRange()) {
         current_row_group_idx_++;
     }
     if (current_row_group_idx_ >= target_row_groups_.size()) {
         next_row_to_read_ = num_rows_;
     } else {
         next_row_to_read_ =
             all_row_group_ranges_[target_row_groups_[current_row_group_idx_].GetRowGroupIndex()]
                 .first;
     }
 }

 Status FileReaderWrapper::SeekToRow(uint64_t row_number) {
     try {
         current_page_filtered_reader_.reset();
         filtered_global_offset_ = 0;

         for (uint64_t i = 0; i < target_row_groups_.size(); i++) {
             if (target_row_groups_[i].IsExcludedByReadRange()) {
                 continue;
             }
             int32_t rg_id = target_row_groups_[i].GetRowGroupIndex();
             uint64_t rg_start = all_row_group_ranges_[rg_id].first;
             uint64_t rg_end = all_row_group_ranges_[rg_id].second;
             if (row_number > rg_start && row_number < rg_end) {
                 return Status::Invalid(
                     fmt::format("seek to row failed. row number {} should not be in the middle of "
                                 "readable range",
                                 row_number));
             }
             if (rg_start >= row_number) {
                 current_row_group_idx_ = i;
                 next_row_to_read_ = rg_start;

                 // Rebuild batch_reader_ for non-page-filtered RGs at/after seek position.
                 std::vector<int32_t> fully_matched_indices;
                 for (uint64_t j = i; j < target_row_groups_.size(); j++) {
                     if (!target_row_groups_[j].IsExcludedByReadRange() &&
                         !target_row_groups_[j].IsPartiallyMatched()) {
                         fully_matched_indices.push_back(target_row_groups_[j].GetRowGroupIndex());
                     }
                 }
                 if (!fully_matched_indices.empty()) {
                     PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_->GetRecordBatchReader(
                         fully_matched_indices, target_column_indices_, &batch_reader_));
                 } else {
                     batch_reader_.reset();
                 }
                 return Status::OK();
             }
         }
         next_row_to_read_ = num_rows_;
         current_row_group_idx_ = target_row_groups_.size();
         return Status::OK();
     }
     PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::SeekToRow")
 }

 Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::NextPageFiltered() {
     int32_t rg_id = target_row_groups_[current_row_group_idx_].GetRowGroupIndex();

     // Construct the per-RG streaming reader on demand.
     if (!current_page_filtered_reader_) {
         const auto& target_rg = target_row_groups_[current_row_group_idx_];
         auto row_group_page_index_reader = GetRowGroupPageIndexReader(rg_id);
         auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges(
             target_rg, target_column_indices_, row_group_page_index_reader,
             file_reader_->parquet_reader());
         bool pre_buffered = !prebuffered_ranges_.empty();
         int64_t max_chunksize = batch_size_ > 0 ? batch_size_ : std::numeric_limits<int64_t>::max();
         PAIMON_ASSIGN_OR_RAISE(
             current_page_filtered_reader_,
             PageFilteredRowGroupReader::ReadFilteredRowGroup(
                 target_rg, target_column_indices_, file_reader_->properties().cache_options(),
                 pre_buffered, page_ranges, max_chunksize, row_group_page_index_reader, pool_,
                 file_reader_.get()));
         current_filtered_row_ranges_ = target_rg.GetRowRanges();
         current_filtered_rg_start_ = all_row_group_ranges_[rg_id].first;
         filtered_global_offset_ = 0;
     }

     std::shared_ptr<arrow::RecordBatch> record_batch;
     PAIMON_RETURN_NOT_OK_FROM_ARROW(current_page_filtered_reader_->ReadNext(&record_batch));

     if (record_batch) {
         auto original_row =
             current_filtered_row_ranges_.MapFilteredIndexToOriginalRow(filtered_global_offset_);
         previous_first_row_ = original_row.has_value() ? current_filtered_rg_start_ +
                                                              static_cast<uint64_t>(*original_row)
                                                        : current_filtered_rg_start_;
         filtered_global_offset_ += record_batch->num_rows();
         // Advance to the next row that survives filtering, or to the row group end when the
         // filtered ranges are exhausted, so that next_row_to_read_ tracks the streaming position.
         auto next_original_row =
             current_filtered_row_ranges_.MapFilteredIndexToOriginalRow(filtered_global_offset_);
         next_row_to_read_ =
             next_original_row.has_value()
                 ? current_filtered_rg_start_ + static_cast<uint64_t>(*next_original_row)
                 : all_row_group_ranges_[rg_id].second;
         return record_batch;
     }

     // RG exhausted — reset and advance.
     current_page_filtered_reader_.reset();
     filtered_global_offset_ = 0;
     AdvanceToNextRowGroup();
     return std::shared_ptr<arrow::RecordBatch>();
 }

 Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::NextFullyMatched() {
     if (!batch_reader_) {
         return std::shared_ptr<arrow::RecordBatch>();
     }

     PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::RecordBatch> record_batch,
                                       batch_reader_->Next());
     if (!record_batch) {
         return std::shared_ptr<arrow::RecordBatch>();
     }
     // Large binary columns (exceed 2GB) may split at different row boundaries. TableBatchReader
     // aligns their chunks by slicing columns, which may leave non-zero child offsets.
     PAIMON_ASSIGN_OR_RAISE(record_batch,
                            ArrowUtils::NormalizeRecordBatchOffsets(record_batch, pool_.get()));

     int32_t rg_id = target_row_groups_[current_row_group_idx_].GetRowGroupIndex();
     uint64_t rg_end = all_row_group_ranges_[rg_id].second;
     int64_t num_rows = record_batch->num_rows();

     previous_first_row_ = next_row_to_read_;
     if (next_row_to_read_ + num_rows < rg_end) {
         next_row_to_read_ += num_rows;
     } else if (next_row_to_read_ + num_rows == rg_end) {
         AdvanceToNextRowGroup();
     } else {
         return Status::Invalid(
             fmt::format("Next failed. next_row_to_read {} + num_rows {} exceeds row group end {}",
                         next_row_to_read_, num_rows, rg_end));
     }
     return record_batch;
 }

 std::shared_ptr<::parquet::RowGroupPageIndexReader> FileReaderWrapper::GetRowGroupPageIndexReader(
     int32_t row_group_index) {
     auto cached = row_group_page_index_readers_.find(row_group_index);
     if (cached != row_group_page_index_readers_.end()) {
         return cached->second;
     }

     std::shared_ptr<::parquet::RowGroupPageIndexReader> row_group_page_index_reader;
     auto page_index_reader = GetPageIndexReader();
     if (page_index_reader) {
         row_group_page_index_reader = page_index_reader->RowGroup(row_group_index);
     }

     // To avoid OOM, limit the number of row group page index readers cached in memory.
     constexpr int32_t kMaxRowGroupPageIndexReaders = 1024;
     if (row_group_page_index_readers_.size() < kMaxRowGroupPageIndexReaders) {
         row_group_page_index_readers_.emplace(row_group_index, row_group_page_index_reader);
     }
     return row_group_page_index_reader;
 }

 Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::Next() {
     try {
         if (PAIMON_UNLIKELY(!reader_initialized_)) {
             PAIMON_RETURN_NOT_OK(PrepareForReading(target_row_groups_, target_column_indices_));
         }

         while (current_row_group_idx_ < target_row_groups_.size()) {
             bool is_partially_matched =
                 target_row_groups_[current_row_group_idx_].IsPartiallyMatched();
             PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::RecordBatch> batch,
                                    is_partially_matched ? NextPageFiltered() : NextFullyMatched());
             if (batch) {
                 return batch;
             } else if (!is_partially_matched) {
                 // Null from fully-matched path means batch_reader_ is globally exhausted.
                 break;
             }
             // current_row_group_idx_ has been advanced in NextPageFiltered() or NextFullyMatched(),
             // loop to try next RG.
         }

         previous_first_row_ = next_row_to_read_;
         return std::shared_ptr<arrow::RecordBatch>();
     }
     PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::Next")
 }

 Result<std::vector<std::pair<uint64_t, uint64_t>>> FileReaderWrapper::GetRowGroupRanges(
     const std::set<int32_t>& row_group_indices) const {
     std::vector<std::pair<uint64_t, uint64_t>> row_group_ranges;
     for (auto row_group_index : row_group_indices) {
         if (static_cast<size_t>(row_group_index) >= all_row_group_ranges_.size()) {
             return Status::Invalid(fmt::format("row group index {} is out of bound {}",
                                                row_group_index, all_row_group_ranges_.size()));
         }
         row_group_ranges.push_back(all_row_group_ranges_[row_group_index]);
     }
     return row_group_ranges;
 }

 Status FileReaderWrapper::PrepareForReadingLazy(
     const std::vector<TargetRowGroup>& target_row_groups,
     const std::vector<int32_t>& column_indices) {
     target_row_groups_ = target_row_groups;
     target_column_indices_ = column_indices;
     reader_initialized_ = false;
     return Status::OK();
 }

 std::vector<::arrow::io::ReadRange> FileReaderWrapper::CollectPreBufferRanges(
     const std::vector<int32_t>& column_indices) {
     std::vector<::arrow::io::ReadRange> ranges;
     auto file_metadata = file_reader_->parquet_reader()->metadata();

     for (const auto& trg : target_row_groups_) {
         if (trg.IsExcludedByReadRange()) continue;

         if (trg.IsPartiallyMatched()) {
             // Page-filtered RGs: only matching page byte ranges.
             auto row_group_page_index_reader = GetRowGroupPageIndexReader(trg.GetRowGroupIndex());
             auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges(
                 trg, column_indices, row_group_page_index_reader, file_reader_->parquet_reader());
             ranges.insert(ranges.end(), std::make_move_iterator(page_ranges.begin()),
                           std::make_move_iterator(page_ranges.end()));
         } else {
             // Fully-matched RGs: entire column chunk ranges.
             auto rg_metadata = file_metadata->RowGroup(trg.GetRowGroupIndex());
             for (int32_t col_idx : column_indices) {
                 auto col_chunk = rg_metadata->ColumnChunk(col_idx);
                 int64_t offset = col_chunk->data_page_offset();
                 if (col_chunk->has_dictionary_page() && col_chunk->dictionary_page_offset() > 0 &&
                     offset > col_chunk->dictionary_page_offset()) {
                     offset = col_chunk->dictionary_page_offset();
                 }
                 ranges.push_back({offset, col_chunk->total_compressed_size()});
             }
         }
     }
     return ranges;
 }

 void FileReaderWrapper::DispatchPreBuffer(std::vector<::arrow::io::ReadRange> ranges) {
     const auto& cache_opts = file_reader_->properties().cache_options();
     ::arrow::io::IOContext io_ctx(pool_.get());
     auto merged_ranges = MergeOverlappingRanges(std::move(ranges));
     try {
         file_reader_->parquet_reader()->PreBufferRanges(merged_ranges, io_ctx, cache_opts);
         prebuffered_ranges_ = std::move(merged_ranges);
     } catch (const std::exception&) {
         prebuffered_ranges_.clear();
     }
 }

 Status FileReaderWrapper::PrepareForReading(const std::vector<TargetRowGroup>& target_row_groups,
                                             const std::vector<int32_t>& column_indices) {
     try {
         target_row_groups_ = target_row_groups;
         target_column_indices_ = column_indices;

         // Partition into fully-matched and page-filtered row groups, skipping excluded ones.
         std::vector<int32_t> fully_matched_row_groups;
         uint64_t active_count = 0;
         for (const auto& trg : target_row_groups_) {
             if (trg.IsExcludedByReadRange()) {
                 continue;
             }
             active_count++;
             if (!trg.IsPartiallyMatched()) {
                 fully_matched_row_groups.push_back(trg.GetRowGroupIndex());
             }
         }

         bool has_partially_matched = fully_matched_row_groups.size() != active_count;

         WaitForPendingPreBuffer();

         // TODO(Yonghao Fang): Neither Paimon nor Arrow manage the size and lifecycle of prebuffered
         // caches. So when a lot of row is needed, there is possibility of OOM due to too much
         // prebuffering. Also, DispatchPreBuffer will drop previous prebuffered ranges by
         // GetRecordBatchReader, which cause IO wastes.

         // Create standard reader for fully-matched row groups.
         if (!fully_matched_row_groups.empty()) {
             PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_->GetRecordBatchReader(
                 fully_matched_row_groups, column_indices, &batch_reader_));
         } else {
             batch_reader_.reset();
         }

         // When page-filtered RGs exist, issue a single PreBuffer covering both kinds.
         // Otherwise GetRecordBatchReader already issued PreBuffer internally.
         if (has_partially_matched) {
             auto all_ranges = CollectPreBufferRanges(column_indices);
             DispatchPreBuffer(std::move(all_ranges));
         }

         // Reset read state. Find the first non-excluded row group.
         uint64_t first_active_idx = 0;
         while (first_active_idx < target_row_groups_.size() &&
                target_row_groups_[first_active_idx].IsExcludedByReadRange()) {
             first_active_idx++;
         }
         if (first_active_idx >= target_row_groups_.size()) {
             next_row_to_read_ = num_rows_;
         } else {
             next_row_to_read_ =
                 all_row_group_ranges_[target_row_groups_[first_active_idx].GetRowGroupIndex()]
                     .first;
         }
         previous_first_row_ = std::numeric_limits<uint64_t>::max();
         current_row_group_idx_ = first_active_idx;
         reader_initialized_ = true;
         return Status::OK();
     }
     PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::PrepareForReading")
 }

 Status FileReaderWrapper::ApplyReadRanges(
     const std::vector<std::pair<uint64_t, uint64_t>>& read_ranges) {
     if (read_ranges.empty()) {
         for (auto& trg : target_row_groups_) {
             trg.SetExcludedByReadRange(true);
         }
         reader_initialized_ = false;
         return Status::OK();
     }
     // Build a set of row group indices whose range matches one of the read ranges.
     std::set<int32_t> matching_rg_indices;
     for (const auto& read_range : read_ranges) {
         for (size_t i = 0; i < all_row_group_ranges_.size(); i++) {
             if (all_row_group_ranges_[i] == read_range) {
                 matching_rg_indices.insert(static_cast<int32_t>(i));
             }
         }
     }
     // Mark each target row group as excluded or not based on the matching set.
     for (auto& trg : target_row_groups_) {
         trg.SetExcludedByReadRange(matching_rg_indices.count(trg.GetRowGroupIndex()) == 0);
     }
     reader_initialized_ = false;
     return Status::OK();
 }

 std::shared_ptr<::parquet::PageIndexReader> FileReaderWrapper::GetPageIndexReader() {
     try {
         return file_reader_->parquet_reader()->GetPageIndexReader();
     } catch (...) {
         // Page index is optional; degrade gracefully if the metadata read throws.
         return nullptr;
     }
 }

 Result<RowRanges> FileReaderWrapper::CalculateFilteredRowRanges(
     int32_t row_group_index, const std::shared_ptr<Predicate>& predicate,
     const std::map<std::string, int32_t>& column_name_to_index) {
     try {
         auto meta_data = file_reader_->parquet_reader()->metadata();
         int64_t row_count = meta_data->RowGroup(row_group_index)->num_rows();

         if (!predicate) {
             return RowRanges::CreateSingle(row_count);
         }

         return ColumnIndexFilter::CalculateRowRanges(predicate,
                                                      GetRowGroupPageIndexReader(row_group_index),
                                                      column_name_to_index, row_count);
     }
     PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::CalculateFilteredRowRanges")
 }

 }  // namespace paimon::parquet
	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you under the Apache License, Version 2.0 (the
	* "License"); you may not use this file except in compliance
	* with the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/

	#include "paimon/format/parquet/file_reader_wrapper.h"

	#include <algorithm>
	#include <cassert>
	#include <cstddef>

	#include "arrow/io/interfaces.h"
	#include "arrow/record_batch.h"
	#include "arrow/util/range.h"
	#include "fmt/format.h"
	#include "paimon/common/utils/arrow/arrow_utils.h"
	#include "paimon/format/parquet/column_index_filter.h"
	#include "paimon/format/parquet/page_filtered_row_group_reader.h"
	#include "paimon/format/parquet/parquet_format_defs.h"
	#include "paimon/macros.h"
	#include "parquet/arrow/reader.h"
	#include "parquet/arrow/schema.h"
	#include "parquet/file_reader.h"
	#include "parquet/metadata.h"
	#include "parquet/page_index.h"

	namespace paimon::parquet {

	namespace {

	// Merge overlapping or adjacent ReadRanges into a minimal set of non-overlapping ranges.
	// PreBufferRanges requires non-overlapping ranges, so this is necessary when combining
	// ranges from multiple sources (page-level ranges, column chunk ranges, etc.).
	std::vector<::arrow::io::ReadRange> MergeOverlappingRanges(
	std::vector<::arrow::io::ReadRange> ranges) {
	if (ranges.empty()) {
	return ranges;
	}

	// Sort by offset
	std::sort(ranges.begin(), ranges.end(),
	[](const ::arrow::io::ReadRange& a, const ::arrow::io::ReadRange& b) {
	return a.offset < b.offset;
	});

	std::vector<::arrow::io::ReadRange> merged;
	merged.push_back(ranges[0]);

	for (size_t i = 1; i < ranges.size(); ++i) {
	auto& last = merged.back();
	const auto& curr = ranges[i];
	// Check if current range overlaps or is adjacent to the last merged range
	int64_t last_end = last.offset + last.length;
	if (curr.offset <= last_end) {
	// Merge: extend the last range if current extends beyond it
	int64_t curr_end = curr.offset + curr.length;
	if (curr_end > last_end) {
	last.length = curr_end - last.offset;
	}
	} else {
	// No overlap, add as new range
	merged.push_back(curr);
	}
	}

	return merged;
	}

	} // namespace

	Result<std::unique_ptr<FileReaderWrapper>> FileReaderWrapper::Create(
	std::unique_ptr<::parquet::arrow::FileReader>&& file_reader, int64_t batch_size,
	std::shared_ptr<::arrow::MemoryPool> pool) {
	try {
	if (file_reader == nullptr) {
	return Status::Invalid("file reader wrapper create failed. file reader is nullptr");
	}
	std::vector<std::pair<uint64_t, uint64_t>> all_row_group_ranges;
	auto meta_data = file_reader->parquet_reader()->metadata();
	// prepare [start_row_idx, end_row_idx) for all row groups
	uint64_t start_row_idx = 0;
	for (int32_t i = 0; i < meta_data->num_row_groups(); i++) {
	uint64_t end_row_idx = start_row_idx + meta_data->RowGroup(i)->num_rows();
	all_row_group_ranges.emplace_back(start_row_idx, end_row_idx);
	start_row_idx = end_row_idx;
	}
	uint64_t num_rows = file_reader->parquet_reader()->metadata()->num_rows();
	if (start_row_idx != num_rows) {
	assert(false);
	return Status::Invalid(fmt::format(
	"unexpected error. row group ranges not match with num rows {}", num_rows));
	}
	std::vector<int32_t> columns_indices =
	arrow::internal::Iota(file_reader->parquet_reader()->metadata()->num_columns());
	auto file_reader_wrapper = std::unique_ptr<FileReaderWrapper>(new FileReaderWrapper(
	std::move(file_reader), all_row_group_ranges, num_rows, batch_size, pool));
	std::vector<TargetRowGroup> all_target_row_groups;
	for (int32_t i = 0; i < file_reader_wrapper->GetNumberOfRowGroups(); i++) {
	all_target_row_groups.emplace_back(/rg_index=/i, /is_partially_matched=/false,
	/ranges=/RowRanges());
	}
	PAIMON_RETURN_NOT_OK(
	file_reader_wrapper->PrepareForReadingLazy(all_target_row_groups, columns_indices));
	return file_reader_wrapper;
	}
	PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::Create")
	}

	FileReaderWrapper::~FileReaderWrapper() {
	WaitForPendingPreBuffer();
	}

	Result<std::shared_ptr<arrow::Schema>> FileReaderWrapper::GetSchema() const {
	try {
	std::shared_ptr<arrow::Schema> file_schema;
	PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_->GetSchema(&file_schema));
	return file_schema;
	}
	PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::GetSchema")
	}

	Status FileReaderWrapper::Close() {
	try {
	if (batch_reader_) {
	PAIMON_RETURN_NOT_OK_FROM_ARROW(batch_reader_->Close());
	}
	return Status::OK();
	}
	PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::Close")
	}

	FileReaderWrapper::FileReaderWrapper(
	std::unique_ptr<::parquet::arrow::FileReader>&& file_reader,
	const std::vector<std::pair<uint64_t, uint64_t>>& all_row_group_ranges, uint64_t num_rows,
	int64_t batch_size, std::shared_ptr<::arrow::MemoryPool> pool)
	: file_reader_(std::move(file_reader)),
	all_row_group_ranges_(all_row_group_ranges),
	pool_(std::move(pool)),
	batch_size_(batch_size),
	num_rows_(num_rows) {}

	void FileReaderWrapper::WaitForPendingPreBuffer() {
	if (!prebuffered_ranges_.empty() && file_reader_) {
	// Wait for all outstanding PreBuffer async reads to complete before destruction.
	// Without this, JindoSDK async pread callbacks may fire after the underlying
	// buffers and memory pool are freed, causing use-after-free crashes.
	auto status =
	file_reader_->parquet_reader()->WhenBufferedRanges(prebuffered_ranges_).status();
	(void)status; // Best-effort; ignore errors during cleanup
	prebuffered_ranges_.clear();
	}
	}

	void FileReaderWrapper::AdvanceToNextRowGroup() {
	current_row_group_idx_++;
	// Skip row groups excluded by read range.
	while (current_row_group_idx_ < target_row_groups_.size() &&
	target_row_groups_[current_row_group_idx_].IsExcludedByReadRange()) {
	current_row_group_idx_++;
	}
	if (current_row_group_idx_ >= target_row_groups_.size()) {
	next_row_to_read_ = num_rows_;
	} else {
	next_row_to_read_ =
	all_row_group_ranges_[target_row_groups_[current_row_group_idx_].GetRowGroupIndex()]
	.first;
	}
	}

	Status FileReaderWrapper::SeekToRow(uint64_t row_number) {
	try {
	current_page_filtered_reader_.reset();
	filtered_global_offset_ = 0;

	for (uint64_t i = 0; i < target_row_groups_.size(); i++) {
	if (target_row_groups_[i].IsExcludedByReadRange()) {
	continue;
	}
	int32_t rg_id = target_row_groups_[i].GetRowGroupIndex();
	uint64_t rg_start = all_row_group_ranges_[rg_id].first;
	uint64_t rg_end = all_row_group_ranges_[rg_id].second;
	if (row_number > rg_start && row_number < rg_end) {
	return Status::Invalid(
	fmt::format("seek to row failed. row number {} should not be in the middle of "
	"readable range",
	row_number));
	}
	if (rg_start >= row_number) {
	current_row_group_idx_ = i;
	next_row_to_read_ = rg_start;

	// Rebuild batch_reader_ for non-page-filtered RGs at/after seek position.
	std::vector<int32_t> fully_matched_indices;
	for (uint64_t j = i; j < target_row_groups_.size(); j++) {
	if (!target_row_groups_[j].IsExcludedByReadRange() &&
	!target_row_groups_[j].IsPartiallyMatched()) {
	fully_matched_indices.push_back(target_row_groups_[j].GetRowGroupIndex());
	}
	}
	if (!fully_matched_indices.empty()) {
	PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_->GetRecordBatchReader(
	fully_matched_indices, target_column_indices_, &batch_reader_));
	} else {
	batch_reader_.reset();
	}
	return Status::OK();
	}
	}
	next_row_to_read_ = num_rows_;
	current_row_group_idx_ = target_row_groups_.size();
	return Status::OK();
	}
	PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::SeekToRow")
	}

	Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::NextPageFiltered() {
	int32_t rg_id = target_row_groups_[current_row_group_idx_].GetRowGroupIndex();

	// Construct the per-RG streaming reader on demand.
	if (!current_page_filtered_reader_) {
	const auto& target_rg = target_row_groups_[current_row_group_idx_];
	auto row_group_page_index_reader = GetRowGroupPageIndexReader(rg_id);
	auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges(
	target_rg, target_column_indices_, row_group_page_index_reader,
	file_reader_->parquet_reader());
	bool pre_buffered = !prebuffered_ranges_.empty();
	int64_t max_chunksize = batch_size_ > 0 ? batch_size_ : std::numeric_limits<int64_t>::max();
	PAIMON_ASSIGN_OR_RAISE(
	current_page_filtered_reader_,
	PageFilteredRowGroupReader::ReadFilteredRowGroup(
	target_rg, target_column_indices_, file_reader_->properties().cache_options(),
	pre_buffered, page_ranges, max_chunksize, row_group_page_index_reader, pool_,
	file_reader_.get()));
	current_filtered_row_ranges_ = target_rg.GetRowRanges();
	current_filtered_rg_start_ = all_row_group_ranges_[rg_id].first;
	filtered_global_offset_ = 0;
	}

	std::shared_ptr<arrow::RecordBatch> record_batch;
	PAIMON_RETURN_NOT_OK_FROM_ARROW(current_page_filtered_reader_->ReadNext(&record_batch));

	if (record_batch) {
	auto original_row =
	current_filtered_row_ranges_.MapFilteredIndexToOriginalRow(filtered_global_offset_);
	previous_first_row_ = original_row.has_value() ? current_filtered_rg_start_ +
	static_cast<uint64_t>(*original_row)
	: current_filtered_rg_start_;
	filtered_global_offset_ += record_batch->num_rows();
	// Advance to the next row that survives filtering, or to the row group end when the
	// filtered ranges are exhausted, so that next_row_to_read_ tracks the streaming position.
	auto next_original_row =
	current_filtered_row_ranges_.MapFilteredIndexToOriginalRow(filtered_global_offset_);
	next_row_to_read_ =
	next_original_row.has_value()
	? current_filtered_rg_start_ + static_cast<uint64_t>(*next_original_row)
	: all_row_group_ranges_[rg_id].second;
	return record_batch;
	}

	// RG exhausted — reset and advance.
	current_page_filtered_reader_.reset();
	filtered_global_offset_ = 0;
	AdvanceToNextRowGroup();
	return std::shared_ptr<arrow::RecordBatch>();
	}

	Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::NextFullyMatched() {
	if (!batch_reader_) {
	return std::shared_ptr<arrow::RecordBatch>();
	}

	PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<arrow::RecordBatch> record_batch,
	batch_reader_->Next());
	if (!record_batch) {
	return std::shared_ptr<arrow::RecordBatch>();
	}
	// Large binary columns (exceed 2GB) may split at different row boundaries. TableBatchReader
	// aligns their chunks by slicing columns, which may leave non-zero child offsets.
	PAIMON_ASSIGN_OR_RAISE(record_batch,
	ArrowUtils::NormalizeRecordBatchOffsets(record_batch, pool_.get()));

	int32_t rg_id = target_row_groups_[current_row_group_idx_].GetRowGroupIndex();
	uint64_t rg_end = all_row_group_ranges_[rg_id].second;
	int64_t num_rows = record_batch->num_rows();

	previous_first_row_ = next_row_to_read_;
	if (next_row_to_read_ + num_rows < rg_end) {
	next_row_to_read_ += num_rows;
	} else if (next_row_to_read_ + num_rows == rg_end) {
	AdvanceToNextRowGroup();
	} else {
	return Status::Invalid(
	fmt::format("Next failed. next_row_to_read {} + num_rows {} exceeds row group end {}",
	next_row_to_read_, num_rows, rg_end));
	}
	return record_batch;
	}

	std::shared_ptr<::parquet::RowGroupPageIndexReader> FileReaderWrapper::GetRowGroupPageIndexReader(
	int32_t row_group_index) {
	auto cached = row_group_page_index_readers_.find(row_group_index);
	if (cached != row_group_page_index_readers_.end()) {
	return cached->second;
	}

	std::shared_ptr<::parquet::RowGroupPageIndexReader> row_group_page_index_reader;
	auto page_index_reader = GetPageIndexReader();
	if (page_index_reader) {
	row_group_page_index_reader = page_index_reader->RowGroup(row_group_index);
	}

	// To avoid OOM, limit the number of row group page index readers cached in memory.
	constexpr int32_t kMaxRowGroupPageIndexReaders = 1024;
	if (row_group_page_index_readers_.size() < kMaxRowGroupPageIndexReaders) {
	row_group_page_index_readers_.emplace(row_group_index, row_group_page_index_reader);
	}
	return row_group_page_index_reader;
	}

	Result<std::shared_ptr<arrow::RecordBatch>> FileReaderWrapper::Next() {
	try {
	if (PAIMON_UNLIKELY(!reader_initialized_)) {
	PAIMON_RETURN_NOT_OK(PrepareForReading(target_row_groups_, target_column_indices_));
	}

	while (current_row_group_idx_ < target_row_groups_.size()) {
	bool is_partially_matched =
	target_row_groups_[current_row_group_idx_].IsPartiallyMatched();
	PAIMON_ASSIGN_OR_RAISE(std::shared_ptr<arrow::RecordBatch> batch,
	is_partially_matched ? NextPageFiltered() : NextFullyMatched());
	if (batch) {
	return batch;
	} else if (!is_partially_matched) {
	// Null from fully-matched path means batch_reader_ is globally exhausted.
	break;
	}
	// current_row_group_idx_ has been advanced in NextPageFiltered() or NextFullyMatched(),
	// loop to try next RG.
	}

	previous_first_row_ = next_row_to_read_;
	return std::shared_ptr<arrow::RecordBatch>();
	}
	PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::Next")
	}

	Result<std::vector<std::pair<uint64_t, uint64_t>>> FileReaderWrapper::GetRowGroupRanges(
	const std::set<int32_t>& row_group_indices) const {
	std::vector<std::pair<uint64_t, uint64_t>> row_group_ranges;
	for (auto row_group_index : row_group_indices) {
	if (static_cast<size_t>(row_group_index) >= all_row_group_ranges_.size()) {
	return Status::Invalid(fmt::format("row group index {} is out of bound {}",
	row_group_index, all_row_group_ranges_.size()));
	}
	row_group_ranges.push_back(all_row_group_ranges_[row_group_index]);
	}
	return row_group_ranges;
	}

	Status FileReaderWrapper::PrepareForReadingLazy(
	const std::vector<TargetRowGroup>& target_row_groups,
	const std::vector<int32_t>& column_indices) {
	target_row_groups_ = target_row_groups;
	target_column_indices_ = column_indices;
	reader_initialized_ = false;
	return Status::OK();
	}

	std::vector<::arrow::io::ReadRange> FileReaderWrapper::CollectPreBufferRanges(
	const std::vector<int32_t>& column_indices) {
	std::vector<::arrow::io::ReadRange> ranges;
	auto file_metadata = file_reader_->parquet_reader()->metadata();

	for (const auto& trg : target_row_groups_) {
	if (trg.IsExcludedByReadRange()) continue;

	if (trg.IsPartiallyMatched()) {
	// Page-filtered RGs: only matching page byte ranges.
	auto row_group_page_index_reader = GetRowGroupPageIndexReader(trg.GetRowGroupIndex());
	auto page_ranges = PageFilteredRowGroupReader::ComputePageRanges(
	trg, column_indices, row_group_page_index_reader, file_reader_->parquet_reader());
	ranges.insert(ranges.end(), std::make_move_iterator(page_ranges.begin()),
	std::make_move_iterator(page_ranges.end()));
	} else {
	// Fully-matched RGs: entire column chunk ranges.
	auto rg_metadata = file_metadata->RowGroup(trg.GetRowGroupIndex());
	for (int32_t col_idx : column_indices) {
	auto col_chunk = rg_metadata->ColumnChunk(col_idx);
	int64_t offset = col_chunk->data_page_offset();
	if (col_chunk->has_dictionary_page() && col_chunk->dictionary_page_offset() > 0 &&
	offset > col_chunk->dictionary_page_offset()) {
	offset = col_chunk->dictionary_page_offset();
	}
	ranges.push_back({offset, col_chunk->total_compressed_size()});
	}
	}
	}
	return ranges;
	}

	void FileReaderWrapper::DispatchPreBuffer(std::vector<::arrow::io::ReadRange> ranges) {
	const auto& cache_opts = file_reader_->properties().cache_options();
	::arrow::io::IOContext io_ctx(pool_.get());
	auto merged_ranges = MergeOverlappingRanges(std::move(ranges));
	try {
	file_reader_->parquet_reader()->PreBufferRanges(merged_ranges, io_ctx, cache_opts);
	prebuffered_ranges_ = std::move(merged_ranges);
	} catch (const std::exception&) {
	prebuffered_ranges_.clear();
	}
	}

	Status FileReaderWrapper::PrepareForReading(const std::vector<TargetRowGroup>& target_row_groups,
	const std::vector<int32_t>& column_indices) {
	try {
	target_row_groups_ = target_row_groups;
	target_column_indices_ = column_indices;

	// Partition into fully-matched and page-filtered row groups, skipping excluded ones.
	std::vector<int32_t> fully_matched_row_groups;
	uint64_t active_count = 0;
	for (const auto& trg : target_row_groups_) {
	if (trg.IsExcludedByReadRange()) {
	continue;
	}
	active_count++;
	if (!trg.IsPartiallyMatched()) {
	fully_matched_row_groups.push_back(trg.GetRowGroupIndex());
	}
	}

	bool has_partially_matched = fully_matched_row_groups.size() != active_count;

	WaitForPendingPreBuffer();

	// TODO(Yonghao Fang): Neither Paimon nor Arrow manage the size and lifecycle of prebuffered
	// caches. So when a lot of row is needed, there is possibility of OOM due to too much
	// prebuffering. Also, DispatchPreBuffer will drop previous prebuffered ranges by
	// GetRecordBatchReader, which cause IO wastes.

	// Create standard reader for fully-matched row groups.
	if (!fully_matched_row_groups.empty()) {
	PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_->GetRecordBatchReader(
	fully_matched_row_groups, column_indices, &batch_reader_));
	} else {
	batch_reader_.reset();
	}

	// When page-filtered RGs exist, issue a single PreBuffer covering both kinds.
	// Otherwise GetRecordBatchReader already issued PreBuffer internally.
	if (has_partially_matched) {
	auto all_ranges = CollectPreBufferRanges(column_indices);
	DispatchPreBuffer(std::move(all_ranges));
	}

	// Reset read state. Find the first non-excluded row group.
	uint64_t first_active_idx = 0;
	while (first_active_idx < target_row_groups_.size() &&
	target_row_groups_[first_active_idx].IsExcludedByReadRange()) {
	first_active_idx++;
	}
	if (first_active_idx >= target_row_groups_.size()) {
	next_row_to_read_ = num_rows_;
	} else {
	next_row_to_read_ =
	all_row_group_ranges_[target_row_groups_[first_active_idx].GetRowGroupIndex()]
	.first;
	}
	previous_first_row_ = std::numeric_limits<uint64_t>::max();
	current_row_group_idx_ = first_active_idx;
	reader_initialized_ = true;
	return Status::OK();
	}
	PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::PrepareForReading")
	}

	Status FileReaderWrapper::ApplyReadRanges(
	const std::vector<std::pair<uint64_t, uint64_t>>& read_ranges) {
	if (read_ranges.empty()) {
	for (auto& trg : target_row_groups_) {
	trg.SetExcludedByReadRange(true);
	}
	reader_initialized_ = false;
	return Status::OK();
	}
	// Build a set of row group indices whose range matches one of the read ranges.
	std::set<int32_t> matching_rg_indices;
	for (const auto& read_range : read_ranges) {
	for (size_t i = 0; i < all_row_group_ranges_.size(); i++) {
	if (all_row_group_ranges_[i] == read_range) {
	matching_rg_indices.insert(static_cast<int32_t>(i));
	}
	}
	}
	// Mark each target row group as excluded or not based on the matching set.
	for (auto& trg : target_row_groups_) {
	trg.SetExcludedByReadRange(matching_rg_indices.count(trg.GetRowGroupIndex()) == 0);
	}
	reader_initialized_ = false;
	return Status::OK();
	}

	std::shared_ptr<::parquet::PageIndexReader> FileReaderWrapper::GetPageIndexReader() {
	try {
	return file_reader_->parquet_reader()->GetPageIndexReader();
	} catch (...) {
	// Page index is optional; degrade gracefully if the metadata read throws.
	return nullptr;
	}
	}

	Result<RowRanges> FileReaderWrapper::CalculateFilteredRowRanges(
	int32_t row_group_index, const std::shared_ptr<Predicate>& predicate,
	const std::map<std::string, int32_t>& column_name_to_index) {
	try {
	auto meta_data = file_reader_->parquet_reader()->metadata();
	int64_t row_count = meta_data->RowGroup(row_group_index)->num_rows();

	if (!predicate) {
	return RowRanges::CreateSingle(row_count);
	}

	return ColumnIndexFilter::CalculateRowRanges(predicate,
	GetRowGroupPageIndexReader(row_group_index),
	column_name_to_index, row_count);
	}
	PAIMON_PARQUET_CATCH_AND_RETURN_STATUS("FileReaderWrapper::CalculateFilteredRowRanges")
	}

	} // namespace paimon::parquet