planner: optimize duplicate removal for common-handle batch point get

Project: pingcap/tidb | Language: Go | Concepts: algorithmic-optimization, data-structures | Difficulty: mid
Link: https://github.com/pingcap/tidb/pull/70389

The problem

When executing batch point-get queries on tables with common handles (composite primary keys), the system needs to deduplicate input values. The previous implementation used repeated in-place deletions from slices, which is inefficient and error-prone, especially when handling large numbers of values or many duplicates/nulls.

Why it's non-trivial

A naive fix might try to remove duplicates and nulls by deleting elements from slices while iterating, but this can easily cause index errors, missed elements, or inefficient O(n^2) behavior. The code must also keep related slices (like values and partition indexes) in sync and avoid leaking references to discarded data.

The approach

This PR rewrites the value and handle pruning logic to rebuild filtered slices in-place, avoiding repeated deletions and ensuring alignment. Here are the key changes:

1. Rebuilding filtered slices instead of deleting in-place

Instead of deleting elements while iterating, the code now appends valid entries to a filtered slice (using Go's slice re-slicing idiom):

-		usedValues := make([]bool, len(p.IndexValues))
-		for i, value := range p.IndexValues {
+		filteredValues := p.IndexValues[:0]
+		for _, value := range p.IndexValues {
 			if types.DatumsContainNull(value) {
 				continue
 			}
 			handle, err := kv.NewCommonHandle(value)
 			if err != nil {
 				return nil, false, err
 			}
 			if dedup.Exist(handle) {
 				continue
 			}
 			dedup.Set(handle, true)
 			handles = append(handles, handle)
-			usedValues[i] = true
+			filteredValues = append(filteredValues, value)
 		}
-		skipped := 0
-		for i, use := range usedValues {
-			if !use {
-				curr := i - skipped
-				p.IndexValues = slices.Delete(p.IndexValues, curr, curr+1)
-				skipped++
-			}
-		}
+		clear(p.IndexValues[len(filteredValues):])
+		p.IndexValues = filteredValues

This code efficiently compacts the slice by reusing its memory and only keeping valid, deduplicated values. The clear call ensures that any leftover elements are zeroed, preventing memory leaks.

2. Partition pruning also uses filtered slices

When partition pruning is needed, the same approach is used to keep the filtered values and handles in sync:

-			skipped := 0
+			filteredHandles := handles[:0]
+			filteredValues = p.IndexValues[:0]
 			partitionsFound := 0
 			for i, idx := range partIdxs {
-				if partIdxs[i] < 0 ||
+				if idx < 0 ||
 					(p.SinglePartition &&
-						partIdxs[i] != p.PartitionIdxs[0]) ||
+						idx != p.PartitionIdxs[0]) ||
 					!isInExplicitPartitions(pi, idx, p.PartitionNames) {
-					curr := i - skipped
-					handles = slices.Delete(handles, curr, curr+1)
-					p.IndexValues = slices.Delete(p.IndexValues, curr, curr+1)
-					skipped++
 					continue
-				} else if !p.SinglePartition {
+				}
+				filteredHandles = append(filteredHandles, handles[i])
+				filteredValues = append(filteredValues, p.IndexValues[i])
+				if !p.SinglePartition {
 					p.PartitionIdxs = append(p.PartitionIdxs, idx)
 				}
 				partitionsFound++
 			}
+			clear(handles[len(filteredHandles):])
+			clear(p.IndexValues[len(filteredValues):])
+			handles = filteredHandles
+			p.IndexValues = filteredValues

This ensures that after pruning, all slices are aligned and contain only the relevant entries.

3. Test coverage for edge cases

A new test file covers cases like empty input, all duplicates, all nulls, and mixed duplicates/nulls:

func TestPruneCommonHandleDuplicateValues(t *testing.T) {
    // ...
    testCases := []struct {
        name     string
        input    [][]types.Datum
        expected []string
    }{
        {name: "empty", input: nil, expected: nil},
        {name: "no duplicates", input: indexValues("a", "b", "c"), expected: []string{"a", "b", "c"}},
        {name: "duplicates and null", input: indexValues("b", "a", "b", nil, "c", "a"), expected: []string{"b", "a", "c"}},
        {name: "all duplicates", input: indexValues("a", "a", "a"), expected: []string{"a"}},
        {name: "only nulls", input: indexValues(nil, nil), expected: nil},
    }
    // ...
}

This ensures the new logic works correctly for all relevant scenarios.

From the review thread

No notable review discussion on this PR.

Takeaways

• When filtering or deduplicating slices, rebuild the slice in-place using a filtered view rather than deleting elements while iterating. This avoids index errors and is more efficient.
• Always clear unused slice elements after compaction to prevent memory leaks or unintended data retention.
• Keep related data structures (like parallel slices) in sync when filtering, and validate their alignment post-processing.