// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed 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 <algorithm>
#include <numeric>
#include <cstddef>
#include <memory>
#include <mutex>
#include <optional>
#include <shared_mutex>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include "TimestampData.h"
#include "TimestampIndex.h"
#include "common/ArrayOffsets.h"
#include "common/EasyAssert.h"
#include "common/Schema.h"
#include "common/TrackingStdAllocator.h"
#include "common/Types.h"
#include "mmap/ChunkedColumn.h"
#include "segcore/AckResponder.h"
#include "segcore/ConcurrentVector.h"
#include <tuple>
#include <type_traits>

namespace milvus::segcore {

constexpr int64_t Unlimited = -1;
constexpr int64_t NoLimit = 0;
// If `bitset_count * 100` > `total_count * BruteForceSelectivity`, we use pk index.
// Otherwise, we use bruteforce to retrieve all the pks and then sort them.
constexpr int64_t BruteForceSelectivity = 10;

using Condition = std::function<bool(int64_t)>;

// Compressed storage for offset -> int64 PK reverse lookup.
// Uses global base + per-block base + bitpacked deltas.
// Block size = 128, O(1) random access within each block.
// Each block stores deltas with a uniform bit width (max bits needed in that block).
class CompressedInt64PkArray {
 public:
    static constexpr int64_t kBlockSize = 128;

    CompressedInt64PkArray() = default;

    void
    build(const int64_t* pks, int64_t num_rows) {
        if (num_rows == 0) {
            return;
        }
        num_rows_ = num_rows;

        // Find global minimum as base
        base_ = pks[0];
        for (int64_t i = 1; i < num_rows; ++i) {
            if (pks[i] < base_) {
                base_ = pks[i];
            }
        }

        // Calculate number of blocks
        int64_t num_blocks = (num_rows + kBlockSize - 1) / kBlockSize;
        block_bit_widths_.reserve(num_blocks);
        block_offsets_.reserve(num_blocks);
        block_bases_.reserve(num_blocks);

        for (int64_t block_id = 0; block_id < num_blocks; ++block_id) {
            int64_t block_start = block_id * kBlockSize;
            int64_t block_end = std::min(block_start + kBlockSize, num_rows);
            int64_t block_count = block_end - block_start;

            // Find block minimum
            int64_t block_min = pks[block_start];
            for (int64_t i = block_start + 1; i < block_end; ++i) {
                if (pks[i] < block_min) {
                    block_min = pks[i];
                }
            }

            // Find max delta to determine bit width
            uint64_t max_delta = 0;
            for (int64_t i = block_start; i < block_end; ++i) {
                uint64_t delta = static_cast<uint64_t>(pks[i] - block_min);
                if (delta > max_delta) {
                    max_delta = delta;
                }
            }

            uint8_t bits = bit_width(max_delta);
            block_bases_.push_back(static_cast<uint64_t>(block_min - base_));
            block_bit_widths_.push_back(bits);
            block_offsets_.push_back(static_cast<uint32_t>(data_.size()));

            // Pack deltas with uniform bit width
            if (bits == 0) {
                // All values identical in this block, no data needed
            } else {
                // Number of bytes needed: ceil(block_count * bits / 8)
                size_t num_bytes =
                    (static_cast<size_t>(block_count) * bits + 7) / 8;
                size_t data_start = data_.size();
                data_.resize(data_start + num_bytes, 0);

                for (int64_t i = 0; i < block_count; ++i) {
                    uint64_t delta =
                        static_cast<uint64_t>(pks[block_start + i] - block_min);
                    pack_value(data_start, i, bits, delta);
                }
            }
        }

        // Shrink to fit
        data_.shrink_to_fit();
        block_offsets_.shrink_to_fit();
        block_bases_.shrink_to_fit();
        block_bit_widths_.shrink_to_fit();
    }

    int64_t
    at(int64_t offset) const {
        AssertInfo(offset >= 0 && offset < num_rows_,
                   "offset out of range: {} not in [0, {})",
                   offset,
                   num_rows_);

        int64_t block_id = offset / kBlockSize;
        int64_t idx_in_block = offset % kBlockSize;
        uint8_t bits = block_bit_widths_[block_id];

        uint64_t delta = 0;
        if (bits > 0) {
            delta = unpack_value(block_offsets_[block_id], idx_in_block, bits);
        }

        return base_ + static_cast<int64_t>(block_bases_[block_id]) +
               static_cast<int64_t>(delta);
    }

    void
    bulk_at(const int64_t* offsets, int64_t count, int64_t* output) const {
        for (int64_t i = 0; i < count; ++i) {
            output[i] = at(offsets[i]);
        }
    }

    size_t
    memory_size() const {
        return sizeof(*this) + data_.capacity() +
               block_offsets_.capacity() * sizeof(uint32_t) +
               block_bases_.capacity() * sizeof(uint64_t) +
               block_bit_widths_.capacity() * sizeof(uint8_t);
    }

    int64_t
    num_rows() const {
        return num_rows_;
    }

    bool
    empty() const {
        return num_rows_ == 0;
    }

 private:
    static uint8_t
    bit_width(uint64_t max_val) {
        if (max_val == 0) {
            return 0;
        }
        return static_cast<uint8_t>(64 - __builtin_clzll(max_val));
    }

    void
    pack_value(size_t data_start, int64_t idx, uint8_t bits, uint64_t value) {
        uint64_t bit_offset = static_cast<uint64_t>(idx) * bits;
        size_t byte_offset = data_start + bit_offset / 8;
        unsigned bit_shift = bit_offset % 8;

        // Low 64 bits of (value << bit_shift)
        uint64_t lo = value << bit_shift;
        for (unsigned b = 0; b < 8 && byte_offset + b < data_.size(); ++b) {
            data_[byte_offset + b] |= static_cast<uint8_t>(lo >> (b * 8));
        }
        // High bits that overflowed past 64-bit boundary
        if (bit_shift > 0 && bit_shift + bits > 64) {
            uint64_t hi = value >> (64 - bit_shift);
            for (unsigned b = 0; byte_offset + 8 + b < data_.size() && b < 8;
                 ++b) {
                data_[byte_offset + 8 + b] |=
                    static_cast<uint8_t>(hi >> (b * 8));
            }
        }
    }

    uint64_t
    unpack_value(uint32_t block_data_offset, int64_t idx, uint8_t bits) const {
        uint64_t bit_offset = static_cast<uint64_t>(idx) * bits;
        size_t byte_offset = block_data_offset + bit_offset / 8;
        unsigned bit_shift = bit_offset % 8;

        // Read low 8 bytes
        uint64_t lo = 0;
        for (unsigned b = 0; b < 8 && byte_offset + b < data_.size(); ++b) {
            lo |= static_cast<uint64_t>(data_[byte_offset + b]) << (b * 8);
        }
        uint64_t result = lo >> bit_shift;

        // Read high bytes if crossing 64-bit boundary
        if (bit_shift > 0 && bit_shift + bits > 64) {
            uint64_t hi = 0;
            for (unsigned b = 0; byte_offset + 8 + b < data_.size() && b < 8;
                 ++b) {
                hi |= static_cast<uint64_t>(data_[byte_offset + 8 + b])
                      << (b * 8);
            }
            result |= hi << (64 - bit_shift);
        }

        uint64_t mask = (bits == 64) ? ~uint64_t(0) : (uint64_t(1) << bits) - 1;
        return result & mask;
    }

 private:
    int64_t base_ = 0;  // Global minimum PK
    int64_t num_rows_ = 0;
    std::vector<uint32_t> block_offsets_;  // Byte offset of each block in data_
    std::vector<uint64_t> block_bases_;    // Block base relative to global base
    std::vector<uint8_t> block_bit_widths_;  // Bit width per block
    std::vector<uint8_t> data_;              // Bitpacked deltas
};

class OffsetMap {
 public:
    virtual ~OffsetMap() = default;

    virtual bool
    contain(const PkType& pk) const = 0;

    virtual std::vector<int64_t>
    find(const PkType& pk) const = 0;

    virtual void
    find_range(const PkType& pk,
               proto::plan::OpType op,
               BitsetTypeView& bitset,
               Condition condition) const = 0;

    virtual void
    insert(const PkType& pk, int64_t offset) = 0;

    virtual void
    seal() = 0;

    virtual bool
    empty() const = 0;

