Merge pull request #1020 from seebs/distinctSet

Distinct operations on set fields

executor.go

@@ -1398,6 +1398,10 @@ func (e *executor) executeDistinctShard(ctx context.Context, qcx *Qcx, index str
 
 	var filter *Row
 	var filterBitmap *roaring.Bitmap
+	// If a filter *is* specified, an empty filter means nothing, and any
+	// filter at all means there's filtering to do. If a filter is *not*
+	// specified, then we don't need to do any filtering. So a nil
+	// filterBitmap (which we get if there's no children) means no filter.
 	if len(c.Children) == 1 {
 		row, err := e.executeBitmapCallShard(ctx, qcx, index, c.Children[0], shard)
 		if err != nil {
@@ -1406,17 +1410,85 @@ func (e *executor) executeDistinctShard(ctx context.Context, qcx *Qcx, index str
 		filter = row
 		if filter != nil && len(filter.segments) > 0 {
 			filterBitmap = filter.segments[0].data
-		} else {
-			filterBitmap = roaring.NewFileBitmap()
+		}
+		// if we had a filter to consider, but it came back empty, we
+		// can go ahead and save time by returning the empty results,
+		// because the filter excluded everything.
+		if filterBitmap == nil || !filterBitmap.Any() {
+			return SignedRow{}, nil
 		}
 	}
 
 	bsig := field.bsiGroup(fieldName)
 	if bsig == nil {
-		return result, nil
+		return executeDistinctShardSet(ctx, qcx, idx, fieldName, shard, filterBitmap)
 	}
-	view := viewBSIGroupPrefix + fieldName
+	return executeDistinctShardBSI(ctx, qcx, idx, fieldName, shard, bsig, filterBitmap)
+}
 
+func executeDistinctShardSet(ctx context.Context, qcx *Qcx, idx *Index, fieldName string, shard uint64, filterBitmap *roaring.Bitmap) (result SignedRow, err error) {
+	index := idx.Name()
+	tx, finisher := qcx.GetTx(Txo{Write: !writable, Index: idx, Shard: shard})
+	defer finisher(&err)
+
+	fragData, _, err := tx.ContainerIterator(index, fieldName, "standard", shard, 0)
+	if err != nil {
+		return SignedRow{}, errors.Wrap(err, "getting fragment data")
+	}
+	// We can't grab the containers "for each row" from the set-type field,
+	// because we don't know how many rows there are, and some of them
+	// might be empty, so really, we're going to iterate through the
+	// containers, and then intersect them with the filter if present.
+	var filter []*roaring.Container
+	if filterBitmap != nil {
+		filter = make([]*roaring.Container, 1<<shardVsContainerExponent)
+		filterIterator, _ := filterBitmap.Containers.Iterator(0)
+		// So let's get these all with a nice convenient 0 offset...
+		for filterIterator.Next() {
+			k, c := filterIterator.Value()
+			if c.N() == 0 {
+				continue
+			}
+			filter[k%(1<<shardVsContainerExponent)] = c
+		}
+	}
+	rows := roaring.NewSliceBitmap()
+	prevRow := ^uint64(0)
+	seenThisRow := false
+	for fragData.Next() {
+		k, c := fragData.Value()
+		row := k >> shardVsContainerExponent
+		if row == prevRow {
+			if seenThisRow {
+				continue
+			}
+		} else {
+			seenThisRow = false
+			prevRow = row
+		}
+		if filterBitmap != nil {
+			if roaring.IntersectionAny(c, filter[k%(1<<shardVsContainerExponent)]) {
+				_, err = rows.Add(row)
+				if err != nil {
+					return SignedRow{}, errors.Wrap(err, "collecting results")
+				}
+				seenThisRow = true
+			}
+		} else if c.N() != 0 {
+			_, err = rows.Add(row)
+			if err != nil {
+				return SignedRow{}, errors.Wrap(err, "recording results")
+			}
+			seenThisRow = true
+		}
+	}
+
+	return SignedRow{Pos: NewRowFromBitmap(rows)}, nil
+}
+
+func executeDistinctShardBSI(ctx context.Context, qcx *Qcx, idx *Index, fieldName string, shard uint64, bsig *bsiGroup, filterBitmap *roaring.Bitmap) (result SignedRow, err error) {
+	view := viewBSIGroupPrefix + fieldName
+	index := idx.Name()
 	depth := uint64(bsig.BitDepth)
 	offset := bsig.Base
 
@@ -1427,7 +1499,7 @@ func (e *executor) executeDistinctShard(ctx context.Context, qcx *Qcx, index str
 	if err != nil {
 		return result, err
 	}
-	if filter != nil {
+	if filterBitmap != nil {
 		existsBitmap = existsBitmap.Intersect(filterBitmap)
 	}
 	if !existsBitmap.Any() {

executor_test.go

@@ -5322,6 +5322,9 @@ func TestExecutor_ForeignIndex(t *testing.T) {
 		pilosa.OptFieldTypeInt(0, math.MaxInt64),
 		pilosa.OptFieldForeignIndex("parent"),
 	)
+	c.CreateField(t, "child", pilosa.IndexOptions{}, "parent_set_id",
+		pilosa.OptFieldForeignIndex("parent"),
+	)
 	c.CreateField(t, "child", pilosa.IndexOptions{}, "color",
 		pilosa.OptFieldKeys(),
 	)
@@ -5347,6 +5350,12 @@ func TestExecutor_ForeignIndex(t *testing.T) {
 			Set(%d, parent_id="one")
 			Set(4, parent_id="twenty-one")
 		`, ShardWidth))
+	c.Query(t, "child", fmt.Sprintf(`
+			Set(1, parent_set_id="one")
+			Set(2, parent_set_id="two")
+			Set(%d, parent_set_id="one")
+			Set(4, parent_set_id="twenty-one")
+		`, ShardWidth))
 
 	// Populate color data.
 	c.Query(t, "child", fmt.Sprintf(`
@@ -5360,6 +5369,10 @@ func TestExecutor_ForeignIndex(t *testing.T) {
 	if !sameStringSlice(distinct.Pos.Keys, []string{"one", "two", "twenty-one"}) {
 		t.Fatalf("unexpected keys: %v", distinct.Pos.Keys)
 	}
+	distinct = c.Query(t, "child", `Distinct(index="child", field="parent_set_id")`).Results[0].(pilosa.SignedRow)
+	if !sameStringSlice(distinct.Pos.Keys, []string{"one", "two", "twenty-one"}) {
+		t.Fatalf("unexpected keys: %v", distinct.Pos.Keys)
+	}
 
 	eq := c.Query(t, "child", `Row(parent_id=="one")`).Results[0].(*pilosa.Row)
 	if !reflect.DeepEqual(eq.Columns(), []uint64{1, ShardWidth}) {
@@ -5375,6 +5388,10 @@ func TestExecutor_ForeignIndex(t *testing.T) {
 	if !reflect.DeepEqual(join.Keys, []string{"one"}) {
 		t.Fatalf("unexpected keys: %v", join.Keys)
 	}
+	join = c.Query(t, "parent", fmt.Sprintf(`Intersect(Row(general=%d), Distinct(Row(color="blue"), index="child", field="parent_set_id"))`, ShardWidth)).Results[0].(*pilosa.Row)
+	if !reflect.DeepEqual(join.Keys, []string{"one"}) {
+		t.Fatalf("unexpected keys: %v", join.Keys)
+	}
 }
 
 // sameStringSlice is a helper function which compares two string
@@ -6430,16 +6447,26 @@ func TestExecutor_BareDistinct(t *testing.T) {
 	c.CreateField(t, "i", pilosa.IndexOptions{}, "ints",
 		pilosa.OptFieldTypeInt(0, math.MaxInt64),
 	)
+	c.CreateField(t, "i", pilosa.IndexOptions{}, "set")
+	c.CreateField(t, "i", pilosa.IndexOptions{}, "filter")
 
 	// Populate integer data.
 	c.Query(t, "i", fmt.Sprintf(`
 			Set(0, ints=1)
 			Set(%d, ints=2)
 		`, ShardWidth))
+	c.Query(t, "i", fmt.Sprintf(`
+		         Set(0, set=1)
+			 Set(1, set=2)
+			 Set(%d, set=2)
+			 Set(0, filter=1)
+			 Set(%d, filter=1)
+	        `, 65537, 65537))
 
 	for _, pql := range []string{
 		`Distinct(field="ints")`,
 		`Distinct(index="i", field="ints")`,
+		`Distinct(Row(filter=1), field="set")`,
 	} {
 		exp := []uint64{1, 2}
 		res := c.Query(t, "i", pql).Results[0].(pilosa.SignedRow)

field.go

@@ -932,7 +932,7 @@ func (f *Field) applyOptions(opt FieldOptions) error {
 		f.options.BitDepth = 0
 		f.options.TimeQuantum = ""
 		f.options.Keys = opt.Keys
-		f.options.ForeignIndex = ""
+		f.options.ForeignIndex = opt.ForeignIndex
 	case FieldTypeInt, FieldTypeDecimal:
 		f.options.Type = opt.Type
 		f.options.CacheType = CacheTypeNone
@@ -978,7 +978,7 @@ func (f *Field) applyOptions(opt FieldOptions) error {
 			f.Close()
 			return errors.Wrap(err, "setting time quantum")
 		}
-		f.options.ForeignIndex = ""
+		f.options.ForeignIndex = opt.ForeignIndex
 	case FieldTypeBool:
 		f.options.Type = FieldTypeBool
 		f.options.CacheType = CacheTypeNone

http/handler.go

@@ -1286,8 +1286,6 @@ func (o *fieldOptions) validate() error {
 			return pilosa.NewBadRequestError(errors.New("max does not apply to field type set"))
 		} else if o.TimeQuantum != nil {
 			return pilosa.NewBadRequestError(errors.New("timeQuantum does not apply to field type set"))
-		} else if o.ForeignIndex != nil {
-			return pilosa.NewBadRequestError(errors.New("set field cannot be a foreign key"))
 		}
 	case pilosa.FieldTypeInt:
 		if o.CacheType != nil {
@@ -1322,8 +1320,6 @@ func (o *fieldOptions) validate() error {
 			return pilosa.NewBadRequestError(errors.New("max does not apply to field type time"))
 		} else if o.TimeQuantum == nil {
 			return pilosa.NewBadRequestError(errors.New("timeQuantum is required for field type time"))
-		} else if o.ForeignIndex != nil {
-			return pilosa.NewBadRequestError(errors.New("time field cannot be a foreign key"))
 		}
 	case pilosa.FieldTypeMutex:
 		if o.CacheType == nil {
@@ -1338,8 +1334,6 @@ func (o *fieldOptions) validate() error {
 			return pilosa.NewBadRequestError(errors.New("max does not apply to field type mutex"))
 		} else if o.TimeQuantum != nil {
 			return pilosa.NewBadRequestError(errors.New("timeQuantum does not apply to field type mutex"))
-		} else if o.ForeignIndex != nil {
-			return pilosa.NewBadRequestError(errors.New("mutex field cannot be a foreign key"))
 		}
 	case pilosa.FieldTypeBool:
 		if o.CacheType != nil {

roaring/roaring.go

@@ -3969,6 +3969,156 @@ func flipRun(b *Container) *Container {
 	return flipBitmap(x)
 }
 
+// IntersectionAny checks whether two containers have any overlap without
+// counting past the first bit found.
+func IntersectionAny(a, b *Container) bool {
+	return intersectionAny(a, b)
+}
+
+func intersectionAny(a, b *Container) bool {
+	an := a.N()
+	if an == 0 {
+		return false
+	}
+	bn := b.N()
+	if bn == 0 {
+		return false
+	}
+	if an+bn > MaxContainerVal+1 {
+		return true
+	}
+	if a.isArray() {
+		if b.isArray() {
+			return intersectionAnyArrayArray(a, b)
+		} else if b.isRun() {
+			return intersectionAnyArrayRun(a, b)
+		} else {
+			return intersectionAnyArrayBitmap(a, b)
+		}
+	} else if a.isRun() {
+		if b.isArray() {
+			return intersectionAnyArrayRun(b, a)
+		} else if b.isRun() {
+			return intersectionAnyRunRun(a, b)
+		} else {
+			return intersectionAnyRunBitmap(a, b)
+		}
+	} else {
+		if b.isArray() {
+			return intersectionAnyArrayBitmap(b, a)
+		} else if b.isRun() {
+			return intersectionAnyRunBitmap(b, a)
+		} else {
+			return intersectionAnyBitmapBitmap(a, b)
+		}
+	}
+}
+
+func intersectionAnyArrayArray(a, b *Container) bool {
+	ca, cb := a.array(), b.array()
+	nb := len(cb)
+	j := 0
+	for _, va := range ca {
+		for cb[j] < va {
+			j++
+			if j >= nb {
+				return false
+			}
+		}
+		if cb[j] == va {
+			return true
+		}
+	}
+	return false
+}
+
+func intersectionAnyArrayRun(a, b *Container) bool {
+	array, runs := a.array(), b.runs()
+	na, nb := len(array), len(runs)
+	for i, j := 0, 0; i < na && j < nb; {
+		va, vb := array[i], runs[j]
+		if va < vb.Start {
+			i++
+		} else if va >= vb.Start && va <= vb.Last {
+			return true
+		} else if va > vb.Last {
+			j++
+		}
+	}
+	return false
+}
+
+func intersectionAnyArrayBitmap(a, b *Container) bool {
+	bitmap := b.bitmap()[:1024]
+	for _, val := range a.array() {
+		i := int(val >> 6)
+		off := val % 64
+		if (bitmap[i]>>off)&1 != 0 {
+			return true
+		}
+	}
+	return false
+}
+
+func intersectionAnyRunRun(a, b *Container) bool {
+	ra, rb := a.runs(), b.runs()
+	na, nb := len(ra), len(rb)
+	for i, j := 0, 0; i < na && j < nb; {
+		va, vb := ra[i], rb[j]
+		if va.Last < vb.Start {
+			// |--va--| |--vb--|
+			i++
+		} else if va.Start > vb.Last {
+			// |--vb--| |--va--|
+			j++
+		} else {
+			// va.Last >= vb.Start, and va.Start <= vb.Last,
+			// means there must be overlap
+			return true
+		}
+	}
+	return false
+}
+
+func intersectionAnyRunBitmap(a, b *Container) bool {
+	bb := b.bitmap()[:1024]
+	runs := a.runs()
+	for _, r := range runs {
+		loWord, loBit := r.Start/64, r.Start%64
+		hiWord, hiBit := r.Last/64, r.Last%64
+		if loBit != 0 {
+			w := bb[loWord]
+			mask := (uint64(1) << loBit) - 1
+			if w&^mask != 0 {
+				return true
+			}
+		}
+		for i := loWord; i < hiWord; i++ {
+			if bb[i] != 0 {
+				return true
+			}
+		}
+		if hiBit != 0 {
+			w := bb[hiWord]
+			mask := (uint64(1) << hiBit) - 1
+			if w&mask != 0 {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+func intersectionAnyBitmapBitmap(a, b *Container) bool {
+	ba, bb := a.bitmap()[:1024], b.bitmap()[:1024]
+	for i, v := range ba {
+		if bb[i]&v != 0 {
+			return true
+		}
+	}
+	return false
+}
+
 func intersectionCount(a, b *Container) int32 {
 	if a.N() == MaxContainerVal+1 {
 		return b.N()

roaring/roaring_container_test.go

@@ -87,6 +87,11 @@ func TestIntersectVariants(t *testing.T) {
 						t.Errorf("intersecting %s[%d] and %s[%d]: container has N %d, count was %d",
 							n1, i1, n2, i2, full.N(), count)
 					}
+					any := intersectionAny(c1, c2)
+					if any != (count != 0) {
+						t.Errorf("intersecting %s[%d] and %s[%d]: any %t, count was %d",
+							n1, i1, n2, i2, any, count)
+					}
 				}
 			}
 		}