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

feat: storage v2 seal segment load (#41567)

Browse Source 
storage v2 chunked seal segment loading is based on caching layer. A
cell unit in storage v2 is a parquet row group in remote object storage,
containing all fields. Therefore, each field needs a proxy to do related
one field operations.

<img width="965" alt="Screenshot 2025-04-28 at 10 59 30"
src="https://github.com/user-attachments/assets/83e93a10-3b1d-4066-ac17-b996d5650416"
/>

related: #39173

---------

Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>
This commit is contained in:
sthuang
2025-04-30 14:22:58 +08:00
committed by GitHub
parent 3bd6268d3c
commit e9442f575d
31 changed files with 2317 additions and 542 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/core/src/common/GroupChunk.h
+101
View File
@@ -0,0 +1,101 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <cstdint>
#include <memory>
#include <utility>
#include "common/Types.h"
#include "common/Chunk.h"
namespace milvus {
// GroupChunk represents a collection of chunks for different fields in a group
class GroupChunk {
public:
GroupChunk() = default;
explicit GroupChunk(
const std::unordered_map<FieldId, std::shared_ptr<Chunk>>& chunks)
: chunks_(chunks) {
}
virtual ~GroupChunk() = default;
// Get the chunk for a specific field
std::shared_ptr<Chunk>
GetChunk(FieldId field_id) const {
auto it = chunks_.find(field_id);
if (it == chunks_.end()) {
return nullptr;
}
return it->second;
}
// Add a chunk for a specific field
void
AddChunk(FieldId field_id, std::shared_ptr<Chunk> chunk) {
if (chunks_.find(field_id) != chunks_.end()) {
PanicInfo(ErrorCode::FieldAlreadyExist,
"Field {} already exists in GroupChunk",
field_id.get());
}
chunks_[field_id] = std::move(chunk);
}
uint64_t
Size() const {
uint64_t total_size = 0;
for (const auto& chunk : chunks_) {
total_size += chunk.second->Size();
}
return total_size;
}
// Get all chunks
const std::unordered_map<FieldId, std::shared_ptr<Chunk>>&
GetChunks() const {
return chunks_;
}
size_t
CellByteSize() const {
size_t total_size = 0;
for (const auto& chunk : chunks_) {
total_size += chunk.second->CellByteSize();
}
return total_size;
}
// Get the number of rows in this group chunk
int64_t
RowNums() const {
if (chunks_.empty()) {
return 0;
}
return chunks_.begin()->second->RowNums();
}
// Check if the chunk for a specific field exists
bool
HasChunk(FieldId field_id) const {
return chunks_.find(field_id) != chunks_.end();
}
private:
std::unordered_map<FieldId, std::shared_ptr<Chunk>> chunks_;
};
} // namespace milvus
internal/core/src/common/Schema.h
+9
View File
@@ -237,6 +237,15 @@ class Schema {
return fields_;
}
const std::unordered_map<FieldId, FieldMeta>
get_field_metas(std::vector<FieldId> field_ids) {
std::unordered_map<FieldId, FieldMeta> field_metas;
for (const auto& field_id : field_ids) {
field_metas.emplace(field_id, operator[](field_id));
}
return field_metas;
}
const std::vector<FieldId>&
get_field_ids() const {
return field_ids_;
internal/core/src/index/SkipIndex.h
+18 -6
View File
@@ -12,7 +12,7 @@
#pragma once
#include <unordered_map>
#include "mmap/ChunkedColumnGroup.h"
#include "common/Types.h"
namespace milvus {
@@ -294,8 +294,8 @@ class SkipIndex {
ProcessStringFieldMetrics(const T& var_column) {
int num_rows = var_column.NumRows();
// find first not null value
int64_t start = 0;
for (int64_t i = start; i < num_rows; i++) {
size_t start = 0;
for (size_t i = start; i < num_rows; i++) {
if (!var_column.IsValid(i)) {
start++;
continue;
@@ -305,15 +305,27 @@ class SkipIndex {
if (start > num_rows - 1) {
return {std::string(), std::string(), num_rows};
}
std::string min_string = var_column.RawAt(start);
// Handle both ChunkedVariableColumn and ProxyChunkColumn cases
std::string min_string;
if constexpr (std::is_same_v<T, ProxyChunkColumn>) {
min_string = var_column.template RawAt<std::string>(start);
} else {
min_string = var_column.RawAt(start);
}
std::string max_string = min_string;
int64_t null_count = start;
for (int64_t i = start; i < num_rows; i++) {
for (size_t i = start; i < num_rows; i++) {
if (!var_column.IsValid(i)) {
null_count++;
continue;
}
const auto& val = var_column.RawAt(i);
std::string val;
if constexpr (std::is_same_v<T, ProxyChunkColumn>) {
val = var_column.template RawAt<std::string>(i);
} else {
val = var_column.RawAt(i);
}
if (val < min_string) {
min_string = val;
}
internal/core/src/mmap/ChunkedColumn.h
+55 -21
View File
@@ -26,7 +26,6 @@
#include <vector>
#include <math.h>
#include "cachinglayer/CacheSlot.h"
#include "cachinglayer/Manager.h"
#include "cachinglayer/Translator.h"
#include "cachinglayer/Utils.h"
@@ -37,6 +36,8 @@
#include "common/Span.h"
#include "common/Array.h"
#include "segcore/storagev1translator/ChunkTranslator.h"
#include "cachinglayer/Translator.h"
#include "mmap/ChunkedColumnInterface.h"
namespace milvus {
@@ -55,7 +56,7 @@ std::pair<size_t, size_t> inline GetChunkIDByOffset(
return {chunk_idx, offset_in_chunk};
}
class ChunkedColumnBase {
class ChunkedColumnBase : public ChunkedColumnInterface {
public:
// memory mode ctor
explicit ChunkedColumnBase(std::unique_ptr<Translator<Chunk>> translator,
@@ -69,14 +70,14 @@ class ChunkedColumnBase {
virtual ~ChunkedColumnBase() = default;
PinWrapper<const char*>
DataOfChunk(int chunk_id) const {
DataOfChunk(int chunk_id) const override {
auto ca = SemiInlineGet(slot_->PinCells({chunk_id}));
auto chunk = ca->get_cell_of(chunk_id);
return PinWrapper<const char*>(ca, chunk->Data());
}
bool
IsValid(size_t offset) const {
IsValid(size_t offset) const override {
if (!nullable_) {
return true;
}
@@ -85,7 +86,7 @@ class ChunkedColumnBase {
}
bool
IsValid(int64_t chunk_id, int64_t offset) const {
IsValid(int64_t chunk_id, int64_t offset) const override {
if (nullable_) {
auto ca =
SemiInlineGet(slot_->PinCells({static_cast<cid_t>(chunk_id)}));
@@ -96,23 +97,23 @@ class ChunkedColumnBase {
}
bool
IsNullable() const {
IsNullable() const override {
return nullable_;
}
size_t
NumRows() const {
NumRows() const override {
return num_rows_;
};
int64_t
num_chunks() const {
num_chunks() const override {
return num_chunks_;
}
// This returns only memory byte size.
size_t
DataByteSize() const {
DataByteSize() const override {
auto size = 0;
for (auto i = 0; i < num_chunks_; i++) {
size += slot_->size_of_cell(i).memory_bytes;
@@ -121,13 +122,13 @@ class ChunkedColumnBase {
}
int64_t
chunk_row_nums(int64_t chunk_id) const {
chunk_row_nums(int64_t chunk_id) const override {
return GetNumRowsUntilChunk(chunk_id + 1) -
GetNumRowsUntilChunk(chunk_id);
}
virtual PinWrapper<SpanBase>
Span(int64_t chunk_id) const {
Span(int64_t chunk_id) const override {
PanicInfo(ErrorCode::Unsupported,
"Span only supported for ChunkedColumn");
}
@@ -136,14 +137,15 @@ class ChunkedColumnBase {
std::pair<std::vector<std::string_view>, FixedVector<bool>>>
StringViews(int64_t chunk_id,
std::optional<std::pair<int64_t, int64_t>> offset_len =
std::nullopt) const {
std::nullopt) const override {
PanicInfo(ErrorCode::Unsupported,
"StringViews only supported for VariableColumn");
}
virtual PinWrapper<std::pair<std::vector<ArrayView>, FixedVector<bool>>>
ArrayViews(int64_t chunk_id,
std::optional<std::pair<int64_t, int64_t>> offset_len) const {
ArrayViews(
int64_t chunk_id,
std::optional<std::pair<int64_t, int64_t>> offset_len) const override {
PanicInfo(ErrorCode::Unsupported,
"ArrayViews only supported for ArrayChunkedColumn");
}
@@ -151,13 +153,13 @@ class ChunkedColumnBase {
virtual PinWrapper<
std::pair<std::vector<std::string_view>, FixedVector<bool>>>
ViewsByOffsets(int64_t chunk_id,
const FixedVector<int32_t>& offsets) const {
const FixedVector<int32_t>& offsets) const override {
PanicInfo(ErrorCode::Unsupported,
"ViewsByOffsets only supported for VariableColumn");
}
std::pair<size_t, size_t>
GetChunkIDByOffset(int64_t offset) const {
GetChunkIDByOffset(int64_t offset) const override {
AssertInfo(offset < num_rows_,
"offset {} is out of range, num_rows: {}",
offset,
@@ -167,24 +169,30 @@ class ChunkedColumnBase {
}
PinWrapper<Chunk*>
GetChunk(int64_t chunk_id) const {
GetChunk(int64_t chunk_id) const override {
auto ca = SemiInlineGet(slot_->PinCells({chunk_id}));
auto chunk = ca->get_cell_of(chunk_id);
return PinWrapper<Chunk*>(ca, chunk);
}
int64_t
GetNumRowsUntilChunk(int64_t chunk_id) const {
GetNumRowsUntilChunk(int64_t chunk_id) const override {
return GetNumRowsUntilChunk()[chunk_id];
}
const std::vector<int64_t>&
GetNumRowsUntilChunk() const {
GetNumRowsUntilChunk() const override {
auto meta = static_cast<milvus::segcore::storagev1translator::CTMeta*>(
slot_->meta());
return meta->num_rows_until_chunk_;
}
virtual const char*
ValueAt(int64_t offset) override {
PanicInfo(ErrorCode::Unsupported,
"ValueAt only supported for ChunkedColumn");
}
protected:
bool nullable_{false};
size_t num_rows_{0};
@@ -202,7 +210,7 @@ class ChunkedColumn : public ChunkedColumnBase {
// TODO(tiered storage 1): this method should be replaced with a bulk access method.
const char*
ValueAt(int64_t offset) {
ValueAt(int64_t offset) override {
auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(offset);
auto ca =
SemiInlineGet(slot_->PinCells({static_cast<cid_t>(chunk_id)}));
@@ -285,7 +293,7 @@ class ChunkedArrayColumn : public ChunkedColumnBase {
// TODO(tiered storage 1): this method should be replaced with a bulk access method.
ScalarArray
RawAt(const int i) const {
PrimitivieRawAt(const int i) const override {
auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(i);
auto ca =
SemiInlineGet(slot_->PinCells({static_cast<cid_t>(chunk_id)}));
@@ -306,4 +314,30 @@ class ChunkedArrayColumn : public ChunkedColumnBase {
ca, static_cast<ArrayChunk*>(chunk)->Views(offset_len));
}
};
inline std::shared_ptr<ChunkedColumnInterface>
MakeChunkedColumnBase(DataType data_type,
std::unique_ptr<Translator<milvus::Chunk>> translator,
const FieldMeta& field_meta) {
if (ChunkedColumnInterface::IsChunkedVariableColumnDataType(data_type)) {
if (data_type == DataType::JSON) {
return std::static_pointer_cast<ChunkedColumnInterface>(
std::make_shared<ChunkedVariableColumn<milvus::Json>>(
std::move(translator), field_meta));
}
return std::static_pointer_cast<ChunkedColumnInterface>(
std::make_shared<ChunkedVariableColumn<std::string>>(
std::move(translator), field_meta));
}
if (ChunkedColumnInterface::IsChunkedArrayColumnDataType(data_type)) {
return std::static_pointer_cast<ChunkedColumnInterface>(
std::make_shared<ChunkedArrayColumn>(std::move(translator),
field_meta));
}
return std::static_pointer_cast<ChunkedColumnInterface>(
std::make_shared<ChunkedColumn>(std::move(translator), field_meta));
}
} // namespace milvus
internal/core/src/mmap/ChunkedColumnGroup.h
+295
View File
@@ -0,0 +1,295 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <folly/io/IOBuf.h>
#include <sys/mman.h>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <memory>
#include <vector>
#include <math.h>
#include "cachinglayer/CacheSlot.h"
#include "cachinglayer/Manager.h"
#include "cachinglayer/Translator.h"
#include "cachinglayer/Utils.h"
#include "common/Array.h"
#include "common/Chunk.h"
#include "common/GroupChunk.h"
#include "common/Common.h"
#include "common/EasyAssert.h"
#include "common/Span.h"
#include "common/Array.h"
#include "mmap/ChunkedColumn.h"
#include "mmap/ChunkedColumnInterface.h"
#include "segcore/storagev2translator/GroupCTMeta.h"
namespace milvus {
using GroupChunkVector = std::vector<std::shared_ptr<GroupChunk>>;
using namespace milvus::cachinglayer;
// ChunkedColumnGroup represents a collection of group chunks
class ChunkedColumnGroup {
public:
explicit ChunkedColumnGroup(
std::unique_ptr<Translator<GroupChunk>> translator)
: slot_(Manager::GetInstance().CreateCacheSlot(std::move(translator))) {
num_chunks_ = slot_->num_cells();
num_rows_ = GetNumRowsUntilChunk().back();
}
virtual ~ChunkedColumnGroup() = default;
// Get the number of group chunks
size_t
num_chunks() const {
return num_chunks_;
}
PinWrapper<GroupChunk*>
GetGroupChunk(int64_t chunk_id) const {
auto ca = SemiInlineGet(slot_->PinCells({chunk_id}));
auto chunk = ca->get_cell_of(chunk_id);
return PinWrapper<GroupChunk*>(ca, chunk);
}
int64_t
NumRows() const {
return num_rows_;
}
int64_t
GetNumRowsUntilChunk(int64_t chunk_id) const {
AssertInfo(chunk_id >= 0 && chunk_id <= num_chunks_,
"chunk_id out of range: " + std::to_string(chunk_id));
return GetNumRowsUntilChunk()[chunk_id];
}
const std::vector<int64_t>&
GetNumRowsUntilChunk() const {
auto meta =
static_cast<milvus::segcore::storagev2translator::GroupCTMeta*>(
slot_->meta());
return meta->num_rows_until_chunk_;
}
protected:
mutable std::shared_ptr<CacheSlot<GroupChunk>> slot_;
size_t num_chunks_{0};
size_t num_rows_{0};
};
class ProxyChunkColumn : public ChunkedColumnInterface {
public:
explicit ProxyChunkColumn(std::shared_ptr<ChunkedColumnGroup> group,
FieldId field_id,
const FieldMeta& field_meta)
: group_(group),
field_id_(field_id),
field_meta_(field_meta),
data_type_(field_meta.get_data_type()) {
}
PinWrapper<const char*>
DataOfChunk(int chunk_id) const override {
auto group_chunk = group_->GetGroupChunk(chunk_id);
auto chunk = group_chunk.get()->GetChunk(field_id_);
return PinWrapper<const char*>(group_chunk, chunk->Data());
}
bool
IsValid(size_t offset) const override {
auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(offset);
auto group_chunk = group_->GetGroupChunk(chunk_id);
auto chunk = group_chunk.get()->GetChunk(field_id_);
return chunk->isValid(offset_in_chunk);
}
bool
IsValid(int64_t chunk_id, int64_t offset) const override {
auto group_chunk = group_->GetGroupChunk(chunk_id);
auto chunk = group_chunk.get()->GetChunk(field_id_);
return chunk->isValid(offset);
}
bool
IsNullable() const override {
return field_meta_.is_nullable();
}
size_t
NumRows() const override {
return group_->NumRows();
}
int64_t
num_chunks() const override {
return group_->num_chunks();
}
size_t
DataByteSize() const override {
size_t total_size = 0;
for (int64_t i = 0; i < num_chunks(); ++i) {
auto group_chunk = group_->GetGroupChunk(i);
auto chunk = group_chunk.get()->GetChunk(field_id_);
total_size += chunk->Size();
}
return total_size;
}
int64_t
chunk_row_nums(int64_t chunk_id) const override {
return group_->GetNumRowsUntilChunk(chunk_id + 1) -
group_->GetNumRowsUntilChunk(chunk_id);
}
PinWrapper<SpanBase>
Span(int64_t chunk_id) const override {
if (!IsChunkedColumnDataType(data_type_)) {
PanicInfo(ErrorCode::Unsupported,
"Span only supported for ChunkedColumn");
}
auto chunk_wrapper = group_->GetGroupChunk(chunk_id);
auto chunk = chunk_wrapper.get()->GetChunk(field_id_);
return PinWrapper<SpanBase>(
chunk_wrapper, static_cast<FixedWidthChunk*>(chunk.get())->Span());
}
PinWrapper<std::pair<std::vector<std::string_view>, FixedVector<bool>>>
StringViews(int64_t chunk_id,
std::optional<std::pair<int64_t, int64_t>> offset_len =
std::nullopt) const override {
if (!IsChunkedVariableColumnDataType(data_type_)) {
PanicInfo(ErrorCode::Unsupported,
"StringViews only supported for ChunkedVariableColumn");
}
auto chunk_wrapper = group_->GetGroupChunk(chunk_id);
auto chunk = chunk_wrapper.get()->GetChunk(field_id_);
return PinWrapper<
std::pair<std::vector<std::string_view>, FixedVector<bool>>>(
chunk_wrapper,
static_cast<StringChunk*>(chunk.get())->StringViews(offset_len));
}
PinWrapper<std::pair<std::vector<ArrayView>, FixedVector<bool>>>
ArrayViews(int64_t chunk_id,
std::optional<std::pair<int64_t, int64_t>> offset_len =
std::nullopt) const override {
if (!IsChunkedArrayColumnDataType(data_type_)) {
PanicInfo(ErrorCode::Unsupported,
"ArrayViews only supported for ChunkedArrayColumn");
}
auto chunk_wrapper = group_->GetGroupChunk(chunk_id);
auto chunk = chunk_wrapper.get()->GetChunk(field_id_);
return PinWrapper<std::pair<std::vector<ArrayView>, FixedVector<bool>>>(
chunk_wrapper,
static_cast<ArrayChunk*>(chunk.get())->Views(offset_len));
}
PinWrapper<std::pair<std::vector<std::string_view>, FixedVector<bool>>>
ViewsByOffsets(int64_t chunk_id,
const FixedVector<int32_t>& offsets) const override {
if (!IsChunkedVariableColumnDataType(data_type_)) {
PanicInfo(
ErrorCode::Unsupported,
"ViewsByOffsets only supported for ChunkedVariableColumn");
}
auto chunk_wrapper = group_->GetGroupChunk(chunk_id);
auto chunk = chunk_wrapper.get()->GetChunk(field_id_);
return PinWrapper<
std::pair<std::vector<std::string_view>, FixedVector<bool>>>(
chunk_wrapper,
static_cast<StringChunk*>(chunk.get())->ViewsByOffsets(offsets));
}
std::pair<size_t, size_t>
GetChunkIDByOffset(int64_t offset) const override {
int64_t current_offset = 0;
for (int64_t i = 0; i < num_chunks(); ++i) {
auto rows = chunk_row_nums(i);
if (current_offset + rows > offset) {
return {i, offset - current_offset};
}
current_offset += rows;
}
return {num_chunks() - 1, chunk_row_nums(num_chunks() - 1) - 1};
}
PinWrapper<Chunk*>
GetChunk(int64_t chunk_id) const override {
auto group_chunk = group_->GetGroupChunk(chunk_id);
auto chunk = group_chunk.get()->GetChunk(field_id_);
return PinWrapper<Chunk*>(group_chunk, chunk.get());
}
int64_t
GetNumRowsUntilChunk(int64_t chunk_id) const override {
return group_->GetNumRowsUntilChunk(chunk_id);
}
const std::vector<int64_t>&
GetNumRowsUntilChunk() const override {
return group_->GetNumRowsUntilChunk();
}
const char*
ValueAt(int64_t offset) override {
auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(offset);
auto chunk = GetChunk(chunk_id);
return chunk.get()->ValueAt(offset_in_chunk);
}
template <typename T>
T
RawAt(size_t i) const {
if (!IsChunkedVariableColumnDataType(data_type_)) {
PanicInfo(ErrorCode::Unsupported,
"RawAt only supported for ChunkedVariableColumn");
}
auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(i);
auto chunk = group_->GetGroupChunk(chunk_id).get()->GetChunk(field_id_);
std::string_view str_view =
static_cast<StringChunk*>(chunk.get())->operator[](offset_in_chunk);
return T(str_view.data(), str_view.size());
}
ScalarArray
PrimitivieRawAt(const int i) const override {
if (!IsChunkedArrayColumnDataType(data_type_)) {
PanicInfo(ErrorCode::Unsupported,
"PrimitivieRawAt only supported for ChunkedArrayColumn");
}
auto [chunk_id, offset_in_chunk] = GetChunkIDByOffset(i);
auto chunk = group_->GetGroupChunk(chunk_id).get()->GetChunk(field_id_);
return static_cast<ArrayChunk*>(chunk.get())
->View(offset_in_chunk)
.output_data();
}
private:
std::shared_ptr<ChunkedColumnGroup> group_;
FieldId field_id_;
const FieldMeta& field_meta_;
DataType data_type_;
};
} // namespace milvus
internal/core/src/mmap/ChunkedColumnInterface.h
+130
View File
@@ -0,0 +1,130 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include "cachinglayer/CacheSlot.h"
#include "common/Chunk.h"
namespace milvus {
using namespace milvus::cachinglayer;
class ChunkedColumnInterface {
public:
virtual ~ChunkedColumnInterface() = default;
// Get raw data pointer of a specific chunk
virtual cachinglayer::PinWrapper<const char*>
DataOfChunk(int chunk_id) const = 0;
// Check if the value at given offset is valid (not null)
virtual bool
IsValid(size_t offset) const = 0;
virtual bool
IsValid(int64_t chunk_id, int64_t offset) const = 0;
// Check if the column can contain null values
virtual bool
IsNullable() const = 0;
// Get total number of rows in the column
virtual size_t
NumRows() const = 0;
// Get total number of chunks in the column
virtual int64_t
num_chunks() const = 0;
// Get total byte size of the column data
virtual size_t
DataByteSize() const = 0;
// Get number of rows in a specific chunk
virtual int64_t
chunk_row_nums(int64_t chunk_id) const = 0;
virtual PinWrapper<SpanBase>
Span(int64_t chunk_id) const = 0;
virtual PinWrapper<
std::pair<std::vector<std::string_view>, FixedVector<bool>>>
StringViews(int64_t chunk_id,
std::optional<std::pair<int64_t, int64_t>> offset_len =
std::nullopt) const = 0;
virtual PinWrapper<std::pair<std::vector<ArrayView>, FixedVector<bool>>>
ArrayViews(int64_t chunk_id,
std::optional<std::pair<int64_t, int64_t>> offset_len) const = 0;
virtual PinWrapper<
std::pair<std::vector<std::string_view>, FixedVector<bool>>>
ViewsByOffsets(int64_t chunk_id,
const FixedVector<int32_t>& offsets) const = 0;
// Convert a global offset to (chunk_id, offset_in_chunk) pair
virtual std::pair<size_t, size_t>
GetChunkIDByOffset(int64_t offset) const = 0;
virtual PinWrapper<Chunk*>
GetChunk(int64_t chunk_id) const = 0;
// Get number of rows before a specific chunk
virtual int64_t
GetNumRowsUntilChunk(int64_t chunk_id) const = 0;
// Get vector of row counts before each chunk
virtual const std::vector<int64_t>&
GetNumRowsUntilChunk() const = 0;
virtual const char*
ValueAt(int64_t offset) = 0;
// Get raw value at given offset for variable length types (string/json)
template <typename T>
T
RawAt(size_t offset) const {
PanicInfo(ErrorCode::Unsupported,
"RawAt only supported for ChunkedVariableColumn");
}
// Get raw value at given offset for array type
virtual ScalarArray
PrimitivieRawAt(int offset) const {
PanicInfo(ErrorCode::Unsupported,
"PrimitivieRawAt only supported for ChunkedArrayColumn");
}
static bool
IsChunkedVariableColumnDataType(DataType data_type) {
return data_type == DataType::STRING ||
data_type == DataType::VARCHAR || data_type == DataType::TEXT ||
data_type == DataType::JSON;
}
static bool
IsChunkedArrayColumnDataType(DataType data_type) {
return data_type == DataType::ARRAY;
}
static bool
IsChunkedColumnDataType(DataType data_type) {
return !IsChunkedVariableColumnDataType(data_type) &&
!IsChunkedArrayColumnDataType(data_type);
}
};
} // namespace milvus
internal/core/src/query/CachedSearchIterator.cpp
+1 -1
View File
@@ -134,7 +134,7 @@ CachedSearchIterator::CachedSearchIterator(
}
CachedSearchIterator::CachedSearchIterator(
ChunkedColumnBase* column,
ChunkedColumnInterface* column,
const dataset::SearchDataset& query_ds,
const SearchInfo& search_info,
const std::map<std::string, std::string>& index_info,
internal/core/src/query/CachedSearchIterator.h
+1 -1
View File
@@ -61,7 +61,7 @@ class CachedSearchIterator {
const milvus::DataType& data_type);
// For sealed segment with chunked data, BF
CachedSearchIterator(ChunkedColumnBase* column,
CachedSearchIterator(ChunkedColumnInterface* column,
const dataset::SearchDataset& dataset,
const SearchInfo& search_info,
const std::map<std::string, std::string>& index_info,
internal/core/src/query/SearchOnSealed.cpp
+1 -1
View File
@@ -86,7 +86,7 @@ SearchOnSealedIndex(const Schema& schema,
void
SearchOnSealedColumn(const Schema& schema,
ChunkedColumnBase* column,
ChunkedColumnInterface* column,
const SearchInfo& search_info,
const std::map<std::string, std::string>& index_info,
const void* query_data,
internal/core/src/query/SearchOnSealed.h
+1 -1
View File
@@ -29,7 +29,7 @@ SearchOnSealedIndex(const Schema& schema,
void
SearchOnSealedColumn(const Schema& schema,
ChunkedColumnBase* column,
ChunkedColumnInterface* column,
const SearchInfo& search_info,
const std::map<std::string, std::string>& index_info,
const void* query_data,
internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp
+338 -153
View File
@@ -31,11 +31,13 @@
#include "common/Array.h"
#include "common/Chunk.h"
#include "common/Consts.h"
#include "common/ChunkWriter.h"
#include "common/EasyAssert.h"
#include "common/FieldMeta.h"
#include "common/Json.h"
#include "common/LoadInfo.h"
#include "common/Schema.h"
#include "common/SystemProperty.h"
#include "common/Tracer.h"
#include "common/Types.h"
#include "google/protobuf/message_lite.h"
@@ -48,9 +50,14 @@
#include "query/SearchOnSealed.h"
#include "segcore/storagev1translator/ChunkTranslator.h"
#include "segcore/storagev1translator/DefaultValueChunkTranslator.h"
#include "segcore/storagev1translator/InsertRecordTranslator.h"
#include "segcore/storagev2translator/GroupChunkTranslator.h"
#include "mmap/ChunkedColumnInterface.h"
#include "mmap/ChunkedColumnGroup.h"
#include "storage/Util.h"
#include "storage/ThreadPools.h"
#include "storage/MmapManager.h"
#include "milvus-storage/format/parquet/file_reader.h"
#include "milvus-storage/filesystem/fs.h"
namespace milvus::segcore {
@@ -186,6 +193,116 @@ ChunkedSegmentSealedImpl::LoadScalarIndex(const LoadIndexInfo& info) {
void
ChunkedSegmentSealedImpl::LoadFieldData(const LoadFieldDataInfo& load_info) {
switch (load_info.storage_version) {
case 2:
load_column_group_data_internal(load_info);
if (fields_.find(TimestampFieldID) != fields_.end()) {
auto timestamp_proxy_column = fields_.at(TimestampFieldID);
auto num_rows =
load_info.field_infos.at(TimestampFieldID.get()).row_count;
std::vector<Timestamp> timestamps(num_rows);
int64_t offset = 0;
for (int i = 0; i < timestamp_proxy_column->num_chunks(); i++) {
auto chunk = timestamp_proxy_column->GetChunk(i);
auto fixed_chunk =
static_cast<FixedWidthChunk*>(chunk.get());
auto span = fixed_chunk->Span();
for (size_t j = 0; j < span.row_count(); j++) {
auto ts = *(int64_t*)((char*)span.data() +
j * span.element_sizeof());
timestamps[offset++] = ts;
}
}
init_timestamp_index(timestamps, num_rows);
system_ready_count_++;
AssertInfo(
offset == num_rows,
"timestamp total row count {} not equal to expected {}",
offset,
num_rows);
}
break;
default:
load_field_data_internal(load_info);
break;
}
}
void
ChunkedSegmentSealedImpl::load_column_group_data_internal(
const LoadFieldDataInfo& load_info) {
size_t num_rows = storage::GetNumRowsForLoadInfo(load_info);
ArrowSchemaPtr arrow_schema = schema_->ConvertToArrowSchema();
for (auto& [id, info] : load_info.field_infos) {
AssertInfo(info.row_count > 0, "The row count of field data is 0");
auto column_group_id = FieldId(id);
auto insert_files = info.insert_files;
storage::SortByPath(insert_files);
auto fs = milvus_storage::ArrowFileSystemSingleton::GetInstance()
.GetArrowFileSystem();
auto file_reader = std::make_shared<milvus_storage::FileRowGroupReader>(
fs, insert_files[0], arrow_schema);
std::shared_ptr<milvus_storage::PackedFileMetadata> metadata =
file_reader->file_metadata();
auto field_id_mapping = metadata->GetFieldIDMapping();
std::vector<milvus_storage::RowGroupMetadataVector> row_group_meta_list;
for (const auto& file : insert_files) {
auto reader =
std::make_shared<milvus_storage::FileRowGroupReader>(fs, file);
row_group_meta_list.push_back(
reader->file_metadata()->GetRowGroupMetadataVector());
}
milvus_storage::FieldIDList field_id_list =
metadata->GetGroupFieldIDList().GetFieldIDList(
column_group_id.get());
std::vector<FieldId> milvus_field_ids;
for (int i = 0; i < field_id_list.size(); ++i) {
milvus_field_ids.push_back(FieldId(field_id_list.Get(i)));
}
auto column_group_info = FieldDataInfo(
column_group_id.get(), num_rows, load_info.mmap_dir_path);
LOG_INFO("segment {} loads column group {} with num_rows {}",
this->get_segment_id(),
column_group_id.get(),
num_rows);
auto field_metas = schema_->get_field_metas(milvus_field_ids);
auto translator =
std::make_unique<storagev2translator::GroupChunkTranslator>(
get_segment_id(),
field_metas,
column_group_info,
insert_files,
info.enable_mmap,
row_group_meta_list,
field_id_list);
auto chunked_column_group =
std::make_shared<ChunkedColumnGroup>(std::move(translator));
// Create ProxyChunkColumn for each field in this column group
for (const auto& field_id : milvus_field_ids) {
auto field_meta = field_metas.at(field_id);
auto column = std::make_shared<ProxyChunkColumn>(
chunked_column_group, field_id, field_meta);
auto data_type = field_meta.get_data_type();
load_field_data_common(
field_id, column, num_rows, data_type, info.enable_mmap, true);
}
}
}
void
ChunkedSegmentSealedImpl::load_field_data_internal(
const LoadFieldDataInfo& load_info) {
size_t num_rows = storage::GetNumRowsForLoadInfo(load_info);
AssertInfo(
!num_rows_.has_value() || num_rows_ == num_rows,
@@ -196,12 +313,7 @@ ChunkedSegmentSealedImpl::LoadFieldData(const LoadFieldDataInfo& load_info) {
auto field_id = FieldId(id);
auto insert_files = info.insert_files;
std::sort(insert_files.begin(),
insert_files.end(),
[](const std::string& a, const std::string& b) {
return std::stol(a.substr(a.find_last_of('/') + 1)) <
std::stol(b.substr(b.find_last_of('/') + 1));
});
storage::SortByPath(insert_files);
auto field_data_info =
FieldDataInfo(field_id.get(), num_rows, load_info.mmap_dir_path);
@@ -211,32 +323,23 @@ ChunkedSegmentSealedImpl::LoadFieldData(const LoadFieldDataInfo& load_info) {
num_rows);
if (SystemProperty::Instance().IsSystem(field_id)) {
if (field_id == RowFieldID) {
// ignore row id field
continue;
}
std::unique_ptr<Translator<InsertRecord<true>>> translator =
std::make_unique<storagev1translator::InsertRecordTranslator>(
this->get_segment_id(),
DataType::INT64,
field_data_info,
schema_,
is_sorted_by_pk_,
insert_files,
this);
insert_record_slot_ =
Manager::GetInstance().CreateCacheSlot(std::move(translator));
deleted_record_ = std::make_unique<DeletedRecord<true>>(
insert_record_slot_,
[this](const PkType& pk, Timestamp timestamp) {
return this->search_pk(pk, timestamp);
},
this->get_segment_id());
{
std::unique_lock lck(mutex_);
update_row_count(num_rows);
}
++system_ready_count_;
auto parallel_degree = static_cast<uint64_t>(
DEFAULT_FIELD_MAX_MEMORY_LIMIT / FILE_SLICE_SIZE);
field_data_info.arrow_reader_channel->set_capacity(parallel_degree *
2);
auto& pool =
ThreadPools::GetThreadPool(milvus::ThreadPoolPriority::MIDDLE);
pool.Submit(LoadArrowReaderFromRemote,
insert_files,
field_data_info.arrow_reader_channel);
LOG_INFO("segment {} submits load field {} task to thread pool",
this->get_segment_id(),
field_id.get());
load_system_field_internal(field_id, field_data_info);
LOG_INFO("segment {} loads system field {} mmap false done",
this->get_segment_id(),
field_id.get());
} else {
auto field_meta = schema_->operator[](field_id);
std::unique_ptr<Translator<milvus::Chunk>> translator =
@@ -247,87 +350,57 @@ ChunkedSegmentSealedImpl::LoadFieldData(const LoadFieldDataInfo& load_info) {
insert_files,
info.enable_mmap);
std::shared_ptr<milvus::ChunkedColumnBase> column{};
auto data_type = field_meta.get_data_type();
switch (data_type) {
case milvus::DataType::STRING:
case milvus::DataType::VARCHAR:
case milvus::DataType::TEXT: {
column =
std::make_shared<ChunkedVariableColumn<std::string>>(
std::move(translator), field_meta);
break;
}
case milvus::DataType::JSON: {
column =
std::make_shared<ChunkedVariableColumn<milvus::Json>>(
std::move(translator), field_meta);
break;
}
case milvus::DataType::ARRAY: {
column = std::make_shared<ChunkedArrayColumn>(
std::move(translator), field_meta);
break;
}
default: {
column = std::make_shared<ChunkedColumn>(
std::move(translator), field_meta);
break;
}
}
auto column = MakeChunkedColumnBase(
data_type, std::move(translator), field_meta);
{
std::unique_lock lck(mutex_);
fields_.emplace(field_id, column);
if (info.enable_mmap) {
mmap_fields_.insert(field_id);
}
if (!num_rows_.has_value()) {
num_rows_ = num_rows;
}
}
if (!info.enable_mmap) {
stats_.mem_size += column->DataByteSize();
if (IsVariableDataType(data_type)) {
if (data_type == milvus::DataType::STRING ||
data_type == milvus::DataType::VARCHAR ||
data_type == milvus::DataType::TEXT) {
auto var_column = std::dynamic_pointer_cast<
ChunkedVariableColumn<std::string>>(column);
LoadStringSkipIndex(field_id, 0, *var_column);
}
// update average row data size
SegmentInternalInterface::set_field_avg_size(
field_id, num_rows, column->DataByteSize());
} else {
auto num_chunk = column->num_chunks();
for (int i = 0; i < num_chunk; ++i) {
auto pw = column->Span(i);
LoadPrimitiveSkipIndex(field_id,
i,
data_type,
pw.get().data(),
pw.get().valid_data(),
pw.get().row_count());
}
}
}
if (generate_interim_index(field_id)) {
std::unique_lock lck(mutex_);
// mmap_fields is useless, no change
fields_.erase(field_id);
set_bit(field_data_ready_bitset_, field_id, false);
} else {
std::unique_lock lck(mutex_);
set_bit(field_data_ready_bitset_, field_id, true);
}
load_field_data_common(
field_id, column, num_rows, data_type, info.enable_mmap, false);
}
}
}
void
ChunkedSegmentSealedImpl::load_system_field_internal(FieldId field_id,
FieldDataInfo& data) {
auto num_rows = data.row_count;
AssertInfo(SystemProperty::Instance().IsSystem(field_id),
"system field is not system field");
auto system_field_type =
SystemProperty::Instance().GetSystemFieldType(field_id);
if (system_field_type == SystemFieldType::Timestamp) {
std::vector<Timestamp> timestamps(num_rows);
int64_t offset = 0;
FieldMeta field_meta(
FieldName(""), FieldId(0), DataType::INT64, false, std::nullopt);
std::shared_ptr<milvus::ArrowDataWrapper> r;
while (data.arrow_reader_channel->pop(r)) {
auto array_vec = read_single_column_batches(r->reader);
auto chunk = create_chunk(field_meta, 1, array_vec);
auto chunk_ptr = static_cast<FixedWidthChunk*>(chunk.get());
std::copy_n(static_cast<const Timestamp*>(chunk_ptr->Span().data()),
chunk_ptr->Span().row_count(),
timestamps.data() + offset);
offset += chunk_ptr->Span().row_count();
}
init_timestamp_index(timestamps, num_rows);
++system_ready_count_;
} else {
AssertInfo(system_field_type == SystemFieldType::RowId,
"System field type of id column is not RowId");
// Consume rowid field data but not really load it
// storage::CollectFieldDataChannel(data.arrow_reader_channel);
std::shared_ptr<milvus::ArrowDataWrapper> r;
while (data.arrow_reader_channel->pop(r)) {
}
}
{
std::unique_lock lck(mutex_);
update_row_count(num_rows);
}
}
void
ChunkedSegmentSealedImpl::LoadDeletedRecord(const LoadDeletedRecordInfo& info) {
AssertInfo(info.row_count > 0, "The row count of deleted record is 0");
@@ -343,7 +416,7 @@ ChunkedSegmentSealedImpl::LoadDeletedRecord(const LoadDeletedRecordInfo& info) {
auto timestamps = reinterpret_cast<const Timestamp*>(info.timestamps);
// step 2: push delete info to delete_record
deleted_record_->LoadPush(pks, timestamps);
deleted_record_.LoadPush(pks, timestamps);
}
void
@@ -493,7 +566,7 @@ ChunkedSegmentSealedImpl::get_row_count() const {
int64_t
ChunkedSegmentSealedImpl::get_deleted_count() const {
std::shared_lock lck(mutex_);
return deleted_record_->size();
return deleted_record_.size();
}
const Schema&
@@ -505,7 +578,7 @@ void
ChunkedSegmentSealedImpl::mask_with_delete(BitsetTypeView& bitset,
int64_t ins_barrier,
Timestamp timestamp) const {
deleted_record_->Query(bitset, ins_barrier, timestamp);
deleted_record_.Query(bitset, ins_barrier, timestamp);
}
void
@@ -695,13 +768,11 @@ ChunkedSegmentSealedImpl::check_search(const query::Plan* plan) const {
std::vector<SegOffset>
ChunkedSegmentSealedImpl::search_pk(const PkType& pk,
Timestamp timestamp) const {
auto ir_accessor = pin_insert_record();
auto ir = ir_accessor->get_cell_of(0);
if (!is_sorted_by_pk_) {
return ir->search_pk(pk, timestamp);
return insert_record_.search_pk(pk, timestamp);
}
return search_sorted_pk(pk, [this, timestamp, ir](int64_t offset) {
return ir->timestamps()[offset] <= timestamp;
return search_sorted_pk(pk, [this, timestamp](int64_t offset) {
return insert_record_.timestamps_[offset] <= timestamp;
});
}
@@ -780,9 +851,7 @@ std::pair<std::vector<OffsetMap::OffsetType>, bool>
ChunkedSegmentSealedImpl::find_first(int64_t limit,
const BitsetType& bitset) const {
if (!is_sorted_by_pk_) {
auto ir_accessor = pin_insert_record();
return ir_accessor->get_cell_of(0)->pk2offset_->find_first(limit,
bitset);
return insert_record_.pk2offset_->find_first(limit, bitset);
}
if (limit == Unlimited || limit == NoLimit) {
limit = num_rows_.value();
@@ -824,13 +893,18 @@ ChunkedSegmentSealedImpl::ChunkedSegmentSealedImpl(
index_ready_bitset_(schema->size()),
binlog_index_bitset_(schema->size()),
scalar_indexings_(schema->size()),
insert_record_slot_(nullptr),
insert_record_(*schema, MAX_ROW_COUNT),
schema_(schema),
id_(segment_id),
col_index_meta_(index_meta),
TEST_skip_index_for_retrieve_(TEST_skip_index_for_retrieve),
is_sorted_by_pk_(is_sorted_by_pk),
deleted_record_(nullptr) {
deleted_record_(
&insert_record_,
[this](const PkType& pk, Timestamp timestamp) {
return this->search_pk(pk, timestamp);
},
segment_id) {
mmap_descriptor_ = std::shared_ptr<storage::MmapChunkDescriptor>(
new storage::MmapChunkDescriptor({segment_id, SegmentType::Sealed}));
auto mcm = storage::MmapManager::GetInstance().GetMmapChunkManager();
@@ -852,14 +926,13 @@ ChunkedSegmentSealedImpl::bulk_subscript(SystemFieldType system_type,
AssertInfo(is_system_field_ready(),
"System field isn't ready when do bulk_insert, segID:{}",
id_);
auto ir_accessor = pin_insert_record();
switch (system_type) {
case SystemFieldType::Timestamp:
AssertInfo(
ir_accessor->get_cell_of(0)->timestamps().num_chunk() == 1,
insert_record_.timestamps_.num_chunk() == 1,
"num chunk of timestamp not equal to 1 for sealed segment");
bulk_subscript_impl<Timestamp>(
ir_accessor->get_cell_of(0)->timestamps().get_chunk_data(0),
this->insert_record_.timestamps_.get_chunk_data(0),
seg_offsets,
count,
static_cast<Timestamp*>(output));
@@ -888,7 +961,7 @@ ChunkedSegmentSealedImpl::bulk_subscript_impl(const void* src_raw,
}
template <typename S, typename T>
void
ChunkedSegmentSealedImpl::bulk_subscript_impl(ChunkedColumnBase* field,
ChunkedSegmentSealedImpl::bulk_subscript_impl(ChunkedColumnInterface* field,
const int64_t* seg_offsets,
int64_t count,
T* dst) {
@@ -904,7 +977,7 @@ ChunkedSegmentSealedImpl::bulk_subscript_impl(ChunkedColumnBase* field,
template <typename S, typename T>
void
ChunkedSegmentSealedImpl::bulk_subscript_ptr_impl(
ChunkedColumnBase* column,
ChunkedColumnInterface* column,
const int64_t* seg_offsets,
int64_t count,
google::protobuf::RepeatedPtrField<T>* dst) {
@@ -918,21 +991,21 @@ ChunkedSegmentSealedImpl::bulk_subscript_ptr_impl(
template <typename T>
void
ChunkedSegmentSealedImpl::bulk_subscript_array_impl(
ChunkedColumnBase* column,
ChunkedColumnInterface* column,
const int64_t* seg_offsets,
int64_t count,
google::protobuf::RepeatedPtrField<T>* dst) {
auto field = reinterpret_cast<ChunkedArrayColumn*>(column);
for (int64_t i = 0; i < count; ++i) {
auto offset = seg_offsets[i];
dst->at(i) = std::move(field->RawAt(offset));
dst->at(i) = std::move(field->PrimitivieRawAt(offset));
}
}
// for dense vector
void
ChunkedSegmentSealedImpl::bulk_subscript_impl(int64_t element_sizeof,
ChunkedColumnBase* field,
ChunkedColumnInterface* field,
const int64_t* seg_offsets,
int64_t count,
void* dst_raw) {
@@ -957,7 +1030,7 @@ ChunkedSegmentSealedImpl::ClearData() {
num_rows_ = std::nullopt;
scalar_indexings_.clear();
vector_indexings_.clear();
insert_record_slot_ = nullptr;
insert_record_.clear();
fields_.clear();
variable_fields_avg_size_.clear();
stats_.mem_size = 0;
@@ -1394,8 +1467,7 @@ ChunkedSegmentSealedImpl::search_ids(const IdArray& id_array,
for (auto& pk : pks) {
std::vector<SegOffset> pk_offsets;
if (!is_sorted_by_pk_) {
auto ir_accessor = pin_insert_record();
pk_offsets = ir_accessor->get_cell_of(0)->search_pk(pk, timestamp);
pk_offsets = insert_record_.search_pk(pk, timestamp);
} else {
pk_offsets = search_pk(pk, timestamp);
}
@@ -1439,15 +1511,13 @@ ChunkedSegmentSealedImpl::Delete(int64_t size,
}
// if insert record is empty (may be only-load meta but not data for lru-cache at go side),
// filtering may cause the deletion lost, skip the filtering to avoid it.
auto ir_accessor = pin_insert_record();
if (!ir_accessor->get_cell_of(0)->empty_pks()) {
auto end =
std::remove_if(ordering.begin(),
ordering.end(),
[&](const std::tuple<Timestamp, PkType>& record) {
return !ir_accessor->get_cell_of(0)->contain(
std::get<1>(record));
});
if (!insert_record_.empty_pks()) {
auto end = std::remove_if(
ordering.begin(),
ordering.end(),
[&](const std::tuple<Timestamp, PkType>& record) {
return !insert_record_.contain(std::get<1>(record));
});
size = end - ordering.begin();
ordering.resize(size);
}
@@ -1466,7 +1536,7 @@ ChunkedSegmentSealedImpl::Delete(int64_t size,
sort_pks[i] = pk;
}
deleted_record_->StreamPush(sort_pks, sort_timestamps.data());
deleted_record_.StreamPush(sort_pks, sort_timestamps.data());
return SegcoreError::success();
}
@@ -1486,9 +1556,7 @@ ChunkedSegmentSealedImpl::LoadSegmentMeta(
for (auto& info : segment_meta.metas()) {
slice_lengths.push_back(info.row_count());
}
auto ir_accessor = pin_insert_record();
ir_accessor->get_cell_of(0)->timestamp_index_.set_length_meta(
std::move(slice_lengths));
insert_record_.timestamp_index_.set_length_meta(std::move(slice_lengths));
PanicInfo(NotImplemented, "unimplemented");
}
@@ -1502,21 +1570,19 @@ void
ChunkedSegmentSealedImpl::mask_with_timestamps(BitsetTypeView& bitset_chunk,
Timestamp timestamp) const {
// TODO change the
auto ir_accessor = pin_insert_record();
auto ir = ir_accessor->get_cell_of(0);
AssertInfo(ir->timestamps().num_chunk() == 1,
AssertInfo(insert_record_.timestamps_.num_chunk() == 1,
"num chunk not equal to 1 for sealed segment");
auto timestamps_data =
(const milvus::Timestamp*)ir->timestamps().get_chunk_data(0);
auto timestamps_data_size = ir->timestamps().get_chunk_size(0);
(const milvus::Timestamp*)insert_record_.timestamps_.get_chunk_data(0);
auto timestamps_data_size = insert_record_.timestamps_.get_chunk_size(0);
AssertInfo(timestamps_data_size == get_row_count(),
fmt::format("Timestamp size not equal to row count: {}, {}",
timestamps_data_size,
get_row_count()));
auto range = ir->timestamp_index_.get_active_range(timestamp);
auto range = insert_record_.timestamp_index_.get_active_range(timestamp);
// range == (size_, size_) and size_ is this->timestamps().size().
// range == (size_, size_) and size_ is this->timestamps_.size().
// it means these data are all useful, we don't need to update bitset_chunk.
// It can be thought of as an OR operation with another bitmask that is all 0s, but it is not necessary to do so.
if (range.first == range.second && range.first == timestamps_data_size) {
@@ -1594,7 +1660,7 @@ ChunkedSegmentSealedImpl::generate_interim_index(const FieldId field_id) {
// if (row_count < field_binlog_config->GetBuildThreshold()) {
// return false;
// }
std::shared_ptr<ChunkedColumnBase> vec_data{};
std::shared_ptr<ChunkedColumnInterface> vec_data{};
{
std::shared_lock lck(mutex_);
vec_data = fields_.at(field_id);
@@ -1656,6 +1722,125 @@ ChunkedSegmentSealedImpl::LazyCheckSchema(const Schema& sch) {
}
}
void
ChunkedSegmentSealedImpl::load_field_data_common(
FieldId field_id,
const std::shared_ptr<ChunkedColumnInterface>& column,
size_t num_rows,
DataType data_type,
bool enable_mmap,
bool is_proxy_column) {
{
std::unique_lock lck(mutex_);
fields_.emplace(field_id, column);
if (enable_mmap) {
mmap_fields_.insert(field_id);
}
}
// system field only needs to emplace column to fields_ map
if (SystemProperty::Instance().IsSystem(field_id)) {
return;
}
if (!enable_mmap) {
stats_.mem_size += column->DataByteSize();
if (IsVariableDataType(data_type)) {
if (IsStringDataType(data_type)) {
if (!is_proxy_column) {
auto var_column = std::dynamic_pointer_cast<
ChunkedVariableColumn<std::string>>(column);
AssertInfo(var_column != nullptr,
"column is not of variable type");
LoadStringSkipIndex(field_id, 0, *var_column);
} else {
auto var_column =
std::dynamic_pointer_cast<ProxyChunkColumn>(column);
AssertInfo(var_column != nullptr,
"column is not of variable type");
LoadStringSkipIndex(field_id, 0, *var_column);
}
}
// update average row data size
SegmentInternalInterface::set_field_avg_size(
field_id, num_rows, column->DataByteSize());
} else {
auto num_chunk = column->num_chunks();
for (int i = 0; i < num_chunk; ++i) {
if (!is_proxy_column) {
auto primitive_column =
std::dynamic_pointer_cast<ChunkedColumn>(column);
AssertInfo(primitive_column != nullptr,
"column is not of primitive type");
auto pw = primitive_column->Span(i);
LoadPrimitiveSkipIndex(field_id,
i,
data_type,
pw.get().data(),
pw.get().valid_data(),
pw.get().row_count());
} else {
auto proxy_column =
std::dynamic_pointer_cast<ProxyChunkColumn>(column);
AssertInfo(proxy_column != nullptr,
"column is not of proxy type");
auto pw = proxy_column->Span(i);
LoadPrimitiveSkipIndex(field_id,
i,
data_type,
pw.get().data(),
pw.get().valid_data(),
pw.get().row_count());
}
}
}
}
// set pks to offset
if (schema_->get_primary_field_id() == field_id && !is_sorted_by_pk_) {
AssertInfo(field_id.get() != -1, "Primary key is -1");
AssertInfo(insert_record_.empty_pks(), "already exists");
insert_record_.insert_pks(data_type, column.get());
insert_record_.seal_pks();
}
if (generate_interim_index(field_id)) {
std::unique_lock lck(mutex_);
// mmap_fields is useless, no change
fields_.erase(field_id);
set_bit(field_data_ready_bitset_, field_id, false);
} else {
std::unique_lock lck(mutex_);
set_bit(field_data_ready_bitset_, field_id, true);
}
{
std::unique_lock lck(mutex_);
update_row_count(num_rows);
}
}
void
ChunkedSegmentSealedImpl::init_timestamp_index(
const std::vector<Timestamp>& timestamps, size_t num_rows) {
TimestampIndex index;
auto min_slice_length = num_rows < 4096 ? 1 : 4096;
auto meta =
GenerateFakeSlices(timestamps.data(), num_rows, min_slice_length);
index.set_length_meta(std::move(meta));
// todo ::opt to avoid copy timestamps from field data
index.build_with(timestamps.data(), num_rows);
// use special index
std::unique_lock lck(mutex_);
AssertInfo(insert_record_.timestamps_.empty(), "already exists");
insert_record_.timestamps_.set_data_raw(
0, timestamps.data(), timestamps.size());
insert_record_.timestamp_index_ = std::move(index);
AssertInfo(insert_record_.timestamps_.num_chunk() == 1,
"num chunk not equal to 1 for sealed segment");
stats_.mem_size += sizeof(Timestamp) * num_rows;
}
void
ChunkedSegmentSealedImpl::Reopen(SchemaPtr sch) {
std::unique_lock lck(mutex_);
internal/core/src/segcore/ChunkedSegmentSealedImpl.h
+35 -24
View File
@@ -73,7 +73,7 @@ class ChunkedSegmentSealedImpl : public SegmentSealed {
bool
Contain(const PkType& pk) const override {
return pin_insert_record()->get_cell_of(0)->contain(pk);
return insert_record_.contain(pk);
}
void
@@ -144,7 +144,12 @@ class ChunkedSegmentSealedImpl : public SegmentSealed {
public:
size_t
GetMemoryUsageInBytes() const override {
return stats_.mem_size.load() + deleted_record_->mem_size();
return stats_.mem_size.load() + deleted_record_.mem_size();
}
InsertRecord<true>&
get_insert_record() override {
return insert_record_;
}
int64_t
@@ -281,12 +286,12 @@ class ChunkedSegmentSealedImpl : public SegmentSealed {
const ConcurrentVector<Timestamp>&
get_timestamps() const override {
return pin_insert_record()->get_cell_of(0)->timestamps_;
return insert_record_.timestamps_;
}
private:
void
LoadSystemFieldInternal(FieldId field_id, FieldDataInfo& data);
load_system_field_internal(FieldId field_id, FieldDataInfo& data);
template <typename S, typename T = S>
static void
@@ -297,28 +302,28 @@ class ChunkedSegmentSealedImpl : public SegmentSealed {
template <typename S, typename T = S>
static void
bulk_subscript_impl(ChunkedColumnBase* field,
bulk_subscript_impl(ChunkedColumnInterface* field,
const int64_t* seg_offsets,
int64_t count,
T* dst_raw);
template <typename S, typename T = S>
static void
bulk_subscript_ptr_impl(ChunkedColumnBase* field,
bulk_subscript_ptr_impl(ChunkedColumnInterface* field,
const int64_t* seg_offsets,
int64_t count,
google::protobuf::RepeatedPtrField<T>* dst_raw);
template <typename T>
static void
bulk_subscript_array_impl(ChunkedColumnBase* column,
bulk_subscript_array_impl(ChunkedColumnInterface* column,
const int64_t* seg_offsets,
int64_t count,
google::protobuf::RepeatedPtrField<T>* dst);
static void
bulk_subscript_impl(int64_t element_sizeof,
ChunkedColumnBase* field,
ChunkedColumnInterface* field,
const int64_t* seg_offsets,
int64_t count,
void* dst_raw);
@@ -335,7 +340,7 @@ class ChunkedSegmentSealedImpl : public SegmentSealed {
void
update_row_count(int64_t row_count) {
num_rows_ = row_count;
deleted_record_->set_sealed_row_count(row_count);
deleted_record_.set_sealed_row_count(row_count);
}
void
@@ -375,11 +380,24 @@ class ChunkedSegmentSealedImpl : public SegmentSealed {
void
fill_empty_field(const FieldMeta& field_meta);
// used only by unit test
std::shared_ptr<CacheSlot<InsertRecord<true>>>
get_insert_record_slot() const override {
return insert_record_slot_;
}
void
init_timestamp_index(const std::vector<Timestamp>& timestamps,
size_t num_rows);
void
load_field_data_internal(const LoadFieldDataInfo& load_info);
void
load_column_group_data_internal(const LoadFieldDataInfo& load_info);
void
load_field_data_common(
FieldId field_id,
const std::shared_ptr<ChunkedColumnInterface>& column,
size_t num_rows,
DataType data_type,
bool enable_mmap,
bool is_proxy_column);
private:
// InsertRecord needs to pin pk column.
@@ -402,24 +420,17 @@ class ChunkedSegmentSealedImpl : public SegmentSealed {
// vector field index
SealedIndexingRecord vector_indexings_;
std::shared_ptr<CellAccessor<InsertRecord<true>>>
pin_insert_record() const {
return insert_record_slot_->PinCells({0})
.via(&folly::InlineExecutor::instance())
.get();
}
// inserted fields data and row_ids, timestamps
mutable std::shared_ptr<CacheSlot<InsertRecord<true>>> insert_record_slot_;
InsertRecord<true> insert_record_;
// deleted pks
mutable std::unique_ptr<DeletedRecord<true>> deleted_record_;
mutable DeletedRecord<true> deleted_record_;
LoadFieldDataInfo field_data_info_;
SchemaPtr schema_;
int64_t id_;
mutable std::unordered_map<FieldId, std::shared_ptr<ChunkedColumnBase>>
mutable std::unordered_map<FieldId, std::shared_ptr<ChunkedColumnInterface>>
fields_;
std::unordered_set<FieldId> mmap_fields_;
internal/core/src/segcore/DeletedRecord.h
+9 -22
View File
@@ -19,9 +19,13 @@
#include <vector>
#include <folly/ConcurrentSkipList.h>
#include "cachinglayer/CacheSlot.h"
#include "AckResponder.h"
#include "common/Schema.h"
#include "common/Types.h"
#include "segcore/Record.h"
#include "segcore/InsertRecord.h"
#include "segcore/SegmentInterface.h"
#include "ConcurrentVector.h"
namespace milvus::segcore {
@@ -48,12 +52,7 @@ static int32_t DELETE_PAIR_SIZE = sizeof(std::pair<Timestamp, Offset>);
template <bool is_sealed = false>
class DeletedRecord {
public:
using InsertRecordType = std::conditional_t<
is_sealed,
std::shared_ptr<milvus::cachinglayer::CacheSlot<InsertRecord<true>>>,
InsertRecord<false>*>;
DeletedRecord(InsertRecordType insert_record,
DeletedRecord(InsertRecord<is_sealed>* insert_record,
std::function<std::vector<SegOffset>(
const PkType& pk, Timestamp timestamp)> search_pk_func,
int64_t segment_id)
@@ -109,18 +108,6 @@ class DeletedRecord {
Timestamp max_timestamp = 0;
SortedDeleteList::Accessor accessor(deleted_lists_);
InsertRecord<is_sealed>* insert_record = nullptr;
std::shared_ptr<
milvus::cachinglayer::CellAccessor<InsertRecord<is_sealed>>>
cell_accessor{nullptr};
if constexpr (is_sealed) {
cell_accessor = insert_record_->PinCells({0})
.via(&folly::InlineExecutor::instance())
.get();
insert_record = cell_accessor->get_cell_of(0);
} else {
insert_record = insert_record_;
}
for (size_t i = 0; i < pks.size(); ++i) {
auto deleted_pk = pks[i];
auto deleted_ts = timestamps[i];
@@ -137,7 +124,7 @@ class DeletedRecord {
}
// if insert record and delete record is same timestamp,
// delete not take effect on this record.
if (deleted_ts == insert_record->timestamps_[row_id]) {
if (deleted_ts == insert_record_->timestamps_[row_id]) {
continue;
}
accessor.insert(std::make_pair(deleted_ts, row_id));
@@ -146,7 +133,7 @@ class DeletedRecord {
deleted_mask_.set(row_id);
} else {
// need to add mask size firstly for growing segment
deleted_mask_.resize(insert_record->size());
deleted_mask_.resize(insert_record_->size());
deleted_mask_.set(row_id);
}
removed_num++;
@@ -335,7 +322,7 @@ class DeletedRecord {
public:
std::atomic<int64_t> n_ = 0;
std::atomic<int64_t> mem_size_ = 0;
InsertRecordType insert_record_;
InsertRecord<is_sealed>* insert_record_;
std::function<std::vector<SegOffset>(const PkType& pk, Timestamp timestamp)>
search_pk_func_;
int64_t segment_id_{0};
internal/core/src/segcore/InsertRecord.h
+5 -2
View File
@@ -334,7 +334,7 @@ struct InsertRecord {
}
void
insert_pks(milvus::DataType data_type, ChunkedColumnBase* data) {
insert_pks(milvus::DataType data_type, ChunkedColumnInterface* data) {
std::lock_guard lck(shared_mutex_);
int64_t offset = 0;
switch (data_type) {
@@ -353,7 +353,10 @@ struct InsertRecord {
case DataType::VARCHAR: {
auto num_chunk = data->num_chunks();
for (int i = 0; i < num_chunk; ++i) {
auto pw = data->StringViews(i);
auto column =
reinterpret_cast<ChunkedVariableColumn<std::string>*>(
data);
auto pw = column->StringViews(i);
auto pks = pw.get().first;
for (auto& pk : pks) {
pk2offset_->insert(std::string(pk), offset++);
internal/core/src/segcore/SegmentGrowingImpl.cpp
+2 -12
View File
@@ -271,12 +271,7 @@ SegmentGrowingImpl::load_field_data_internal(const LoadFieldDataInfo& infos) {
for (auto& [id, info] : infos.field_infos) {
auto field_id = FieldId(id);
auto insert_files = info.insert_files;
std::sort(insert_files.begin(),
insert_files.end(),
[](const std::string& a, const std::string& b) {
return std::stol(a.substr(a.find_last_of('/') + 1)) <
std::stol(b.substr(b.find_last_of('/') + 1));
});
storage::SortByPath(insert_files);
auto channel = std::make_shared<FieldDataChannel>();
auto& pool =
@@ -412,12 +407,7 @@ SegmentGrowingImpl::load_column_group_data_internal(
for (auto& [id, info] : infos.field_infos) {
auto column_group_id = FieldId(id);
auto insert_files = info.insert_files;
std::sort(insert_files.begin(),
insert_files.end(),
[](const std::string& a, const std::string& b) {
return std::stol(a.substr(a.find_last_of('/') + 1)) <
std::stol(b.substr(b.find_last_of('/') + 1));
});
storage::SortByPath(insert_files);
auto fs = milvus_storage::ArrowFileSystemSingleton::GetInstance()
.GetArrowFileSystem();
auto file_reader = std::make_shared<milvus_storage::FileRowGroupReader>(
internal/core/src/segcore/SegmentInterface.h
+1
View File
@@ -32,6 +32,7 @@
#include "common/BitsetView.h"
#include "common/QueryResult.h"
#include "common/QueryInfo.h"
#include "mmap/ChunkedColumnInterface.h"
#include "query/Plan.h"
#include "query/PlanNode.h"
#include "pb/schema.pb.h"
internal/core/src/segcore/SegmentSealed.h
+3 -4
View File
@@ -19,7 +19,6 @@
#include "pb/segcore.pb.h"
#include "segcore/InsertRecord.h"
#include "segcore/SegmentInterface.h"
#include "cachinglayer/CacheSlot.h"
#include "segcore/Types.h"
namespace milvus::segcore {
@@ -48,6 +47,9 @@ class SegmentSealed : public SegmentInternalInterface {
LoadTextIndex(FieldId field_id,
std::unique_ptr<index::TextMatchIndex> index) = 0;
virtual InsertRecord<true>&
get_insert_record() = 0;
virtual index::IndexBase*
GetJsonIndex(FieldId field_id, std::string path) const override {
JSONIndexKey key;
@@ -104,9 +106,6 @@ class SegmentSealed : public SegmentInternalInterface {
index->second->JsonCastType() == DataType::DOUBLE);
}
virtual std::shared_ptr<CacheSlot<InsertRecord<true>>>
get_insert_record_slot() const = 0;
protected:
struct JSONIndexKey {
FieldId field_id;
internal/core/src/segcore/storagev1translator/InsertRecordTranslator.cpp
-139
View File
@@ -1,139 +0,0 @@
#include "segcore/storagev1translator/InsertRecordTranslator.h"
#include <memory>
#include <vector>
#include <string>
#include "fmt/core.h"
#include "cachinglayer/Utils.h"
#include "common/ChunkWriter.h"
#include "common/Types.h"
#include "common/SystemProperty.h"
#include "segcore/Utils.h"
#include "storage/ThreadPools.h"
namespace milvus::segcore::storagev1translator {
InsertRecordTranslator::InsertRecordTranslator(
int64_t segment_id,
DataType data_type,
FieldDataInfo field_data_info,
SchemaPtr schema,
bool is_sorted_by_pk,
std::vector<std::string> insert_files,
ChunkedSegmentSealedImpl* chunked_segment)
: segment_id_(segment_id),
data_type_(data_type),
key_(fmt::format("seg_{}_ir_f_{}", segment_id, field_data_info.field_id)),
field_data_info_(field_data_info),
schema_(schema),
is_sorted_by_pk_(is_sorted_by_pk),
insert_files_(insert_files),
chunked_segment_(chunked_segment),
meta_(milvus::cachinglayer::StorageType::MEMORY) {
}
size_t
InsertRecordTranslator::num_cells() const {
return 1;
}
milvus::cachinglayer::cid_t
InsertRecordTranslator::cell_id_of(milvus::cachinglayer::uid_t uid) const {
return 0;
}
milvus::cachinglayer::ResourceUsage
InsertRecordTranslator::estimated_byte_size_of_cell(
milvus::cachinglayer::cid_t cid) const {
return {0, 0};
}
const std::string&
InsertRecordTranslator::key() const {
return key_;
}
std::vector<std::pair<milvus::cachinglayer::cid_t,
std::unique_ptr<milvus::segcore::InsertRecord<true>>>>
InsertRecordTranslator::get_cells(
const std::vector<milvus::cachinglayer::cid_t>& cids) {
AssertInfo(cids.size() == 1 && cids[0] == 0,
"InsertRecordTranslator only supports single cell");
FieldId fid = FieldId(field_data_info_.field_id);
auto parallel_degree =
static_cast<uint64_t>(DEFAULT_FIELD_MAX_MEMORY_LIMIT / FILE_SLICE_SIZE);
// TODO(tiered storage 4): we should phase out this thread pool and use folly executor.
auto& pool = ThreadPools::GetThreadPool(milvus::ThreadPoolPriority::MIDDLE);
pool.Submit(LoadArrowReaderFromRemote,
insert_files_,
field_data_info_.arrow_reader_channel);
LOG_INFO("segment {} submits load field {} task to thread pool",
segment_id_,
field_data_info_.field_id);
auto num_rows = field_data_info_.row_count;
AssertInfo(milvus::SystemProperty::Instance().IsSystem(fid),
"system field is not system field");
auto system_field_type =
milvus::SystemProperty::Instance().GetSystemFieldType(fid);
AssertInfo(system_field_type == SystemFieldType::Timestamp,
"system field is not timestamp");
std::vector<Timestamp> timestamps(num_rows);
int64_t offset = 0;
FieldMeta field_meta(
FieldName(""), FieldId(0), DataType::INT64, false, std::nullopt);
std::shared_ptr<milvus::ArrowDataWrapper> r;
while (field_data_info_.arrow_reader_channel->pop(r)) {
arrow::ArrayVector array_vec = read_single_column_batches(r->reader);
auto chunk = create_chunk(field_meta, 1, array_vec);
auto chunk_ptr = static_cast<FixedWidthChunk*>(chunk.get());
std::copy_n(static_cast<const Timestamp*>(chunk_ptr->Span().data()),
chunk_ptr->Span().row_count(),
timestamps.data() + offset);
offset += chunk_ptr->Span().row_count();
}
TimestampIndex index;
auto min_slice_length = num_rows < 4096 ? 1 : 4096;
auto meta =
GenerateFakeSlices(timestamps.data(), num_rows, min_slice_length);
index.set_length_meta(std::move(meta));
// todo ::opt to avoid copy timestamps from field data
index.build_with(timestamps.data(), num_rows);
std::unique_ptr<milvus::segcore::InsertRecord<true>> ir =
std::make_unique<milvus::segcore::InsertRecord<true>>(*schema_,
MAX_ROW_COUNT);
// use special index
AssertInfo(ir->timestamps_.empty(), "already exists");
ir->timestamps_.set_data_raw(0, timestamps.data(), timestamps.size());
ir->timestamp_index_ = std::move(index);
AssertInfo(ir->timestamps_.num_chunk() == 1,
"num chunk not equal to 1 for sealed segment");
chunked_segment_->stats_.mem_size += sizeof(Timestamp) * num_rows;
auto pk_field_id = schema_->get_primary_field_id();
AssertInfo(pk_field_id.has_value(),
"primary key field not found in schema");
auto pk_field_meta = schema_->operator[](pk_field_id.value());
// set pks to offset
if (!is_sorted_by_pk_) {
AssertInfo(ir->empty_pks(), "already exists");
auto it = chunked_segment_->fields_.find(pk_field_id.value());
AssertInfo(it != chunked_segment_->fields_.end(),
"primary key field not found in segment");
ir->insert_pks(pk_field_meta.get_data_type(), it->second.get());
ir->seal_pks();
}
std::vector<std::pair<milvus::cachinglayer::cid_t,
std::unique_ptr<milvus::segcore::InsertRecord<true>>>>
cells;
cells.emplace_back(0, std::move(ir));
return cells;
}
} // namespace milvus::segcore::storagev1translator
internal/core/src/segcore/storagev1translator/InsertRecordTranslator.h
-72
View File
@@ -1,72 +0,0 @@
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// 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 <string>
#include <vector>
#include "cachinglayer/Translator.h"
#include "cachinglayer/Utils.h"
#include "common/Schema.h"
#include "segcore/InsertRecord.h"
#include "segcore/ChunkedSegmentSealedImpl.h"
namespace milvus::segcore::storagev1translator {
class InsertRecordTranslator : public milvus::cachinglayer::Translator<
milvus::segcore::InsertRecord<true>> {
public:
InsertRecordTranslator(int64_t segment_id,
DataType data_type,
FieldDataInfo field_data_info,
SchemaPtr schema,
bool is_sorted_by_pk,
std::vector<std::string> insert_files,
ChunkedSegmentSealedImpl* chunked_segment);
size_t
num_cells() const override;
milvus::cachinglayer::cid_t
cell_id_of(milvus::cachinglayer::uid_t uid) const override;
milvus::cachinglayer::ResourceUsage
estimated_byte_size_of_cell(milvus::cachinglayer::cid_t cid) const override;
const std::string&
key() const override;
// each calling of this will trigger a new download.
std::vector<std::pair<milvus::cachinglayer::cid_t,
std::unique_ptr<milvus::segcore::InsertRecord<true>>>>
get_cells(const std::vector<milvus::cachinglayer::cid_t>& cids) override;
// InsertRecord does not have meta.
milvus::cachinglayer::Meta*
meta() override {
return &meta_;
}
private:
std::unique_ptr<milvus::segcore::InsertRecord<true>>
load_column_in_memory() const;
int64_t segment_id_;
std::string key_;
DataType data_type_;
FieldDataInfo field_data_info_;
std::vector<std::string> insert_files_;
mutable size_t estimated_byte_size_of_cell_;
SchemaPtr schema_;
bool is_sorted_by_pk_;
ChunkedSegmentSealedImpl* chunked_segment_;
milvus::cachinglayer::Meta meta_;
};
} // namespace milvus::segcore::storagev1translator
internal/core/src/segcore/storagev2translator/GroupCTMeta.h
+30
View File
@@ -0,0 +1,30 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include "cachinglayer/Translator.h"
namespace milvus::segcore::storagev2translator {
struct GroupCTMeta : public milvus::cachinglayer::Meta {
std::vector<int64_t> num_rows_until_chunk_;
std::vector<int64_t> chunk_memory_size_;
GroupCTMeta(milvus::cachinglayer::StorageType storage_type)
: milvus::cachinglayer::Meta(storage_type) {
}
};
} // namespace milvus::segcore::storagev2translator
internal/core/src/segcore/storagev2translator/GroupChunkTranslator.cpp
+274
View File
@@ -0,0 +1,274 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include "segcore/storagev2translator/GroupChunkTranslator.h"
#include "segcore/storagev2translator/GroupCTMeta.h"
#include "common/GroupChunk.h"
#include "mmap/Types.h"
#include "common/Types.h"
#include "milvus-storage/common/metadata.h"
#include "milvus-storage/filesystem/fs.h"
#include "storage/ThreadPools.h"
#include "segcore/memory_planner.h"
#include <string>
#include <vector>
#include <unordered_map>
#include "arrow/type.h"
#include "arrow/type_fwd.h"
#include "cachinglayer/Utils.h"
#include "common/ChunkWriter.h"
namespace milvus::segcore::storagev2translator {
GroupChunkTranslator::GroupChunkTranslator(
int64_t segment_id,
const std::unordered_map<FieldId, FieldMeta>& field_metas,
FieldDataInfo column_group_info,
std::vector<std::string> insert_files,
bool use_mmap,
std::vector<milvus_storage::RowGroupMetadataVector>& row_group_meta_list,
milvus_storage::FieldIDList field_id_list)
: segment_id_(segment_id),
key_(fmt::format("seg_{}_cg_{}", segment_id, column_group_info.field_id)),
field_metas_(field_metas),
column_group_info_(column_group_info),
insert_files_(insert_files),
use_mmap_(use_mmap),
row_group_meta_list_(row_group_meta_list),
field_id_list_(field_id_list),
meta_(use_mmap ? milvus::cachinglayer::StorageType::DISK
: milvus::cachinglayer::StorageType::MEMORY) {
AssertInfo(insert_files_.size() == row_group_meta_list_.size(),
"Number of insert files must match number of row group metas");
meta_.num_rows_until_chunk_.push_back(0);
for (const auto& row_group_meta : row_group_meta_list_) {
for (int i = 0; i < row_group_meta.size(); ++i) {
meta_.num_rows_until_chunk_.push_back(
meta_.num_rows_until_chunk_.back() +
row_group_meta.Get(i).row_num());
meta_.chunk_memory_size_.push_back(
row_group_meta.Get(i).memory_size());
}
}
AssertInfo(
meta_.num_rows_until_chunk_.back() == column_group_info_.row_count,
fmt::format("data lost while loading column group {}: found "
"num rows {} but expected {}",
column_group_info_.field_id,
meta_.num_rows_until_chunk_.back(),
column_group_info_.row_count));
}
GroupChunkTranslator::~GroupChunkTranslator() {
for (auto chunk : group_chunks_) {
if (chunk != nullptr) {
// let the GroupChunk to be deleted by the unique_ptr
auto chunk_ptr = std::unique_ptr<GroupChunk>(chunk);
}
}
}
size_t
GroupChunkTranslator::num_cells() const {
return meta_.chunk_memory_size_.size();
}
milvus::cachinglayer::cid_t
GroupChunkTranslator::cell_id_of(milvus::cachinglayer::uid_t uid) const {
return uid;
}
milvus::cachinglayer::ResourceUsage
GroupChunkTranslator::estimated_byte_size_of_cell(
milvus::cachinglayer::cid_t cid) const {
auto [file_idx, row_group_idx] = get_file_and_row_group_index(cid);
auto& row_group_meta = row_group_meta_list_[file_idx].Get(row_group_idx);
return {static_cast<int64_t>(row_group_meta.memory_size()), 0};
}
const std::string&
GroupChunkTranslator::key() const {
return key_;
}
std::pair<size_t, size_t>
GroupChunkTranslator::get_file_and_row_group_index(
milvus::cachinglayer::cid_t cid) const {
size_t file_idx = 0;
size_t remaining_cid = cid;
for (; file_idx < row_group_meta_list_.size(); ++file_idx) {
const auto& file_metas = row_group_meta_list_[file_idx];
if (remaining_cid < file_metas.size()) {
return {file_idx, remaining_cid};
}
remaining_cid -= file_metas.size();
}
return {0, 0}; // Default to first file and first row group if not found
}
std::vector<std::pair<cachinglayer::cid_t, std::unique_ptr<milvus::GroupChunk>>>
GroupChunkTranslator::get_cells(const std::vector<cachinglayer::cid_t>& cids) {
std::vector<std::pair<milvus::cachinglayer::cid_t,
std::unique_ptr<milvus::GroupChunk>>>
cells;
cells.reserve(cids.size());
// Create row group lists for requested cids
std::vector<std::vector<int64_t>> row_group_lists;
row_group_lists.reserve(insert_files_.size());
for (size_t i = 0; i < insert_files_.size(); ++i) {
row_group_lists.emplace_back();
}
for (auto cid : cids) {
auto [file_idx, row_group_idx] = get_file_and_row_group_index(cid);
row_group_lists[file_idx].push_back(row_group_idx);
}
auto parallel_degree =
static_cast<uint64_t>(DEFAULT_FIELD_MAX_MEMORY_LIMIT / FILE_SLICE_SIZE);
auto strategy =
std::make_unique<ParallelDegreeSplitStrategy>(parallel_degree);
auto& pool = ThreadPools::GetThreadPool(milvus::ThreadPoolPriority::MIDDLE);
auto fs = milvus_storage::ArrowFileSystemSingleton::GetInstance()
.GetArrowFileSystem();
auto load_future = pool.Submit([&]() {
return LoadWithStrategy(insert_files_,
column_group_info_.arrow_reader_channel,
DEFAULT_FIELD_MAX_MEMORY_LIMIT,
std::move(strategy),
row_group_lists);
});
LOG_INFO("segment {} submits load fields {} task to thread pool",
segment_id_,
field_id_list_.ToString());
if (!use_mmap_) {
load_column_group_in_memory();
} else {
load_column_group_in_mmap();
}
for (auto cid : cids) {
AssertInfo(group_chunks_[cid] != nullptr,
"GroupChunkTranslator::get_cells failed to load cell {} of "
"CacheSlot {}.",
cid,
key_);
cells.emplace_back(
cid, std::unique_ptr<milvus::GroupChunk>(group_chunks_[cid]));
group_chunks_[cid] = nullptr;
}
return cells;
}
void
GroupChunkTranslator::load_column_group_in_memory() {
std::vector<size_t> row_counts(field_id_list_.size(), 0);
std::shared_ptr<milvus::ArrowDataWrapper> r;
std::vector<std::string> files;
std::vector<size_t> file_offsets;
while (column_group_info_.arrow_reader_channel->pop(r)) {
for (const auto& table : r->arrow_tables) {
process_batch(table, files, file_offsets, row_counts);
}
}
}
void
GroupChunkTranslator::load_column_group_in_mmap() {
std::vector<std::string> files;
std::vector<size_t> file_offsets;
std::vector<size_t> row_counts;
// Initialize files and offsets
for (size_t i = 0; i < field_id_list_.size(); ++i) {
auto field_id = field_id_list_.Get(i);
auto filepath =
std::filesystem::path(column_group_info_.mmap_dir_path) /
std::to_string(segment_id_) / std::to_string(field_id);
auto dir = filepath.parent_path();
std::filesystem::create_directories(dir);
files.push_back(filepath.string());
file_offsets.push_back(0);
row_counts.push_back(0);
}
std::shared_ptr<milvus::ArrowDataWrapper> r;
while (column_group_info_.arrow_reader_channel->pop(r)) {
for (const auto& table : r->arrow_tables) {
process_batch(table, files, file_offsets, row_counts);
}
}
}
void
GroupChunkTranslator::process_batch(const std::shared_ptr<arrow::Table>& table,
const std::vector<std::string>& files,
std::vector<size_t>& file_offsets,
std::vector<size_t>& row_counts) {
// Create chunks for each field in this batch
std::unordered_map<FieldId, std::shared_ptr<Chunk>> chunks;
// Iterate through field_id_list to get field_id and create chunk
for (size_t i = 0; i < field_id_list_.size(); ++i) {
auto field_id = field_id_list_.Get(i);
auto fid = milvus::FieldId(field_id);
if (fid == RowFieldID) {
// ignore row id field
continue;
}
auto it = field_metas_.find(fid);
AssertInfo(it != field_metas_.end(),
"Field id not found in field_metas");
const auto& field_meta = it->second;
const arrow::ArrayVector& array_vec = table->column(i)->chunks();
auto dim =
IsVectorDataType(field_meta.get_data_type()) &&
!IsSparseFloatVectorDataType(field_meta.get_data_type())
? field_meta.get_dim()
: 1;
std::unique_ptr<Chunk> chunk;
if (!use_mmap_) {
// Memory mode
chunk = create_chunk(field_meta, dim, array_vec);
} else {
// Mmap mode
int flags = O_RDWR;
if (file_offsets[i] == 0) {
// First write to this file, create and truncate
flags |= O_CREAT | O_TRUNC;
}
auto file = File::Open(files[i], flags);
// should seek to the file offset before writing
file.Seek(file_offsets[i], SEEK_SET);
chunk =
create_chunk(field_meta, dim, file, file_offsets[i], array_vec);
file_offsets[i] += chunk->Size();
}
row_counts[i] += chunk->RowNums();
chunks[fid] = std::move(chunk);
}
// Create GroupChunk from chunks and store in results
auto group_chunk = std::make_unique<milvus::GroupChunk>(chunks);
group_chunks_.emplace_back(group_chunk.release());
}
} // namespace milvus::segcore::storagev2translator
internal/core/src/segcore/storagev2translator/GroupChunkTranslator.h
+104
View File
@@ -0,0 +1,104 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <string>
#include <vector>
#include <memory>
#include <unordered_map>
#include "cachinglayer/Translator.h"
#include "cachinglayer/Utils.h"
#include "milvus-storage/common/metadata.h"
#include "mmap/Types.h"
#include "common/Types.h"
#include "common/GroupChunk.h"
#include "segcore/ChunkedSegmentSealedImpl.h"
#include "segcore/InsertRecord.h"
#include "segcore/storagev2translator/GroupCTMeta.h"
namespace milvus::segcore::storagev2translator {
class GroupChunkTranslator
: public milvus::cachinglayer::Translator<milvus::GroupChunk> {
public:
GroupChunkTranslator(
int64_t segment_id,
const std::unordered_map<FieldId, FieldMeta>& field_metas,
FieldDataInfo column_group_info,
std::vector<std::string> insert_files,
bool use_mmap,
std::vector<milvus_storage::RowGroupMetadataVector>&
row_group_meta_list,
milvus_storage::FieldIDList field_id_list);
~GroupChunkTranslator() override;
size_t
num_cells() const override;
milvus::cachinglayer::cid_t
cell_id_of(milvus::cachinglayer::uid_t uid) const override;
milvus::cachinglayer::ResourceUsage
estimated_byte_size_of_cell(milvus::cachinglayer::cid_t cid) const override;
const std::string&
key() const override;
std::vector<std::pair<milvus::cachinglayer::cid_t,
std::unique_ptr<milvus::GroupChunk>>>
get_cells(const std::vector<milvus::cachinglayer::cid_t>& cids) override;
std::pair<size_t, size_t>
get_file_and_row_group_index(milvus::cachinglayer::cid_t cid) const;
milvus::cachinglayer::Meta*
meta() override {
return &meta_;
}
private:
void
load_column_group_in_memory();
void
load_column_group_in_mmap();
void
process_batch(const std::shared_ptr<arrow::Table>& table,
const std::vector<std::string>& files,
std::vector<size_t>& file_offsets,
std::vector<size_t>& row_counts);
int64_t segment_id_;
std::string key_;
std::unordered_map<FieldId, FieldMeta> field_metas_;
FieldDataInfo column_group_info_;
std::vector<std::string> insert_files_;
std::vector<milvus_storage::RowGroupMetadataVector>& row_group_meta_list_;
milvus_storage::FieldIDList field_id_list_;
SchemaPtr schema_;
bool is_sorted_by_pk_;
ChunkedSegmentSealedImpl* chunked_segment_;
std::unique_ptr<milvus::segcore::InsertRecord<true>> ir_;
GroupCTMeta meta_;
std::vector<milvus::GroupChunk*> group_chunks_;
int64_t timestamp_offet_;
bool use_mmap_;
};
} // namespace milvus::segcore::storagev2translator
internal/core/unittest/CMakeLists.txt
+3
View File
@@ -100,6 +100,9 @@ set(MILVUS_TEST_FILES
test_growing_storage_v2.cpp
test_memory_planner.cpp
test_storage_v2_index_raw_data.cpp
test_chunked_column_group.cpp
test_group_chunk_translator.cpp
test_chunked_segment_storage_v2.cpp
)
if ( INDEX_ENGINE STREQUAL "cardinal" )
internal/core/unittest/test_cachinglayer/cachinglayer_test_utils.h
+67
View File
@@ -18,7 +18,9 @@
#include "cachinglayer/Translator.h"
#include "common/Chunk.h"
#include "common/GroupChunk.h"
#include "segcore/storagev1translator/ChunkTranslator.h"
#include "segcore/storagev2translator/GroupChunkTranslator.h"
#include "cachinglayer/lrucache/DList.h"
namespace milvus {
@@ -91,6 +93,71 @@ class TestChunkTranslator : public Translator<milvus::Chunk> {
std::vector<std::unique_ptr<Chunk>> chunks_;
};
class TestGroupChunkTranslator : public Translator<milvus::GroupChunk> {
public:
TestGroupChunkTranslator(std::vector<int64_t> num_rows_per_chunk,
std::string key,
std::vector<std::unique_ptr<GroupChunk>>&& chunks)
: Translator<milvus::GroupChunk>(),
num_cells_(num_rows_per_chunk.size()),
chunks_(std::move(chunks)),
meta_(
segcore::storagev2translator::GroupCTMeta(StorageType::MEMORY)) {
meta_.num_rows_until_chunk_.reserve(num_cells_ + 1);
meta_.num_rows_until_chunk_.push_back(0);
for (int i = 0; i < num_cells_; ++i) {
meta_.num_rows_until_chunk_.push_back(
meta_.num_rows_until_chunk_[i] + num_rows_per_chunk[i]);
}
key_ = key;
}
~TestGroupChunkTranslator() override {
}
size_t
num_cells() const override {
return num_cells_;
}
cid_t
cell_id_of(uid_t uid) const override {
return uid;
}
ResourceUsage
estimated_byte_size_of_cell(cid_t cid) const override {
return {0, 0};
}
const std::string&
key() const override {
return key_;
}
Meta*
meta() override {
return &meta_;
}
std::vector<std::pair<cid_t, std::unique_ptr<milvus::GroupChunk>>>
get_cells(const std::vector<cid_t>& cids) override {
std::vector<std::pair<cid_t, std::unique_ptr<milvus::GroupChunk>>> res;
res.reserve(cids.size());
for (auto cid : cids) {
AssertInfo(cid < chunks_.size() && chunks_[cid] != nullptr,
"TestGroupChunkTranslator assumes no eviction.");
res.emplace_back(cid, std::move(chunks_[cid]));
}
return res;
}
private:
size_t num_cells_;
std::vector<std::unique_ptr<GroupChunk>> chunks_;
std::string key_;
segcore::storagev2translator::GroupCTMeta meta_;
};
namespace cachinglayer::internal {
class DListTestFriend {
public:
internal/core/unittest/test_chunked_column_group.cpp
+273
View File
@@ -0,0 +1,273 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <gtest/gtest.h>
#include <memory>
#include <vector>
#include <arrow/array.h>
#include <arrow/io/memory.h>
#include <parquet/arrow/reader.h>
#include <cstring>
#include "common/Chunk.h"
#include "common/GroupChunk.h"
#include "common/FieldMeta.h"
#include "common/Types.h"
#include "mmap/ChunkedColumnGroup.h"
#include "test_cachinglayer/cachinglayer_test_utils.h"
#include "test_utils/storage_test_utils.h"
#include "storage/Event.h"
#include "storage/Util.h"
#include "common/ChunkWriter.h"
#include "common/FieldData.h"
using namespace milvus;
using namespace milvus::storage;
std::shared_ptr<Chunk>
create_chunk(const FixedVector<int64_t>& data) {
auto field_data =
milvus::storage::CreateFieldData(storage::DataType::INT64);
field_data->FillFieldData(data.data(), data.size());
storage::InsertEventData event_data;
auto payload_reader =
std::make_shared<milvus::storage::PayloadReader>(field_data);
event_data.payload_reader = payload_reader;
auto ser_data = event_data.Serialize();
auto buffer = std::make_shared<arrow::io::BufferReader>(
ser_data.data() + 2 * sizeof(milvus::Timestamp),
ser_data.size() - 2 * sizeof(milvus::Timestamp));
parquet::arrow::FileReaderBuilder reader_builder;
auto s = reader_builder.Open(buffer);
EXPECT_TRUE(s.ok());
std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
s = reader_builder.Build(&arrow_reader);
EXPECT_TRUE(s.ok());
std::shared_ptr<::arrow::RecordBatchReader> rb_reader;
s = arrow_reader->GetRecordBatchReader(&rb_reader);
EXPECT_TRUE(s.ok());
FieldMeta field_meta(FieldName("a"),
milvus::FieldId(1),
DataType::INT64,
false,
std::nullopt);
arrow::ArrayVector array_vec = read_single_column_batches(rb_reader);
return create_chunk(field_meta, 1, array_vec);
}
// Helper function to create chunks for string data
std::shared_ptr<Chunk>
create_chunk(const FixedVector<std::string>& data) {
auto field_data =
milvus::storage::CreateFieldData(storage::DataType::VARCHAR);
field_data->FillFieldData(data.data(), data.size());
storage::InsertEventData event_data;
auto payload_reader =
std::make_shared<milvus::storage::PayloadReader>(field_data);
event_data.payload_reader = payload_reader;
auto ser_data = event_data.Serialize();
auto buffer = std::make_shared<arrow::io::BufferReader>(
ser_data.data() + 2 * sizeof(milvus::Timestamp),
ser_data.size() - 2 * sizeof(milvus::Timestamp));
parquet::arrow::FileReaderBuilder reader_builder;
auto s = reader_builder.Open(buffer);
EXPECT_TRUE(s.ok());
std::unique_ptr<parquet::arrow::FileReader> arrow_reader;
s = reader_builder.Build(&arrow_reader);
EXPECT_TRUE(s.ok());
std::shared_ptr<::arrow::RecordBatchReader> rb_reader;
s = arrow_reader->GetRecordBatchReader(&rb_reader);
EXPECT_TRUE(s.ok());
FieldMeta field_meta(FieldName("b"),
milvus::FieldId(2),
DataType::STRING,
false,
std::nullopt);
arrow::ArrayVector array_vec = read_single_column_batches(rb_reader);
return create_chunk(field_meta, 1, array_vec);
}
// Test fixture for chunk tests
class ChunkedColumnGroupTest : public ::testing::Test {
protected:
ChunkedColumnGroupTest()
: int64_field_meta(
FieldName(""), FieldId(0), DataType::INT64, false, std::nullopt),
string_field_meta(FieldName(""),
FieldId(0),
DataType::STRING,
false,
std::nullopt) {
}
void
SetUp() override {
// Create test data
int64_data = {1, 2, 3, 4, 5};
string_data = {"a", "b", "c", "d", "e"};
// Create field metadata
int64_field_meta = FieldMeta(FieldName("int64_field"),
milvus::FieldId(1),
DataType::INT64,
false,
std::nullopt);
string_field_meta = FieldMeta(FieldName("string_field"),
milvus::FieldId(2),
DataType::STRING,
false,
std::nullopt);
// Create chunks
int64_chunk = std::move(create_chunk(int64_data));
string_chunk = std::move(create_chunk(string_data));
}
FixedVector<int64_t> int64_data;
FixedVector<std::string> string_data;
FieldMeta int64_field_meta;
FieldMeta string_field_meta;
std::shared_ptr<Chunk> int64_chunk;
std::shared_ptr<Chunk> string_chunk;
};
TEST_F(ChunkedColumnGroupTest, GroupChunk) {
std::unordered_map<FieldId, std::shared_ptr<Chunk>> chunks;
chunks[FieldId(1)] = int64_chunk;
chunks[FieldId(2)] = string_chunk;
auto group_chunk = std::make_unique<GroupChunk>(chunks);
// Has chunk
EXPECT_EQ(group_chunk->RowNums(), 5);
EXPECT_TRUE(group_chunk->HasChunk(FieldId(1)));
EXPECT_TRUE(group_chunk->HasChunk(FieldId(2)));
EXPECT_FALSE(group_chunk->HasChunk(FieldId(3)));
// Get chunk
auto retrieved_int64_chunk = group_chunk->GetChunk(FieldId(1));
auto retrieved_string_chunk = group_chunk->GetChunk(FieldId(2));
EXPECT_EQ(retrieved_int64_chunk->RowNums(), 5);
EXPECT_EQ(retrieved_string_chunk->RowNums(), 5);
// Get all chunks
auto all_chunks = group_chunk->GetChunks();
EXPECT_EQ(all_chunks.size(), 2);
EXPECT_EQ(all_chunks[FieldId(1)], int64_chunk);
EXPECT_EQ(all_chunks[FieldId(2)], string_chunk);
// Add chunk
auto new_int64_chunk = create_chunk(FixedVector<int64_t>{6, 7, 8, 9, 10});
EXPECT_NO_THROW(group_chunk->AddChunk(FieldId(3), new_int64_chunk));
EXPECT_TRUE(group_chunk->HasChunk(FieldId(3)));
EXPECT_EQ(group_chunk->GetChunk(FieldId(3))->RowNums(), 5);
auto another_int64_chunk =
create_chunk(FixedVector<int64_t>{11, 12, 13, 14, 15});
EXPECT_THROW(group_chunk->AddChunk(FieldId(3), another_int64_chunk),
std::exception);
// Size
uint64_t expected_size =
int64_chunk->Size() + string_chunk->Size() + new_int64_chunk->Size();
EXPECT_EQ(group_chunk->Size(), expected_size);
// Cell byte size
uint64_t expected_cell_size = int64_chunk->CellByteSize() +
string_chunk->CellByteSize() +
new_int64_chunk->CellByteSize();
EXPECT_EQ(group_chunk->CellByteSize(), expected_cell_size);
// Test empty group chunk
auto empty_group_chunk = std::make_unique<GroupChunk>();
EXPECT_EQ(empty_group_chunk->RowNums(), 0);
EXPECT_EQ(empty_group_chunk->Size(), 0);
EXPECT_EQ(empty_group_chunk->CellByteSize(), 0);
}
TEST_F(ChunkedColumnGroupTest, ChunkedColumnGroup) {
std::unordered_map<FieldId, std::shared_ptr<Chunk>> chunks;
chunks[FieldId(1)] = int64_chunk;
chunks[FieldId(2)] = string_chunk;
auto group_chunk = std::make_unique<GroupChunk>(chunks);
std::vector<std::unique_ptr<GroupChunk>> group_chunks;
group_chunks.push_back(std::move(group_chunk));
auto translator = std::make_unique<TestGroupChunkTranslator>(
std::vector<int64_t>{5}, "test_key", std::move(group_chunks));
auto column_group =
std::make_shared<ChunkedColumnGroup>(std::move(translator));
// basic properties
EXPECT_EQ(column_group->num_chunks(), 1);
EXPECT_EQ(column_group->NumRows(), 5);
// Get group chunk
auto retrieved_group_chunk = column_group->GetGroupChunk(0);
EXPECT_NE(retrieved_group_chunk.get(), nullptr);
EXPECT_EQ(retrieved_group_chunk.get()->RowNums(), 5);
// GetNumRowsUntilChunk
EXPECT_EQ(column_group->GetNumRowsUntilChunk(0), 0);
EXPECT_EQ(column_group->GetNumRowsUntilChunk(1), 5);
// GetNumRowsUntilChunk vector
const auto& rows_until_chunk = column_group->GetNumRowsUntilChunk();
EXPECT_EQ(rows_until_chunk.size(), 2);
EXPECT_EQ(rows_until_chunk[0], 0);
EXPECT_EQ(rows_until_chunk[1], 5);
// boundary conditions
EXPECT_THROW(column_group->GetNumRowsUntilChunk(100), std::exception); // Out of range
}
TEST_F(ChunkedColumnGroupTest, ProxyChunkColumn) {
std::unordered_map<FieldId, std::shared_ptr<Chunk>> chunks;
chunks[FieldId(1)] = int64_chunk;
chunks[FieldId(2)] = string_chunk;
auto group_chunk = std::make_unique<GroupChunk>(chunks);
std::vector<std::unique_ptr<GroupChunk>> group_chunks;
group_chunks.push_back(std::move(group_chunk));
auto translator = std::make_unique<TestGroupChunkTranslator>(
std::vector<int64_t>{5}, "test_key", std::move(group_chunks));
auto column_group =
std::make_shared<ChunkedColumnGroup>(std::move(translator));
// Test int64 proxy
auto proxy_int64 = std::make_shared<ProxyChunkColumn>(
column_group, FieldId(1), int64_field_meta);
EXPECT_EQ(proxy_int64->NumRows(), 5);
EXPECT_EQ(proxy_int64->num_chunks(), 1);
EXPECT_FALSE(proxy_int64->IsNullable());
EXPECT_NE(proxy_int64->DataOfChunk(0).get(), nullptr);
EXPECT_NE(proxy_int64->ValueAt(0), nullptr);
EXPECT_TRUE(proxy_int64->IsValid(0));
EXPECT_TRUE(proxy_int64->IsValid(0, 0));
// Test string proxy
auto proxy_string = std::make_shared<ProxyChunkColumn>(
column_group, FieldId(2), string_field_meta);
EXPECT_EQ(proxy_string->NumRows(), 5);
EXPECT_EQ(proxy_string->num_chunks(), 1);
EXPECT_FALSE(proxy_string->IsNullable());
}
internal/core/unittest/test_chunked_segment.cpp
+14 -64
View File
@@ -179,21 +179,7 @@ class TestChunkSegment : public testing::TestWithParam<bool> {
void
SetUp() override {
bool pk_is_string = GetParam();
auto schema = std::make_shared<Schema>();
auto int64_fid = schema->AddDebugField("int64", DataType::INT64, true);
auto pk_fid = schema->AddDebugField(
"pk", pk_is_string ? DataType::VARCHAR : DataType::INT64, false);
auto str_fid =
schema->AddDebugField("string1", DataType::VARCHAR, true);
auto str2_fid =
schema->AddDebugField("string2", DataType::VARCHAR, true);
schema->AddField(FieldName("ts"),
TimestampFieldID,
DataType::INT64,
true,
std::nullopt);
schema->set_primary_field_id(pk_fid);
auto schema = segcore::GenChunkedSegmentTestSchema(pk_is_string);
segment = segcore::CreateSealedSegment(
schema,
nullptr,
@@ -203,55 +189,17 @@ class TestChunkSegment : public testing::TestWithParam<bool> {
true);
test_data_count = 10000;
auto arrow_i64_field = arrow::field(
"int64",
arrow::int64(),
true,
arrow::key_value_metadata({milvus_storage::ARROW_FIELD_ID_KEY},
{std::to_string(100)}));
auto arrow_pk_field = arrow::field(
"pk",
pk_is_string ? arrow::utf8() : arrow::int64(),
false,
arrow::key_value_metadata({milvus_storage::ARROW_FIELD_ID_KEY},
{std::to_string(101)}));
auto arrow_ts_field = arrow::field(
"ts",
arrow::int64(),
true,
arrow::key_value_metadata({milvus_storage::ARROW_FIELD_ID_KEY},
{std::to_string(1)}));
auto arrow_str_field = arrow::field(
"string1",
arrow::utf8(),
true,
arrow::key_value_metadata({milvus_storage::ARROW_FIELD_ID_KEY},
{std::to_string(102)}));
auto arrow_str2_field = arrow::field(
"string2",
arrow::utf8(),
true,
arrow::key_value_metadata({milvus_storage::ARROW_FIELD_ID_KEY},
{std::to_string(103)}));
std::vector<std::shared_ptr<arrow::Field>> arrow_fields = {
arrow_ts_field,
arrow_str2_field,
arrow_str_field,
arrow_pk_field,
arrow_i64_field,
};
auto expected_arrow_schema =
std::make_shared<arrow::Schema>(arrow_fields);
ASSERT_EQ(schema->ConvertToArrowSchema()->ToString(),
expected_arrow_schema->ToString());
std::vector<FieldId> field_ids = {
int64_fid, pk_fid, TimestampFieldID, str_fid, str2_fid};
fields = {{"int64", int64_fid},
{"pk", pk_fid},
schema->get_field_id(FieldName("int64")),
schema->get_field_id(FieldName("pk")),
TimestampFieldID,
schema->get_field_id(FieldName("string1")),
schema->get_field_id(FieldName("string2"))};
fields = {{"int64", schema->get_field_id(FieldName("int64"))},
{"pk", schema->get_field_id(FieldName("pk"))},
{"ts", TimestampFieldID},
{"string1", str_fid},
{"string2", str2_fid}};
{"string1", schema->get_field_id(FieldName("string1"))},
{"string2", schema->get_field_id(FieldName("string2"))}};
int start_id = 0;
chunk_num = 2;
@@ -273,11 +221,13 @@ class TestChunkSegment : public testing::TestWithParam<bool> {
std::vector<int64_t> test_data(test_data_count);
std::iota(test_data.begin(), test_data.end(), start_id);
for (int i = 0; i < arrow_fields.size(); i++) {
for (int i = 0; i < field_ids.size(); i++) {
// if i < 3, this is system field, thus must be int64.
// other fields include a pk field and 2 string fields. pk field is string if pk_is_string is true.
auto datatype =
i < 3 && (field_ids[i] != pk_fid || !pk_is_string)
i < 3 && (field_ids[i] !=
schema->get_field_id(FieldName("pk")) ||
!pk_is_string)
? DataType::INT64
: DataType::VARCHAR;
FieldDataPtr field_data{nullptr};
internal/core/unittest/test_chunked_segment_storage_v2.cpp
+310
View File
@@ -0,0 +1,310 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <folly/Conv.h>
#include <arrow/record_batch.h>
#include <arrow/util/key_value_metadata.h>
#include <arrow/table_builder.h>
#include <arrow/type_fwd.h>
#include <gtest/gtest.h>
#include <algorithm>
#include <cstdint>
#include <iostream>
#include <memory>
#include <string>
#include <vector>
#include "common/Consts.h"
#include "common/FieldDataInterface.h"
#include "common/Schema.h"
#include "common/Types.h"
#include "expr/ITypeExpr.h"
#include "index/IndexFactory.h"
#include "index/IndexInfo.h"
#include "index/Meta.h"
#include "milvus-storage/common/constants.h"
#include "milvus-storage/filesystem/fs.h"
#include "milvus-storage/packed/writer.h"
#include "milvus-storage/format/parquet/file_reader.h"
#include "pb/plan.pb.h"
#include "pb/schema.pb.h"
#include "query/ExecPlanNodeVisitor.h"
#include "segcore/SegcoreConfig.h"
#include "segcore/SegmentSealed.h"
#include "segcore/ChunkedSegmentSealedImpl.h"
#include "segcore/SegmentGrowing.h"
#include "segcore/SegmentGrowingImpl.h"
#include "segcore/Utils.h"
#include "segcore/memory_planner.h"
#include "segcore/Types.h"
#include "test_utils/DataGen.h"
using namespace milvus;
using namespace milvus::segcore;
using namespace milvus::segcore::storagev1translator;
class TestChunkSegmentStorageV2 : public testing::TestWithParam<bool> {
protected:
void
SetUp() override {
bool pk_is_string = GetParam();
auto schema = segcore::GenChunkedSegmentTestSchema(pk_is_string);
segment = segcore::CreateSealedSegment(
schema,
nullptr,
-1,
segcore::SegcoreConfig::default_config(),
false,
true);
// Initialize file system
auto conf = milvus_storage::ArrowFileSystemConfig();
conf.storage_type = "local";
conf.root_path = "/tmp";
milvus_storage::ArrowFileSystemSingleton::GetInstance().Init(conf);
auto fs = milvus_storage::ArrowFileSystemSingleton::GetInstance()
.GetArrowFileSystem();
// Prepare paths and column groups
std::vector<std::string> paths = {"/tmp/0.parquet", "/tmp/1.parquet"};
std::vector<std::vector<int>> column_groups = {
{0, 1, 2}, {3, 4}}; // short columns and long columns
auto writer_memory = 16 * 1024 * 1024;
auto storage_config = milvus_storage::StorageConfig();
// Create writer
milvus_storage::PackedRecordBatchWriter writer(
fs,
paths,
schema->ConvertToArrowSchema(),
storage_config,
column_groups,
writer_memory);
// Generate and write data
int64_t row_count = 0;
int start_id = 0;
std::vector<std::string> str_data;
for (int i = 0; i < test_data_count * chunk_num; i++) {
str_data.push_back("test" + std::to_string(i));
}
std::sort(str_data.begin(), str_data.end());
fields = {{"int64", schema->get_field_id(FieldName("int64"))},
{"pk", schema->get_field_id(FieldName("pk"))},
{"ts", TimestampFieldID},
{"string1", schema->get_field_id(FieldName("string1"))},
{"string2", schema->get_field_id(FieldName("string2"))}};
auto arrow_schema = schema->ConvertToArrowSchema();
for (int chunk_id = 0; chunk_id < chunk_num;
chunk_id++, start_id += test_data_count) {
std::vector<int64_t> test_data(test_data_count);
std::iota(test_data.begin(), test_data.end(), start_id);
// Create arrow arrays for each field
std::vector<std::shared_ptr<arrow::Array>> arrays;
for (int i = 0; i < arrow_schema->fields().size(); i++) {
if (arrow_schema->fields()[i]->type()->id() ==
arrow::Type::INT64) {
arrow::Int64Builder builder;
auto status =
builder.AppendValues(test_data.data(), test_data_count);
EXPECT_TRUE(status.ok());
std::shared_ptr<arrow::Array> array;
status = builder.Finish(&array);
EXPECT_TRUE(status.ok());
arrays.push_back(array);
} else {
arrow::StringBuilder builder;
std::vector<std::string> str_values;
for (int j = 0; j < test_data_count; j++) {
str_values.push_back(str_data[start_id + j]);
}
auto status = builder.AppendValues(str_values);
EXPECT_TRUE(status.ok());
std::shared_ptr<arrow::Array> array;
status = builder.Finish(&array);
EXPECT_TRUE(status.ok());
arrays.push_back(array);
}
}
// Create record batch
auto record_batch = arrow::RecordBatch::Make(
schema->ConvertToArrowSchema(), test_data_count, arrays);
row_count += test_data_count;
EXPECT_TRUE(writer.Write(record_batch).ok());
}
EXPECT_TRUE(writer.Close().ok());
LoadFieldDataInfo load_info;
load_info.field_infos.emplace(
int64_t(0),
FieldBinlogInfo{int64_t(0),
static_cast<int64_t>(row_count),
std::vector<int64_t>(chunk_num * test_data_count),
false,
std::vector<std::string>({paths[0]})});
load_info.field_infos.emplace(
int64_t(1),
FieldBinlogInfo{int64_t(1),
static_cast<int64_t>(row_count),
std::vector<int64_t>(chunk_num * test_data_count),
false,
std::vector<std::string>({paths[1]})});
load_info.mmap_dir_path = "";
load_info.storage_version = 2;
segment->LoadFieldData(load_info);
}
segcore::SegmentSealedUPtr segment;
int chunk_num = 2;
int test_data_count = 10000;
std::unordered_map<std::string, FieldId> fields;
};
INSTANTIATE_TEST_SUITE_P(TestChunkSegmentStorageV2,
TestChunkSegmentStorageV2,
testing::Bool());
TEST_P(TestChunkSegmentStorageV2, TestTermExpr) {
bool pk_is_string = GetParam();
// query int64 expr
std::vector<proto::plan::GenericValue> filter_data;
for (int i = 1; i <= 10; ++i) {
proto::plan::GenericValue v;
v.set_int64_val(i);
filter_data.push_back(v);
}
auto term_filter_expr = std::make_shared<expr::TermFilterExpr>(
expr::ColumnInfo(fields.at("int64"), milvus::DataType::INT64),
filter_data);
BitsetType final;
auto plan = std::make_shared<plan::FilterBitsNode>(DEFAULT_PLANNODE_ID,
term_filter_expr);
final = query::ExecuteQueryExpr(
plan, segment.get(), chunk_num * test_data_count, MAX_TIMESTAMP);
ASSERT_EQ(10, final.count());
std::vector<proto::plan::GenericValue> filter_str_data;
for (int i = 1; i <= 10; ++i) {
proto::plan::GenericValue v;
v.set_string_val("test" + std::to_string(i));
filter_str_data.push_back(v);
}
// query pk expr
auto pk_term_filter_expr = std::make_shared<expr::TermFilterExpr>(
expr::ColumnInfo(
fields.at("pk"),
pk_is_string ? milvus::DataType::VARCHAR : milvus::DataType::INT64),
pk_is_string ? filter_str_data : filter_data);
plan = std::make_shared<plan::FilterBitsNode>(DEFAULT_PLANNODE_ID,
pk_term_filter_expr);
final = query::ExecuteQueryExpr(
plan, segment.get(), chunk_num * test_data_count, MAX_TIMESTAMP);
ASSERT_EQ(10, final.count());
// query pk in second chunk
std::vector<proto::plan::GenericValue> filter_data2;
proto::plan::GenericValue v;
if (pk_is_string) {
v.set_string_val("test" + std::to_string(test_data_count + 1));
} else {
v.set_int64_val(test_data_count + 1);
}
filter_data2.push_back(v);
pk_term_filter_expr = std::make_shared<expr::TermFilterExpr>(
expr::ColumnInfo(
fields.at("pk"),
pk_is_string ? milvus::DataType::VARCHAR : milvus::DataType::INT64),
filter_data2);
plan = std::make_shared<plan::FilterBitsNode>(DEFAULT_PLANNODE_ID,
pk_term_filter_expr);
final = query::ExecuteQueryExpr(
plan, segment.get(), chunk_num * test_data_count, MAX_TIMESTAMP);
ASSERT_EQ(1, final.count());
}
TEST_P(TestChunkSegmentStorageV2, TestCompareExpr) {
srand(time(NULL));
bool pk_is_string = GetParam();
milvus::DataType pk_data_type =
pk_is_string ? milvus::DataType::VARCHAR : milvus::DataType::INT64;
auto expr = std::make_shared<expr::CompareExpr>(
pk_is_string ? fields.at("string1") : fields.at("int64"),
fields.at("pk"),
pk_data_type,
pk_data_type,
proto::plan::OpType::Equal);
auto plan =
std::make_shared<plan::FilterBitsNode>(DEFAULT_PLANNODE_ID, expr);
BitsetType final = query::ExecuteQueryExpr(
plan, segment.get(), chunk_num * test_data_count, MAX_TIMESTAMP);
ASSERT_EQ(chunk_num * test_data_count, final.count());
expr = std::make_shared<expr::CompareExpr>(fields.at("string1"),
fields.at("string2"),
milvus::DataType::VARCHAR,
milvus::DataType::VARCHAR,
proto::plan::OpType::Equal);
plan = std::make_shared<plan::FilterBitsNode>(DEFAULT_PLANNODE_ID, expr);
final = query::ExecuteQueryExpr(
plan, segment.get(), chunk_num * test_data_count, MAX_TIMESTAMP);
ASSERT_EQ(chunk_num * test_data_count, final.count());
// test with inverted index
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 = rand();
index_meta.index_version = rand();
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::INVERTED_INDEX_TYPE;
auto index = index::IndexFactory::GetInstance().CreateScalarIndex(
create_index_info, file_manager_ctx);
std::vector<int64_t> data(test_data_count * chunk_num);
auto pw = segment->chunk_data<int64_t>(fid, 0);
auto d = pw.get();
std::copy(
d.data(), d.data() + test_data_count, data.begin() + test_data_count);
index->BuildWithRawDataForUT(data.size(), data.data());
segcore::LoadIndexInfo load_index_info;
load_index_info.index = std::move(index);
load_index_info.field_id = fid.get();
segment->LoadIndex(load_index_info);
expr = std::make_shared<expr::CompareExpr>(
pk_is_string ? fields.at("string1") : fields.at("int64"),
fields.at("pk"),
pk_data_type,
pk_data_type,
proto::plan::OpType::Equal);
plan = std::make_shared<plan::FilterBitsNode>(DEFAULT_PLANNODE_ID, expr);
final = query::ExecuteQueryExpr(
plan, segment.get(), chunk_num * test_data_count, MAX_TIMESTAMP);
ASSERT_EQ(chunk_num * test_data_count, final.count());
}
internal/core/unittest/test_delete_record.cpp
+4 -4
View File
@@ -38,12 +38,12 @@ TEST(DeleteMVCC, common_case) {
auto pks = dataset.get_col<int64_t>(pk);
auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);
ASSERT_EQ(c, segment->get_real_count());
auto insert_record = segment->get_insert_record_slot();
auto& insert_record = segment->get_insert_record();
auto segment_ptr = segment.get();
DeletedRecord<true> delete_record(
insert_record,
[segment_ptr](const PkType& pk, Timestamp timestamp) {
return segment_ptr->search_pk(pk, timestamp);
&insert_record,
[&insert_record](const PkType& pk, Timestamp timestamp) {
return insert_record.search_pk(pk, timestamp);
},
0);
delete_record.set_sealed_row_count(c);
internal/core/unittest/test_group_chunk_translator.cpp
+162
View File
@@ -0,0 +1,162 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <folly/Conv.h>
#include <arrow/record_batch.h>
#include <arrow/util/key_value_metadata.h>
#include <gtest/gtest.h>
#include <cstdint>
#include "arrow/type_fwd.h"
#include "common/Schema.h"
#include "common/Types.h"
#include "gtest/gtest.h"
#include "milvus-storage/filesystem/fs.h"
#include "milvus-storage/packed/writer.h"
#include "milvus-storage/format/parquet/file_reader.h"
#include "test_utils/DataGen.h"
#include "segcore/storagev2translator/GroupChunkTranslator.h"
#include <memory>
#include <string>
#include <vector>
#include <filesystem>
using namespace milvus;
using namespace milvus::segcore;
using namespace milvus::segcore::storagev2translator;
class TestGroupChunkTranslator : public ::testing::TestWithParam<bool> {
void
SetUp() override {
auto conf = milvus_storage::ArrowFileSystemConfig();
conf.storage_type = "local";
conf.root_path = path_;
milvus_storage::ArrowFileSystemSingleton::GetInstance().Init(conf);
fs_ = milvus_storage::ArrowFileSystemSingleton::GetInstance()
.GetArrowFileSystem();
schema_ = CreateTestSchema();
arrow_schema_ = schema_->ConvertToArrowSchema();
int64_t per_batch = 1000;
int64_t n_batch = 3;
int64_t dim = 128;
// Write data to storage v2
paths_ = std::vector<std::string>{path_ + "/19530.parquet"};
auto column_groups = std::vector<std::vector<int>>{
{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}};
auto writer_memory = 16 * 1024 * 1024;
auto storage_config = milvus_storage::StorageConfig();
milvus_storage::PackedRecordBatchWriter writer(fs_,
paths_,
arrow_schema_,
storage_config,
column_groups,
writer_memory);
int64_t total_rows = 0;
for (int64_t i = 0; i < n_batch; i++) {
auto dataset = DataGen(schema_, per_batch);
auto record_batch =
ConvertToArrowRecordBatch(dataset, dim, arrow_schema_);
total_rows += record_batch->num_rows();
EXPECT_TRUE(writer.Write(record_batch).ok());
}
EXPECT_TRUE(writer.Close().ok());
}
protected:
~TestGroupChunkTranslator() {
if (GetParam()) { // if use_mmap is true
std::string mmap_dir = std::to_string(segment_id_);
if (std::filesystem::exists(mmap_dir)) {
std::filesystem::remove_all(mmap_dir);
}
}
}
SchemaPtr schema_;
milvus_storage::ArrowFileSystemPtr fs_;
std::shared_ptr<arrow::Schema> arrow_schema_;
std::string path_ = "/tmp";
std::vector<std::string> paths_;
int64_t segment_id_ = 0;
};
TEST_P(TestGroupChunkTranslator, TestWithMmap) {
auto use_mmap = GetParam();
std::unordered_map<FieldId, FieldMeta> field_metas = schema_->get_fields();
auto column_group_info = FieldDataInfo(0, 3000, "");
// Get row group metadata
std::vector<milvus_storage::RowGroupMetadataVector> row_group_meta_list;
auto fr =
std::make_shared<milvus_storage::FileRowGroupReader>(fs_, paths_[0]);
auto field_id_list =
fr->file_metadata()->GetGroupFieldIDList().GetFieldIDList(0);
row_group_meta_list.push_back(
fr->file_metadata()->GetRowGroupMetadataVector());
GroupChunkTranslator translator(segment_id_,
field_metas,
column_group_info,
paths_,
use_mmap,
row_group_meta_list,
field_id_list);
// num cells
EXPECT_EQ(translator.num_cells(), row_group_meta_list[0].size());
// cell id of
for (size_t i = 0; i < translator.num_cells(); ++i) {
EXPECT_EQ(translator.cell_id_of(i), i);
}
// key
EXPECT_EQ(translator.key(), "seg_0_cg_0");
// estimated byte size
for (size_t i = 0; i < translator.num_cells(); ++i) {
auto [file_idx, row_group_idx] =
translator.get_file_and_row_group_index(i);
auto& row_group_meta = row_group_meta_list[file_idx].Get(row_group_idx);
auto usage = translator.estimated_byte_size_of_cell(i);
EXPECT_EQ(usage.memory_bytes,
static_cast<int64_t>(row_group_meta.memory_size()));
}
// getting cells
std::vector<cachinglayer::cid_t> cids = {0, 1};
auto cells = translator.get_cells(cids);
EXPECT_EQ(cells.size(), cids.size());
// Verify mmap directory and files if in mmap mode
if (use_mmap) {
std::string mmap_dir = std::to_string(segment_id_);
EXPECT_TRUE(std::filesystem::exists(mmap_dir));
// Verify each field has a corresponding file
for (size_t i = 0; i < field_id_list.size(); ++i) {
auto field_id = field_id_list.Get(i);
std::string field_file = mmap_dir + "/" + std::to_string(field_id);
EXPECT_TRUE(std::filesystem::exists(field_file));
}
}
}
INSTANTIATE_TEST_SUITE_P(TestGroupChunkTranslator,
TestGroupChunkTranslator,
testing::Bool());
internal/core/unittest/test_sealed.cpp
+1 -6
View File
@@ -1308,12 +1308,7 @@ TEST(Sealed, DeleteCount) {
auto dataset = DataGen(schema, N);
auto segment = CreateSealedWithFieldDataLoaded(schema, dataset);
auto insert_record_slot = segment->get_insert_record_slot();
auto ca = milvus::cachinglayer::SemiInlineGet(
insert_record_slot->PinCells({0}));
auto insert_record = ca->get_cell_of(0);
insert_record->seal_pks();
segment->get_insert_record().seal_pks();
int64_t c = 10;
ASSERT_EQ(segment->get_deleted_count(), 0);
internal/core/unittest/test_utils/DataGen.h
+70 -9
View File
@@ -336,15 +336,59 @@ GenerateRandomSparseFloatVector(size_t rows,
return tensor;
}
inline GeneratedData DataGen(SchemaPtr schema,
int64_t N,
uint64_t seed = 42,
uint64_t ts_offset = 0,
int repeat_count = 1,
int array_len = 10,
bool random_pk = false,
bool random_val = true,
bool random_valid = false) {
inline SchemaPtr CreateTestSchema() {
auto schema = std::make_shared<milvus::Schema>();
auto bool_field =
schema->AddDebugField("bool", milvus::DataType::BOOL, true);
auto int8_field =
schema->AddDebugField("int8", milvus::DataType::INT8, true);
auto int16_field =
schema->AddDebugField("int16", milvus::DataType::INT16, true);
auto int32_field =
schema->AddDebugField("int32", milvus::DataType::INT32, true);
auto int64_field = schema->AddDebugField("int64", milvus::DataType::INT64);
auto float_field =
schema->AddDebugField("float", milvus::DataType::FLOAT, true);
auto double_field =
schema->AddDebugField("double", milvus::DataType::DOUBLE, true);
auto varchar_field =
schema->AddDebugField("varchar", milvus::DataType::VARCHAR, true);
auto json_field =
schema->AddDebugField("json", milvus::DataType::JSON, true);
auto int_array_field = schema->AddDebugField(
"int_array", milvus::DataType::ARRAY, milvus::DataType::INT8, true);
auto long_array_field = schema->AddDebugField(
"long_array", milvus::DataType::ARRAY, milvus::DataType::INT64, true);
auto bool_array_field = schema->AddDebugField(
"bool_array", milvus::DataType::ARRAY, milvus::DataType::BOOL, true);
auto string_array_field = schema->AddDebugField("string_array",
milvus::DataType::ARRAY,
milvus::DataType::VARCHAR,
true);
auto double_array_field = schema->AddDebugField("double_array",
milvus::DataType::ARRAY,
milvus::DataType::DOUBLE,
true);
auto float_array_field = schema->AddDebugField(
"float_array", milvus::DataType::ARRAY, milvus::DataType::FLOAT, true);
auto vec = schema->AddDebugField("embeddings",
milvus::DataType::VECTOR_FLOAT,
128,
knowhere::metric::L2);
schema->set_primary_field_id(int64_field);
return schema;
}
inline GeneratedData
DataGen(SchemaPtr schema,
int64_t N,
uint64_t seed = 42,
uint64_t ts_offset = 0,
int repeat_count = 1,
int array_len = 10,
bool random_pk = false,
bool random_val = true,
bool random_valid = false) {
using std::vector;
std::default_random_engine random(seed);
std::normal_distribution<> distr(0, 1);
@@ -1422,4 +1466,21 @@ gen_all_data_types_schema() {
return schema;
}
inline SchemaPtr
GenChunkedSegmentTestSchema(bool pk_is_string) {
auto schema = std::make_shared<Schema>();
auto int64_fid = schema->AddDebugField("int64", DataType::INT64, true);
auto pk_fid = schema->AddDebugField(
"pk", pk_is_string ? DataType::VARCHAR : DataType::INT64, false);
auto str_fid = schema->AddDebugField("string1", DataType::VARCHAR, true);
auto str2_fid = schema->AddDebugField("string2", DataType::VARCHAR, true);
schema->AddField(FieldName("ts"),
TimestampFieldID,
DataType::INT64,
false,
std::nullopt);
schema->set_primary_field_id(pk_fid);
return schema;
}
} // namespace milvus::segcore