From 51c104d05615fca88b1d472b71f970050c60b903 Mon Sep 17 00:00:00 2001 From: Prateek Gaur Date: Sun, 4 Oct 2026 00:10:07 +0000 Subject: [PATCH 01/11] Choose ALP scaling parameters by encoded size ALP turns a float into an integer by multiplying it by 10^exponent and then by 10^-factor. This adds that arithmetic and the search that picks the pair for each vector. The search estimates the size of every pair on a sample of 256 values, counting the packed width of the values that convert and the full cost of those that do not, and breaks ties toward the larger exponent as the ALP paper does. After eight vectors it only tries the five pairs those vectors chose most often. The arithmetic follows the reference implementation's rounding and order of multiplication exactly. Rounding decides which values become exceptions, so a writer that differs by one bit writes a different page. --- parquet/internal/encoding/alp_params.go | 154 +++++++++++++ parquet/internal/encoding/alp_params_test.go | 179 +++++++++++++++ parquet/internal/encoding/alp_scaling.go | 214 ++++++++++++++++++ parquet/internal/encoding/alp_scaling_test.go | 206 +++++++++++++++++ 4 files changed, 753 insertions(+) create mode 100644 parquet/internal/encoding/alp_params.go create mode 100644 parquet/internal/encoding/alp_params_test.go create mode 100644 parquet/internal/encoding/alp_scaling.go create mode 100644 parquet/internal/encoding/alp_scaling_test.go diff --git a/parquet/internal/encoding/alp_params.go b/parquet/internal/encoding/alp_params.go new file mode 100644 index 000000000..60e90c8aa --- /dev/null +++ b/parquet/internal/encoding/alp_params.go @@ -0,0 +1,154 @@ +// 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 + for _, v := range values { + e, ok := alpTryEncode(v, exponent, factor) + 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 + } + candidate := alpEncodingParams{exponent: e, factor: f} + if alpBetterParams(candidate, sizeBits, best, bestBits) { + best, bestBits = candidate, sizeBits + } + } + } + return best +} + +// alpFindBestParamsWithPresets returns whichever shortlisted scaling is cheapest +// for the given values. The shortlist is ordered by how often each scaling won +// during sampling, so the search stops once alpPresetGiveUpAfter of them in a +// row fail to improve on the best. +func alpFindBestParamsWithPresets[T alpFloat](values []T, presets []alpEncodingParams) alpEncodingParams { + best, bestBits := presets[0], int64(math.MaxInt64) + notImproved := 0 + + for _, p := range presets { + // Only the size counts here. Every shortlisted scaling already encoded + // enough values during sampling to be worth trying. + sizeBits, _ := alpEstimateSizeBits(values, p.exponent, p.factor) + if sizeBits >= bestBits { + notImproved++ + if notImproved == alpPresetGiveUpAfter { + break + } + continue + } + best, bestBits = p, sizeBits + notImproved = 0 + } + return best +} diff --git a/parquet/internal/encoding/alp_params_test.go b/parquet/internal/encoding/alp_params_test.go new file mode 100644 index 000000000..5ae231a14 --- /dev/null +++ b/parquet/internal/encoding/alp_params_test.go @@ -0,0 +1,179 @@ +// 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 + +import ( + "math" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// alpCountTestExceptions counts the values a scaling cannot encode. +func alpCountTestExceptions[T float32 | float64](values []T, params alpEncodingParams) int { + count := 0 + for _, v := range values { + if alpDecode[T](alpEncode(v, params.exponent, params.factor), params.exponent, params.factor) != v { + count++ + } + } + return count +} + +// assertCheapestParams fails if any exponent and factor pair estimates smaller +// than the one the search returned. +func assertCheapestParams[T float32 | float64](t *testing.T, values []T, got alpEncodingParams) { + t.Helper() + gotBits, _ := alpEstimateSizeBits(values, got.exponent, got.factor) + for e := range alpNumExponents[T]() { + for f := range e + 1 { + bits, numEncodable := alpEstimateSizeBits(values, e, f) + if numEncodable < 2 { + continue + } + assert.LessOrEqual(t, gotBits, bits, + "exponent %d factor %d estimates %d bits, below the chosen %d/%d at %d", + e, f, bits, got.exponent, got.factor, gotBits) + } + } +} + +func TestALPFindBestFloat32Params(t *testing.T) { + values := []float32{1.23, 4.56, 7.89, 10.11, 12.13} + params := alpFindBestParams(values) + assert.Equal(t, 0, alpCountTestExceptions(values, params), "two decimal places always encode") + assertCheapestParams(t, values, params) + assert.Equal(t, 2, params.exponent-params.factor, "the values scale by 10^2") +} + +// TestALPFindBestFloat64Params covers a case where the exponent is not the count +// of decimal places. Scaling 4.56 by 10^2 gives an integer that decodes to +// 4.5600000000000005, so the search has to reach a larger pair that does not +// lose the last bit. +func TestALPFindBestFloat64Params(t *testing.T) { + values := []float64{1.23, 4.56, 7.89, 10.11, 12.13} + params := alpFindBestParams(values) + assert.Equal(t, 0, alpCountTestExceptions(values, params)) + assertCheapestParams(t, values, params) + assert.Equal(t, 2, params.exponent-params.factor, "the values scale by 10^2") +} + +func TestALPFindBestParamsWithExceptions(t *testing.T) { + values := []float32{ + 1.23, 4.56, float32(math.Inf(1)), 7.89, + alpNegZero[float32](), 10.11, + } + params := alpFindBestParams(values) + assert.Equal(t, 2, alpCountTestExceptions(values, params), + "the infinity and the negative zero are the only exceptions") + assertCheapestParams(t, values, params) +} + +// TestALPFindBestParamsPrefersNarrowPacking pins the reason the search weighs +// estimated size rather than exceptions alone: both pairs below encode every +// value, and the one with the smaller factor packs 30 bits wider. +func TestALPFindBestParamsPrefersNarrowPacking(t *testing.T) { + values := []float64{1.23, 4.56, 7.89} + wide, _ := alpEstimateSizeBits(values, 14, 0) + narrow, _ := alpEstimateSizeBits(values, 14, 12) + assert.Less(t, narrow, wide) + + params := alpFindBestParams(values) + assert.Equal(t, 0, alpCountTestExceptions(values, params)) + assertCheapestParams(t, values, params) +} + +func TestALPFindBestFloat32ParamsWithPresets(t *testing.T) { + values := []float32{1.23, 4.56, 7.89} + presets := []alpEncodingParams{{3, 0}, {2, 0}, {1, 0}} + params := alpFindBestParamsWithPresets(values, presets) + assert.Equal(t, 0, alpCountTestExceptions(values, params)) + // 10^3 encodes the values too, but three digits wider than it has to. + assert.Equal(t, 2, params.exponent) + assert.Equal(t, 0, params.factor) +} + +func TestALPFindBestFloat64ParamsWithPresets(t *testing.T) { + values := []float64{1.23, 4.56, 7.89} + presets := []alpEncodingParams{{2, 0}, {14, 12}, {16, 12}} + params := alpFindBestParamsWithPresets(values, presets) + assert.Equal(t, 0, alpCountTestExceptions(values, params)) + assert.Equal(t, 14, params.exponent) + assert.Equal(t, 12, params.factor) +} + +// TestALPFindBestParamsWithPresetsStopsImproving covers the early exit: the +// shortlist is ordered by how often each scaling won, so once several in a row +// have failed to beat the best, the rest are not worth estimating. +func TestALPFindBestParamsWithPresetsStopsImproving(t *testing.T) { + values := []float32{1.23, 4.56, 7.89} + presets := []alpEncodingParams{{2, 0}, {3, 0}, {4, 0}, {5, 0}, {6, 0}} + require.GreaterOrEqual(t, len(presets), alpPresetGiveUpAfter+1) + + params := alpFindBestParamsWithPresets(values, presets) + assert.Equal(t, 2, params.exponent) +} + +func TestALPFindBestParamsIntegerData(t *testing.T) { + values := []float32{1.0, 2.0, 3.0, 100.0, 200.0} + params := alpFindBestParams(values) + assert.Equal(t, 0, alpCountTestExceptions(values, params)) + assert.Equal(t, params.exponent, params.factor, + "integers need no scaling, so the exponent and the factor cancel") +} + +// TestALPFindBestParamsAllExceptions covers a vector no scaling can encode. The +// search returns the pair that leaves the values as they are, and the encoder +// writes all of them verbatim. +func TestALPFindBestParamsAllExceptions(t *testing.T) { + inf := float32(math.Inf(1)) + negZero := alpNegZero[float32]() + values := []float32{inf, float32(math.Inf(-1)), negZero, inf} + + params := alpFindBestParams(values) + assert.Equal(t, len(values), alpCountTestExceptions(values, params)) + assert.Equal(t, alpNumExponents[float32]()-1, params.exponent) + assert.Equal(t, params.exponent, params.factor) +} + +func TestALPBetterParams(t *testing.T) { + a := alpEncodingParams{exponent: 3, factor: 1} + b := alpEncodingParams{exponent: 2, factor: 2} + assert.True(t, alpBetterParams(a, 100, b, 200), "fewer bits wins") + assert.False(t, alpBetterParams(a, 200, b, 100)) + assert.True(t, alpBetterParams(a, 100, b, 100), "a tie goes to the larger exponent") + assert.True(t, alpBetterParams(alpEncodingParams{exponent: 2, factor: 2}, 100, + alpEncodingParams{exponent: 2, factor: 1}, 100), "then to the larger factor") +} + +func TestALPSample(t *testing.T) { + values := make([]float64, 4*alpSamplerSamples) + for i := range values { + values[i] = float64(i) + } + + sample := alpSample(values, nil) + assert.Len(t, sample, alpSamplerSamples) + assert.Equal(t, float64(0), sample[0]) + assert.Equal(t, float64(4), sample[1], "the stride spreads the sample over the vector") + + // Fewer values than the sample size means every value is sampled. + short := alpSample(values[:10], nil) + assert.Equal(t, values[:10], short) + assert.Empty(t, alpSample(values[:0], nil)) +} diff --git a/parquet/internal/encoding/alp_scaling.go b/parquet/internal/encoding/alp_scaling.go new file mode 100644 index 000000000..6b591055b --- /dev/null +++ b/parquet/internal/encoding/alp_scaling.go @@ -0,0 +1,214 @@ +// 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 implementations must use identical scaling arithmetic. Rounding decides +// whether a value is packed or stored verbatim as an exception, so even a +// one-bit difference can change the page contents. + +import ( + "math" + + "github.com/apache/arrow-go/v18/parquet" +) + +const ( + alpMagicFloat32 = float32(12582912.0) // 2^22 + 2^23 + alpMagicFloat64 = float64(6755399441055744.0) // 2^51 + 2^52 + + // The largest and smallest values of each type that convert to an integer + // of the matching width. They stop short of the integer limits because + // neither type can represent every integer that close to them. + alpEncodingUpperFloat32 = float32(2147483520.0) + alpEncodingLowerFloat32 = float32(-2147483520.0) + alpEncodingUpperFloat64 = float64(9223372036854774784.0) + alpEncodingLowerFloat64 = float64(-9223372036854774784.0) +) + +type alpFloat interface { + float32 | float64 +} + +// The exponent and factor index these tables, whose lengths define the search +// space. A float32 multiplier above 1e10 pushes every value outside int32. +// +// Scaling multiplies by a tabulated 10^-i rather than dividing by 10^i. The two +// differ by up to one unit in the last place, and every ALP implementation has +// to pick the same one or they decode each other's pages to different values. +var ( + alpFloatPow10 = [...]float32{ + 1e0, 1e1, 1e2, 1e3, 1e4, 1e5, 1e6, 1e7, 1e8, 1e9, 1e10, + } + alpFloatNegPow10 = [...]float32{ + 1e0, 1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7, 1e-8, 1e-9, 1e-10, + } + alpDoublePow10 = [...]float64{ + 1e0, 1e1, 1e2, 1e3, 1e4, 1e5, 1e6, 1e7, 1e8, 1e9, + 1e10, 1e11, 1e12, 1e13, 1e14, 1e15, 1e16, 1e17, 1e18, + } + alpDoubleNegPow10 = [...]float64{ + 1e0, 1e-1, 1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7, 1e-8, 1e-9, + 1e-10, 1e-11, 1e-12, 1e-13, 1e-14, 1e-15, 1e-16, 1e-17, 1e-18, + } +) + +func alpMagic[T alpFloat]() T { + var z T + switch any(z).(type) { + case float32: + return T(alpMagicFloat32) + case float64: + return T(alpMagicFloat64) + } + panic("parquet: ALP: unsupported float type") +} + +func alpPow10[T alpFloat](i int) T { + var z T + switch any(z).(type) { + case float32: + return T(alpFloatPow10[i]) + case float64: + return T(alpDoublePow10[i]) + } + panic("parquet: ALP: unsupported float type") +} + +func alpNegPow10[T alpFloat](i int) T { + var z T + switch any(z).(type) { + case float32: + return T(alpFloatNegPow10[i]) + case float64: + return T(alpDoubleNegPow10[i]) + } + panic("parquet: ALP: unsupported float type") +} + +func alpNumExponents[T alpFloat]() int { + var z T + switch any(z).(type) { + case float32: + return len(alpFloatPow10) + case float64: + return len(alpDoublePow10) + } + panic("parquet: ALP: unsupported float type") +} + +// The frame of reference and exception values use the physical type's width. +func alpEncodedBytes[T alpFloat]() int { + var z T + switch any(z).(type) { + case float32: + return 4 + case float64: + return 8 + } + panic("parquet: ALP: unsupported float type") +} + +// alpEncodingLimits returns the scaled range that converts to an integer. The +// bounds stop short of the integer limits because floats cannot represent every +// integer near them. +func alpEncodingLimits[T alpFloat]() (lo, hi T) { + var z T + switch any(z).(type) { + case float32: + return T(alpEncodingLowerFloat32), T(alpEncodingUpperFloat32) + case float64: + return T(alpEncodingLowerFloat64), T(alpEncodingUpperFloat64) + } + panic("parquet: ALP: unsupported float type") +} + +func alpParquetType[T alpFloat]() parquet.Type { + var z T + switch any(z).(type) { + case float32: + return parquet.Types.Float + case float64: + return parquet.Types.Double + } + panic("parquet: ALP: unsupported float type") +} + +func alpFloatToBits[T alpFloat](v T) uint64 { + switch f := any(v).(type) { + case float32: + return uint64(math.Float32bits(f)) + case float64: + return math.Float64bits(f) + } + panic("parquet: ALP: unsupported float type") +} + +func alpFloatFromBits[T alpFloat](b uint64) T { + var z T + switch any(z).(type) { + case float32: + return T(math.Float32frombits(uint32(b))) + case float64: + return T(math.Float64frombits(b)) + } + panic("parquet: ALP: unsupported float type") +} + +// Adding and subtracting this constant rounds as in the reference +// implementation without a library call. +func alpFastRound[T alpFloat](v T) int64 { + magic := alpMagic[T]() + if v >= 0 { + return int64((v + magic) - magic) + } + return int64((v - magic) + magic) +} + +// NaNs, infinities, out-of-range values, and negative zero cannot round-trip +// through ALP's integer representation. +func alpImpossibleToEncode[T alpFloat](v T) bool { + lo, hi := alpEncodingLimits[T]() + f := float64(v) + return math.IsNaN(f) || v > hi || v < lo || (f == 0 && math.Signbit(f)) +} + +// Unencodable values use a sentinel that cannot decode to the original value, +// so the caller records them as exceptions. +func alpNumberToInt[T alpFloat](v T) int64 { + if alpImpossibleToEncode(v) { + _, hi := alpEncodingLimits[T]() + return int64(hi) + } + return alpFastRound(v) +} + +func alpEncode[T alpFloat](v T, exponent, factor int) int64 { + return alpNumberToInt(v * alpPow10[T](exponent) * alpNegPow10[T](factor)) +} + +// The multiplication order is part of the format; reassociation can change the +// final bit. +func alpDecode[T alpFloat](encoded int64, exponent, factor int) T { + return T(encoded) * alpPow10[T](factor) * alpNegPow10[T](exponent) +} + +// alpTryEncode reports whether v survives the required round trip. Values that +// fail are stored verbatim as exceptions. +func alpTryEncode[T alpFloat](v T, exponent, factor int) (encoded int64, ok bool) { + encoded = alpEncode(v, exponent, factor) + return encoded, alpDecode[T](encoded, exponent, factor) == v +} diff --git a/parquet/internal/encoding/alp_scaling_test.go b/parquet/internal/encoding/alp_scaling_test.go new file mode 100644 index 000000000..0545714a6 --- /dev/null +++ b/parquet/internal/encoding/alp_scaling_test.go @@ -0,0 +1,206 @@ +// 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 + +import ( + "math" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestALPFastRoundFloat32(t *testing.T) { + tests := []struct { + name string + input float32 + expected int64 + }{ + {"positive integer", 5.0, 5}, + {"positive round up", 2.7, 3}, + {"positive round down", 2.3, 2}, + {"positive half", 2.5, 2}, // round-to-even via magic trick + {"negative integer", -5.0, -5}, + {"negative round up", -2.3, -2}, + {"negative round down", -2.7, -3}, + {"zero", 0.0, 0}, + {"small positive", 0.1, 0}, + {"large integer", 12345.0, 12345}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := alpFastRound(tt.input) + assert.Equal(t, tt.expected, result, "alpFastRound(%v) = %d, want %d", tt.input, result, tt.expected) + }) + } +} + +func TestALPFastRoundFloat64(t *testing.T) { + tests := []struct { + name string + input float64 + expected int64 + }{ + {"positive integer", 5.0, 5}, + {"positive round up", 2.7, 3}, + {"positive round down", 2.3, 2}, + {"negative integer", -5.0, -5}, + {"negative round up", -2.3, -2}, + {"negative round down", -2.7, -3}, + {"zero", 0.0, 0}, + {"small positive", 0.1, 0}, + {"large integer", 123456789.0, 123456789}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + result := alpFastRound(tt.input) + assert.Equal(t, tt.expected, result, "alpFastRound(%v) = %d, want %d", tt.input, result, tt.expected) + }) + } +} + +// alpNegZero returns negative zero, which ALP has to store as an exception +// because the round trip would drop the sign. +func alpNegZero[T alpFloat]() T { + return T(math.Copysign(0, -1)) +} + +func alpIsException[T alpFloat](t *testing.T, v T, exponent, factor int) bool { + t.Helper() + _, ok := alpTryEncode(v, exponent, factor) + return !ok +} + +func TestALPTryEncodeFloat32Exceptions(t *testing.T) { + assert.True(t, alpIsException(t, float32(math.NaN()), 2, 0), "NaN has no integer form") + assert.True(t, alpIsException(t, float32(math.Inf(1)), 2, 0), "+Inf has no integer form") + assert.True(t, alpIsException(t, float32(math.Inf(-1)), 2, 0), "-Inf has no integer form") + assert.True(t, alpIsException(t, alpNegZero[float32](), 2, 0), "-0.0 would come back as +0.0") + assert.True(t, alpIsException(t, float32(1e30), 10, 0), "1e30 * 10^10 overflows int32") + + assert.False(t, alpIsException(t, float32(1.23), 2, 0), "1.23 at 10^2 encodes to 123") + assert.False(t, alpIsException(t, float32(0.0), 2, 0)) + assert.False(t, alpIsException(t, float32(-1.0), 2, 0)) + assert.False(t, alpIsException(t, float32(42.0), 0, 0), "42.0 is already an integer") +} + +func TestALPTryEncodeFloat64Exceptions(t *testing.T) { + assert.True(t, alpIsException(t, math.NaN(), 2, 0), "NaN has no integer form") + assert.True(t, alpIsException(t, math.Inf(1), 2, 0), "+Inf has no integer form") + assert.True(t, alpIsException(t, math.Inf(-1), 2, 0), "-Inf has no integer form") + assert.True(t, alpIsException(t, alpNegZero[float64](), 2, 0), "-0.0 would come back as +0.0") + assert.True(t, alpIsException(t, 1e300, 18, 0), "1e300 * 10^18 overflows int64") + + assert.False(t, alpIsException(t, 1.23, 2, 0), "1.23 at 10^2 encodes to 123") + assert.False(t, alpIsException(t, 0.0, 2, 0)) + assert.False(t, alpIsException(t, -1.0, 2, 0)) + assert.False(t, alpIsException(t, 42.0, 0, 0), "42.0 is already an integer") +} + +func TestALPTryEncodeReturnsTheEncodedValue(t *testing.T) { + e, ok := alpTryEncode(1.23, 2, 0) + require.True(t, ok) + assert.EqualValues(t, 123, e, "1.23 at 10^2 encodes to 123") + + e, ok = alpTryEncode(1.5, 3, 2) + require.True(t, ok) + assert.EqualValues(t, 15, e, "10^3 * 10^-2 scales by ten") +} + +func TestALPEncodeDecodeFloat32(t *testing.T) { + tests := []struct { + name string + value float32 + exponent int + factor int + encoded int64 + }{ + {"integer no scaling", 42.0, 0, 0, 42}, + {"one decimal", 1.5, 1, 0, 15}, + {"two decimals", 1.23, 2, 0, 123}, + {"negative", -1.23, 2, 0, -123}, + {"zero", 0.0, 5, 0, 0}, + {"with factor", 1230.0, 5, 2, 1230000}, // 1230 * 10^5 / 10^2 = 1230 * 1000 + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + encoded := alpEncode(tt.value, tt.exponent, tt.factor) + assert.Equal(t, tt.encoded, encoded, "encode(%v, exp=%d, fac=%d)", tt.value, tt.exponent, tt.factor) + + decoded := alpDecode[float32](encoded, tt.exponent, tt.factor) + assert.Equal(t, math.Float32bits(tt.value), math.Float32bits(decoded), + "decode(%d, exp=%d, fac=%d) = %v, want %v", encoded, tt.exponent, tt.factor, decoded, tt.value) + }) + } +} + +func TestALPEncodeDecodeFloat64(t *testing.T) { + tests := []struct { + name string + value float64 + exponent int + factor int + encoded int64 + }{ + {"integer no scaling", 42.0, 0, 0, 42}, + {"one decimal", 1.5, 1, 0, 15}, + {"two decimals", 1.23, 2, 0, 123}, + {"negative", -1.23, 2, 0, -123}, + {"zero", 0.0, 5, 0, 0}, + {"with factor", 1230.0, 5, 2, 1230000}, // 1230 * 10^5 / 10^2 = 1230 * 1000 + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + encoded := alpEncode(tt.value, tt.exponent, tt.factor) + assert.Equal(t, tt.encoded, encoded, "encode(%v, exp=%d, fac=%d)", tt.value, tt.exponent, tt.factor) + + decoded := alpDecode[float64](encoded, tt.exponent, tt.factor) + assert.Equal(t, math.Float64bits(tt.value), math.Float64bits(decoded), + "decode(%d, exp=%d, fac=%d) = %v, want %v", encoded, tt.exponent, tt.factor, decoded, tt.value) + }) + } +} + +func TestALPPowerTables(t *testing.T) { + assert.Equal(t, float32(1.0), alpPow10[float32](0)) + assert.Equal(t, float32(100.0), alpPow10[float32](2)) + assert.Equal(t, float32(1e10), alpPow10[float32](alpNumExponents[float32]()-1)) + assert.Equal(t, float32(1.0), alpNegPow10[float32](0)) + assert.Equal(t, float32(1e-10), alpNegPow10[float32](alpNumExponents[float32]()-1)) + + assert.Equal(t, float64(1.0), alpPow10[float64](0)) + assert.Equal(t, float64(1e18), alpPow10[float64](alpNumExponents[float64]()-1)) + assert.Equal(t, float64(1e-18), alpNegPow10[float64](alpNumExponents[float64]()-1)) +} + +// TestALPScalingMultipliesByANegativePower guards the one arithmetic choice +// that two implementations can disagree on. Dividing by 10^exponent and +// multiplying by the tabulated 10^-exponent differ by a unit in the last place, +// so another implementation could decode the page to a different bit pattern. +func TestALPScalingMultipliesByANegativePower(t *testing.T) { + const value = float32(1016.10999) + const exponent, factor = 6, 0 + + encoded := alpEncode(value, exponent, factor) + assert.Equal(t, value, alpDecode[float32](encoded, exponent, factor)) + assert.NotEqual(t, value, float32(encoded)/(alpPow10[float32](exponent)/alpPow10[float32](factor)), + "dividing has stopped losing the last bit, so this test no longer guards anything") +} From ce27eef85966df17136f792f681a382e14d3df86 Mon Sep 17 00:00:00 2001 From: Arnav Balyan Date: Sun, 4 Oct 2026 00:11:40 +0000 Subject: [PATCH 02/11] Encode and decode ALP float pages This adds the ALP page format for FLOAT and DOUBLE columns. A page has a seven-byte header, one offset per vector, and vectors of 1024 values. Each vector stores its scaling, a frame of reference, the bit-packed differences from that frame, and the values that did not survive the scaling as exceptions patched in by position. Packing uses the existing BitWriter and BitReader. The differences are computed in unsigned arithmetic, so a DOUBLE vector spanning most of the int64 range packs at 64 bits instead of overflowing. The decoder accepts any vector size a page declares, from 8 to 32768 values, so it reads pages from writers that do not use 1024. It checks every offset, bit width and exception position against the page before reading, and returns an error when the page holds fewer values than the page header claims instead of a short read the column reader would retry forever. ALP takes encoding id 10. The generated Thrift bindings have no checked-in source, so the new value is added to them by hand. Co-authored-by: Prateek Gaur --- parquet/internal/encoding/alp_encoding.go | 608 ++++++++ .../internal/encoding/alp_encoding_test.go | 1263 +++++++++++++++++ parquet/internal/encoding/encoding_test.go | 11 + parquet/internal/encoding/typed_encoder.go | 4 + parquet/internal/gen-go/parquet/parquet.go | 4 + parquet/types.go | 3 + 6 files changed, 1893 insertions(+) create mode 100644 parquet/internal/encoding/alp_encoding.go create mode 100644 parquet/internal/encoding/alp_encoding_test.go diff --git a/parquet/internal/encoding/alp_encoding.go b/parquet/internal/encoding/alp_encoding.go new file mode 100644 index 000000000..23c6114d3 --- /dev/null +++ b/parquet/internal/encoding/alp_encoding.go @@ -0,0 +1,608 @@ +// 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 uint32. 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, rather than ALP for a sparse column + alpIntegerEncodingFOR = 0 // frame of reference plus bit packing + + alpHeaderSize = 7 + + // This encoder writes 1,024-value vectors. 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 by every vector, so encoding allocates nothing per vector. + encodedScratch []int64 + deltaScratch []uint64 + excPositions []uint16 + excValues []T + sampleScratch []T + + vectorsProcessed int + cachedPresets []alpEncodingParams + presetCounts map[alpEncodingParams]int +} + +func newALPEncoder[T alpFloat](e format.Encoding, descr *schema.Column, mem memory.Allocator) *alpEncoder[T] { + return &alpEncoder[T]{ + encoder: newEncoderBase(e, descr, mem), + vectorSize: alpDefaultVectorSize, + logVectorSize: alpDefaultLogVectorSize, + vectorBuf: make([]T, alpDefaultVectorSize), + encodedScratch: make([]int64, alpDefaultVectorSize), + deltaScratch: make([]uint64, alpDefaultVectorSize), + sampleScratch: make([]T, 0, alpSamplerSamples), + presetCounts: make(map[alpEncodingParams]int), + } +} + +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) { + nbuf := make([]T, len(in)) + nvalid := spacedCompress(in, nbuf, validBits, validBitsOffset) + enc.Put(nbuf[: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. + minValue := int64(math.MaxInt64) + for i, v := range values { + e, ok := alpTryEncode(v, params.exponent, params.factor) + 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 unencoded buffered values at their physical +// width. +func (enc *alpEncoder[T]) EstimatedDataEncodedSize() int64 { + return int64(len(enc.encodedBuf)) + int64(enc.bufCount*alpEncodedBytes[T]()) +} + +// 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 := make([]byte, 0, alpHeaderSize+offsetArraySize) + header = append(header, 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.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. +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) + dec.numVectors = (dec.totalCount + dec.vectorSize - 1) / dec.vectorSize + 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:] + dec.decodedValues = slices.Grow(dec.decodedValues[:0], dec.vectorSize)[:dec.vectorSize] + dec.deltas = slices.Grow(dec.deltas[:0], dec.vectorSize)[:dec.vectorSize] + + return nil +} + +// vectorLength returns how many values the given vector holds. Only the last +// one can be short. +func (dec *alpDecoder[T]) vectorLength(vectorIdx int) int { + if vectorIdx < dec.numVectors-1 { + return dec.vectorSize + } + if lastLen := dec.totalCount % dec.vectorSize; lastLen != 0 { + return lastLen + } + return dec.vectorSize +} + +// decodeVector decodes one vector into dec.decodedValues. It checks every size +// and position it reads, since the page may be truncated or corrupt. +func (dec *alpDecoder[T]) decodeVector(vectorIdx int) error { + vectorLen := dec.vectorLength(vectorIdx) + valueBytes := alpEncodedBytes[T]() + + // On 32-bit platforms, converting a large uint32 offset can make pos + // negative. + pos := int(dec.vectorOffsets[vectorIdx]) + if pos < 0 || pos+alpVectorHeaderSize[T]() > len(dec.bodyData) { + return fmt.Errorf("parquet: ALP vector %d starts past the end of the page", vectorIdx) + } + + exponent := int(dec.bodyData[pos]) + factor := int(dec.bodyData[pos+1]) + numExceptions := int(binary.LittleEndian.Uint16(dec.bodyData[pos+2:])) + frameOfRef := alpReadInt[T](dec.bodyData[pos+4:]) + bitWidth := int(dec.bodyData[pos+4+valueBytes]) + pos += alpVectorHeaderSize[T]() + + if exponent >= alpNumExponents[T]() || factor > exponent { + return fmt.Errorf("parquet: invalid ALP exponent %d and factor %d", exponent, factor) + } + if numExceptions > vectorLen { + return fmt.Errorf("parquet: ALP vector %d claims %d exceptions in %d values", + vectorIdx, numExceptions, vectorLen) + } + if bitWidth > valueBytes*8 { + return fmt.Errorf("parquet: invalid ALP bit width: %d", bitWidth) + } + + packedSize := alpPackedSize(vectorLen, bitWidth) + exceptionSize := numExceptions * (2 + valueBytes) + if pos+packedSize+exceptionSize > len(dec.bodyData) { + return fmt.Errorf("parquet: ALP vector %d runs past the end of the page", vectorIdx) + } + + deltas := dec.deltas[:vectorLen] + if err := alpUnpackBits(dec.bitReader, dec.packedReader, dec.bodyData[pos:pos+packedSize], deltas, bitWidth); err != nil { + return err + } + pos += packedSize + + for i, d := range deltas { + dec.decodedValues[i] = alpDecode[T](int64(d)+frameOfRef, exponent, factor) + } + + // Positions come first, then the values, so read both in one pass instead + // of copying the positions out. + valuePos := pos + numExceptions*2 + for e := range numExceptions { + idx := binary.LittleEndian.Uint16(dec.bodyData[pos+e*2:]) + if int(idx) >= vectorLen { + return fmt.Errorf("parquet: ALP exception position %d is outside vector %d", idx, vectorIdx) + } + dec.decodedValues[idx] = alpReadFloat[T](dec.bodyData[valuePos+e*valueBytes:]) + } + + return nil +} + +// Decode fills out with the remaining requested values, decoding vectors as +// needed. +func (dec *alpDecoder[T]) Decode(out []T) (int, error) { + toRead := min(len(out), dec.nvals) + + read := 0 + // Stop at the last value the page holds, which is not always the end of the + // last vector: that one is short whenever the count is not a whole number of + // vectors. + for read < toRead && dec.currentIndex < dec.totalCount { + vectorIdx := dec.currentIndex / dec.vectorSize + if vectorIdx != dec.currentVectorIndex { + if err := dec.decodeVector(vectorIdx); err != nil { + dec.nvals -= read + return read, err + } + dec.currentVectorIndex = vectorIdx + } + + indexInVector := dec.currentIndex % dec.vectorSize + vectorLen := dec.vectorLength(vectorIdx) + n := copy(out[read:toRead], dec.decodedValues[indexInVector:vectorLen]) + read += n + dec.currentIndex += n + } + + dec.nvals -= read + if read < toRead { + // The page header and the ALP header disagree about how many values the + // page holds. Report it rather than return a short read, which the column + // reader would retry forever, since a read of nothing leaves it where it + // was. + return read, fmt.Errorf("parquet: ALP page holds %d values, %d more were requested", + dec.totalCount, toRead-read) + } + return read, nil +} + +// DecodeSpaced leaves a slot for every null. +func (dec *alpDecoder[T]) DecodeSpaced(out []T, nullCount int, validBits []byte, validBitsOffset int64) (int, error) { + toRead := len(out) - nullCount + valuesRead, err := dec.Decode(out[:toRead]) + if err != nil { + return valuesRead, err + } + if valuesRead != toRead { + return valuesRead, errors.New("parquet: number of values / definition levels read did not match") + } + + return spacedExpand(out, nullCount, validBits, validBitsOffset), nil +} + +// Discard skips n values without decoding intervening vectors. +func (dec *alpDecoder[T]) Discard(n int) (int, error) { + n = min(n, dec.nvals) + dec.nvals -= n + dec.currentIndex += n + return n, nil +} + +type ( + // ALPFloat32Encoder encodes float32 values with ALP. + ALPFloat32Encoder = alpEncoder[float32] + // ALPFloat64Encoder encodes float64 values with ALP. + ALPFloat64Encoder = alpEncoder[float64] + // ALPFloat32Decoder decodes ALP-encoded float32 values. + ALPFloat32Decoder = alpDecoder[float32] + // ALPFloat64Decoder decodes ALP-encoded float64 values. + ALPFloat64Decoder = alpDecoder[float64] +) diff --git a/parquet/internal/encoding/alp_encoding_test.go b/parquet/internal/encoding/alp_encoding_test.go new file mode 100644 index 000000000..b7ab0378e --- /dev/null +++ b/parquet/internal/encoding/alp_encoding_test.go @@ -0,0 +1,1263 @@ +// 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 + +import ( + "bytes" + "encoding/binary" + "fmt" + "math" + "slices" + "testing" + + "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" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// alpTestUnpackBits unpacks with readers of its own, since the decoder's are +// tied to a page. +func alpTestUnpackBits(packed []byte, out []uint64, bitWidth int) error { + reader := bytes.NewReader(nil) + return alpUnpackBits(utils.NewBitReader(reader), reader, packed, out, bitWidth) +} + +func TestALPPackUnpackBits(t *testing.T) { + tests := []struct { + name string + values []uint64 + bitWidth int + }{ + {"1-bit values", []uint64{0, 1, 0, 1, 1, 0, 1, 0}, 1}, + {"3-bit values", []uint64{0, 1, 2, 3, 4, 5, 6, 7}, 3}, + {"4-bit values", []uint64{0, 1, 5, 10, 15, 3, 7, 12}, 4}, + {"8-bit values", []uint64{0, 128, 255, 1, 42, 100, 200, 50}, 8}, + {"16-bit values", []uint64{0, 1000, 50000, 65535, 32768, 1, 256, 2048}, 16}, + {"32-bit values", []uint64{0, math.MaxUint32, 12345, 67890}, 32}, + {"48-bit values", []uint64{0, 1 << 47, 12345678901234, 1}, 48}, + {"64-bit values", []uint64{0, math.MaxUint64, 12345, 67890}, 64}, + {"zero bit width", []uint64{0, 0, 0, 0}, 0}, + {"single value", []uint64{42}, 6}, + // A full vector exercises the vectorized unpack path, which handles + // whole groups of 32 values and leaves the remainder to the scalar one. + {"full vector", func() []uint64 { + v := make([]uint64, alpDefaultVectorSize) + for i := range v { + v[i] = uint64(i % 256) + } + return v + }(), 8}, + {"partial group", func() []uint64 { + v := make([]uint64, 100) + for i := range v { + v[i] = uint64(i * 1000) + } + return v + }(), 20}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + packed := alpAppendPackedBits(nil, tt.values, tt.bitWidth) + assert.Len(t, packed, alpPackedSize(len(tt.values), tt.bitWidth)) + + unpacked := make([]uint64, len(tt.values)) + require.NoError(t, alpTestUnpackBits(packed, unpacked, tt.bitWidth)) + assert.Equal(t, tt.values, unpacked, "pack/unpack round-trip failed") + }) + } +} + +// TestALPPackBitsKnownOutput pins the byte layout, so a change of bit packer +// cannot silently break pages written by an earlier one. +func TestALPPackBitsKnownOutput(t *testing.T) { + packed := alpAppendPackedBits(nil, []uint64{0, 1, 2, 3}, 2) + require.Len(t, packed, 1) + assert.Equal(t, byte(0xE4), packed[0]) +} + +// TestALPUnpackBitsShortInput requires truncated packed data to return an error. +// The two lengths cover the scalar path for a partial group of 32 values and the +// vectorized path for a whole group. +func TestALPUnpackBitsShortInput(t *testing.T) { + for _, count := range []int{8, 64} { + out := make([]uint64, count) + assert.Error(t, alpTestUnpackBits([]byte{0xFF}, out, 8), "count=%d", count) + } +} + +func newFloat32Column() *schema.Column { + return schema.NewColumn(schema.NewFloat32Node("float32", parquet.Repetitions.Required, -1), 0, 0) +} + +func newFloat64Column() *schema.Column { + return schema.NewColumn(schema.NewFloat64Node("float64", parquet.Repetitions.Required, -1), 0, 0) +} + +func alpFloat32RoundTrip(t *testing.T, values []float32) { + t.Helper() + col := newFloat32Column() + mem := memory.DefaultAllocator + + enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + dec := newALPDecoder[float32](format.Encoding_ALP, col) + err = dec.SetData(len(values), buf.Bytes()) + require.NoError(t, err) + + out := make([]float32, len(values)) + n, err := dec.Decode(out) + require.NoError(t, err) + assert.Equal(t, len(values), n) + + for i, v := range values { + assert.Equal(t, math.Float32bits(v), math.Float32bits(out[i]), + "index %d: got %v (bits %08x), want %v (bits %08x)", + i, out[i], math.Float32bits(out[i]), v, math.Float32bits(v)) + } +} + +func alpFloat64RoundTrip(t *testing.T, values []float64) { + t.Helper() + col := newFloat64Column() + mem := memory.DefaultAllocator + + enc := newALPEncoder[float64](format.Encoding_ALP, col, mem) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + dec := newALPDecoder[float64](format.Encoding_ALP, col) + err = dec.SetData(len(values), buf.Bytes()) + require.NoError(t, err) + + out := make([]float64, len(values)) + n, err := dec.Decode(out) + require.NoError(t, err) + assert.Equal(t, len(values), n) + + for i, v := range values { + assert.Equal(t, math.Float64bits(v), math.Float64bits(out[i]), + "index %d: got %v (bits %016x), want %v (bits %016x)", + i, out[i], math.Float64bits(out[i]), v, math.Float64bits(v)) + } +} + +func TestALPFloat32MonetaryRoundTrip(t *testing.T) { + values := make([]float32, 2000) + for i := range values { + values[i] = float32(i) * 0.01 + } + alpFloat32RoundTrip(t, values) +} + +func TestALPFloat64MonetaryRoundTrip(t *testing.T) { + values := make([]float64, 2000) + for i := range values { + values[i] = float64(i) * 0.01 + } + alpFloat64RoundTrip(t, values) +} + +func TestALPFloat32IntegerRoundTrip(t *testing.T) { + values := make([]float32, 1500) + for i := range values { + values[i] = float32(i) + } + alpFloat32RoundTrip(t, values) +} + +func TestALPFloat64IntegerRoundTrip(t *testing.T) { + values := make([]float64, 1500) + for i := range values { + values[i] = float64(i) + } + alpFloat64RoundTrip(t, values) +} + +func TestALPFloat32ExceptionsRoundTrip(t *testing.T) { + negZero := alpNegZero[float32]() + inf := float32(math.Inf(1)) + ninf := float32(math.Inf(-1)) + + values := []float32{ + 1.23, 4.56, inf, 7.89, negZero, 10.11, + ninf, 12.13, 14.15, 16.17, + } + alpFloat32RoundTrip(t, values) +} + +func TestALPFloat64ExceptionsRoundTrip(t *testing.T) { + negZero := alpNegZero[float64]() + inf := math.Inf(1) + ninf := math.Inf(-1) + + values := []float64{ + 1.23, 4.56, inf, 7.89, negZero, 10.11, + ninf, 12.13, 14.15, 16.17, + } + alpFloat64RoundTrip(t, values) +} + +func TestALPFloat32AllExceptionsRoundTrip(t *testing.T) { + negZero := alpNegZero[float32]() + inf := float32(math.Inf(1)) + ninf := float32(math.Inf(-1)) + + values := []float32{inf, ninf, negZero, inf, ninf, negZero, inf, ninf} + alpFloat32RoundTrip(t, values) +} + +func TestALPFloat64AllExceptionsRoundTrip(t *testing.T) { + negZero := alpNegZero[float64]() + inf := math.Inf(1) + ninf := math.Inf(-1) + + values := []float64{inf, ninf, negZero, inf, ninf, negZero, inf, ninf} + alpFloat64RoundTrip(t, values) +} + +func TestALPFloat32SingleValueRoundTrip(t *testing.T) { + alpFloat32RoundTrip(t, []float32{42.0}) + alpFloat32RoundTrip(t, []float32{0.0}) + alpFloat32RoundTrip(t, []float32{-1.23}) +} + +func TestALPFloat64SingleValueRoundTrip(t *testing.T) { + alpFloat64RoundTrip(t, []float64{42.0}) + alpFloat64RoundTrip(t, []float64{0.0}) + alpFloat64RoundTrip(t, []float64{-1.23}) +} + +func TestALPFloat32PartialVectorRoundTrip(t *testing.T) { + for _, count := range []int{1, 2, 3, 7, 100, 500, 1023} { + t.Run(fmt.Sprintf("count_%d", count), func(t *testing.T) { + values := make([]float32, count) + for i := range values { + values[i] = float32(i) * 0.1 + } + alpFloat32RoundTrip(t, values) + }) + } +} + +func TestALPFloat64PartialVectorRoundTrip(t *testing.T) { + for _, count := range []int{1, 2, 3, 7, 100, 500, 1023} { + t.Run(fmt.Sprintf("count_%d", count), func(t *testing.T) { + values := make([]float64, count) + for i := range values { + values[i] = float64(i) * 0.1 + } + alpFloat64RoundTrip(t, values) + }) + } +} + +func TestALPFloat32ExactVectorSizeRoundTrip(t *testing.T) { + values := make([]float32, 1024) + for i := range values { + values[i] = float32(i) * 0.01 + } + alpFloat32RoundTrip(t, values) +} + +func TestALPFloat64ExactVectorSizeRoundTrip(t *testing.T) { + values := make([]float64, 1024) + for i := range values { + values[i] = float64(i) * 0.01 + } + alpFloat64RoundTrip(t, values) +} + +func TestALPFloat32MultipleVectorsRoundTrip(t *testing.T) { + for _, count := range []int{1025, 2048, 3000, 5000, 10000} { + t.Run(fmt.Sprintf("count_%d", count), func(t *testing.T) { + values := make([]float32, count) + for i := range values { + values[i] = float32(i) * 0.01 + } + alpFloat32RoundTrip(t, values) + }) + } +} + +func TestALPFloat64MultipleVectorsRoundTrip(t *testing.T) { + for _, count := range []int{1025, 2048, 3000, 5000, 10000} { + t.Run(fmt.Sprintf("count_%d", count), func(t *testing.T) { + values := make([]float64, count) + for i := range values { + values[i] = float64(i) * 0.01 + } + alpFloat64RoundTrip(t, values) + }) + } +} + +func TestALPFloat32ConstantValueRoundTrip(t *testing.T) { + values := make([]float32, 1500) + for i := range values { + values[i] = 3.14 + } + alpFloat32RoundTrip(t, values) +} + +func TestALPFloat64ConstantValueRoundTrip(t *testing.T) { + values := make([]float64, 1500) + for i := range values { + values[i] = 3.14 + } + alpFloat64RoundTrip(t, values) +} + +func TestALPFloat32ZeroValuesRoundTrip(t *testing.T) { + values := make([]float32, 1024) + alpFloat32RoundTrip(t, values) +} + +func TestALPFloat64ZeroValuesRoundTrip(t *testing.T) { + values := make([]float64, 1024) + alpFloat64RoundTrip(t, values) +} + +func TestALPFloat32NegativeValuesRoundTrip(t *testing.T) { + values := make([]float32, 1500) + for i := range values { + values[i] = -float32(i) * 0.01 + } + alpFloat32RoundTrip(t, values) +} + +func TestALPFloat64NegativeValuesRoundTrip(t *testing.T) { + values := make([]float64, 1500) + for i := range values { + values[i] = -float64(i) * 0.01 + } + alpFloat64RoundTrip(t, values) +} + +func TestALPFloat32ScientificRoundTrip(t *testing.T) { + values := make([]float32, 2000) + for i := range values { + values[i] = float32(i)*0.001 + 273.15 // Kelvin temperatures + } + alpFloat32RoundTrip(t, values) +} + +func TestALPFloat64ScientificRoundTrip(t *testing.T) { + values := make([]float64, 2000) + for i := range values { + values[i] = float64(i)*0.001 + 273.15 + } + alpFloat64RoundTrip(t, values) +} + +func TestALPFloat32ProgressiveDecoding(t *testing.T) { + values := make([]float32, 3000) + for i := range values { + values[i] = float32(i) * 0.01 + } + + col := newFloat32Column() + mem := memory.DefaultAllocator + + enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + dec := newALPDecoder[float32](format.Encoding_ALP, col) + err = dec.SetData(len(values), buf.Bytes()) + require.NoError(t, err) + + chunkSize := 100 + totalRead := 0 + for totalRead < len(values) { + remaining := len(values) - totalRead + toRead := chunkSize + if toRead > remaining { + toRead = remaining + } + out := make([]float32, toRead) + n, err := dec.Decode(out) + require.NoError(t, err) + assert.Equal(t, toRead, n) + + for i := 0; i < n; i++ { + assert.Equal(t, math.Float32bits(values[totalRead+i]), math.Float32bits(out[i]), + "index %d: got %v, want %v", totalRead+i, out[i], values[totalRead+i]) + } + totalRead += n + } + assert.Equal(t, len(values), totalRead) +} + +func TestALPFloat64ProgressiveDecoding(t *testing.T) { + values := make([]float64, 3000) + for i := range values { + values[i] = float64(i) * 0.01 + } + + col := newFloat64Column() + mem := memory.DefaultAllocator + + enc := newALPEncoder[float64](format.Encoding_ALP, col, mem) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + dec := newALPDecoder[float64](format.Encoding_ALP, col) + err = dec.SetData(len(values), buf.Bytes()) + require.NoError(t, err) + + chunkSize := 100 + totalRead := 0 + for totalRead < len(values) { + remaining := len(values) - totalRead + toRead := chunkSize + if toRead > remaining { + toRead = remaining + } + out := make([]float64, toRead) + n, err := dec.Decode(out) + require.NoError(t, err) + assert.Equal(t, toRead, n) + + for i := 0; i < n; i++ { + assert.Equal(t, math.Float64bits(values[totalRead+i]), math.Float64bits(out[i]), + "index %d: got %v, want %v", totalRead+i, out[i], values[totalRead+i]) + } + totalRead += n + } + assert.Equal(t, len(values), totalRead) +} + +func TestALPFloat32Discard(t *testing.T) { + values := make([]float32, 3000) + for i := range values { + values[i] = float32(i) * 0.01 + } + + col := newFloat32Column() + mem := memory.DefaultAllocator + + enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + dec := newALPDecoder[float32](format.Encoding_ALP, col) + err = dec.SetData(len(values), buf.Bytes()) + require.NoError(t, err) + + discarded, err := dec.Discard(1500) + require.NoError(t, err) + assert.Equal(t, 1500, discarded) + + out := make([]float32, 1500) + n, err := dec.Decode(out) + require.NoError(t, err) + assert.Equal(t, 1500, n) + + for i := 0; i < n; i++ { + assert.Equal(t, math.Float32bits(values[1500+i]), math.Float32bits(out[i]), + "index %d: got %v, want %v", 1500+i, out[i], values[1500+i]) + } +} + +func TestALPFloat64Discard(t *testing.T) { + values := make([]float64, 3000) + for i := range values { + values[i] = float64(i) * 0.01 + } + + col := newFloat64Column() + mem := memory.DefaultAllocator + + enc := newALPEncoder[float64](format.Encoding_ALP, col, mem) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + dec := newALPDecoder[float64](format.Encoding_ALP, col) + err = dec.SetData(len(values), buf.Bytes()) + require.NoError(t, err) + + discarded, err := dec.Discard(1500) + require.NoError(t, err) + assert.Equal(t, 1500, discarded) + + out := make([]float64, 1500) + n, err := dec.Decode(out) + require.NoError(t, err) + assert.Equal(t, 1500, n) + + for i := 0; i < n; i++ { + assert.Equal(t, math.Float64bits(values[1500+i]), math.Float64bits(out[i]), + "index %d: got %v, want %v", 1500+i, out[i], values[1500+i]) + } +} + +func TestALPDiscardMoreThanAvailable(t *testing.T) { + values := []float32{1.0, 2.0, 3.0} + + col := newFloat32Column() + mem := memory.DefaultAllocator + + enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + dec := newALPDecoder[float32](format.Encoding_ALP, col) + err = dec.SetData(len(values), buf.Bytes()) + require.NoError(t, err) + + discarded, err := dec.Discard(100) + require.NoError(t, err) + assert.Equal(t, 3, discarded, "should only discard available values") +} + +func TestALPFloat32MultiplePages(t *testing.T) { + col := newFloat32Column() + mem := memory.DefaultAllocator + + dec := newALPDecoder[float32](format.Encoding_ALP, col) + + for page := 0; page < 5; page++ { + values := make([]float32, 500) + for i := range values { + values[i] = float32(page*500+i) * 0.01 + } + + enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + + err = dec.SetData(len(values), buf.Bytes()) + require.NoError(t, err) + + out := make([]float32, len(values)) + n, err := dec.Decode(out) + require.NoError(t, err) + assert.Equal(t, len(values), n) + + for i, v := range values { + assert.Equal(t, math.Float32bits(v), math.Float32bits(out[i]), + "page %d, index %d: mismatch", page, i) + } + + buf.Release() + } +} + +func TestALPFloat64MultiplePages(t *testing.T) { + col := newFloat64Column() + mem := memory.DefaultAllocator + + dec := newALPDecoder[float64](format.Encoding_ALP, col) + + for page := 0; page < 5; page++ { + values := make([]float64, 500) + for i := range values { + values[i] = float64(page*500+i) * 0.01 + } + + enc := newALPEncoder[float64](format.Encoding_ALP, col, mem) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + + err = dec.SetData(len(values), buf.Bytes()) + require.NoError(t, err) + + out := make([]float64, len(values)) + n, err := dec.Decode(out) + require.NoError(t, err) + assert.Equal(t, len(values), n) + + for i, v := range values { + assert.Equal(t, math.Float64bits(v), math.Float64bits(out[i]), + "page %d, index %d: mismatch", page, i) + } + + buf.Release() + } +} + +func TestALPEncoderReset(t *testing.T) { + col := newFloat32Column() + mem := memory.DefaultAllocator + + enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + + values1 := []float32{1.23, 4.56, 7.89} + enc.Put(values1) + buf1, err := enc.FlushValues() + require.NoError(t, err) + defer buf1.Release() + + enc.Reset() + values2 := []float32{10.11, 12.13} + enc.Put(values2) + buf2, err := enc.FlushValues() + require.NoError(t, err) + defer buf2.Release() + + dec := newALPDecoder[float32](format.Encoding_ALP, col) + err = dec.SetData(len(values1), buf1.Bytes()) + require.NoError(t, err) + out1 := make([]float32, len(values1)) + n, err := dec.Decode(out1) + require.NoError(t, err) + assert.Equal(t, len(values1), n) + for i, v := range values1 { + assert.Equal(t, math.Float32bits(v), math.Float32bits(out1[i])) + } + + err = dec.SetData(len(values2), buf2.Bytes()) + require.NoError(t, err) + out2 := make([]float32, len(values2)) + n, err = dec.Decode(out2) + require.NoError(t, err) + assert.Equal(t, len(values2), n) + for i, v := range values2 { + assert.Equal(t, math.Float32bits(v), math.Float32bits(out2[i])) + } +} + +func TestALPEncoderType(t *testing.T) { + col32 := newFloat32Column() + col64 := newFloat64Column() + mem := memory.DefaultAllocator + + enc32 := newALPEncoder[float32](format.Encoding_ALP, col32, mem) + assert.Equal(t, parquet.Types.Float, enc32.Type()) + + enc64 := newALPEncoder[float64](format.Encoding_ALP, col64, mem) + assert.Equal(t, parquet.Types.Double, enc64.Type()) +} + +func TestALPDecoderType(t *testing.T) { + col32 := newFloat32Column() + col64 := newFloat64Column() + + dec32 := &alpDecoder[float32]{decoder: newDecoderBase(format.Encoding_ALP, col32)} + assert.Equal(t, parquet.Types.Float, dec32.Type()) + + dec64 := &alpDecoder[float64]{decoder: newDecoderBase(format.Encoding_ALP, col64)} + assert.Equal(t, parquet.Types.Double, dec64.Type()) +} + +func TestALPEncoderPutSpaced(t *testing.T) { + col := newFloat32Column() + mem := memory.DefaultAllocator + + values := []float32{1.0, 2.0, 0.0, 4.0, 0.0, 6.0, 7.0, 0.0} + validBits := []byte{0b01101011} // bits 0,1,3,5,6 set + + enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + enc.PutSpaced(values, validBits, 0) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + expected := []float32{1.0, 2.0, 4.0, 6.0, 7.0} + + dec := newALPDecoder[float32](format.Encoding_ALP, col) + err = dec.SetData(len(expected), buf.Bytes()) + require.NoError(t, err) + + out := make([]float32, len(expected)) + n, err := dec.Decode(out) + require.NoError(t, err) + assert.Equal(t, len(expected), n) + + for i, v := range expected { + assert.Equal(t, math.Float32bits(v), math.Float32bits(out[i]), + "index %d: got %v, want %v", i, out[i], v) + } +} + +func TestALPHeaderFormat(t *testing.T) { + col := newFloat32Column() + mem := memory.DefaultAllocator + + values := make([]float32, 100) + for i := range values { + values[i] = float32(i) + } + + enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + data := buf.Bytes() + require.GreaterOrEqual(t, len(data), alpHeaderSize) + + assert.Equal(t, byte(alpCompressionMode), data[0], "compression mode") + assert.Equal(t, byte(alpIntegerEncodingFOR), data[1], "integer encoding") + assert.Equal(t, byte(alpDefaultLogVectorSize), data[2], "log vector size") + assert.Equal(t, uint32(100), binary.LittleEndian.Uint32(data[3:]), "element count") +} + +func TestALPHeaderFormatMultipleVectors(t *testing.T) { + col := newFloat64Column() + mem := memory.DefaultAllocator + + values := make([]float64, 3000) + for i := range values { + values[i] = float64(i) * 0.01 + } + + enc := newALPEncoder[float64](format.Encoding_ALP, col, mem) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + data := buf.Bytes() + require.GreaterOrEqual(t, len(data), alpHeaderSize) + assert.Equal(t, uint32(3000), binary.LittleEndian.Uint32(data[3:])) + + numVectors := 3 + offsetArraySize := numVectors * 4 + require.GreaterOrEqual(t, len(data), alpHeaderSize+offsetArraySize) + + // Offsets count from the start of the array, so the first vector begins + // where the array ends and each later one begins after its predecessor. + firstOffset := binary.LittleEndian.Uint32(data[alpHeaderSize:]) + assert.Equal(t, uint32(offsetArraySize), firstOffset) + + for i := 1; i < numVectors; i++ { + offset := binary.LittleEndian.Uint32(data[alpHeaderSize+i*4:]) + prevOffset := binary.LittleEndian.Uint32(data[alpHeaderSize+(i-1)*4:]) + assert.Greater(t, offset, prevOffset, + "offset[%d]=%d should be greater than offset[%d]=%d", i, offset, i-1, prevOffset) + } +} + +// alpTestPage returns a page whose header is valid, for tests that corrupt one +// field at a time. +func alpTestPage(numElements uint32) []byte { + data := make([]byte, 100) + data[0] = alpCompressionMode + data[1] = alpIntegerEncodingFOR + data[2] = alpDefaultLogVectorSize + binary.LittleEndian.PutUint32(data[3:], numElements) + return data +} + +func TestALPDecoderInvalidData(t *testing.T) { + col := newFloat32Column() + dec := newALPDecoder[float32](format.Encoding_ALP, col) + + err := dec.SetData(10, []byte{1, 2, 3}) + assert.Error(t, err) + assert.Contains(t, err.Error(), "too short") +} + +func TestALPDecoderInvalidHeaderFields(t *testing.T) { + tests := []struct { + name string + offset int + value byte + errMsg string + }{ + {"compression mode", 0, 99, "compression mode"}, + {"integer encoding", 1, 99, "integer encoding"}, + {"log vector size too large", 2, alpMaxLogVectorSize + 1, "log vector size"}, + {"log vector size too small", 2, alpMinLogVectorSize - 1, "log vector size"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + col := newFloat32Column() + dec := newALPDecoder[float32](format.Encoding_ALP, col) + + data := alpTestPage(1) + data[tt.offset] = tt.value + + err := dec.SetData(1, data) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errMsg) + }) + } +} + +func TestALPDecoderTruncatedOffsetArray(t *testing.T) { + col := newFloat32Column() + dec := newALPDecoder[float32](format.Encoding_ALP, col) + + data := alpTestPage(1024)[:alpHeaderSize+2] // not enough for a 4-byte offset + + err := dec.SetData(1024, data) + assert.Error(t, err) + assert.Contains(t, err.Error(), "too short") +} + +// TestALPFloat64WidestPacking covers the widest packing a page can carry. The +// values spread over most of what an int64 holds, so their differences need all +// 64 bits and the difference between the two ends overflows a signed subtraction. +func TestALPFloat64WidestPacking(t *testing.T) { + // Every value is an odd multiple of 2^53, which a float64 holds exactly. + // The two ends of the run differ by 1023 * 2^54. + values := make([]float64, alpDefaultVectorSize) + for i := range values { + values[i] = float64(int64(i-len(values)/2)*(1<<54) + (1 << 53)) + } + + enc := newALPEncoder[float64](format.Encoding_ALP, newFloat64Column(), memory.DefaultAllocator) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + // One offset follows the header, and it holds the size of the offset array. + page := buf.Bytes() + vectorPos := alpHeaderSize + int(binary.LittleEndian.Uint32(page[alpHeaderSize:])) + assert.Equal(t, byte(64), page[vectorPos+alpVectorHeaderSize[float64]()-1], "bit width") + assert.Zero(t, binary.LittleEndian.Uint16(page[vectorPos+2:]), "exception count") + + alpFloat64RoundTrip(t, values) +} + +// alpEncodeTestPage encodes values into a single-vector page and returns the +// byte offset of that vector. +func alpEncodeTestPage(t *testing.T, values []float32) (page []byte, vectorPos int) { + t.Helper() + + enc := newALPEncoder[float32](format.Encoding_ALP, newFloat32Column(), memory.DefaultAllocator) + enc.Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + t.Cleanup(buf.Release) + + page = append([]byte(nil), buf.Bytes()...) + // One offset follows the header, and it holds the size of the offset array. + vectorPos = alpHeaderSize + int(binary.LittleEndian.Uint32(page[alpHeaderSize:])) + return page, vectorPos +} + +// TestALPDecoderCorruptVector covers the checks a vector header gets. Each case +// writes one field a writer could not have produced, which the decoder has to +// report rather than read past the page or scale by a table it has no entry for. +func TestALPDecoderCorruptVector(t *testing.T) { + values := make([]float32, 100) + for i := range values { + values[i] = float32(i) * 1.25 + } + + tests := []struct { + name string + corrupt func(page []byte, vectorPos int) []byte + errMsg string + }{ + { + name: "exponent past the power table", + corrupt: func(page []byte, vectorPos int) []byte { + page[vectorPos] = byte(alpNumExponents[float32]()) + return page + }, + errMsg: "invalid ALP exponent", + }, + { + name: "factor above the exponent", + corrupt: func(page []byte, vectorPos int) []byte { + page[vectorPos+1] = page[vectorPos] + 1 + return page + }, + errMsg: "invalid ALP exponent", + }, + { + name: "more exceptions than values", + corrupt: func(page []byte, vectorPos int) []byte { + binary.LittleEndian.PutUint16(page[vectorPos+2:], uint16(len(values)+1)) + return page + }, + errMsg: "exceptions in", + }, + { + name: "bit width past the integer width", + corrupt: func(page []byte, vectorPos int) []byte { + page[vectorPos+alpVectorHeaderSize[float32]()-1] = 33 + return page + }, + errMsg: "invalid ALP bit width", + }, + { + name: "vector longer than the page", + corrupt: func(page []byte, vectorPos int) []byte { + return page[:len(page)-1] + }, + errMsg: "past the end of the page", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + page, vectorPos := alpEncodeTestPage(t, values) + page = tt.corrupt(page, vectorPos) + + dec := newALPDecoder[float32](format.Encoding_ALP, newFloat32Column()) + require.NoError(t, dec.SetData(len(values), page)) + + _, err := dec.Decode(make([]float32, len(values))) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.errMsg) + }) + } +} + +// TestALPDecoderExceptionPositionOutsideVector covers the one check that needs a +// vector with an exception in it: the position that patches the value back in +// has to land inside the vector. +func TestALPDecoderExceptionPositionOutsideVector(t *testing.T) { + values := []float32{1.25, 2.5, float32(math.Inf(1)), 3.75} + page, vectorPos := alpEncodeTestPage(t, values) + + // The positions follow the packed deltas, which the vector's own bit width + // sizes. + headerSize := alpVectorHeaderSize[float32]() + bitWidth := int(page[vectorPos+headerSize-1]) + positionPos := vectorPos + headerSize + alpPackedSize(len(values), bitWidth) + require.Equal(t, 1, int(binary.LittleEndian.Uint16(page[vectorPos+2:])), "one exception") + binary.LittleEndian.PutUint16(page[positionPos:], uint16(len(values))) + + dec := newALPDecoder[float32](format.Encoding_ALP, newFloat32Column()) + require.NoError(t, dec.SetData(len(values), page)) + + _, err := dec.Decode(make([]float32, len(values))) + require.Error(t, err) + assert.Contains(t, err.Error(), "exception position") +} + +func TestALPPresetCacheBuilds(t *testing.T) { + col := newFloat32Column() + mem := memory.DefaultAllocator + + enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + + values := make([]float32, 9000) + for i := range values { + values[i] = float32(i) * 0.01 + } + enc.Put(values) + + assert.NotNil(t, enc.cachedPresets, "preset cache should be built after 8 vectors") + assert.True(t, len(enc.cachedPresets) > 0, "should have at least one preset") + assert.True(t, len(enc.cachedPresets) <= alpMaxPresetCombinations, "should not exceed max presets") + + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + dec := newALPDecoder[float32](format.Encoding_ALP, col) + err = dec.SetData(len(values), buf.Bytes()) + require.NoError(t, err) + + out := make([]float32, len(values)) + n, err := dec.Decode(out) + require.NoError(t, err) + assert.Equal(t, len(values), n) + + for i, v := range values { + assert.Equal(t, math.Float32bits(v), math.Float32bits(out[i]), + "index %d: got %v, want %v", i, out[i], v) + } +} + +func TestALPEstimatedDataEncodedSize(t *testing.T) { + col := newFloat32Column() + mem := memory.DefaultAllocator + + enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + + assert.Equal(t, int64(0), enc.EstimatedDataEncodedSize()) + + values := make([]float32, 100) + for i := range values { + values[i] = float32(i) * 0.01 + } + enc.Put(values) + size := enc.EstimatedDataEncodedSize() + assert.True(t, size > 0, "estimated size should be positive after Put") +} + +func TestALPFloat32IncrementalPut(t *testing.T) { + col := newFloat32Column() + mem := memory.DefaultAllocator + + enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + allValues := make([]float32, 2500) + for i := range allValues { + allValues[i] = float32(i) * 0.01 + } + + offset := 0 + chunks := []int{100, 200, 500, 724, 976} // total = 2500 + for _, chunk := range chunks { + enc.Put(allValues[offset : offset+chunk]) + offset += chunk + } + + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + dec := newALPDecoder[float32](format.Encoding_ALP, col) + err = dec.SetData(len(allValues), buf.Bytes()) + require.NoError(t, err) + + out := make([]float32, len(allValues)) + n, err := dec.Decode(out) + require.NoError(t, err) + assert.Equal(t, len(allValues), n) + + for i, v := range allValues { + assert.Equal(t, math.Float32bits(v), math.Float32bits(out[i]), + "index %d: got %v, want %v", i, out[i], v) + } +} + +func TestALPFloat32MixedDataPatterns(t *testing.T) { + var values []float32 + + for i := 0; i < 500; i++ { + values = append(values, float32(i)*0.01) + } + for i := 0; i < 500; i++ { + values = append(values, float32(i)) + } + for i := 0; i < 500; i++ { + values = append(values, float32(i)*0.001+100.0) + } + values = append(values, float32(math.Inf(1))) + values = append(values, float32(math.Inf(-1))) + values = append(values, alpNegZero[float32]()) + for i := 0; i < 497; i++ { + values = append(values, float32(i)*0.1) + } + + alpFloat32RoundTrip(t, values) +} + +func TestALPFloat64MixedDataPatterns(t *testing.T) { + var values []float64 + + for i := 0; i < 500; i++ { + values = append(values, float64(i)*0.01) + } + for i := 0; i < 500; i++ { + values = append(values, float64(i)) + } + for i := 0; i < 500; i++ { + values = append(values, float64(i)*0.001+100.0) + } + values = append(values, math.Inf(1)) + values = append(values, math.Inf(-1)) + values = append(values, alpNegZero[float64]()) + for i := 0; i < 497; i++ { + values = append(values, float64(i)*0.1) + } + + alpFloat64RoundTrip(t, values) +} + +func TestALPFloat32ViaPublicAPI(t *testing.T) { + col := newFloat32Column() + mem := memory.DefaultAllocator + + values := make([]float32, 2000) + for i := range values { + values[i] = float32(i) * 0.01 + } + + enc := NewEncoder(parquet.Types.Float, parquet.Encodings.ALP, false, col, mem) + assert.Equal(t, parquet.Encodings.ALP, enc.Encoding()) + assert.Equal(t, parquet.Types.Float, enc.Type()) + + enc.(Float32Encoder).Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + dec := NewDecoder(parquet.Types.Float, parquet.Encodings.ALP, col, mem) + assert.Equal(t, parquet.Encodings.ALP, dec.Encoding()) + assert.Equal(t, parquet.Types.Float, dec.Type()) + + dec.SetData(len(values), buf.Bytes()) + out := make([]float32, len(values)) + n, err := dec.(Float32Decoder).Decode(out) + require.NoError(t, err) + assert.Equal(t, len(values), n) + + for i, v := range values { + assert.Equal(t, math.Float32bits(v), math.Float32bits(out[i]), + "index %d: mismatch", i) + } +} + +func TestALPFloat64ViaPublicAPI(t *testing.T) { + col := newFloat64Column() + mem := memory.DefaultAllocator + + values := make([]float64, 2000) + for i := range values { + values[i] = float64(i) * 0.01 + } + + enc := NewEncoder(parquet.Types.Double, parquet.Encodings.ALP, false, col, mem) + assert.Equal(t, parquet.Encodings.ALP, enc.Encoding()) + assert.Equal(t, parquet.Types.Double, enc.Type()) + + enc.(Float64Encoder).Put(values) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + dec := NewDecoder(parquet.Types.Double, parquet.Encodings.ALP, col, mem) + assert.Equal(t, parquet.Encodings.ALP, dec.Encoding()) + assert.Equal(t, parquet.Types.Double, dec.Type()) + + dec.SetData(len(values), buf.Bytes()) + out := make([]float64, len(values)) + n, err := dec.(Float64Decoder).Decode(out) + require.NoError(t, err) + assert.Equal(t, len(values), n) + + for i, v := range values { + assert.Equal(t, math.Float64bits(v), math.Float64bits(out[i]), + "index %d: mismatch", i) + } +} + +func TestALPPanicsForNonFloatTypes(t *testing.T) { + col := schema.NewColumn(schema.NewInt32Node("int32", parquet.Repetitions.Required, -1), 0, 0) + mem := memory.DefaultAllocator + + assert.Panics(t, func() { + NewEncoder(parquet.Types.Int32, parquet.Encodings.ALP, false, col, mem) + }, "ALP should panic for non-float encoder") + + assert.Panics(t, func() { + NewDecoder(parquet.Types.Int32, parquet.Encodings.ALP, col, mem) + }, "ALP should panic for non-float decoder") +} + +// alpDecodePage decodes a whole page, so that a test can check what an encoder +// wrote without going through a file. +func alpDecodePage[T alpFloat](t *testing.T, page Buffer, numValues int) []T { + t.Helper() + dec := newALPDecoder[T](format.Encoding_ALP, nil) + require.NoError(t, dec.SetData(numValues, page.Bytes())) + + out := make([]T, numValues) + n, err := dec.Decode(out) + require.NoError(t, err) + require.Equal(t, numValues, n) + return out +} + +// A column writes a new page every time its buffered size passes the page size, +// and calls FlushValues for each one without resetting the encoder in between. +// A page therefore has to hold its own values and no others. +func TestALPEncoderFlushValuesStartsANewPage(t *testing.T) { + enc := newALPEncoder[float64](format.Encoding_ALP, nil, memory.DefaultAllocator) + + first := make([]float64, 2*alpDefaultVectorSize) + for i := range first { + first[i] = float64(i) / 100 + } + enc.Put(first) + page, err := enc.FlushValues() + require.NoError(t, err) + got := alpDecodePage[float64](t, page, len(first)) + page.Release() + assert.Equal(t, first, got) + + second := []float64{9.5, 8.25, 7.125} + enc.Put(second) + assert.EqualValues(t, len(second)*8, enc.EstimatedDataEncodedSize(), + "the flushed page still counts towards the next one") + + page, err = enc.FlushValues() + require.NoError(t, err) + got = alpDecodePage[float64](t, page, len(second)) + page.Release() + assert.Equal(t, second, got) +} + +// The shortlist describes the column, so it outlives a page. A second page of +// the same values encodes to the same bytes as the first. +func TestALPEncoderKeepsPresetsAcrossPages(t *testing.T) { + enc := newALPEncoder[float64](format.Encoding_ALP, nil, memory.DefaultAllocator) + + values := make([]float64, 10*alpDefaultVectorSize) + for i := range values { + values[i] = float64(i%997) / 100 + } + + enc.Put(values) + first, err := enc.FlushValues() + require.NoError(t, err) + require.NotNil(t, enc.cachedPresets, "ten vectors is past the sampling threshold") + presets := slices.Clone(enc.cachedPresets) + + enc.Put(values) + second, err := enc.FlushValues() + require.NoError(t, err) + + assert.Equal(t, presets, enc.cachedPresets) + assert.Equal(t, first.Bytes(), second.Bytes()) + first.Release() + second.Release() +} + +// Reset returns the encoder to its initial state, shortlist included. +func TestALPEncoderResetClearsPresets(t *testing.T) { + enc := newALPEncoder[float64](format.Encoding_ALP, nil, memory.DefaultAllocator) + + values := make([]float64, 10*alpDefaultVectorSize) + for i := range values { + values[i] = float64(i) / 100 + } + enc.Put(values) + require.NotNil(t, enc.cachedPresets) + + enc.Reset() + assert.Nil(t, enc.cachedPresets) + assert.Empty(t, enc.presetCounts) + assert.Zero(t, enc.vectorsProcessed) + assert.Zero(t, enc.EstimatedDataEncodedSize()) +} diff --git a/parquet/internal/encoding/encoding_test.go b/parquet/internal/encoding/encoding_test.go index 1589c1332..b1e34cd23 100644 --- a/parquet/internal/encoding/encoding_test.go +++ b/parquet/internal/encoding/encoding_test.go @@ -449,6 +449,17 @@ func (b *BaseEncodingTestSuite) TestByteStreamSplitRoundTrip() { } } +func (b *BaseEncodingTestSuite) TestALPRoundTrip() { + b.initData(10000, 1) + + switch b.typ { + case reflect.TypeOf(float32(0)), reflect.TypeOf(float64(0)): + b.checkRoundTrip(parquet.Encodings.ALP) + default: + b.Panics(func() { b.checkRoundTrip(parquet.Encodings.ALP) }) + } +} + func (b *BaseEncodingTestSuite) TestSpacedRoundTrip() { exec := func(vals, repeats int, validBitsOffset int64, nullProb float64) { b.Run(fmt.Sprintf("%d vals %d repeats %d offset %0.3f null", vals, repeats, validBitsOffset, 1-nullProb), func() { diff --git a/parquet/internal/encoding/typed_encoder.go b/parquet/internal/encoding/typed_encoder.go index 735ee9252..0fc9aa851 100644 --- a/parquet/internal/encoding/typed_encoder.go +++ b/parquet/internal/encoding/typed_encoder.go @@ -133,6 +133,8 @@ func (floatEncodingTraits[T]) Encoder(e format.Encoding, useDict bool, descr *sc encoder: newEncoderBase(e, descr, mem), }, } + case format.Encoding_ALP: + return newALPEncoder[T](e, descr, mem) default: panic("unimplemented encoding for float types: " + e.String()) } @@ -286,6 +288,8 @@ func (floatDecoderTraits[T]) Decoder(e parquet.Encoding, descr *schema.Column, u return &PlainDecoder[T]{decoder: newDecoderBase(format.Encoding(e), descr)} case parquet.Encodings.ByteStreamSplit: return &ByteStreamSplitDecoder[T]{decoder: newDecoderBase(format.Encoding(e), descr)} + case parquet.Encodings.ALP: + return newALPDecoder[T](format.Encoding(e), descr) default: panic("unimplemented encoding for float types: " + e.String()) } diff --git a/parquet/internal/gen-go/parquet/parquet.go b/parquet/internal/gen-go/parquet/parquet.go index 54547a264..e6c06fc55 100644 --- a/parquet/internal/gen-go/parquet/parquet.go +++ b/parquet/internal/gen-go/parquet/parquet.go @@ -454,6 +454,7 @@ const ( Encoding_DELTA_BYTE_ARRAY Encoding = 7 Encoding_RLE_DICTIONARY Encoding = 8 Encoding_BYTE_STREAM_SPLIT Encoding = 9 + Encoding_ALP Encoding = 10 ) var knownEncodingValues = []Encoding{ @@ -466,6 +467,7 @@ var knownEncodingValues = []Encoding{ Encoding_DELTA_BYTE_ARRAY, Encoding_RLE_DICTIONARY, Encoding_BYTE_STREAM_SPLIT, + Encoding_ALP, } func EncodingValues() iter.Seq[Encoding] { @@ -489,6 +491,7 @@ func (p Encoding) String() string { case Encoding_DELTA_BYTE_ARRAY: return "DELTA_BYTE_ARRAY" case Encoding_RLE_DICTIONARY: return "RLE_DICTIONARY" case Encoding_BYTE_STREAM_SPLIT: return "BYTE_STREAM_SPLIT" + case Encoding_ALP: return "ALP" } return "" } @@ -504,6 +507,7 @@ func EncodingFromString(s string) (Encoding, error) { case "DELTA_BYTE_ARRAY": return Encoding_DELTA_BYTE_ARRAY, nil case "RLE_DICTIONARY": return Encoding_RLE_DICTIONARY, nil case "BYTE_STREAM_SPLIT": return Encoding_BYTE_STREAM_SPLIT, nil + case "ALP": return Encoding_ALP, nil } return Encoding(0), fmt.Errorf("not a valid Encoding string") } diff --git a/parquet/types.go b/parquet/types.go index c1bd4372a..ff47666a9 100644 --- a/parquet/types.go +++ b/parquet/types.go @@ -328,6 +328,8 @@ var ( DeltaBinaryPacked Encoding DeltaLengthByteArray Encoding ByteStreamSplit Encoding + // ALP encodes FLOAT and DOUBLE columns. + ALP Encoding }{ Plain: Encoding(format.Encoding_PLAIN), PlainDict: Encoding(format.Encoding_PLAIN_DICTIONARY), @@ -338,6 +340,7 @@ var ( DeltaBinaryPacked: Encoding(format.Encoding_DELTA_BINARY_PACKED), DeltaLengthByteArray: Encoding(format.Encoding_DELTA_LENGTH_BYTE_ARRAY), ByteStreamSplit: Encoding(format.Encoding_BYTE_STREAM_SPLIT), + ALP: Encoding(format.Encoding_ALP), } // ColumnOrders contains constants for the Column Ordering fields From 65dc2da6ba5d01d3da474f918d7537df3970e504 Mon Sep 17 00:00:00 2001 From: Prateek Gaur Date: Sun, 4 Oct 2026 00:12:30 +0000 Subject: [PATCH 03/11] Integrate ALP with Parquet readers and writers parquet-format marks ALP as a preview encoding, and a reader without it cannot open a file that uses it. Writing ALP therefore takes an explicit opt-in, WithALPEncoding for every column or WithALPEncodingFor for one, and NewWriterProperties panics when a column asks for ALP without it, as it does for any unusable encoding. The opt-in permits ALP but does not select it. A column still chooses it with WithEncoding or WithEncodingFor, and needs dictionary encoding off, because a writer that abandons its dictionary falls back to PLAIN. The column reader returns an error for an ALP page on a column that is not FLOAT or DOUBLE, where the decoder traits would otherwise panic. The tests write and read columns of several pages with and without nulls, under both data page versions. --- parquet/file/alp_test.go | 230 +++++++++++++++++++++++ parquet/file/column_reader.go | 8 + parquet/file/column_writer_test.go | 12 ++ parquet/file/file_writer_test.go | 218 +++++++++++++++++++++ parquet/reader_writer_properties_test.go | 53 ++++++ parquet/types.go | 3 +- parquet/writer_properties.go | 86 +++++++++ 7 files changed, 609 insertions(+), 1 deletion(-) create mode 100644 parquet/file/alp_test.go diff --git a/parquet/file/alp_test.go b/parquet/file/alp_test.go new file mode 100644 index 000000000..f06507b70 --- /dev/null +++ b/parquet/file/alp_test.go @@ -0,0 +1,230 @@ +// 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 file_test + +import ( + "bytes" + "math" + "math/rand" + "testing" + + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/compress" + "github.com/apache/arrow-go/v18/parquet/file" + "github.com/apache/arrow-go/v18/parquet/internal/encoding" + "github.com/apache/arrow-go/v18/parquet/schema" + "github.com/stretchr/testify/require" +) + +// alpWriteReadDouble writes one optional or required double column with ALP and +// reads it back. A nil defLevels writes a required column. +func alpWriteReadDouble(t *testing.T, values []float64, defLevels []int16, opts ...parquet.WriterProperty) ([]float64, []int16) { + t.Helper() + + rep := parquet.Repetitions.Required + if defLevels != nil { + rep = parquet.Repetitions.Optional + } + opts = append([]parquet.WriterProperty{ + parquet.WithEncoding(parquet.Encodings.ALP), + parquet.WithALPEncoding(true), + parquet.WithDictionaryDefault(false), + }, opts...) + + field, err := schema.NewPrimitiveNode("value", rep, parquet.Types.Double, -1, -1) + require.NoError(t, err) + sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, schema.FieldList{field}, 0) + require.NoError(t, err) + + sink := encoding.NewBufferWriter(0, memory.DefaultAllocator) + writer := file.NewParquetWriter(sink, sc, file.WithWriterProps(parquet.NewWriterProperties(opts...))) + rgw := writer.AppendRowGroup() + cw, err := rgw.NextColumn() + require.NoError(t, err) + _, err = cw.(*file.Float64ColumnChunkWriter).WriteBatch(values, defLevels, nil) + require.NoError(t, err) + require.NoError(t, cw.Close()) + require.NoError(t, rgw.Close()) + require.NoError(t, writer.Close()) + + reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes())) + require.NoError(t, err) + defer reader.Close() + + cr, err := reader.RowGroup(0).Column(0) + require.NoError(t, err) + r := cr.(*file.Float64ColumnChunkReader) + + rows := int(reader.NumRows()) + out := make([]float64, rows) + var outDef []int16 + if defLevels != nil { + outDef = make([]int16, rows) + } + for read := 0; read < rows; { + _, n, err := r.ReadBatch(int64(rows), out[read:], outDef, nil) + require.NoError(t, err) + require.NoError(t, r.Err()) + if n == 0 { + break + } + read += n + } + return out, outDef +} + +// An all-null page contains levels but no encoded values. PLAIN is the control. +func TestALPAllNullPage(t *testing.T) { + for _, enc := range []parquet.Encoding{parquet.Encodings.Plain, parquet.Encodings.ALP} { + t.Run(enc.String(), func(t *testing.T) { + const rows = 8 + out, def := alpWriteReadDouble(t, nil, make([]int16, rows), + parquet.WithEncoding(enc)) + require.Len(t, out, rows) + for i, d := range def { + require.Zero(t, d, "row %d should be null", i) + } + }) + } +} + +// Scattered nulls exercise PutSpaced and DecodeSpaced. +func TestALPNullsScattered(t *testing.T) { + const rows = 3000 + values := make([]float64, 0, rows) + defLevels := make([]int16, rows) + for i := range rows { + if i%3 == 0 { + continue // definition level 0: this row is null + } + defLevels[i] = 1 + values = append(values, float64(i)/100) + } + + got, _ := alpWriteReadDouble(t, values, defLevels) + require.Equal(t, values, got[:len(values)]) +} + +// These distributions cover exact decimals, integer overflow, NaNs, +// infinities, and subnormals. +func TestALPRandomRoundtrip(t *testing.T) { + rng := rand.New(rand.NewSource(7)) + pools := []struct { + name string + draw func() float64 + }{ + {"normal", func() float64 { return rng.NormFloat64() }}, + {"two decimal places", func() float64 { return float64(rng.Intn(1_000_000)) / 100 }}, + {"very large", func() float64 { return rng.NormFloat64() * 1e300 }}, + {"very small", func() float64 { return rng.NormFloat64() * 1e-300 }}, + {"random bits", func() float64 { return math.Float64frombits(rng.Uint64()) }}, + {"whole numbers", func() float64 { return float64(rng.Int63()) }}, + } + + for _, pool := range pools { + t.Run(pool.name, func(t *testing.T) { + values := make([]float64, 5000) + for i := range values { + values[i] = pool.draw() + } + + got, _ := alpWriteReadDouble(t, values, nil) + for i := range values { + // Compare the bit patterns, so that a NaN or a negative zero has + // to come back exactly as it went in. + if math.Float64bits(got[i]) != math.Float64bits(values[i]) { + t.Fatalf("index %d: got %v (%016x), want %v (%016x)", + i, got[i], math.Float64bits(got[i]), values[i], math.Float64bits(values[i])) + } + } + }) + } +} + +// Data page v2 stores levels separately from encoded values. +func TestALPDataPageV2(t *testing.T) { + values := []float64{1.25, 2.5, 3.75, 100, 0.001} + got, _ := alpWriteReadDouble(t, values, nil, parquet.WithDataPageVersion(parquet.DataPageV2)) + require.Equal(t, values, got) +} + +// alpSinglePageReader returns one page followed by the end of the chunk. +type alpSinglePageReader struct { + page file.Page + handed bool +} + +func (r *alpSinglePageReader) SetMaxPageHeaderSize(int) {} +func (r *alpSinglePageReader) Page() file.Page { return r.page } +func (r *alpSinglePageReader) Next() bool { + if r.handed { + return false + } + r.handed = true + return r.page != nil +} + +func (r *alpSinglePageReader) Err() error { return nil } +func (r *alpSinglePageReader) Reset(parquet.BufferedReader, int64, compress.Compression, *file.CryptoContext) { +} + +func (r *alpSinglePageReader) GetDictionaryPage() (*file.DictionaryPage, error) { return nil, nil } +func (r *alpSinglePageReader) SeekToPageWithRow(int64) error { return nil } +func (r *alpSinglePageReader) Close() error { return nil } + +// Reject ALP before non-floating-point decoder traits can panic. +func TestALPOnIntegerColumnIsRejected(t *testing.T) { + node, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Int32, -1, -1) + require.NoError(t, err) + descr := schema.NewColumn(node, 0, 0) + + body := memory.NewBufferBytes(make([]byte, 64)) + page := file.NewDataPageV1(body, 4, parquet.Encodings.ALP, + parquet.Encodings.RLE, parquet.Encodings.RLE, int32(body.Len())) + + r := file.NewColumnReader(descr, &alpSinglePageReader{page: page}, memory.DefaultAllocator, nil) + total, read, err := r.(*file.Int32ColumnChunkReader).ReadBatch(4, make([]int32, 4), nil, nil) + require.NoError(t, err) + require.Zero(t, total) + require.Zero(t, read) + // ReadBatch stops rather than returning the failure, as it does for any page + // it cannot open, so the reason is on the reader. + require.ErrorContains(t, r.Err(), "only float and double support ALP encoding") +} + +// A mismatched page count must return an error instead of a zero-length read, +// which the column reader would retry indefinitely. +func TestALPPageClaimingTooManyValues(t *testing.T) { + node, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Double, -1, -1) + require.NoError(t, err) + descr := schema.NewColumn(node, 0, 0) + + enc := encoding.NewEncoder(parquet.Types.Double, parquet.Encodings.ALP, false, descr, memory.DefaultAllocator) + enc.(encoding.Float64Encoder).Put([]float64{1.5, 2.5}) + buf, err := enc.FlushValues() + require.NoError(t, err) + + body := memory.NewBufferBytes(buf.Bytes()) + page := file.NewDataPageV1(body, 8, parquet.Encodings.ALP, + parquet.Encodings.RLE, parquet.Encodings.RLE, int32(body.Len())) + + r := file.NewColumnReader(descr, &alpSinglePageReader{page: page}, memory.DefaultAllocator, nil) + _, read, err := r.(*file.Float64ColumnChunkReader).ReadBatch(8, make([]float64, 8), nil, nil) + require.ErrorContains(t, err, "ALP page holds 2 values") + require.Equal(t, 2, read, "the values the page does hold still come back") +} diff --git a/parquet/file/column_reader.go b/parquet/file/column_reader.go index a296915dc..d85d7b61a 100644 --- a/parquet/file/column_reader.go +++ b/parquet/file/column_reader.go @@ -499,6 +499,14 @@ func (c *columnChunkReader) initDataDecoder(page DataPage, lvlByteLen int64) err format.Encoding_BYTE_STREAM_SPLIT: c.curDecoder = c.decoderTraits.Decoder(parquet.Encoding(encoding), c.descr, false, c.mem) c.decoders[encoding] = c.curDecoder + case format.Encoding_ALP: + // Reject before building a decoder: the traits for a non-float type + // have no ALP case and panic on one. + if pt := c.descr.PhysicalType(); pt != parquet.Types.Float && pt != parquet.Types.Double { + return fmt.Errorf("parquet: only float and double support ALP encoding, got %s", pt) + } + c.curDecoder = c.decoderTraits.Decoder(parquet.Encoding(encoding), c.descr, false, c.mem) + c.decoders[encoding] = c.curDecoder case format.Encoding_RLE_DICTIONARY: return errors.New("parquet: dictionary page must be before data page") default: diff --git a/parquet/file/column_writer_test.go b/parquet/file/column_writer_test.go index 37491a529..e0f293dff 100644 --- a/parquet/file/column_writer_test.go +++ b/parquet/file/column_writer_test.go @@ -358,6 +358,9 @@ func (p *PrimitiveWriterTestSuite) buildWriter(_ int64, columnProps parquet.Colu } else { opts = append(opts, parquet.WithDictionaryDefault(false), parquet.WithEncoding(columnProps.Encoding)) } + if columnProps.Encoding == parquet.Encodings.ALP { + opts = append(opts, parquet.WithALPEncoding(true)) + } opts = append(opts, parquet.WithMaxStatsSize(columnProps.MaxStatsSize), parquet.WithStats(columnProps.StatsEnabled)) p.props = parquet.NewWriterProperties(opts...) @@ -571,6 +574,15 @@ func (p *PrimitiveWriterTestSuite) TestRequiredByteStreamSplit() { } } +func (p *PrimitiveWriterTestSuite) TestRequiredALP() { + switch p.Typ { + case reflect.TypeOf(float32(0)), reflect.TypeOf(float64(0)): + p.testRequiredWithEncoding(parquet.Encodings.ALP) + default: + p.Panics(func() { p.testRequiredWithEncoding(parquet.Encodings.ALP) }) + } +} + func (p *PrimitiveWriterTestSuite) TestRequiredDictionary() { p.testRequiredWithEncoding(parquet.Encodings.PlainDict) } diff --git a/parquet/file/file_writer_test.go b/parquet/file/file_writer_test.go index 870807793..1a061b11a 100644 --- a/parquet/file/file_writer_test.go +++ b/parquet/file/file_writer_test.go @@ -611,6 +611,224 @@ func TestBatchedByteStreamSplitFileRoundtrip(t *testing.T) { require.NoError(t, rdr.Close()) } +func TestALPFileRoundtrip(t *testing.T) { + for _, tc := range []struct { + name string + values []float64 + }{ + { + name: "monetary", + values: []float64{1.23, 4.56, 7.89, 10.11, 12.13, 14.15, 16.17, 18.19}, + }, + { + name: "large_vector", + values: func() []float64 { + v := make([]float64, 2048) + for i := range v { + v[i] = 0.125 + float64(i)*0.25 + } + return v + }(), + }, + { + name: "integers", + values: []float64{1, 2, 3, 100, 200, 300, 1000, 2000}, + }, + { + name: "special_values", + values: []float64{0, math.Copysign(0, -1), 1.5, -1.5, math.Inf(1), math.Inf(-1), math.NaN()}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + props := parquet.NewWriterProperties( + parquet.WithEncoding(parquet.Encodings.ALP), + parquet.WithALPEncoding(true), + parquet.WithDictionaryDefault(false), + ) + + field, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Double, -1, -1) + require.NoError(t, err) + + sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, schema.FieldList{field}, 0) + require.NoError(t, err) + + sink := encoding.NewBufferWriter(0, memory.DefaultAllocator) + writer := file.NewParquetWriter(sink, sc, file.WithWriterProps(props)) + + rgw := writer.AppendRowGroup() + cw, err := rgw.NextColumn() + require.NoError(t, err) + + f64Writer, ok := cw.(*file.Float64ColumnChunkWriter) + require.True(t, ok) + + nVals, err := f64Writer.WriteBatch(tc.values, nil, nil) + require.NoError(t, err) + require.EqualValues(t, len(tc.values), nVals) + + require.NoError(t, cw.Close()) + require.NoError(t, rgw.Close()) + require.NoError(t, writer.Close()) + + reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes())) + require.NoError(t, err) + + require.Equal(t, 1, reader.NumRowGroups()) + require.EqualValues(t, len(tc.values), reader.NumRows()) + + rgr := reader.RowGroup(0) + cr, err := rgr.Column(0) + require.NoError(t, err) + + f64Reader, ok := cr.(*file.Float64ColumnChunkReader) + require.True(t, ok) + + output := make([]float64, len(tc.values)) + total, read, err := f64Reader.ReadBatch(int64(len(tc.values)), output, nil, nil) + require.NoError(t, err) + require.EqualValues(t, len(tc.values), total) + require.EqualValues(t, len(tc.values), read) + + // Compare the bit patterns, so that a NaN or a negative zero has to come + // back exactly as it went in. + for i := range tc.values { + assert.Equal(t, math.Float64bits(tc.values[i]), math.Float64bits(output[i]), + "index %d: got %v (bits %016x), want %v (bits %016x)", + i, output[i], math.Float64bits(output[i]), tc.values[i], math.Float64bits(tc.values[i])) + } + + require.NoError(t, reader.Close()) + }) + } +} + +// A column whose values outgrow the data page size is written as several pages. +// Each one carries its own ALP header, offset array and vectors, so a page that +// repeated an earlier page's values would decode to the wrong numbers here. +func TestALPFileRoundtripManyPages(t *testing.T) { + const numValues = 40000 + + values := make([]float64, numValues) + for i := range values { + values[i] = float64(i) / 100 + } + + props := parquet.NewWriterProperties( + parquet.WithEncoding(parquet.Encodings.ALP), + parquet.WithALPEncoding(true), + parquet.WithDictionaryDefault(false), + // Small enough that 1024-value vectors fill a page every few vectors. + parquet.WithDataPageSize(16*1024), + ) + + field, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Double, -1, -1) + require.NoError(t, err) + + sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, schema.FieldList{field}, 0) + require.NoError(t, err) + + sink := encoding.NewBufferWriter(0, memory.DefaultAllocator) + writer := file.NewParquetWriter(sink, sc, file.WithWriterProps(props)) + + rgw := writer.AppendRowGroup() + cw, err := rgw.NextColumn() + require.NoError(t, err) + + f64Writer, ok := cw.(*file.Float64ColumnChunkWriter) + require.True(t, ok) + + nVals, err := f64Writer.WriteBatch(values, nil, nil) + require.NoError(t, err) + require.EqualValues(t, numValues, nVals) + + require.NoError(t, cw.Close()) + require.NoError(t, rgw.Close()) + require.NoError(t, writer.Close()) + + reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes())) + require.NoError(t, err) + defer reader.Close() + + rgr := reader.RowGroup(0) + cr, err := rgr.Column(0) + require.NoError(t, err) + + f64Reader, ok := cr.(*file.Float64ColumnChunkReader) + require.True(t, ok) + + // Read in batches that do not divide the vector size, so a batch boundary + // lands inside a vector and inside a page. + output := make([]float64, numValues) + for read := 0; read < numValues; { + _, n, err := f64Reader.ReadBatch(700, output[read:], nil, nil) + require.NoError(t, err) + require.NotZero(t, n, "read stalled at %d of %d values", read, numValues) + read += int(n) + } + require.Equal(t, values, output) +} + +func TestALPFloat32FileRoundtrip(t *testing.T) { + values := []float32{1.23, 4.56, 7.89, 0, -1.5, 100.0, 0.001} + + props := parquet.NewWriterProperties( + parquet.WithEncoding(parquet.Encodings.ALP), + parquet.WithALPEncoding(true), + parquet.WithDictionaryDefault(false), + ) + + field, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Float, -1, -1) + require.NoError(t, err) + + sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, schema.FieldList{field}, 0) + require.NoError(t, err) + + sink := encoding.NewBufferWriter(0, memory.DefaultAllocator) + writer := file.NewParquetWriter(sink, sc, file.WithWriterProps(props)) + + rgw := writer.AppendRowGroup() + cw, err := rgw.NextColumn() + require.NoError(t, err) + + f32Writer, ok := cw.(*file.Float32ColumnChunkWriter) + require.True(t, ok) + + nVals, err := f32Writer.WriteBatch(values, nil, nil) + require.NoError(t, err) + require.EqualValues(t, len(values), nVals) + + require.NoError(t, cw.Close()) + require.NoError(t, rgw.Close()) + require.NoError(t, writer.Close()) + + reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes())) + require.NoError(t, err) + + require.Equal(t, 1, reader.NumRowGroups()) + require.EqualValues(t, len(values), reader.NumRows()) + + rgr := reader.RowGroup(0) + cr, err := rgr.Column(0) + require.NoError(t, err) + + f32Reader, ok := cr.(*file.Float32ColumnChunkReader) + require.True(t, ok) + + output := make([]float32, len(values)) + total, read, err := f32Reader.ReadBatch(int64(len(values)), output, nil, nil) + require.NoError(t, err) + require.EqualValues(t, len(values), total) + require.EqualValues(t, len(values), read) + + for i := range values { + assert.Equal(t, math.Float32bits(values[i]), math.Float32bits(output[i]), + "index %d: got %v (bits %08x), want %v (bits %08x)", + i, output[i], math.Float32bits(output[i]), values[i], math.Float32bits(values[i])) + } + + require.NoError(t, reader.Close()) +} + func TestLZ4RawFileRoundtrip(t *testing.T) { input := []int64{ -1, 0, 1, 2, 3, 4, 5, 123456789, -123456789, diff --git a/parquet/reader_writer_properties_test.go b/parquet/reader_writer_properties_test.go index ee8ac78f2..4729323cc 100644 --- a/parquet/reader_writer_properties_test.go +++ b/parquet/reader_writer_properties_test.go @@ -93,3 +93,56 @@ func TestReaderPropsGetStreamInsufficient(t *testing.T) { _, err := props.GetStream(rdr, 12, 15) assert.Error(t, err) } + +func TestWriterPropALPEncodingOptIn(t *testing.T) { + assert.False(t, parquet.NewWriterProperties().ALPEncodingEnabled(), + "ALP is a preview encoding, so it stays off by default") + + assert.Panics(t, func() { + parquet.NewWriterProperties(parquet.WithEncoding(parquet.Encodings.ALP)) + }, "the default encoding cannot be ALP without the flag") + assert.Panics(t, func() { + parquet.NewWriterProperties(parquet.WithEncodingFor("price", parquet.Encodings.ALP)) + }, "a column cannot ask for ALP without the flag") + + // The flag may be given before or after the encoding it allows. + for _, opts := range [][]parquet.WriterProperty{ + {parquet.WithALPEncoding(true), parquet.WithEncodingFor("price", parquet.Encodings.ALP)}, + {parquet.WithEncodingFor("price", parquet.Encodings.ALP), parquet.WithALPEncoding(true)}, + } { + props := parquet.NewWriterProperties(opts...) + assert.True(t, props.ALPEncodingEnabled()) + assert.Equal(t, parquet.Encodings.ALP, props.EncodingFor("price")) + assert.Equal(t, parquet.Encodings.Plain, props.EncodingFor("other"), + "the flag allows ALP without selecting it") + } +} + +func TestWriterPropALPEncodingPerColumn(t *testing.T) { + assert.Panics(t, func() { + parquet.NewWriterProperties( + parquet.WithALPEncodingFor("price", true), + parquet.WithEncodingFor("volume", parquet.Encodings.ALP)) + }, "permission for one column does not carry to another") + + props := parquet.NewWriterProperties( + parquet.WithALPEncodingFor("price", true), + parquet.WithEncodingFor("price", parquet.Encodings.ALP)) + assert.True(t, props.ALPEncodingEnabled(), + "one permitted column is enough to report the encoding as allowed") + assert.True(t, props.ALPEncodingEnabledFor("price")) + assert.False(t, props.ALPEncodingEnabledFor("volume")) + + // A column may also be taken out of a permission given to every column. + props = parquet.NewWriterProperties( + parquet.WithALPEncoding(true), + parquet.WithALPEncodingFor("volume", false)) + assert.True(t, props.ALPEncodingEnabledFor("price")) + assert.False(t, props.ALPEncodingEnabledFor("volume")) + assert.Panics(t, func() { + parquet.NewWriterProperties( + parquet.WithALPEncoding(true), + parquet.WithALPEncodingFor("volume", false), + parquet.WithEncodingFor("volume", parquet.Encodings.ALP)) + }, "the column-specific setting wins over the default one") +} diff --git a/parquet/types.go b/parquet/types.go index ff47666a9..135de94ba 100644 --- a/parquet/types.go +++ b/parquet/types.go @@ -328,7 +328,8 @@ var ( DeltaBinaryPacked Encoding DeltaLengthByteArray Encoding ByteStreamSplit Encoding - // ALP encodes FLOAT and DOUBLE columns. + // ALP encodes FLOAT and DOUBLE columns. parquet-format marks it as a + // preview feature, so writing it requires WithALPEncoding. ALP Encoding }{ Plain: Encoding(format.Encoding_PLAIN), diff --git a/parquet/writer_properties.go b/parquet/writer_properties.go index 9cedc8737..205da8644 100644 --- a/parquet/writer_properties.go +++ b/parquet/writer_properties.go @@ -61,6 +61,9 @@ const ( DefaultAdaptiveBloomFilterEnabled = false DefaultBloomFilterCandidates = 5 + // DefaultALPEnabled is false because ALP is a preview encoding. + DefaultALPEnabled = false + minimumBloomFilterBytes = 32 maximumBloomFilterBytes = 128 * 1024 * 1024 ) @@ -84,6 +87,8 @@ type ColumnProperties struct { AdaptiveBloomFilterEnabled bool BloomFilterCandidates int BloomFilterNDV int64 + // ALPEnabled permits this column to use the preview ALP encoding. + ALPEnabled bool } // DefaultColumnProperties returns the default properties which get utilized for writing. @@ -114,6 +119,7 @@ func DefaultColumnProperties() ColumnProperties { BloomFilterFPP: DefaultBloomFilterFPP, AdaptiveBloomFilterEnabled: DefaultAdaptiveBloomFilterEnabled, BloomFilterCandidates: DefaultBloomFilterCandidates, + ALPEnabled: DefaultALPEnabled, } } @@ -134,6 +140,7 @@ type writerPropConfig struct { bloomFilterEnabled map[string]bool adaptiveBloomFilterEnabled map[string]bool numBloomFilterCandidates map[string]int + alpEnabled map[string]bool } // WriterProperty is used as the options for building a writer properties instance @@ -283,6 +290,39 @@ func WithEncodingPath(path ColumnPath, encoding Encoding) WriterProperty { return WithEncodingFor(path.String(), encoding) } +// WithALPEncoding permits all columns to use the preview ALP encoding. Readers +// without ALP support cannot read files that use it, so ALP is disabled by +// default. +// +// The option grants permission rather than selecting ALP. A column still asks for +// the encoding through WithEncoding or WithEncodingFor, and asking for it without +// this option panics, as asking for any unusable encoding does. +// +// A column reaches ALP only with dictionary encoding turned off. A dictionary +// takes precedence over the encoding a column asks for, and a writer that drops +// the dictionary falls back to PLAIN rather than to ALP, so pair the request with +// WithDictionaryDefault(false) or WithDictionaryFor(path, false). +// +// WithALPEncodingFor overrides this setting for one column. +func WithALPEncoding(enabled bool) WriterProperty { + return func(cfg *writerPropConfig) { + cfg.wr.defColumnProps.ALPEnabled = enabled + } +} + +// WithALPEncodingFor permits or rejects ALP for one column. See WithALPEncoding +// for the requirements that apply when ALP is enabled. +func WithALPEncodingFor(path string, enabled bool) WriterProperty { + return func(cfg *writerPropConfig) { + cfg.alpEnabled[path] = enabled + } +} + +// WithALPEncodingPath is the same as WithALPEncodingFor but takes a ColumnPath. +func WithALPEncodingPath(path ColumnPath, enabled bool) WriterProperty { + return WithALPEncodingFor(path.String(), enabled) +} + // WithCompression specifies the default compression type to use for column writing. func WithCompression(codec compress.Compression) WriterProperty { return func(cfg *writerPropConfig) { @@ -597,6 +637,7 @@ func NewWriterProperties(opts ...WriterProperty) *WriterProperties { bloomFilterEnabled: make(map[string]bool), adaptiveBloomFilterEnabled: make(map[string]bool), numBloomFilterCandidates: make(map[string]int), + alpEnabled: make(map[string]bool), } for _, o := range opts { o(&cfg) @@ -660,9 +701,54 @@ func NewWriterProperties(opts ...WriterProperty) *WriterProperties { get(key).BloomFilterCandidates = value } + for key, value := range cfg.alpEnabled { + get(key).ALPEnabled = 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)") + } + 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)") + } + } + return cfg.wr } +// ALPEncodingEnabled reports whether any column may be written with the ALP +// encoding. +func (w *WriterProperties) ALPEncodingEnabled() bool { + if w.defColumnProps.ALPEnabled { + return true + } + for _, props := range w.columnProps { + if props.ALPEnabled { + return true + } + } + return false +} + +// 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()) +} + // FileEncryptionProperties returns the current encryption properties that were // used to create the writer properties. func (w *WriterProperties) FileEncryptionProperties() *FileEncryptionProperties { From 42d9c3b6cb9329f7edbe807036560b142e052ff1 Mon Sep 17 00:00:00 2001 From: Prateek Gaur Date: Sun, 4 Oct 2026 00:13:04 +0000 Subject: [PATCH 04/11] Test ALP on a shared fixture and through Arrow A round trip through one implementation cannot show that it reads what other writers produce. This moves parquet-testing to the commit that adds alp_extended.zstd.parquet, whose ALP columns use vectors of 32, 1024 and 4096 values next to PLAIN copies of the same data, and checks that every ALP value decodes to the bits of its PLAIN copy. A second test writes an Arrow table with ALP columns and reads it back across compression codecs, several pages and row groups, nulls, and values that can only be stored as exceptions. --- parquet-testing | 2 +- parquet/file/file_reader_test.go | 146 +++++++++++++++++++++ parquet/pqarrow/alp_test.go | 216 +++++++++++++++++++++++++++++++ 3 files changed, 363 insertions(+), 1 deletion(-) create mode 100644 parquet/pqarrow/alp_test.go diff --git a/parquet-testing b/parquet-testing index 92d45b075..09f3cdbde 160000 --- a/parquet-testing +++ b/parquet-testing @@ -1 +1 @@ -Subproject commit 92d45b0752487a4b55fb7f1581c8126ee3e73b0d +Subproject commit 09f3cdbde45302f0f0c689c950e465e98a9df960 diff --git a/parquet/file/file_reader_test.go b/parquet/file/file_reader_test.go index 9fa38c3cf..4d81dce56 100644 --- a/parquet/file/file_reader_test.go +++ b/parquet/file/file_reader_test.go @@ -24,6 +24,7 @@ import ( "encoding/csv" "fmt" "io" + "math" "os" "path" "testing" @@ -1043,3 +1044,148 @@ func TestListColumns(t *testing.T) { } } } + +// alpConformanceFile contains matching ALP and PLAIN columns published by +// apache/parquet-testing. Its ALP columns use vectors of 32, 1024, and 4096 +// values, so decoding cannot assume the writer's default vector size. +// +// The fixture independently covers NaN payloads, infinities, negative zero, +// subnormals, nulls, constant and all-exception vectors, and full-width frames. +// See parquet-testing/data/README.md. +const alpConformanceFile = "alp_extended.zstd.parquet" + +// floatBitsValue keeps NaN payloads and distinguishes negative zero from zero. +type floatBitsValue struct { + null bool + bits uint64 +} + +func TestReadALPConformanceFile(t *testing.T) { + dir := os.Getenv("PARQUET_TEST_DATA") + if dir == "" { + dir = "../../parquet-testing/data" + t.Log("PARQUET_TEST_DATA not set, using ../../parquet-testing/data") + } + require.DirExists(t, dir) + + props := parquet.NewReaderProperties(memory.DefaultAllocator) + rdr, err := file.OpenParquetFile(path.Join(dir, alpConformanceFile), false, + file.WithReadProps(props)) + require.NoError(t, err) + defer rdr.Close() + + const numRows = 9032 + require.EqualValues(t, numRows, rdr.MetaData().GetNumRows()) + + for _, tc := range []struct { + alp, plain string + }{ + {"float_alp_1024", "float_plain"}, + {"float_alp_4096", "float_plain"}, + {"float_alp_32", "float_plain"}, + {"double_alp_1024", "double_plain"}, + {"double_alp_4096", "double_plain"}, + {"double_alp_32", "double_plain"}, + } { + t.Run(tc.alp, func(t *testing.T) { + alpCol := alpColumnIndex(t, rdr, tc.alp) + + // A decoder that fell back to another encoding would satisfy the + // comparison below, so check that the pages really are ALP. + for rg := range rdr.NumRowGroups() { + chunk, err := rdr.RowGroup(rg).MetaData().ColumnChunk(alpCol) + require.NoError(t, err) + assert.Contains(t, chunk.Encodings(), parquet.Encodings.ALP, + "row group %d of %s should be ALP encoded", rg, tc.alp) + } + + want := readFloatColumn(t, rdr, alpColumnIndex(t, rdr, tc.plain)) + got := readFloatColumn(t, rdr, alpCol) + require.Len(t, want, numRows) + require.Len(t, got, numRows) + + for i := range want { + if want[i] != got[i] { + t.Fatalf("row %d: got %s, want %s", i, got[i], want[i]) + } + } + }) + } +} + +func (v floatBitsValue) String() string { + if v.null { + return "null" + } + return fmt.Sprintf("%#016x", v.bits) +} + +// alpColumnIndex returns the index of the named leaf column. +func alpColumnIndex(t *testing.T, rdr *file.Reader, name string) int { + t.Helper() + + schema := rdr.MetaData().Schema + for c := range schema.NumColumns() { + if schema.Column(c).Name() == name { + return c + } + } + t.Fatalf("no column named %s", name) + return -1 +} + +// readFloatColumn reads one column of every row group, one entry per row. Float +// bit patterns are widened to 64 bits so that both physical types compare the +// same way. +func readFloatColumn(t *testing.T, rdr *file.Reader, col int) []floatBitsValue { + t.Helper() + + maxDef := rdr.MetaData().Schema.Column(col).MaxDefinitionLevel() + var values []floatBitsValue + for rg := range rdr.NumRowGroups() { + cr, err := rdr.RowGroup(rg).Column(col) + require.NoError(t, err) + + switch r := cr.(type) { + case *file.Float32ColumnChunkReader: + values = appendFloatValues(t, values, maxDef, r.ReadBatch, + func(v float32) uint64 { return uint64(math.Float32bits(v)) }) + case *file.Float64ColumnChunkReader: + values = appendFloatValues(t, values, maxDef, r.ReadBatch, math.Float64bits) + default: + t.Fatalf("column %d: unexpected reader %T", col, cr) + } + } + return values +} + +// appendFloatValues drains one column chunk, which ReadBatch may return one page +// at a time. Nulls take a definition level but not a value, so the two advance +// independently. +func appendFloatValues[T float32 | float64](t *testing.T, values []floatBitsValue, maxDef int16, + readBatch func(int64, []T, []int16, []int16) (int64, int, error), + bits func(T) uint64) []floatBitsValue { + t.Helper() + + const batchSize = 1024 + batch := make([]T, batchSize) + defLvls := make([]int16, batchSize) + for { + total, valuesRead, err := readBatch(batchSize, batch, defLvls, nil) + require.NoError(t, err) + if total == 0 { + return values + } + + taken := 0 + for i := range int(total) { + if defLvls[i] < maxDef { + values = append(values, floatBitsValue{null: true}) + continue + } + values = append(values, floatBitsValue{bits: bits(batch[taken])}) + taken++ + } + require.Equal(t, valuesRead, taken, "read %d values for %d non-null levels", valuesRead, taken) + } +} diff --git a/parquet/pqarrow/alp_test.go b/parquet/pqarrow/alp_test.go new file mode 100644 index 000000000..e8ac8ffa8 --- /dev/null +++ b/parquet/pqarrow/alp_test.go @@ -0,0 +1,216 @@ +// 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 pqarrow_test + +import ( + "bytes" + "context" + "math" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/compress" + "github.com/apache/arrow-go/v18/parquet/file" + "github.com/apache/arrow-go/v18/parquet/pqarrow" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// This spans several vectors and three row groups in the test below. +const alpNumRows = 5000 + +// alpColumns mixes ordinary decimals, exceptions, and nulls. Nulls consume a +// definition level without consuming an encoded value. +func alpColumns() (floats []float32, doubles []float64, valid []bool) { + floats = make([]float32, alpNumRows) + doubles = make([]float64, alpNumRows) + valid = make([]bool, alpNumRows) + + for i := range floats { + floats[i] = float32(i%1000) / 100 + doubles[i] = float64(i%100000) / 1000 + valid[i] = i%100 != 7 + } + // A vector of one repeated value packs at a bit width of zero, and a vector + // of full-mantissa values is nothing but exceptions. + for i := 1024; i < 2048; i++ { + floats[i], doubles[i] = 7.77, 7.77 + } + for i := 2048; i < 3072; i++ { + floats[i] = float32(math.Pi) * float32(i) + doubles[i] = math.Pi * float64(i) + } + // Values that cannot round-trip through ALP's decimal representation. The + // NaNs carry payloads that must return bit for bit. + for _, s := range []struct { + row int + bits uint32 + dbl uint64 + }{ + {10, 0x7FC00000, 0x7FF8000000000000}, // quiet NaN + {11, 0x7FC0DEAD, 0x7FF800DEADBEEF00}, // quiet NaN with a payload + {12, 0xFFC00001, 0xFFF8000000000001}, // negative quiet NaN + {13, 0x7F800000, 0x7FF0000000000000}, // +Inf + {14, 0xFF800000, 0xFFF0000000000000}, // -Inf + {15, 0x80000000, 0x8000000000000000}, // -0.0 + {16, 0x00000001, 0x0000000000000001}, // subnormal + {17, 0x7F7FFFFF, 0x7FEFFFFFFFFFFFFF}, // largest finite + } { + floats[s.row] = math.Float32frombits(s.bits) + doubles[s.row] = math.Float64frombits(s.dbl) + valid[s.row] = true + } + return floats, doubles, valid +} + +// alpTable holds the columns as an Arrow table, ready to write. +func alpTable(t *testing.T, mem memory.Allocator) arrow.Table { + t.Helper() + + floats, doubles, valid := alpColumns() + schema := arrow.NewSchema([]arrow.Field{ + {Name: "f", Type: arrow.PrimitiveTypes.Float32, Nullable: true}, + {Name: "d", Type: arrow.PrimitiveTypes.Float64, Nullable: true}, + }, nil) + + fb := array.NewFloat32Builder(mem) + defer fb.Release() + fb.AppendValues(floats, valid) + farr := fb.NewArray() + defer farr.Release() + + db := array.NewFloat64Builder(mem) + defer db.Release() + db.AppendValues(doubles, valid) + darr := db.NewArray() + defer darr.Release() + + rec := array.NewRecordBatch(schema, []arrow.Array{farr, darr}, alpNumRows) + defer rec.Release() + return array.NewTableFromRecords(schema, []arrow.RecordBatch{rec}) +} + +// TestALPArrowRoundtrip covers the application path through the Arrow writer +// and reader, including compression, pages, row groups, nulls, and exceptions. +func TestALPArrowRoundtrip(t *testing.T) { + for _, codec := range []compress.Compression{ + compress.Codecs.Uncompressed, + compress.Codecs.Snappy, + compress.Codecs.Zstd, + } { + t.Run(codec.String(), func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + tbl := alpTable(t, mem) + defer tbl.Release() + + props := parquet.NewWriterProperties( + parquet.WithAllocator(mem), + parquet.WithEncoding(parquet.Encodings.ALP), + // A dictionary would win before the encoding is ever asked for. + parquet.WithDictionaryDefault(false), + parquet.WithALPEncoding(true), + parquet.WithCompression(codec), + // Small pages, so a column chunk holds more than one of them. + parquet.WithDataPageSize(4096), + ) + + var buf bytes.Buffer + // The chunk size leaves three row groups, the last of them short. + require.NoError(t, pqarrow.WriteTable(tbl, &buf, 2048, props, pqarrow.DefaultWriterProps())) + + rdr, err := file.NewParquetReader(bytes.NewReader(buf.Bytes()), + file.WithReadProps(parquet.NewReaderProperties(mem))) + require.NoError(t, err) + defer rdr.Close() + + require.Equal(t, 3, rdr.NumRowGroups()) + for rg := range rdr.NumRowGroups() { + for c := range rdr.MetaData().Schema.NumColumns() { + chunk, err := rdr.RowGroup(rg).MetaData().ColumnChunk(c) + require.NoError(t, err) + assert.Contains(t, chunk.Encodings(), parquet.Encodings.ALP, + "row group %d column %d should be ALP encoded", rg, c) + } + } + + fr, err := pqarrow.NewFileReader(rdr, pqarrow.ArrowReadProperties{BatchSize: 137}, mem) + require.NoError(t, err) + got, err := fr.ReadTable(context.Background()) + require.NoError(t, err) + defer got.Release() + + require.EqualValues(t, alpNumRows, got.NumRows()) + floats, doubles, valid := alpColumns() + assertALPFloat32Column(t, got.Column(0), floats, valid) + assertALPFloat64Column(t, got.Column(1), doubles, valid) + }) + } +} + +// assertALPFloat32Column compares a column against the values written, on bit +// patterns rather than on values: a NaN does not equal itself, and -0.0 equals +// 0.0 although the two are different values. +func assertALPFloat32Column(t *testing.T, col *arrow.Column, want []float32, valid []bool) { + t.Helper() + + row := 0 + for _, chunk := range col.Data().Chunks() { + values := chunk.(*array.Float32) + for i := range values.Len() { + if !valid[row] { + require.Truef(t, values.IsNull(i), "row %d should be null", row) + row++ + continue + } + require.Falsef(t, values.IsNull(i), "row %d should hold a value", row) + if math.Float32bits(values.Value(i)) != math.Float32bits(want[row]) { + t.Fatalf("row %d: got %#08x, want %#08x", + row, math.Float32bits(values.Value(i)), math.Float32bits(want[row])) + } + row++ + } + } + require.Equal(t, alpNumRows, row) +} + +func assertALPFloat64Column(t *testing.T, col *arrow.Column, want []float64, valid []bool) { + t.Helper() + + row := 0 + for _, chunk := range col.Data().Chunks() { + values := chunk.(*array.Float64) + for i := range values.Len() { + if !valid[row] { + require.Truef(t, values.IsNull(i), "row %d should be null", row) + row++ + continue + } + require.Falsef(t, values.IsNull(i), "row %d should hold a value", row) + if math.Float64bits(values.Value(i)) != math.Float64bits(want[row]) { + t.Fatalf("row %d: got %#016x, want %#016x", + row, math.Float64bits(values.Value(i)), math.Float64bits(want[row])) + } + row++ + } + } + require.Equal(t, alpNumRows, row) +} From a017cacb210a849bc5e3a530ae0292750c889058 Mon Sep 17 00:00:00 2001 From: Prateek Gaur Date: Sun, 4 Oct 2026 00:13:55 +0000 Subject: [PATCH 05/11] Benchmark ALP encoding and decoding The exception rate drives what ALP costs, so the benchmarks encode and decode FLOAT and DOUBLE columns at several sizes with no exceptions, one exception in ten values, and nothing but exceptions. A batched DOUBLE decode reads 137 values at a time, so batches end partway through a vector the way column reader calls do. --- .../encoding/alp_encoding_benchmark_test.go | 185 ++++++++++++++++++ 1 file changed, 185 insertions(+) create mode 100644 parquet/internal/encoding/alp_encoding_benchmark_test.go diff --git a/parquet/internal/encoding/alp_encoding_benchmark_test.go b/parquet/internal/encoding/alp_encoding_benchmark_test.go new file mode 100644 index 000000000..a138b5f40 --- /dev/null +++ b/parquet/internal/encoding/alp_encoding_benchmark_test.go @@ -0,0 +1,185 @@ +// 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_test + +import ( + "fmt" + "math" + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/internal/encoding" +) + +// ALP performance depends on its exception rate. The benchmarks cover no +// exceptions, one exception per ten values, and all exceptions. +var alpBenchmarkCases = []struct { + name string + exceptionInterval int +}{ + {"no exceptions", 0}, + {"10% exceptions", 10}, + {"all exceptions", 1}, +} + +func alpBenchmarkValues[T float32 | float64](size, exceptionInterval int) []T { + values := make([]T, size) + for i := range values { + values[i] = T(i%100000) / 1000 + if exceptionInterval != 0 && i%exceptionInterval == 0 { + values[i] = T(math.Pi) * T(i+1) + } + } + return values +} + +func BenchmarkALPEncodingFloat32(b *testing.B) { + for _, benchmark := range alpBenchmarkCases { + b.Run(benchmark.name, func(b *testing.B) { + for sz := MINSIZE; sz < MAXSIZE+1; sz *= 2 { + b.Run(fmt.Sprintf("len %d", sz), func(b *testing.B) { + values := alpBenchmarkValues[float32](sz, benchmark.exceptionInterval) + encoder := encoding.NewEncoder(parquet.Types.Float, parquet.Encodings.ALP, + false, nil, memory.DefaultAllocator).(encoding.Float32Encoder) + b.ResetTimer() + b.SetBytes(int64(len(values) * arrow.Float32SizeBytes)) + for b.Loop() { + encoder.Put(values) + buf, _ := encoder.FlushValues() + buf.Release() + } + }) + } + }) + } +} + +func BenchmarkALPEncodingFloat64(b *testing.B) { + for _, benchmark := range alpBenchmarkCases { + b.Run(benchmark.name, func(b *testing.B) { + for sz := MINSIZE; sz < MAXSIZE+1; sz *= 2 { + b.Run(fmt.Sprintf("len %d", sz), func(b *testing.B) { + values := alpBenchmarkValues[float64](sz, benchmark.exceptionInterval) + encoder := encoding.NewEncoder(parquet.Types.Double, parquet.Encodings.ALP, + false, nil, memory.DefaultAllocator).(encoding.Float64Encoder) + b.ResetTimer() + b.SetBytes(int64(len(values) * arrow.Float64SizeBytes)) + for b.Loop() { + encoder.Put(values) + buf, _ := encoder.FlushValues() + buf.Release() + } + }) + } + }) + } +} + +func BenchmarkALPDecodingFloat32(b *testing.B) { + for _, benchmark := range alpBenchmarkCases { + b.Run(benchmark.name, func(b *testing.B) { + for sz := MINSIZE; sz < MAXSIZE+1; sz *= 2 { + b.Run(fmt.Sprintf("len %d", sz), func(b *testing.B) { + values := alpBenchmarkValues[float32](sz, benchmark.exceptionInterval) + output := make([]float32, sz) + encoder := encoding.NewEncoder(parquet.Types.Float, parquet.Encodings.ALP, + false, nil, memory.DefaultAllocator).(encoding.Float32Encoder) + encoder.Put(values) + buf, _ := encoder.FlushValues() + defer buf.Release() + + decoder := encoding.NewDecoder(parquet.Types.Float, parquet.Encodings.ALP, + nil, memory.DefaultAllocator).(encoding.Float32Decoder) + b.ResetTimer() + b.SetBytes(int64(len(values) * arrow.Float32SizeBytes)) + for b.Loop() { + if err := decoder.SetData(sz, buf.Bytes()); err != nil { + b.Fatal(err) + } + if _, err := decoder.Decode(output); err != nil { + b.Fatal(err) + } + } + }) + } + }) + } +} + +func BenchmarkALPDecodingFloat64(b *testing.B) { + for _, benchmark := range alpBenchmarkCases { + b.Run(benchmark.name, func(b *testing.B) { + for sz := MINSIZE; sz < MAXSIZE+1; sz *= 2 { + b.Run(fmt.Sprintf("len %d", sz), func(b *testing.B) { + values := alpBenchmarkValues[float64](sz, benchmark.exceptionInterval) + output := make([]float64, sz) + encoder := encoding.NewEncoder(parquet.Types.Double, parquet.Encodings.ALP, + false, nil, memory.DefaultAllocator).(encoding.Float64Encoder) + encoder.Put(values) + buf, _ := encoder.FlushValues() + defer buf.Release() + + decoder := encoding.NewDecoder(parquet.Types.Double, parquet.Encodings.ALP, + nil, memory.DefaultAllocator).(encoding.Float64Decoder) + b.ResetTimer() + b.SetBytes(int64(len(values) * arrow.Float64SizeBytes)) + for b.Loop() { + if err := decoder.SetData(sz, buf.Bytes()); err != nil { + b.Fatal(err) + } + if _, err := decoder.Decode(output); err != nil { + b.Fatal(err) + } + } + }) + } + }) + } +} + +// Batched decoding exercises calls that stop partway through a vector. +func BenchmarkALPDecodingFloat64Batched(b *testing.B) { + const batchSize = 137 + for sz := MINSIZE; sz < MAXSIZE+1; sz *= 2 { + b.Run(fmt.Sprintf("len %d", sz), func(b *testing.B) { + values := alpBenchmarkValues[float64](sz, 10) + output := make([]float64, batchSize) + encoder := encoding.NewEncoder(parquet.Types.Double, parquet.Encodings.ALP, + false, nil, memory.DefaultAllocator).(encoding.Float64Encoder) + encoder.Put(values) + buf, _ := encoder.FlushValues() + defer buf.Release() + + decoder := encoding.NewDecoder(parquet.Types.Double, parquet.Encodings.ALP, + nil, memory.DefaultAllocator).(encoding.Float64Decoder) + b.ResetTimer() + b.SetBytes(int64(len(values) * arrow.Float64SizeBytes)) + for b.Loop() { + if err := decoder.SetData(sz, buf.Bytes()); err != nil { + b.Fatal(err) + } + for left := sz; left > 0; left -= batchSize { + if _, err := decoder.Decode(output[:min(left, batchSize)]); err != nil { + b.Fatal(err) + } + } + } + }) + } +} From 4fd6624e299baf6b1d2d82d4e71a0e341a23f282 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Thu, 29 Oct 2026 12:27:15 -0400 Subject: [PATCH 06/11] perf(parquet): hoist ALP scaling constants and add a vector size option - Look the power-of-ten constants up once per vector (alpScaling) instead of once per value; encode is ~35-40% and decode ~30-44% faster. - Add the frame of reference back in the encoded integer's width, as the spec requires for FLOAT. - Reuse scratch for PutSpaced and the page header. - Add WithALPVectorSize/WithALPVectorSizeFor/WithALPVectorSizePath (power of two, 8-32768, default 1024) and apply it in the float column writers. - Document ALP in the package docs. --- parquet/doc.go | 6 ++ parquet/file/column_writer_types.gen.go | 12 +++ parquet/file/column_writer_types.gen.go.tmpl | 8 ++ parquet/file/file_writer_test.go | 58 +++++++++++++ parquet/internal/encoding/alp_encoding.go | 87 +++++++++++++++---- .../internal/encoding/alp_encoding_test.go | 52 +++++++++++ parquet/internal/encoding/alp_params.go | 3 +- parquet/internal/encoding/alp_params_test.go | 3 +- parquet/internal/encoding/alp_scaling.go | 78 ++++++++++------- parquet/internal/encoding/alp_scaling_test.go | 34 +++++--- parquet/reader_writer_properties_test.go | 22 +++++ parquet/writer_properties.go | 71 +++++++++++++++ 12 files changed, 368 insertions(+), 66 deletions(-) diff --git a/parquet/doc.go b/parquet/doc.go index edccc3c0f..13dbdb7d6 100644 --- a/parquet/doc.go +++ b/parquet/doc.go @@ -70,6 +70,12 @@ // // - Byte Stream Split (Float, Double, Int32, Int64, FixedLenByteArray) // +// - ALP, Adaptive Lossless floating-Point (Float, Double). The Parquet +// format marks ALP as a preview encoding, so writing it is opt-in: enable +// it with WithALPEncoding, select it with WithEncoding or WithEncodingFor, +// and turn off dictionary encoding for those columns. Reading needs no +// option. +// // Tip: Some platforms don't necessarily support all kinds of encodings. If you're not // sure what to use, just use Plain and Dictionary encoding. // diff --git a/parquet/file/column_writer_types.gen.go b/parquet/file/column_writer_types.gen.go index 906ecdbe9..f766fb09b 100644 --- a/parquet/file/column_writer_types.gen.go +++ b/parquet/file/column_writer_types.gen.go @@ -858,6 +858,12 @@ type Float32ColumnChunkWriter struct { func NewFloat32ColumnChunkWriter(meta *metadata.ColumnChunkMetaDataBuilder, pager PageWriter, useDict bool, enc parquet.Encoding, props *parquet.WriterProperties) *Float32ColumnChunkWriter { ret := &Float32ColumnChunkWriter{columnWriter: newColumnWriterBase(meta, pager, useDict, enc, props)} ret.currentEncoder = encoding.Float32EncoderTraits.Encoder(format.Encoding(enc), useDict, meta.Descr(), props.Allocator()) + if alp, ok := ret.currentEncoder.(encoding.ALPEncoder); ok { + // The writer properties have already validated the size. + if err := alp.SetVectorSize(props.ALPVectorSizeFor(meta.Descr().Path())); err != nil { + panic(err) + } + } if useDict && ret.bloomFilter != nil { ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking() } @@ -1129,6 +1135,12 @@ type Float64ColumnChunkWriter struct { func NewFloat64ColumnChunkWriter(meta *metadata.ColumnChunkMetaDataBuilder, pager PageWriter, useDict bool, enc parquet.Encoding, props *parquet.WriterProperties) *Float64ColumnChunkWriter { ret := &Float64ColumnChunkWriter{columnWriter: newColumnWriterBase(meta, pager, useDict, enc, props)} ret.currentEncoder = encoding.Float64EncoderTraits.Encoder(format.Encoding(enc), useDict, meta.Descr(), props.Allocator()) + if alp, ok := ret.currentEncoder.(encoding.ALPEncoder); ok { + // The writer properties have already validated the size. + if err := alp.SetVectorSize(props.ALPVectorSizeFor(meta.Descr().Path())); err != nil { + panic(err) + } + } if useDict && ret.bloomFilter != nil { ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking() } diff --git a/parquet/file/column_writer_types.gen.go.tmpl b/parquet/file/column_writer_types.gen.go.tmpl index c9e1ad7ab..c2b5fbdb4 100644 --- a/parquet/file/column_writer_types.gen.go.tmpl +++ b/parquet/file/column_writer_types.gen.go.tmpl @@ -49,6 +49,14 @@ func New{{.Name}}ColumnChunkWriter(meta *metadata.ColumnChunkMetaDataBuilder, pa {{- end}} ret := &{{.Name}}ColumnChunkWriter{columnWriter: newColumnWriterBase(meta, pager, useDict, enc, props)} ret.currentEncoder = encoding.{{.Name}}EncoderTraits.Encoder(format.Encoding(enc), useDict, meta.Descr(), props.Allocator()) +{{- if or (eq .Name "Float32") (eq .Name "Float64")}} + if alp, ok := ret.currentEncoder.(encoding.ALPEncoder); ok { + // The writer properties have already validated the size. + if err := alp.SetVectorSize(props.ALPVectorSizeFor(meta.Descr().Path())); err != nil { + panic(err) + } + } +{{- end}} if useDict && ret.bloomFilter != nil { ret.currentEncoder.(encoding.DictEncoder).EnableDictionaryReferenceTracking() } diff --git a/parquet/file/file_writer_test.go b/parquet/file/file_writer_test.go index 1a061b11a..834b2b3d8 100644 --- a/parquet/file/file_writer_test.go +++ b/parquet/file/file_writer_test.go @@ -702,6 +702,64 @@ func TestALPFileRoundtrip(t *testing.T) { } } +// The vector size is a writer property, so it has to reach each column's +// encoder: the pages of a column must declare the size set for it, and still +// read back. +func TestALPFileVectorSize(t *testing.T) { + values := make([]float64, 5000) + for i := range values { + values[i] = float64(i%1234) / 100 + } + + props := parquet.NewWriterProperties( + parquet.WithEncoding(parquet.Encodings.ALP), + parquet.WithALPEncoding(true), + parquet.WithDictionaryDefault(false), + parquet.WithALPVectorSizeFor("small", 32), + parquet.WithALPVectorSizeFor("large", 4096), + ) + var fields schema.FieldList + for _, name := range []string{"small", "default", "large"} { + fields = append(fields, schema.NewFloat64Node(name, parquet.Repetitions.Required, -1)) + } + sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, fields, -1) + require.NoError(t, err) + + var buf bytes.Buffer + writer := file.NewParquetWriter(&buf, sc, file.WithWriterProps(props)) + rgw := writer.AppendRowGroup() + for range fields { + cw, err := rgw.NextColumn() + require.NoError(t, err) + _, err = cw.(*file.Float64ColumnChunkWriter).WriteBatch(values, nil, nil) + require.NoError(t, err) + require.NoError(t, cw.Close()) + } + require.NoError(t, rgw.Close()) + require.NoError(t, writer.Close()) + + reader, err := file.NewParquetReader(bytes.NewReader(buf.Bytes())) + require.NoError(t, err) + defer reader.Close() + + for col, logSize := range []byte{5, 10, 12} { + pages, err := reader.RowGroup(0).GetColumnPageReader(col) + require.NoError(t, err) + require.True(t, pages.Next(), "column %d has a page", col) + // A required column has no levels, so the page data starts with the + // ALP header, whose third byte is log2 of the vector size. + assert.Equal(t, logSize, pages.Page().Data()[2], "column %d", col) + + cr, err := reader.RowGroup(0).Column(col) + require.NoError(t, err) + got := make([]float64, len(values)) + total, _, err := cr.(*file.Float64ColumnChunkReader).ReadBatch(int64(len(values)), got, nil, nil) + require.NoError(t, err) + require.EqualValues(t, len(values), total) + assert.Equal(t, values, got, "column %d", col) + } +} + // A column whose values outgrow the data page size is written as several pages. // Each one carries its own ALP header, offset array and vectors, so a page that // repeated an earlier page's values would decode to the wrong numbers here. diff --git a/parquet/internal/encoding/alp_encoding.go b/parquet/internal/encoding/alp_encoding.go index 23c6114d3..4590bfa6e 100644 --- a/parquet/internal/encoding/alp_encoding.go +++ b/parquet/internal/encoding/alp_encoding.go @@ -156,12 +156,15 @@ type alpEncoder[T alpFloat] struct { encodedBuf []byte vectorSizes []int - // Scratch reused by every vector, so encoding allocates nothing per vector. + // Scratch reused across vectors, calls, and pages, so steady-state + // encoding allocates nothing. encodedScratch []int64 deltaScratch []uint64 excPositions []uint16 excValues []T sampleScratch []T + spacedScratch []T + pageHeader []byte vectorsProcessed int cachedPresets []alpEncodingParams @@ -169,16 +172,38 @@ type alpEncoder[T alpFloat] struct { } func newALPEncoder[T alpFloat](e format.Encoding, descr *schema.Column, mem memory.Allocator) *alpEncoder[T] { - return &alpEncoder[T]{ - encoder: newEncoderBase(e, descr, mem), - vectorSize: alpDefaultVectorSize, - logVectorSize: alpDefaultLogVectorSize, - vectorBuf: make([]T, alpDefaultVectorSize), - encodedScratch: make([]int64, alpDefaultVectorSize), - deltaScratch: make([]uint64, alpDefaultVectorSize), - sampleScratch: make([]T, 0, alpSamplerSamples), - presetCounts: make(map[alpEncodingParams]int), + 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< 1< 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 { @@ -203,9 +228,9 @@ func (enc *alpEncoder[T]) Put(in []T) { // PutSpaced drops the null slots of in and buffers the values that remain. func (enc *alpEncoder[T]) PutSpaced(in []T, validBits []byte, validBitsOffset int64) { - nbuf := make([]T, len(in)) - nvalid := spacedCompress(in, nbuf, validBits, validBitsOffset) - enc.Put(nbuf[:nvalid]) + 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. @@ -235,9 +260,10 @@ func (enc *alpEncoder[T]) encodeVector(values []T) { // 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 := alpTryEncode(v, params.exponent, params.factor) + e, ok := s.tryEncode(v) if !ok { enc.excPositions = append(enc.excPositions, uint16(i)) enc.excValues = append(enc.excValues, v) @@ -341,8 +367,7 @@ func (enc *alpEncoder[T]) FlushValues() (Buffer, error) { offsetArraySize := len(enc.vectorSizes) * 4 enc.sink.Reserve(alpHeaderSize + offsetArraySize + len(enc.encodedBuf)) - header := make([]byte, 0, alpHeaderSize+offsetArraySize) - header = append(header, alpCompressionMode, alpIntegerEncodingFOR, byte(enc.logVectorSize)) + header := append(enc.pageHeader[:0], alpCompressionMode, alpIntegerEncodingFOR, byte(enc.logVectorSize)) header = binary.LittleEndian.AppendUint32(header, uint32(enc.totalCount)) offset := uint32(offsetArraySize) @@ -350,6 +375,7 @@ func (enc *alpEncoder[T]) FlushValues() (Buffer, error) { header = binary.LittleEndian.AppendUint32(header, offset) offset += uint32(size) } + enc.pageHeader = header enc.sink.Write(header) enc.sink.Write(enc.encodedBuf) @@ -364,7 +390,7 @@ func (enc *alpEncoder[T]) resetPage() { } // Reset returns the encoder to its state when new, shortlist included, so that -// it can encode an unrelated column. +// it can encode an unrelated column. The vector size stays as set. func (enc *alpEncoder[T]) Reset() { enc.encoder.Reset() enc.resetPage() @@ -518,8 +544,20 @@ func (dec *alpDecoder[T]) decodeVector(vectorIdx int) error { } pos += packedSize - for i, d := range deltas { - dec.decodedValues[i] = alpDecode[T](int64(d)+frameOfRef, exponent, factor) + // The deltas were computed in the encoded integer's width with wrapping + // arithmetic, so add the frame of reference back the same way. + s := newALPScaling[T](exponent, factor) + out := dec.decodedValues[:len(deltas)] + if valueBytes == 4 { + ref := uint32(frameOfRef) + for i, d := range deltas { + out[i] = s.decode(int64(int32(uint32(d) + ref))) + } + } else { + ref := uint64(frameOfRef) + for i, d := range deltas { + out[i] = s.decode(int64(d + ref)) + } } // Positions come first, then the values, so read both in one pass instead @@ -606,3 +644,14 @@ type ( // ALPFloat64Decoder decodes ALP-encoded float64 values. ALPFloat64Decoder = alpDecoder[float64] ) + +// ALPEncoder is implemented by the ALP encoders, so that a column writer can +// set the vector size without knowing the physical type. +type ALPEncoder interface { + SetVectorSize(size int) error +} + +var ( + _ ALPEncoder = (*ALPFloat32Encoder)(nil) + _ ALPEncoder = (*ALPFloat64Encoder)(nil) +) diff --git a/parquet/internal/encoding/alp_encoding_test.go b/parquet/internal/encoding/alp_encoding_test.go index b7ab0378e..b8c4b3f04 100644 --- a/parquet/internal/encoding/alp_encoding_test.go +++ b/parquet/internal/encoding/alp_encoding_test.go @@ -21,6 +21,7 @@ import ( "encoding/binary" "fmt" "math" + "math/bits" "slices" "testing" @@ -1261,3 +1262,54 @@ func TestALPEncoderResetClearsPresets(t *testing.T) { assert.Zero(t, enc.vectorsProcessed) assert.Zero(t, enc.EstimatedDataEncodedSize()) } + +// Every vector size the format allows must round-trip, including partial last +// vectors and the largest size, whose exception positions reach 32767. +func TestALPEncoderVectorSizes(t *testing.T) { + for _, size := range []int{8, 32, 1024, 4096, 1 << 15} { + t.Run(fmt.Sprint(size), func(t *testing.T) { + values := make([]float64, 2*size+size/2+3) + for i := range values { + values[i] = float64(i%1000) / 100 + } + // Exceptions at both ends of every vector. + for i := 0; i < len(values); i += size { + values[i] = math.Pi + values[min(i+size, len(values))-1] = math.NaN() + } + + enc := newALPEncoder[float64](format.Encoding_ALP, nil, memory.DefaultAllocator) + require.NoError(t, enc.SetVectorSize(size)) + enc.Put(values) + page, err := enc.FlushValues() + require.NoError(t, err) + defer page.Release() + + data := page.Bytes() + assert.Equal(t, byte(bits.TrailingZeros(uint(size))), data[2], "log vector size") + numVectors := (len(values) + size - 1) / size + assert.EqualValues(t, numVectors*4, binary.LittleEndian.Uint32(data[alpHeaderSize:]), + "the first vector starts after %d offsets", numVectors) + + got := alpDecodePage[float64](t, page, len(values)) + for i, v := range values { + require.Equal(t, math.Float64bits(v), math.Float64bits(got[i]), "index %d", i) + } + }) + } +} + +func TestALPEncoderSetVectorSizeRejects(t *testing.T) { + enc := newALPEncoder[float32](format.Encoding_ALP, nil, memory.DefaultAllocator) + for _, size := range []int{0, 4, 1000, 1 << 16} { + assert.Error(t, enc.SetVectorSize(size), "size %d", size) + } + + enc.Put([]float32{1.5}) + assert.Error(t, enc.SetVectorSize(32), "values written at 1024 cannot move to 32") + + page, err := enc.FlushValues() + require.NoError(t, err) + page.Release() + require.NoError(t, enc.SetVectorSize(32), "a flushed page leaves nothing buffered") +} diff --git a/parquet/internal/encoding/alp_params.go b/parquet/internal/encoding/alp_params.go index 60e90c8aa..7da63a0dd 100644 --- a/parquet/internal/encoding/alp_params.go +++ b/parquet/internal/encoding/alp_params.go @@ -67,8 +67,9 @@ func alpExceptionBits[T alpFloat]() int { 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 := alpTryEncode(v, exponent, factor) + e, ok := s.tryEncode(v) if !ok { numExceptions++ continue diff --git a/parquet/internal/encoding/alp_params_test.go b/parquet/internal/encoding/alp_params_test.go index 5ae231a14..ce77038ad 100644 --- a/parquet/internal/encoding/alp_params_test.go +++ b/parquet/internal/encoding/alp_params_test.go @@ -27,8 +27,9 @@ import ( // alpCountTestExceptions counts the values a scaling cannot encode. func alpCountTestExceptions[T float32 | float64](values []T, params alpEncodingParams) int { count := 0 + s := newALPScaling[T](params.exponent, params.factor) for _, v := range values { - if alpDecode[T](alpEncode(v, params.exponent, params.factor), params.exponent, params.factor) != v { + if _, ok := s.tryEncode(v); !ok { count++ } } diff --git a/parquet/internal/encoding/alp_scaling.go b/parquet/internal/encoding/alp_scaling.go index 6b591055b..7474a1bfa 100644 --- a/parquet/internal/encoding/alp_scaling.go +++ b/parquet/internal/encoding/alp_scaling.go @@ -168,47 +168,61 @@ func alpFloatFromBits[T alpFloat](b uint64) T { panic("parquet: ALP: unsupported float type") } -// Adding and subtracting this constant rounds as in the reference -// implementation without a library call. -func alpFastRound[T alpFloat](v T) int64 { - magic := alpMagic[T]() - if v >= 0 { - return int64((v + magic) - magic) - } - return int64((v - magic) + magic) -} +// alpScaling holds the constants of one exponent and factor pair. Encoding and +// decoding look them up once per vector rather than once per value. +type alpScaling[T alpFloat] struct { + // Encoding multiplies by 10^exponent then 10^-factor; decoding multiplies + // by 10^factor then 10^-exponent. + encExp, encFac T + decFac, decExp T -// NaNs, infinities, out-of-range values, and negative zero cannot round-trip -// through ALP's integer representation. -func alpImpossibleToEncode[T alpFloat](v T) bool { - lo, hi := alpEncodingLimits[T]() - f := float64(v) - return math.IsNaN(f) || v > hi || v < lo || (f == 0 && math.Signbit(f)) + lo, hi, magic T } -// Unencodable values use a sentinel that cannot decode to the original value, -// so the caller records them as exceptions. -func alpNumberToInt[T alpFloat](v T) int64 { - if alpImpossibleToEncode(v) { - _, hi := alpEncodingLimits[T]() - return int64(hi) +func newALPScaling[T alpFloat](exponent, factor int) alpScaling[T] { + lo, hi := alpEncodingLimits[T]() + return alpScaling[T]{ + encExp: alpPow10[T](exponent), + encFac: alpNegPow10[T](factor), + decFac: alpPow10[T](factor), + decExp: alpNegPow10[T](exponent), + lo: lo, + hi: hi, + magic: alpMagic[T](), } - return alpFastRound(v) } -func alpEncode[T alpFloat](v T, exponent, factor int) int64 { - return alpNumberToInt(v * alpPow10[T](exponent) * alpNegPow10[T](factor)) +// round adds and subtracts a magic constant, which rounds to the nearest +// integer without a library call. +func (s *alpScaling[T]) round(v T) int64 { + if v >= 0 { + return int64((v + s.magic) - s.magic) + } + return int64((v - s.magic) + s.magic) +} + +// encode returns the integer v scales to. NaNs, infinities, out-of-range +// values, and negative zero cannot round-trip through an integer, so they get a +// sentinel that cannot decode to the original value. +func (s *alpScaling[T]) encode(v T) int64 { + // The conversion keeps the product from fusing with the addition in round, + // so every platform rounds the same product. + scaled := T(v * s.encExp * s.encFac) + if scaled != scaled || scaled > s.hi || scaled < s.lo || (scaled == 0 && math.Signbit(float64(scaled))) { + return int64(s.hi) + } + return s.round(scaled) } -// The multiplication order is part of the format; reassociation can change the -// final bit. -func alpDecode[T alpFloat](encoded int64, exponent, factor int) T { - return T(encoded) * alpPow10[T](factor) * alpNegPow10[T](exponent) +// decode is the normative ALP decoding. The multiplication order is part of +// the format; reassociation can change the final bit. +func (s *alpScaling[T]) decode(encoded int64) T { + return T(encoded) * s.decFac * s.decExp } -// alpTryEncode reports whether v survives the required round trip. Values that +// tryEncode reports whether v survives the required round trip. Values that // fail are stored verbatim as exceptions. -func alpTryEncode[T alpFloat](v T, exponent, factor int) (encoded int64, ok bool) { - encoded = alpEncode(v, exponent, factor) - return encoded, alpDecode[T](encoded, exponent, factor) == v +func (s *alpScaling[T]) tryEncode(v T) (encoded int64, ok bool) { + encoded = s.encode(v) + return encoded, s.decode(encoded) == v } diff --git a/parquet/internal/encoding/alp_scaling_test.go b/parquet/internal/encoding/alp_scaling_test.go index 0545714a6..bd1bf0d97 100644 --- a/parquet/internal/encoding/alp_scaling_test.go +++ b/parquet/internal/encoding/alp_scaling_test.go @@ -44,8 +44,9 @@ func TestALPFastRoundFloat32(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - result := alpFastRound(tt.input) - assert.Equal(t, tt.expected, result, "alpFastRound(%v) = %d, want %d", tt.input, result, tt.expected) + s := newALPScaling[float32](0, 0) + result := s.round(tt.input) + assert.Equal(t, tt.expected, result, "round(%v) = %d, want %d", tt.input, result, tt.expected) }) } } @@ -69,8 +70,9 @@ func TestALPFastRoundFloat64(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - result := alpFastRound(tt.input) - assert.Equal(t, tt.expected, result, "alpFastRound(%v) = %d, want %d", tt.input, result, tt.expected) + s := newALPScaling[float64](0, 0) + result := s.round(tt.input) + assert.Equal(t, tt.expected, result, "round(%v) = %d, want %d", tt.input, result, tt.expected) }) } } @@ -83,7 +85,8 @@ func alpNegZero[T alpFloat]() T { func alpIsException[T alpFloat](t *testing.T, v T, exponent, factor int) bool { t.Helper() - _, ok := alpTryEncode(v, exponent, factor) + s := newALPScaling[T](exponent, factor) + _, ok := s.tryEncode(v) return !ok } @@ -114,11 +117,13 @@ func TestALPTryEncodeFloat64Exceptions(t *testing.T) { } func TestALPTryEncodeReturnsTheEncodedValue(t *testing.T) { - e, ok := alpTryEncode(1.23, 2, 0) + s := newALPScaling[float64](2, 0) + e, ok := s.tryEncode(1.23) require.True(t, ok) assert.EqualValues(t, 123, e, "1.23 at 10^2 encodes to 123") - e, ok = alpTryEncode(1.5, 3, 2) + s = newALPScaling[float64](3, 2) + e, ok = s.tryEncode(1.5) require.True(t, ok) assert.EqualValues(t, 15, e, "10^3 * 10^-2 scales by ten") } @@ -141,10 +146,11 @@ func TestALPEncodeDecodeFloat32(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - encoded := alpEncode(tt.value, tt.exponent, tt.factor) + s := newALPScaling[float32](tt.exponent, tt.factor) + encoded := s.encode(tt.value) assert.Equal(t, tt.encoded, encoded, "encode(%v, exp=%d, fac=%d)", tt.value, tt.exponent, tt.factor) - decoded := alpDecode[float32](encoded, tt.exponent, tt.factor) + decoded := s.decode(encoded) assert.Equal(t, math.Float32bits(tt.value), math.Float32bits(decoded), "decode(%d, exp=%d, fac=%d) = %v, want %v", encoded, tt.exponent, tt.factor, decoded, tt.value) }) @@ -169,10 +175,11 @@ func TestALPEncodeDecodeFloat64(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - encoded := alpEncode(tt.value, tt.exponent, tt.factor) + s := newALPScaling[float64](tt.exponent, tt.factor) + encoded := s.encode(tt.value) assert.Equal(t, tt.encoded, encoded, "encode(%v, exp=%d, fac=%d)", tt.value, tt.exponent, tt.factor) - decoded := alpDecode[float64](encoded, tt.exponent, tt.factor) + decoded := s.decode(encoded) assert.Equal(t, math.Float64bits(tt.value), math.Float64bits(decoded), "decode(%d, exp=%d, fac=%d) = %v, want %v", encoded, tt.exponent, tt.factor, decoded, tt.value) }) @@ -199,8 +206,9 @@ func TestALPScalingMultipliesByANegativePower(t *testing.T) { const value = float32(1016.10999) const exponent, factor = 6, 0 - encoded := alpEncode(value, exponent, factor) - assert.Equal(t, value, alpDecode[float32](encoded, exponent, factor)) + s := newALPScaling[float32](exponent, factor) + encoded := s.encode(value) + assert.Equal(t, value, s.decode(encoded)) assert.NotEqual(t, value, float32(encoded)/(alpPow10[float32](exponent)/alpPow10[float32](factor)), "dividing has stopped losing the last bit, so this test no longer guards anything") } diff --git a/parquet/reader_writer_properties_test.go b/parquet/reader_writer_properties_test.go index 4729323cc..949c2355c 100644 --- a/parquet/reader_writer_properties_test.go +++ b/parquet/reader_writer_properties_test.go @@ -146,3 +146,25 @@ func TestWriterPropALPEncodingPerColumn(t *testing.T) { parquet.WithEncodingFor("volume", parquet.Encodings.ALP)) }, "the column-specific setting wins over the default one") } + +func TestWriterPropALPVectorSize(t *testing.T) { + props := parquet.NewWriterProperties() + assert.Equal(t, parquet.DefaultALPVectorSize, props.ALPVectorSizeFor("price")) + + props = parquet.NewWriterProperties( + parquet.WithALPVectorSize(4096), + parquet.WithALPVectorSizeFor("price", 32), + parquet.WithEncodingFor("volume", parquet.Encodings.Plain)) + assert.Equal(t, 32, props.ALPVectorSizeFor("price")) + assert.Equal(t, 4096, props.ALPVectorSizeFor("volume"), + "a column configured for something else keeps the default size") + assert.Equal(t, 4096, props.ALPVectorSizeFor("other")) + + for _, size := range []int{0, 4, 1000, 1 << 16} { + assert.Panics(t, func() { parquet.NewWriterProperties(parquet.WithALPVectorSize(size)) }, "size %d", size) + assert.Panics(t, func() { parquet.NewWriterProperties(parquet.WithALPVectorSizeFor("price", size)) }, "size %d", size) + } + for _, size := range []int{8, 1 << 15} { + assert.NotPanics(t, func() { parquet.NewWriterProperties(parquet.WithALPVectorSize(size)) }, "size %d", size) + } +} diff --git a/parquet/writer_properties.go b/parquet/writer_properties.go index 205da8644..8c51c4174 100644 --- a/parquet/writer_properties.go +++ b/parquet/writer_properties.go @@ -17,6 +17,7 @@ package parquet import ( + "fmt" "math" "github.com/apache/arrow-go/v18/arrow" @@ -63,6 +64,12 @@ const ( // DefaultALPEnabled is false because ALP is a preview encoding. DefaultALPEnabled = false + // DefaultALPVectorSize is how many values an ALP vector holds unless + // WithALPVectorSize says otherwise. + DefaultALPVectorSize = 1024 + + minALPVectorSize = 1 << 3 + maxALPVectorSize = 1 << 15 minimumBloomFilterBytes = 32 maximumBloomFilterBytes = 128 * 1024 * 1024 @@ -89,6 +96,8 @@ type ColumnProperties struct { BloomFilterNDV int64 // ALPEnabled permits this column to use the preview ALP encoding. ALPEnabled bool + // ALPVectorSize is how many values each ALP vector of this column holds. + ALPVectorSize int } // DefaultColumnProperties returns the default properties which get utilized for writing. @@ -120,6 +129,7 @@ func DefaultColumnProperties() ColumnProperties { AdaptiveBloomFilterEnabled: DefaultAdaptiveBloomFilterEnabled, BloomFilterCandidates: DefaultBloomFilterCandidates, ALPEnabled: DefaultALPEnabled, + ALPVectorSize: DefaultALPVectorSize, } } @@ -141,6 +151,7 @@ type writerPropConfig struct { adaptiveBloomFilterEnabled map[string]bool numBloomFilterCandidates map[string]int alpEnabled map[string]bool + alpVectorSize map[string]int } // WriterProperty is used as the options for building a writer properties instance @@ -323,6 +334,33 @@ func WithALPEncodingPath(path ColumnPath, enabled bool) WriterProperty { return WithALPEncodingFor(path.String(), enabled) } +// WithALPVectorSize sets how many values each ALP vector holds for all +// columns. It must be a power of two from 8 to 32768; the default is +// DefaultALPVectorSize. Larger vectors amortize the per-vector header, while +// smaller ones let each vector pick the scaling that suits its values. +// +// It only affects columns written with the ALP encoding, and readers honor +// whatever size a page declares. +func WithALPVectorSize(size int) WriterProperty { + return func(cfg *writerPropConfig) { + cfg.wr.defColumnProps.ALPVectorSize = size + } +} + +// WithALPVectorSizeFor sets the ALP vector size for one column. See +// WithALPVectorSize. +func WithALPVectorSizeFor(path string, size int) WriterProperty { + return func(cfg *writerPropConfig) { + cfg.alpVectorSize[path] = size + } +} + +// WithALPVectorSizePath is the same as WithALPVectorSizeFor but takes a +// ColumnPath. +func WithALPVectorSizePath(path ColumnPath, size int) WriterProperty { + return WithALPVectorSizeFor(path.String(), size) +} + // WithCompression specifies the default compression type to use for column writing. func WithCompression(codec compress.Compression) WriterProperty { return func(cfg *writerPropConfig) { @@ -638,6 +676,7 @@ func NewWriterProperties(opts ...WriterProperty) *WriterProperties { adaptiveBloomFilterEnabled: make(map[string]bool), numBloomFilterCandidates: make(map[string]int), alpEnabled: make(map[string]bool), + alpVectorSize: make(map[string]int), } for _, o := range opts { o(&cfg) @@ -704,17 +743,22 @@ func NewWriterProperties(opts ...WriterProperty) *WriterProperties { 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)") } + checkALPVectorSize(path, props.ALPVectorSize) } return cfg.wr @@ -749,6 +793,33 @@ 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 + } + panic(fmt.Sprintf("parquet: %s must be a power of two from %d to %d, got %d", + target, minALPVectorSize, maxALPVectorSize, size)) +} + // FileEncryptionProperties returns the current encryption properties that were // used to create the writer properties. func (w *WriterProperties) FileEncryptionProperties() *FileEncryptionProperties { From 9977f1b716a4129b1842e3dc270d2a3f6c2b2cee Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Thu, 29 Oct 2026 13:13:08 -0400 Subject: [PATCH 07/11] fix(parquet): address ALP review findings - Treat a FLOAT value whose rounded integer leaves int32 as an exception. Rounding a scaled value just under 2^31 could reach 2^31, which readers rebuild in int32 with the sign flipped (or a 33-bit packing width). - Avoid int overflow on 32-bit platforms when counting vectors and bounds-checking vector offsets in the decoder. - List the encodings the data pages actually used in the column chunk metadata when there is no dictionary, so a column that asked for ALP but fell back to PLAIN is not labelled ALP. - Point the docs at per-column ALP selection, make ALPEncodingEnabled report the default like the other getters, and add ALPVectorSize. - Count vector offsets in EstimatedDataEncodedSize; fix stale comments. --- parquet/doc.go | 6 +-- parquet/file/file_writer_test.go | 41 +++++++++++++++++++ parquet/internal/encoding/alp_encoding.go | 27 ++++++------ .../internal/encoding/alp_encoding_test.go | 25 +++++++++++ parquet/internal/encoding/alp_scaling.go | 16 ++++++++ parquet/metadata/column_chunk.go | 13 +++++- parquet/reader_writer_properties_test.go | 4 +- parquet/writer_properties.go | 28 ++++++------- 8 files changed, 125 insertions(+), 35 deletions(-) diff --git a/parquet/doc.go b/parquet/doc.go index 13dbdb7d6..aaddbbb12 100644 --- a/parquet/doc.go +++ b/parquet/doc.go @@ -72,9 +72,9 @@ // // - ALP, Adaptive Lossless floating-Point (Float, Double). The Parquet // format marks ALP as a preview encoding, so writing it is opt-in: enable -// it with WithALPEncoding, select it with WithEncoding or WithEncodingFor, -// and turn off dictionary encoding for those columns. Reading needs no -// option. +// it with WithALPEncoding, select it per column with WithEncodingFor or +// WithEncodingPath, and turn off dictionary encoding for those columns. +// Reading needs no option. // // Tip: Some platforms don't necessarily support all kinds of encodings. If you're not // sure what to use, just use Plain and Dictionary encoding. diff --git a/parquet/file/file_writer_test.go b/parquet/file/file_writer_test.go index 834b2b3d8..c8f6cf6c2 100644 --- a/parquet/file/file_writer_test.go +++ b/parquet/file/file_writer_test.go @@ -702,6 +702,47 @@ func TestALPFileRoundtrip(t *testing.T) { } } +// A column that asks for ALP but keeps its dictionary falls back to PLAIN when +// the dictionary outgrows its limit before the first page. The footer must then +// list PLAIN, not ALP: a reader without ALP support may reject a file whose +// footer names it, though no page uses it. +func TestALPRequestedButDictionaryDroppedListsPlain(t *testing.T) { + values := make([]float64, 20000) + for i := range values { + values[i] = float64(i) / 7 + } + props := parquet.NewWriterProperties( + parquet.WithALPEncoding(true), + parquet.WithEncodingFor("value", parquet.Encodings.ALP), + parquet.WithDictionaryPageSizeLimit(1024)) + sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, schema.FieldList{ + schema.NewFloat64Node("value", parquet.Repetitions.Required, -1), + }, -1) + require.NoError(t, err) + + var buf bytes.Buffer + writer := file.NewParquetWriter(&buf, sc, file.WithWriterProps(props)) + rgw := writer.AppendRowGroup() + cw, err := rgw.NextColumn() + require.NoError(t, err) + _, err = cw.(*file.Float64ColumnChunkWriter).WriteBatch(values, nil, nil) + require.NoError(t, err) + require.NoError(t, cw.Close()) + require.NoError(t, rgw.Close()) + require.NoError(t, writer.Close()) + + reader, err := file.NewParquetReader(bytes.NewReader(buf.Bytes())) + require.NoError(t, err) + defer reader.Close() + chunk, err := reader.RowGroup(0).MetaData().ColumnChunk(0) + require.NoError(t, err) + for _, s := range chunk.EncodingStats() { + require.NotEqual(t, parquet.Encodings.ALP, s.Encoding, "no page should be ALP") + } + assert.NotContains(t, chunk.Encodings(), parquet.Encodings.ALP) + assert.Contains(t, chunk.Encodings(), parquet.Encodings.Plain) +} + // The vector size is a writer property, so it has to reach each column's // encoder: the pages of a column must declare the size set for it, and still // read back. diff --git a/parquet/internal/encoding/alp_encoding.go b/parquet/internal/encoding/alp_encoding.go index 4590bfa6e..a47e15b7c 100644 --- a/parquet/internal/encoding/alp_encoding.go +++ b/parquet/internal/encoding/alp_encoding.go @@ -25,7 +25,7 @@ package encoding // 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 uint32. An array +// 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: // @@ -55,14 +55,14 @@ import ( const ( // The first two header bytes select the ALP variants. The decoder rejects // variants other than the ones this encoder writes. - alpCompressionMode = 0 // ALP, rather than ALP for a sparse column + alpCompressionMode = 0 // ALP; other values are reserved for variants such as ALP-RD alpIntegerEncodingFOR = 0 // frame of reference plus bit packing alpHeaderSize = 7 - // This encoder writes 1,024-value vectors. The format allows anything from - // 2^alpMinLogVectorSize to 2^alpMaxLogVectorSize, and the decoder honors - // what a page declares. + // 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 @@ -156,8 +156,7 @@ type alpEncoder[T alpFloat] struct { encodedBuf []byte vectorSizes []int - // Scratch reused across vectors, calls, and pages, so steady-state - // encoding allocates nothing. + // Scratch reused across vectors, calls, and pages. encodedScratch []int64 deltaScratch []uint64 excPositions []uint16 @@ -344,10 +343,10 @@ func (enc *alpEncoder[T]) buildPresetCache() { } } -// EstimatedDataEncodedSize counts unencoded buffered values at their physical -// width. +// EstimatedDataEncodedSize counts the encoded vectors and their offsets, and +// unencoded buffered values at their physical width. func (enc *alpEncoder[T]) EstimatedDataEncodedSize() int64 { - return int64(len(enc.encodedBuf)) + int64(enc.bufCount*alpEncodedBytes[T]()) + return int64(len(enc.encodedBuf) + 4*len(enc.vectorSizes) + enc.bufCount*alpEncodedBytes[T]()) } // FlushValues writes one page and clears the values it wrote, so that the next @@ -463,7 +462,8 @@ func (dec *alpDecoder[T]) SetData(nvals int, data []byte) error { dec.vectorSize = 1 << logVectorSize dec.totalCount = int(numElements) - dec.numVectors = (dec.totalCount + dec.vectorSize - 1) / dec.vectorSize + // 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 @@ -508,9 +508,10 @@ func (dec *alpDecoder[T]) decodeVector(vectorIdx int) error { valueBytes := alpEncodedBytes[T]() // On 32-bit platforms, converting a large uint32 offset can make pos - // negative. + // negative, and adding to a large one can overflow, so compare against + // the length less the header instead. pos := int(dec.vectorOffsets[vectorIdx]) - if pos < 0 || pos+alpVectorHeaderSize[T]() > len(dec.bodyData) { + if pos < 0 || pos > len(dec.bodyData)-alpVectorHeaderSize[T]() { return fmt.Errorf("parquet: ALP vector %d starts past the end of the page", vectorIdx) } diff --git a/parquet/internal/encoding/alp_encoding_test.go b/parquet/internal/encoding/alp_encoding_test.go index b8c4b3f04..49452743b 100644 --- a/parquet/internal/encoding/alp_encoding_test.go +++ b/parquet/internal/encoding/alp_encoding_test.go @@ -1313,3 +1313,28 @@ func TestALPEncoderSetVectorSizeRejects(t *testing.T) { page.Release() require.NoError(t, enc.SetVectorSize(32), "a flushed page leaves nothing buffered") } + +// A FLOAT value scaled to just under 2^31 can round up to 2^31, which a vector +// cannot hold: readers rebuild its integers in int32 and would flip the sign. +func TestALPFloat32RoundingPastInt32IsAnException(t *testing.T) { + s := newALPScaling[float32](9, 6) + _, ok := s.tryEncode(2147483.5) + assert.False(t, ok, "2147483.5 at (9,6) rounds to 2^31") + + alpFloat32RoundTrip(t, []float32{-24.875, 2147483.5, 12}) + alpFloat32RoundTrip(t, []float32{-24.875, -2147483.5, 12}) +} + +// Page fields near the int32 limits must be rejected, not overflow int on +// 32-bit platforms. +func TestALPDecoderHugeFields(t *testing.T) { + dec := newALPDecoder[float32](format.Encoding_ALP, newFloat32Column()) + assert.Error(t, dec.SetData(1, alpTestPage(math.MaxInt32)), + "the vector count must not overflow") + + page, _ := alpEncodeTestPage(t, []float32{1.5, 2.5}) + binary.LittleEndian.PutUint32(page[alpHeaderSize:], math.MaxInt32-4) + require.NoError(t, dec.SetData(2, page)) + _, err := dec.Decode(make([]float32, 2)) + assert.ErrorContains(t, err, "past the end of the page") +} diff --git a/parquet/internal/encoding/alp_scaling.go b/parquet/internal/encoding/alp_scaling.go index 7474a1bfa..8ff42e614 100644 --- a/parquet/internal/encoding/alp_scaling.go +++ b/parquet/internal/encoding/alp_scaling.go @@ -177,10 +177,17 @@ type alpScaling[T alpFloat] struct { decFac, decExp T lo, hi, magic T + // The range of the integer type a vector stores: int32 for FLOAT, int64 + // for DOUBLE. + intMin, intMax int64 } func newALPScaling[T alpFloat](exponent, factor int) alpScaling[T] { lo, hi := alpEncodingLimits[T]() + intMin, intMax := int64(math.MinInt64), int64(math.MaxInt64) + if alpEncodedBytes[T]() == 4 { + intMin, intMax = math.MinInt32, math.MaxInt32 + } return alpScaling[T]{ encExp: alpPow10[T](exponent), encFac: alpNegPow10[T](factor), @@ -189,6 +196,8 @@ func newALPScaling[T alpFloat](exponent, factor int) alpScaling[T] { lo: lo, hi: hi, magic: alpMagic[T](), + intMin: intMin, + intMax: intMax, } } @@ -222,7 +231,14 @@ func (s *alpScaling[T]) decode(encoded int64) T { // tryEncode reports whether v survives the required round trip. Values that // fail are stored verbatim as exceptions. +// +// Rounding a scaled value just under the limit can still reach 2^31, which a +// FLOAT vector cannot store: readers rebuild its integers in int32. Such a +// value is an exception even though it would round-trip in int64. func (s *alpScaling[T]) tryEncode(v T) (encoded int64, ok bool) { encoded = s.encode(v) + if encoded < s.intMin || encoded > s.intMax { + return encoded, false + } return encoded, s.decode(encoded) == v } diff --git a/parquet/metadata/column_chunk.go b/parquet/metadata/column_chunk.go index c5ae5029a..4825ef818 100644 --- a/parquet/metadata/column_chunk.go +++ b/parquet/metadata/column_chunk.go @@ -22,7 +22,9 @@ import ( "errors" "fmt" "io" + "maps" "reflect" + "slices" "github.com/apache/arrow-go/v18/arrow/memory" "github.com/apache/arrow-go/v18/parquet" @@ -405,14 +407,21 @@ func (c *ColumnChunkMetaDataBuilder) Finish(info ChunkMetaInfo, hasDict, dictFal } else { thriftEncodings = append(thriftEncodings, format.Encoding(c.props.DictionaryPageEncoding())) } - } else { // no dictionary + } else if len(encStats.DataEncodingStats) > 0 { + // List what the data pages used rather than what was asked for: a + // writer whose dictionary was dropped before its first page writes + // PLAIN pages whatever encoding the properties name. + for _, e := range slices.Sorted(maps.Keys(encStats.DataEncodingStats)) { + thriftEncodings = append(thriftEncodings, format.Encoding(e)) + } + } else { // no pages to go by thriftEncodings = append(thriftEncodings, format.Encoding(c.props.EncodingFor(c.column.Path()))) } thriftEncodings = append(thriftEncodings, format.Encoding(parquet.Encodings.RLE)) // Only PLAIN encoding is supported for fallback in V1 // TODO(zeroshade): Use user specified encoding for V2 - if dictFallback { + if dictFallback && !slices.Contains(thriftEncodings, format.Encoding_PLAIN) { thriftEncodings = append(thriftEncodings, format.Encoding_PLAIN) } c.chunk.MetaData.Encodings = thriftEncodings diff --git a/parquet/reader_writer_properties_test.go b/parquet/reader_writer_properties_test.go index 949c2355c..c41b12b5d 100644 --- a/parquet/reader_writer_properties_test.go +++ b/parquet/reader_writer_properties_test.go @@ -128,8 +128,8 @@ func TestWriterPropALPEncodingPerColumn(t *testing.T) { props := parquet.NewWriterProperties( parquet.WithALPEncodingFor("price", true), parquet.WithEncodingFor("price", parquet.Encodings.ALP)) - assert.True(t, props.ALPEncodingEnabled(), - "one permitted column is enough to report the encoding as allowed") + assert.False(t, props.ALPEncodingEnabled(), + "like the other getters, the unqualified one reports the default") assert.True(t, props.ALPEncodingEnabledFor("price")) assert.False(t, props.ALPEncodingEnabledFor("volume")) diff --git a/parquet/writer_properties.go b/parquet/writer_properties.go index 8c51c4174..a523ddb31 100644 --- a/parquet/writer_properties.go +++ b/parquet/writer_properties.go @@ -306,8 +306,12 @@ func WithEncodingPath(path ColumnPath, encoding Encoding) WriterProperty { // default. // // The option grants permission rather than selecting ALP. A column still asks for -// the encoding through WithEncoding or WithEncodingFor, and asking for it without -// this option panics, as asking for any unusable encoding does. +// the encoding through WithEncodingFor or WithEncodingPath, and asking for it +// without this option panics, as asking for any unusable encoding does. +// +// ALP supports only FLOAT and DOUBLE columns. WithEncoding(Encodings.ALP) sets +// it for every column, so writing any other column without a dictionary +// (BOOLEAN and INT96 never use one) panics; prefer selecting it per column. // // A column reaches ALP only with dictionary encoding turned off. A dictionary // takes precedence over the encoding a column asks for, and a writer that drops @@ -764,19 +768,13 @@ func NewWriterProperties(opts ...WriterProperty) *WriterProperties { return cfg.wr } -// ALPEncodingEnabled reports whether any column may be written with the ALP -// encoding. -func (w *WriterProperties) ALPEncodingEnabled() bool { - if w.defColumnProps.ALPEnabled { - return true - } - for _, props := range w.columnProps { - if props.ALPEnabled { - return true - } - } - return false -} +// 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. From 1218254b5efafb445181f6864dee6503cb8c939f Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Thu, 29 Oct 2026 13:13:53 -0400 Subject: [PATCH 08/11] fix(parquet/metadata): list RLE once when a data page used it --- parquet/metadata/column_chunk.go | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/parquet/metadata/column_chunk.go b/parquet/metadata/column_chunk.go index 4825ef818..1a2ee0dc7 100644 --- a/parquet/metadata/column_chunk.go +++ b/parquet/metadata/column_chunk.go @@ -418,7 +418,11 @@ func (c *ColumnChunkMetaDataBuilder) Finish(info ChunkMetaInfo, hasDict, dictFal thriftEncodings = append(thriftEncodings, format.Encoding(c.props.EncodingFor(c.column.Path()))) } - thriftEncodings = append(thriftEncodings, format.Encoding(parquet.Encodings.RLE)) + // RLE for the repetition and definition levels, unless a data page already + // listed it. + if !slices.Contains(thriftEncodings, format.Encoding_RLE) { + thriftEncodings = append(thriftEncodings, format.Encoding_RLE) + } // Only PLAIN encoding is supported for fallback in V1 // TODO(zeroshade): Use user specified encoding for V2 if dictFallback && !slices.Contains(thriftEncodings, format.Encoding_PLAIN) { From a87963c66b646b9329a154dbfab8b35f2cbf7f91 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Thu, 29 Oct 2026 13:34:54 -0400 Subject: [PATCH 09/11] test(parquet): tighten and consolidate ALP tests - Cover V2 pages with nulls and compression, and the spaced write/read paths, which the existing tests named but did not reach. - Prove pages are ALP via data-page encoding stats, and check the many-pages test really writes several pages. - Merge the float32/float64 copy-paste round trips into one generic table test; drop tests that only restated types or duplicated the shared encoding suite; move reader-error tests to column_reader_test.go on testutils.MockPageReader; gather file-level ALP tests in alp_test.go. - Use checked allocators, and report allocations in the benchmarks. --- parquet/file/alp_test.go | 534 +++++++++++++--- parquet/file/column_reader_test.go | 67 ++ parquet/file/file_reader_test.go | 9 +- parquet/file/file_writer_test.go | 317 ---------- .../encoding/alp_encoding_benchmark_test.go | 10 + .../internal/encoding/alp_encoding_test.go | 583 ++++++------------ parquet/pqarrow/alp_test.go | 12 +- parquet/reader_writer_properties_test.go | 32 +- 8 files changed, 748 insertions(+), 816 deletions(-) diff --git a/parquet/file/alp_test.go b/parquet/file/alp_test.go index f06507b70..d4d5568a6 100644 --- a/parquet/file/alp_test.go +++ b/parquet/file/alp_test.go @@ -22,102 +22,232 @@ import ( "math/rand" "testing" + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/bitutil" "github.com/apache/arrow-go/v18/arrow/memory" "github.com/apache/arrow-go/v18/parquet" "github.com/apache/arrow-go/v18/parquet/compress" "github.com/apache/arrow-go/v18/parquet/file" - "github.com/apache/arrow-go/v18/parquet/internal/encoding" + "github.com/apache/arrow-go/v18/parquet/metadata" "github.com/apache/arrow-go/v18/parquet/schema" "github.com/stretchr/testify/require" ) -// alpWriteReadDouble writes one optional or required double column with ALP and -// reads it back. A nil defLevels writes a required column. -func alpWriteReadDouble(t *testing.T, values []float64, defLevels []int16, opts ...parquet.WriterProperty) ([]float64, []int16) { +// writeALPFile writes one row group and returns a reader owned by the caller. +// The callback writes each column; options can override the ALP defaults. +func writeALPFile(t *testing.T, mem memory.Allocator, fields schema.FieldList, writeColumn func(file.ColumnChunkWriter), opts ...parquet.WriterProperty) *file.Reader { t.Helper() - rep := parquet.Repetitions.Required - if defLevels != nil { - rep = parquet.Repetitions.Optional - } opts = append([]parquet.WriterProperty{ + parquet.WithAllocator(mem), parquet.WithEncoding(parquet.Encodings.ALP), parquet.WithALPEncoding(true), parquet.WithDictionaryDefault(false), }, opts...) - - field, err := schema.NewPrimitiveNode("value", rep, parquet.Types.Double, -1, -1) - require.NoError(t, err) - sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, schema.FieldList{field}, 0) + sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, fields, -1) require.NoError(t, err) - sink := encoding.NewBufferWriter(0, memory.DefaultAllocator) - writer := file.NewParquetWriter(sink, sc, file.WithWriterProps(parquet.NewWriterProperties(opts...))) + var sink bytes.Buffer + writer := file.NewParquetWriter(&sink, sc, file.WithWriterProps(parquet.NewWriterProperties(opts...))) + defer writer.Close() rgw := writer.AppendRowGroup() - cw, err := rgw.NextColumn() - require.NoError(t, err) - _, err = cw.(*file.Float64ColumnChunkWriter).WriteBatch(values, defLevels, nil) - require.NoError(t, err) - require.NoError(t, cw.Close()) + for range fields { + cw, err := rgw.NextColumn() + require.NoError(t, err) + writeColumn(cw) + require.NoError(t, cw.Close()) + } require.NoError(t, rgw.Close()) require.NoError(t, writer.Close()) - reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes())) + reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes()), + file.WithReadProps(parquet.NewReaderProperties(mem))) require.NoError(t, err) - defer reader.Close() + require.Equal(t, 1, reader.NumRowGroups()) + return reader +} + +// readALPDoubleColumn returns packed non-null values and, for optional columns, +// definition levels. Row and value offsets advance independently across batches. +func readALPDoubleColumn(t *testing.T, reader *file.Reader, col, batchSize int) ([]float64, []int16) { + t.Helper() - cr, err := reader.RowGroup(0).Column(0) + cr, err := reader.RowGroup(0).Column(col) require.NoError(t, err) + defer cr.Close() r := cr.(*file.Float64ColumnChunkReader) rows := int(reader.NumRows()) out := make([]float64, rows) var outDef []int16 - if defLevels != nil { + if cr.Descriptor().MaxDefinitionLevel() > 0 { outDef = make([]int16, rows) + for i := range outDef { + outDef[i] = -1 + } } - for read := 0; read < rows; { - _, n, err := r.ReadBatch(int64(rows), out[read:], outDef, nil) + valuesRead := 0 + for rowsRead := 0; rowsRead < rows; { + var def []int16 + if outDef != nil { + def = outDef[rowsRead:] + } + batch := min(batchSize, rows-rowsRead) + total, n, err := r.ReadBatch(int64(batch), out[valuesRead:], def, nil) require.NoError(t, err) require.NoError(t, r.Err()) - if n == 0 { - break + require.EqualValues(t, batch, total, "read stalled at row %d", rowsRead) + rowsRead += int(total) + valuesRead += n + } + require.False(t, r.HasNext()) + require.NoError(t, r.Err()) + return out[:valuesRead], outDef +} + +func assertALPDataPageEncoding(t *testing.T, chunk *metadata.ColumnChunkMetaData) { + t.Helper() + + dataPageEntries := 0 + for _, stat := range chunk.EncodingStats() { + if stat.PageType == file.PageTypeDataPage || stat.PageType == file.PageTypeDataPageV2 { + dataPageEntries++ + require.Equal(t, parquet.Encodings.ALP, stat.Encoding) } - read += n } - return out, outDef + require.Positive(t, dataPageEntries, "missing data-page encoding statistics") } // An all-null page contains levels but no encoded values. PLAIN is the control. func TestALPAllNullPage(t *testing.T) { for _, enc := range []parquet.Encoding{parquet.Encodings.Plain, parquet.Encodings.ALP} { t.Run(enc.String(), func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + const rows = 8 - out, def := alpWriteReadDouble(t, nil, make([]int16, rows), - parquet.WithEncoding(enc)) - require.Len(t, out, rows) - for i, d := range def { - require.Zero(t, d, "row %d should be null", i) + defLevels := make([]int16, rows) + reader := writeALPFile(t, mem, schema.FieldList{ + schema.NewFloat64Node("value", parquet.Repetitions.Optional, -1), + }, func(cw file.ColumnChunkWriter) { + n, err := cw.(*file.Float64ColumnChunkWriter).WriteBatch(nil, defLevels, nil) + require.NoError(t, err) + require.Zero(t, n) + }, parquet.WithEncoding(enc)) + defer reader.Close() + require.EqualValues(t, rows, reader.NumRows()) + if enc == parquet.Encodings.ALP { + chunk, err := reader.RowGroup(0).MetaData().ColumnChunk(0) + require.NoError(t, err) + assertALPDataPageEncoding(t, chunk) } + + out, def := readALPDoubleColumn(t, reader, 0, rows) + require.Empty(t, out) + require.Equal(t, defLevels, def) }) } } -// Scattered nulls exercise PutSpaced and DecodeSpaced. +// Keep the packed path covered alongside PutSpaced and DecodeSpaced. The latter +// is reached through RecordReader; column ReadBatch only returns packed values. func TestALPNullsScattered(t *testing.T) { - const rows = 3000 - values := make([]float64, 0, rows) + const ( + rows = 3000 + validBitsOffset = 3 + ) + values := make([]float64, rows) + packed := make([]float64, 0, rows) defLevels := make([]int16, rows) + validBits := make([]byte, bitutil.BytesForBits(rows+validBitsOffset)) for i := range rows { if i%3 == 0 { - continue // definition level 0: this row is null + values[i] = -9999 // A null slot must not enter the encoded values. + continue } defLevels[i] = 1 - values = append(values, float64(i)/100) + bitutil.SetBit(validBits, i+validBitsOffset) + values[i] = float64(i) / 100 + switch i { + case 1: + values[i] = math.Copysign(0, -1) + case 2: + values[i] = math.Float64frombits(0x7ff8000000001234) + case 4: + values[i] = math.Inf(1) + case 5: + values[i] = math.SmallestNonzeroFloat64 + } + packed = append(packed, values[i]) } - got, _ := alpWriteReadDouble(t, values, defLevels) - require.Equal(t, values, got[:len(values)]) + for _, tc := range []struct { + name string + spaced bool + }{ + {"packed", false}, + {"spaced", true}, + } { + t.Run(tc.name, func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + reader := writeALPFile(t, mem, schema.FieldList{ + schema.NewFloat64Node("value", parquet.Repetitions.Optional, -1), + }, func(cw file.ColumnChunkWriter) { + w := cw.(*file.Float64ColumnChunkWriter) + if tc.spaced { + n, err := w.WriteBatchSpacedWithError(values, defLevels, nil, validBits, validBitsOffset) + require.NoError(t, err) + require.EqualValues(t, rows, n) + } else { + n, err := w.WriteBatch(packed, defLevels, nil) + require.NoError(t, err) + require.EqualValues(t, len(packed), n) + } + }) + defer reader.Close() + require.EqualValues(t, rows, reader.NumRows()) + chunk, err := reader.RowGroup(0).MetaData().ColumnChunk(0) + require.NoError(t, err) + assertALPDataPageEncoding(t, chunk) + + gotPacked, gotDef := readALPDoubleColumn(t, reader, 0, 513) + require.Equal(t, defLevels, gotDef) + require.Len(t, gotPacked, len(packed)) + for i := range packed { + require.Equal(t, math.Float64bits(packed[i]), math.Float64bits(gotPacked[i]), "value %d", i) + } + if !tc.spaced { + return + } + + pages, err := reader.RowGroup(0).GetColumnPageReader(0) + require.NoError(t, err) + rr := file.NewRecordReader(reader.MetaData().Schema.Column(0), + file.LevelInfo{NullSlotUsage: 1, DefLevel: 1}, arrow.PrimitiveTypes.Float64, mem, reader.BufferPool()) + defer rr.Release() + rr.SetPageReader(pages) + require.NoError(t, rr.Reserve(rows)) + n, err := rr.ReadRecords(rows) + require.NoError(t, err) + require.EqualValues(t, rows, n) + require.Equal(t, rows, rr.ValuesWritten()) + require.EqualValues(t, rows-len(packed), rr.NullCount()) + require.EqualValues(t, rows, rr.LevelsPos()) + require.Equal(t, defLevels, rr.DefLevels()[:rows]) + + got := arrow.Float64Traits.CastFromBytes(rr.Values())[:rows] + valid := rr.ReleaseValidBits() + defer valid.Release() + for i, def := range defLevels { + require.Equal(t, def == 1, bitutil.BitIsSet(valid.Bytes(), i), "row %d validity", i) + if def == 1 { + require.Equal(t, math.Float64bits(values[i]), math.Float64bits(got[i]), "row %d", i) + } + } + }) + } } // These distributions cover exact decimals, integer overflow, NaNs, @@ -138,12 +268,29 @@ func TestALPRandomRoundtrip(t *testing.T) { for _, pool := range pools { t.Run(pool.name, func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + values := make([]float64, 5000) for i := range values { values[i] = pool.draw() } - got, _ := alpWriteReadDouble(t, values, nil) + reader := writeALPFile(t, mem, schema.FieldList{ + schema.NewFloat64Node("value", parquet.Repetitions.Required, -1), + }, func(cw file.ColumnChunkWriter) { + n, err := cw.(*file.Float64ColumnChunkWriter).WriteBatch(values, nil, nil) + require.NoError(t, err) + require.EqualValues(t, len(values), n) + }) + defer reader.Close() + require.EqualValues(t, len(values), reader.NumRows()) + chunk, err := reader.RowGroup(0).MetaData().ColumnChunk(0) + require.NoError(t, err) + assertALPDataPageEncoding(t, chunk) + + got, _ := readALPDoubleColumn(t, reader, 0, len(values)) + require.Len(t, got, len(values)) for i := range values { // Compare the bit patterns, so that a NaN or a negative zero has // to come back exactly as it went in. @@ -158,73 +305,280 @@ func TestALPRandomRoundtrip(t *testing.T) { // Data page v2 stores levels separately from encoded values. func TestALPDataPageV2(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + values := []float64{1.25, 2.5, 3.75, 100, 0.001} - got, _ := alpWriteReadDouble(t, values, nil, parquet.WithDataPageVersion(parquet.DataPageV2)) + defLevels := []int16{1, 0, 1, 1, 0, 1, 0, 1} + reader := writeALPFile(t, mem, schema.FieldList{ + schema.NewFloat64Node("value", parquet.Repetitions.Optional, -1), + }, func(cw file.ColumnChunkWriter) { + n, err := cw.(*file.Float64ColumnChunkWriter).WriteBatch(values, defLevels, nil) + require.NoError(t, err) + require.EqualValues(t, len(values), n) + }, parquet.WithDataPageVersion(parquet.DataPageV2), parquet.WithCompression(compress.Codecs.Zstd)) + defer reader.Close() + require.EqualValues(t, len(defLevels), reader.NumRows()) + chunk, err := reader.RowGroup(0).MetaData().ColumnChunk(0) + require.NoError(t, err) + assertALPDataPageEncoding(t, chunk) + require.Equal(t, compress.Codecs.Zstd, chunk.Compression()) + + pages, err := reader.RowGroup(0).GetColumnPageReader(0) + require.NoError(t, err) + defer pages.Close() + require.True(t, pages.Next()) + require.Equal(t, file.PageTypeDataPageV2, pages.Page().Type()) + page := pages.Page().(*file.DataPageV2) + require.Equal(t, parquet.Encodings.ALP, parquet.Encoding(page.Encoding())) + require.True(t, page.IsCompressed()) + require.EqualValues(t, len(defLevels), page.NumRows()) + require.EqualValues(t, len(defLevels)-len(values), page.NumNulls()) + require.Positive(t, page.DefinitionLevelByteLen()) + require.False(t, pages.Next()) + require.NoError(t, pages.Err()) + + got, gotDef := readALPDoubleColumn(t, reader, 0, len(defLevels)) + require.Equal(t, defLevels, gotDef) require.Equal(t, values, got) } -// alpSinglePageReader returns one page followed by the end of the chunk. -type alpSinglePageReader struct { - page file.Page - handed bool +func TestALPFileRoundtrip(t *testing.T) { + for _, tc := range []struct { + name string + values []float64 + }{ + { + name: "monetary", + values: []float64{1.23, 4.56, 7.89, 10.11, 12.13, 14.15, 16.17, 18.19}, + }, + { + name: "large_vector", + values: func() []float64 { + v := make([]float64, 2048) + for i := range v { + v[i] = 0.125 + float64(i)*0.25 + } + return v + }(), + }, + { + name: "integers", + values: []float64{1, 2, 3, 100, 200, 300, 1000, 2000}, + }, + { + name: "special_values", + values: []float64{0, math.Copysign(0, -1), 1.5, -1.5, math.Inf(1), math.Inf(-1), math.NaN()}, + }, + } { + t.Run(tc.name, func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + reader := writeALPFile(t, mem, schema.FieldList{ + schema.NewFloat64Node("value", parquet.Repetitions.Required, -1), + }, func(cw file.ColumnChunkWriter) { + n, err := cw.(*file.Float64ColumnChunkWriter).WriteBatch(tc.values, nil, nil) + require.NoError(t, err) + require.EqualValues(t, len(tc.values), n) + }) + defer reader.Close() + require.EqualValues(t, len(tc.values), reader.NumRows()) + chunk, err := reader.RowGroup(0).MetaData().ColumnChunk(0) + require.NoError(t, err) + assertALPDataPageEncoding(t, chunk) + + output, _ := readALPDoubleColumn(t, reader, 0, len(tc.values)) + require.Len(t, output, len(tc.values)) + + // Compare the bit patterns, so that a NaN or a negative zero has to come + // back exactly as it went in. + for i := range tc.values { + require.Equal(t, math.Float64bits(tc.values[i]), math.Float64bits(output[i]), + "index %d: got %v (bits %016x), want %v (bits %016x)", + i, output[i], math.Float64bits(output[i]), tc.values[i], math.Float64bits(tc.values[i])) + } + }) + } } -func (r *alpSinglePageReader) SetMaxPageHeaderSize(int) {} -func (r *alpSinglePageReader) Page() file.Page { return r.page } -func (r *alpSinglePageReader) Next() bool { - if r.handed { - return false +// A column that asks for ALP but keeps its dictionary falls back to PLAIN when +// the dictionary outgrows its limit before the first page. The footer must then +// list PLAIN, not ALP: a reader without ALP support may reject a file whose +// footer names it, though no page uses it. +func TestALPRequestedButDictionaryDroppedListsPlain(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + values := make([]float64, 20000) + for i := range values { + values[i] = float64(i) / 7 + } + reader := writeALPFile(t, mem, schema.FieldList{ + schema.NewFloat64Node("value", parquet.Repetitions.Required, -1), + }, func(cw file.ColumnChunkWriter) { + n, err := cw.(*file.Float64ColumnChunkWriter).WriteBatch(values, nil, nil) + require.NoError(t, err) + require.EqualValues(t, len(values), n) + }, parquet.WithDictionaryDefault(true), + parquet.WithEncodingFor("value", parquet.Encodings.ALP), + parquet.WithDictionaryPageSizeLimit(1024)) + defer reader.Close() + chunk, err := reader.RowGroup(0).MetaData().ColumnChunk(0) + require.NoError(t, err) + dataPageEntries := 0 + for _, s := range chunk.EncodingStats() { + require.NotEqual(t, parquet.Encodings.ALP, s.Encoding, "no page should be ALP") + if s.PageType == file.PageTypeDataPage || s.PageType == file.PageTypeDataPageV2 { + dataPageEntries++ + require.Equal(t, parquet.Encodings.Plain, s.Encoding) + } } - r.handed = true - return r.page != nil + require.Positive(t, dataPageEntries) + // This test also checks that the footer does not advertise an unused codec. + require.NotContains(t, chunk.Encodings(), parquet.Encodings.ALP) + require.Contains(t, chunk.Encodings(), parquet.Encodings.Plain) + + got, _ := readALPDoubleColumn(t, reader, 0, len(values)) + require.Equal(t, values, got) } -func (r *alpSinglePageReader) Err() error { return nil } -func (r *alpSinglePageReader) Reset(parquet.BufferedReader, int64, compress.Compression, *file.CryptoContext) { +// The vector size is a writer property, so it has to reach each column's +// encoder: the pages of a column must declare the size set for it, and still +// read back. +func TestALPFileVectorSize(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + values := make([]float64, 5000) + for i := range values { + values[i] = float64(i%1234) / 100 + } + + var fields schema.FieldList + for _, name := range []string{"small", "default", "large"} { + fields = append(fields, schema.NewFloat64Node(name, parquet.Repetitions.Required, -1)) + } + reader := writeALPFile(t, mem, fields, func(cw file.ColumnChunkWriter) { + n, err := cw.(*file.Float64ColumnChunkWriter).WriteBatch(values, nil, nil) + require.NoError(t, err) + require.EqualValues(t, len(values), n) + }, parquet.WithALPVectorSizeFor("small", 32), + parquet.WithALPVectorSizeFor("large", 4096)) + defer reader.Close() + + for col, logSize := range []byte{5, 10, 12} { + chunk, err := reader.RowGroup(0).MetaData().ColumnChunk(col) + require.NoError(t, err) + assertALPDataPageEncoding(t, chunk) + + func() { + pages, err := reader.RowGroup(0).GetColumnPageReader(col) + require.NoError(t, err) + defer pages.Close() + pageCount := 0 + for pages.Next() { + page := pages.Page() + require.Equal(t, file.PageTypeDataPage, page.Type()) + require.Equal(t, parquet.Encodings.ALP, parquet.Encoding(page.Encoding())) + // Required columns have no levels before the ALP header. + require.GreaterOrEqual(t, len(page.Data()), 3) + require.Equal(t, logSize, page.Data()[2], "column %d", col) + pageCount++ + } + require.NoError(t, pages.Err()) + require.Positive(t, pageCount) + }() + + got, _ := readALPDoubleColumn(t, reader, col, len(values)) + require.Equal(t, values, got, "column %d", col) + } } -func (r *alpSinglePageReader) GetDictionaryPage() (*file.DictionaryPage, error) { return nil, nil } -func (r *alpSinglePageReader) SeekToPageWithRow(int64) error { return nil } -func (r *alpSinglePageReader) Close() error { return nil } +// A column whose values outgrow the data page size is written as several pages. +// Each one carries its own ALP header, offset array and vectors, so a page that +// repeated an earlier page's values would decode to the wrong numbers here. +func TestALPFileRoundtripManyPages(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) -// Reject ALP before non-floating-point decoder traits can panic. -func TestALPOnIntegerColumnIsRejected(t *testing.T) { - node, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Int32, -1, -1) - require.NoError(t, err) - descr := schema.NewColumn(node, 0, 0) + const numValues = 40000 - body := memory.NewBufferBytes(make([]byte, 64)) - page := file.NewDataPageV1(body, 4, parquet.Encodings.ALP, - parquet.Encodings.RLE, parquet.Encodings.RLE, int32(body.Len())) + values := make([]float64, numValues) + for i := range values { + values[i] = float64(i) / 100 + } - r := file.NewColumnReader(descr, &alpSinglePageReader{page: page}, memory.DefaultAllocator, nil) - total, read, err := r.(*file.Int32ColumnChunkReader).ReadBatch(4, make([]int32, 4), nil, nil) + reader := writeALPFile(t, mem, schema.FieldList{ + schema.NewFloat64Node("value", parquet.Repetitions.Required, -1), + }, func(cw file.ColumnChunkWriter) { + n, err := cw.(*file.Float64ColumnChunkWriter).WriteBatch(values, nil, nil) + require.NoError(t, err) + require.EqualValues(t, numValues, n) + }, parquet.WithDataPageSize(16*1024)) + defer reader.Close() + + require.EqualValues(t, numValues, reader.NumRows()) + chunk, err := reader.RowGroup(0).MetaData().ColumnChunk(0) require.NoError(t, err) - require.Zero(t, total) - require.Zero(t, read) - // ReadBatch stops rather than returning the failure, as it does for any page - // it cannot open, so the reason is on the reader. - require.ErrorContains(t, r.Err(), "only float and double support ALP encoding") -} + assertALPDataPageEncoding(t, chunk) -// A mismatched page count must return an error instead of a zero-length read, -// which the column reader would retry indefinitely. -func TestALPPageClaimingTooManyValues(t *testing.T) { - node, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Double, -1, -1) + pages, err := reader.RowGroup(0).GetColumnPageReader(0) require.NoError(t, err) - descr := schema.NewColumn(node, 0, 0) + defer pages.Close() + dataPages := 0 + for pages.Next() { + require.Equal(t, file.PageTypeDataPage, pages.Page().Type()) + require.Equal(t, parquet.Encodings.ALP, parquet.Encoding(pages.Page().Encoding())) + dataPages++ + } + require.NoError(t, pages.Err()) + require.Greater(t, dataPages, 1, "the roundtrip must span multiple data pages") - enc := encoding.NewEncoder(parquet.Types.Double, parquet.Encodings.ALP, false, descr, memory.DefaultAllocator) - enc.(encoding.Float64Encoder).Put([]float64{1.5, 2.5}) - buf, err := enc.FlushValues() + // Read in batches that do not divide the vector size, so a batch boundary + // lands inside a vector and inside a page. + output, _ := readALPDoubleColumn(t, reader, 0, 700) + require.Equal(t, values, output) +} + +func TestALPFloat32FileRoundtrip(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + values := []float32{1.23, 4.56, 7.89, 0, -1.5, 100.0, 0.001} + + reader := writeALPFile(t, mem, schema.FieldList{ + schema.NewFloat32Node("value", parquet.Repetitions.Required, -1), + }, func(cw file.ColumnChunkWriter) { + n, err := cw.(*file.Float32ColumnChunkWriter).WriteBatch(values, nil, nil) + require.NoError(t, err) + require.EqualValues(t, len(values), n) + }) + defer reader.Close() + require.EqualValues(t, len(values), reader.NumRows()) + + rgr := reader.RowGroup(0) + chunk, err := rgr.MetaData().ColumnChunk(0) + require.NoError(t, err) + assertALPDataPageEncoding(t, chunk) + cr, err := rgr.Column(0) require.NoError(t, err) + defer cr.Close() - body := memory.NewBufferBytes(buf.Bytes()) - page := file.NewDataPageV1(body, 8, parquet.Encodings.ALP, - parquet.Encodings.RLE, parquet.Encodings.RLE, int32(body.Len())) + f32Reader, ok := cr.(*file.Float32ColumnChunkReader) + require.True(t, ok) - r := file.NewColumnReader(descr, &alpSinglePageReader{page: page}, memory.DefaultAllocator, nil) - _, read, err := r.(*file.Float64ColumnChunkReader).ReadBatch(8, make([]float64, 8), nil, nil) - require.ErrorContains(t, err, "ALP page holds 2 values") - require.Equal(t, 2, read, "the values the page does hold still come back") + output := make([]float32, len(values)) + total, read, err := f32Reader.ReadBatch(int64(len(values)), output, nil, nil) + require.NoError(t, err) + require.EqualValues(t, len(values), total) + require.EqualValues(t, len(values), read) + require.False(t, f32Reader.HasNext()) + require.NoError(t, f32Reader.Err()) + + for i := range values { + require.Equal(t, math.Float32bits(values[i]), math.Float32bits(output[i]), + "index %d: got %v (bits %08x), want %v (bits %08x)", + i, output[i], math.Float32bits(output[i]), values[i], math.Float32bits(values[i])) + } } diff --git a/parquet/file/column_reader_test.go b/parquet/file/column_reader_test.go index 2786f9e35..a0e3557d5 100644 --- a/parquet/file/column_reader_test.go +++ b/parquet/file/column_reader_test.go @@ -34,6 +34,7 @@ import ( "github.com/apache/arrow-go/v18/internal/utils" "github.com/apache/arrow-go/v18/parquet" "github.com/apache/arrow-go/v18/parquet/file" + "github.com/apache/arrow-go/v18/parquet/internal/encoding" "github.com/apache/arrow-go/v18/parquet/internal/testutils" "github.com/apache/arrow-go/v18/parquet/pqarrow" "github.com/apache/arrow-go/v18/parquet/schema" @@ -854,6 +855,72 @@ func TestFullSeekRow(t *testing.T) { } } +// Reject ALP before non-floating-point decoder traits can panic. +func TestALPOnIntegerColumnIsRejected(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + node, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Int32, -1, -1) + require.NoError(t, err) + descr := schema.NewColumn(node, 0, 0) + + body := memory.NewResizableBuffer(mem) + body.Resize(64) + clear(body.Bytes()) + page := file.NewDataPageV1(body, 4, parquet.Encodings.ALP, + parquet.Encodings.RLE, parquet.Encodings.RLE, int32(body.Len())) + defer page.Release() + + pages := new(testutils.MockPageReader) + pages.Test(t) + pages.TestData().Set("pages", []file.Page{page}) + pages.On("Err").Return((error)(nil)) + r := file.NewColumnReader(descr, pages, mem, nil) + defer r.Close() + total, read, err := r.(*file.Int32ColumnChunkReader).ReadBatch(4, make([]int32, 4), nil, nil) + require.NoError(t, err) + require.Zero(t, total) + require.Zero(t, read) + // ReadBatch stops rather than returning the failure, as it does for any page + // it cannot open, so the reason is on the reader. + require.ErrorContains(t, r.Err(), "only float and double support ALP encoding") +} + +// A mismatched page count must return an error instead of a zero-length read, +// which the column reader would retry indefinitely. +func TestALPPageClaimingTooManyValues(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + node, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Double, -1, -1) + require.NoError(t, err) + descr := schema.NewColumn(node, 0, 0) + + enc := encoding.NewEncoder(parquet.Types.Double, parquet.Encodings.ALP, false, descr, mem) + defer enc.Release() + enc.(encoding.Float64Encoder).Put([]float64{1.5, 2.5}) + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + body := memory.NewBufferBytes(buf.Bytes()) + page := file.NewDataPageV1(body, 8, parquet.Encodings.ALP, + parquet.Encodings.RLE, parquet.Encodings.RLE, int32(body.Len())) + defer page.Release() + + pages := new(testutils.MockPageReader) + pages.Test(t) + pages.TestData().Set("pages", []file.Page{page}) + pages.On("Err").Return((error)(nil)) + r := file.NewColumnReader(descr, pages, mem, nil) + defer r.Close() + values := make([]float64, 8) + _, read, err := r.(*file.Float64ColumnChunkReader).ReadBatch(8, values, nil, nil) + require.ErrorContains(t, err, "ALP page holds 2 values") + require.Equal(t, 2, read, "the values the page does hold still come back") + require.Equal(t, []float64{1.5, 2.5}, values[:read]) +} + func BenchmarkReadInt32Column(b *testing.B) { // generate parquet with RLE-dictionary encoded int32 column tempdir := b.TempDir() diff --git a/parquet/file/file_reader_test.go b/parquet/file/file_reader_test.go index 2530ce447..e33b3c652 100644 --- a/parquet/file/file_reader_test.go +++ b/parquet/file/file_reader_test.go @@ -1068,7 +1068,10 @@ func TestReadALPConformanceFile(t *testing.T) { } require.DirExists(t, dir) - props := parquet.NewReaderProperties(memory.DefaultAllocator) + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + props := parquet.NewReaderProperties(mem) rdr, err := file.OpenParquetFile(path.Join(dir, alpConformanceFile), false, file.WithReadProps(props)) require.NoError(t, err) @@ -1095,8 +1098,7 @@ func TestReadALPConformanceFile(t *testing.T) { for rg := range rdr.NumRowGroups() { chunk, err := rdr.RowGroup(rg).MetaData().ColumnChunk(alpCol) require.NoError(t, err) - assert.Contains(t, chunk.Encodings(), parquet.Encodings.ALP, - "row group %d of %s should be ALP encoded", rg, tc.alp) + assertALPDataPageEncoding(t, chunk) } want := readFloatColumn(t, rdr, alpColumnIndex(t, rdr, tc.plain)) @@ -1145,6 +1147,7 @@ func readFloatColumn(t *testing.T, rdr *file.Reader, col int) []floatBitsValue { for rg := range rdr.NumRowGroups() { cr, err := rdr.RowGroup(rg).Column(col) require.NoError(t, err) + defer cr.Close() switch r := cr.(type) { case *file.Float32ColumnChunkReader: diff --git a/parquet/file/file_writer_test.go b/parquet/file/file_writer_test.go index c8f6cf6c2..870807793 100644 --- a/parquet/file/file_writer_test.go +++ b/parquet/file/file_writer_test.go @@ -611,323 +611,6 @@ func TestBatchedByteStreamSplitFileRoundtrip(t *testing.T) { require.NoError(t, rdr.Close()) } -func TestALPFileRoundtrip(t *testing.T) { - for _, tc := range []struct { - name string - values []float64 - }{ - { - name: "monetary", - values: []float64{1.23, 4.56, 7.89, 10.11, 12.13, 14.15, 16.17, 18.19}, - }, - { - name: "large_vector", - values: func() []float64 { - v := make([]float64, 2048) - for i := range v { - v[i] = 0.125 + float64(i)*0.25 - } - return v - }(), - }, - { - name: "integers", - values: []float64{1, 2, 3, 100, 200, 300, 1000, 2000}, - }, - { - name: "special_values", - values: []float64{0, math.Copysign(0, -1), 1.5, -1.5, math.Inf(1), math.Inf(-1), math.NaN()}, - }, - } { - t.Run(tc.name, func(t *testing.T) { - props := parquet.NewWriterProperties( - parquet.WithEncoding(parquet.Encodings.ALP), - parquet.WithALPEncoding(true), - parquet.WithDictionaryDefault(false), - ) - - field, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Double, -1, -1) - require.NoError(t, err) - - sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, schema.FieldList{field}, 0) - require.NoError(t, err) - - sink := encoding.NewBufferWriter(0, memory.DefaultAllocator) - writer := file.NewParquetWriter(sink, sc, file.WithWriterProps(props)) - - rgw := writer.AppendRowGroup() - cw, err := rgw.NextColumn() - require.NoError(t, err) - - f64Writer, ok := cw.(*file.Float64ColumnChunkWriter) - require.True(t, ok) - - nVals, err := f64Writer.WriteBatch(tc.values, nil, nil) - require.NoError(t, err) - require.EqualValues(t, len(tc.values), nVals) - - require.NoError(t, cw.Close()) - require.NoError(t, rgw.Close()) - require.NoError(t, writer.Close()) - - reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes())) - require.NoError(t, err) - - require.Equal(t, 1, reader.NumRowGroups()) - require.EqualValues(t, len(tc.values), reader.NumRows()) - - rgr := reader.RowGroup(0) - cr, err := rgr.Column(0) - require.NoError(t, err) - - f64Reader, ok := cr.(*file.Float64ColumnChunkReader) - require.True(t, ok) - - output := make([]float64, len(tc.values)) - total, read, err := f64Reader.ReadBatch(int64(len(tc.values)), output, nil, nil) - require.NoError(t, err) - require.EqualValues(t, len(tc.values), total) - require.EqualValues(t, len(tc.values), read) - - // Compare the bit patterns, so that a NaN or a negative zero has to come - // back exactly as it went in. - for i := range tc.values { - assert.Equal(t, math.Float64bits(tc.values[i]), math.Float64bits(output[i]), - "index %d: got %v (bits %016x), want %v (bits %016x)", - i, output[i], math.Float64bits(output[i]), tc.values[i], math.Float64bits(tc.values[i])) - } - - require.NoError(t, reader.Close()) - }) - } -} - -// A column that asks for ALP but keeps its dictionary falls back to PLAIN when -// the dictionary outgrows its limit before the first page. The footer must then -// list PLAIN, not ALP: a reader without ALP support may reject a file whose -// footer names it, though no page uses it. -func TestALPRequestedButDictionaryDroppedListsPlain(t *testing.T) { - values := make([]float64, 20000) - for i := range values { - values[i] = float64(i) / 7 - } - props := parquet.NewWriterProperties( - parquet.WithALPEncoding(true), - parquet.WithEncodingFor("value", parquet.Encodings.ALP), - parquet.WithDictionaryPageSizeLimit(1024)) - sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, schema.FieldList{ - schema.NewFloat64Node("value", parquet.Repetitions.Required, -1), - }, -1) - require.NoError(t, err) - - var buf bytes.Buffer - writer := file.NewParquetWriter(&buf, sc, file.WithWriterProps(props)) - rgw := writer.AppendRowGroup() - cw, err := rgw.NextColumn() - require.NoError(t, err) - _, err = cw.(*file.Float64ColumnChunkWriter).WriteBatch(values, nil, nil) - require.NoError(t, err) - require.NoError(t, cw.Close()) - require.NoError(t, rgw.Close()) - require.NoError(t, writer.Close()) - - reader, err := file.NewParquetReader(bytes.NewReader(buf.Bytes())) - require.NoError(t, err) - defer reader.Close() - chunk, err := reader.RowGroup(0).MetaData().ColumnChunk(0) - require.NoError(t, err) - for _, s := range chunk.EncodingStats() { - require.NotEqual(t, parquet.Encodings.ALP, s.Encoding, "no page should be ALP") - } - assert.NotContains(t, chunk.Encodings(), parquet.Encodings.ALP) - assert.Contains(t, chunk.Encodings(), parquet.Encodings.Plain) -} - -// The vector size is a writer property, so it has to reach each column's -// encoder: the pages of a column must declare the size set for it, and still -// read back. -func TestALPFileVectorSize(t *testing.T) { - values := make([]float64, 5000) - for i := range values { - values[i] = float64(i%1234) / 100 - } - - props := parquet.NewWriterProperties( - parquet.WithEncoding(parquet.Encodings.ALP), - parquet.WithALPEncoding(true), - parquet.WithDictionaryDefault(false), - parquet.WithALPVectorSizeFor("small", 32), - parquet.WithALPVectorSizeFor("large", 4096), - ) - var fields schema.FieldList - for _, name := range []string{"small", "default", "large"} { - fields = append(fields, schema.NewFloat64Node(name, parquet.Repetitions.Required, -1)) - } - sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, fields, -1) - require.NoError(t, err) - - var buf bytes.Buffer - writer := file.NewParquetWriter(&buf, sc, file.WithWriterProps(props)) - rgw := writer.AppendRowGroup() - for range fields { - cw, err := rgw.NextColumn() - require.NoError(t, err) - _, err = cw.(*file.Float64ColumnChunkWriter).WriteBatch(values, nil, nil) - require.NoError(t, err) - require.NoError(t, cw.Close()) - } - require.NoError(t, rgw.Close()) - require.NoError(t, writer.Close()) - - reader, err := file.NewParquetReader(bytes.NewReader(buf.Bytes())) - require.NoError(t, err) - defer reader.Close() - - for col, logSize := range []byte{5, 10, 12} { - pages, err := reader.RowGroup(0).GetColumnPageReader(col) - require.NoError(t, err) - require.True(t, pages.Next(), "column %d has a page", col) - // A required column has no levels, so the page data starts with the - // ALP header, whose third byte is log2 of the vector size. - assert.Equal(t, logSize, pages.Page().Data()[2], "column %d", col) - - cr, err := reader.RowGroup(0).Column(col) - require.NoError(t, err) - got := make([]float64, len(values)) - total, _, err := cr.(*file.Float64ColumnChunkReader).ReadBatch(int64(len(values)), got, nil, nil) - require.NoError(t, err) - require.EqualValues(t, len(values), total) - assert.Equal(t, values, got, "column %d", col) - } -} - -// A column whose values outgrow the data page size is written as several pages. -// Each one carries its own ALP header, offset array and vectors, so a page that -// repeated an earlier page's values would decode to the wrong numbers here. -func TestALPFileRoundtripManyPages(t *testing.T) { - const numValues = 40000 - - values := make([]float64, numValues) - for i := range values { - values[i] = float64(i) / 100 - } - - props := parquet.NewWriterProperties( - parquet.WithEncoding(parquet.Encodings.ALP), - parquet.WithALPEncoding(true), - parquet.WithDictionaryDefault(false), - // Small enough that 1024-value vectors fill a page every few vectors. - parquet.WithDataPageSize(16*1024), - ) - - field, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Double, -1, -1) - require.NoError(t, err) - - sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, schema.FieldList{field}, 0) - require.NoError(t, err) - - sink := encoding.NewBufferWriter(0, memory.DefaultAllocator) - writer := file.NewParquetWriter(sink, sc, file.WithWriterProps(props)) - - rgw := writer.AppendRowGroup() - cw, err := rgw.NextColumn() - require.NoError(t, err) - - f64Writer, ok := cw.(*file.Float64ColumnChunkWriter) - require.True(t, ok) - - nVals, err := f64Writer.WriteBatch(values, nil, nil) - require.NoError(t, err) - require.EqualValues(t, numValues, nVals) - - require.NoError(t, cw.Close()) - require.NoError(t, rgw.Close()) - require.NoError(t, writer.Close()) - - reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes())) - require.NoError(t, err) - defer reader.Close() - - rgr := reader.RowGroup(0) - cr, err := rgr.Column(0) - require.NoError(t, err) - - f64Reader, ok := cr.(*file.Float64ColumnChunkReader) - require.True(t, ok) - - // Read in batches that do not divide the vector size, so a batch boundary - // lands inside a vector and inside a page. - output := make([]float64, numValues) - for read := 0; read < numValues; { - _, n, err := f64Reader.ReadBatch(700, output[read:], nil, nil) - require.NoError(t, err) - require.NotZero(t, n, "read stalled at %d of %d values", read, numValues) - read += int(n) - } - require.Equal(t, values, output) -} - -func TestALPFloat32FileRoundtrip(t *testing.T) { - values := []float32{1.23, 4.56, 7.89, 0, -1.5, 100.0, 0.001} - - props := parquet.NewWriterProperties( - parquet.WithEncoding(parquet.Encodings.ALP), - parquet.WithALPEncoding(true), - parquet.WithDictionaryDefault(false), - ) - - field, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.Float, -1, -1) - require.NoError(t, err) - - sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, schema.FieldList{field}, 0) - require.NoError(t, err) - - sink := encoding.NewBufferWriter(0, memory.DefaultAllocator) - writer := file.NewParquetWriter(sink, sc, file.WithWriterProps(props)) - - rgw := writer.AppendRowGroup() - cw, err := rgw.NextColumn() - require.NoError(t, err) - - f32Writer, ok := cw.(*file.Float32ColumnChunkWriter) - require.True(t, ok) - - nVals, err := f32Writer.WriteBatch(values, nil, nil) - require.NoError(t, err) - require.EqualValues(t, len(values), nVals) - - require.NoError(t, cw.Close()) - require.NoError(t, rgw.Close()) - require.NoError(t, writer.Close()) - - reader, err := file.NewParquetReader(bytes.NewReader(sink.Bytes())) - require.NoError(t, err) - - require.Equal(t, 1, reader.NumRowGroups()) - require.EqualValues(t, len(values), reader.NumRows()) - - rgr := reader.RowGroup(0) - cr, err := rgr.Column(0) - require.NoError(t, err) - - f32Reader, ok := cr.(*file.Float32ColumnChunkReader) - require.True(t, ok) - - output := make([]float32, len(values)) - total, read, err := f32Reader.ReadBatch(int64(len(values)), output, nil, nil) - require.NoError(t, err) - require.EqualValues(t, len(values), total) - require.EqualValues(t, len(values), read) - - for i := range values { - assert.Equal(t, math.Float32bits(values[i]), math.Float32bits(output[i]), - "index %d: got %v (bits %08x), want %v (bits %08x)", - i, output[i], math.Float32bits(output[i]), values[i], math.Float32bits(values[i])) - } - - require.NoError(t, reader.Close()) -} - func TestLZ4RawFileRoundtrip(t *testing.T) { input := []int64{ -1, 0, 1, 2, 3, 4, 5, 123456789, -123456789, diff --git a/parquet/internal/encoding/alp_encoding_benchmark_test.go b/parquet/internal/encoding/alp_encoding_benchmark_test.go index a138b5f40..6972d7a0e 100644 --- a/parquet/internal/encoding/alp_encoding_benchmark_test.go +++ b/parquet/internal/encoding/alp_encoding_benchmark_test.go @@ -57,6 +57,8 @@ func BenchmarkALPEncodingFloat32(b *testing.B) { values := alpBenchmarkValues[float32](sz, benchmark.exceptionInterval) encoder := encoding.NewEncoder(parquet.Types.Float, parquet.Encodings.ALP, false, nil, memory.DefaultAllocator).(encoding.Float32Encoder) + defer encoder.Release() + b.ReportAllocs() b.ResetTimer() b.SetBytes(int64(len(values) * arrow.Float32SizeBytes)) for b.Loop() { @@ -78,6 +80,8 @@ func BenchmarkALPEncodingFloat64(b *testing.B) { values := alpBenchmarkValues[float64](sz, benchmark.exceptionInterval) encoder := encoding.NewEncoder(parquet.Types.Double, parquet.Encodings.ALP, false, nil, memory.DefaultAllocator).(encoding.Float64Encoder) + defer encoder.Release() + b.ReportAllocs() b.ResetTimer() b.SetBytes(int64(len(values) * arrow.Float64SizeBytes)) for b.Loop() { @@ -100,12 +104,14 @@ func BenchmarkALPDecodingFloat32(b *testing.B) { output := make([]float32, sz) encoder := encoding.NewEncoder(parquet.Types.Float, parquet.Encodings.ALP, false, nil, memory.DefaultAllocator).(encoding.Float32Encoder) + defer encoder.Release() encoder.Put(values) buf, _ := encoder.FlushValues() defer buf.Release() decoder := encoding.NewDecoder(parquet.Types.Float, parquet.Encodings.ALP, nil, memory.DefaultAllocator).(encoding.Float32Decoder) + b.ReportAllocs() b.ResetTimer() b.SetBytes(int64(len(values) * arrow.Float32SizeBytes)) for b.Loop() { @@ -131,12 +137,14 @@ func BenchmarkALPDecodingFloat64(b *testing.B) { output := make([]float64, sz) encoder := encoding.NewEncoder(parquet.Types.Double, parquet.Encodings.ALP, false, nil, memory.DefaultAllocator).(encoding.Float64Encoder) + defer encoder.Release() encoder.Put(values) buf, _ := encoder.FlushValues() defer buf.Release() decoder := encoding.NewDecoder(parquet.Types.Double, parquet.Encodings.ALP, nil, memory.DefaultAllocator).(encoding.Float64Decoder) + b.ReportAllocs() b.ResetTimer() b.SetBytes(int64(len(values) * arrow.Float64SizeBytes)) for b.Loop() { @@ -162,12 +170,14 @@ func BenchmarkALPDecodingFloat64Batched(b *testing.B) { output := make([]float64, batchSize) encoder := encoding.NewEncoder(parquet.Types.Double, parquet.Encodings.ALP, false, nil, memory.DefaultAllocator).(encoding.Float64Encoder) + defer encoder.Release() encoder.Put(values) buf, _ := encoder.FlushValues() defer buf.Release() decoder := encoding.NewDecoder(parquet.Types.Double, parquet.Encodings.ALP, nil, memory.DefaultAllocator).(encoding.Float64Decoder) + b.ReportAllocs() b.ResetTimer() b.SetBytes(int64(len(values) * arrow.Float64SizeBytes)) for b.Loop() { diff --git a/parquet/internal/encoding/alp_encoding_test.go b/parquet/internal/encoding/alp_encoding_test.go index 49452743b..3bc979de4 100644 --- a/parquet/internal/encoding/alp_encoding_test.go +++ b/parquet/internal/encoding/alp_encoding_test.go @@ -113,266 +113,76 @@ func newFloat64Column() *schema.Column { return schema.NewColumn(schema.NewFloat64Node("float64", parquet.Repetitions.Required, -1), 0, 0) } -func alpFloat32RoundTrip(t *testing.T, values []float32) { +func alpRoundTrip[T alpFloat](t *testing.T, values []T) { t.Helper() - col := newFloat32Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) - enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + enc := newALPEncoder[T](format.Encoding_ALP, nil, mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) defer buf.Release() - dec := newALPDecoder[float32](format.Encoding_ALP, col) - err = dec.SetData(len(values), buf.Bytes()) - require.NoError(t, err) - - out := make([]float32, len(values)) - n, err := dec.Decode(out) - require.NoError(t, err) - assert.Equal(t, len(values), n) - + out := alpDecodePage[T](t, buf, len(values)) for i, v := range values { - assert.Equal(t, math.Float32bits(v), math.Float32bits(out[i]), - "index %d: got %v (bits %08x), want %v (bits %08x)", - i, out[i], math.Float32bits(out[i]), v, math.Float32bits(v)) + switch want := any(v).(type) { + case float32: + assert.Equal(t, math.Float32bits(want), math.Float32bits(float32(out[i])), + "index %d: got %v, want %v", i, out[i], v) + case float64: + assert.Equal(t, math.Float64bits(want), math.Float64bits(float64(out[i])), + "index %d: got %v, want %v", i, out[i], v) + } } } -func alpFloat64RoundTrip(t *testing.T, values []float64) { +func testALPFloatRoundTrip[T alpFloat](t *testing.T) { t.Helper() - col := newFloat64Column() - mem := memory.DefaultAllocator - - enc := newALPEncoder[float64](format.Encoding_ALP, col, mem) - enc.Put(values) - buf, err := enc.FlushValues() - require.NoError(t, err) - defer buf.Release() - - dec := newALPDecoder[float64](format.Encoding_ALP, col) - err = dec.SetData(len(values), buf.Bytes()) - require.NoError(t, err) - - out := make([]float64, len(values)) - n, err := dec.Decode(out) - require.NoError(t, err) - assert.Equal(t, len(values), n) - - for i, v := range values { - assert.Equal(t, math.Float64bits(v), math.Float64bits(out[i]), - "index %d: got %v (bits %016x), want %v (bits %016x)", - i, out[i], math.Float64bits(out[i]), v, math.Float64bits(v)) - } -} - -func TestALPFloat32MonetaryRoundTrip(t *testing.T) { - values := make([]float32, 2000) - for i := range values { - values[i] = float32(i) * 0.01 - } - alpFloat32RoundTrip(t, values) -} - -func TestALPFloat64MonetaryRoundTrip(t *testing.T) { - values := make([]float64, 2000) - for i := range values { - values[i] = float64(i) * 0.01 - } - alpFloat64RoundTrip(t, values) -} - -func TestALPFloat32IntegerRoundTrip(t *testing.T) { - values := make([]float32, 1500) - for i := range values { - values[i] = float32(i) - } - alpFloat32RoundTrip(t, values) -} - -func TestALPFloat64IntegerRoundTrip(t *testing.T) { - values := make([]float64, 1500) - for i := range values { - values[i] = float64(i) - } - alpFloat64RoundTrip(t, values) -} - -func TestALPFloat32ExceptionsRoundTrip(t *testing.T) { - negZero := alpNegZero[float32]() - inf := float32(math.Inf(1)) - ninf := float32(math.Inf(-1)) - - values := []float32{ - 1.23, 4.56, inf, 7.89, negZero, 10.11, - ninf, 12.13, 14.15, 16.17, - } - alpFloat32RoundTrip(t, values) -} - -func TestALPFloat64ExceptionsRoundTrip(t *testing.T) { - negZero := alpNegZero[float64]() - inf := math.Inf(1) - ninf := math.Inf(-1) - - values := []float64{ - 1.23, 4.56, inf, 7.89, negZero, 10.11, - ninf, 12.13, 14.15, 16.17, - } - alpFloat64RoundTrip(t, values) -} - -func TestALPFloat32AllExceptionsRoundTrip(t *testing.T) { - negZero := alpNegZero[float32]() - inf := float32(math.Inf(1)) - ninf := float32(math.Inf(-1)) - - values := []float32{inf, ninf, negZero, inf, ninf, negZero, inf, ninf} - alpFloat32RoundTrip(t, values) -} - -func TestALPFloat64AllExceptionsRoundTrip(t *testing.T) { - negZero := alpNegZero[float64]() - inf := math.Inf(1) - ninf := math.Inf(-1) - - values := []float64{inf, ninf, negZero, inf, ninf, negZero, inf, ninf} - alpFloat64RoundTrip(t, values) -} - -func TestALPFloat32SingleValueRoundTrip(t *testing.T) { - alpFloat32RoundTrip(t, []float32{42.0}) - alpFloat32RoundTrip(t, []float32{0.0}) - alpFloat32RoundTrip(t, []float32{-1.23}) -} - -func TestALPFloat64SingleValueRoundTrip(t *testing.T) { - alpFloat64RoundTrip(t, []float64{42.0}) - alpFloat64RoundTrip(t, []float64{0.0}) - alpFloat64RoundTrip(t, []float64{-1.23}) -} - -func TestALPFloat32PartialVectorRoundTrip(t *testing.T) { - for _, count := range []int{1, 2, 3, 7, 100, 500, 1023} { - t.Run(fmt.Sprintf("count_%d", count), func(t *testing.T) { - values := make([]float32, count) - for i := range values { - values[i] = float32(i) * 0.1 - } - alpFloat32RoundTrip(t, values) - }) - } -} - -func TestALPFloat64PartialVectorRoundTrip(t *testing.T) { - for _, count := range []int{1, 2, 3, 7, 100, 500, 1023} { - t.Run(fmt.Sprintf("count_%d", count), func(t *testing.T) { - values := make([]float64, count) - for i := range values { - values[i] = float64(i) * 0.1 - } - alpFloat64RoundTrip(t, values) - }) - } -} - -func TestALPFloat32ExactVectorSizeRoundTrip(t *testing.T) { - values := make([]float32, 1024) - for i := range values { - values[i] = float32(i) * 0.01 - } - alpFloat32RoundTrip(t, values) -} - -func TestALPFloat64ExactVectorSizeRoundTrip(t *testing.T) { - values := make([]float64, 1024) - for i := range values { - values[i] = float64(i) * 0.01 - } - alpFloat64RoundTrip(t, values) -} - -func TestALPFloat32MultipleVectorsRoundTrip(t *testing.T) { - for _, count := range []int{1025, 2048, 3000, 5000, 10000} { - t.Run(fmt.Sprintf("count_%d", count), func(t *testing.T) { - values := make([]float32, count) - for i := range values { - values[i] = float32(i) * 0.01 - } - alpFloat32RoundTrip(t, values) - }) + negZero := alpNegZero[T]() + inf, ninf := T(math.Inf(1)), T(math.Inf(-1)) + exceptions := []T{1.23, 4.56, inf, 7.89, negZero, 10.11, ninf, 12.13, 14.15, 16.17} + allExceptions := []T{inf, ninf, negZero, inf, ninf, negZero, inf, ninf} + tests := []struct { + name string + counts []int + value func(int) T + }{ + {"Monetary", []int{2000}, func(i int) T { return T(i) * 0.01 }}, + {"Integer", []int{1500}, func(i int) T { return T(i) }}, + {"Exceptions", []int{len(exceptions)}, func(i int) T { return exceptions[i] }}, + {"AllExceptions", []int{len(allExceptions)}, func(i int) T { return allExceptions[i] }}, + {"SingleValue/Positive", []int{1}, func(int) T { return 42 }}, + {"SingleValue/Zero", []int{1}, func(int) T { return 0 }}, + {"SingleValue/Negative", []int{1}, func(int) T { return -1.23 }}, + {"PartialVector", []int{1, 2, 3, 7, 100, 500, 1023}, func(i int) T { return T(i) * 0.1 }}, + {"ExactVectorSize", []int{1024}, func(i int) T { return T(i) * 0.01 }}, + {"MultipleVectors", []int{1025, 2048, 3000, 5000, 10000}, func(i int) T { return T(i) * 0.01 }}, + {"ConstantValue", []int{1500}, func(int) T { return 3.14 }}, + {"ZeroValues", []int{1024}, func(int) T { return 0 }}, + {"NegativeValues", []int{1500}, func(i int) T { return -T(i) * 0.01 }}, + {"Scientific", []int{2000}, func(i int) T { return T(i)*0.001 + 273.15 }}, // Kelvin temperatures } -} -func TestALPFloat64MultipleVectorsRoundTrip(t *testing.T) { - for _, count := range []int{1025, 2048, 3000, 5000, 10000} { - t.Run(fmt.Sprintf("count_%d", count), func(t *testing.T) { - values := make([]float64, count) - for i := range values { - values[i] = float64(i) * 0.01 + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + for _, count := range tt.counts { + t.Run(fmt.Sprintf("count_%d", count), func(t *testing.T) { + values := make([]T, count) + for i := range values { + values[i] = tt.value(i) + } + alpRoundTrip(t, values) + }) } - alpFloat64RoundTrip(t, values) }) } } -func TestALPFloat32ConstantValueRoundTrip(t *testing.T) { - values := make([]float32, 1500) - for i := range values { - values[i] = 3.14 - } - alpFloat32RoundTrip(t, values) -} - -func TestALPFloat64ConstantValueRoundTrip(t *testing.T) { - values := make([]float64, 1500) - for i := range values { - values[i] = 3.14 - } - alpFloat64RoundTrip(t, values) -} - -func TestALPFloat32ZeroValuesRoundTrip(t *testing.T) { - values := make([]float32, 1024) - alpFloat32RoundTrip(t, values) -} - -func TestALPFloat64ZeroValuesRoundTrip(t *testing.T) { - values := make([]float64, 1024) - alpFloat64RoundTrip(t, values) -} - -func TestALPFloat32NegativeValuesRoundTrip(t *testing.T) { - values := make([]float32, 1500) - for i := range values { - values[i] = -float32(i) * 0.01 - } - alpFloat32RoundTrip(t, values) -} - -func TestALPFloat64NegativeValuesRoundTrip(t *testing.T) { - values := make([]float64, 1500) - for i := range values { - values[i] = -float64(i) * 0.01 - } - alpFloat64RoundTrip(t, values) -} - -func TestALPFloat32ScientificRoundTrip(t *testing.T) { - values := make([]float32, 2000) - for i := range values { - values[i] = float32(i)*0.001 + 273.15 // Kelvin temperatures - } - alpFloat32RoundTrip(t, values) -} - -func TestALPFloat64ScientificRoundTrip(t *testing.T) { - values := make([]float64, 2000) - for i := range values { - values[i] = float64(i)*0.001 + 273.15 - } - alpFloat64RoundTrip(t, values) +func TestALPFloatRoundTrip(t *testing.T) { + t.Run("Float32", testALPFloatRoundTrip[float32]) + t.Run("Float64", testALPFloatRoundTrip[float64]) } func TestALPFloat32ProgressiveDecoding(t *testing.T) { @@ -382,9 +192,11 @@ func TestALPFloat32ProgressiveDecoding(t *testing.T) { } col := newFloat32Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) @@ -423,9 +235,11 @@ func TestALPFloat64ProgressiveDecoding(t *testing.T) { } col := newFloat64Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) enc := newALPEncoder[float64](format.Encoding_ALP, col, mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) @@ -464,9 +278,11 @@ func TestALPFloat32Discard(t *testing.T) { } col := newFloat32Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) @@ -498,9 +314,11 @@ func TestALPFloat64Discard(t *testing.T) { } col := newFloat64Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) enc := newALPEncoder[float64](format.Encoding_ALP, col, mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) @@ -529,9 +347,11 @@ func TestALPDiscardMoreThanAvailable(t *testing.T) { values := []float32{1.0, 2.0, 3.0} col := newFloat32Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) @@ -548,7 +368,8 @@ func TestALPDiscardMoreThanAvailable(t *testing.T) { func TestALPFloat32MultiplePages(t *testing.T) { col := newFloat32Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) dec := newALPDecoder[float32](format.Encoding_ALP, col) @@ -559,9 +380,11 @@ func TestALPFloat32MultiplePages(t *testing.T) { } enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) + defer buf.Release() err = dec.SetData(len(values), buf.Bytes()) require.NoError(t, err) @@ -575,14 +398,13 @@ func TestALPFloat32MultiplePages(t *testing.T) { assert.Equal(t, math.Float32bits(v), math.Float32bits(out[i]), "page %d, index %d: mismatch", page, i) } - - buf.Release() } } func TestALPFloat64MultiplePages(t *testing.T) { col := newFloat64Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) dec := newALPDecoder[float64](format.Encoding_ALP, col) @@ -593,9 +415,11 @@ func TestALPFloat64MultiplePages(t *testing.T) { } enc := newALPEncoder[float64](format.Encoding_ALP, col, mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) + defer buf.Release() err = dec.SetData(len(values), buf.Bytes()) require.NoError(t, err) @@ -609,16 +433,16 @@ func TestALPFloat64MultiplePages(t *testing.T) { assert.Equal(t, math.Float64bits(v), math.Float64bits(out[i]), "page %d, index %d: mismatch", page, i) } - - buf.Release() } } func TestALPEncoderReset(t *testing.T) { col := newFloat32Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + defer enc.Release() values1 := []float32{1.23, 4.56, 7.89} enc.Put(values1) @@ -655,37 +479,16 @@ func TestALPEncoderReset(t *testing.T) { } } -func TestALPEncoderType(t *testing.T) { - col32 := newFloat32Column() - col64 := newFloat64Column() - mem := memory.DefaultAllocator - - enc32 := newALPEncoder[float32](format.Encoding_ALP, col32, mem) - assert.Equal(t, parquet.Types.Float, enc32.Type()) - - enc64 := newALPEncoder[float64](format.Encoding_ALP, col64, mem) - assert.Equal(t, parquet.Types.Double, enc64.Type()) -} - -func TestALPDecoderType(t *testing.T) { - col32 := newFloat32Column() - col64 := newFloat64Column() - - dec32 := &alpDecoder[float32]{decoder: newDecoderBase(format.Encoding_ALP, col32)} - assert.Equal(t, parquet.Types.Float, dec32.Type()) - - dec64 := &alpDecoder[float64]{decoder: newDecoderBase(format.Encoding_ALP, col64)} - assert.Equal(t, parquet.Types.Double, dec64.Type()) -} - func TestALPEncoderPutSpaced(t *testing.T) { col := newFloat32Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) values := []float32{1.0, 2.0, 0.0, 4.0, 0.0, 6.0, 7.0, 0.0} validBits := []byte{0b01101011} // bits 0,1,3,5,6 set enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + defer enc.Release() enc.PutSpaced(values, validBits, 0) buf, err := enc.FlushValues() require.NoError(t, err) @@ -710,7 +513,8 @@ func TestALPEncoderPutSpaced(t *testing.T) { func TestALPHeaderFormat(t *testing.T) { col := newFloat32Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) values := make([]float32, 100) for i := range values { @@ -718,6 +522,7 @@ func TestALPHeaderFormat(t *testing.T) { } enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) @@ -734,7 +539,8 @@ func TestALPHeaderFormat(t *testing.T) { func TestALPHeaderFormatMultipleVectors(t *testing.T) { col := newFloat64Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) values := make([]float64, 3000) for i := range values { @@ -742,6 +548,7 @@ func TestALPHeaderFormatMultipleVectors(t *testing.T) { } enc := newALPEncoder[float64](format.Encoding_ALP, col, mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) @@ -838,7 +645,10 @@ func TestALPFloat64WidestPacking(t *testing.T) { values[i] = float64(int64(i-len(values)/2)*(1<<54) + (1 << 53)) } - enc := newALPEncoder[float64](format.Encoding_ALP, newFloat64Column(), memory.DefaultAllocator) + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + enc := newALPEncoder[float64](format.Encoding_ALP, newFloat64Column(), mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) @@ -850,7 +660,7 @@ func TestALPFloat64WidestPacking(t *testing.T) { assert.Equal(t, byte(64), page[vectorPos+alpVectorHeaderSize[float64]()-1], "bit width") assert.Zero(t, binary.LittleEndian.Uint16(page[vectorPos+2:]), "exception count") - alpFloat64RoundTrip(t, values) + alpRoundTrip(t, values) } // alpEncodeTestPage encodes values into a single-vector page and returns the @@ -858,11 +668,14 @@ func TestALPFloat64WidestPacking(t *testing.T) { func alpEncodeTestPage(t *testing.T, values []float32) (page []byte, vectorPos int) { t.Helper() - enc := newALPEncoder[float32](format.Encoding_ALP, newFloat32Column(), memory.DefaultAllocator) + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + enc := newALPEncoder[float32](format.Encoding_ALP, newFloat32Column(), mem) + defer enc.Release() enc.Put(values) buf, err := enc.FlushValues() require.NoError(t, err) - t.Cleanup(buf.Release) + defer buf.Release() page = append([]byte(nil), buf.Bytes()...) // One offset follows the header, and it holds the size of the offset array. @@ -965,9 +778,11 @@ func TestALPDecoderExceptionPositionOutsideVector(t *testing.T) { func TestALPPresetCacheBuilds(t *testing.T) { col := newFloat32Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + defer enc.Release() values := make([]float32, 9000) for i := range values { @@ -998,28 +813,84 @@ func TestALPPresetCacheBuilds(t *testing.T) { } } -func TestALPEstimatedDataEncodedSize(t *testing.T) { - col := newFloat32Column() - mem := memory.DefaultAllocator - - enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) +func testALPEstimatedDataEncodedSize[T alpFloat](t *testing.T, valueWidth int) { + t.Helper() + tests := []struct { + name string + count int + value func(int) T + allExceptions bool + }{ + {"Empty", 0, func(int) T { return 0 }, false}, + {"BufferedMonetary", 100, func(i int) T { return T(i) * 0.01 }, false}, + {"BufferedDecimals", 100, func(i int) T { return T(i) / 4 }, false}, + {"IntegerVector", alpDefaultVectorSize, func(i int) T { return T(i) }, false}, + {"ConstantVector", alpDefaultVectorSize, func(int) T { return 3.14 }, false}, + {"MonetaryVectors", 2 * alpDefaultVectorSize, func(i int) T { return T(i) * 0.01 }, false}, + {"DecimalsWithBufferedTail", alpDefaultVectorSize + 100, func(i int) T { return T(i) / 4 }, false}, + {"AllExceptionsVector", alpDefaultVectorSize, func(int) T { return T(math.Inf(1)) }, true}, + {"AllExceptionsBuffered", 100, func(int) T { return T(math.Inf(1)) }, true}, + {"AllExceptionsWithBufferedTail", alpDefaultVectorSize + 100, func(int) T { return T(math.Inf(1)) }, true}, + } - assert.Equal(t, int64(0), enc.EstimatedDataEncodedSize()) + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + enc := newALPEncoder[T](format.Encoding_ALP, nil, mem) + defer enc.Release() + assert.Zero(t, enc.EstimatedDataEncodedSize()) - values := make([]float32, 100) - for i := range values { - values[i] = float32(i) * 0.01 + values := make([]T, tt.count) + for i := range values { + values[i] = tt.value(i) + } + enc.Put(values) + estimate := enc.EstimatedDataEncodedSize() + buf, err := enc.FlushValues() + require.NoError(t, err) + defer buf.Release() + + const pageHeaderSize = 7 + require.GreaterOrEqual(t, buf.Len(), pageHeaderSize) + payloadSize := int64(buf.Len() - pageHeaderSize) + pending := tt.count % alpDefaultVectorSize + switch { + case pending == 0: + assert.Equal(t, payloadSize, estimate, "completed vectors have an exact size") + case tt.allExceptions: + // A pending all-exception vector needs one 4-byte offset, + // exponent/factor (2), exception count (2), frame of reference, + // bit width (1), and one 2-byte position per exception. + // There are no packed delta bits; the estimate already counts + // the verbatim values at their physical width. + overhead := int64(4 + 2 + 2 + valueWidth + 1 + 2*pending) + assert.Equal(t, payloadSize, estimate+overhead) + default: + assert.GreaterOrEqual(t, estimate, payloadSize, + "buffered decimal values should not exceed their raw-width estimate") + } + assert.Zero(t, enc.EstimatedDataEncodedSize(), "flushing starts a new page") + }) } - enc.Put(values) - size := enc.EstimatedDataEncodedSize() - assert.True(t, size > 0, "estimated size should be positive after Put") +} + +func TestALPEstimatedDataEncodedSize(t *testing.T) { + t.Run("Float32", func(t *testing.T) { + testALPEstimatedDataEncodedSize[float32](t, 4) + }) + t.Run("Float64", func(t *testing.T) { + testALPEstimatedDataEncodedSize[float64](t, 8) + }) } func TestALPFloat32IncrementalPut(t *testing.T) { col := newFloat32Column() - mem := memory.DefaultAllocator + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) enc := newALPEncoder[float32](format.Encoding_ALP, col, mem) + defer enc.Release() allValues := make([]float32, 2500) for i := range allValues { allValues[i] = float32(i) * 0.01 @@ -1070,7 +941,7 @@ func TestALPFloat32MixedDataPatterns(t *testing.T) { values = append(values, float32(i)*0.1) } - alpFloat32RoundTrip(t, values) + alpRoundTrip(t, values) } func TestALPFloat64MixedDataPatterns(t *testing.T) { @@ -1092,84 +963,13 @@ func TestALPFloat64MixedDataPatterns(t *testing.T) { values = append(values, float64(i)*0.1) } - alpFloat64RoundTrip(t, values) + alpRoundTrip(t, values) } -func TestALPFloat32ViaPublicAPI(t *testing.T) { - col := newFloat32Column() - mem := memory.DefaultAllocator - - values := make([]float32, 2000) - for i := range values { - values[i] = float32(i) * 0.01 - } - - enc := NewEncoder(parquet.Types.Float, parquet.Encodings.ALP, false, col, mem) - assert.Equal(t, parquet.Encodings.ALP, enc.Encoding()) - assert.Equal(t, parquet.Types.Float, enc.Type()) - - enc.(Float32Encoder).Put(values) - buf, err := enc.FlushValues() - require.NoError(t, err) - defer buf.Release() - - dec := NewDecoder(parquet.Types.Float, parquet.Encodings.ALP, col, mem) - assert.Equal(t, parquet.Encodings.ALP, dec.Encoding()) - assert.Equal(t, parquet.Types.Float, dec.Type()) - - dec.SetData(len(values), buf.Bytes()) - out := make([]float32, len(values)) - n, err := dec.(Float32Decoder).Decode(out) - require.NoError(t, err) - assert.Equal(t, len(values), n) - - for i, v := range values { - assert.Equal(t, math.Float32bits(v), math.Float32bits(out[i]), - "index %d: mismatch", i) - } -} - -func TestALPFloat64ViaPublicAPI(t *testing.T) { - col := newFloat64Column() - mem := memory.DefaultAllocator - - values := make([]float64, 2000) - for i := range values { - values[i] = float64(i) * 0.01 - } - - enc := NewEncoder(parquet.Types.Double, parquet.Encodings.ALP, false, col, mem) - assert.Equal(t, parquet.Encodings.ALP, enc.Encoding()) - assert.Equal(t, parquet.Types.Double, enc.Type()) - - enc.(Float64Encoder).Put(values) - buf, err := enc.FlushValues() - require.NoError(t, err) - defer buf.Release() - - dec := NewDecoder(parquet.Types.Double, parquet.Encodings.ALP, col, mem) - assert.Equal(t, parquet.Encodings.ALP, dec.Encoding()) - assert.Equal(t, parquet.Types.Double, dec.Type()) - - dec.SetData(len(values), buf.Bytes()) - out := make([]float64, len(values)) - n, err := dec.(Float64Decoder).Decode(out) - require.NoError(t, err) - assert.Equal(t, len(values), n) - - for i, v := range values { - assert.Equal(t, math.Float64bits(v), math.Float64bits(out[i]), - "index %d: mismatch", i) - } -} - -func TestALPPanicsForNonFloatTypes(t *testing.T) { +func TestALPDecoderPanicsForNonFloatType(t *testing.T) { col := schema.NewColumn(schema.NewInt32Node("int32", parquet.Repetitions.Required, -1), 0, 0) - mem := memory.DefaultAllocator - - assert.Panics(t, func() { - NewEncoder(parquet.Types.Int32, parquet.Encodings.ALP, false, col, mem) - }, "ALP should panic for non-float encoder") + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) assert.Panics(t, func() { NewDecoder(parquet.Types.Int32, parquet.Encodings.ALP, col, mem) @@ -1194,7 +994,10 @@ func alpDecodePage[T alpFloat](t *testing.T, page Buffer, numValues int) []T { // and calls FlushValues for each one without resetting the encoder in between. // A page therefore has to hold its own values and no others. func TestALPEncoderFlushValuesStartsANewPage(t *testing.T) { - enc := newALPEncoder[float64](format.Encoding_ALP, nil, memory.DefaultAllocator) + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + enc := newALPEncoder[float64](format.Encoding_ALP, nil, mem) + defer enc.Release() first := make([]float64, 2*alpDefaultVectorSize) for i := range first { @@ -1203,8 +1006,8 @@ func TestALPEncoderFlushValuesStartsANewPage(t *testing.T) { enc.Put(first) page, err := enc.FlushValues() require.NoError(t, err) + defer page.Release() got := alpDecodePage[float64](t, page, len(first)) - page.Release() assert.Equal(t, first, got) second := []float64{9.5, 8.25, 7.125} @@ -1214,15 +1017,18 @@ func TestALPEncoderFlushValuesStartsANewPage(t *testing.T) { page, err = enc.FlushValues() require.NoError(t, err) + defer page.Release() got = alpDecodePage[float64](t, page, len(second)) - page.Release() assert.Equal(t, second, got) } // The shortlist describes the column, so it outlives a page. A second page of // the same values encodes to the same bytes as the first. func TestALPEncoderKeepsPresetsAcrossPages(t *testing.T) { - enc := newALPEncoder[float64](format.Encoding_ALP, nil, memory.DefaultAllocator) + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + enc := newALPEncoder[float64](format.Encoding_ALP, nil, mem) + defer enc.Release() values := make([]float64, 10*alpDefaultVectorSize) for i := range values { @@ -1232,22 +1038,25 @@ func TestALPEncoderKeepsPresetsAcrossPages(t *testing.T) { enc.Put(values) first, err := enc.FlushValues() require.NoError(t, err) + defer first.Release() require.NotNil(t, enc.cachedPresets, "ten vectors is past the sampling threshold") presets := slices.Clone(enc.cachedPresets) enc.Put(values) second, err := enc.FlushValues() require.NoError(t, err) + defer second.Release() assert.Equal(t, presets, enc.cachedPresets) assert.Equal(t, first.Bytes(), second.Bytes()) - first.Release() - second.Release() } // Reset returns the encoder to its initial state, shortlist included. func TestALPEncoderResetClearsPresets(t *testing.T) { - enc := newALPEncoder[float64](format.Encoding_ALP, nil, memory.DefaultAllocator) + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + enc := newALPEncoder[float64](format.Encoding_ALP, nil, mem) + defer enc.Release() values := make([]float64, 10*alpDefaultVectorSize) for i := range values { @@ -1278,7 +1087,10 @@ func TestALPEncoderVectorSizes(t *testing.T) { values[min(i+size, len(values))-1] = math.NaN() } - enc := newALPEncoder[float64](format.Encoding_ALP, nil, memory.DefaultAllocator) + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + enc := newALPEncoder[float64](format.Encoding_ALP, nil, mem) + defer enc.Release() require.NoError(t, enc.SetVectorSize(size)) enc.Put(values) page, err := enc.FlushValues() @@ -1300,7 +1112,10 @@ func TestALPEncoderVectorSizes(t *testing.T) { } func TestALPEncoderSetVectorSizeRejects(t *testing.T) { - enc := newALPEncoder[float32](format.Encoding_ALP, nil, memory.DefaultAllocator) + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + enc := newALPEncoder[float32](format.Encoding_ALP, nil, mem) + defer enc.Release() for _, size := range []int{0, 4, 1000, 1 << 16} { assert.Error(t, enc.SetVectorSize(size), "size %d", size) } @@ -1321,8 +1136,8 @@ func TestALPFloat32RoundingPastInt32IsAnException(t *testing.T) { _, ok := s.tryEncode(2147483.5) assert.False(t, ok, "2147483.5 at (9,6) rounds to 2^31") - alpFloat32RoundTrip(t, []float32{-24.875, 2147483.5, 12}) - alpFloat32RoundTrip(t, []float32{-24.875, -2147483.5, 12}) + alpRoundTrip(t, []float32{-24.875, 2147483.5, 12}) + alpRoundTrip(t, []float32{-24.875, -2147483.5, 12}) } // Page fields near the int32 limits must be rejected, not overflow int on diff --git a/parquet/pqarrow/alp_test.go b/parquet/pqarrow/alp_test.go index e8ac8ffa8..836aa9d6d 100644 --- a/parquet/pqarrow/alp_test.go +++ b/parquet/pqarrow/alp_test.go @@ -28,6 +28,7 @@ import ( "github.com/apache/arrow-go/v18/parquet" "github.com/apache/arrow-go/v18/parquet/compress" "github.com/apache/arrow-go/v18/parquet/file" + format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet" "github.com/apache/arrow-go/v18/parquet/pqarrow" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -147,8 +148,15 @@ func TestALPArrowRoundtrip(t *testing.T) { for c := range rdr.MetaData().Schema.NumColumns() { chunk, err := rdr.RowGroup(rg).MetaData().ColumnChunk(c) require.NoError(t, err) - assert.Contains(t, chunk.Encodings(), parquet.Encodings.ALP, - "row group %d column %d should be ALP encoded", rg, c) + hasDataPages := false + for _, stats := range chunk.EncodingStats() { + if stats.PageType == format.PageType_DATA_PAGE || stats.PageType == format.PageType_DATA_PAGE_V2 { + hasDataPages = true + assert.Equal(t, parquet.Encodings.ALP, stats.Encoding, + "row group %d column %d should be ALP encoded", rg, c) + } + } + assert.True(t, hasDataPages, "row group %d column %d has no data page encoding statistics", rg, c) } } diff --git a/parquet/reader_writer_properties_test.go b/parquet/reader_writer_properties_test.go index c41b12b5d..8985c0e07 100644 --- a/parquet/reader_writer_properties_test.go +++ b/parquet/reader_writer_properties_test.go @@ -95,9 +95,6 @@ func TestReaderPropsGetStreamInsufficient(t *testing.T) { } func TestWriterPropALPEncodingOptIn(t *testing.T) { - assert.False(t, parquet.NewWriterProperties().ALPEncodingEnabled(), - "ALP is a preview encoding, so it stays off by default") - assert.Panics(t, func() { parquet.NewWriterProperties(parquet.WithEncoding(parquet.Encodings.ALP)) }, "the default encoding cannot be ALP without the flag") @@ -111,7 +108,6 @@ func TestWriterPropALPEncodingOptIn(t *testing.T) { {parquet.WithEncodingFor("price", parquet.Encodings.ALP), parquet.WithALPEncoding(true)}, } { props := parquet.NewWriterProperties(opts...) - assert.True(t, props.ALPEncodingEnabled()) assert.Equal(t, parquet.Encodings.ALP, props.EncodingFor("price")) assert.Equal(t, parquet.Encodings.Plain, props.EncodingFor("other"), "the flag allows ALP without selecting it") @@ -125,20 +121,19 @@ func TestWriterPropALPEncodingPerColumn(t *testing.T) { parquet.WithEncodingFor("volume", parquet.Encodings.ALP)) }, "permission for one column does not carry to another") - props := parquet.NewWriterProperties( - parquet.WithALPEncodingFor("price", true), - parquet.WithEncodingFor("price", parquet.Encodings.ALP)) - assert.False(t, props.ALPEncodingEnabled(), - "like the other getters, the unqualified one reports the default") - assert.True(t, props.ALPEncodingEnabledFor("price")) - assert.False(t, props.ALPEncodingEnabledFor("volume")) + assert.NotPanics(t, func() { + parquet.NewWriterProperties( + parquet.WithALPEncodingFor("price", true), + parquet.WithEncodingFor("price", parquet.Encodings.ALP)) + }) // A column may also be taken out of a permission given to every column. - props = parquet.NewWriterProperties( - parquet.WithALPEncoding(true), - parquet.WithALPEncodingFor("volume", false)) - assert.True(t, props.ALPEncodingEnabledFor("price")) - assert.False(t, props.ALPEncodingEnabledFor("volume")) + assert.NotPanics(t, func() { + parquet.NewWriterProperties( + parquet.WithALPEncoding(true), + parquet.WithALPEncodingFor("volume", false), + parquet.WithEncodingFor("price", parquet.Encodings.ALP)) + }) assert.Panics(t, func() { parquet.NewWriterProperties( parquet.WithALPEncoding(true), @@ -148,10 +143,7 @@ func TestWriterPropALPEncodingPerColumn(t *testing.T) { } func TestWriterPropALPVectorSize(t *testing.T) { - props := parquet.NewWriterProperties() - assert.Equal(t, parquet.DefaultALPVectorSize, props.ALPVectorSizeFor("price")) - - props = parquet.NewWriterProperties( + props := parquet.NewWriterProperties( parquet.WithALPVectorSize(4096), parquet.WithALPVectorSizeFor("price", 32), parquet.WithEncodingFor("volume", parquet.Encodings.Plain)) From 2befc7298e0233e6b7b48afde2e4e4971aa9873c Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Thu, 29 Oct 2026 13:39:46 -0400 Subject: [PATCH 10/11] fix(parquet): estimate a buffered ALP vector at its all-exception size EstimatedDataEncodedSize counted buffered values at their raw width, so a vector of NaNs or infinities (each stored with a 2-byte position, plus the vector header and offset) was underestimated and pages could overshoot the page size. --- parquet/internal/encoding/alp_encoding.go | 13 +++++++--- .../internal/encoding/alp_encoding_test.go | 24 +++++++++---------- 2 files changed, 21 insertions(+), 16 deletions(-) diff --git a/parquet/internal/encoding/alp_encoding.go b/parquet/internal/encoding/alp_encoding.go index a47e15b7c..79a7143f6 100644 --- a/parquet/internal/encoding/alp_encoding.go +++ b/parquet/internal/encoding/alp_encoding.go @@ -343,10 +343,17 @@ func (enc *alpEncoder[T]) buildPresetCache() { } } -// EstimatedDataEncodedSize counts the encoded vectors and their offsets, and -// unencoded buffered values at their physical width. +// EstimatedDataEncodedSize counts the encoded vectors and their offsets, plus +// the vector still being buffered sized as if every value were an exception: +// its offset, header, and each value with its 2-byte position. That is exact +// for the worst common case, NaNs, infinities and other unencodable values, and +// above the size of a vector that packs well. func (enc *alpEncoder[T]) EstimatedDataEncodedSize() int64 { - return int64(len(enc.encodedBuf) + 4*len(enc.vectorSizes) + enc.bufCount*alpEncodedBytes[T]()) + size := len(enc.encodedBuf) + 4*len(enc.vectorSizes) + if enc.bufCount > 0 { + size += 4 + alpVectorHeaderSize[T]() + enc.bufCount*(alpEncodedBytes[T]()+2) + } + return int64(size) } // FlushValues writes one page and clears the values it wrote, so that the next diff --git a/parquet/internal/encoding/alp_encoding_test.go b/parquet/internal/encoding/alp_encoding_test.go index 3bc979de4..29c83b2de 100644 --- a/parquet/internal/encoding/alp_encoding_test.go +++ b/parquet/internal/encoding/alp_encoding_test.go @@ -813,7 +813,7 @@ func TestALPPresetCacheBuilds(t *testing.T) { } } -func testALPEstimatedDataEncodedSize[T alpFloat](t *testing.T, valueWidth int) { +func testALPEstimatedDataEncodedSize[T alpFloat](t *testing.T) { t.Helper() tests := []struct { name string @@ -859,16 +859,11 @@ func testALPEstimatedDataEncodedSize[T alpFloat](t *testing.T, valueWidth int) { case pending == 0: assert.Equal(t, payloadSize, estimate, "completed vectors have an exact size") case tt.allExceptions: - // A pending all-exception vector needs one 4-byte offset, - // exponent/factor (2), exception count (2), frame of reference, - // bit width (1), and one 2-byte position per exception. - // There are no packed delta bits; the estimate already counts - // the verbatim values at their physical width. - overhead := int64(4 + 2 + 2 + valueWidth + 1 + 2*pending) - assert.Equal(t, payloadSize, estimate+overhead) + assert.Equal(t, payloadSize, estimate, + "a pending vector is estimated as if every value were an exception") default: assert.GreaterOrEqual(t, estimate, payloadSize, - "buffered decimal values should not exceed their raw-width estimate") + "a pending vector must not be underestimated") } assert.Zero(t, enc.EstimatedDataEncodedSize(), "flushing starts a new page") }) @@ -877,10 +872,10 @@ func testALPEstimatedDataEncodedSize[T alpFloat](t *testing.T, valueWidth int) { func TestALPEstimatedDataEncodedSize(t *testing.T) { t.Run("Float32", func(t *testing.T) { - testALPEstimatedDataEncodedSize[float32](t, 4) + testALPEstimatedDataEncodedSize[float32](t) }) t.Run("Float64", func(t *testing.T) { - testALPEstimatedDataEncodedSize[float64](t, 8) + testALPEstimatedDataEncodedSize[float64](t) }) } @@ -1012,8 +1007,11 @@ func TestALPEncoderFlushValuesStartsANewPage(t *testing.T) { second := []float64{9.5, 8.25, 7.125} enc.Put(second) - assert.EqualValues(t, len(second)*8, enc.EstimatedDataEncodedSize(), - "the flushed page still counts towards the next one") + fresh := newALPEncoder[float64](format.Encoding_ALP, nil, mem) + defer fresh.Release() + fresh.Put(second) + assert.Equal(t, fresh.EstimatedDataEncodedSize(), enc.EstimatedDataEncodedSize(), + "the flushed page no longer counts towards the next one") page, err = enc.FlushValues() require.NoError(t, err) From b7c6e12364269f7ff01e5a642b3eea7bf7dd2c87 Mon Sep 17 00:00:00 2001 From: Matt Topol Date: Thu, 29 Oct 2026 13:44:40 -0400 Subject: [PATCH 11/11] fix(parquet): bound a buffered ALP vector's estimate by its largest size A vector can pack every slot at full width and also store exceptions, so sizing the pending vector as all exceptions still undercounted mixed vectors of NaNs and wide integers. Count both. --- parquet/internal/encoding/alp_encoding.go | 12 +++-- .../internal/encoding/alp_encoding_test.go | 49 +++++++++++-------- 2 files changed, 36 insertions(+), 25 deletions(-) diff --git a/parquet/internal/encoding/alp_encoding.go b/parquet/internal/encoding/alp_encoding.go index 79a7143f6..00ae0818f 100644 --- a/parquet/internal/encoding/alp_encoding.go +++ b/parquet/internal/encoding/alp_encoding.go @@ -344,14 +344,16 @@ func (enc *alpEncoder[T]) buildPresetCache() { } // EstimatedDataEncodedSize counts the encoded vectors and their offsets, plus -// the vector still being buffered sized as if every value were an exception: -// its offset, header, and each value with its 2-byte position. That is exact -// for the worst common case, NaNs, infinities and other unencodable values, and -// above the size of a vector that packs well. +// 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 { - size += 4 + alpVectorHeaderSize[T]() + enc.bufCount*(alpEncodedBytes[T]()+2) + width := alpEncodedBytes[T]() + size += 4 + alpVectorHeaderSize[T]() + enc.bufCount*(2*width+2) } return int64(size) } diff --git a/parquet/internal/encoding/alp_encoding_test.go b/parquet/internal/encoding/alp_encoding_test.go index 29c83b2de..a9147601d 100644 --- a/parquet/internal/encoding/alp_encoding_test.go +++ b/parquet/internal/encoding/alp_encoding_test.go @@ -815,22 +815,35 @@ func TestALPPresetCacheBuilds(t *testing.T) { func testALPEstimatedDataEncodedSize[T alpFloat](t *testing.T) { t.Helper() + lo, hi := alpEncodingLimits[T]() tests := []struct { - name string - count int - value func(int) T - allExceptions bool + name string + count int + value func(int) T }{ - {"Empty", 0, func(int) T { return 0 }, false}, - {"BufferedMonetary", 100, func(i int) T { return T(i) * 0.01 }, false}, - {"BufferedDecimals", 100, func(i int) T { return T(i) / 4 }, false}, - {"IntegerVector", alpDefaultVectorSize, func(i int) T { return T(i) }, false}, - {"ConstantVector", alpDefaultVectorSize, func(int) T { return 3.14 }, false}, - {"MonetaryVectors", 2 * alpDefaultVectorSize, func(i int) T { return T(i) * 0.01 }, false}, - {"DecimalsWithBufferedTail", alpDefaultVectorSize + 100, func(i int) T { return T(i) / 4 }, false}, - {"AllExceptionsVector", alpDefaultVectorSize, func(int) T { return T(math.Inf(1)) }, true}, - {"AllExceptionsBuffered", 100, func(int) T { return T(math.Inf(1)) }, true}, - {"AllExceptionsWithBufferedTail", alpDefaultVectorSize + 100, func(int) T { return T(math.Inf(1)) }, true}, + {"Empty", 0, func(int) T { return 0 }}, + {"BufferedMonetary", 100, func(i int) T { return T(i) * 0.01 }}, + {"BufferedDecimals", 100, func(i int) T { return T(i) / 4 }}, + {"IntegerVector", alpDefaultVectorSize, func(i int) T { return T(i) }}, + {"ConstantVector", alpDefaultVectorSize, func(int) T { return 3.14 }}, + {"MonetaryVectors", 2 * alpDefaultVectorSize, func(i int) T { return T(i) * 0.01 }}, + {"DecimalsWithBufferedTail", alpDefaultVectorSize + 100, func(i int) T { return T(i) / 4 }}, + {"AllExceptionsVector", alpDefaultVectorSize, func(int) T { return T(math.Inf(1)) }}, + {"AllExceptionsBuffered", 100, func(int) T { return T(math.Inf(1)) }}, + {"AllExceptionsWithBufferedTail", alpDefaultVectorSize + 100, func(int) T { return T(math.Inf(1)) }}, + // Half NaNs, half integers near both ends of the integer range: the + // encoder packs the integers at nearly full width and also stores + // every NaN as an exception, the largest a vector gets. + {"MixedWideBuffered", 100, func(i int) T { + switch i % 4 { + case 0, 2: + return T(math.NaN()) + case 1: + return T(math.Trunc(float64(lo) * 0.97)) + default: + return T(math.Trunc(float64(hi) * 0.97)) + } + }}, } for _, tt := range tests { @@ -855,13 +868,9 @@ func testALPEstimatedDataEncodedSize[T alpFloat](t *testing.T) { require.GreaterOrEqual(t, buf.Len(), pageHeaderSize) payloadSize := int64(buf.Len() - pageHeaderSize) pending := tt.count % alpDefaultVectorSize - switch { - case pending == 0: + if pending == 0 { assert.Equal(t, payloadSize, estimate, "completed vectors have an exact size") - case tt.allExceptions: - assert.Equal(t, payloadSize, estimate, - "a pending vector is estimated as if every value were an exception") - default: + } else { assert.GreaterOrEqual(t, estimate, payloadSize, "a pending vector must not be underestimated") }