    using OffsetType = int64_t;
    // TODO: in fact, we can retrieve the pk here. Not sure which way is more efficient.
    virtual std::pair<std::vector<OffsetMap::OffsetType>, bool>
    find_first_n(int64_t limit, const BitsetTypeView& bitset) const = 0;

    // Element-level version of find_first_n.
    // Find first N elements that pass filter, ordered by PK then element_index.
    // Returns:
    //   - vector of unique doc_offsets (no duplicates)
    //   - vector of element_indices per doc (element_indices[i] for doc_offsets[i])
    //   - has_more flag indicating if there are more results
    virtual std::
        tuple<std::vector<int64_t>, std::vector<std::vector<int32_t>>, bool>
        find_first_n_element(
            int64_t limit,
            const BitsetTypeView& element_bitset,
            const IArrayOffsets* array_offsets,
            const std::optional<QueryIteratorCursor>& cursor) const = 0;

    virtual void
    clear() = 0;

    virtual size_t
    memory_size() const = 0;

    // Returns true if the implementation stores nothing per-row (e.g.
    // VirtualPKOffsetMap derives pk→offset via bit-extract). Callers that
    // fall back to full pk-column scans (GetAllChunks + two-pointers) can
    // short-circuit to find_range() when this is true — O(M) not O(M+N).
    virtual bool
    is_zero_storage() const {
        return false;
    }
};

template <typename T>
class OffsetOrderedMap : public OffsetMap {
 public:
    using OrderedMap = std::map<
        T,
        std::vector<int64_t>,
        std::less<>,
        TrackingStdAllocator<std::pair<const T, std::vector<int64_t>>>>;

    bool
    contain(const PkType& pk) const override {
        std::shared_lock<std::shared_mutex> lck(mtx_);

        return map_.find(std::get<T>(pk)) != map_.end();
    }

    std::vector<int64_t>
    find(const PkType& pk) const override {
        std::shared_lock<std::shared_mutex> lck(mtx_);

        auto offset_vector = map_.find(std::get<T>(pk));
        return offset_vector != map_.end() ? offset_vector->second
                                           : std::vector<int64_t>();
    }

    void
    find_range(const PkType& pk,
               proto::plan::OpType op,
               BitsetTypeView& bitset,
               Condition condition) const override {
        std::shared_lock<std::shared_mutex> lck(mtx_);
        const T& target = std::get<T>(pk);

        if (op == proto::plan::OpType::Equal) {
            auto it = map_.find(target);
            if (it != map_.end()) {
                for (auto offset : it->second) {
                    if (condition(offset) && offset < bitset.size()) {
                        bitset[offset] = true;
                    }
                }
            }
        } else if (op == proto::plan::OpType::GreaterEqual) {
            auto it = map_.lower_bound(target);
            for (; it != map_.end(); ++it) {
                for (auto offset : it->second) {
                    if (condition(offset) && offset < bitset.size()) {
                        bitset[offset] = true;
                    }
                }
            }
        } else if (op == proto::plan::OpType::GreaterThan) {
            auto it = map_.upper_bound(target);
            for (; it != map_.end(); ++it) {
                for (auto offset : it->second) {
                    if (condition(offset) && offset < bitset.size()) {
                        bitset[offset] = true;
                    }
                }
            }
        } else if (op == proto::plan::OpType::LessEqual) {
            auto it = map_.upper_bound(target);
            for (auto ptr = map_.begin(); ptr != it; ++ptr) {
                for (auto offset : ptr->second) {
                    if (condition(offset) && offset < bitset.size()) {
                        bitset[offset] = true;
                    }
                }
            }
        } else if (op == proto::plan::OpType::LessThan) {
            auto it = map_.lower_bound(target);
            for (auto ptr = map_.begin(); ptr != it; ++ptr) {
                for (auto offset : ptr->second) {
                    if (condition(offset) && offset < bitset.size()) {
                        bitset[offset] = true;
                    }
                }
            }
        } else {
            ThrowInfo(ErrorCode::Unsupported,
                      fmt::format("unsupported op type {}", op));
        }
    }

    void
    insert(const PkType& pk, int64_t offset) override {
        std::unique_lock<std::shared_mutex> lck(mtx_);

        map_[std::get<T>(pk)].emplace_back(offset);
    }

    void
    seal() override {
        ThrowInfo(
            NotImplemented,
            "OffsetOrderedMap used for growing segment could not be sealed.");
    }

    bool
    empty() const override {
        std::shared_lock<std::shared_mutex> lck(mtx_);

        return map_.empty();
    }

    std::pair<std::vector<OffsetMap::OffsetType>, bool>
    find_first_n(int64_t limit, const BitsetTypeView& bitset) const override {
        std::shared_lock<std::shared_mutex> lck(mtx_);

        if (limit == Unlimited || limit == NoLimit) {
            limit = map_.size();
        }

        // TODO: we can't retrieve pk by offset very conveniently.
        //      Selectivity should be done outside.
        return find_first_n_by_index(limit, bitset);
    }

    std::tuple<std::vector<int64_t>, std::vector<std::vector<int32_t>>, bool>
    find_first_n_element(
        int64_t limit,
        const BitsetTypeView& element_bitset,
        const IArrayOffsets* array_offsets,
        const std::optional<QueryIteratorCursor>& cursor) const override {
        std::shared_lock<std::shared_mutex> lck(mtx_);

        if (limit == Unlimited || limit == NoLimit) {
            limit = static_cast<int64_t>(element_bitset.size());
        }

        return find_first_n_element_by_index(
            limit, element_bitset, array_offsets, cursor);
    }

    void
    clear() override {
        std::unique_lock<std::shared_mutex> lck(mtx_);
        map_.clear();
    }

    size_t
    memory_size() const override {
        std::shared_lock<std::shared_mutex> lck(mtx_);
        return map_.get_allocator().total_allocated();
    }

 private:
    std::pair<std::vector<OffsetMap::OffsetType>, bool>
    find_first_n_by_index(int64_t limit, const BitsetTypeView& bitset) const {
        int64_t hit_num = 0;  // avoid counting the number everytime.
        auto size = bitset.size();
        int64_t cnt = size - bitset.count();
        auto more_hit_than_limit = cnt > limit;
        limit = std::min(limit, cnt);
        std::vector<int64_t> seg_offsets;
        seg_offsets.reserve(limit);
        auto it = map_.begin();
        for (; hit_num < limit && it != map_.end(); it++) {
            // Offsets in the growing segment are ordered by timestamp,
            // so traverse from back to front to obtain the latest offset.
            for (int i = it->second.size() - 1; i >= 0; --i) {
                auto seg_offset = it->second[i];
                if (seg_offset >= size) {
                    // Frequently concurrent insert/query will cause this case.
                    continue;
                }

                if (!bitset[seg_offset]) {
                    seg_offsets.push_back(seg_offset);
                    hit_num++;
                    // PK hit, no need to continue traversing offsets with the same PK.
                    break;
                }
            }
        }
        return {seg_offsets, more_hit_than_limit && it != map_.end()};
    }

