milvus/client/row/data.go

// 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 row

import (
	"encoding/json"
	"fmt"
	"reflect"
	"strconv"

	"github.com/cockroachdb/errors"

	"github.com/milvus-io/milvus/client/v2/column"
	"github.com/milvus-io/milvus/client/v2/entity"
)

const (
	// MilvusTag struct tag const for milvus row based struct
	MilvusTag = `milvus`

	// MilvusSkipTagValue struct tag const for skip this field.
	MilvusSkipTagValue = `-`

	// MilvusTagSep struct tag const for attribute separator
	MilvusTagSep = `;`

	// MilvusTagName struct tag const for field name
	MilvusTagName = `NAME`

	// VectorDimTag struct tag const for vector dimension
	VectorDimTag = `DIM`

	// VectorTypeTag struct tag const for binary vector type
	VectorTypeTag = `VECTOR_TYPE`

	// MilvusPrimaryKey struct tag const for primary key indicator
	MilvusPrimaryKey = `PRIMARY_KEY`

	// MilvusAutoID struct tag const for auto id indicator
	MilvusAutoID = `AUTO_ID`

	// MilvusMaxLength struct tag const for max length
	MilvusMaxLength = `MAX_LENGTH`

	// DimMax dimension max value
	DimMax = 65535
)

func AnyToColumns(rows []interface{}, schemas ...*entity.Schema) ([]column.Column, error) {
	rowsLen := len(rows)
	if rowsLen == 0 {
		return []column.Column{}, errors.New("0 length column")
	}

	var sch *entity.Schema
	var err error
	// if schema not provided, try to parse from row
	if len(schemas) == 0 {
		sch, err = ParseSchema(rows[0])
		if err != nil {
			return []column.Column{}, err
		}
	} else {
		// use first schema provided
		sch = schemas[0]
	}

	isDynamic := sch.EnableDynamicField
	var dynamicCol *column.ColumnJSONBytes

	nameColumns := make(map[string]column.Column)
	for _, field := range sch.Fields {
		// skip auto id pk field
		if field.PrimaryKey && field.AutoID {
			continue
		}
		switch field.DataType {
		case entity.FieldTypeBool:
			data := make([]bool, 0, rowsLen)
			col := column.NewColumnBool(field.Name, data)
			nameColumns[field.Name] = col
		case entity.FieldTypeInt8:
			data := make([]int8, 0, rowsLen)
			col := column.NewColumnInt8(field.Name, data)
			nameColumns[field.Name] = col
		case entity.FieldTypeInt16:
			data := make([]int16, 0, rowsLen)
			col := column.NewColumnInt16(field.Name, data)
			nameColumns[field.Name] = col
		case entity.FieldTypeInt32:
			data := make([]int32, 0, rowsLen)
			col := column.NewColumnInt32(field.Name, data)
			nameColumns[field.Name] = col
		case entity.FieldTypeInt64:
			data := make([]int64, 0, rowsLen)
			col := column.NewColumnInt64(field.Name, data)
			nameColumns[field.Name] = col
		case entity.FieldTypeFloat:
			data := make([]float32, 0, rowsLen)
			col := column.NewColumnFloat(field.Name, data)
			nameColumns[field.Name] = col
		case entity.FieldTypeDouble:
			data := make([]float64, 0, rowsLen)
			col := column.NewColumnDouble(field.Name, data)
			nameColumns[field.Name] = col
		case entity.FieldTypeString, entity.FieldTypeVarChar:
			data := make([]string, 0, rowsLen)
			col := column.NewColumnVarChar(field.Name, data)
			nameColumns[field.Name] = col
		case entity.FieldTypeJSON:
			data := make([][]byte, 0, rowsLen)
			col := column.NewColumnJSONBytes(field.Name, data)
			nameColumns[field.Name] = col
		case entity.FieldTypeArray:
			col := NewArrayColumn(field)
			if col == nil {
				return nil, errors.Newf("unsupported element type %s for Array", field.ElementType.String())
			}
			nameColumns[field.Name] = col
		case entity.FieldTypeFloatVector:
			data := make([][]float32, 0, rowsLen)
			dimStr, has := field.TypeParams[entity.TypeParamDim]
			if !has {
				return []column.Column{}, errors.New("vector field with no dim")
			}
			dim, err := strconv.ParseInt(dimStr, 10, 64)
			if err != nil {
				return []column.Column{}, fmt.Errorf("vector field with bad format dim: %s", err.Error())
			}
			col := column.NewColumnFloatVector(field.Name, int(dim), data)
			nameColumns[field.Name] = col
		case entity.FieldTypeBinaryVector:
			data := make([][]byte, 0, rowsLen)
			dim, err := field.GetDim()
			if err != nil {
				return []column.Column{}, err
			}
			col := column.NewColumnBinaryVector(field.Name, int(dim), data)
			nameColumns[field.Name] = col
		case entity.FieldTypeFloat16Vector:
			data := make([][]byte, 0, rowsLen)
			dim, err := field.GetDim()
			if err != nil {
				return []column.Column{}, err
			}
			col := column.NewColumnFloat16Vector(field.Name, int(dim), data)
			nameColumns[field.Name] = col
		case entity.FieldTypeBFloat16Vector:
			data := make([][]byte, 0, rowsLen)
			dim, err := field.GetDim()
			if err != nil {
				return []column.Column{}, err
			}
			col := column.NewColumnBFloat16Vector(field.Name, int(dim), data)
			nameColumns[field.Name] = col
		case entity.FieldTypeSparseVector:
			data := make([]entity.SparseEmbedding, 0, rowsLen)
			col := column.NewColumnSparseVectors(field.Name, data)
			nameColumns[field.Name] = col
		}
	}

	if isDynamic {
		dynamicCol = column.NewColumnJSONBytes("", make([][]byte, 0, rowsLen)).WithIsDynamic(true)
	}

	for _, row := range rows {
		// collection schema name need not to be same, since receiver could has other names
		v := reflect.ValueOf(row)
		set, err := reflectValueCandi(v)
		if err != nil {
			return nil, err
		}

		for idx, field := range sch.Fields {
			// skip dynamic field if visible
			if isDynamic && field.IsDynamic {
				continue
			}
			// skip auto id pk field
			if field.PrimaryKey && field.AutoID {
				// remove pk field from candidates set, avoid adding it into dynamic column
				delete(set, field.Name)
				continue
			}
			column, ok := nameColumns[field.Name]
			if !ok {
				return nil, fmt.Errorf("expected unhandled field %s", field.Name)
			}

			candi, ok := set[field.Name]
			if !ok {
				return nil, fmt.Errorf("row %d does not has field %s", idx, field.Name)
			}
			err := column.AppendValue(candi.v.Interface())
			if err != nil {
				return nil, err
			}
			delete(set, field.Name)
		}

		if isDynamic {
			m := make(map[string]interface{})
			for name, candi := range set {
				m[name] = candi.v.Interface()
			}
			bs, err := json.Marshal(m)
			if err != nil {
				return nil, fmt.Errorf("failed to marshal dynamic field %w", err)
			}
			err = dynamicCol.AppendValue(bs)
			if err != nil {
				return nil, fmt.Errorf("failed to append value to dynamic field %w", err)
			}
		}
	}
	columns := make([]column.Column, 0, len(nameColumns))
	for _, column := range nameColumns {
		columns = append(columns, column)
	}
	if isDynamic {
		columns = append(columns, dynamicCol)
	}
	return columns, nil
}

