move to using roaring iterators for UnmarshalBinary

The new roaring iterator used for the remap and importroaring
things could also be used for unmarshalling roaring streams,
and it's a slightly simpler design that doesn't need two passes
through the data. This patch cleans that up a bit, makes it work
better with ops logs, and uses that instead. It appears to
noticably but not immensely reduce the time imports take, but it
also gets us back down to one thing parsing roaring formats.

There are a couple of subtle changes to errors we were testing
for in various tests, and one of the fuzz tests goes away because
it was actually itself an erroneous error message -- it was reporting
the header of a roaring file as an invalid op because the op log
reader was running on the header for roaring files with zero
containers. Oops.

ctl/check_test.go

@@ -93,8 +93,9 @@ func TestCheckCommand_Run(t *testing.T) {
 		t.Fatalf("copy: %v", err)
 	}
 
-	if !strings.HasPrefix(err.Error(), "checking bitmap: unmarshalling: reading roaring header:") {
-		t.Fatalf("expect error: invalid roaring file, actual: '%s'", err)
+	expectedPrefix := "checking bitmap: unmarshalling: unknown roaring magic number 12849"
+	if !strings.HasPrefix(err.Error(), expectedPrefix) {
+		t.Fatalf("expect error: '%s...', actual: '%s'", expectedPrefix, err)
 	}
 	//	Todo: need correct roaring file for happy path
 }

ctl/inspect_test.go

@@ -41,8 +41,9 @@ func TestInspectCommand_Run(t *testing.T) {
 	file.Close()
 	cm.Path = file.Name()
 	err = cm.Run(context.Background())
-	if err != nil && err.Error() != "unmarshalling: reading roaring header: did not find expected serialCookie in header" {
-		t.Fatalf("can't run command: %v", err)
+	expectedError := "unmarshalling: unknown roaring magic number 12849"
+	if err != nil && err.Error() != expectedError {
+		t.Fatalf("expected error '%s', got '%v'", expectedError, err)
 	}
 
 	w.Close()

fragment.go

@@ -399,6 +399,8 @@ func (f *fragment) openStorage(unmarshalData bool) error {
 				}
 			}()
 		}
+		// set the preference for mapping based on whether the data's mmapped
+		f.storage.PreferMapping(newStorageData != nil)
 		// so we have a problem here: if this fails, it's unclear whether
 		// *either* or *both* of old and new storage data might be in use.
 		// So we call the thing that should unconditionally unmap both of them...

roaring/fuzz_test.go

