github.com/wrgl/wrgl@v0.14.0/pkg/index/ordered_hash_set_test.go (about)

     1  // SPDX-License-Identifier: Apache-2.0
     2  // Copyright © 2022 Wrangle Ltd
     3  
     4  package index
     5  
     6  import (
     7  	"bytes"
     8  	"os"
     9  	"testing"
    10  
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  	"github.com/wrgl/wrgl/pkg/objects"
    14  	"github.com/wrgl/wrgl/pkg/testutils"
    15  )
    16  
    17  func makeBytesMatrix(n, m int) [][]byte {
    18  	rows := make([][]byte, n)
    19  	for i := 0; i < len(rows); i++ {
    20  		rows[i] = testutils.SecureRandomBytes(m)
    21  	}
    22  	return rows
    23  }
    24  
    25  func TestOrderedHashSet(t *testing.T) {
    26  	otherHashes := makeBytesMatrix(10, 16)
    27  	for c, rows := range [][][]byte{
    28  		{},
    29  		{testutils.SecureRandomBytes(32)},
    30  		{testutils.SecureRandomBytes(32), testutils.SecureRandomBytes(32)},
    31  		{
    32  			append([]byte{0}, testutils.SecureRandomBytes(31)...),
    33  			append([]byte{128}, testutils.SecureRandomBytes(31)...),
    34  			append([]byte{255}, testutils.SecureRandomBytes(31)...),
    35  		},
    36  		{
    37  			append([]byte{128}, testutils.SecureRandomBytes(31)...),
    38  			append([]byte{0}, testutils.SecureRandomBytes(31)...),
    39  			append([]byte{255}, testutils.SecureRandomBytes(31)...),
    40  			append([]byte{128}, testutils.SecureRandomBytes(31)...),
    41  			append([]byte{255}, testutils.SecureRandomBytes(31)...),
    42  			append([]byte{0}, testutils.SecureRandomBytes(31)...),
    43  		},
    44  		makeBytesMatrix(1024, 32),
    45  	} {
    46  		buf := bytes.NewBuffer(nil)
    47  		w := NewOrderedHashSetWriter(buf, rows)
    48  		require.NoError(t, w.Flush())
    49  
    50  		i, err := NewOrderedHashSet(objects.NopCloser(bytes.NewReader(buf.Bytes())))
    51  		require.NoError(t, err)
    52  		for j, row := range rows {
    53  			k, err := i.IndexOf(row[:16])
    54  			require.NoError(t, err, "case %d", c)
    55  			assert.Equal(t, j, k, "case %d", c)
    56  		}
    57  		for _, s := range otherHashes {
    58  			k, err := i.IndexOf(s)
    59  			require.NoError(t, err, "case %d", c)
    60  			assert.Equal(t, -1, k, "case %d", c)
    61  		}
    62  	}
    63  }
    64  
    65  func BenchmarkOrderedHashSet1M(b *testing.B) {
    66  	n := 1024 * 1024
    67  	rows := makeBytesMatrix(n, 32)
    68  	f, err := testutils.TempFile("", "test_hash_index")
    69  	require.NoError(b, err)
    70  	defer os.Remove(f.Name())
    71  	w := NewOrderedHashSetWriter(f, rows)
    72  	require.NoError(b, w.Flush())
    73  	require.NoError(b, f.Close())
    74  
    75  	f, err = os.Open(f.Name())
    76  	require.NoError(b, err)
    77  	defer f.Close()
    78  	i, err := NewOrderedHashSet(f)
    79  	require.NoError(b, err)
    80  	b.ResetTimer()
    81  	for j := 0; j < b.N; j++ {
    82  		k, err := i.IndexOf(rows[j%n])
    83  		require.NoError(b, err)
    84  		assert.Equal(b, j%n, k)
    85  	}
    86  }