Support import column attrs in client

client.go

@@ -70,6 +70,7 @@ type InternalClient interface {
 	SendMessage(ctx context.Context, uri *URI, msg []byte) error
 	RetrieveShardFromURI(ctx context.Context, index, field, view string, shard uint64, uri URI) (io.ReadCloser, error)
 	ImportRoaring(ctx context.Context, uri *URI, index, field string, shard uint64, remote bool, req *ImportRoaringRequest) error
+	ImportColumnAttrs(ctx context.Context, uri *URI, index string, req *ImportColumnAttrsRequest) error
 }
 
 //===============
@@ -137,6 +138,11 @@ func (n nopInternalClient) ImportValue2(ctx context.Context, req *ImportValueReq
 func (n nopInternalClient) ImportRoaring(ctx context.Context, uri *URI, index, field string, shard uint64, remote bool, req *ImportRoaringRequest) error {
 	return nil
 }
+
+func (n nopInternalClient) ImportColumnAttrs(ctx context.Context, uri *URI, index string, req *ImportColumnAttrsRequest) error {
+	return nil
+}
+
 func (n nopInternalClient) EnsureIndex(ctx context.Context, name string, options IndexOptions) error {
 	return nil
 }

encoding/proto/proto.go

@@ -225,6 +225,14 @@ func (Serializer) Unmarshal(buf []byte, m pilosa.Message) error {
 		}
 		decodeImportRoaringRequest(msg, mt)
 		return nil
+	case *pilosa.ImportColumnAttrsRequest:
+		msg := &internal.ImportColumnAttrsRequest{}
+		err := proto.Unmarshal(buf, msg)
+		if err != nil {
+			return errors.Wrap(err, "unmarshaling ImportColumnAttrsRequest")
+		}
+		decodeImportColumnAttrsRequest(msg, mt)
+		return nil
 	case *pilosa.ImportResponse:
 		msg := &internal.ImportResponse{}
 		err := proto.Unmarshal(buf, msg)
@@ -318,6 +326,8 @@ func encodeToProto(m pilosa.Message) proto.Message {
 		return encodeImportValueRequest(mt)
 	case *pilosa.ImportRoaringRequest:
 		return encodeImportRoaringRequest(mt)
+	case *pilosa.ImportColumnAttrsRequest:
+		return encodeImportColumnAttrsRequest(mt)
 	case *pilosa.ImportResponse:
 		return encodeImportResponse(mt)
 	case *pilosa.BlockDataRequest:
@@ -395,6 +405,16 @@ func encodeImportRoaringRequest(m *pilosa.ImportRoaringRequest) *internal.Import
 	}
 }
 
+func encodeImportColumnAttrsRequest(m *pilosa.ImportColumnAttrsRequest) *internal.ImportColumnAttrsRequest {
+	return &internal.ImportColumnAttrsRequest{
+		Index:     m.Index,
+		Shard:     m.Shard,
+		AttrKey:   m.AttrKey,
+		AttrVals:  m.AttrVals,
+		ColumnIDs: m.ColumnIDs,
+	}
+}
+
 func encodeQueryRequest(m *pilosa.QueryRequest) *internal.QueryRequest {
 	r := &internal.QueryRequest{
 		Query:           m.Query,
@@ -1010,6 +1030,14 @@ func decodeImportRoaringRequest(pb *internal.ImportRoaringRequest, m *pilosa.Imp
 	m.Views = views
 }
 
+func decodeImportColumnAttrsRequest(pb *internal.ImportColumnAttrsRequest, m *pilosa.ImportColumnAttrsRequest) {
+	m.Index = pb.Index
+	m.Shard = pb.Shard
+	m.AttrKey = pb.AttrKey
+	m.AttrVals = pb.AttrVals
+	m.ColumnIDs = pb.ColumnIDs
+}
+
 func decodeImportResponse(pb *internal.ImportResponse, m *pilosa.ImportResponse) {
 	m.Err = pb.Err
 }

http/client.go

@@ -696,6 +696,56 @@ func (c *InternalClient) ImportRoaring(ctx context.Context, uri *pilosa.URI, ind
 	return nil
 }
 
+// ImportColumnAttrs does bulk import of column attrs
+func (c *InternalClient) ImportColumnAttrs(ctx context.Context, uri *pilosa.URI, index string, req *pilosa.ImportColumnAttrsRequest) error {
+	span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.ImportRoaring")
+	defer span.Finish()
+
+	if index == "" {
+		return pilosa.ErrIndexRequired
+	}
+	if uri == nil {
+		uri = c.defaultURI
+	}
+
+	url := fmt.Sprintf("%s/index/%s/import-column-attrs", uri, index)
+
+	// Marshal data to protobuf.
+	data, err := c.serializer.Marshal(req)
+	if err != nil {
+		return errors.Wrap(err, "marshal import-column-attrs request")
+	}
+
+	// Generate HTTP request.
+	httpReq, err := http.NewRequest("POST", url, bytes.NewBuffer(data))
+	if err != nil {
+		return errors.Wrap(err, "creating request")
+	}
+	httpReq.Header.Set("Content-Type", "application/x-protobuf")
+	httpReq.Header.Set("Accept", "application/x-protobuf")
+	httpReq.Header.Set("User-Agent", "pilosa/"+pilosa.Version)
+
+	// Execute request against the host.
+	resp, err := c.executeRequest(httpReq.WithContext(ctx))
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	dec := json.NewDecoder(resp.Body)
+	rbody := &pilosa.ImportResponse{}
+	err = dec.Decode(rbody)
+	// Decode can return EOF when no error occurred. helpful!
+	if err != nil && err != io.EOF {
+		return errors.Wrap(err, "decoding response body")
+	}
+	if rbody.Err != "" {
+		return errors.Wrap(errors.New(rbody.Err), "importing roaring")
+	}
+
+	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 {
 	span, ctx := tracing.StartSpanFromContext(ctx, "InternalClient.ExportCSV")

http/client_test.go

@@ -22,6 +22,7 @@ import (
 	"fmt"
 	gohttp "net/http"
 	"reflect"
+	"strconv"
 	"testing"
 	"time"
 
@@ -394,6 +395,60 @@ func TestClient_Import(t *testing.T) {
 	}
 }
 
+// Ensure client can bulk import column attrs.
+func TestClient_ImportColumnAttrs(t *testing.T) {
+	cluster := test.MustNewCluster(t, 2)
+	for _, c := range cluster {
+		c.Config.Cluster.ReplicaN = 2
+	}
+	err := cluster.Start()
+	if err != nil {
+		t.Fatalf("starting cluster: %v", err)
+	}
+	defer cluster.Close()
+
+	ctx := context.Background()
+	_, err = cluster[0].API.CreateIndex(ctx, "i", pilosa.IndexOptions{})
+	if err != nil {
+		t.Fatalf("creating index: %v", err)
+	}
+	_, err = cluster[0].API.CreateField(ctx, "i", "f", pilosa.OptFieldTypeSet(pilosa.CacheTypeRanked, 100))
+	if err != nil {
+		t.Fatalf("creating field: %v", err)
+	}
+	_, err = cluster[0].API.Query(ctx, &pilosa.QueryRequest{Index: "i", Query: "Set(0, f=0) Set(1, f=0) Set(2, f=0) Set(3, f=0) Set(4, f=0)"})
+	if err != nil {
+		t.Fatalf("querying: %v", err)
+	}
+
+	attrKey := "k"
+	// Send import request.
+	host := cluster[0].URL()
+	c := MustNewClient(host, http.GetHTTPClient(nil))
+	colAttrsReq := makeImportColumnAttrsRequest("i", 0, attrKey)
+	if err := c.ImportColumnAttrs(ctx, &cluster[1].API.Node().URI, "i", colAttrsReq); err != nil {
+		t.Fatal(err)
+	}
+
+	// Verify data.
+	pql := "Options(Row(f=0), columnAttrs=true)"
+	res, err := cluster[1].API.Query(ctx, &pilosa.QueryRequest{Index: "i", Query: pql})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(res.ColumnAttrSets) != 5 {
+		t.Fatal("incorrect number of column attrs set")
+	}
+
+	for _, v := range res.ColumnAttrSets {
+		attrVal := attrFun(v.ID)
+		if attrVal != v.Attrs[attrKey] {
+			t.Fatal(err)
+		}
+	}
+
+}
+
 // Ensure client can bulk import data.
 func TestClient_ImportRoaring(t *testing.T) {
 	cluster := test.MustNewCluster(t, 2)
@@ -1195,3 +1250,23 @@ func makeImportRoaringRequest(clear bool, viewData string) *pilosa.ImportRoaring
 		},
 	}
 }
+
+func attrFun(id uint64) string {
+	return strconv.FormatInt(int64(id), 10)
+}
+
+func makeImportColumnAttrsRequest(index string, shard int64, attrKey string) *pilosa.ImportColumnAttrsRequest {
+	colIDs := make([]uint64, 0, 5)
+	attrVals := make([]string, 0, 5)
+	for n := uint64(0); n < 5; n++ {
+		colIDs = append(colIDs, n)
+		attrVals = append(attrVals, attrFun(n))
+	}
+	return &pilosa.ImportColumnAttrsRequest{
+		Index:     index,
+		Shard:     shard,
+		AttrKey:   attrKey,
+		ColumnIDs: colIDs,
+		AttrVals:  attrVals,
+	}
+}

internal/public.proto

@@ -130,3 +130,11 @@ message ImportRoaringRequest {
 	bool Clear = 1;
 	repeated ImportRoaringRequestView views = 2;
 }
+
+message ImportColumnAttrsRequest {
+	string Index = 1;
+	int64 Shard = 2;
+	string AttrKey = 3;
+	repeated string AttrVals = 4;
+	repeated uint64 ColumnIDs = 5;
+}