ivfflat memory escape to heap by cpegeric · Pull Request #23826 · matrixorigin/matrixone

cpegeric · 2026-03-09T10:48:00Z

What type of PR is this?

Which issue(s) this PR fixes:

What this PR does / why we need it:

fix the memory escape to heap in kmeans, brute force index
use sample() in ivf_create
remove get slow centroid stats in ivf search

qodo-code-review · 2026-03-09T10:48:21Z

ⓘ You are approaching your monthly quota for Qodo. Upgrade your plan

Review Summary by Qodo

🐞 Bug fix

Walkthroughs

Description

• Fix memory escape to heap in kmeans and brute force index
• Use object pools for temporary allocations in vector operations
• Replace RAND() with SAMPLE() function for IVF index creation
• Remove centroid statistics tracking and small centroid threshold logic
• Migrate to math/rand/v2 and use stack-allocated temporary buffers

Diagram

flowchart LR
  A["Memory Allocations"] -->|"Use sync.Pool"| B["Brute Force Index"]
  A -->|"Use sync.Pool"| C["Product L2 Operator"]
  A -->|"Use malloc allocator"| D["KMeans Clustering"]
  E["IVF Index Creation"] -->|"Replace RAND with SAMPLE"| F["Sampling"]
  G["IVF Search"] -->|"Remove centroid stats"| H["Simplified Logic"]
  I["Random Generation"] -->|"Migrate to math/rand/v2"| J["New RNG API"]

File Changes

1. pkg/common/util/unsafe.go ✨ Enhancement +5/-0

Add UnsafeSizeOf generic helper function

pkg/common/util/unsafe.go

2. pkg/frontend/variables.go ⚙️ Configuration changes +0/-8

Remove ivf_small_centroid_threshold system variable

pkg/frontend/variables.go

3. pkg/sql/colexec/productl2/product_l2.go 🐞 Bug fix +104/-40

Implement object pools for vector allocations

pkg/sql/colexec/productl2/product_l2.go

View more (11)

4. pkg/sql/colexec/productl2/types.go 🐞 Bug fix +8/-0

Add pool-based fields for center and null vectors

pkg/sql/colexec/productl2/types.go

5. pkg/sql/colexec/table_function/ivf_create.go ✨ Enhancement +11/-2

Replace RAND with SAMPLE and add logging

pkg/sql/colexec/table_function/ivf_create.go

6. pkg/vectorindex/brute_force/brute_force.go 🐞 Bug fix +102/-13

Implement sync.Pool for dataset and result allocations

pkg/vectorindex/brute_force/brute_force.go

7. pkg/vectorindex/brute_force/cpu.go ✨ Enhancement +2/-1

Add nthread parameter to NewBruteForceIndex

pkg/vectorindex/brute_force/cpu.go

8. pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go 🐞 Bug fix +93/-23

Use malloc allocator and pre-allocated buffers in kmeans

pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go

9. pkg/vectorindex/ivfflat/kmeans/elkans/clusterer_bench_test.go 🧪 Tests +2/-2

Update to math/rand/v2 API

pkg/vectorindex/ivfflat/kmeans/elkans/clusterer_bench_test.go

10. pkg/vectorindex/ivfflat/kmeans/elkans/clusterer_test.go 🧪 Tests +17/-9

Update tests for new random and allocation behavior

pkg/vectorindex/ivfflat/kmeans/elkans/clusterer_test.go

11. pkg/vectorindex/ivfflat/kmeans/elkans/initializer.go 🐞 Bug fix +9/-11

Migrate to math/rand/v2 and remove stored rand field

pkg/vectorindex/ivfflat/kmeans/elkans/initializer.go

12. pkg/vectorindex/ivfflat/kmeans/elkans/initializer_test.go 🧪 Tests +5/-8

Update expected values for new random initialization

pkg/vectorindex/ivfflat/kmeans/elkans/initializer_test.go

