github.com/parquet-go/parquet-go@v0.21.1-0.20240501160520-b3c3a0c3ed6f/internal/bitpack/unpack_test.go (about)

     1  package bitpack_test
     2  
     3  import (
     4  	"fmt"
     5  	"math/rand"
     6  	"reflect"
     7  	"testing"
     8  
     9  	"github.com/parquet-go/parquet-go/internal/bitpack"
    10  )
    11  
    12  const (
    13  	blockSize = 128
    14  )
    15  
    16  func TestUnpackInt32(t *testing.T) {
    17  	for bitWidth := uint(1); bitWidth <= 32; bitWidth++ {
    18  		t.Run(fmt.Sprintf("bitWidth=%d", bitWidth), func(t *testing.T) {
    19  			block := [blockSize]int32{}
    20  			bitMask := int32(bitWidth<<1) - 1
    21  
    22  			prng := rand.New(rand.NewSource(0))
    23  			for i := range block {
    24  				block[i] = prng.Int31() & bitMask
    25  			}
    26  
    27  			size := (blockSize * bitWidth) / 8
    28  			buf := make([]byte, size+bitpack.PaddingInt32)
    29  			bitpack.PackInt32(buf, block[:], bitWidth)
    30  
    31  			src := buf[:size]
    32  			dst := make([]int32, blockSize)
    33  
    34  			for n := 1; n <= blockSize; n++ {
    35  				for i := range dst {
    36  					dst[i] = 0
    37  				}
    38  
    39  				bitpack.UnpackInt32(dst[:n], src, bitWidth)
    40  
    41  				if !reflect.DeepEqual(block[:n], dst[:n]) {
    42  					t.Fatalf("values mismatch for length=%d\nwant: %v\ngot:  %v", n, block[:n], dst[:n])
    43  				}
    44  			}
    45  		})
    46  	}
    47  }
    48  
    49  func TestUnpackInt64(t *testing.T) {
    50  	for bitWidth := uint(1); bitWidth <= 63; bitWidth++ {
    51  		t.Run(fmt.Sprintf("bitWidth=%d", bitWidth), func(t *testing.T) {
    52  			block := [blockSize]int64{}
    53  			bitMask := int64(bitWidth<<1) - 1
    54  
    55  			prng := rand.New(rand.NewSource(0))
    56  			for i := range block {
    57  				block[i] = prng.Int63() & bitMask
    58  			}
    59  
    60  			size := (blockSize * bitWidth) / 8
    61  			buf := make([]byte, size+bitpack.PaddingInt64)
    62  			bitpack.PackInt64(buf, block[:], bitWidth)
    63  
    64  			src := buf[:size]
    65  			dst := make([]int64, blockSize)
    66  
    67  			for n := 1; n <= blockSize; n++ {
    68  				for i := range dst {
    69  					dst[i] = 0
    70  				}
    71  
    72  				bitpack.UnpackInt64(dst[:n], src, bitWidth)
    73  
    74  				if !reflect.DeepEqual(block[:n], dst[:n]) {
    75  					t.Fatalf("values mismatch for length=%d\nwant: %v\ngot:  %v", n, block[:n], dst[:n])
    76  				}
    77  			}
    78  		})
    79  	}
    80  }
    81  
    82  func BenchmarkUnpackInt32(b *testing.B) {
    83  	for bitWidth := uint(1); bitWidth <= 32; bitWidth++ {
    84  		block := [blockSize]int32{}
    85  		buf := [4*blockSize + bitpack.PaddingInt32]byte{}
    86  		bitpack.PackInt32(buf[:], block[:], bitWidth)
    87  
    88  		b.Run(fmt.Sprintf("bitWidth=%d", bitWidth), func(b *testing.B) {
    89  			dst := block[:]
    90  			src := buf[:]
    91  
    92  			for i := 0; i < b.N; i++ {
    93  				bitpack.UnpackInt32(dst, src, bitWidth)
    94  			}
    95  
    96  			b.SetBytes(4 * blockSize)
    97  		})
    98  	}
    99  }
   100  
   101  func BenchmarkUnpackInt64(b *testing.B) {
   102  	for bitWidth := uint(1); bitWidth <= 64; bitWidth++ {
   103  		block := [blockSize]int64{}
   104  		buf := [8*blockSize + bitpack.PaddingInt64]byte{}
   105  		bitpack.PackInt64(buf[:], block[:], bitWidth)
   106  
   107  		b.Run(fmt.Sprintf("bitWidth=%d", bitWidth), func(b *testing.B) {
   108  			dst := block[:]
   109  			src := buf[:]
   110  
   111  			for i := 0; i < b.N; i++ {
   112  				bitpack.UnpackInt64(dst, src, bitWidth)
   113  			}
   114  
   115  			b.SetBytes(4 * blockSize)
   116  		})
   117  	}
   118  }