Merge pull request #1098 from niaow/topn-v3

Add TopK with perpendicular BSI bitmaps

bsi.go

@@ -0,0 +1,120 @@
+// Copyright 2020 Pilosa Corp.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package pilosa
+
+import "math/bits"
+
+// bsiData contains BSI-structured data.
+type bsiData []*Row
+
+// insert a value for a column in the BSI data.
+func (bsi *bsiData) insert(column uint64, value uint64) {
+	data := *bsi
+	for value != 0 {
+		bit := bits.TrailingZeros64(value)
+		value &^= 1 << bit
+
+		for len(data) <= bit {
+			data = append(data, NewRow())
+		}
+
+		data[bit].SetBit(column)
+	}
+	*bsi = data
+}
+
+// pivotDescending loops over nonzero BSI values in descending order.
+// For each value, the provided function is called with the value and a slice of the associated columns.
+func (bsi bsiData) pivotDescending(filter *Row, branch uint64, limit, offset *uint64, fn func(uint64, ...uint64)) {
+	// This "pivot" algorithm works by treating the BSI data as a tree.
+	// Each branch of this tree corresponds to a power-of-2-sized range of BSI values.
+	// Each range is subdivided into 2 ranges of half size, which form lower branches.
+	// Eventually, a range of width 1 cannot be subdivided and forms a leaf.
+	// At each branch and leaf, there is a bitmap of all columns within the corresponding range.
+	// The lower branches are formed as a difference or intersect of the upper branch's bitmap with the BSI bit that subdivides the range.
+	// This function uses a depth-first search over this virtual tree.
+
+	switch {
+	case !filter.Any():
+		// There are no remaining data.
+
+	case offset != nil && *offset >= filter.Count():
+		// Skip this entire branch.
+		*offset -= filter.Count()
+
+	case limit != nil && *limit == 0:
+		// The limit has been reached.
+		// No more data is necessary.
+
+	case len(bsi) == 0:
+		// This is a leaf node.
+		cols := filter.Columns()
+		if offset != nil {
+			cols = cols[*offset:]
+			*offset = 0
+		}
+		if limit != nil {
+			if *limit < uint64(len(cols)) {
+				cols = cols[:*limit]
+			}
+			*limit -= uint64(len(cols))
+		}
+		fn(branch, cols...)
+
+	default:
+		// Pivot over the highest bit.
+		upperBranch, lowerBranch := branch|(1<<uint(len(bsi)-1)), branch
+		splitBit := bsi[len(bsi)-1]
+		lowerBits := bsi[:len(bsi)-1]
+		lowerBits.pivotDescending(filter.Intersect(splitBit), upperBranch, limit, offset, fn)
+		lowerBits.pivotDescending(filter.Difference(splitBit), lowerBranch, limit, offset, fn)
+	}
+}
+
+/*
+// distribution generates a BSI histogram for the input.
+// TODO: I forgot what I was going to use this for.
+// Could probbably use this for:
+// - quartile queries
+// - TopN on int
+func (bsi bsiData) distribution(filter *Row) bsiData {
+	var dist bsiData
+	bsi.pivotDescending(filter, 0, nil, nil, func(count uint64, values ...uint64) {
+		dist.insert(count, uint64(len(values)))
+	})
+	return dist
+}
+*/
+
+// addBSI adds BSI values together.
+func addBSI(x, y bsiData) bsiData {
+	if len(x) > len(y) {
+		x, y = y, x
+	}
+	carry := NewRow()
+	out := make(bsiData, 0, len(y))
+	for i, v := range x {
+		out = append(out, v.Xor(y[i]).Xor(carry))
+		carry = v.Intersect(y[i]).Union(v.Intersect(carry), y[i].Intersect(carry))
+	}
+	for _, v := range y[len(x):] {
+		out = append(out, v.Xor(carry))
+		carry = v.Intersect(carry)
+	}
+	if carry.Any() {
+		out = append(out, carry)
+	}
+	return out
+}

encoding/proto/proto.go

@@ -312,6 +312,14 @@ func (s Serializer) Unmarshal(buf []byte, m pilosa.Message) error {
 		}
 		s.decodeAtomicRecord(msg, mt)
 		return nil
+	case *[]*pilosa.Row:
+		msg := &internal.RowMatrix{}
+		err := proto.Unmarshal(buf, msg)
+		if err != nil {
+			return errors.Wrap(err, "unmarshaling RowMatrix")
+		}
+		*mt = s.decodeRowMatrix(msg)
+		return nil
 	default:
 		panic(fmt.Sprintf("unhandled pilosa.Message of type %T: %#v", mt, m))
 	}
