Skip to content

refactor(datacell): extract variable record layouts #2737

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
LHT129 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

refactor(datacell): extract variable record layouts #2737

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 2 additions & 4 deletions src/datacell/multi_vector_datacell.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -114,12 +115,9 @@ class MultiVectorDataCell : public FlattenInterface {

private:
std::shared_ptr<Quantizer<QuantTmpl>> quantizer_{nullptr};
std::shared_ptr<BasicIO<IOTmpl>> io_{nullptr};

Allocator* const allocator_{nullptr};
std::shared_ptr<MemoryBlockIO> offset_io_{nullptr};
uint64_t current_offset_{0};
std::mutex current_offset_mutex_;
VariableRecordLayout<HeaderLengthLocationPolicy, MemoryBlockIO, IOTmpl> layout_{};

uint32_t multi_vector_dim_{0};
MetricType metric_{MetricType::METRIC_TYPE_L2SQR};
Expand Down
108 changes: 66 additions & 42 deletions src/datacell/multi_vector_datacell.inl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

#include <algorithm>
#include <cstring>
#include <limits>

#include "common.h"
#include "multi_vector_datacell.h"
Expand All @@ -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<uint64_t>(dimension) * value_size;
if (values_per_token != 0 &&
token_count > (std::numeric_limits<uint64_t>::max() - header_size) / values_per_token) {
throw VsagException(ErrorType::INVALID_ARGUMENT,
"MultiVectorDataCell: record size overflow");
}
return header_size + static_cast<uint64_t>(token_count) * values_per_token;
}

template <typename QuantTmpl, typename IOTmpl>
MultiVectorDataCell<QuantTmpl, IOTmpl>::MultiVectorDataCell(
const QuantizerParamPtr& quantization_param,
Expand All @@ -36,9 +50,12 @@ MultiVectorDataCell<QuantTmpl, IOTmpl>::MultiVectorDataCell(
this->quantizer_ = std::make_shared<QuantTmpl>(quantization_param, common_param);
this->backend_ =
QuantizerDistanceBackend<QuantTmpl>::Get(static_cast<const QuantTmpl&>(*this->quantizer_));
this->io_ = std::make_shared<IOTmpl>(io_param, common_param);
this->offset_io_ =
auto io = std::make_shared<IOTmpl>(io_param, common_param);
auto offset_io =
std::make_shared<MemoryBlockIO>(Options::Instance().block_size_limit(), allocator_);
layout_.SetIO(std::move(offset_io), std::move(io));
layout_.SetLocationPolicy(
HeaderLengthLocationPolicy{static_cast<uint64_t>(multi_vector_dim_) * sizeof(float)});
this->max_capacity_ = 0;
this->code_size_ = 0;
}
Expand Down Expand Up @@ -68,23 +85,16 @@ MultiVectorDataCell<QuantTmpl, IOTmpl>::InsertVector(const void* vector, InnerId
}
}

const uint64_t vector_bytes = static_cast<uint64_t>(multi_vector->len_) *
static_cast<uint64_t>(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<const uint8_t*>(&old_offset),
sizeof(old_offset),
static_cast<uint64_t>(idx) * sizeof(old_offset));
io_->Write(codes.data, code_size, old_offset);
}

template <typename QuantTmpl, typename IOTmpl>
Expand Down Expand Up @@ -116,7 +126,7 @@ MultiVectorDataCell<QuantTmpl, IOTmpl>::Resize(InnerIdType new_capacity) {
if (new_capacity <= this->max_capacity_) {
return;
}
this->offset_io_->Resize(static_cast<uint64_t>(new_capacity) * sizeof(uint64_t));
layout_.ResizeLocations(new_capacity);
this->max_capacity_ = new_capacity;
}

