emkornfield commented on a change in pull request #9671: URL: https://github.com/apache/arrow/pull/9671#discussion_r598848234
########## File path: go/parquet/internal/utils/bit_block_counter.go ########## @@ -0,0 +1,260 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package utils + +import ( + "math" + "math/bits" + "unsafe" + + "github.com/apache/arrow/go/arrow/bitutil" + "github.com/apache/arrow/go/arrow/endian" +) + +var toLEFunc func(uint64) uint64 + +func init() { + if endian.IsBigEndian { + toLEFunc = bits.ReverseBytes64 + } else { + toLEFunc = func(in uint64) uint64 { return in } + } +} + +func loadWord(byt []byte) uint64 { + return toLEFunc(*(*uint64)(unsafe.Pointer(&byt[0]))) +} + +func shiftWord(current, next uint64, shift int64) uint64 { + if shift == 0 { + return current + } + return (current >> shift) | (next << (64 - shift)) +} + +// BitBlockCount is returned by the various bit block counter utilities +// in order to return a length of bits and the population count of that +// slice of bits. +type BitBlockCount struct { + Len int16 + Popcnt int16 +} + +// NoneSet returns true if ALL the bits were 0 in this set, ie: Popcnt == 0 +func (b BitBlockCount) NoneSet() bool { + return b.Popcnt == 0 +} + +// AllSet returns true if ALL the bits were 1 in this set, ie: Popcnt == Len +func (b BitBlockCount) AllSet() bool { + return b.Len == b.Popcnt +} + +// BitBlockCounter is a utility for grabbing chunks of a bitmap at a time and efficiently +// counting the number of bits which are 1. +type BitBlockCounter struct { + bitmap []byte + bitsRemaining int64 + bitOffset int8 +} + +const ( + WordBits int64 = 64 + FourWordsBits int64 = WordBits * 4 +) + +// NewBitBlockCounter returns a BitBlockCounter for the passed bitmap starting at startOffset +// of length nbits. +func NewBitBlockCounter(bitmap []byte, startOffset, nbits int64) *BitBlockCounter { + return &BitBlockCounter{ + bitmap: bitmap[startOffset/8:], + bitsRemaining: nbits, + bitOffset: int8(startOffset % 8), + } +} + +// getBlockSlow is for returning a block of the requested size when there aren't +// enough bits remaining to do a full word computation. +func (b *BitBlockCounter) getBlockSlow(blockSize int64) BitBlockCount { + runlen := int16(Min(b.bitsRemaining, blockSize)) + popcnt := int16(bitutil.CountSetBits(b.bitmap, int(b.bitOffset), int(runlen))) + b.bitsRemaining -= int64(runlen) + b.bitmap = b.bitmap[runlen/8:] + return BitBlockCount{runlen, popcnt} +} + +// Return the next run of available bits, usually 256. The returned +// pair contains the size of run and the number of true values. The last +// block will have a length less than 256 if the bitmap length is not a +// multiple of 256, and will return 0-length blocks in subsequent +// invocations. +func (b *BitBlockCounter) NextFourWords() BitBlockCount { + if b.bitsRemaining == 0 { + return BitBlockCount{0, 0} + } + + totalPopcnt := 0 + if b.bitOffset == 0 { + // if we're aligned at 0 bitoffset, then we can easily just jump from + // word to word nice and easy. + if b.bitsRemaining < FourWordsBits { + return b.getBlockSlow(FourWordsBits) + } + totalPopcnt += bits.OnesCount64(loadWord(b.bitmap)) + totalPopcnt += bits.OnesCount64(loadWord(b.bitmap[8:])) + totalPopcnt += bits.OnesCount64(loadWord(b.bitmap[16:])) + totalPopcnt += bits.OnesCount64(loadWord(b.bitmap[24:])) + } else { + // When the offset is > 0, we need there to be a word beyond the last + // aligned word in the bitmap for the bit shifting logic. + if b.bitsRemaining < 5*FourWordsBits-int64(b.bitOffset) { + return b.getBlockSlow(FourWordsBits) + } + + current := loadWord(b.bitmap) + next := loadWord(b.bitmap[8:]) + totalPopcnt += bits.OnesCount64(shiftWord(current, next, int64(b.bitOffset))) + + current = next + next = loadWord(b.bitmap[16:]) + totalPopcnt += bits.OnesCount64(shiftWord(current, next, int64(b.bitOffset))) + + current = next + next = loadWord(b.bitmap[24:]) + totalPopcnt += bits.OnesCount64(shiftWord(current, next, int64(b.bitOffset))) + + current = next + next = loadWord(b.bitmap[32:]) + totalPopcnt += bits.OnesCount64(shiftWord(current, next, int64(b.bitOffset))) + } + b.bitmap = b.bitmap[bitutil.BytesForBits(FourWordsBits):] + b.bitsRemaining -= FourWordsBits + return BitBlockCount{256, int16(totalPopcnt)} +} + +// Return the next run of available bits, usually 64. The returned +// pair contains the size of run and the number of true values. The last +// block will have a length less than 64 if the bitmap length is not a +// multiple of 64, and will return 0-length blocks in subsequent +// invocations. +func (b *BitBlockCounter) NextWord() BitBlockCount { + if b.bitsRemaining == 0 { + return BitBlockCount{0, 0} + } + popcnt := 0 + if b.bitOffset == 0 { + if b.bitsRemaining < WordBits { + return b.getBlockSlow(WordBits) + } + popcnt = bits.OnesCount64(loadWord(b.bitmap)) + } else { + // When the offset is > 0, we need there to be a word beyond the last + // aligned word in the bitmap for the bit shifting logic. + if b.bitsRemaining < 2*WordBits-int64(b.bitOffset) { Review comment: I assume no spacing with the multiple is inline with Go style guide? (it reads strangely to me) -- 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. For queries about this service, please contact Infrastructure at: us...@infra.apache.org