@@ -542,6 +550,9 @@ func (s Serializer) encodeQueryResponse(m *pilosa.QueryResponse) *internal.Query
 		case pilosa.PairField:
 			pb.Results[i].Type = queryResultTypePairField
 			pb.Results[i].PairField = s.encodePairField(result)
+		case []*pilosa.Row:
+			pb.Results[i].Type = queryResultTypeRowMatrix
+			pb.Results[i].RowMatrix = s.encodeRowMatrix(result)
 		case nil:
 			pb.Results[i].Type = queryResultTypeNil
 		default:
@@ -900,6 +911,15 @@ func (s Serializer) encodeAtomicRecord(msg *pilosa.AtomicRecord) *internal.Atomi
 	return ar
 }
 
+func (s Serializer) encodeRowMatrix(msg []*pilosa.Row) *internal.RowMatrix {
+	rows := make([]*internal.Row, len(msg))
+	for i, r := range msg {
+		rows[i] = s.encodeRow(r)
+	}
+
+	return &internal.RowMatrix{Rows: rows}
+}
+
 func (s Serializer) encodeTransaction(trns *pilosa.Transaction) *internal.Transaction {
 	if trns == nil {
 		return nil
@@ -1341,6 +1361,14 @@ func (s Serializer) decodeAtomicRecord(pb *internal.AtomicRecord, m *pilosa.Atom
 	}
 }
 
+func (s Serializer) decodeRowMatrix(pb *internal.RowMatrix) []*pilosa.Row {
+	rows := make([]*pilosa.Row, len(pb.Rows))
+	for i, r := range pb.Rows {
+		rows[i] = s.decodeRow(r)
+	}
+	return rows
+}
+
 func decodeTransaction(pb *internal.Transaction, trns *pilosa.Transaction) {
 	trns.ID = pb.ID
 	trns.Active = pb.Active
@@ -1364,6 +1392,7 @@ const (
 	queryResultTypeRowIdentifiers
 	queryResultTypePair
 	queryResultTypePairField
+	queryResultTypeRowMatrix
 	queryResultTypeSignedRow
 	queryResultTypeExtractedIDMatrix
 	queryResultTypeExtractedTable
@@ -1401,6 +1430,8 @@ func (s Serializer) decodeQueryResult(pb *internal.QueryResult) interface{} {
 		return s.decodeExtractedIDMatrix(pb.ExtractedIDMatrix)
 	case queryResultTypeExtractedTable:
 		return s.decodeExtractedTable(pb.ExtractedTable)
+	case queryResultTypeRowMatrix:
+		return s.decodeRowMatrix(pb.RowMatrix)
 	}
 	panic(fmt.Sprintf("unknown type: %d", pb.Type))
 }

executor.go

@@ -324,6 +324,12 @@ func (e *executor) safeCopy(resp QueryResponse) (out QueryResponse) {
 			// so does not contain bitmap material, and
 			// should not need to be cloned.
 			out.Results = append(out.Results, x)
+		case []*Row:
+			safe := make([]*Row, len(x))
+			for i, v := range x {
+				safe[i] = v.Clone()
+			}
+			out.Results = append(out.Results, safe)
 		default:
 			panic(fmt.Sprintf("handle %T here", v))
 		}
@@ -703,6 +709,9 @@ func (e *executor) executeCall(ctx context.Context, qcx *Qcx, index string, c *p
 	case "SetColumnAttrs":
 		statFn()
 		return nil, e.executeSetColumnAttrs(ctx, qcx, index, c, opt)
+	case "TopK":
+		statFn()
+		return e.executeTopK(ctx, qcx, index, c, shards, opt)
 	case "TopN":
 		statFn()
 		return e.executeTopN(ctx, qcx, index, c, shards, opt)
@@ -1808,6 +1817,126 @@ func (e *executor) executeMaxRowShard(ctx context.Context, qcx *Qcx, index strin
 	}, nil
 }
 
+func (e *executor) executeTopK(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shards []uint64, opt *execOptions) (interface{}, error) {
+	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeTopK")
+	defer span.Finish()
+
+	mapFn := func(ctx context.Context, shard uint64) (_ interface{}, err error) {
+		return e.executeTopKShard(ctx, qcx, index, c, shard)
+	}
+
+	reduceFn := func(ctx context.Context, prev, v interface{}) interface{} {
+		x, _ := prev.([]*Row)
+		y, _ := v.([]*Row)
+		return ([]*Row)(addBSI(x, y))
+	}
+
+	other, err := e.mapReduce(ctx, index, shards, c, opt, mapFn, reduceFn)
+	if err != nil {
+		return nil, err
+	}
+	results, _ := other.([]*Row)
+
+	if opt.Remote {
+		return results, nil
+	}
+
+	k, hasK, err := c.UintArg("k")
+	if err != nil {
+		return nil, errors.Wrap(err, "fetching k")
+	}
+
+	var limit *uint64
+	if hasK {
+		limit = &k
+	}
+
+	var dst []Pair
+	bsiData(results).pivotDescending(NewRow().Union(results...), 0, limit, nil, func(count uint64, ids ...uint64) {
+		for _, id := range ids {
+			dst = append(dst, Pair{
+				ID:    id,
+				Count: count,
+			})
+		}
+	})
+
+	fieldName, hasFieldName, err := c.StringArg("_field")
+	if err != nil {
+		return nil, errors.Wrap(err, "fetching TopK field")
+	} else if !hasFieldName {
+		return nil, errors.New("missing field in TopK")
+	}
+
+	return &PairsField{
+		Pairs: dst,
+		Field: fieldName,
+	}, nil
+}
+
+func (e *executor) executeTopKShard(ctx context.Context, qcx *Qcx, index string, c *pql.Call, shard uint64) (_ []*Row, err0 error) {
+	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeTopKShard")
+	defer span.Finish()
+
+	// Look up the index.
+	idx := e.Holder.Index(index)
+	if idx == nil {
+		return nil, ErrIndexNotFound
+	}
+
+	// Look up the field.
+	fieldName, hasFieldName, err := c.StringArg("_field")
+	if err != nil {
+		return nil, errors.Wrap(err, "fetching TopK field")
+	} else if !hasFieldName {
+		return nil, errors.New("missing field in TopK")
+	}
+	f := idx.Field(fieldName)
+	if f == nil {
+		return nil, ErrFieldNotFound
+	}
+
+	// Fetch the filter.
+	var filterBitmap *Row
+	if filter, hasFilter, err := c.CallArg("filter"); err != nil {
+		return nil, err
+	} else if hasFilter {
+		filterBitmap, err = e.executeBitmapCallShard(ctx, qcx, index, filter, shard)
+		if err != nil {
+			return nil, err
+		}
+		if !filterBitmap.Any() {
+			return []*Row(nil), nil
+		}
+	}
+
+	tx, finisher, err := qcx.GetTx(Txo{Write: !writable, Index: idx, Shard: shard})
+	if err != nil {
+		return nil, err
+	}
+	defer finisher(&err0)
+
+	ftype := f.Type()
+	switch ftype {
+	case FieldTypeSet, FieldTypeTime:
+		return e.executeTopKShardSet(ctx, tx, filterBitmap, index, fieldName, shard)
+	default:
+		return nil, errors.Errorf("field type %q is not yet supported by TopK", ftype)
+	}
+}
+
+func (e *executor) executeTopKShardSet(ctx context.Context, tx Tx, filter *Row, index, field string, shard uint64) ([]*Row, error) {
+	span, ctx := tracing.StartSpanFromContext(ctx, "Executor.executeTopKShardSet")
+	defer span.Finish()
+
+	f := e.Holder.fragment(index, field, viewStandard, shard)
+	if f == nil {
+		return nil, nil
+	}
+
+	return f.cardinalityBSISet(ctx, tx, filter)
+}
+
 // executeTopN executes a TopN() call.
 // This first performs the TopN() to determine the top results and then
 // requeries to retrieve the full counts for each of the top results.

executor_test.go

@@ -1063,6 +1063,37 @@ func TestExecutor_Execute_SetRowAttrs(t *testing.T) {
 	})
 }
 
