unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to OSSpedia
algorithmic-optimizationdata-structures

enhance: resolve the destination column by index in AppendFieldData

Project: milvus-io/milvus | Language: Go | Concepts: algorithmic-optimization, data-structures | Difficulty: mid
Link: https://github.com/milvus-io/milvus/pull/52270

The problem

The AppendFieldData function was inefficient: every time it appended a row, it rebuilt a map from field ID to destination column, even though the destination columns (dst) are usually in the same order as the source columns (src). This wasted CPU cycles, especially when appending many rows, because the mapping was recomputed for every single row.

Why it's non-trivial

A naive fix would just use index-based access (dst[i]), but that assumes dst and src are always in the same order and shape. In reality, some callers might have permuted or partially filled destination slices, so the code must still handle those cases safely—falling back to the map lookup only when necessary.

The approach

The new approach optimizes for the common case (parallel dst and src), but preserves correctness for edge cases. Here are the key changes:

1. Remove unconditional map creation.
Previously, the function always built a map from field ID to destination column, even if it wasn't needed:

-dstMap := make(map[int64]*schemapb.FieldData)
-for _, fieldData := range dst {
-	if fieldData != nil {
-		dstMap[fieldData.FieldId] = fieldData
-	}
-}

This code built a map for every row, even when dst and src were index-aligned.

2. Use fast-path index lookup when possible.
Now, for each source column, the code first checks if the destination column at the same index matches the field ID:

if i < len(dst) && dst[i] != nil && dst[i].FieldId == fieldData.FieldId {
    dstFieldData, ok = dst[i], true
} else {
    // fallback to map
}

This is the fast path: if dst and src are parallel, it avoids the map entirely.

3. Lazily build the map only if needed.
If the fast path fails (e.g., columns are permuted), the code builds the map just once and uses it for the remaining lookups:

if dstMap == nil {
    dstMap = make(map[int64]*schemapb.FieldData, len(dst))
    for _, fd := range dst {
        if fd != nil {
            dstMap[fd.FieldId] = fd
        }
    }
}
dstFieldData, ok = dstMap[fieldData.FieldId]

This ensures the map is only built if the index-based lookup fails, saving work in the common case.

4. Defensive bounds checking and panic for broken callers.
If a caller passes a dst that's too short, the code panics with a clear error message:

if i < len(dst) {
    dst[i] = dstFieldData
} else {
    panic(fmt.Sprintf("AppendFieldData: dst has %d columns, src has %d; "+
        "callers must size dst to at least len(src)", len(dst), len(src)))
}

This prevents silent data loss and makes contract violations obvious.

5. Comprehensive tests.
A new test file covers all scenarios:

  • Fast path (parallel columns)
  • Lazy creation (nil destination entries)
  • Permuted columns (forces map fallback)
  • Panic on too-short destination

Example test for permuted columns:

t.Run("dst permuted relative to src (forces the map fallback)", func(t *testing.T) {
    src := newSrc()
    prepared := PrepareResultFieldData(src, rows)
    // Reverse dst so index i no longer corresponds to src[i].
    dst := []*schemapb.FieldData{prepared[2], prepared[1], prepared[0]}
    for r := int64(0); r < rows; r++ {
        AppendFieldData(dst, src, r)
    }
    // Data must still land in the column with the matching FieldId.
    ...
})

From the review thread

No notable review discussion on this PR.

Takeaways

  • Optimize for the common case, but preserve correctness for edge cases. Fast paths should be used when assumptions hold, but code must gracefully handle less common input shapes.
  • Avoid redundant work inside tight loops. Building data structures (like maps) should be done only when truly needed, not on every iteration.
  • Fail fast and loudly on contract violations. Defensive programming (e.g., panicking with a clear message when an invariant is broken) helps surface bugs early and prevents silent data corruption.
  • Thoroughly test all code paths, including edge cases and error handling. Comprehensive tests ensure both performance gains and correctness are maintained.