Merge pull request #1601 from travisturner/import-key-values

Import key values

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 {

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
 }

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()

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]
 

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)
 	}

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.

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.