Repository: segmentio/parquet-go
Pull request #127 (2022-04-07 16:29:03)

add parquet.Search
This PR adds a new `parquet.Search` function, which intends to help leverage column indexes to find the page that may contain a value. Since this is the main purpose of the column index, I think it makes sense to provide this feature in the library.

There are two cases:
* the index is in ascending order, in which case we can do a binary search on the min page values
* the index is not ordered, in which case we perform a linear scan for the first match

Technically we could potentially support another case to leverage descending ordering on column indexes, but I don't have any use cases for it at this time.


Patch:
diff --git a/column_index.go b/column_index.go
index d23387b..c927a54 100644
--- a/column_index.go
+++ b/column_index.go
@@ -34,6 +34,57 @@ type ColumnIndex interface {
 	IsDescending() bool
 }
 
+// NewColumnIndex constructs a ColumnIndex instance from the given parquet
+// format column index. The kind argument configures the type of values
+func NewColumnIndex(kind Kind, index *format.ColumnIndex) ColumnIndex {
+	return &formatColumnIndex{
+		kind:  kind,
+		index: index,
+	}
+}
+
+type formatColumnIndex struct {
+	kind  Kind
+	index *format.ColumnIndex
+}
+
+func (f *formatColumnIndex) NumPages() int {
+	return len(f.index.MinValues)
+}
+
+func (f *formatColumnIndex) NullCount(i int) int64 {
+	if len(f.index.NullCounts) > 0 {
+		return f.index.NullCounts[0]
+	}
+	return 0
+}
+
+func (f *formatColumnIndex) NullPage(i int) bool {
+	return len(f.index.NullPages) > 0 && f.index.NullPages[i]
+}
+
+func (f *formatColumnIndex) MinValue(i int) Value {
+	if f.NullPage(i) {
+		return Value{}
+	}
+	return f.kind.Value(f.index.MinValues[i])
+}
+
+func (f *formatColumnIndex) MaxValue(i int) Value {
+	if f.NullPage(i) {
+		return Value{}
+	}
+	return f.kind.Value(f.index.MaxValues[i])
+}
+
+func (f *formatColumnIndex) IsAscending() bool {
+	return f.index.BoundaryOrder == format.Ascending
+}
+
+func (f *formatColumnIndex) IsDescending() bool {
+	return f.index.BoundaryOrder == format.Descending
+}
+
 type emptyColumnIndex struct{}
 
 func (emptyColumnIndex) NumPages() int       { return 0 }
@@ -46,17 +97,43 @@ func (emptyColumnIndex) IsDescending() bool  { return false }
 
 type fileColumnIndex struct{ chunk *fileColumnChunk }
 
-func (i fileColumnIndex) NumPages() int         { return len(i.chunk.columnIndex.NullPages) }
-func (i fileColumnIndex) NullCount(j int) int64 { return i.chunk.columnIndex.NullCounts[j] }
-func (i fileColumnIndex) NullPage(j int) bool   { return i.chunk.columnIndex.NullPages[j] }
-func (i fileColumnIndex) MinValue(j int) Value  { return i.makeValue(i.chunk.columnIndex.MinValues[j]) }
-func (i fileColumnIndex) MaxValue(j int) Value  { return i.makeValue(i.chunk.columnIndex.MaxValues[j]) }
+func (i fileColumnIndex) NumPages() int {
+	return len(i.chunk.columnIndex.NullPages)
+}
+
+func (i fileColumnIndex) NullCount(j int) int64 {
+	if len(i.chunk.columnIndex.NullCounts) > 0 {
+		return i.chunk.columnIndex.NullCounts[j]
+	}
+	return 0
+}
+
+func (i fileColumnIndex) NullPage(j int) bool {
+	return len(i.chunk.columnIndex.NullPages) > 0 && i.chunk.columnIndex.NullPages[j]
+}
+
+func (i fileColumnIndex) MinValue(j int) Value {
+	if i.NullPage(j) {
+		return Value{}
+	}
+	return i.makeValue(i.chunk.columnIndex.MinValues[j])
+}
+
+func (i fileColumnIndex) MaxValue(j int) Value {
+	if i.NullPage(j) {
+		return Value{}
+	}
+	return i.makeValue(i.chunk.columnIndex.MaxValues[j])
+}
+
 func (i fileColumnIndex) IsAscending() bool {
 	return i.chunk.columnIndex.BoundaryOrder == format.Ascending
 }
+
 func (i fileColumnIndex) IsDescending() bool {
 	return i.chunk.columnIndex.BoundaryOrder == format.Descending
 }
+
 func (i *fileColumnIndex) makeValue(b []byte) Value {
 	return i.chunk.column.typ.Kind().Value(b)
 }
