implement alternate groupByIterator using fragment rowIterator

doesn't re-intersect the same rows for every record

executor.go

@@ -943,7 +943,7 @@ func mergeGroupCounts(gc, other []GroupCount) []GroupCount {
 }
 
 func (e *executor) executeGroupByShard(_ context.Context, index string, c *pql.Call, shard uint64, childRows []RowIDs) ([]GroupCount, error) {
-	iter := newGroupByIterator(childRows, c.Children, index, shard, e.Holder)
+	iter := newGroupByIterator2(childRows, c.Children, index, shard, e.Holder)
 	if iter == nil {
 		return []GroupCount{}, nil
 	}
@@ -2596,10 +2596,122 @@ func filterWithRows(rows []uint64) rowFilter {
 	}
 }
 
+type groupByIterator2 struct {
+	rowIters []*rowIterator
+	rows     []struct {
+		row *Row
+		id  uint64
+	}
+	done   bool
+	fields []FieldRow
+}
+
+func newGroupByIterator2(rowIDs []RowIDs, children []*pql.Call, index string, shard uint64, holder *Holder) *groupByIterator2 {
+	gbi := &groupByIterator2{
+		rowIters: make([]*rowIterator, len(children)),
+		rows: make([]struct {
+			row *Row
+			id  uint64
+		}, len(children)),
+		fields: make([]FieldRow, len(children)),
+	}
+
+	ignorePrev := false
+	for i, call := range children {
+		fieldName := call.Args["field"].(string) // this has already been validated by this point
+		gbi.fields[i].Field = fieldName
+		// Fetch fragment.
+		frag := holder.fragment(index, fieldName, viewStandard, shard)
+		if frag == nil { // this means this whole shard doesn't have all it needs to continue
+			return nil
+		}
+		filters := []rowFilter{}
+		if len(rowIDs[i]) > 0 {
+			filters = append(filters, filterWithRows(rowIDs[i]))
+		}
+		gbi.rowIters[i] = frag.rowIterator(i != 0, filters...)
+
+		prev, hasPrev, err := call.UintArg("previous")
+		if err != nil {
+			panic("getting prev")
+		} else if hasPrev && !ignorePrev {
+			if i == len(children)-1 {
+				prev += 1
+			}
+			gbi.rowIters[i].Seek(prev)
+		}
+		nextRow, rowID, wrapped := gbi.rowIters[i].Next()
+		if nextRow == nil {
+			gbi.done = true
+			return gbi
+		}
+		gbi.rows[i].row = nextRow
+		gbi.rows[i].id = rowID
+		if hasPrev && rowID != prev {
+			ignorePrev = true
+		}
+		if wrapped {
+			for j := i - 1; j >= 0; j-- {
+				nextRow, rowID, wrapped := gbi.rowIters[j].Next()
+				if nextRow == nil {
+					gbi.done = true
+					return gbi
+				}
+				gbi.rows[j].row = nextRow
+				gbi.rows[j].id = rowID
+				if !wrapped {
+					break
+				}
+			}
+		}
+	}
+
+	for i := 1; i < len(gbi.rows); i++ {
+		gbi.rows[i].row = gbi.rows[i].row.Intersect(gbi.rows[i-1].row)
+	}
+
+	return gbi
+}
+
+func (gbi *groupByIterator2) nextAtIdx(i int) {
+	nr, rowID, wrapped := gbi.rowIters[i].Next()
+	if nr == nil {
+		gbi.done = true
+		return
+	}
+	if wrapped && i != 0 {
+		gbi.nextAtIdx(i - 1)
+	}
+	if i != 0 {
+		gbi.rows[i].row = nr.Intersect(gbi.rows[i-1].row)
+	} else {
+		gbi.rows[i].row = nr
+	}
+	gbi.rows[i].id = rowID
+}
+
+func (gbi *groupByIterator2) Next() (ret ppi, done bool) {
+	if gbi.done {
+		return ret, true
+	}
+	ret.row = gbi.rows[len(gbi.rows)-1].row
+	ret.gCnt.Count = ret.row.Count()
+	ret.gCnt.Group = make([]FieldRow, len(gbi.rows))
+	copy(ret.gCnt.Group, gbi.fields)
+	for i, r := range gbi.rows {
+		ret.gCnt.Group[i].RowID = r.id
+	}
+
+	// set up for next call
+	gbi.nextAtIdx(len(gbi.rows) - 1)
+
+	return ret, false
+}
+
 type groupByIterator struct {
 	fragments []*fragment
 	current   []uint64
-	rows      []RowIDs
+	rowIDs    []RowIDs
 	fields    []FieldRow
 }
 
