diff --git a/executor.go b/executor.go index 928080ea6..712eafd33 100644 --- a/executor.go +++ b/executor.go @@ -341,6 +341,8 @@ func (e *executor) executeCall(ctx context.Context, index string, c *pql.Call, s return e.executeGroupBy(ctx, index, c, shards, opt) case "Options": return e.executeOptionsCall(ctx, index, c, shards, opt) + case "IncludesColumn": + return e.executeIncludesColumnCall(ctx, index, c, shards, opt) default: e.Holder.Stats.CountWithCustomTags(c.Name, 1, 1.0, []string{indexTag}) return e.executeBitmapCall(ctx, index, c, shards, opt) @@ -411,6 +413,58 @@ func (e *executor) executeOptionsCall(ctx context.Context, index string, c *pql. return e.executeCall(ctx, index, c.Children[0], shards, optCopy) } +// executeIncludesColumnCall executes an IncludesColumn() call. +func (e *executor) executeIncludesColumnCall(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *execOptions) (bool, error) { + // Get the shard containing the column, since that's the only + // shard that needs to execute this query. + var shard uint64 + col, ok, err := c.UintArg("column") + if err != nil { + return false, errors.Wrap(err, "getting column from args") + } else if !ok { + return false, errors.New("IncludesColumn call must specify a column") + } + shard = col / ShardWidth + + // If shard is not in shards, bail early. + if !uint64InSlice(shard, shards) { + return false, nil + } + + // Execute calls in bulk on each remote node and merge. + mapFn := func(shard uint64) (interface{}, error) { + return e.executeIncludesColumnCallShard(ctx, index, c, shard, col) + } + + // Merge returned results at coordinating node. + reduceFn := func(prev, v interface{}) interface{} { + other, _ := prev.(bool) + return other || v.(bool) + } + + result, err := e.mapReduce(ctx, index, []uint64{shard}, c, opt, mapFn, reduceFn) + if err != nil { + return false, err + } + return result.(bool), nil +} + +// executeIncludesColumnCallShard +func (e *executor) executeIncludesColumnCallShard(ctx context.Context, index string, c *pql.Call, shard uint64, column uint64) (bool, error) { + span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeIncludesColumnCallShard") + defer span.Finish() + + if len(c.Children) == 1 { + row, err := e.executeBitmapCallShard(ctx, index, c.Children[0], shard) + if err != nil { + return false, errors.Wrap(err, "executing bitmap call") + } + return row.Includes(column), nil + } + + return false, errors.New("IncludesColumn call must specify a row query") +} + // executeSum executes a Sum() call. func (e *executor) executeSum(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *execOptions) (ValCount, error) { span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeSum") diff --git a/executor_test.go b/executor_test.go index a280be964..383ea21a8 100644 --- a/executor_test.go +++ b/executor_test.go @@ -4090,3 +4090,60 @@ func TestExecutor_Execute_Shift(t *testing.T) { } }) } + +func TestExecutor_Execute_IncludesColumn(t *testing.T) { + t.Run("results", func(t *testing.T) { + c := test.MustRunCluster(t, 1) + defer c.Close() + hldr := test.Holder{Holder: c[0].Server.Holder()} + hldr.SetBit("i", "general", 10, 1) + hldr.SetBit("i", "general", 10, ShardWidth) + hldr.SetBit("i", "general", 10, 2*ShardWidth) + + for i, tt := range []struct { + col uint64 + expIncluded bool + }{ + {1, true}, + {2, false}, + {ShardWidth, true}, + {ShardWidth + 1, false}, + {2 * ShardWidth, true}, + {(2 * ShardWidth) + 1, false}, + } { + t.Run(fmt.Sprint(i), func(t *testing.T) { + if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: fmt.Sprintf("IncludesColumn(Row(general=10), column=%d)", tt.col)}); err != nil { + t.Fatal(err) + } else if tt.expIncluded && !res.Results[0].(bool) { + t.Fatalf("expected to find column: %d", tt.col) + } else if !tt.expIncluded && res.Results[0].(bool) { + t.Fatalf("did not expect to find column: %d", tt.col) + } + }) + } + }) + t.Run("errors", func(t *testing.T) { + c := test.MustRunCluster(t, 1) + defer c.Close() + hldr := test.Holder{Holder: c[0].Server.Holder()} + hldr.SetBit("i", "general", 10, 1) + + t.Run("no column", func(t *testing.T) { + expErr := "IncludesColumn call must specify a column" + if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `IncludesColumn(Row(general=10))`}); err == nil { + t.Fatalf("expected to get an error") + } else if !strings.Contains(err.Error(), expErr) { + t.Fatalf("expected error: %s, but got: %s", expErr, err.Error()) + } + }) + + t.Run("no row query", func(t *testing.T) { + expErr := "IncludesColumn call must specify a row query" + if _, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `IncludesColumn(column=1)`}); err == nil { + t.Fatalf("expected to get an error") + } else if !strings.Contains(err.Error(), expErr) { + t.Fatalf("expected error: %s, but got: %s", expErr, err.Error()) + } + }) + }) +} diff --git a/row.go b/row.go index 90e19d4b1..519416420 100644 --- a/row.go +++ b/row.go @@ -326,6 +326,21 @@ func (r *Row) Columns() []uint64 { return a } +// Includes returns true if the row contains the given column. +func (r *Row) Includes(col uint64) bool { + // TODO: improve the efficiency of this method by + // performing the column filter at the bitmap level + // rather than iterating through the results here. + for i := range r.segments { + for _, c := range r.segments[i].Columns() { + if c == col { + return true + } + } + } + return false +} + // rowSegment holds a subset of a row. // This could point to a mmapped roaring bitmap or an in-memory bitmap. The // width of the segment will always match the shard width. diff --git a/row_test.go b/row_test.go index a1fd7de4b..175faa4e7 100644 --- a/row_test.go +++ b/row_test.go @@ -123,5 +123,18 @@ func TestRow_IsEmpty(t *testing.T) { if !res.IsEmpty() { t.Fatal("Result Should Be Empty\n") } - +} + +func TestRow_Includes(t *testing.T) { + row := pilosa.NewRow(0, 2*ShardWidth) + + if !row.Includes(0) { + t.Fatal("row should include 0") + } + if row.Includes(1) { + t.Fatal("row should not include 1") + } + if !row.Includes(2 * ShardWidth) { + t.Fatalf("row should include %d", 2*ShardWidth) + } }