From aba67364b1be3a3cc88e5ba01463aa2dba59f8a8 Mon Sep 17 00:00:00 2001 From: Alan Bernstein Date: Fri, 29 Nov 2019 12:47:38 -0600 Subject: [PATCH] Add support for importing column attrs --- api.go | 27 +++++++++++ api_test.go | 121 ++++++++++++++++++++++++++++++++++++++++++++++++ handler.go | 10 ++++ http/handler.go | 45 ++++++++++++++++++ 4 files changed, 203 insertions(+) diff --git a/api.go b/api.go index 328d4d07c..8c212d061 100644 --- a/api.go +++ b/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 { diff --git a/api_test.go b/api_test.go index d412def1a..3ba15c4a0 100644 --- a/api_test.go +++ b/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{ diff --git a/handler.go b/handler.go index e99dbed69..19fcbac44 100644 --- a/handler.go +++ b/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 { diff --git a/http/handler.go b/http/handler.go index b9e905a81..c575afd31 100644 --- a/http/handler.go +++ b/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.