fix: Fix filtering by partition key fails for importing data

internal/proxy/impl.go

@@ -6100,7 +6100,11 @@
 				resp.Status = merr.Status(err)
 				return resp, nil
 			}
-			partitionIDs = lo.Values(partitions)
+			_, partitionIDs, err = typeutil.RearrangePartitionsForPartitionKey(partitions)
+			if err != nil {
+				resp.Status = merr.Status(err)
+				return resp, nil
+			}
 		} else {
 			if req.GetPartitionName() == "" {
 				req.PartitionName = Params.CommonCfg.DefaultPartitionName.GetValue()

internal/proxy/msg_pack.go

@@ -231,7 +231,7 @@ func repackInsertDataWithPartitionKey(ctx context.Context,
 	}
 
 	channel2RowOffsets := assignChannelsByPK(result.IDs, channelNames, insertMsg)
-	partitionNames, err := getDefaultPartitionNames(ctx, insertMsg.GetDbName(), insertMsg.CollectionName)
+	partitionNames, err := getDefaultPartitionsInPartitionKeyMode(ctx, insertMsg.GetDbName(), insertMsg.CollectionName)
 	if err != nil {
 		log.Warn("get default partition names failed in partition key mode",
 			zap.String("collectionName", insertMsg.CollectionName),

internal/proxy/task_test.go

@@ -3135,7 +3135,7 @@ func TestCreateCollectionTaskWithPartitionKey(t *testing.T) {
 		// check default partitions
 		err = InitMetaCache(ctx, rc, nil, nil)
 		assert.NoError(t, err)
-		partitionNames, err := getDefaultPartitionNames(ctx, "", task.CollectionName)
+		partitionNames, err := getDefaultPartitionsInPartitionKeyMode(ctx, "", task.CollectionName)
 		assert.NoError(t, err)
 		assert.Equal(t, task.GetNumPartitions(), int64(len(partitionNames)))
 

internal/proxy/util.go

@@ -1395,7 +1395,7 @@ func hasParitionKeyModeField(schema *schemapb.CollectionSchema) bool {
 	return false
 }
 
-// getDefaultPartitionNames only used in partition key mode
+// getDefaultPartitionsInPartitionKeyMode only used in partition key mode
 func getDefaultPartitionsInPartitionKeyMode(ctx context.Context, dbName string, collectionName string) ([]string, error) {
 	partitions, err := globalMetaCache.GetPartitions(ctx, dbName, collectionName)
 	if err != nil {
@@ -1411,32 +1411,6 @@ func getDefaultPartitionsInPartitionKeyMode(ctx context.Context, dbName string,
 	return partitionNames, nil
 }
 
-// getDefaultPartitionNames only used in partition key mode
-func getDefaultPartitionNames(ctx context.Context, dbName string, collectionName string) ([]string, error) {
-	partitions, err := globalMetaCache.GetPartitions(ctx, dbName, collectionName)
-	if err != nil {
-		return nil, err
-	}
-
-	// Make sure the order of the partition names got every time is the same
-	partitionNames := make([]string, len(partitions))
-	for partitionName := range partitions {
-		splits := strings.Split(partitionName, "_")
-		if len(splits) < 2 {
-			err = fmt.Errorf("bad default partion name in partition ket mode: %s", partitionName)
-			return nil, err
-		}
-		index, err := strconv.ParseInt(splits[len(splits)-1], 10, 64)
-		if err != nil {
-			return nil, err
-		}
-
-		partitionNames[index] = partitionName
-	}
-
-	return partitionNames, nil
-}
-
 func assignChannelsByPK(pks *schemapb.IDs, channelNames []string, insertMsg *msgstream.InsertMsg) map[string][]int {
 	insertMsg.HashValues = typeutil.HashPK2Channels(pks, channelNames)
 

tests/integration/import/partition_key_test.go

@@ -0,0 +1,215 @@
+// 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 importv2
+
+import (
+	"context"
+	"fmt"
+	"math/rand"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/golang/protobuf/proto"
+	"github.com/samber/lo"
+	"go.uber.org/zap"
+
+	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
+	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
+	"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
+	"github.com/milvus-io/milvus/internal/proto/internalpb"
+	"github.com/milvus-io/milvus/pkg/common"
+	"github.com/milvus-io/milvus/pkg/log"
+	"github.com/milvus-io/milvus/pkg/util/funcutil"
+	"github.com/milvus-io/milvus/pkg/util/merr"
+	"github.com/milvus-io/milvus/pkg/util/metric"
+	"github.com/milvus-io/milvus/tests/integration"
+)
+
+func (s *BulkInsertSuite) TestImportWithPartitionKey() {
+	const (
+		rowCount = 10000
+	)
+
+	c := s.Cluster
+	ctx, cancel := context.WithTimeout(c.GetContext(), 60*time.Second)
+	defer cancel()
+
+	collectionName := "TestBulkInsert_WithPartitionKey_" + funcutil.GenRandomStr()
+
+	schema := integration.ConstructSchema(collectionName, dim, true, &schemapb.FieldSchema{
+		FieldID:      100,
+		Name:         integration.Int64Field,
+		IsPrimaryKey: true,
+		DataType:     schemapb.DataType_Int64,
+		AutoID:       true,
+	}, &schemapb.FieldSchema{
+		FieldID:  101,
+		Name:     integration.FloatVecField,
+		DataType: schemapb.DataType_FloatVector,
+		TypeParams: []*commonpb.KeyValuePair{
+			{
+				Key:   common.DimKey,
+				Value: fmt.Sprintf("%d", dim),
+			},
+		},
+	}, &schemapb.FieldSchema{
+		FieldID:  102,
+		Name:     integration.VarCharField,
+		DataType: schemapb.DataType_VarChar,
+		TypeParams: []*commonpb.KeyValuePair{
+			{
+				Key:   common.MaxLengthKey,
+				Value: fmt.Sprintf("%d", 256),
+			},
+		},
+		IsPartitionKey: true,
+	})
+	marshaledSchema, err := proto.Marshal(schema)
+	s.NoError(err)
+
+	createCollectionStatus, err := c.Proxy.CreateCollection(ctx, &milvuspb.CreateCollectionRequest{
+		DbName:         "",
+		CollectionName: collectionName,
+		Schema:         marshaledSchema,
+		ShardsNum:      common.DefaultShardsNum,
+	})
+	s.NoError(err)
+	s.Equal(int32(0), createCollectionStatus.GetCode())
+
+	// create index
+	createIndexStatus, err := c.Proxy.CreateIndex(ctx, &milvuspb.CreateIndexRequest{
+		CollectionName: collectionName,
+		FieldName:      integration.FloatVecField,
+		IndexName:      "_default",
+		ExtraParams:    integration.ConstructIndexParam(dim, integration.IndexFaissIvfFlat, metric.L2),
+	})
+	s.NoError(err)
+	s.Equal(int32(0), createIndexStatus.GetCode())
+
+	s.WaitForIndexBuilt(ctx, collectionName, integration.FloatVecField)
+
+	// import
+	var files []*internalpb.ImportFile
+	err = os.MkdirAll(c.ChunkManager.RootPath(), os.ModePerm)
+	s.NoError(err)
+
+	filePath := fmt.Sprintf("/tmp/test_%d.parquet", rand.Int())
+	insertData, err := GenerateParquetFileAndReturnInsertData(filePath, schema, rowCount)
+	s.NoError(err)
+	defer os.Remove(filePath)
+	files = []*internalpb.ImportFile{
+		{
+			Paths: []string{
+				filePath,
+			},
+		},
+	}
+
+	importResp, err := c.Proxy.ImportV2(ctx, &internalpb.ImportRequest{
+		CollectionName: collectionName,
+		Files:          files,
+	})
+	s.NoError(err)
+	s.Equal(int32(0), importResp.GetStatus().GetCode())
+	log.Info("Import result", zap.Any("importResp", importResp))
+
+	jobID := importResp.GetJobID()
+	err = WaitForImportDone(ctx, c, jobID)
+	s.NoError(err)
+
+	// load
+	loadStatus, err := c.Proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
+		CollectionName: collectionName,
+	})
+	s.NoError(err)
+	s.Equal(commonpb.ErrorCode_Success, loadStatus.GetErrorCode())
+	s.WaitForLoad(ctx, collectionName)
+
+	segments, err := c.MetaWatcher.ShowSegments()
+	s.NoError(err)
+	s.NotEmpty(segments)
+	log.Info("Show segments", zap.Any("segments", segments))
+
+	// load refresh
+	loadStatus, err = c.Proxy.LoadCollection(ctx, &milvuspb.LoadCollectionRequest{
+		CollectionName: collectionName,
+		Refresh:        true,
+	})
+	s.NoError(err)
+	s.Equal(commonpb.ErrorCode_Success, loadStatus.GetErrorCode())
+	s.WaitForLoadRefresh(ctx, "", collectionName)
+
+	// query partition key, TermExpr
+	queryNum := 10
+	partitionKeyData := insertData.Data[int64(102)].GetRows().([]string)
+	queryData := partitionKeyData[:queryNum]
+	strs := lo.Map(queryData, func(str string, _ int) string {
+		return fmt.Sprintf("\"%s\"", str)
+	})
+	str := strings.Join(strs, `,`)
+	expr := fmt.Sprintf("%s in [%v]", integration.VarCharField, str)
+	queryResult, err := c.Proxy.Query(ctx, &milvuspb.QueryRequest{
+		CollectionName: collectionName,
+		Expr:           expr,
+		OutputFields:   []string{integration.VarCharField},
+	})
+	err = merr.CheckRPCCall(queryResult, err)
+	s.NoError(err)
+	for _, data := range queryResult.GetFieldsData() {
+		if data.GetType() == schemapb.DataType_VarChar {
+			resData := data.GetScalars().GetStringData().GetData()
+			s.Equal(queryNum, len(resData))
+			s.ElementsMatch(resData, queryData)
+		}
+	}
+
+	// query partition key, CmpOp 1
+	expr = fmt.Sprintf("%s >= 0", integration.Int64Field)
+	queryResult, err = c.Proxy.Query(ctx, &milvuspb.QueryRequest{
+		CollectionName: collectionName,
+		Expr:           expr,
+		OutputFields:   []string{integration.VarCharField},
+	})
+	err = merr.CheckRPCCall(queryResult, err)
+	s.NoError(err)
+	for _, data := range queryResult.GetFieldsData() {
+		if data.GetType() == schemapb.DataType_VarChar {
+			resData := data.GetScalars().GetStringData().GetData()
+			s.Equal(rowCount, len(resData))
+			s.ElementsMatch(resData, partitionKeyData)
+		}
+	}
+
+	// query partition key, CmpOp 2
+	target := partitionKeyData[rand.Intn(rowCount)]
+	expr = fmt.Sprintf("%s == \"%s\"", integration.VarCharField, target)
+	queryResult, err = c.Proxy.Query(ctx, &milvuspb.QueryRequest{
+		CollectionName: collectionName,
+		Expr:           expr,
+		OutputFields:   []string{integration.VarCharField},
+	})
+	err = merr.CheckRPCCall(queryResult, err)
+	s.NoError(err)
+	for _, data := range queryResult.GetFieldsData() {
+		if data.GetType() == schemapb.DataType_VarChar {
+			resData := data.GetScalars().GetStringData().GetData()
+			s.Equal(1, len(resData))
+			s.Equal(resData[0], target)
+		}
+	}
+}

tests/integration/import/util_test.go

@@ -46,33 +46,38 @@ import (
 const dim = 128
 
 func GenerateParquetFile(filePath string, schema *schemapb.CollectionSchema, numRows int) error {
+	_, err := GenerateParquetFileAndReturnInsertData(filePath, schema, numRows)
+	return err
+}
+
+func GenerateParquetFileAndReturnInsertData(filePath string, schema *schemapb.CollectionSchema, numRows int) (*storage.InsertData, error) {
 	w, err := os.OpenFile(filePath, os.O_RDWR|os.O_CREATE, 0o666)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	pqSchema, err := pq.ConvertToArrowSchema(schema)
 	if err != nil {
-		return err
+		return nil, err
 	}
 	fw, err := pqarrow.NewFileWriter(pqSchema, w, parquet.NewWriterProperties(parquet.WithMaxRowGroupLength(int64(numRows))), pqarrow.DefaultWriterProps())
 	if err != nil {
-		return err
+		return nil, err
 	}
 	defer fw.Close()
 
 	insertData, err := testutil.CreateInsertData(schema, numRows)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	columns, err := testutil.BuildArrayData(schema, insertData)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	recordBatch := array.NewRecord(pqSchema, columns, int64(numRows))
-	return fw.Write(recordBatch)
+	return insertData, fw.Write(recordBatch)
 }
 
 func GenerateNumpyFiles(cm storage.ChunkManager, schema *schemapb.CollectionSchema, rowCount int) (*internalpb.ImportFile, error) {