github.com/milvus-io/milvus-sdk-go/v2@v2.4.1/examples/insert/insert.go (about) 1 package main 2 3 import ( 4 "context" 5 "encoding/csv" 6 "fmt" 7 "log" 8 "os" 9 "strconv" 10 "strings" 11 "time" 12 13 "github.com/milvus-io/milvus-sdk-go/v2/client" 14 "github.com/milvus-io/milvus-sdk-go/v2/entity" 15 ) 16 17 func main() { 18 // Milvus instance proxy address, may verify in your env/settings 19 milvusAddr := `localhost:19530` 20 21 // setup context for client creation, use 2 seconds here 22 ctx := context.Background() 23 ctx, cancel := context.WithTimeout(ctx, 2*time.Second) 24 defer cancel() 25 26 c, err := client.NewClient(ctx, client.Config{ 27 Address: milvusAddr, 28 }) 29 if err != nil { 30 // handling error and exit, to make example simple here 31 log.Fatal("failed to connect to milvus:", err.Error()) 32 } 33 // in a main func, remember to close the client 34 defer c.Close() 35 36 // here is the collection name we use in this example 37 collectionName := `gosdk_insert_example` 38 39 has, err := c.HasCollection(ctx, collectionName) 40 if err != nil { 41 log.Fatal("failed to check whether collection exists:", err.Error()) 42 } 43 if has { 44 // collection with same name exist, clean up mess 45 _ = c.DropCollection(ctx, collectionName) 46 } 47 48 // define collection schema, see film.csv 49 schema := entity.NewSchema().WithName(collectionName).WithDescription("this is the example collection for insert and search"). 50 WithField(entity.NewField().WithName("ID").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)). 51 WithField(entity.NewField().WithName("Year").WithDataType(entity.FieldTypeInt32)). 52 WithField(entity.NewField().WithName("Vector").WithDataType(entity.FieldTypeFloatVector).WithDim(8)) 53 54 err = c.CreateCollection(ctx, schema, entity.DefaultShardNumber) // only 1 shard 55 if err != nil { 56 log.Fatal("failed to create collection:", err.Error()) 57 } 58 59 films, err := loadFilmCSV() 60 if err != nil { 61 log.Fatal("failed to load film data csv:", err.Error()) 62 } 63 64 // row-base covert to column-base 65 ids := make([]int64, 0, len(films)) 66 years := make([]int32, 0, len(films)) 67 vectors := make([][]float32, 0, len(films)) 68 // string field is not supported yet 69 idTitle := make(map[int64]string) 70 for idx, film := range films { 71 ids = append(ids, film.ID) 72 idTitle[film.ID] = film.Title 73 years = append(years, film.Year) 74 vectors = append(vectors, films[idx].Vector[:]) // prevent same vector 75 } 76 idColumn := entity.NewColumnInt64("ID", ids) 77 yearColumn := entity.NewColumnInt32("Year", years) 78 vectorColumn := entity.NewColumnFloatVector("Vector", 8, vectors) 79 80 // insert into default partition 81 _, err = c.Insert(ctx, collectionName, "", idColumn, yearColumn, vectorColumn) 82 if err != nil { 83 log.Fatal("failed to insert film data:", err.Error()) 84 } 85 log.Println("insert completed") 86 ctx, cancel = context.WithTimeout(context.Background(), time.Second*120) 87 defer cancel() 88 err = c.Flush(ctx, collectionName, false) 89 if err != nil { 90 log.Fatal("failed to flush collection:", err.Error()) 91 } 92 log.Println("flush completed") 93 94 // load collection with async=false 95 err = c.LoadCollection(ctx, collectionName, false) 96 if err != nil { 97 log.Fatal("failed to load collection:", err.Error()) 98 } 99 log.Println("load collection completed") 100 101 searchFilm := films[0] // use first fim to search 102 vector := entity.FloatVector(searchFilm.Vector[:]) 103 // Use flat search param 104 sp, _ := entity.NewIndexFlatSearchParam() 105 sr, err := c.Search(ctx, collectionName, []string{}, "Year > 1990", []string{"ID"}, []entity.Vector{vector}, "Vector", 106 entity.L2, 10, sp) 107 if err != nil { 108 log.Fatal("fail to search collection:", err.Error()) 109 } 110 for _, result := range sr { 111 var idColumn *entity.ColumnInt64 112 for _, field := range result.Fields { 113 if field.Name() == "ID" { 114 c, ok := field.(*entity.ColumnInt64) 115 if ok { 116 idColumn = c 117 } 118 } 119 } 120 if idColumn == nil { 121 log.Fatal("result field not math") 122 } 123 for i := 0; i < result.ResultCount; i++ { 124 id, err := idColumn.ValueByIdx(i) 125 if err != nil { 126 log.Fatal(err.Error()) 127 } 128 title := idTitle[id] 129 fmt.Printf("file id: %d title: %s scores: %f\n", id, title, result.Scores[i]) 130 } 131 } 132 133 // clean up 134 _ = c.DropCollection(ctx, collectionName) 135 } 136 137 type film struct { 138 ID int64 139 Title string 140 Year int32 141 Vector [8]float32 // fix length array 142 } 143 144 func loadFilmCSV() ([]film, error) { 145 f, err := os.Open("../films.csv") // assume you are in examples/insert folder, if not, please change the path 146 if err != nil { 147 return []film{}, err 148 } 149 r := csv.NewReader(f) 150 raw, err := r.ReadAll() 151 if err != nil { 152 return []film{}, err 153 } 154 films := make([]film, 0, len(raw)) 155 for _, line := range raw { 156 if len(line) < 4 { // insuffcient column 157 continue 158 } 159 fi := film{} 160 // ID 161 v, err := strconv.ParseInt(line[0], 10, 64) 162 if err != nil { 163 continue 164 } 165 fi.ID = v 166 // Title 167 fi.Title = line[1] 168 // Year 169 v, err = strconv.ParseInt(line[2], 10, 64) 170 if err != nil { 171 continue 172 } 173 fi.Year = int32(v) 174 // Vector 175 vectorStr := strings.ReplaceAll(line[3], "[", "") 176 vectorStr = strings.ReplaceAll(vectorStr, "]", "") 177 parts := strings.Split(vectorStr, ",") 178 if len(parts) != 8 { // dim must be 8 179 continue 180 } 181 for idx, part := range parts { 182 part = strings.TrimSpace(part) 183 v, err := strconv.ParseFloat(part, 32) 184 if err != nil { 185 continue 186 } 187 fi.Vector[idx] = float32(v) 188 } 189 films = append(films, fi) 190 } 191 return films, nil 192 }