    std::tuple<std::vector<int64_t>, std::vector<std::vector<int32_t>>, bool>
    find_first_n_element_by_index(
        int64_t limit,
        const BitsetTypeView& element_bitset,
        const IArrayOffsets* array_offsets,
        const std::optional<QueryIteratorCursor>& cursor) const {
        std::vector<int64_t> doc_offsets;
        std::vector<std::vector<int32_t>> element_indices;

        int64_t hit_num = 0;
        auto element_size = static_cast<int64_t>(element_bitset.size());
        // Clamp limit to the actual number of matching elements,
        // same as find_first_n_by_index does for doc-level queries.
        int64_t cnt = element_size - element_bitset.count();
        auto more_hit_than_limit = cnt > limit;
        limit = std::min(limit, cnt);

        // Traverse map_ in PK order; for each PK visit offsets from back to
        // front to obtain the latest offset, and only use the first (newest)
        // offset that has matching elements (same as find_first_n_by_index)
        // to avoid returning stale versions.
        // Element ranges are resolved in windows: gather the next
        // kRangeBatchSize candidate doc offsets in traversal order and
        // resolve them with ONE CopyRowElementRanges call (one lock-free
        // published snapshot for the whole window on growing) instead of one
        // virtual call per doc. Ranges prefetched for docs the traversal ends up
        // skipping (older offsets of an already-hit PK, or docs past the
        // limit) are simply unused — the lookups are pure.
        constexpr int64_t kRangeBatchSize = 64;
        int32_t batch_offsets[kRangeBatchSize];
        std::pair<int32_t, int32_t> batch_ranges[kRangeBatchSize];
        const typename OrderedMap::value_type* batch_pks[kRangeBatchSize];

        std::vector<int32_t> matching_indices;
        // Gather cursor over the flattened (pk, offset) traversal.
        auto gather_it = map_.begin();
        int64_t gather_idx =
            gather_it == map_.end()
                ? -1
                : static_cast<int64_t>(gather_it->second.size()) - 1;
        // PK entry whose remaining (older) offsets must be skipped — the
        // batched equivalent of the per-doc inner-loop `break`.
        const typename OrderedMap::value_type* skip_pk = nullptr;
        while (hit_num < limit && gather_it != map_.end()) {
            // Gather the next window of candidate docs.
            int64_t n = 0;
            while (n < kRangeBatchSize && gather_it != map_.end()) {
                if (gather_idx < 0) {
                    ++gather_it;
                    if (gather_it == map_.end()) {
                        break;
                    }
                    gather_idx =
                        static_cast<int64_t>(gather_it->second.size()) - 1;
                    continue;
                }
                batch_offsets[n] =
                    static_cast<int32_t>(gather_it->second[gather_idx]);
                batch_pks[n] = &*gather_it;
                ++n;
                --gather_idx;
            }
            array_offsets->CopyRowElementRanges(batch_offsets, n, batch_ranges);

            for (int64_t k = 0; k < n && hit_num < limit; ++k) {
                const auto* pk_entry = batch_pks[k];
                if (pk_entry == skip_pk) {
                    continue;
                }
                const int64_t doc_offset = batch_offsets[k];

                // Get element range for this doc
                auto [first_elem, last_elem] = batch_ranges[k];

                // Collect all matching element indices for this doc
                matching_indices.clear();
                for (int64_t elem_id = first_elem;
                     elem_id < last_elem && hit_num < limit;
                     ++elem_id) {
                    if (elem_id >= element_size) {
                        continue;
                    }
                    if (is_skipped_by_cursor(
                            pk_entry->first, elem_id - first_elem, cursor)) {
                        continue;
                    }
                    if (!element_bitset[elem_id]) {  // 0 means pass filter
                        matching_indices.push_back(
                            static_cast<int32_t>(elem_id - first_elem));
                        hit_num++;
                    }
                }

                // Only add doc if it has matching elements
                if (!matching_indices.empty()) {
                    doc_offsets.push_back(doc_offset);
                    element_indices.push_back(std::move(matching_indices));
                    // PK hit, no need to continue traversing older offsets with the same PK.
                    skip_pk = pk_entry;
                } else if (is_cursor_pk(pk_entry->first, cursor)) {
                    // The cursor applies to the newest visible row for this PK.
                    // Do not fall through to older offsets of the same PK.
                    skip_pk = pk_entry;
                }
            }
        }

        bool has_more = more_hit_than_limit && hit_num >= limit;
        return {std::move(doc_offsets), std::move(element_indices), has_more};
    }

    bool
    is_skipped_by_cursor(
        const T& pk,
        int64_t element_offset,
        const std::optional<QueryIteratorCursor>& cursor) const {
        return is_cursor_pk(pk, cursor) &&
               element_offset <= cursor->last_element_offset;
    }

    bool
    is_cursor_pk(const T& pk,
                 const std::optional<QueryIteratorCursor>& cursor) const {
        if (!cursor.has_value()) {
            return false;
        }
        auto last_pk = std::get_if<T>(&cursor->last_pk);
        return last_pk != nullptr && *last_pk == pk;
    }

 private:
    OrderedMap map_;
    mutable std::shared_mutex mtx_;
};

template <typename T>
class OffsetOrderedArray : public OffsetMap {
 public:
    bool
    contain(const PkType& pk) const override {
        const T& target = std::get<T>(pk);
        auto it =
            std::lower_bound(array_.begin(),
                             array_.end(),
                             target,
                             [](const std::pair<T, int64_t>& elem,
                                const T& value) { return elem.first < value; });

        return it != array_.end() && it->first == target;
    }

    std::vector<int64_t>
    find(const PkType& pk) const override {
        check_search();

        const T& target = std::get<T>(pk);
        auto it =
            std::lower_bound(array_.begin(),
                             array_.end(),
                             target,
                             [](const std::pair<T, int64_t>& elem,
                                const T& value) { return elem.first < value; });

        std::vector<int64_t> offset_vector;
        for (; it != array_.end() && it->first == target; ++it) {
            offset_vector.push_back(it->second);
        }

        return offset_vector;
    }

    void
    find_range(const PkType& pk,
               proto::plan::OpType op,
               BitsetTypeView& bitset,
               Condition condition) const override {
        check_search();
        auto lower_bound_comp = [](const std::pair<T, int64_t>& elem,
                                   const T& value) {
            return elem.first < value;
        };
        auto upper_bound_comp = [](const T& value,
                                   const std::pair<T, int64_t>& elem) {
            return value < elem.first;
        };

        const T& target = std::get<T>(pk);
        if (op == proto::plan::OpType::Equal) {
            auto it = std::lower_bound(
                array_.begin(), array_.end(), target, lower_bound_comp);
            for (; it != array_.end() && it->first == target; ++it) {
                if (condition(it->second)) {
                    bitset[it->second] = true;
                }
            }
        } else if (op == proto::plan::OpType::GreaterEqual) {
            auto it = std::lower_bound(
                array_.begin(), array_.end(), target, lower_bound_comp);
            for (; it < array_.end(); ++it) {
                if (condition(it->second)) {
                    bitset[it->second] = true;
                }
            }
        } else if (op == proto::plan::OpType::GreaterThan) {
            auto it = std::upper_bound(
                array_.begin(), array_.end(), target, upper_bound_comp);
            for (; it < array_.end(); ++it) {
                if (condition(it->second)) {
                    bitset[it->second] = true;
                }
            }
        } else if (op == proto::plan::OpType::LessEqual) {
            auto it = std::upper_bound(
                array_.begin(), array_.end(), target, upper_bound_comp);
            for (auto ptr = array_.begin(); ptr < it; ++ptr) {
                if (condition(ptr->second)) {
                    bitset[ptr->second] = true;
                }
            }
        } else if (op == proto::plan::OpType::LessThan) {
            auto it = std::lower_bound(
                array_.begin(), array_.end(), target, lower_bound_comp);
            for (auto ptr = array_.begin(); ptr < it; ++ptr) {
                if (condition(ptr->second)) {
                    bitset[ptr->second] = true;
                }
            }
        } else {
            ThrowInfo(ErrorCode::Unsupported,
                      fmt::format("unsupported op type {}", op));
        }
    }

    void
    insert(const PkType& pk, int64_t offset) override {
        if (is_sealed) {
            ThrowInfo(Unsupported,
                      "OffsetOrderedArray could not insert after seal");
        }
        array_.push_back(
            std::make_pair(std::get<T>(pk), static_cast<int32_t>(offset)));
    }

    void
    seal() override {
        sort(array_.begin(), array_.end());
        is_sealed = true;
    }

    bool
    empty() const override {
        return array_.empty();
    }

    std::pair<std::vector<OffsetMap::OffsetType>, bool>
    find_first_n(int64_t limit, const BitsetTypeView& bitset) const override {
        check_search();

        if (limit == Unlimited || limit == NoLimit) {
            limit = array_.size();
        }

        // TODO: we can't retrieve pk by offset very conveniently.
        //      Selectivity should be done outside.
        return find_first_n_by_index(limit, bitset);
    }

    std::tuple<std::vector<int64_t>, std::vector<std::vector<int32_t>>, bool>
    find_first_n_element(
        int64_t limit,
        const BitsetTypeView& element_bitset,
        const IArrayOffsets* array_offsets,
        const std::optional<QueryIteratorCursor>& cursor) const override {
        check_search();

        if (limit == Unlimited || limit == NoLimit) {
            limit = static_cast<int64_t>(element_bitset.size());
        }

        return find_first_n_element_by_index(
            limit, element_bitset, array_offsets, cursor);
    }

    void
    clear() override {
        array_.clear();
        is_sealed = false;
    }

    size_t
    memory_size() const override {
        return sizeof(std::pair<T, int32_t>) * array_.capacity();
    }

