Merge pull request #1761 from tgruben/shift-op

Shift operator

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(<ROW_CALL>, [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:**

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")

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)
+		}
+	})
+}

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
 

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

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)
+		}
+	}
+}

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) }

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()