github.com/grafana/pyroscope@v1.18.0/pkg/segmentwriter/memdb/profiles.go (about)

     1  package memdb
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  
     7  	"github.com/parquet-go/parquet-go"
     8  
     9  	v1 "github.com/grafana/pyroscope/pkg/phlaredb/schemas/v1"
    10  	"github.com/grafana/pyroscope/pkg/util/build"
    11  )
    12  
    13  const segmentsParquetWriteBufferSize = 32 << 10
    14  
    15  func WriteProfiles(metrics *HeadMetrics, profiles []v1.InMemoryProfile) ([]byte, error) {
    16  	buf := &bytes.Buffer{}
    17  	w := parquet.NewGenericWriter[*v1.Profile](
    18  		buf,
    19  		parquet.PageBufferSize(segmentsParquetWriteBufferSize),
    20  		parquet.CreatedBy("github.com/grafana/pyroscope/", build.Version, build.Revision),
    21  		v1.ProfilesSchema,
    22  	)
    23  
    24  	if _, err := parquet.CopyRows(w, v1.NewInMemoryProfilesRowReader(profiles)); err != nil {
    25  		return nil, fmt.Errorf("failed to write profile rows to parquet table: %w", err)
    26  	}
    27  	if err := w.Close(); err != nil {
    28  		return nil, fmt.Errorf("failed to close parquet table: %w", err)
    29  	}
    30  
    31  	metrics.writtenProfileSegments.WithLabelValues("success").Inc()
    32  	res := buf.Bytes()
    33  	metrics.writtenProfileSegmentsBytes.Observe(float64(len(res)))
    34  	metrics.rowsWritten.WithLabelValues("profiles").Add(float64(len(profiles)))
    35  	return res, nil
    36  }