 private:
    std::pair<std::vector<OffsetMap::OffsetType>, bool>
    find_first_n_by_index(int64_t limit, const BitsetTypeView& bitset) const {
        int64_t hit_num = 0;  // avoid counting the number everytime.
        auto size = bitset.size();
        int64_t cnt = size - bitset.count();
        auto more_hit_than_limit = cnt > limit;
        limit = std::min(limit, cnt);
        std::vector<int64_t> seg_offsets;
        seg_offsets.reserve(limit);
        auto it = array_.begin();
        for (; hit_num < limit && it != array_.end(); it++) {
            auto seg_offset = it->second;
            if (seg_offset >= size) {
                // In fact, this case won't happen on sealed segments.
                continue;
            }

            if (!bitset[seg_offset]) {
                seg_offsets.push_back(seg_offset);
                hit_num++;
            }
        }
        return {seg_offsets, more_hit_than_limit && it != array_.end()};
    }

    std::tuple<std::vector<int64_t>, std::vector<std::vector<int32_t>>, bool>
    find_first_n_element_by_index(
        int64_t limit,
        const BitsetTypeView& element_bitset,
        const IArrayOffsets* array_offsets,
        const std::optional<QueryIteratorCursor>& cursor) const {
        std::vector<int64_t> doc_offsets;
        std::vector<std::vector<int32_t>> element_indices;

        int64_t hit_num = 0;
        auto element_size = static_cast<int64_t>(element_bitset.size());
        // Clamp limit to the actual number of matching elements,
        // same as find_first_n_by_index does for doc-level queries.
        int64_t cnt = element_size - element_bitset.count();
        auto more_hit_than_limit = cnt > limit;
        limit = std::min(limit, cnt);

        // Traverse array_ in PK order (already sorted). Element ranges are
        // resolved in windows of kRangeBatchSize docs with one
        // CopyRowElementRanges call each instead of one virtual call per
        // doc; ranges prefetched past an early limit-exit are unused (the
        // lookups are pure).
        constexpr int64_t kRangeBatchSize = 64;
        int32_t batch_rows[kRangeBatchSize];
        std::pair<int32_t, int32_t> batch_ranges[kRangeBatchSize];

        std::vector<int32_t> matching_indices;
        const int64_t total = static_cast<int64_t>(array_.size());
        for (int64_t base = 0; base < total && hit_num < limit;
             base += kRangeBatchSize) {
            const int64_t n = std::min(kRangeBatchSize, total - base);
            for (int64_t k = 0; k < n; ++k) {
                batch_rows[k] = array_[base + k].second;
            }
            array_offsets->CopyRowElementRanges(batch_rows, n, batch_ranges);

            for (int64_t k = 0; k < n && hit_num < limit; ++k) {
                const auto& entry = array_[base + k];
                auto doc_offset = entry.second;

                // Get element range for this doc
                auto [first_elem, last_elem] = batch_ranges[k];

                // Collect all matching element indices for this doc
                matching_indices.clear();
                for (int64_t elem_id = first_elem;
                     elem_id < last_elem && hit_num < limit;
                     ++elem_id) {
                    if (elem_id >= element_size) {
                        continue;
                    }
                    if (is_skipped_by_cursor(
                            entry.first, elem_id - first_elem, cursor)) {
                        continue;
                    }
                    if (!element_bitset[elem_id]) {  // 0 means pass filter
                        matching_indices.push_back(
                            static_cast<int32_t>(elem_id - first_elem));
                        hit_num++;
                    }
                }

                // Only add doc if it has matching elements
                if (!matching_indices.empty()) {
                    doc_offsets.push_back(doc_offset);
                    element_indices.push_back(std::move(matching_indices));
                }
            }
        }

        bool has_more = more_hit_than_limit && hit_num >= limit;
        return {std::move(doc_offsets), std::move(element_indices), has_more};
    }

    bool
    is_skipped_by_cursor(
        const T& pk,
        int64_t element_offset,
        const std::optional<QueryIteratorCursor>& cursor) const {
        if (!cursor.has_value()) {
            return false;
        }
        auto last_pk = std::get_if<T>(&cursor->last_pk);
        return last_pk != nullptr && *last_pk == pk &&
               element_offset <= cursor->last_element_offset;
    }

    void
    check_search() const {
        AssertInfo(is_sealed,
                   "OffsetOrderedArray could not search before seal");
    }

 private:
    bool is_sealed = false;
    std::vector<std::pair<T, int32_t>> array_;
};

// VirtualPKOffsetMap is a zero-storage OffsetMap for external collections.
// Virtual PK = (truncated_segment_id << 32) | offset, so pk→offset is a
// simple bit-extract: offset = pk & 0xFFFFFFFF. No data structure needed.
// This replaces OffsetOrderedArray for external tables, saving ~17 GB for
// 1B rows (which would otherwise store 1B (pk, offset) pairs).
class VirtualPKOffsetMap : public OffsetMap {
 public:
    explicit VirtualPKOffsetMap(int64_t segment_id, int64_t num_rows)
        : truncated_segment_id_(segment_id & 0xFFFFFFFF),
          shifted_segment_id_(truncated_segment_id_ << 32),
          num_rows_(num_rows) {
    }

    bool
    contain(const PkType& pk) const override {
        int64_t vpk = std::get<int64_t>(pk);
        int64_t seg = static_cast<int64_t>(static_cast<uint64_t>(vpk) >> 32);
        if (seg != truncated_segment_id_)
            return false;
        int64_t offset = vpk & 0xFFFFFFFF;
        return offset >= 0 && offset < num_rows_;
    }

    std::vector<int64_t>
    find(const PkType& pk) const override {
        int64_t vpk = std::get<int64_t>(pk);
        int64_t seg = static_cast<int64_t>(static_cast<uint64_t>(vpk) >> 32);
        if (seg != truncated_segment_id_)
            return {};
        int64_t offset = vpk & 0xFFFFFFFF;
        if (offset >= 0 && offset < num_rows_)
            return {offset};
        return {};
    }

    void
    find_range(const PkType& pk,
               proto::plan::OpType op,
               BitsetTypeView& bitset,
               Condition condition) const override {
        int64_t target = std::get<int64_t>(pk);
        // Virtual PKs for this segment are [base, base+num_rows_)
        // where base = shifted_segment_id_
        int64_t base = shifted_segment_id_;
        int64_t lo = 0, hi = num_rows_;  // offset range

        // Convert target to offset space
        int64_t target_offset = target - base;

        switch (op) {
            case proto::plan::OpType::Equal: {
                if (target_offset >= 0 && target_offset < num_rows_) {
                    if (condition(target_offset)) {
                        bitset[target_offset] = true;
                    }
                }
                break;
            }
            case proto::plan::OpType::GreaterEqual:
                lo = std::max(int64_t(0), target_offset);
                for (int64_t i = lo; i < hi; i++) {
                    if (condition(i))
                        bitset[i] = true;
                }
                break;
            case proto::plan::OpType::GreaterThan:
                lo = std::max(int64_t(0), target_offset + 1);
                for (int64_t i = lo; i < hi; i++) {
                    if (condition(i))
                        bitset[i] = true;
                }
                break;
            case proto::plan::OpType::LessEqual:
                hi = std::min(num_rows_, target_offset + 1);
                for (int64_t i = 0; i < hi; i++) {
                    if (condition(i))
                        bitset[i] = true;
                }
                break;
            case proto::plan::OpType::LessThan:
                hi = std::min(num_rows_, target_offset);
                for (int64_t i = 0; i < hi; i++) {
                    if (condition(i))
                        bitset[i] = true;
                }
                break;
            default:
                ThrowInfo(ErrorCode::Unsupported,
                          fmt::format("unsupported op type {}", op));
        }
    }

    void
    insert(const PkType& pk, int64_t offset) override {
        // No-op: virtual PK mapping is computed, not stored
    }

    void
    seal() override {
        // No-op
    }

    bool
    empty() const override {
        return num_rows_ == 0;
    }

    std::pair<std::vector<OffsetMap::OffsetType>, bool>
    find_first_n(int64_t limit, const BitsetTypeView& bitset) const override {
        // Virtual PKs are sorted by offset (since pk = base | offset),
        // so iterating offsets 0..N-1 gives PK-sorted order.
        auto size = static_cast<int64_t>(bitset.size());
        int64_t cnt = size - bitset.count();
        if (limit == Unlimited || limit == NoLimit) {
            limit = cnt;
        }
        auto more_hit_than_limit = cnt > limit;
        limit = std::min(limit, cnt);
        std::vector<int64_t> seg_offsets;
        seg_offsets.reserve(limit);
        for (int64_t i = 0;
             i < size && static_cast<int64_t>(seg_offsets.size()) < limit;
             i++) {
            if (!bitset[i]) {
                seg_offsets.push_back(i);
            }
        }
        return {seg_offsets,
                more_hit_than_limit &&
                    static_cast<int64_t>(seg_offsets.size()) >= limit};
    }

