Merge pull request #22 from weaviate/improve_expand

improvement: introduce expandConditionally method

Author: antas-marcin

.github/workflows/go.yml

@@ -19,7 +19,7 @@ jobs:
     - name: Set up Go
       uses: actions/setup-go@v3
       with:
-        go-version: 1.19
+        go-version: 1.24
         cache: true
 
     - name: Test

bitmap.go

@@ -178,15 +178,12 @@ func (ra *Bitmap) expandKeys(bySize uint64) uint64 {
 	ra.scootRight(curSize, bySize)
 	ra.keys = uint16To64SliceUnsafe(ra.data[:curSize+bySize])
 	ra.keys.setNodeSize(int(curSize + bySize))
-	
+
 	// All containers have moved to the right by bySize bytes.
 	// Update their offsets.
 	n := ra.keys
-	for i := 0; i < n.maxKeys(); i++ {
-		val := n.val(i)
-		if val > 0 {
-			n.setAt(valOffset(i), val+uint64(bySize))
-		}
+	for i := 0; i < n.numKeys(); i++ {
+		n.setAt(valOffset(i), n.val(i)+uint64(bySize))
 	}
 	return bySize
 }

bitmap_opt.go

@@ -130,7 +130,7 @@ func andNotContainers(a, b, res *Bitmap, optBuf []uint16) {
 	bi, bn := 0, b.keys.numKeys()
 
 	akToAc := map[uint64][]uint16{}
-	sizeContainers := uint64(0)
+	sizeContainers := 0
 	newKeys := 0
 
 	for ai < an && bi < bn {
@@ -154,7 +154,7 @@ func andNotContainers(a, b, res *Bitmap, optBuf []uint16) {
 			ac := a.getContainer(off)
 			if getCardinality(ac) > 0 {
 				akToAc[ak] = ac
-				sizeContainers += uint64(len(ac))
+				sizeContainers += len(ac)
 				newKeys++
 			}
 			ai++
@@ -168,18 +168,13 @@ func andNotContainers(a, b, res *Bitmap, optBuf []uint16) {
 		if getCardinality(ac) > 0 {
 			ak := a.keys.key(ai)
 			akToAc[ak] = ac
-			sizeContainers += uint64(len(ac))
+			sizeContainers += len(ac)
 			newKeys++
 		}
 	}
 
 	if sizeContainers > 0 {
-		// 2x uint64 = 8x uint16; for key and offset; twice as much as actually needed
-		sizeKeys := 8 * uint64(newKeys) * 2
-		// ensure enough space for new containers and keys,
-		// allocate required memory just once to avoid copying underlying data slice multiple times
-		res.expandNoLengthChange(sizeContainers + sizeKeys)
-		res.expandKeys(sizeKeys)
+		res.expandConditionally(newKeys, sizeContainers)
 
 		for ak, ac := range akToAc {
 			// create a new container and update the key offset to this container.
@@ -290,7 +285,7 @@ func orContainers(a, b, res *Bitmap, buf []uint16) {
 
 	akToAc := map[uint64][]uint16{}
 	bkToBc := map[uint64][]uint16{}
-	sizeContainers := uint64(0)
+	sizeContainers := 0
 	newKeys := 0
 
 	for ai < an && bi < bn {
@@ -318,7 +313,7 @@ func orContainers(a, b, res *Bitmap, buf []uint16) {
 			ac := a.getContainer(off)
 			if getCardinality(ac) > 0 {
 				akToAc[ak] = ac
-				sizeContainers += uint64(len(ac))
+				sizeContainers += len(ac)
 				newKeys++
 			}
 			ai++
@@ -327,7 +322,7 @@ func orContainers(a, b, res *Bitmap, buf []uint16) {
 			bc := b.getContainer(off)
 			if getCardinality(bc) > 0 {
 				bkToBc[bk] = bc
-				sizeContainers += uint64(len(bc))
+				sizeContainers += len(bc)
 				newKeys++
 			}
 			bi++
@@ -339,7 +334,7 @@ func orContainers(a, b, res *Bitmap, buf []uint16) {
 		if getCardinality(ac) > 0 {
 			ak := a.keys.key(ai)
 			akToAc[ak] = ac
-			sizeContainers += uint64(len(ac))
+			sizeContainers += len(ac)
 			newKeys++
 		}
 	}
@@ -349,18 +344,13 @@ func orContainers(a, b, res *Bitmap, buf []uint16) {
 		if getCardinality(bc) > 0 {
 			bk := b.keys.key(bi)
 			bkToBc[bk] = bc
-			sizeContainers += uint64(len(bc))
+			sizeContainers += len(bc)
 			newKeys++
 		}
 	}
 
 	if sizeContainers > 0 {
-		// 2x uint64 = 8x uint16; for key and offset; twice as much as actually needed
-		sizeKeys := 8 * uint64(newKeys) * 2
-		// ensure enough space for new containers and keys,
-		// allocate required memory just once avoid copying underlying data slice multiple times
-		res.expandNoLengthChange(sizeContainers + sizeKeys)
-		res.expandKeys(sizeKeys)
+		res.expandConditionally(newKeys, sizeContainers)
 
 		for ak, ac := range akToAc {
 			// create a new container and update the key offset to this container.
@@ -395,7 +385,7 @@ func orContainersInRange(a, b *Bitmap, bi, bn int, buf []uint16) {
 	// copy containers from b to a all at once
 	// expanding underlying data slice and keys subslice once
 	bkToBc := map[uint64][]uint16{}
-	sizeContainers := uint64(0)
+	sizeContainers := 0
 	newKeys := 0
 
 	for ai < an && bi < bn {
@@ -435,7 +425,7 @@ func orContainersInRange(a, b *Bitmap, bi, bn int, buf []uint16) {
 			bc := b.getContainer(off)
 			if getCardinality(bc) > 0 {
 				bkToBc[bk] = bc
-				sizeContainers += uint64(len(bc))
+				sizeContainers += len(bc)
 				newKeys++
 			}
 			bi++
@@ -447,23 +437,13 @@ func orContainersInRange(a, b *Bitmap, bi, bn int, buf []uint16) {
 		if getCardinality(bc) > 0 {
 			bk := b.keys.key(bi)
 			bkToBc[bk] = bc
-			sizeContainers += uint64(len(bc))
+			sizeContainers += len(bc)
 			newKeys++
 		}
 	}
 
 	if sizeContainers > 0 {
-		if a.keysExpansionRequired(newKeys) {
-			// 2x uint64 = 8x uint16; for key and offset; twice as much as actually needed
-			sizeKeys := 8 * uint64(newKeys) * 2
-
-			// ensure enough space for new containers and keys,
-			// allocate required memory just once to avoid copying underlying data slice multiple times
-			a.expandNoLengthChange(sizeContainers + sizeKeys)
-			a.expandKeys(sizeKeys)
-		} else {
-			a.expandNoLengthChange(sizeContainers)
-		}
+		a.expandConditionally(newKeys, sizeContainers)
 
 		for bk, bc := range bkToBc {
 			// create a new container and update the key offset to this container.
@@ -498,36 +478,26 @@ func (ra *Bitmap) OrConc(bm *Bitmap, maxConcurrency int) *Bitmap {
 		return ra
 	}
 
-	var totalSizeContainers uint64
 	var totalNewKeys int
+	var totalSizeContainers int
 	var allKeys []uint64
 	var allContainers [][]uint16
 	lock := new(sync.Mutex)
 	inlineVsMutateLock := new(sync.RWMutex)
 	callback := func(bi, bj, _ int) {
 		buf := make([]uint16, maxContainerSize)
-		sizeContainers, newKeys, keys, containers := orContainersInRangeConc(ra, bm, bi, bj, buf, inlineVsMutateLock)
+		newKeys, sizeContainers, keys, containers := orContainersInRangeConc(ra, bm, bi, bj, buf, inlineVsMutateLock)
 
 		lock.Lock()
-		totalSizeContainers += sizeContainers
 		totalNewKeys += newKeys
+		totalSizeContainers += sizeContainers
 		allKeys = append(allKeys, keys...)
 		allContainers = append(allContainers, containers...)
 		lock.Unlock()
 	}
 	concurrentlyInRanges(numContainers, concurrency, callback)
 	if totalSizeContainers > 0 {
-		if ra.keysExpansionRequired(totalNewKeys) {
-			// 2x uint64 = 8x uint16; for key and offset; twice as much as actually needed
-			totalSizeKeys := 8 * uint64(totalNewKeys) * 2
-
-			// ensure enough space for new containers and keys,
-			// allocate required memory just once to avoid copying underlying data slice multiple times
-			ra.expandNoLengthChange(totalSizeContainers + totalSizeKeys)
-			ra.expandKeys(totalSizeKeys)
-		} else {
-			ra.expandNoLengthChange(totalSizeContainers)
-		}
+		ra.expandConditionally(totalNewKeys, totalSizeContainers)
 
 		for i, container := range allContainers {
 			// create a new container and update the key offset to this container.
@@ -541,7 +511,7 @@ func (ra *Bitmap) OrConc(bm *Bitmap, maxConcurrency int) *Bitmap {
 }
 
 func orContainersInRangeConc(a, b *Bitmap, bi, bn int, buf []uint16, inlineVsMutateLock *sync.RWMutex,
-) (sizeContainers uint64, newKeys int, bKeys []uint64, bContainers [][]uint16) {
+) (newKeys, sizeContainers int, bKeys []uint64, bContainers [][]uint16) {
 	bk := b.keys.key(bi)
 	ai := a.keys.search(bk)
 	an := a.keys.numKeys()
@@ -597,7 +567,7 @@ func orContainersInRangeConc(a, b *Bitmap, bi, bn int, buf []uint16, inlineVsMut
 			if getCardinality(bc) > 0 {
 				bKeys = append(bKeys, bk)
 				bContainers = append(bContainers, bc)
-				sizeContainers += uint64(len(bc))
+				sizeContainers += len(bc)
 				newKeys++
 			}
 			bi++
@@ -610,7 +580,7 @@ func orContainersInRangeConc(a, b *Bitmap, bi, bn int, buf []uint16, inlineVsMut
 			bk := b.keys.key(bi)
 			bKeys = append(bKeys, bk)
 			bContainers = append(bContainers, bc)
-			sizeContainers += uint64(len(bc))
+			sizeContainers += len(bc)
 			newKeys++
 		}
 	}
@@ -914,15 +884,7 @@ func (ra *Bitmap) FillUp(maxX uint64) {
 	}
 
 	// calculate required memory to allocate and expand underlying slice
-	containersLen := uint64(requiredContainersCount * maxContainerSize)
-	if ra.keysExpansionRequired(requiredContainersCount) {
-		// 2x uint64 = 8x uint16; for key and offset; twice as much as actually needed
-		keysLen := uint64(requiredContainersCount*2*4) * 2
-		ra.expandNoLengthChange(containersLen + keysLen)
-		ra.expandKeys(keysLen)
-	} else {
-		ra.expandNoLengthChange(containersLen)
-	}
+	ra.expandConditionally(requiredContainersCount, requiredContainersCount*maxContainerSize)
 
 	var onesBitmap bitmap
 	if containerIdx < maxContainersCount {
@@ -1069,6 +1031,65 @@ func (b bitmap) fillWithOnes() {
 	}
 }
 
-func (ra *Bitmap) keysExpansionRequired(newKeys int) bool {
-	return ra.keys.numKeys()+newKeys >= ra.keys.maxKeys()
+func (ra *Bitmap) expandConditionally(newKeys int, sizeContainers int) {
+	cp := cap(ra.data)
+	ln := len(ra.data)
+
+	sizeKeys := 8 * newKeys // 2x uint64 (key+offset) = 8x uint16
+	if ra.keys.numKeys()+newKeys < ra.keys.maxKeys() {
+		if ln+sizeContainers <= cp {
+			// keys and containers fit. nothing to do
+			return
+		}
+		// keys fit, containers do not. expand slice to make room for containers and only new keys
+	} else {
+		if ln+sizeKeys+sizeContainers <= cp {
+			curNumKeys := ra.keys.numKeys()
+
+			// keys do not fit, containers do. just move containers to make room for keys
+			if curNumKeys > newKeys {
+				// make room for up to curNumKeys additional keys
+				sizeKeys = 8 * curNumKeys // 2x uint64 (key+offset) = 8x uint16
+				if left := cp - ln - sizeContainers; left < sizeKeys {
+					sizeKeys = left / 8 * 8
+				}
+			}
+
+			// new containers will fit. just move containers to make room for keys
+			curSizeKeys := ra.keys.size()
+			newSizeKeys := curSizeKeys + sizeKeys
+
+			ra.data = ra.data[:ln+sizeKeys]
+			n := copy(ra.data[newSizeKeys:], ra.data[curSizeKeys:])
+			ra.memMoved += n
+			Memclr(ra.data[curSizeKeys:newSizeKeys]) // Zero out the space in the middle.
+
+			ra.keys = uint16To64SliceUnsafe(ra.data[:newSizeKeys])
+			ra.keys.setNodeSize(newSizeKeys)
+			for i := 0; i < curNumKeys; i++ {
+				ra.keys.setAt(valOffset(i), ra.keys.val(i)+uint64(sizeKeys))
+			}
+			return
+		}
+		// neither keys nor containers fit. expand slice to make room for containers and more keys
+		sizeKeys = 8 * max(newKeys, ra.keys.numKeys()) // 2x uint64 (key+offset) = 8x uint16
+	}
+
+	// expand 2x (or up to sizeKeys+sizeNewContainers if 2x is too little)
+	growBy := max(cp, sizeKeys+sizeContainers)
+	out := make([]uint16, ln+sizeKeys, cp+growBy)
+
+	curSizeKeys := ra.keys.size()
+	newSizeKeys := curSizeKeys + sizeKeys
+	copy(out, ra.data[:curSizeKeys])
+	copy(out[newSizeKeys:], ra.data[curSizeKeys:])
+	ra.data = out
+	ra._ptr = nil // Allow Go to GC whatever this was pointing to.
+	// Re-reference ra.keys correctly because underlying array has changed.
+
+	ra.keys = uint16To64SliceUnsafe(ra.data[:newSizeKeys])
+	ra.keys.setNodeSize(newSizeKeys)
+	for i := 0; i < ra.keys.numKeys(); i++ {
+		ra.keys.setAt(valOffset(i), ra.keys.val(i)+uint64(sizeKeys))
+	}
 }