+func TestExecutor_Execute_TopK(t *testing.T) {
+	c := test.MustRunCluster(t, 2)
+	defer c.Close()
+
+	// Load some test data into a set field.
+	c.CreateField(t, "i", pilosa.IndexOptions{TrackExistence: true}, "f")
+	c.ImportBits(t, "i", "f", [][2]uint64{
+		{0, 0},
+		{0, 1},
+		{0, ShardWidth + 2},
+		{10, 2},
+		{10, ShardWidth},
+		{10, 2 * ShardWidth},
+		{10, ShardWidth + 1},
+		{20, ShardWidth},
+	})
+
+	// Execute query.
+	if result, err := c.GetNode(0).API.Query(context.Background(), &pilosa.QueryRequest{Index: "i", Query: `TopK(f, k=2)`}); err != nil {
+		t.Fatal(err)
+	} else if !reflect.DeepEqual(result.Results, []interface{}{&pilosa.PairsField{
+		Pairs: []pilosa.Pair{
+			{ID: 10, Count: 4},
+			{ID: 0, Count: 3},
+		},
+		Field: "f",
+	}}) {
+		t.Fatalf("unexpected result: %s", spew.Sdump(result))
+	}
+}
+
 // Ensure a TopN() query can be executed.
 func TestExecutor_Execute_TopN(t *testing.T) {
 	t.Run("RowIDColumnID", func(t *testing.T) {

fragment.go

@@ -1775,6 +1775,36 @@ func (f *fragment) forEachBit(tx Tx, fn func(rowID, columnID uint64) error) erro
 	})
 }
 
