github.com/apache/arrow/go/v14@v14.0.1/parquet/file/page_writer.go

// 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 (
	"bytes"
	"sync"

	"github.com/apache/arrow/go/v14/arrow/memory"
	"github.com/apache/arrow/go/v14/parquet"
	"github.com/apache/arrow/go/v14/parquet/compress"
	"github.com/apache/arrow/go/v14/parquet/internal/encoding"
	"github.com/apache/arrow/go/v14/parquet/internal/encryption"
	format "github.com/apache/arrow/go/v14/parquet/internal/gen-go/parquet"
	"github.com/apache/arrow/go/v14/parquet/internal/thrift"
	"github.com/apache/arrow/go/v14/parquet/internal/utils"
	"github.com/apache/arrow/go/v14/parquet/metadata"
	libthrift "github.com/apache/thrift/lib/go/thrift"
	"golang.org/x/xerrors"
)

// PageWriter is the interface for both serialized and buffered page writers
type PageWriter interface {
	// Closes the current page, flushing any buffered data pages/dictionary pages
	// based on the input parameters. Subsequent calls have no effect.
	Close(hasDict, fallback bool) error
	// Write the provided datapage out to the underlying writer
	WriteDataPage(page DataPage) (int64, error)
	// Write the provided dictionary page out to the underlying writer
	WriteDictionaryPage(page *DictionaryPage) (int64, error)
	// returns true if there is a configured compressor for the data
	HasCompressor() bool
	// use the configured compressor and writer properties to compress the data in src
	// using the buffer buf. Returns the slice of the compressed bytes which may be
	// the bytes in the provided buffer
	Compress(buf *bytes.Buffer, src []byte) []byte
	// Allow reuse of the pagewriter object by resetting it using these values instead
	// of having to create a new object.
	Reset(sink utils.WriterTell, codec compress.Compression, compressionLevel int, metadata *metadata.ColumnChunkMetaDataBuilder, rgOrdinal, columnOrdinal int16, metaEncryptor, dataEncryptor encryption.Encryptor) error
}

type serializedPageWriter struct {
	mem      memory.Allocator
	metaData *metadata.ColumnChunkMetaDataBuilder
	sink     utils.WriterTell

	nvalues           int64
	dictPageOffset    int64
	dataPageOffset    int64
	totalUncompressed int64
	totalCompressed   int64
	pageOrdinal       int16
	rgOrdinal         int16
	columnOrdinal     int16

	compressLevel int
	compressor    compress.Codec
	metaEncryptor encryption.Encryptor
	dataEncryptor encryption.Encryptor
	encryptionBuf bytes.Buffer

	dataPageAAD       []byte
	dataPageHeaderAAD []byte

	dictEncodingStats map[parquet.Encoding]int32
	dataEncodingStats map[parquet.Encoding]int32

	thriftSerializer *thrift.Serializer
}

func createSerializedPageWriter(sink utils.WriterTell, codec compress.Compression, compressionLevel int, metadata *metadata.ColumnChunkMetaDataBuilder, rowGroupOrdinal, columnChunkOrdinal int16, mem memory.Allocator, metaEncryptor, dataEncryptor encryption.Encryptor) (PageWriter, error) {
	var (
		compressor compress.Codec
		err        error
	)
	if codec != compress.Codecs.Uncompressed {
		compressor, err = compress.GetCodec(codec)
		if err != nil {
			return nil, err
		}
	}

	pgwriter := &serializedPageWriter{
		sink:              sink,
		compressor:        compressor,
		compressLevel:     compressionLevel,
		metaData:          metadata,
		rgOrdinal:         rowGroupOrdinal,
		columnOrdinal:     columnChunkOrdinal,
		mem:               mem,
		metaEncryptor:     metaEncryptor,
		dataEncryptor:     dataEncryptor,
		dictEncodingStats: make(map[parquet.Encoding]int32),
		dataEncodingStats: make(map[parquet.Encoding]int32),
		thriftSerializer:  thrift.NewThriftSerializer(),
	}
	if metaEncryptor != nil || dataEncryptor != nil {
		pgwriter.initEncryption()
	}
	return pgwriter, nil
}

