diff --git a/executor.go b/executor.go index b765a5ac9..a5c9b4391 100644 --- a/executor.go +++ b/executor.go @@ -1531,6 +1531,11 @@ func (e *executor) executeNotShard(ctx context.Context, index string, c *pql.Cal // 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 { @@ -1542,7 +1547,7 @@ func (e *executor) executeShiftShard(ctx context.Context, index string, c *pql.C return nil, err } - return row.Shift(), nil + return row.Shift(n) } // executeCount executes a count() call. diff --git a/executor_test.go b/executor_test.go index 47b8ec9b3..3d17a14e9 100644 --- a/executor_test.go +++ b/executor_test.go @@ -3724,25 +3724,26 @@ func TestExecutor_Execute_Shift(t *testing.T) { 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))`}); err != nil { + 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)))`}); err != nil { + 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))`}); err != nil { + 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) @@ -3753,18 +3754,31 @@ func TestExecutor_Execute_Shift(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-1) - if res, err := c[0].API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `Shift(Row(general=10))`}); err != nil { + 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, []uint64{ShardWidth}) { - t.Fatalf("unexpected columns: %+v", columns) + } 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, []uint64{ShardWidth + 1}) { - t.Fatalf("unexpected columns: %+v", columns) + } 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) } }) @@ -3778,16 +3792,15 @@ func TestExecutor_Execute_Shift(t *testing.T) { 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))`}); err != nil { + 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)))`}); err != nil { + 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 a0568dbc2..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,30 +757,36 @@ func (b *Bitmap) Xor(other *Bitmap) *Bitmap { return output } -func (b *Bitmap) Shift() *Bitmap { +// 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) - last := false + lastCarry := false lastKey := uint64(0) for iiter.Next() { ki, ci := iiter.Value() o, carry := shift(ci) - if last { + if lastCarry { o.add(0) } if o.n > 0 { output.Containers.Put(ki, o) } - last = carry + lastCarry = carry lastKey = ki } - if last { //handle the overflow + // 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 + return output, nil } // removeEmptyContainers deletes all containers that have a count of zero. @@ -3387,15 +3397,19 @@ func xorBitmapBitmap(a, b *Container) *Container { return output } -func shift(a *Container) (*Container, bool) { - if a.isArray() { - return shiftArray(a) - } else if a.isRun() { - return shiftRun(a) +// 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(a) + return shiftBitmap(c) } +// shiftArray is an array-specific implementation of shift(). func shiftArray(a *Container) (*Container, bool) { statsHit("shift/Array") carry := false @@ -3404,7 +3418,7 @@ func shiftArray(a *Container) (*Container, bool) { output.array = output.array[:0] output.n = a.n for _, v := range a.array { - if v+1 == 0 { //overflow + if v+1 == 0 { // overflow carry = true output.n -= 1 } else { @@ -3414,6 +3428,7 @@ func shiftArray(a *Container) (*Container, bool) { return output, carry } +// shiftBitmap is a bitmap-specific implementation of shift(). func shiftBitmap(a *Container) (*Container, bool) { statsHit("shift/Bitmap") carry := false @@ -3421,17 +3436,15 @@ func shiftBitmap(a *Container) (*Container, bool) { output.bitmap = make([]uint64, len(a.bitmap)) output.bitmap = output.bitmap[:0] output.n = a.n - lastcarry := false - for i, v := range a.bitmap { + lastCarry := false + for _, v := range a.bitmap { carry = (v & (1 << 63)) != 0 v = v << 1 - if i != 0 { - if lastcarry { - v |= 1 - } + if lastCarry { + v |= 1 } output.bitmap = append(output.bitmap, v) - lastcarry = carry + lastCarry = carry } if carry { output.n -= 1 @@ -3439,6 +3452,7 @@ func shiftBitmap(a *Container) (*Container, bool) { return output, carry } +// shiftRun is a run-specific implementation of shift(). func shiftRun(a *Container) (*Container, bool) { statsHit("shift/Run") carry := false @@ -3446,11 +3460,11 @@ func shiftRun(a *Container) (*Container, bool) { output.runs = make([]interval16, len(a.runs)) output.runs = output.runs[:0] for _, v := range a.runs { - if v.start+1 == 0 { + 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 { + } else if v.last+1 == 0 { // final run ends on container edge v.start += 1 carry = true output.n -= 1 diff --git a/roaring/roaring_internal_test.go b/roaring/roaring_internal_test.go index 2dda414b3..5ca0edb2b 100644 --- a/roaring/roaring_internal_test.go +++ b/roaring/roaring_internal_test.go @@ -3346,9 +3346,12 @@ func TestShiftArray(t *testing.T) { for i, test := range tests { a.array = test.array a.n = int32(len(a.array)) - ret, _ := shiftArray(a) - if !reflect.DeepEqual(ret.array, test.exp) { - t.Fatalf("test #%v expected %v, but got %v", i, test.exp, ret.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) } } } @@ -3378,9 +3381,12 @@ func TestShiftBitmap(t *testing.T) { for i, test := range tests { a.bitmap = test.bitmap a.n = 1 - ret, _ := shiftBitmap(a) - if !reflect.DeepEqual(ret.bitmap, test.exp) { - t.Fatalf("test #%v expected %v, but got %v", i, test.exp, ret.bitmap) + 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) } } } @@ -3422,9 +3428,12 @@ func TestShiftRun(t *testing.T) { for i, test := range tests { a.runs = test.runs a.n = test.n - ret, c := shiftRun(a) - if !reflect.DeepEqual(ret.runs, test.exp) && c == test.carry && ret.n == test.en { - t.Fatalf("test #%v expected %v, but got %v %d", i, test.exp, ret.runs, ret.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 2a0eddf58..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,14 +168,30 @@ func (r *Row) Difference(other *Row) *Row { return &Row{segments: segments} } -// Shift returns the bitwise shift of r by 1 bit. -func (r *Row) Shift() *Row { - var segments []rowSegment - for _, segment := range r.segments { - segments = append(segments, *segment.Shift()) +// 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 } - return &Row{segments: segments} + 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. @@ -352,15 +369,18 @@ func (s *rowSegment) Xor(other *rowSegment) *rowSegment { } // Shift returns s shifted by 1 bit. -func (s *rowSegment) Shift() *rowSegment { +func (s *rowSegment) Shift() (*rowSegment, error) { //TODO deal with overflow - data := s.data.Shift() + 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.