// 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 "format_v2/parquet/parquet_reader.h"

#include <algorithm>

#include <map>

#include <memory>

#include <optional>

#include <ranges>

#include <utility>

#include <vector>

#include "core/assert_cast.h"

#include "core/block/block.h"

#include "core/data_type/data_type_array.h"

#include "core/data_type/data_type_factory.hpp"

#include "core/data_type/data_type_map.h"

#include "core/data_type/data_type_nullable.h"

#include "core/data_type/data_type_struct.h"

#include "format_v2/column_mapper.h"

#include "format_v2/parquet/parquet_column_schema.h"

#include "format_v2/parquet/parquet_file_context.h"

#include "format_v2/parquet/parquet_scan.h"

#include "format_v2/parquet/parquet_statistics.h"

#include "format_v2/parquet/reader/column_reader.h"

#include "runtime/runtime_state.h"

namespace doris::format::parquet {

struct ParquetReaderScanState {

    ParquetFileContext file_context;

    std::vector<std::unique_ptr<ParquetColumnSchema>> file_schema;

    RowGroupScanPlan scan_plan;

    ParquetScanScheduler scheduler;

    const cctz::time_zone* timezone = nullptr;

    bool enable_bloom_filter = false;

};

DataTypePtr nullable_like_original(const DataTypePtr& type, DataTypePtr nested_type) {

    return type != nullptr && type->is_nullable() ? make_nullable(nested_type) : nested_type;

}

int timestamp_tz_scale(const ParquetTypeDescriptor& type_descriptor) {

    switch (type_descriptor.time_unit) {

    case ParquetTimeUnit::MILLIS:

        return 3;

    case ParquetTimeUnit::MICROS:

    case ParquetTimeUnit::UNKNOWN:

    default:

        return 6;

    }

}

bool should_map_to_timestamp_tz(const ParquetColumnSchema& column_schema) {

    const auto& type_descriptor = column_schema.type_descriptor;

    return type_descriptor.is_timestamp && type_descriptor.timestamp_is_adjusted_to_utc;

}

DataTypePtr apply_timestamp_tz_mapping(ParquetColumnSchema* column_schema) {

    DORIS_CHECK(column_schema != nullptr);

    if (column_schema->kind == ParquetColumnSchemaKind::PRIMITIVE) {

        if (should_map_to_timestamp_tz(*column_schema)) {

            const bool nullable =

                    column_schema->type != nullptr && column_schema->type->is_nullable();

            const auto scale = timestamp_tz_scale(column_schema->type_descriptor);

            column_schema->type = DataTypeFactory::instance().create_data_type(TYPE_TIMESTAMPTZ,

                                                                               nullable, 0, scale);

            column_schema->type_descriptor.doris_type = column_schema->type;

        }

        return column_schema->type;

    }

    std::vector<DataTypePtr> child_types;

    child_types.reserve(column_schema->children.size());

    for (auto& child : column_schema->children) {

        child_types.push_back(apply_timestamp_tz_mapping(child.get()));

    }

    if (column_schema->kind == ParquetColumnSchemaKind::LIST) {

        DORIS_CHECK(child_types.size() == 1);

        column_schema->type = nullable_like_original(

                column_schema->type, std::make_shared<DataTypeArray>(child_types[0]));

    } else if (column_schema->kind == ParquetColumnSchemaKind::MAP) {

        DORIS_CHECK(child_types.size() == 2);

        column_schema->type = nullable_like_original(

                column_schema->type, std::make_shared<DataTypeMap>(make_nullable(child_types[0]),

                                                                   make_nullable(child_types[1])));

    } else if (column_schema->kind == ParquetColumnSchemaKind::STRUCT) {

        Strings child_names;

        child_names.reserve(column_schema->children.size());

        for (const auto& child : column_schema->children) {

            child_names.push_back(child->name);

        }

        column_schema->type = nullable_like_original(

                column_schema->type, std::make_shared<DataTypeStruct>(child_types, child_names));

    }

    return column_schema->type;

}

static Status find_projected_minmax_leaf(const ParquetColumnSchema& column_schema,

                                         const format::LocalColumnIndex& projection,

                                         const ParquetColumnSchema** leaf_schema) {

    DORIS_CHECK(leaf_schema != nullptr);

    if (projection.project_all_children || projection.children.empty()) {

        if (column_schema.leaf_column_id < 0) {

            return Status::NotSupported(

                    "Parquet aggregate pushdown only supports primitive column {}",

                    column_schema.name);

        }

        if (column_schema.max_repetition_level > 0) {

            return Status::NotSupported(

                    "Parquet aggregate pushdown does not support repeated column {}",

                    column_schema.name);

        }

        *leaf_schema = &column_schema;

        return Status::OK();

    }

    if (projection.children.size() != 1) {

        return Status::NotSupported(

                "Parquet aggregate pushdown only supports a single nested leaf under column {}",

                column_schema.name);

    }

    const auto& child_projection = projection.children[0];

    const auto child_schema_it =

            std::ranges::find_if(column_schema.children, [&](const auto& child_schema) {

                return child_schema->local_id == child_projection.local_id();

            });

    if (child_schema_it != column_schema.children.end()) {

        return find_projected_minmax_leaf(**child_schema_it, child_projection, leaf_schema);

    }

    return Status::InvalidArgument("Invalid parquet aggregate projection local id {} for column {}",

                                   child_projection.local_id(), column_schema.name);

}

void ParquetReader::_fill_column_definition(const ParquetColumnSchema& column_schema,

                                            format::ColumnDefinition* field) const {

    if (column_schema.parquet_field_id >= 0) {

        field->identifier = Field::create_field<TYPE_INT>(column_schema.parquet_field_id);

    } else {

        field->identifier = Field::create_field<TYPE_STRING>(column_schema.name);

    }

    field->local_id = column_schema.local_id;

    field->name = column_schema.name;

    field->type = column_schema.type != nullptr && !column_schema.type->is_nullable()

                          ? make_nullable(column_schema.type)

                          : column_schema.type;

    field->children.clear();

    field->children.reserve(column_schema.children.size());

    for (const auto& child : column_schema.children) {

        format::ColumnDefinition child_field;

        _fill_column_definition(*child, &child_field);

        field->children.push_back(std::move(child_field));

    }

}

ParquetReader::ParquetReader(std::shared_ptr<io::FileSystemProperties>& system_properties,

                             std::unique_ptr<io::FileDescription>& file_description,

                             std::shared_ptr<io::IOContext> io_ctx, RuntimeProfile* profile,

                             std::optional<format::GlobalRowIdContext> global_rowid_context,

                             bool enable_mapping_timestamp_tz)

