diff --git a/executor.go b/executor.go index a0b0b75b4..0d7241ea5 100644 --- a/executor.go +++ b/executor.go @@ -3851,11 +3851,13 @@ func (e *executor) executeClearRowShard(ctx context.Context, tx Tx, index string // executeSetRow executes a Store() call. func (e *executor) executeSetRow(ctx context.Context, tx Tx, indexName string, c *pql.Call, shards []uint64, opt *execOptions) (bool, error) { - // Ensure the field type supports Store(). + // Parse arguments. fieldName, err := c.FieldArg() if err != nil { return false, errors.New("field required for Store()") } + argKey, argKeyed := c.Args[fieldName].(string) + field := e.Holder.Field(indexName, fieldName) if field == nil { // Find index. @@ -3865,17 +3867,34 @@ func (e *executor) executeSetRow(ctx context.Context, tx Tx, indexName string, c } // Create field. - field, err = index.CreateField(fieldName, OptFieldTypeSet(CacheTypeNone, 0)) + opts := []FieldOption{OptFieldTypeSet(CacheTypeNone, 0)} + if argKeyed { + opts = append(opts, OptFieldKeys()) + } + field, err = index.CreateField(fieldName, opts...) if err != nil { // We wrap these because we want to indicate that it wasn't found, // but also the problem we encountered trying to create it. return false, newNotFoundError(errors.Wrap(err, "creating field")) } } + // Ensure the field type supports Store(). if field.Type() != FieldTypeSet { return false, fmt.Errorf("can't Store() on a %s field", field.Type()) } + if argKeyed { + // Translate the argument key. + if !field.Keys() { + return false, errors.New("can't Store() a key into an unkeyed field") + } + id, err := field.TranslateStore().TranslateKey(argKey) + if err != nil { + return false, errors.Wrap(err, "translating a key for Store()") + } + c.Args[fieldName] = id + } + // Execute calls in bulk on each remote node and merge. mapFn := func(ctx context.Context, shard uint64) (interface{}, error) { return e.executeSetRowShard(ctx, tx, indexName, c, shard) diff --git a/executor_test.go b/executor_test.go index 3abaebf46..5af2712be 100644 --- a/executor_test.go +++ b/executor_test.go @@ -4291,6 +4291,54 @@ func TestExecutor_Execute_SetRow(t *testing.T) { t.Fatalf("unexpected columns: %+v", bits) } }) + t.Run("Set_Keyed", func(t *testing.T) { + c := test.MustRunCluster(t, 1) + defer c.Close() + hldr := c.GetHolder(0) + index := hldr.MustCreateIndexIfNotExists("i", pilosa.IndexOptions{TrackExistence: true}) + if _, err := index.CreateField("f", pilosa.OptFieldTypeDefault(), pilosa.OptFieldKeys()); err != nil { + t.Fatal(err) + } + + // Set bits. + if _, err := c.GetNode(0).API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Set(1, f="a")`}); err != nil { + t.Fatal(err) + } + + if res, err := c.GetNode(0).API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Row(f="a")`}); err != nil { + t.Fatal(err) + } else if bits := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(bits, []uint64{1}) { + t.Fatalf("unexpected columns: %+v", bits) + } + + // Store row a into a different row. + if res, err := c.GetNode(0).API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Store(Row(f="a"), f="b")`}); err != nil { + t.Fatal(err) + } else if res := res.Results[0].(bool); !res { + t.Fatalf("unexpected set row result: %+v", res) + } + + // Ensure the row was populated. + if res, err := c.GetNode(0).API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Row(f="b")`}); err != nil { + t.Fatal(err) + } else if bits := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(bits, []uint64{1}) { + t.Fatalf("unexpected columns: %+v", bits) + } + + // Store row 10 into a table which doesn't exist. + if res, err := c.GetNode(0).API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Store(Row(f="a"), nonexistent="c")`}); err != nil { + t.Fatal(err) + } else if res := res.Results[0].(bool); !res { + t.Fatalf("unexpected set row result: %+v", res) + } + + // Ensure the row was populated. + if res, err := c.GetNode(0).API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Row(nonexistent="c")`}); err != nil { + t.Fatal(err) + } else if bits := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(bits, []uint64{1}) { + t.Fatalf("unexpected columns: %+v", bits) + } + }) } func benchmarkExistence(nn bool, b *testing.B) {