(db/snapshot): Add new TypeNameField and update some ut (#3100)

* (db/snapshot): Add TypeNameField
Signed-off-by: Npeng.xu <peng.xu@zilliz.com>

* (db/snapshot): add test for DropAllIndexOperation
Signed-off-by: Npeng.xu <peng.xu@zilliz.com>
Co-authored-by: NJin Hai <hai.jin@zilliz.com>

core/src/db/meta/MetaNames.h

@@ -29,6 +29,7 @@ extern const char* F_NAME;
 extern const char* F_PARAMS;
 extern const char* F_SIZE;
 extern const char* F_ROW_COUNT;
+extern const char* F_TYPE_NAME;
 
 ////////////////////////////////////////////////////////
 // Table names

core/src/db/meta/MetaResourceAttrs.cpp

@@ -34,6 +34,7 @@ const char* F_NAME = snapshot::NameField::Name;
 const char* F_PARAMS = snapshot::ParamsField::Name;
 const char* F_SIZE = snapshot::SizeField::Name;
 const char* F_ROW_COUNT = snapshot::RowCountField::Name;
+const char* F_TYPE_NAME = snapshot::TypeNameField::Name;
 
 const char* TABLE_COLLECTION = snapshot::Collection::Name;
 const char* TABLE_COLLECTION_COMMIT = snapshot::CollectionCommit::Name;
@@ -72,7 +73,8 @@ ResourceAttrMapOf(const std::string& table, std::vector<std::string>& attrs) {
         {snapshot::FieldCommit::Name,
          {F_COLLECTON_ID, F_FIELD_ID, F_MAPPINGS, F_ID, F_LSN, F_STATE, F_CREATED_ON, F_UPDATED_ON}},
         {snapshot::FieldElement::Name,
-         {F_COLLECTON_ID, F_FIELD_ID, F_NAME, F_FTYPE, F_PARAMS, F_ID, F_LSN, F_STATE, F_CREATED_ON, F_UPDATED_ON}},
+         {F_COLLECTON_ID, F_FIELD_ID, F_NAME, F_FTYPE, F_TYPE_NAME, F_PARAMS, F_ID, F_LSN, F_STATE, F_CREATED_ON,
+          F_UPDATED_ON}},
     };
 
     if (ResourceAttrMap.find(table) == ResourceAttrMap.end()) {

core/src/db/meta/MetaResourceAttrs.h

@@ -152,6 +152,10 @@ AttrValue2Str(typename ResourceContext<ResourceT>::ResPtr src, const std::string
         auto name_field = std::dynamic_pointer_cast<snapshot::NameField>(src);
         str_value = name_field->GetName();
         str2str(str_value, value);
+    } else if (F_TYPE_NAME == attr) {
+        auto type_name_field = std::dynamic_pointer_cast<snapshot::TypeNameField>(src);
+        str_value = type_name_field->GetTypeName();
+        str2str(str_value, value);
     } else if (F_PARAMS == attr) {
         auto params_field = std::dynamic_pointer_cast<snapshot::ParamsField>(src);
         json_value = params_field->GetParams();

core/src/db/meta/MetaSession.h

@@ -304,6 +304,14 @@ MetaSession::Select(const std::string& field, const std::vector<U>& values,
             }
         }
 
+        auto tn_p = std::dynamic_pointer_cast<snapshot::TypeNameField>(resource);
+        if (tn_p != nullptr) {
+            iter = raw.find(F_TYPE_NAME);
+            if (iter != raw.end()) {
+                tn_p->SetTypeName(iter->second);
+            }
+        }
+
         resources.push_back(std::move(resource));
     }
 

core/src/db/meta/backend/MySqlEngine.cpp

