From c172ca0680a3e04aea20f49a06de7148b549008b Mon Sep 17 00:00:00 2001
From: Matt Jaffee <matthew.jaffee@gmail.com>
Date: Mon, 8 Oct 2018 16:16:39 -0500
Subject: [PATCH] combine fragment.rows and rowsForColumn with generalized
 filter

use filter funcs with closures for state instead of methods on structs. seems a
bit cleaner.
---
 executor.go               | 65 ++++++++++++++++-------------
 executor_internal_test.go | 18 ++++++++
 fragment.go               | 87 +++++++++------------------------------
 fragment_internal_test.go |  6 +--
 4 files changed, 78 insertions(+), 98 deletions(-)

diff --git a/executor.go b/executor.go
index 349604139..3d8618b9d 100644
--- a/executor.go
+++ b/executor.go
@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"github.com/pilosa/pilosa/pql"
+	"github.com/pilosa/pilosa/roaring"
 	"github.com/pkg/errors"
 )
 
@@ -1013,19 +1014,18 @@ func (e *executor) executeRowsShard(ctx context.Context, index string, c *pql.Ca
 	}
 
 	filters := []rowFilter{}
-	if limit, hasLimit, err := c.UintArg("limit"); err != nil {
-		return nil, errors.Wrap(err, "getting limit")
-	} else if hasLimit {
-		filters = append(filters, (&filterWithLimit{limit: limit}).filter)
-	}
-
 	if columnID, ok, err := c.UintArg("column"); err != nil {
 		return nil, err
 	} else if ok {
-		return frag.rowsForColumn(start, columnID, filters...), nil
-	} else {
-		return frag.rows(start, filters...), nil
+		filters = append(filters, filterColumn(columnID))
 	}
+	if limit, hasLimit, err := c.UintArg("limit"); err != nil {
+		return nil, errors.Wrap(err, "getting limit")
+	} else if hasLimit {
+		filters = append(filters, filterWithLimit(limit))
+	}
+
+	return frag.rows(start, filters...), nil
 }
 
 // getGroupByFilterFunction returns a rowFilter based on the
@@ -1064,10 +1064,10 @@ func getGroupByFilterFunction(fieldDirective string) (ret []rowFilter, err error
 		}
 	}
 	if hasOffset {
-		ret = append(ret, (&filterWithOffset{offset: offset}).filter)
+		ret = append(ret, filterWithOffset(offset))
 	}
 	if hasLimit {
-		ret = append(ret, (&filterWithLimit{limit: limit}).filter)
+		ret = append(ret, filterWithLimit(limit))
 	}
 	return ret, nil
 }
@@ -2319,22 +2319,31 @@ func isString(v interface{}) bool {
 	return ok
 }
 
