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/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.