Repository: segmentio/parquet-go
Pull request #388 (2022-10-25 23:08:49)

search: add comments to search.go
I was going through this function to check my understanding of page indexes while thinking about granules in TraceDB.

I believe that granules in other databases like clickhouse and frostdb are the equivlant of column pages in parquet.  The page index is the sparse index over the pages (granules) wherein the data being searched for is contained.

This check on the binary search mostly confirms that suspicion, if I'm not mistaken.  Please correct me if I'm misunderstaneding something.

Also, as part of this, I found what is possibly an optimization to avoid skipping ahead a page, by comparing the current value to the max in the current page, instead of letting the binary search check the next page, then falling through to the final if-block.

It may or may not be a bug also, if, for example, the following happens:

you have pages:

    [1,2,3]
    [3,4,5]
    [6,7,8]

Searching for 3 in this case would not return the first page and instead return the second, missing values of 3 in the first page.

I'll have to look at writing a test for this, but the current search test will pass creating this scenario without altering the test somewhat (since it's only looking for presence, not the expected index)

achille-roussel: > My understanding is that the current algorithm returns early if it finds a page with minValue == searchValue

OK, my mistake.

> How is this different from what's implemented in this PR?

The difference would be that we would only do one comparison per iteration: first loop compares min values, second compares max values.

The change you are proposing is good, I just want to explore whether there are opportunities for simplification 👍 
abraithwaite: > The difference would be that we would only do one comparison per iteration: first loop compares min values, second compares max values.

Gotcha. I did think about this given the existing code was only doing one bounds check (on `index.MinValue`) a bit and couldn't figure out a way to handle this kind of case:

```
search: 5
[1,2,3]
[4,5,6]
[7,8,9]
```

