github.com/segmentio/parquet-go@v0.0.0-20230712180008-5d42db8f0d47/allocator.go (about)

     1  package parquet
     2  
     3  import "github.com/segmentio/parquet-go/internal/unsafecast"
     4  
     5  type allocator struct{ buffer []byte }
     6  
     7  func (a *allocator) makeBytes(n int) []byte {
     8  	if free := cap(a.buffer) - len(a.buffer); free < n {
     9  		newCap := 2 * cap(a.buffer)
    10  		if newCap == 0 {
    11  			newCap = 4096
    12  		}
    13  		for newCap < n {
    14  			newCap *= 2
    15  		}
    16  		a.buffer = make([]byte, 0, newCap)
    17  	}
    18  
    19  	i := len(a.buffer)
    20  	j := len(a.buffer) + n
    21  	a.buffer = a.buffer[:j]
    22  	return a.buffer[i:j:j]
    23  }
    24  
    25  func (a *allocator) copyBytes(v []byte) []byte {
    26  	b := a.makeBytes(len(v))
    27  	copy(b, v)
    28  	return b
    29  }
    30  
    31  func (a *allocator) copyString(v string) string {
    32  	b := a.makeBytes(len(v))
    33  	copy(b, v)
    34  	return unsafecast.BytesToString(b)
    35  }
    36  
    37  func (a *allocator) reset() {
    38  	a.buffer = a.buffer[:0]
    39  }
    40  
    41  // rowAllocator is a memory allocator used to make a copy of rows referencing
    42  // memory buffers that parquet-go does not have ownership of.
    43  //
    44  // This type is used in the implementation of various readers and writers that
    45  // need to capture rows passed to the ReadRows/WriteRows methods. Copies to a
    46  // local buffer is necessary in those cases to repect the reader/writer
    47  // contracts that do not allow the implementations to retain the rows they
    48  // are passed as arguments.
    49  //
    50  // See: RowBuffer, DedupeRowReader, DedupeRowWriter
    51  type rowAllocator struct{ allocator }
    52  
    53  func (a *rowAllocator) capture(row Row) {
    54  	for i, v := range row {
    55  		switch v.Kind() {
    56  		case ByteArray, FixedLenByteArray:
    57  			row[i].ptr = unsafecast.AddressOfBytes(a.copyBytes(v.byteArray()))
    58  		}
    59  	}
    60  }