Add probe command for query function

models/collection.go

@@ -39,6 +39,24 @@ type CollectionHistory struct {
 	Dropped bool
 }
 
+func (c *Collection) GetPKField() (FieldSchema, bool) {
+	for _, field := range c.Schema.Fields {
+		if field.IsPrimaryKey {
+			return field, true
+		}
+	}
+	return FieldSchema{}, false
+}
+
+func (c *Collection) GetVectorField() (FieldSchema, bool) {
+	for _, field := range c.Schema.Fields {
+		if field.DataType == DataTypeBinaryVector || field.DataType == DataTypeFloatVector {
+			return field, true
+		}
+	}
+	return FieldSchema{}, false
+}
+
 // newCollectionFromBase fetchs common base information form proto objects
 func newCollectionFromBase[collectionBase interface {
 	GetID() int64

models/collection_schema.go

@@ -1,5 +1,10 @@
 package models
 
+import (
+	"errors"
+	"strconv"
+)
+
 type CollectionSchema struct {
 	Name   string
 	Fields []FieldSchema
@@ -15,6 +20,19 @@ type FieldSchema struct {
 	Properties   map[string]string
 }
 
+func (fs *FieldSchema) GetDim() (int64, error) {
+	if fs.DataType != DataTypeFloatVector && fs.DataType != DataTypeBinaryVector {
+		return -1, errors.New("field is not vector")
+	}
+	raw, ok := fs.Properties["dim"]
+	if !ok {
+		return -1, errors.New("field not has dim property")
+	}
+
+	dim, err := strconv.ParseInt(raw, 10, 64)
+	return dim, err
+}
+
 func newSchemaFromBase[schemaBase interface {
 	GetName() string
 }](schema schemaBase) CollectionSchema {