// 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/rowid_fetcher.h"

#include <brpc/callback.h>

#include <butil/endpoint.h>

#include <fmt/format.h>

#include <gen_cpp/data.pb.h>

#include <gen_cpp/internal_service.pb.h>

#include <gen_cpp/olap_file.pb.h>

#include <gen_cpp/types.pb.h>

#include <glog/logging.h>

#include <stddef.h>

#include <stdint.h>

#include <algorithm>

#include <cstdint>

#include <limits>

#include <memory>

#include <mutex>

#include <ostream>

#include <string>

#include <unordered_map>

#include <utility>

#include <vector>

#include "bthread/countdown_event.h"

#include "cloud/cloud_storage_engine.h"

#include "cloud/cloud_tablet.h"

#include "cloud/cloud_tablet_mgr.h"

#include "cloud/config.h"

#include "common/config.h"

#include "common/consts.h"

#include "common/exception.h"

#include "common/signal_handler.h"

#include "core/assert_cast.h"

#include "core/block/block.h" // Block

#include "core/column/column.h"

#include "core/column/column_nullable.h"

#include "core/column/column_string.h"

#include "core/data_type/data_type_struct.h"

#include "core/data_type_serde/data_type_serde.h"

#include "core/string_ref.h"

#include "core/types.h"

#include "exec/operator/file_scan_operator.h"

#include "exec/scan/file_scanner.h"

#include "exec/scan/file_scanner_v2.h"

#include "format/orc/vorc_reader.h"

#include "format/parquet/vparquet_reader.h"

#include "format_v2/table/lance_reader.h"

#include "io/io_common.h"

#include "runtime/descriptors.h"

#include "runtime/exec_env.h"      // ExecEnv

#include "runtime/fragment_mgr.h"  // FragmentMgr

#include "runtime/runtime_state.h" // RuntimeState

#include "runtime/workload_group/workload_group_manager.h"

#include "semaphore"

#include "storage/olap_common.h"

#include "storage/rowset/beta_rowset.h"

#include "storage/segment/column_reader.h"

#include "storage/storage_engine.h"

#include "storage/tablet/tablet_fwd.h"

#include "storage/tablet/tablet_schema.h"

#include "storage/tablet_info.h" // DorisNodesInfo

#include "storage/utils.h"

#include "util/brpc_client_cache.h" // BrpcClientCache

#include "util/defer_op.h"

#include "util/jsonb/serialize.h"

#include "util/pretty_printer.h"