Only checking Min on the first loop would skip to page [7,8,9] and would return a bad result (not in columnChunk), if I'm not mistaken.
abraithwaite: Took the feedback regarding multiple checks on every iteration and implemented it, while still not regressing on the bug found previously (#389).

Patch:
diff --git a/search.go b/search.go
index bf23c52..db0ec8c 100644
--- a/search.go
+++ b/search.go
@@ -75,33 +75,65 @@ func Find(index ColumnIndex, value Value, cmp func(Value, Value) int) int {
 
 func binarySearch(index ColumnIndex, value Value, cmp func(Value, Value) int) int {
 	n := index.NumPages()
-	i := 0
-	j := n
+	curIdx := 0
+	topIdx := n
 
-	for (j - i) > 1 {
-		k := ((j - i) / 2) + i
-		c := cmp(value, index.MinValue(k))
+	// while there's at least one more page to check
+	for (topIdx - curIdx) > 1 {
+
+		// nextIdx is set to halfway between curIdx and topIdx
+		nextIdx := ((topIdx - curIdx) / 2) + curIdx
+
+		smallerThanMin := cmp(value, index.MinValue(nextIdx))
 
 		switch {
-		case c < 0:
-			j = k
-		case c > 0:
-			i = k
-		default:
-			return k
+		// search below pages[nextIdx]
+		case smallerThanMin < 0:
+			topIdx = nextIdx
+		// search pages[nextIdx] and above
+		case smallerThanMin > 0:
+			curIdx = nextIdx
+		case smallerThanMin == 0:
+			// this case is hit when winValue == value of nextIdx
+			// we must check below this index to find if there's
+			// another page before this.
+			// e.g. searching for first page 3 is in:
+			// [1,2,3]
+			// [3,4,5]
+			// [6,7,8]
+
+			// if the page proceeding this has a maxValue matching the value we're
+			// searching, continue the search.
+			// otherwise, we can return early
+			//
+			// cases covered by else block
+			// if cmp(value, index.MaxValue(nextIdx-1)) < 0: the value is only in this page
+			// if cmp(value, index.MaxValue(nextIdx-1)) > 0: we've got a sorting problem with overlapping pages
+			//
+			// bounds check not needed for nextIdx-1 because nextIdx is guaranteed to be at least curIdx + 1
+			// line 82 & 85 above
+			if cmp(value, index.MaxValue(nextIdx-1)) == 0 {
+				topIdx = nextIdx
+			} else {
+				return nextIdx
+			}
 		}
 	}
 
-	if i < n {
-		min := index.MinValue(i)
-		max := index.MaxValue(i)
+	// last page check, if it wasn't explicitly found above
+	if curIdx < n {
+
+		// check pages[curIdx] for value
+		min := index.MinValue(curIdx)
+		max := index.MaxValue(curIdx)
 
-		if cmp(value, min) < 0 || cmp(max, value) < 0 {
-			i = n
+		// if value is not in pages[curIdx], then it's not in this columnChunk
+		if cmp(value, min) < 0 || cmp(value, max) > 0 {
+			curIdx = n
 		}
 	}
 
-	return i
+	return curIdx
 }
 
 func linearSearch(index ColumnIndex, value Value, cmp func(Value, Value) int) int {


Test patch:
diff --git a/search_test.go b/search_test.go
index 76f40ac..e10fd99 100644
--- a/search_test.go
+++ b/search_test.go
@@ -1,8 +1,6 @@
 package parquet_test
 
 import (
-	"fmt"
-	"math"
 	"testing"
 
 	"github.com/segmentio/parquet-go"
@@ -34,11 +32,24 @@ func TestSearchBinary(t *testing.T) {
 	testSearch(t, [][]int32{
 		{0, 1, 2, 3, 4, 5, 6, 7, 8, 9},
 		{10, 10, 10, 10},
-		{21, 22, 23, 24, 25},
-		{30},
-		{31},
+		{21, 22, 24, 25, 30},
+		{30, 30},
+		{30, 31},
 		{32},
 		{42, 43, 44, 45, 46, 47, 48, 49},
+	}, [][]int{
+		{10, 1},
+		{0, 0},
+		{9, 0},
+		// non-existant, but would be in this page
+		{23, 2},
+		// ensure we find the first page
+		{30, 2},
+		{31, 4},
+		// out of bounds
+		{99, 7},
+		// out of bounds
+		{-1, 7},
 	})
 }
 
@@ -47,12 +58,23 @@ func TestSearchLinear(t *testing.T) {
 		{10, 10, 10, 10},
 		{0, 1, 2, 3, 4, 5, 6, 7, 8, 9},
 		{21, 22, 23, 24, 25},
-		{19, 18, 17, 16, 15, 14, 13, 12, 11},
+		{19, 18, 17, 16, 14, 13, 12, 11},
 		{42, 43, 44, 45, 46, 47, 48, 49},
+	}, [][]int{
+		{10, 0},
+		{0, 1},
+		{9, 1},
+		{48, 4},
+		// non-existant, but could be in this page
+		{15, 3},
+		// out of bounds
+		{99, 5},
+		// out of bounds
+		{-1, 5},
 	})
 }
 
-func testSearch(t *testing.T, pages [][]int32) {
+func testSearch(t *testing.T, pages [][]int32, expectIndex [][]int) {
 	indexer := parquet.Int32Type.NewColumnIndexer(0)
 
 	for _, values := range pages {
@@ -77,22 +99,12 @@ func testSearch(t *testing.T, pages [][]int32) {
 	formatIndex := indexer.ColumnIndex()
 	columnIndex := parquet.NewColumnIndex(parquet.Int32, &formatIndex)
 
-	for i, values := range pages {
-		t.Run(fmt.Sprintf("page#%02d", i), func(t *testing.T) {
-			for _, value := range values {
-				v := parquet.ValueOf(value)
-				j := parquet.Search(columnIndex, v, parquet.Int32Type)
-
-				if i != j {
-					t.Errorf("searching for value %v: got=%d want=%d", v, j, i)
-				}
-			}
+	for _, values := range expectIndex {
+		v := parquet.ValueOf(values[0])
+		j := parquet.Search(columnIndex, v, parquet.Int32Type)
 
-			for _, test := range []int32{math.MinInt32, math.MaxInt32} {
-				if page := parquet.Search(columnIndex, parquet.ValueOf(test), parquet.Int32Type); page != len(pages) {
-					t.Errorf("search for non-existing value %v: got=%d want=%d", test, page, len(pages))
-				}
-			}
-		})
+		if values[1] != j {
+			t.Errorf("searching for value %v: got=%d want=%d", v, j, values[1])
+		}
 	}
 }