    std::tuple<std::vector<int64_t>, std::vector<std::vector<int32_t>>, bool>
    find_first_n_element(
        int64_t limit,
        const BitsetTypeView& element_bitset,
        const IArrayOffsets* array_offsets,
        const std::optional<QueryIteratorCursor>& cursor) const override {
        // External tables don't support array fields, but implement
        // the interface for completeness.
        auto element_size = static_cast<int64_t>(element_bitset.size());
        if (limit == Unlimited || limit == NoLimit) {
            limit = element_size;
        }
        int64_t cnt = element_size - element_bitset.count();
        auto more_hit_than_limit = cnt > limit;
        limit = std::min(limit, cnt);

        std::vector<int64_t> doc_offsets;
        std::vector<std::vector<int32_t>> element_indices;
        int64_t hit_num = 0;

        // Docs are scanned sequentially, so resolve element ranges with one
        // contiguous CopyRowElementStarts call per window of kRowBatchSize
        // rows ([starts[i], starts[i+1]) is row (base + i)'s range) instead
        // of one virtual call per row.
        constexpr int64_t kRowBatchSize = 64;
        int32_t batch_starts[kRowBatchSize + 1];

        for (int64_t base = 0; base < num_rows_ && hit_num < limit;
             base += kRowBatchSize) {
            const int64_t n = std::min(kRowBatchSize, num_rows_ - base);
            array_offsets->CopyRowElementStarts(base, n, batch_starts);
            for (int64_t i = 0; i < n && hit_num < limit; i++) {
                const int64_t doc = base + i;
                const int32_t first_elem = batch_starts[i];
                const int32_t last_elem = batch_starts[i + 1];
                std::vector<int32_t> matching;
                for (int64_t e = first_elem; e < last_elem && hit_num < limit;
                     e++) {
                    if (e < element_size && !element_bitset[e] &&
                        !is_skipped_by_cursor(doc, e - first_elem, cursor)) {
                        matching.push_back(
                            static_cast<int32_t>(e - first_elem));
                        hit_num++;
                    }
                }
                if (!matching.empty()) {
                    doc_offsets.push_back(doc);
                    element_indices.push_back(std::move(matching));
                }
            }
        }
        return {std::move(doc_offsets),
                std::move(element_indices),
                more_hit_than_limit && hit_num >= limit};
    }

    void
    clear() override {
        // No-op
    }

    size_t
    memory_size() const override {
        return sizeof(VirtualPKOffsetMap);
    }

    bool
    is_skipped_by_cursor(
        int64_t doc,
        int64_t element_offset,
        const std::optional<QueryIteratorCursor>& cursor) const {
        if (!cursor.has_value()) {
            return false;
        }
        auto last_pk = std::get_if<int64_t>(&cursor->last_pk);
        return last_pk != nullptr && *last_pk == (shifted_segment_id_ | doc) &&
               element_offset <= cursor->last_element_offset;
    }

    bool
    is_zero_storage() const override {
        return true;
    }

 private:
    int64_t truncated_segment_id_;
    int64_t shifted_segment_id_;
    int64_t num_rows_;
};

class InsertRecordSealed {
 public:
    InsertRecordSealed(
        const Schema& schema,
        const int64_t size_per_chunk,
        const storage::MmapChunkDescriptorPtr mmap_descriptor = nullptr) {
        std::optional<FieldId> pk_field_id = schema.get_primary_field_id();
        // for sealed segment, only pk field is added.
        for (auto& field : schema) {
            auto field_id = field.first;
            if (pk_field_id.has_value() && pk_field_id.value() == field_id) {
                auto& field_meta = field.second;
                AssertInfo(!field_meta.is_nullable(),
                           "Primary key should not be nullable");
                switch (field_meta.get_data_type()) {
                    case DataType::INT64: {
                        pk2offset_ =
                            std::make_unique<OffsetOrderedArray<int64_t>>();
                        is_int64_pk_ = true;
                        break;
                    }
                    case DataType::VARCHAR: {
                        pk2offset_ =
                            std::make_unique<OffsetOrderedArray<std::string>>();
                        is_int64_pk_ = false;
                        break;
                    }
                    default: {
                        ThrowInfo(DataTypeInvalid,
                                  fmt::format("unsupported pk type",
                                              field_meta.get_data_type()));
                    }
                }
            }
        }
    }

    ~InsertRecordSealed() {
        if (estimated_memory_size_ > 0) {
            cachinglayer::Manager::GetInstance().RefundLoadedResource(
                {static_cast<int64_t>(estimated_memory_size_), 0});
            estimated_memory_size_ = 0;
        }
    }

    bool
    contain(const PkType& pk) const {
        return pk2offset_->contain(pk);
    }

    // Returns true when the owned pk2offset_ is zero-storage (currently only
    // VirtualPKOffsetMap). Sealed search_pks uses this to short-circuit the
    // sorted two-pointers path.
    bool
    pk2offset_is_zero_storage() const {
        return pk2offset_ != nullptr && pk2offset_->is_zero_storage();
    }

    std::vector<SegOffset>
    search_pk(const PkType& pk,
              Timestamp timestamp,
              bool include_same_ts = true) const {
        std::shared_lock lck(shared_mutex_);
        std::vector<SegOffset> res_offsets;
        auto offset_iter = pk2offset_->find(pk);
        auto timestamp_hit =
            include_same_ts ? [](const Timestamp& ts1,
                                 const Timestamp& ts2) { return ts1 <= ts2; }
                            : [](const Timestamp& ts1, const Timestamp& ts2) {
                                  return ts1 < ts2;
                              };
        for (auto offset : offset_iter) {
            if (timestamp_hit(timestamps_[offset], timestamp)) {
                res_offsets.emplace_back(offset);
            }
        }
        return res_offsets;
    }

    void
    search_pk_range(const PkType& pk,
                    proto::plan::OpType op,
                    BitsetTypeView& bitset) const {
        pk2offset_->find_range(
            pk, op, bitset, [](int64_t offset) { return true; });
    }

    void
    search_pk_range(const PkType& pk,
                    Timestamp timestamp,
                    proto::plan::OpType op,
                    BitsetTypeView& bitset) const {
        auto condition = [this, timestamp](int64_t offset) {
            return timestamps_[offset] <= timestamp;
        };
        pk2offset_->find_range(pk, op, bitset, condition);
    }

    void
    search_pk_binary_range(const PkType& lower_pk,
                           bool lower_inclusive,
                           const PkType& upper_pk,
                           bool upper_inclusive,
                           BitsetTypeView& bitset) const {
        auto lower_op = lower_inclusive ? proto::plan::OpType::GreaterEqual
                                        : proto::plan::OpType::GreaterThan;
        auto upper_op = upper_inclusive ? proto::plan::OpType::LessEqual
                                        : proto::plan::OpType::LessThan;

        BitsetType upper_result(bitset.size());
        auto upper_view = upper_result.view();

        // values >= lower_pk (or > lower_pk if not inclusive)
        pk2offset_->find_range(
            lower_pk, lower_op, bitset, [](int64_t offset) { return true; });

        // values <= upper_pk (or < upper_pk if not inclusive)
        pk2offset_->find_range(
            upper_pk, upper_op, upper_view, [](int64_t offset) {
                return true;
            });

        bitset &= upper_result;
    }

    void
    insert_pks(milvus::DataType data_type, ChunkedColumnInterface* data) {
        std::lock_guard lck(shared_mutex_);
        int64_t offset = 0;
        switch (data_type) {
            case DataType::INT64: {
                auto num_chunk = data->num_chunks();
                std::vector<int64_t> chunk_ids(num_chunk);
                std::iota(chunk_ids.begin(), chunk_ids.end(), 0);
                data->PrefetchChunks(nullptr, chunk_ids);
                for (int i = 0; i < num_chunk; ++i) {
                    auto pw = data->DataOfChunk(nullptr, i);
                    auto pks = reinterpret_cast<const int64_t*>(pw.get());
                    auto chunk_num_rows = data->chunk_row_nums(i);
                    for (int j = 0; j < chunk_num_rows; ++j) {
                        pk2offset_->insert(pks[j], offset++);
                    }
                }
                break;
            }
            case DataType::VARCHAR: {
                auto num_chunk = data->num_chunks();
                for (int i = 0; i < num_chunk; ++i) {
                    auto column =
                        reinterpret_cast<ChunkedVariableColumn<std::string>*>(
                            data);
                    auto pw = column->StringViews(nullptr, i);
                    auto pks = pw.get().first;
                    for (auto& pk : pks) {
                        pk2offset_->insert(std::string(pk), offset++);
                    }
                }
                break;
            }
            default: {
                ThrowInfo(DataTypeInvalid,
                          fmt::format("unsupported primary key data type",
                                      data_type));
            }
        }
    }

