Copilot commented on code in PR #1289:
URL: https://github.com/apache/arrow-go/pull/1289#discussion_r4528728923
##########
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(fmt.Sprintf("parquet: ALP encoding for column %q
requires WithALPEncoding(true) or WithALPEncodingFor(%q, true)",
+ path, path))
+ }
+ checkALPVectorSize(path, props.ALPVectorSize)
+ }
+
return cfg.wr
}
+// ALPEncodingEnabled reports whether columns may use the ALP encoding by
+// default. Use ALPEncodingEnabledFor for a specific column.
+func (w *WriterProperties) ALPEncodingEnabled() bool { return
w.defColumnProps.ALPEnabled }
+
+// ALPVectorSize returns the default number of values in each ALP vector. Use
+// ALPVectorSizeFor for a specific column.
+func (w *WriterProperties) ALPVectorSize() int { return
w.defColumnProps.ALPVectorSize }
+
+// ALPEncodingEnabledFor reports whether the given column may be written with
the
+// ALP encoding, or the default value if the column was not separately
specified.
+func (w *WriterProperties) ALPEncodingEnabledFor(path string) bool {
+ if p, ok := w.columnProps[path]; ok {
+ return p.ALPEnabled
+ }
+ return w.defColumnProps.ALPEnabled
+}
+
+// ALPEncodingEnabledPath is the same as ALPEncodingEnabledFor but takes a
+// ColumnPath object.
+func (w *WriterProperties) ALPEncodingEnabledPath(path ColumnPath) bool {
+ return w.ALPEncodingEnabledFor(path.String())
+}
+
+// ALPVectorSizeFor returns how many values each ALP vector of the given column
+// holds, or the default if the column was not separately specified.
+func (w *WriterProperties) ALPVectorSizeFor(path string) int {
+ if p, ok := w.columnProps[path]; ok {
+ return p.ALPVectorSize
+ }
+ return w.defColumnProps.ALPVectorSize
+}
+
+// ALPVectorSizePath is the same as ALPVectorSizeFor but takes a ColumnPath
+// object.
+func (w *WriterProperties) ALPVectorSizePath(path ColumnPath) int {
+ return w.ALPVectorSizeFor(path.String())
+}
+
+func checkALPVectorSize(path string, size int) {
+ if size >= minALPVectorSize && size <= maxALPVectorSize &&
size&(size-1) == 0 {
+ return
+ }
+ target := "the ALP vector size"
+ if path != "" {
+ target += " of column " + path
Review Comment:
The panic message for per-column vector sizes concatenates the raw `path`
without quoting/escaping (unlike nearby ALP opt-in panics that use `%q`). For
nested/complex paths this can be ambiguous. Consider using formatting that
quotes the path (e.g., building `target` with `%q` or including `path` as its
own `%q` argument) to make the error message unambiguous.
##########
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 {
+ sizeBits, numEncodable := alpEstimateSizeBits(values,
e, f)
+ // A scaling that turns almost everything into an
exception tells us
+ // nothing about how wide the rest would pack.
+ if numEncodable < 2 {
+ continue
+ }
Review Comment:
This uses `for i := range <int>` (integer range loops), which requires Go
1.22+. If the repo’s minimum supported Go version is below 1.22, this won’t
compile. Either (a) rewrite these loops to classic `for i := 0; i < n; i++`
form, or (b) bump the module’s Go version accordingly and ensure CI/tooling
reflects it.
##########
parquet/internal/encoding/alp_encoding.go:
##########
@@ -0,0 +1,673 @@
+// 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)
+ }
+ // nvals counts the page's nulls too, so a smaller element count is
valid,
+ // but a page can't hold more non-null values than it has values.
+ if int64(numElements) > int64(nvals) {
+ return fmt.Errorf("parquet: ALP page claims %d values, but the
page holds at most %d",
+ numElements, nvals)
+ }
+
+ 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
+
+ offsetArraySize := dec.numVectors * 4
+ if len(data) < alpHeaderSize+offsetArraySize {
+ return fmt.Errorf("parquet: ALP data too short for offset
array: need %d, have %d",
+ alpHeaderSize+offsetArraySize, len(data))
+ }
+
+ // Reuse these buffers across pages. Only the offset count varies by
page.
+ dec.vectorOffsets = slices.Grow(dec.vectorOffsets[:0],
dec.numVectors)[:dec.numVectors]
+ for i := range dec.vectorOffsets {
+ dec.vectorOffsets[i] =
binary.LittleEndian.Uint32(data[alpHeaderSize+i*4:])
+ }
+
+ // Vector offsets count from the start of the offset array, so keep the
page
+ // from there rather than from the start of the data.
+ dec.bodyData = data[alpHeaderSize:]
Review Comment:
The decoder currently trusts the vector offset table without validating that
offsets actually point past the offset array (e.g., first offset should be >=
offsetArraySize) and are monotonically increasing / within bounds. A corrupted
offset of 0 would point into the offset array itself; this can lead to decoding
garbage without necessarily triggering an error. Consider validating in SetData
that: (1) each offset is >= uint32(offsetArraySize), (2) offsets are
non-decreasing (ideally strictly increasing unless empty vectors are allowed),
and (3) each offset is <= len(dec.bodyData)-alpVectorHeaderSize[T](). This
makes the 'bounds-checks every size and offset' claim true in practice and
avoids silent corruption on malformed pages.
--
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]
