segcore impl doen

Signed-off-by: SpadeA <[email protected]>

Author: SpadeA-Tang

internal/core/src/common/ArrayOffsets.cpp

@@ -0,0 +1,259 @@
+// Licensed to the LF AI & Data foundation 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 "ArrayOffsets.h"
+#include "segcore/SegmentInterface.h"
+#include "log/Log.h"
+#include "common/EasyAssert.h"
+
+namespace milvus {
+
+namespace {
+ArrayOffsetsSealed
+BuildArrayOffsetsSealedFromColumn(
+    const segcore::SegmentInternalInterface* segment,
+    const FieldMeta& field_meta,
+    int64_t row_count) {
+    FieldId field_id = field_meta.get_id();
+    auto data_type = field_meta.get_data_type();
+
+    ArrayOffsetsSealed result;
+    result.doc_count = row_count;
+    result.doc_to_element_range_.resize(row_count);
+
+    auto temp_op_ctx = std::make_unique<OpContext>();
+    auto op_ctx_ptr = temp_op_ctx.get();
+
+    int64_t num_chunks = segment->num_chunk(field_id);
+    int64_t current_doc_id = 0;
+
+    if (data_type == DataType::VECTOR_ARRAY) {
+        for (int64_t chunk_id = 0; chunk_id < num_chunks; ++chunk_id) {
+            auto pin_wrapper = segment->chunk_view<VectorArrayView>(
+                op_ctx_ptr, field_id, chunk_id);
+            const auto& [vector_array_views, valid_flags] = pin_wrapper.get();
+
+            for (size_t i = 0; i < vector_array_views.size(); ++i) {
+                int64_t array_len = 0;
+                if (valid_flags.empty() || valid_flags[i]) {
+                    array_len = vector_array_views[i].length();
+                }
+
+                // Record the start position for this doc
+                int64_t elem_start = result.element_info.size();
+
+                // Add (doc_id, element_index) for each element
+                for (int32_t j = 0; j < array_len; ++j) {
+                    result.element_info.emplace_back(current_doc_id, j);
+                }
+
+                int64_t elem_end = result.element_info.size();
+                result.doc_to_element_range_[current_doc_id] = {elem_start,
+                                                                elem_end};
+
+                current_doc_id++;
+            }
+        }
+    } else {
+        for (int64_t chunk_id = 0; chunk_id < num_chunks; ++chunk_id) {
+            auto pin_wrapper =
+                segment->chunk_view<ArrayView>(op_ctx_ptr, field_id, chunk_id);
+            const auto& [array_views, valid_flags] = pin_wrapper.get();
+
+            for (size_t i = 0; i < array_views.size(); ++i) {
+                int64_t array_len = 0;
+                if (valid_flags.empty() || valid_flags[i]) {
+                    array_len = array_views[i].length();
+                }
+
+                // Record the start position for this doc
+                int64_t elem_start = result.element_info.size();
+
+                // Add (doc_id, element_index) for each element
+                for (int32_t j = 0; j < array_len; ++j) {
+                    result.element_info.emplace_back(current_doc_id, j);
+                }
+
+                int64_t elem_end = result.element_info.size();
+                result.doc_to_element_range_[current_doc_id] = {elem_start,
+                                                                elem_end};
+
+                current_doc_id++;
+            }
+        }
+    }
+
+    AssertInfo(current_doc_id == row_count,
+               "Document count mismatch: expected {}, got {}",
+               row_count,
+               current_doc_id);
+
+    return result;
+}
+
+}  // anonymous namespace
+
+std::pair<int64_t, int64_t>
+ArrayOffsetsSealed::ElementIDToDoc(int64_t elem_id) const {
+    assert(elem_id >= 0 && elem_id < GetTotalElementCount());
+
+    const auto& [doc_id, elem_idx] = element_info[elem_id];
+    return {doc_id, elem_idx};
+}
+
+std::pair<int64_t, int64_t>
+ArrayOffsetsSealed::DocIDToElementID(int64_t doc_id) const {
+    assert(doc_id >= 0 && doc_id <= doc_count);
+
+    // If doc_id equals doc_count, return the total element count
+    // This represents "past the last document"
+    if (doc_id >= doc_count) {
+        return {GetTotalElementCount(), GetTotalElementCount()};
+    }
+
+    return doc_to_element_range_[doc_id];
+}
+
+std::pair<int64_t, int64_t>
+ArrayOffsetsGrowing::ElementIDToDoc(int64_t elem_id) const {
+    std::shared_lock lock(mutex_);
+    assert(elem_id >= 0 &&
+           elem_id < static_cast<int64_t>(element_info_.size()));
+    const auto& [doc_id, elem_idx] = element_info_[elem_id];
+    return {doc_id, elem_idx};
+}
+
+std::pair<int64_t, int64_t>
+ArrayOffsetsGrowing::DocIDToElementID(int64_t doc_id) const {
+    std::shared_lock lock(mutex_);
+    assert(doc_id >= 0 && doc_id <= committed_doc_count_);
+
+    // If doc_id equals or exceeds committed_doc_count_, return the total element count
+    // This represents "past the last document"
+    if (doc_id >= committed_doc_count_) {
+        return {static_cast<int64_t>(element_info_.size()),
+                static_cast<int64_t>(element_info_.size())};
+    }
+
+    return doc_to_element_range_[doc_id];
+}
+
+ArrayOffsetsSealed
+ArrayOffsetsSealed::BuildFromSegment(const void* segment,
+                                     const FieldMeta& field_meta) {
+    auto seg = static_cast<const segcore::SegmentInternalInterface*>(segment);
+
+    int64_t row_count = seg->get_row_count();
+    if (row_count == 0) {
+        LOG_INFO(
+            "ArrayOffsetsSealed::BuildFromSegment: empty segment for struct "
+            "'{}'",
+            field_meta.get_name().get());
+        ArrayOffsetsSealed result;
+        result.doc_count = 0;
+        return result;
+    }
+
+    ArrayOffsetsSealed result =
+        BuildArrayOffsetsSealedFromColumn(seg, field_meta, row_count);
+
+    int64_t total_elements = result.GetTotalElementCount();
+
+    LOG_INFO(
+        "ArrayOffsetsSealed::BuildFromSegment: struct_name='{}', "
+        "field_id={}, row_count={}, total_elements={}",
+        field_meta.get_name().get(),
+        field_meta.get_id().get(),
+        row_count,
+        total_elements);
+
+    return result;
+}
+
+void
+ArrayOffsetsGrowing::Insert(int64_t doc_id_start,
+                            const int32_t* array_lengths,
+                            int64_t count) {
+    std::unique_lock lock(mutex_);
+
+    for (int64_t i = 0; i < count; ++i) {
+        int64_t doc_id = doc_id_start + i;
+        int32_t array_len = array_lengths[i];
+
+        if (doc_id == committed_doc_count_) {
+            // Record the start position for this doc
+            int64_t elem_start = element_info_.size();
+
+            for (int32_t j = 0; j < array_len; ++j) {
+                element_info_.emplace_back(static_cast<int32_t>(doc_id), j);
+            }
+
+            // Record the end position for this doc
+            int64_t elem_end = element_info_.size();
+
+            // Ensure doc_to_element_range_ is large enough
+            if (static_cast<int64_t>(doc_to_element_range_.size()) <= doc_id) {
+                doc_to_element_range_.resize(doc_id + 1);
+            }
+            doc_to_element_range_[doc_id] = {elem_start, elem_end};
+
+            committed_doc_count_++;
+
+            // Try to drain pending documents
+            DrainPendingDocs();
+        } else {
+            // Cache this document for later
+            pending_docs_[doc_id] = {doc_id, array_len};
+        }
+    }
+}
+
+void
+ArrayOffsetsGrowing::DrainPendingDocs() {
+    while (true) {
+        auto it = pending_docs_.find(committed_doc_count_);
+        if (it == pending_docs_.end()) {
+            break;
+        }
+
+        // Commit this pending document
+        const auto& pending = it->second;
+
+        // Record the start position for this doc
+        int64_t elem_start = element_info_.size();
+
+        for (int32_t j = 0; j < pending.array_len; ++j) {
+            element_info_.emplace_back(static_cast<int32_t>(pending.doc_id), j);
+        }
+
+        // Record the end position for this doc
+        int64_t elem_end = element_info_.size();
+
+        // Ensure doc_to_element_range_ is large enough
+        if (static_cast<int64_t>(doc_to_element_range_.size()) <=
+            pending.doc_id) {
+            doc_to_element_range_.resize(pending.doc_id + 1);
+        }
+        doc_to_element_range_[pending.doc_id] = {elem_start, elem_end};
+
+        committed_doc_count_++;
+
+        // Remove from pending
+        pending_docs_.erase(it);
+    }
+}
+
+}  // namespace milvus