fix: clear import segments NumOfRows when import task retries (#49986)

issue: #49985

## Summary

When a DataNode running an import task fails or returns `Retry`,
DataCoord resets the task to `Pending` and reschedules it on another
node. The preallocated segment IDs are reused on retry, but
`PickSegment` on the DataNode picks **randomly** whenever more than one
segment exists for a `(vchannel, partition)`
(https://github.com/milvus-io/milvus/blob/master/internal/datanode/importv2/util.go#L150-L162),
so the retried attempt can write to a different subset of the
preallocated segments than the failed attempt did.

`NumOfRows` updates from progress reports were monotonic-only and were
never cleared across attempts. As a result, segments touched by the
failed attempt but skipped by the retried attempt kept stale `NumOfRows`
without any insert binlogs being persisted (DataCoord only persists
binlogs on `Completed`,
https://github.com/milvus-io/milvus/blob/master/internal/datacoord/import_task_import.go#L209-L234).
Sort compaction later read this stale `NumOfRows` via
`originSegment.GetNumOfRows()`
(https://github.com/milvus-io/milvus/blob/master/internal/datacoord/import_util.go#L899),
found zero binlog rows for the orphan segment, and failed with
`unexpected row count` or EOF
(https://github.com/milvus-io/milvus/blob/master/internal/datanode/compactor/sort_compaction.go#L335-L343).

## Fix

- New operator `ResetImportingSegmentRows(segmentIDs ...int64)` in
`internal/datacoord/meta.go` — clears `NumOfRows` and `MaxRowNum` only
on segments still in `Importing` state.
- In `(*importTask).QueryTaskOnWorker`, call it on every preallocated
segment **before** transitioning the task back to `Pending`. If the
reset fails, the state transition is skipped so the next polling tick
retries the whole sequence.
- Orphan segments touched by the failed attempt but skipped by the
retried attempt now have `NumOfRows = 0` after the retry completes, so
`createSortCompactionTask`'s zero-rows guard drops them instead of
feeding stale counts into sort compaction.

## Test plan

- [x]
`TestImportTask_QueryTaskOnWorker/QueryImport_rpc_failed_resets_NumOfRows`
— new subtest verifies that `NumOfRows` on every preallocated segment is
reset to 0 and the task transitions to `Pending` when QueryImport
returns an error.
- [x] Existing `TestImportTask_QueryTaskOnWorker` subtests still pass
(the existing `QueryImport rpc failed` subtest's meta was extended with
`segments: NewSegmentsInfo()` so the reset code path runs without
nil-deref).
- [x] `golangci-lint run ./internal/datacoord/` is clean.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

internal/datacoord/import_task_import.go

@@ -163,6 +163,18 @@ func (t *importTask) QueryTaskOnWorker(cluster session.Cluster) {
 	}
 	resp, err := cluster.QueryImport(t.GetNodeID(), req)
 	if err != nil || resp.GetState() == datapb.ImportTaskStateV2_Retry {
+		// Clear partial progress recorded from the failed attempt. Otherwise,
+		// if the retried attempt's PickSegment lands on a different subset of
+		// the preallocated segments, the segments it skips will keep stale
+		// NumOfRows without insert binlogs — causing sort compaction to fail
+		// with "unexpected row count" or EOF.
+		if segmentIDs := t.GetSegmentIDs(); len(segmentIDs) > 0 {
+			if resetErr := t.meta.UpdateSegmentsInfo(context.TODO(), ResetImportingSegmentRows(segmentIDs...)); resetErr != nil {
+				log.Warn("failed to reset import segment row counts on retry",
+					WrapTaskLog(t, zap.Error(resetErr))...)
+				return
+			}
+		}
 		updateErr := t.importMeta.UpdateTask(context.TODO(), t.GetTaskID(), UpdateState(datapb.ImportTaskStateV2_Pending))
 		if updateErr != nil {
 			log.Warn("failed to update import task state to pending", WrapTaskLog(t, zap.Error(updateErr))...)

internal/datacoord/import_task_import_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 
+	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
 	"github.com/milvus-io/milvus/internal/datacoord/allocator"
 	"github.com/milvus-io/milvus/internal/datacoord/session"
 	"github.com/milvus-io/milvus/internal/metastore/mocks"
@@ -271,8 +272,11 @@ func TestImportTask_QueryTaskOnWorker(t *testing.T) {
 			State:        datapb.ImportTaskStateV2_InProgress,
 		}
 		task := &importTask{
-			alloc:      nil,
-			meta:       &meta{collections: typeutil.NewConcurrentMap[UniqueID, *collectionInfo]()},
+			alloc: nil,
+			meta: &meta{
+				collections: typeutil.NewConcurrentMap[UniqueID, *collectionInfo](),
+				segments:    NewSegmentsInfo(),
+			},
 			importMeta: im,
 			tr:         timerecord.NewTimeRecorder(""),
 		}
@@ -289,6 +293,66 @@ func TestImportTask_QueryTaskOnWorker(t *testing.T) {
 		assert.Equal(t, datapb.ImportTaskStateV2_InProgress, task.GetState())
 	})
 
+	t.Run("QueryImport rpc failed resets NumOfRows", func(t *testing.T) {
+		catalog := mocks.NewDataCoordCatalog(t)
+		catalog.EXPECT().ListImportJobs(mock.Anything).Return(nil, nil)
+		catalog.EXPECT().ListPreImportTasks(mock.Anything).Return(nil, nil)
+		catalog.EXPECT().ListImportTasks(mock.Anything).Return(nil, nil)
+		catalog.EXPECT().SaveImportTask(mock.Anything, mock.Anything).Return(nil)
+
+		im, err := NewImportMeta(context.TODO(), catalog, nil, nil)
+		assert.NoError(t, err)
+
+		taskProto := &datapb.ImportTaskV2{
+			JobID:        1,
+			TaskID:       2,
+			CollectionID: 3,
+			SegmentIDs:   []int64{5, 6},
+			NodeID:       7,
+			State:        datapb.ImportTaskStateV2_InProgress,
+		}
+		segCatalog := mocks.NewDataCoordCatalog(t)
+		segCatalog.EXPECT().AlterSegments(mock.Anything, mock.Anything).Return(nil)
+		task := &importTask{
+			alloc: nil,
+			meta: &meta{
+				catalog:     segCatalog,
+				collections: typeutil.NewConcurrentMap[UniqueID, *collectionInfo](),
+				segments:    NewSegmentsInfo(),
+			},
+			importMeta: im,
+			tr:         timerecord.NewTimeRecorder(""),
+		}
+		task.task.Store(taskProto)
+		err = im.AddTask(context.TODO(), task)
+		assert.NoError(t, err)
+
+		task.meta.segments.SetSegment(5, &SegmentInfo{
+			SegmentInfo: &datapb.SegmentInfo{
+				ID:        5,
+				State:     commonpb.SegmentState_Importing,
+				NumOfRows: 100,
+				MaxRowNum: 100,
+			},
+		})
+		task.meta.segments.SetSegment(6, &SegmentInfo{
+			SegmentInfo: &datapb.SegmentInfo{
+				ID:        6,
+				State:     commonpb.SegmentState_Importing,
+				NumOfRows: 50,
+				MaxRowNum: 50,
+			},
+		})
+
+		cluster := session.NewMockCluster(t)
+		cluster.EXPECT().QueryImport(mock.Anything, mock.Anything).Return(nil, errors.New("mock err"))
+		task.QueryTaskOnWorker(cluster)
+
+		assert.Equal(t, datapb.ImportTaskStateV2_Pending, task.GetState())
+		assert.EqualValues(t, 0, task.meta.segments.GetSegment(5).GetNumOfRows())
+		assert.EqualValues(t, 0, task.meta.segments.GetSegment(6).GetNumOfRows())
+	})
+
 	t.Run("import failed", func(t *testing.T) {
 		catalog := mocks.NewDataCoordCatalog(t)
 		catalog.EXPECT().ListImportJobs(mock.Anything).Return(nil, nil)

internal/datacoord/meta.go

@@ -1399,6 +1399,35 @@ func UpdateImportedRows(segmentID int64, rows int64) UpdateOperator {
 	}
 }
 
+// ResetImportingSegmentRows clears NumOfRows and MaxRowNum on each given
+// segment that is still in the Importing state. Used to discard partial
+// progress reported by a failed import attempt before the task is rescheduled,
+// so segments the retried attempt skips do not keep stale row counts that
+// would break sort compaction.
+func ResetImportingSegmentRows(segmentIDs ...int64) UpdateOperator {
+	return func(modPack *updateSegmentPack) bool {
+		anyReset := false
+		for _, segmentID := range segmentIDs {
+			segment := modPack.Get(segmentID)
+			if segment == nil {
+				log.Ctx(context.TODO()).Warn("meta update: reset importing segment rows failed - segment not found",
+					zap.Int64("segmentID", segmentID))
+				continue
+			}
+			if segment.GetState() != commonpb.SegmentState_Importing {
+				log.Ctx(context.TODO()).Warn("meta update: reset importing segment rows skipped - segment not in Importing state",
+					zap.Int64("segmentID", segmentID),
+					zap.String("state", segment.GetState().String()))
+				continue
+			}
+			segment.NumOfRows = 0
+			segment.MaxRowNum = 0
+			anyReset = true
+		}
+		return anyReset
+	}
+}
+
 func UpdateIsImporting(segmentID int64, isImporting bool) UpdateOperator {
 	return func(modPack *updateSegmentPack) bool {
 		segment := modPack.Get(segmentID)