diff --git a/s2/encode_test.go b/s2/encode_test.go index fcf04191fa..a1ad5ac591 100644 --- a/s2/encode_test.go +++ b/s2/encode_test.go @@ -351,6 +351,43 @@ func TestIndex(t *testing.T) { } } +func BenchmarkIndexFind(b *testing.B) { + fatalErr := func(t testing.TB, err error) { + if err != nil { + t.Fatal(err) + } + } + for blocks := 1; blocks <= 65536; blocks *= 2 { + if blocks == 65536 { + blocks = 65535 + } + + var index Index + index.reset(100) + index.TotalUncompressed = int64(blocks) * 100 + index.TotalCompressed = int64(blocks) * 100 + for i := 0; i < blocks; i++ { + err := index.add(int64(i*100), int64(i*100)) + fatalErr(b, err) + } + + rng := rand.New(rand.NewSource(0xabeefcafe)) + b.Run(fmt.Sprintf("blocks-%d", len(index.info)), func(b *testing.B) { + b.ResetTimer() + b.ReportAllocs() + const prime4bytes = 2654435761 + rng2 := rng.Int63() + for i := 0; i < b.N; i++ { + rng2 = ((rng2 + prime4bytes) * prime4bytes) >> 32 + // Find offset: + _, _, err := index.Find(rng2 % (int64(blocks) * 100)) + fatalErr(b, err) + } + }) + } + +} + func TestWriterPadding(t *testing.T) { n := 100 if testing.Short() { diff --git a/s2/index.go b/s2/index.go index fd857682e4..7b24a0060b 100644 --- a/s2/index.go +++ b/s2/index.go @@ -10,6 +10,7 @@ import ( "encoding/json" "fmt" "io" + "sort" ) const ( @@ -100,6 +101,15 @@ func (i *Index) Find(offset int64) (compressedOff, uncompressedOff int64, err er if offset > i.TotalUncompressed { return 0, 0, io.ErrUnexpectedEOF } + if len(i.info) > 200 { + n := sort.Search(len(i.info), func(n int) bool { + return i.info[n].uncompressedOffset > offset + }) + if n == 0 { + n = 1 + } + return i.info[n-1].compressedOffset, i.info[n-1].uncompressedOffset, nil + } for _, info := range i.info { if info.uncompressedOffset > offset { break