roaring import: fix api doc, refactor/rename, test merge

api.go

@@ -297,8 +297,8 @@ func (api *API) Field(_ context.Context, indexName, fieldName string) (*Field, e
 
 // ImportRoaringBytes is a low level interface for importing data to Pilosa when
 // extremely high throughput is desired. The data must be encoded in a
-// particular way which may be unintuitive (discussed below), and overwrites any
-// existing data in the particular fragment into which it is being loaded.
+// particular way which may be unintuitive (discussed below). The data is merged
+// with existing data.
 //
 // It takes as input a roaring bitmap which it uses as the data for the
 // indicated index, field, and shard. The bitmap may be encoded according to the

fragment.go

@@ -44,15 +44,19 @@ import (
 )
 
 const (
-	// ShardWidth is the number of column IDs in a shard.
-	ShardWidth = 1048576
+	// ShardWidth is the number of column IDs in a shard. It must be a power of 2 greater than or equal to 16.
+	shardWidthExponent = 20
+	ShardWidth         = 1 << shardWidthExponent
 
-	// containersPerRowSegment is dependent upon ShardWidth,
-	// and it represents the number of containers per shard row
-	// (or rowSegment). Since containers are set in roaring
-	// to be 2^16, then this const should be ShardWidth / 2^16.
-	// It is represented as the exponent n of 2^n.
-	containersPerRowSegment = 4
+	// shardVsContainerExponent is the power of 2 of ShardWith minus the power
+	// of two of roaring container width (which is 16).
+	// 2^shardVsContainerExponent is the number of containers in a shard row.
+	//
+	// It is represented in this rather awkward way because calculating the row
+	// which a given container is in means dividing by the number of rows per
+	// container which is performantly expressed as a right shift by this
+	// exponent.
+	shardVsContainerExponent = shardWidthExponent - 16
 
 	// snapshotExt is the file extension used for an in-process snapshot.
 	snapshotExt = ".snapshotting"
@@ -1432,21 +1436,21 @@ func (f *fragment) importRoaringBytes(roaringBytes []byte) error {
 	}
 
 	// get a list of keys in order to update the cache
-	i, _ := bm.Containers.Iterator(0)
-	set := make([]uint64, 0)
+	iter, _ := bm.Containers.Iterator(0)
+	rowsInData := make([]uint64, 0)
 	var lastRow uint64 = math.MaxUint64
 
-	for i.Next() {
-		key, _ := i.Value()
+	for iter.Next() {
+		key, _ := iter.Value()
 
 		// virtual row for the current container
-		vRow := key >> containersPerRowSegment
+		vRow := key >> shardVsContainerExponent
 
 		// skip dups
 		if vRow == lastRow {
 			continue
 		}
-			set = append(set, vRow)
+		rowsInData = append(rowsInData, vRow)
 		lastRow = vRow
 	}
 
@@ -1454,7 +1458,7 @@ func (f *fragment) importRoaringBytes(roaringBytes []byte) error {
 		bm = f.storage.Union(bm)
 	}
 
-	for _,rowID := range set {
+	for _, rowID := range rowsInData {
 		n := bm.CountRange(rowID*ShardWidth, (rowID+1)*ShardWidth)
 		f.cache.BulkAdd(rowID, n)
 	}
@@ -1493,6 +1497,8 @@ func (f *fragment) snapshot() error {
 	return snapshot(f, f.storage)
 }
 
+// snapshot writes the fragment f with bm as the data. It is unprotected, and
+// f.mu must be locked when calling it.
 func snapshot(f *fragment, bm *roaring.Bitmap) error {
 
 	f.Logger.Printf("fragment: snapshotting %s/%s/%s/%d", f.index, f.field, f.view, f.shard)
@@ -1763,7 +1769,7 @@ func (f *fragment) rows() []uint64 {
 		key, _ := i.Value()
 
 		// virtual row for the current container
-		vRow := key >> containersPerRowSegment
+		vRow := key >> shardVsContainerExponent
 
 		// skip dups
 		if vRow == lastRow {
@@ -1792,7 +1798,7 @@ func (f *fragment) rowsForColumn(columnID uint64) []uint64 {
 		key, c := i.Value()
 
 		// virtual row for the current container
-		vRow := key >> containersPerRowSegment
+		vRow := key >> shardVsContainerExponent
 
 		// column container key for virtual row
 		colKey = ((vRow * ShardWidth) + colID) >> 16

fragment_internal_test.go

@@ -17,14 +17,17 @@ package pilosa
 import (
 	"bytes"
 	"flag"
+	"fmt"
 	"io/ioutil"
 	"math"
 	"reflect"
+	"sort"
 	"testing"
 	"testing/quick"
 
 	"github.com/davecgh/go-spew/spew"
 	"github.com/pilosa/pilosa/pql"
+	"github.com/pilosa/pilosa/roaring"
 )
 
 // Test flags
@@ -1397,3 +1400,93 @@ func TestFragment_RowsIteration(t *testing.T) {
 		}
 	})
 }
+
+// Test Importing roaring data.
+func TestFragment_RoaringImport(t *testing.T) {
+	tests := [][][]uint64{
+		{
+			[]uint64{0},
+			[]uint64{1},
+		},
+		{
+			[]uint64{0, 65535, 65536, 65537, 65538, 65539, 130000},
+			[]uint64{1000, 67000, 130000},
+		},
+		{
+			[]uint64{0, 65535, 65536, 65537, 65538, 65539, 130000},
+			[]uint64{0, 65535, 65536, 65537, 65538, 65539, 130000},
+		},
+		{
+			[]uint64{0, 65535, 65536, 1<<20 + 1, 1<<20*2 + 1},
+			[]uint64{1, 1<<20 + 2, 1<<20*2 + 2},
+		},
+	}
+
+	for i, test := range tests {
+		t.Run(fmt.Sprintf("importroaring%d", i), func(t *testing.T) {
+			f := mustOpenFragment("i", "f", viewStandard, 0, "")
+			defer f.Close()
+			for num, input := range test {
+				buf := &bytes.Buffer{}
+				bm := roaring.NewBitmap(input...)
+				_, err := bm.WriteTo(buf)
+				if err != nil {
+					t.Fatalf("writing to buffer: %v", err)
+				}
+				f.importRoaringBytes(buf.Bytes())
+				exp := calcExpected(test[:num+1]...)
+				for row, expCols := range exp {
+					cols := f.row(uint64(row)).Columns()
+					t.Logf("\nrow: %d\n  exp:%v\n  got:%v", row, expCols, cols)
+					if !reflect.DeepEqual(cols, expCols) {
+						t.Fatalf("input%d, row %d\n  exp:%v\n  got:%v", num, row, expCols, cols)
+					}
+				}
+			}
+		})
+	}
+}
+
+// calcExpected takes a number of slices of uint64 represented data to be added
+// to a fragment. It calculates which rows and bits set in each row would be
+// expected after importing that data, and returns a [][]uint64 where the index
+// into the first slice is row id.
+func calcExpected(inputs ...[]uint64) [][]uint64 {
+	// create map of row id to set of column values in that row.
+	rows := make(map[uint64]map[uint64]struct{})
+	var maxrow uint64 = 0
+	for _, input := range inputs {
+		for _, val := range input {
+			row := val / ShardWidth
+			if row > maxrow {
+				maxrow = row
+			}
+			m, ok := rows[row]
+			if !ok {
+				m = make(map[uint64]struct{})
+				rows[row] = m
+			}
+			m[val%ShardWidth] = struct{}{}
+		}
+	}
+
+	// initialize ret slice
+	ret := make([][]uint64, maxrow+1)
+	for i, _ := range ret {
+		ret[i] = make([]uint64, 0)
+	}
+
+	// populate ret slices from rows map
+	for row, vals := range rows {
+		for val, _ := range vals {
+			ret[row] = append(ret[row], val)
+		}
+	}
+
+	// sort ret slices
+	for _, slice := range ret {
+		sort.Slice(slice, func(i int, j int) bool { return slice[i] < slice[j] })
+	}
+
+	return ret
+}