namespace doris {

#include "common/compile_check_begin.h"

namespace {

void set_topn_lazy_materialization_file_cache_stats(

        const io::FileCacheStatistics& stats, PTopNLazyMaterializationFileCacheStats* pstats) {

    pstats->set_local_io_count(stats.num_local_io_total);

    pstats->set_local_io_bytes(stats.bytes_read_from_local);

    pstats->set_remote_io_count(stats.num_remote_io_total);

    pstats->set_remote_io_bytes(stats.bytes_read_from_remote);

    pstats->set_skip_cache_io_count(stats.num_skip_cache_io_total);

    pstats->set_write_cache_bytes(stats.bytes_write_into_cache);

    pstats->set_local_io_time(stats.local_io_timer);

    pstats->set_remote_io_time(stats.remote_io_timer);

    pstats->set_write_cache_io_time(stats.write_cache_io_timer);

}

} // namespace

Status RowIDFetcher::init() {

    DorisNodesInfo nodes_info;

    nodes_info.setNodes(_fetch_option.t_fetch_opt.nodes_info);

    for (auto [node_id, node_info] : nodes_info.nodes_info()) {

        auto client = ExecEnv::GetInstance()->brpc_internal_client_cache()->get_client(

                node_info.host, node_info.brpc_port);

        if (!client) {

            LOG(WARNING) << "Get rpc stub failed, host=" << node_info.host

                         << ", port=" << node_info.brpc_port;

            return Status::InternalError("RowIDFetcher failed to init rpc client, host={}, port={}",

                                         node_info.host, node_info.brpc_port);

        }

        _stubs.push_back(client);

    }

    return Status::OK();

}

PMultiGetRequest RowIDFetcher::_init_fetch_request(const ColumnString& row_locs) const {

    PMultiGetRequest mget_req;

    _fetch_option.desc->to_protobuf(mget_req.mutable_desc());

    for (SlotDescriptor* slot : _fetch_option.desc->slots()) {

        // ignore rowid

        if (slot->col_name() == BeConsts::ROWID_COL) {

            continue;

        }

        slot->to_protobuf(mget_req.add_slots());

    }

    for (size_t i = 0; i < row_locs.size(); ++i) {

        PRowLocation row_loc;

        StringRef row_id_rep = row_locs.get_data_at(i);

        // TODO: When transferring data between machines with different byte orders (endianness),

        // not performing proper handling may lead to issues in parsing and exchanging the data.

        auto location = reinterpret_cast<const GlobalRowLoacation*>(row_id_rep.data);

        row_loc.set_tablet_id(location->tablet_id);

        row_loc.set_rowset_id(location->row_location.rowset_id.to_string());

        row_loc.set_segment_id(location->row_location.segment_id);

        row_loc.set_ordinal_id(location->row_location.row_id);

        *mget_req.add_row_locs() = std::move(row_loc);

    }

    // Set column desc

    for (const TColumn& tcolumn : _fetch_option.t_fetch_opt.column_desc) {

        TabletColumn column(tcolumn);

        column.to_schema_pb(mget_req.add_column_desc());

    }

    PUniqueId& query_id = *mget_req.mutable_query_id();

    query_id.set_hi(_fetch_option.runtime_state->query_id().hi);

    query_id.set_lo(_fetch_option.runtime_state->query_id().lo);

    mget_req.set_be_exec_version(_fetch_option.runtime_state->be_exec_version());

    mget_req.set_fetch_row_store(_fetch_option.t_fetch_opt.fetch_row_store);

    return mget_req;

}

Status RowIDFetcher::_merge_rpc_results(const PMultiGetRequest& request,

                                        const std::vector<PMultiGetResponse>& rsps,

                                        const std::vector<brpc::Controller>& cntls,

                                        Block* output_block,

                                        std::vector<PRowLocation>* rows_id) const {

    output_block->clear();

    for (const auto& cntl : cntls) {

        if (cntl.Failed()) {

            LOG(WARNING) << "Failed to fetch meet rpc error:" << cntl.ErrorText()

                         << ", host:" << cntl.remote_side();

            return Status::InternalError(cntl.ErrorText());

        }

    }

    DataTypeSerDeSPtrs serdes;

    std::unordered_map<uint32_t, uint32_t> col_uid_to_idx;

    std::vector<std::string> default_values;

    default_values.resize(_fetch_option.desc->slots().size());

    auto merge_function = [&](const PMultiGetResponse& resp) {

        Status st(Status::create(resp.status()));

        if (!st.ok()) {

            LOG(WARNING) << "Failed to fetch " << st.to_string();

            return st;

        }

        for (const PRowLocation& row_id : resp.row_locs()) {

            rows_id->push_back(row_id);

        }

        // Merge binary rows

        if (request.fetch_row_store()) {

            CHECK(resp.row_locs().size() == resp.binary_row_data_size());

            if (output_block->is_empty_column()) {

                *output_block = Block(_fetch_option.desc->slots(), 1);

            }

            if (serdes.empty() && col_uid_to_idx.empty()) {

                serdes = create_data_type_serdes(_fetch_option.desc->slots());

                for (int i = 0; i < _fetch_option.desc->slots().size(); ++i) {

                    col_uid_to_idx[_fetch_option.desc->slots()[i]->col_unique_id()] = i;

                    default_values[i] = _fetch_option.desc->slots()[i]->col_default_value();

                }

            }

            auto output_columns_guard = output_block->mutate_columns_scoped();

            MutableColumns& output_columns = output_columns_guard.mutable_columns();

            for (int i = 0; i < resp.binary_row_data_size(); ++i) {

                RETURN_IF_ERROR(JsonbSerializeUtil::jsonb_to_columns(

                        serdes, resp.binary_row_data(i).data(), resp.binary_row_data(i).size(),

                        col_uid_to_idx, output_columns, default_values, {}));

            }

            return Status::OK();

        }

        // Merge partial blocks

        Block partial_block;

        [[maybe_unused]] size_t uncompressed_size = 0;

        [[maybe_unused]] int64_t uncompressed_time = 0;

        RETURN_IF_ERROR(

                partial_block.deserialize(resp.block(), &uncompressed_size, &uncompressed_time));

        if (partial_block.is_empty_column()) {

            return Status::OK();

        }

        CHECK(resp.row_locs().size() == partial_block.rows());

        if (output_block->is_empty_column()) {

            output_block->swap(partial_block);

        } else if (partial_block.columns() != output_block->columns()) {

            return Status::Error<ErrorCode::INTERNAL_ERROR>(

                    "Merge block not match, self:[{}], input:[{}], ", output_block->dump_types(),

                    partial_block.dump_types());

        } else {

            for (int i = 0; i < output_block->columns(); ++i) {

                auto column_guard = output_block->mutate_column_scoped(i);

                MutableColumnPtr& column = column_guard.mutable_column();

                column->insert_range_from(

                        *partial_block.get_by_position(i).column->convert_to_full_column_if_const(),

                        0, partial_block.rows());

            }

        }

        return Status::OK();

    };

    for (const auto& resp : rsps) {

        RETURN_IF_ERROR(merge_function(resp));

    }

    return Status::OK();

}

Status RowIDFetcher::fetch(const ColumnPtr& column_row_ids, Block* res_block) {

    CHECK(!_stubs.empty());

    PMultiGetRequest mget_req = _init_fetch_request(

            assert_cast<const ColumnString&>(*remove_nullable(column_row_ids).get()));

    std::vector<PMultiGetResponse> resps(_stubs.size());

    std::vector<brpc::Controller> cntls(_stubs.size());

    bthread::CountdownEvent counter(cast_set<int>(_stubs.size()));

    for (size_t i = 0; i < _stubs.size(); ++i) {

        cntls[i].set_timeout_ms(_fetch_option.runtime_state->execution_timeout() * 1000);

        auto callback = brpc::NewCallback(fetch_callback, &counter);

        _stubs[i]->multiget_data(&cntls[i], &mget_req, &resps[i], callback);

    }

    counter.wait();

    // Merge

    std::vector<PRowLocation> rows_locs;

    rows_locs.reserve(rows_locs.size());

    RETURN_IF_ERROR(_merge_rpc_results(mget_req, resps, cntls, res_block, &rows_locs));

    if (rows_locs.size() < column_row_ids->size()) {

        return Status::InternalError("Miss matched return row loc count {}, expected {}, input {}",

                                     rows_locs.size(), res_block->rows(), column_row_ids->size());

    }

    // Final sort by row_ids sequence, since row_ids is already sorted if need

    std::map<GlobalRowLoacation, size_t> positions;

    for (size_t i = 0; i < rows_locs.size(); ++i) {

        RowsetId rowset_id;

        rowset_id.init(rows_locs[i].rowset_id());

        GlobalRowLoacation grl(rows_locs[i].tablet_id(), rowset_id,

                               cast_set<uint32_t>(rows_locs[i].segment_id()),

                               cast_set<uint32_t>(rows_locs[i].ordinal_id()));

        positions[grl] = i;

    };

    // TODO remove this warning code

    if (positions.size() < rows_locs.size()) {

        LOG(WARNING) << "cwntains duplicated row entry";

    }

    IColumn::Permutation permutation;

    permutation.reserve(column_row_ids->size());

    for (size_t i = 0; i < column_row_ids->size(); ++i) {

        auto location =

                reinterpret_cast<const GlobalRowLoacation*>(column_row_ids->get_data_at(i).data);

        permutation.push_back(positions[*location]);

    }

    for (size_t i = 0; i < res_block->columns(); ++i) {

        res_block->get_by_position(i).column =

                res_block->get_by_position(i).column->permute(permutation, permutation.size());

    }

    // Check row consistency

    RETURN_IF_CATCH_EXCEPTION(res_block->check_number_of_rows());

    VLOG_DEBUG << "dump block:" << res_block->dump_data(0, 10);

    return Status::OK();

}

struct IteratorKey {