    // Build only the compressed offset->pk mapping (no pk2offset_ index).
    // Used by sorted segments where pk2offset_ is not needed.
    void
    build_offset2pk(milvus::DataType data_type, ChunkedColumnInterface* data) {
        if (data_type != DataType::INT64 || !is_int64_pk_) {
            return;
        }
        std::lock_guard lck(shared_mutex_);
        int64_t total_rows = data->NumRows();
        std::vector<int64_t> all_pks;
        all_pks.reserve(total_rows);

        auto num_chunk = data->num_chunks();
        // Prefetch all chunks so that subsequent sequential
        // DataOfChunk() calls hit the cache instead of blocking
        // on HIGH pool downloads one at a time.
        std::vector<int64_t> chunk_ids(num_chunk);
        std::iota(chunk_ids.begin(), chunk_ids.end(), 0);
        data->PrefetchChunks(nullptr, chunk_ids);

        for (int i = 0; i < num_chunk; ++i) {
            auto pw = data->DataOfChunk(nullptr, i);
            auto pks = reinterpret_cast<const int64_t*>(pw.get());
            auto chunk_num_rows = data->chunk_row_nums(i);
            for (int j = 0; j < chunk_num_rows; ++j) {
                all_pks.push_back(pks[j]);
            }
        }

        offset2pk_ = std::make_unique<CompressedInt64PkArray>();
        offset2pk_->build(all_pks.data(), all_pks.size());

        size_t mem = offset2pk_->memory_size();
        cachinglayer::Manager::GetInstance().ChargeLoadedResource(
            {static_cast<int64_t>(mem), 0});
        estimated_memory_size_ += mem;
    }

    bool
    empty_pks() const {
        std::shared_lock lck(shared_mutex_);
        return pk2offset_->empty();
    }

    void
    seal_pks() {
        std::lock_guard lck(shared_mutex_);
        pk2offset_->seal();
        // update estimated memory size to caching layer
        // offset2pk_ memory is charged separately in build_offset2pk()
        size_t total_size = pk2offset_->memory_size();
        cachinglayer::Manager::GetInstance().ChargeLoadedResource(
            {static_cast<int64_t>(total_size), 0});
        estimated_memory_size_ += total_size;
    }

    // Replace the pk2offset_ with a VirtualPKOffsetMap for external tables.
    // This avoids materializing the pk→offset mapping (saves ~17 GB for 1B rows).
    void
    set_virtual_pk_offset_map(int64_t segment_id, int64_t num_rows) {
        std::lock_guard lck(shared_mutex_);
        pk2offset_ = std::make_unique<VirtualPKOffsetMap>(segment_id, num_rows);
        is_int64_pk_ = true;
    }

    // Constant mode: all timestamps return the same value.
    // Used by external tables where all rows are always visible (ts=0).
    // Saves ~8 GB for 1B-row external tables.
    void
    init_timestamps_constant(int64_t num_rows, Timestamp value = 0) {
        std::lock_guard lck(shared_mutex_);
        timestamps_.InitConstant(num_rows, value);
        // TimestampIndex is not needed — external tables have no TTL or
        // time-travel. Callers that check timestamps_.empty() will see false
        // (total_size_ > 0), so the fast paths work correctly.
    }

    // Own mode: takes ownership of timestamp data (StorageV1 / multi-chunk)
    void
    init_timestamps_from_owned(std::vector<Timestamp> data,
                               TimestampIndex timestamp_index) {
        std::lock_guard lck(shared_mutex_);
        size_t count = data.size();
        timestamps_.InitFromOwnedData(std::move(data));
        timestamp_index_ = std::move(timestamp_index);
        size_t ts_data_size = count * sizeof(Timestamp);
        size_t ts_index_size = timestamp_index_.memory_size();
        cachinglayer::Manager::GetInstance().ChargeLoadedResource(
            {static_cast<int64_t>(ts_data_size + ts_index_size), 0});
        estimated_memory_size_ += ts_data_size + ts_index_size;
    }

    const TimestampData&
    timestamps() const {
        return timestamps_;
    }

    void
    clear() {
        timestamps_.clear();
        timestamp_index_ = TimestampIndex();
        if (pk2offset_) {
            pk2offset_->clear();
        }
        offset2pk_.reset();

        reserved = 0;
        if (estimated_memory_size_ > 0) {
            cachinglayer::Manager::GetInstance().RefundLoadedResource(
                {static_cast<int64_t>(estimated_memory_size_), 0});
            estimated_memory_size_ = 0;
        }
    }

    // Returns true if PK type is int64 and offset2pk is available
    bool
    has_int64_pk_index() const {
        return is_int64_pk_ && offset2pk_ != nullptr;
    }

    // Get PK by offset. Only valid when has_int64_pk_index() returns true.
    int64_t
    get_int64_pk_by_offset(int64_t offset) const {
        AssertInfo(
            has_int64_pk_index(),
            "get_int64_pk_by_offset requires int64 PK with offset2pk index");
        return offset2pk_->at(offset);
    }

    // Bulk get PKs by offsets. Only valid when has_int64_pk_index() returns true.
    void
    bulk_get_int64_pks_by_offsets(const int64_t* offsets,
                                  int64_t count,
                                  int64_t* output) const {
        AssertInfo(has_int64_pk_index(),
                   "bulk_get_int64_pks_by_offsets requires int64 PK with "
                   "offset2pk index");
        offset2pk_->bulk_at(offsets, count, output);
    }

 public:
    TimestampData timestamps_;
    std::atomic<int64_t> reserved = 0;
    // used for timestamps index of sealed segment
    TimestampIndex timestamp_index_;
    // pks to row offset
    std::unique_ptr<OffsetMap> pk2offset_;
    // offset to pk (compressed), only available for int64 PK
    std::unique_ptr<CompressedInt64PkArray> offset2pk_;
    // whether PK type is int64
    bool is_int64_pk_ = false;
    // estimated memory size of InsertRecord, only used for sealed segment
    int64_t estimated_memory_size_{0};

 protected:
    mutable std::shared_mutex shared_mutex_{};
};

class InsertRecordGrowing {
 public:
    InsertRecordGrowing(
        const Schema& schema,
        const int64_t size_per_chunk,
        const storage::MmapChunkDescriptorPtr mmap_descriptor = nullptr)
        : timestamps_(size_per_chunk), row_ids_(size_per_chunk) {
        std::optional<FieldId> pk_field_id = schema.get_primary_field_id();
        for (auto& field : schema) {
            auto field_id = field.first;
            auto& field_meta = field.second;
            if (pk_field_id.has_value() && pk_field_id.value() == field_id) {
                AssertInfo(!field_meta.is_nullable(),
                           "Primary key should not be nullable");
                switch (field_meta.get_data_type()) {
                    case DataType::INT64: {
                        pk2offset_ =
                            std::make_unique<OffsetOrderedMap<int64_t>>();
                        break;
                    }
                    case DataType::VARCHAR: {
                        pk2offset_ =
                            std::make_unique<OffsetOrderedMap<std::string>>();
                        break;
                    }
                    default: {
                        ThrowInfo(DataTypeInvalid,
                                  fmt::format("unsupported pk type: {}",
                                              field_meta.get_data_type()));
                    }
                }
            }
            append_field_meta(
                field_id, field_meta, size_per_chunk, mmap_descriptor);
        }
    }

    bool
    contain(const PkType& pk) const {
        return pk2offset_->contain(pk);
    }

