Merge pull request #1943 from molecula/fb1216

[FB-1216] Improve sum aggregate performance

fragment.go

@@ -770,50 +770,31 @@ func (f *fragment) setValueBase(txOrig Tx, columnID uint64, bitDepth uint64, val
 // sum returns the sum of a given bsiGroup as well as the number of columns involved.
 // A bitmap can be passed in to optionally filter the computed columns.
 func (f *fragment) sum(tx Tx, filter *Row, bitDepth uint64) (sum int64, count uint64, err error) {
-	// Compute count based on the existence row.
-	consider, err := f.row(tx, bsiExistsBit)
-	if err != nil {
-		return sum, count, err
-	} else if filter != nil {
-		consider = consider.Intersect(filter)
-	}
-	count = consider.Count()
-
-	// Get negative set
-	nrow, err := f.row(tx, bsiSignBit)
-	if err != nil {
-		return sum, count, err
-	}
-
-	// Filter negative set
-	nrow = consider.Intersect(nrow)
-
-	// Get postive set
-	prow := consider.Difference(nrow)
-
-	// Compute the sum based on the bit count of each row multiplied by the
-	// place value of each row. For example, 10 bits in the 1's place plus
-	// 4 bits in the 2's place plus 3 bits in the 4's place equals a total
-	// sum of 30:
-	//
-	//   10*(2^0) + 4*(2^1) + 3*(2^2) = 30
-	//
-	// Execute once for positive numbers and once for negative. Subtract the
-	// negative sum from the positive sum.
-	for i := uint64(0); i < bitDepth; i++ {
-		row, err := f.row(tx, uint64(bsiOffsetBit+i))
-		if err != nil {
-			return sum, count, err
+	// If there's a provided filter, but it has no contents for this particular
+	// shard, we're done and can return early. If there's no provided filter,
+	// though, we want to run with no-filter, as opposed to an empty filter.
+	var filterData *roaring.Bitmap
+	if filter != nil {
+		for _, seg := range filter.segments {
+			if seg.shard == f.shard {
+				filterData = seg.data
+				break
+			}
 		}
-
-		psum := int64((1 << i) * row.intersectionCount(prow))
-		nsum := int64((1 << i) * row.intersectionCount(nrow))
-
-		// Squash to reduce the possibility of overflow.
-		sum += psum - nsum
+		// if filter is empty, we're done
+		if filterData == nil {
+			return 0, 0, nil
+		}
+	}
+	bsiFilt := roaring.NewBitmapBSICountFilter(filterData)
+	err = tx.ApplyFilter(f.index(), f.field(), f.view(), f.shard, 0, bsiFilt)
+	if err != nil && err != io.EOF {
+		return sum, count, errors.Wrap(err, "finding existing positions")
 	}
 
-	return sum, count, nil
+	c32, sum := bsiFilt.Total()
+
+	return sum, uint64(c32), nil
 }
 
 // min returns the min of a given bsiGroup as well as the number of columns involved.

rbf/cursorx.go

@@ -177,7 +177,7 @@ func intoContainer(l leafCell, tx *Tx, replacing *roaring.Container, target []by
 	case ContainerTypeBitmapPtr:
 		_, bm, _ := tx.leafCellBitmap(toPgno(cpMaybe))
 		cloneMaybe := bm
-		c = roaring.RemakeContainerBitmap(replacing, cloneMaybe)
+		c = roaring.RemakeContainerBitmapN(replacing, cloneMaybe, int32(l.BitN))
 	case ContainerTypeBitmap:
 		c = roaring.RemakeContainerBitmapN(replacing, toArray64(cpMaybe), int32(l.BitN))
 	case ContainerTypeRLE:
@@ -216,9 +216,9 @@ func toContainer(l leafCell, tx *Tx) (c *roaring.Container) {
 	case ContainerTypeBitmapPtr:
 		_, bm, _ := tx.leafCellBitmap(toPgno(cpMaybe))
 		cloneMaybe := bm
-		c = roaring.NewContainerBitmap(-1, cloneMaybe)
+		c = roaring.NewContainerBitmap(l.BitN, cloneMaybe)
 	case ContainerTypeBitmap:
-		c = roaring.NewContainerBitmap(-1, toArray64(cpMaybe))
+		c = roaring.NewContainerBitmap(l.BitN, toArray64(cpMaybe))
 	case ContainerTypeRLE:
 		c = roaring.NewContainerRun(toInterval16(cpMaybe))
 	}

rbf/tx.go

@@ -1258,6 +1258,22 @@ func (tx *Tx) ContainerIterator(name string, key uint64) (citer roaring.Containe
 	return &containerIterator{cursor: c}, exact, nil
 }
 
+// Shared pool for in-memory database pages.
+// These are used before being flushed to disk.
+var containerFilterPool = &sync.Pool{}
+
+func getContainerFilter(c *Cursor, filter roaring.BitmapFilter, tx *Tx) *containerFilter {
+	existing := containerFilterPool.Get()
+	if existing == nil {
+		return &containerFilter{cursor: c, filter: filter, tx: tx}
+	}
+	f := existing.(*containerFilter)
+	f.cursor = c
+	f.filter = filter
+	f.tx = tx
+	return f
+}
+
 func (tx *Tx) ApplyFilter(name string, key uint64, filter roaring.BitmapFilter) (err error) {
 	tx.mu.RLock()
 	defer tx.mu.RUnlock()
@@ -1273,7 +1289,7 @@ func (tx *Tx) ApplyFilter(name string, key uint64, filter roaring.BitmapFilter)
 	if err != nil {
 		return err
 	}
-	f := containerFilter{cursor: c, filter: filter, tx: tx}
+	f := getContainerFilter(c, filter, tx)
 	defer f.Close()
 	return f.Apply()
 }
@@ -1615,6 +1631,8 @@ type containerFilter struct {
 
 func (s *containerFilter) Close() {
 	s.cursor.Close()
+	s.cursor = nil
+	containerFilterPool.Put(s)
 }
 
 func (s *containerFilter) Apply() (err error) {

roaring/filter.go

@@ -879,3 +879,136 @@ func ApplyFilterToIterator(filter BitmapFilter, iter ContainerIterator) error {
 	}
 	return nil
 }
+
+// BitmapBSICountFilter gives counts of values in each value-holding row
+// of a BSI field, constrained by a filter. The first row of the data is
+// taken to be an existence bit, which is intersected into the filter to
+// constrain it, and the second is used as a sign bit. The rows after that
+// are treated as value rows, and their counts of bits, overlapping with
+// positive and negative bits in the sign rows, are returned to a callback
+// function.
+//
+// The total counts of positions evaluated are returned with a row count
+// of ^uint64(0) prior to row counts.
+type BitmapBSICountFilter struct {
+	containers  []*Container
+	positive    []*Container
+	negative    []*Container
+	nextOffsets []uint64
+	count       int32
+	psum, nsum  uint64
+}
+
+func (b *BitmapBSICountFilter) Total() (count int32, total int64) {
+	return b.count, int64(b.psum) - int64(b.nsum)
+}
+
+func (b *BitmapBSICountFilter) ConsiderKey(key FilterKey, n int32) FilterResult {
+	pos := key & keyMask
+	if b.containers[pos] == nil || n == 0 {
+		return key.RejectUntilOffset(b.nextOffsets[pos])
+	}
+	return key.NeedData()
+}
+
+func (b *BitmapBSICountFilter) ConsiderData(key FilterKey, data *Container) FilterResult {
+	pos := key & keyMask
+	filter := b.containers[pos]
+	if filter == nil {
+		key.RejectUntilOffset(b.nextOffsets[pos])
+	}
+	row := uint64(key >> rowExponent) // row count within the fragment
+	// How do we translate the filter and existence bit into actionable things?
+	// Assume the sign row is empty. We want positive values for anything in
+	// the intersection of the filter and the positive bits. If the sign row
+	// isn't empty, we want positive values for that intersection, less the
+	// sign row, and negative for the intersection of the filter/positive and
+	// the sign bits. So we can just stash the intermediate filter+existence
+	// as positive, then split it up if we have sign bits, which we often don't.
+	setup := false
+	switch row {
+	case 0: // existence bit
+		b.positive[pos] = intersect(b.containers[pos], data)
+		if b.positive[pos] == data {
+			b.positive[pos] = b.positive[pos].Clone()
+		}
+		b.count += int32(b.positive[pos].N())
+		setup = true
+	case 1: // sign bit
+		// split into negative/positive components. doesn't affect total
+		// count.
+		b.negative[pos] = intersect(b.positive[pos], data)
+		if b.negative[pos] == data {
+			b.negative[pos] = b.negative[pos].Clone()
+		}
+		b.positive[pos] = difference(b.positive[pos], data)
+		setup = true
+	}
+	// if we were doing setup (first two rows), we're done
+	if setup {
+		return key.MatchOneUntilOffset(b.nextOffsets[pos])
+	}
+	// helpful reminder: a nil container is a valid empty container, and
+	// intersectionCount knows this.
+	pcount := intersectionCount(b.positive[pos], data)
+	ncount := intersectionCount(b.negative[pos], data)
+	b.psum += (uint64(pcount) << (row - 2))
+	b.nsum += (uint64(ncount) << (row - 2))
+	return key.MatchOneUntilOffset(b.nextOffsets[pos])
+}
+
+// NewBitmapBSICountFilter creates a BitmapBSICountFilter, used for tasks
+// like computing the sum of a BSI field matching a given filter.
+//
+// The input filter is assumed to represent one "row" of a shard's data,
+// which is to say, a range of up to rowWidth consecutive containers starting
+// at some multiple of rowWidth. We coerce that to the 0..rowWidth range
+// because offset-within-row is what we care about.
+func NewBitmapBSICountFilter(filter *Bitmap) *BitmapBSICountFilter {
+	containers := make([]*Container, rowWidth*3)
+	b := &BitmapBSICountFilter{
+		containers:  containers[:rowWidth],
+		positive:    containers[rowWidth : rowWidth*2],
+		negative:    containers[rowWidth*2 : rowWidth*3],
+		nextOffsets: make([]uint64, rowWidth),
+	}
+	if filter == nil {
+		for i := range b.containers {
+			b.containers[i] = NewContainerRun([]Interval16{{Start: 0, Last: 65535}})
+			b.nextOffsets[i] = uint64(i+1) % rowWidth
+		}
+		return b
+	}
+	count := 0
+	iter, _ := filter.Containers.Iterator(0)
+	last := uint64(0)
+	for iter.Next() {
+		k, v := iter.Value()
+		// Coerce container key into the 0-rowWidth range we'll be
+		// using to compare against containers within each row.
+		k = k & keyMask
+		b.containers[k] = v
+		last = k
+		count++
+	}
+	// if there's only one container, we need to populate everything with
+	// its position.
+	if count == 1 {
+		for i := range b.containers {
+			b.nextOffsets[i] = last
+		}
+	} else {
+		// Point each container at the offset of the next valid container.
+		// With sparse bitmaps this will potentially make skipping faster.
+		for i := range b.containers {
+			if b.containers[i] != nil {
+				for int(last) != i {
+					b.nextOffsets[last] = uint64(i)
+					last = (last + 1) % rowWidth
+				}
+			}
+		}
+	}
+
+	return b
+}