    int64_t tablet_id;

    RowsetId rowset_id;

    uint64_t segment_id;

    int slot_id;

    // unordered map std::equal_to

    bool operator==(const IteratorKey& rhs) const {

        return tablet_id == rhs.tablet_id && rowset_id == rhs.rowset_id &&

               segment_id == rhs.segment_id && slot_id == rhs.slot_id;

    }

};

struct SegKey {

    int64_t tablet_id;

    RowsetId rowset_id;

    uint64_t segment_id;

    // unordered map std::equal_to

    bool operator==(const SegKey& rhs) const {

        return tablet_id == rhs.tablet_id && rowset_id == rhs.rowset_id &&

               segment_id == rhs.segment_id;

    }

};

struct HashOfSegKey {

    size_t operator()(const SegKey& key) const {

        size_t seed = 0;

        seed = HashUtil::hash64(&key.tablet_id, sizeof(key.tablet_id), seed);

        seed = HashUtil::hash64(&key.rowset_id.hi, sizeof(key.rowset_id.hi), seed);

        seed = HashUtil::hash64(&key.rowset_id.mi, sizeof(key.rowset_id.mi), seed);

        seed = HashUtil::hash64(&key.rowset_id.lo, sizeof(key.rowset_id.lo), seed);

        seed = HashUtil::hash64(&key.segment_id, sizeof(key.segment_id), seed);

        return seed;

    }

};

struct HashOfIteratorKey {

    size_t operator()(const IteratorKey& key) const {

        size_t seed = 0;

        seed = HashUtil::hash64(&key.tablet_id, sizeof(key.tablet_id), seed);

        seed = HashUtil::hash64(&key.rowset_id.hi, sizeof(key.rowset_id.hi), seed);

        seed = HashUtil::hash64(&key.rowset_id.mi, sizeof(key.rowset_id.mi), seed);

        seed = HashUtil::hash64(&key.rowset_id.lo, sizeof(key.rowset_id.lo), seed);

        seed = HashUtil::hash64(&key.segment_id, sizeof(key.segment_id), seed);

        seed = HashUtil::hash64(&key.slot_id, sizeof(key.slot_id), seed);

        return seed;

    }

};

struct IteratorItem {

    std::unique_ptr<ColumnIterator> iterator;

    SegmentSharedPtr segment;

    // for holding the reference of storage read options to avoid use after release

    StorageReadOptions storage_read_options;

};

static void set_slot_access_paths(const SlotDescriptor& slot, const TabletSchema& schema,

                                  StorageReadOptions& storage_read_options) {

    int32_t unique_id = slot.col_unique_id();

    const int field_index =

            unique_id >= 0 ? schema.field_index(unique_id) : schema.field_index(slot.col_name());

    if (field_index >= 0) {

        const auto& column = schema.column(field_index);

        unique_id = column.unique_id() >= 0 ? column.unique_id() : column.parent_unique_id();

    }

    if (unique_id < 0) {

        return;

    }

    if (!slot.all_access_paths().empty()) {

        storage_read_options.all_access_paths[unique_id] = slot.all_access_paths();

    }

    if (!slot.predicate_access_paths().empty()) {

        storage_read_options.predicate_access_paths[unique_id] = slot.predicate_access_paths();

    }

}

struct SegItem {

    BaseTabletSPtr tablet;

    BetaRowsetSharedPtr rowset;

    // for holding the reference of segment to avoid use after release

    SegmentSharedPtr segment;

};

// Groups all row_ids belonging to the same segment for batched reading.

// Position index tracks where each row_id originated in the original request,

// so results can be scattered back to the correct output positions.

struct DorisFormatReadBatch {

    std::shared_ptr<FileMapping> file_mapping;

    // (row_id, index_in_request) pairs for all rows in this segment.

