1/**
2 * Copyright 2021 Alibaba, Inc. and its affiliates. All Rights Reserved.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#define private public
18#include "server/write_request_builder.h"
19#undef private
20#include <gtest/gtest.h>
21#include "common/error_code.h"
22
23using namespace ::proxima::be;
24using namespace ::proxima::be::server;
25
26class WriteRequestBuilderTest : public testing::Test {
27 protected:
28 void SetUp() {}
29
30 void TearDown() {}
31};
32
33void fill_collection_info(proto::WriteRequest &request,
34 meta::CollectionMetaPtr *meta,
35 agent::ColumnOrderMapPtr *order_map,
36 bool with_repo = true, bool is_bytes = false) {
37 request.set_request_id("00000000");
38 if (with_repo) {
39 request.set_magic_number(140140140);
40 }
41 int dim = 4;
42 request.set_collection_name("mytest");
43 auto *row_meta = request.mutable_row_meta();
44 auto *column_meta = row_meta->add_index_column_metas();
45 column_meta->set_column_name("field1");
46 column_meta->set_data_type(proto::DataType::DT_VECTOR_FP32);
47 column_meta->set_dimension(dim);
48 ;
49 row_meta->add_forward_column_names("forward_f1");
50 row_meta->add_forward_column_names("forward_f2");
51 auto *row1 = request.add_rows();
52 row1->set_primary_key(1000);
53 row1->set_operation_type(::proxima::be::proto::OP_INSERT);
54 if (with_repo) {
55 auto *ctx = row1->mutable_lsn_context();
56 ctx->set_lsn(1);
57 ctx->set_context("binlog:123");
58 }
59 if (!is_bytes) {
60 row1->mutable_index_column_values()->add_values()->set_string_value(
61 "[1,2,3,4]");
62 } else {
63 std::vector<float> vectors = {1.0, 2.0, 3.0, 4.0};
64 row1->mutable_index_column_values()->add_values()->set_bytes_value(
65 (const void *)vectors.data(), vectors.size() * sizeof(float));
66 }
67 row1->mutable_forward_column_values()->add_values()->set_float_value(10.0);
68 row1->mutable_forward_column_values()->add_values()->set_int32_value(20);
69
70 *meta = std::make_shared<meta::CollectionMeta>();
71 (*meta)->set_name("mytest");
72
73 if (with_repo) {
74 meta::RepositoryBasePtr repo =
75 std::make_shared<meta::DatabaseRepositoryMeta>();
76 repo->set_name("mytest");
77 (*meta)->set_repository(repo);
78 }
79
80 (*meta)->mutable_forward_columns()->emplace_back("forward_f1");
81 (*meta)->mutable_forward_columns()->emplace_back("forward_f2");
82
83 auto *index_columns = (*meta)->mutable_index_columns();
84
85 meta::ColumnMetaPtr column_meta1 = std::make_shared<meta::ColumnMeta>();
86 column_meta1->set_index_type(IndexTypes::PROXIMA_GRAPH_INDEX);
87 column_meta1->set_name("field1");
88 column_meta1->set_data_type(DataTypes::VECTOR_FP32);
89 column_meta1->set_dimension(4);
90
91 index_columns->emplace_back(column_meta1);
92
93 *order_map = std::make_shared<agent::ColumnOrderMap>();
94 (*order_map)->add_column_order(**meta);
95}
96
97int create_write_request(const proto::WriteRequest &request,
98 const meta::CollectionMetaPtr &meta,
99 agent::ColumnOrderMapPtr &order_map,
100 agent::WriteRequest *write_request) {
101 auto column_order = order_map->get_column_order(request.collection_name());
102
103 return WriteRequestBuilder::build(*meta, *column_order, request,
104 write_request);
105}
106
107TEST_F(WriteRequestBuilderTest, TestCreateSuccessWithProxyWrite) {
108 proto::WriteRequest request;
109 meta::CollectionMetaPtr meta;
110 agent::ColumnOrderMapPtr order_map;
111 fill_collection_info(request, &meta, &order_map, true);
112 agent::WriteRequest write_request;
113 int ret = create_write_request(request, meta, order_map, &write_request);
114 ASSERT_EQ(ret, 0);
115 ASSERT_EQ(write_request.magic_number(), 140140140);
116 ASSERT_EQ(write_request.collection_name(), "mytest");
117 auto &records = write_request.get_collection_dataset();
118 ASSERT_EQ(records.size(), (size_t)1);
119 ASSERT_EQ(records[0]->size(), (size_t)1);
120 ASSERT_EQ(records[0]->schema_revision_, 0);
121 auto &raw_data = records[0]->get(0);
122 ASSERT_EQ(raw_data.primary_key, 1000);
123 ASSERT_EQ(raw_data.operation_type, OperationTypes::INSERT);
124 ASSERT_EQ(raw_data.lsn, 1);
125 ASSERT_EQ(raw_data.lsn_context, "binlog:123");
126 ::proxima::be::proto::GenericValueList forward_list;
127 forward_list.add_values()->set_float_value(10.0);
128 forward_list.add_values()->set_int32_value(20);
129 std::string expected_forward_str;
130 forward_list.SerializeToString(&expected_forward_str);
131 ASSERT_EQ(raw_data.forward_data, expected_forward_str);
132 auto &column_datas = raw_data.column_datas;
133 ASSERT_EQ(column_datas.size(), (size_t)1);
134 ASSERT_EQ(column_datas[0].column_name, "field1");
135 ASSERT_EQ(column_datas[0].data_type, DataTypes::VECTOR_FP32);
136 ASSERT_EQ(column_datas[0].dimension, (uint32_t)4);
137 ASSERT_EQ(column_datas[0].data.size(), (size_t)16);
138 const float *vector_data = (const float *)column_datas[0].data.data();
139 ASSERT_FLOAT_EQ(vector_data[0], 1.0);
140 ASSERT_FLOAT_EQ(vector_data[1], 2.0);
141 ASSERT_FLOAT_EQ(vector_data[2], 3.0);
142 ASSERT_FLOAT_EQ(vector_data[3], 4.0);
143}
144
145TEST_F(WriteRequestBuilderTest, TestCreateSuccessWithDirectWrite) {
146 proto::WriteRequest request;
147 meta::CollectionMetaPtr meta;
148 agent::ColumnOrderMapPtr order_map;
149 fill_collection_info(request, &meta, &order_map, false);
150 agent::WriteRequest write_request;
151 int ret = create_write_request(request, meta, order_map, &write_request);
152 ASSERT_EQ(ret, 0);
153 ASSERT_EQ(write_request.magic_number(), 0);
154 ASSERT_EQ(write_request.collection_name(), "mytest");
155 auto &records = write_request.get_collection_dataset();
156 ASSERT_EQ(records.size(), (size_t)1);
157 auto &record = records[0];
158 ASSERT_EQ(record->schema_revision_, 0);
159 auto &raw_data = record->get(0);
160 ASSERT_EQ(raw_data.primary_key, 1000);
161 ASSERT_EQ(raw_data.lsn_check, false);
162 ASSERT_EQ(raw_data.operation_type, OperationTypes::INSERT);
163 ::proxima::be::proto::GenericValueList forward_list;
164 forward_list.add_values()->set_float_value(10.0);
165 forward_list.add_values()->set_int32_value(20);
166 std::string expected_forward_str;
167 forward_list.SerializeToString(&expected_forward_str);
168 ASSERT_EQ(raw_data.forward_data, expected_forward_str);
169 auto &column_datas = raw_data.column_datas;
170 ASSERT_EQ(column_datas.size(), (size_t)1);
171 ASSERT_EQ(column_datas[0].column_name, "field1");
172 ASSERT_EQ(column_datas[0].data_type, DataTypes::VECTOR_FP32);
173 ASSERT_EQ(column_datas[0].dimension, (uint32_t)4);
174 ASSERT_EQ(column_datas[0].data.size(), (size_t)16);
175 const float *vector_data = (const float *)column_datas[0].data.data();
176 ASSERT_FLOAT_EQ(vector_data[0], 1.0);
177 ASSERT_FLOAT_EQ(vector_data[1], 2.0);
178 ASSERT_FLOAT_EQ(vector_data[2], 3.0);
179 ASSERT_FLOAT_EQ(vector_data[3], 4.0);
180}
181
182TEST_F(WriteRequestBuilderTest, TestCreateFailedWithDimensionMismatched) {
183 proto::WriteRequest request;
184 meta::CollectionMetaPtr meta;
185 agent::ColumnOrderMapPtr order_map;
186 fill_collection_info(request, &meta, &order_map, true);
187 request.mutable_row_meta()->mutable_index_column_metas(0)->set_dimension(100);
188 agent::WriteRequest write_request;
189 int ret = create_write_request(request, meta, order_map, &write_request);
190 ASSERT_EQ(ret, ErrorCode_InvalidWriteRequest);
191}
192
193TEST_F(WriteRequestBuilderTest, TestCreateFailedWithValidate) {
194 proto::WriteRequest request;
195 meta::CollectionMetaPtr meta;
196 agent::ColumnOrderMapPtr order_map;
197 fill_collection_info(request, &meta, &order_map, true);
198 agent::WriteRequest write_request;
199 proto::WriteRequest tmp_request;
200 int ret = create_write_request(tmp_request, meta, order_map, &write_request);
201 ASSERT_EQ(ret, ErrorCode_InvalidWriteRequest);
202}
203
204TEST_F(WriteRequestBuilderTest, TestCreateFailedWithBuildProxyRequest) {
205 proto::WriteRequest request;
206 meta::CollectionMetaPtr meta;
207 agent::ColumnOrderMapPtr order_map;
208 agent::WriteRequest write_request;
209 fill_collection_info(request, &meta, &order_map, true);
210 request.mutable_rows(0)->clear_lsn_context();
211 int ret = create_write_request(request, meta, order_map, &write_request);
212 ASSERT_EQ(ret, ErrorCode_EmptyLsnContext);
213}
214
215TEST_F(WriteRequestBuilderTest, TestCreateFailedWithBuildDirectRequest) {
216 proto::WriteRequest request;
217 meta::CollectionMetaPtr meta;
218 agent::ColumnOrderMapPtr order_map;
219 agent::WriteRequest write_request;
220 fill_collection_info(request, &meta, &order_map, false);
221 request.mutable_rows(0)
222 ->mutable_index_column_values()
223 ->mutable_values(0)
224 ->set_string_value("invalid vector");
225 int ret = create_write_request(request, meta, order_map, &write_request);
226 ASSERT_EQ(ret, ErrorCode_MismatchedDimension);
227}
228
229TEST_F(WriteRequestBuilderTest, TestGetIndexAndForwardModeWithFullMatch) {
230 proto::WriteRequest request;
231 meta::CollectionMetaPtr meta;
232 agent::ColumnOrderMapPtr order_map;
233 agent::WriteRequest write_request;
234 fill_collection_info(request, &meta, &order_map, false);
235 bool index_full_match;
236 bool forward_full_match;
237 WriteRequestBuilder::get_index_and_forward_mode(
238 request, *meta, &index_full_match, &forward_full_match);
239 ASSERT_TRUE(index_full_match);
240 ASSERT_TRUE(forward_full_match);
241}
242
243TEST_F(WriteRequestBuilderTest, TestGetIndexAndForwardModeWithNotFullMatch) {
244 proto::WriteRequest request;
245 meta::CollectionMetaPtr meta;
246 agent::ColumnOrderMapPtr order_map;
247 agent::WriteRequest write_request;
248 fill_collection_info(request, &meta, &order_map, false);
249 request.mutable_row_meta()->set_forward_column_names(0, "invalid");
250 auto *index_meta = request.mutable_row_meta()->mutable_index_column_metas(0);
251 index_meta->set_column_name("invalid");
252
253 bool index_full_match;
254 bool forward_full_match;
255 WriteRequestBuilder::get_index_and_forward_mode(
256 request, *meta, &index_full_match, &forward_full_match);
257 ASSERT_FALSE(index_full_match);
258 ASSERT_FALSE(forward_full_match);
259}
260
261TEST_F(WriteRequestBuilderTest, TestValidateRequest) {
262 proto::WriteRequest request;
263 meta::CollectionMetaPtr meta;
264 agent::ColumnOrderMapPtr order_map;
265 agent::WriteRequest write_request;
266 fill_collection_info(request, &meta, &order_map, true);
267 agent::ColumnOrderPtr column_order = order_map->get_column_order("mytest");
268 int ret = WriteRequestBuilder::validate_request(request, *meta, *column_order,
269 true, true);
270 ASSERT_EQ(ret, 0);
271}
272
273TEST_F(WriteRequestBuilderTest, TestValidateRequestFailedWithIndexColumnSize) {
274 proto::WriteRequest request;
275 meta::CollectionMetaPtr meta;
276 agent::ColumnOrderMapPtr order_map;
277 agent::WriteRequest write_request;
278 fill_collection_info(request, &meta, &order_map, true);
279 agent::ColumnOrderPtr column_order = order_map->get_column_order("mytest");
280 request.mutable_row_meta()->add_index_column_metas();
281 int ret = WriteRequestBuilder::validate_request(request, *meta, *column_order,
282 true, true);
283 ASSERT_EQ(ret, ErrorCode_InvalidWriteRequest);
284}
285
286TEST_F(WriteRequestBuilderTest, TestValidateRequestFailedWithIndexColumnName) {
287 proto::WriteRequest request;
288 meta::CollectionMetaPtr meta;
289 agent::ColumnOrderMapPtr order_map;
290 agent::WriteRequest write_request;
291 fill_collection_info(request, &meta, &order_map, true);
292 agent::ColumnOrderPtr column_order = order_map->get_column_order("mytest");
293 auto *index_meta = request.mutable_row_meta()->mutable_index_column_metas(0);
294 index_meta->set_column_name("invalid");
295
296 int ret = WriteRequestBuilder::validate_request(request, *meta, *column_order,
297 false, true);
298 ASSERT_EQ(ret, ErrorCode_InvalidWriteRequest);
299}
300
301TEST_F(WriteRequestBuilderTest,
302 TestValidateRequestFailedWithForwardColumnSize) {
303 proto::WriteRequest request;
304 meta::CollectionMetaPtr meta;
305 agent::ColumnOrderMapPtr order_map;
306 agent::WriteRequest write_request;
307 fill_collection_info(request, &meta, &order_map, true);
308 agent::ColumnOrderPtr column_order = order_map->get_column_order("mytest");
309 request.mutable_row_meta()->add_forward_column_names("invalid");
310 int ret = WriteRequestBuilder::validate_request(request, *meta, *column_order,
311 true, false);
312 ASSERT_EQ(ret, ErrorCode_InvalidWriteRequest);
313}
314
315TEST_F(WriteRequestBuilderTest,
316 TestValidateRequestFailedWithForwardColumnName) {
317 proto::WriteRequest request;
318 meta::CollectionMetaPtr meta;
319 agent::ColumnOrderMapPtr order_map;
320 agent::WriteRequest write_request;
321 fill_collection_info(request, &meta, &order_map, true);
322 agent::ColumnOrderPtr column_order = order_map->get_column_order("mytest");
323 request.mutable_row_meta()->set_forward_column_names(0, "invalid");
324 int ret = WriteRequestBuilder::validate_request(request, *meta, *column_order,
325 true, false);
326 ASSERT_EQ(ret, ErrorCode_InvalidWriteRequest);
327}
328
329TEST_F(WriteRequestBuilderTest,
330 TestValidateRequestFailedWithIndexColumnSizeZero) {
331 proto::WriteRequest request;
332 meta::CollectionMetaPtr meta;
333 agent::ColumnOrderMapPtr order_map;
334 agent::WriteRequest write_request;
335 fill_collection_info(request, &meta, &order_map, true);
336 agent::ColumnOrderPtr column_order = order_map->get_column_order("mytest");
337 request.mutable_row_meta()->clear_index_column_metas();
338 int ret = WriteRequestBuilder::validate_request(request, *meta, *column_order,
339 true, false);
340 ASSERT_EQ(ret, ErrorCode_InvalidWriteRequest);
341}
342
343TEST_F(WriteRequestBuilderTest,
344 TestValidateRequestFailedWithIndexColumnSizeMismatched) {
345 proto::WriteRequest request;
346 meta::CollectionMetaPtr meta;
347 agent::ColumnOrderMapPtr order_map;
348 agent::WriteRequest write_request;
349 fill_collection_info(request, &meta, &order_map, true);
350 agent::ColumnOrderPtr column_order = order_map->get_column_order("mytest");
351 request.mutable_rows(0)->clear_index_column_values();
352 int ret = WriteRequestBuilder::validate_request(request, *meta, *column_order,
353 true, false);
354 ASSERT_EQ(ret, ErrorCode_InvalidWriteRequest);
355}
356
357TEST_F(WriteRequestBuilderTest,
358 TestValidateRequestFailedWithForwardColumnSizeMismatched) {
359 proto::WriteRequest request;
360 meta::CollectionMetaPtr meta;
361 agent::ColumnOrderMapPtr order_map;
362 agent::WriteRequest write_request;
363 fill_collection_info(request, &meta, &order_map, true);
364 agent::ColumnOrderPtr column_order = order_map->get_column_order("mytest");
365 request.mutable_rows(0)->clear_forward_column_values();
366 int ret = WriteRequestBuilder::validate_request(request, *meta, *column_order,
367 true, false);
368 ASSERT_EQ(ret, ErrorCode_InvalidWriteRequest);
369}
370
371TEST_F(WriteRequestBuilderTest, TestBuildForwardsDataWithFull) {
372 proto::WriteRequest request;
373 meta::CollectionMetaPtr meta;
374 agent::ColumnOrderMapPtr order_map;
375 fill_collection_info(request, &meta, &order_map, true);
376 agent::ColumnOrderPtr column_order = order_map->get_column_order("mytest");
377
378 auto *row = request.mutable_rows(0);
379 auto &row_meta = request.row_meta();
380 index::CollectionDataset::RowData row_data;
381 bool forward_full_match = true;
382 int ret = WriteRequestBuilder::build_forwards_data(
383 *row, row_meta, *column_order, *meta, forward_full_match, &row_data);
384 ASSERT_EQ(ret, 0);
385}
386
387TEST_F(WriteRequestBuilderTest, TestBuildForwardsDataWithNotFull) {
388 proto::WriteRequest request;
389 meta::CollectionMetaPtr meta;
390 agent::ColumnOrderMapPtr order_map;
391 fill_collection_info(request, &meta, &order_map, true);
392 agent::ColumnOrderPtr column_order = order_map->get_column_order("mytest");
393
394 auto *row = request.mutable_rows(0);
395 auto &row_meta = request.row_meta();
396 index::CollectionDataset::RowData row_data;
397 bool forward_full_match = false;
398 int ret = WriteRequestBuilder::build_forwards_data(
399 *row, row_meta, *column_order, *meta, forward_full_match, &row_data);
400 ASSERT_EQ(ret, 0);
401}
402
403TEST_F(WriteRequestBuilderTest, TestBuildForwardsDataWithInvalidForwardColumn) {
404 proto::WriteRequest request;
405 meta::CollectionMetaPtr meta;
406 agent::ColumnOrderMapPtr order_map;
407 fill_collection_info(request, &meta, &order_map, true);
408 agent::ColumnOrderPtr column_order = order_map->get_column_order("mytest");
409
410 auto *row = request.mutable_rows(0);
411 auto *row_meta = request.mutable_row_meta();
412 row_meta->set_forward_column_names(0, "invalid");
413 index::CollectionDataset::RowData row_data;
414 bool forward_full_match = false;
415 int ret = WriteRequestBuilder::build_forwards_data(
416 *row, *row_meta, *column_order, *meta, forward_full_match, &row_data);
417 ASSERT_EQ(ret, ErrorCode_InvalidWriteRequest);
418}
419
420TEST_F(WriteRequestBuilderTest, TestBuildIndexesDataWithFull) {
421 proto::WriteRequest request;
422 meta::CollectionMetaPtr meta;
423 agent::ColumnOrderMapPtr order_map;
424 fill_collection_info(request, &meta, &order_map, true);
425
426 auto *row = request.mutable_rows(0);
427 auto &row_meta = request.row_meta();
428 index::CollectionDataset::RowData row_data;
429 bool index_full_match = true;
430 int ret = WriteRequestBuilder::build_indexes_data(
431 *row, row_meta, *meta, index_full_match, &row_data);
432 ASSERT_EQ(ret, 0);
433}
434
435TEST_F(WriteRequestBuilderTest, TestBuildIndexesDataWithNotFull) {
436 proto::WriteRequest request;
437 meta::CollectionMetaPtr meta;
438 agent::ColumnOrderMapPtr order_map;
439 fill_collection_info(request, &meta, &order_map, true);
440
441 auto *row = request.mutable_rows(0);
442 auto &row_meta = request.row_meta();
443 index::CollectionDataset::RowData row_data;
444 bool index_full_match = false;
445 int ret = WriteRequestBuilder::build_indexes_data(
446 *row, row_meta, *meta, index_full_match, &row_data);
447 ASSERT_EQ(ret, 0);
448}
449
450TEST_F(WriteRequestBuilderTest,
451 TestBuildIndexesDataWithIndexColumnNameInvalid) {
452 proto::WriteRequest request;
453 meta::CollectionMetaPtr meta;
454 agent::ColumnOrderMapPtr order_map;
455 fill_collection_info(request, &meta, &order_map, true);
456
457 auto *row = request.mutable_rows(0);
458 auto *row_meta = request.mutable_row_meta();
459 auto *index_meta = request.mutable_row_meta()->mutable_index_column_metas(0);
460 index_meta->set_column_name("invalid");
461
462 index::CollectionDataset::RowData row_data;
463 bool index_full_match = false;
464 int ret = WriteRequestBuilder::build_indexes_data(
465 *row, *row_meta, *meta, index_full_match, &row_data);
466 ASSERT_EQ(ret, ErrorCode_MismatchedIndexColumn);
467}
468
469TEST_F(WriteRequestBuilderTest,
470 TestBuildIndexesDataWithIndexColumnTypeInvalid) {
471 proto::WriteRequest request;
472 meta::CollectionMetaPtr meta;
473 agent::ColumnOrderMapPtr order_map;
474 fill_collection_info(request, &meta, &order_map, true);
475
476 auto *row = request.mutable_rows(0);
477 auto *row_meta = request.mutable_row_meta();
478 row->mutable_index_column_values()->mutable_values(0)->set_float_value(
479 123.456);
480
481 index::CollectionDataset::RowData row_data;
482 bool index_full_match = false;
483 int ret = WriteRequestBuilder::build_indexes_data(
484 *row, *row_meta, *meta, index_full_match, &row_data);
485 ASSERT_EQ(ret, ErrorCode_MismatchedIndexColumn);
486}
487
488TEST_F(WriteRequestBuilderTest, TestBuildIndexesDataWithIndexTypeBytes) {
489 proto::WriteRequest request;
490 meta::CollectionMetaPtr meta;
491 agent::ColumnOrderMapPtr order_map;
492 fill_collection_info(request, &meta, &order_map, true, true);
493
494 auto *row = request.mutable_rows(0);
495 auto *row_meta = request.mutable_row_meta();
496 index::CollectionDataset::RowData row_data;
497 bool index_full_match = false;
498 int ret = WriteRequestBuilder::build_indexes_data(
499 *row, *row_meta, *meta, index_full_match, &row_data);
500 ASSERT_EQ(ret, 0);
501}
502