diff --git a/api.go b/api.go index 708ca61ea..14ab78969 100644 --- a/api.go +++ b/api.go @@ -664,10 +664,26 @@ func (api *API) ImportValue(_ context.Context, req *ImportValueRequest) error { return errors.Wrap(err, "validating api method") } + index := api.holder.Index(req.Index) + if index == nil { + return newNotFoundError(ErrIndexNotFound) + } + field, err := api.indexField(req.Index, req.Field, req.Shard) if err != nil { return errors.Wrap(err, "getting field") } + + // Translate column keys. + if index.Keys() { + if len(req.ColumnIDs) != 0 { + return errors.New("column ids cannot be used because index uses string keys") + } + if req.ColumnIDs, err = api.holder.translateFile.TranslateColumnsToUint64(index.Name(), req.ColumnKeys); err != nil { + return errors.Wrap(err, "translating columns") + } + } + // Import into fragment. err = field.importValue(req.ColumnIDs, req.Values) if err != nil { diff --git a/client.go b/client.go index 0f64ba8fc..9c5016646 100644 --- a/client.go +++ b/client.go @@ -42,6 +42,7 @@ type InternalClient interface { EnsureIndex(ctx context.Context, name string, options IndexOptions) error EnsureField(ctx context.Context, indexName string, fieldName string) error ImportValue(ctx context.Context, index, field string, shard uint64, vals []FieldValue) error + ImportValueK(ctx context.Context, index, field string, vals []FieldValue) error ExportCSV(ctx context.Context, index, field string, shard uint64, w io.Writer) error CreateField(ctx context.Context, index, field string) error FragmentBlocks(ctx context.Context, uri *URI, index, field string, shard uint64) ([]FragmentBlock, error) @@ -114,6 +115,9 @@ func (n nopInternalClient) EnsureField(ctx context.Context, indexName string, fi func (n nopInternalClient) ImportValue(ctx context.Context, index, field string, shard uint64, vals []FieldValue) error { return nil } +func (n nopInternalClient) ImportValueK(ctx context.Context, index, field string, vals []FieldValue) error { + return nil +} func (n nopInternalClient) ExportCSV(ctx context.Context, index, field string, shard uint64, w io.Writer) error { return nil } diff --git a/ctl/import.go b/ctl/import.go index e8d6c8f37..2969b581b 100644 --- a/ctl/import.go +++ b/ctl/import.go @@ -340,7 +340,7 @@ func (cmd *ImportCommand) bufferValues(ctx context.Context, useColumnKeys bool, // If we've reached the buffer size then import FieldValues. if len(a) == cmd.BufferSize { - if err := cmd.importValues(ctx, a); err != nil { + if err := cmd.importValues(ctx, useColumnKeys, a); err != nil { return err } a = a[:0] @@ -348,13 +348,22 @@ func (cmd *ImportCommand) bufferValues(ctx context.Context, useColumnKeys bool, } // If there are still values in the buffer then flush them. - return cmd.importValues(ctx, a) + return cmd.importValues(ctx, useColumnKeys, a) } // importValues sends batches of FieldValues to the server. -func (cmd *ImportCommand) importValues(ctx context.Context, vals []pilosa.FieldValue) error { +func (cmd *ImportCommand) importValues(ctx context.Context, useColumnKeys bool, vals []pilosa.FieldValue) error { logger := log.New(cmd.Stderr, "", log.LstdFlags) + // If keys are used, all values are sent to the primary translate store (i.e. coordinator). + if useColumnKeys { + logger.Printf("importing keyed values: n=%d", len(vals)) + if err := cmd.client.ImportValueK(ctx, cmd.Index, cmd.Field, vals); err != nil { + return errors.Wrap(err, "importing keys") + } + return nil + } + // Group vals by shard. logger.Printf("grouping %d vals", len(vals)) valsByShard := http.FieldValues(vals).GroupByShard() diff --git a/ctl/import_test.go b/ctl/import_test.go index c2f7893f4..cd75a5d0a 100644 --- a/ctl/import_test.go +++ b/ctl/import_test.go @@ -128,6 +128,33 @@ func TestImportCommand_RunKeys(t *testing.T) { } } +// Ensure that integer import with keys runs. +func TestImportCommand_RunValueKeys(t *testing.T) { + buf := bytes.Buffer{} + stdin, stdout, stderr := GetIO(buf) + cm := NewImportCommand(stdin, stdout, stderr) + file, err := ioutil.TempFile("", "import-key.csv") + file.Write([]byte("foo1,2\nfoo3,4\nfoo5,6")) + ctx := context.Background() + if err != nil { + t.Fatal(err) + } + + cmd := test.MustRunCluster(t, 1)[0] + cm.Host = cmd.API.Node().URI.HostPort() + + http.DefaultClient.Do(MustNewHTTPRequest("POST", "http://"+cm.Host+"/index/i", strings.NewReader(`{"options":{"keys": true}}`))) + http.DefaultClient.Do(MustNewHTTPRequest("POST", "http://"+cm.Host+"/index/i/field/f", strings.NewReader(`{"options":{"type": "int", "min": 0, "max": 100}}`))) + + cm.Index = "i" + cm.Field = "f" + cm.Paths = []string{file.Name()} + err = cm.Run(ctx) + if err != nil { + t.Fatalf("Import Run with keys doesn't work: %s", err) + } +} + func TestImportCommand_InvalidFile(t *testing.T) { cmd := test.MustRunCluster(t, 1)[0] diff --git a/encoding/proto/proto.go b/encoding/proto/proto.go index fcdde95f9..298683e92 100644 --- a/encoding/proto/proto.go +++ b/encoding/proto/proto.go @@ -890,6 +890,7 @@ func decodeRow(pr *internal.Row) *pilosa.Row { r := pilosa.NewRow() r.Attrs = decodeAttrs(pr.Attrs) + r.Keys = pr.Keys for _, v := range pr.Columns { r.SetBit(v) } diff --git a/http/client.go b/http/client.go index e5c270cbb..0bfdbdddd 100644 --- a/http/client.go +++ b/http/client.go @@ -473,6 +473,38 @@ func (c *InternalClient) ImportValue(ctx context.Context, index, field string, s return nil } +// 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) + } + + // Get the coordinator node; all bits are sent to the + // primary translate store (i.e. coordinator). + nodes, err := c.Nodes(ctx) + if err != nil { + return fmt.Errorf("getting nodes: %s", err) + } + coord := getCoordinatorNode(nodes) + if coord == nil { + return fmt.Errorf("could not find the coordinator node") + } + + // Import to node. + if err := c.importNode(ctx, coord, index, field, buf); err != nil { + return fmt.Errorf("import node: host=%s, err=%s", coord.URI, err) + } + + return nil +} + // marshalImportValuePayload marshalls the import parameters into a protobuf byte slice. func (c *InternalClient) marshalImportValuePayload(index, field string, shard uint64, vals []pilosa.FieldValue) ([]byte, error) { // Separate row and column IDs to reduce allocations. diff --git a/http/client_test.go b/http/client_test.go index 80c4d0425..702860c23 100644 --- a/http/client_test.go +++ b/http/client_test.go @@ -129,11 +129,6 @@ func TestClient_MultiNode(t *testing.T) { Remote: false, } - _, err = client[0].Query(context.Background(), "i", queryRequest) - if err != nil { - t.Fatal(err) - } - result, err := client[0].Query(context.Background(), "i", queryRequest) if err != nil { t.Fatal(err) @@ -376,6 +371,56 @@ func TestClient_ImportKeys(t *testing.T) { } }) }) + + t.Run("IntegerFieldSingleNode", func(t *testing.T) { + cmd := test.MustRunCluster(t, 1)[0] + host := cmd.URL() + holder := cmd.Server.Holder() + hldr := test.Holder{Holder: holder} + + fldName := "f" + + // Load bitmap into cache to ensure cache gets updated. + index := hldr.MustCreateIndexIfNotExists("i", pilosa.IndexOptions{Keys: true}) + field, err := index.CreateFieldIfNotExists(fldName, pilosa.OptFieldTypeInt(-100, 100)) + if err != nil { + t.Fatal(err) + } + + // Send import request. + c := MustNewClient(host, http.GetHTTPClient(nil)) + if err := c.ImportValue(context.Background(), "i", "f", 0, []pilosa.FieldValue{ + {ColumnKey: "col1", Value: -10}, + {ColumnKey: "col2", Value: 20}, + {ColumnKey: "col3", Value: 40}, + }); err != nil { + t.Fatal(err) + } + + // Verify Sum. + sum, cnt, err := field.Sum(nil, fldName) + if err != nil { + t.Fatal(err) + } + if sum != 50 || cnt != 3 { + t.Fatalf("unexpected values: got sum=%v, count=%v; expected sum=50, cnt=3", sum, cnt) + } + + // Verify Range + queryRequest := &pilosa.QueryRequest{ + Query: fmt.Sprintf(`Range(%s>10)`, fldName), + Remote: false, + } + + result, err := c.Query(context.Background(), "i", queryRequest) + if err != nil { + t.Fatal(err) + } + + if !reflect.DeepEqual(result.Results[0].(*pilosa.Row).Keys, []string{"col2", "col3"}) { + t.Fatalf("unexpected column keys: %s", spew.Sdump(result)) + } + }) } // Ensure client can bulk import value data.