    std::vector<std::pair<segment_v2::rowid_t, size_t>> row_ids_with_positions;

};

static void scatter_scan_blocks_to_result_block(

        const std::vector<std::pair<size_t, size_t>>& row_id_block_idx,

        std::vector<Block>& scan_blocks, Block& result_block) {

    for (size_t column_id = 0; column_id < result_block.columns(); ++column_id) {

        auto column_guard = result_block.mutate_column_scoped(column_id);

        MutableColumnPtr& dst_col = column_guard.mutable_column();

        for (const auto& [pos_block, block_idx] : row_id_block_idx) {

            DCHECK(scan_blocks.size() > pos_block);

            DCHECK(scan_blocks[pos_block].columns() > column_id);

            const auto& src_col = *scan_blocks[pos_block].get_by_position(column_id).column;

            uint32_t scan_position = cast_set<uint32_t>(block_idx);

            dst_col->insert_indices_from(src_col, &scan_position, &scan_position + 1);

        }

    }

}

Status RowIdStorageReader::read_by_rowids(const PMultiGetRequest& request,

                                          PMultiGetResponse* response) {

    // read from storage engine row id by row id

    OlapReaderStatistics stats;

    Block result_block;

    int64_t acquire_tablet_ms = 0;

    int64_t acquire_rowsets_ms = 0;

    int64_t acquire_segments_ms = 0;

    int64_t lookup_row_data_ms = 0;

    // init desc

    std::vector<SlotDescriptor> slots;

    slots.reserve(request.slots().size());

    for (const auto& pslot : request.slots()) {

        slots.push_back(SlotDescriptor(pslot));

    }

    // init read schema

    TabletSchema full_read_schema;

    for (const ColumnPB& column_pb : request.column_desc()) {

        full_read_schema.append_column(TabletColumn(column_pb));

    }

    std::unordered_map<IteratorKey, IteratorItem, HashOfIteratorKey> iterator_map;

    // read row by row

    for (int i = 0; i < request.row_locs_size(); ++i) {

        const auto& row_loc = request.row_locs(i);

        MonotonicStopWatch watch;

        watch.start();

        BaseTabletSPtr tablet = scope_timer_run(

                [&]() {

                    auto res = ExecEnv::get_tablet(row_loc.tablet_id(), nullptr, true);

                    return !res.has_value() ? nullptr

                                            : std::dynamic_pointer_cast<BaseTablet>(res.value());

                },

                &acquire_tablet_ms);

        RowsetId rowset_id;

        rowset_id.init(row_loc.rowset_id());

        if (!tablet) {

            continue;

        }

        // We ensured it's rowset is not released when init Tablet reader param, rowset->update_delayed_expired_timestamp();

        BetaRowsetSharedPtr rowset = std::static_pointer_cast<BetaRowset>(scope_timer_run(

                [&]() {

                    return ExecEnv::GetInstance()->storage_engine().get_quering_rowset(rowset_id);

                },

                &acquire_rowsets_ms));

        if (!rowset) {

            LOG(INFO) << "no such rowset " << rowset_id;

            continue;

        }

        size_t row_size = 0;

        Defer _defer([&]() {

            LOG_EVERY_N(INFO, 100)

                    << "multiget_data single_row, cost(us):" << watch.elapsed_time() / 1000

                    << ", row_size:" << row_size;

            *response->add_row_locs() = row_loc;

        });

        // TODO: supoort session variable enable_page_cache and disable_file_cache if necessary.

        SegmentCacheHandle segment_cache;

        RETURN_IF_ERROR(scope_timer_run(

                [&]() {

                    return SegmentLoader::instance()->load_segments(rowset, &segment_cache, true);

                },

                &acquire_segments_ms));

        // find segment

        auto it = std::find_if(segment_cache.get_segments().cbegin(),

                               segment_cache.get_segments().cend(),

                               [&row_loc](const segment_v2::SegmentSharedPtr& seg) {

                                   return seg->id() == row_loc.segment_id();

                               });

        if (it == segment_cache.get_segments().end()) {

            continue;

        }

        segment_v2::SegmentSharedPtr segment = *it;

        GlobalRowLoacation row_location(row_loc.tablet_id(), rowset->rowset_id(),

                                        cast_set<uint32_t>(row_loc.segment_id()),

                                        cast_set<uint32_t>(row_loc.ordinal_id()));

        // fetch by row store, more effcient way

        if (request.fetch_row_store()) {

            if (!tablet->tablet_schema()->has_row_store_for_all_columns()) {

                return Status::InternalError("Tablet {} does not have row store for all columns",

                                             tablet->tablet_id());

            }

            RowLocation loc(rowset_id, segment->id(), cast_set<uint32_t>(row_loc.ordinal_id()));

            std::string* value = response->add_binary_row_data();

            RETURN_IF_ERROR(scope_timer_run(

                    [&]() { return tablet->lookup_row_data({}, loc, rowset, stats, *value); },

                    &lookup_row_data_ms));

            row_size = value->size();

            continue;

        }

        // fetch by column store

        if (result_block.is_empty_column()) {

            result_block = Block(slots, request.row_locs().size());

        }

        VLOG_DEBUG << "Read row location "

                   << fmt::format("{}, {}, {}, {}", row_location.tablet_id,

                                  row_location.row_location.rowset_id.to_string(),

                                  row_location.row_location.segment_id,

                                  row_location.row_location.row_id);

        for (int x = 0; x < slots.size(); ++x) {

            std::vector<segment_v2::rowid_t> row_ids {

                    static_cast<segment_v2::rowid_t>(row_loc.ordinal_id())};

            // The scoped mutation republishes a potentially replaced V2 column when the reader

            // transfers ownership of its first assembled batch.

            auto column_guard = result_block.mutate_column_scoped(x);

            MutableColumnPtr& column = column_guard.mutable_column();

            IteratorKey iterator_key {.tablet_id = tablet->tablet_id(),

                                      .rowset_id = rowset_id,

                                      .segment_id = row_loc.segment_id(),

                                      .slot_id = slots[x].id()};

            IteratorItem& iterator_item = iterator_map[iterator_key];

            if (iterator_item.segment == nullptr) {

                // hold the reference

                iterator_map[iterator_key].segment = segment;

                iterator_item.storage_read_options.stats = &stats;

                iterator_item.storage_read_options.io_ctx.reader_type = ReaderType::READER_QUERY;

            }

            segment = iterator_item.segment;

            set_slot_access_paths(slots[x], full_read_schema, iterator_item.storage_read_options);

            RETURN_IF_ERROR(segment->seek_and_read_by_rowid(

                    full_read_schema, &slots[x], row_ids, column,

                    iterator_item.storage_read_options, iterator_item.iterator));

        }

    }

    // serialize block if not empty

    if (!result_block.is_empty_column()) {

        VLOG_DEBUG << "dump block:" << result_block.dump_data(0, 10)

                   << ", be_exec_version:" << request.be_exec_version();

        [[maybe_unused]] size_t compressed_size = 0;

        [[maybe_unused]] size_t uncompressed_size = 0;

        [[maybe_unused]] int64_t compress_time = 0;

        int be_exec_version = request.has_be_exec_version() ? request.be_exec_version() : 0;

        RETURN_IF_ERROR(result_block.serialize(be_exec_version, response->mutable_block(),

                                               &uncompressed_size, &compressed_size, &compress_time,

                                               segment_v2::CompressionTypePB::LZ4));

    }

    LOG(INFO) << "Query stats: "

              << fmt::format(

                         "query_id:{}, "

                         "hit_cached_pages:{}, total_pages_read:{}, compressed_bytes_read:{}, "

                         "io_latency:{}ns, "

                         "uncompressed_bytes_read:{},"

                         "bytes_read:{},"

                         "acquire_tablet_ms:{}, acquire_rowsets_ms:{}, acquire_segments_ms:{}, "

                         "lookup_row_data_ms:{}",

                         print_id(request.query_id()), stats.cached_pages_num,

                         stats.total_pages_num, stats.compressed_bytes_read, stats.io_ns,

                         stats.uncompressed_bytes_read, stats.bytes_read, acquire_tablet_ms,

                         acquire_rowsets_ms, acquire_segments_ms, lookup_row_data_ms);

    return Status::OK();

}

Status RowIdStorageReader::read_by_rowids(const PMultiGetRequestV2& request,

                                          PMultiGetResponseV2* response) {

    if (request.request_block_descs_size()) {

        auto tquery_id = ((UniqueId)request.query_id()).to_thrift();

        // todo: use mutableBlock instead of block

        std::vector<Block> result_blocks(request.request_block_descs_size());

        OlapReaderStatistics stats;

        int64_t acquire_tablet_ms = 0;

        int64_t acquire_rowsets_ms = 0;

        int64_t acquire_segments_ms = 0;

        int64_t lookup_row_data_ms = 0;

        int64_t external_init_reader_avg_ms = 0;

        int64_t external_get_block_avg_ms = 0;

        size_t external_scan_range_cnt = 0;

        const auto file_cache_miss_policy =

                request.file_cache_remote_only_on_miss()

                        ? io::FileCacheMissPolicy::REMOTE_ONLY_ON_MISS

                        : io::FileCacheMissPolicy::READ_THROUGH_AND_WRITE_BACK;

        // Add counters for different file mapping types

        std::unordered_map<FileMappingType, int64_t> file_type_counts;

        auto id_file_map =

                ExecEnv::GetInstance()->get_id_manager()->get_id_file_map(request.query_id());

        // if id_file_map is null, means the BE not have scan range, just return ok

        if (!id_file_map) {

            // padding empty block to response

            LOG(INFO) << "id_file_map not found for query_id: " << print_id(request.query_id());

            for (int i = 0; i < request.request_block_descs_size(); ++i) {

                response->add_blocks();

            }

            return Status::OK();

        }

        for (int i = 0; i < request.request_block_descs_size(); ++i) {

            const auto& request_block_desc = request.request_block_descs(i);

            PMultiGetBlockV2* pblock = response->add_blocks();

            if (request_block_desc.row_id_size() >= 1) {

                if (request_block_desc.file_id_size() != request_block_desc.row_id_size()) {

                    return Status::InvalidArgument(

                            "Row-id fetch has mismatched file_id and row_id counts: "

                            "file_ids={}, row_ids={}",

                            request_block_desc.file_id_size(), request_block_desc.row_id_size());

                }

                const auto row_location_version = request_block_desc.row_location_version();

                switch (row_location_version) {

                case static_cast<uint32_t>(ROW_VERSION::FILE_LOCAL_ROW_ID):

                case static_cast<uint32_t>(ROW_VERSION::LANCE_DATASET_ROW_ID):

                    break;

                default:

                    return Status::NotSupported("unsupported row location version in fetch: {}",

                                                row_location_version);

                }

                // Rows in this request block belong to the same relation. The first mapping is

                // sufficient to select the internal/external fetch path; individual file IDs

                // are still used below to group the actual reads.

                const auto first_file_id = request_block_desc.file_id(0);

                auto first_file_mapping = id_file_map->get_file_mapping(first_file_id);

                if (!first_file_mapping) {

                    return Status::InternalError(

                            "Backend:{} file mapping not found, query_id: {}, file_id: {}",

                            BackendOptions::get_localhost(), print_id(request.query_id()),

                            first_file_id);

                }

                if (first_file_mapping->type != FileMappingType::EXTERNAL &&

                    row_location_version ==

                            static_cast<uint32_t>(ROW_VERSION::LANCE_DATASET_ROW_ID)) {

                    return Status::InvalidArgument(

                            "Lance dataset row IDs require an external file mapping");

                }

                file_type_counts[first_file_mapping->type] += request_block_desc.row_id_size();

                // prepare slots to build block

                std::vector<SlotDescriptor> slots;

                slots.reserve(request_block_desc.slots_size());

                for (const auto& pslot : request_block_desc.slots()) {

                    slots.push_back(SlotDescriptor(pslot));

                }

                try {

                    if (first_file_mapping->type == FileMappingType::INTERNAL) {

                        RETURN_IF_ERROR(read_batch_doris_format_row(

                                request_block_desc, id_file_map, slots, tquery_id, result_blocks[i],

                                stats, &acquire_tablet_ms, &acquire_rowsets_ms,

                                &acquire_segments_ms, &lookup_row_data_ms, file_cache_miss_policy));

                    } else {

                        RETURN_IF_ERROR(read_batch_external_row(

                                request.wg_id(), request_block_desc, id_file_map, slots,

                                first_file_mapping, tquery_id, result_blocks[i],

                                pblock->mutable_profile(), &external_init_reader_avg_ms,

                                &external_get_block_avg_ms, &external_scan_range_cnt));

                    }

                } catch (const Exception& e) {

                    return Status::Error<false>(e.code(), "Row id fetch failed because {}",

                                                e.what());

                }

            }

            [[maybe_unused]] size_t compressed_size = 0;

            [[maybe_unused]] size_t uncompressed_size = 0;

            [[maybe_unused]] int64_t compress_time = 0;

            int be_exec_version = request.has_be_exec_version() ? request.be_exec_version() : 0;

            RETURN_IF_ERROR(result_blocks[i].serialize(

                    be_exec_version, pblock->mutable_block(), &uncompressed_size, &compressed_size,

                    &compress_time, segment_v2::CompressionTypePB::LZ4));

        }

        // Build file type statistics string

        std::string file_type_stats;

        for (const auto& [type, count] : file_type_counts) {

            if (!file_type_stats.empty()) {

                file_type_stats += ", ";

            }

            file_type_stats += fmt::format("{}:{}", type, count);

        }

        LOG(INFO) << "Query stats: "

                  << fmt::format(

                             "query_id:{}, "

                             "Internal table:"

                             "hit_cached_pages:{}, total_pages_read:{}, compressed_bytes_read:{}, "

                             "io_latency:{}ns, uncompressed_bytes_read:{}, bytes_read:{}, "

                             "acquire_tablet_ms:{}, acquire_rowsets_ms:{}, acquire_segments_ms:{}, "

                             "lookup_row_data_ms:{}, file_types:[{}]; "

                             "External table : init_reader_ms:{}, get_block_ms:{}, "

                             "external_scan_range_cnt:{}",

                             print_id(request.query_id()), stats.cached_pages_num,

                             stats.total_pages_num, stats.compressed_bytes_read, stats.io_ns,

                             stats.uncompressed_bytes_read, stats.bytes_read, acquire_tablet_ms,

                             acquire_rowsets_ms, acquire_segments_ms, lookup_row_data_ms,

                             file_type_stats, external_init_reader_avg_ms,

                             external_get_block_avg_ms, external_scan_range_cnt);

        set_topn_lazy_materialization_file_cache_stats(

                stats.file_cache_stats,

                response->mutable_topn_lazy_materialization_file_cache_stats());

    }

    if (request.has_gc_id_map() && request.gc_id_map()) {

        ExecEnv::GetInstance()->get_id_manager()->remove_id_file_map(request.query_id());

    }

    return Status::OK();

}

Status RowIdStorageReader::read_batch_doris_format_row(

        const PRequestBlockDesc& request_block_desc, std::shared_ptr<IdFileMap> id_file_map,

        std::vector<SlotDescriptor>& slots, const TUniqueId& query_id, Block& result_block,

        OlapReaderStatistics& stats, int64_t* acquire_tablet_ms, int64_t* acquire_rowsets_ms,

        int64_t* acquire_segments_ms, int64_t* lookup_row_data_ms,

        io::FileCacheMissPolicy file_cache_miss_policy) {

    if (result_block.is_empty_column()) [[likely]] {

        result_block = Block(slots, request_block_desc.row_id_size());

    }

    TabletSchema full_read_schema;

    for (const ColumnPB& column_pb : request_block_desc.column_descs()) {

        full_read_schema.append_column(TabletColumn(column_pb));

    }

    std::unordered_map<IteratorKey, IteratorItem, HashOfIteratorKey> iterator_map;

    std::unordered_map<SegKey, SegItem, HashOfSegKey> seg_map;

    std::string row_store_buffer;

    RowStoreReadStruct row_store_read_struct(row_store_buffer);

    if (request_block_desc.fetch_row_store()) {

        for (int i = 0; i < request_block_desc.slots_size(); ++i) {

            row_store_read_struct.serdes.emplace_back(slots[i].get_data_type_ptr()->get_serde());

            row_store_read_struct.col_uid_to_idx[slots[i].col_unique_id()] = i;

            row_store_read_struct.default_values.emplace_back(slots[i].col_default_value());

        }

    }

    // Phase 1: Group all row_ids by their (tablet_id, rowset_id, segment_id) key.

    // Unlike the old code which only batched adjacent rows with the same file_id,

    // this merges non-contiguous same-segment requests into a single batch,

    // maximizing the number of rows read per seek_and_read_by_rowid call.

    std::vector<DorisFormatReadBatch> scan_batches;

    std::unordered_map<SegKey, size_t, HashOfSegKey> batch_idx_by_seg;

    // (batch_idx, position_in_batch) for each row in the original request.

    std::vector<std::pair<size_t, size_t>> row_id_block_idx(request_block_desc.row_id_size());

    for (int j = 0; j < request_block_desc.row_id_size(); ++j) {

        auto file_id = request_block_desc.file_id(j);

        const auto row_id = request_block_desc.row_id(j);

        if (row_id > std::numeric_limits<uint32_t>::max()) {

            return Status::InvalidArgument("internal row id exceeds uint32 range: {}", row_id);

        }

        auto file_mapping = id_file_map->get_file_mapping(file_id);

        if (!file_mapping) {

            return Status::InternalError(

                    "Backend:{} file_mapping not found, query_id: {}, file_id: {}",

                    BackendOptions::get_localhost(), print_id(query_id), file_id);

        }

        // Derive segment key and group by it — rows from the same segment are batched together

        // even if they are interleaved with rows from other segments in the request.

        auto [tablet_id, rowset_id, segment_id] = file_mapping->get_doris_format_info();

        SegKey seg_key {.tablet_id = tablet_id, .rowset_id = rowset_id, .segment_id = segment_id};

        auto [it, inserted] = batch_idx_by_seg.emplace(seg_key, scan_batches.size());

        if (inserted) {

            // First time seeing this segment, create a new batch for it.

            scan_batches.emplace_back();

            scan_batches.back().file_mapping = file_mapping;

        }

        // Record (row_id, original_request_index) for later sorting and scattering.

        scan_batches[it->second].row_ids_with_positions.emplace_back(static_cast<uint32_t>(row_id),

                                                                     j);

    }

    // Phase 2: For each segment, sort row_ids ascending (required by ColumnIterator),

    // deduplicate, then read all rows in a single batch call.

    std::vector<Block> scan_blocks(scan_batches.size());

    for (size_t batch_idx = 0; batch_idx < scan_batches.size(); ++batch_idx) {

        auto& scan_batch = scan_batches[batch_idx];

        auto& row_ids_with_positions = scan_batch.row_ids_with_positions;

        std::sort(row_ids_with_positions.begin(), row_ids_with_positions.end(),

                  [](const auto& lhs, const auto& rhs) { return lhs.first < rhs.first; });

        // Column iterators read rowids monotonically. Deduplicate consecutive identical row_ids

        // (different file_ids may map to the same row), then scatter rows back to their original

        // request positions.

        std::vector<uint32_t> row_ids;

        row_ids.reserve(row_ids_with_positions.size());

        // Also builds the scatter map: row_id_block_idx[original_request_idx] ->

        // (batch_idx, deduplicated_position_in_batch).

        for (const auto& [row_id, result_idx] : row_ids_with_positions) {

            if (row_ids.empty() || row_ids.back() != row_id) {

                row_ids.emplace_back(row_id);

            }

            row_id_block_idx[result_idx] = std::make_pair(batch_idx, row_ids.size() - 1);

        }

        scan_blocks[batch_idx] = Block(slots, row_ids.size());

        RETURN_IF_ERROR(read_doris_format_row(

                id_file_map, scan_batch.file_mapping, row_ids, slots, full_read_schema,

                row_store_read_struct, stats, acquire_tablet_ms, acquire_rowsets_ms,

                acquire_segments_ms, lookup_row_data_ms, seg_map, iterator_map,

                file_cache_miss_policy, scan_blocks[batch_idx]));

    }

    scatter_scan_blocks_to_result_block(row_id_block_idx, scan_blocks, result_block);

    return Status::OK();

}

const std::string RowIdStorageReader::ScannersRunningTimeProfile = "ScannersRunningTime";

const std::string RowIdStorageReader::InitReaderAvgTimeProfile = "InitReaderAvgTime";

const std::string RowIdStorageReader::GetBlockAvgTimeProfile = "GetBlockAvgTime";

const std::string RowIdStorageReader::FileReadLinesProfile = "FileReadLines";

const std::string RowIdStorageReader::LanceDatasetOpenTimeProfile = "LanceDatasetOpenTime";

const std::string RowIdStorageReader::LanceRowIdTakeReadTimeProfile = "LanceRowIdTakeReadTime";

const std::string RowIdStorageReader::LanceArrowToDorisBlockTimeProfile =

