diff --git a/internal/proxy/validate_util.go b/internal/proxy/validate_util.go index 904a2c5c1b..1ee77b3cf1 100644 --- a/internal/proxy/validate_util.go +++ b/internal/proxy/validate_util.go @@ -529,12 +529,7 @@ func (v *validateUtil) fillWithDefaultValue(field *schemapb.FieldData, fieldSche return merr.WrapErrParameterInvalid(numRows, len(field.GetValidData()), msg) } defaultValue := fieldSchema.GetDefaultValue().GetStringData() - geomT, err := wkt.Unmarshal(defaultValue) - if err != nil { - log.Warn("invalid default value for geometry field", zap.Error(err)) - return merr.WrapErrParameterInvalidMsg("invalid default value for geometry field") - } - defaultValueWkbBytes, err := wkb.Marshal(geomT, wkb.NDR) + defaultValueWkbBytes, err := common.ConvertWKTToWKB(defaultValue) if err != nil { log.Warn("invalid default value for geometry field", zap.Error(err)) return merr.WrapErrParameterInvalidMsg("invalid default value for geometry field") diff --git a/internal/util/importutilv2/csv/row_parser.go b/internal/util/importutilv2/csv/row_parser.go index cb566e060f..1ec7226c81 100644 --- a/internal/util/importutilv2/csv/row_parser.go +++ b/internal/util/importutilv2/csv/row_parser.go @@ -23,9 +23,6 @@ import ( "github.com/cockroachdb/errors" "github.com/samber/lo" - "github.com/twpayne/go-geom/encoding/wkb" - "github.com/twpayne/go-geom/encoding/wkbcommon" - "github.com/twpayne/go-geom/encoding/wkt" "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" "github.com/milvus-io/milvus/internal/json" @@ -289,11 +286,7 @@ func (r *rowParser) parseEntity(field *schemapb.FieldSchema, obj string) (any, e } return []byte(obj), nil case schemapb.DataType_Geometry: - geomT, err := wkt.Unmarshal(obj) - if err != nil { - return nil, r.wrapTypeError(obj, field) - } - wkbValue, err := wkb.Marshal(geomT, wkb.NDR, wkbcommon.WKBOptionEmptyPointHandling(wkbcommon.EmptyPointHandlingNaN)) + wkbValue, err := pkgcommon.ConvertWKTToWKB(obj) if err != nil { return nil, r.wrapTypeError(obj, field) } diff --git a/internal/util/importutilv2/json/row_parser.go b/internal/util/importutilv2/json/row_parser.go index 8de4b0045b..b9c889bf64 100644 --- a/internal/util/importutilv2/json/row_parser.go +++ b/internal/util/importutilv2/json/row_parser.go @@ -22,9 +22,6 @@ import ( "github.com/cockroachdb/errors" "github.com/samber/lo" - "github.com/twpayne/go-geom/encoding/wkb" - "github.com/twpayne/go-geom/encoding/wkbcommon" - "github.com/twpayne/go-geom/encoding/wkt" "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" "github.com/milvus-io/milvus/internal/json" @@ -428,11 +425,7 @@ func (r *rowParser) parseEntity(fieldID int64, obj any) (any, error) { return nil, r.wrapTypeError(obj, fieldID) } - geomT, err := wkt.Unmarshal(wktValue) - if err != nil { - return nil, r.wrapTypeError(wktValue, fieldID) - } - wkbValue, err := wkb.Marshal(geomT, wkb.NDR, wkbcommon.WKBOptionEmptyPointHandling(wkbcommon.EmptyPointHandlingNaN)) + wkbValue, err := pkgcommon.ConvertWKTToWKB(wktValue) if err != nil { return nil, r.wrapTypeError(wktValue, fieldID) } diff --git a/internal/util/importutilv2/numpy/field_reader.go b/internal/util/importutilv2/numpy/field_reader.go index d26beed31d..b785c286a5 100644 --- a/internal/util/importutilv2/numpy/field_reader.go +++ b/internal/util/importutilv2/numpy/field_reader.go @@ -30,6 +30,7 @@ import ( "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" "github.com/milvus-io/milvus/internal/json" "github.com/milvus-io/milvus/internal/util/importutilv2/common" + pkgcommon "github.com/milvus-io/milvus/pkg/v2/common" "github.com/milvus-io/milvus/pkg/v2/util/merr" "github.com/milvus-io/milvus/pkg/v2/util/parameterutil" "github.com/milvus-io/milvus/pkg/v2/util/typeutil" @@ -170,6 +171,22 @@ func (c *FieldReader) Next(count int64) (any, error) { if err != nil { return nil, err } + case schemapb.DataType_Geometry: + var strs []string + strs, err = c.ReadString(readCount) + if err != nil { + return nil, err + } + byteArr := make([][]byte, 0) + for _, wktValue := range strs { + wkbValue, err := pkgcommon.ConvertWKTToWKB(wktValue) + if err != nil { + return nil, err + } + byteArr = append(byteArr, wkbValue) + } + data = byteArr + c.readPosition += int(readCount) case schemapb.DataType_JSON: var strs []string strs, err = c.ReadString(readCount) diff --git a/internal/util/importutilv2/numpy/util.go b/internal/util/importutilv2/numpy/util.go index 889ee18994..49343c6959 100644 --- a/internal/util/importutilv2/numpy/util.go +++ b/internal/util/importutilv2/numpy/util.go @@ -223,7 +223,7 @@ func validateHeader(npyReader *npy.Reader, field *schemapb.FieldSchema, dim int) if shape[1] != dim/8 { return wrapDimError(shape[1]*8, dim, field) } - case schemapb.DataType_VarChar, schemapb.DataType_JSON: + case schemapb.DataType_VarChar, schemapb.DataType_JSON, schemapb.DataType_Geometry: if len(shape) != 1 { return wrapShapeError(len(shape), 1, field) } diff --git a/internal/util/importutilv2/parquet/field_reader.go b/internal/util/importutilv2/parquet/field_reader.go index d309c66cae..1c97d3b0f7 100644 --- a/internal/util/importutilv2/parquet/field_reader.go +++ b/internal/util/importutilv2/parquet/field_reader.go @@ -25,9 +25,6 @@ import ( "github.com/apache/arrow/go/v17/parquet/pqarrow" "github.com/cockroachdb/errors" "github.com/samber/lo" - "github.com/twpayne/go-geom/encoding/wkb" - "github.com/twpayne/go-geom/encoding/wkbcommon" - "github.com/twpayne/go-geom/encoding/wkt" "golang.org/x/exp/constraints" "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" @@ -35,6 +32,7 @@ import ( "github.com/milvus-io/milvus/internal/storage" "github.com/milvus-io/milvus/internal/util/importutilv2/common" "github.com/milvus-io/milvus/internal/util/nullutil" + pkgcommon "github.com/milvus-io/milvus/pkg/v2/common" "github.com/milvus-io/milvus/pkg/v2/util/merr" "github.com/milvus-io/milvus/pkg/v2/util/parameterutil" "github.com/milvus-io/milvus/pkg/v2/util/typeutil" @@ -728,11 +726,7 @@ func ReadNullableGeometryData(pcr *FieldReader, count int64) (any, []bool, error wkbValues = append(wkbValues, []byte(nil)) continue } - geomT, err := wkt.Unmarshal(wktValue) - if err != nil { - return nil, nil, err - } - wkbValue, err := wkb.Marshal(geomT, wkb.NDR, wkbcommon.WKBOptionEmptyPointHandling(wkbcommon.EmptyPointHandlingNaN)) + wkbValue, err := pkgcommon.ConvertWKTToWKB(wktValue) if err != nil { return nil, nil, err } @@ -753,11 +747,7 @@ func ReadGeometryData(pcr *FieldReader, count int64) (any, error) { wkbValues := make([][]byte, 0) for _, wktValue := range data.([]string) { - geomT, err := wkt.Unmarshal(wktValue) - if err != nil { - return nil, err - } - wkbValue, err := wkb.Marshal(geomT, wkb.NDR, wkbcommon.WKBOptionEmptyPointHandling(wkbcommon.EmptyPointHandlingNaN)) + wkbValue, err := pkgcommon.ConvertWKTToWKB(wktValue) if err != nil { return nil, err } diff --git a/internal/util/nullutil/nullutil.go b/internal/util/nullutil/nullutil.go index a14b969517..0d7d3a95fb 100644 --- a/internal/util/nullutil/nullutil.go +++ b/internal/util/nullutil/nullutil.go @@ -20,6 +20,7 @@ import ( "fmt" "github.com/milvus-io/milvus-proto/go-api/v2/schemapb" + "github.com/milvus-io/milvus/pkg/v2/common" "github.com/milvus-io/milvus/pkg/v2/util/merr" ) @@ -55,6 +56,8 @@ func GetDefaultValue(field *schemapb.FieldSchema) (any, error) { return field.GetDefaultValue().GetDoubleData(), nil case schemapb.DataType_String, schemapb.DataType_VarChar: return field.GetDefaultValue().GetStringData(), nil + case schemapb.DataType_Geometry: + return common.ConvertWKTToWKB(field.GetDefaultValue().GetStringData()) default: msg := fmt.Sprintf("type (%s) not support default_value", field.GetDataType().String()) return nil, merr.WrapErrParameterInvalidMsg(msg) diff --git a/pkg/common/common.go b/pkg/common/common.go index f2f0706b02..0cbd3ba033 100644 --- a/pkg/common/common.go +++ b/pkg/common/common.go @@ -25,6 +25,9 @@ import ( "github.com/cockroachdb/errors" "github.com/samber/lo" + "github.com/twpayne/go-geom/encoding/wkb" + "github.com/twpayne/go-geom/encoding/wkbcommon" + "github.com/twpayne/go-geom/encoding/wkt" "go.uber.org/zap" "github.com/milvus-io/milvus-proto/go-api/v2/commonpb" @@ -482,3 +485,11 @@ func IsAllowInsertAutoID(kvs ...*commonpb.KeyValuePair) (bool, bool) { } return false, false } + +func ConvertWKTToWKB(wktStr string) ([]byte, error) { + geomT, err := wkt.Unmarshal(wktStr) + if err != nil { + return nil, err + } + return wkb.Marshal(geomT, wkb.NDR, wkbcommon.WKBOptionEmptyPointHandling(wkbcommon.EmptyPointHandlingNaN)) +}