perf(parquet/pqarrow): write FixedSizeBinary values directly

parquet/file/fixed_len_byte_array_column_writer_arrow.go

@@ -0,0 +1,161 @@
+// 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 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
+	}
+	_, ok := w.pageStatistics.(*metadata.FixedLenByteArrayStatistics)
+	return ok
+}
+
+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, 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, numValues 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)
+		stats.IncNulls(numValues - numSpaced)
+	}
+	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, batch-toWrite)
+		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, batch, writeBits, writeBitsOffset)
+		if err := w.commitWriteAndCheckPageLimit(batch, info.numSpaced()); err != nil {
+			panic(err)
+		}
+		valueOffset += info.numSpaced()
+		w.checkDictionarySizeLimit()
+	})
+	return valueOffset, nil
+}

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")
+				}
+			})
+		}
+	}
+}

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

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())
+}