@@ -2607,7 +2719,7 @@ func newGroupByIterator(rowIDs []RowIDs, children []*pql.Call, index string, sha
 	gbi := &groupByIterator{
 		fragments: make([]*fragment, len(rowIDs)),
 		current:   make([]uint64, len(rowIDs)),
-		rows:      rowIDs,
+		rowIDs:    rowIDs,
 		fields:    make([]FieldRow, len(rowIDs)),
 	}
 	for i, call := range children {
@@ -2645,16 +2757,16 @@ func (gbi *groupByIterator) Next() (ret ppi, done bool) {
 		}
 		frag := gbi.fragments[i]
 		filters := []rowFilter{}
-		if len(gbi.rows[i]) > 0 {
-			filters = append(filters, filterWithRows(gbi.rows[i]))
+		if len(gbi.rowIDs[i]) > 0 {
+			filters = append(filters, filterWithRows(gbi.rowIDs[i]))
 		}
 		filters = append(filters, filterWithLimit(1))
 		rowIDs := frag.rows(gbi.current[i], filters...)
 		if len(rowIDs) == 0 && i != 0 { // wrap around
 			wrap = true
 			gbi.current[i] = 0
-			if len(gbi.rows[i]) > 0 {
-				gbi.current[i] = gbi.rows[i][0]
+			if len(gbi.rowIDs[i]) > 0 {
+				gbi.current[i] = gbi.rowIDs[i][0]
 			}
 			rowIDs = frag.rows(gbi.current[i], filters...)
 		}

fragment.go

@@ -2013,6 +2013,42 @@ func (f *fragment) rows(start uint64, filters ...rowFilter) []uint64 {
 	return rows
 }
 
+type rowIterator struct {
+	f      *fragment
+	rowIDs []uint64
+	cur    int
+	wrap   bool
+}
+
+func (f *fragment) rowIterator(wrap bool, filters ...rowFilter) *rowIterator {
+	return &rowIterator{
+		f:      f,
+		rowIDs: f.rows(0, filters...), // TODO: this may be memory intensive in high cardinality cases
+		wrap:   wrap,
+	}
+}
+
+func (ri *rowIterator) Seek(rowID uint64) {
+	idx := sort.Search(len(ri.rowIDs), func(i int) bool {
+		return ri.rowIDs[i] >= rowID
+	})
+	ri.cur = idx
+}
+
+func (ri *rowIterator) Next() (r *Row, rowID uint64, wrapped bool) {
+	if ri.cur >= len(ri.rowIDs) {
+		if !ri.wrap || len(ri.rowIDs) == 0 {
+			return nil, 0, true
+		}
+		ri.Seek(0)
+		wrapped = true
+	}
+	rowID = ri.rowIDs[ri.cur]
+	r = ri.f.row(rowID)
+	ri.cur += 1
+	return r, rowID, wrapped
+}
+
 // FragmentBlock represents info about a subsection of the rows in a block.
 // This is used for comparing data in remote blocks for active anti-entropy.
 type FragmentBlock struct {

fragment_internal_test.go

@@ -1819,3 +1819,122 @@ func calcExpected(inputs ...[]uint64) [][]uint64 {
 
 	return ret
 }
+
+func TestFragmentRowIterator(t *testing.T) {
+	t.Run("basic", func(t *testing.T) {
+		f := mustOpenFragment("i", "f", "v", 0, CacheTypeRanked)
+		f.mustSetBits(0, 0)
+		f.mustSetBits(1, 0)
+		f.mustSetBits(2, 0)
+		f.mustSetBits(3, 0)
+
+		iter := f.rowIterator(false)
+		for i := uint64(0); i < 4; i++ {
+			row, id, wrapped := iter.Next()
+			if id != i {
+				t.Fatalf("expected row %d but got %d", i, id)
+			}
+			if wrapped != false {
+				t.Fatalf("shouldn't have wrapped")
+			}
+			if !reflect.DeepEqual(row.Columns(), []uint64{0}) {
+				t.Fatalf("got wrong columns back on iteration %d - should just be 0 but %v", i, row.Columns())
+			}
+		}
+		row, id, wrapped := iter.Next()
+		if row != nil {
+			t.Fatalf("row should be nil after iterator is exhausted, got %v", row.Columns())
+		}
+		if id != 0 {
+			t.Fatalf("id should be 0 after iterator is exhausted, got %d", id)
+		}
+		if wrapped != true {
+			t.Fatalf("wrapped should be true after iterator is exhausted")
+		}
+		f.Close()
+	})
+
+	t.Run("skipped rows", func(t *testing.T) {
+		f := mustOpenFragment("i", "f", "v", 0, CacheTypeRanked)
+		f.mustSetBits(1, 0)
+		f.mustSetBits(3, 0)
+		f.mustSetBits(5, 0)
+		f.mustSetBits(7, 0)
+
+		iter := f.rowIterator(false)
+		for i := uint64(1); i < 8; i += 2 {
+			row, id, wrapped := iter.Next()
+			if id != i {
+				t.Fatalf("expected row %d but got %d", i, id)
+			}
+			if wrapped != false {
+				t.Fatalf("shouldn't have wrapped")
+			}
+			if !reflect.DeepEqual(row.Columns(), []uint64{0}) {
+				t.Fatalf("got wrong columns back on iteration %d - should just be 0 but %v", i, row.Columns())
+			}
+		}
+		row, id, wrapped := iter.Next()
+		if row != nil {
+			t.Fatalf("row should be nil after iterator is exhausted, got %v", row.Columns())
+		}
+		if id != 0 {
+			t.Fatalf("id should be 0 after iterator is exhausted, got %d", id)
+		}
+		if wrapped != true {
+			t.Fatalf("wrapped should be true after iterator is exhausted")
+		}
+		f.Close()
+	})
+
+	t.Run("basic wrapped", func(t *testing.T) {
+		f := mustOpenFragment("i", "f", "v", 0, CacheTypeRanked)
+		f.mustSetBits(0, 0)
+		f.mustSetBits(1, 0)
+		f.mustSetBits(2, 0)
+		f.mustSetBits(3, 0)
+
+		iter := f.rowIterator(true)
+		for i := uint64(0); i < 5; i++ {
+			row, id, wrapped := iter.Next()
+			if id != i%4 {
+				t.Fatalf("expected row %d but got %d", i%4, id)
+			}
+			if wrapped && i < 4 {
+				t.Fatalf("shouldn't have wrapped")
+			} else if !wrapped && i >= 4 {
+				t.Fatalf("should have wrapped")
+			}
+			if !reflect.DeepEqual(row.Columns(), []uint64{0}) {
+				t.Fatalf("got wrong columns back on iteration %d - should just be 0 but %v", i, row.Columns())
+			}
+		}
+		f.Close()
+	})
+
+	t.Run("skipped rows wrapped", func(t *testing.T) {
+		f := mustOpenFragment("i", "f", "v", 0, CacheTypeRanked)
+		f.mustSetBits(1, 0)
+		f.mustSetBits(3, 0)
+		f.mustSetBits(5, 0)
+		f.mustSetBits(7, 0)
+
+		iter := f.rowIterator(true)
+		for i := uint64(1); i < 10; i += 2 {
+			row, id, wrapped := iter.Next()
+			if id != i%8 {
+				t.Errorf("expected row %d but got %d", i%8, id)
+			}
+			if wrapped && i < 8 {
+				t.Errorf("shouldn't have wrapped")
+			} else if !wrapped && i >= 8 {
+				t.Errorf("should have wrapped")
+			}
+			if !reflect.DeepEqual(row.Columns(), []uint64{0}) {
+				t.Fatalf("got wrong columns back on iteration %d - should just be 0 but %v", i, row.Columns())
+			}
+		}
+		f.Close()
+	})
+
+}