github.com/fraugster/parquet-go@v0.12.0/hybrid_test.go (about)

     1  package goparquet
     2  
     3  import (
     4  	"bytes"
     5  	"math"
     6  	"math/rand"
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/stretchr/testify/assert"
    12  )
    13  
    14  func buildData(bitWidth int, l int) []int32 {
    15  	if bitWidth > 32 || bitWidth < 0 {
    16  		panic("wrong bitwidth")
    17  	}
    18  
    19  	max := int32(math.Pow(2, float64(bitWidth)))
    20  	res := make([]int32, l)
    21  	for i := 0; i < l; i++ {
    22  		if bitWidth == 0 {
    23  			res[i] = 0
    24  		} else if bitWidth < 31 {
    25  			res[i] = rand.Int31n(max)
    26  		} else {
    27  			res[i] = rand.Int31()
    28  		}
    29  	}
    30  
    31  	return res
    32  }
    33  
    34  func TestHybrid(t *testing.T) {
    35  	for i := 0; i < 32; i++ {
    36  		data := &bytes.Buffer{}
    37  		enc := newHybridEncoder(i)
    38  		assert.NoError(t, enc.initSize(data))
    39  		to1 := buildData(i, 8*10240+5)
    40  		assert.NoError(t, enc.encode(to1))
    41  
    42  		to2 := buildData(i, 1000)
    43  		assert.NoError(t, enc.encode(to2))
    44  
    45  		assert.NoError(t, enc.Close())
    46  
    47  		buf2 := bytes.NewReader(data.Bytes())
    48  		dec := newHybridDecoder(i)
    49  		assert.NoError(t, dec.initSize(buf2))
    50  		var toR []int32
    51  		total := len(to1) + len(to2)
    52  		for j := 0; j < total; j++ {
    53  			d, err := dec.next()
    54  			if err != nil {
    55  				break
    56  			}
    57  			toR = append(toR, d)
    58  		}
    59  		assert.Equal(t, toR, append(to1, to2...))
    60  	}
    61  }
    62  
    63  func TestOnlyOne(t *testing.T) {
    64  	data := &packedArray{}
    65  	data.reset(1)
    66  	for i := int32(0); i < 1000; i++ {
    67  		data.appendSingle(i)
    68  	}
    69  	data.flush()
    70  
    71  	buf := &bytes.Buffer{}
    72  	require.NoError(t, encodeLevelsV1(buf, 1, data))
    73  	read := make([]int32, 1000)
    74  	dec := newHybridDecoder(1)
    75  	require.NoError(t, dec.initSize(bytes.NewReader(buf.Bytes())))
    76  	require.NoError(t, decodeInt32(dec, read))
    77  	require.Equal(t, data.toArray(), read)
    78  }