// 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.

#pragma once

#include <stddef.h>

#include <stdint.h>

#include <limits>

#include <memory>

#include <string>

#include <tuple>

#include <unordered_map>

#include <utility>

#include <vector>

#include "common/status.h"

#include "core/column/column.h"

#include "exprs/vexpr_fwd.h"

#include "format/parquet/parquet_common.h"

#include "format/parquet/vparquet_column_reader.h"

#include "format/table/table_format_reader.h"

#include "io/fs/file_reader_writer_fwd.h"

#include "storage/id_manager.h"

#include "storage/utils.h"

namespace cctz {

class time_zone;

} // namespace cctz

namespace doris {

class ObjectPool;

class RowDescriptor;

class RuntimeState;

class SlotDescriptor;

class TupleDescriptor;

namespace io {

struct IOContext;

} // namespace io

class Block;

class FieldDescriptor;

} // namespace doris

namespace tparquet {

class ColumnMetaData;

class OffsetIndex;

class RowGroup;

} // namespace tparquet

namespace doris::segment_v2 {

class RowIdColumnIteratorV2;

}

namespace doris {

#include "common/compile_check_begin.h"

// TODO: we need to determine it by test.

class RowGroupReader : public ProfileCollector {

public:

    struct IcebergRowIdParams {

        bool enabled = false;

        std::string file_path;

        int32_t partition_spec_id = 0;

        std::string partition_data_json;

        int row_id_column_pos = -1;

    };

    std::shared_ptr<TableSchemaChangeHelper::Node> _table_info_node_ptr;

    static const std::vector<int64_t> NO_DELETE;

    struct RowGroupIndex {

        int32_t row_group_id;

        int64_t first_row;

        int64_t last_row;

        RowGroupIndex(int32_t id, int64_t first, int64_t last)

                : row_group_id(id), first_row(first), last_row(last) {}

    };

    // table name

    struct LazyReadContext {

        // all conjuncts: in sql, join runtime filter, topn runtime filter.

        VExprContextSPtrs conjuncts;

        // ParquetReader::set_fill_columns(xxx, xxx) will set these two members

        std::unordered_map<std::string, std::tuple<std::string, const SlotDescriptor*>>

                fill_partition_columns;

        std::unordered_map<std::string, VExprContextSPtr> fill_missing_columns;

        phmap::flat_hash_map<int, std::vector<std::shared_ptr<ColumnPredicate>>>

                slot_id_to_predicates;

        bool can_lazy_read = false;

        // block->rows() returns the number of rows of the first column,

        // so we should check and resize the first column

        bool resize_first_column = true;

        std::vector<std::string> all_read_columns;

        // include predicate_partition_columns & predicate_missing_columns

        std::vector<uint32_t> all_predicate_col_ids;

        // save slot_id to find dict filter column name, because expr column name may

        // be different with parquet column name

        // std::pair<std::vector<col_name>, std::vector<slot_id>>

        std::pair<std::vector<std::string>, std::vector<int>> predicate_columns;

        std::vector<std::string> lazy_read_columns;

        std::unordered_map<std::string, std::tuple<std::string, const SlotDescriptor*>>

                predicate_partition_columns;

        // lazy read partition columns or all partition columns

        std::unordered_map<std::string, std::tuple<std::string, const SlotDescriptor*>>

                partition_columns;

        std::unordered_map<std::string, VExprContextSPtr> predicate_missing_columns;

        VExprContextSPtrs missing_columns_conjuncts;

        // lazy read missing columns or all missing columns

        std::unordered_map<std::string, VExprContextSPtr> missing_columns;

        // should turn off filtering by page index, lazy read and dict filter if having complex type

        bool has_complex_type = false;

    };

    /**

     * Support row-level delete in iceberg:

     * https://iceberg.apache.org/spec/#position-delete-files

     */

    struct PositionDeleteContext {

        // the filtered rows in current row group

        const std::vector<int64_t>& delete_rows;

        // the first row id of current row group in parquet file

        const int64_t first_row_id;

        // the number of rows in current row group

        const int64_t num_rows;

        const int64_t last_row_id;

        // current row id to read in the row group

        int64_t current_row_id;

        // start index in delete_rows

        const int64_t start_index;

        // end index in delete_rows

        const int64_t end_index;