-type filterWithOffset struct {
-	offset uint64
-}
-
-func (fo *filterWithOffset) filter(rowID uint64) (bool, bool) {
-	return rowID >= fo.offset, false
-}
-
-type filterWithLimit struct {
-	limit uint64
-}
-
-func (fl *filterWithLimit) filter(rowID uint64) (bool, bool) { // nolint: unparam
-	if fl.limit > 0 {
-		fl.limit--
-		return true, false
+func filterWithOffset(offset uint64) rowFilter {
+	return func(rowID, key uint64, c *roaring.Container) (include, done bool) {
+		return rowID >= offset, false
+	}
+}
+
+// filterWithLimit returns a filter which will only allow a limited number of
+// rows to be returned. It should be applied last so that it is only called (and
+// therefore only updates its internal state) if the row is being included by
+// every other filter.
+func filterWithLimit(limit uint64) rowFilter {
+	return func(rowID, key uint64, c *roaring.Container) (include, done bool) {
+		if limit > 0 {
+			limit--
+			return true, false
+		}
+		return false, true
+	}
+}
+
+func filterColumn(col uint64) rowFilter {
+	return func(rowID, key uint64, c *roaring.Container) (include, done bool) {
+		colID := col % ShardWidth
+		colKey := ((rowID * ShardWidth) + colID) >> 16
+		colVal := uint16(colID & 0xFFFF) // columnID within the container
+		return colKey == key && c.Contains(colVal), false
 	}
-	return false, true
 }
diff --git a/executor_internal_test.go b/executor_internal_test.go
index e90948561..990509128 100644
--- a/executor_internal_test.go
+++ b/executor_internal_test.go
@@ -120,3 +120,21 @@ func isInt(a interface{}) bool {
 		return false
 	}
 }
+
+func TestFilterWithLimit(t *testing.T) {
+	f := filterWithLimit(5)
+
+	for i := uint64(0); i < 5; i++ {
+		include, done := f(i, i*(1<<shardVsContainerExponent), nil)
+		if done {
+			t.Fatalf("limit filter ended early on iteration %d", i)
+		}
+		if !include {
+			t.Fatalf("limit filter should always include until done")
+		}
+	}
+	inc, done := f(5, 5*(1<<shardVsContainerExponent)+1, nil)
+	if !done {
+		t.Fatalf("limit filter should have been done, but got inc: %v done: %v", inc, done)
+	}
+}
diff --git a/fragment.go b/fragment.go
index 6f6b376a5..a738b0829 100644
--- a/fragment.go
+++ b/fragment.go
@@ -1761,16 +1761,21 @@ func (f *fragment) readCacheFromArchive(r io.Reader) error {
 	return nil
 }
 
-// rowFilter is a filter function which takes a rowID
-// and determines if that row should be included in
-// the result set. Additionally, it signals whether
-// to halt processing any more rows. The two bool
-// returned are (1) include row, (2) break further
-// processing.
-type rowFilter func(rowID uint64) (bool, bool)
+// rowFilter is a function signature for controlling iteration over containers
+// in a fragment. It will be invoked on each container found and returns two
+// booleans. The first is whether the row this container is in should be
+// included or skipped, and the second is whether to stop processing or
+// continue.
+type rowFilter func(rowID, key uint64, c *roaring.Container) (include, done bool)
 
-// rows returns all rows by calling rowsWithFilter()
-// with a completely unrestrictive filter.
+// rows returns all rows starting from 'start'. Filters will be applied in
+// order. All filters must return true to include the row. Once a row is
+// included, further containers in that row will be skipped. So, for a row to be
+// included, there must be one container in that row where all filters return
+// true. For a row to be skipped, at least one filter must return false for each
+// container in that row (it need not be the same filter for each). Any filter
+// returning done == true will cause processing to stop and the rows accumulated
+// so far will be returned.
 func (f *fragment) rows(start uint64, filters ...rowFilter) []uint64 {
 	startKey := rowToKey(start)
 	i, _ := f.storage.Containers.Iterator(startKey)
@@ -1779,7 +1784,7 @@ func (f *fragment) rows(start uint64, filters ...rowFilter) []uint64 {
 
 	// Loop over the existing containers.
 	for i.Next() {
-		key, _ := i.Value()
+		key, c := i.Value()
 
 		// virtual row for the current container
 		vRow := key >> shardVsContainerExponent
@@ -1790,71 +1795,19 @@ func (f *fragment) rows(start uint64, filters ...rowFilter) []uint64 {
 		}
 
 		// apply filters
-		addRow := true
+		addRow, done := true, false
 		for _, filter := range filters {
-			add, done := filter(vRow)
+			addRow, done = filter(vRow, key, c)
 			if done {
 				return rows
 			}
-			addRow = add && addRow
 			if !addRow {
 				break
 			}
 		}
 		if addRow {
-			rows = append(rows, vRow)
-		}
-
-		lastRow = vRow
-	}
-	return rows
-}
-
-func (f *fragment) rowsForColumn(start, columnID uint64, filters ...rowFilter) []uint64 {
-	if columnID/ShardWidth != f.shard {
-		panic(fmt.Sprintln("fragment.rowsForColumn should never be called with a columnID which is not in the fragment's shard",
-			columnID, columnID/ShardWidth, f.shard))
-	}
-
-	startKey := rowToKey(start)
-	i, _ := f.storage.Containers.Iterator(startKey)
-	rows := make([]uint64, 0)
-
-	colID := columnID % ShardWidth
-	colVal := uint16(colID & 0xFFFF) // columnID within the container
-
-	var colKey uint64
-
-	// Loop over the existing containers.
-	for i.Next() {
-		key, c := i.Value()
-
-		// virtual row for the current container
-		vRow := key >> shardVsContainerExponent
-
-		// column container key for virtual row
-		colKey = ((vRow * ShardWidth) + colID) >> 16
-
-		if colKey != key {
-			continue
-		}
-
-		// apply filter
-		if c.Contains(colVal) {
-			addRow := true
-			for _, filter := range filters {
-				add, done := filter(vRow)
-				if done {
-					return rows
-				}
-				addRow = add && addRow
-				if !addRow {
-					break
-				}
-			}
-			if addRow {
-				rows = append(rows, vRow)
-			}
+			lastRow = vRow
+			rows = append(rows, key>>shardVsContainerExponent)
 		}
 	}
 	return rows
@@ -2140,7 +2093,7 @@ func newRowsVector(f *fragment) *rowsVector {
 // Additionally, it returns true if a value was found,
 // otherwise it returns false.
 func (v *rowsVector) Get(colID uint64) (uint64, bool) {
-	rows := v.f.rowsForColumn(0, colID)
+	rows := v.f.rows(0, filterColumn(colID))
 	if len(rows) == 1 {
 		return rows[0], true
 	}
diff --git a/fragment_internal_test.go b/fragment_internal_test.go
index f9d0bd999..36651fee0 100644
--- a/fragment_internal_test.go
+++ b/fragment_internal_test.go
@@ -1346,7 +1346,7 @@ func TestFragment_RowsIteration(t *testing.T) {
 			t.Fatalf("Do not match %v %v", expectedAll, ids)
 		}
 
-		ids = f.rowsForColumn(0, 1)
+		ids = f.rows(0, filterColumn(1))
 		if !reflect.DeepEqual(expectedOdd, ids) {
 			t.Fatalf("Do not match %v %v", expectedOdd, ids)
 		}
@@ -1370,7 +1370,7 @@ func TestFragment_RowsIteration(t *testing.T) {
 			t.Fatalf("Do not match %v %v", expected, ids)
 		}
 
-		ids = f.rowsForColumn(0, 66000)
+		ids = f.rows(0, filterColumn(66000))
 		if !reflect.DeepEqual(expected, ids) {
 			t.Fatalf("Do not match %v %v", expected, ids)
 		}
@@ -1392,7 +1392,7 @@ func TestFragment_RowsIteration(t *testing.T) {
 				if !reflect.DeepEqual(expectedRows, ids) {
 					t.Fatalf("Do not match %v %v", expectedRows, ids)
 				}
-				ids = f.rowsForColumn(0, c)
+				ids = f.rows(0, filterColumn(c))
 				if !reflect.DeepEqual(expectedRows, ids) {
 					t.Fatalf("Do not match %v %v", expectedRows, ids)
 				}