Repository: segmentio/parquet-go
Pull request #218 (2022-06-08 07:14:22)

Optimize delta prefix search
This PR fixes a somewhat broken implementation of the binary search of prefix lengths in the DELTA_BYTE_ARRAY encoding.

```
name                                old time/op  new time/op  delta
BinarySearchPrefixLength/size=10    14.7ns ± 0%  11.3ns ± 2%  -23.21%  (p=0.000 n=9+10)
BinarySearchPrefixLength/size=100   41.1ns ± 2%  19.6ns ± 6%  -52.34%  (p=0.000 n=10+10)
BinarySearchPrefixLength/size=1000  88.3ns ± 1%  27.6ns ± 4%  -68.71%  (p=0.000 n=9+10)
```


Patch:
diff --git a/encoding/delta/byte_array.go b/encoding/delta/byte_array.go
index 98b9a3e..ed34c1a 100644
--- a/encoding/delta/byte_array.go
+++ b/encoding/delta/byte_array.go
@@ -3,6 +3,7 @@ package delta
 import (
 	"bytes"
 	"fmt"
+	"sort"
 
 	"github.com/segmentio/parquet-go/encoding"
 	"github.com/segmentio/parquet-go/encoding/plain"
@@ -10,7 +11,7 @@ import (
 )
 
 const (
-	maxLinearSearchPrefixLength = 128 // arbitrary
+	maxLinearSearchPrefixLength = 64 // arbitrary
 )
 
 type ByteArrayEncoding struct {
@@ -55,7 +56,7 @@ func (e *ByteArrayEncoding) EncodeByteArray(dst, src []byte) ([]byte, error) {
 		if len(v) <= maxLinearSearchPrefixLength {
 			p = linearSearchPrefixLength(lastValue, v)
 		} else {
-			p = binarySearchPrefixLength(len(lastValue)/2, lastValue, v)
+			p = binarySearchPrefixLength(lastValue, v)
 		}
 
 		prefix.values = append(prefix.values, int32(p))
@@ -229,17 +230,12 @@ func linearSearchPrefixLength(base, data []byte) (n int) {
 	return n
 }
 
-func binarySearchPrefixLength(max int, base, data []byte) int {
-	for len(base) > 0 {
-		if bytes.HasPrefix(data, base[:max]) {
-			if max == len(base) {
-				return max
-			}
-			max += (len(base)-max)/2 + 1
-		} else {
-			base = base[:max-1]
-			max /= 2
-		}
+func binarySearchPrefixLength(base, data []byte) int {
+	n := len(base)
+	if n > len(data) {
+		n = len(data)
 	}
-	return 0
+	return sort.Search(n, func(i int) bool {
+		return !bytes.Equal(base[:i+1], data[:i+1])
+	})
 }
diff --git a/page.go b/page.go
index 987344a..c78c3fe 100644
--- a/page.go
+++ b/page.go
@@ -1,6 +1,7 @@
 package parquet
 
 import (
+	"bytes"
 	"encoding/binary"
 	"fmt"
 	"io"
@@ -1012,7 +1013,7 @@ func (page *byteArrayPage) min() (min []byte) {
 		for i := 4 + len(min); i < len(page.values); {
 			v := page.valueAt(uint32(i))
 
-			if string(v) < string(min) {
+			if bytes.Compare(v, min) < 0 {
 				min = v
 			}
 
@@ -1030,7 +1031,7 @@ func (page *byteArrayPage) max() (max []byte) {
 		for i := 4 + len(max); i < len(page.values); {
 			v := page.valueAt(uint32(i))
 
-			if string(v) > string(max) {
+			if bytes.Compare(v, max) > 0 {
 				max = v
 			}
 
@@ -1050,9 +1051,9 @@ func (page *byteArrayPage) bounds() (min, max []byte) {
 			v := page.valueAt(uint32(i))
 
 			switch {
-			case string(v) < string(min):
+			case bytes.Compare(v, min) < 0:
 				min = v
-			case string(v) > string(max):
+			case bytes.Compare(v, max) > 0:
 				max = v
 			}
 


Test patch:
diff --git a/encoding/delta/byte_array_test.go b/encoding/delta/byte_array_test.go
index 571a3da..b994d0c 100644
--- a/encoding/delta/byte_array_test.go
+++ b/encoding/delta/byte_array_test.go
@@ -2,6 +2,7 @@ package delta
 
 import (
 	"bytes"
+	"fmt"
 	"testing"
 )
 
@@ -11,7 +12,7 @@ func TestLinearSearchPrefixLength(t *testing.T) {
 
 func TestBinarySearchPrefixLength(t *testing.T) {
 	testSearchPrefixLength(t, func(base, data []byte) int {
-		return binarySearchPrefixLength(len(base)/2, base, data)
+		return binarySearchPrefixLength(base, data)
 	})
 }
 
@@ -128,20 +129,27 @@ func testSearchPrefixLength(t *testing.T, prefixLength func(base, data []byte) i
 	}
 }
 
-func BenchmarkLinearSearchkPrefixLength(b *testing.B) {
+func BenchmarkLinearSearchPrefixLength(b *testing.B) {
 	benchmarkSearchPrefixLength(b, linearSearchPrefixLength)
 }
 
-func BenchmarkBinarySearchkPrefixLength(b *testing.B) {
+func BenchmarkBinarySearchPrefixLength(b *testing.B) {
 	benchmarkSearchPrefixLength(b, func(base, data []byte) int {
-		return binarySearchPrefixLength(len(base)/2, base, data)
+		return binarySearchPrefixLength(base, data)
 	})
 }
 
 func benchmarkSearchPrefixLength(b *testing.B, prefixLength func(base, data []byte) int) {
-	value := bytes.Repeat([]byte("0123456789"), 100)
+	buffer := bytes.Repeat([]byte("0123456789"), 100)
 
-	for i := 0; i < b.N; i++ {
-		_ = prefixLength(value, value)
+	for _, size := range []int{10, 100, 1000} {
+		b.Run(fmt.Sprintf("size=%d", size), func(b *testing.B) {
+			base := buffer[:size]
+			data := buffer[:size/2]
+
+			for i := 0; i < b.N; i++ {
+				_ = prefixLength(base, data)
+			}
+		})
 	}
 }