    std::vector<SegOffset>
    search_pk(const PkType& pk,
              Timestamp timestamp,
              bool include_same_ts = true) const {
        std::shared_lock<std::shared_mutex> lck(shared_mutex_);
        std::vector<SegOffset> res_offsets;
        auto offset_iter = pk2offset_->find(pk);
        auto timestamp_hit =
            include_same_ts ? [](const Timestamp& ts1,
                                 const Timestamp& ts2) { return ts1 <= ts2; }
                            : [](const Timestamp& ts1, const Timestamp& ts2) {
                                  return ts1 < ts2;
                              };
        for (auto offset : offset_iter) {
            if (timestamp_hit(timestamps_[offset], timestamp)) {
                res_offsets.emplace_back(offset);
            }
        }
        return res_offsets;
    }

    void
    search_pk_range(const PkType& pk,
                    proto::plan::OpType op,
                    BitsetTypeView& bitset) const {
        pk2offset_->find_range(
            pk, op, bitset, [](int64_t offset) { return true; });
    }

    void
    search_pk_range(const PkType& pk,
                    Timestamp timestamp,
                    proto::plan::OpType op,
                    BitsetTypeView& bitset) const {
        auto condition = [this, timestamp](int64_t offset) {
            return timestamps_[offset] <= timestamp;
        };
        pk2offset_->find_range(pk, op, bitset, condition);
    }

    void
    insert_pks(const std::vector<FieldDataPtr>& field_datas) {
        std::lock_guard lck(shared_mutex_);
        int64_t offset = 0;
        for (auto& data : field_datas) {
            int64_t row_count = data->get_num_rows();
            auto data_type = data->get_data_type();
            switch (data_type) {
                case DataType::INT64: {
                    for (int i = 0; i < row_count; ++i) {
                        pk2offset_->insert(
                            *static_cast<const int64_t*>(data->RawValue(i)),
                            offset++);
                    }
                    break;
                }
                case DataType::VARCHAR: {
                    for (int i = 0; i < row_count; ++i) {
                        pk2offset_->insert(
                            *static_cast<const std::string*>(data->RawValue(i)),
                            offset++);
                    }
                    break;
                }
                default: {
                    ThrowInfo(DataTypeInvalid,
                              fmt::format("unsupported primary key data type",
                                          data_type));
                }
            }
        }
    }

    bool
    empty_pks() const {
        std::shared_lock lck(shared_mutex_);
        return pk2offset_->empty();
    }

    void
    seal_pks() {
        std::lock_guard lck(shared_mutex_);
        pk2offset_
            ->seal();  // will throw for growing map, consistent with previous behavior
    }

    const ConcurrentVector<Timestamp>&
    timestamps() const {
        return timestamps_;
    }

    const ConcurrentVector<int64_t>&
    row_ids() const {
        return row_ids_;
    }

    void
    clear() {
        timestamps_.clear();
        row_ids_.clear();
        timestamp_index_ = TimestampIndex();
        if (pk2offset_) {
            pk2offset_->clear();
        }
        reserved = 0;
        {
            std::unique_lock<std::shared_mutex> lck(field_map_mutex_);
            data_.clear();
            valid_data_.clear();
        }
        ack_responder_.clear();
    }