// NewPageWriter returns a page writer using either the buffered or serialized implementations
func NewPageWriter(sink utils.WriterTell, codec compress.Compression, compressionLevel int, metadata *metadata.ColumnChunkMetaDataBuilder, rowGroupOrdinal, columnChunkOrdinal int16, mem memory.Allocator, buffered bool, metaEncryptor, dataEncryptor encryption.Encryptor) (PageWriter, error) {
	if buffered {
		return newBufferedPageWriter(sink, codec, compressionLevel, metadata, rowGroupOrdinal, columnChunkOrdinal, mem, metaEncryptor, dataEncryptor)
	}
	return createSerializedPageWriter(sink, codec, compressionLevel, metadata, rowGroupOrdinal, columnChunkOrdinal, mem, metaEncryptor, dataEncryptor)
}

// Reset allows reusing the pagewriter object instead of creating a new one.
func (pw *serializedPageWriter) Reset(sink utils.WriterTell, codec compress.Compression, compressionLevel int, metadata *metadata.ColumnChunkMetaDataBuilder, rowGroupOrdinal, columnChunkOrdinal int16, metaEncryptor, dataEncryptor encryption.Encryptor) error {
	var (
		compressor compress.Codec
		err        error
	)
	if codec != compress.Codecs.Uncompressed {
		compressor, err = compress.GetCodec(codec)
		if err != nil {
			return err
		}
	}

	pw.sink = sink
	pw.compressor = compressor
	pw.compressLevel = compressionLevel
	pw.metaData = metadata
	pw.rgOrdinal = rowGroupOrdinal
	pw.columnOrdinal = columnChunkOrdinal
	pw.metaEncryptor = metaEncryptor
	pw.dataEncryptor = dataEncryptor
	pw.dictEncodingStats = make(map[parquet.Encoding]int32)
	pw.dataEncodingStats = make(map[parquet.Encoding]int32)

	pw.nvalues = 0
	pw.dictPageOffset = 0
	pw.dataPageOffset = 0
	pw.totalUncompressed = 0
	pw.totalCompressed = 0
	pw.pageOrdinal = 0

	if metaEncryptor != nil || dataEncryptor != nil {
		pw.initEncryption()
	}
	return nil
}

func (pw *serializedPageWriter) initEncryption() {
	if pw.dataEncryptor != nil {
		pw.dataPageAAD = []byte(encryption.CreateModuleAad(pw.dataEncryptor.FileAad(), encryption.DataPageModule, pw.rgOrdinal, pw.columnOrdinal, -1))
	}
	if pw.metaEncryptor != nil {
		pw.dataPageHeaderAAD = []byte(encryption.CreateModuleAad(pw.metaEncryptor.FileAad(), encryption.DataPageHeaderModule, pw.rgOrdinal, pw.columnOrdinal, -1))
	}
}

func (pw *serializedPageWriter) updateEncryption(moduleType int8) error {
	switch moduleType {
	case encryption.ColumnMetaModule:
		pw.metaEncryptor.UpdateAad(encryption.CreateModuleAad(pw.metaEncryptor.FileAad(), moduleType, pw.rgOrdinal, pw.columnOrdinal, -1))
	case encryption.DataPageModule:
		encryption.QuickUpdatePageAad(pw.dataPageAAD, pw.pageOrdinal)
		pw.dataEncryptor.UpdateAad(string(pw.dataPageAAD))
	case encryption.DataPageHeaderModule:
		encryption.QuickUpdatePageAad(pw.dataPageHeaderAAD, pw.pageOrdinal)
		pw.metaEncryptor.UpdateAad(string(pw.dataPageHeaderAAD))
	case encryption.DictPageHeaderModule:
		pw.metaEncryptor.UpdateAad(encryption.CreateModuleAad(pw.metaEncryptor.FileAad(), moduleType, pw.rgOrdinal, pw.columnOrdinal, -1))
	case encryption.DictPageModule:
		pw.dataEncryptor.UpdateAad(encryption.CreateModuleAad(pw.dataEncryptor.FileAad(), moduleType, pw.rgOrdinal, pw.columnOrdinal, -1))
	default:
		return xerrors.New("unknown module type in updateencryption")
	}
	return nil
}

