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

     1  package bitpack
     2  
     3  import (
     4  	"encoding/binary"
     5  )
     6  
     7  // PackInt32 packs values from src to dst, each value is packed into the given
     8  // bit width regardless of how many bits are needed to represent it.
     9  //
    10  // The function panics if dst is too short to hold the bit packed values.
    11  func PackInt32(dst []byte, src []int32, bitWidth uint) {
    12  	assertPack(dst, len(src), bitWidth)
    13  	packInt32(dst, src, bitWidth)
    14  }
    15  
    16  func packInt32(dst []byte, src []int32, bitWidth uint) {
    17  	n := ByteCount(uint(len(src)) * bitWidth)
    18  	b := dst[:n]
    19  
    20  	for i := range b {
    21  		b[i] = 0
    22  	}
    23  
    24  	bitMask := uint32(1<<bitWidth) - 1
    25  	bitOffset := uint(0)
    26  
    27  	for _, value := range src {
    28  		i := bitOffset / 32
    29  		j := bitOffset % 32
    30  
    31  		lo := binary.LittleEndian.Uint32(dst[(i+0)*4:])
    32  		hi := binary.LittleEndian.Uint32(dst[(i+1)*4:])
    33  
    34  		lo |= (uint32(value) & bitMask) << j
    35  		hi |= (uint32(value) >> (32 - j))
    36  
    37  		binary.LittleEndian.PutUint32(dst[(i+0)*4:], lo)
    38  		binary.LittleEndian.PutUint32(dst[(i+1)*4:], hi)
    39  
    40  		bitOffset += bitWidth
    41  	}
    42  }
    43  
    44  // PackInt64 packs values from src to dst, each value is packed into the given
    45  // bit width regardless of how many bits are needed to represent it.
    46  //
    47  // The function panics if dst is too short to hold the bit packed values.
    48  func PackInt64(dst []byte, src []int64, bitWidth uint) {
    49  	assertPack(dst, len(src), bitWidth)
    50  	packInt64(dst, src, bitWidth)
    51  }
    52  
    53  func packInt64(dst []byte, src []int64, bitWidth uint) {
    54  	n := ByteCount(uint(len(src)) * bitWidth)
    55  	b := dst[:n]
    56  
    57  	for i := range b {
    58  		b[i] = 0
    59  	}
    60  
    61  	bitMask := uint64(1<<bitWidth) - 1
    62  	bitOffset := uint(0)
    63  
    64  	for _, value := range src {
    65  		i := bitOffset / 64
    66  		j := bitOffset % 64
    67  
    68  		lo := binary.LittleEndian.Uint64(dst[(i+0)*8:])
    69  		hi := binary.LittleEndian.Uint64(dst[(i+1)*8:])
    70  
    71  		lo |= (uint64(value) & bitMask) << j
    72  		hi |= (uint64(value) >> (64 - j))
    73  
    74  		binary.LittleEndian.PutUint64(dst[(i+0)*8:], lo)
    75  		binary.LittleEndian.PutUint64(dst[(i+1)*8:], hi)
    76  
    77  		bitOffset += bitWidth
    78  	}
    79  }
    80  
    81  func assertPack(dst []byte, count int, bitWidth uint) {
    82  	_ = dst[:ByteCount(bitWidth*uint(count))]
    83  }