        "LanceArrowToDorisBlockTime";

const std::string RowIdStorageReader::LanceRowIdFetchTotalTimeProfile = "LanceRowIdFetchTotalTime";

const std::map<std::string, TUnit::type> RowIdStorageReader::LanceFetchCountersProfile = {

        {"LanceRowIdFetchRows", TUnit::UNIT},

        {"LanceRowIdFetchCalls", TUnit::UNIT},

        {"LanceDataCacheBytesReadFromCache", TUnit::BYTES},

        {"LanceDataCacheBytesReadFromRemote", TUnit::BYTES},

};

const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseLocalIOCount =

        "TopNLazyMaterializationSecondPhaseLocalIOCount";

const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseLocalIOBytes =

        "TopNLazyMaterializationSecondPhaseLocalIOBytes";

const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseRemoteIOCount =

        "TopNLazyMaterializationSecondPhaseRemoteIOCount";

const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseRemoteIOBytes =

        "TopNLazyMaterializationSecondPhaseRemoteIOBytes";

const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseSkipCacheIOCount =

        "TopNLazyMaterializationSecondPhaseSkipCacheIOCount";

const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseWriteCacheBytes =

        "TopNLazyMaterializationSecondPhaseWriteCacheBytes";

const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseLocalIOTime =

        "TopNLazyMaterializationSecondPhaseLocalIOTime";

const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseRemoteIOTime =

