diff --git a/docs/query-language.md b/docs/query-language.md index 12d39647f..1f405e144 100644 --- a/docs/query-language.md +++ b/docs/query-language.md @@ -792,6 +792,33 @@ Options(Row(f1=10), shards=[0, 2]) {"attrs":{},"columns":[100, 2097152]} ``` +#### Row Constant + +**Spec:** + +``` +ConstRow(columns=<[]COLUMN>) +``` + +**Description:** + +`ConstRow` provides a constant bitmap value that can be used in place of a `Row` call. +The columns can be specified as integer IDs or strings. + +**Result Type:** row value columns. + +e.g. `{"attrs":{},"columns":[10, 20]}` + +**Examples:** + +Filter specified columns to only those with a bit set in row 1 of the field `stargazer` (repositories that are starred by user 1): +```request +Intersect(ConstRow(columns=[10, 20, 30]), Row(stargazer=1)) +``` +```response +{"attrs":{},"columns":[10, 20]} +``` + #### Rows **Spec:** diff --git a/executor.go b/executor.go index a973e9565..89fa88504 100644 --- a/executor.go +++ b/executor.go @@ -523,7 +523,6 @@ func (e *executor) execute(ctx context.Context, tx Tx, index string, q *pql.Quer } // preprocessQuery expands any calls that need preprocessing. -// So far, this only needs to process UnionRows. func (e *executor) preprocessQuery(ctx context.Context, tx Tx, index string, c *pql.Call, shards []uint64, opt *execOptions) (*pql.Call, error) { switch c.Name { case "UnionRows": @@ -605,6 +604,51 @@ func (e *executor) preprocessQuery(ctx context.Context, tx Tx, index string, c * Children: rows, }, nil + case "ConstRow": + // Fetch user-provided columns list. + cols, _ := c.Args["columns"].([]interface{}) + var ids []uint64 + var keys []string + for _, c := range cols { + switch c := c.(type) { + case uint64: + ids = append(ids, c) + case int64: + ids = append(ids, uint64(c)) + case string: + keys = append(keys, c) + default: + return nil, errors.Errorf("invalid column identifier %v of type %T", c, c) + } + } + + // Translate keys to IDs. + if len(keys) > 0 { + keyIDs, err := e.Cluster.translateIndexKeys(ctx, index, keys) + if err != nil { + return nil, errors.Wrap(err, "translating column IDs in ConstRow") + } + ids = append(ids, keyIDs...) + } + + // Split IDs by shard. + shardSet := make(map[uint64][]uint64) + for _, id := range ids { + shardSet[id/ShardWidth] = append(shardSet[id/ShardWidth], id) + } + + // Convert ID sets to per-shard Row objects. + precomputed := make(map[uint64]interface{}) + for _, s := range shards { + precomputed[s] = NewRow(shardSet[s]...) + } + + // Generate a precomputed call with the data. + return &pql.Call{ + Name: "Precomputed", + Precomputed: precomputed, + }, nil + default: // Recurse through child calls. out := make([]*pql.Call, len(c.Children)) diff --git a/executor_test.go b/executor_test.go index b83ef7c32..48b38ade6 100644 --- a/executor_test.go +++ b/executor_test.go @@ -61,6 +61,25 @@ func getTempDirString() (td *string) { return td } +func TestExecutor_Execute_ConstRow(t *testing.T) { + c := test.MustRunCluster(t, 2) + defer c.Close() + + c.CreateField(t, "i", pilosa.IndexOptions{}, "h") + c.ImportBits(t, "i", "h", [][2]uint64{ + {1, 2}, + {3, 4}, + {5, 6}, + }) + + resp := c.Query(t, "i", `ConstRow(columns=[2,6])`) + expect := []uint64{2, 6} + got := resp.Results[0].(*pilosa.Row).Columns() + if !reflect.DeepEqual(expect, got) { + t.Errorf("expected %v but got %v", expect, got) + } +} + // Ensure a row query can be executed. func TestExecutor_Execute_Row(t *testing.T) { t.Run("RowIDColumnID", func(t *testing.T) { diff --git a/pql/ast.go b/pql/ast.go index 839a3d4d4..04621ed83 100644 --- a/pql/ast.go +++ b/pql/ast.go @@ -396,6 +396,13 @@ var callInfoByFunc = map[string]callInfo{ "Extract": {allowUnknown: false}, "Xor": {allowUnknown: false}, + "ConstRow": { + allowUnknown: false, + prototypes: map[string]interface{}{ + "columns": []interface{}{}, + }, + }, + // things that take _field "TopN": allowUnderField, // special cases: