xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: clarify segment-local row ID handling and tests and fixes bugs (#432)

This commit is contained in:
Qinren Zhou 2026-06-03 15:20:45 +08:00 committed by GitHub
parent 95e5ad5105
commit dbea635019
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 839 additions and 784 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
src/db/index/column/vector_column/combined_vector_column_indexer.cc
View File

@ -116,6 +116,18 @@ Result<IndexResults::Ptr> CombinedVectorColumnIndexer::Search(
}
}
std::unique_ptr<vector_column_params::GroupByParams> group_by;
if (query_params.group_by) {
auto group_by_func = query_params.group_by->group_by;
auto block_offset = block_offsets_[i];
group_by = std::make_unique<vector_column_params::GroupByParams>(
query_params.group_by->group_topk, query_params.group_by->group_count,
[group_by_func = std::move(group_by_func),
block_offset](uint64_t block_doc_id) {
return group_by_func(block_doc_id + block_offset);
});
}
vector_column_params::QueryParams modified_query_params{
query_params.data_type,
query_params.dimension,
@ -123,12 +135,7 @@ Result<IndexResults::Ptr> CombinedVectorColumnIndexer::Search(
filter,
query_params.fetch_vector,
query_params.query_params,
query_params.group_by
? std::make_unique<vector_column_params::GroupByParams>(
query_params.group_by->group_topk,
query_params.group_by->group_count,
query_params.group_by->group_by)
: nullptr,
std::move(group_by),
{},
need_refine ? std::shared_ptr<vector_column_params::RefinerParam>(
new vector_column_params::RefinerParam{
@ -270,4 +277,4 @@ CombinedVectorColumnIndexer::Fetch(uint32_t segment_doc_id) const {
return indexer->Fetch(target_block_doc_id);
}
} // namespace zvec
} // namespace zvec

3
src/db/index/column/vector_column/combined_vector_column_indexer.h
View File

@ -38,7 +38,6 @@ class CombinedVectorColumnIndexer {
const vector_column_params::VectorData &vector_data,
const vector_column_params::QueryParams &query_params);
// doc_id is segment local id
virtual Result<vector_column_params::VectorDataBuffer> Fetch(
uint32_t segment_doc_id) const;
@ -81,4 +80,4 @@ class CombinedVectorColumnIndexer {
uint64_t min_doc_id_{0};
};
} // namespace zvec
} // namespace zvec

265
src/db/index/segment/segment.cc
View File

