implement Store() in the executor (i.e. setRow())

executor.go

@@ -184,6 +184,8 @@ func (e *executor) executeCall(ctx context.Context, index string, c *pql.Call, s
 		return e.executeClearBit(ctx, index, c, opt)
 	case "ClearRow":
 		return e.executeClearRow(ctx, index, c, shards, opt)
+	case "Store":
+		return e.executeSetRow(ctx, index, c, shards, opt)
 	case "Count":
 		e.Holder.Stats.CountWithCustomTags(c.Name, 1, 1.0, []string{indexTag})
 		return e.executeCount(ctx, index, c, shards, opt)
@@ -1213,6 +1215,98 @@ func (e *executor) executeClearRowShard(_ context.Context, index string, c *pql.
 	return changed, nil
 }
 
+// executeSetRow executes a SetRow() call.
+func (e *executor) executeSetRow(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *execOptions) (bool, error) {
+	// Ensure the field type supports SetRow().
+	fieldName, err := c.FieldArg()
+	if err != nil {
+		return false, errors.New("SetRow() argument required: field")
+	}
+	field := e.Holder.Field(index, fieldName)
+	if field == nil {
+		return false, ErrFieldNotFound
+	}
+
+	switch field.Type() {
+	case FieldTypeSet:
+		// These field types support SetRow().
+	default:
+		return false, fmt.Errorf("SetRow() is not supported on %s field types", field.Type())
+	}
+
+	// Execute calls in bulk on each remote node and merge.
+	mapFn := func(shard uint64) (interface{}, error) {
+		return e.executeSetRowShard(ctx, index, c, shard)
+	}
+
+	// Merge returned results at coordinating node.
+	reduceFn := func(prev, v interface{}) interface{} {
+		val := v.(bool)
+		if prev == nil {
+			return val
+		}
+		return val || prev.(bool)
+	}
+
+	result, err := e.mapReduce(ctx, index, shards, c, opt, mapFn, reduceFn)
+	return result.(bool), err
+}
+
+// executeSetRowShard executes a SetRow() call for a single shard.
+func (e *executor) executeSetRowShard(ctx context.Context, index string, c *pql.Call, shard uint64) (bool, error) {
+	fieldName, err := c.FieldArg()
+	if err != nil {
+		return false, errors.New("SetRow() argument required: field")
+	}
+
+	// Read fields using labels.
+	rowID, ok, err := c.UintArg(fieldName)
+	if err != nil {
+		return false, fmt.Errorf("reading SetRow() row: %v", err)
+	} else if !ok {
+		return false, fmt.Errorf("SetRow() row argument '%v' required", rowLabel)
+	}
+
+	field := e.Holder.Field(index, fieldName)
+	if field == nil {
+		return false, ErrFieldNotFound
+	}
+
+	// Retrieve source row.
+	var src *Row
+	if len(c.Children) == 1 {
+		row, err := e.executeBitmapCallShard(ctx, index, c.Children[0], shard)
+		if err != nil {
+			return false, errors.Wrap(err, "getting source row")
+		}
+		src = row
+	} else {
+		return false, errors.New("SetRow() requires a source row")
+	}
+
+	// Set the row on the standard view.
+	changed := false
+	fragment := e.Holder.fragment(index, fieldName, viewStandard, shard)
+	if fragment == nil {
+		// Since the destination fragment doesn't exist, create one.
+		view, err := field.createViewIfNotExists(viewStandard)
+		if err != nil {
+			return false, errors.Wrap(err, "creating view")
+		}
+		fragment, err = view.createFragmentIfNotExists(shard)
+		if err != nil {
+			return false, errors.Wrapf(err, "creating fragment: %d", shard)
+		}
+	}
+	cleared, err := fragment.setRow(src, rowID)
+	if err != nil {
+		return false, errors.Wrapf(err, "setting row %d on view %s shard %d", rowID, viewStandard, shard)
+	}
+	changed = changed || cleared
+
+	return changed, nil
+}
+
 // executeSet executes a Set() call.
 func (e *executor) executeSet(ctx context.Context, index string, c *pql.Call, opt *execOptions) (bool, error) {
 	// Read colID.

executor_test.go

@@ -1913,6 +1913,145 @@ func TestExecutor_Execute_ClearRow(t *testing.T) {
 	})
 }
 