        // current index in delete_rows

        int64_t index;

        const bool has_filter;

        PositionDeleteContext(const std::vector<int64_t>& delete_rows, const int64_t num_rows,

                              const int64_t first_row_id, const int64_t start_index,

                              const int64_t end_index)

                : delete_rows(delete_rows),

                  first_row_id(first_row_id),

                  num_rows(num_rows),

                  last_row_id(first_row_id + num_rows),

                  current_row_id(first_row_id),

                  start_index(start_index),

                  end_index(end_index),

                  index(start_index),

                  has_filter(end_index > start_index) {}

        PositionDeleteContext(const int64_t num_rows, const int64_t first_row)

                : PositionDeleteContext(NO_DELETE, num_rows, first_row, 0, 0) {}

        PositionDeleteContext(const PositionDeleteContext& filter) = default;

    };

    RowGroupReader(io::FileReaderSPtr file_reader, const std::vector<std::string>& read_columns,

                   const int32_t row_group_id, const tparquet::RowGroup& row_group,

                   const cctz::time_zone* ctz, io::IOContext* io_ctx,

                   const PositionDeleteContext& position_delete_ctx,

                   const LazyReadContext& lazy_read_ctx, RuntimeState* state,

                   const std::set<uint64_t>& column_ids,

                   const std::set<uint64_t>& filter_column_ids);

    ~RowGroupReader();

    Status init(const FieldDescriptor& schema, RowRanges& row_ranges,

                std::unordered_map<int, tparquet::OffsetIndex>& col_offsets,

                const TupleDescriptor* tuple_descriptor, const RowDescriptor* row_descriptor,

                const std::unordered_map<std::string, int>* colname_to_slot_id,

                const VExprContextSPtrs* not_single_slot_filter_conjuncts,

                const std::unordered_map<int, VExprContextSPtrs>* slot_id_to_filter_conjuncts);

    Status next_batch(Block* block, size_t batch_size, size_t* read_rows, bool* batch_eof);

    int64_t lazy_read_filtered_rows() const { return _lazy_read_filtered_rows; }

    int64_t predicate_filter_time() const { return _predicate_filter_time; }

    int64_t dict_filter_rewrite_time() const { return _dict_filter_rewrite_time; }

    int64_t condition_cache_filtered_rows() const { return _condition_cache_filtered_rows; }

    ParquetColumnReader::ColumnStatistics merged_column_statistics();

    void set_remaining_rows(int64_t rows) { _remaining_rows = rows; }

    int64_t get_remaining_rows() { return _remaining_rows; }

    // Filters read_ranges by removing row chunks whose condition cache granules are all-false.

    // Pure algorithm, exposed as static for testability.

    static RowRanges filter_ranges_by_cache(const RowRanges& read_ranges,

                                            const std::vector<bool>& cache, int64_t first_row,

                                            int64_t base_granule = 0);

    void set_row_id_column_iterator(

            const std::pair<std::shared_ptr<segment_v2::RowIdColumnIteratorV2>, int>&

                    iterator_pair) {

        _row_id_column_iterator_pair = iterator_pair;

    }

    void set_iceberg_rowid_params(const IcebergRowIdParams& params) {

        _iceberg_rowid_params = params;

    }

    void set_current_row_group_idx(RowGroupIndex row_group_idx) {

        _current_row_group_idx = row_group_idx;

    }

    void set_col_name_to_block_idx(

            std::unordered_map<std::string, uint32_t>* col_name_to_block_idx) {

        _col_name_to_block_idx = col_name_to_block_idx;

    }

    void set_condition_cache_context(std::shared_ptr<ConditionCacheContext> ctx) {

        _condition_cache_ctx = std::move(ctx);

    }

protected:

    void _collect_profile_before_close() override {

        if (_file_reader != nullptr) {

            _file_reader->collect_profile_before_close();

        }

    }

private:

    Status _read_empty_batch(size_t batch_size, size_t* read_rows, bool* batch_eof,

                             bool* modify_row_ids);

    Status _read_column_data(Block* block, const std::vector<std::string>& columns,

                             size_t batch_size, size_t* read_rows, bool* batch_eof,

                             FilterMap& filter_map);

    Status _do_lazy_read(Block* block, size_t batch_size, size_t* read_rows, bool* batch_eof);

