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  }