Merge pull request #5 from parquet-go/optimize-pack-purego

optimize purego bit packing

Author: fpetkovski

pack.go

@@ -1,9 +1,5 @@
 package bitpack
 
-import (
-	"encoding/binary"
-)
-
 // PackInt32 packs values from src to dst, each value is packed into the given
 // bit width regardless of how many bits are needed to represent it.
 //
@@ -13,34 +9,6 @@ func PackInt32(dst []byte, src []int32, bitWidth uint) {
 	packInt32(dst, src, bitWidth)
 }
 
-func packInt32(dst []byte, src []int32, bitWidth uint) {
-	n := ByteCount(uint(len(src)) * bitWidth)
-	b := dst[:n]
-
-	for i := range b {
-		b[i] = 0
-	}
-
-	bitMask := uint32(1<<bitWidth) - 1
-	bitOffset := uint(0)
-
-	for _, value := range src {
-		i := bitOffset / 32
-		j := bitOffset % 32
-
-		lo := binary.LittleEndian.Uint32(dst[(i+0)*4:])
-		hi := binary.LittleEndian.Uint32(dst[(i+1)*4:])
-
-		lo |= (uint32(value) & bitMask) << j
-		hi |= (uint32(value) >> (32 - j))
-
-		binary.LittleEndian.PutUint32(dst[(i+0)*4:], lo)
-		binary.LittleEndian.PutUint32(dst[(i+1)*4:], hi)
-
-		bitOffset += bitWidth
-	}
-}
-
 // PackInt64 packs values from src to dst, each value is packed into the given
 // bit width regardless of how many bits are needed to represent it.
 //
@@ -50,34 +18,6 @@ func PackInt64(dst []byte, src []int64, bitWidth uint) {
 	packInt64(dst, src, bitWidth)
 }
 
-func packInt64(dst []byte, src []int64, bitWidth uint) {
-	n := ByteCount(uint(len(src)) * bitWidth)
-	b := dst[:n]
-
-	for i := range b {
-		b[i] = 0
-	}
-
-	bitMask := uint64(1<<bitWidth) - 1
-	bitOffset := uint(0)
-
-	for _, value := range src {
-		i := bitOffset / 64
-		j := bitOffset % 64
-
-		lo := binary.LittleEndian.Uint64(dst[(i+0)*8:])
-		hi := binary.LittleEndian.Uint64(dst[(i+1)*8:])
-
-		lo |= (uint64(value) & bitMask) << j
-		hi |= (uint64(value) >> (64 - j))
-
-		binary.LittleEndian.PutUint64(dst[(i+0)*8:], lo)
-		binary.LittleEndian.PutUint64(dst[(i+1)*8:], hi)
-
-		bitOffset += bitWidth
-	}
-}
-
 func assertPack(dst []byte, count int, bitWidth uint) {
 	_ = dst[:ByteCount(bitWidth*uint(count))]
 }

pack_arm64.go

@@ -0,0 +1,31 @@
+//go:build !purego
+
+package bitpack
+
+//go:noescape
+func packInt32ARM64(dst []byte, src []int32, bitWidth uint)
+
+//go:noescape
+func packInt32NEON(dst []byte, src []int32, bitWidth uint)
+
+//go:noescape
+func packInt64ARM64(dst []byte, src []int64, bitWidth uint)
+
+//go:noescape
+func packInt64NEON(dst []byte, src []int64, bitWidth uint)
+
+func packInt32(dst []byte, src []int32, bitWidth uint) {
+	if bitWidth <= 8 {
+		packInt32NEON(dst, src, bitWidth)
+	} else {
+		packInt32ARM64(dst, src, bitWidth)
+	}
+}
+
+func packInt64(dst []byte, src []int64, bitWidth uint) {
+	if bitWidth <= 8 {
+		packInt64NEON(dst, src, bitWidth)
+	} else {
+		packInt64ARM64(dst, src, bitWidth)
+	}
+}