func NewArrayColumn(f *entity.Field) column.Column {
	switch f.ElementType {
	case entity.FieldTypeBool:
		return column.NewColumnBoolArray(f.Name, nil)

	case entity.FieldTypeInt8:
		return column.NewColumnInt8Array(f.Name, nil)

	case entity.FieldTypeInt16:
		return column.NewColumnInt16Array(f.Name, nil)

	case entity.FieldTypeInt32:
		return column.NewColumnInt32Array(f.Name, nil)

	case entity.FieldTypeInt64:
		return column.NewColumnInt64Array(f.Name, nil)

	case entity.FieldTypeFloat:
		return column.NewColumnFloatArray(f.Name, nil)

	case entity.FieldTypeDouble:
		return column.NewColumnDoubleArray(f.Name, nil)

	case entity.FieldTypeVarChar:
		return column.NewColumnVarCharArray(f.Name, nil)

	default:
		return nil
	}
}

type fieldCandi struct {
	name    string
	v       reflect.Value
	options map[string]string
}

func reflectValueCandi(v reflect.Value) (map[string]fieldCandi, error) {
	if v.Kind() == reflect.Ptr {
		v = v.Elem()
	}

	result := make(map[string]fieldCandi)
	switch v.Kind() {
	case reflect.Map: // map[string]any
		iter := v.MapRange()
		for iter.Next() {
			key := iter.Key().String()
			result[key] = fieldCandi{
				name: key,
				v:    iter.Value(),
			}
		}
		return result, nil
	case reflect.Struct:
		for i := 0; i < v.NumField(); i++ {
			ft := v.Type().Field(i)
			name := ft.Name
			tag, ok := ft.Tag.Lookup(MilvusTag)

			settings := make(map[string]string)
			if ok {
				if tag == MilvusSkipTagValue {
					continue
				}
				settings = ParseTagSetting(tag, MilvusTagSep)
				fn, has := settings[MilvusTagName]
				if has {
					// overwrite column to tag name
					name = fn
				}
			}
			_, ok = result[name]
			// duplicated
			if ok {
				return nil, fmt.Errorf("column has duplicated name: %s when parsing field: %s", name, ft.Name)
			}

			v := v.Field(i)
			if v.Kind() == reflect.Array {
				v = v.Slice(0, v.Len())
			}

			result[name] = fieldCandi{
				name:    name,
				v:       v,
				options: settings,
			}
		}

		return result, nil
	default:
		return nil, fmt.Errorf("unsupport row type: %s", v.Kind().String())
	}
}