github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/query/util.go (about)

     1  package query
     2  
     3  import (
     4  	"slices"
     5  	"strings"
     6  
     7  	"github.com/colega/zeropool"
     8  	"github.com/parquet-go/parquet-go"
     9  )
    10  
    11  func GetColumnIndexByPath(root *parquet.Column, s string) (index, depth int) {
    12  	colSelector := strings.Split(s, ".")
    13  	n := root
    14  	for len(colSelector) > 0 {
    15  		n = n.Column(colSelector[0])
    16  		if n == nil {
    17  			return -1, -1
    18  		}
    19  
    20  		colSelector = colSelector[1:]
    21  		depth++
    22  	}
    23  
    24  	return n.Index(), depth
    25  }
    26  
    27  func HasColumn(root *parquet.Column, s string) bool {
    28  	index, _ := GetColumnIndexByPath(root, s)
    29  	return index >= 0
    30  }
    31  
    32  func RowGroupBoundaries(groups []parquet.RowGroup) []int64 {
    33  	b := make([]int64, len(groups))
    34  	var o int64
    35  	for i := range b {
    36  		o += groups[i].NumRows()
    37  		b[i] = o
    38  	}
    39  	return b
    40  }
    41  
    42  func SplitRows(rows, groups []int64) [][]int64 {
    43  	switch len(groups) {
    44  	case 0:
    45  		return nil
    46  	case 1:
    47  		return [][]int64{rows}
    48  	}
    49  	// Sanity check: max row must be less than
    50  	// the number of rows in the last group.
    51  	if rows[len(rows)-1] >= groups[len(groups)-1] {
    52  		panic(ErrSeekOutOfRange)
    53  	}
    54  	split := make([][]int64, len(groups))
    55  	var j, r int
    56  	maxRow := groups[j]
    57  	for i, rn := range rows {
    58  		if rn < maxRow {
    59  			continue
    60  		}
    61  		split[j], rows = rows[:i-r], rows[i-r:]
    62  		r = i
    63  		// Find matching group.
    64  		for x, v := range groups[j:] {
    65  			if rn >= v {
    66  				continue
    67  			}
    68  			j += x
    69  			break
    70  		}
    71  		maxRow = groups[j]
    72  	}
    73  	// Last bit.
    74  	split[j] = rows
    75  	// Subtract group offset from the row numbers,
    76  	// which makes them local to the group.
    77  	for i, g := range split[1:] {
    78  		offset := groups[i]
    79  		for n := range g {
    80  			g[n] -= offset
    81  		}
    82  	}
    83  	return split
    84  }
    85  
    86  var parquetValuesPool = zeropool.New(func() []parquet.Value { return nil })
    87  
    88  func CloneParquetValues(values []parquet.Value) []parquet.Value {
    89  	p := parquetValuesPool.Get()
    90  	p = slices.Grow(p, len(values))
    91  	p = p[:len(values)]
    92  	for i, v := range values {
    93  		p[i] = v.Clone()
    94  	}
    95  	return p
    96  }
    97  
    98  func ReleaseParquetValues(b [][]parquet.Value) {
    99  	for _, s := range b {
   100  		if cap(s) > 0 {
   101  			parquetValuesPool.Put(s)
   102  		}
   103  	}
   104  }
   105  
   106  var uint64valuesPool = zeropool.New(func() []uint64 { return nil })
   107  
   108  func CloneUint64ParquetValues(values []parquet.Value) []uint64 {
   109  	uint64s := uint64valuesPool.Get()
   110  	if l := len(values); cap(uint64s) < l {
   111  		uint64s = make([]uint64, 0, 2*l)
   112  	}
   113  	uint64s = uint64s[:len(values)]
   114  	for i, v := range values {
   115  		uint64s[i] = v.Uint64()
   116  	}
   117  	return uint64s
   118  }
   119  
   120  func ReleaseUint64Values(b [][]uint64) {
   121  	for _, s := range b {
   122  		if len(s) > 0 {
   123  			uint64valuesPool.Put(s)
   124  		}
   125  	}
   126  }
   127  
   128  var uint32valuesPool = zeropool.New(func() []uint32 { return nil })
   129  
   130  func CloneUint32ParquetValues(values []parquet.Value) []uint32 {
   131  	uint32s := uint32valuesPool.Get()
   132  	if l := len(values); cap(uint32s) < l {
   133  		uint32s = make([]uint32, 0, 2*l)
   134  	}
   135  	uint32s = uint32s[:len(values)]
   136  	for i, v := range values {
   137  		uint32s[i] = v.Uint32()
   138  	}
   139  	return uint32s
   140  }
   141  
   142  func ReleaseUint32Values(b [][]uint32) {
   143  	for _, s := range b {
   144  		if len(s) > 0 {
   145  			uint32valuesPool.Put(s)
   146  		}
   147  	}
   148  }