From 9e8d00b58f8c636459bce6de1ead069262868e60 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Sat, 29 Aug 2026 15:27:43 +0200 Subject: [PATCH 1/2] perf(parquet/pqarrow): write fixed-size binary values directly --- ...ixed_len_byte_array_column_writer_arrow.go | 154 ++++++++++++++ .../encoding/fixed_len_byte_array_encoder.go | 43 ++++ .../fixed_len_byte_array_encoder_test.go | 26 +++ .../metadata/fixed_len_byte_array_arrow.go | 144 +++++++++++++ .../fixed_len_byte_array_arrow_test.go | 69 +++++++ parquet/pqarrow/encode_arrow.go | 15 ++ .../pqarrow/fixed_size_binary_bench_test.go | 90 +++++++++ parquet/pqarrow/fixed_size_binary_test.go | 189 ++++++++++++++++++ 8 files changed, 730 insertions(+) create mode 100644 parquet/file/fixed_len_byte_array_column_writer_arrow.go create mode 100644 parquet/metadata/fixed_len_byte_array_arrow.go create mode 100644 parquet/metadata/fixed_len_byte_array_arrow_test.go create mode 100644 parquet/pqarrow/fixed_size_binary_bench_test.go create mode 100644 parquet/pqarrow/fixed_size_binary_test.go diff --git a/parquet/file/fixed_len_byte_array_column_writer_arrow.go b/parquet/file/fixed_len_byte_array_column_writer_arrow.go new file mode 100644 index 000000000..6c6fcbfdb --- /dev/null +++ b/parquet/file/fixed_len_byte_array_column_writer_arrow.go @@ -0,0 +1,154 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package file + +import ( + "fmt" + + "github.com/apache/arrow-go/v18/internal/utils" + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/metadata" +) + +type fixedLenByteArrayArrowEncoder interface { + PutArrow([]byte) + PutArrowSpaced([]byte, []byte, int64) +} + +// SupportsArrowValues reports whether the active fixed-length byte-array +// encoder can consume values directly from an Arrow value buffer. +func (w *FixedLenByteArrayColumnChunkWriter) SupportsArrowValues() bool { + if _, ok := w.currentEncoder.(fixedLenByteArrayArrowEncoder); !ok { + return false + } + if w.pageStatistics == nil { + return true + } + _, ok := w.pageStatistics.(*metadata.FixedLenByteArrayStatistics) + return ok +} + +func (w *FixedLenByteArrayColumnChunkWriter) writeArrowValues(values []byte, byteWidth int) { + w.currentEncoder.(fixedLenByteArrayArrowEncoder).PutArrow(values) + if stats, ok := w.pageStatistics.(*metadata.FixedLenByteArrayStatistics); ok { + stats.UpdateFromArrowFixedWidth(values, byteWidth, 0) + } + if w.bloomFilter != nil && w.currentEncoder.Encoding() != parquet.Encodings.PlainDict { + metadata.InsertArrowFixedLenHashes(w.bloomFilter, values, byteWidth) + } +} + +func (w *FixedLenByteArrayColumnChunkWriter) writeArrowValuesSpaced(values []byte, byteWidth int, numRead int64, validBits []byte, validBitsOffset int64) { + enc := w.currentEncoder.(fixedLenByteArrayArrowEncoder) + numSpaced := int64(len(values) / byteWidth) + if numSpaced == numRead { + enc.PutArrow(values) + } else { + enc.PutArrowSpaced(values, validBits, validBitsOffset) + } + + if stats, ok := w.pageStatistics.(*metadata.FixedLenByteArrayStatistics); ok { + stats.UpdateFromArrowFixedWidthSpaced(values, byteWidth, validBits, validBitsOffset, numSpaced-numRead) + } + if w.bloomFilter != nil && w.currentEncoder.Encoding() != parquet.Encodings.PlainDict { + metadata.InsertSpacedArrowFixedLenHashes(w.bloomFilter, numRead, values, byteWidth, validBits, validBitsOffset) + } +} + +// WriteBatchArrow writes fixed-length byte-array values directly from an Arrow +// value buffer. The buffer contains typeLength bytes per value. +func (w *FixedLenByteArrayColumnChunkWriter) WriteBatchArrow(values []byte, defLevels, repLevels []int16) (valueOffset int64, err error) { + defer func() { + if r := recover(); r != nil { + err = utils.FormatRecoveredError("unknown error type", r) + } + }() + if !w.SupportsArrowValues() { + return 0, fmt.Errorf("parquet: current fixed-length byte-array encoder does not support Arrow values") + } + + typeLen := int(w.descr.TypeLength()) + if typeLen <= 0 || len(values)%typeLen != 0 { + return 0, fmt.Errorf("parquet: Arrow fixed-length values are not aligned to the type length") + } + length := len(values) / typeLen + if defLevels != nil { + length = len(defLevels) + } + if length == 0 { + return 0, nil + } + + w.doBatches(int64(length), repLevels, func(offset, batch int64) { + toWrite := w.writeLevels(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) + start := int(valueOffset) * typeLen + end := int(valueOffset+toWrite) * typeLen + w.writeArrowValues(values[start:end], typeLen) + if err := w.commitWriteAndCheckPageLimit(batch, toWrite); err != nil { + panic(err) + } + valueOffset += toWrite + w.checkDictionarySizeLimit() + }) + return valueOffset, nil +} + +// WriteBatchSpacedArrow writes fixed-length byte-array values directly from an +// Arrow value buffer while using validBits to skip null values. +func (w *FixedLenByteArrayColumnChunkWriter) WriteBatchSpacedArrow(values []byte, defLevels, repLevels []int16, validBits []byte, validBitsOffset int64) (valueOffset int64, err error) { + defer func() { + if r := recover(); r != nil { + err = utils.FormatRecoveredError("unknown error type", r) + } + }() + if !w.SupportsArrowValues() { + return 0, fmt.Errorf("parquet: current fixed-length byte-array encoder does not support Arrow values") + } + + typeLen := int(w.descr.TypeLength()) + if typeLen <= 0 || len(values)%typeLen != 0 { + return 0, fmt.Errorf("parquet: Arrow fixed-length values are not aligned to the type length") + } + length := len(defLevels) + if defLevels == nil { + length = len(values) / typeLen + } + if length == 0 { + return 0, nil + } + + w.doBatches(int64(length), repLevels, func(offset, batch int64) { + info := w.maybeCalculateValidityBits(levelSliceOrNil(defLevels, offset, batch), batch) + w.writeLevelsSpaced(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) + + start := int(valueOffset) * typeLen + end := int(valueOffset+info.numSpaced()) * typeLen + writeBits := validBits + writeBitsOffset := validBitsOffset + valueOffset + if w.bitsBuffer != nil { + writeBits = w.bitsBuffer.Bytes() + writeBitsOffset = 0 + } + w.writeArrowValuesSpaced(values[start:end], typeLen, info.batchNum, writeBits, writeBitsOffset) + if err := w.commitWriteAndCheckPageLimit(batch, info.numSpaced()); err != nil { + panic(err) + } + valueOffset += info.numSpaced() + w.checkDictionarySizeLimit() + }) + return valueOffset, nil +} diff --git a/parquet/internal/encoding/fixed_len_byte_array_encoder.go b/parquet/internal/encoding/fixed_len_byte_array_encoder.go index 190802eec..cb22f0d5f 100644 --- a/parquet/internal/encoding/fixed_len_byte_array_encoder.go +++ b/parquet/internal/encoding/fixed_len_byte_array_encoder.go @@ -56,6 +56,20 @@ func (enc *PlainFixedLenByteArrayEncoder) Put(in []parquet.FixedLenByteArray) { } } +// PutArrow writes fixed-width values already laid out in an Arrow value buffer. +// The buffer contains typeLen bytes for each value. +func (enc *PlainFixedLenByteArrayEncoder) PutArrow(values []byte) { + if len(values) == 0 { + return + } + if enc.typeLen <= 0 || len(values)%enc.typeLen != 0 { + panic("parquet: Arrow fixed-length values are not aligned to the type length") + } + + enc.sink.Reserve(len(values)) + enc.sink.UnsafeWrite(values) +} + func (enc *PlainFixedLenByteArrayEncoder) Release() { enc.encoder.Release() enc.zeroValue = nil @@ -82,6 +96,35 @@ func (enc *PlainFixedLenByteArrayEncoder) PutSpaced(in []parquet.FixedLenByteArr } } +// PutArrowSpaced writes fixed-width values from an Arrow value buffer while +// skipping values whose corresponding validity bits are unset. +func (enc *PlainFixedLenByteArrayEncoder) PutArrowSpaced(values []byte, validBits []byte, validBitsOffset int64) { + if validBits == nil { + enc.PutArrow(values) + return + } + if enc.typeLen <= 0 || len(values)%enc.typeLen != 0 { + panic("parquet: Arrow fixed-length values are not aligned to the type length") + } + + nvalues := int64(len(values) / enc.typeLen) + if enc.bitSetReader == nil { + enc.bitSetReader = bitutils.NewSetBitRunReader(validBits, validBitsOffset, nvalues) + } else { + enc.bitSetReader.Reset(validBits, validBitsOffset, nvalues) + } + + for { + run := enc.bitSetReader.NextRun() + if run.Length == 0 { + break + } + start := int(run.Pos) * enc.typeLen + end := int(run.Pos+run.Length) * enc.typeLen + enc.PutArrow(values[start:end]) + } +} + // Type returns the underlying physical type this encoder works with, Fixed Length byte arrays. func (PlainFixedLenByteArrayEncoder) Type() parquet.Type { return parquet.Types.FixedLenByteArray diff --git a/parquet/internal/encoding/fixed_len_byte_array_encoder_test.go b/parquet/internal/encoding/fixed_len_byte_array_encoder_test.go index 1edee3156..9c6240d57 100644 --- a/parquet/internal/encoding/fixed_len_byte_array_encoder_test.go +++ b/parquet/internal/encoding/fixed_len_byte_array_encoder_test.go @@ -102,3 +102,29 @@ func TestPlainFixedLenByteArrayEncoder_ReusesZeroValue(t *testing.T) { encoder.Put([]parquet.FixedLenByteArray{nil}) require.Same(t, &zeroValue[0], &encoder.zeroValue[0]) } + +func TestPlainFixedLenByteArrayEncoder_PutArrow(t *testing.T) { + sink := NewPooledBufferWriter(0) + elem := schema.NewFixedLenByteArrayNode("test", parquet.Repetitions.Required, 4, 0) + descr := schema.NewColumn(elem, 0, 0) + encoder := &PlainFixedLenByteArrayEncoder{ + encoder: encoder{ + descr: descr, + typeLen: 4, + sink: sink, + }, + } + defer encoder.Release() + + values := []byte("abcdefghijklmnop") + encoder.PutArrow(values) + require.Equal(t, values, sink.Bytes()) + + sink.Reset(0) + encoder.PutArrowSpaced(values, []byte{0b0101}, 0) + require.Equal(t, []byte("abcdijkl"), sink.Bytes()) + + sink.Reset(0) + encoder.PutArrowSpaced(values, []byte{0b1010}, 1) + require.Equal(t, []byte("abcdijkl"), sink.Bytes()) +} diff --git a/parquet/metadata/fixed_len_byte_array_arrow.go b/parquet/metadata/fixed_len_byte_array_arrow.go new file mode 100644 index 000000000..88bae9b3a --- /dev/null +++ b/parquet/metadata/fixed_len_byte_array_arrow.go @@ -0,0 +1,144 @@ +// 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 metadata + +import ( + "github.com/apache/arrow-go/v18/internal/bitutils" + "github.com/apache/arrow-go/v18/parquet" +) + +func arrowFixedLenValueCount(values []byte, byteWidth int) int64 { + if byteWidth <= 0 || len(values)%byteWidth != 0 { + panic("parquet: Arrow fixed-length values are not aligned to the type length") + } + return int64(len(values) / byteWidth) +} + +// UpdateFromArrowFixedWidth updates fixed-length byte-array statistics from +// values laid out in an Arrow value buffer. +func (s *FixedLenByteArrayStatistics) UpdateFromArrowFixedWidth(values []byte, byteWidth int, numNull int64) { + nvalues := arrowFixedLenValueCount(values, byteWidth) + s.IncNulls(numNull) + s.nvalues += nvalues + if nvalues == 0 { + return + } + + min, max := s.defaultMin(), s.defaultMax() + for offset := 0; offset < len(values); offset += byteWidth { + value := parquet.FixedLenByteArray(values[offset : offset+byteWidth]) + min = s.minval(min, value) + max = s.maxval(max, value) + } + s.SetMinMax(min, max) +} + +// UpdateFromArrowFixedWidthSpaced updates fixed-length byte-array statistics +// from an Arrow value buffer whose null positions are described by validBits. +func (s *FixedLenByteArrayStatistics) UpdateFromArrowFixedWidthSpaced(values []byte, byteWidth int, validBits []byte, validBitsOffset, numNull int64) { + nvalues := arrowFixedLenValueCount(values, byteWidth) + if validBits == nil { + s.UpdateFromArrowFixedWidth(values, byteWidth, numNull) + return + } + + s.IncNulls(numNull) + s.nvalues += nvalues - numNull + if nvalues == 0 || nvalues == numNull { + return + } + + if s.bitSetReader == nil { + s.bitSetReader = bitutils.NewSetBitRunReader(validBits, validBitsOffset, nvalues) + } else { + s.bitSetReader.Reset(validBits, validBitsOffset, nvalues) + } + + min, max := s.defaultMin(), s.defaultMax() + for { + run := s.bitSetReader.NextRun() + if run.Length == 0 { + break + } + for pos := run.Pos; pos < run.Pos+run.Length; pos++ { + start := int(pos) * byteWidth + value := parquet.FixedLenByteArray(values[start : start+byteWidth]) + min = s.minval(min, value) + max = s.maxval(max, value) + } + } + s.SetMinMax(min, max) +} + +// InsertArrowFixedLenHashes inserts hashes for fixed-length values laid out in +// an Arrow value buffer. +func InsertArrowFixedLenHashes(b BloomFilterBuilder, values []byte, byteWidth int) { + if len(values) == 0 { + return + } + arrowFixedLenValueCount(values, byteWidth) + + h := b.Hasher() + var ( + byteBatch [bloomFilterHashBatchSize][]byte + hashBatch [bloomFilterHashBatchSize]uint64 + ) + for offset := 0; offset < len(values); offset += bloomFilterHashBatchSize * byteWidth { + end := min(offset+bloomFilterHashBatchSize*byteWidth, len(values)) + n := (end - offset) / byteWidth + for i := 0; i < n; i++ { + start := offset + i*byteWidth + byteBatch[i] = values[start : start+byteWidth] + } + b.InsertBulk(sum64s(h, byteBatch[:n], hashBatch[:n])) + } +} + +// InsertSpacedArrowFixedLenHashes inserts hashes for valid fixed-length values +// from an Arrow value buffer. +func InsertSpacedArrowFixedLenHashes(b BloomFilterBuilder, numValid int64, values []byte, byteWidth int, validBits []byte, validBitsOffset int64) { + if numValid == 0 { + return + } + if validBits == nil { + InsertArrowFixedLenHashes(b, values, byteWidth) + return + } + + nvalues := arrowFixedLenValueCount(values, byteWidth) + h := b.Hasher() + var ( + byteBatch [bloomFilterHashBatchSize][]byte + hashBatch [bloomFilterHashBatchSize]uint64 + ) + setReader := bitutils.NewSetBitRunReader(validBits, validBitsOffset, nvalues) + for { + run := setReader.NextRun() + if run.Length == 0 { + break + } + for pos := run.Pos; pos < run.Pos+run.Length; pos += bloomFilterHashBatchSize { + end := min(pos+int64(bloomFilterHashBatchSize), run.Pos+run.Length) + n := int(end - pos) + for i := 0; i < n; i++ { + start := int(pos+int64(i)) * byteWidth + byteBatch[i] = values[start : start+byteWidth] + } + b.InsertBulk(sum64s(h, byteBatch[:n], hashBatch[:n])) + } + } +} diff --git a/parquet/metadata/fixed_len_byte_array_arrow_test.go b/parquet/metadata/fixed_len_byte_array_arrow_test.go new file mode 100644 index 000000000..6ad5779d9 --- /dev/null +++ b/parquet/metadata/fixed_len_byte_array_arrow_test.go @@ -0,0 +1,69 @@ +// 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 metadata + +import ( + "testing" + + "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/schema" + "github.com/stretchr/testify/require" +) + +func TestFixedLenByteArrayStatisticsUpdateFromArrowFixedWidth(t *testing.T) { + node, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Optional, parquet.Types.FixedLenByteArray, -1, 3) + require.NoError(t, err) + descr := schema.NewColumn(node, 0, 0) + stats := NewStatistics(descr, memory.DefaultAllocator).(*FixedLenByteArrayStatistics) + + values := []byte("bbb" + "aaa" + "ccc" + "abc") + stats.UpdateFromArrowFixedWidth(values, 3, 0) + + require.Equal(t, int64(4), stats.NumValues()) + require.Equal(t, int64(0), stats.NullCount()) + require.Equal(t, []byte("aaa"), []byte(stats.Min())) + require.Equal(t, []byte("ccc"), []byte(stats.Max())) + + validBits := []byte{0b1101} + stats.Reset() + stats.UpdateFromArrowFixedWidthSpaced(values, 3, validBits, 0, 1) + require.Equal(t, int64(3), stats.NumValues()) + require.Equal(t, int64(1), stats.NullCount()) + require.Equal(t, []byte("abc"), []byte(stats.Min())) + require.Equal(t, []byte("ccc"), []byte(stats.Max())) +} + +func TestInsertArrowFixedLenHashes(t *testing.T) { + values := []byte("aaaa" + "bbbb" + "cccc" + "dddd") + parquetValues := []parquet.FixedLenByteArray{ + values[0:4], values[4:8], values[8:12], values[12:16], + } + + bloom := newBatchRecordingBloomFilter(xxhasher{}) + InsertArrowFixedLenHashes(bloom, values, 4) + require.Equal(t, GetHashes(xxhasher{}, parquetValues), flattenHashBatches(bloom.batches)) + + validBits := make([]byte, bitutil.BytesForBits(6)) + bitutil.SetBit(validBits, 2) + bitutil.SetBit(validBits, 4) + valid := []parquet.FixedLenByteArray{parquetValues[1], parquetValues[3]} + bloom = newBatchRecordingBloomFilter(xxhasher{}) + InsertSpacedArrowFixedLenHashes(bloom, 2, values, 4, validBits, 1) + require.Equal(t, GetHashes(xxhasher{}, valid), flattenHashBatches(bloom.batches)) +} diff --git a/parquet/pqarrow/encode_arrow.go b/parquet/pqarrow/encode_arrow.go index 8b49fb954..3824aa982 100644 --- a/parquet/pqarrow/encode_arrow.go +++ b/parquet/pqarrow/encode_arrow.go @@ -686,6 +686,21 @@ func writeDenseArrow(ctx *arrowWriteContext, cw file.ColumnChunkWriter, leafArr case *file.FixedLenByteArrayColumnChunkWriter: switch dt := leafArr.DataType().(type) { case *arrow.FixedSizeBinaryType: + if wr.SupportsArrowValues() { + buffer := leafArr.Data().Buffers()[1] + var valueBuf []byte + if buffer != nil { + start := leafArr.Data().Offset() * dt.ByteWidth + end := start + leafArr.Len()*dt.ByteWidth + valueBuf = buffer.Bytes()[start:end] + } + if !maybeParentNulls && noNulls { + _, err = wr.WriteBatchArrow(valueBuf, defLevels, repLevels) + } else { + _, err = wr.WriteBatchSpacedArrow(valueBuf, defLevels, repLevels, leafArr.NullBitmapBytes(), int64(leafArr.Data().Offset())) + } + return err + } data := make([]parquet.FixedLenByteArray, leafArr.Len()) for idx := range data { data[idx] = leafArr.(*array.FixedSizeBinary).Value(idx) diff --git a/parquet/pqarrow/fixed_size_binary_bench_test.go b/parquet/pqarrow/fixed_size_binary_bench_test.go new file mode 100644 index 000000000..035438e15 --- /dev/null +++ b/parquet/pqarrow/fixed_size_binary_bench_test.go @@ -0,0 +1,90 @@ +// 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" + "strconv" + "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/pqarrow" +) + +func benchmarkFixedSizeBinaryTable(mem memory.Allocator, n, byteWidth int, nullable bool) (arrow.Table, int64) { + builder := array.NewFixedSizeBinaryBuilder(mem, &arrow.FixedSizeBinaryType{ByteWidth: byteWidth}) + builder.Reserve(n) + for i := 0; i < n; i++ { + if nullable && i%10 == 0 { + builder.AppendNull() + continue + } + + value := make([]byte, byteWidth) + copy(value, strconv.AppendInt(nil, int64(i), 10)) + builder.Append(value) + } + arr := builder.NewArray() + builder.Release() + + sch := arrow.NewSchema([]arrow.Field{{ + Name: "value", + Type: &arrow.FixedSizeBinaryType{ByteWidth: byteWidth}, + Nullable: nullable, + }}, nil) + col := arrow.NewColumnFromArr(sch.Field(0), arr) + arr.Release() + tbl := array.NewTable(sch, []arrow.Column{col}, int64(n)) + col.Release() + return tbl, int64(n * byteWidth) +} + +func BenchmarkWriteArrowFixedSizeBinary(b *testing.B) { + const ( + n = 64 * 1024 + byteWidth = 16 + ) + mem := memory.DefaultAllocator + + for _, nullable := range []bool{false, true} { + tbl, inputBytes := benchmarkFixedSizeBinaryTable(mem, n, byteWidth, nullable) + b.Run("nullable="+strconv.FormatBool(nullable), func(b *testing.B) { + defer tbl.Release() + for _, stats := range []bool{false, true} { + b.Run("stats="+strconv.FormatBool(stats), func(b *testing.B) { + props := parquet.NewWriterProperties( + parquet.WithDictionaryDefault(false), + parquet.WithStats(stats), + parquet.WithCompression(compress.Codecs.Uncompressed), + ) + b.SetBytes(inputBytes) + b.ReportAllocs() + for b.Loop() { + var buf bytes.Buffer + if err := pqarrow.WriteTable(tbl, &buf, int64(n), props, pqarrow.DefaultWriterProps()); err != nil { + b.Fatal(err) + } + } + }) + } + }) + } +} diff --git a/parquet/pqarrow/fixed_size_binary_test.go b/parquet/pqarrow/fixed_size_binary_test.go new file mode 100644 index 000000000..c56b3f86d --- /dev/null +++ b/parquet/pqarrow/fixed_size_binary_test.go @@ -0,0 +1,189 @@ +// 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 ( + "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/pqarrow" + "github.com/stretchr/testify/require" +) + +func fixedSizeBinaryTable(t *testing.T, values [][]byte, byteWidth int) arrow.Table { + t.Helper() + mem := memory.DefaultAllocator + builder := array.NewFixedSizeBinaryBuilder(mem, &arrow.FixedSizeBinaryType{ByteWidth: byteWidth}) + for _, value := range values { + if value == nil { + builder.AppendNull() + } else { + builder.Append(value) + } + } + arr := builder.NewArray() + builder.Release() + + field := arrow.Field{ + Name: "value", + Type: &arrow.FixedSizeBinaryType{ByteWidth: byteWidth}, + Nullable: true, + } + sch := arrow.NewSchema([]arrow.Field{field}, nil) + col := arrow.NewColumnFromArr(field, arr) + arr.Release() + tbl := array.NewTable(sch, []arrow.Column{col}, int64(len(values))) + col.Release() + return tbl +} + +func TestWriteArrowFixedSizeBinaryDirect(t *testing.T) { + tbl := fixedSizeBinaryTable(t, [][]byte{ + {0x03, 0x02, 0x01}, + nil, + {0x09, 0x08, 0x07}, + {0x06, 0x05, 0x04}, + {0x0c, 0x0b, 0x0a}, + {0x0f, 0x0e, 0x0d}, + }, 3) + defer tbl.Release() + + for _, encoding := range []parquet.Encoding{parquet.Encodings.Plain, parquet.Encodings.ByteStreamSplit} { + t.Run(encoding.String(), func(t *testing.T) { + writerProps := parquet.NewWriterProperties( + parquet.WithDictionaryDefault(false), + parquet.WithEncodingFor("value", encoding), + parquet.WithStats(true), + parquet.WithBatchSize(2), + parquet.WithDataPageSize(16), + parquet.WithPageIndexEnabled(true), + parquet.WithBloomFilterEnabledFor("value", true), + parquet.WithBloomFilterNDVFor("value", tbl.NumRows()), + ) + data := writeParquetTable(t, tbl, tbl.NumRows(), writerProps) + got := readParquetTable(t, data, pqarrow.ArrowReadProperties{}) + defer got.Release() + assertTableColumnsEqual(t, tbl, got) + }) + } +} + +func TestWriteArrowFixedSizeBinaryDirectWithSlice(t *testing.T) { + full := fixedSizeBinaryTable(t, [][]byte{ + {0x00, 0x01, 0x02, 0x03}, + {0x04, 0x05, 0x06, 0x07}, + {0x08, 0x09, 0x0a, 0x0b}, + nil, + {0x10, 0x11, 0x12, 0x13}, + {0x14, 0x15, 0x16, 0x17}, + }, 4) + defer full.Release() + + sliced := array.NewSlice(full.Column(0).Data().Chunk(0), 1, 5) + defer sliced.Release() + field := arrow.Field{Name: "value", Type: &arrow.FixedSizeBinaryType{ByteWidth: 4}, Nullable: true} + sch := arrow.NewSchema([]arrow.Field{field}, nil) + col := arrow.NewColumnFromArr(field, sliced) + tbl := array.NewTable(sch, []arrow.Column{col}, int64(sliced.Len())) + col.Release() + defer tbl.Release() + + writerProps := parquet.NewWriterProperties( + parquet.WithDictionaryDefault(false), + parquet.WithBatchSize(2), + parquet.WithStats(true), + ) + data := writeParquetTable(t, tbl, tbl.NumRows(), writerProps) + got := readParquetTable(t, data, pqarrow.ArrowReadProperties{}) + defer got.Release() + assertTableColumnsEqual(t, tbl, got) +} + +func TestWriteArrowFixedSizeBinaryDirectNestedList(t *testing.T) { + mem := memory.DefaultAllocator + dtype := &arrow.FixedSizeBinaryType{ByteWidth: 3} + builder := array.NewListBuilder(mem, dtype) + values := builder.ValueBuilder().(*array.FixedSizeBinaryBuilder) + + builder.Append(true) + values.Append([]byte("aaa")) + values.Append([]byte("bbb")) + builder.AppendNull() + builder.Append(true) + builder.Append(true) + values.Append([]byte("ccc")) + values.AppendNull() + arr := builder.NewListArray() + builder.Release() + + field := arrow.Field{Name: "value", Type: arr.DataType(), Nullable: true} + sch := arrow.NewSchema([]arrow.Field{field}, nil) + col := arrow.NewColumnFromArr(field, arr) + arr.Release() + tbl := array.NewTable(sch, []arrow.Column{col}, 4) + col.Release() + defer tbl.Release() + + writerProps := parquet.NewWriterProperties( + parquet.WithDictionaryDefault(false), + parquet.WithBatchSize(2), + parquet.WithStats(true), + parquet.WithPageIndexEnabled(true), + ) + data := writeParquetTable(t, tbl, tbl.NumRows(), writerProps) + got := readParquetTable(t, data, pqarrow.ArrowReadProperties{}) + defer got.Release() + assertTableColumnsEqual(t, tbl, got) +} + +func TestWriteArrowFixedSizeBinaryDictionaryFallbackPath(t *testing.T) { + tbl := fixedSizeBinaryTable(t, [][]byte{ + []byte("foo!"), + []byte("bar!"), + nil, + []byte("foo!"), + []byte("baz!"), + }, 4) + defer tbl.Release() + + writerProps := parquet.NewWriterProperties( + parquet.WithDictionaryDefault(true), + parquet.WithStats(true), + ) + data := writeParquetTable(t, tbl, tbl.NumRows(), writerProps) + got := readParquetTable(t, data, pqarrow.ArrowReadProperties{}) + defer got.Release() + assertTableColumnsEqual(t, tbl, got) +} + +func TestWriteArrowFixedSizeBinaryAllNull(t *testing.T) { + tbl := fixedSizeBinaryTable(t, [][]byte{nil, nil, nil, nil}, 8) + defer tbl.Release() + + writerProps := parquet.NewWriterProperties( + parquet.WithDictionaryDefault(false), + parquet.WithStats(true), + ) + data := writeParquetTable(t, tbl, tbl.NumRows(), writerProps) + got := readParquetTable(t, data, pqarrow.ArrowReadProperties{}) + defer got.Release() + assertTableColumnsEqual(t, tbl, got) + require.Equal(t, tbl.NumRows(), got.NumRows()) +} From 387c41db1cba95d7ef785d73ac76595e2cc414b3 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Sun, 30 Aug 2026 22:50:33 +0200 Subject: [PATCH 2/2] fix(parquet): preserve fixed-width write limits and null counts --- ...ixed_len_byte_array_column_writer_arrow.go | 19 ++++-- ...len_byte_array_column_writer_arrow_test.go | 62 +++++++++++++++++ parquet/pqarrow/fixed_size_binary_test.go | 68 +++++++++++++++++++ 3 files changed, 143 insertions(+), 6 deletions(-) create mode 100644 parquet/file/fixed_len_byte_array_column_writer_arrow_test.go diff --git a/parquet/file/fixed_len_byte_array_column_writer_arrow.go b/parquet/file/fixed_len_byte_array_column_writer_arrow.go index 6c6fcbfdb..8acb31b19 100644 --- a/parquet/file/fixed_len_byte_array_column_writer_arrow.go +++ b/parquet/file/fixed_len_byte_array_column_writer_arrow.go @@ -30,11 +30,17 @@ type fixedLenByteArrayArrowEncoder interface { } // SupportsArrowValues reports whether the active fixed-length byte-array -// encoder can consume values directly from an Arrow value buffer. +// encoder and configured batch size support writing an Arrow value buffer directly. func (w *FixedLenByteArrayColumnChunkWriter) SupportsArrowValues() bool { if _, ok := w.currentEncoder.(fixedLenByteArrayArrowEncoder); !ok { return false } + typeLen := int64(w.descr.TypeLength()) + batchSize := w.props.WriteBatchSize() + const maxSafeBatchDataSize int64 = 1 << 30 + if typeLen <= 0 || batchSize <= 0 || batchSize > max(1, maxSafeBatchDataSize/(typeLen+4)) { + return false + } if w.pageStatistics == nil { return true } @@ -42,17 +48,17 @@ func (w *FixedLenByteArrayColumnChunkWriter) SupportsArrowValues() bool { return ok } -func (w *FixedLenByteArrayColumnChunkWriter) writeArrowValues(values []byte, byteWidth int) { +func (w *FixedLenByteArrayColumnChunkWriter) writeArrowValues(values []byte, byteWidth int, numNulls int64) { w.currentEncoder.(fixedLenByteArrayArrowEncoder).PutArrow(values) if stats, ok := w.pageStatistics.(*metadata.FixedLenByteArrayStatistics); ok { - stats.UpdateFromArrowFixedWidth(values, byteWidth, 0) + stats.UpdateFromArrowFixedWidth(values, byteWidth, numNulls) } if w.bloomFilter != nil && w.currentEncoder.Encoding() != parquet.Encodings.PlainDict { metadata.InsertArrowFixedLenHashes(w.bloomFilter, values, byteWidth) } } -func (w *FixedLenByteArrayColumnChunkWriter) writeArrowValuesSpaced(values []byte, byteWidth int, numRead int64, validBits []byte, validBitsOffset int64) { +func (w *FixedLenByteArrayColumnChunkWriter) writeArrowValuesSpaced(values []byte, byteWidth int, numRead, numValues int64, validBits []byte, validBitsOffset int64) { enc := w.currentEncoder.(fixedLenByteArrayArrowEncoder) numSpaced := int64(len(values) / byteWidth) if numSpaced == numRead { @@ -63,6 +69,7 @@ func (w *FixedLenByteArrayColumnChunkWriter) writeArrowValuesSpaced(values []byt if stats, ok := w.pageStatistics.(*metadata.FixedLenByteArrayStatistics); ok { stats.UpdateFromArrowFixedWidthSpaced(values, byteWidth, validBits, validBitsOffset, numSpaced-numRead) + stats.IncNulls(numValues - numSpaced) } if w.bloomFilter != nil && w.currentEncoder.Encoding() != parquet.Encodings.PlainDict { metadata.InsertSpacedArrowFixedLenHashes(w.bloomFilter, numRead, values, byteWidth, validBits, validBitsOffset) @@ -97,7 +104,7 @@ func (w *FixedLenByteArrayColumnChunkWriter) WriteBatchArrow(values []byte, defL toWrite := w.writeLevels(batch, levelSliceOrNil(defLevels, offset, batch), levelSliceOrNil(repLevels, offset, batch)) start := int(valueOffset) * typeLen end := int(valueOffset+toWrite) * typeLen - w.writeArrowValues(values[start:end], typeLen) + w.writeArrowValues(values[start:end], typeLen, batch-toWrite) if err := w.commitWriteAndCheckPageLimit(batch, toWrite); err != nil { panic(err) } @@ -143,7 +150,7 @@ func (w *FixedLenByteArrayColumnChunkWriter) WriteBatchSpacedArrow(values []byte writeBits = w.bitsBuffer.Bytes() writeBitsOffset = 0 } - w.writeArrowValuesSpaced(values[start:end], typeLen, info.batchNum, writeBits, writeBitsOffset) + w.writeArrowValuesSpaced(values[start:end], typeLen, info.batchNum, batch, writeBits, writeBitsOffset) if err := w.commitWriteAndCheckPageLimit(batch, info.numSpaced()); err != nil { panic(err) } diff --git a/parquet/file/fixed_len_byte_array_column_writer_arrow_test.go b/parquet/file/fixed_len_byte_array_column_writer_arrow_test.go new file mode 100644 index 000000000..54eca951b --- /dev/null +++ b/parquet/file/fixed_len_byte_array_column_writer_arrow_test.go @@ -0,0 +1,62 @@ +// 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 + +import ( + "fmt" + "testing" + + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/apache/arrow-go/v18/parquet" + "github.com/apache/arrow-go/v18/parquet/internal/encoding" + "github.com/apache/arrow-go/v18/parquet/schema" + "github.com/stretchr/testify/require" +) + +func TestFixedLenByteArraySupportsArrowValuesBatchLimit(t *testing.T) { + for _, tc := range []struct { + byteWidth int32 + limit int64 + }{ + {byteWidth: 3, limit: 153391689}, + {byteWidth: 1 << 20, limit: 1023}, + {byteWidth: 1 << 30, limit: 1}, + } { + byteWidth, limit := tc.byteWidth, tc.limit + for _, batchSize := range []int64{-1, 0, limit, limit + 1} { + t.Run(fmt.Sprintf("width-%d/batch-%d", byteWidth, batchSize), func(t *testing.T) { + node, err := schema.NewPrimitiveNode("value", parquet.Repetitions.Required, parquet.Types.FixedLenByteArray, -1, byteWidth) + require.NoError(t, err) + descr := schema.NewColumn(node, 0, 0) + enc := encoding.NewEncoder(parquet.Types.FixedLenByteArray, parquet.Encodings.Plain, false, descr, memory.DefaultAllocator) + defer enc.Release() + writer := &FixedLenByteArrayColumnChunkWriter{columnWriter: columnWriter{ + descr: descr, + props: parquet.NewWriterProperties(parquet.WithBatchSize(batchSize)), + currentEncoder: enc, + }} + require.Equal(t, batchSize == limit, writer.SupportsArrowValues()) + if batchSize != limit { + _, err := writer.WriteBatchArrow(nil, nil, nil) + require.ErrorContains(t, err, "does not support Arrow values") + _, err = writer.WriteBatchSpacedArrow(nil, nil, nil, nil, 0) + require.ErrorContains(t, err, "does not support Arrow values") + } + }) + } + } +} diff --git a/parquet/pqarrow/fixed_size_binary_test.go b/parquet/pqarrow/fixed_size_binary_test.go index c56b3f86d..58c0c36e2 100644 --- a/parquet/pqarrow/fixed_size_binary_test.go +++ b/parquet/pqarrow/fixed_size_binary_test.go @@ -17,12 +17,15 @@ package pqarrow_test import ( + "bytes" + "fmt" "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/file" "github.com/apache/arrow-go/v18/parquet/pqarrow" "github.com/stretchr/testify/require" ) @@ -187,3 +190,68 @@ func TestWriteArrowFixedSizeBinaryAllNull(t *testing.T) { assertTableColumnsEqual(t, tbl, got) require.Equal(t, tbl.NumRows(), got.NumRows()) } + +func TestWriteArrowFixedSizeBinaryNestedNullStatistics(t *testing.T) { + for _, pageVersion := range []parquet.DataPageVersion{parquet.DataPageV1, parquet.DataPageV2} { + for _, nullableValues := range []bool{false, true} { + t.Run(fmt.Sprintf("page-%d/nullable-values-%t", pageVersion, nullableValues), func(t *testing.T) { + builder := array.NewListBuilder(memory.DefaultAllocator, &arrow.FixedSizeBinaryType{ByteWidth: 3}) + defer builder.Release() + values := builder.ValueBuilder().(*array.FixedSizeBinaryBuilder) + builder.Append(true) + values.Append([]byte("aaa")) + builder.AppendNull() + builder.Append(true) + builder.Append(true) + values.Append([]byte("zzz")) + nullCount := int64(2) + if nullableValues { + values.AppendNull() + nullCount++ + } + arr := builder.NewListArray() + defer arr.Release() + field := arrow.Field{Name: "value", Type: arr.DataType(), Nullable: true} + column := arrow.NewColumnFromArr(field, arr) + defer column.Release() + tbl := array.NewTable(arrow.NewSchema([]arrow.Field{field}, nil), []arrow.Column{column}, int64(arr.Len())) + defer tbl.Release() + props := parquet.NewWriterProperties(parquet.WithDictionaryDefault(false), + parquet.WithDataPageVersion(pageVersion), parquet.WithBatchSize(2)) + data := writeParquetTable(t, tbl, tbl.NumRows(), props) + reader, err := file.NewParquetReader(bytes.NewReader(data)) + require.NoError(t, err) + defer reader.Close() + chunk, err := reader.MetaData().RowGroup(0).ColumnChunk(0) + require.NoError(t, err) + stats, err := chunk.Statistics() + require.NoError(t, err) + require.Equal(t, nullCount, stats.NullCount()) + require.Equal(t, int64(2), stats.NumValues()) + require.Equal(t, []byte("aaa"), stats.EncodeMin()) + require.Equal(t, []byte("zzz"), stats.EncodeMax()) + }) + } + } +} + +func TestWriteArrowFixedSizeBinaryBatchSizeFallback(t *testing.T) { + for _, batchSize := range []int64{-1, 0, (1<<30)/7 + 1} { + for _, nullable := range []bool{false, true} { + t.Run(fmt.Sprintf("batch-%d/nullable-%t", batchSize, nullable), func(t *testing.T) { + values := [][]byte{[]byte("aaa"), []byte("zzz")} + if nullable { + values = append(values, nil) + } + tbl := fixedSizeBinaryTable(t, values, 3) + defer tbl.Release() + props := parquet.NewWriterProperties(parquet.WithDictionaryDefault(false), + parquet.WithBatchSize(batchSize)) + data := writeParquetTable(t, tbl, tbl.NumRows(), props) + got := readParquetTable(t, data, pqarrow.ArrowReadProperties{}) + defer got.Release() + assertTableColumnsEqual(t, tbl, got) + }) + } + } +}