func (pw *serializedPageWriter) Close(hasDict, fallback bool) error {
	if pw.metaEncryptor != nil {
		pw.updateEncryption(encryption.ColumnMetaModule)
	}

	chunkInfo := metadata.ChunkMetaInfo{
		NumValues:        pw.nvalues,
		DictPageOffset:   pw.dictPageOffset,
		IndexPageOffset:  -1,
		DataPageOffset:   pw.dataPageOffset,
		CompressedSize:   pw.totalCompressed,
		UncompressedSize: pw.totalUncompressed,
	}
	encodingStats := metadata.EncodingStats{
		DictEncodingStats: pw.dictEncodingStats,
		DataEncodingStats: pw.dataEncodingStats,
	}
	pw.metaData.Finish(chunkInfo, hasDict, fallback, encodingStats, pw.metaEncryptor)
	_, err := pw.metaData.WriteTo(pw.sink)
	return err
}

func (pw *serializedPageWriter) Compress(buf *bytes.Buffer, src []byte) []byte {
	maxCompressed := pw.compressor.CompressBound(int64(len(src)))
	buf.Grow(int(maxCompressed))
	return pw.compressor.EncodeLevel(buf.Bytes(), src, pw.compressLevel)
}

var dataPageV1HeaderPool = sync.Pool{
	New: func() interface{} { return format.NewDataPageHeader() },
}

func (pw *serializedPageWriter) setDataPageHeader(pageHdr *format.PageHeader, page *DataPageV1) {
	pageHdr.Type = format.PageType_DATA_PAGE
	hdr := dataPageV1HeaderPool.Get().(*format.DataPageHeader)
	hdr.NumValues = page.nvals
	hdr.Encoding = page.encoding
	hdr.DefinitionLevelEncoding = page.defLvlEncoding
	hdr.RepetitionLevelEncoding = page.repLvlEncoding
	hdr.Statistics = page.statistics.ToThrift()
	pageHdr.DataPageHeader = hdr
	pageHdr.DataPageHeaderV2 = nil
	pageHdr.DictionaryPageHeader = nil
}

var dataPageV2HeaderPool = sync.Pool{
	New: func() interface{} { return format.NewDataPageHeaderV2() },
}

func (pw *serializedPageWriter) setDataPageV2Header(pageHdr *format.PageHeader, page *DataPageV2) {
	pageHdr.Type = format.PageType_DATA_PAGE_V2
	hdr := dataPageV2HeaderPool.Get().(*format.DataPageHeaderV2)
	hdr.NumValues = page.nvals
	hdr.NumNulls = page.nulls
	hdr.NumRows = page.nrows
	hdr.Encoding = page.encoding
	hdr.DefinitionLevelsByteLength = page.defLvlByteLen
	hdr.RepetitionLevelsByteLength = page.repLvlByteLen
	hdr.IsCompressed = page.compressed
	hdr.Statistics = page.statistics.ToThrift()
	pageHdr.DataPageHeaderV2 = hdr
	pageHdr.DataPageHeader = nil
	pageHdr.DictionaryPageHeader = nil
}

func (pw *serializedPageWriter) HasCompressor() bool          { return pw.compressor != nil }
func (pw *serializedPageWriter) NumValues() int64             { return pw.nvalues }
func (pw *serializedPageWriter) DictionaryPageOffset() int64  { return pw.dictPageOffset }
func (pw *serializedPageWriter) DataPageoffset() int64        { return pw.dataPageOffset }
func (pw *serializedPageWriter) TotalCompressedSize() int64   { return pw.totalCompressed }
func (pw *serializedPageWriter) TotalUncompressedSize() int64 { return pw.totalUncompressed }