    Status _rebuild_filter_map(FilterMap& filter_map,

                               DorisUniqueBufferPtr<uint8_t>& filter_map_data,

                               size_t pre_read_rows) const;

    Status _fill_partition_columns(

            Block* block, size_t rows,

            const std::unordered_map<std::string, std::tuple<std::string, const SlotDescriptor*>>&

                    partition_columns);

    Status _fill_missing_columns(

            Block* block, size_t rows,

            const std::unordered_map<std::string, VExprContextSPtr>& missing_columns);

    Status _build_pos_delete_filter(size_t read_rows);

    Status _filter_block(Block* block, int column_to_keep,

                         const std::vector<uint32_t>& columns_to_filter);

    Status _filter_block_internal(Block* block, const std::vector<uint32_t>& columns_to_filter,

                                  const IColumn::Filter& filter);

    bool _can_filter_by_dict(int slot_id, const tparquet::ColumnMetaData& column_metadata);

    bool is_dictionary_encoded(const tparquet::ColumnMetaData& column_metadata);

    Status _rewrite_dict_predicates();

    Status _rewrite_dict_conjuncts(std::vector<int32_t>& dict_codes, int slot_id, bool is_nullable);

    Status _convert_dict_cols_to_string_cols(Block* block);

    void _filter_read_ranges_by_condition_cache();

    void _mark_condition_cache_granules(const uint8_t* filter_data, size_t num_rows,

                                        int64_t batch_seq_start);

    Status _get_current_batch_row_id(size_t read_rows);

    Status _fill_row_id_columns(Block* block, size_t read_rows, bool is_current_row_ids);

    Status _append_iceberg_rowid_column(Block* block, size_t read_rows, bool is_current_row_ids);

    io::FileReaderSPtr _file_reader;

    std::unordered_map<std::string, std::unique_ptr<ParquetColumnReader>>

            _column_readers; // table_column_name

    const std::vector<std::string>& _read_table_columns;

    const int32_t _row_group_id;

    const tparquet::RowGroup& _row_group_meta;

    int64_t _remaining_rows;

    const cctz::time_zone* _ctz = nullptr;

    io::IOContext* _io_ctx = nullptr;

    PositionDeleteContext _position_delete_ctx;

    // merge the row ranges generated from page index and position delete.

    RowRanges _read_ranges;

    const LazyReadContext& _lazy_read_ctx;

    int64_t _lazy_read_filtered_rows = 0;

    int64_t _predicate_filter_time = 0;

    int64_t _dict_filter_rewrite_time = 0;

    int64_t _condition_cache_filtered_rows = 0;

    // If continuous batches are skipped, we can cache them to skip a whole page

    size_t _cached_filtered_rows = 0;

    std::shared_ptr<ConditionCacheContext> _condition_cache_ctx;

    std::unique_ptr<IColumn::Filter> _pos_delete_filter_ptr;

    int64_t _total_read_rows = 0;

    const TupleDescriptor* _tuple_descriptor = nullptr;

    const RowDescriptor* _row_descriptor = nullptr;

    const std::unordered_map<std::string, int>* _col_name_to_slot_id = nullptr;

    const std::unordered_map<int, VExprContextSPtrs>* _slot_id_to_filter_conjuncts = nullptr;

    VExprContextSPtrs _dict_filter_conjuncts;

    VExprContextSPtrs _filter_conjuncts;

    // std::pair<col_name, slot_id>

    std::vector<std::pair<std::string, int>> _dict_filter_cols;

    RuntimeState* _state = nullptr;

    std::shared_ptr<ObjectPool> _obj_pool;

    const std::set<uint64_t>& _column_ids;

    const std::set<uint64_t>& _filter_column_ids;

    bool _is_row_group_filtered = false;

    RowGroupIndex _current_row_group_idx {0, 0, 0};

    std::pair<std::shared_ptr<segment_v2::RowIdColumnIteratorV2>, int>

            _row_id_column_iterator_pair = {nullptr, -1};

    std::vector<rowid_t> _current_batch_row_ids;

    std::unordered_map<std::string, uint32_t>* _col_name_to_block_idx = nullptr;

    IcebergRowIdParams _iceberg_rowid_params;

};

#include "common/compile_check_end.h"

} // namespace doris