+// Ensure a row can be set.
+func TestExecutor_Execute_SetRow(t *testing.T) {
+	t.Run("Set_NewRow", func(t *testing.T) {
+		c := test.MustRunCluster(t, 1)
+		defer c.Close()
+		hldr := test.Holder{Holder: c[0].Server.Holder()}
+		index := hldr.MustCreateIndexIfNotExists("i", pilosa.IndexOptions{TrackExistence: true})
+		if _, err := index.CreateField("f", pilosa.OptFieldTypeDefault()); err != nil {
+			t.Fatal(err)
+		}
+		if _, err := index.CreateField("tmp", pilosa.OptFieldTypeDefault()); err != nil {
+			t.Fatal(err)
+		}
+
+		// Set bits.
+		if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `` +
+			fmt.Sprintf("Set(%d, f=%d)\n", 3, 10) +
+			fmt.Sprintf("Set(%d, f=%d)\n", ShardWidth-1, 10) +
+			fmt.Sprintf("Set(%d, f=%d)\n", ShardWidth+1, 10),
+		}); err != nil {
+			t.Fatal(err)
+		}
+
+		if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Row(f=10)`}); err != nil {
+			t.Fatal(err)
+		} else if bits := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(bits, []uint64{3, ShardWidth - 1, ShardWidth + 1}) {
+			t.Fatalf("unexpected columns: %+v", bits)
+		}
+
+		// Store row 10 into a different row.
+		if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Store(Row(f=10), tmp=20)`}); 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[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Row(tmp=20)`}); err != nil {
+			t.Fatal(err)
+		} else if bits := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(bits, []uint64{3, ShardWidth - 1, ShardWidth + 1}) {
+			t.Fatalf("unexpected columns: %+v", bits)
+		}
+	})
+	t.Run("Set_NoSource", func(t *testing.T) {
+		c := test.MustRunCluster(t, 1)
+		defer c.Close()
+		hldr := test.Holder{Holder: c[0].Server.Holder()}
+		index := hldr.MustCreateIndexIfNotExists("i", pilosa.IndexOptions{TrackExistence: true})
+		_, err := index.CreateField("f", pilosa.OptFieldTypeDefault())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Set bits.
+		if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `` +
+			fmt.Sprintf("Set(%d, f=%d)\n", 3, 10) +
+			fmt.Sprintf("Set(%d, f=%d)\n", ShardWidth-1, 10) +
+			fmt.Sprintf("Set(%d, f=%d)\n", ShardWidth+1, 10),
+		}); err != nil {
+			t.Fatal(err)
+		}
+
+		if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Row(f=10)`}); err != nil {
+			t.Fatal(err)
+		} else if bits := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(bits, []uint64{3, ShardWidth - 1, ShardWidth + 1}) {
+			t.Fatalf("unexpected columns: %+v", bits)
+		}
+
+		// Store row 9 (which doesn't exist) into a different row.
+		if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Store(Row(f=9), f=20)`}); 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[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Row(f=20)`}); err != nil {
+			t.Fatal(err)
+		} else if bits := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(bits, []uint64{}) {
+			t.Fatalf("unexpected columns: %+v", bits)
+		}
+
+		// Store row 9 (which doesn't exist) into a row that does exist.
+		if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Store(Row(f=9), f=10)`}); 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[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Row(f=10)`}); err != nil {
+			t.Fatal(err)
+		} else if bits := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(bits, []uint64{}) {
+			t.Fatalf("unexpected columns: %+v", bits)
+		}
+	})
+	t.Run("Set_ExistingDestination", func(t *testing.T) {
+		c := test.MustRunCluster(t, 1)
+		defer c.Close()
+		hldr := test.Holder{Holder: c[0].Server.Holder()}
+		index := hldr.MustCreateIndexIfNotExists("i", pilosa.IndexOptions{TrackExistence: true})
+		_, err := index.CreateField("f", pilosa.OptFieldTypeDefault())
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		// Set bits.
+		if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `` +
+			fmt.Sprintf("Set(%d, f=%d)\n", 3, 10) +
+			fmt.Sprintf("Set(%d, f=%d)\n", ShardWidth-1, 10) +
+			fmt.Sprintf("Set(%d, f=%d)\n", ShardWidth+1, 10) +
+			fmt.Sprintf("Set(%d, f=%d)\n", 1, 20) +
+			fmt.Sprintf("Set(%d, f=%d)\n", ShardWidth+1, 20),
+		}); err != nil {
+			t.Fatal(err)
+		}
+
+		if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Row(f=20)`}); err != nil {
+			t.Fatal(err)
+		} else if bits := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(bits, []uint64{1, ShardWidth + 1}) {
+			t.Fatalf("unexpected columns: %+v", bits)
+		}
+
+		// Store row 10 into an existing row.
+		if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Store(Row(f=10), f=20)`}); 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[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Row(f=20)`}); err != nil {
+			t.Fatal(err)
+		} else if bits := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(bits, []uint64{3, ShardWidth - 1, ShardWidth + 1}) {
+			t.Fatalf("unexpected columns: %+v", bits)
+		}
+	})
+}
+
 func benchmarkExistence(nn bool, b *testing.B) {
 	c := test.MustRunCluster(b, 1)
 	defer c.Close()

pql/pql.peg

@@ -11,6 +11,7 @@ Call <-  'Set' {p.startCall("Set")} open col comma args (comma timestamp)? close
        / 'SetColumnAttrs' {p.startCall("SetColumnAttrs")} open col comma args close {p.endCall()}
        / 'Clear' {p.startCall("Clear")} open col comma args close {p.endCall()}
        / 'ClearRow' {p.startCall("ClearRow")} open arg close {p.endCall()}
+       / 'Store' {p.startCall("Store")} open Call comma arg close {p.endCall()}
        / 'TopN' {p.startCall("TopN")} open posfield (comma allargs)? close {p.endCall()}
        / 'Range' {p.startCall("Range")} open (timerange / conditional / arg) close {p.endCall()}
        / < IDENT > { p.startCall(buffer[begin:end] ) } open allargs comma? close { p.endCall() }