github.com/fraugster/parquet-go@v0.12.0/hybrid_encoder.go (about)

     1  package goparquet
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  	"io"
     7  )
     8  
     9  type hybridEncoder struct {
    10  	w io.Writer
    11  
    12  	left       []int32
    13  	original   io.Writer
    14  	bitWidth   int
    15  	unpackerFn pack8int32Func
    16  
    17  	data *packedArray
    18  }
    19  
    20  func newHybridEncoder(bitWidth int) *hybridEncoder {
    21  	p := &packedArray{}
    22  	return &hybridEncoder{
    23  		bitWidth:   bitWidth,
    24  		unpackerFn: pack8Int32FuncByWidth[bitWidth],
    25  		data:       p,
    26  	}
    27  }
    28  
    29  func (he *hybridEncoder) init(w io.Writer) error {
    30  	he.w = w
    31  	he.left = nil
    32  	he.original = nil
    33  
    34  	he.data.reset(he.bitWidth)
    35  	return nil
    36  }
    37  
    38  func (he *hybridEncoder) initSize(w io.Writer) error {
    39  	_ = he.init(&bytes.Buffer{})
    40  	he.original = w
    41  
    42  	return nil
    43  }
    44  
    45  func (he *hybridEncoder) write(items ...[]byte) error {
    46  	for i := range items {
    47  		if err := writeFull(he.w, items[i]); err != nil {
    48  			return err
    49  		}
    50  	}
    51  
    52  	return nil
    53  }
    54  
    55  func (he *hybridEncoder) bpEncode() error {
    56  	// If the bit width is zero, no need to write any
    57  	if he.bitWidth == 0 {
    58  		return nil
    59  	}
    60  	l := he.data.count
    61  	if x := l % 8; x != 0 {
    62  		l += 8 - x
    63  	}
    64  
    65  	header := ((l / 8) << 1) | 1
    66  	buf := make([]byte, 4) // big enough for int
    67  	cnt := binary.PutUvarint(buf, uint64(header))
    68  
    69  	return he.write(buf[:cnt], he.data.data)
    70  }
    71  
    72  func (he *hybridEncoder) encode(data []int32) error {
    73  	for i := range data {
    74  		he.data.appendSingle(data[i])
    75  	}
    76  
    77  	return nil
    78  }
    79  
    80  func (he *hybridEncoder) encodePacked(data *packedArray) error {
    81  	he.data.appendArray(data)
    82  
    83  	return nil
    84  }
    85  
    86  func (he *hybridEncoder) flush() error {
    87  	he.data.flush()
    88  	return he.bpEncode()
    89  }
    90  
    91  func (he *hybridEncoder) Close() error {
    92  	if he.bitWidth == 0 {
    93  		return nil
    94  	}
    95  	if err := he.flush(); err != nil {
    96  		return err
    97  	}
    98  
    99  	if he.original != nil {
   100  		data := he.w.(*bytes.Buffer).Bytes()
   101  		var size = uint32(len(data))
   102  		if err := binary.Write(he.original, binary.LittleEndian, size); err != nil {
   103  			return err
   104  		}
   105  		return writeFull(he.original, data)
   106  	}
   107  
   108  	return nil
   109  }