123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195 |
- // Licensed to the LF AI & Data foundation under one
- // or more contributor license agreements. See the NOTICE file
- // distributed with this work for additional information
- // regarding copyright ownership. The ASF licenses this file
- // to you under the Apache License, Version 2.0 (the
- // "License"); you may not use this file except in compliance
- // with the License. You may obtain a copy of the License at
- //
- // http://www.apache.org/licenses/LICENSE-2.0
- //
- // Unless required by applicable law or agreed to in writing, software
- // distributed under the License is distributed on an "AS IS" BASIS,
- // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- // See the License for the specific language governing permissions and
- // limitations under the License.
- package client
- import (
- "reflect"
- "runtime/debug"
- "github.com/cockroachdb/errors"
- "github.com/milvus-io/milvus/client/v2/column"
- "github.com/milvus-io/milvus/client/v2/entity"
- "github.com/milvus-io/milvus/client/v2/row"
- )
- // ResultSet is struct for search result set.
- type ResultSet struct {
- // internal schema for unmarshaling
- sch *entity.Schema
- ResultCount int // the returning entry count
- GroupByValue column.Column
- IDs column.Column // auto generated id, can be mapped to the columns from `Insert` API
- Fields DataSet // output field data
- Scores []float32 // distance to the target vector
- Err error // search error if any
- }
- // GetColumn returns column with provided field name.
- func (rs *ResultSet) GetColumn(fieldName string) column.Column {
- for _, column := range rs.Fields {
- if column.Name() == fieldName {
- return column
- }
- }
- return nil
- }
- // Unmarshal puts dataset into receiver in row based way.
- // `receiver` shall be a slice of pointer of model struct
- // eg, []*Records, in which type `Record` defines the row data.
- // note that distance/score is not unmarshaled here.
- func (sr *ResultSet) Unmarshal(receiver any) (err error) {
- err = sr.Fields.Unmarshal(receiver)
- if err != nil {
- return err
- }
- return sr.fillPKEntry(receiver)
- }
- func (sr *ResultSet) fillPKEntry(receiver any) (err error) {
- defer func() {
- if x := recover(); x != nil {
- err = errors.Newf("failed to unmarshal result set: %v, stack: %v", x, string(debug.Stack()))
- }
- }()
- rr := reflect.ValueOf(receiver)
- if rr.Kind() == reflect.Ptr {
- if rr.IsNil() && rr.CanAddr() {
- rr.Set(reflect.New(rr.Type().Elem()))
- }
- rr = rr.Elem()
- }
- rt := rr.Type()
- rv := rr
- switch rt.Kind() {
- case reflect.Slice:
- pkField := sr.sch.PKField()
- et := rt.Elem()
- for et.Kind() == reflect.Ptr {
- et = et.Elem()
- }
- candidates := row.ParseCandidate(et)
- candi, ok := candidates[pkField.Name]
- if !ok {
- // pk field not found in struct, skip
- return nil
- }
- for i := 0; i < sr.IDs.Len(); i++ {
- row := rv.Index(i)
- for row.Kind() == reflect.Ptr {
- row = row.Elem()
- }
- val, err := sr.IDs.Get(i)
- if err != nil {
- return err
- }
- row.Field(candi).Set(reflect.ValueOf(val))
- }
- rr.Set(rv)
- default:
- return errors.Newf("receiver need to be slice or array but get %v", rt.Kind())
- }
- return nil
- }
- // DataSet is an alias type for column slice.
- // Returned by query API.
- type DataSet []column.Column
- // Len returns the row count of dataset.
- // if there is no column, it shall return 0.
- func (ds DataSet) Len() int {
- if len(ds) == 0 {
- return 0
- }
- return ds[0].Len()
- }
- // Unmarshal puts dataset into receiver in row based way.
- // `receiver` shall be a slice of pointer of model struct
- // eg, []*Records, in which type `Record` defines the row data.
- func (ds DataSet) Unmarshal(receiver any) (err error) {
- defer func() {
- if x := recover(); x != nil {
- err = errors.Newf("failed to unmarshal result set: %v, stack: %v", x, string(debug.Stack()))
- }
- }()
- rr := reflect.ValueOf(receiver)
- if rr.Kind() == reflect.Ptr {
- if rr.IsNil() && rr.CanAddr() {
- rr.Set(reflect.New(rr.Type().Elem()))
- }
- rr = rr.Elem()
- }
- rt := rr.Type()
- rv := rr
- switch rt.Kind() {
- // TODO maybe support Array and just fill data
- // case reflect.Array:
- case reflect.Slice:
- et := rt.Elem()
- if et.Kind() != reflect.Ptr {
- return errors.Newf("receiver must be slice of pointers but get: %v", et.Kind())
- }
- for et.Kind() == reflect.Ptr {
- et = et.Elem()
- }
- for i := 0; i < ds.Len(); i++ {
- data := reflect.New(et)
- err := ds.fillData(data.Elem(), et, i)
- if err != nil {
- return err
- }
- rv = reflect.Append(rv, data)
- }
- rr.Set(rv)
- default:
- return errors.Newf("receiver need to be slice or array but get %v", rt.Kind())
- }
- return nil
- }
- func (ds DataSet) fillData(data reflect.Value, dataType reflect.Type, idx int) error {
- m := row.ParseCandidate(dataType)
- for i := 0; i < len(ds); i++ {
- name := ds[i].Name()
- fidx, ok := m[name]
- if !ok {
- // if target is not found, the behavior here is to ignore the column
- // `strict` mode could be added in the future to return error if any column missing
- continue
- }
- val, err := ds[i].Get(idx)
- if err != nil {
- return err
- }
- // TODO check datatype, return error here instead of reflect panicking & recover
- data.Field(fidx).Set(reflect.ValueOf(val))
- }
- return nil
- }
|