@@ -25,53 +25,50 @@ func TestUnmarshalBinary(t *testing.T) {
 	}{
 		{ // Checks for the zero containers situation
 			cr:       []byte(":0\x00\x00\x01\x00\x00\x000000"), //":000000"
-			expected: "reading roaring header: malformed bitmap, key-cardinality slice overruns buffer at 12",
-		},
-		{ // Checks for int overflow
-			cr: []byte("<0\x000\x00\x00\x00\x00000000000000" +
-				"0"), //"<000000000000000"
-			expected: "unmarshaling as pilosa roaring: unknown op type: 48",
+			expected: "reading official header: malformed bitmap, key-cardinality slice overruns buffer at 12",
 		},
 		{ // The next 5 check for malformed bitmaps
 			cr: []byte("<0\x0000000000000000000" +
 				"\x00\x00\xec\x00\x03\x00\x00\x00\xec000"), //"<000000000000000000ÏÏ000"
-			expected: "unmarshaling as pilosa roaring: malformed bitmap, key-cardinality not provided for 67372036 containers",
+			expected: "insufficient data for header + offsets: want 12935430920 bytes, got 32",
 		},
 		{
 			cr: []byte("<0\x00\x02\x00\x00\x00\\f\x01\xb5\x8d\x009\v\x01\x00\x00\x00\x00" +
 				"\x00\x00e\x04\x00\x00\x00\x04\xfd\x00\x01\x00"), //"<0\f
µç9
e˝"
-			expected: "unmarshaling as pilosa roaring: malformed bitmap, key-cardinality not provided for 128625322 containers",
+			expected: "insufficient data for header + offsets: want 24696061960 bytes, got 32",
 		},
 		{
 			cr:       []byte("<0\x00\x02\x00\x00\x00&x.field safe"), //"<0&x.field safe"
-			expected: "unmarshaling as pilosa roaring: malformed bitmap, key-cardinality not provided for 53127850 containers",
+			expected: "insufficient data for header + offsets: want 10200547336 bytes, got 20",
 		},
 		{
 			cr: []byte("<0\x00\x00\x14\x00\x00\x00\x80\xffp\x05_ 4\x114089" +
 				"\x00\x00\xff\x000\x00\x02\x00\x00\x00\x00\xff\u007f\x00\x00\x01\x10\x00\x00j" +
 				"\x02\x00\x00$\x04_\x00\xff\u007f\xff062616163\x00" + //"<0ġp_ 44089ˇ0ˇj$_ˇˇ0626161630ø¸ad$j√"
 				"0\x00\x02\x00\x01\xbf\x00\x04\x00\xfcad$\x00\x00j\x10\x00\x00\xc3"),
-			expected: "unmarshaling as pilosa roaring: malformed bitmap, key-cardinality not provided for 1 containers",
+			expected: "insufficient data for header + offsets: want 328 bytes, got 80",
 		},
 		{ // 0 containers because the container is partially formed, but not fully (ie. 3/12 = 0)
 			cr:       []byte("<0\x00\x02\x03\x00\x00\x00쳫\v\x00d9\v\x00\x009\v"), //<0쳫d99
-			expected: "unmarshaling as pilosa roaring: malformed bitmap, key-cardinality not provided for 0 containers",
+			expected: "insufficient data for header + offsets: want 56 bytes, got 20",
 		},
 		{ // Checks for incomplete offset in readWithRuns
 			cr:       []byte(";0\x00\x00\v00000"), //";0000000"
-			expected: "reading offsets from official roaring format: offset incomplete: len=10",
+			expected: "container 0/1, expect run length at 9/10 bytes",
 		},
 		{ // Checks for incomplete offset in readOffsets
 			cr: []byte(":0\x00\x00\x03\x00\x00\x00000000000000" +
 				"\x00"), //:0000000000000
-			expected: "reading offsets from official roaring format: offset incomplete: len=1",
+			expected: "insufficient data for offsets (need 12 bytes, found 1)",
 		},
 	}
 
 	for _, crash := range confirmedCrashers {
 		err := b.UnmarshalBinary(crash.cr)
-		if err.Error() != crash.expected {
-			t.Errorf("Expected: %s, Got: %s", crash.expected, err)
+		if err == nil {
+			t.Errorf("expected: %s, got: no error", crash.expected)
+		} else if err.Error() != crash.expected {
+			t.Errorf("expected: %s, got: %s", crash.expected, err)
 		}
 	}
 

roaring/roaring.go

@@ -147,6 +147,8 @@ type Bitmap struct {
 
 	// User-defined flags.
 	Flags byte
+	// should we try to keep things mapped?
+	preferMapping bool
 
 	// Number of bit change operations written to the writer. Some operations
 	// contain multiple values, so "ops" represents the number of distinct
@@ -1125,7 +1127,11 @@ func (b *Bitmap) writeToUnoptimized(w io.Writer) (n int64, err error) {
 // bitmap and yield information about containers, including type, size, and
 // the location of their data structures.
 type roaringIterator interface {
+	// Next yields the information about the next container
 	Next() (key uint64, cType byte, n int, length int, pointer *uint16, err error)
+	// Remaining yields the bytes left over past the end of the roaring data,
+	// which is typically an ops log in our case.
+	Remaining() []byte
 }
 
 // baseRoaringIterator holds values used by both Pilosa and official Roaring
@@ -1142,6 +1148,7 @@ type baseRoaringIterator struct {
 	currentLen        int
 	currentPointer    *uint16
 	currentDataOffset uint32
+	lastDataOffset    int64
 	lastErr           error
 }
 
@@ -1189,10 +1196,14 @@ func newOfficialRoaringIterator(data []byte) (*officialRoaringIterator, error) {
 	r.headers = data[headerOffset:offsetOffset]
 	// note: offsets are only actually used with the no-run headers.
 	if r.haveRuns {
-		// start out pointed at where the offsets would have been.
 		r.currentDataOffset = uint32(offsetOffset)
 	} else {
+		if len(r.data) < offsetOffset+int(r.keys*4) {
+			return nil, fmt.Errorf("insufficient data for offsets (need %d bytes, found %d)",
+				r.keys*4, len(r.data)-offsetOffset)
+		}
 		r.offsets = data[offsetOffset : offsetOffset+int(r.keys*4)]
+		r.currentDataOffset = uint32(offsetOffset)
 	}
 	// set key to -1; user should call Next first.
 	r.currentIdx = -1
@@ -1212,6 +1223,12 @@ func newPilosaRoaringIterator(data []byte) (*pilosaRoaringIterator, error) {
 	r.keys = int64(binary.LittleEndian.Uint32(data[3+1 : 8]))
 	// it could happen
 	if r.keys == 0 {
+		// special case: what if we have zero containers, but a valid ops log after them?
+		// set currentDataOffset so that Done will set lastDataOffset and Remaining() will
+		// work.
+		if len(data) > headerBaseSize {
+			r.currentDataOffset = headerBaseSize
+		}
 		// not an error, exactly. it's valid and well-formed, we just have nothing to do
 		r.Done(io.EOF)
 		return r, nil
@@ -1227,6 +1244,10 @@ func newPilosaRoaringIterator(data []byte) (*pilosaRoaringIterator, error) {
 	offsetEnd := offsetStart + (r.keys * 4)
 	r.headers = data[headerStart:headerEnd]
 	r.offsets = data[offsetStart:offsetEnd]
+	// if there's no containers, we want to act as though data started at the end
+	// of the list of offsets, which was also empty, so we don't think the entire thing
+	// is actually a malformed op
+	r.currentDataOffset = uint32(offsetEnd)
 	// set key to -1; user should call Next first.
 	r.currentIdx = -1
 	r.currentKey = ^uint64(0)
@@ -1257,9 +1278,17 @@ func (r *baseRoaringIterator) Done(err error) {
 	r.currentN = 0
 	r.currentLen = 0
 	r.currentPointer = nil
+	r.lastDataOffset = int64(r.currentDataOffset)
 	r.currentDataOffset = 0
 }
 
+func (r *baseRoaringIterator) Remaining() []byte {
+	if r.lastDataOffset == 0 {
+		return nil
+	}
+	return r.data[r.lastDataOffset:]
+}
+
 func (r *pilosaRoaringIterator) Next() (key uint64, cType byte, n int, length int, pointer *uint16, err error) {
 	if r.currentIdx >= r.keys {
 		// we're already done
@@ -1306,6 +1335,7 @@ func (r *pilosaRoaringIterator) Next() (key uint64, cType byte, n int, length in
 			r.currentIdx, r.keys, r.currentKey, r.currentDataOffset, size, len(r.data)))
 		return r.Current()
 	}
+	r.currentDataOffset += uint32(size)
 	r.lastErr = nil
 	return r.Current()
 }
@@ -1333,6 +1363,11 @@ func (r *officialRoaringIterator) Next() (key uint64, cType byte, n int, length
 	// a run container keeps its data after an initial 2 byte length header
 	var runCount uint16
 	if r.currentType == containerRun {
+		if int(r.currentDataOffset)+2 > len(r.data) {
+			r.Done(fmt.Errorf("container %d/%d, expect run length at %d/%d bytes",
+				r.currentIdx, r.keys, r.currentDataOffset, len(r.data)))
+			return r.Current()
+		}
 		runCount = binary.LittleEndian.Uint16(r.data[r.currentDataOffset : r.currentDataOffset+runCountHeaderSize])
 		r.currentDataOffset += 2
 	}
@@ -1554,78 +1589,62 @@ func (b *Bitmap) ImportRoaringBits(data []byte, clear bool, log bool, rowSize ui
 		err = b.writeOp(&op)
 	}
 	return changed, rowSet, err
+}
 
+func (b *Bitmap) PreferMapping(preferred bool) {
+	b.preferMapping = preferred
 }
 
 // unmarshalPilosaRoaring treats data as being encoded in Pilosa's 64 bit
 // roaring format and decodes it into b.
-func (b *Bitmap) unmarshalPilosaRoaring(data []byte) error {
-	if len(data) < headerBaseSize {
-		return errors.New("data too small")
+func (b *Bitmap) UnmarshalBinary(data []byte) (err error) {
+	if data == nil {
+		return errors.New("no roaring bitmap provided")
+	}
+	var itr roaringIterator
+	var itrKey uint64
+	var itrCType byte
+	var itrN int
+	var itrLen int
+	var itrPointer *uint16
+	var itrErr error
+
+	itr, err = newRoaringIterator(data)
+	if err != nil {
+		return err
+	}
+	if itr == nil {
+		return errors.New("failed to create roaring iterator, but don't know why")
 	}
 
-	// Verify the first two bytes are a valid MagicNumber, and second two bytes match current storageVersion.
-	fileMagic := uint32(binary.LittleEndian.Uint16(data[0:2]))
-	fileVersion := uint32(data[2])
-	b.Flags = data[3]
-	if fileMagic != MagicNumber {
-		return fmt.Errorf("invalid roaring file, magic number %v is incorrect", fileMagic)
-	}
+	b.Containers.Reset()
 
-	if fileVersion != storageVersion {
-		return fmt.Errorf("wrong roaring version, file is v%d, server requires v%d", fileVersion, storageVersion)
-	}
-
-	// Read key count in bytes sizeof(cookie)+sizeof(flag):(sizeof(cookie)+sizeof(uint32)).
-	keyN := binary.LittleEndian.Uint32(data[3+1 : 8])
-	if uint32(len(data)) < headerBaseSize+keyN*12 {
-		return fmt.Errorf("malformed bitmap, key-cardinality not provided for %d containers", int(keyN)/12)
-	}
-
-	headerSize := headerBaseSize
-	b.Containers.ResetN(int(keyN))
-	// Descriptive header section: Read container keys and cardinalities.
-	for i, buf := 0, data[headerSize:]; i < int(keyN); i, buf = i+1, buf[12:] {
-		b.Containers.PutContainerValues(
-			binary.LittleEndian.Uint64(buf[0:8]),
-			byte(binary.LittleEndian.Uint16(buf[8:10])),
-			int(binary.LittleEndian.Uint16(buf[10:12]))+1,
-			true)
-	}
-	opsOffset := headerSize + int(keyN)*12
-
-	// Read container offsets and attach data.
-	citer, _ := b.Containers.Iterator(0)
-	for i, buf := 0, data[opsOffset:]; i < int(keyN); i, buf = i+1, buf[4:] {
-		offset := binary.LittleEndian.Uint32(buf[0:4])
-		// Verify the offset is within the bounds of the input data.
-		if int(offset) >= len(data) {
-			return fmt.Errorf("offset out of bounds: off=%d, len=%d", offset, len(data))
+	itrKey, itrCType, itrN, itrLen, itrPointer, itrErr = itr.Next()
+	for itrErr == nil {
+		newC := &Container{
+			typeID:  itrCType,
+			n:       int32(itrN),
+			len:     int32(itrLen),
+			cap:     int32(itrLen),
+			pointer: itrPointer,
+			flags:   flagMapped,
 		}
-
-		// Map byte slice directly to the container data.
-		citer.Next()
-		_, c := citer.Value()
-		// this shouldn't happen, since we don't normally store nils.
-		if c == nil {
-			continue
-		}
-		switch c.typ() {
-		case containerRun:
-			runCount := binary.LittleEndian.Uint16(data[offset : offset+runCountHeaderSize])
-			c.setRuns((*[0xFFFFFFF]interval16)(unsafe.Pointer(&data[offset+runCountHeaderSize]))[:runCount:runCount])
-			opsOffset = int(offset) + runCountHeaderSize + len(c.runs())*interval16Size
-		case containerArray:
-			c.setArray((*[0xFFFFFFF]uint16)(unsafe.Pointer(&data[offset]))[:c.N():c.N()])
-			opsOffset = int(offset) + len(c.array())*2 // sizeof(uint32)
-		case containerBitmap:
-			c.setBitmap((*[0xFFFFFFF]uint64)(unsafe.Pointer(&data[offset]))[:bitmapN:bitmapN])
-			opsOffset = int(offset) + len(c.bitmap())*8 // sizeof(uint64)
+		if !b.preferMapping {
+			newC.unmapOrClone()
 		}
+		b.Containers.Put(itrKey, newC)
+		itrKey, itrCType, itrN, itrLen, itrPointer, itrErr = itr.Next()
+	}
+	// note: if we get a non-EOF err, it's possible that we made SOME
+	// changes but didn't log them. I don't have a good solution to this.
+	if itrErr != io.EOF {
+		return itrErr
 	}
 
 	// Read ops log until the end of the file.
-	buf := data[opsOffset:]
+	b.ops = 0
+	b.opN = 0
+	buf := itr.Remaining()
 	for {
 		// Exit when there are no more ops to parse.
 		if len(buf) == 0 {
@@ -1648,7 +1667,6 @@ func (b *Bitmap) unmarshalPilosaRoaring(data []byte) error {
 		// Move the buffer forward.
 		buf = buf[opr.size():]
 	}
-
 	return nil
 }
 
@@ -5138,56 +5156,7 @@ func readOfficialHeader(buf []byte) (size uint32, containerTyper func(index uint
 	return size, containerTyper, header, pos, haveRuns, err
 }
 
-// UnmarshalBinary decodes b from a binary-encoded byte slice. data can be in
-// either official roaring format or Pilosa's roaring format.
-func (b *Bitmap) UnmarshalBinary(data []byte) error {
-	if data == nil {
-		// Nothing to unmarshal
-		return nil
-	}
-	statsHit("Bitmap/UnmarshalBinary")
-	b.opN = 0 // reset opN since we're reading new data.
-	fileMagic := uint32(binary.LittleEndian.Uint16(data[0:2]))
-	if fileMagic == MagicNumber { // if pilosa roaring
-		return errors.Wrap(b.unmarshalPilosaRoaring(data), "unmarshaling as pilosa roaring")
-	}
-
-	keyN, containerTyper, header, pos, haveRuns, err := readOfficialHeader(data)
-	if err != nil {
-		return errors.Wrap(err, "reading roaring header")
-	}
-	// Only the Pilosa roaring format has flags. The official Roaring format
-	// hasn't got space in its header for flags.
-	b.Flags = 0
-
-	b.Containers.ResetN(int(keyN))
-	// Descriptive header section: Read container keys and cardinalities.
-	for i, buf := uint(0), data[header:]; i < uint(keyN); i, buf = i+1, buf[4:] {
-		card := int(binary.LittleEndian.Uint16(buf[2:4])) + 1
-		b.Containers.PutContainerValues(
-			uint64(binary.LittleEndian.Uint16(buf[0:2])),
-			containerTyper(i, card), /// container type voodo with isRunBitmap
-			card,
-			true)
-	}
-
-	// Read container offsets and attach data.
-	if haveRuns {
-		err := readWithRuns(b, data, pos, keyN)
-		if err != nil {
-			return errors.Wrap(err, "reading offsets from official roaring format")
-		}
-	} else {
-		err := readOffsets(b, data, pos, keyN)
-		if err != nil {
-			return errors.Wrap(err, "reading offsets from official roaring format")
-		}
-	}
-	return nil
-}
-
 func readOffsets(b *Bitmap, data []byte, pos int, keyN uint32) error {
-
 	citer, _ := b.Containers.Iterator(0)
 	for i, buf := 0, data[pos:]; i < int(keyN); i, buf = i+1, buf[4:] {
 		// Verify the offset is fully formed

roaring/roaring_internal_test.go

@@ -3409,11 +3409,11 @@ func TestUnmarshalRoaringWithErrors(t *testing.T) {
 	}{
 		{ // Runs a bitmap without runs and no containers through the official roaring
 			hexString:     "3A30000000000000",
-			expectedError: "reading roaring header: malformed bitmap, key-cardinality slice overruns buffer at 8",
+			expectedError: "reading official header: malformed bitmap, key-cardinality slice overruns buffer at 8",
 		},
 		{ // Runs a bitmap with runs and no containers through the official roaring
 			hexString:     "3B30000000000000",
-			expectedError: "reading roaring header: malformed bitmap, key-cardinality slice overruns buffer at 9",
+			expectedError: "reading official header: malformed bitmap, key-cardinality slice overruns buffer at 9",
 		},
 		{ // Runs a bitmap in the Pilosa format through the Pilosa roaring
 			hexString: "3C30000000000000",