Add support for importing column attrs

api.go

@@ -1170,6 +1170,33 @@ func (api *API) ImportValue(ctx context.Context, req *ImportValueRequest, opts .
 
 }
 
+func (api *API) ImportColumnAttrs(ctx context.Context, req *ImportColumnAttrsRequest, opts ...ImportOption) error {
+	span, _ := tracing.StartSpanFromContext(ctx, "API.ImportColumnAttrs")
+	defer span.Finish()
+
+	index, err := api.Index(ctx, req.Index)
+	if err != nil {
+		return errors.Wrap(err, "getting index")
+	}
+
+	if err := api.validateShardOwnership(req.Index, uint64(req.Shard)); err != nil {
+		return errors.Wrap(err, "validating shard ownership")
+	}
+
+	bulkAttrs := make(map[uint64]map[string]interface{})
+	for n := 0; n < len(req.ColumnIDs); n++ {
+		bulkAttrs[uint64(req.ColumnIDs[n])] = map[string]interface{}{req.AttrKey: req.AttrVals[n]}
+	}
+	if err := index.ColumnAttrStore().SetBulkAttrs(bulkAttrs); err != nil {
+		return err
+	}
+
+	if err != nil {
+		api.server.logger.Printf("import error: index=%s, shard=%d, len(columns)=%d, err=%s", req.Index, req.Shard, len(req.ColumnIDs), err)
+	}
+	return errors.Wrap(err, "importing column attrs")
+}
+
 func importExistenceColumns(index *Index, columnIDs []uint64) error {
 	ef := index.existenceField()
 	if ef == nil {

api_test.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"math"
 	"reflect"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -30,6 +31,126 @@ import (
 	"github.com/pilosa/pilosa/v2/test"
 )
 
+// attrFun defines a mapping from columnID -> attr value
+func attrFun(id uint64) string {
+	//return fmt.Sprintf("%x", md5.Sum([]byte(strconv.FormatInt(int64(id), 10))))
+	return strconv.FormatInt(int64(id), 10)
+}
+
+func TestAPI_ImportColumnAttrs(t *testing.T) {
+	/*
+	   columns seconds
+	   100      1.150
+	   1000     1.568
+	   10000    5.156
+	   100000  38.179
+	*/
+	c := test.MustRunCluster(t, 2,
+		[]server.CommandOption{
+			server.OptCommandServerOptions(
+				pilosa.OptServerNodeID("node0"),
+				pilosa.OptServerClusterHasher(&offsetModHasher{}),
+			)},
+		[]server.CommandOption{
+			server.OptCommandServerOptions(
+				pilosa.OptServerNodeID("node1"),
+				pilosa.OptServerClusterHasher(&offsetModHasher{}),
+			)},
+	)
+	defer c.Close()
+
+	m0 := c[0]
+	m1 := c[1]
+	t.Run("ImportColumnAttrs", func(t *testing.T) {
+		ctx := context.Background()
+		index := "i"
+		field := "f"
+		attrKey := "columnid-md5"
+
+		_, err := m0.API.CreateIndex(ctx, index, pilosa.IndexOptions{})
+		if err != nil {
+			t.Fatalf("creating index: %v", err)
+		}
+		_, err = m0.API.CreateField(ctx, index, field)
+		if err != nil {
+			t.Fatalf("creating field: %v", err)
+		}
+
+		// Generate some attrs for two shards
+		columnIDs0 := make([]uint64, 0, 100)
+		attrVals0 := make([]string, 0, 100)
+		columnIDs1 := make([]uint64, 0, 100)
+		attrVals1 := make([]string, 0, 100)
+		for n := 0; n < 1000000; n += 10000 {
+			columnIDs0 = append(columnIDs0, uint64(n))
+			md50 := attrFun(uint64(n))
+			attrVals0 = append(attrVals0, md50)
+			setPql0 := fmt.Sprintf("Set(%d, %s=0) ", n, field)
+			m0.API.Query(ctx, &pilosa.QueryRequest{Index: index, Query: setPql0})
+
+			columnIDs1 = append(columnIDs1, uint64(n+ShardWidth))
+			md51 := attrFun(uint64(n + ShardWidth))
+			attrVals1 = append(attrVals1, md51)
+			setPql1 := fmt.Sprintf("Set(%d, %s=0) ", n+ShardWidth, field)
+			m1.API.Query(ctx, &pilosa.QueryRequest{Index: index, Query: setPql1})
+		}
+
+		// send shard0 to node1
+		req := &pilosa.ImportColumnAttrsRequest{
+			AttrKey:   attrKey,
+			ColumnIDs: columnIDs0,
+			AttrVals:  attrVals0,
+			Shard:     0,
+			Index:     index,
+		}
+
+		if err := m1.API.ImportColumnAttrs(ctx, req); err != nil {
+			t.Fatal(err)
+		}
+
+		// send shard1 to node0
+		req = &pilosa.ImportColumnAttrsRequest{
+			AttrKey:   attrKey,
+			ColumnIDs: columnIDs1,
+			AttrVals:  attrVals1,
+			Shard:     1,
+			Index:     index,
+		}
+
+		if err := m0.API.ImportColumnAttrs(ctx, req); err != nil {
+			t.Fatal(err)
+		}
+
+		// Query node0.
+		pql := fmt.Sprintf("Options(Row(%s=0), columnAttrs=true)", field)
+		res, err := m0.API.Query(ctx, &pilosa.QueryRequest{Index: index, Query: pql})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		for _, v := range res.ColumnAttrSets {
+			attrVal := attrFun(v.ID)
+			if attrVal != v.Attrs[attrKey] {
+				t.Fatal(err)
+			}
+		}
+		// Query node1.
+		pql = fmt.Sprintf("Options(Row(%s=0), columnAttrs=true)", field)
+		res, err = m1.API.Query(ctx, &pilosa.QueryRequest{Index: index, Query: pql})
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		for _, v := range res.ColumnAttrSets {
+			attrVal := attrFun(v.ID)
+			if attrVal != v.Attrs[attrKey] {
+				t.Fatal(err)
+			}
+		}
+
+	})
+}
+
 func TestAPI_Import(t *testing.T) {
 	c := test.MustRunCluster(t, 2,
 		[]server.CommandOption{

handler.go

@@ -147,6 +147,16 @@ func (i *ImportValueRequest) Validate() error {
 	return nil
 }
 
+// ImportColumnAttrsRequest describes the import request structure
+// for a ColumnAttr import
+type ImportColumnAttrsRequest struct {
+	AttrKey   string
+	ColumnIDs []uint64
+	AttrVals  []string
+	Shard     int64
+	Index     string
+}
+
 // ImportRequest describes the import request structure
 // for an import.
 type ImportRequest struct {

http/handler.go

@@ -286,6 +286,7 @@ func newRouter(handler *Handler) *mux.Router {
 	router.HandleFunc("/index/{index}", handler.handlePostIndex).Methods("POST").Name("PostIndex")
 	router.HandleFunc("/index/{index}", handler.handleDeleteIndex).Methods("DELETE").Name("DeleteIndex")
 	//router.HandleFunc("/index/{index}/field", handler.handleGetFields).Methods("GET") // Not implemented.
+	router.HandleFunc("/index/{index}/import-column-attrs", handler.handlePostImportColumnAttrs).Methods("POST").Name("PostImportColumnAttrs")
 	router.HandleFunc("/index/{index}/field/{field}", handler.handlePostField).Methods("POST").Name("PostField")
 	router.HandleFunc("/index/{index}/field/{field}", handler.handleDeleteField).Methods("DELETE").Name("DeleteField")
 	router.HandleFunc("/index/{index}/field/{field}/import", handler.handlePostImport).Methods("POST").Name("PostImport")
@@ -1621,6 +1622,50 @@ func GetHTTPClient(t *tls.Config) *http.Client {
 	return &http.Client{Transport: transport}
 }
 
+// handlePostImportColumnAttrs
+func (h *Handler) handlePostImportColumnAttrs(w http.ResponseWriter, r *http.Request) {
+	// Verify that request is only communicating over protobufs.
+	if r.Header.Get("Content-Type") != "application/x-protobuf" {
+		http.Error(w, "Unsupported media type", http.StatusUnsupportedMediaType)
+		return
+	} else if r.Header.Get("Accept") != "application/x-protobuf" {
+		http.Error(w, "Not acceptable", http.StatusNotAcceptable)
+		return
+	}
+
+	opts := []pilosa.ImportOption{}
+
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	req := &pilosa.ImportColumnAttrsRequest{}
+	if err := h.api.Serializer.Unmarshal(body, req); err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	if err := h.api.ImportColumnAttrs(r.Context(), req, opts...); err != nil {
+		http.Error(w, err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	// Marshal response object.
+	buf, e := h.api.Serializer.Marshal(&pilosa.ImportResponse{Err: ""})
+	if e != nil {
+		http.Error(w, fmt.Sprintf("marshal import-column-attrs response"), http.StatusInternalServerError)
+		return
+	}
+
+	// Write response.
+	_, err = w.Write(buf)
+	if err != nil {
+		h.logger.Printf("writing import-column-attrs response: %v", err)
+	}
+}
+
 // handlPostRoaringImport
 func (h *Handler) handlePostImportRoaring(w http.ResponseWriter, r *http.Request) {
 	// Verify that request is only communicating over protobufs.