diff --git a/search.go b/search.go
new file mode 100644
index 0000000..758b1b3
--- /dev/null
+++ b/search.go
@@ -0,0 +1,113 @@
+package parquet
+
+// CompareNullsFirst constructs a comparison function which assumes that null
+// values are smaller than all other values.
+func CompareNullsFirst(cmp func(Value, Value) int) func(Value, Value) int {
+	return func(a, b Value) int {
+		switch {
+		case a.IsNull():
+			if b.IsNull() {
+				return 0
+			}
+			return -1
+		case b.IsNull():
+			return +1
+		default:
+			return cmp(a, b)
+		}
+	}
+}
+
+// CompareNullsLast constructs a comparison function which assumes that null
+// values are greater than all other values.
+func CompareNullsLast(cmp func(Value, Value) int) func(Value, Value) int {
+	return func(a, b Value) int {
+		switch {
+		case a.IsNull():
+			if b.IsNull() {
+				return 0
+			}
+			return +1
+		case b.IsNull():
+			return -1
+		default:
+			return cmp(a, b)
+		}
+	}
+}
+
+// Search is like Find, but uses the default ordering of the given type.
+func Search(index ColumnIndex, value Value, typ Type) int {
+	return Find(index, value, CompareNullsLast(typ.Compare))
+}
+
+// Find uses the column index passed as argument to find the page that the
+// given value is expected to be found in.
+//
+// The function returns the index of the first page that might contain the
+// value. If the function determines that the value does not exist in the
+// index, NumPages is returned.
+//
+// The comparison function passed as last argument is used to determine the
+// relative order of values. This should generally be the Compare method of
+// the column type, but can sometimes be customized to modify how null values
+// are interpreted, for example:
+//
+//	pageIndex := parquet.Find(columnIndex, value,
+//		parquet.CompareNullsFirst(typ.Compare),
+//	)
+//
+func Find(index ColumnIndex, value Value, cmp func(Value, Value) int) int {
+	switch {
+	case index.IsAscending():
+		return binarySearch(index, value, cmp)
+	default:
+		return linearSearch(index, value, cmp)
+	}
+}
+
+func binarySearch(index ColumnIndex, value Value, cmp func(Value, Value) int) int {
+	n := index.NumPages()
+	i := 0
+	j := n
+
+	for (j - i) > 1 {
+		k := ((j - i) / 2) + i
+		c := cmp(value, index.MinValue(k))
+
+		switch {
+		case c < 0:
+			j = k
+		case c > 0:
+			i = k
+		default:
+			return k
+		}
+	}
+
+	if i < n {
+		min := index.MinValue(i)
+		max := index.MaxValue(i)
+
+		if cmp(value, min) < 0 || cmp(max, value) < 0 {
+			i = n
+		}
+	}
+
+	return i
+}
+
+func linearSearch(index ColumnIndex, value Value, cmp func(Value, Value) int) int {
+	n := index.NumPages()
+
+	for i := 0; i < n; i++ {
+		min := index.MinValue(i)
+		max := index.MaxValue(i)
+
+		if cmp(min, value) <= 0 && cmp(value, max) <= 0 {
+			return i
+		}
+	}
+
+	return n
+}


Test patch:
diff --git a/search_test.go b/search_test.go
new file mode 100644
index 0000000..76f40ac
--- /dev/null
+++ b/search_test.go
@@ -0,0 +1,98 @@
+package parquet_test
+
+import (
+	"fmt"
+	"math"
+	"testing"
+
+	"github.com/segmentio/parquet-go"
+)
+
+func assertCompare(t *testing.T, a, b parquet.Value, cmp func(parquet.Value, parquet.Value) int, want int) {
+	if got := cmp(a, b); got != want {
+		t.Errorf("compare(%v, %v): got=%d want=%d", a, b, got, want)
+	}
+}
+
+func TestCompareNullsFirst(t *testing.T) {
+	cmp := parquet.CompareNullsFirst(parquet.Int32Type.Compare)
+	assertCompare(t, parquet.Value{}, parquet.Value{}, cmp, 0)
+	assertCompare(t, parquet.Value{}, parquet.ValueOf(int32(0)), cmp, -1)
+	assertCompare(t, parquet.ValueOf(int32(0)), parquet.Value{}, cmp, +1)
+	assertCompare(t, parquet.ValueOf(int32(0)), parquet.ValueOf(int32(1)), cmp, -1)
+}
+
+func TestCompareNullsLast(t *testing.T) {
+	cmp := parquet.CompareNullsLast(parquet.Int32Type.Compare)
+	assertCompare(t, parquet.Value{}, parquet.Value{}, cmp, 0)
+	assertCompare(t, parquet.Value{}, parquet.ValueOf(int32(0)), cmp, +1)
+	assertCompare(t, parquet.ValueOf(int32(0)), parquet.Value{}, cmp, -1)
+	assertCompare(t, parquet.ValueOf(int32(0)), parquet.ValueOf(int32(1)), cmp, -1)
+}
+
+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},
+		{32},
+		{42, 43, 44, 45, 46, 47, 48, 49},
+	})
+}
+
+func TestSearchLinear(t *testing.T) {
+	testSearch(t, [][]int32{
+		{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},
+		{42, 43, 44, 45, 46, 47, 48, 49},
+	})
+}
+
+func testSearch(t *testing.T, pages [][]int32) {
+	indexer := parquet.Int32Type.NewColumnIndexer(0)
+
+	for _, values := range pages {
+		min := values[0]
+		max := values[0]
+
+		for _, v := range values[1:] {
+			switch {
+			case v < min:
+				min = v
+			case v > max:
+				max = v
+			}
+		}
+
+		indexer.IndexPage(int64(len(values)), 0,
+			parquet.ValueOf(min),
+			parquet.ValueOf(max),
+		)
+	}
+
+	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 _, 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))
+				}
+			}
+		})
+	}
+}