    void
    append_field_meta(
        FieldId field_id,
        const FieldMeta& field_meta,
        int64_t size_per_chunk,
        const storage::MmapChunkDescriptorPtr mmap_descriptor = nullptr) {
        if (field_meta.is_nullable()) {
            this->append_valid_data(field_id, size_per_chunk);
        }
        const milvus::storage::MmapConfig mmap_config =
            storage::MmapManager::GetInstance().GetMmapConfig();
        // todo: @cqy123456, scalar and vector mmap should be depend on GetScalarFieldEnableMmap()/ GetVectorFieldEnableMmap()
        storage::MmapChunkDescriptorPtr scalar_mmap_descriptor =
            mmap_config.GetEnableGrowingMmap() ? mmap_descriptor : nullptr;
        storage::MmapChunkDescriptorPtr vec_mmap_descriptor =
            mmap_config.GetEnableGrowingMmap() ? mmap_descriptor : nullptr;
        storage::MmapChunkDescriptorPtr dense_vec_mmap_descriptor = nullptr;
        // todo: @cqy123456, remove all condition and select of dense_vec_mmap_descriptor later
        {
            const auto& segcore_config =
                milvus::segcore::SegcoreConfig::default_config();
            auto enable_intermin_index =
                segcore_config.get_enable_interim_segment_index();
            auto dense_vec_intermin_index_type =
                segcore_config.get_dense_vector_intermin_index_type();
            auto enable_mmap_field_mmap =
                mmap_config.GetVectorFieldEnableMmap();
            bool enable_dense_vec_mmap =
                enable_intermin_index && enable_mmap_field_mmap &&
                (dense_vec_intermin_index_type ==
                 knowhere::IndexEnum::INDEX_FAISS_SCANN_DVR);
            dense_vec_mmap_descriptor =
                enable_dense_vec_mmap || mmap_config.GetEnableGrowingMmap()
                    ? mmap_descriptor
                    : nullptr;
        }
        if (field_meta.is_vector()) {
            if (field_meta.get_data_type() == DataType::VECTOR_FLOAT) {
                this->append_data<FloatVector>(field_id,
                                               field_meta.get_dim(),
                                               size_per_chunk,
                                               dense_vec_mmap_descriptor);
                return;
            } else if (field_meta.get_data_type() == DataType::VECTOR_BINARY) {
                this->append_data<BinaryVector>(field_id,
                                                field_meta.get_dim(),
                                                size_per_chunk,
                                                dense_vec_mmap_descriptor);
                return;
            } else if (field_meta.get_data_type() == DataType::VECTOR_FLOAT16) {
                this->append_data<Float16Vector>(field_id,
                                                 field_meta.get_dim(),
                                                 size_per_chunk,
                                                 dense_vec_mmap_descriptor);
                return;
            } else if (field_meta.get_data_type() ==
                       DataType::VECTOR_BFLOAT16) {
                this->append_data<BFloat16Vector>(field_id,
                                                  field_meta.get_dim(),
                                                  size_per_chunk,
                                                  dense_vec_mmap_descriptor);
                return;
            } else if (field_meta.get_data_type() ==
                       DataType::VECTOR_SPARSE_U32_F32) {
                this->append_data<SparseFloatVector>(
                    field_id, size_per_chunk, vec_mmap_descriptor);
                return;
            } else if (field_meta.get_data_type() == DataType::VECTOR_INT8) {
                this->append_data<Int8Vector>(field_id,
                                              field_meta.get_dim(),
                                              size_per_chunk,
                                              dense_vec_mmap_descriptor);
                return;
            } else if (field_meta.get_data_type() == DataType::VECTOR_ARRAY) {
                this->append_data<VectorArray>(field_id,
                                               field_meta.get_dim(),
                                               size_per_chunk,
                                               dense_vec_mmap_descriptor);
                return;
            } else {
                ThrowInfo(DataTypeInvalid,
                          fmt::format("unsupported vector type",
                                      field_meta.get_data_type()));
            }
        }
        switch (field_meta.get_data_type()) {
            case DataType::BOOL: {
                this->append_data<bool>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            case DataType::INT8: {
                this->append_data<int8_t>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            case DataType::INT16: {
                this->append_data<int16_t>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            case DataType::INT32: {
                this->append_data<int32_t>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            case DataType::INT64: {
                this->append_data<int64_t>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            case DataType::FLOAT: {
                this->append_data<float>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            case DataType::DOUBLE: {
                this->append_data<double>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            case DataType::TIMESTAMPTZ: {
                this->append_data<int64_t>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            case DataType::STRING:
            case DataType::VARCHAR:
            case DataType::TEXT: {
                this->append_data<std::string>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            case DataType::JSON: {
                this->append_data<Json>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            case DataType::ARRAY: {
                this->append_data<Array>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            case DataType::GEOMETRY: {
                this->append_data<std::string>(
                    field_id, size_per_chunk, scalar_mmap_descriptor);
                return;
            }
            default: {
                ThrowInfo(DataTypeInvalid,
                          fmt::format("unsupported scalar type",
                                      field_meta.get_data_type()));
            }
        }
    }

    void
    insert_pk(const PkType& pk, int64_t offset) {
        std::lock_guard lck(shared_mutex_);
        pk2offset_->insert(pk, offset);
    }

    // get data without knowing the type
    VectorBase*
    get_data_base(FieldId field_id) const {
        // Guard the field map against concurrent structural modification
        // (e.g. append_field_meta during schema evolution rehashing data_).
        std::shared_lock<std::shared_mutex> lck(field_map_mutex_);
        return get_data_base_unlocked(field_id);
    }

    // get field data in given type, const version
    template <typename Type>
    const ConcurrentVector<Type>*
    get_data(FieldId field_id) const {
        auto base_ptr = get_data_base(field_id);
        auto ptr = dynamic_cast<const ConcurrentVector<Type>*>(base_ptr);
        Assert(ptr);
        return ptr;
    }

    // get field data in given type, non-const version
    template <typename Type>
    ConcurrentVector<Type>*
    get_data(FieldId field_id) {
        auto base_ptr = get_data_base(field_id);
        auto ptr = dynamic_cast<ConcurrentVector<Type>*>(base_ptr);
        Assert(ptr);
        return ptr;
    }

    ThreadSafeValidDataPtr
    get_valid_data(FieldId field_id) const {
        // Guard the field map against concurrent structural modification
        // (e.g. append_field_meta during schema evolution rehashing valid_data_).
        std::shared_lock<std::shared_mutex> lck(field_map_mutex_);
        return get_valid_data_unlocked(field_id);
    }

    bool
    is_data_exist(FieldId field_id) const {
        std::shared_lock<std::shared_mutex> lck(field_map_mutex_);
        return data_.find(field_id) != data_.end();
    }

    // Field ids that have materialized columns. The exact set the flush
    // layout must be trimmed to: non-materialized function outputs are
    // absent here.
    std::vector<int64_t>
    get_data_field_ids() const {
        std::shared_lock<std::shared_mutex> lck(field_map_mutex_);
        std::vector<int64_t> ids;
        ids.reserve(data_.size());
        for (const auto& [field_id, entry] : data_) {
            // An allocated-but-empty column (e.g. a function output the
            // replayed older-era inserts never filled) is not materialized.
            if (!entry || entry->empty()) {
                continue;
            }
            ids.push_back(field_id.get());
        }
        return ids;
    }

    bool
    is_valid_data_exist(FieldId field_id) const {
        std::shared_lock<std::shared_mutex> lck(field_map_mutex_);
        return is_valid_data_exist_unlocked(field_id);
    }

    SpanBase
    get_span_base(FieldId field_id, int64_t chunk_id) const {
        // Take the shared lock once for the whole span resolution instead of
        // re-locking inside each of get_data_base / is_valid_data_exist /
        // get_valid_data. This both removes the redundant lock churn on this
        // hot path and gives a consistent snapshot of the field map across the
        // three lookups.
        std::shared_lock<std::shared_mutex> lck(field_map_mutex_);
        auto data = get_data_base_unlocked(field_id);
        if (is_valid_data_exist_unlocked(field_id)) {
            auto size = data->get_chunk_size(chunk_id);
            auto element_offset = data->get_element_offset(chunk_id);
            return SpanBase(data->get_chunk_data(chunk_id),
                            get_valid_data_unlocked(field_id)->get_chunk_data(
                                element_offset),
                            size,
                            data->get_element_size());
        }
        return data->get_span_base(chunk_id);
    }

    // append a column of scalar type
    void
    append_valid_data(FieldId field_id, int64_t size_per_chunk) {
        auto valid_data = std::make_shared<ThreadSafeValidData>(size_per_chunk);
        std::unique_lock<std::shared_mutex> lck(field_map_mutex_);
        valid_data_.emplace(field_id, std::move(valid_data));
    }

    // append a column of vector type
    template <typename VectorType>
    void
    append_data(FieldId field_id,
                int64_t dim,
                int64_t size_per_chunk,
                const storage::MmapChunkDescriptorPtr mmap_descriptor) {
        static_assert(std::is_base_of_v<VectorTrait, VectorType>);
        bool use_mapping_storage = is_valid_data_exist(field_id);
        // Resolve valid_data and build the column before taking the write lock
        // so the shared-lock readers above are not nested inside the unique lock.
        auto column = std::make_unique<ConcurrentVector<VectorType>>(
            dim,
            size_per_chunk,
            mmap_descriptor,
            use_mapping_storage ? get_valid_data(field_id) : nullptr,
            use_mapping_storage);
        std::unique_lock<std::shared_mutex> lck(field_map_mutex_);
        data_.emplace(field_id, std::move(column));
    }

    // append a column of scalar or sparse float vector type
    template <typename Type>
    void
    append_data(FieldId field_id,
                int64_t size_per_chunk,
                const storage::MmapChunkDescriptorPtr mmap_descriptor) {
        static_assert(IsScalar<Type> || IsSparse<Type>);
        bool use_mapping_storage = is_valid_data_exist(field_id);
        // Resolve valid_data and build the column before taking the write lock
        // so the shared-lock readers above are not nested inside the unique lock.
        std::unique_ptr<ConcurrentVector<Type>> column;
        if constexpr (IsSparse<Type>) {
            column = std::make_unique<ConcurrentVector<Type>>(
                size_per_chunk,
                mmap_descriptor,
                use_mapping_storage ? get_valid_data(field_id) : nullptr,
                use_mapping_storage);
        } else {
            column = std::make_unique<ConcurrentVector<Type>>(
                size_per_chunk,
                mmap_descriptor,
                use_mapping_storage ? get_valid_data(field_id) : nullptr);
        }
        std::unique_lock<std::shared_mutex> lck(field_map_mutex_);
        data_.emplace(field_id, std::move(column));
    }

    void
    drop_field_data(FieldId field_id) {
        std::unique_lock<std::shared_mutex> lck(field_map_mutex_);
        data_.erase(field_id);
        valid_data_.erase(field_id);
    }

    int64_t
    row_count() const {
        return ack_responder_.GetAck();
    }

    bool
    empty() const {
        return pk2offset_->empty();
    }

 public:
    ConcurrentVector<Timestamp> timestamps_;
    ConcurrentVector<int64_t> row_ids_;
    std::atomic<int64_t> reserved = 0;
    TimestampIndex timestamp_index_;
    std::unique_ptr<OffsetMap> pk2offset_;

    // used for preInsert of growing segment
    AckResponder ack_responder_;

 private:
    // Unlocked field-map lookups. Callers MUST already hold field_map_mutex_
    // (at least shared). These let multi-lookup hot paths such as
    // get_span_base resolve under a single lock acquisition.
    VectorBase*
    get_data_base_unlocked(FieldId field_id) const {
        AssertInfo(data_.find(field_id) != data_.end(),
                   "Cannot find field_data with field_id: " +
                       std::to_string(field_id.get()));
        AssertInfo(data_.at(field_id) != nullptr,
                   "data_ at i is null" + std::to_string(field_id.get()));
        return data_.at(field_id).get();
    }

    ThreadSafeValidDataPtr
    get_valid_data_unlocked(FieldId field_id) const {
        AssertInfo(valid_data_.find(field_id) != valid_data_.end(),
                   "Cannot find valid_data with field_id: " +
                       std::to_string(field_id.get()));
        AssertInfo(valid_data_.at(field_id) != nullptr,
                   "valid_data_ at i is null" + std::to_string(field_id.get()));
        return valid_data_.at(field_id);
    }

    bool
    is_valid_data_exist_unlocked(FieldId field_id) const {
        return valid_data_.find(field_id) != valid_data_.end();
    }

    std::unordered_map<FieldId, std::unique_ptr<VectorBase>> data_{};
    std::unordered_map<FieldId, ThreadSafeValidDataPtr> valid_data_{};
    mutable std::shared_mutex shared_mutex_{};
    // Protects the structure of data_ / valid_data_ against concurrent
    // rehash: structural writes (append_*/drop/clear during schema evolution)
    // take it unique, lookups (get_data_base/get_valid_data/...) take it shared.
    // Kept separate from shared_mutex_ so frequent reads do not contend with
    // pk inserts.
    //
    // PRECONDITION: this lock only protects the map structure, NOT element
    // lifetime. The raw VectorBase* returned by get_data_base() escapes the
    // shared lock, so callers rely on no concurrent erase()/clear() of the
    // entry they hold (append-only mutation is safe). Today that holds because
    // drop_field_data() has no callers and clear() is never called
    // concurrently with reads; whoever adds drop-field support (e.g. schema
    // evolution dropping fields) must also fence element lifetime, this lock
    // alone will not save the escaped pointers.
    mutable std::shared_mutex field_map_mutex_{};
};

// Keep the original template API via alias
template <bool is_sealed>
using InsertRecord =
    std::conditional_t<is_sealed, InsertRecordSealed, InsertRecordGrowing>;

}  // namespace milvus::segcore