@@ -48,6 +48,7 @@ static const MetaField MetaUpdatedOnField = MetaField(F_UPDATED_ON, "BIGINT", "N
 static const MetaField MetaParamsField = MetaField(F_PARAMS, "JSON", "NOT NULL");
 static const MetaField MetaSizeField = MetaField(F_SIZE, "BIGINT", "NOT NULL");
 static const MetaField MetaRowCountField = MetaField(F_ROW_COUNT, "BIGINT", "NOT NULL");
+static const MetaField MetaTypeNameField = MetaField(F_TYPE_NAME, "VARCHAR(255)", "NOT NULL");
 
 // Environment schema
 static const MetaSchema COLLECTION_SCHEMA(TABLE_COLLECTION, {MetaIdField, MetaNameField, MetaLSNField, MetaParamsField,
@@ -99,8 +100,8 @@ static const MetaSchema FIELDCOMMIT_SCHEMA(TABLE_FIELD_COMMIT,
 
 static const MetaSchema FIELDELEMENT_SCHEMA(TABLE_FIELD_ELEMENT,
                                             {MetaIdField, MetaCollectionIdField, MetaFieldIdField, MetaNameField,
-                                             MetaFtypeField, MetaParamsField, MetaLSNField, MetaStateField,
-                                             MetaCreatedOnField, MetaUpdatedOnField});
+                                             MetaFtypeField, MetaTypeNameField, MetaParamsField, MetaLSNField,
+                                             MetaStateField, MetaCreatedOnField, MetaUpdatedOnField});
 
 }  // namespace
 

core/src/db/meta/backend/SqliteEngine.cpp

@@ -44,6 +44,7 @@ static const MetaField MetaUpdatedOnField = MetaField(F_UPDATED_ON, "BIGINT", "N
 static const MetaField MetaParamsField = MetaField(F_PARAMS, "VARCHAR(255)", "NOT NULL");
 static const MetaField MetaSizeField = MetaField(F_SIZE, "BIGINT", "NOT NULL");
 static const MetaField MetaRowCountField = MetaField(F_ROW_COUNT, "BIGINT", "NOT NULL");
+static const MetaField MetaTypeNameField = MetaField(F_TYPE_NAME, "VARCHAR(255)", "NOT NULL");
 
 // Environment schema
 static const MetaSchema COLLECTION_SCHEMA(TABLE_COLLECTION, {MetaIdField, MetaNameField, MetaLSNField, MetaParamsField,
@@ -117,8 +118,8 @@ static const MetaSchema FIELDCOMMIT_SCHEMA(TABLE_FIELD_COMMIT,
 
 static const MetaSchema FIELDELEMENT_SCHEMA(TABLE_FIELD_ELEMENT,
                                             {MetaIdField, MetaCollectionIdField, MetaFieldIdField, MetaNameField,
-                                             MetaFtypeField, MetaParamsField, MetaLSNField, MetaStateField,
-                                             MetaCreatedOnField, MetaUpdatedOnField});
+                                             MetaFtypeField, MetaTypeNameField, MetaParamsField, MetaLSNField,
+                                             MetaStateField, MetaCreatedOnField, MetaUpdatedOnField});
 
 /////////////////////////////////////////////////////
 static AttrsMapList* QueryData = nullptr;

core/src/db/snapshot/CompoundOperations.cpp

@@ -921,18 +921,19 @@ CreateCollectionOperation::DoExecute(StorePtr store) {
         AddStepWithLsn(*field, c_context_.lsn, f_ctx_p);
         MappingT element_ids = {};
         FieldElementPtr raw_element;
-        status = store->CreateResource<FieldElement>(
-            FieldElement(collection->GetID(), field->GetID(), DEFAULT_RAW_DATA_NAME, FieldElementType::FET_RAW),
-            raw_element);
+        status =
+            store->CreateResource<FieldElement>(FieldElement(collection->GetID(), field->GetID(), DEFAULT_RAW_DATA_NAME,
+                                                             FieldElementType::FET_RAW, DEFAULT_RAW_DATA_NAME),
+                                                raw_element);
         auto fe_ctx_p = ResourceContextBuilder<FieldElement>().SetOp(meta::oUpdate).CreatePtr();
         AddStepWithLsn(*raw_element, c_context_.lsn, fe_ctx_p);
         element_ids.insert(raw_element->GetID());
         for (auto& element_schema : field_elements) {
             FieldElementPtr element;
-            status =
-                store->CreateResource<FieldElement>(FieldElement(collection->GetID(), field->GetID(),
-                                                                 element_schema->GetName(), element_schema->GetFtype()),
-                                                    element);
+            status = store->CreateResource<FieldElement>(
+                FieldElement(collection->GetID(), field->GetID(), element_schema->GetName(), element_schema->GetFtype(),
+                             element_schema->GetTypeName()),
+                element);
             auto t_fe_ctx_p = ResourceContextBuilder<FieldElement>().SetOp(meta::oUpdate).CreatePtr();
             AddStepWithLsn(*element, c_context_.lsn, t_fe_ctx_p);
             element_ids.insert(element->GetID());

core/src/db/snapshot/Resources.cpp

@@ -186,12 +186,13 @@ FieldCommit::FieldCommit(ID_TYPE collection_id, ID_TYPE field_id, const MappingT
 }
 
 FieldElement::FieldElement(ID_TYPE collection_id, ID_TYPE field_id, const std::string& name, FTYPE_TYPE ftype,
-                           const json& params, ID_TYPE id, LSN_TYPE lsn, State state, TS_TYPE created_on,
-                           TS_TYPE updated_on)
+                           const std::string& type_name, const json& params, ID_TYPE id, LSN_TYPE lsn, State state,
+                           TS_TYPE created_on, TS_TYPE updated_on)
     : CollectionIdField(collection_id),
       FieldIdField(field_id),
       NameField(name),
       FtypeField(ftype),
+      TypeNameField(type_name),
       ParamsField(params),
       IdField(id),
       LsnField(lsn),

core/src/db/snapshot/Resources.h

@@ -349,6 +349,27 @@ class SegmentIdField {
     ID_TYPE segment_id_;
 };
 
+class TypeNameField {
+ public:
+    static constexpr const char* Name = "TypeName";
+
+    explicit TypeNameField(std::string val) : type_name_(std::move(val)) {
+    }
+
+    const std::string&
+    GetTypeName() const {
+        return type_name_;
+    }
+
+    void
+    SetTypeName(const std::string& val) {
+        type_name_ = val;
+    }
+
+ protected:
+    std::string type_name_;
+};
+
 class NameField {
  public:
     static constexpr const char* Name = "name";
@@ -709,6 +730,7 @@ class FieldElement : public BaseResource<FieldElement>,
                      public FieldIdField,
                      public NameField,
                      public FtypeField,
+                     public TypeNameField,
                      public ParamsField,
                      public IdField,
                      public LsnField,
@@ -723,8 +745,9 @@ class FieldElement : public BaseResource<FieldElement>,
     using VecT = std::vector<Ptr>;
     static constexpr const char* Name = "FieldElement";
     FieldElement(ID_TYPE collection_id, ID_TYPE field_id, const std::string& name, FTYPE_TYPE ftype,
-                 const json& params = JEmpty, ID_TYPE id = 0, LSN_TYPE lsn = 0, State status = PENDING,
-                 TS_TYPE created_on = GetMicroSecTimeStamp(), TS_TYPE UpdatedOnField = GetMicroSecTimeStamp());
+                 const std::string& type_name = "", const json& params = JEmpty, ID_TYPE id = 0, LSN_TYPE lsn = 0,
+                 State status = PENDING, TS_TYPE created_on = GetMicroSecTimeStamp(),
+                 TS_TYPE UpdatedOnField = GetMicroSecTimeStamp());
 };
 
 using FieldElementPtr = FieldElement::Ptr;

core/src/db/snapshot/Snapshot.cpp

@@ -19,13 +19,11 @@ namespace snapshot {
 
 void
 Snapshot::RefAll() {
-    /* std::cout << this << " RefAll SS=" << GetID() << " SS RefCnt=" << ref_count() << std::endl; */
     std::apply([this](auto&... resource) { ((DoRef(resource)), ...); }, resources_);
 }
 
 void
 Snapshot::UnRefAll() {
-    /* std::cout << this << " UnRefAll SS=" << GetID() << " SS RefCnt=" << ref_count() << std::endl; */
     std::apply([this](auto&... resource) { ((DoUnRef(resource)), ...); }, resources_);
 }
 
@@ -249,6 +247,7 @@ Snapshot::ToString() const {
         for (auto& fe_id : fc_m) {
             auto fe = GetResource<FieldElement>(fe_id);
             ss << "\n\tFieldElement: id=" << fe_id << ",name=" << fe->GetName() << " CID=" << fe->GetCollectionId();
+            ss << ",ftype=" << fe->GetFtype() << ",typename=" << fe->GetTypeName();
         }
     }
 

core/unittest/db/test_snapshot.cpp

@@ -541,31 +541,41 @@ TEST_F(SnapshotTest, IndexTest) {
     sf_collector->Iterate();
     auto prev_total = sf_collector->segment_files_.size();
 
+    decltype(sf_context) osf_context;
+    SFContextBuilder(osf_context, ss, {sf_context.field_element_name});
+    std::vector<SegmentFileContext> sfs_context = {sf_context, osf_context};
+
     auto new_total = 0;
     auto partitions = ss->GetResources<Partition>();
     for (auto& kv : partitions) {
         num = RandomInt(2, 5);
         auto row_cnt = 1024;
         for (auto i = 0; i < num; ++i) {
-            ASSERT_TRUE(CreateSegment(ss, kv.first, next_lsn(), sf_context, row_cnt).ok());
+            ASSERT_TRUE(CreateSegment(ss, kv.first, next_lsn(), sfs_context, row_cnt).ok());
         }
-        new_total += num;
+        new_total += num * sfs_context.size();
     }
 
     status = Snapshots::GetInstance().GetSnapshot(ss, ss->GetName());
     ASSERT_TRUE(status.ok());
+    std::cout << ss->ToString() << std::endl;
 
     sf_collector = std::make_shared<SegmentFileCollector>(ss, filter2);
     sf_collector->Iterate();
     auto total = sf_collector->segment_files_.size();
     ASSERT_EQ(total, prev_total + new_total);
 
-    auto field_element_id = ss->GetFieldElementId(sf_context.field_name,
-            sf_context.field_element_name);
-    ASSERT_NE(field_element_id, 0);
+    std::set<ID_TYPE> fe_ids;
+
+    for (auto ctx : sfs_context) {
+        auto field_element_id = ss->GetFieldElementId(ctx.field_name,
+                ctx.field_element_name);
+        ASSERT_NE(field_element_id, 0);
+        fe_ids.insert(field_element_id);
+    }
 
     auto filter3 = [&](SegmentFile::Ptr segment_file) -> bool {
-        return segment_file->GetFieldElementId() == field_element_id;
+        return fe_ids.find(segment_file->GetFieldElementId()) != fe_ids.end();
     };
     sf_collector = std::make_shared<SegmentFileCollector>(ss, filter3);
     sf_collector->Iterate();
@@ -573,14 +583,9 @@ TEST_F(SnapshotTest, IndexTest) {
 
     OperationContext d_a_i_ctx;
     d_a_i_ctx.lsn = next_lsn();
-    d_a_i_ctx.stale_field_elements.push_back(ss->GetResource<FieldElement>(field_element_id));
-
-    FieldElement::Ptr fe;
-    status = ss->GetFieldElement(sf_context.field_name, sf_context.field_element_name,
-            fe);
-
-    ASSERT_TRUE(status.ok());
-    ASSERT_EQ(fe, d_a_i_ctx.stale_field_elements[0]);
+    for (auto fe_id : fe_ids) {
+        d_a_i_ctx.stale_field_elements.push_back(ss->GetResource<FieldElement>(fe_id));
+    }
 
     auto drop_all_index_op = std::make_shared<DropAllIndexOperation>(d_a_i_ctx, ss);
     status = drop_all_index_op->Push();
@@ -605,14 +610,18 @@ TEST_F(SnapshotTest, IndexTest) {
     /*     } */
     /* } */
 
+    std::cout << ss->ToString() << std::endl;
     sf_collector = std::make_shared<SegmentFileCollector>(ss, filter2);
     sf_collector->Iterate();
     ASSERT_EQ(sf_collector->segment_files_.size(), total - specified_segment_files_cnt);
+    std::cout << "sf size = " << sf_collector->segment_files_.size() << std::endl;
+    std::cout << "total = " << total << std::endl;
+    std::cout << "specified_segment_files_cnt = " << specified_segment_files_cnt << std::endl;
 
     {
         auto& field_elements = ss->GetResources<FieldElement>();
         for (auto& kv : field_elements) {
-            ASSERT_NE(kv.second->GetID(), field_element_id);
+            ASSERT_EQ(fe_ids.find(kv.second->GetID()), fe_ids.end());
         }
     }
 }

core/unittest/db/utils.h

@@ -98,10 +98,15 @@ RandomInt(int start, int end) {
 }
 
 inline void
-SFContextBuilder(SegmentFileContext& ctx, ScopedSnapshotT sss) {
+SFContextBuilder(SegmentFileContext& ctx, ScopedSnapshotT sss,
+        const std::set<std::string>& exclude_field_element_names = {}) {
     auto field = sss->GetResources<Field>().begin()->second;
     ctx.field_name = field->GetName();
     for (auto& kv : sss->GetResources<FieldElement>()) {
+        auto name = kv.second->GetName();
+        if (exclude_field_element_names.find(name) != exclude_field_element_names.end()) {
+            continue;
+        }
         ctx.field_element_name = kv.second->GetName();
         break;
     }
@@ -286,6 +291,30 @@ CreateSegment(ScopedSnapshotT ss, ID_TYPE partition_id, LSN_TYPE lsn, const Segm
     return op->GetSnapshot(ss);
 }
 
+inline Status
+CreateSegment(ScopedSnapshotT ss, ID_TYPE partition_id, LSN_TYPE lsn,
+        const std::vector<SegmentFileContext>& sfs_context,
+        SIZE_TYPE row_cnt) {
+    OperationContext context;
+    context.lsn = lsn;
+    context.prev_partition = ss->GetResource<Partition>(partition_id);
+    auto op = std::make_shared<NewSegmentOperation>(context, ss);
+    SegmentPtr new_seg;
+    STATUS_CHECK(op->CommitNewSegment(new_seg));
+    for (auto& sf_context : sfs_context) {
+        SegmentFilePtr seg_file;
+        auto nsf_context = sf_context;
+        nsf_context.segment_id = new_seg->GetID();
+        nsf_context.partition_id = new_seg->GetPartitionId();
+        STATUS_CHECK(op->CommitNewSegmentFile(nsf_context, seg_file));
+        seg_file->SetSize(row_cnt * 10);
+    }
+    op->CommitRowCount(row_cnt);
+    STATUS_CHECK(op->Push());
+
+    return op->GetSnapshot(ss);
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 class BaseTest : public ::testing::Test {
  protected: