From 67830b74cf82e1e0393ae0d2242cc25ce99bf755 Mon Sep 17 00:00:00 2001 From: Seebs Date: Tue, 25 Jun 2019 16:57:08 -0500 Subject: [PATCH] allow importRoaring to work with official format roaring I didn't think of this, because we don't use it much in the client. This is a bit hairy because really official roaring is two fairly different formats, one with runs and one without. --- roaring/roaring.go | 157 +++++++++++++++++++++++++++++++++++++++------ 1 file changed, 139 insertions(+), 18 deletions(-) diff --git a/roaring/roaring.go b/roaring/roaring.go index 174ef39a4..0d61020e8 100644 --- a/roaring/roaring.go +++ b/roaring/roaring.go @@ -1124,7 +1124,11 @@ func (b *Bitmap) writeToUnoptimized(w io.Writer) (n int64, err error) { // roaringIterator represents something which can iterate through a roaring // bitmap and yield information about containers, including type, size, and // the location of their data structures. -type roaringIterator struct { +type roaringIterator interface { + Next() (key uint64, cType byte, n int, length int, pointer *uint16, err error) +} + +type baseRoaringIterator struct { data []byte keys int64 headers []byte @@ -1139,20 +1143,58 @@ type roaringIterator struct { lastErr error } -func newRoaringIterator(data []byte) (*roaringIterator, error) { - if len(data) < headerBaseSize { - return nil, errors.New("invalid data: not long enough to be a roaring header") +type pilosaRoaringIterator struct { + baseRoaringIterator +} + +type officialRoaringIterator struct { + baseRoaringIterator + containerTyper func(index uint, card int) byte + haveRuns bool +} + +func newOfficialRoaringIterator(data []byte) (*officialRoaringIterator, error) { + r := &officialRoaringIterator{} + r.data = data + + // share code with the existing unmarshal code + var offsetOffset, headerOffset int + var err error + var keys uint32 + + // we ignore the flags, since we don't have to process them for anything. + keys, r.containerTyper, headerOffset, offsetOffset, _, r.haveRuns, err = readOfficialHeader(data) + if err != nil { + return nil, fmt.Errorf("reading official header: %v", err) } - // Verify the first two bytes are a valid MagicNumber, and second two bytes match current storageVersion. - fileMagic := uint32(binary.LittleEndian.Uint16(data[0:2])) + if keys == 0 { + // not an error, exactly. it's valid and well-formed, we just have nothing to do + r.Done(io.EOF) + return r, nil + } + r.keys = int64(keys) + 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 { + r.offsets = data[offsetOffset : offsetOffset+int(r.keys*4)] + } + // set key to -1; user should call Next first. + r.currentIdx = -1 + r.currentKey = ^uint64(0) + r.lastErr = errors.New("tried to read iterator without calling Next first") + return r, nil +} + +func newPilosaRoaringIterator(data []byte) (*pilosaRoaringIterator, error) { fileVersion := uint32(data[2]) - if fileMagic != MagicNumber { - return nil, fmt.Errorf("invalid roaring file, magic number %v is incorrect", fileMagic) - } if fileVersion != storageVersion { return nil, fmt.Errorf("wrong roaring version, file is v%d, server requires v%d", fileVersion, storageVersion) } - r := &roaringIterator{data: data} + r := &pilosaRoaringIterator{} + r.data = data // Read key count in bytes sizeof(cookie)+sizeof(flag):(sizeof(cookie)+sizeof(uint32)). r.keys = int64(binary.LittleEndian.Uint32(data[3+1 : 8])) // it could happen @@ -1179,8 +1221,23 @@ func newRoaringIterator(data []byte) (*roaringIterator, error) { return r, nil } +func newRoaringIterator(data []byte) (roaringIterator, error) { + if len(data) < headerBaseSize { + return nil, errors.New("invalid data: not long enough to be a roaring header") + } + // Verify the first two bytes are a valid MagicNumber, and second two bytes match current storageVersion. + fileMagic := uint32(binary.LittleEndian.Uint16(data[0:2])) + switch fileMagic { + case serialCookie, serialCookieNoRunContainer: + return newOfficialRoaringIterator(data) + case MagicNumber: + return newPilosaRoaringIterator(data) + } + return nil, fmt.Errorf("unknown roaring magic number %d", fileMagic) +} + // Done marks the iterator as complete, recording err as the reason why -func (r *roaringIterator) Done(err error) { +func (r *baseRoaringIterator) Done(err error) { r.lastErr = err r.currentKey = ^uint64(0) r.currentType = 0 @@ -1190,7 +1247,7 @@ func (r *roaringIterator) Done(err error) { r.currentDataOffset = 0 } -func (r *roaringIterator) Next() (key uint64, cType byte, n int, length int, pointer *uint16, err error) { +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 return r.Current() @@ -1206,8 +1263,11 @@ func (r *roaringIterator) Next() (key uint64, cType byte, n int, length int, poi r.currentType = byte(binary.LittleEndian.Uint16(header[8:10])) r.currentN = int(binary.LittleEndian.Uint16(header[10:12])) + 1 r.currentDataOffset = binary.LittleEndian.Uint32(r.offsets[r.currentIdx*4:]) + // a run container keeps its data after an initial 2 byte length header + var runCount uint16 if r.currentType == containerRun { + runCount = binary.LittleEndian.Uint16(r.data[r.currentDataOffset : r.currentDataOffset+runCountHeaderSize]) r.currentDataOffset += 2 } if r.currentDataOffset > uint32(len(r.data)) || r.currentDataOffset < headerBaseSize { @@ -1225,7 +1285,7 @@ func (r *roaringIterator) Next() (key uint64, cType byte, n int, length int, poi r.currentLen = 1024 size = 8192 case containerRun: - r.currentLen = int(*((*uint16)(unsafe.Pointer(&r.data[r.currentDataOffset-2])))) + r.currentLen = int(runCount) size = r.currentLen * 4 } if int64(r.currentDataOffset)+int64(size) > int64(len(r.data)) { @@ -1237,7 +1297,70 @@ func (r *roaringIterator) Next() (key uint64, cType byte, n int, length int, poi return r.Current() } -func (r *roaringIterator) Current() (key uint64, cType byte, n int, length int, pointer *uint16, err error) { +func (r *officialRoaringIterator) Next() (key uint64, cType byte, n int, length int, pointer *uint16, err error) { + if r.currentIdx >= r.keys { + // we're already done + return r.Current() + } + r.currentIdx++ + if r.currentIdx == r.keys { + // this is the last key. transition state to the finalized state + r.Done(io.EOF) + return r.Current() + } + header := r.headers[r.currentIdx*4:] + r.currentKey = uint64(binary.LittleEndian.Uint16(header[0:2])) + r.currentN = int(binary.LittleEndian.Uint16(header[2:4])) + 1 + r.currentType = r.containerTyper(uint(r.currentIdx), r.currentN) + // with runs, we can't actually look up offsets; the format just stores + // things sequentially. so we have to actually track the offset in that case. + if !r.haveRuns { + r.currentDataOffset = binary.LittleEndian.Uint32(r.offsets[r.currentIdx*4:]) + } + // a run container keeps its data after an initial 2 byte length header + var runCount uint16 + if r.currentType == containerRun { + runCount = binary.LittleEndian.Uint16(r.data[r.currentDataOffset : r.currentDataOffset+runCountHeaderSize]) + r.currentDataOffset += 2 + } + if r.currentDataOffset > uint32(len(r.data)) || r.currentDataOffset < headerBaseSize { + r.Done(fmt.Errorf("container %d/%d, key %d, had offset %d, maximum %d", + r.currentIdx, r.keys, r.currentKey, r.currentDataOffset, len(r.data))) + return r.Current() + } + r.currentPointer = (*uint16)(unsafe.Pointer(&r.data[r.currentDataOffset])) + var size int + switch r.currentType { + case containerArray: + r.currentLen = r.currentN + size = r.currentLen * 2 + case containerBitmap: + r.currentLen = 1024 + size = 8192 + case containerRun: + // official format stores runs as start/len, we want to convert, but since + // they might be mmapped, we can't write to that memory + newRuns := make([]interval16, runCount) + oldRuns := (*[65536]interval16)(unsafe.Pointer(r.currentPointer))[:runCount:runCount] + copy(newRuns, oldRuns) + for i := range newRuns { + newRuns[i].last += newRuns[i].start + } + r.currentPointer = (*uint16)(unsafe.Pointer(&newRuns[0])) + r.currentLen = int(runCount) + size = r.currentLen * 4 + } + if int64(r.currentDataOffset)+int64(size) > int64(len(r.data)) { + r.Done(fmt.Errorf("container %d/%d, key %d, had offset %d+%d size, maximum %d", + 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() +} + +func (r *baseRoaringIterator) Current() (key uint64, cType byte, n int, length int, pointer *uint16, err error) { return r.currentKey, r.currentType, r.currentN, r.currentLen, r.currentPointer, r.lastErr } @@ -1250,13 +1373,11 @@ func (r *roaringIterator) Current() (key uint64, cType byte, n int, length int, // Regardless, after this function runs, no containers have // mapped storage which does not refer to data; either they got mapped // to the new storage, or storage was allocated for them. -// -// Data should be in the Pilosa roaring format. func (b *Bitmap) RemapRoaringStorage(data []byte) (mappedAny bool, returnErr error) { if b.Containers == nil { return false, nil } - var itr *roaringIterator + var itr roaringIterator var err error var itrKey uint64 var itrCType byte @@ -1331,7 +1452,7 @@ func (b *Bitmap) ImportRoaringBits(data []byte, clear bool, log bool, rowSize ui if data == nil { return 0, nil, errors.New("no roaring bitmap provided") } - var itr *roaringIterator + var itr roaringIterator var itrKey uint64 var itrCType byte var itrN int