        "TopNLazyMaterializationSecondPhaseRemoteIOTime";

const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseWriteCacheIOTime =

        "TopNLazyMaterializationSecondPhaseWriteCacheIOTime";

const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseRowsRead =

        "TopNLazyMaterializationSecondPhaseRowsRead";

const std::string RowIdStorageReader::TopNLazyMaterializationSecondPhaseSegmentsRead =

        "TopNLazyMaterializationSecondPhaseSegmentsRead";

bool RowIdStorageReader::should_use_file_scanner_v2(const TQueryOptions& query_options,

                                                    const TFileScanRangeParams& scan_params,

                                                    const TFileRangeDesc& range) {

    const auto format_type =

            range.__isset.format_type ? range.format_type : scan_params.format_type;

    // Phase two inherits the query options, including the Thrift presence bit. Reuse phase one's

    // policy so disabling V2 (or an older payload omitting the option) also keeps row fetches on V1.

    return FileScanLocalState::should_use_file_scanner_v2(query_options, false, scan_params) &&

           (format_type == TFileFormatType::FORMAT_PARQUET ||

            format_type == TFileFormatType::FORMAT_ORC) &&

           FileScannerV2::is_supported(scan_params, range);

}

TFileRangeDesc RowIdStorageReader::build_external_fetch_range(const TFileRangeDesc& source_range) {

    // Rows were selected after delete filtering. Preserve the original path and row lineage

    // needed by virtual columns, and do not mutate the FileMapping shared by other fetches.

    auto range = source_range;

    range.table_format_params.iceberg_params.__set_delete_files({});

    range.table_format_params.transactional_hive_params = TTransactionalHiveDesc {};

    return range;

}

TFileScanRangeParams RowIdStorageReader::build_external_scan_params(

        const TFileScanRangeParams& source_params, const TFileRangeDesc& range,

        const std::vector<SlotDescriptor>& scan_slots,

        const std::vector<uint32_t>& scan_column_idxs) {

    DORIS_CHECK(scan_slots.size() == scan_column_idxs.size());

    auto params = source_params;

    params.required_slots.clear();

    params.column_idxs.clear();

    params.slot_name_to_schema_pos.clear();

    const std::set partition_names(range.columns_from_path_keys.begin(),

                                   range.columns_from_path_keys.end());

    for (size_t slot_idx = 0; slot_idx < scan_slots.size(); ++slot_idx) {

        const auto& slot = scan_slots[slot_idx];

        const auto column_idx = scan_column_idxs[slot_idx];

        TFileScanSlotInfo slot_info;

        slot_info.__set_slot_id(slot.id());

        // Hive V2 checks the Thrift presence bit before trusting is_file_slot. Without it,

        // partition columns consume physical file indexes and invalidate the rebuilt projection.

        bool is_file_slot = !partition_names.contains(slot.col_name());

        if (source_params.__isset.column_name_to_category) {

            // Lazy metadata slots may be absent from phase one's required_slots and have new

            // slot IDs here. The pinned schema's name map preserves their original categories.

            const auto it = source_params.column_name_to_category.find(slot.col_name());

            const auto category = it != source_params.column_name_to_category.end()

                                          ? it->second

                                          : TColumnCategory::REGULAR;

            slot_info.__set_category(category);

            is_file_slot =

                    category == TColumnCategory::REGULAR || category == TColumnCategory::GENERATED;

        }

        slot_info.__set_is_file_slot(is_file_slot);

        if (is_file_slot) {

            params.column_idxs.emplace_back(column_idx);

        }

        params.default_value_of_src_slot.emplace(slot.id(), TExpr {});

        params.required_slots.emplace_back(slot_info);

        params.slot_name_to_schema_pos.emplace(slot.col_name(), column_idx);

    }

    return params;

}

Status RowIdStorageReader::read_lance_rows_by_row_ids(

        const TFileRangeDesc& scan_range_desc, const std::vector<uint64_t>& row_ids,

        const std::vector<SlotDescriptor>& slots, RuntimeState* runtime_state,

        RuntimeProfile* runtime_profile, const TFileScanRangeParams& scan_params, Block* block,

        RowIdStorageReader::ExternalFetchStatistics* fetch_statistics) {

    // Unlike Parquet/ORC, a Lance row ID is a native uint64 ID in a fixed dataset snapshot,

    // rather than a row ordinal interpreted by a reader for one physical file range. Phase-two

    // fetch therefore opens that dataset snapshot and uses dataset-level take_rows directly. It

    // does not create a FileScanner or rescan the fragment split that produced the row ID.

    std::vector<format::ColumnDefinition> projected_columns;

    projected_columns.reserve(slots.size());

    for (const auto& slot : slots) {

        format::ColumnDefinition column;

        column.identifier = Field::create_field<TYPE_STRING>(slot.col_name());

        column.name = slot.col_name();

        column.type = slot.get_data_type_ptr();

        projected_columns.emplace_back(std::move(column));

    }

    format::lance::LanceTableReader reader;

    auto lance_scan_params = scan_params;

    RETURN_IF_ERROR(scope_timer_run(

            [&]() {

                return reader.init({

                        .projected_columns = projected_columns,

                        .conjuncts = {},

                        .format = format::FileFormat::LANCE,

                        .scan_params = &lance_scan_params,

                        .io_ctx = nullptr,

                        .runtime_state = runtime_state,