furyhawk/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-07-21 10:15:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: align expression cursor with scalar index raw data (#50233)

Browse Source 
issue: #50197

---------

Signed-off-by: Shawn Wang <shawn.wang@zilliz.com>
This commit is contained in:
sparknack
2026-06-16 16:34:23 +08:00
committed by GitHub
parent 5cd7cbd3f5
commit 778f53fa70
8 changed files with 592 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/core/src/exec/expression/ColumnExpr.cpp
+7 -13
View File
@@ -98,20 +98,19 @@ PhyColumnExpr::DoEval(OffsetVector* input) {
TargetBitmapView valid_res(res_vec->GetValidRawData(), real_batch_size);
valid_res.set();
auto data_barrier = segment_chunk_reader_.segment_->num_chunk_data(
expr_->GetColumn().field_id_);
int64_t processed_rows = 0;
const auto size_per_chunk = segment_chunk_reader_.SizePerChunk();
for (auto i = 0; i < real_batch_size; ++i) {
auto offset = (*input)[i];
auto [chunk_id,
chunk_offset] = [&]() -> std::pair<int64_t, int64_t> {
if (segment_chunk_reader_.segment_->type() ==
SegmentType::Growing) {
const auto size_per_chunk =
segment_chunk_reader_.SizePerChunk();
return {offset / size_per_chunk, offset % size_per_chunk};
} else if (segment_chunk_reader_.segment_->is_chunked() &&
data_barrier > 0) {
segment_chunk_reader_.segment_->num_chunk_data(
expr_->GetColumn().field_id_) > 0) {
return segment_chunk_reader_.segment_->get_chunk_by_offset(
expr_->GetColumn().field_id_, offset);
} else {
@@ -122,8 +121,7 @@ PhyColumnExpr::DoEval(OffsetVector* input) {
expr_->GetColumn().data_type_,
expr_->GetColumn().field_id_,
chunk_id,
data_barrier,
pinned_index_);
PinnedIndexForRawLookup());
auto chunk_data_by_offset = cda(chunk_offset);
if (!chunk_data_by_offset.has_value()) {
valid_res[processed_rows] = false;
@@ -153,7 +151,7 @@ PhyColumnExpr::DoEval(OffsetVector* input) {
expr_->GetColumn().field_id_,
current_chunk_id_,
current_chunk_pos_,
pinned_index_);
PinnedIndexForRawLookup());
for (int i = 0; i < real_batch_size; ++i) {
auto data = cda();
if (!data.has_value()) {
@@ -175,9 +173,6 @@ PhyColumnExpr::DoEval(OffsetVector* input) {
TargetBitmapView valid_res(res_vec->GetValidRawData(), real_batch_size);
valid_res.set();
auto data_barrier = segment_chunk_reader_.segment_->num_chunk_data(
expr_->GetColumn().field_id_);
int64_t processed_rows = 0;
for (int64_t chunk_id = current_chunk_id_; chunk_id < num_chunk_;
++chunk_id) {
@@ -190,8 +185,7 @@ PhyColumnExpr::DoEval(OffsetVector* input) {
expr_->GetColumn().data_type_,
expr_->GetColumn().field_id_,
chunk_id,
data_barrier,
pinned_index_);
PinnedIndexForRawLookup());
for (int i = chunk_id == current_chunk_id_ ? current_chunk_pos_ : 0;
i < chunk_size;
internal/core/src/exec/expression/ColumnExpr.h
+34 -11
View File
@@ -59,13 +59,16 @@ class PhyColumnExpr : public Expr {
auto& field_meta = schema[expr_->GetColumn().field_id_];
pinned_index_ = PinIndex(op_ctx_, segment, field_meta);
is_indexed_ = pinned_index_.size() > 0;
use_index_data_ =
is_indexed_ &&
segment->HasRawData(expr_->GetColumn().field_id_.get());
if (segment->is_chunked()) {
num_chunk_ =
is_indexed_
use_index_data_
? pinned_index_.size()
: segment->num_chunk_data(expr_->GetColumn().field_id_);
} else {
num_chunk_ = is_indexed_
num_chunk_ = use_index_data_
? pinned_index_.size()
: upper_div(segment_chunk_reader_.active_count_,
segment_chunk_reader_.SizePerChunk());
@@ -83,12 +86,16 @@ class PhyColumnExpr : public Expr {
MoveCursor() override {
if (!has_offset_input_) {
if (segment_chunk_reader_.segment_->is_chunked()) {
segment_chunk_reader_.MoveCursorForMultipleChunk(
current_chunk_id_,
current_chunk_pos_,
expr_->GetColumn().field_id_,
num_chunk_,
batch_size_);
if (use_index_data_) {
MoveCursorForIndexed();
} else {
segment_chunk_reader_.MoveCursorForMultipleChunk(
current_chunk_id_,
current_chunk_pos_,
expr_->GetColumn().field_id_,
num_chunk_,
batch_size_);
}
} else {
segment_chunk_reader_.MoveCursorForSingleChunk(
current_chunk_id_,
@@ -99,13 +106,21 @@ class PhyColumnExpr : public Expr {
}
}
void
MoveCursorForIndexed() {
current_chunk_pos_ = current_chunk_pos_ + batch_size_ >=
segment_chunk_reader_.active_count_
? segment_chunk_reader_.active_count_
: current_chunk_pos_ + batch_size_;
}
private:
int64_t
GetCurrentRows() const {
if (segment_chunk_reader_.segment_->is_chunked()) {
auto current_rows =
is_indexed_ && segment_chunk_reader_.segment_->type() ==
SegmentType::Sealed
use_index_data_ && segment_chunk_reader_.segment_->type() ==
SegmentType::Sealed
? current_chunk_pos_
: segment_chunk_reader_.segment_->num_rows_until_chunk(
expr_->GetColumn().field_id_, current_chunk_id_) +
@@ -148,8 +163,16 @@ class PhyColumnExpr : public Expr {
return false;
}
private:
segcore::PinnedIndexView
PinnedIndexForRawLookup() const {
if (!use_index_data_) {
return {};
}
return {pinned_index_.data(), pinned_index_.size()};
}
bool is_indexed_;
bool use_index_data_;
int64_t num_chunk_{0};
int64_t current_chunk_id_{0};
internal/core/src/exec/expression/CompareExpr.cpp
+16 -25
View File
@@ -61,9 +61,10 @@ PhyCompareFilterExpr::ExecCompareExprDispatcher(OpType op, EvalCtx& context) {
TargetBitmapView res(res_vec->GetRawData(), real_batch_size);
TargetBitmapView valid_res(res_vec->GetValidRawData(), real_batch_size);
auto left_data_barrier = segment_chunk_reader_.segment_->num_chunk_data(
expr_->left_field_id_);
auto right_data_barrier =
auto left_raw_data_chunk_count =
segment_chunk_reader_.segment_->num_chunk_data(
expr_->left_field_id_);
auto right_raw_data_chunk_count =
segment_chunk_reader_.segment_->num_chunk_data(
expr_->right_field_id_);
@@ -72,13 +73,13 @@ PhyCompareFilterExpr::ExecCompareExprDispatcher(OpType op, EvalCtx& context) {
for (auto i = 0; i < real_batch_size; ++i) {
auto offset = (*input)[i];
auto get_chunk_id_and_offset =
[&](const FieldId field,
const int64_t data_barrier) -> std::pair<int64_t, int64_t> {
[&](const FieldId field, const int64_t raw_data_chunk_count)
-> std::pair<int64_t, int64_t> {
if (segment_chunk_reader_.segment_->type() ==
SegmentType::Growing) {
return {offset / size_per_chunk, offset % size_per_chunk};
} else if (segment_chunk_reader_.segment_->is_chunked() &&
data_barrier > 0) {
raw_data_chunk_count > 0) {
return segment_chunk_reader_.segment_->get_chunk_by_offset(
field, offset);
} else {
@@ -87,21 +88,19 @@ PhyCompareFilterExpr::ExecCompareExprDispatcher(OpType op, EvalCtx& context) {
};
auto [left_chunk_id, left_chunk_offset] =
get_chunk_id_and_offset(left_field_, left_data_barrier);
auto [right_chunk_id, right_chunk_offset] =
get_chunk_id_and_offset(right_field_, right_data_barrier);
get_chunk_id_and_offset(left_field_, left_raw_data_chunk_count);
auto [right_chunk_id, right_chunk_offset] = get_chunk_id_and_offset(
right_field_, right_raw_data_chunk_count);
auto left = segment_chunk_reader_.GetChunkDataAccessor(
expr_->left_data_type_,
expr_->left_field_id_,
left_chunk_id,
left_data_barrier,
pinned_index_left_);
LeftPinnedIndexForRawLookup());
auto right = segment_chunk_reader_.GetChunkDataAccessor(
expr_->right_data_type_,
expr_->right_field_id_,
right_chunk_id,
right_data_barrier,
pinned_index_right_);
RightPinnedIndexForRawLookup());
auto left_opt = left(left_chunk_offset);
auto right_opt = right(right_chunk_offset);
if (!left_opt.has_value() || !right_opt.has_value()) {
@@ -136,13 +135,13 @@ PhyCompareFilterExpr::ExecCompareExprDispatcher(OpType op, EvalCtx& context) {
expr_->left_field_id_,
left_current_chunk_id_,
left_current_chunk_pos_,
pinned_index_left_);
LeftPinnedIndexForRawLookup());
auto right = segment_chunk_reader_.GetMultipleChunkDataAccessor(
expr_->right_data_type_,
expr_->right_field_id_,
right_current_chunk_id_,
right_current_chunk_pos_,
pinned_index_right_);
RightPinnedIndexForRawLookup());
for (int i = 0; i < real_batch_size; ++i) {
auto left_value = left(), right_value = right();
if (!left_value.has_value() || !right_value.has_value()) {
@@ -168,12 +167,6 @@ PhyCompareFilterExpr::ExecCompareExprDispatcher(OpType op, EvalCtx& context) {
TargetBitmapView valid_res(res_vec->GetValidRawData(), real_batch_size);
valid_res.set();
auto left_data_barrier = segment_chunk_reader_.segment_->num_chunk_data(
expr_->left_field_id_);
auto right_data_barrier =
segment_chunk_reader_.segment_->num_chunk_data(
expr_->right_field_id_);
int64_t processed_rows = 0;
for (int64_t chunk_id = current_chunk_id_; chunk_id < num_chunk_;
++chunk_id) {
@@ -186,14 +179,12 @@ PhyCompareFilterExpr::ExecCompareExprDispatcher(OpType op, EvalCtx& context) {
expr_->left_data_type_,
expr_->left_field_id_,
chunk_id,
left_data_barrier,
pinned_index_left_);
LeftPinnedIndexForRawLookup());
auto right = segment_chunk_reader_.GetChunkDataAccessor(
expr_->right_data_type_,
expr_->right_field_id_,
chunk_id,
right_data_barrier,
pinned_index_right_);
RightPinnedIndexForRawLookup());
for (int i = chunk_id == current_chunk_id_ ? current_chunk_pos_ : 0;
i < chunk_size;
internal/core/src/exec/expression/CompareExpr.h
+28 -8
View File
@@ -165,27 +165,30 @@ class PhyCompareFilterExpr : public Expr {
pinned_index_right_ = PinIndex(op_ctx_, segment, right_field_meta);
is_left_indexed_ = pinned_index_left_.size() > 0;
is_right_indexed_ = pinned_index_right_.size() > 0;
left_use_index_data_ =
is_left_indexed_ && segment->HasRawData(left_field_.get());
right_use_index_data_ =
is_right_indexed_ && segment->HasRawData(right_field_.get());
if (segment->is_chunked()) {
left_num_chunk_ =
is_left_indexed_ ? pinned_index_left_.size()
left_use_index_data_ ? pinned_index_left_.size()
: segment->type() == SegmentType::Growing
? upper_div(segment_chunk_reader_.active_count_,
segment_chunk_reader_.SizePerChunk())
: segment->num_chunk_data(left_field_);
right_num_chunk_ =
is_right_indexed_ ? pinned_index_right_.size()
right_use_index_data_ ? pinned_index_right_.size()
: segment->type() == SegmentType::Growing
? upper_div(segment_chunk_reader_.active_count_,
segment_chunk_reader_.SizePerChunk())
: segment->num_chunk_data(right_field_);
num_chunk_ = left_num_chunk_;
} else {
num_chunk_ = is_left_indexed_
num_chunk_ = left_use_index_data_
? pinned_index_left_.size()
: upper_div(segment_chunk_reader_.active_count_,
segment_chunk_reader_.SizePerChunk());
}
AssertInfo(
batch_size_ > 0,
fmt::format("expr batch size should greater than zero, but now: {}",
@@ -206,7 +209,7 @@ class PhyCompareFilterExpr : public Expr {
MoveCursor() override {
if (!has_offset_input_) {
if (segment_chunk_reader_.segment_->is_chunked()) {
if (is_left_indexed_) {
if (left_use_index_data_) {
MoveCursorForIndexed(left_current_chunk_pos_);
} else {
segment_chunk_reader_.MoveCursorForMultipleChunk(
@@ -216,7 +219,7 @@ class PhyCompareFilterExpr : public Expr {
left_num_chunk_,
batch_size_);
}
if (is_right_indexed_) {
if (right_use_index_data_) {
MoveCursorForIndexed(right_current_chunk_pos_);
} else {
segment_chunk_reader_.MoveCursorForMultipleChunk(
@@ -257,12 +260,27 @@ class PhyCompareFilterExpr : public Expr {
}
private:
segcore::PinnedIndexView
LeftPinnedIndexForRawLookup() const {
if (!left_use_index_data_) {
return {};
}
return {pinned_index_left_.data(), pinned_index_left_.size()};
}
segcore::PinnedIndexView
RightPinnedIndexForRawLookup() const {
if (!right_use_index_data_) {
return {};
}
return {pinned_index_right_.data(), pinned_index_right_.size()};
}
int64_t
GetCurrentRows() {
if (segment_chunk_reader_.segment_->is_chunked()) {
auto current_rows =
is_left_indexed_ && segment_chunk_reader_.segment_->HasRawData(
left_field_.get())
left_use_index_data_
? left_current_chunk_pos_
: segment_chunk_reader_.segment_->num_rows_until_chunk(
left_field_, left_current_chunk_id_) +
@@ -577,6 +595,8 @@ class PhyCompareFilterExpr : public Expr {
const FieldId right_field_;
bool is_left_indexed_;
bool is_right_indexed_;
bool left_use_index_data_;
bool right_use_index_data_;
int64_t num_chunk_{0};
int64_t left_num_chunk_{0};
int64_t right_num_chunk_{0};
internal/core/src/mmap/ChunkedColumnGroup.h
+10
View File
@@ -76,6 +76,10 @@ class ChunkedColumnGroup {
PinWrapper<GroupChunk*>
GetGroupChunk(milvus::OpContext* op_ctx, int64_t chunk_id) const {
AssertInfo(
chunk_id >= 0 && chunk_id < num_chunks_,
"[StorageV2] chunk_id out of range: " + std::to_string(chunk_id) +
", num_chunks: " + std::to_string(num_chunks_));
auto ca = SemiInlineGet(slot_->PinCells(op_ctx, {chunk_id}));
auto chunk = ca->get_cell_of(chunk_id);
return PinWrapper<GroupChunk*>(std::move(ca), chunk);
@@ -84,6 +88,12 @@ class ChunkedColumnGroup {
std::shared_ptr<CellAccessor<GroupChunk>>
GetGroupChunks(milvus::OpContext* op_ctx,
const std::vector<int64_t>& chunk_ids) {
for (auto chunk_id : chunk_ids) {
AssertInfo(chunk_id >= 0 && chunk_id < num_chunks_,
"[StorageV2] chunk_id out of range: " +
std::to_string(chunk_id) +
", num_chunks: " + std::to_string(num_chunks_));
}
return SemiInlineGet(slot_->PinCells(op_ctx, chunk_ids));
}
internal/core/src/segcore/ChunkedSegmentSealedStorageV2Test.cpp
+378
View File
@@ -44,6 +44,8 @@
#include "common/Span.h"
#include "common/Types.h"
#include "common/protobuf_utils.h"
#include "exec/expression/EvalCtx.h"
#include "exec/expression/Expr.h"
#include "expr/ITypeExpr.h"
#include "filemanager/InputStream.h"
#include "gtest/gtest.h"
@@ -51,6 +53,7 @@
#include "index/IndexFactory.h"
#include "index/IndexInfo.h"
#include "index/Meta.h"
#include "index/ScalarIndex.h"
#include "milvus-storage/common/config.h"
#include "milvus-storage/filesystem/fs.h"
#include "milvus-storage/packed/writer.h"
@@ -60,6 +63,7 @@
#include "query/ExecPlanNodeVisitor.h"
#include "segcore/ChunkedSegmentSealedImpl.h"
#include "segcore/SegcoreConfig.h"
#include "segcore/SegmentChunkReader.h"
#include "segcore/SegmentSealed.h"
#include "segcore/Types.h"
#include "segcore/storagev1translator/ChunkTranslator.h"
@@ -72,6 +76,98 @@ using namespace milvus;
using namespace milvus::segcore;
using namespace milvus::segcore::storagev1translator;
namespace {
class RawLookupOnlyIndex : public index::ScalarIndex<int64_t> {
public:
RawLookupOnlyIndex() : index::ScalarIndex<int64_t>("raw_lookup_only") {
}
index::ScalarIndexType
GetIndexType() const override {
return index::ScalarIndexType::STLSORT;
}
void
Build(size_t, const int64_t*, const bool* = nullptr) override {
}
const TargetBitmap
In(size_t, const int64_t*) override {
return {};
}
const TargetBitmap
NotIn(size_t, const int64_t*) override {
return {};
}
const TargetBitmap
IsNull() override {
return {};
}
TargetBitmap
IsNotNull() override {
return {};
}
const TargetBitmap
Range(const int64_t&, OpType) override {
return {};
}
const TargetBitmap
Range(const int64_t&, bool, const int64_t&, bool) override {
return {};
}
std::optional<int64_t>
Reverse_Lookup(size_t offset) const override {
last_lookup_offset = offset;
return static_cast<int64_t>(offset);
}
void
Build(const Config& = {}) override {
}
BinarySet
Serialize(const Config& = {}) override {
return {};
}
void
Load(const BinarySet&, const Config& = {}) override {
}
void
Load(milvus::tracer::TraceContext, const Config& = {}) override {
}
int64_t
Count() override {
return 0;
}
int64_t
Size() override {
return 0;
}
index::IndexStatsPtr
Upload(const Config& = {}) override {
return nullptr;
}
const bool
HasRawData() const override {
return true;
}
mutable size_t last_lookup_offset = 0;
};
} // namespace
class TestChunkSegmentStorageV2 : public testing::TestWithParam<bool> {
protected:
segcore::SegmentSealedUPtr
@@ -230,6 +326,78 @@ class TestChunkSegmentStorageV2 : public testing::TestWithParam<bool> {
ASSERT_TRUE(status.ok());
}
int64_t
RowCount() const {
return chunk_num * test_data_count;
}
void
LoadInt64ScalarIndex(const std::string& index_type) {
auto fid = fields.at("int64");
auto file_manager_ctx = storage::FileManagerContext();
file_manager_ctx.fieldDataMeta.field_schema.set_data_type(
milvus::proto::schema::Int64);
file_manager_ctx.fieldDataMeta.field_schema.set_fieldid(fid.get());
file_manager_ctx.fieldDataMeta.field_id = fid.get();
milvus::storage::IndexMeta index_meta;
index_meta.field_id = fid.get();
index_meta.build_id = 1000 + fid.get();
index_meta.index_version = 2000 + fid.get();
file_manager_ctx.indexMeta = index_meta;
index::CreateIndexInfo create_index_info;
create_index_info.field_type = milvus::DataType::INT64;
create_index_info.index_type = index_type;
auto index = index::IndexFactory::GetInstance().CreateScalarIndex(
create_index_info, file_manager_ctx);
std::vector<int64_t> data(RowCount());
std::iota(data.begin(), data.end(), 0);
index->BuildWithRawDataForUT(data.size(), data.data());
segcore::LoadIndexInfo load_index_info;
load_index_info.index_params = GenIndexParams(index.get());
load_index_info.cache_index =
CreateTestCacheIndex("int64_scalar_index", std::move(index));
load_index_info.field_id = fid.get();
segment->LoadIndex(load_index_info);
}
void
LoadString1ScalarIndex(const std::string& index_type) {
auto fid = fields.at("string1");
auto file_manager_ctx = storage::FileManagerContext();
file_manager_ctx.fieldDataMeta.field_schema.set_data_type(
milvus::proto::schema::VarChar);
file_manager_ctx.fieldDataMeta.field_schema.set_fieldid(fid.get());
file_manager_ctx.fieldDataMeta.field_id = fid.get();
milvus::storage::IndexMeta index_meta;
index_meta.field_id = fid.get();
index_meta.build_id = 1000 + fid.get();
index_meta.index_version = 2000 + fid.get();
file_manager_ctx.indexMeta = index_meta;
index::CreateIndexInfo create_index_info;
create_index_info.field_type = milvus::DataType::VARCHAR;
create_index_info.index_type = index_type;
auto index = index::IndexFactory::GetInstance().CreateScalarIndex(
create_index_info, file_manager_ctx);
std::vector<std::string> data;
data.reserve(RowCount());
for (int64_t i = 0; i < RowCount(); ++i) {
data.push_back("test" + std::to_string(i));
}
index->BuildWithRawDataForUT(data.size(), data.data());
segcore::LoadIndexInfo load_index_info;
load_index_info.index_params = GenIndexParams(index.get());
load_index_info.cache_index =
CreateTestCacheIndex("string1_scalar_index", std::move(index));
load_index_info.field_id = fid.get();
segment->LoadIndex(load_index_info);
}
segcore::SegmentSealedUPtr segment;
SchemaPtr schema_;
LoadFieldDataInfo load_info_;
@@ -371,6 +539,216 @@ TEST_P(TestChunkSegmentStorageV2, TestCompareExpr) {
ASSERT_EQ(chunk_num * test_data_count, final.count());
}
TEST_P(TestChunkSegmentStorageV2, TestColumnExprWithScalarIndexRawData) {
LoadInt64ScalarIndex(index::ASCENDING_SORT);
ASSERT_TRUE(segment->HasRawData(fields.at("int64").get()));
auto query_config = std::make_shared<exec::QueryConfig>(
std::unordered_map<std::string, std::string>{
{exec::QueryConfig::kExprEvalBatchSize, "4096"}});
exec::QueryContext query_context("column_expr_scalar_index_raw_data",
segment.get(),
RowCount(),
MAX_TIMESTAMP,
0,
0,
query::PlanOptions(),
query_config);
exec::ExecContext exec_context(&query_context);
std::vector<expr::TypedExprPtr> exprs{std::make_shared<expr::ColumnExpr>(
expr::ColumnInfo(fields.at("int64"), milvus::DataType::INT64))};
exec::ExprSet expr_set(exprs, &exec_context);
exec::EvalCtx eval_context(&exec_context);
int64_t offset = 0;
while (offset < RowCount()) {
std::vector<VectorPtr> results;
expr_set.Eval(eval_context, results);
ASSERT_EQ(1, results.size());
auto column = std::dynamic_pointer_cast<ColumnVector>(results[0]);
ASSERT_NE(column, nullptr);
auto expected_batch_size = std::min<int64_t>(4096, RowCount() - offset);
ASSERT_EQ(expected_batch_size, column->size());
auto values = column->RawAsValues<int64_t>();
for (int64_t i = 0; i < expected_batch_size; ++i) {
ASSERT_TRUE(column->ValidAt(i));
ASSERT_EQ(offset + i, values[i]);
}
offset += expected_batch_size;
}
}
TEST_P(TestChunkSegmentStorageV2,
TestChunkDataAccessorFallsBackWhenPinnedIndexViewIsEmpty) {
SegmentChunkReader reader(nullptr, segment.get(), RowCount());
auto accessor = reader.GetChunkDataAccessor(
milvus::DataType::INT64, fields.at("int64"), 0, {});
auto value = accessor(7);
ASSERT_TRUE(value.has_value());
ASSERT_EQ(7, segcore::get_from_variant<int64_t>(value));
}
TEST_P(TestChunkSegmentStorageV2,
TestChunkDataAccessorUsesGlobalOffsetForFieldLevelScalarIndex) {
auto raw_lookup_index = std::make_unique<RawLookupOnlyIndex>();
std::vector<PinWrapper<const index::IndexBase*>> pinned_indexes;
pinned_indexes.emplace_back(raw_lookup_index.get());
SegmentChunkReader reader(nullptr, segment.get(), RowCount());
auto accessor = reader.GetChunkDataAccessor(
milvus::DataType::INT64,
fields.at("int64"),
1,
{pinned_indexes.data(), pinned_indexes.size()});
auto expected_offset =
segment->num_rows_until_chunk(fields.at("int64"), 1) + 7;
auto value = accessor(7);
ASSERT_TRUE(value.has_value());
ASSERT_EQ(expected_offset, segcore::get_from_variant<int64_t>(value));
ASSERT_EQ(expected_offset, raw_lookup_index->last_lookup_offset);
}
TEST_P(TestChunkSegmentStorageV2,
TestChunkDataAccessorThrowsWhenPinnedIndexAndRawDataAreUnavailable) {
LoadString1ScalarIndex(index::INVERTED_INDEX_TYPE);
segment->DropFieldData(fields.at("string1"));
ASSERT_FALSE(segment->HasRawData(fields.at("string1").get()));
ASSERT_EQ(0, segment->num_chunk_data(fields.at("string1")));
SegmentChunkReader reader(nullptr, segment.get(), RowCount());
EXPECT_THROW(reader.GetChunkDataAccessor(
milvus::DataType::VARCHAR, fields.at("string1"), 0, {}),
SegcoreError);
}
TEST_P(TestChunkSegmentStorageV2,
TestColumnExprOffsetInputFallsBackWhenScalarIndexHasNoRawData) {
LoadInt64ScalarIndex(index::INVERTED_INDEX_TYPE);
ASSERT_FALSE(segment->HasRawData(fields.at("int64").get()));
auto query_config = std::make_shared<exec::QueryConfig>(
std::unordered_map<std::string, std::string>{
{exec::QueryConfig::kExprEvalBatchSize, "4096"}});
exec::QueryContext query_context("column_expr_offset_input",
segment.get(),
RowCount(),
MAX_TIMESTAMP,
0,
0,
query::PlanOptions(),
query_config);
exec::ExecContext exec_context(&query_context);
std::vector<expr::TypedExprPtr> exprs{std::make_shared<expr::ColumnExpr>(
expr::ColumnInfo(fields.at("int64"), milvus::DataType::INT64))};
exec::ExprSet expr_set(exprs, &exec_context);
exec::OffsetVector offsets;
offsets.push_back(7);
offsets.push_back(7000);
exec::EvalCtx eval_context(&exec_context, &offsets);
std::vector<VectorPtr> results;
expr_set.Eval(eval_context, results);
ASSERT_EQ(1, results.size());
auto column = std::dynamic_pointer_cast<ColumnVector>(results[0]);
ASSERT_NE(column, nullptr);
ASSERT_EQ(offsets.size(), column->size());
auto values = column->RawAsValues<int64_t>();
for (size_t i = 0; i < offsets.size(); ++i) {
ASSERT_TRUE(column->ValidAt(i));
ASSERT_EQ(offsets[i], values[i]);
}
}
TEST_P(TestChunkSegmentStorageV2,
TestColumnExprOffsetInputThrowsWhenIndexAndRawDataAreUnavailable) {
LoadString1ScalarIndex(index::INVERTED_INDEX_TYPE);
segment->DropFieldData(fields.at("string1"));
ASSERT_FALSE(segment->HasRawData(fields.at("string1").get()));
ASSERT_EQ(0, segment->num_chunk_data(fields.at("string1")));
auto query_config = std::make_shared<exec::QueryConfig>(
std::unordered_map<std::string, std::string>{
{exec::QueryConfig::kExprEvalBatchSize, "4096"}});
exec::QueryContext query_context("column_expr_offset_input_no_raw_data",
segment.get(),
RowCount(),
MAX_TIMESTAMP,
0,
0,
query::PlanOptions(),
query_config);
exec::ExecContext exec_context(&query_context);
std::vector<expr::TypedExprPtr> exprs{std::make_shared<expr::ColumnExpr>(
expr::ColumnInfo(fields.at("string1"), milvus::DataType::VARCHAR))};
exec::ExprSet expr_set(exprs, &exec_context);
exec::OffsetVector offsets;
offsets.push_back(0);
exec::EvalCtx eval_context(&exec_context, &offsets);
std::vector<VectorPtr> results;
EXPECT_THROW(expr_set.Eval(eval_context, results), SegcoreError);
}
TEST_P(TestChunkSegmentStorageV2,
TestCompareExprSkippedCursorWithScalarIndexWithoutRawData) {
LoadInt64ScalarIndex(index::INVERTED_INDEX_TYPE);
ASSERT_FALSE(segment->HasRawData(fields.at("int64").get()));
proto::plan::GenericValue threshold;
threshold.set_int64_val(12000);
auto range_expr = std::make_shared<expr::UnaryRangeFilterExpr>(
expr::ColumnInfo(fields.at("int64"), milvus::DataType::INT64),
proto::plan::OpType::GreaterEqual,
threshold);
auto right_field = GetParam() ? fields.at("int64") : fields.at("pk");
auto compare_expr =
std::make_shared<expr::CompareExpr>(fields.at("int64"),
right_field,
milvus::DataType::INT64,
milvus::DataType::INT64,
proto::plan::OpType::Equal);
auto conjunct_expr = std::make_shared<expr::LogicalBinaryExpr>(
expr::LogicalBinaryExpr::OpType::And, range_expr, compare_expr);
auto plan = std::make_shared<plan::FilterBitsNode>(DEFAULT_PLANNODE_ID,
conjunct_expr);
auto query_config = std::make_shared<exec::QueryConfig>(
std::unordered_map<std::string, std::string>{
{exec::QueryConfig::kExprEvalBatchSize, "6000"}});
auto query_context =
std::make_shared<exec::QueryContext>(DEAFULT_QUERY_ID,
segment.get(),
RowCount(),
MAX_TIMESTAMP,
0,
0,
query::PlanOptions(),
query_config);
auto plan_fragment = plan::PlanFragment(plan);
auto row =
query::ExecPlanNodeVisitor::ExecuteTask(plan_fragment, query_context);
ASSERT_NE(row, nullptr);
ASSERT_EQ(row->childrens().size(), 1);
auto col_vec = std::dynamic_pointer_cast<ColumnVector>(row->childrens()[0]);
ASSERT_NE(col_vec, nullptr);
BitsetTypeView view(col_vec->GetRawData(), col_vec->size());
BitsetType final(view);
final.flip();
ASSERT_EQ(RowCount() - threshold.int64_val(), final.count());
}
// Test DropFieldData behavior based on parquet file structure.
// In this test setup, the parquet files are organized as:
// - paths[0] contains columns {0, 4, 3} = int64, ts, string2 (multi-field column group)
internal/core/src/segcore/SegmentChunkReader.cpp
+106 -50
View File
@@ -31,19 +31,43 @@
#include "storage/Types.h"
namespace milvus::segcore {
namespace {
std::pair<const index::IndexBase*, int64_t>
GetIndexAndBaseOffset(const SegmentInternalInterface* segment,
FieldId field_id,
int chunk_id,
PinnedIndexView pinned_index) {
if (pinned_index.empty()) {
return {nullptr, 0};
}
if (chunk_id >= 0 && segment->type() == SegmentType::Sealed &&
segment->is_chunked() && pinned_index.size() == 1) {
auto base_offset =
chunk_id == 0 ? 0
: segment->num_rows_until_chunk(field_id, chunk_id);
return {pinned_index[0].get(), base_offset};
}
if (chunk_id >= 0 && static_cast<size_t>(chunk_id) < pinned_index.size()) {
return {pinned_index[static_cast<size_t>(chunk_id)].get(), 0};
}
return {nullptr, 0};
}
} // namespace
template <typename T>
MultipleChunkDataAccessor
SegmentChunkReader::GetMultipleChunkDataAccessor(
FieldId field_id,
int64_t& current_chunk_id,
int64_t& current_chunk_pos,
const std::vector<PinWrapper<const index::IndexBase*>>& pinned_index)
const {
PinnedIndexView pinned_index) const {
const index::IndexBase* index = nullptr;
if (current_chunk_id < pinned_index.size()) {
index = pinned_index[current_chunk_id].get();
}
if (index) {
auto index_ptr = dynamic_cast<const index::ScalarIndex<T>*>(index);
if (index_ptr->HasRawData()) {
@@ -61,6 +85,12 @@ SegmentChunkReader::GetMultipleChunkDataAccessor(
};
}
}
auto num_chunks = segment_->num_chunk_data(field_id);
AssertInfo(current_chunk_id < num_chunks,
"field {} cursor chunk_id {} exceeds num_chunks {}",
field_id.get(),
current_chunk_id,
num_chunks);
// pw is captured by value, each time we need to access a new chunk, we need to
// pin a new Chunk.
auto pw = segment_->chunk_data<T>(op_ctx_, field_id, current_chunk_id);
@@ -69,12 +99,18 @@ SegmentChunkReader::GetMultipleChunkDataAccessor(
auto chunk_valid_data = chunk_info.valid_data();
auto current_chunk_size = segment_->chunk_size(field_id, current_chunk_id);
return [=,
this,
pw = std::move(pw),
&current_chunk_id,
&current_chunk_pos]() mutable -> const data_access_type {
if (current_chunk_pos >= current_chunk_size) {
current_chunk_id++;
current_chunk_pos = 0;
AssertInfo(current_chunk_id < num_chunks,
"field {} cursor chunk_id {} exceeds num_chunks {}",
field_id.get(),
current_chunk_id,
num_chunks);
// the old chunk will be unpinned, pw will now pin the new chunk.
pw = segment_->chunk_data<T>(op_ctx_, field_id, current_chunk_id);
chunk_data = pw.get().data();
@@ -96,13 +132,11 @@ SegmentChunkReader::GetMultipleChunkDataAccessor<std::string>(
FieldId field_id,
int64_t& current_chunk_id,
int64_t& current_chunk_pos,
const std::vector<PinWrapper<const index::IndexBase*>>& pinned_index)
const {
PinnedIndexView pinned_index) const {
const index::IndexBase* index = nullptr;
if (current_chunk_id < pinned_index.size()) {
index = pinned_index[current_chunk_id].get();
}
if (index) {
auto index_ptr =
dynamic_cast<const index::ScalarIndex<std::string>*>(index);
@@ -120,6 +154,12 @@ SegmentChunkReader::GetMultipleChunkDataAccessor<std::string>(
};
}
}
auto num_chunks = segment_->num_chunk_data(field_id);
AssertInfo(current_chunk_id < num_chunks,
"field {} cursor chunk_id {} exceeds num_chunks {}",
field_id.get(),
current_chunk_id,
num_chunks);
if (segment_->type() == SegmentType::Growing &&
!storage::MmapManager::GetInstance()
.GetMmapConfig()
@@ -137,12 +177,18 @@ SegmentChunkReader::GetMultipleChunkDataAccessor<std::string>(
chunk_data,
chunk_valid_data,
current_chunk_size,
num_chunks,
// pw = std::move(pw),
&current_chunk_id,
&current_chunk_pos]() mutable -> const data_access_type {
if (current_chunk_pos >= current_chunk_size) {
current_chunk_id++;
current_chunk_pos = 0;
AssertInfo(current_chunk_id < num_chunks,
"field {} cursor chunk_id {} exceeds num_chunks {}",
field_id.get(),
current_chunk_id,
num_chunks);
pw = segment_->chunk_data<std::string>(
op_ctx_, field_id, current_chunk_id);
chunk_data = pw.get().data();
@@ -163,12 +209,18 @@ SegmentChunkReader::GetMultipleChunkDataAccessor<std::string>(
auto current_chunk_size =
segment_->chunk_size(field_id, current_chunk_id);
return [=,
this,
pw = std::move(pw),
&current_chunk_id,
&current_chunk_pos]() mutable -> const data_access_type {
if (current_chunk_pos >= current_chunk_size) {
current_chunk_id++;
current_chunk_pos = 0;
AssertInfo(current_chunk_id < num_chunks,
"field {} cursor chunk_id {} exceeds num_chunks {}",
field_id.get(),
current_chunk_id,
num_chunks);
pw = segment_->chunk_view<std::string_view>(
op_ctx_, field_id, current_chunk_id);
current_chunk_size =
@@ -192,8 +244,7 @@ SegmentChunkReader::GetMultipleChunkDataAccessor(
FieldId field_id,
int64_t& current_chunk_id,
int64_t& current_chunk_pos,
const std::vector<PinWrapper<const index::IndexBase*>>& pinned_index)
const {
PinnedIndexView pinned_index) const {
switch (data_type) {
case DataType::BOOL:
return GetMultipleChunkDataAccessor<bool>(
@@ -230,25 +281,30 @@ SegmentChunkReader::GetMultipleChunkDataAccessor(
template <typename T>
ChunkDataAccessor
SegmentChunkReader::GetChunkDataAccessor(
FieldId field_id,
int chunk_id,
int data_barrier,
const std::vector<PinWrapper<const index::IndexBase*>>& pinned_index)
const {
if (chunk_id >= data_barrier) {
auto index = pinned_index[chunk_id].get();
auto index_ptr = dynamic_cast<const index::ScalarIndex<T>*>(index);
if (index->HasRawData()) {
return [index_ptr](int i) mutable -> const data_access_type {
auto raw = index_ptr->Reverse_Lookup(i);
SegmentChunkReader::GetChunkDataAccessor(FieldId field_id,
int chunk_id,
PinnedIndexView pinned_index) const {
auto index_and_base_offset =
GetIndexAndBaseOffset(segment_, field_id, chunk_id, pinned_index);
auto index = index_and_base_offset.first;
auto base_offset = index_and_base_offset.second;
auto index_ptr = dynamic_cast<const index::ScalarIndex<T>*>(index);
if (index_ptr != nullptr && index_ptr->HasRawData()) {
return
[index_ptr, base_offset](int i) mutable -> const data_access_type {
auto raw = index_ptr->Reverse_Lookup(base_offset + i);
if (!raw.has_value()) {
return std::nullopt;
}
return raw.value();
};
}
}
auto num_chunks = segment_->num_chunk_data(field_id);
AssertInfo(chunk_id >= 0 && chunk_id < num_chunks,
"field {} chunk_id {} exceeds raw data chunks {}",
field_id.get(),
chunk_id,
num_chunks);
auto pw = segment_->chunk_data<T>(op_ctx_, field_id, chunk_id);
return [pw = std::move(pw)](int i) mutable -> const data_access_type {
auto chunk_info = pw.get();
@@ -264,25 +320,29 @@ SegmentChunkReader::GetChunkDataAccessor(
template <>
ChunkDataAccessor
SegmentChunkReader::GetChunkDataAccessor<std::string>(
FieldId field_id,
int chunk_id,
int data_barrier,
const std::vector<PinWrapper<const index::IndexBase*>>& pinned_index)
const {
if (chunk_id >= data_barrier) {
auto index = pinned_index[chunk_id].get();
auto index_ptr =
dynamic_cast<const index::ScalarIndex<std::string>*>(index);
if (index_ptr->HasRawData()) {
return [index_ptr](int i) mutable -> const data_access_type {
auto raw = index_ptr->Reverse_Lookup(i);
FieldId field_id, int chunk_id, PinnedIndexView pinned_index) const {
auto index_and_base_offset =
GetIndexAndBaseOffset(segment_, field_id, chunk_id, pinned_index);
auto index = index_and_base_offset.first;
auto base_offset = index_and_base_offset.second;
auto index_ptr =
dynamic_cast<const index::ScalarIndex<std::string>*>(index);
if (index_ptr != nullptr && index_ptr->HasRawData()) {
return
[index_ptr, base_offset](int i) mutable -> const data_access_type {
auto raw = index_ptr->Reverse_Lookup(base_offset + i);
if (!raw.has_value()) {
return std::nullopt;
}
return raw.value();
};
}
}
auto num_chunks = segment_->num_chunk_data(field_id);
AssertInfo(chunk_id >= 0 && chunk_id < num_chunks,
"field {} chunk_id {} exceeds raw data chunks {}",
field_id.get(),
chunk_id,
num_chunks);
if (segment_->type() == SegmentType::Growing &&
!storage::MmapManager::GetInstance()
.GetMmapConfig()
@@ -312,40 +372,36 @@ SegmentChunkReader::GetChunkDataAccessor<std::string>(
}
ChunkDataAccessor
SegmentChunkReader::GetChunkDataAccessor(
DataType data_type,
FieldId field_id,
int chunk_id,
int data_barrier,
const std::vector<PinWrapper<const index::IndexBase*>>& pinned_index)
const {
SegmentChunkReader::GetChunkDataAccessor(DataType data_type,
FieldId field_id,
int chunk_id,
PinnedIndexView pinned_index) const {
switch (data_type) {
case DataType::BOOL:
return GetChunkDataAccessor<bool>(
field_id, chunk_id, data_barrier, pinned_index);
return GetChunkDataAccessor<bool>(field_id, chunk_id, pinned_index);
case DataType::INT8:
return GetChunkDataAccessor<int8_t>(
field_id, chunk_id, data_barrier, pinned_index);
field_id, chunk_id, pinned_index);
case DataType::INT16:
return GetChunkDataAccessor<int16_t>(
field_id, chunk_id, data_barrier, pinned_index);
field_id, chunk_id, pinned_index);
case DataType::INT32:
return GetChunkDataAccessor<int32_t>(
field_id, chunk_id, data_barrier, pinned_index);
field_id, chunk_id, pinned_index);
case DataType::TIMESTAMPTZ:
case DataType::INT64:
return GetChunkDataAccessor<int64_t>(
field_id, chunk_id, data_barrier, pinned_index);
field_id, chunk_id, pinned_index);
case DataType::FLOAT:
return GetChunkDataAccessor<float>(
field_id, chunk_id, data_barrier, pinned_index);
field_id, chunk_id, pinned_index);
case DataType::DOUBLE:
return GetChunkDataAccessor<double>(
field_id, chunk_id, data_barrier, pinned_index);
field_id, chunk_id, pinned_index);
case DataType::VARCHAR:
case DataType::TEXT: {
return GetChunkDataAccessor<std::string>(
field_id, chunk_id, data_barrier, pinned_index);
field_id, chunk_id, pinned_index);
}
default:
ThrowInfo(DataTypeInvalid, "unsupported data type: {}", data_type);
internal/core/src/segcore/SegmentChunkReader.h
+13 -19
View File
@@ -21,6 +21,7 @@
#include <string>
#include <vector>
#include <boost/core/span.hpp>
#include "boost/variant/variant.hpp"
#include "cachinglayer/CacheSlot.h"
#include "common/OpContext.h"
@@ -43,6 +44,7 @@ using data_access_type = std::optional<boost::variant<bool,
using ChunkDataAccessor = std::function<const data_access_type(int)>;
using MultipleChunkDataAccessor = std::function<const data_access_type()>;
using PinnedIndexView = boost::span<const PinWrapper<const index::IndexBase*>>;
// Helper to extract a value of type T from data_access_type.
// For std::string, handles both std::string and std::string_view in the variant.
@@ -102,21 +104,17 @@ class SegmentChunkReader {
}
MultipleChunkDataAccessor
GetMultipleChunkDataAccessor(
DataType data_type,
FieldId field_id,
int64_t& current_chunk_id,
int64_t& current_chunk_pos,
const std::vector<PinWrapper<const index::IndexBase*>>& pinned_index)
const;
GetMultipleChunkDataAccessor(DataType data_type,
FieldId field_id,
int64_t& current_chunk_id,
int64_t& current_chunk_pos,
PinnedIndexView pinned_index) const;
ChunkDataAccessor
GetChunkDataAccessor(DataType data_type,
FieldId field_id,
int chunk_id,
int data_barrier,
const std::vector<PinWrapper<const index::IndexBase*>>&
pinned_index) const;
PinnedIndexView pinned_index) const;
void
MoveCursorForMultipleChunk(int64_t& current_chunk_id,
@@ -178,20 +176,16 @@ class SegmentChunkReader {
private:
template <typename T>
MultipleChunkDataAccessor
GetMultipleChunkDataAccessor(
FieldId field_id,
int64_t& current_chunk_id,
int64_t& current_chunk_pos,
const std::vector<PinWrapper<const index::IndexBase*>>& pinned_index)
const;
GetMultipleChunkDataAccessor(FieldId field_id,
int64_t& current_chunk_id,
int64_t& current_chunk_pos,
PinnedIndexView pinned_index) const;
template <typename T>
ChunkDataAccessor
GetChunkDataAccessor(FieldId field_id,
int chunk_id,
int data_barrier,
const std::vector<PinWrapper<const index::IndexBase*>>&
pinned_index) const;
PinnedIndexView pinned_index) const;
const int64_t size_per_chunk_;
milvus::OpContext* op_ctx_;