zeroshade commented on code in PR #1289: URL: https://github.com/apache/arrow-go/pull/1289#discussion_r4528457733
########## parquet/internal/encoding/alp_encoding.go: ########## @@ -0,0 +1,667 @@ +// 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 encoding + +// ALP (Adaptive Lossless Floating-Point) encoding for FLOAT and DOUBLE columns. +// +// ALP multiplies each value by a power of ten, rounds the product to an integer, +// and stores bit-packed differences from the smallest integer. A value that does +// not survive that round trip is stored verbatim as an exception. alp_scaling.go +// holds that arithmetic and alp_params.go chooses the scaling; this file writes +// and reads the bytes. +// +// A page begins with a 7-byte header: compression mode, integer encoding, log2 +// of the vector size, then the value count as a little-endian int32. An array +// of little-endian uint32 offsets follows, one per vector, each counted from the +// start of that array. Every vector then holds: +// +// exponent (1 byte), factor (1 byte), exception count (uint16) +// frame of reference (int32 for FLOAT, int64 for DOUBLE), bit width (1 byte) +// the bit-packed differences from the frame of reference +// the exception positions (uint16 each), then the exception values +// +// Reference: https://dl.acm.org/doi/10.1145/3626717 + +import ( + "bytes" + "encoding/binary" + "errors" + "fmt" + "math" + "math/bits" + "slices" + + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet" + "github.com/apache/arrow-go/v18/parquet/internal/utils" + "github.com/apache/arrow-go/v18/parquet/schema" +) + +const ( + // The first two header bytes select the ALP variants. The decoder rejects + // variants other than the ones this encoder writes. + alpCompressionMode = 0 // ALP; other values are reserved for variants such as ALP-RD + alpIntegerEncodingFOR = 0 // frame of reference plus bit packing + + alpHeaderSize = 7 + + // Encoders write 1,024-value vectors unless SetVectorSize says otherwise. + // The format allows anything from 2^alpMinLogVectorSize to + // 2^alpMaxLogVectorSize, and the decoder honors what a page declares. + alpDefaultVectorSize = 1024 + alpDefaultLogVectorSize = 10 + alpMinLogVectorSize = 3 + alpMaxLogVectorSize = 15 +) + +func alpVectorHeaderSize[T alpFloat]() int { + return 1 + 1 + 2 + alpEncodedBytes[T]() + 1 +} + +func alpPackedSize(count, bitWidth int) int { + return (count*bitWidth + 7) / 8 +} + +// alpAppendPackedBits appends values to dst, packed at bitWidth bits each, +// least significant bit first. +func alpAppendPackedBits(dst []byte, values []uint64, bitWidth int) []byte { + if bitWidth == 0 || len(values) == 0 { + return dst + } + + // BitWriter flushes eight bytes at a time, so it needs room past the last + // packed byte. The return drops those bytes again. + start, size := len(dst), alpPackedSize(len(values), bitWidth) + dst = slices.Grow(dst, size+8)[:start+size+8] + clear(dst[start:]) + + bw := utils.NewBitWriter(utils.NewWriterAtBuffer(dst[start:])) + for _, v := range values { + bw.WriteValue(v, uint(bitWidth)) + } + bw.Flush(true) + return dst[:start+size] +} + +// alpUnpackBits reverses alpAppendPackedBits into out, which must hold as many +// values as were packed. The reader is reset onto packed, so one reader serves +// every vector of a page. +func alpUnpackBits(br *utils.BitReader, packed *bytes.Reader, data []byte, out []uint64, bitWidth int) error { + if bitWidth == 0 { + clear(out) + return nil + } + + packed.Reset(data) + br.Reset(packed) + n, err := br.GetBatch(uint(bitWidth), out) + if err != nil { + return err + } + if n != len(out) { + return fmt.Errorf("parquet: ALP unpacked %d of %d values", n, len(out)) + } + return nil +} + +func alpAppendInt[T alpFloat](dst []byte, v uint64) []byte { + if alpEncodedBytes[T]() == 4 { + return binary.LittleEndian.AppendUint32(dst, uint32(v)) + } + return binary.LittleEndian.AppendUint64(dst, v) +} + +func alpReadInt[T alpFloat](src []byte) int64 { + if alpEncodedBytes[T]() == 4 { + return int64(int32(binary.LittleEndian.Uint32(src))) + } + return int64(binary.LittleEndian.Uint64(src)) +} + +func alpReadFloat[T alpFloat](src []byte) T { + if alpEncodedBytes[T]() == 4 { + return alpFloatFromBits[T](uint64(binary.LittleEndian.Uint32(src))) + } + return alpFloatFromBits[T](binary.LittleEndian.Uint64(src)) +} + +// alpEncoder encodes a float column with ALP, a vector at a time. It buffers +// values until it holds a whole vector, so the values of a page reach the sink +// only when the caller flushes it. +type alpEncoder[T alpFloat] struct { + encoder + + vectorSize int + logVectorSize int + + vectorBuf []T + bufCount int + totalCount int + + encodedBuf []byte + vectorSizes []int + + // Scratch reused across vectors, calls, and pages. + encodedScratch []int64 + deltaScratch []uint64 + excPositions []uint16 + excValues []T + sampleScratch []T + spacedScratch []T + pageHeader []byte + + vectorsProcessed int + cachedPresets []alpEncodingParams + presetCounts map[alpEncodingParams]int +} + +func newALPEncoder[T alpFloat](e format.Encoding, descr *schema.Column, mem memory.Allocator) *alpEncoder[T] { + enc := &alpEncoder[T]{ + encoder: newEncoderBase(e, descr, mem), + sampleScratch: make([]T, 0, alpSamplerSamples), + presetCounts: make(map[alpEncodingParams]int), + } + enc.setVectorSize(alpDefaultLogVectorSize) + return enc +} + +func (enc *alpEncoder[T]) setVectorSize(logSize int) { + enc.logVectorSize = logSize + enc.vectorSize = 1 << logSize + enc.vectorBuf = make([]T, enc.vectorSize) + enc.encodedScratch = make([]int64, enc.vectorSize) + enc.deltaScratch = make([]uint64, enc.vectorSize) +} + +// SetVectorSize sets how many values each vector holds: a power of two from +// 8 to 32768, 1024 by default. It must be called before the first value is +// put, and Reset keeps it. +func (enc *alpEncoder[T]) SetVectorSize(size int) error { + if size < 1<<alpMinLogVectorSize || size > 1<<alpMaxLogVectorSize || size&(size-1) != 0 { + return fmt.Errorf("parquet: ALP vector size must be a power of two from %d to %d, got %d", + 1<<alpMinLogVectorSize, 1<<alpMaxLogVectorSize, size) + } + if enc.bufCount > 0 || len(enc.vectorSizes) > 0 { + return errors.New("parquet: ALP vector size cannot change once values are buffered") + } + if size != enc.vectorSize { + enc.setVectorSize(bits.TrailingZeros(uint(size))) + } + return nil +} + +func (enc *alpEncoder[T]) Type() parquet.Type { + return alpParquetType[T]() +} + +// Put buffers values and encodes complete vectors. FlushValues handles the +// final partial vector. +func (enc *alpEncoder[T]) Put(in []T) { + for len(in) > 0 { + n := copy(enc.vectorBuf[enc.bufCount:], in) + enc.bufCount += n + enc.totalCount += n + in = in[n:] + + if enc.bufCount == enc.vectorSize { + enc.encodeVector(enc.vectorBuf[:enc.bufCount]) + enc.bufCount = 0 + } + } +} + +// PutSpaced drops the null slots of in and buffers the values that remain. +func (enc *alpEncoder[T]) PutSpaced(in []T, validBits []byte, validBitsOffset int64) { + enc.spacedScratch = slices.Grow(enc.spacedScratch[:0], len(in))[:len(in)] + nvalid := spacedCompress(in, enc.spacedScratch, validBits, validBitsOffset) + enc.Put(enc.spacedScratch[:nvalid]) +} + +// The first alpSamplerVectors vectors seed the preset list. +func (enc *alpEncoder[T]) findParams(values []T) alpEncodingParams { + sample := alpSample(values, enc.sampleScratch[:0]) + + if len(enc.cachedPresets) > 0 { + return alpFindBestParamsWithPresets(sample, enc.cachedPresets) + } + + params := alpFindBestParams(sample) + enc.presetCounts[params]++ + return params +} + +func (enc *alpEncoder[T]) encodeVector(values []T) { + params := enc.findParams(values) + + enc.vectorsProcessed++ + if enc.cachedPresets == nil && enc.vectorsProcessed >= alpSamplerVectors { + enc.buildPresetCache() + } + + encoded := enc.encodedScratch[:len(values)] + enc.excPositions = enc.excPositions[:0] + enc.excValues = enc.excValues[:0] + + // Encoding a value also determines whether it is an exception, so collect + // both results in one pass. + s := newALPScaling[T](params.exponent, params.factor) + minValue := int64(math.MaxInt64) + for i, v := range values { + e, ok := s.tryEncode(v) + if !ok { + enc.excPositions = append(enc.excPositions, uint16(i)) + enc.excValues = append(enc.excValues, v) + continue + } + encoded[i] = e + minValue = min(e, minValue) + } + + if len(enc.excValues) == len(values) { + // Nothing encoded, so nothing constrains the frame of reference. + minValue = 0 + } + // An exception still takes up a slot among the packed integers. Filling the + // slot with the frame of reference itself costs no bits, so the exceptions + // cannot widen the packing. + for _, pos := range enc.excPositions { + encoded[pos] = minValue + } + + // The difference between two int64 values can exceed int64. Compute it in + // unsigned arithmetic; the decoder adds the wrapped difference back. + deltas := enc.deltaScratch[:len(values)] + var maxDelta uint64 + for i, e := range encoded { + deltas[i] = uint64(e) - uint64(minValue) + maxDelta = max(deltas[i], maxDelta) + } + bitWidth := bits.Len64(maxDelta) + + startLen := len(enc.encodedBuf) + enc.encodedBuf = append(enc.encodedBuf, byte(params.exponent), byte(params.factor)) + enc.encodedBuf = binary.LittleEndian.AppendUint16(enc.encodedBuf, uint16(len(enc.excPositions))) + enc.encodedBuf = alpAppendInt[T](enc.encodedBuf, uint64(minValue)) + enc.encodedBuf = append(enc.encodedBuf, byte(bitWidth)) + enc.encodedBuf = alpAppendPackedBits(enc.encodedBuf, deltas, bitWidth) + + for _, pos := range enc.excPositions { + enc.encodedBuf = binary.LittleEndian.AppendUint16(enc.encodedBuf, pos) + } + for _, v := range enc.excValues { + enc.encodedBuf = alpAppendInt[T](enc.encodedBuf, alpFloatToBits(v)) + } + + enc.vectorSizes = append(enc.vectorSizes, len(enc.encodedBuf)-startLen) +} + +// buildPresetCache keeps the scaling pairs the sampled vectors chose most often, +// ordered by how often they won. +func (enc *alpEncoder[T]) buildPresetCache() { + type preset struct { + params alpEncodingParams + count int + } + + sorted := make([]preset, 0, len(enc.presetCounts)) + for k, v := range enc.presetCounts { + sorted = append(sorted, preset{k, v}) + } + // The counts come out of a map, so the exponent and the factor break ties: + // without them the shortlist, and the encoded bytes, would differ between + // two runs over the same input. + slices.SortStableFunc(sorted, func(a, b preset) int { + switch { + case a.count != b.count: + return b.count - a.count + case a.params.exponent != b.params.exponent: + return b.params.exponent - a.params.exponent + default: + return b.params.factor - a.params.factor + } + }) + + numPresets := min(len(sorted), alpMaxPresetCombinations) + enc.cachedPresets = make([]alpEncodingParams, numPresets) + for i, p := range sorted[:numPresets] { + enc.cachedPresets[i] = p.params + } +} + +// EstimatedDataEncodedSize counts the encoded vectors and their offsets, plus +// an upper bound for the vector still being buffered: its offset and header, +// every slot packed at the full integer width, and every value also stored as +// an exception with its 2-byte position. A vector that mixes exceptions with +// wide packed values can approach that bound, so nothing smaller is safe; the +// overestimate is limited to one vector. +func (enc *alpEncoder[T]) EstimatedDataEncodedSize() int64 { + size := len(enc.encodedBuf) + 4*len(enc.vectorSizes) + if enc.bufCount > 0 { + width := alpEncodedBytes[T]() + size += 4 + alpVectorHeaderSize[T]() + enc.bufCount*(2*width+2) + } + return int64(size) +} + +// FlushValues writes one page and clears the values it wrote, so that the next +// Put starts a new page. The scaling shortlist survives, because it describes +// the column rather than the page. +// +// A page always carries the header, even when it holds no values: a column of +// optional values can produce a page whose rows are all null, and a reader +// rejects an empty page rather than reading it as no values. +func (enc *alpEncoder[T]) FlushValues() (Buffer, error) { + if enc.bufCount > 0 { + enc.encodeVector(enc.vectorBuf[:enc.bufCount]) + enc.bufCount = 0 + } + defer enc.resetPage() + + offsetArraySize := len(enc.vectorSizes) * 4 + enc.sink.Reserve(alpHeaderSize + offsetArraySize + len(enc.encodedBuf)) + + header := append(enc.pageHeader[:0], alpCompressionMode, alpIntegerEncodingFOR, byte(enc.logVectorSize)) + header = binary.LittleEndian.AppendUint32(header, uint32(enc.totalCount)) + + offset := uint32(offsetArraySize) + for _, size := range enc.vectorSizes { + header = binary.LittleEndian.AppendUint32(header, offset) + offset += uint32(size) + } + enc.pageHeader = header + + enc.sink.Write(header) + enc.sink.Write(enc.encodedBuf) + return enc.sink.Finish(), nil +} + +// resetPage drops the encoded vectors of the page just written. +func (enc *alpEncoder[T]) resetPage() { + enc.totalCount = 0 + enc.encodedBuf = enc.encodedBuf[:0] + enc.vectorSizes = enc.vectorSizes[:0] +} + +// Reset returns the encoder to its state when new, shortlist included, so that +// it can encode an unrelated column. The vector size stays as set. +func (enc *alpEncoder[T]) Reset() { + enc.encoder.Reset() + enc.resetPage() + enc.bufCount = 0 + enc.vectorsProcessed = 0 + enc.cachedPresets = nil + clear(enc.presetCounts) +} + +// alpDecoder decodes one ALP vector at a time and keeps it across partial reads. +type alpDecoder[T alpFloat] struct { + decoder + + vectorSize int + totalCount int + numVectors int + + vectorOffsets []uint32 + bodyData []byte + + // The vector the decoder holds, and how far through the page it has read. + currentVectorIndex int + currentIndex int + decodedValues []T + deltas []uint64 + + // One reader pair, reset onto each vector's packed deltas in turn. + packedReader *bytes.Reader + bitReader *utils.BitReader +} + +func newALPDecoder[T alpFloat](e format.Encoding, descr *schema.Column) *alpDecoder[T] { + packed := bytes.NewReader(nil) + return &alpDecoder[T]{ + decoder: newDecoderBase(e, descr), + packedReader: packed, + bitReader: utils.NewBitReader(packed), + } +} + +func (dec *alpDecoder[T]) Type() parquet.Type { + return alpParquetType[T]() +} + +// SetData points the decoder at one page and reads its header. It rejects +// unsupported ALP variants and truncated vector-offset tables. +func (dec *alpDecoder[T]) SetData(nvals int, data []byte) error { + if len(data) < alpHeaderSize { + return fmt.Errorf("parquet: ALP data too short for header: %d bytes", len(data)) + } + + compressionMode := data[0] + integerEncoding := data[1] + logVectorSize := data[2] + numElements := int32(binary.LittleEndian.Uint32(data[3:7])) + + if compressionMode != alpCompressionMode { + return fmt.Errorf("parquet: unsupported ALP compression mode: %d", compressionMode) + } + if integerEncoding != alpIntegerEncodingFOR { + return fmt.Errorf("parquet: unsupported ALP integer encoding: %d", integerEncoding) + } + if logVectorSize < alpMinLogVectorSize || logVectorSize > alpMaxLogVectorSize { + return fmt.Errorf("parquet: invalid ALP log vector size: %d, must be between %d and %d", + logVectorSize, alpMinLogVectorSize, alpMaxLogVectorSize) + } + if numElements < 0 { + return fmt.Errorf("parquet: invalid ALP element count: %d", numElements) + } + + dec.vectorSize = 1 << logVectorSize + dec.totalCount = int(numElements) + // Rounding up by adding vectorSize-1 first could overflow a 32-bit int. + dec.numVectors = dec.totalCount/dec.vectorSize + min(dec.totalCount%dec.vectorSize, 1) + dec.currentIndex = 0 + dec.currentVectorIndex = -1 + dec.nvals = nvals Review Comment: Partly applied in 85c4a970. `SetData` now rejects a header that claims more values than the page holds. Requiring the counts to be equal would be wrong, because `nvals` is the page's value count including nulls, so a nullable page legitimately carries fewer ALP elements. `TestALPDecoderElementCountExceedsPage` covers both directions. I left the offset-table validation lazy. Each vector's offset and size are already bounds-checked in `decodeVector` before anything is read, so a corrupt page fails with an error either way, and checking every offset up front would cost an extra pass for each page. ########## parquet/internal/encoding/alp_params.go: ########## @@ -0,0 +1,155 @@ +// 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 encoding + +// The encoder samples each vector and chooses the exponent and factor that +// minimize its encoded size. After the initial sampling window, it searches +// only the combinations that won. + +import ( + "math" + "math/bits" +) + +const ( + // alpSamplerVectors is how many vectors get the full exponent and factor + // search before the encoder falls back to the combinations those vectors + // chose. alpMaxPresetCombinations caps how many it keeps. + alpSamplerVectors = 8 + alpMaxPresetCombinations = 5 + + // alpSamplerSamples is how many values of a vector the search looks at, and + // alpPresetGiveUpAfter is how many shortlisted scalings may fail to improve + // on the best before the per-vector search stops looking. + alpSamplerSamples = 256 + alpPresetGiveUpAfter = 4 +) + +// alpEncodingParams is the scaling one vector is encoded with: ALP multiplies +// each value by 10^exponent, then by 10^-factor. +type alpEncodingParams struct { + exponent int + factor int +} + +// alpSample appends up to alpSamplerSamples evenly spaced values to out. +func alpSample[T alpFloat](values []T, out []T) []T { + stride := max(1, (len(values)+alpSamplerSamples-1)/alpSamplerSamples) + for i := 0; i < len(values); i += stride { + out = append(out, values[i]) + } + return out +} + +// alpExceptionBits is what one exception costs: the value at its full width, +// plus the position that patches it back in. +func alpExceptionBits[T alpFloat]() int { + return 8*alpEncodedBytes[T]() + 16 +} + +// alpEstimateSizeBits estimates the encoded size for a parameter pair and +// reports how many values it can encode. The cost includes the packed integers +// at the width left by the frame of reference and every exception. +func alpEstimateSizeBits[T alpFloat](values []T, exponent, factor int) (sizeBits int64, numEncodable int) { + minEncoded, maxEncoded := int64(math.MaxInt64), int64(math.MinInt64) + numExceptions := 0 + s := newALPScaling[T](exponent, factor) + for _, v := range values { + e, ok := s.tryEncode(v) + if !ok { + numExceptions++ + continue + } + minEncoded, maxEncoded = min(minEncoded, e), max(maxEncoded, e) + } + + sizeBits = int64(numExceptions) * int64(alpExceptionBits[T]()) + numEncodable = len(values) - numExceptions + if numEncodable == 0 { + return sizeBits, 0 + } + // The encoded range can exceed int64, so compute its width in unsigned + // arithmetic. + width := bits.Len64(uint64(maxEncoded) - uint64(minEncoded)) + return sizeBits + int64(len(values))*int64(width), numEncodable +} + +// For equal-size estimates, the ALP paper prefers the larger exponent, +// then the larger factor (Afroozeh et al., SIGMOD 2023, section 3.1.2). +func alpBetterParams(a alpEncodingParams, aBits int64, b alpEncodingParams, bBits int64) bool { + switch { + case aBits != bBits: + return aBits < bBits + case a.exponent != b.exponent: + return a.exponent > b.exponent + default: + return a.factor > b.factor + } +} + +// alpFindBestParams returns the cheapest scaling over every exponent and factor +// pair. The factor never exceeds the exponent, because scaling by a net power +// below one would move the decimal point the wrong way. +func alpFindBestParams[T alpFloat](values []T) alpEncodingParams { + maxExponent := alpNumExponents[T]() - 1 + // The fallback treats the values as integers: an exponent and a factor of + // equal size cancel, so whatever is already integral packs and the rest + // become exceptions. It stands when no scaling encodes enough to compare. + best := alpEncodingParams{exponent: maxExponent, factor: maxExponent} + bestBits := int64(math.MaxInt64) + + for e := range maxExponent + 1 { + for f := range e + 1 { Review Comment: No change needed: `go.mod` declares `go 1.26.0`, so range-over-int is available everywhere this module builds. ########## parquet/writer_properties.go: ########## @@ -660,9 +744,80 @@ func NewWriterProperties(opts ...WriterProperty) *WriterProperties { get(key).BloomFilterCandidates = value } + for key, value := range cfg.alpEnabled { + get(key).ALPEnabled = value + } + for key, value := range cfg.alpVectorSize { + get(key).ALPVectorSize = value + } + + // The check runs here rather than in the options, because an option that + // selects ALP may be given before the one that allows it. + if cfg.wr.defColumnProps.Encoding == Encodings.ALP && !cfg.wr.defColumnProps.ALPEnabled { + panic("parquet: ALP encoding requires WithALPEncoding(true)") + } + checkALPVectorSize("", cfg.wr.defColumnProps.ALPVectorSize) + for path, props := range cfg.wr.columnProps { + if props.Encoding == Encodings.ALP && !props.ALPEnabled { + panic("parquet: ALP encoding for column " + path + + " requires WithALPEncoding(true) or WithALPEncodingFor(" + path + ", true)") Review Comment: Done in 85c4a970. The message now formats the path with `%q` in both places. -- 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]
