diff --git a/docs/query-language.md b/docs/query-language.md index 82ffecd2a..4dd838727 100644 --- a/docs/query-language.md +++ b/docs/query-language.md @@ -594,6 +594,34 @@ Count(Row(stargazer=1)) * Result is the number of repositories that user 1 has starred. +#### Shift +**Spec:** + +``` +Shift(, [n=UINT]) +``` + +**Description:** + +Returns the row specified by `ROW_CALL` shifted by `n` bits. + +**Result Type:** object with attrs and columns + +attrs will always be empty + +**Examples:** + +Query all columns with a bit set in row 1 of the field `stargazer` +and shift the result by 2: +```request +Shift(Row(stargazer=1), n=2) +``` +```response +{"attrs":{},"columns":[12, 22]} +``` + +* columns are the repositories which user 1 has starred shifted by 2 bits. + #### TopN **Spec:** diff --git a/executor.go b/executor.go index b17cf95e2..a5c9b4391 100644 --- a/executor.go +++ b/executor.go @@ -558,6 +558,8 @@ func (e *executor) executeBitmapCallShard(ctx context.Context, index string, c * return e.executeXorShard(ctx, index, c, shard) case "Not": return e.executeNotShard(ctx, index, c, shard) + case "Shift": + return e.executeShiftShard(ctx, index, c, shard) default: return nil, fmt.Errorf("unknown call: %s", c.Name) } @@ -1527,6 +1529,27 @@ func (e *executor) executeNotShard(ctx context.Context, index string, c *pql.Cal return existenceRow.Difference(row), nil } +// executeShiftShard executes a shift() call for a local shard. +func (e *executor) executeShiftShard(ctx context.Context, index string, c *pql.Call, shard uint64) (*Row, error) { + n, _, err := c.IntArg("n") + if err != nil { + return nil, fmt.Errorf("executeShiftShard: %v", err) + } + + if len(c.Children) == 0 { + return nil, errors.New("Shift() requires an input row") + } else if len(c.Children) > 1 { + return nil, errors.New("Shift() only accepts a single row input") + } + + row, err := e.executeBitmapCallShard(ctx, index, c.Children[0], shard) + if err != nil { + return nil, err + } + + return row.Shift(n) +} + // executeCount executes a count() call. func (e *executor) executeCount(ctx context.Context, index string, c *pql.Call, shards []uint64, opt *execOptions) (uint64, error) { span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeCount") diff --git a/executor_test.go b/executor_test.go index 910ab14ba..3d17a14e9 100644 --- a/executor_test.go +++ b/executor_test.go @@ -3716,3 +3716,91 @@ func runCallTest(t *testing.T, writeQuery string, readQueries []string, indexOpt return responses } + +func TestExecutor_Execute_Shift(t *testing.T) { + t.Run("Shift Bit 0", 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, 0) + + if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Shift(Row(general=10), n=1)`}); err != nil { + t.Fatal(err) + } else if columns := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, []uint64{1}) { + t.Fatalf("unexpected columns: %+v", columns) + } + + if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Shift(Shift(Row(general=10), n=1), n=1)`}); err != nil { + t.Fatal(err) + } else if columns := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, []uint64{2}) { + t.Fatalf("unexpected columns: %+v", columns) + } + }) + + t.Run("Shift container boundary", 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, 65535) + + if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Shift(Row(general=10), n=1)`}); err != nil { + t.Fatal(err) + } else if columns := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, []uint64{65536}) { + t.Fatalf("unexpected columns: %+v", columns) + } + }) + + t.Run("Shift shard boundary", func(t *testing.T) { + c := test.MustRunCluster(t, 1) + defer c.Close() + hldr := test.Holder{Holder: c[0].Server.Holder()} + + orig := []uint64{1, ShardWidth - 1, ShardWidth + 1} + shift1 := []uint64{2, ShardWidth, ShardWidth + 2} + shift2 := []uint64{3, ShardWidth + 1, ShardWidth + 3} + + for _, bit := range orig { + hldr.SetBit("i", "general", 10, bit) + } + + if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Shift(Row(general=10), n=1)`}); err != nil { + t.Fatal(err) + } else if columns := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, shift1) { + t.Fatalf("unexpected shift by 1: expected: %+v, but got: %+v", shift1, columns) + } + + if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Shift(Row(general=10), n=2)`}); err != nil { + t.Fatal(err) + } else if columns := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, shift2) { + t.Fatalf("unexpected shift by 2: expected: %+v, but got: %+v", shift2, columns) + } + + if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Shift(Shift(Row(general=10)))`}); err != nil { + t.Fatal(err) + } else if columns := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, orig) { + t.Fatalf("unexpected shift by 0: expected: %+v, but got: %+v", orig, columns) + } + }) + + t.Run("Shift shard boundary no create", 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, ShardWidth-2) //shardwidth -1 + hldr.SetBit("i", "general", 10, ShardWidth-1) //shardwidth + hldr.SetBit("i", "general", 10, ShardWidth) //shardwidth +1 + hldr.SetBit("i", "general", 10, ShardWidth+2) //shardwidth +3 + + exp := []uint64{ShardWidth - 1, ShardWidth, ShardWidth + 1, ShardWidth + 3} + if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Shift(Row(general=10), n=1)`}); err != nil { + t.Fatal(err) + } else if columns := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, exp) { + t.Fatalf("unexpected columns: %+v", columns) + } + if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Shift(Shift(Row(general=10), n=1), n=1)`}); err != nil { + t.Fatal(err) + } else if columns := res.Results[0].(*pilosa.Row).Columns(); !reflect.DeepEqual(columns, []uint64{ShardWidth, ShardWidth + 1, ShardWidth + 2, ShardWidth + 4}) { + t.Fatalf("unexpected columns: \n%+v\n%+v", columns, exp) + } + }) +} diff --git a/roaring/roaring.go b/roaring/roaring.go index 2fe76e8ef..abc4e4d2f 100644 --- a/roaring/roaring.go +++ b/roaring/roaring.go @@ -52,16 +52,20 @@ const ( // bitmapN is the number of values in a container.bitmap. bitmapN = (1 << 16) / 64 - //containerArray indicates a container of bit position values + // containerArray indicates a container of bit position values containerArray = byte(1) - //containerBitmap indicates a container of bits packed in a uint64 array block + // containerBitmap indicates a container of bits packed in a uint64 array block containerBitmap = byte(2) - //containerRun indicates a container of run encoded bits + // containerRun indicates a container of run encoded bits containerRun = byte(3) maxContainerVal = 0xffff + + // maxContainerKey is the key representing the last container in a full row. + // It is the full bitmap space (2^64) divided by container width (2^16). + maxContainerKey = (1 << 48) - 1 ) type Containers interface { @@ -753,6 +757,38 @@ func (b *Bitmap) Xor(other *Bitmap) *Bitmap { return output } +// Shift shifts the contents of b by 1. +func (b *Bitmap) Shift(n int) (*Bitmap, error) { + if n != 1 { + return nil, errors.New("cannot shift by a value other than 1") + } + output := NewBitmap() + iiter, _ := b.Containers.Iterator(0) + lastCarry := false + lastKey := uint64(0) + for iiter.Next() { + ki, ci := iiter.Value() + o, carry := shift(ci) + if lastCarry { + o.add(0) + } + if o.n > 0 { + output.Containers.Put(ki, o) + } + lastCarry = carry + lastKey = ki + } + // As long as the carry wasn't from the max container, + // append a new container and add the carried bit. + if lastCarry && lastKey != maxContainerKey { + extra := NewContainer() + extra.add(0) + output.Containers.Put(lastKey+1, extra) + } + + return output, nil +} + // removeEmptyContainers deletes all containers that have a count of zero. func (b *Bitmap) removeEmptyContainers() { citer, _ := b.Containers.Iterator(0) @@ -3361,6 +3397,88 @@ func xorBitmapBitmap(a, b *Container) *Container { return output } +// shift() shifts the contents of c by one. It returns +// the new container and a bool indicating whether a +// carry bit was shifted out. +func shift(c *Container) (*Container, bool) { + if c.isArray() { + return shiftArray(c) + } else if c.isRun() { + return shiftRun(c) + } + return shiftBitmap(c) +} + +// shiftArray is an array-specific implementation of shift(). +func shiftArray(a *Container) (*Container, bool) { + statsHit("shift/Array") + carry := false + output := &Container{containerType: containerArray} + output.array = make([]uint16, len(a.array)) + output.array = output.array[:0] + output.n = a.n + for _, v := range a.array { + if v+1 == 0 { // overflow + carry = true + output.n -= 1 + } else { + output.array = append(output.array, v+1) + } + } + return output, carry +} + +// shiftBitmap is a bitmap-specific implementation of shift(). +func shiftBitmap(a *Container) (*Container, bool) { + statsHit("shift/Bitmap") + carry := false + output := &Container{containerType: containerBitmap} + output.bitmap = make([]uint64, len(a.bitmap)) + output.bitmap = output.bitmap[:0] + output.n = a.n + lastCarry := false + for _, v := range a.bitmap { + carry = (v & (1 << 63)) != 0 + v = v << 1 + if lastCarry { + v |= 1 + } + output.bitmap = append(output.bitmap, v) + lastCarry = carry + } + if carry { + output.n -= 1 + } + return output, carry +} + +// shiftRun is a run-specific implementation of shift(). +func shiftRun(a *Container) (*Container, bool) { + statsHit("shift/Run") + carry := false + output := &Container{containerType: containerRun} + output.runs = make([]interval16, len(a.runs)) + output.runs = output.runs[:0] + for _, v := range a.runs { + if v.start+1 == 0 { // final run was 1 bit on container edge + carry = true + output.n -= 1 + break + } else if v.last+1 == 0 { // final run ends on container edge + v.start += 1 + carry = true + output.n -= 1 + } else { + v.start += 1 + v.last += 1 + carry = false + } + output.runs = append(output.runs, v) + } + + return output, carry +} + // opType represents a type of operation. type opType uint8 diff --git a/roaring/roaring_helpers_test.go b/roaring/roaring_helpers_test.go index cd22978fb..f99d922e8 100644 --- a/roaring/roaring_helpers_test.go +++ b/roaring/roaring_helpers_test.go @@ -106,6 +106,24 @@ func bitmapFirstBitSet() []uint64 { return bitmap } +func bitmapSecondBitSet() []uint64 { + bitmap := make([]uint64, bitmapN) + bitmap[0] = 0x0000000000000002 + return bitmap +} + +func bitmapLastBitFirstRowSet() []uint64 { + bitmap := make([]uint64, bitmapN) + bitmap[0] = 0x8000000000000000 + return bitmap +} + +func bitmapFirstBitSecoundRowSet() []uint64 { + bitmap := make([]uint64, bitmapN) + bitmap[1] = 0x0000000000000001 + return bitmap +} + func bitmapLastBitSet() []uint64 { bitmap := make([]uint64, bitmapN) bitmap[bitmapN-1] = 0x8000000000000000 diff --git a/roaring/roaring_internal_test.go b/roaring/roaring_internal_test.go index ff45e6c86..5ca0edb2b 100644 --- a/roaring/roaring_internal_test.go +++ b/roaring/roaring_internal_test.go @@ -3317,3 +3317,123 @@ func TestEquals(t *testing.T) { } } */ +func TestShiftArray(t *testing.T) { + a := &Container{ + containerType: containerArray, + } + tests := []struct { + array []uint16 + exp []uint16 + }{ + { + array: []uint16{1}, + exp: []uint16{2}, + }, + { + array: []uint16{}, + exp: []uint16{}, + }, + { + array: []uint16{1, 2, 3, 4, 5, 11, 12}, + exp: []uint16{2, 3, 4, 5, 6, 12, 13}, + }, + { + array: []uint16{65535}, + exp: []uint16{}, + }, + } + + for i, test := range tests { + a.array = test.array + a.n = int32(len(a.array)) + ret1, _ := shift(a) // test generic shift function + ret2, _ := shiftArray(a) // test array-specific shift function + if !reflect.DeepEqual(ret1.array, test.exp) { + t.Fatalf("test #%v shift() expected %v, but got %v", i, test.exp, ret1.array) + } else if !reflect.DeepEqual(ret2.array, test.exp) { + t.Fatalf("test #%v shiftArray() expected %v, but got %v", i, test.exp, ret2.array) + } + } +} + +func TestShiftBitmap(t *testing.T) { + a := &Container{ + containerType: containerBitmap, + } + tests := []struct { + bitmap []uint64 + exp []uint64 + }{ + { + bitmap: bitmapFirstBitSet(), + exp: bitmapSecondBitSet(), + }, + { + bitmap: bitmapLastBitSet(), + exp: bitmapEmpty(), + }, + { + bitmap: bitmapLastBitFirstRowSet(), + exp: bitmapFirstBitSecoundRowSet(), + }, + } + + for i, test := range tests { + a.bitmap = test.bitmap + a.n = 1 + ret1, _ := shift(a) // test generic shift function + ret2, _ := shiftBitmap(a) // test bitmap-specific shift function + if !reflect.DeepEqual(ret1.bitmap, test.exp) { + t.Fatalf("test #%v shift() expected %v, but got %v", i, test.exp, ret1.bitmap) + } else if !reflect.DeepEqual(ret2.bitmap, test.exp) { + t.Fatalf("test #%v shiftBitmap() expected %v, but got %v", i, test.exp, ret2.bitmap) + } + } +} +func TestShiftRun(t *testing.T) { + a := &Container{ + containerType: containerRun, + } + + tests := []struct { + runs []interval16 + n int32 + en int32 + exp []interval16 + carry bool + }{ + { + runs: []interval16{{start: 5, last: 10}}, + n: 5, + en: 5, + exp: []interval16{{start: 6, last: 11}}, + carry: false, + }, + { + runs: []interval16{{start: 5, last: 65535}}, + n: 65530, + en: 65529, + exp: []interval16{{start: 6, last: 65535}}, + carry: true, + }, + { + runs: []interval16{{start: 65535, last: 65535}}, + n: 1, + en: 0, + exp: []interval16{}, + carry: true, + }, + } + + for i, test := range tests { + a.runs = test.runs + a.n = test.n + ret1, c1 := shift(a) // test generic shift function + ret2, c2 := shiftRun(a) // test run-specific shift function + if !reflect.DeepEqual(ret1.runs, test.exp) && c1 == test.carry && ret1.n == test.en { + t.Fatalf("test #%v shift() expected %v, but got %v %d", i, test.exp, ret1.runs, ret1.n) + } else if !reflect.DeepEqual(ret2.runs, test.exp) && c2 == test.carry && ret2.n == test.en { + t.Fatalf("test #%v shiftRun() expected %v, but got %v %d", i, test.exp, ret2.runs, ret2.n) + } + } +} diff --git a/roaring/roaring_test.go b/roaring/roaring_test.go index b165042a9..a4333d0fc 100644 --- a/roaring/roaring_test.go +++ b/roaring/roaring_test.go @@ -985,6 +985,18 @@ func TestBitmap_IntersectionCount_Mixed(t *testing.T) { } } +func TestBitmap_Shift(t *testing.T) { + var max uint64 = math.MaxUint64 + bm1 := roaring.NewFileBitmap(0, 1, 2, 3, 4, 5, 6, 7, 9, 10, 65536, max) + bm2 := roaring.NewFileBitmap(1, 2, 3, 4, 5, 6, 7, 8, 10, 11, 65537) + + if got, err := bm1.Shift(1); err != nil { + t.Fatal(err) + } else if !reflect.DeepEqual(got.Slice(), bm2.Slice()) { + t.Fatalf("unexpected bitmap: expected %v, but got %v", bm2.Slice(), got.Slice()) + } +} + func TestBitmap_Quick_Array1(t *testing.T) { testBitmapQuick(t, 1000, 1000, 2000) } func TestBitmap_Quick_Array2(t *testing.T) { testBitmapQuick(t, 10000, 0, 1000) } func TestBitmap_Quick_Bitmap1(t *testing.T) { testBitmapQuick(t, 10000, 0, 10000) } diff --git a/row.go b/row.go index 255c3e847..0a8fc0acd 100644 --- a/row.go +++ b/row.go @@ -19,6 +19,7 @@ import ( "sort" "github.com/pilosa/pilosa/roaring" + "github.com/pkg/errors" ) // Row is a set of integers (the associated columns), and attributes which are @@ -167,6 +168,32 @@ func (r *Row) Difference(other *Row) *Row { return &Row{segments: segments} } +// Shift returns the bitwise shift of r by n bits. +// Currently only positive shift values are supported. +func (r *Row) Shift(n int64) (*Row, error) { + if n < 0 { + return nil, errors.New("cannot shift by negative values") + } else if n == 0 { + return r, nil + } + + work := r + var segments []rowSegment + for i := int64(0); i < n; i++ { + segments = segments[:0] + for _, segment := range work.segments { + shifted, err := segment.Shift() + if err != nil { + return nil, errors.Wrap(err, "shifting row segment") + } + segments = append(segments, *shifted) + } + work = &Row{segments: segments} + } + + return work, nil +} + // SetBit sets the i-th column of the row. func (r *Row) SetBit(i uint64) (changed bool) { return r.createSegmentIfNotExists(i / ShardWidth).SetBit(i) @@ -341,6 +368,21 @@ func (s *rowSegment) Xor(other *rowSegment) *rowSegment { } } +// Shift returns s shifted by 1 bit. +func (s *rowSegment) Shift() (*rowSegment, error) { + //TODO deal with overflow + data, err := s.data.Shift(1) + if err != nil { + return nil, errors.Wrap(err, "shifting roaring data") + } + + return &rowSegment{ + data: *data, + shard: s.shard, + n: data.Count(), + }, nil +} + // SetBit sets the i-th column of the row. func (s *rowSegment) SetBit(i uint64) (changed bool) { s.ensureWritable()