Expand All @@ -135,14 +145,24 @@ MultiVectorDataCell<QuantTmpl, IOTmpl>::GetMetricType() {
template <typename QuantTmpl, typename IOTmpl>
const uint8_t*
MultiVectorDataCell<QuantTmpl, IOTmpl>::GetCodesById(InnerIdType id, bool& need_release) const {
uint64_t offset = 0;
offset_io_->Read(sizeof(offset), static_cast<uint64_t>(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<uint64_t>(len) * multi_vector_dim_ * sizeof(float);
if (not layout_.Payload().Read(offset, sizeof(len), reinterpret_cast<uint8_t*>(&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<uint8_t*>(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;
}
Expand All @@ -164,9 +184,9 @@ void
MultiVectorDataCell<QuantTmpl, IOTmpl>::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);
}

Expand All @@ -175,9 +195,11 @@ void
MultiVectorDataCell<QuantTmpl, IOTmpl>::Deserialize(lvalue_or_rvalue<StreamReader> 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<QuantTmpl>::Get(static_cast<const QuantTmpl&>(*this->quantizer_));
Expand Down Expand Up @@ -213,19 +235,16 @@ MultiVectorDataCell<QuantTmpl, IOTmpl>::Query(float* result_dists,
// Step 1: Read all offsets (offset_io_ is MemoryBlockIO, in-memory, fast)
std::vector<uint64_t> offsets(id_count);
for (InnerIdType i = 0; i < id_count; ++i) {
bool ok = offset_io_->Read(sizeof(uint64_t),
static_cast<uint64_t>(idx[i]) * sizeof(uint64_t),
reinterpret_cast<uint8_t*>(&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<uint32_t> lens(id_count);
std::vector<uint64_t> len_sizes(id_count, sizeof(uint32_t));
if (!this->io_->MultiRead(reinterpret_cast<uint8_t*>(lens.data()),
len_sizes.data(),
offsets.data(),
static_cast<uint64_t>(id_count))) {
if (!layout_.Payload().MultiRead(offsets.data(),
len_sizes.data(),
static_cast<uint64_t>(id_count),
reinterpret_cast<uint8_t*>(lens.data()))) {
throw VsagException(ErrorType::READ_ERROR,
"MultiVectorDataCell: failed to read token counts");
}
Expand All @@ -234,13 +253,21 @@ MultiVectorDataCell<QuantTmpl, IOTmpl>::Query(float* result_dists,
std::vector<uint64_t> 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<uint64_t>(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<uint64_t>::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<uint64_t>(id_count))) {
if (!layout_.Payload().MultiRead(
offsets.data(), data_sizes.data(), static_cast<uint64_t>(id_count), all_codes.data)) {
Comment thread
LHT129 marked this conversation as resolved.
throw VsagException(ErrorType::READ_ERROR,
"MultiVectorDataCell: failed to read token data");
}
Expand All @@ -261,10 +288,7 @@ template <typename QuantTmpl, typename IOTmpl>
uint64_t
MultiVectorDataCell<QuantTmpl, IOTmpl>::GetMemoryUsage() const {
uint64_t memory = sizeof(MultiVectorDataCell<QuantTmpl, IOTmpl>);
memory += this->offset_io_->size_;
if (IOTmpl::InMemory) {
memory += this->io_->GetMemoryUsage();
}
memory += layout_.GetMemoryUsage();
memory += sizeof(QuantTmpl);
return memory;
}
Expand Down
27 changes: 10 additions & 17 deletions src/datacell/sparse_vector_datacell.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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<uint64_t>(new_capacity - total_count_) * max_code_size_ + current_offset_;
this->io_->Resize(io_size);
this->offset_io_->Resize(static_cast<uint64_t>(new_capacity) * sizeof(DocLocation));
std::lock_guard lock(mutex_);
uint64_t io_size = static_cast<uint64_t>(new_capacity - total_count_) * max_code_size_ +
layout_.GetNextOffset();
layout_.ReservePayload(io_size);
layout_.ResizeLocations(new_capacity);
this->max_capacity_ = new_capacity;
}

Expand Down Expand Up @@ -149,12 +150,12 @@ class SparseVectorDataCell : public FlattenInterface {

inline void
SetIO(std::shared_ptr<BasicIO<IOTmpl>> 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
Expand Down Expand Up @@ -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
Expand All @@ -217,14 +213,11 @@ class SparseVectorDataCell : public FlattenInterface {
static constexpr uint32_t SERIALIZE_FORMAT_VERSION_V2 = 2;

std::shared_ptr<Quantizer<QuantTmpl>> quantizer_{nullptr};
std::shared_ptr<BasicIO<IOTmpl>> io_{nullptr};
QueryIOStrategy query_io_strategy_{QueryIOStrategy::MULTI_READ};

Allocator* const allocator_{nullptr};
std::shared_ptr<MemoryBlockIO> offset_io_{nullptr};
uint64_t current_offset_{0};
VariableRecordLayout<OffsetAndLengthLocationPolicy, MemoryBlockIO, IOTmpl> layout_{};
uint64_t max_code_size_{0};
std::mutex current_offset_mutex_;
};

} // namespace vsag
Expand Down
Loading
Loading