13. pkg/vectorindex/ivfflat/search.go 🐞 Bug fix +23/-94

Remove centroid stats and simplify bloom filter logic

pkg/vectorindex/ivfflat/search.go

14. pkg/vectorindex/ivfflat/search_test.go 🧪 Tests +0/-55

Remove tests for merged centroids functionality

pkg/vectorindex/ivfflat/search_test.go

qodo-code-review · 2026-03-09T10:48:23Z

Code Review by Qodo

🐞 Bugs (5) 📘 Rule violations (0) 📎 Requirement gaps (0)

1. Centroid scratch aliases dataset 🐞 Bug ✓ Correctness

Description

In Elkan kmeans, km.nextCentroids is swapped with km.centroids and then reused as the mutable
scratch buffer, but InitCentroids returns centroid slices that alias km.vectorList entries. This
makes recalculateCentroids overwrite the original training vectors in later iterations, yielding
incorrect centroids and corrupting the input dataset.

Code

pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go[R287-292]

+		newCentroids := km.recalculateCentroids(ctx, rnd, km.nextCentroids, km.membersCount) // step 4

-		km.updateBounds(ctx, newCentroids) // step 5 and 6
+		km.updateBounds(ctx, newCentroids, km.centroidShiftDist) // step 5 and 6

-		km.centroids = newCentroids // step 7
+		km.centroids, km.nextCentroids = newCentroids, km.centroids // step 7

Evidence
elkansCluster computes newCentroids into km.nextCentroids and then swaps km.nextCentroids to the
previous km.centroids; because initial km.centroids elements are direct references to km.vectorList
vectors, the next iteration’s recalculateCentroids zeroing/writes will mutate those underlying
vectors.
pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go[275-292]
pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go[552-560]
pkg/vectorindex/ivfflat/kmeans/elkans/initializer.go[44-63]
pkg/vectorindex/ivfflat/kmeans/elkans/initializer.go[93-100]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

### Issue description
`InitCentroids` returns centroid slices that alias `km.vectorList` vectors. After the first iteration, `elkansCluster` swaps `km.nextCentroids` with `km.centroids`, causing the next call to `recalculateCentroids` to use those aliased slices as the writable scratch buffer and overwrite the training data.

### Issue Context
`recalculateCentroids` explicitly zeroes and accumulates into the provided `newCentroids` buffer.

### Fix Focus Areas
- pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go[275-292]
- pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go[552-560]
- pkg/vectorindex/ivfflat/kmeans/elkans/initializer.go[44-63]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

2. Go slice stored in C 🐞 Bug ⛯ Reliability

Description

NewKMeans allocates km.nextCentroids backing storage via a C allocator and unsafe slice casts, but
recalculateCentroids can assign a Go-heap slice (randVector) into newCentroids[c] for empty
clusters. This stores an untracked Go pointer in non-Go memory, risking GC use-after-free and memory
corruption.

Code

pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go[R578-583]

			randVector := make([]T, len(km.vectorList[0]))
			for l := range randVector {
-				randVector[l] = T(km.rand.Float32())
+				randVector[l] = T(rnd.Float32())
			}
			newCentroids[c] = randVector

Evidence
nextCentroids is backed by memory returned from malloc.NewCAllocator().Allocate(...) and cast to
[][]T. Writing a Go-allocated slice header (randVector) into that non-Go backing array makes the
GC unaware of the reference to Go heap memory.
pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go[107-173]
pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go[548-583]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

### Issue description
Empty-cluster handling assigns a Go-heap slice (`randVector`) into `newCentroids[c]`, but `newCentroids` can be backed by C-allocated memory. This stores a Go pointer in non-Go memory and can lead to GC unsafety and memory corruption.

### Issue Context
`nextCentroids` is built from `malloc.NewCAllocator().Allocate` memory and converted with `UnsafeSliceCastToLength`.

### Fix Focus Areas
- pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go[107-173]
- pkg/vectorindex/ivfflat/kmeans/elkans/clusterer.go[571-588]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

