initial support for bulk importing standard roaring files per shard

api.go

@@ -1,4 +1,3 @@
-// Copyright 2017 Pilosa Corp.
 //
 // Licensed under the Apache License, Version 2.0 (the "License");
 // you may not use this file except in compliance with the License.
@@ -24,6 +23,7 @@ import (
 	"io/ioutil"
 	"strconv"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/pilosa/pilosa/pql"
@@ -294,6 +294,42 @@ func (api *API) Field(_ context.Context, indexName, fieldName string) (*Field, e
 	return field, nil
 }
 
+// ImportRoaringBytes fast loading of standard roaring format
+func (api *API) ImportRoaringBytes(ctx context.Context, roaringBytes []byte, indexName, fieldName string, shard uint64, forward bool) (err error) {
+	if err = api.validate(apiField); err != nil {
+		err = errors.Wrap(err, "validating api method")
+		return
+	}
+	nodes := api.cluster.shardNodes(indexName, shard)
+	var wg sync.WaitGroup
+
+	for _, node := range nodes {
+		if node.ID == api.server.nodeID {
+			field := api.holder.Field(indexName, fieldName)
+			if field == nil {
+				err = newNotFoundError(ErrFieldNotFound)
+				return
+			}
+			wg.Add(1)
+			go func(node *Node) {
+				err = field.importRoaringBytes(roaringBytes, shard)
+				wg.Done()
+			}(node)
+		} else if forward {
+			wg.Add(1)
+			//forward it on
+			go func(node *Node) {
+				//execute on node
+				err = api.server.defaultClient.ImportRoaringBytes(ctx, node, indexName, fieldName, shard, roaringBytes, forward)
+				wg.Done()
+			}(node)
+		}
+
+	}
+	wg.Wait()
+	return
+}
+
 // DeleteField removes the named field from the named index. If the index is not
 // found, an error is returned. If the field is not found, it is ignored and no
 // action is taken.

client.go

@@ -51,6 +51,7 @@ type InternalClient interface {
 	RowAttrDiff(ctx context.Context, uri *URI, index, field string, blks []AttrBlock) (map[uint64]map[string]interface{}, error)
 	SendMessage(ctx context.Context, uri *URI, msg []byte) error
 	RetrieveShardFromURI(ctx context.Context, index, field string, shard uint64, uri URI) (io.ReadCloser, error)
+	ImportRoaringBytes(ctx context.Context, node *Node, index, field string, shard uint64, roaringBytes []byte, forward bool) error
 }
 
 //===============
@@ -106,6 +107,9 @@ func (n nopInternalClient) Import(ctx context.Context, index, field string, shar
 func (n nopInternalClient) ImportK(ctx context.Context, index, field string, bits []Bit) error {
 	return nil
 }
+func (n nopInternalClient) ImportRoaringBytes(ctx context.Context, node *Node, index, field string, shard uint64, roaringBytes []byte, forward bool) error {
+	return nil
+}
 func (n nopInternalClient) EnsureIndex(ctx context.Context, name string, options IndexOptions) error {
 	return nil
 }

field.go

@@ -1055,6 +1055,26 @@ func (f *Field) importValue(columnIDs []uint64, values []int64) error {
 	return nil
 }
 
+func (f *Field) importRoaringBytes(stdRoaringBytes []byte, shard uint64) error {
+	viewName := viewStandard
+
+	view, err := f.createViewIfNotExists(viewName)
+	if err != nil {
+		return errors.Wrap(err, "creating view")
+	}
+
+	frag, err := view.CreateFragmentIfNotExists(shard)
+	if err != nil {
+		return errors.Wrap(err, "creating fragment")
+	}
+
+	if err := frag.importRoaringBytes(stdRoaringBytes); err != nil {
+		return err
+	}
+
+	return nil
+}
+
 type fieldSlice []*Field
 
 func (p fieldSlice) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }

fragment.go

@@ -1327,69 +1327,66 @@ func (f *fragment) mergeBlock(id int, data []pairSet) (sets, clears []pairSet, e
 func (f *fragment) bulkImport(rowIDs, columnIDs []uint64) error {
 	f.mu.Lock()
 	defer f.mu.Unlock()
+
 	// Verify that there are an equal number of row ids and column ids.
 	if len(rowIDs) != len(columnIDs) {
 		return fmt.Errorf("mismatch of row/column len: %d != %d", len(rowIDs), len(columnIDs))
 	}
 
 	// Disconnect op writer so we don't append updates.
-	f.storage.OpWriter = nil
-
+	localBitmap := roaring.NewBitmap()
+	localBitmap.OpWriter = nil
 	// Process every bit.
 	// If an error occurs then reopen the storage.
 	lastID := uint64(0)
-	if err := func() error {
-		set := make(map[uint64]struct{})
-		for i := range rowIDs {
-			rowID, columnID := rowIDs[i], columnIDs[i]
+	set := make(map[uint64]struct{})
+	for i := range rowIDs {
+		rowID, columnID := rowIDs[i], columnIDs[i]
 
-			// Determine the position of the bit in the storage.
-			pos, err := f.pos(rowID, columnID)
-			if err != nil {
-				return errors.Wrap(err, "getting bit pos")
-			}
-
-			// Write to storage.
-			_, err = f.storage.Add(pos)
-			if err != nil {
-				return errors.Wrap(err, "writing")
-			}
-			// Reduce the StatsD rate for high volume stats
-			f.stats.Count("ImportBit", 1, 0.0001)
-			// import optimization to avoid linear foreach calls
-			// slight risk of concurrent cache counter being off but
-			// no real danger
-			if i == 0 || rowID != lastID {
-				lastID = rowID
-				set[rowID] = struct{}{}
-			}
-
-			// Invalidate block checksum.
-			delete(f.checksums, int(rowID/HashBlockSize))
+		// Determine the position of the bit in the storage.
+		pos, err := f.pos(rowID, columnID)
+		if err != nil {
+			return err
 		}
 
-		// Update cache counts for all rows.
-		for rowID := range set {
-			// Import should ALWAYS have row() load a new row from fragment.storage
-			// because the row that's in rowCache hasn't been updated with
-			// this import's data.
-			f.cache.BulkAdd(rowID, f.unprotectedRow(rowID, false, false).Count())
+		// Write to storage.
+		_, err = localBitmap.Add(pos)
+		if err != nil {
+			return err
+		}
+		// Reduce the StatsD rate for high volume stats
+		f.stats.Count("ImportBit", 1, 0.0001)
+		// import optimization to avoid linear foreach calls
+		// slight risk of concurrent cache counter being off but
+		// no real danger
+		if i == 0 || rowID != lastID {
+			lastID = rowID
+			set[rowID] = struct{}{}
 		}
 
-		f.cache.Invalidate()
-		return nil
-	}(); err != nil {
-		_ = f.closeStorage()
-		_ = f.openStorage()
-		return err
+		// Invalidate block checksum.
+		delete(f.checksums, int(rowID/HashBlockSize))
 	}
 
-	// Write the storage to disk and reload.
-	if err := f.snapshot(); err != nil {
-		return errors.Wrap(err, "snapshotting")
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	//f.storage.Unmmap()
+	// Update cache counts for all rows.
+	var results *roaring.Bitmap
+	if f.storage.Count() > 0 {
+		results = f.storage.Union(localBitmap)
+	} else {
+		results = localBitmap
 	}
 
-	return nil
+	for rowID := range set {
+		n := results.CountRange(rowID*ShardWidth, (rowID+1)*ShardWidth)
+		f.cache.BulkAdd(rowID, n)
+	}
+
+	f.cache.Invalidate()
+	return snapshot(f, results)
 }
 
 // importValue bulk imports a set of range-encoded values.
@@ -1425,6 +1422,19 @@ func (f *fragment) importValue(columnIDs, values []uint64, bitDepth uint) error
 	return nil
 }
 
+// importRoaringBytes imports from standard roaring data format defined at
+// https://github.com/RoaringBitmap/RoaringFormatSpec
+func (f *fragment) importRoaringBytes(roaringBytes []byte) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	bm, err := roaring.UnmarshalStandardRoaring(roaringBytes)
+	if err != nil {
+		return err
+	}
+	err = snapshot(f, bm)
+	return err
+}
+
 // incrementOpN increase the operation count by one.
 // If the count exceeds the maximum allowed then a snapshot is performed.
 func (f *fragment) incrementOpN() error {
@@ -1452,6 +1462,11 @@ func track(start time.Time, message string, stats StatsClient, logger Logger) {
 }
 
 func (f *fragment) snapshot() error {
+	return snapshot(f, f.storage)
+}
+
+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)
 	completeMessage := fmt.Sprintf("fragment: snapshot complete %s/%s/%s/%d", f.index, f.field, f.view, f.shard)
 	start := time.Now()
@@ -1467,7 +1482,7 @@ func (f *fragment) snapshot() error {
 
 	// Write storage to snapshot.
 	bw := bufio.NewWriter(file)
-	if _, err := f.storage.WriteTo(bw); err != nil {
+	if _, err := bm.WriteTo(bw); err != nil {
 		return fmt.Errorf("snapshot write to: %s", err)
 	}
 

http/client.go

@@ -475,12 +475,6 @@ func (c *InternalClient) ImportValue(ctx context.Context, index, field string, s
 
 // ImportValueK bulk imports keyed field values to a host.
 func (c *InternalClient) ImportValueK(ctx context.Context, index, field string, vals []pilosa.FieldValue) error {
-	if index == "" {
-		return pilosa.ErrIndexRequired
-	} else if field == "" {
-		return pilosa.ErrFieldRequired
-	}
-
 	buf, err := c.marshalImportValuePayload(index, field, 0, vals)
 	if err != nil {
 		return fmt.Errorf("Error Creating Payload: %s", err)
@@ -527,6 +521,46 @@ func (c *InternalClient) marshalImportValuePayload(index, field string, shard ui
 	return buf, nil
 }
 
+// ImportRoaringBytes fast import of raw bits in roaring standard format
+func (c *InternalClient) ImportRoaringBytes(ctx context.Context, node *pilosa.Node, index, field string, shard uint64, roaringBytes []byte, forward bool) error {
+	if index == "" {
+		return pilosa.ErrIndexRequired
+	} else if field == "" {
+		return pilosa.ErrFieldRequired
+	}
+
+	endpoint := fmt.Sprintf("/index/%s/field/%s/importroaring/%d", index, field, shard)
+
+	// Create URL.
+	u := nodePathToURL(node, endpoint)
+	if forward {
+		v := url.Values{}
+		v.Set("noforward", "y")
+		u.RawQuery = v.Encode()
+	}
+	// Generate HTTP request.
+	req, err := http.NewRequest("POST", u.String(), bytes.NewBuffer(roaringBytes))
+	if err != nil {
+		return errors.Wrap(err, "creating request")
+	}
+	req.Header.Set("Accept", "application/x-binary")
+	req.Header.Set("User-Agent", "pilosa/"+pilosa.Version)
+
+	// Execute request against the host.
+	resp, err := c.httpClient.Do(req.WithContext(ctx))
+	if err != nil {
+		return errors.Wrap(err, "executing request")
+	}
+	defer resp.Body.Close()
+
+	// Validate status code.
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("invalid status: %d", resp.StatusCode)
+	}
+
+	return nil
+}
+
 // ExportCSV bulk exports data for a single shard from a host to CSV format.
 func (c *InternalClient) ExportCSV(ctx context.Context, index, field string, shard uint64, w io.Writer) error {
 	if index == "" {

http/handler.go

@@ -177,6 +177,7 @@ func (h *Handler) populateValidators() {
 	h.validators["GetFragmentData"] = queryValidationSpecRequired("index", "field", "shard")
 	h.validators["PostFragmentData"] = queryValidationSpecRequired("index", "field", "shard")
 	h.validators["GetFragmentBlocks"] = queryValidationSpecRequired("index", "field", "view", "shard")
+	h.validators["ImportRoaringBitmap"] = queryValidationSpecRequired("index", "field", "shard")
 }
 
 func (h *Handler) queryArgValidator(next http.Handler) http.Handler {
@@ -217,6 +218,7 @@ func newRouter(handler *Handler) *mux.Router {
 	router.HandleFunc("/index/{index}/field/{field}", handler.handlePostField).Methods("POST")
 	router.HandleFunc("/index/{index}/field/{field}", handler.handleDeleteField).Methods("DELETE")
 	router.HandleFunc("/index/{index}/field/{field}/import", handler.handlePostImport).Methods("POST")
+	router.HandleFunc("/index/{index}/field/{field}/importroaring/{shard}", handler.handlePostRoaringImport).Methods("POST")
 	router.HandleFunc("/index/{index}/query", handler.handlePostQuery).Methods("POST").Name("PostQuery")
 	router.HandleFunc("/info", handler.handleGetInfo).Methods("GET")
 	router.HandleFunc("/recalculate-caches", handler.handleRecalculateCaches).Methods("POST")
@@ -1427,3 +1429,49 @@ func GetHTTPClient(t *tls.Config) *http.Client {
 	}
 	return &http.Client{Transport: transport}
 }
+
+// handlPostRoaringImport
+func (h *Handler) handlePostRoaringImport(w http.ResponseWriter, r *http.Request) {
+	// Verify that request is only communicating over protobufs.
+	if r.Header.Get("Content-Type") != "application/x-binary" {
+		http.Error(w, "Unsupported media type", http.StatusUnsupportedMediaType)
+		return
+	} /*else if validHeaderAcceptJSON(r.Header) {
+		http.Error(w, "Not acceptable", http.StatusNotAcceptable)
+		return
+	}*/
+
+	indexName := mux.Vars(r)["index"]
+	fieldName := mux.Vars(r)["field"]
+	shardName := mux.Vars(r)["shard"]
+	_, noForward := r.URL.Query()["noforward"]
+
+	// Read entire body.
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	shard, err := strconv.ParseUint(shardName, 10, 64)
+	if err != nil {
+		http.Error(w, "shard should be an unsigned integer", http.StatusBadRequest)
+		return
+	}
+	//TODO give meaningful stats for import
+	err = h.api.ImportRoaringBytes(r.Context(), body, indexName, fieldName, shard, !noForward)
+
+	// Marshal response object.
+	msg := string("")
+	if err != nil {
+		msg = err.Error()
+	}
+	buf, e := h.api.Serializer.Marshal(&pilosa.ImportResponse{Err: msg})
+	if e != nil {
+		http.Error(w, fmt.Sprintf("marshal import response"), http.StatusInternalServerError)
+		return
+	}
+
+	// Write response.
+	w.Write(buf)
+}

roaring/roaring.go

@@ -3408,7 +3408,7 @@ func readStandardHeader(buf []byte) (size uint32, containerTyper func(index uint
 	return
 }
 
-func unmarshalStandardRoaring(data []byte) (*Bitmap, error) {
+func UnmarshalStandardRoaring(data []byte) (*Bitmap, error) {
 
 	b := NewBitmap()
 	keyN, containerTyper, header, pos, err, haveRuns := readStandardHeader(data)
@@ -3482,10 +3482,7 @@ func readWithRuns(b *Bitmap, data []byte, pos int, keyN uint32) (err error) {
 			c.runs = (*[0xFFFFFFF]interval16)(unsafe.Pointer(&data[pos+runCountHeaderSize]))[:runCount]
 
 			for o := range c.runs { //need to convert to start:end vs start:length :(
-				r := c.runs[o]
-				r.last = r.start + r.last
-				c.runs[o] = r
-
+				c.runs[o].last = c.runs[o].start + c.runs[o].last
 			}
 			pos += int((runCount * interval16Size) + runCountHeaderSize)
 		case containerArray: