fix: align growing source storage v3 flush (#50738)

issue: #50697

Summary:

1. Align StorageV3 column group layout for growing-source flush by
reusing the segment currentSplit/manifest layout, passing a schema-based
split pattern to C++, and using the schema_based writer policy instead
of always writing a single column group.

2. Project growing flush output to the target segment layout by passing
AllowedFieldIDs to C++, so ordinary fields, vector fields, text fields,
and BM25 output fields are all filtered to the old segment layout when
appending to an existing StorageV3 manifest.

3. Stop retrying non-retryable layout mismatches and preserve V3 layout
metadata, so column count/group mismatches do not loop forever and
recovery/balance can restore currentSplit for future layout-compatible
appends.

Signed-off-by: luzhang <luzhang@zilliz.com>
Co-authored-by: luzhang <luzhang@zilliz.com>
This commit is contained in:
zhagnlu
2026-06-27 19:58:26 +08:00
committed by GitHub
co-authored by luzhang
parent 0909c274bc
commit 2310c93b15
25 changed files with 1139 additions and 149 deletions
+9
View File
@@ -26,6 +26,7 @@ import (
type Params struct {
StorageVersion int64 `json:"storage_version,omitempty"`
StorageFormat string `json:"storage_format,omitempty"`
BinLogMaxSize uint64 `json:"binlog_max_size,omitempty"`
UseMergeSort bool `json:"use_merge_sort,omitempty"`
MaxSegmentMergeSort int `json:"max_segment_merge_sort,omitempty"`
@@ -46,6 +47,7 @@ func GenParams() Params {
}
return Params{
StorageVersion: storageVersion,
StorageFormat: paramtable.Get().DataNodeCfg.StorageFormat.GetValue(),
BinLogMaxSize: paramtable.Get().DataNodeCfg.BinLogMaxSize.GetAsUint64(),
UseMergeSort: paramtable.Get().DataNodeCfg.UseMergeSort.GetAsBool(),
MaxSegmentMergeSort: paramtable.Get().DataNodeCfg.MaxSegmentMergeSort.GetAsInt(),
@@ -60,6 +62,13 @@ func GenParams() Params {
}
}
func (p Params) GetStorageFormat() string {
if p.StorageFormat != "" {
return p.StorageFormat
}
return paramtable.Get().DataNodeCfg.StorageFormat.GetValue()
}
func GenerateJSONParams(schema *schemapb.CollectionSchema) (string, error) {
compactionParams := GenParams()
// TEXT fields require at least V3 manifest storage for LOB support.
+1
View File
@@ -41,6 +41,7 @@ func TestGetJSONParams(t *testing.T) {
assert.NoError(t, err)
assert.Equal(t, Params{
StorageVersion: storageVersion,
StorageFormat: paramtable.Get().DataNodeCfg.StorageFormat.GetValue(),
BinLogMaxSize: paramtable.Get().DataNodeCfg.BinLogMaxSize.GetAsUint64(),
UseMergeSort: paramtable.Get().DataNodeCfg.UseMergeSort.GetAsBool(),
MaxSegmentMergeSort: paramtable.Get().DataNodeCfg.MaxSegmentMergeSort.GetAsInt(),
@@ -238,8 +238,28 @@ class FlushGrowingSegmentTest : public ::testing::Test {
auto buffer = buffer_result.ValueOrDie();
return ParseBM25StatsBlob(buffer->data(), buffer->size());
}
std::string
SerializeSchemaBlob(const SchemaPtr& schema) {
auto schema_proto = schema->ToProto();
std::string blob;
EXPECT_TRUE(schema_proto.SerializeToString(&blob));
EXPECT_FALSE(blob.empty());
return blob;
}
void
SetFlushSchema(CFlushConfig& config, const std::string& schema_blob) {
config.schema_blob = schema_blob.data();
config.schema_length = static_cast<int64_t>(schema_blob.size());
}
};
#define C_FLUSH_CONFIG_WITH_SCHEMA(config, schema) \
auto config##_schema_blob = SerializeSchemaBlob(schema); \
CFlushConfig config{}; \
SetFlushSchema(config, config##_schema_blob)
// test basic flush with scalar fields
TEST_F(FlushGrowingSegmentTest, BasicFlushScalarFields) {
// create schema with scalar fields
@@ -264,7 +284,7 @@ TEST_F(FlushGrowingSegmentTest, BasicFlushScalarFields) {
dataset.raw_);
// prepare flush config
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -312,7 +332,7 @@ TEST_F(FlushGrowingSegmentTest, FlushAllowsStaleReadVersionOverwrite) {
std::string segment_path = test_dir_ + "/segment_stale_read_version";
CFlushConfig first_config{};
C_FLUSH_CONFIG_WITH_SCHEMA(first_config, schema);
first_config.segment_path = segment_path.c_str();
first_config.read_version = -1;
first_config.retry_limit = 3;
@@ -324,7 +344,7 @@ TEST_F(FlushGrowingSegmentTest, FlushAllowsStaleReadVersionOverwrite) {
ASSERT_EQ(first_result.num_rows, 50);
auto acknowledged_version = first_result.committed_version;
CFlushConfig orphan_config{};
C_FLUSH_CONFIG_WITH_SCHEMA(orphan_config, schema);
orphan_config.segment_path = segment_path.c_str();
orphan_config.read_version = acknowledged_version;
orphan_config.retry_limit = 3;
@@ -336,7 +356,7 @@ TEST_F(FlushGrowingSegmentTest, FlushAllowsStaleReadVersionOverwrite) {
ASSERT_EQ(orphan_result.num_rows, 25);
ASSERT_GT(orphan_result.committed_version, acknowledged_version);
CFlushConfig retry_config{};
C_FLUSH_CONFIG_WITH_SCHEMA(retry_config, schema);
retry_config.segment_path = segment_path.c_str();
retry_config.read_version = acknowledged_version;
retry_config.retry_limit = 3;
@@ -379,7 +399,7 @@ TEST_F(FlushGrowingSegmentTest, FlushUsesWriterFormatFromConfig) {
dataset.timestamps_.data(),
dataset.raw_);
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_writer_format";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -422,7 +442,7 @@ TEST_F(FlushGrowingSegmentTest, FlushWithVectorFields) {
dataset.raw_);
// prepare flush config
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_vec";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -466,7 +486,7 @@ TEST_F(FlushGrowingSegmentTest, FlushWithStringFields) {
dataset.raw_);
// prepare flush config
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_str";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -510,7 +530,7 @@ TEST_F(FlushGrowingSegmentTest, FlushPartialRange) {
dataset.raw_);
// prepare flush config
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_partial";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -559,7 +579,7 @@ TEST_F(FlushGrowingSegmentTest, FlushWithTextColumnConfig) {
dataset.raw_);
// prepare flush config with TEXT column
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_text";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -610,7 +630,7 @@ TEST_F(FlushGrowingSegmentTest, FlushWithNullableFields) {
dataset.raw_);
// prepare flush config
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_nullable";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -709,7 +729,7 @@ TEST_F(FlushGrowingSegmentTest, FlushOrdinaryFieldSemanticsRoundTrip) {
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_semantics";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -826,7 +846,7 @@ TEST_F(FlushGrowingSegmentTest, FlushTimestamptzPartialRangeRoundTrip) {
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_timestamptz";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -919,7 +939,7 @@ TEST_F(FlushGrowingSegmentTest, FlushNullableScalarTypesPartialRangeRoundTrip) {
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_nullable_scalars";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -1022,7 +1042,7 @@ TEST_F(FlushGrowingSegmentTest, FlushVectorArrayRoundTrip) {
dataset.timestamps_.data(),
dataset.raw_);
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_vector_array";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -1106,7 +1126,7 @@ TEST_F(FlushGrowingSegmentTest, FlushNullableVectorArrayRoundTrip) {
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_nullable_vector_array";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -1187,7 +1207,7 @@ TEST_F(FlushGrowingSegmentTest, FlushVectorArrayElementTypesRoundTrip) {
dataset.timestamps_.data(),
dataset.raw_);
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path =
test_dir_ + "/segment_vector_array_" + test_case.segment_suffix;
config.segment_path = segment_path.c_str();
@@ -1265,7 +1285,7 @@ TEST_F(FlushGrowingSegmentTest, FlushStringAndTextRoundTrip) {
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_string_text";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -1379,7 +1399,7 @@ TEST_F(FlushGrowingSegmentTest, FlushArrayElementTypesRoundTrip) {
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_arrays";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -1498,7 +1518,7 @@ TEST_F(FlushGrowingSegmentTest, FlushNullableFloatVectorKeepsCompactMapping) {
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_nullable_vec";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -1566,7 +1586,7 @@ TEST_F(FlushGrowingSegmentTest, FlushNullableInt8VectorKeepsCompactMapping) {
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_int8_vec";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -1669,7 +1689,7 @@ TEST_F(FlushGrowingSegmentTest, FlushNullableFixedWidthVectorTypesRoundTrip) {
segment->Insert(
0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path =
test_dir_ + "/segment_nullable_" + test_case.segment_suffix;
config.segment_path = segment_path.c_str();
@@ -1777,7 +1797,7 @@ TEST_F(FlushGrowingSegmentTest, FlushNullableSparseVectorKeepsCompactMapping) {
segment->PreInsert(N);
segment->Insert(0, N, row_ids.data(), timestamps.data(), insert_data.get());
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_sparse_vec";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -1866,7 +1886,7 @@ TEST_F(FlushGrowingSegmentTest, FlushBM25StatsRangeAndCompoundManifest) {
int64_t first_bm25_stats_log_ids[] = {1001};
int64_t second_bm25_stats_log_ids[] = {1002};
CFlushConfig first_config{};
C_FLUSH_CONFIG_WITH_SCHEMA(first_config, schema);
first_config.segment_path = segment_path.c_str();
first_config.read_version = -1;
first_config.retry_limit = 3;
@@ -1889,7 +1909,7 @@ TEST_F(FlushGrowingSegmentTest, FlushBM25StatsRangeAndCompoundManifest) {
EXPECT_EQ(first_stats->rows_with_token[10], 2);
EXPECT_EQ(first_stats->rows_with_token[20], 1);
CFlushConfig second_config{};
C_FLUSH_CONFIG_WITH_SCHEMA(second_config, schema);
second_config.segment_path = segment_path.c_str();
second_config.read_version = first_result.committed_version;
second_config.retry_limit = 3;
@@ -1982,7 +2002,7 @@ TEST_F(FlushGrowingSegmentTest, FlushEmptyRange) {
dataset.raw_);
// prepare flush config
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_empty";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -2028,7 +2048,7 @@ TEST_F(FlushGrowingSegmentTest, FlushLargeDataMultipleChunks) {
dataset.raw_);
// prepare flush config
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_large";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -2076,7 +2096,7 @@ TEST_F(FlushGrowingSegmentTest, FlushMultipleTextColumns) {
dataset.raw_);
// prepare flush config with multiple TEXT columns
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_multi_text";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -2129,7 +2149,7 @@ TEST_F(FlushGrowingSegmentTest, FlushWithBoolField) {
dataset.raw_);
// prepare flush config
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_bool";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -2177,7 +2197,7 @@ TEST_F(FlushGrowingSegmentTest, FlushAllNumericTypes) {
dataset.raw_);
// prepare flush config
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_numeric";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -2220,7 +2240,7 @@ TEST_F(FlushGrowingSegmentTest, FlushDifferentVectorTypes) {
dataset.timestamps_.data(),
dataset.raw_);
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_fp16";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -2259,7 +2279,7 @@ TEST_F(FlushGrowingSegmentTest, FlushDifferentVectorTypes) {
dataset.timestamps_.data(),
dataset.raw_);
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_bf16";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -2298,7 +2318,7 @@ TEST_F(FlushGrowingSegmentTest, FlushDifferentVectorTypes) {
dataset.timestamps_.data(),
dataset.raw_);
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_binary";
config.segment_path = segment_path.c_str();
config.read_version = -1;
@@ -2341,7 +2361,7 @@ TEST_F(FlushGrowingSegmentTest, FlushSealedSegmentFails) {
ASSERT_NE(segment, nullptr);
// prepare flush config
CFlushConfig config{};
C_FLUSH_CONFIG_WITH_SCHEMA(config, schema);
std::string segment_path = test_dir_ + "/segment_sealed";
config.segment_path = segment_path.c_str();
config.read_version = -1;
+73 -6
View File
@@ -202,6 +202,17 @@ ParseReopenSchema(const void* schema_blob,
return schema;
}
milvus::SchemaPtr
ParseFlushSchema(const void* schema_blob, const int64_t schema_length) {
AssertInfo(schema_blob != nullptr, "flush schema is null");
AssertInfo(schema_length > 0, "flush schema length must be positive");
milvus::proto::schema::CollectionSchema collection_schema;
auto suc = collection_schema.ParseFromArray(schema_blob, schema_length);
AssertInfo(suc, "parse flush schema proto failed");
return milvus::Schema::ParseFrom(collection_schema);
}
CFuture*
AsyncReopenSegment(CTraceContext c_trace,
CSegmentInterface c_segment,
@@ -1749,6 +1760,18 @@ FlushGrowingSegmentData(CSegmentInterface c_segment,
milvus::UnexpectedError,
"invalid BM25 config: bm25_stats_log_ids is null");
}
if (config->num_allowed_fields > 0 &&
config->allowed_field_ids == nullptr) {
return milvus::FailureCStatus(
milvus::UnexpectedError,
"invalid allowed field config: allowed_field_ids is null");
}
if (config->schema_blob == nullptr || config->schema_length <= 0) {
return milvus::FailureCStatus(
milvus::UnexpectedError,
"invalid flush schema config: schema_blob is null or "
"schema_length is not positive");
}
// no data to flush
if (start_offset == end_offset) {
@@ -1765,8 +1788,11 @@ FlushGrowingSegmentData(CSegmentInterface c_segment,
"segment is not a growing segment");
}
// get schema from segment
auto& schema = growing_segment->get_schema();
// Use the schema selected by the flush task. The growing segment's
// runtime schema may be advanced by concurrent LazyCheckSchema/Reopen.
auto flush_schema =
ParseFlushSchema(config->schema_blob, config->schema_length);
const auto& schema = *flush_schema;
auto& insert_record = growing_segment->get_insert_record();
int64_t total_rows = end_offset - start_offset;
@@ -1777,6 +1803,10 @@ FlushGrowingSegmentData(CSegmentInterface c_segment,
bm25_stats_log_ids[config->bm25_field_ids[i]] =
config->bm25_stats_log_ids[i];
}
std::unordered_set<int64_t> allowed_field_ids;
for (size_t i = 0; i < config->num_allowed_fields; i++) {
allowed_field_ids.insert(config->allowed_field_ids[i]);
}
std::unordered_map<int64_t, BM25StatsAccumulator> bm25_stats;
// Use get_field_ids() (ordered vector) instead of get_fields() (unordered_map)
@@ -1814,6 +1844,11 @@ FlushGrowingSegmentData(CSegmentInterface c_segment,
if (field_id == RowFieldID) {
continue; // skip RowID system field
}
if (!allowed_field_ids.empty() && field_id != TimestampFieldID &&
allowed_field_ids.find(field_id.get()) ==
allowed_field_ids.end()) {
continue;
}
const auto& field_meta = schema[field_id];
@@ -1930,14 +1965,34 @@ FlushGrowingSegmentData(CSegmentInterface c_segment,
}
// set required properties for ColumnGroupPolicy
// use single column group policy (all columns in one group)
writer_config.properties[PROPERTY_WRITER_POLICY] =
std::string(LOON_COLUMN_GROUP_POLICY_SINGLE);
if (config->schema_based_pattern &&
config->schema_based_pattern[0] != '\0') {
milvus_storage::api::SetValue(
writer_config.properties,
PROPERTY_WRITER_POLICY,
LOON_COLUMN_GROUP_POLICY_SCHEMA_BASED);
milvus_storage::api::SetValue(writer_config.properties,
PROPERTY_WRITER_SCHEMA_BASE_PATTERNS,
config->schema_based_pattern);
if (config->schema_based_formats &&
config->schema_based_formats[0] != '\0') {
milvus_storage::api::SetValue(
writer_config.properties,
PROPERTY_WRITER_SCHEMA_BASE_FORMATS,
config->schema_based_formats);
}
} else {
milvus_storage::api::SetValue(writer_config.properties,
PROPERTY_WRITER_POLICY,
LOON_COLUMN_GROUP_POLICY_SINGLE);
}
auto writer_format =
config->writer_format && config->writer_format[0] != '\0'
? std::string(config->writer_format)
: std::string(LOON_FORMAT_PARQUET);
writer_config.properties[PROPERTY_WRITER_FORMAT] = writer_format;
milvus_storage::api::SetValue(writer_config.properties,
PROPERTY_WRITER_FORMAT,
writer_format.c_str());
// add TEXT column configs
for (size_t i = 0; i < config->num_text_columns; i++) {
@@ -1946,6 +2001,18 @@ FlushGrowingSegmentData(CSegmentInterface c_segment,
if (config->text_lob_paths && config->text_lob_paths[i]) {
text_config.lob_base_path = config->text_lob_paths[i];
}
if (config->text_inline_threshold > 0) {
text_config.inline_threshold =
static_cast<size_t>(config->text_inline_threshold);
}
if (config->text_max_lob_file_bytes > 0) {
text_config.max_lob_file_bytes =
static_cast<size_t>(config->text_max_lob_file_bytes);
}
if (config->text_flush_threshold_bytes > 0) {
text_config.flush_threshold_bytes =
static_cast<size_t>(config->text_flush_threshold_bytes);
}
text_config.properties = writer_config.properties;
writer_config.lob_columns[text_config.field_id] = text_config;
}
+15 -6
View File
@@ -300,13 +300,22 @@ ExprResCacheEraseSegment(int64_t segment_id);
* @brief Configuration for flushing growing segment data to storage.
*/
typedef struct CFlushConfig {
const char* segment_path; // base path for segment manifest and data
int64_t read_version; // version to read (-1 = latest)
uint32_t retry_limit; // retry limit for commit
const char* writer_format; // writer.format
const char* segment_path; // base path for segment manifest and data
const void* schema_blob; // serialized CollectionSchema for this flush task
int64_t schema_length; // length of schema_blob in bytes
int64_t read_version; // version to read (-1 = latest)
uint32_t retry_limit; // retry limit for commit
const char* writer_format; // writer.format
const char* schema_based_pattern; // writer.split.schema_based.patterns
const char* schema_based_formats; // writer.split.schema_based.formats
int64_t* allowed_field_ids; // projected field IDs to flush
size_t num_allowed_fields; // number of projected fields
// TEXT column configurations
int64_t* text_field_ids; // array of TEXT field IDs
const char** text_lob_paths; // array of LOB paths for each TEXT field
int64_t* text_field_ids; // array of TEXT field IDs
const char** text_lob_paths; // array of LOB paths for each TEXT field
int64_t text_inline_threshold;
int64_t text_max_lob_file_bytes;
int64_t text_flush_threshold_bytes;
size_t num_text_columns; // number of TEXT columns
int64_t* bm25_field_ids; // array of BM25 sparse output field IDs
int64_t* bm25_stats_log_ids; // array of BM25 stats log IDs
@@ -1070,6 +1070,7 @@ func (t *clusteringCompactionTask) getWriterOpts() []storage.RwOption {
storage.WithBufferSize(t.bufferSize),
storage.WithStorageConfig(t.compactionParams.StorageConfig),
storage.WithUseLoonFFI(t.compactionParams.UseLoonFFI),
storage.WithWriterFormat(t.compactionParams.GetStorageFormat()),
}
if t.lobContext != nil && t.lobContext.ShouldRewriteAnyField() {
@@ -170,6 +170,7 @@ func (t *mixCompactionTask) buildWriterOptions(ctx context.Context) []storage.Rw
writerOpts := []storage.RwOption{
storage.WithStorageConfig(t.compactionParams.StorageConfig),
storage.WithUseLoonFFI(t.compactionParams.UseLoonFFI),
storage.WithWriterFormat(t.compactionParams.GetStorageFormat()),
}
if t.lobContext != nil && t.lobContext.ShouldRewriteAnyField() {
@@ -180,6 +180,7 @@ func (t *sortCompactionTask) sortSegment(ctx context.Context) (*datapb.Compactio
storage.WithVersion(t.storageVersion),
storage.WithStorageConfig(t.compactionParams.StorageConfig),
storage.WithUseLoonFFI(t.useLoonFFI),
storage.WithWriterFormat(t.compactionParams.GetStorageFormat()),
}
if t.lobContext != nil && t.lobContext.HasReuseAllFields() {
writerOpts = append(writerOpts, storage.WithTextRefsAsBinary())
+152 -30
View File
@@ -21,6 +21,7 @@ import (
"fmt"
"path"
"sort"
"strings"
"sync"
"github.com/cockroachdb/errors"
@@ -32,6 +33,7 @@ import (
"github.com/milvus-io/milvus/internal/allocator"
"github.com/milvus-io/milvus/internal/flushcommon/metacache"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/storagecommon"
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/metrics"
@@ -42,19 +44,28 @@ import (
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/retry"
"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
type GrowingFlushConfig struct {
SegmentBasePath string
PartitionBasePath string
CollectionID int64
PartitionID int64
TextFieldIDs []int64
TextLobPaths []string
BM25FieldIDs []int64
BM25StatsLogIDs []int64
WriteMergedBM25Stats bool
ReadVersion int64
SegmentBasePath string
PartitionBasePath string
CollectionID int64
PartitionID int64
Schema *schemapb.CollectionSchema
TextFieldIDs []int64
TextLobPaths []string
TextInlineThreshold int64
TextMaxLobFileBytes int64
TextFlushThresholdBytes int64
BM25FieldIDs []int64
BM25StatsLogIDs []int64
WriteMergedBM25Stats bool
ReadVersion int64
WriterFormat string
SchemaBasedPattern string
SchemaBasedFormats string
AllowedFieldIDs []int64
}
type GrowingFlushResult struct {
@@ -320,11 +331,12 @@ type GrowingSourceSyncTask struct {
schema *schemapb.CollectionSchema
source GrowingFlushSource
chunkManager storage.ChunkManager
allocator allocator.Interface
manifestPath string
flushedSize int64
bm25Stats map[int64]*storage.BM25Stats
chunkManager storage.ChunkManager
allocator allocator.Interface
manifestPath string
flushedSize int64
insertBinlogs map[int64]*datapb.FieldBinlog
bm25Stats map[int64]*storage.BM25Stats
committedManifestPath string
committedBM25Stats map[int64]*storage.BM25Stats
@@ -541,6 +553,10 @@ func (t *GrowingSourceSyncTask) Run(ctx context.Context) (err error) {
return merr.WrapErrServiceInternalMsg("growing source target offset is behind flushed rows, flushedRows=%d targetOffset=%d segmentID=%d",
segment.FlushedRows(), t.targetOffset, t.segmentID)
}
columnGroups, err := t.getColumnGroups(segment)
if err != nil {
return err
}
if t.committedManifestPath != "" {
t.manifestPath = t.committedManifestPath
t.bm25Stats = t.committedBM25Stats
@@ -553,7 +569,7 @@ func (t *GrowingSourceSyncTask) Run(ctx context.Context) (err error) {
if t.source.CurrentOffset() < t.targetOffset {
return merr.WrapErrServiceInternalMsg("growing flush source is behind target offset, current=%d target=%d", t.source.CurrentOffset(), t.targetOffset)
}
config, err := t.buildFlushConfig(segment)
config, err := t.buildFlushConfig(segment, columnGroups)
if err != nil {
return err
}
@@ -574,6 +590,9 @@ func (t *GrowingSourceSyncTask) Run(ctx context.Context) (err error) {
}
}
t.flushedSize = expectedRows
if expectedRows > 0 && len(columnGroups) > 0 {
t.insertBinlogs = buildV3ColumnGroupFieldBinlogs(columnGroups, 0, 0, 0, nil, nil, nil, nil, nil)
}
if t.metaWriter != nil {
if err := t.metaWriter.UpdateGrowingSourceSync(ctx, t); err != nil {
@@ -591,6 +610,9 @@ func (t *GrowingSourceSyncTask) Run(ctx context.Context) (err error) {
if len(t.bm25Stats) > 0 {
actions = append(actions, metacache.MergeBm25Stats(t.bm25Stats))
}
if len(columnGroups) > 0 {
actions = append(actions, metacache.UpdateCurrentSplit(columnGroups))
}
if t.IsFlush() {
actions = append(actions, metacache.UpdateState(commonpb.SegmentState_Flushed))
}
@@ -612,18 +634,45 @@ func (t *GrowingSourceSyncTask) Run(ctx context.Context) (err error) {
return nil
}
func (t *GrowingSourceSyncTask) buildFlushConfig(segment *metacache.SegmentInfo) (*GrowingFlushConfig, error) {
func (t *GrowingSourceSyncTask) getColumnGroups(segment *metacache.SegmentInfo) ([]storagecommon.ColumnGroup, error) {
return resolveColumnGroups(segment, t.schema, t.segmentID, func() map[int64]storagecommon.ColumnStats {
return map[int64]storagecommon.ColumnStats{}
}), nil
}
func (t *GrowingSourceSyncTask) schemaBasedPattern(columnGroups []storagecommon.ColumnGroup) (string, error) {
if len(columnGroups) == 0 {
return "", nil
}
arrowSchema, err := storage.ConvertToArrowSchema(t.schema, true)
if err != nil {
return "", merr.WrapErrServiceInternal(
fmt.Sprintf("can not convert collection schema %s to arrow schema: %s", t.schema.GetName(), err.Error()))
}
schemaBasedPattern, err := packed.SchemaBasedPattern(arrowSchema, columnGroups)
if err != nil {
return "", merr.WrapErrServiceInternal(
fmt.Sprintf("can not build schema based writer pattern %s", err.Error()))
}
return schemaBasedPattern, nil
}
func (t *GrowingSourceSyncTask) buildFlushConfig(segment *metacache.SegmentInfo, columnGroups []storagecommon.ColumnGroup) (*GrowingFlushConfig, error) {
segmentBasePath := path.Join(t.chunkManager.RootPath(), common.SegmentInsertLogPath,
metautil.JoinIDPath(t.collectionID, t.partitionID, t.segmentID))
partitionBasePath := path.Join(t.chunkManager.RootPath(), common.SegmentInsertLogPath,
metautil.JoinIDPath(t.collectionID, t.partitionID))
allowedFieldIDs, allowedFieldSet := allowedFieldsFromColumnGroups(columnGroups)
var textFieldIDs []int64
var textLobPaths []string
var bm25FieldIDs []int64
var bm25StatsLogIDs []int64
if t.schema != nil {
for _, field := range t.schema.GetFields() {
for _, field := range typeutil.GetAllFieldSchemas(t.schema) {
if !fieldAllowed(allowedFieldSet, field.GetFieldID()) {
continue
}
if field.GetDataType() == schemapb.DataType_Text {
fieldID := field.GetFieldID()
textFieldIDs = append(textFieldIDs, fieldID)
@@ -632,7 +681,10 @@ func (t *GrowingSourceSyncTask) buildFlushConfig(segment *metacache.SegmentInfo)
}
for _, function := range t.schema.GetFunctions() {
if function.GetType() == schemapb.FunctionType_BM25 && len(function.GetOutputFieldIds()) > 0 {
bm25FieldIDs = append(bm25FieldIDs, function.GetOutputFieldIds()[0])
outputFieldID := function.GetOutputFieldIds()[0]
if fieldAllowed(allowedFieldSet, outputFieldID) {
bm25FieldIDs = append(bm25FieldIDs, outputFieldID)
}
}
}
}
@@ -643,24 +695,94 @@ func (t *GrowingSourceSyncTask) buildFlushConfig(segment *metacache.SegmentInfo)
return nil, err
}
}
writerFormat := paramtable.Get().DataNodeCfg.StorageFormat.GetValue()
schemaBasedPattern, err := t.schemaBasedPattern(columnGroups)
if err != nil {
return nil, err
}
readVersion, err := growingSourceReadVersion(segment.ManifestPath(), columnGroups)
if err != nil {
return nil, err
}
schemaBasedFormats := strings.Join(storagecommon.ColumnGroupFormats(columnGroups, writerFormat), ",")
return &GrowingFlushConfig{
SegmentBasePath: segmentBasePath,
PartitionBasePath: partitionBasePath,
CollectionID: t.collectionID,
PartitionID: t.partitionID,
TextFieldIDs: textFieldIDs,
TextLobPaths: textLobPaths,
BM25FieldIDs: bm25FieldIDs,
BM25StatsLogIDs: bm25StatsLogIDs,
WriteMergedBM25Stats: t.IsFlush() && t.level != datapb.SegmentLevel_L0 && t.schema != nil && hasBM25Function(t.schema),
ReadVersion: manifestVersion(segment.ManifestPath()),
SegmentBasePath: segmentBasePath,
PartitionBasePath: partitionBasePath,
CollectionID: t.collectionID,
PartitionID: t.partitionID,
Schema: t.schema,
TextFieldIDs: textFieldIDs,
TextLobPaths: textLobPaths,
TextInlineThreshold: paramtable.Get().DataNodeCfg.TextInlineThreshold.GetAsInt64(),
TextMaxLobFileBytes: paramtable.Get().DataNodeCfg.TextMaxLobFileBytes.GetAsInt64(),
TextFlushThresholdBytes: paramtable.Get().DataNodeCfg.TextFlushThresholdBytes.GetAsInt64(),
BM25FieldIDs: bm25FieldIDs,
BM25StatsLogIDs: bm25StatsLogIDs,
WriteMergedBM25Stats: t.IsFlush() && t.level != datapb.SegmentLevel_L0 && t.schema != nil && hasBM25Function(t.schema),
ReadVersion: readVersion,
WriterFormat: writerFormat,
SchemaBasedPattern: schemaBasedPattern,
SchemaBasedFormats: schemaBasedFormats,
AllowedFieldIDs: allowedFieldIDs,
}, nil
}
func growingSourceReadVersion(manifestPath string, columnGroups []storagecommon.ColumnGroup) (int64, error) {
if manifestPath == "" {
return packed.ManifestEarliest, nil
}
_, version, err := packedManifestVersion(manifestPath)
if err != nil {
return 0, err
}
if version == packed.ManifestEarliest {
return version, nil
}
for _, columnGroup := range columnGroups {
if columnGroup.Format == "" {
return 0, merr.WrapErrDataIntegrityMsg("column group %d fields %v missing format for existing manifest %s",
columnGroup.GroupID, columnGroup.Fields, manifestPath)
}
}
return version, nil
}
func allowedFieldsFromColumnGroups(columnGroups []storagecommon.ColumnGroup) ([]int64, map[int64]struct{}) {
if len(columnGroups) == 0 {
return nil, nil
}
allowed := make(map[int64]struct{})
for _, group := range columnGroups {
for _, fieldID := range group.Fields {
allowed[fieldID] = struct{}{}
}
}
if len(allowed) == 0 {
return nil, nil
}
allowedFieldIDs := lo.Keys(allowed)
sort.Slice(allowedFieldIDs, func(i, j int) bool {
return allowedFieldIDs[i] < allowedFieldIDs[j]
})
return allowedFieldIDs, allowed
}
func fieldAllowed(allowed map[int64]struct{}, fieldID int64) bool {
if len(allowed) == 0 {
return true
}
_, ok := allowed[fieldID]
return ok
}
func (t *GrowingSourceSyncTask) allocBM25StatsLogIDs(count int) ([]int64, error) {
return t.allocLogIDs(count, "bm25 stats")
}
func (t *GrowingSourceSyncTask) allocLogIDs(count int, purpose string) ([]int64, error) {
if t.allocator == nil {
return nil, merr.WrapErrServiceInternal("id allocator is nil when allocating bm25 stats log ids")
return nil, merr.WrapErrServiceInternal(fmt.Sprintf("id allocator is nil when allocating %s ids", purpose))
}
ids := make([]int64, count)
for i := range ids {
@@ -18,6 +18,8 @@ import (
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
type fakeAllocator struct {
@@ -63,6 +65,14 @@ func (s *fakeCommitGrowingFlushSource) CommitGrowingFlush(targetOffset int64) {
}
func TestGrowingSourceSyncTaskBuildFlushConfigBM25(t *testing.T) {
paramtable.Get().Init(paramtable.NewBaseTable())
paramtable.Get().Save(paramtable.Get().DataNodeCfg.TextInlineThreshold.Key, "12345")
paramtable.Get().Save(paramtable.Get().DataNodeCfg.TextMaxLobFileBytes.Key, "67890")
paramtable.Get().Save(paramtable.Get().DataNodeCfg.TextFlushThresholdBytes.Key, "23456")
defer paramtable.Get().Reset(paramtable.Get().DataNodeCfg.TextInlineThreshold.Key)
defer paramtable.Get().Reset(paramtable.Get().DataNodeCfg.TextMaxLobFileBytes.Key)
defer paramtable.Get().Reset(paramtable.Get().DataNodeCfg.TextFlushThresholdBytes.Key)
segmentID := int64(1)
schema := &schemapb.CollectionSchema{
Name: "bm25",
@@ -97,10 +107,16 @@ func TestGrowingSourceSyncTaskBuildFlushConfigBM25(t *testing.T) {
WithAllocator(&fakeAllocator{next: 500}).
WithFlush()
config, err := task.buildFlushConfig(segment)
columnGroups, err := task.getColumnGroups(segment)
require.NoError(t, err)
config, err := task.buildFlushConfig(segment, columnGroups)
require.NoError(t, err)
require.Equal(t, schema, config.Schema)
require.Equal(t, []int64{101}, config.TextFieldIDs)
require.Equal(t, []string{"/root/insert_log/3/2/lobs/101"}, config.TextLobPaths)
require.EqualValues(t, 12345, config.TextInlineThreshold)
require.EqualValues(t, 67890, config.TextMaxLobFileBytes)
require.EqualValues(t, 23456, config.TextFlushThresholdBytes)
require.Equal(t, []int64{102}, config.BM25FieldIDs)
require.Equal(t, []int64{500}, config.BM25StatsLogIDs)
require.True(t, config.WriteMergedBM25Stats)
@@ -121,7 +137,9 @@ func TestGrowingSourceSyncTaskBuildFlushConfigStartsFromEarliestManifest(t *test
WithSegmentID(1).
WithChunkManager(cm)
config, err := task.buildFlushConfig(segment)
columnGroups, err := task.getColumnGroups(segment)
require.NoError(t, err)
config, err := task.buildFlushConfig(segment, columnGroups)
require.NoError(t, err)
require.EqualValues(t, packed.ManifestEarliest, config.ReadVersion)
}
@@ -144,7 +162,9 @@ func TestGrowingSourceSyncTaskBuildFlushConfigBM25AllocatorError(t *testing.T) {
WithChunkManager(cm).
WithAllocator(&fakeAllocator{err: errors.New("alloc failed")})
_, err := task.buildFlushConfig(segment)
columnGroups, err := task.getColumnGroups(segment)
require.NoError(t, err)
_, err = task.buildFlushConfig(segment, columnGroups)
require.ErrorContains(t, err, "alloc failed")
}
@@ -165,10 +185,232 @@ func TestGrowingSourceSyncTaskBuildFlushConfigBM25RequiresAllocator(t *testing.T
WithSchema(schema).
WithChunkManager(cm)
_, err := task.buildFlushConfig(segment)
columnGroups, err := task.getColumnGroups(segment)
require.NoError(t, err)
_, err = task.buildFlushConfig(segment, columnGroups)
require.ErrorContains(t, err, "id allocator is nil")
}
func TestGrowingSourceSyncTaskBuildFlushConfigUsesCurrentSplitPattern(t *testing.T) {
segmentID := int64(1)
schema := &schemapb.CollectionSchema{
Name: "text",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_Text},
{FieldID: 102, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}},
},
}
cm := mock_storage.NewMockChunkManager(t)
cm.EXPECT().RootPath().Return("/root").Maybe()
segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: segmentID,
PartitionID: 2,
StorageVersion: storage.StorageV3,
ManifestPath: `{"ver":7,"base_path":"/root/insert_log/3/2/1"}`,
Binlogs: []*datapb.FieldBinlog{
{FieldID: 0, ChildFields: []int64{100}, Format: "parquet"},
{FieldID: 101, ChildFields: []int64{101}, Format: "vortex"},
{FieldID: 102, ChildFields: []int64{102}, Format: "parquet"},
},
}, pkoracle.NewBloomFilterSet(), nil)
task := NewGrowingSourceSyncTask().
WithCollectionID(3).
WithPartitionID(2).
WithSegmentID(segmentID).
WithSchema(schema).
WithChunkManager(cm)
columnGroups, err := task.getColumnGroups(segment)
require.NoError(t, err)
config, err := task.buildFlushConfig(segment, columnGroups)
require.NoError(t, err)
require.Equal(t, "100,101,102", config.SchemaBasedPattern)
require.Equal(t, "parquet,vortex,parquet", config.SchemaBasedFormats)
require.Equal(t, []int64{100, 101, 102}, config.AllowedFieldIDs)
require.Equal(t, columnGroups, segment.GetCurrentSplit())
require.NotEmpty(t, config.WriterFormat)
}
func TestGrowingSourceSyncTaskBuildFlushConfigRequiresCurrentSplitFormatForExistingManifest(t *testing.T) {
segmentID := int64(1)
schema := &schemapb.CollectionSchema{
Name: "text",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_Text},
},
}
cm := mock_storage.NewMockChunkManager(t)
cm.EXPECT().RootPath().Return("/root").Maybe()
segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: segmentID,
PartitionID: 2,
StorageVersion: storage.StorageV3,
ManifestPath: packed.MarshalManifestPath("/root/insert_log/3/2/1", 7),
Binlogs: []*datapb.FieldBinlog{
{FieldID: 0, ChildFields: []int64{100}},
{FieldID: 101, ChildFields: []int64{101}, Format: "parquet"},
},
}, pkoracle.NewBloomFilterSet(), nil)
task := NewGrowingSourceSyncTask().
WithCollectionID(3).
WithPartitionID(2).
WithSegmentID(segmentID).
WithSchema(schema).
WithChunkManager(cm)
columnGroups, err := task.getColumnGroups(segment)
require.NoError(t, err)
_, err = task.buildFlushConfig(segment, columnGroups)
require.ErrorIs(t, err, merr.ErrDataIntegrity)
require.ErrorContains(t, err, "missing format")
}
func TestGrowingSourceSyncTaskBuildFlushConfigAllowsMissingFormatForEarliestManifest(t *testing.T) {
segmentID := int64(1)
schema := &schemapb.CollectionSchema{
Name: "text",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_Text},
},
}
cm := mock_storage.NewMockChunkManager(t)
cm.EXPECT().RootPath().Return("/root").Maybe()
segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: segmentID,
PartitionID: 2,
StorageVersion: storage.StorageV3,
ManifestPath: packed.MarshalManifestPath("/root/insert_log/3/2/1", packed.ManifestEarliest),
Binlogs: []*datapb.FieldBinlog{
{FieldID: 0, ChildFields: []int64{100}},
{FieldID: 101, ChildFields: []int64{101}},
},
}, pkoracle.NewBloomFilterSet(), nil)
task := NewGrowingSourceSyncTask().
WithCollectionID(3).
WithPartitionID(2).
WithSegmentID(segmentID).
WithSchema(schema).
WithChunkManager(cm)
columnGroups, err := task.getColumnGroups(segment)
require.NoError(t, err)
config, err := task.buildFlushConfig(segment, columnGroups)
require.NoError(t, err)
require.EqualValues(t, packed.ManifestEarliest, config.ReadVersion)
require.NotEmpty(t, config.SchemaBasedFormats)
}
func TestGrowingSourceSyncTaskBuildFlushConfigProjectsToCurrentSplit(t *testing.T) {
segmentID := int64(1)
schema := &schemapb.CollectionSchema{
Name: "text",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_Text},
{FieldID: 102, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}},
{FieldID: 103, Name: "added_scalar", DataType: schemapb.DataType_Int64, Nullable: true},
{FieldID: 104, Name: "added_text", DataType: schemapb.DataType_Text, Nullable: true},
{FieldID: 105, Name: "added_sparse", DataType: schemapb.DataType_SparseFloatVector},
},
Functions: []*schemapb.FunctionSchema{
{Type: schemapb.FunctionType_BM25, OutputFieldIds: []int64{105}},
},
}
cm := mock_storage.NewMockChunkManager(t)
cm.EXPECT().RootPath().Return("/root").Maybe()
segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: segmentID,
PartitionID: 2,
StorageVersion: storage.StorageV3,
ManifestPath: `{"ver":7,"base_path":"/root/insert_log/3/2/1"}`,
Binlogs: []*datapb.FieldBinlog{
{FieldID: 100, ChildFields: []int64{100}, Format: "parquet"},
{FieldID: 101, ChildFields: []int64{101}, Format: "vortex"},
{FieldID: 102, ChildFields: []int64{102}, Format: "parquet"},
},
}, pkoracle.NewBloomFilterSet(), nil)
task := NewGrowingSourceSyncTask().
WithCollectionID(3).
WithPartitionID(2).
WithSegmentID(segmentID).
WithSchema(schema).
WithChunkManager(cm).
WithAllocator(&fakeAllocator{next: 500}).
WithFlush()
columnGroups, err := task.getColumnGroups(segment)
require.NoError(t, err)
config, err := task.buildFlushConfig(segment, columnGroups)
require.NoError(t, err)
require.Equal(t, schema, config.Schema)
require.Equal(t, []int64{100, 101, 102}, config.AllowedFieldIDs)
require.Equal(t, []int64{101}, config.TextFieldIDs)
require.Equal(t, []string{"/root/insert_log/3/2/lobs/101"}, config.TextLobPaths)
require.Empty(t, config.BM25FieldIDs)
require.Empty(t, config.BM25StatsLogIDs)
require.True(t, config.WriteMergedBM25Stats)
}
func TestGrowingSourceSyncTaskBuildsColumnGroupBinlogs(t *testing.T) {
segmentID := int64(1)
schema := &schemapb.CollectionSchema{
Name: "text",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_Text},
{FieldID: 102, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: "dim", Value: "4"}}},
},
}
mc := metacache.NewMockMetaCache(t)
segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: segmentID,
PartitionID: 2,
State: commonpb.SegmentState_Growing,
StorageVersion: storage.StorageV3,
}, pkoracle.NewBloomFilterSet(), nil)
mc.EXPECT().GetSegmentByID(segmentID).Return(segment, true)
mc.EXPECT().UpdateSegments(mock.Anything, mock.Anything).Run(func(action metacache.SegmentAction, filters ...metacache.SegmentFilter) {
action(segment)
}).Return()
cm := mock_storage.NewMockChunkManager(t)
cm.EXPECT().RootPath().Return("/root").Maybe()
metaWriter := NewMockMetaWriter(t)
metaWriter.EXPECT().UpdateGrowingSourceSync(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, task *GrowingSourceSyncTask) error {
require.Len(t, task.insertBinlogs, 3)
require.Empty(t, task.insertBinlogs[0].GetBinlogs())
require.Equal(t, []int64{100}, task.insertBinlogs[0].GetChildFields())
return nil
})
task := NewGrowingSourceSyncTask().
WithCollectionID(3).
WithPartitionID(2).
WithSegmentID(segmentID).
WithChannelName("ch").
WithStartPosition(&msgpb.MsgPosition{Timestamp: 100}).
WithCheckpoint(&msgpb.MsgPosition{Timestamp: 200}).
WithBatchRows(10).
WithTargetOffset(10).
WithMetaCache(mc).
WithMetaWriter(metaWriter).
WithSchema(schema).
WithChunkManager(cm).
WithSource(&fakeCommitGrowingFlushSource{})
require.NoError(t, task.Run(context.Background()))
require.Len(t, segment.GetCurrentSplit(), 3)
require.Equal(t, []int64{100}, segment.GetCurrentSplit()[0].Fields)
}
func TestGrowingSourceSyncTaskCommitRetainedSourceOnlyOnFinalization(t *testing.T) {
run := func(t *testing.T, finalize func(*GrowingSourceSyncTask), expectCommit bool) {
segmentID := int64(1)
+36 -12
View File
@@ -168,6 +168,16 @@ func (b *brokerMetaWriter) UpdateGrowingSourceSync(ctx context.Context, task *Gr
return merr.WrapErrSegmentNotFound(task.segmentID)
}
insertFieldBinlogs := segment.Binlogs()
if len(task.insertBinlogs) > 0 {
// Growing-source V3 insert binlogs are column-group skeletons. They
// fully describe the current layout, so do not append them to older
// skeletons from previous incremental syncs.
insertFieldBinlogs = storage.SortFieldBinlogs(task.insertBinlogs)
}
statsFieldBinlogs := segment.Statslogs()
deltaFieldBinlogs := segment.Deltalogs()
bm25FieldBinlogs := segment.Bm25logs()
startPos := task.startPositions()
checkPoints := []*datapb.CheckPoint{{
SegmentID: task.segmentID,
@@ -181,6 +191,10 @@ func (b *brokerMetaWriter) UpdateGrowingSourceSync(ctx context.Context, task *Gr
mlog.Int64("ParitionID", task.partitionID),
mlog.Any("startPos", startPos),
mlog.Any("checkPoints", checkPoints),
mlog.Int("binlogNum", lo.SumBy(insertFieldBinlogs, func(fBinlog *datapb.FieldBinlog) int { return len(fBinlog.GetBinlogs()) })),
mlog.Int("statslogNum", lo.SumBy(statsFieldBinlogs, func(fBinlog *datapb.FieldBinlog) int { return len(fBinlog.GetBinlogs()) })),
mlog.Int("deltalogNum", lo.SumBy(deltaFieldBinlogs, func(fBinlog *datapb.FieldBinlog) int { return len(fBinlog.GetBinlogs()) })),
mlog.Int("bm25logNum", lo.SumBy(bm25FieldBinlogs, func(fBinlog *datapb.FieldBinlog) int { return len(fBinlog.GetBinlogs()) })),
mlog.String("manifestPath", task.manifestPath),
mlog.String("vChannelName", task.channelName),
)
@@ -191,18 +205,22 @@ func (b *brokerMetaWriter) UpdateGrowingSourceSync(ctx context.Context, task *Gr
commonpbutil.WithMsgID(0),
commonpbutil.WithSourceID(b.serverID),
),
SegmentID: task.segmentID,
CollectionID: task.collectionID,
PartitionID: task.partitionID,
CheckPoints: checkPoints,
StartPositions: startPos,
Flushed: task.IsFlush(),
Dropped: task.IsDrop(),
Channel: task.channelName,
SegLevel: task.level,
StorageVersion: storage.StorageV3,
WithFullBinlogs: false,
ManifestPath: task.manifestPath,
SegmentID: task.segmentID,
CollectionID: task.collectionID,
PartitionID: task.partitionID,
Field2BinlogPaths: insertFieldBinlogs,
Field2StatslogPaths: statsFieldBinlogs,
Field2Bm25LogPaths: bm25FieldBinlogs,
Deltalogs: deltaFieldBinlogs,
CheckPoints: checkPoints,
StartPositions: startPos,
Flushed: task.IsFlush(),
Dropped: task.IsDrop(),
Channel: task.channelName,
SegLevel: task.level,
StorageVersion: storage.StorageV3,
WithFullBinlogs: true,
ManifestPath: task.manifestPath,
}
err := retry.Handle(ctx, func() (bool, error) {
@@ -229,6 +247,12 @@ func (b *brokerMetaWriter) UpdateGrowingSourceSync(ctx context.Context, task *Gr
task.metacache.UpdateSegments(metacache.SetStartPosRecorded(true), metacache.WithSegmentIDs(lo.Map(startPos, func(pos *datapb.SegmentStartPosition, _ int) int64 {
return pos.GetSegmentID()
})...))
task.metacache.UpdateSegments(metacache.MergeSegmentAction(
metacache.UpdateBinlogs(insertFieldBinlogs),
metacache.UpdateStatslogs(statsFieldBinlogs),
metacache.UpdateDeltalogs(deltaFieldBinlogs),
metacache.UpdateBm25logs(bm25FieldBinlogs),
), metacache.WithSegmentIDs(task.segmentID))
return nil
}
@@ -12,6 +12,8 @@ import (
"github.com/milvus-io/milvus/internal/flushcommon/broker"
"github.com/milvus-io/milvus/internal/flushcommon/metacache"
"github.com/milvus-io/milvus/internal/flushcommon/metacache/pkoracle"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/storagecommon"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
@@ -102,6 +104,151 @@ func (s *MetaWriterSuite) TestReturnError() {
s.Error(err)
}
func (s *MetaWriterSuite) TestGrowingSourceSyncPersistsColumnGroupBinlogs() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
seg := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: 1,
CollectionID: 3,
PartitionID: 2,
StorageVersion: storage.StorageV3,
Statslogs: []*datapb.FieldBinlog{{
FieldID: 100,
Binlogs: []*datapb.Binlog{{
LogID: 31,
LogPath: "stats/100/31",
}},
}},
Deltalogs: []*datapb.FieldBinlog{{
FieldID: 0,
Binlogs: []*datapb.Binlog{{
LogID: 41,
LogPath: "delta/41",
}},
}},
Bm25Statslogs: []*datapb.FieldBinlog{{
FieldID: 102,
Binlogs: []*datapb.Binlog{{
LogID: 51,
LogPath: "bm25/102/51",
}},
}},
}, pkoracle.NewBloomFilterSet(), nil)
metacache.UpdateNumOfRows(5)(seg)
s.metacache.EXPECT().GetSegmentByID(int64(1)).Return(seg, true)
s.metacache.EXPECT().GetSegmentsBy(mock.Anything, mock.Anything, mock.Anything).Return(nil)
s.metacache.EXPECT().UpdateSegments(mock.Anything, mock.Anything).Run(func(action metacache.SegmentAction, filters ...metacache.SegmentFilter) {
action(seg)
}).Return()
columnGroups := []storagecommon.ColumnGroup{
{GroupID: 0, Fields: []int64{100}, Format: "parquet"},
{GroupID: 101, Fields: []int64{101}, Format: "parquet"},
}
task := NewGrowingSourceSyncTask().
WithCollectionID(3).
WithPartitionID(2).
WithSegmentID(1).
WithChannelName("ch").
WithStartPosition(&msgpb.MsgPosition{Timestamp: 100}).
WithCheckpoint(&msgpb.MsgPosition{Timestamp: 200}).
WithBatchRows(10).
WithTargetOffset(15).
WithMetaCache(s.metacache)
task.manifestPath = "manifest"
task.insertBinlogs = buildV3ColumnGroupFieldBinlogs(columnGroups, 0, 0, 0, nil, nil, nil, nil, nil)
s.broker.EXPECT().SaveBinlogPaths(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, req *datapb.SaveBinlogPathsRequest) error {
s.True(req.GetWithFullBinlogs())
s.EqualValues(storage.StorageV3, req.GetStorageVersion())
s.Equal("manifest", req.GetManifestPath())
s.Len(req.GetField2BinlogPaths(), 2)
s.EqualValues(0, req.GetField2BinlogPaths()[0].GetFieldID())
s.EqualValues([]int64{100}, req.GetField2BinlogPaths()[0].GetChildFields())
s.Equal("parquet", req.GetField2BinlogPaths()[0].GetFormat())
s.Empty(req.GetField2BinlogPaths()[0].GetBinlogs())
s.EqualValues(101, req.GetField2BinlogPaths()[1].GetFieldID())
s.EqualValues([]int64{101}, req.GetField2BinlogPaths()[1].GetChildFields())
s.Len(req.GetField2StatslogPaths(), 1)
s.EqualValues(100, req.GetField2StatslogPaths()[0].GetFieldID())
s.Equal("stats/100/31", req.GetField2StatslogPaths()[0].GetBinlogs()[0].GetLogPath())
s.Len(req.GetDeltalogs(), 1)
s.Equal("delta/41", req.GetDeltalogs()[0].GetBinlogs()[0].GetLogPath())
s.Len(req.GetField2Bm25LogPaths(), 1)
s.EqualValues(102, req.GetField2Bm25LogPaths()[0].GetFieldID())
s.Equal("bm25/102/51", req.GetField2Bm25LogPaths()[0].GetBinlogs()[0].GetLogPath())
return nil
})
err := s.writer.UpdateGrowingSourceSync(ctx, task)
s.NoError(err)
s.Len(seg.Binlogs(), 2)
s.EqualValues([]int64{100}, seg.Binlogs()[0].GetChildFields())
s.Len(seg.Statslogs(), 1)
s.Len(seg.Deltalogs(), 1)
s.Len(seg.Bm25logs(), 1)
}
func (s *MetaWriterSuite) TestGrowingSourceSyncDoesNotAccumulateColumnGroupSkeletons() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
columnGroups := []storagecommon.ColumnGroup{
{GroupID: 0, Fields: []int64{100}, Format: "parquet"},
{GroupID: 101, Fields: []int64{101}, Format: "vortex"},
}
existingSkeleton := storage.SortFieldBinlogs(buildV3ColumnGroupFieldBinlogs(columnGroups, 0, 0, 0, nil, nil, nil, nil, nil))
seg := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: 1,
CollectionID: 3,
PartitionID: 2,
StorageVersion: storage.StorageV3,
Binlogs: existingSkeleton,
}, pkoracle.NewBloomFilterSet(), nil)
metacache.UpdateNumOfRows(5)(seg)
s.metacache.EXPECT().GetSegmentByID(int64(1)).Return(seg, true)
s.metacache.EXPECT().GetSegmentsBy(mock.Anything, mock.Anything, mock.Anything).Return(nil)
s.metacache.EXPECT().UpdateSegments(mock.Anything, mock.Anything).Run(func(action metacache.SegmentAction, filters ...metacache.SegmentFilter) {
action(seg)
}).Return()
task := NewGrowingSourceSyncTask().
WithCollectionID(3).
WithPartitionID(2).
WithSegmentID(1).
WithChannelName("ch").
WithStartPosition(&msgpb.MsgPosition{Timestamp: 100}).
WithCheckpoint(&msgpb.MsgPosition{Timestamp: 200}).
WithBatchRows(10).
WithTargetOffset(15).
WithMetaCache(s.metacache)
task.manifestPath = "manifest"
task.insertBinlogs = buildV3ColumnGroupFieldBinlogs(columnGroups, 0, 0, 0, nil, nil, nil, nil, nil)
s.broker.EXPECT().SaveBinlogPaths(mock.Anything, mock.Anything).RunAndReturn(
func(ctx context.Context, req *datapb.SaveBinlogPathsRequest) error {
s.True(req.GetWithFullBinlogs())
s.Equal("manifest", req.GetManifestPath())
s.Len(req.GetField2BinlogPaths(), 2)
s.EqualValues(0, req.GetField2BinlogPaths()[0].GetFieldID())
s.EqualValues([]int64{100}, req.GetField2BinlogPaths()[0].GetChildFields())
s.Equal("parquet", req.GetField2BinlogPaths()[0].GetFormat())
s.EqualValues(101, req.GetField2BinlogPaths()[1].GetFieldID())
s.EqualValues([]int64{101}, req.GetField2BinlogPaths()[1].GetChildFields())
s.Equal("vortex", req.GetField2BinlogPaths()[1].GetFormat())
return nil
})
err := s.writer.UpdateGrowingSourceSync(ctx, task)
s.NoError(err)
s.Len(seg.Binlogs(), 2)
s.EqualValues(0, seg.Binlogs()[0].GetFieldID())
s.EqualValues(101, seg.Binlogs()[1].GetFieldID())
}
func (s *MetaWriterSuite) TestGrowingSourceSyncMetaErrorsReturnError() {
testCases := []struct {
name string
+53 -15
View File
@@ -291,27 +291,65 @@ func (bw *BulkPackWriterV3) writeInserts(ctx context.Context, pack *SyncPack, ba
return result
}
logs = make(map[int64]*datapb.FieldBinlog)
for _, columnGroup := range bw.columnGroups {
logs = buildV3ColumnGroupFieldBinlogs(
bw.columnGroups,
w.GetWrittenRowNum(),
tsFrom,
tsTo,
func(columnGroupID int64) int64 { return int64(w.GetColumnGroupWrittenCompressed(columnGroupID)) },
func(columnGroupID int64) int64 { return int64(w.GetColumnGroupWrittenUncompressed(columnGroupID)) },
nil,
w.GetWrittenPaths,
getFieldNullCounts,
)
return logs, files, nil
}
func buildV3ColumnGroupFieldBinlogs(
columnGroups []storagecommon.ColumnGroup,
entriesNum int64,
tsFrom uint64,
tsTo uint64,
compressedSize func(columnGroupID int64) int64,
memorySize func(columnGroupID int64) int64,
logID func(columnGroupID int64) int64,
logPath func(columnGroupID int64) string,
fieldNullCounts func(columnGroup storagecommon.ColumnGroup) map[int64]int64,
) map[int64]*datapb.FieldBinlog {
logs := make(map[int64]*datapb.FieldBinlog, len(columnGroups))
for _, columnGroup := range columnGroups {
columnGroupID := columnGroup.GroupID
logs[columnGroupID] = &datapb.FieldBinlog{
fieldBinlog := &datapb.FieldBinlog{
FieldID: columnGroupID,
ChildFields: columnGroup.Fields,
Format: columnGroup.Format,
Binlogs: []*datapb.Binlog{
{
LogSize: int64(w.GetColumnGroupWrittenCompressed(columnGroup.GroupID)),
MemorySize: int64(w.GetColumnGroupWrittenUncompressed(columnGroup.GroupID)),
LogPath: w.GetWrittenPaths(columnGroupID),
EntriesNum: w.GetWrittenRowNum(),
TimestampFrom: tsFrom,
TimestampTo: tsTo,
FieldNullCounts: getFieldNullCounts(columnGroup),
},
},
}
if entriesNum > 0 {
binlog := &datapb.Binlog{
EntriesNum: entriesNum,
TimestampFrom: tsFrom,
TimestampTo: tsTo,
}
if compressedSize != nil {
binlog.LogSize = compressedSize(columnGroupID)
}
if memorySize != nil {
binlog.MemorySize = memorySize(columnGroupID)
}
if logID != nil {
binlog.LogID = logID(columnGroupID)
}
if logPath != nil {
binlog.LogPath = logPath(columnGroupID)
}
if fieldNullCounts != nil {
binlog.FieldNullCounts = fieldNullCounts(columnGroup)
}
fieldBinlog.Binlogs = []*datapb.Binlog{binlog}
}
logs[columnGroupID] = fieldBinlog
}
return logs, files, nil
return logs
}
func (bw *BulkPackWriterV3) resolveInsertWriterFormats() (string, []string, error) {
+14 -6
View File
@@ -213,17 +213,21 @@ func (t *SyncTask) Run(ctx context.Context) (err error) {
}
func (t *SyncTask) getColumnGroups(segmentInfo *metacache.SegmentInfo) []storagecommon.ColumnGroup {
// column group only needed for storage v2 segment
return resolveColumnGroups(segmentInfo, t.schema, t.segmentID, t.calcColumnStats)
}
func resolveColumnGroups(segmentInfo *metacache.SegmentInfo, schema *schemapb.CollectionSchema, segmentID int64, calcColumnStats func() map[int64]storagecommon.ColumnStats) []storagecommon.ColumnGroup {
// column group only needed for storage v2/v3 segments
if segmentInfo.GetStorageVersion() != storage.StorageV2 && segmentInfo.GetStorageVersion() != storage.StorageV3 {
return nil
}
// empty pack
if len(t.pack.insertData) == 0 && t.schema == nil {
if schema == nil {
return nil
}
allFields := typeutil.GetAllFieldSchemas(t.schema)
allFields := typeutil.GetAllFieldSchemas(schema)
// use previous split if already exists
if currentSplit := segmentInfo.GetCurrentSplit(); currentSplit != nil {
@@ -232,7 +236,7 @@ func (t *SyncTask) getColumnGroups(segmentInfo *metacache.SegmentInfo) []storage
if len(cg.Fields) == 0 {
result := storagecommon.SplitColumns(allFields, map[int64]storagecommon.ColumnStats{}, storagecommon.NewSelectedDataTypePolicy(), storagecommon.NewRemanentShortPolicy(-1))
result = storagecommon.FillColumnGroupFormats(result, paramtable.Get().DataNodeCfg.StorageFormat.GetValue())
mlog.Info(context.TODO(), "use legacy split policy", mlog.FieldSegmentID(t.segmentID), mlog.Stringers("columnGroups", result))
mlog.Info(context.TODO(), "use legacy split policy", mlog.FieldSegmentID(segmentID), mlog.Stringers("columnGroups", result))
return result
}
}
@@ -253,9 +257,13 @@ func (t *SyncTask) getColumnGroups(segmentInfo *metacache.SegmentInfo) []storage
}
policies := storagecommon.DefaultPolicies()
result := storagecommon.SplitColumns(allFields, t.calcColumnStats(), policies...)
stats := map[int64]storagecommon.ColumnStats{}
if calcColumnStats != nil {
stats = calcColumnStats()
}
result := storagecommon.SplitColumns(allFields, stats, policies...)
result = storagecommon.FillColumnGroupFormats(result, paramtable.Get().DataNodeCfg.StorageFormat.GetValue())
mlog.Info(context.TODO(), "sync new split columns", mlog.FieldSegmentID(t.segmentID), mlog.Stringers("columnGroups", result))
mlog.Info(context.TODO(), "sync new split columns", mlog.FieldSegmentID(segmentID), mlog.Stringers("columnGroups", result))
return result
}
@@ -20,6 +20,7 @@ import (
"github.com/milvus-io/milvus/internal/flushcommon/metacache/pkoracle"
"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
@@ -838,13 +839,12 @@ func (s *L0WriteBufferSuite) TestBufferDataGrowingSourceMode() {
return &syncmgr.GrowingFlushResult{ManifestPath: "manifest-retry", NumRows: 10}, nil
},
}
errorHandlerCalled := false
var errorHandlerCalls atomic.Int32
wb, err := NewL0WriteBuffer(s.channelName, metacache, s.syncMgr, &writeBufferOption{
idAllocator: s.allocator,
growingSourceRetryInterval: time.Hour,
errorHandler: func(error) {
errorHandlerCalled = true
panic("growing source task should not call writebuffer error handler directly")
errorHandlerCalls.Add(1)
},
growingSourceResolver: func(segmentID int64, targetOffset int64, _ *msgpb.MsgPosition) (syncmgr.GrowingFlushSource, syncmgr.GrowingSourceState) {
return source, syncmgr.GrowingSourceUsable
@@ -862,7 +862,6 @@ func (s *L0WriteBufferSuite) TestBufferDataGrowingSourceMode() {
done <- textTask
}()
err := textTask.Run(ctx)
textTask.HandleError(err)
for _, callback := range callbacks {
if cbErr := callback(err); cbErr != nil {
return struct{}{}, cbErr
@@ -883,7 +882,7 @@ func (s *L0WriteBufferSuite) TestBufferDataGrowingSourceMode() {
s.ErrorContains(conc.AwaitAll(futures...), "mock growing source flush error")
firstTask := <-done
s.False(errorHandlerCalled)
s.EqualValues(1, errorHandlerCalls.Load())
progress, ok := l0wb.growingSourceProgress[int64(1010)]
s.True(ok)
s.EqualValues(1, progress.failureCount)
@@ -899,6 +898,7 @@ func (s *L0WriteBufferSuite) TestBufferDataGrowingSourceMode() {
s.Require().Len(futures, 1)
s.NoError(conc.AwaitAll(futures...))
secondTask := <-done
s.EqualValues(1, errorHandlerCalls.Load())
s.EqualValues(10, secondTask.BatchRows())
segment, ok = metacache.GetSegmentByID(1010)
s.True(ok)
@@ -907,6 +907,74 @@ func (s *L0WriteBufferSuite) TestBufferDataGrowingSourceMode() {
s.Equal("manifest-retry", segment.ManifestPath())
})
s.Run("source_task_layout_mismatch_is_non_retryable", func() {
textSchema := s.textSchema()
metacache := s.newTextRealMetaCache(textSchema)
source := fakeGrowingFlushSource{
flushFunc: func(context.Context, int64, int64, *syncmgr.GrowingFlushConfig) (*syncmgr.GrowingFlushResult, error) {
return nil, fmt.Errorf("Invalid: Column group size mismatch: existing has 10 groups, but appended has 1 groups: segcore error[segcoreCode=2001]")
},
}
var errorHandlerCalls atomic.Int32
wb, err := NewL0WriteBuffer(s.channelName, metacache, s.syncMgr, &writeBufferOption{
idAllocator: s.allocator,
growingSourceRetryInterval: time.Hour,
errorHandler: func(error) {
errorHandlerCalls.Add(1)
},
growingSourceResolver: func(segmentID int64, targetOffset int64, _ *msgpb.MsgPosition) (syncmgr.GrowingFlushSource, syncmgr.GrowingSourceState) {
return source, syncmgr.GrowingSourceUsable
},
})
s.NoError(err)
done := make(chan struct{})
s.syncMgr.EXPECT().SyncData(mock.Anything, mock.AnythingOfType("*syncmgr.GrowingSourceSyncTask"), mock.Anything).
RunAndReturn(func(ctx context.Context, task syncmgr.Task, callbacks ...func(error) error) (*conc.Future[struct{}], error) {
textTask := task.(*syncmgr.GrowingSourceSyncTask)
textTask.WithChunkManager(storage.NewLocalChunkManager(objectstorage.RootPath(s.T().TempDir())))
return conc.Go(func() (struct{}, error) {
defer close(done)
err := textTask.Run(ctx)
for _, callback := range callbacks {
if cbErr := callback(err); cbErr != nil {
return struct{}{}, cbErr
}
}
return struct{}{}, err
}), nil
}).Once()
_, msg := s.composeTextInsertMsg(1014, 10)
insertData, err := PrepareInsert(textSchema, s.pkSchema, []*msgstream.InsertMsg{msg})
s.NoError(err)
err = wb.BufferData(insertData, nil, &msgpb.MsgPosition{Timestamp: 100}, &msgpb.MsgPosition{Timestamp: 200}, 100)
s.NoError(err)
l0wb := wb.(*l0WriteBuffer)
futures := l0wb.syncSegments(context.Background(), []int64{1014})
s.Require().Len(futures, 1)
s.ErrorContains(conc.AwaitAll(futures...), "Column group size mismatch")
<-done
s.EqualValues(1, errorHandlerCalls.Load())
progress, ok := l0wb.growingSourceProgress[int64(1014)]
s.True(ok)
s.EqualValues(1, progress.failureCount)
s.Contains(progress.lastFailure, "Column group size mismatch")
s.True(progress.nonRetryableFailure)
s.False(l0wb.growingSourceRetryScheduled)
segments, retry := l0wb.getGrowingSourceSegmentsToRetry()
s.Empty(segments)
s.False(retry)
segment, ok := metacache.GetSegmentByID(1014)
s.True(ok)
s.EqualValues(0, segment.FlushedRows())
s.EqualValues(0, segment.SyncingRows())
s.EqualValues(10, segment.BufferRows())
})
s.Run("source_task_row_count_mismatch_retries_without_meta_update", func() {
textSchema := s.textSchema()
metacache := s.newTextRealMetaCache(textSchema)
@@ -916,14 +984,13 @@ func (s *L0WriteBufferSuite) TestBufferDataGrowingSourceMode() {
},
}
metaWriter := syncmgr.NewMockMetaWriter(s.T())
errorHandlerCalled := false
var errorHandlerCalls atomic.Int32
wb, err := NewL0WriteBuffer(s.channelName, metacache, s.syncMgr, &writeBufferOption{
idAllocator: s.allocator,
metaWriter: metaWriter,
growingSourceRetryInterval: time.Hour,
errorHandler: func(error) {
errorHandlerCalled = true
panic("growing source row count mismatch should not call writebuffer error handler directly")
errorHandlerCalls.Add(1)
},
growingSourceResolver: func(segmentID int64, targetOffset int64, _ *msgpb.MsgPosition) (syncmgr.GrowingFlushSource, syncmgr.GrowingSourceState) {
return source, syncmgr.GrowingSourceUsable
@@ -959,7 +1026,7 @@ func (s *L0WriteBufferSuite) TestBufferDataGrowingSourceMode() {
s.ErrorContains(conc.AwaitAll(futures...), "row count mismatch")
<-done
s.False(errorHandlerCalled)
s.EqualValues(1, errorHandlerCalls.Load())
progress, ok := l0wb.growingSourceProgress[int64(1011)]
s.True(ok)
s.EqualValues(1, progress.failureCount)
@@ -976,9 +1043,9 @@ func (s *L0WriteBufferSuite) TestBufferDataGrowingSourceMode() {
var ranges [][2]int64
source := fakeGrowingFlushSource{
currentOffset: 15,
flushFunc: func(_ context.Context, startOffset, endOffset int64, _ *syncmgr.GrowingFlushConfig) (*syncmgr.GrowingFlushResult, error) {
flushFunc: func(_ context.Context, startOffset, endOffset int64, config *syncmgr.GrowingFlushConfig) (*syncmgr.GrowingFlushResult, error) {
ranges = append(ranges, [2]int64{startOffset, endOffset})
return &syncmgr.GrowingFlushResult{ManifestPath: fmt.Sprintf("manifest-%d", endOffset), NumRows: endOffset - startOffset}, nil
return &syncmgr.GrowingFlushResult{ManifestPath: packed.MarshalManifestPath(config.SegmentBasePath, endOffset), NumRows: endOffset - startOffset}, nil
},
}
wb, err := NewL0WriteBuffer(s.channelName, metacache, s.syncMgr, &writeBufferOption{
@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"path"
"strings"
"sync"
"time"
@@ -100,15 +101,16 @@ type checkpointCandidates struct {
}
type growingSourceProgress struct {
segmentID int64
targetOffset int64
syncingOffset int64
syncing bool
pendingFlush bool
pendingCommitted *growingSourcePendingCommittedFlush
batches []growingSourceProgressBatch
failureCount int64
lastFailure string
segmentID int64
targetOffset int64
syncingOffset int64
syncing bool
pendingFlush bool
pendingCommitted *growingSourcePendingCommittedFlush
nonRetryableFailure bool
batches []growingSourceProgressBatch
failureCount int64
lastFailure string
}
type growingSourcePendingCommittedFlush struct {
@@ -177,6 +179,19 @@ func (p *growingSourceProgress) failSync(err error) {
}
}
func (p *growingSourceProgress) markNonRetryableFailure() {
p.nonRetryableFailure = true
}
func isGrowingSourceLayoutMismatch(err error) bool {
if err == nil {
return false
}
msg := err.Error()
return strings.Contains(msg, "Column count mismatch") ||
strings.Contains(msg, "Column group size mismatch")
}
func cloneBM25StatsMap(stats map[int64]*storage.BM25Stats) map[int64]*storage.BM25Stats {
if len(stats) == 0 {
return nil
@@ -656,6 +671,9 @@ func (wb *writeBufferBase) warnGrowingSourceFallback(segmentID int64, targetOffs
}
func (wb *writeBufferBase) growingSourceProgressSyncable(segmentID int64, progress *growingSourceProgress, rollbackFlushing bool, markSealedFlushing bool) (bool, bool) {
if progress.nonRetryableFailure {
return false, false
}
if progress.syncing {
if segment, ok := wb.metaCache.GetSegmentByID(segmentID); ok &&
(segment.State() == commonpb.SegmentState_Sealed || segment.State() == commonpb.SegmentState_Flushing) {
@@ -927,7 +945,15 @@ func (wb *writeBufferBase) submitSyncTasks(ctx context.Context, syncTasks []sync
progress.failSync(err)
wb.rollbackGrowingSourceSyncTaskLocked(growingSourceTask)
wb.observeGrowingSourceSyncFailureLocked(growingSourceTask.SegmentID(), progress)
wb.scheduleGrowingSourceRetryLocked()
if isGrowingSourceLayoutMismatch(err) {
progress.markNonRetryableFailure()
mlog.Error(ctx, "growing-source source sync failed with non-retryable layout mismatch",
mlog.Int64("segmentID", growingSourceTask.SegmentID()),
mlog.Int64("targetOffset", progress.targetOffset),
mlog.String("lastFailure", progress.lastFailure))
} else {
wb.scheduleGrowingSourceRetryLocked()
}
} else {
if growingSourceTask.IsFlush() {
progress.pendingFlush = false
@@ -1020,6 +1046,9 @@ func (wb *writeBufferBase) growingSourceProgressSelectedByPolicy(ts typeutil.Tim
if progress == nil {
return false
}
if progress.nonRetryableFailure {
return false
}
if progress.pendingFlush {
return true
}
@@ -1442,6 +1471,7 @@ func (wb *writeBufferBase) getGrowingSourceSyncTask(ctx context.Context, segment
WithMetaWriter(wb.metaWriter).
WithSchema(wb.metaCache.GetSchema(schemaTimestamp)).
WithAllocator(wb.allocator).
WithFailureCallback(wb.errHandler).
// Same as above: keep the critical write path retrying despite the
// retry.Do InputError short-circuit.
WithWriteRetryOptions(retry.AttemptAlways(), retry.MaxSleepTime(10*time.Second),
@@ -6,6 +6,7 @@ import (
"testing"
"time"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
@@ -698,6 +699,15 @@ func (s *WriteBufferSuite) TestGrowingSourceProgressSelectedByPolicy() {
s.True(selected)
})
s.Run("non_retryable_failure", func() {
selected := s.wb.growingSourceProgressSelectedByPolicy(recentTs, 1007, &growingSourceProgress{
segmentID: 1007,
pendingFlush: true,
nonRetryableFailure: true,
})
s.False(selected)
})
s.Run("sealed_segment", func() {
segment := metacache.NewSegmentInfo(&datapb.SegmentInfo{
ID: 1002,
@@ -765,6 +775,39 @@ func (s *WriteBufferSuite) TestGrowingSourceProgressSelectedByPolicy() {
})
}
func (s *WriteBufferSuite) TestGrowingSourceLayoutMismatch() {
s.True(isGrowingSourceLayoutMismatch(errors.New("flush growing source data: Invalid: Column count mismatch at index 0: existing has 21 columns, but appended has 44 columns: segcore error[segcoreCode=2001]")))
s.True(isGrowingSourceLayoutMismatch(errors.New("flush growing source data: Invalid: Column group size mismatch: existing has 10 groups, but appended has 1 groups: segcore error[segcoreCode=2001]")))
s.False(isGrowingSourceLayoutMismatch(errors.New("flush growing source data: mock transient error")))
s.False(isGrowingSourceLayoutMismatch(nil))
}
func (s *WriteBufferSuite) TestGrowingSourceProgressSyncableSkipsNonRetryableFailure() {
syncable, retry := s.wb.growingSourceProgressSyncable(1001, &growingSourceProgress{
segmentID: 1001,
nonRetryableFailure: true,
batches: []growingSourceProgressBatch{
{endPosition: &msgpb.MsgPosition{Timestamp: 100}, endOffset: 10},
},
}, false, true)
s.False(syncable)
s.False(retry)
}
func (s *WriteBufferSuite) TestGrowingSourceProgressRetrySkipsNonRetryableFailure() {
s.wb.growingSourceProgress[1001] = &growingSourceProgress{
segmentID: 1001,
nonRetryableFailure: true,
batches: []growingSourceProgressBatch{
{endPosition: &msgpb.MsgPosition{Timestamp: 100}, endOffset: 10},
},
}
segments, retry := s.wb.getGrowingSourceSegmentsToRetry()
s.Empty(segments)
s.False(retry)
}
func (s *WriteBufferSuite) TestDropPartitions() {
wb, err := newWriteBufferBase(s.channelName, s.metacache, s.syncMgr, &writeBufferOption{
pkStatsFactory: func(vchannel *datapb.SegmentInfo) pkoracle.PkStat {
@@ -545,16 +545,24 @@ func (s *delegatorGrowingFlushSource) CurrentOffset() int64 {
func (s *delegatorGrowingFlushSource) FlushGrowingData(ctx context.Context, startOffset, endOffset int64, config *syncmgr.GrowingFlushConfig) (*syncmgr.GrowingFlushResult, error) {
result, err := s.segment.FlushData(ctx, startOffset, endOffset, &segments.FlushConfig{
SegmentBasePath: config.SegmentBasePath,
PartitionBasePath: config.PartitionBasePath,
CollectionID: config.CollectionID,
PartitionID: config.PartitionID,
TextFieldIDs: config.TextFieldIDs,
TextLobPaths: config.TextLobPaths,
BM25FieldIDs: config.BM25FieldIDs,
BM25StatsLogIDs: config.BM25StatsLogIDs,
WriteMergedBM25Stats: config.WriteMergedBM25Stats,
ReadVersion: config.ReadVersion,
SegmentBasePath: config.SegmentBasePath,
PartitionBasePath: config.PartitionBasePath,
CollectionID: config.CollectionID,
PartitionID: config.PartitionID,
Schema: config.Schema,
TextFieldIDs: config.TextFieldIDs,
TextLobPaths: config.TextLobPaths,
TextInlineThreshold: config.TextInlineThreshold,
TextMaxLobFileBytes: config.TextMaxLobFileBytes,
TextFlushThresholdBytes: config.TextFlushThresholdBytes,
BM25FieldIDs: config.BM25FieldIDs,
BM25StatsLogIDs: config.BM25StatsLogIDs,
WriteMergedBM25Stats: config.WriteMergedBM25Stats,
ReadVersion: config.ReadVersion,
WriterFormat: config.WriterFormat,
SchemaBasedPattern: config.SchemaBasedPattern,
SchemaBasedFormats: config.SchemaBasedFormats,
AllowedFieldIDs: config.AllowedFieldIDs,
})
if err != nil || result == nil {
return nil, err
@@ -21,13 +21,37 @@ import (
"testing"
"time"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/flushcommon/syncmgr"
"github.com/milvus-io/milvus/internal/querynodev2/segments"
)
func TestDelegatorGrowingFlushSourcePassesTaskSchema(t *testing.T) {
ctx := context.Background()
schema := &schemapb.CollectionSchema{Name: "task-schema"}
segment := segments.NewMockSegment(t)
segment.EXPECT().
FlushData(ctx, int64(3), int64(7), mock.Anything).
RunAndReturn(func(_ context.Context, startOffset int64, endOffset int64, config *segments.FlushConfig) (*segments.FlushResult, error) {
require.EqualValues(t, 3, startOffset)
require.EqualValues(t, 7, endOffset)
require.True(t, config.Schema == schema)
return &segments.FlushResult{ManifestPath: "manifest", NumRows: 4}, nil
})
source := &delegatorGrowingFlushSource{segment: segment}
result, err := source.FlushGrowingData(ctx, 3, 7, &syncmgr.GrowingFlushConfig{
Schema: schema,
})
require.NoError(t, err)
require.Equal(t, "manifest", result.ManifestPath)
require.EqualValues(t, 4, result.NumRows)
}
func TestDelegatorGrowingSourceProviderCloseWaitsForSourceRelease(t *testing.T) {
segmentManager := segments.NewMockSegmentManager(t)
segment := segments.NewMockSegment(t)
+44
View File
@@ -1656,12 +1656,30 @@ func (s *LocalSegment) FlushData(ctx context.Context, startOffset, endOffset int
if startOffset == endOffset {
return nil, nil
}
if config == nil {
return nil, merr.WrapErrServiceInternalMsg("flush config is nil")
}
if config.Schema == nil {
return nil, merr.WrapErrServiceInternalMsg("flush schema is nil")
}
schemaBlob, err := proto.Marshal(config.Schema)
if err != nil {
return nil, merr.WrapErrServiceInternalMsg("marshal flush schema failed: %s", err.Error())
}
if len(schemaBlob) == 0 {
return nil, merr.WrapErrServiceInternalMsg("marshal flush schema returned empty blob")
}
// build C flush config
var cConfig C.CFlushConfig
cSegmentPath := C.CString(config.SegmentBasePath)
defer C.free(unsafe.Pointer(cSegmentPath))
cConfig.segment_path = cSegmentPath
cSchemaBlob := C.CBytes(schemaBlob)
defer C.free(cSchemaBlob)
cConfig.schema_blob = cSchemaBlob
cConfig.schema_length = C.int64_t(len(schemaBlob))
cConfig.read_version = C.int64_t(config.ReadVersion)
cConfig.retry_limit = C.uint32_t(3)
@@ -1671,6 +1689,29 @@ func (s *LocalSegment) FlushData(ctx context.Context, startOffset, endOffset int
defer C.free(unsafe.Pointer(cWriterFormat))
cConfig.writer_format = cWriterFormat
}
schemaBasedPattern := config.SchemaBasedPattern
if schemaBasedPattern != "" {
cSchemaBasedPattern := C.CString(schemaBasedPattern)
defer C.free(unsafe.Pointer(cSchemaBasedPattern))
cConfig.schema_based_pattern = cSchemaBasedPattern
}
schemaBasedFormats := config.SchemaBasedFormats
if schemaBasedFormats != "" {
cSchemaBasedFormats := C.CString(schemaBasedFormats)
defer C.free(unsafe.Pointer(cSchemaBasedFormats))
cConfig.schema_based_formats = cSchemaBasedFormats
}
numAllowedFields := len(config.AllowedFieldIDs)
if numAllowedFields > 0 {
cAllowedFieldIDs := (*C.int64_t)(C.malloc(C.size_t(numAllowedFields) * C.size_t(unsafe.Sizeof(C.int64_t(0)))))
defer C.free(unsafe.Pointer(cAllowedFieldIDs))
allowedFieldIDSlice := unsafe.Slice(cAllowedFieldIDs, numAllowedFields)
for i, fieldID := range config.AllowedFieldIDs {
allowedFieldIDSlice[i] = C.int64_t(fieldID)
}
cConfig.allowed_field_ids = cAllowedFieldIDs
cConfig.num_allowed_fields = C.size_t(numAllowedFields)
}
// populate TEXT column configs
// All arrays must be C-allocated to avoid "Go pointer to unpinned Go pointer" panic
@@ -1693,6 +1734,9 @@ func (s *LocalSegment) FlushData(ctx context.Context, startOffset, endOffset int
cConfig.text_field_ids = cFieldIDs
cConfig.text_lob_paths = cLobPathPtrs
cConfig.text_inline_threshold = C.int64_t(config.TextInlineThreshold)
cConfig.text_max_lob_file_bytes = C.int64_t(config.TextMaxLobFileBytes)
cConfig.text_flush_threshold_bytes = C.int64_t(config.TextFlushThresholdBytes)
cConfig.num_text_columns = C.size_t(numTextCols)
} else {
cConfig.text_field_ids = nil
@@ -20,6 +20,7 @@ import (
"context"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
pkoracle "github.com/milvus-io/milvus/internal/querynodev2/pkoracle"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/util/segcore"
@@ -151,11 +152,18 @@ type FlushConfig struct {
CollectionID int64
// Partition ID
PartitionID int64
// Schema is the flush-task schema for the offset range being flushed.
// It must not be inferred from the mutable growing segment runtime schema.
Schema *schemapb.CollectionSchema
// TEXT column field IDs
TextFieldIDs []int64
// LOB base paths for each TEXT field (same order as TextFieldIDs)
// Format: {partition_path}/lobs/{field_id}
TextLobPaths []string
// TEXT LOB writer thresholds, aligned with DataNode writer settings.
TextInlineThreshold int64
TextMaxLobFileBytes int64
TextFlushThresholdBytes int64
// BM25 output sparse vector field IDs whose stats should be collected for
// the flushed offset range.
BM25FieldIDs []int64
@@ -172,6 +180,15 @@ type FlushConfig struct {
// it may be resolved from the acknowledged manifest to keep appends
// compatible with existing column groups.
WriterFormat string
// SchemaBasedPattern is passed as writer.split.schema_based.patterns.
// When set, C++ uses schema_based writer policy instead of single.
SchemaBasedPattern string
// SchemaBasedFormats is passed as writer.split.schema_based.formats.
// It preserves per-column-group formats when appending to existing manifests.
SchemaBasedFormats string
// AllowedFieldIDs limits growing flush output to fields compatible with
// the target segment layout.
AllowedFieldIDs []int64
}
// FlushResult contains the result of flushing segment data.
@@ -436,6 +436,7 @@ func (suite *SegmentSuite) TestFlushData() {
PartitionBasePath: suite.rootPath + "/partition",
CollectionID: suite.collectionID,
PartitionID: suite.partitionID,
Schema: suite.collection.Schema(),
}
rowNum := suite.growing.RowNum()
@@ -462,6 +463,7 @@ func (suite *SegmentSuite) TestFlushDataSealedSegmentFails() {
PartitionBasePath: suite.rootPath + "/partition",
CollectionID: suite.collectionID,
PartitionID: suite.partitionID,
Schema: suite.collection.Schema(),
}
// sealed segment should fail
@@ -478,6 +480,7 @@ func (suite *SegmentSuite) TestFlushDataInvalidOffsets() {
PartitionBasePath: suite.rootPath + "/partition",
CollectionID: suite.collectionID,
PartitionID: suite.partitionID,
Schema: suite.collection.Schema(),
}
// Test negative start offset
@@ -499,6 +502,7 @@ func (suite *SegmentSuite) TestFlushDataEmptyRange() {
PartitionBasePath: suite.rootPath + "/partition",
CollectionID: suite.collectionID,
PartitionID: suite.partitionID,
Schema: suite.collection.Schema(),
}
// empty range (start == end) should return nil, nil
@@ -515,6 +519,7 @@ func (suite *SegmentSuite) TestFlushDataPartialRange() {
PartitionBasePath: suite.rootPath + "/partition_partial",
CollectionID: suite.collectionID,
PartitionID: suite.partitionID,
Schema: suite.collection.Schema(),
}
rowNum := suite.growing.RowNum()
+11 -1
View File
@@ -67,6 +67,7 @@ type packedBinlogRecordWriterBase struct {
columnGroups []storagecommon.ColumnGroup
storageConfig *indexpb.StorageConfig
storagePluginContext *indexcgopb.StoragePluginContext
writerFormat string
// basePath is the segment data root, populated by initWriters. The
// underlying packed batch writers do not return a manifest path; the
// caller builds the manifest update against this base path.
@@ -165,7 +166,10 @@ func (pw *packedBinlogRecordWriterBase) GetRowNum() int64 {
}
func (pw *packedBinlogRecordWriterBase) fillV3ColumnGroupFormats() (string, []string) {
writerFormat := paramtable.Get().DataNodeCfg.StorageFormat.GetValue()
writerFormat := pw.writerFormat
if writerFormat == "" {
writerFormat = paramtable.Get().DataNodeCfg.StorageFormat.GetValue()
}
pw.columnGroups = storagecommon.FillColumnGroupFormats(pw.columnGroups, writerFormat)
return writerFormat, storagecommon.ColumnGroupFormats(pw.columnGroups, writerFormat)
}
@@ -339,6 +343,7 @@ func newPackedBinlogRecordWriter(collectionID, partitionID, segmentID UniqueID,
blobsWriter ChunkedBlobsWriter, allocator allocator.Interface, maxRowNum int64, bufferSize, multiPartUploadSize int64, columnGroups []storagecommon.ColumnGroup,
storageConfig *indexpb.StorageConfig,
storagePluginContext *indexcgopb.StoragePluginContext,
writerFormat string,
) (*PackedBinlogRecordWriter, error) {
arrowSchema, err := ConvertToArrowSchema(schema, true)
if err != nil {
@@ -360,6 +365,7 @@ func newPackedBinlogRecordWriter(collectionID, partitionID, segmentID UniqueID,
columnGroups: columnGroups,
storageConfig: storageConfig,
storagePluginContext: storagePluginContext,
writerFormat: writerFormat,
tsFrom: typeutil.MaxTimestamp,
tsTo: 0,
ttlFieldID: getTTLFieldID(schema),
@@ -559,6 +565,7 @@ func newPackedManifestRecordWriter(collectionID, partitionID, segmentID UniqueID
storageConfig *indexpb.StorageConfig,
storagePluginContext *indexcgopb.StoragePluginContext,
textRefsAsBinary bool,
writerFormat string,
) (*PackedManifestRecordWriter, error) {
arrowSchema, err := ConvertToArrowSchema(schema, true)
if err != nil {
@@ -580,6 +587,7 @@ func newPackedManifestRecordWriter(collectionID, partitionID, segmentID UniqueID
columnGroups: columnGroups,
storageConfig: storageConfig,
storagePluginContext: storagePluginContext,
writerFormat: writerFormat,
tsFrom: typeutil.MaxTimestamp,
tsTo: 0,
ttlFieldID: getTTLFieldID(schema),
@@ -736,6 +744,7 @@ func NewPackedTextManifestRecordWriter(
columnGroups []storagecommon.ColumnGroup,
storageConfig *indexpb.StorageConfig,
textColumnConfigs []packed.TextColumnConfig,
writerFormat string,
) (*PackedTextManifestRecordWriter, error) {
arrowSchema, err := ConvertToArrowSchema(schema, true)
if err != nil {
@@ -756,6 +765,7 @@ func NewPackedTextManifestRecordWriter(
multiPartUploadSize: multiPartUploadSize,
columnGroups: columnGroups,
storageConfig: storageConfig,
writerFormat: writerFormat,
tsFrom: typeutil.MaxTimestamp,
tsTo: 0,
ttlFieldID: getTTLFieldID(schema),
@@ -43,7 +43,7 @@ func TestPackedManifestRecordWriter_CloseWithoutWrite(t *testing.T) {
w, err := newPackedManifestRecordWriter(1, 2, 3, schema,
ChunkedBlobsWriter(func(_ []*Blob) error { return nil }),
allocator.NewLocalAllocator(1, 1<<20),
1024, 0, 0, nil, cfg, nil, false)
1024, 0, 0, nil, cfg, nil, false, "")
require.NoError(t, err)
// No Write before Close. The internal `writer` field stays nil so
@@ -71,7 +71,7 @@ func TestPackedTextManifestRecordWriter_CloseWithoutWrite(t *testing.T) {
w, err := NewPackedTextManifestRecordWriter(1, 2, 3, schema,
ChunkedBlobsWriter(func(_ []*Blob) error { return nil }),
allocator.NewLocalAllocator(1, 1<<20),
1024, 0, 0, nil, cfg, nil)
1024, 0, 0, nil, cfg, nil, "")
require.NoError(t, err)
// No Write before Close. The text writer's nil-handling path must
@@ -118,7 +118,7 @@ func TestPackedManifestRecordWriter_FillsV3ColumnGroupFormats(t *testing.T) {
w, err := newPackedManifestRecordWriter(1, 2, 3, schema,
ChunkedBlobsWriter(func(_ []*Blob) error { return nil }),
allocator.NewLocalAllocator(1, 1<<20),
1024, 0, 0, columnGroups, cfg, nil, false)
1024, 0, 0, columnGroups, cfg, nil, false, "")
require.NoError(t, err)
require.NoError(t, w.initWriters(nil))
@@ -171,7 +171,7 @@ func TestPackedManifestRecordWriter_TextRefsUseBinarySchema(t *testing.T) {
w, err := newPackedManifestRecordWriter(1, 2, 3, schema,
ChunkedBlobsWriter(func(_ []*Blob) error { return nil }),
allocator.NewLocalAllocator(1, 1<<20),
1024, 0, 0, columnGroups, cfg, nil, true)
1024, 0, 0, columnGroups, cfg, nil, true, "")
require.NoError(t, err)
err = w.initWriters(nil)
require.NoError(t, err)
@@ -179,6 +179,48 @@ func TestPackedManifestRecordWriter_TextRefsUseBinarySchema(t *testing.T) {
require.Equal(t, arrow.BINARY, gotSchema.Field(2).Type.ID())
}
func TestPackedManifestRecordWriter_UsesExplicitWriterFormat(t *testing.T) {
params := paramtable.Get()
require.NoError(t, params.Save(params.DataNodeCfg.StorageFormat.Key, "vortex"))
defer params.Reset(params.DataNodeCfg.StorageFormat.Key)
schema := &schemapb.CollectionSchema{Fields: []*schemapb.FieldSchema{
{FieldID: common.TimeStampField, DataType: schemapb.DataType_Int64},
{FieldID: common.RowIDField, DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, DataType: schemapb.DataType_Int64},
}}
columnGroups := []storagecommon.ColumnGroup{
{GroupID: 0, Columns: []int{0, 1}, Fields: []int64{common.TimeStampField, common.RowIDField}},
{GroupID: 101, Columns: []int{2}, Fields: []int64{101}},
}
cfg := &indexpb.StorageConfig{StorageType: "local", RootPath: t.TempDir()}
var gotWriterFormat string
var gotSchemaBasedFormats []string
patch := mockey.Mock(newPackedRecordBatchWriter).To(
func(_ string, _ *schemapb.CollectionSchema, _, _ int64, _ []storagecommon.ColumnGroup,
_ *indexpb.StorageConfig, _ *indexcgopb.StoragePluginContext, validatePK bool, textRefsAsBinary bool,
writerFormat string, schemaBasedFormats []string,
) (*packedRecordBatchWriter, error) {
assert.True(t, validatePK)
assert.False(t, textRefsAsBinary)
gotWriterFormat = writerFormat
gotSchemaBasedFormats = append([]string(nil), schemaBasedFormats...)
return &packedRecordBatchWriter{}, nil
}).Build()
defer patch.UnPatch()
w, err := newPackedManifestRecordWriter(1, 2, 3, schema,
ChunkedBlobsWriter(func(_ []*Blob) error { return nil }),
allocator.NewLocalAllocator(1, 1<<20),
1024, 0, 0, columnGroups, cfg, nil, false, "parquet")
require.NoError(t, err)
require.NoError(t, w.initWriters(nil))
assert.Equal(t, "parquet", gotWriterFormat)
assert.Equal(t, []string{"parquet", "parquet"}, gotSchemaBasedFormats)
}
func TestPackedTextManifestRecordWriter_FillsV3ColumnGroupFormats(t *testing.T) {
params := paramtable.Get()
require.NoError(t, params.Save(params.DataNodeCfg.StorageFormat.Key, "vortex"))
@@ -212,7 +254,7 @@ func TestPackedTextManifestRecordWriter_FillsV3ColumnGroupFormats(t *testing.T)
w, err := NewPackedTextManifestRecordWriter(1, 2, 3, schema,
ChunkedBlobsWriter(func(_ []*Blob) error { return nil }),
allocator.NewLocalAllocator(1, 1<<20),
1024, 0, 0, columnGroups, cfg, nil)
1024, 0, 0, columnGroups, cfg, nil, "")
require.NoError(t, err)
require.NoError(t, w.initWriters(nil))
+10
View File
@@ -76,6 +76,7 @@ type rwOptions struct {
textColumnConfigs []packed.TextColumnConfig // TEXT column configurations for REWRITE_ALL mode
textRefsAsBinary bool // TEXT columns already contain encoded LOB refs and should be copied as-is
externalReader packed.ExternalReaderContext
writerFormat string
}
func (o *rwOptions) validate() error {
@@ -214,6 +215,12 @@ func WithTextRefsAsBinary() RwOption {
}
}
func WithWriterFormat(format string) RwOption {
return func(options *rwOptions) {
options.writerFormat = format
}
}
func makeBlobsReader(ctx context.Context, binlogs []*datapb.FieldBinlog, downloader downloaderFn) (ChunkedBlobsReader, error) {
if len(binlogs) == 0 {
return func() ([]*Blob, error) {
@@ -445,6 +452,7 @@ func NewBinlogRecordWriter(ctx context.Context, collectionID, partitionID, segme
rwOptions.bufferSize, rwOptions.multiPartUploadSize, rwOptions.columnGroups,
rwOptions.storageConfig,
pluginContext,
rwOptions.writerFormat,
)
case StorageV3:
// if TEXT column configs are provided, use the text writer with TEXT column support
@@ -454,6 +462,7 @@ func NewBinlogRecordWriter(ctx context.Context, collectionID, partitionID, segme
rwOptions.bufferSize, rwOptions.multiPartUploadSize, rwOptions.columnGroups,
rwOptions.storageConfig,
rwOptions.textColumnConfigs,
rwOptions.writerFormat,
)
}
return newPackedManifestRecordWriter(collectionID, partitionID, segmentID, schema,
@@ -462,6 +471,7 @@ func NewBinlogRecordWriter(ctx context.Context, collectionID, partitionID, segme
rwOptions.storageConfig,
pluginContext,
rwOptions.textRefsAsBinary,
rwOptions.writerFormat,
)
}
return nil, merr.WrapErrServiceInternalMsg("unsupported storage version %d", rwOptions.version)