github.com/parquet-go/parquet-go@v0.21.1-0.20240501160520-b3c3a0c3ed6f/allocator.go (about) 1 package parquet 2 3 import "github.com/parquet-go/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 }