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/doc.go b/parquet/doc.go index edccc3c0f..aaddbbb12 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 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/alp_test.go b/parquet/file/alp_test.go new file mode 100644 index 000000000..d4d5568a6 --- /dev/null +++ b/parquet/file/alp_test.go @@ -0,0 +1,584 @@ +// 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" + "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/metadata" + "github.com/apache/arrow-go/v18/parquet/schema" + "github.com/stretchr/testify/require" +) + +// 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() + + opts = append([]parquet.WriterProperty{ + parquet.WithAllocator(mem), + parquet.WithEncoding(parquet.Encodings.ALP), + parquet.WithALPEncoding(true), + parquet.WithDictionaryDefault(false), + }, opts...) + sc, err := schema.NewGroupNode("test", parquet.Repetitions.Required, fields, -1) + require.NoError(t, err) + + var sink bytes.Buffer + writer := file.NewParquetWriter(&sink, sc, file.WithWriterProps(parquet.NewWriterProperties(opts...))) + defer writer.Close() + rgw := writer.AppendRowGroup() + 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()), + file.WithReadProps(parquet.NewReaderProperties(mem))) + require.NoError(t, err) + 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(col) + require.NoError(t, err) + defer cr.Close() + r := cr.(*file.Float64ColumnChunkReader) + + rows := int(reader.NumRows()) + out := make([]float64, rows) + var outDef []int16 + if cr.Descriptor().MaxDefinitionLevel() > 0 { + outDef = make([]int16, rows) + for i := range outDef { + outDef[i] = -1 + } + } + 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()) + 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) + } + } + 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 + 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) + }) + } +} + +// 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 + 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 { + values[i] = -9999 // A null slot must not enter the encoded values. + continue + } + defLevels[i] = 1 + 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]) + } + + 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, +// 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) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + values := make([]float64, 5000) + for i := range values { + values[i] = pool.draw() + } + + 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. + 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) { + mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) + defer mem.AssertSize(t, 0) + + values := []float64{1.25, 2.5, 3.75, 100, 0.001} + 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) +} + +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])) + } + }) + } +} + +// 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) + } + } + 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) +} + +// 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) + } +} + +// 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) + + const numValues = 40000 + + values := make([]float64, numValues) + for i := range values { + values[i] = float64(i) / 100 + } + + 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) + assertALPDataPageEncoding(t, chunk) + + pages, err := reader.RowGroup(0).GetColumnPageReader(0) + require.NoError(t, err) + 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") + + // 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() + + 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) + 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.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_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/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/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_reader_test.go b/parquet/file/file_reader_test.go index 2dd63e390..e33b3c652 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" @@ -1044,6 +1045,154 @@ 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) + + 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) + 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) + assertALPDataPageEncoding(t, chunk) + } + + 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) + defer cr.Close() + + 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) + } +} + func TestEncryptedFileCipherAndAAD(t *testing.T) { const footerKey = "0123456789012345" values := []int32{3, 7, 11} diff --git a/parquet/internal/encoding/alp_encoding.go b/parquet/internal/encoding/alp_encoding.go new file mode 100644 index 000000000..00ae0818f --- /dev/null +++ b/parquet/internal/encoding/alp_encoding.go @@ -0,0 +1,667 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package encoding + +// ALP (Adaptive Lossless Floating-Point) encoding for FLOAT and DOUBLE columns. +// +// ALP multiplies each value by a power of ten, rounds the product to an integer, +// and stores bit-packed differences from the smallest integer. A value that does +// not survive that round trip is stored verbatim as an exception. alp_scaling.go +// holds that arithmetic and alp_params.go chooses the scaling; this file writes +// and reads the bytes. +// +// A page begins with a 7-byte header: compression mode, integer encoding, log2 +// of the vector size, then the value count as a little-endian int32. An array +// of little-endian uint32 offsets follows, one per vector, each counted from the +// start of that array. Every vector then holds: +// +// exponent (1 byte), factor (1 byte), exception count (uint16) +// frame of reference (int32 for FLOAT, int64 for DOUBLE), bit width (1 byte) +// the bit-packed differences from the frame of reference +// the exception positions (uint16 each), then the exception values +// +// Reference: https://dl.acm.org/doi/10.1145/3626717 + +import ( + "bytes" + "encoding/binary" + "errors" + "fmt" + "math" + "math/bits" + "slices" + + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + format "github.com/apache/arrow-go/v18/parquet/internal/gen-go/parquet" + "github.com/apache/arrow-go/v18/parquet/internal/utils" + "github.com/apache/arrow-go/v18/parquet/schema" +) + +const ( + // The first two header bytes select the ALP variants. The decoder rejects + // variants other than the ones this encoder writes. + alpCompressionMode = 0 // ALP; other values are reserved for variants such as ALP-RD + alpIntegerEncodingFOR = 0 // frame of reference plus bit packing + + alpHeaderSize = 7 + + // Encoders write 1,024-value vectors unless SetVectorSize says otherwise. + // The format allows anything from 2^alpMinLogVectorSize to + // 2^alpMaxLogVectorSize, and the decoder honors what a page declares. + alpDefaultVectorSize = 1024 + alpDefaultLogVectorSize = 10 + alpMinLogVectorSize = 3 + alpMaxLogVectorSize = 15 +) + +func alpVectorHeaderSize[T alpFloat]() int { + return 1 + 1 + 2 + alpEncodedBytes[T]() + 1 +} + +func alpPackedSize(count, bitWidth int) int { + return (count*bitWidth + 7) / 8 +} + +// alpAppendPackedBits appends values to dst, packed at bitWidth bits each, +// least significant bit first. +func alpAppendPackedBits(dst []byte, values []uint64, bitWidth int) []byte { + if bitWidth == 0 || len(values) == 0 { + return dst + } + + // BitWriter flushes eight bytes at a time, so it needs room past the last + // packed byte. The return drops those bytes again. + start, size := len(dst), alpPackedSize(len(values), bitWidth) + dst = slices.Grow(dst, size+8)[:start+size+8] + clear(dst[start:]) + + bw := utils.NewBitWriter(utils.NewWriterAtBuffer(dst[start:])) + for _, v := range values { + bw.WriteValue(v, uint(bitWidth)) + } + bw.Flush(true) + return dst[:start+size] +} + +// alpUnpackBits reverses alpAppendPackedBits into out, which must hold as many +// values as were packed. The reader is reset onto packed, so one reader serves +// every vector of a page. +func alpUnpackBits(br *utils.BitReader, packed *bytes.Reader, data []byte, out []uint64, bitWidth int) error { + if bitWidth == 0 { + clear(out) + return nil + } + + packed.Reset(data) + br.Reset(packed) + n, err := br.GetBatch(uint(bitWidth), out) + if err != nil { + return err + } + if n != len(out) { + return fmt.Errorf("parquet: ALP unpacked %d of %d values", n, len(out)) + } + return nil +} + +func alpAppendInt[T alpFloat](dst []byte, v uint64) []byte { + if alpEncodedBytes[T]() == 4 { + return binary.LittleEndian.AppendUint32(dst, uint32(v)) + } + return binary.LittleEndian.AppendUint64(dst, v) +} + +func alpReadInt[T alpFloat](src []byte) int64 { + if alpEncodedBytes[T]() == 4 { + return int64(int32(binary.LittleEndian.Uint32(src))) + } + return int64(binary.LittleEndian.Uint64(src)) +} + +func alpReadFloat[T alpFloat](src []byte) T { + if alpEncodedBytes[T]() == 4 { + return alpFloatFromBits[T](uint64(binary.LittleEndian.Uint32(src))) + } + return alpFloatFromBits[T](binary.LittleEndian.Uint64(src)) +} + +// alpEncoder encodes a float column with ALP, a vector at a time. It buffers +// values until it holds a whole vector, so the values of a page reach the sink +// only when the caller flushes it. +type alpEncoder[T alpFloat] struct { + encoder + + vectorSize int + logVectorSize int + + vectorBuf []T + bufCount int + totalCount int + + encodedBuf []byte + vectorSizes []int + + // Scratch reused across vectors, calls, and pages. + encodedScratch []int64 + deltaScratch []uint64 + excPositions []uint16 + excValues []T + sampleScratch []T + spacedScratch []T + pageHeader []byte + + vectorsProcessed int + cachedPresets []alpEncodingParams + presetCounts map[alpEncodingParams]int +} + +func newALPEncoder[T alpFloat](e format.Encoding, descr *schema.Column, mem memory.Allocator) *alpEncoder[T] { + enc := &alpEncoder[T]{ + encoder: newEncoderBase(e, descr, mem), + sampleScratch: make([]T, 0, alpSamplerSamples), + presetCounts: make(map[alpEncodingParams]int), + } + enc.setVectorSize(alpDefaultLogVectorSize) + return enc +} + +func (enc *alpEncoder[T]) setVectorSize(logSize int) { + enc.logVectorSize = logSize + enc.vectorSize = 1 << logSize + enc.vectorBuf = make([]T, enc.vectorSize) + enc.encodedScratch = make([]int64, enc.vectorSize) + enc.deltaScratch = make([]uint64, enc.vectorSize) +} + +// SetVectorSize sets how many values each vector holds: a power of two from +// 8 to 32768, 1024 by default. It must be called before the first value is +// put, and Reset keeps it. +func (enc *alpEncoder[T]) SetVectorSize(size int) error { + if size < 1< 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 { + return alpParquetType[T]() +} + +// Put buffers values and encodes complete vectors. FlushValues handles the +// final partial vector. +func (enc *alpEncoder[T]) Put(in []T) { + for len(in) > 0 { + n := copy(enc.vectorBuf[enc.bufCount:], in) + enc.bufCount += n + enc.totalCount += n + in = in[n:] + + if enc.bufCount == enc.vectorSize { + enc.encodeVector(enc.vectorBuf[:enc.bufCount]) + enc.bufCount = 0 + } + } +} + +// PutSpaced drops the null slots of in and buffers the values that remain. +func (enc *alpEncoder[T]) PutSpaced(in []T, validBits []byte, validBitsOffset int64) { + enc.spacedScratch = slices.Grow(enc.spacedScratch[:0], len(in))[:len(in)] + nvalid := spacedCompress(in, enc.spacedScratch, validBits, validBitsOffset) + enc.Put(enc.spacedScratch[:nvalid]) +} + +// The first alpSamplerVectors vectors seed the preset list. +func (enc *alpEncoder[T]) findParams(values []T) alpEncodingParams { + sample := alpSample(values, enc.sampleScratch[:0]) + + if len(enc.cachedPresets) > 0 { + return alpFindBestParamsWithPresets(sample, enc.cachedPresets) + } + + params := alpFindBestParams(sample) + enc.presetCounts[params]++ + return params +} + +func (enc *alpEncoder[T]) encodeVector(values []T) { + params := enc.findParams(values) + + enc.vectorsProcessed++ + if enc.cachedPresets == nil && enc.vectorsProcessed >= alpSamplerVectors { + enc.buildPresetCache() + } + + encoded := enc.encodedScratch[:len(values)] + enc.excPositions = enc.excPositions[:0] + enc.excValues = enc.excValues[:0] + + // Encoding a value also determines whether it is an exception, so collect + // both results in one pass. + s := newALPScaling[T](params.exponent, params.factor) + minValue := int64(math.MaxInt64) + for i, v := range values { + e, ok := s.tryEncode(v) + if !ok { + enc.excPositions = append(enc.excPositions, uint16(i)) + enc.excValues = append(enc.excValues, v) + continue + } + encoded[i] = e + minValue = min(e, minValue) + } + + if len(enc.excValues) == len(values) { + // Nothing encoded, so nothing constrains the frame of reference. + minValue = 0 + } + // An exception still takes up a slot among the packed integers. Filling the + // slot with the frame of reference itself costs no bits, so the exceptions + // cannot widen the packing. + for _, pos := range enc.excPositions { + encoded[pos] = minValue + } + + // The difference between two int64 values can exceed int64. Compute it in + // unsigned arithmetic; the decoder adds the wrapped difference back. + deltas := enc.deltaScratch[:len(values)] + var maxDelta uint64 + for i, e := range encoded { + deltas[i] = uint64(e) - uint64(minValue) + maxDelta = max(deltas[i], maxDelta) + } + bitWidth := bits.Len64(maxDelta) + + startLen := len(enc.encodedBuf) + enc.encodedBuf = append(enc.encodedBuf, byte(params.exponent), byte(params.factor)) + enc.encodedBuf = binary.LittleEndian.AppendUint16(enc.encodedBuf, uint16(len(enc.excPositions))) + enc.encodedBuf = alpAppendInt[T](enc.encodedBuf, uint64(minValue)) + enc.encodedBuf = append(enc.encodedBuf, byte(bitWidth)) + enc.encodedBuf = alpAppendPackedBits(enc.encodedBuf, deltas, bitWidth) + + for _, pos := range enc.excPositions { + enc.encodedBuf = binary.LittleEndian.AppendUint16(enc.encodedBuf, pos) + } + for _, v := range enc.excValues { + enc.encodedBuf = alpAppendInt[T](enc.encodedBuf, alpFloatToBits(v)) + } + + enc.vectorSizes = append(enc.vectorSizes, len(enc.encodedBuf)-startLen) +} + +// buildPresetCache keeps the scaling pairs the sampled vectors chose most often, +// ordered by how often they won. +func (enc *alpEncoder[T]) buildPresetCache() { + type preset struct { + params alpEncodingParams + count int + } + + sorted := make([]preset, 0, len(enc.presetCounts)) + for k, v := range enc.presetCounts { + sorted = append(sorted, preset{k, v}) + } + // The counts come out of a map, so the exponent and the factor break ties: + // without them the shortlist, and the encoded bytes, would differ between + // two runs over the same input. + slices.SortStableFunc(sorted, func(a, b preset) int { + switch { + case a.count != b.count: + return b.count - a.count + case a.params.exponent != b.params.exponent: + return b.params.exponent - a.params.exponent + default: + return b.params.factor - a.params.factor + } + }) + + numPresets := min(len(sorted), alpMaxPresetCombinations) + enc.cachedPresets = make([]alpEncodingParams, numPresets) + for i, p := range sorted[:numPresets] { + enc.cachedPresets[i] = p.params + } +} + +// EstimatedDataEncodedSize counts the encoded vectors and their offsets, plus +// an upper bound for the vector still being buffered: its offset and header, +// every slot packed at the full integer width, and every value also stored as +// an exception with its 2-byte position. A vector that mixes exceptions with +// wide packed values can approach that bound, so nothing smaller is safe; the +// overestimate is limited to one vector. +func (enc *alpEncoder[T]) EstimatedDataEncodedSize() int64 { + size := len(enc.encodedBuf) + 4*len(enc.vectorSizes) + if enc.bufCount > 0 { + width := alpEncodedBytes[T]() + size += 4 + alpVectorHeaderSize[T]() + enc.bufCount*(2*width+2) + } + return int64(size) +} + +// FlushValues writes one page and clears the values it wrote, so that the next +// Put starts a new page. The scaling shortlist survives, because it describes +// the column rather than the page. +// +// A page always carries the header, even when it holds no values: a column of +// optional values can produce a page whose rows are all null, and a reader +// rejects an empty page rather than reading it as no values. +func (enc *alpEncoder[T]) FlushValues() (Buffer, error) { + if enc.bufCount > 0 { + enc.encodeVector(enc.vectorBuf[:enc.bufCount]) + enc.bufCount = 0 + } + defer enc.resetPage() + + offsetArraySize := len(enc.vectorSizes) * 4 + enc.sink.Reserve(alpHeaderSize + offsetArraySize + len(enc.encodedBuf)) + + header := append(enc.pageHeader[:0], alpCompressionMode, alpIntegerEncodingFOR, byte(enc.logVectorSize)) + header = binary.LittleEndian.AppendUint32(header, uint32(enc.totalCount)) + + offset := uint32(offsetArraySize) + for _, size := range enc.vectorSizes { + header = binary.LittleEndian.AppendUint32(header, offset) + offset += uint32(size) + } + enc.pageHeader = header + + enc.sink.Write(header) + enc.sink.Write(enc.encodedBuf) + return enc.sink.Finish(), nil +} + +// resetPage drops the encoded vectors of the page just written. +func (enc *alpEncoder[T]) resetPage() { + enc.totalCount = 0 + enc.encodedBuf = enc.encodedBuf[:0] + enc.vectorSizes = enc.vectorSizes[:0] +} + +// Reset returns the encoder to its state when new, shortlist included, so that +// it can encode an unrelated column. The vector size stays as set. +func (enc *alpEncoder[T]) Reset() { + enc.encoder.Reset() + enc.resetPage() + enc.bufCount = 0 + enc.vectorsProcessed = 0 + enc.cachedPresets = nil + clear(enc.presetCounts) +} + +// alpDecoder decodes one ALP vector at a time and keeps it across partial reads. +type alpDecoder[T alpFloat] struct { + decoder + + vectorSize int + totalCount int + numVectors int + + vectorOffsets []uint32 + bodyData []byte + + // The vector the decoder holds, and how far through the page it has read. + currentVectorIndex int + currentIndex int + decodedValues []T + deltas []uint64 + + // One reader pair, reset onto each vector's packed deltas in turn. + packedReader *bytes.Reader + bitReader *utils.BitReader +} + +func newALPDecoder[T alpFloat](e format.Encoding, descr *schema.Column) *alpDecoder[T] { + packed := bytes.NewReader(nil) + return &alpDecoder[T]{ + decoder: newDecoderBase(e, descr), + packedReader: packed, + bitReader: utils.NewBitReader(packed), + } +} + +func (dec *alpDecoder[T]) Type() parquet.Type { + return alpParquetType[T]() +} + +// SetData points the decoder at one page and reads its header. It rejects +// unsupported ALP variants and truncated vector-offset tables. +func (dec *alpDecoder[T]) SetData(nvals int, data []byte) error { + if len(data) < alpHeaderSize { + return fmt.Errorf("parquet: ALP data too short for header: %d bytes", len(data)) + } + + compressionMode := data[0] + integerEncoding := data[1] + logVectorSize := data[2] + numElements := int32(binary.LittleEndian.Uint32(data[3:7])) + + if compressionMode != alpCompressionMode { + return fmt.Errorf("parquet: unsupported ALP compression mode: %d", compressionMode) + } + if integerEncoding != alpIntegerEncodingFOR { + return fmt.Errorf("parquet: unsupported ALP integer encoding: %d", integerEncoding) + } + if logVectorSize < alpMinLogVectorSize || logVectorSize > alpMaxLogVectorSize { + return fmt.Errorf("parquet: invalid ALP log vector size: %d, must be between %d and %d", + logVectorSize, alpMinLogVectorSize, alpMaxLogVectorSize) + } + if numElements < 0 { + return fmt.Errorf("parquet: invalid ALP element count: %d", numElements) + } + + dec.vectorSize = 1 << logVectorSize + dec.totalCount = int(numElements) + // Rounding up by adding vectorSize-1 first could overflow a 32-bit int. + dec.numVectors = dec.totalCount/dec.vectorSize + min(dec.totalCount%dec.vectorSize, 1) + dec.currentIndex = 0 + dec.currentVectorIndex = -1 + dec.nvals = nvals + + 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, 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 > len(dec.bodyData)-alpVectorHeaderSize[T]() { + 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 + + // 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 + // 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] +) + +// 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_benchmark_test.go b/parquet/internal/encoding/alp_encoding_benchmark_test.go new file mode 100644 index 000000000..6972d7a0e --- /dev/null +++ b/parquet/internal/encoding/alp_encoding_benchmark_test.go @@ -0,0 +1,195 @@ +// 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) + defer encoder.Release() + b.ReportAllocs() + 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) + defer encoder.Release() + b.ReportAllocs() + 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) + 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() { + 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) + 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() { + 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) + 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() { + 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) + } + } + } + }) + } +} diff --git a/parquet/internal/encoding/alp_encoding_test.go b/parquet/internal/encoding/alp_encoding_test.go new file mode 100644 index 000000000..a9147601d --- /dev/null +++ b/parquet/internal/encoding/alp_encoding_test.go @@ -0,0 +1,1162 @@ +// 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" + "math/bits" + "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 alpRoundTrip[T alpFloat](t *testing.T, values []T) { + t.Helper() + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + 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() + + out := alpDecodePage[T](t, buf, len(values)) + for i, v := range values { + 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 testALPFloatRoundTrip[T alpFloat](t *testing.T) { + t.Helper() + 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 + } + + 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) + }) + } + }) + } +} + +func TestALPFloatRoundTrip(t *testing.T) { + t.Run("Float32", testALPFloatRoundTrip[float32]) + t.Run("Float64", testALPFloatRoundTrip[float64]) +} + +func TestALPFloat32ProgressiveDecoding(t *testing.T) { + values := make([]float32, 3000) + for i := range values { + values[i] = float32(i) * 0.01 + } + + col := newFloat32Column() + 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) + 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.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) + 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.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) + 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.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) + 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.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) + 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.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + 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) + 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) + + 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) + } + } +} + +func TestALPFloat64MultiplePages(t *testing.T) { + col := newFloat64Column() + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + 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) + 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) + + 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) + } + } +} + +func TestALPEncoderReset(t *testing.T) { + col := newFloat32Column() + 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) + 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 TestALPEncoderPutSpaced(t *testing.T) { + col := newFloat32Column() + 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) + 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.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + values := make([]float32, 100) + for i := range values { + values[i] = float32(i) + } + + 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() + + 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.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + values := make([]float64, 3000) + for i := range values { + values[i] = float64(i) * 0.01 + } + + 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() + + 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)) + } + + 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) + 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") + + alpRoundTrip(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() + + 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) + defer 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.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 { + 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 alpFloat](t *testing.T) { + t.Helper() + lo, hi := alpEncodingLimits[T]() + tests := []struct { + name string + count int + value func(int) T + }{ + {"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 { + 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([]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 + if pending == 0 { + assert.Equal(t, payloadSize, estimate, "completed vectors have an exact size") + } else { + assert.GreaterOrEqual(t, estimate, payloadSize, + "a pending vector must not be underestimated") + } + assert.Zero(t, enc.EstimatedDataEncodedSize(), "flushing starts a new page") + }) + } +} + +func TestALPEstimatedDataEncodedSize(t *testing.T) { + t.Run("Float32", func(t *testing.T) { + testALPEstimatedDataEncodedSize[float32](t) + }) + t.Run("Float64", func(t *testing.T) { + testALPEstimatedDataEncodedSize[float64](t) + }) +} + +func TestALPFloat32IncrementalPut(t *testing.T) { + col := newFloat32Column() + 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 + } + + 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) + } + + alpRoundTrip(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) + } + + alpRoundTrip(t, values) +} + +func TestALPDecoderPanicsForNonFloatType(t *testing.T) { + col := schema.NewColumn(schema.NewInt32Node("int32", parquet.Repetitions.Required, -1), 0, 0) + mem := memory.NewCheckedAllocator(memory.DefaultAllocator) + defer mem.AssertSize(t, 0) + + 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) { + 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 { + first[i] = float64(i) / 100 + } + enc.Put(first) + page, err := enc.FlushValues() + require.NoError(t, err) + defer page.Release() + got := alpDecodePage[float64](t, page, len(first)) + assert.Equal(t, first, got) + + second := []float64{9.5, 8.25, 7.125} + enc.Put(second) + 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) + defer page.Release() + got = alpDecodePage[float64](t, page, len(second)) + 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) { + 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 { + values[i] = float64(i%997) / 100 + } + + 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()) +} + +// Reset returns the encoder to its initial state, shortlist included. +func TestALPEncoderResetClearsPresets(t *testing.T) { + 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 { + 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()) +} + +// 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() + } + + 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() + 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) { + 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) + } + + 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") +} + +// 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") + + 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 +// 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_params.go b/parquet/internal/encoding/alp_params.go new file mode 100644 index 000000000..7da63a0dd --- /dev/null +++ b/parquet/internal/encoding/alp_params.go @@ -0,0 +1,155 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package encoding + +// The encoder samples each vector and chooses the exponent and factor that +// minimize its encoded size. After the initial sampling window, it searches +// only the combinations that won. + +import ( + "math" + "math/bits" +) + +const ( + // alpSamplerVectors is how many vectors get the full exponent and factor + // search before the encoder falls back to the combinations those vectors + // chose. alpMaxPresetCombinations caps how many it keeps. + alpSamplerVectors = 8 + alpMaxPresetCombinations = 5 + + // alpSamplerSamples is how many values of a vector the search looks at, and + // alpPresetGiveUpAfter is how many shortlisted scalings may fail to improve + // on the best before the per-vector search stops looking. + alpSamplerSamples = 256 + alpPresetGiveUpAfter = 4 +) + +// alpEncodingParams is the scaling one vector is encoded with: ALP multiplies +// each value by 10^exponent, then by 10^-factor. +type alpEncodingParams struct { + exponent int + factor int +} + +// alpSample appends up to alpSamplerSamples evenly spaced values to out. +func alpSample[T alpFloat](values []T, out []T) []T { + stride := max(1, (len(values)+alpSamplerSamples-1)/alpSamplerSamples) + for i := 0; i < len(values); i += stride { + out = append(out, values[i]) + } + return out +} + +// alpExceptionBits is what one exception costs: the value at its full width, +// plus the position that patches it back in. +func alpExceptionBits[T alpFloat]() int { + return 8*alpEncodedBytes[T]() + 16 +} + +// alpEstimateSizeBits estimates the encoded size for a parameter pair and +// reports how many values it can encode. The cost includes the packed integers +// at the width left by the frame of reference and every exception. +func alpEstimateSizeBits[T alpFloat](values []T, exponent, factor int) (sizeBits int64, numEncodable int) { + minEncoded, maxEncoded := int64(math.MaxInt64), int64(math.MinInt64) + numExceptions := 0 + s := newALPScaling[T](exponent, factor) + for _, v := range values { + e, ok := s.tryEncode(v) + if !ok { + numExceptions++ + continue + } + minEncoded, maxEncoded = min(minEncoded, e), max(maxEncoded, e) + } + + sizeBits = int64(numExceptions) * int64(alpExceptionBits[T]()) + numEncodable = len(values) - numExceptions + if numEncodable == 0 { + return sizeBits, 0 + } + // The encoded range can exceed int64, so compute its width in unsigned + // arithmetic. + width := bits.Len64(uint64(maxEncoded) - uint64(minEncoded)) + return sizeBits + int64(len(values))*int64(width), numEncodable +} + +// For equal-size estimates, the ALP paper prefers the larger exponent, +// then the larger factor (Afroozeh et al., SIGMOD 2023, section 3.1.2). +func alpBetterParams(a alpEncodingParams, aBits int64, b alpEncodingParams, bBits int64) bool { + switch { + case aBits != bBits: + return aBits < bBits + case a.exponent != b.exponent: + return a.exponent > b.exponent + default: + return a.factor > b.factor + } +} + +// alpFindBestParams returns the cheapest scaling over every exponent and factor +// pair. The factor never exceeds the exponent, because scaling by a net power +// below one would move the decimal point the wrong way. +func alpFindBestParams[T alpFloat](values []T) alpEncodingParams { + maxExponent := alpNumExponents[T]() - 1 + // The fallback treats the values as integers: an exponent and a factor of + // equal size cancel, so whatever is already integral packs and the rest + // become exceptions. It stands when no scaling encodes enough to compare. + best := alpEncodingParams{exponent: maxExponent, factor: maxExponent} + bestBits := int64(math.MaxInt64) + + for e := range maxExponent + 1 { + for f := range e + 1 { + sizeBits, numEncodable := alpEstimateSizeBits(values, e, f) + // A scaling that turns almost everything into an exception tells us + // nothing about how wide the rest would pack. + if numEncodable < 2 { + continue + } + 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..ce77038ad --- /dev/null +++ b/parquet/internal/encoding/alp_params_test.go @@ -0,0 +1,180 @@ +// 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 + s := newALPScaling[T](params.exponent, params.factor) + for _, v := range values { + if _, ok := s.tryEncode(v); !ok { + 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..8ff42e614 --- /dev/null +++ b/parquet/internal/encoding/alp_scaling.go @@ -0,0 +1,244 @@ +// 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") +} + +// 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 + + 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), + decFac: alpPow10[T](factor), + decExp: alpNegPow10[T](exponent), + lo: lo, + hi: hi, + magic: alpMagic[T](), + intMin: intMin, + intMax: intMax, + } +} + +// 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) +} + +// 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 +} + +// 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/internal/encoding/alp_scaling_test.go b/parquet/internal/encoding/alp_scaling_test.go new file mode 100644 index 000000000..bd1bf0d97 --- /dev/null +++ b/parquet/internal/encoding/alp_scaling_test.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 + +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) { + 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) + }) + } +} + +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) { + 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) + }) + } +} + +// 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() + s := newALPScaling[T](exponent, factor) + _, ok := s.tryEncode(v) + 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) { + 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") + + 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") +} + +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) { + 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 := 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) + }) + } +} + +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) { + 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 := 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) + }) + } +} + +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 + + 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/internal/encoding/encoding_test.go b/parquet/internal/encoding/encoding_test.go index a7009b2df..9c5e0404a 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 b4861e085..3e53f6ca4 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/metadata/column_chunk.go b/parquet/metadata/column_chunk.go index c5ae5029a..1a2ee0dc7 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,25 @@ 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)) + // 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 { + if dictFallback && !slices.Contains(thriftEncodings, format.Encoding_PLAIN) { thriftEncodings = append(thriftEncodings, format.Encoding_PLAIN) } c.chunk.MetaData.Encodings = thriftEncodings diff --git a/parquet/pqarrow/alp_test.go b/parquet/pqarrow/alp_test.go new file mode 100644 index 000000000..836aa9d6d --- /dev/null +++ b/parquet/pqarrow/alp_test.go @@ -0,0 +1,224 @@ +// 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" + 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" +) + +// 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) + 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) + } + } + + 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) +} diff --git a/parquet/reader_writer_properties_test.go b/parquet/reader_writer_properties_test.go index ee8ac78f2..8985c0e07 100644 --- a/parquet/reader_writer_properties_test.go +++ b/parquet/reader_writer_properties_test.go @@ -93,3 +93,70 @@ func TestReaderPropsGetStreamInsufficient(t *testing.T) { _, err := props.GetStream(rdr, 12, 15) assert.Error(t, err) } + +func TestWriterPropALPEncodingOptIn(t *testing.T) { + 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.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") + + 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. + 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), + parquet.WithALPEncodingFor("volume", false), + parquet.WithEncodingFor("volume", parquet.Encodings.ALP)) + }, "the column-specific setting wins over the default one") +} + +func TestWriterPropALPVectorSize(t *testing.T) { + 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/types.go b/parquet/types.go index c1bd4372a..135de94ba 100644 --- a/parquet/types.go +++ b/parquet/types.go @@ -328,6 +328,9 @@ var ( DeltaBinaryPacked Encoding DeltaLengthByteArray Encoding ByteStreamSplit Encoding + // 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), PlainDict: Encoding(format.Encoding_PLAIN_DICTIONARY), @@ -338,6 +341,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 diff --git a/parquet/writer_properties.go b/parquet/writer_properties.go index 9cedc8737..a523ddb31 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" @@ -61,6 +62,15 @@ const ( DefaultAdaptiveBloomFilterEnabled = false DefaultBloomFilterCandidates = 5 + // 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 ) @@ -84,6 +94,10 @@ type ColumnProperties struct { AdaptiveBloomFilterEnabled bool BloomFilterCandidates int 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. @@ -114,6 +128,8 @@ func DefaultColumnProperties() ColumnProperties { BloomFilterFPP: DefaultBloomFilterFPP, AdaptiveBloomFilterEnabled: DefaultAdaptiveBloomFilterEnabled, BloomFilterCandidates: DefaultBloomFilterCandidates, + ALPEnabled: DefaultALPEnabled, + ALPVectorSize: DefaultALPVectorSize, } } @@ -134,6 +150,8 @@ type writerPropConfig struct { bloomFilterEnabled map[string]bool 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 @@ -283,6 +301,70 @@ 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 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 +// 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) +} + +// 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) { @@ -597,6 +679,8 @@ 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), + alpVectorSize: make(map[string]int), } for _, o := range opts { o(&cfg) @@ -660,9 +744,80 @@ func NewWriterProperties(opts ...WriterProperty) *WriterProperties { get(key).BloomFilterCandidates = value } + for key, value := range cfg.alpEnabled { + get(key).ALPEnabled = value + } + for key, value := range cfg.alpVectorSize { + get(key).ALPVectorSize = value + } + + // The check runs here rather than in the options, because an option that + // selects ALP may be given before the one that allows it. + if cfg.wr.defColumnProps.Encoding == Encodings.ALP && !cfg.wr.defColumnProps.ALPEnabled { + panic("parquet: ALP encoding requires WithALPEncoding(true)") + } + checkALPVectorSize("", cfg.wr.defColumnProps.ALPVectorSize) + for path, props := range cfg.wr.columnProps { + if props.Encoding == Encodings.ALP && !props.ALPEnabled { + panic("parquet: ALP encoding for column " + path + + " requires WithALPEncoding(true) or WithALPEncodingFor(" + path + ", true)") + } + checkALPVectorSize(path, props.ALPVectorSize) + } + return cfg.wr } +// ALPEncodingEnabled reports whether columns may use the ALP encoding by +// default. Use ALPEncodingEnabledFor for a specific column. +func (w *WriterProperties) ALPEncodingEnabled() bool { return w.defColumnProps.ALPEnabled } + +// ALPVectorSize returns the default number of values in each ALP vector. Use +// ALPVectorSizeFor for a specific column. +func (w *WriterProperties) ALPVectorSize() int { return w.defColumnProps.ALPVectorSize } + +// ALPEncodingEnabledFor reports whether the given column may be written with the +// ALP encoding, or the default value if the column was not separately specified. +func (w *WriterProperties) ALPEncodingEnabledFor(path string) bool { + if p, ok := w.columnProps[path]; ok { + return p.ALPEnabled + } + return w.defColumnProps.ALPEnabled +} + +// ALPEncodingEnabledPath is the same as ALPEncodingEnabledFor but takes a +// ColumnPath object. +func (w *WriterProperties) ALPEncodingEnabledPath(path ColumnPath) bool { + return w.ALPEncodingEnabledFor(path.String()) +} + +// ALPVectorSizeFor returns how many values each ALP vector of the given column +// holds, or the default if the column was not separately specified. +func (w *WriterProperties) ALPVectorSizeFor(path string) int { + if p, ok := w.columnProps[path]; ok { + return p.ALPVectorSize + } + return w.defColumnProps.ALPVectorSize +} + +// ALPVectorSizePath is the same as ALPVectorSizeFor but takes a ColumnPath +// object. +func (w *WriterProperties) ALPVectorSizePath(path ColumnPath) int { + return w.ALPVectorSizeFor(path.String()) +} + +func checkALPVectorSize(path string, size int) { + if size >= minALPVectorSize && size <= maxALPVectorSize && size&(size-1) == 0 { + return + } + target := "the ALP vector size" + if path != "" { + target += " of column " + path + } + 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 {