// 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 "exec/scan/parallel_scanner_builder.h"

#include <cstddef>

#include "cloud/cloud_storage_engine.h"

#include "cloud/cloud_tablet_hotspot.h"

#include "cloud/config.h"

#include "common/status.h"

#include "exec/operator/olap_scan_operator.h"

#include "exec/scan/olap_scanner.h"

#include "io/io_common.h"

#include "runtime/query_context.h"

#include "storage/rowset/beta_rowset.h"

#include "storage/segment/segment_loader.h"

#include "storage/tablet/base_tablet.h"

namespace doris {

namespace {

io::FileCacheStatistics take_initial_file_cache_stats(

        std::unordered_map<int64_t, io::FileCacheStatistics>* preload_stats, int64_t tablet_id) {

    auto it = preload_stats->find(tablet_id);

    if (it == preload_stats->end()) {

        return {};

    }

    auto stats = std::move(it->second);

    preload_stats->erase(it);

    return stats;

}

io::IOContext create_preload_io_context(RuntimeState* state, OlapReaderStatistics* preload_stats) {

    io::IOContext io_ctx;

    io_ctx.reader_type = ReaderType::READER_QUERY;

    io_ctx.file_cache_stats = preload_stats ? &preload_stats->file_cache_stats : nullptr;

    if (state == nullptr) {

        return io_ctx;

    }

    io_ctx.query_id = &state->query_id();

    io_ctx.read_file_cache = state->query_options().enable_file_cache;

    io_ctx.is_disposable = state->query_options().disable_file_cache;

    if (auto* query_ctx = state->get_query_ctx(); query_ctx != nullptr) {

        io_ctx.remote_scan_cache_write_limiter = query_ctx->remote_scan_cache_write_limiter();

    }

    return io_ctx;

}

} // namespace

Status ParallelScannerBuilder::build_scanners(std::list<ScannerSPtr>& scanners) {

    RETURN_IF_ERROR(_load());

    if (_scan_parallelism_by_per_segment) {

        return _build_scanners_by_per_segment(scanners);

    } else if (_is_dup_mow_key) {

        // Default strategy for DUP/MOW tables: split by rowids within segments

        return _build_scanners_by_rowid(scanners);

    } else {

        // TODO: support to split by key range

        return Status::NotSupported("split by key range not supported yet.");

    }

}

Status ParallelScannerBuilder::_build_scanners_by_rowid(std::list<ScannerSPtr>& scanners) {

    DCHECK_GE(_rows_per_scanner, _min_rows_per_scanner);

    for (auto&& [tablet, version] : _tablets) {

        DCHECK(_all_read_sources.contains(tablet->tablet_id()));

        auto& entire_read_source = _all_read_sources[tablet->tablet_id()];

        if (config::is_cloud_mode()) {

            // FIXME(plat1ko): Avoid pointer cast

            ExecEnv::GetInstance()->storage_engine().to_cloud().tablet_hotspot().count(*tablet);

        }

        // `rs_splits` in `entire read source` will be devided into several partitial read sources

        // to build several parallel scanners, based on segment rows number. All the partitial read sources

        // share the same delete predicates from their corresponding entire read source.

        TabletReadSource partitial_read_source;

        int64_t rows_collected = 0;

        for (auto& rs_split : entire_read_source.rs_splits) {

            auto reader = rs_split.rs_reader;

            auto rowset = reader->rowset();

            const auto rowset_id = rowset->rowset_id();

            const auto& segments_rows = _all_segments_rows[rowset_id];

            if (rowset->num_rows() == 0) {

                continue;

            }

            int64_t segment_start = 0;

            auto split = RowSetSplits(reader->clone());

            for (size_t i = 0; i != segments_rows.size(); ++i) {

                const size_t rows_of_segment = segments_rows[i];

                RowRanges row_ranges;

                int64_t offset_in_segment = 0;

                // try to split large segments into RowRanges

                while (offset_in_segment < rows_of_segment) {

                    const int64_t remaining_rows = rows_of_segment - offset_in_segment;

                    auto rows_need = _rows_per_scanner - rows_collected;

                    // 0.9: try to avoid splitting the segments into excessively small parts.

                    if (rows_need >= remaining_rows * 9 / 10) {

                        rows_need = remaining_rows;

                    }

                    DCHECK_LE(rows_need, remaining_rows);

                    // RowRange stands for range: [From, To), From is inclusive, To is exclusive.

                    row_ranges.add({offset_in_segment,

                                    offset_in_segment + static_cast<int64_t>(rows_need)});

                    rows_collected += rows_need;

                    offset_in_segment += rows_need;

                    // If collected enough rows, build a new scanner

                    if (rows_collected >= _rows_per_scanner) {

                        split.segment_offsets.first = segment_start,

                        split.segment_offsets.second = i + 1;

                        split.segment_row_ranges.emplace_back(std::move(row_ranges));

                        DCHECK_EQ(split.segment_offsets.second - split.segment_offsets.first,

                                  split.segment_row_ranges.size());

                        partitial_read_source.rs_splits.emplace_back(std::move(split));

                        scanners.emplace_back(_build_scanner(

                                tablet, version, _key_ranges,

                                {.rs_splits = std::move(partitial_read_source.rs_splits),

                                 .delete_predicates = entire_read_source.delete_predicates,

                                 .delete_bitmap = entire_read_source.delete_bitmap},

                                take_initial_file_cache_stats(&_tablet_preload_file_cache_stats,

                                                              tablet->tablet_id())));

                        partitial_read_source = {};

                        split = RowSetSplits(reader->clone());

                        row_ranges = RowRanges();

                        segment_start = offset_in_segment < rows_of_segment ? i : i + 1;

                        rows_collected = 0;

                    }

                }

                // The non-empty `row_ranges` means there are some rows left in this segment not added into `split`.

                if (!row_ranges.is_empty()) {

                    DCHECK_GT(rows_collected, 0);

                    DCHECK_EQ(row_ranges.to(), rows_of_segment);

                    split.segment_row_ranges.emplace_back(std::move(row_ranges));

                }

            }

            DCHECK_LE(rows_collected, _rows_per_scanner);

            if (rows_collected > 0) {

                split.segment_offsets.first = segment_start;

                split.segment_offsets.second = segments_rows.size();

                DCHECK_GT(split.segment_offsets.second, split.segment_offsets.first);

                DCHECK_EQ(split.segment_row_ranges.size(),

                          split.segment_offsets.second - split.segment_offsets.first);

                partitial_read_source.rs_splits.emplace_back(std::move(split));

            }

        } // end `for (auto& rowset : rowsets)`

        DCHECK_LE(rows_collected, _rows_per_scanner);

        if (rows_collected > 0) {

            DCHECK_GT(partitial_read_source.rs_splits.size(), 0);

#ifndef NDEBUG

            for (auto& split : partitial_read_source.rs_splits) {

                DCHECK(split.rs_reader != nullptr);

                DCHECK_LT(split.segment_offsets.first, split.segment_offsets.second);

                DCHECK_EQ(split.segment_row_ranges.size(),

                          split.segment_offsets.second - split.segment_offsets.first);

            }

#endif

            scanners.emplace_back(

                    _build_scanner(tablet, version, _key_ranges,

                                   {.rs_splits = std::move(partitial_read_source.rs_splits),

                                    .delete_predicates = entire_read_source.delete_predicates,

                                    .delete_bitmap = entire_read_source.delete_bitmap},

                                   take_initial_file_cache_stats(&_tablet_preload_file_cache_stats,

                                                                 tablet->tablet_id())));

        }

    }

    return Status::OK();

}

// Build scanners so that each segment is exclusively scanned by a single scanner.

// This guarantees the number of scanners equals the number of segments across all rowsets

// for the involved tablets. It preserves delete predicates and key ranges, and clones

// RowsetReader per scanner to avoid sharing between scanners.

Status ParallelScannerBuilder::_build_scanners_by_per_segment(std::list<ScannerSPtr>& scanners) {

    DCHECK_GE(_rows_per_scanner, _min_rows_per_scanner);

    for (auto&& [tablet, version] : _tablets) {

        DCHECK(_all_read_sources.contains(tablet->tablet_id()));

        auto& entire_read_source = _all_read_sources[tablet->tablet_id()];

        if (config::is_cloud_mode()) {

            // FIXME(plat1ko): Avoid pointer cast

            ExecEnv::GetInstance()->storage_engine().to_cloud().tablet_hotspot().count(*tablet);

        }

        // For each RowSet split in the read source, split by segment id and build

        // one scanner per segment. Keep delete predicates shared.

        for (auto& rs_split : entire_read_source.rs_splits) {

            auto reader = rs_split.rs_reader;

            auto rowset = reader->rowset();

            const auto rowset_id = rowset->rowset_id();

            const auto& segments_rows = _all_segments_rows[rowset_id];

            if (segments_rows.empty() || rowset->num_rows() == 0) {

                continue;

            }

            // Build scanners for [i, i+1) segment range, without row-range slicing.

            for (int64_t i = 0; i < rowset->num_segments(); ++i) {

                RowSetSplits split(reader->clone());

                split.segment_offsets.first = i;

                split.segment_offsets.second = i + 1;

                // No row-ranges slicing; scan whole segment i.

                DCHECK_GE(split.segment_offsets.second, split.segment_offsets.first + 1);

                TabletReadSource partitial_read_source;

                partitial_read_source.rs_splits.emplace_back(std::move(split));

                scanners.emplace_back(_build_scanner(

                        tablet, version, _key_ranges,

                        {.rs_splits = std::move(partitial_read_source.rs_splits),

                         .delete_predicates = entire_read_source.delete_predicates,

                         .delete_bitmap = entire_read_source.delete_bitmap},

                        take_initial_file_cache_stats(&_tablet_preload_file_cache_stats,

                                                      tablet->tablet_id())));

            }

        }

    }

    return Status::OK();

}

/**

 * Load rowsets of each tablet with specified version, segments of each rowset.

 */

Status ParallelScannerBuilder::_load() {

    _total_rows = 0;

    size_t idx = 0;

    bool enable_segment_cache = _state->query_options().__isset.enable_segment_cache

                                        ? _state->query_options().enable_segment_cache

                                        : true;

    for (auto&& [tablet, version] : _tablets) {

        const auto tablet_id = tablet->tablet_id();

        _all_read_sources[tablet_id] = _read_sources[idx];

        const auto& read_source = _all_read_sources[tablet_id];

        for (auto& rs_split : read_source.rs_splits) {

            auto rowset = rs_split.rs_reader->rowset();

            RETURN_IF_ERROR(rowset->load());

            const auto rowset_id = rowset->rowset_id();

            auto beta_rowset = std::dynamic_pointer_cast<BetaRowset>(rowset);

            std::vector<uint32_t> segment_rows;

            OlapReaderStatistics preload_stats;

            auto preload_io_ctx = create_preload_io_context(_state, &preload_stats);

            RETURN_IF_ERROR(beta_rowset->get_segment_num_rows(&segment_rows, enable_segment_cache,

                                                              &preload_stats, &preload_io_ctx));

            _tablet_preload_file_cache_stats[tablet_id].merge_from(preload_stats.file_cache_stats);

            auto segment_count = rowset->num_segments();

            for (int64_t i = 0; i != segment_count; i++) {

                _all_segments_rows[rowset_id].emplace_back(segment_rows[i]);

            }

            _total_rows += rowset->num_rows();

        }

        idx++;

    }

    _rows_per_scanner = _total_rows / _max_scanners_count;

    _rows_per_scanner = std::max<size_t>(_rows_per_scanner, _min_rows_per_scanner);

    return Status::OK();

}

std::shared_ptr<OlapScanner> ParallelScannerBuilder::_build_scanner(

        BaseTabletSPtr tablet, int64_t version, const std::vector<OlapScanRange*>& key_ranges,

        TabletReadSource&& read_source, io::FileCacheStatistics&& initial_file_cache_stats) {

    OlapScanner::Params params {

            .state = _state,

            .profile = _scanner_profile.get(),

            .key_ranges = key_ranges,

            .tablet = std::move(tablet),

            .version = version,

            .read_source = std::move(read_source),

            .initial_file_cache_stats = std::move(initial_file_cache_stats),

            .limit = _limit,

            .aggregation = _is_preaggregation,

            .read_row_binlog = false,

            .binlog_scan_type = TBinlogScanType::NONE,

            .start_tso = std::nullopt,

            .end_tso = std::nullopt,

    };

    return OlapScanner::create_shared(_parent, std::move(params));

}

} // namespace doris