3. GPU build API mismatch 🐞 Bug ✓ Correctness

Description

The !gpu brute_force.NewBruteForceIndex signature was changed to add an nthread parameter and call
sites were updated, but the gpu-tagged implementation still uses the old 4-argument signature.
Building with -tags gpu will fail due to this signature mismatch.

Code

pkg/vectorindex/brute_force/cpu.go[R25-31]

func NewBruteForceIndex[T types.RealNumbers](dataset [][]T,
	dimension uint,
	m metric.MetricType,
-	elemsz uint) (cache.VectorIndexSearchIf, error) {
+	elemsz uint,
+	nthread uint) (cache.VectorIndexSearchIf, error) {

	return NewCpuBruteForceIndex[T](dataset, dimension, m, elemsz)

Evidence
The cpu build variant now requires 5 arguments, and non-tagged call sites (e.g., ivfflat/search.go)
pass 5 arguments; under -tags gpu, the gpu build variant defines a 4-argument function with the
same name, causing compile-time mismatch.
pkg/vectorindex/brute_force/cpu.go[25-32]
pkg/vectorindex/brute_force/gpu.go[45-49]
pkg/vectorindex/ivfflat/search.go[239-244]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

### Issue description
CPU and GPU build variants expose different `NewBruteForceIndex` function signatures, breaking `-tags gpu` builds.

### Issue Context
Non-tagged call sites now pass 5 args.

### Fix Focus Areas
- pkg/vectorindex/brute_force/cpu.go[25-32]
- pkg/vectorindex/brute_force/gpu.go[45-60]
- pkg/vectorindex/ivfflat/search.go[238-244]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

View more (1)

4. Stats count ignores version 🐞 Bug ✓ Correctness

Description

IvfflatSearchIndex.LoadStats now counts all rows in the entries table without filtering by
idx.Version. This makes Meta.DataSize inconsistent with version-scoped entry queries and can
mis-size bloom filters and other heuristics.

Code

pkg/vectorindex/ivfflat/search.go[R78-83]

+	logutil.Infof("IVFFLAT START: gets data size")
+	sql := fmt.Sprintf("SELECT COUNT(1) FROM `%s`.`%s`",
		tblcfg.DbName, tblcfg.EntriesTable,
-		catalog.SystemSI_IVFFLAT_TblCol_Entries_version,
-		idx.Version,
-		catalog.SystemSI_IVFFLAT_TblCol_Entries_id,
	)

	res, err := runSql(sqlproc, sql)

Evidence
Entries are clearly versioned (other queries include entries_version = idx.Version), but
LoadStats’ COUNT query no longer filters by version, so it overcounts when multiple versions exist.
pkg/vectorindex/ivfflat/search.go[78-95]
pkg/vectorindex/ivfflat/search.go[158-165]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

### Issue description
`LoadStats` counts entries across all versions, but subsequent operations are version-filtered, making `Meta.DataSize` incorrect.

### Issue Context
The entries table is versioned and queries elsewhere use `entries_version = idx.Version`.

### Fix Focus Areas
- pkg/vectorindex/ivfflat/search.go[72-97]
- pkg/vectorindex/ivfflat/search.go[154-166]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

5. Pools pin batch buffers 🐞 Bug ➹ Performance

Description

productl2 returns pooled [][]float32/[][]float64 slices to sync.Pool without clearing elements or
shrinking length, even though each element is a zero-copy slice over batch bytes. This can retain
references to large batch buffers in the pool and prevent timely memory reclamation.

Code

pkg/sql/colexec/productl2/product_l2.go[R290-304]

+	if ctr.centersF32 != nil {
+		put1D(&pool2DF32, ctr.centersF32)
+		ctr.centersF32 = nil
+	}
+	if ctr.centersF64 != nil {
+		put1D(&pool2DF64, ctr.centersF64)
+		ctr.centersF64 = nil
+	}
+	if ctr.nullvecF32 != nil {
+		put1D(&pool1DF32, ctr.nullvecF32)
+		ctr.nullvecF32 = nil
+	}
+	if ctr.nullvecF64 != nil {
+		put1D(&pool1DF64, ctr.nullvecF64)
+		ctr.nullvecF64 = nil

Evidence
Centroid/probe vectors come from types.BytesToArray, which is a zero-copy unsafe cast over the
input bytes. Those inner slices are stored into the pooled outer slice and then the outer slice
pointer is returned to sync.Pool as-is, keeping those references alive while the pooled object
remains reachable.
pkg/sql/colexec/productl2/product_l2.go[152-160]
pkg/sql/colexec/productl2/product_l2.go[226-248]
pkg/sql/colexec/productl2/product_l2.go[285-305]
pkg/container/types/encoding.go[62-72]

Agent prompt

The issue below was found during a code review. Follow the provided context and guidance below and implement a solution

### Issue description
Pooled outer slices keep references to zero-copy inner slices backed by batch bytes, pinning those buffers until overwritten.

### Issue Context
`BytesToArray`/`DecodeSlice` is a zero-copy unsafe cast, so retaining the resulting slices retains the underlying batch buffers.

### Fix Focus Areas
- pkg/sql/colexec/productl2/product_l2.go[135-173]
- pkg/sql/colexec/productl2/product_l2.go[219-248]
- pkg/sql/colexec/productl2/product_l2.go[285-306]
- pkg/container/types/encoding.go[62-72]

ⓘ Copy this prompt and use it to remediate the issue with your preferred AI generation tools

ⓘ The new review experience is currently in Beta. Learn more

This reverts commit 50564fb.

mergify · 2026-03-13T10:44:22Z

Merge Queue Status

✅ Entered queue — 2026-03-13 10:44 UTC · Rule: main
✅ Checks started · in-place
🚫 Left the queue — 2026-03-13 11:07 UTC · at 5a05e1497bce58deddaccd4b08c26275f71a6a2e

This pull request spent 23 minutes 49 seconds in the queue, with no time running CI.

Required conditions to merge

Reason

The merge conditions cannot be satisfied due to failing checks

Hint

You may have to fix your CI before adding the pull request to the queue again.
If you update this pull request, to fix the CI, it will automatically be requeued once the queue conditions match again.
If you think this was a flaky issue instead, you can requeue the pull request, without updating it, by posting a @mergifyio queue comment.

mergify · 2026-03-14T09:28:27Z

Merge Queue Status

✅ Entered queue — 2026-03-14 09:28 UTC · Rule: main
✅ Checks started · in-place
🚫 Left the queue — 2026-03-14 09:50 UTC · at d1e82709164b9c17c689a32e7bbed9d6d47e29bd

This pull request spent 21 minutes 58 seconds in the queue, with no time running CI.

Required conditions to merge

Reason

The merge conditions cannot be satisfied due to failing checks

Hint

You may have to fix your CI before adding the pull request to the queue again.
If you update this pull request, to fix the CI, it will automatically be requeued once the queue conditions match again.
If you think this was a flaky issue instead, you can requeue the pull request, without updating it, by posting a @mergifyio queue comment.

mergify · 2026-03-16T10:48:55Z

Merge Queue Status

✅ Entered queue — 2026-03-16 10:48 UTC · Rule: main
✅ Checks passed · in-place
✅ Merged — 2026-03-16 10:48 UTC · at dd54724e6ca4ef34432a1f06c58a17097b51872c

This pull request spent 13 seconds in the queue, with no time running CI.

Required conditions to merge

cherry pick

7d9fb3a

cpegeric requested review from XuPeng-SH, aunjgr, fengttt and ouyuanning as code owners March 9, 2026 10:48

cpegeric had a problem deploying to ci March 9, 2026 10:48 — with GitHub Actions Error

cpegeric temporarily deployed to ci March 9, 2026 10:48 — with GitHub Actions Inactive

cpegeric had a problem deploying to ci March 9, 2026 10:48 — with GitHub Actions Error

cpegeric temporarily deployed to ci March 9, 2026 10:48 — with GitHub Actions Inactive

matrix-meow added the size/L Denotes a PR that changes [500,999] lines label Mar 9, 2026

Merge branch 'main' into ivf_escapeheap

09d02c5

mergify Bot added the kind/bug Something isn't working label Mar 9, 2026

mergify Bot had a problem deploying to ci March 9, 2026 10:49 Failure

mergify Bot had a problem deploying to ci March 9, 2026 10:49 Error

mergify Bot temporarily deployed to ci March 9, 2026 10:49 Inactive

mergify Bot had a problem deploying to ci March 9, 2026 10:49 Failure

mergify Bot temporarily deployed to ci March 9, 2026 10:49 Inactive

mergify Bot had a problem deploying to ci March 9, 2026 10:49 Error

cpegeric had a problem deploying to ci March 9, 2026 12:18 — with GitHub Actions Error

cpegeric added 11 commits March 9, 2026 12:27

go fmt

7cc5bb3

sca

caa06ba

revise test

deb4202

default to use go brute force index

473642f

remove partial

84ceb5f

go fmt

e09b50f

cleanup malloc

4b8cc4c

fast max heap

88fc7d1

go fmt

cce3bd0

go fmt

7d3de29

Merge branch 'main' into ivf_escapeheap

43751b2

ouyuanning approved these changes Mar 11, 2026

View reviewed changes

heni02 approved these changes Mar 11, 2026

View reviewed changes

aunjgr approved these changes Mar 11, 2026

View reviewed changes

cpegeric added 4 commits March 11, 2026 15:42

split by dataset

50564fb

Revert "split by dataset"

f996a50

This reverts commit 50564fb.

return error instead of panic when OOM

01edfc2

set deallocators to null after deallocate all memories

fae851d

XuPeng-SH approved these changes Mar 12, 2026

View reviewed changes

fengttt approved these changes Mar 12, 2026

View reviewed changes

Merge branch 'main' into ivf_escapeheap

5a05e14

cpegeric and others added 2 commits March 13, 2026 16:59

error message and better malloc

0edeb56

Merge branch 'main' into ivf_escapeheap

d1e8270

Merge branch 'main' into ivf_escapeheap

dd54724

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

ivfflat memory escape to heap#23826

ivfflat memory escape to heap#23826
mergify[bot] merged 26 commits intomatrixorigin:mainfrom
cpegeric:ivf_escapeheap

cpegeric commented Mar 9, 2026

Uh oh!

qodo-code-review Bot commented Mar 9, 2026

Uh oh!

qodo-code-review Bot commented Mar 9, 2026 •

edited

Loading

Uh oh!

mergify Bot commented Mar 13, 2026 •

edited

Loading

Uh oh!

mergify Bot commented Mar 14, 2026 •

edited

Loading

Uh oh!

mergify Bot commented Mar 16, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Conversation

cpegeric commented Mar 9, 2026

What type of PR is this?

Which issue(s) this PR fixes:

What this PR does / why we need it:

Uh oh!

qodo-code-review Bot commented Mar 9, 2026

Review Summary by Qodo

Walkthroughs

File Changes

Uh oh!

qodo-code-review Bot commented Mar 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review by Qodo

Uh oh!

mergify Bot commented Mar 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge Queue Status

Reason

Hint

Uh oh!

mergify Bot commented Mar 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge Queue Status

Reason

Hint

Uh oh!

mergify Bot commented Mar 16, 2026

Merge Queue Status

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

qodo-code-review Bot commented Mar 9, 2026 •

edited

Loading

mergify Bot commented Mar 13, 2026 •

edited

Loading

mergify Bot commented Mar 14, 2026 •

edited

Loading