github.com/vc42/parquet-go@v0.0.0-20240320194221-1a9adb5f23f5/writer_go18_test.go (about)

     1  //go:build go1.18
     2  
     3  package parquet_test
     4  
     5  import (
     6  	"io"
     7  	"math/rand"
     8  	"reflect"
     9  	"testing"
    10  
    11  	"github.com/vc42/parquet-go"
    12  )
    13  
    14  func BenchmarkGenericWriter(b *testing.B) {
    15  	benchmarkGenericWriter[benchmarkRowType](b)
    16  	benchmarkGenericWriter[booleanColumn](b)
    17  	benchmarkGenericWriter[int32Column](b)
    18  	benchmarkGenericWriter[int64Column](b)
    19  	benchmarkGenericWriter[floatColumn](b)
    20  	benchmarkGenericWriter[doubleColumn](b)
    21  	benchmarkGenericWriter[byteArrayColumn](b)
    22  	benchmarkGenericWriter[fixedLenByteArrayColumn](b)
    23  	benchmarkGenericWriter[stringColumn](b)
    24  	benchmarkGenericWriter[indexedStringColumn](b)
    25  	benchmarkGenericWriter[uuidColumn](b)
    26  	benchmarkGenericWriter[mapColumn](b)
    27  	benchmarkGenericWriter[decimalColumn](b)
    28  	benchmarkGenericWriter[contact](b)
    29  	benchmarkGenericWriter[paddedBooleanColumn](b)
    30  	benchmarkGenericWriter[optionalInt32Column](b)
    31  	benchmarkGenericWriter[repeatedInt32Column](b)
    32  }
    33  
    34  func benchmarkGenericWriter[Row generator[Row]](b *testing.B) {
    35  	var model Row
    36  	b.Run(reflect.TypeOf(model).Name(), func(b *testing.B) {
    37  		prng := rand.New(rand.NewSource(0))
    38  		rows := make([]Row, benchmarkNumRows)
    39  		for i := range rows {
    40  			rows[i] = rows[i].generate(prng)
    41  		}
    42  
    43  		b.Run("go1.17", func(b *testing.B) {
    44  			writer := parquet.NewWriter(io.Discard, parquet.SchemaOf(rows[0]))
    45  			i := 0
    46  			benchmarkRowsPerSecond(b, func() int {
    47  				for j := 0; j < benchmarkRowsPerStep; j++ {
    48  					if err := writer.Write(&rows[i]); err != nil {
    49  						b.Fatal(err)
    50  					}
    51  				}
    52  
    53  				i += benchmarkRowsPerStep
    54  				i %= benchmarkNumRows
    55  
    56  				if i == 0 {
    57  					writer.Close()
    58  					writer.Reset(io.Discard)
    59  				}
    60  				return benchmarkRowsPerStep
    61  			})
    62  		})
    63  
    64  		b.Run("go1.18", func(b *testing.B) {
    65  			writer := parquet.NewGenericWriter[Row](io.Discard)
    66  			i := 0
    67  			benchmarkRowsPerSecond(b, func() int {
    68  				n, err := writer.Write(rows[i : i+benchmarkRowsPerStep])
    69  				if err != nil {
    70  					b.Fatal(err)
    71  				}
    72  
    73  				i += benchmarkRowsPerStep
    74  				i %= benchmarkNumRows
    75  
    76  				if i == 0 {
    77  					writer.Close()
    78  					writer.Reset(io.Discard)
    79  				}
    80  				return n
    81  			})
    82  		})
    83  	})
    84  }