implement a BSI-aware filter to avoid OffsetRange calls in fragment.sum

We don't really need to fully extract every row, we just need counts.
This naive approach uses logic similar to BitmapBitmapFilter, but tweaks
it so that we can intercept the existence and sign bit rows, work with
an optional filter, and yield a sum. We accumulate the statistics
internally, rather than using a callback, because I tried to make it
work with a callback and it was a complete mess.

Note the fancy check for container reuse in the BSI Count filter.
This is because intersection(full container, X) is just the original
X, *not* a copy, but in this case we need a copy because RBF
ApplyFilter will in fact reuse a single container's storage for
each consecutive container.

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.

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
+}