diff --git a/executor.go b/executor.go index 29da40b27..235000938 100644 --- a/executor.go +++ b/executor.go @@ -118,17 +118,6 @@ func (e *executor) Execute(ctx context.Context, index string, q *pql.Query, shar resp.Results = results - // Translate response objects from ids to keys, if necessary. - // No need to translate a remote call. - if !opt.Remote { - for i := range results { - results[i], err = e.translateResult(index, idx, q.Calls[i], results[i]) - if err != nil { - return resp, err - } - } - } - // Fill column attributes if requested. if opt.ColumnAttrs { // Consolidate all column ids across all calls. @@ -148,8 +137,8 @@ func (e *executor) Execute(ctx context.Context, index string, q *pql.Query, shar } // Translate column attributes, if necessary. - if e.Holder.translateFile != nil { - for _, col := range resp.ColumnAttrSets { + if idx.Keys() { + for _, col := range columnAttrSets { v, err := e.Holder.translateFile.TranslateColumnToString(index, col.ID) if err != nil { return resp, err @@ -161,6 +150,17 @@ func (e *executor) Execute(ctx context.Context, index string, q *pql.Query, shar resp.ColumnAttrSets = columnAttrSets } + // Translate response objects from ids to keys, if necessary. + // No need to translate a remote call. + if !opt.Remote { + for i := range results { + results[i], err = e.translateResult(index, idx, q.Calls[i], results[i]) + if err != nil { + return resp, err + } + } + } + return resp, nil } @@ -1765,9 +1765,16 @@ func (e *executor) mapperLocal(ctx context.Context, shards []uint64, mapFn mapFu } } +var translateCallCol = map[string]struct{}{ + "Set": struct{}{}, + "Clear": struct{}{}, + "Row": struct{}{}, + "SetColumnAttrs": struct{}{}, +} + func (e *executor) translateCall(index string, idx *Index, c *pql.Call) error { var colKey, rowKey, fieldName string - if c.Name == "Set" || c.Name == "Clear" || c.Name == "Row" { + if _, ok := translateCallCol[c.Name]; ok { // Positional args in new PQL syntax require special handling here. colKey = "_" + columnLabel fieldName, _ = c.FieldArg() diff --git a/executor_test.go b/executor_test.go index d6d529efe..5a5f68ff9 100644 --- a/executor_test.go +++ b/executor_test.go @@ -1509,6 +1509,36 @@ func TestExecutor_QueryCall(t *testing.T) { } }) + t.Run("columnAttrsWithKeys", func(t *testing.T) { + c := test.MustRunCluster(t, 1) + defer c.Close() + hldr := test.Holder{Holder: c[0].Server.Holder()} + + // Set columns for rows 0, 10, & 20 across two shards. + if idx, err := hldr.CreateIndex("i", pilosa.IndexOptions{Keys: true}); err != nil { + t.Fatal(err) + } else if _, err := idx.CreateField("f", pilosa.OptFieldKeys()); err != nil { + t.Fatal(err) + } else if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: ` + Set("one-hundred", f="ten") + SetColumnAttrs("one-hundred", foo="bar") + `}); err != nil { + t.Fatal(err) + } + + targetColAttrSets := []*pilosa.ColumnAttrSet{ + {Key: "one-hundred", Attrs: map[string]interface{}{"foo": "bar"}}, + } + + if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Options(Row(f="ten"), columnAttrs=true)`}); err != nil { + t.Fatal(err) + } else if keys := res.Results[0].(*pilosa.Row).Keys; !reflect.DeepEqual(keys, []string{"one-hundred"}) { + t.Fatalf("unexpected keys: %+v", keys) + } else if attrs := res.ColumnAttrSets; !reflect.DeepEqual(attrs, targetColAttrSets) { + t.Fatalf("unexpected attrs: %s", spew.Sdump(attrs)) + } + }) + t.Run("shards", func(t *testing.T) { c := test.MustRunCluster(t, 1) defer c.Close()