func (pw *serializedPageWriter) WriteDictionaryPage(page *DictionaryPage) (int64, error) {
	uncompressed := len(page.Data())

	var data []byte
	if pw.HasCompressor() {
		var buffer bytes.Buffer
		data = pw.Compress(&buffer, page.Data())
		// data = buffer.Bytes()
	} else {
		data = page.Data()
	}

	dictPageHeader := &format.DictionaryPageHeader{
		NumValues: page.NumValues(),
		Encoding:  page.Encoding(),
		IsSorted:  libthrift.BoolPtr(page.IsSorted()),
	}

	if pw.dataEncryptor != nil {
		pw.updateEncryption(encryption.DictPageModule)
		pw.encryptionBuf.Reset()
		pw.encryptionBuf.Grow(pw.dataEncryptor.CiphertextSizeDelta() + len(data))
		pw.dataEncryptor.Encrypt(&pw.encryptionBuf, data)
		data = pw.encryptionBuf.Bytes()
	}

	pageHdr := pageHeaderPool.Get().(*format.PageHeader)
	defer pageHeaderPool.Put(pageHdr)
	pageHdr.Type = format.PageType_DICTIONARY_PAGE
	pageHdr.UncompressedPageSize = int32(uncompressed)
	pageHdr.CompressedPageSize = int32(len(data))
	pageHdr.DictionaryPageHeader = dictPageHeader
	pageHdr.DataPageHeader = nil
	pageHdr.DataPageHeaderV2 = nil

	startPos := pw.sink.Tell()
	if pw.dictPageOffset == 0 {
		pw.dictPageOffset = int64(startPos)
	}

	if pw.metaEncryptor != nil {
		if err := pw.updateEncryption(encryption.DictPageHeaderModule); err != nil {
			return 0, err
		}
	}
	headerSize, err := pw.thriftSerializer.Serialize(pageHdr, pw.sink, pw.metaEncryptor)
	if err != nil {
		return 0, err
	}
	written, err := pw.sink.Write(data)
	if err != nil {
		return 0, err
	}

	written += headerSize

	pw.totalUncompressed += int64(uncompressed + headerSize)
	pw.totalCompressed = int64(written)
	pw.dictEncodingStats[parquet.Encoding(page.encoding)]++
	return int64(written), nil
}

var pageHeaderPool = sync.Pool{
	New: func() interface{} {
		return format.NewPageHeader()
	},
}

func (pw *serializedPageWriter) WriteDataPage(page DataPage) (int64, error) {
	uncompressed := page.UncompressedSize()
	data := page.Data()

	if pw.dataEncryptor != nil {
		if err := pw.updateEncryption(encryption.DataPageModule); err != nil {
			return 0, err
		}
		pw.encryptionBuf.Reset()
		pw.encryptionBuf.Grow(pw.dataEncryptor.CiphertextSizeDelta() + len(data))
		pw.dataEncryptor.Encrypt(&pw.encryptionBuf, data)
		data = pw.encryptionBuf.Bytes()
	}

	pageHdr := pageHeaderPool.Get().(*format.PageHeader)
	defer pageHeaderPool.Put(pageHdr)
	pageHdr.UncompressedPageSize = uncompressed
	pageHdr.CompressedPageSize = int32(len(data))

	switch dpage := page.(type) {
	case *DataPageV1:
		pw.setDataPageHeader(pageHdr, dpage)
		defer dataPageV1HeaderPool.Put(pageHdr.DataPageHeader)
	case *DataPageV2:
		pw.setDataPageV2Header(pageHdr, dpage)
		defer dataPageV2HeaderPool.Put(pageHdr.DataPageHeaderV2)
	default:
		return 0, xerrors.New("parquet: unexpected page type")
	}

	startPos := pw.sink.Tell()
	if pw.pageOrdinal == 0 {
		pw.dataPageOffset = int64(startPos)
	}

	if pw.metaEncryptor != nil {
		if err := pw.updateEncryption(encryption.DataPageHeaderModule); err != nil {
			return 0, err
		}
	}
	headerSize, err := pw.thriftSerializer.Serialize(pageHdr, pw.sink, pw.metaEncryptor)
	if err != nil {
		return 0, err
	}
	written, err := pw.sink.Write(data)
	if err != nil {
		return int64(written), err
	}
	written += headerSize

	pw.totalUncompressed += int64(uncompressed) + int64(headerSize)
	pw.totalCompressed += int64(written)
	pw.nvalues += int64(page.NumValues())
	pw.dataEncodingStats[parquet.Encoding(page.Encoding())]++
	pw.pageOrdinal++
	return int64(written), nil
}