        : FileReader(system_properties, file_description, io_ctx, profile),

          _global_rowid_context(global_rowid_context),

          _enable_mapping_timestamp_tz(enable_mapping_timestamp_tz) {}

ParquetReader::~ParquetReader() = default;

Status ParquetReader::init(RuntimeState* state) {

    RETURN_IF_ERROR(format::FileReader::init(state));

    if (_profile != nullptr) {

        COUNTER_UPDATE(_parquet_profile.file_reader_create_time,

                       _reader_statistics.file_reader_create_time);

        COUNTER_UPDATE(_parquet_profile.open_file_num, _reader_statistics.open_file_num);

    }

    _state = std::make_unique<ParquetReaderScanState>();

    _state->enable_bloom_filter =

            state != nullptr && state->query_options().enable_parquet_filter_by_bloom_filter;

    if (state != nullptr) {

        _state->timezone = &state->timezone_obj();

        _state->scheduler.set_timezone(&state->timezone_obj());

        _state->scheduler.set_enable_strict_mode(state->enable_strict_mode());

    }

    // Open parquet file and parse metadata to get file schema.

    RETURN_IF_ERROR(_state->file_context.open(_tracing_file_reader, _io_ctx.get()));

    // Build file schema from parquet metadata.

    // A file reader may expose raw file identifiers, such as Parquet field_id, through ColumnDefinition::identifier

    RETURN_IF_ERROR(

            build_parquet_column_schema(*_state->file_context.schema, &_state->file_schema));

    if (_enable_mapping_timestamp_tz) {

        for (auto& column_schema : _state->file_schema) {

            apply_timestamp_tz_mapping(column_schema.get());

        }

    }

    return Status::OK();

}

Status ParquetReader::get_schema(std::vector<format::ColumnDefinition>* file_schema) const {

    if (file_schema == nullptr) {

        return Status::InvalidArgument("file_schema is null");

    }

    file_schema->clear();

    if (_state == nullptr || _state->file_context.schema == nullptr) {

        return Status::Uninitialized("ParquetReader is not open");

    }

    file_schema->reserve(_state->file_schema.size());

    for (size_t column_idx = 0; column_idx < _state->file_schema.size(); ++column_idx) {

        format::ColumnDefinition field;

        _fill_column_definition(*_state->file_schema[column_idx], &field);

        DORIS_CHECK(field.local_id == static_cast<int32_t>(column_idx));

        file_schema->push_back(std::move(field));

    }

    if (_global_rowid_context.has_value()) {

        file_schema->push_back(format::global_rowid_column_definition());

    }

    return Status::OK();

}

std::unique_ptr<format::TableColumnMapper> ParquetReader::create_column_mapper(

        format::TableColumnMapperOptions options) const {

    return std::make_unique<format::ParquetColumnMapper>(std::move(options));

}

Status ParquetReader::open(std::shared_ptr<format::FileScanRequest> request) {

    if (_state == nullptr || _state->file_context.metadata == nullptr ||

        _state->file_context.schema == nullptr) {

        return Status::Uninitialized("ParquetReader is not open");

    }

    auto request_snapshot = request;

    DORIS_CHECK(request_snapshot != nullptr);

    RETURN_IF_ERROR(format::FileReader::open(std::move(request)));

    const int num_fields = static_cast<int>(_state->file_schema.size());

    for (const auto& column_filter : request_snapshot->column_predicate_filters) {

        const auto file_column_id = column_filter.effective_file_column_id();

        if (!file_column_id.is_valid() || file_column_id.value() >= num_fields) {

            return Status::InvalidArgument("Invalid parquet filter top-level local id {}",

                                           file_column_id.value());

        }

    }

    // `local_positions.empty()` means all columns are needed by table reader

    // TODO(gabriel): It will happen only for TVF `select *` query.

    if (request_snapshot->local_positions.empty()) {

        for (const auto& col : request_snapshot->predicate_columns) {

            request_snapshot->local_positions.emplace(col.column_id(),

                                                      format::LocalIndex(col.column_id().value()));

        }

        for (const auto& col : request_snapshot->non_predicate_columns) {

            request_snapshot->local_positions.emplace(col.column_id(),

                                                      format::LocalIndex(col.column_id().value()));

        }

    }

    for (const auto& col : request_snapshot->predicate_columns) {

        DORIS_CHECK(request_snapshot->local_positions.count(col.column_id()) > 0);

        const auto local_id = col.local_id();

        if (local_id == format::ROW_POSITION_COLUMN_ID ||

            local_id == format::GLOBAL_ROWID_COLUMN_ID) {

            continue;

        }

        DORIS_CHECK(local_id >= 0 && local_id < num_fields);

    }

    for (const auto& col : request_snapshot->non_predicate_columns) {

        DORIS_CHECK(request_snapshot->local_positions.count(col.column_id()) > 0);

        const auto local_id = col.local_id();

        if (local_id == format::ROW_POSITION_COLUMN_ID ||

            local_id == format::GLOBAL_ROWID_COLUMN_ID) {

            continue;

        }

        DORIS_CHECK(local_id >= 0 && local_id < num_fields);

    }

    RowGroupScanPlan row_group_plan;

    ParquetScanRange scan_range;

    scan_range.start_offset = _file_description->range_start_offset;

    scan_range.size = _file_description->range_size;

    scan_range.file_size = _file_description->file_size;

    // Get selected ranges in row groups according to metadata (Row-Group level index and Page Index including Zonemap, Dictionary, Bloom Filter).

    RETURN_IF_ERROR(plan_parquet_row_groups(

            *_state->file_context.metadata, _state->file_context.file_reader.get(),

            _state->file_schema, *request_snapshot, scan_range, _state->enable_bloom_filter,

            &row_group_plan, _state->timezone));

    if (_profile != nullptr) {

        _parquet_profile.update_pruning_stats(row_group_plan.pruning_stats);

    }

    _state->scan_plan = row_group_plan;

    _state->scheduler.set_page_skip_profile(_parquet_profile.page_skip_profile());

    _state->scheduler.set_global_rowid_context(_global_rowid_context);

    _state->scheduler.set_scan_profile(_parquet_profile.scan_profile());

    _state->scheduler.set_plan(std::move(row_group_plan));

    _eof = _state->scheduler.empty();

    return Status::OK();

}

Status ParquetReader::get_block(Block* file_block, size_t* rows, bool* eof) {

    if (_state == nullptr || _state->file_context.file_reader == nullptr ||

        _state->file_context.schema == nullptr) {

        return Status::Uninitialized("ParquetReader is not open");

    }

    *rows = 0;

    if (_eof) {

        *eof = true;

        return Status::OK();

    }

    auto request_snapshot = _request;

    if (request_snapshot == nullptr) {

        return Status::Cancelled("ParquetReader is closed");

    }

    const auto predicate_filtered_rows_before = _state->scheduler.predicate_filtered_rows();

    RETURN_IF_ERROR(_state->scheduler.read_next_batch(_state->file_context, _state->file_schema,

                                                      *request_snapshot, file_block, rows, eof));

    if (_io_ctx != nullptr) {

        _io_ctx->predicate_filtered_rows +=

                _state->scheduler.predicate_filtered_rows() - predicate_filtered_rows_before;

    }

    _eof = *eof;

    return Status::OK();

}

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