+// cardinalityBSISet constructs a perpendicular BSI bitmap containing the cardinality of each specified row in a set field.
+func (f *fragment) cardinalityBSISet(ctx context.Context, tx Tx, filter *Row) ([]*Row, error) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+
+	// Fetch row IDs.
+	rowIDs, err := f.unprotectedRows(ctx, tx, 0)
+	if err != nil {
+		return nil, err
+	}
+
+	// Count the bits in each row.
+	var out bsiData
+	for _, id := range rowIDs {
+		row, err := f.unprotectedRow(tx, id)
+		if err != nil {
+			return nil, err
+		}
+		var count uint64
+		if filter != nil {
+			count = row.intersectionCount(filter)
+		} else {
+			count = row.Count()
+		}
+		out.insert(id, count)
+	}
+
+	return out, nil
+}
+
 // top returns the top rows from the fragment.
 // If opt.Src is specified then only rows which intersect src are returned.
 // If opt.FilterValues exist then the row attribute specified by field is matched.
@@ -3100,11 +3130,16 @@ func (f *fragment) unprotectedRows(ctx context.Context, tx Tx, start uint64, fil
 	var lastRow uint64 = math.MaxUint64
 
 	// Loop over the existing containers.
+	var k uint16
 	for i.Next() {
-		// caller doesn't need a result anymore.
-		if err := ctx.Err(); err != nil {
-			return nil, err
+		if k == 0 {
+			if err := ctx.Err(); err != nil {
+				// caller doesn't need a result anymore.
+				return nil, err
+			}
 		}
+		k++
+
 		key, c := i.Value()
 
 		// virtual row for the current container

internal/public.proto

@@ -9,6 +9,10 @@ message Row {
 	bytes Roaring = 4;
 }
 
+message RowMatrix {
+	repeated Row Rows = 1;
+}
+
 message SignedRow {
 	Row Pos = 1;
 	Row Neg = 2;
@@ -161,6 +165,7 @@ message QueryResult {
 	PairField PairField = 12;
 	ExtractedIDMatrix ExtractedIDMatrix = 13;
 	ExtractedTable ExtractedTable = 14;
+	RowMatrix RowMatrix = 15;
 }
 
 message ImportRequest {

pql/ast.go

@@ -416,6 +416,15 @@ var callInfoByFunc = map[string]callInfo{
 		callType: PrecallGlobal,
 	},
 
+	"TopK": {
+		allowUnknown: false,
+		prototypes: map[string]interface{}{
+			"_field": "",
+			"k":      int64(0),
+			"filter": nil,
+		},
+	},
+
 	// things that take _field
 	"TopN": allowUnderField,
 	// special cases:

pql/pql.peg

@@ -13,6 +13,7 @@ Call <-  "Set" {p.startCall("Set")} open col comma args (comma timestamp)? close
        / "ClearRow" {p.startCall("ClearRow")} open arg close {p.endCall()}
        / "Store" {p.startCall("Store")} open Call comma arg close {p.endCall()}
        / "TopN" {p.startCall("TopN")} open posfield (comma allargs)? close {p.endCall()}
+       / "TopK" {p.startCall("TopK")} open posfield (comma allargs)? close {p.endCall()}
        / "Rows" {p.startCall("Rows")} open posfield (comma allargs)? close {p.endCall()}
        / "Range" {p.startCall("Range")} open field eq value comma 'from='? {p.addField("from")} timestampfmt {p.addVal(text)} comma 'to='? sp {p.addField("to")} timestampfmt {p.addVal(text)} close {p.endCall()}
        / < IDENT > { p.startCall(text) } open allargs comma? close { p.endCall() }