zeroshade commented on code in PR #1404:
URL: https://github.com/apache/arrow-go/pull/1404#discussion_r4362291200
##########
parquet/internal/utils/bit_reader.go:
##########
@@ -619,24 +620,35 @@ func (b *BitReader) GetBatchLevels(bits uint, out
[]int16, maxLevel int16) (int,
if _, err := b.reader.Seek(b.byteoffset, io.SeekStart); err != nil {
return i, maxCount, err
}
- for i < length {
- unpackSize := utils.Min(buflen, length-i)
- numUnpacked, err := unpack32(b.reader,
b.unpackBuf[:unpackSize], int(bits))
-
- for k, val := range b.unpackBuf[:numUnpacked] {
- level := int16(val)
- out[i+k] = level
- if level == maxLevel {
- maxCount++
- }
- }
- i += numUnpacked
- b.byteoffset += int64(numUnpacked * int(bits) / 8)
+ if bits == 1 || bits == 2 {
Review Comment:
On amd64 the 2-bit case comes out slower than the existing unpack path.
Pinned, 5×300 ms runs of `BenchmarkPackedLevelDecode`: width=2 alternating
n=65536 30.1–30.5 µs before vs 32.5–33.9 µs after; mostly-defined n=65536
30.9–31.3 vs 32.8–33.2; n=1024 cases 4–5% slower. 1-bit is 1.2–1.4× faster.
Could this be `if bits == 1` unless the 2-bit loop can be made faster on amd64
too?
##########
parquet/internal/utils/bit_reader.go:
##########
@@ -658,6 +670,74 @@ func (b *BitReader) GetBatchLevels(bits uint, out []int16,
maxLevel int16) (int,
return i, maxCount, nil
}
+// getBatchLevelsNarrow consumes complete 32-value groups of packed 1-bit or
+// 2-bit levels. It must be called with a byte-aligned reader; the caller
decodes
+// any trailing values with next. Only complete groups are written on short
reads.
+func (b *BitReader) getBatchLevelsNarrow(width uint, out []int16, maxLevel
int16) (int, int64, error) {
+ const valuesPerGroup = 32
+ groupBytes := 4 * int(width)
+ toDecode := len(out) / valuesPerGroup * valuesPerGroup
+ packed := arrow.Uint32Traits.CastToBytes(b.unpackBuf[:])
+ decoded := 0
+ var maxCount int64
+
+ for decoded < toDecode {
+ batch := min(toDecode-decoded, buflen)
+ batchBytes := batch * int(width) / 8
+ nread, err := io.ReadFull(b.reader, packed[:batchBytes])
+ complete := nread / groupBytes * valuesPerGroup
+ completeBytes := complete * int(width) / 8
+ dst := out[decoded : decoded+complete]
+
+ if width == 1 {
+ for _, value := range packed[:completeBytes] {
+ dst[0] = int16(value & 1)
+ dst[1] = int16((value >> 1) & 1)
+ dst[2] = int16((value >> 2) & 1)
+ dst[3] = int16((value >> 3) & 1)
+ dst[4] = int16((value >> 4) & 1)
+ dst[5] = int16((value >> 5) & 1)
+ dst[6] = int16((value >> 6) & 1)
+ dst[7] = int16(value >> 7)
+ switch maxLevel {
+ case 0:
+ maxCount += int64(8 -
mathbits.OnesCount8(value))
+ case 1:
+ maxCount +=
int64(mathbits.OnesCount8(value))
+ }
+ dst = dst[8:]
+ }
+ } else {
+ // Each 2-bit value occupies one of four pairs in the
byte.
Review Comment:
This is the loop that regresses on amd64 (see the comment above). If you
keep it, please include amd64 numbers alongside the M1 ones.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