    if (_state == nullptr) {

        return;

    }

    _state->scheduler.set_condition_cache_context(std::move(ctx));

    if (_io_ctx != nullptr) {

        // Condition-cache HIT filters row ranges before batch reading, so skipped rows never belong

        // to a later get_block() batch. Report the plan-level skipped rows at the same point where

        // the scan plan is rewritten.

        _io_ctx->condition_cache_filtered_rows += _state->scheduler.condition_cache_filtered_rows();

    }

}

int64_t ParquetReader::get_total_rows() const {

    if (_state == nullptr) {

        return 0;

    }

    int64_t rows = 0;

    for (const auto& row_group_plan : _state->scan_plan.row_groups) {

        rows += row_group_plan.row_group_rows;

    }

    return rows;

}

Status ParquetReader::get_aggregate_result(const format::FileAggregateRequest& request,

                                           format::FileAggregateResult* result) {

    DORIS_CHECK(result != nullptr);

    if (_state == nullptr || _state->file_context.metadata == nullptr ||

        _state->file_context.schema == nullptr) {

        return Status::Uninitialized("ParquetReader is not open");

    }

    result->count = 0;

    result->columns.clear();

    if (request.agg_type != TPushAggOp::type::COUNT &&

        request.agg_type != TPushAggOp::type::MINMAX) {

        return Status::NotSupported("Unsupported parquet aggregate pushdown type {}",

                                    request.agg_type);

    }

    // Aggregate row count in all selected row groups. For MIN/MAX aggregate, this is used to determine whether there is no row group selected.

    for (const auto& row_group_plan : _state->scan_plan.row_groups) {

        auto row_group_metadata =

                _state->file_context.metadata->RowGroup(row_group_plan.row_group_id);

        DORIS_CHECK(row_group_metadata != nullptr);

        result->count += row_group_metadata->num_rows();

    }

    if (request.agg_type == TPushAggOp::type::COUNT) {

        return Status::OK();

    }

    result->columns.resize(request.columns.size());

    for (size_t request_column_idx = 0; request_column_idx < request.columns.size();

         ++request_column_idx) {

        const auto file_column_id = request.columns[request_column_idx].projection.local_id();

        if (file_column_id < 0 ||

            file_column_id >= static_cast<int32_t>(_state->file_schema.size())) {

            return Status::InvalidArgument("Invalid parquet aggregate column id {}",

                                           file_column_id);

        }

        const auto& column_schema = _state->file_schema[file_column_id];

        DORIS_CHECK(column_schema != nullptr);

        const ParquetColumnSchema* leaf_schema = nullptr;

        RETURN_IF_ERROR(find_projected_minmax_leaf(

                *column_schema, request.columns[request_column_idx].projection, &leaf_schema));

        DORIS_CHECK(leaf_schema != nullptr);

        auto& aggregate_column = result->columns[request_column_idx];

        aggregate_column.projection = request.columns[request_column_idx].projection;

        for (const auto& row_group_plan : _state->scan_plan.row_groups) {

            auto row_group_metadata =

                    _state->file_context.metadata->RowGroup(row_group_plan.row_group_id);

            DORIS_CHECK(row_group_metadata != nullptr);

            auto column_chunk = row_group_metadata->ColumnChunk(leaf_schema->leaf_column_id);

            DORIS_CHECK(column_chunk != nullptr);

            const auto statistics = ParquetStatisticsUtils::TransformColumnStatistics(

                    *leaf_schema, column_chunk->statistics(), _state->timezone);

            if (!statistics.has_min_max) {

                return Status::NotSupported("Missing parquet min/max statistics for column {}",

                                            leaf_schema->name);

            }

            if (!aggregate_column.has_min || statistics.min_value < aggregate_column.min_value) {

                aggregate_column.min_value = statistics.min_value;

                aggregate_column.has_min = true;

            }

            if (!aggregate_column.has_max || aggregate_column.max_value < statistics.max_value) {

                aggregate_column.max_value = statistics.max_value;

                aggregate_column.has_max = true;

            }

        }

        if (!aggregate_column.has_min || !aggregate_column.has_max) {

            return Status::NotSupported("No parquet row group selected for min/max pushdown");

        }

    }

    return Status::OK();

}

Status ParquetReader::close() {

    if (_state != nullptr) {

        RETURN_IF_ERROR(_state->file_context.close());

    }

    return FileReader::close();

}

void ParquetReader::_init_profile() {

    _parquet_profile.init(_profile);

}

} // namespace doris::format::parquet