type bufferedPageWriter struct {
	finalSink          utils.WriterTell
	inMemSink          *encoding.BufferWriter
	metadata           *metadata.ColumnChunkMetaDataBuilder
	pager              *serializedPageWriter
	hasDictionaryPages bool
}

func newBufferedPageWriter(sink utils.WriterTell, codec compress.Compression, compressionLevel int, metadata *metadata.ColumnChunkMetaDataBuilder, rgOrdinal, columnOrdinal int16, mem memory.Allocator, metaEncryptor, dataEncryptor encryption.Encryptor) (PageWriter, error) {
	wr := &bufferedPageWriter{
		finalSink:          sink,
		metadata:           metadata,
		hasDictionaryPages: false,
		inMemSink:          encoding.NewBufferWriter(0, mem),
	}
	pager, err := createSerializedPageWriter(wr.inMemSink, codec, compressionLevel, metadata, rgOrdinal, columnOrdinal, mem, metaEncryptor, dataEncryptor)
	if err != nil {
		return nil, err
	}
	wr.pager = pager.(*serializedPageWriter)
	return wr, nil
}

func (bw *bufferedPageWriter) Reset(sink utils.WriterTell, codec compress.Compression, compressionLevel int, metadata *metadata.ColumnChunkMetaDataBuilder, rgOrdinal, columnOrdinal int16, metaEncryptor, dataEncryptor encryption.Encryptor) error {
	bw.finalSink = sink
	bw.metadata = metadata
	bw.hasDictionaryPages = false
	bw.inMemSink.Reset(0)

	return bw.pager.Reset(bw.inMemSink, codec, compressionLevel, metadata, rgOrdinal, columnOrdinal, metaEncryptor, dataEncryptor)
}

func (bw *bufferedPageWriter) WriteDictionaryPage(page *DictionaryPage) (int64, error) {
	bw.hasDictionaryPages = true
	return bw.pager.WriteDictionaryPage(page)
}

func (bw *bufferedPageWriter) Close(hasDict, fallback bool) error {
	if bw.pager.metaEncryptor != nil {
		bw.pager.updateEncryption(encryption.ColumnMetaModule)
	}

	position := bw.finalSink.Tell()
	dictOffset := int64(0)
	if bw.hasDictionaryPages {
		dictOffset = bw.pager.DictionaryPageOffset() + position
	}

	chunkInfo := metadata.ChunkMetaInfo{
		NumValues:        bw.pager.NumValues(),
		DictPageOffset:   dictOffset,
		IndexPageOffset:  -1,
		DataPageOffset:   bw.pager.DataPageoffset() + position,
		CompressedSize:   bw.pager.TotalCompressedSize(),
		UncompressedSize: bw.pager.TotalUncompressedSize(),
	}
	encodingStats := metadata.EncodingStats{
		DictEncodingStats: bw.pager.dictEncodingStats,
		DataEncodingStats: bw.pager.dataEncodingStats,
	}
	bw.metadata.Finish(chunkInfo, hasDict, fallback, encodingStats, bw.pager.metaEncryptor)
	bw.metadata.WriteTo(bw.inMemSink)

	buf := bw.inMemSink.Finish()
	defer buf.Release()
	_, err := bw.finalSink.Write(buf.Bytes())
	return err
}

func (bw *bufferedPageWriter) WriteDataPage(page DataPage) (int64, error) {
	return bw.pager.WriteDataPage(page)
}

func (bw *bufferedPageWriter) HasCompressor() bool {
	return bw.pager.HasCompressor()
}

func (bw *bufferedPageWriter) Compress(buf *bytes.Buffer, src []byte) []byte {
	return bw.pager.Compress(buf, src)
}