diff --git a/src/datacell/multi_vector_datacell.h b/src/datacell/multi_vector_datacell.h index 189b386886..80cf608a99 100644 --- a/src/datacell/multi_vector_datacell.h +++ b/src/datacell/multi_vector_datacell.h @@ -18,6 +18,7 @@ #include "flatten_interface.h" #include "io/common/basic_io.h" #include "io/memory_block_io/memory_block_io.h" +#include "layout/variable_record_layout.h" #include "quantization/multi_vector_computer.h" #include "vsag/dataset.h" @@ -114,12 +115,9 @@ class MultiVectorDataCell : public FlattenInterface { private: std::shared_ptr> quantizer_{nullptr}; - std::shared_ptr> io_{nullptr}; Allocator* const allocator_{nullptr}; - std::shared_ptr offset_io_{nullptr}; - uint64_t current_offset_{0}; - std::mutex current_offset_mutex_; + VariableRecordLayout layout_{}; uint32_t multi_vector_dim_{0}; MetricType metric_{MetricType::METRIC_TYPE_L2SQR}; diff --git a/src/datacell/multi_vector_datacell.inl b/src/datacell/multi_vector_datacell.inl index b5fea2f1a5..d609ecba07 100644 --- a/src/datacell/multi_vector_datacell.inl +++ b/src/datacell/multi_vector_datacell.inl @@ -17,6 +17,7 @@ #include #include +#include #include "common.h" #include "multi_vector_datacell.h" @@ -25,6 +26,19 @@ namespace vsag { +inline uint64_t +GetMultiVectorCodeSize(uint32_t token_count, uint32_t dimension) { + constexpr uint64_t header_size = sizeof(uint32_t); + constexpr uint64_t value_size = sizeof(float); + const uint64_t values_per_token = static_cast(dimension) * value_size; + if (values_per_token != 0 && + token_count > (std::numeric_limits::max() - header_size) / values_per_token) { + throw VsagException(ErrorType::INVALID_ARGUMENT, + "MultiVectorDataCell: record size overflow"); + } + return header_size + static_cast(token_count) * values_per_token; +} + template MultiVectorDataCell::MultiVectorDataCell( const QuantizerParamPtr& quantization_param, @@ -36,9 +50,12 @@ MultiVectorDataCell::MultiVectorDataCell( this->quantizer_ = std::make_shared(quantization_param, common_param); this->backend_ = QuantizerDistanceBackend::Get(static_cast(*this->quantizer_)); - this->io_ = std::make_shared(io_param, common_param); - this->offset_io_ = + auto io = std::make_shared(io_param, common_param); + auto offset_io = std::make_shared(Options::Instance().block_size_limit(), allocator_); + layout_.SetIO(std::move(offset_io), std::move(io)); + layout_.SetLocationPolicy( + HeaderLengthLocationPolicy{static_cast(multi_vector_dim_) * sizeof(float)}); this->max_capacity_ = 0; this->code_size_ = 0; } @@ -68,23 +85,16 @@ MultiVectorDataCell::InsertVector(const void* vector, InnerId } } - const uint64_t vector_bytes = static_cast(multi_vector->len_) * - static_cast(multi_vector_dim_) * sizeof(float); - const uint64_t code_size = sizeof(uint32_t) + vector_bytes; + const uint64_t code_size = GetMultiVectorCodeSize(multi_vector->len_, multi_vector_dim_); + const uint64_t vector_bytes = code_size - sizeof(uint32_t); ByteBuffer codes(code_size, allocator_); std::memcpy(codes.data, &multi_vector->len_, sizeof(uint32_t)); std::memcpy(codes.data + sizeof(uint32_t), multi_vector->vectors_, vector_bytes); - uint64_t old_offset = 0; { - std::lock_guard lock(current_offset_mutex_); - old_offset = current_offset_; - current_offset_ += code_size; + std::lock_guard lock(mutex_); + layout_.Write(idx, codes.data, code_size); } - offset_io_->Write(reinterpret_cast(&old_offset), - sizeof(old_offset), - static_cast(idx) * sizeof(old_offset)); - io_->Write(codes.data, code_size, old_offset); } template @@ -116,7 +126,7 @@ MultiVectorDataCell::Resize(InnerIdType new_capacity) { if (new_capacity <= this->max_capacity_) { return; } - this->offset_io_->Resize(static_cast(new_capacity) * sizeof(uint64_t)); + layout_.ResizeLocations(new_capacity); this->max_capacity_ = new_capacity; } @@ -135,14 +145,24 @@ MultiVectorDataCell::GetMetricType() { template const uint8_t* MultiVectorDataCell::GetCodesById(InnerIdType id, bool& need_release) const { - uint64_t offset = 0; - offset_io_->Read(sizeof(offset), static_cast(id) * sizeof(offset), (uint8_t*)&offset); + const uint64_t offset = layout_.ReadLocation(id); uint32_t len = 0; - io_->Read(sizeof(len), offset, (uint8_t*)&len); - uint64_t read_size = - sizeof(uint32_t) + static_cast(len) * multi_vector_dim_ * sizeof(float); + if (not layout_.Payload().Read(offset, sizeof(len), reinterpret_cast(&len))) { + throw VsagException(ErrorType::READ_ERROR, + "MultiVectorDataCell: failed to read token count"); + } + const uint64_t read_size = GetMultiVectorCodeSize(len, multi_vector_dim_); + const uint64_t payload_size = layout_.Payload().GetByteSize(); + if (offset > payload_size || read_size > payload_size - offset) { + throw VsagException(ErrorType::READ_ERROR, + "MultiVectorDataCell: token data range exceeds payload"); + } auto* codes = static_cast(allocator_->Allocate(read_size)); - io_->Read(read_size, offset, codes); + if (not layout_.Payload().Read(offset, read_size, codes)) { + allocator_->Deallocate(codes); + throw VsagException(ErrorType::READ_ERROR, + "MultiVectorDataCell: failed to read token data"); + } need_release = true; return codes; } @@ -164,9 +184,9 @@ void MultiVectorDataCell::Serialize(StreamWriter& writer) { FlattenInterface::Serialize(writer); StreamWriter::WriteObj(writer, multi_vector_dim_); - StreamWriter::WriteObj(writer, current_offset_); - this->offset_io_->Serialize(writer); - this->io_->Serialize(writer); + StreamWriter::WriteObj(writer, layout_.GetNextOffset()); + layout_.Locations().Serialize(writer); + layout_.Payload().Serialize(writer); this->quantizer_->Serialize(writer); } @@ -175,9 +195,11 @@ void MultiVectorDataCell::Deserialize(lvalue_or_rvalue reader) { FlattenInterface::Deserialize(reader); StreamReader::ReadObj(reader, multi_vector_dim_); - StreamReader::ReadObj(reader, current_offset_); - this->offset_io_->Deserialize(reader); - this->io_->Deserialize(reader); + uint64_t current_offset = 0; + StreamReader::ReadObj(reader, current_offset); + layout_.SetNextOffset(current_offset); + layout_.Locations().Deserialize(reader); + layout_.Payload().Deserialize(reader); this->quantizer_->Deserialize(reader); this->backend_ = QuantizerDistanceBackend::Get(static_cast(*this->quantizer_)); @@ -213,19 +235,16 @@ MultiVectorDataCell::Query(float* result_dists, // Step 1: Read all offsets (offset_io_ is MemoryBlockIO, in-memory, fast) std::vector offsets(id_count); for (InnerIdType i = 0; i < id_count; ++i) { - bool ok = offset_io_->Read(sizeof(uint64_t), - static_cast(idx[i]) * sizeof(uint64_t), - reinterpret_cast(&offsets[i])); - CHECK_ARGUMENT(ok, "MultiVectorDataCell: failed to read offset"); + offsets[i] = layout_.ReadLocation(idx[i]); } // Step 2: Batch read all token counts via MultiRead (async IO) std::vector lens(id_count); std::vector len_sizes(id_count, sizeof(uint32_t)); - if (!this->io_->MultiRead(reinterpret_cast(lens.data()), - len_sizes.data(), - offsets.data(), - static_cast(id_count))) { + if (!layout_.Payload().MultiRead(offsets.data(), + len_sizes.data(), + static_cast(id_count), + reinterpret_cast(lens.data()))) { throw VsagException(ErrorType::READ_ERROR, "MultiVectorDataCell: failed to read token counts"); } @@ -234,13 +253,21 @@ MultiVectorDataCell::Query(float* result_dists, std::vector data_sizes(id_count); uint64_t total_size = 0; for (InnerIdType i = 0; i < id_count; ++i) { - data_sizes[i] = - sizeof(uint32_t) + static_cast(lens[i]) * multi_vector_dim_ * sizeof(float); + data_sizes[i] = GetMultiVectorCodeSize(lens[i], multi_vector_dim_); + const uint64_t payload_size = layout_.Payload().GetByteSize(); + if (offsets[i] > payload_size || data_sizes[i] > payload_size - offsets[i]) { + throw VsagException(ErrorType::READ_ERROR, + "MultiVectorDataCell: token data range exceeds payload"); + } + if (data_sizes[i] > std::numeric_limits::max() - total_size) { + throw VsagException(ErrorType::INVALID_ARGUMENT, + "MultiVectorDataCell: batch record size overflow"); + } total_size += data_sizes[i]; } ByteBuffer all_codes(total_size, this->allocator_); - if (!this->io_->MultiRead( - all_codes.data, data_sizes.data(), offsets.data(), static_cast(id_count))) { + if (!layout_.Payload().MultiRead( + offsets.data(), data_sizes.data(), static_cast(id_count), all_codes.data)) { throw VsagException(ErrorType::READ_ERROR, "MultiVectorDataCell: failed to read token data"); } @@ -261,10 +288,7 @@ template uint64_t MultiVectorDataCell::GetMemoryUsage() const { uint64_t memory = sizeof(MultiVectorDataCell); - memory += this->offset_io_->size_; - if (IOTmpl::InMemory) { - memory += this->io_->GetMemoryUsage(); - } + memory += layout_.GetMemoryUsage(); memory += sizeof(QuantTmpl); return memory; } diff --git a/src/datacell/sparse_vector_datacell.h b/src/datacell/sparse_vector_datacell.h index 166e695439..b2d2738190 100644 --- a/src/datacell/sparse_vector_datacell.h +++ b/src/datacell/sparse_vector_datacell.h @@ -22,6 +22,7 @@ #include "inner_string_params.h" #include "io/common/basic_io.h" #include "io/memory_block_io/memory_block_io.h" +#include "layout/variable_record_layout.h" #include "quantization/sparse_quantization/sparse_quantizer.h" #include "vsag/dataset.h" @@ -86,11 +87,11 @@ class SparseVectorDataCell : public FlattenInterface { if (new_capacity <= this->max_capacity_) { return; } - std::scoped_lock lock(mutex_, current_offset_mutex_); - uint64_t io_size = - static_cast(new_capacity - total_count_) * max_code_size_ + current_offset_; - this->io_->Resize(io_size); - this->offset_io_->Resize(static_cast(new_capacity) * sizeof(DocLocation)); + std::lock_guard lock(mutex_); + uint64_t io_size = static_cast(new_capacity - total_count_) * max_code_size_ + + layout_.GetNextOffset(); + layout_.ReservePayload(io_size); + layout_.ResizeLocations(new_capacity); this->max_capacity_ = new_capacity; } @@ -149,12 +150,12 @@ class SparseVectorDataCell : public FlattenInterface { inline void SetIO(std::shared_ptr> io) { - this->io_ = io; + layout_.Payload().SetIO(std::move(io)); } void InitIO(const IOParamPtr& io_param) override { - this->io_->InitIO(io_param); + layout_.Payload().InitIO(io_param); } uint64_t @@ -188,12 +189,7 @@ class SparseVectorDataCell : public FlattenInterface { // Packed so each entry is exactly 12 bytes on disk and in the offset_io_ // buffer. The unpacked layout would round sizeof up to 16 due to the // uint64 alignment requirement, wasting 33% of the offset table. -#pragma pack(push, 1) - struct DocLocation { - uint64_t offset{0}; - uint32_t size{0}; - }; -#pragma pack(pop) + using DocLocation = OffsetAndLengthLocationPolicy::Entry; static_assert(sizeof(DocLocation) == 12, "DocLocation must be 12 bytes on disk"); // Legacy on-disk layout: kept for backward-compatible deserialization of indexes @@ -217,14 +213,11 @@ class SparseVectorDataCell : public FlattenInterface { static constexpr uint32_t SERIALIZE_FORMAT_VERSION_V2 = 2; std::shared_ptr> quantizer_{nullptr}; - std::shared_ptr> io_{nullptr}; QueryIOStrategy query_io_strategy_{QueryIOStrategy::MULTI_READ}; Allocator* const allocator_{nullptr}; - std::shared_ptr offset_io_{nullptr}; - uint64_t current_offset_{0}; + VariableRecordLayout layout_{}; uint64_t max_code_size_{0}; - std::mutex current_offset_mutex_; }; } // namespace vsag diff --git a/src/datacell/sparse_vector_datacell.inl b/src/datacell/sparse_vector_datacell.inl index 6e34b3866e..e7b26a706e 100644 --- a/src/datacell/sparse_vector_datacell.inl +++ b/src/datacell/sparse_vector_datacell.inl @@ -39,18 +39,14 @@ SparseVectorDataCell::query(float* result_dists, const auto load_location = [this](InnerIdType id) { DocLocation location{}; - const bool read_ok = offset_io_->Read(sizeof(location), - static_cast(id) * sizeof(location), - reinterpret_cast(&location)); - CHECK_ARGUMENT(read_ok, "SparseVectorDataCell failed to read document location"); - return location; + return layout_.ReadLocation(id); }; std::shared_lock lock(mutex_); const auto compute_direct = [&](const DocLocation& location, InnerIdType result_index) { bool need_release = false; - const auto* codes = io_->Read(location.size, location.offset, need_release); + const auto* codes = layout_.Read(location, need_release); if (codes == nullptr) { throw VsagException(ErrorType::READ_ERROR, "SparseVectorDataCell failed to read vector codes"); @@ -114,7 +110,7 @@ SparseVectorDataCell::query(float* result_dists, ranges.reserve(id_count); for (uint64_t i = 0; i < static_cast(id_count); ++i) { const auto& location = locations[i].location; - const uint64_t location_end = location.offset + location.size; + const uint64_t location_end = location.offset + location.length; if (not ranges.empty()) { auto& range = ranges.back(); const uint64_t range_end = range.offset + range.size; @@ -127,7 +123,7 @@ SparseVectorDataCell::query(float* result_dists, continue; } } - ranges.push_back({location.offset, location.size, i, i}); + ranges.push_back({location.offset, location.length, i, i}); } const uint64_t range_count = ranges.size(); @@ -143,7 +139,8 @@ SparseVectorDataCell::query(float* result_dists, } Vector scratch(scratch_size, query_allocator); - if (not io_->MultiRead(scratch.data(), read_sizes.data(), read_offsets.data(), range_count)) { + if (not layout_.Payload().MultiRead( + read_offsets.data(), read_sizes.data(), range_count, scratch.data())) { throw VsagException(ErrorType::READ_ERROR, "SparseVectorDataCell failed to read vector-code batch"); } @@ -176,13 +173,15 @@ SparseVectorDataCell::Deserialize(lvalue_or_rvalueio_->Deserialize(reader); - this->offset_io_->Deserialize(reader); + uint64_t current_offset = 0; + StreamReader::ReadObj(reader, current_offset); + layout_.SetNextOffset(current_offset); + layout_.Payload().Deserialize(reader); + layout_.Locations().Deserialize(reader); } else { // Legacy 32-bit format. The uint32 we just read is the old current_offset_. - current_offset_ = static_cast(maybe_sentinel); - this->io_->Deserialize(reader); + layout_.SetNextOffset(static_cast(maybe_sentinel)); + layout_.Payload().Deserialize(reader); // Legacy offset_io_ holds an array of 8-byte LegacyDocLocation records. We // load them and expand each entry to the new 12-byte DocLocation in memory // so the rest of the code can use a single internal representation. @@ -195,7 +194,7 @@ SparseVectorDataCell::Deserialize(lvalue_or_rvalueoffset_io_->Resize(doc_count * sizeof(DocLocation)); + layout_.ResizeLocations(doc_count); if (doc_count > 0) { constexpr uint64_t BATCH = 4096; Vector legacy_batch(allocator_); @@ -212,11 +211,10 @@ SparseVectorDataCell::Deserialize(lvalue_or_rvalue(legacy_batch[i].offset); - new_batch[i].size = legacy_batch[i].size; + new_batch[i].length = legacy_batch[i].size; } - this->offset_io_->Write(reinterpret_cast(new_batch.data()), - batch * sizeof(DocLocation), - cursor * sizeof(DocLocation)); + layout_.Locations().WriteRange( + cursor, reinterpret_cast(new_batch.data()), batch); cursor += batch; remaining -= batch; } @@ -235,9 +233,9 @@ SparseVectorDataCell::Serialize(StreamWriter& writer) { const uint32_t version = SERIALIZE_FORMAT_VERSION_V2; StreamWriter::WriteObj(writer, sentinel); StreamWriter::WriteObj(writer, version); - StreamWriter::WriteObj(writer, current_offset_); - this->io_->Serialize(writer); - this->offset_io_->Serialize(writer); + StreamWriter::WriteObj(writer, layout_.GetNextOffset()); + layout_.Payload().Serialize(writer); + layout_.Locations().Serialize(writer); this->quantizer_->Serialize(writer); } @@ -279,22 +277,11 @@ SparseVectorDataCell::InsertVector(const void* vector, InnerI } Vector codes(code_size, allocator_); quantizer_->EncodeOne((const float*)vector, codes.data()); - DocLocation location; { - std::scoped_lock lock(mutex_, current_offset_mutex_); + std::lock_guard lock(mutex_); total_count_ = std::max(total_count_, idx + 1); max_code_size_ = std::max(max_code_size_, code_size); - const auto required_size = current_offset_ + code_size; - if (required_size > this->io_->size_) { - this->io_->Resize(required_size); - } - location.offset = current_offset_; - location.size = static_cast(code_size); - current_offset_ += code_size; - offset_io_->Write(reinterpret_cast(&location), - sizeof(location), - static_cast(idx) * sizeof(location)); - io_->Write(codes.data(), code_size, location.offset); + layout_.Write(idx, codes.data(), code_size); } } @@ -315,12 +302,7 @@ template const uint8_t* SparseVectorDataCell::get_codes_by_id_no_lock(InnerIdType id, bool& need_release) const { - DocLocation location{}; - const bool read_ok = offset_io_->Read(sizeof(location), - static_cast(id) * sizeof(location), - reinterpret_cast(&location)); - CHECK_ARGUMENT(read_ok, "SparseVectorDataCell failed to read document location"); - const auto* codes = io_->Read(location.size, location.offset, need_release); + const auto* codes = layout_.Read(id, need_release); if (codes == nullptr) { throw VsagException(ErrorType::READ_ERROR, "SparseVectorDataCell failed to read vector codes"); @@ -371,7 +353,7 @@ SparseVectorDataCell::GetSparseVectorByInnerId( template void SparseVectorDataCell::Release(const uint8_t* data) const { - io_->Release(data); + layout_.Release(data); } template @@ -430,7 +412,7 @@ SparseVectorDataCell::SparseVectorDataCell( this->quantizer_ = std::make_shared(quantization_param, common_param); this->backend_ = QuantizerDistanceBackend::Get(static_cast(*this->quantizer_)); - this->io_ = std::make_shared(io_param, common_param); + auto io = std::make_shared(io_param, common_param); const auto& io_type = io_param->GetTypeName(); if (io_type == IO_TYPE_VALUE_MEMORY_IO || io_type == IO_TYPE_VALUE_BLOCK_MEMORY_IO) { this->query_io_strategy_ = QueryIOStrategy::DIRECT_READ; @@ -439,8 +421,9 @@ SparseVectorDataCell::SparseVectorDataCell( } else { this->query_io_strategy_ = QueryIOStrategy::MULTI_READ; } - this->offset_io_ = + auto offset_io = std::make_shared(Options::Instance().block_size_limit(), allocator_); + layout_.SetIO(std::move(offset_io), std::move(io)); this->max_code_size_ = std::max( sizeof(uint32_t), (static_cast(common_param.dim_) * 2 + 1) * sizeof(uint32_t)); this->max_capacity_ = 0; @@ -451,10 +434,7 @@ template uint64_t SparseVectorDataCell::GetMemoryUsage() const { uint64_t memory = sizeof(SparseVectorDataCell); - memory += this->offset_io_->size_; - if (IOTmpl::InMemory) { - memory += this->io_->GetMemoryUsage(); - } + memory += layout_.GetMemoryUsage(); memory += sizeof(QuantTmpl); return memory; } diff --git a/src/layout/byte_range_layout.h b/src/layout/byte_range_layout.h new file mode 100644 index 0000000000..73a2783d66 --- /dev/null +++ b/src/layout/byte_range_layout.h @@ -0,0 +1,119 @@ +// Copyright 2024-present the vsag project +// +// 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 +#include + +#include "index_common_param.h" +#include "io/common/basic_io.h" +#include "storage/stream_reader.h" +#include "storage/stream_writer.h" + +namespace vsag { + +/** Organizes opaque bytes addressed by an explicit byte offset and length. */ +template +class ByteRangeLayout { +public: + using IOType = IOTmpl; + static constexpr bool InMemory = IOTmpl::InMemory; + + ByteRangeLayout() = default; + + ByteRangeLayout(const IOParamPtr& io_param, const IndexCommonParam& common_param) + : io_(std::make_shared(io_param, common_param)) { + } + + void + SetIO(std::shared_ptr> io) { + io_ = std::move(io); + } + + void + Write(uint64_t offset, const uint8_t* data, uint64_t length) { + io_->Write(data, length, offset); + } + + bool + Read(uint64_t offset, uint64_t length, uint8_t* data) const { + return io_->Read(length, offset, data); + } + + [[nodiscard]] const uint8_t* + Read(uint64_t offset, uint64_t length, bool& need_release) const { + return io_->Read(length, offset, need_release); + } + + bool + MultiRead(uint64_t* offsets, uint64_t* lengths, uint64_t count, uint8_t* data) const { + return io_->MultiRead(data, lengths, offsets, count); + } + + void + Release(const uint8_t* data) const { + if (data != nullptr) { + io_->Release(data); + } + } + + void + Prefetch(uint64_t offset, uint64_t length) { + io_->Prefetch(offset, length); + } + + void + Resize(uint64_t byte_size) { + io_->Resize(byte_size); + } + + void + Shrink(uint64_t byte_size) { + io_->Shrink(byte_size); + } + + void + InitIO(const IOParamPtr& io_param) { + io_->InitIO(io_param); + } + + void + Serialize(StreamWriter& writer) { + io_->Serialize(writer); + } + + void + Deserialize(lvalue_or_rvalue reader) { + io_->Deserialize(reader); + } + + [[nodiscard]] uint64_t + GetMemoryUsage() const { + if constexpr (InMemory) { + return io_->GetMemoryUsage(); + } + return 0; + } + + [[nodiscard]] uint64_t + GetByteSize() const { + return io_->size_; + } + +private: + std::shared_ptr> io_{nullptr}; +}; + +} // namespace vsag diff --git a/src/layout/byte_range_layout_test.cpp b/src/layout/byte_range_layout_test.cpp new file mode 100644 index 0000000000..5cb714cef5 --- /dev/null +++ b/src/layout/byte_range_layout_test.cpp @@ -0,0 +1,74 @@ +// Copyright 2024-present the vsag project +// +// 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. + +#include "layout/byte_range_layout.h" + +#include +#include +#include + +#include "impl/allocator/safe_allocator.h" +#include "io/memory_io/memory_io.h" +#include "unittest.h" + +namespace vsag { + +TEST_CASE("ByteRangeLayout addresses opaque byte ranges", "[ut][ByteRangeLayout]") { + IndexCommonParam common_param; + common_param.allocator_ = SafeAllocator::FactoryDefaultAllocator(); + ByteRangeLayout layout(nullptr, common_param); + layout.Resize(16); + + const std::array first{1, 2, 3, 4}; + const std::array second{7, 8, 9}; + layout.Write(2, first.data(), first.size()); + layout.Write(10, second.data(), second.size()); + + std::array output{}; + REQUIRE(layout.Read(2, output.size(), output.data())); + REQUIRE(output == first); + + std::array offsets{10, 2}; + std::array lengths{3, 4}; + std::array batch{}; + REQUIRE(layout.MultiRead(offsets.data(), lengths.data(), offsets.size(), batch.data())); + const std::array expected{7, 8, 9, 1, 2, 3, 4}; + REQUIRE(batch == expected); +} + +TEST_CASE("ByteRangeLayout preserves underlying IO serialization", "[ut][ByteRangeLayout]") { + IndexCommonParam common_param; + common_param.allocator_ = SafeAllocator::FactoryDefaultAllocator(); + ByteRangeLayout source(nullptr, common_param); + const std::array bytes{2, 4, 6, 8, 10}; + source.Resize(bytes.size()); + source.Write(0, bytes.data(), bytes.size()); + + std::stringstream stream; + IOStreamWriter writer(stream); + source.Serialize(writer); + stream.seekg(0, std::ios::beg); + + ByteRangeLayout restored(nullptr, common_param); + IOStreamReader reader(stream); + restored.Deserialize(reader); + bool need_release = true; + const auto* result = restored.Read(0, bytes.size(), need_release); + REQUIRE(result != nullptr); + REQUIRE_FALSE(need_release); + REQUIRE(std::memcmp(result, bytes.data(), bytes.size()) == 0); + REQUIRE(restored.GetMemoryUsage() >= bytes.size()); +} + +} // namespace vsag diff --git a/src/layout/variable_record_layout.h b/src/layout/variable_record_layout.h new file mode 100644 index 0000000000..411159e25a --- /dev/null +++ b/src/layout/variable_record_layout.h @@ -0,0 +1,241 @@ +// Copyright 2024-present the vsag project +// +// 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 +#include +#include +#include +#include + +#include "layout/byte_range_layout.h" +#include "layout/fixed_layout.h" +#include "vsag_exception.h" + +namespace vsag { + +/** Default location policy for records whose offset and length are persisted. */ +struct OffsetAndLengthLocationPolicy { +#pragma pack(push, 1) + struct Entry { + uint64_t offset{0}; + uint32_t length{0}; + }; +#pragma pack(pop) + + static uint64_t + GetOffset(const Entry& entry) { + return entry.offset; + } + + template + static uint64_t + GetLength(const Entry& entry, const PayloadLayout&) { + return entry.length; + } + + static Entry + Make(uint64_t offset, uint64_t length) { + if (length > std::numeric_limits::max()) { + throw VsagException(ErrorType::INVALID_ARGUMENT, + "variable record length exceeds uint32_t limit"); + } + return {offset, static_cast(length)}; + } +}; + +/** Location policy for records whose byte length is derived from a uint32 header. */ +struct HeaderLengthLocationPolicy { + using Entry = uint64_t; + + uint64_t bytes_per_element{0}; + + static uint64_t + GetOffset(const Entry& entry) { + return entry; + } + + template + uint64_t + GetLength(const Entry& entry, const PayloadLayout& payload) const { + if (bytes_per_element == 0) { + throw VsagException(ErrorType::INVALID_ARGUMENT, + "variable record bytes per element must be positive"); + } + uint32_t element_count = 0; + if (not payload.Read( + entry, sizeof(element_count), reinterpret_cast(&element_count))) { + throw VsagException(ErrorType::READ_ERROR, + "failed to read variable record length header"); + } + if (element_count > + (std::numeric_limits::max() - sizeof(element_count)) / bytes_per_element) { + throw VsagException(ErrorType::INVALID_ARGUMENT, + "variable record header length overflow"); + } + return sizeof(element_count) + static_cast(element_count) * bytes_per_element; + } + + static Entry + Make(uint64_t offset, uint64_t) { + return offset; + } +}; + +/** Maps a logical record ID to an opaque variable-length payload. */ +template +class VariableRecordLayout { +public: + using LocationEntry = typename LocationPolicy::Entry; + using LocationLayout = FixedLayout; + using PayloadLayout = ByteRangeLayout; + + VariableRecordLayout() = default; + + VariableRecordLayout(std::shared_ptr> location_io, + std::shared_ptr> payload_io) { + SetIO(std::move(location_io), std::move(payload_io)); + } + + void + SetIO(std::shared_ptr> location_io, + std::shared_ptr> payload_io) { + locations_.SetCodeSize(sizeof(LocationEntry)); + locations_.SetIO(std::move(location_io)); + payload_.SetIO(std::move(payload_io)); + } + + void + SetLocationPolicy(LocationPolicy policy) { + location_policy_ = std::move(policy); + } + + void + Write(InnerIdType id, const uint8_t* data, uint64_t length) { + std::lock_guard lock(append_mutex_); + const uint64_t offset = next_offset_; + if (length > std::numeric_limits::max() - offset) { + throw VsagException(ErrorType::INVALID_ARGUMENT, + "variable record payload offset overflow"); + } + const uint64_t required_size = offset + length; + if (required_size > payload_.GetByteSize()) { + payload_.Resize(required_size); + } + payload_.Write(offset, data, length); + const auto location = location_policy_.Make(offset, length); + locations_.Write(id, reinterpret_cast(&location)); + next_offset_ = required_size; + } + + [[nodiscard]] LocationEntry + ReadLocation(InnerIdType id) const { + LocationEntry location{}; + if (not locations_.Read(id, reinterpret_cast(&location))) { + throw VsagException(ErrorType::READ_ERROR, "failed to read variable record location"); + } + return location; + } + + [[nodiscard]] uint64_t + GetRecordLength(const LocationEntry& location) const { + return location_policy_.GetLength(location, payload_); + } + + [[nodiscard]] const uint8_t* + Read(InnerIdType id, bool& need_release) const { + const auto location = ReadLocation(id); + return Read(location, need_release); + } + + [[nodiscard]] const uint8_t* + Read(const LocationEntry& location, bool& need_release) const { + return payload_.Read( + location_policy_.GetOffset(location), GetRecordLength(location), need_release); + } + + bool + MultiRead(const LocationEntry* locations, + uint64_t count, + uint8_t* data, + Allocator* allocator) const { + Vector offsets(count, allocator); + Vector lengths(count, allocator); + for (uint64_t i = 0; i < count; ++i) { + offsets[i] = location_policy_.GetOffset(locations[i]); + lengths[i] = GetRecordLength(locations[i]); + } + return payload_.MultiRead(offsets.data(), lengths.data(), count, data); + } + + void + Release(const uint8_t* data) const { + payload_.Release(data); + } + + void + ResizeLocations(uint64_t capacity) { + locations_.Resize(capacity); + } + + void + ReservePayload(uint64_t byte_size) { + payload_.Resize(byte_size); + } + + void + SetNextOffset(uint64_t offset) { + next_offset_ = offset; + } + + [[nodiscard]] uint64_t + GetNextOffset() const { + return next_offset_; + } + + LocationLayout& + Locations() { + return locations_; + } + + PayloadLayout& + Payload() { + return payload_; + } + + [[nodiscard]] const PayloadLayout& + Payload() const { + return payload_; + } + + [[nodiscard]] const LocationLayout& + Locations() const { + return locations_; + } + + [[nodiscard]] uint64_t + GetMemoryUsage() const { + return locations_.GetMemoryUsage() + payload_.GetMemoryUsage(); + } + +private: + LocationLayout locations_{}; + PayloadLayout payload_{}; + uint64_t next_offset_{0}; + LocationPolicy location_policy_{}; + std::mutex append_mutex_; +}; + +} // namespace vsag diff --git a/src/layout/variable_record_layout_test.cpp b/src/layout/variable_record_layout_test.cpp new file mode 100644 index 0000000000..01eec5c1a7 --- /dev/null +++ b/src/layout/variable_record_layout_test.cpp @@ -0,0 +1,74 @@ +// Copyright 2024-present the vsag project +// +// 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. + +#include "layout/variable_record_layout.h" + +#include +#include + +#include "impl/allocator/safe_allocator.h" +#include "io/memory_block_io/memory_block_io.h" +#include "io/memory_io/memory_io.h" +#include "unittest.h" +#include "vsag/options.h" + +namespace vsag { + +TEST_CASE("VariableRecordLayout maps ids to appended records", "[ut][VariableRecordLayout]") { + IndexCommonParam common_param; + common_param.allocator_ = SafeAllocator::FactoryDefaultAllocator(); + auto location_io = std::make_shared(Options::Instance().block_size_limit(), + common_param.allocator_.get()); + auto payload_io = std::make_shared(IOParamPtr{}, common_param); + VariableRecordLayout layout(location_io, + payload_io); + layout.ResizeLocations(4); + + const std::array first{1, 3, 5}; + const std::array second{2, 4, 6, 8, 10}; + layout.Write(2, first.data(), first.size()); + layout.Write(0, second.data(), second.size()); + + const auto first_location = layout.ReadLocation(2); + const auto second_location = layout.ReadLocation(0); + REQUIRE(first_location.offset == 0); + REQUIRE(first_location.length == first.size()); + REQUIRE(second_location.offset == first.size()); + REQUIRE(second_location.length == second.size()); + REQUIRE(layout.GetNextOffset() == first.size() + second.size()); + + bool need_release = true; + const auto* record = layout.Read(0, need_release); + REQUIRE(record != nullptr); + REQUIRE_FALSE(need_release); + REQUIRE(std::memcmp(record, second.data(), second.size()) == 0); + + const std::array locations{second_location, first_location}; + std::array batch{}; + REQUIRE(layout.MultiRead( + locations.data(), locations.size(), batch.data(), common_param.allocator_.get())); + const std::array expected{2, 4, 6, 8, 10, 1, 3, 5}; + REQUIRE(batch == expected); +} + +TEST_CASE("HeaderLengthLocationPolicy rejects an unset element width", + "[ut][VariableRecordLayout]") { + IndexCommonParam common_param; + common_param.allocator_ = SafeAllocator::FactoryDefaultAllocator(); + ByteRangeLayout payload(nullptr, common_param); + HeaderLengthLocationPolicy policy; + REQUIRE_THROWS_AS(policy.GetLength(0, payload), VsagException); +} + +} // namespace vsag