@ -49,7 +49,6 @@
#include "db/index/column/fts_column/fts_rocksdb_merge.h"
#include "db/index/column/fts_column/fts_types.h"
#include "db/index/column/inverted_column/inverted_indexer.h"
#include "db/index/column/vector_column/engine_helper.hpp"
#include "db/index/column/vector_column/vector_column_indexer.h"
#include "db/index/column/vector_column/vector_column_params.h"
#include "db/index/common/index_filter.h"
@ -61,11 +60,7 @@
#include "db/index/storage/mmap_forward_store.h"
#include "db/index/storage/store_helper.h"
#include "db/index/storage/wal/wal_file.h"
#include "zvec/ailego/container/params.h"
#include "zvec/core/framework/index_factory.h"
#include "zvec/core/framework/index_meta.h"
#include "zvec/core/framework/index_provider.h"
#include "zvec/core/framework/index_reformer.h"
#include "column_merging_reader.h"
#include "sql_expr_parser.h"
@ -225,10 +220,10 @@ class SegmentImpl : public Segment,
Status destroy() override;
TablePtr fetch(const std::vector<std::string> &columns,
const std::vector<int> &indices) const override;
const std::vector<int> &segment_doc_ids) const override;
ExecBatchPtr fetch(const std::vector<std::string> &columns,
int index) const override;
int segment_doc_id) const override;
RecordBatchReaderPtr scan(
const std::vector<std::string> &columns) const override;
@ -302,7 +297,7 @@ class SegmentImpl : public Segment,
Status append_wal(const Doc &doc);
Status update_version(uint32_t delete_snapshot_path_suffix);
Result<uint64_t> get_global_doc_id(uint32_t local_id) const;
Result<uint64_t> get_global_doc_id(uint32_t segment_doc_id) const;
BlockID allocate_block_id();
@ -323,12 +318,12 @@ class SegmentImpl : public Segment,
TablePtr fetch_normal(const std::vector<std::string> &columns,
const std::shared_ptr<arrow::Schema> &result_schema,
const std::vector<int> &indices) const;
const std::vector<int> &segment_doc_ids) const;
// For performance tuning
TablePtr fetch_perf(const std::vector<std::string> &columns,
const std::shared_ptr<arrow::Schema> &result_schema,
const std::vector<int> &indices) const;
const std::vector<int> &segment_doc_ids) const;
void fresh_persist_chunked_array();
@ -416,7 +411,6 @@ class SegmentImpl : public Segment,
class SegmentImpl::CombinedRecordBatchReader : public arrow::RecordBatchReader {
public:
CombinedRecordBatchReader(
std::shared_ptr<const SegmentImpl> segment,
std::vector<std::shared_ptr<arrow::RecordBatchReader>> readers,
const std::vector<std::string> &columns);
@ -427,14 +421,12 @@ class SegmentImpl::CombinedRecordBatchReader : public arrow::RecordBatchReader {
arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch> *batch) override;
private:
std::shared_ptr<const SegmentImpl> segment_;
std::vector<std::shared_ptr<arrow::RecordBatchReader>> readers_;
std::vector<uint64_t> offsets_;
std::shared_ptr<arrow::Schema> projected_schema_;
bool need_local_doc_id_ = false;
bool emit_segment_row_id_ = false;
size_t current_reader_index_;
size_t local_doc_id_;
int local_doc_id_col_index_ = -1;
uint64_t next_segment_row_id_to_emit_;
int segment_row_id_output_col_index_ = -1;
};
////////////////////////////////////////////////////////////////////////////////////
@ -1433,7 +1425,7 @@ Doc::Ptr SegmentImpl::Fetch(
const auto &block_offsets =
get_persist_block_offsets(BlockType::VECTOR_INDEX, field->name());
auto block_offset = block_offsets[block_idx];
auto local_row = segment_doc_id - block_offset;
auto block_doc_id = segment_doc_id - block_offset;
auto column_name = field->name();
auto iter = vector_indexers_.find(column_name);
@ -1445,12 +1437,12 @@ Doc::Ptr SegmentImpl::Fetch(
continue;
}
auto vector_indexer = vector_indexers[block_idx];
auto fetch_result = vector_indexer->Fetch(local_row);
auto fetch_result = vector_indexer->Fetch(block_doc_id);
if (!fetch_result) {
LOG_ERROR(
"vector indexer fetch failed, local_row: %d, block_idx: %d, "
"vector indexer fetch failed, block_doc_id: %d, block_idx: %d, "
"segment_doc_id: %d",
local_row, block_idx, segment_doc_id);
block_doc_id, block_idx, segment_doc_id);
return nullptr;
}
const auto &vector_buffer = fetch_result.value();
@ -1472,18 +1464,18 @@ Doc::Ptr SegmentImpl::Fetch(
p_block_offsets.empty()
? 0
: p_block_offsets.back() + p_block_metas.back().doc_count_;
int local_row = segment_doc_id - mem_block_offset;
int block_doc_id = segment_doc_id - mem_block_offset;
auto column_name = field->name();
auto iter = memory_vector_indexers_.find(column_name);
if (iter != memory_vector_indexers_.end()) {
auto vector_indexer = iter->second;
auto fetch_result = vector_indexer->Fetch(local_row);
auto fetch_result = vector_indexer->Fetch(block_doc_id);
if (!fetch_result.has_value()) {
LOG_ERROR(
"vector indexer fetch failed, column: %s, doc_count: %lu, "
"mem_block_offset: %d, local_row: %d",
"mem_block_offset: %d, block_doc_id: %d",
field->name().c_str(), vector_indexer->doc_count(),
mem_block_offset, local_row);
mem_block_offset, block_doc_id);
continue;
}
const auto &vector_buffer = fetch_result.value();
@ -2339,17 +2331,17 @@ bool SegmentImpl::validate(const std::vector<std::string> &columns) const {
TablePtr SegmentImpl::fetch_perf(
const std::vector<std::string> &columns,
const std::shared_ptr<arrow::Schema> &result_schema,
const std::vector<int> &indices) const {
const std::vector<int> &segment_doc_ids) const {
std::vector<std::shared_ptr<arrow::ChunkedArray>> chunk_arrays;
chunk_arrays.resize(columns.size());
bool need_local_doc_id = false;
size_t local_doc_id_col_index = 0;
bool has_segment_row_id_column = false;
size_t segment_row_id_col_index = 0;
for (size_t i = 0; i < columns.size(); ++i) {
if (columns[i] == LOCAL_ROW_ID) {
need_local_doc_id = true;
local_doc_id_col_index = i;
has_segment_row_id_column = true;
segment_row_id_col_index = i;
chunk_arrays[i] = nullptr;
continue;
}
@ -2358,18 +2350,19 @@ TablePtr SegmentImpl::fetch_perf(
std::vector<std::shared_ptr<arrow::Array>> result_arrays(columns.size());
std::vector<std::pair<int64_t, int64_t>> indices_in_table;
for (const auto &target_index : indices) {
// Parallel to segment_doc_ids: each pair is (chunk_index, row_index_in_chunk)
std::vector<std::pair<int64_t, int64_t>> chunk_row_indices_for_ids;
for (const auto segment_doc_id : segment_doc_ids) {
auto it = std::upper_bound(chunk_offsets_.begin(), chunk_offsets_.end(),
target_index);
segment_doc_id);
if (it == chunk_offsets_.begin()) {
LOG_ERROR("Target index %d is out of bounds", target_index);
LOG_ERROR("Segment doc ID %d is out of bounds", segment_doc_id);
return nullptr;
}
int chunk_index =
static_cast<int>(std::distance(chunk_offsets_.begin(), it) - 1);
int64_t index_in_chunk = target_index - chunk_offsets_[chunk_index];
indices_in_table.emplace_back(chunk_index, index_in_chunk);
int64_t row_index_in_chunk = segment_doc_id - chunk_offsets_[chunk_index];
chunk_row_indices_for_ids.emplace_back(chunk_index, row_index_in_chunk);
}
for (size_t i = 0; i < columns.size(); ++i) {
@ -2378,8 +2371,8 @@ TablePtr SegmentImpl::fetch_perf(
}
const auto &source_column = chunk_arrays[i];
std::shared_ptr<arrow::Array> array;
auto status =
BuildArrayFromIndicesWithType(source_column, indices_in_table, &array);
auto status = BuildArrayFromIndicesWithType(
source_column, chunk_row_indices_for_ids, &array);
if (!status.ok()) {
LOG_ERROR("BuildArrayFromIndices failed: %s", status.ToString().c_str());
return nullptr;
@ -2387,11 +2380,11 @@ TablePtr SegmentImpl::fetch_perf(
result_arrays[i] = array;
}
if (need_local_doc_id) {
if (has_segment_row_id_column) {
std::vector<uint64_t> values;
values.reserve(indices.size());
for (const auto idx : indices) {
values.push_back(idx);
values.reserve(segment_doc_ids.size());
for (const auto segment_doc_id : segment_doc_ids) {
values.push_back(segment_doc_id);
}
arrow::UInt64Builder builder;
@ -2406,26 +2399,26 @@ TablePtr SegmentImpl::fetch_perf(
LOG_ERROR("Failed to finish builder: %s", s.message().c_str());
return nullptr;
}
result_arrays[local_doc_id_col_index] = array;
result_arrays[segment_row_id_col_index] = array;
}
return arrow::Table::Make(result_schema, result_arrays,
static_cast<int64_t>(indices.size()));
static_cast<int64_t>(segment_doc_ids.size()));
}
TablePtr SegmentImpl::fetch_normal(
const std::vector<std::string> &columns,
const std::shared_ptr<arrow::Schema> &result_schema,
const std::vector<int> &indices) const {
const std::vector<int> &segment_doc_ids) const {
// Store scalars per column: column_index -> (output_row, scalar)
std::vector<std::vector<std::pair<int, std::shared_ptr<arrow::Scalar>>>>
column_results(columns.size());
// Collect local_doc_id values if needed
std::vector<std::pair<int, uint64_t>> local_doc_id_values;
// Collect segment-local row IDs when LOCAL_ROW_ID is requested.
std::vector<std::pair<int, uint64_t>> segment_row_id_values;
// Group fetch requests by block: block_index -> {column -> [(output_row,
// local_row)]}
// Group fetch requests by block:
// block_index -> {column -> [(output_row, block_row)]}
// block_index >= 0: persisted store
// block_index == -1: memory store
std::map<int, std::map<std::string, std::vector<std::pair<int, int>>>>
@ -2436,26 +2429,27 @@ TablePtr SegmentImpl::fetch_normal(
const auto &block_offsets = get_persist_block_offsets(BlockType::SCALAR);
const auto &block_metas = get_persist_block_metas(BlockType::SCALAR);
// Phase 1: Map each (doc_id, column) to its block and local row
for (int output_row = 0; output_row < static_cast<int>(indices.size());
++output_row) {
int doc_id = indices[output_row];
// Phase 1: Map each (segment_doc_id, column) to its block and block-local
// row.
for (int output_row = 0;
output_row < static_cast<int>(segment_doc_ids.size()); ++output_row) {
int segment_doc_id = segment_doc_ids[output_row];
for (size_t col_index = 0; col_index < columns.size(); ++col_index) {
const std::string &col = columns[col_index];
if (col == LOCAL_ROW_ID) {
local_doc_id_values.emplace_back(output_row, doc_id);
segment_row_id_values.emplace_back(output_row, segment_doc_id);
continue;
}
int offset_idx = -1;
int block_index =
find_persist_block_id(BlockType::SCALAR, doc_id, col, &offset_idx);
int block_index = find_persist_block_id(BlockType::SCALAR, segment_doc_id,
col, &offset_idx);
int local_row = -1;
int block_row = -1;
if (block_index != -1 && offset_idx > -1 &&
offset_idx < static_cast<int>(block_offsets.size())) {
local_row = doc_id - block_offsets[offset_idx];
block_request_map[block_index][col].emplace_back(output_row, local_row);
block_row = segment_doc_id - block_offsets[offset_idx];
block_request_map[block_index][col].emplace_back(output_row, block_row);
continue;
}
@ -2467,15 +2461,17 @@ TablePtr SegmentImpl::fetch_normal(
: block_offsets.back() + block_metas.back().doc_count_;
const auto &mem_block = segment_meta_->writing_forward_block().value();
if (mem_offset <= doc_id &&
doc_id < mem_offset + static_cast<int>(mem_block.doc_count_)) {
local_row = doc_id - mem_offset;
block_request_map[-1][col].emplace_back(output_row, local_row);
if (mem_offset <= segment_doc_id &&
segment_doc_id <
mem_offset + static_cast<int>(mem_block.doc_count_)) {
block_row = segment_doc_id - mem_offset;
block_request_map[-1][col].emplace_back(output_row, block_row);
continue;
}
}
LOG_ERROR("Document ID %d not found in segment %d", doc_id, meta()->id());
LOG_ERROR("Segment doc ID %d not found in segment %d", segment_doc_id,
meta()->id());
return nullptr;
}
}
@ -2483,7 +2479,7 @@ TablePtr SegmentImpl::fetch_normal(
// Phase 2: Execute batched fetch per block
for (const auto &[block_index, col_to_rows] : block_request_map) {
std::vector<std::string> fetch_columns;
std::vector<int> fetch_local_rows;
std::vector<int> fetch_block_rows;
std::vector<std::pair<int, int>>
output_to_result_index; // (output_row, result_pos)
@ -2493,20 +2489,20 @@ TablePtr SegmentImpl::fetch_normal(
}
// all column has same output size, here just take first column
for (const auto &[output_row, local_row] :
for (const auto &[output_row, block_row] :
col_to_rows.at(fetch_columns[0])) {
fetch_local_rows.push_back(local_row);
fetch_block_rows.push_back(block_row);
output_to_result_index.emplace_back(
output_row, static_cast<int>(fetch_local_rows.size() - 1));
output_row, static_cast<int>(fetch_block_rows.size() - 1));
}
std::shared_ptr<arrow::Table> block_table;
if (block_index >= 0 &&
block_index < static_cast<int>(persist_stores_.size())) {
block_table =
persist_stores_[block_index]->fetch(fetch_columns, fetch_local_rows);
persist_stores_[block_index]->fetch(fetch_columns, fetch_block_rows);
} else if (block_index == -1 && memory_store_) {
block_table = memory_store_->fetch(fetch_columns, fetch_local_rows);
block_table = memory_store_->fetch(fetch_columns, fetch_block_rows);
}
if (!block_table || block_table->num_rows() == 0) {
@ -2530,7 +2526,7 @@ TablePtr SegmentImpl::fetch_normal(
}
auto flat_array = flat_array_res.ValueOrDie();
for (size_t j = 0; j < fetch_local_rows.size(); ++j) {
for (size_t j = 0; j < fetch_block_rows.size(); ++j) {
auto scalar_result = flat_array->GetScalar(j);
if (!scalar_result.ok()) continue;
int output_row = output_to_result_index[j].first;
@ -2543,14 +2539,14 @@ TablePtr SegmentImpl::fetch_normal(
// Phase 3: Construct result arrays
std::vector<std::shared_ptr<arrow::Array>> result_arrays(columns.size());
bool need_local_doc_id = false;
size_t local_doc_id_col_index = -1;
bool has_segment_row_id_column = false;
size_t segment_row_id_col_index = -1;
for (size_t col_index = 0; col_index < columns.size(); ++col_index) {
const std::string &col = columns[col_index];
if (col == LOCAL_ROW_ID) {
need_local_doc_id = true;
local_doc_id_col_index = col_index;
has_segment_row_id_column = true;
segment_row_id_col_index = col_index;
continue;
}
@ -2558,7 +2554,7 @@ TablePtr SegmentImpl::fetch_normal(
std::sort(result_vec.begin(), result_vec.end());
std::vector<std::shared_ptr<arrow::Scalar>> ordered_scalars;
for (int i = 0; i < static_cast<int>(indices.size()); ++i) {
for (int i = 0; i < static_cast<int>(segment_doc_ids.size()); ++i) {
auto it = std::find_if(
result_vec.begin(), result_vec.end(),
[i](const std::pair<int, std::shared_ptr<arrow::Scalar>> &p) {
@ -2582,13 +2578,13 @@ TablePtr SegmentImpl::fetch_normal(
}
}
// Add LOCAL_ROW_ID array if requested
if (need_local_doc_id) {
std::sort(local_doc_id_values.begin(), local_doc_id_values.end());
// Add segment-local values for the LOCAL_ROW_ID column.
if (has_segment_row_id_column) {
std::sort(segment_row_id_values.begin(), segment_row_id_values.end());
std::vector<uint64_t> values;
values.reserve(local_doc_id_values.size());
for (const auto &[row, id] : local_doc_id_values) {
values.push_back(id);
values.reserve(segment_row_id_values.size());
for (const auto &[row, segment_row_id] : segment_row_id_values) {
values.push_back(segment_row_id);
}
arrow::UInt64Builder builder;
@ -2603,7 +2599,7 @@ TablePtr SegmentImpl::fetch_normal(
LOG_ERROR("Failed to finish builder: %s", s.message().c_str());
return nullptr;
}
result_arrays[local_doc_id_col_index] = std::move(array);
result_arrays[segment_row_id_col_index] = std::move(array);
}
// Wrap arrays into ChunkedArray and build final table
@ -2614,11 +2610,11 @@ TablePtr SegmentImpl::fetch_normal(
}
return arrow::Table::Make(result_schema, result_columns,
static_cast<int64_t>(indices.size()));
static_cast<int64_t>(segment_doc_ids.size()));
}
TablePtr SegmentImpl::fetch(const std::vector<std::string> &columns,
const std::vector<int> &indices) const {
const std::vector<int> &segment_doc_ids) const {
if (!validate(columns)) {
return nullptr;
}
@ -2649,8 +2645,8 @@ TablePtr SegmentImpl::fetch(const std::vector<std::string> &columns,
auto result_schema = std::make_shared<arrow::Schema>(fields);
// Early return for empty indices
if (indices.empty()) {
// Early return for empty segment doc IDs.
if (segment_doc_ids.empty()) {
arrow::ArrayVector empty_arrays;
for (const auto &field : fields) {
empty_arrays.push_back(arrow::MakeEmptyArray(field->type()).ValueOrDie());
@ -2664,13 +2660,13 @@ TablePtr SegmentImpl::fetch(const std::vector<std::string> &columns,
}
if (use_fetch_perf_) {
return fetch_perf(columns, result_schema, indices);
return fetch_perf(columns, result_schema, segment_doc_ids);
}
return fetch_normal(columns, result_schema, indices);
return fetch_normal(columns, result_schema, segment_doc_ids);
}
ExecBatchPtr SegmentImpl::fetch(const std::vector<std::string> &columns,
int doc_id) const {
int segment_doc_id) const {
if (columns.empty()) {
LOG_ERROR("Empty columns");
return nullptr;
@ -2709,12 +2705,12 @@ ExecBatchPtr SegmentImpl::fetch(const std::vector<std::string> &columns,
if (is_in_single_persist_store) {
int offset_idx = -1;
int block_index = find_persist_block_id(BlockType::SCALAR, doc_id,
int block_index = find_persist_block_id(BlockType::SCALAR, segment_doc_id,
columns[0], &offset_idx);
if (block_index != -1 && offset_idx > -1 &&
offset_idx < static_cast<int>(block_offsets.size())) {
int local_row = doc_id - block_offsets[offset_idx];
return persist_stores_[block_index]->fetch(columns, local_row);
int block_row = segment_doc_id - block_offsets[offset_idx];
return persist_stores_[block_index]->fetch(columns, block_row);
}
// Check memory store
@ -2725,14 +2721,15 @@ ExecBatchPtr SegmentImpl::fetch(const std::vector<std::string> &columns,
: block_offsets.back() + block_metas.back().doc_count_;
const auto &mem_block = segment_meta_->writing_forward_block().value();
if (mem_offset <= doc_id &&
doc_id < mem_offset + static_cast<int>(mem_block.doc_count_)) {
int local_row = doc_id - mem_offset;
return memory_store_->fetch(columns, local_row);
if (mem_offset <= segment_doc_id &&
segment_doc_id <
mem_offset + static_cast<int>(mem_block.doc_count_)) {
int block_row = segment_doc_id - mem_offset;
return memory_store_->fetch(columns, block_row);
}
}
} else {
auto table = fetch(columns, std::vector<int>{doc_id});
auto table = fetch(columns, std::vector<int>{segment_doc_id});
if (table) {
std::vector<arrow::Datum> datums;
for (const auto &col : table->columns()) {
@ -2749,7 +2746,7 @@ ExecBatchPtr SegmentImpl::fetch(const std::vector<std::string> &columns,
}
}
LOG_ERROR("Document ID %d not found in persist segment", doc_id);
LOG_ERROR("Segment doc ID %d not found in persist segment", segment_doc_id);
return nullptr;
}
@ -2767,6 +2764,8 @@ RecordBatchReaderPtr SegmentImpl::scan(
std::map<std::pair<int64_t, int64_t>,
std::vector<std::shared_ptr<arrow::ipc::RecordBatchReader>>>
block_groups;
bool emit_segment_row_id =
std::find(columns.begin(), columns.end(), LOCAL_ROW_ID) != columns.end();
for (size_t i = 0; i < scalar_blocks.size() && i < persist_stores_.size();
++i) {
@ -2779,6 +2778,9 @@ RecordBatchReaderPtr SegmentImpl::scan(
interested_cols.push_back(col);
}
}
if (interested_cols.empty() && emit_segment_row_id) {
interested_cols.push_back(GLOBAL_DOC_ID);
}
if (interested_cols.empty()) {
continue;
@ -2794,7 +2796,16 @@ RecordBatchReaderPtr SegmentImpl::scan(
}
if (memory_store_ && memory_store_->num_rows() > 0) {
auto reader = memory_store_->scan(columns);
std::vector<std::string> memory_scan_columns;
for (const auto &col : columns) {
if (col != LOCAL_ROW_ID) {
memory_scan_columns.push_back(col);
}
}
if (memory_scan_columns.empty()) {
memory_scan_columns.push_back(GLOBAL_DOC_ID);
}
auto reader = memory_store_->scan(memory_scan_columns);
if (reader) {
auto &mem_block = segment_meta_->writing_forward_block().value();
auto key = std::make_pair(mem_block.min_doc_id(), mem_block.max_doc_id());
@ -2834,8 +2845,8 @@ RecordBatchReaderPtr SegmentImpl::scan(
}
}
return std::make_shared<CombinedRecordBatchReader>(
shared_from_this(), std::move(merged_readers), columns);
return std::make_shared<CombinedRecordBatchReader>(std::move(merged_readers),
columns);
}
@ -2844,13 +2855,11 @@ RecordBatchReaderPtr SegmentImpl::scan(
////////////////////////////////////////////////////////////////////////////////////
SegmentImpl::CombinedRecordBatchReader::CombinedRecordBatchReader(
std::shared_ptr<const SegmentImpl> segment,
std::vector<std::shared_ptr<arrow::RecordBatchReader>> readers,
const std::vector<std::string> &columns)
: segment_(segment),
readers_(std::move(readers)),
: readers_(std::move(readers)),
current_reader_index_(0),
local_doc_id_(0) {
next_segment_row_id_to_emit_(0) {
if (!readers_.empty()) {
auto schema = readers_[0]->schema();
std::vector<std::shared_ptr<arrow::Field>> selected_fields;
@ -2859,8 +2868,8 @@ SegmentImpl::CombinedRecordBatchReader::CombinedRecordBatchReader(
if (col_name == LOCAL_ROW_ID) {
selected_fields.push_back(
arrow::field(LOCAL_ROW_ID, arrow::uint64(), false));
need_local_doc_id_ = true;
local_doc_id_col_index_ = static_cast<int>(i);
emit_segment_row_id_ = true;
segment_row_id_output_col_index_ = static_cast<int>(i);
} else {
if (auto field = schema->GetFieldByName(col_name); field) {
selected_fields.push_back(field);
@ -2869,17 +2878,6 @@ SegmentImpl::CombinedRecordBatchReader::CombinedRecordBatchReader(
}
projected_schema_ = arrow::schema(selected_fields);
auto segment_meta = segment_->meta();
const auto &blocks = segment_meta->persisted_blocks();
for (const auto &block : blocks) {
if (block.type() != BlockType::SCALAR) continue;
offsets_.push_back(block.min_doc_id_);
}
if (segment_meta->has_writing_forward_block()) {
const auto &mem_block = segment_meta->writing_forward_block().value();
offsets_.push_back(mem_block.min_doc_id_);
}
}
}
@ -2899,24 +2897,25 @@ arrow::Status SegmentImpl::CombinedRecordBatchReader::ReadNext(
return status;
}
if (need_local_doc_id_ && *batch) {
if (emit_segment_row_id_ && *batch) {
auto num_rows = (*batch)->num_rows();
arrow::UInt64Builder builder;
ARROW_RETURN_NOT_OK(builder.Reserve(num_rows));
for (int64_t i = 0; i < num_rows; ++i) {
builder.UnsafeAppend(local_doc_id_++);
builder.UnsafeAppend(next_segment_row_id_to_emit_++);
}
std::shared_ptr<arrow::Array> local_id_array;
ARROW_RETURN_NOT_OK(builder.Finish(&local_id_array));
std::shared_ptr<arrow::Array> segment_row_id_array;
ARROW_RETURN_NOT_OK(builder.Finish(&segment_row_id_array));
auto result =
(*batch)->AddColumn(local_doc_id_col_index_,
(*batch)->AddColumn(segment_row_id_output_col_index_,
projected_schema_->GetFieldByName(LOCAL_ROW_ID),
std::move(local_id_array));
if (result.ok()) {
*batch = std::move(result.ValueOrDie());
std::move(segment_row_id_array));
if (!result.ok()) {
return result.status();
}
*batch = std::move(result.ValueOrDie());
}
if (*batch) {
@ -2924,9 +2923,6 @@ arrow::Status SegmentImpl::CombinedRecordBatchReader::ReadNext(
}
current_reader_index_++;
if (current_reader_index_ < readers_.size()) {
local_doc_id_ = offsets_[current_reader_index_];
}
}
*batch = nullptr;
@ -4404,14 +4400,13 @@ BlockID SegmentImpl::allocate_block_id() {
return block_id_allocator_.fetch_add(1);
}
Result<uint64_t> SegmentImpl::get_global_doc_id(uint32_t local_id) const {
Result<uint64_t> SegmentImpl::get_global_doc_id(uint32_t segment_doc_id) const {
std::lock_guard lock(seg_mtx_);
if (local_id >= doc_ids_.size()) {
if (segment_doc_id >= doc_ids_.size()) {
return tl::make_unexpected(
Status::InvalidArgument("local_id out of range"));
Status::InvalidArgument("segment_doc_id out of range"));
}
// global doc_id
return doc_ids_[local_id];
return doc_ids_[segment_doc_id];
}
@ -4679,4 +4674,4 @@ Result<std::vector<fts::FtsResult>> SegmentImpl::fts_search(
return std::move(ret.value());
}
} // namespace zvec
} // namespace zvec

31
src/db/index/segment/segment.h
View File

@ -64,9 +64,10 @@ class Segment {
virtual SegmentMeta::Ptr meta() const = 0;
// Count documents visible to an optional global-doc-ID filter.
virtual uint64_t doc_count(const IndexFilter::Ptr filter = nullptr) = 0;
// for collection
// ---- Schema and index mutation -----------------------------------------
virtual Status add_column(FieldSchema::Ptr column_schema,
const std::string &expression,
const AddColumnOptions &options) = 0;
@ -84,7 +85,6 @@ class Segment {
std::unordered_map<std::string, VectorColumnIndexer::Ptr>
*quant_vector_indexers) = 0;
// defined in segment.h cause it needs to access block_id generator
virtual Status create_vector_index(
const std::string &column, const IndexParams::Ptr &index_params,
int concurrency, SegmentMeta::Ptr *new_segment_meta,
@ -111,13 +111,11 @@ class Segment {
virtual bool all_vector_index_ready() const = 0;
// defined in segment.h cause it needs to access block_id generator
virtual Status create_scalar_index(
const std::vector<std::string> &columns,
const IndexParams::Ptr &index_params, SegmentMeta::Ptr *new_segment_meta,
InvertedIndexer::Ptr *new_scalar_indexer) = 0;
// defined in segment.h cause it needs to access block_id generator
virtual Status drop_scalar_index(
const std::vector<std::string> &columns,
SegmentMeta::Ptr *new_segment_meta,
@ -127,6 +125,7 @@ class Segment {
const CollectionSchema &schema, const SegmentMeta::Ptr &segment_meta,
const InvertedIndexer::Ptr &scalar_indexer) = 0;
// ---- Data operations ----------------------------------------------------
virtual Status Insert(Doc &doc) = 0;
virtual Status Upsert(Doc &doc) = 0;
@ -142,53 +141,51 @@ class Segment {
&output_fields = std::nullopt,
bool include_vector = true) = 0;
// for sqlengine
virtual TablePtr fetch(const std::vector<std::string> &columns,
const std::vector<int> &indices) const = 0;
const std::vector<int> &segment_doc_ids) const = 0;
virtual ExecBatchPtr fetch(const std::vector<std::string> &columns,
int index) const = 0;
int segment_doc_id) const = 0;
// caller should hold segment shared_ptr for segment handle the indexer's
// lifetime
// Keep Segment alive while consuming the returned reader.
virtual RecordBatchReaderPtr scan(
const std::vector<std::string> &columns) const = 0;
// caller hold segment shared_ptr for segment handle the indexer's lifetime
// ---- Index accessors ----------------------------------------------------
// Keep Segment alive while using returned indexers.
virtual CombinedVectorColumnIndexer::Ptr get_combined_vector_indexer(
const std::string &field_name) const = 0;
// caller hold segment shared_ptr for segment handle the indexer's lifetime
virtual CombinedVectorColumnIndexer::Ptr get_quant_combined_vector_indexer(
const std::string &field_name) const = 0;
// caller hold segment shared_ptr for segment handle the indexer's lifetime
virtual std::vector<VectorColumnIndexer::Ptr> get_vector_indexer(
const std::string &field_name) const = 0;
virtual std::vector<VectorColumnIndexer::Ptr> get_quant_vector_indexer(
const std::string &field_name) const = 0;
// caller hold segment shared_ptr for segment handle the indexer's lifetime
virtual InvertedColumnIndexer::Ptr get_scalar_indexer(
const std::string &field_name) const = 0;
// caller hold segment shared_ptr for segment handle the indexer's lifetime
virtual fts::FtsColumnIndexerPtr get_fts_indexer(
const std::string &field_name) const = 0;
// ---- Index queries and filters -----------------------------------------
virtual Result<std::vector<fts::FtsResult>> fts_search(
const std::string &field_name, const fts::FtsAstNode &ast,
const fts::FtsQueryParams &params) = 0;
// Returned filter is evaluated with segment-local row IDs. It translates the
// local row ID to a global doc ID before consulting the delete store.
virtual const IndexFilter::Ptr get_filter() = 0;
// for others
// ---- Persistence and lifecycle -----------------------------------------
virtual Status flush() = 0;
virtual Status dump() = 0;
// only mark need_destroyed
virtual Status destroy() = 0;
};
} // namespace zvec
} // namespace zvec

234
tests/db/index/segment/segment_row_id_test.cc Normal file
View File

@ -0,0 +1,234 @@
// Copyright 2025-present the zvec 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 <cstdint>
#include <memory>
#include <string>
#include <vector>
#include <arrow/array.h>
#include <arrow/record_batch.h>
#include <gtest/gtest.h>
#include "db/common/constants.h"
#include "segment_test_fixture.h"
using namespace zvec;
namespace {
struct LocalRowIdProjection {
std::vector<std::string> columns;
int row_id_column_index;
};
const std::vector<LocalRowIdProjection> &LocalRowIdProjections() {
static const std::vector<LocalRowIdProjection> projections = {
{{LOCAL_ROW_ID, "id", "name"}, 0},
{{"id", LOCAL_ROW_ID, "name"}, 1},
{{"id", "name", LOCAL_ROW_ID}, 2},
};
return projections;
}
void ExpectSingleRowLocalRowID(const ExecBatchPtr &batch, int column_index,
uint64_t expected) {
ASSERT_TRUE(batch != nullptr);
EXPECT_EQ(batch->length, 1);
EXPECT_EQ(batch->values.size(), 3);
auto local_row_id_scalar = batch->values[column_index].scalar();
ASSERT_TRUE(local_row_id_scalar != nullptr);
auto local_row_id_value =
std::dynamic_pointer_cast<arrow::UInt64Scalar>(local_row_id_scalar);
ASSERT_TRUE(local_row_id_value != nullptr);
EXPECT_EQ(local_row_id_value->value, expected);
}
void ExpectScanLocalRowIDColumn(const RecordBatchReaderPtr &reader,
int column_index, uint32_t expected_rows) {
ASSERT_TRUE(reader != nullptr);
ASSERT_TRUE(reader->schema() != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
uint32_t total_doc = 0;
while (true) {
auto status = reader->ReadNext(&batch);
ASSERT_TRUE(status.ok()) << status.ToString();
if (batch == nullptr) break;
ASSERT_GT(batch->num_columns(), column_index);
EXPECT_EQ(batch->column(column_index)->type()->id(), arrow::Type::UINT64);
EXPECT_EQ(batch->column_name(column_index), LOCAL_ROW_ID);
total_doc += batch->num_rows();
}
EXPECT_EQ(total_doc, expected_rows);
}
void ExpectFetchedLocalRowIDs(
const TablePtr &table, int column_index,
const std::vector<int> &expected_segment_doc_ids) {
ASSERT_TRUE(table != nullptr);
EXPECT_EQ(table->num_columns(), 3);
EXPECT_EQ(table->num_rows(),
static_cast<int64_t>(expected_segment_doc_ids.size()));
auto field = table->schema()->field(column_index);
EXPECT_EQ(field->name(), LOCAL_ROW_ID);
auto id_column = table->column(column_index);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(id_column->chunk(0));
ASSERT_TRUE(id_array != nullptr);
std::vector<uint64_t> actual_ids;
actual_ids.reserve(id_array->length());
for (int64_t i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
std::vector<uint64_t> expected_u64_ids(expected_segment_doc_ids.begin(),
expected_segment_doc_ids.end());
EXPECT_EQ(actual_ids, expected_u64_ids)
<< "LOCAL_ROW_ID values don't match expected order";
}
} // namespace
TEST_P(SegmentTest, FetchSingleRowWithLocalRowIDInRequestedPosition) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_,
options_, 0, 10);
ASSERT_TRUE(segment != nullptr);
for (const auto &projection : LocalRowIdProjections()) {
ExecBatchPtr batch = segment->fetch(projection.columns, 4);
ExpectSingleRowLocalRowID(batch, projection.row_id_column_index, 4);
}
}
TEST_P(SegmentTest, ScanAndFetchPreserveLocalRowIDColumnPosition) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 0, id_map_, delete_store_, version_manager_,
options_, 0, 10);
ASSERT_TRUE(segment != nullptr);
std::vector<int> segment_doc_ids = {0, 3, 6, 1, 0};
for (const auto &projection : LocalRowIdProjections()) {
auto reader = segment->scan(projection.columns);
ExpectScanLocalRowIDColumn(reader, projection.row_id_column_index, 10);
auto table = segment->fetch(projection.columns, segment_doc_ids);
ExpectFetchedLocalRowIDs(table, projection.row_id_column_index,
segment_doc_ids);
}
}
TEST_P(SegmentTest, ScanLocalRowIDIsSegmentLocal) {
options_.max_buffer_size_ = 1 * 1024;
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 100, id_map_, delete_store_, version_manager_,
options_, 100, 25);
ASSERT_TRUE(segment != nullptr);
auto reader = segment->scan({LOCAL_ROW_ID, "id"});
ASSERT_TRUE(reader != nullptr);
std::shared_ptr<arrow::RecordBatch> batch;
std::vector<uint64_t> actual_ids;
while (true) {
auto status = reader->ReadNext(&batch);
ASSERT_TRUE(status.ok()) << status.ToString();
if (batch == nullptr) break;
ASSERT_EQ(batch->num_columns(), 2);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(batch->column(0));
ASSERT_TRUE(id_array != nullptr);
for (int64_t i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
}
std::vector<uint64_t> expected_ids;
for (uint64_t i = 0; i < 25; ++i) {
expected_ids.push_back(i);
}
EXPECT_EQ(actual_ids, expected_ids);
}
TEST_P(SegmentTest, ScanOnlyLocalRowIDDoesNotExposeGlobalDocID) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 100, id_map_, delete_store_, version_manager_,
options_, 100, 10);
ASSERT_TRUE(segment != nullptr);
auto reader = segment->scan({LOCAL_ROW_ID});
ASSERT_TRUE(reader != nullptr);
ASSERT_TRUE(reader->schema() != nullptr);
ASSERT_EQ(reader->schema()->num_fields(), 1);
EXPECT_EQ(reader->schema()->field(0)->name(), LOCAL_ROW_ID);
std::shared_ptr<arrow::RecordBatch> batch;
std::vector<uint64_t> actual_ids;
while (true) {
auto status = reader->ReadNext(&batch);
ASSERT_TRUE(status.ok()) << status.ToString();
if (batch == nullptr) break;
ASSERT_EQ(batch->num_columns(), 1);
EXPECT_EQ(batch->column_name(0), LOCAL_ROW_ID);
auto id_array =
std::dynamic_pointer_cast<arrow::UInt64Array>(batch->column(0));
ASSERT_TRUE(id_array != nullptr);
for (int64_t i = 0; i < id_array->length(); ++i) {
actual_ids.push_back(id_array->Value(i));
}
}
std::vector<uint64_t> expected_ids;
for (uint64_t i = 0; i < 10; ++i) {
expected_ids.push_back(i);
}
EXPECT_EQ(actual_ids, expected_ids);
}
TEST_P(SegmentTest, DocCountDeleteFilterWithNonZeroGlobalDocID) {
auto segment = test::TestHelper::CreateSegmentWithDoc(
col_path_, *schema_, 0, 100, id_map_, delete_store_, version_manager_,
options_, 0, 10);
ASSERT_TRUE(segment != nullptr);
auto status = segment->Delete("pk_5");
EXPECT_TRUE(status.ok()) << "Delete by pk failed: " << status.message();
status = segment->Delete(103);
EXPECT_TRUE(status.ok()) << "Delete by global doc id failed: "
<< status.message();
EXPECT_EQ(segment->doc_count(), 10);
EXPECT_EQ(segment->doc_count(delete_store_->make_filter()), 8);
}
INSTANTIATE_TEST_SUITE_P(MMapTest, SegmentTest, testing::Values(true, false));

962
tests/db/index/segment/segment_test.cc
View File

File diff suppressed because it is too large Load Diff

105
tests/db/index/segment/segment_test_fixture.h Normal file
View File

@ -0,0 +1,105 @@
// Copyright 2025-present the zvec 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 <memory>
#include <stdexcept>
#include <string>
#include <gtest/gtest.h>
#include <zvec/ailego/buffer/block_eviction_queue.h>
#include "db/common/file_helper.h"
#include "db/index/common/delete_store.h"
#include "db/index/common/id_map.h"
#include "db/index/common/version_manager.h"
#include "utils/utils.h"
#include "zvec/db/options.h"
namespace zvec {
class SegmentTest : public testing::TestWithParam<bool> {
protected:
void SetUp() override {
ailego::LoggerBroker::SetLevel(ailego::Logger::LEVEL_WARN);
zvec::ailego::MemoryLimitPool::get_instance().init(MIN_MEMORY_LIMIT_BYTES);
FileHelper::RemoveDirectory(col_path_);
FileHelper::CreateDirectory(col_path_);
auto id_map_path = FileHelper::MakeFilePath(col_path_, FileID::ID_FILE, 0);
id_map_ = IDMap::CreateAndOpen(col_name_, id_map_path, true, false);
if (id_map_ == nullptr) {
throw std::runtime_error("Failed to create id map");
}
delete_store_ = std::make_shared<DeleteStore>(col_name_);
schema_ =
test::TestHelper::CreateSchemaWithScalarIndex(false, false, col_name_);
schema_->add_field(
std::make_shared<FieldSchema>("id", DataType::INT32, false));
schema_->add_field(
std::make_shared<FieldSchema>("name", DataType::STRING, false));
schema_->add_field(
std::make_shared<FieldSchema>("age", DataType::UINT32, false));
schema_->add_field(
std::make_shared<FieldSchema>("binary", DataType::BINARY, false));
schema_->add_field(std::make_shared<FieldSchema>(
"array_binary", DataType::ARRAY_BINARY, false));
bool enable_mmap = GetParam();
Version version;
version.set_schema(*schema_);
version.set_enable_mmap(enable_mmap);
auto version_manager_tmp = VersionManager::Create(col_path_, version);
if (!version_manager_tmp.has_value()) {
throw std::runtime_error("Failed to create version manager");
}
version_manager_ = version_manager_tmp.value();
options_.read_only_ = false;
options_.enable_mmap_ = enable_mmap;
options_.max_buffer_size_ = 64 * 1024 * 1024;
}
void TearDown() override {
id_map_.reset();
delete_store_.reset();
version_manager_.reset();
FileHelper::RemoveDirectory(col_path_);
}
public:
std::string GetColPath() {
return col_path_;
}
protected:
std::string col_name_ = "test_segment";
std::string col_path_ = "./test_collection";
IDMap::Ptr id_map_;
DeleteStore::Ptr delete_store_;
VersionManager::Ptr version_manager_;
CollectionSchema::Ptr schema_;
SegmentOptions options_;
};
} // namespace zvec

2
tests/db/sqlengine/mock_segment.h
View File

@ -333,7 +333,7 @@ class MockSegment : public Segment {
}
ExecBatchPtr fetch(const std::vector<std::string> &columns,
int index) const override {
int segment_doc_id) const override {
LOG_ERROR("Not implemented");
return nullptr;
}