modifying _format_data_for_milvus based on PR comment

Manisha Sudhir · Manisha Sudhir · commit 97c47890c821 · 2023-08-16T18:54:29.000-07:00
diff --git a/sdk/python/feast/expediagroup/vectordb/milvus_online_store.py b/sdk/python/feast/expediagroup/vectordb/milvus_online_store.py
@@ -94,9 +94,9 @@ def online_write_batch(
     ) -> None:
         with MilvusConnectionManager(config.online_store):
             try:
-                entities = self._format_data_for_milvus(data)
+                rows = self._format_data_for_milvus(data)
                 collection_to_load_data = Collection(table.name)
-                collection_to_load_data.insert(entities)
+                collection_to_load_data.insert(rows)
                 #  The flush call will seal any remaining segments and send them for indexing
                 collection_to_load_data.flush()
             except Exception as e:
@@ -257,29 +257,11 @@ def _format_data_for_milvus(self, feast_data):
             feature = []
             for feature_name, val in values.items():
                 val_type = val.WhichOneof("val")
+                value = getattr(val, val_type)
                 if val_type == "float_list_val":
-                    float_list = val.float_list_val.val
-                    final_float_list = np.array(float_list)
-                    feature.append(final_float_list)
-                # TODO: Test out binary lists
-                if val_type == "biinary_list_val":
-                    binary_list = val.binary_list_val.val
-                    feature.append(binary_list)
-                if val_type == "string_val":
-                    string_val = val.string_val
-                    feature.append(string_val)
-                if val_type == "int32_val":
-                    int32_val = val.int32_val
-                    feature.append(int32_val)
-                if val_type == "int64_val":
-                    int64_val = val.int64_val
-                    feature.append(int64_val)
-                if val_type == "bytes_val":
-                    bytes_val = val.bytes_val
-                    feature.append(bytes_val)
-                if val_type == "float_val":
-                    float_val = val.float_val
-                    feature.append(float_val)
+                    value = np.array(value.val)
+                # TODO: Check binary vector conversion
+                feature.append(value)
             milvus_data.append(feature)
 
         transformed_data = [list(item) for item in zip(*milvus_data)]
diff --git a/sdk/python/tests/expediagroup/test_milvus_online_store.py b/sdk/python/tests/expediagroup/test_milvus_online_store.py
@@ -150,7 +150,7 @@ def create_n_customer_test_samples_milvus(self, n=10):
                 ),
                 {
                     "avg_orders_day": ValueProto(
-                        float_list_val=FloatList(val=[1, 2, 3, 4, 5])
+                        float_list_val=FloatList(val=[1.0, 2.1, 3.3, 4.0, 5.0])
                     ),
                     "name": ValueProto(string_val="John"),
                     "age": ValueProto(int64_val=3),

Original file line number	Diff line number	Diff line change
`@@ -150,7 +150,7 @@ def create_n_customer_test_samples_milvus(self, n=10):`
`150`	`150`	`),`
`151`	`151`	`{`
`152`	`152`	`"avg_orders_day": ValueProto(`
`153`		`- float_list_val=FloatList(val=[1, 2, 3, 4, 5])`
	`153`	`+ float_list_val=FloatList(val=[1.0, 2.1, 3.3, 4.0, 5.0])`
`154`	`154`	`),`
`155`	`155`	`"name": ValueProto(string_val="John"),`
`156`	`156`	`"age": ValueProto(int64_val=3),`