Browse Source

Fix a bug for partition key (#503)

Signed-off-by: yhmo <yihua.mo@zilliz.com>
groot 2 years ago
parent
commit
2361ffcdf8

+ 10 - 1
src/main/java/io/milvus/param/ParamUtils.java

@@ -188,13 +188,17 @@ public class ParamUtils {
         MsgBase msgBase = MsgBase.newBuilder().setMsgType(MsgType.Insert).build();
         InsertRequest.Builder insertBuilder = InsertRequest.newBuilder()
                 .setCollectionName(collectionName)
-                .setPartitionName(partitionName)
                 .setBase(msgBase)
                 .setNumRows(requestParam.getRowCount());
 
         // gen fieldData
         // make sure the field order must be consisted with collection schema
+        boolean isPartitionKeyEnabled = false;
         for (FieldType fieldType : fieldTypes) {
+            if (fieldType.isPartitionKey()) {
+                isPartitionKeyEnabled = true;
+            }
+
             boolean found = false;
             for (InsertParam.Field field : fields) {
                 if (field.getName().equals(fieldType.getName())) {
@@ -216,6 +220,11 @@ public class ParamUtils {
             }
         }
 
+        // set partition name only when there is no partition key field
+        if (!isPartitionKeyEnabled) {
+            insertBuilder.setPartitionName(partitionName);
+        }
+
         // gen request
         return insertBuilder.build();
     }

+ 1 - 0
src/main/java/io/milvus/param/dml/InsertParam.java

@@ -73,6 +73,7 @@ public class InsertParam {
 
         /**
          * Set partition name (Optional).
+         * This partition name will be ignored if the collection has a partition key field.
          *
          * @param partitionName partition name
          * @return <code>Builder</code>

+ 29 - 0
src/main/java/io/milvus/response/DescCollResponseWrapper.java

@@ -114,6 +114,34 @@ public class DescCollResponseWrapper {
         return null;
     }
 
+    /**
+     * Get whether the collection dynamic field is enabled
+     *
+     * @return boolean
+     */
+    public boolean isDynamicFieldEnabled() {
+        CollectionSchema schema = response.getSchema();
+        return schema.getEnableDynamicField();
+    }
+
+    /**
+     * Get the partition key field.
+     * Return null if the partition key field doesn't exist.
+     *
+     * @return {@link FieldType} schema of the partition key field
+     */
+    public FieldType getPartitionKeyField() {
+        CollectionSchema schema = response.getSchema();
+        for (int i = 0; i < schema.getFieldsCount(); ++i) {
+            FieldSchema field = schema.getFields(i);
+            if (field.getIsPartitionKey()) {
+                return ParamUtils.ConvertField(field);
+            }
+        }
+
+        return null;
+    }
+
     /**
      * Construct a <code>String</code> by {@link DescCollResponseWrapper} instance.
      *
@@ -129,6 +157,7 @@ public class DescCollResponseWrapper {
                 ", createdUtcTimestamp:" + getCreatedUtcTimestamp() +
                 ", aliases:" + getAliases().toString() +
                 ", fields:" + getFields().toString() +
+                ", isDynamicFieldEnabled:" + isDynamicFieldEnabled() +
                 '}';
     }
 }