351 lines
9.4 KiB
Go
351 lines
9.4 KiB
Go
package storageflux
|
|
|
|
//go:generate env GO111MODULE=on go run github.com/benbjohnson/tmpl -data=@types.tmpldata table.gen.go.tmpl
|
|
|
|
import (
|
|
"errors"
|
|
"sync/atomic"
|
|
|
|
"github.com/influxdata/flux"
|
|
"github.com/influxdata/flux/array"
|
|
"github.com/influxdata/flux/arrow"
|
|
"github.com/influxdata/flux/execute"
|
|
"github.com/influxdata/flux/memory"
|
|
"github.com/influxdata/influxdb/v2/models"
|
|
)
|
|
|
|
type table struct {
|
|
bounds execute.Bounds
|
|
key flux.GroupKey
|
|
cols []flux.ColMeta
|
|
|
|
// cache of the tags on the current series.
|
|
// len(tags) == len(colMeta)
|
|
tags [][]byte
|
|
defs [][]byte
|
|
|
|
done chan struct{}
|
|
|
|
colBufs *colReader
|
|
empty bool
|
|
|
|
err error
|
|
|
|
cancelled, used int32
|
|
cache *tagsCache
|
|
alloc memory.Allocator
|
|
}
|
|
|
|
func newTable(
|
|
done chan struct{},
|
|
bounds execute.Bounds,
|
|
key flux.GroupKey,
|
|
cols []flux.ColMeta,
|
|
defs [][]byte,
|
|
cache *tagsCache,
|
|
alloc memory.Allocator,
|
|
) table {
|
|
return table{
|
|
done: done,
|
|
bounds: bounds,
|
|
key: key,
|
|
tags: make([][]byte, len(cols)),
|
|
defs: defs,
|
|
cols: cols,
|
|
cache: cache,
|
|
alloc: alloc,
|
|
}
|
|
}
|
|
|
|
func (t *table) Key() flux.GroupKey { return t.key }
|
|
func (t *table) Cols() []flux.ColMeta { return t.cols }
|
|
func (t *table) Err() error { return t.err }
|
|
func (t *table) Empty() bool { return t.empty }
|
|
|
|
func (t *table) Cancel() {
|
|
atomic.StoreInt32(&t.cancelled, 1)
|
|
}
|
|
|
|
func (t *table) isCancelled() bool {
|
|
return atomic.LoadInt32(&t.cancelled) != 0
|
|
}
|
|
|
|
func (t *table) init(advance func() bool) {
|
|
t.empty = !advance() && t.err == nil
|
|
}
|
|
|
|
func (t *table) do(f func(flux.ColReader) error, advance func() bool) error {
|
|
// Mark this table as having been used. If this doesn't
|
|
// succeed, then this has already been invoked somewhere else.
|
|
if !atomic.CompareAndSwapInt32(&t.used, 0, 1) {
|
|
return errors.New("table already used")
|
|
}
|
|
defer t.closeDone()
|
|
|
|
// If an error occurred during initialization, that is
|
|
// returned here.
|
|
if t.err != nil {
|
|
return t.err
|
|
}
|
|
|
|
if !t.Empty() {
|
|
t.err = f(t.colBufs)
|
|
t.colBufs.Release()
|
|
|
|
for !t.isCancelled() && t.err == nil && advance() {
|
|
t.err = f(t.colBufs)
|
|
t.colBufs.Release()
|
|
}
|
|
t.colBufs = nil
|
|
}
|
|
|
|
return t.err
|
|
}
|
|
|
|
func (t *table) Done() {
|
|
// Mark the table as having been used. If this has already
|
|
// been done, then nothing needs to be done.
|
|
if atomic.CompareAndSwapInt32(&t.used, 0, 1) {
|
|
defer t.closeDone()
|
|
}
|
|
|
|
if t.colBufs != nil {
|
|
t.colBufs.Release()
|
|
t.colBufs = nil
|
|
}
|
|
}
|
|
|
|
// allocateBuffer will allocate a suitable buffer for the
|
|
// table implementations to use. If the existing buffer
|
|
// is not used anymore, then it may be reused.
|
|
//
|
|
// The allocated buffer can be accessed at colBufs or
|
|
// through the returned colReader.
|
|
func (t *table) allocateBuffer(l int) *colReader {
|
|
if t.colBufs == nil || atomic.LoadInt64(&t.colBufs.refCount) > 0 {
|
|
// The current buffer is still being used so we should
|
|
// generate a new one.
|
|
t.colBufs = &colReader{
|
|
key: t.key,
|
|
colMeta: t.cols,
|
|
cols: make([]array.Array, len(t.cols)),
|
|
}
|
|
}
|
|
t.colBufs.refCount = 1
|
|
t.colBufs.l = l
|
|
return t.colBufs
|
|
}
|
|
|
|
type colReader struct {
|
|
refCount int64
|
|
|
|
key flux.GroupKey
|
|
colMeta []flux.ColMeta
|
|
cols []array.Array
|
|
l int
|
|
}
|
|
|
|
func (cr *colReader) Retain() {
|
|
atomic.AddInt64(&cr.refCount, 1)
|
|
}
|
|
func (cr *colReader) Release() {
|
|
if atomic.AddInt64(&cr.refCount, -1) == 0 {
|
|
for _, col := range cr.cols {
|
|
col.Release()
|
|
}
|
|
}
|
|
}
|
|
|
|
func (cr *colReader) Key() flux.GroupKey { return cr.key }
|
|
func (cr *colReader) Cols() []flux.ColMeta { return cr.colMeta }
|
|
func (cr *colReader) Len() int { return cr.l }
|
|
|
|
func (cr *colReader) Bools(j int) *array.Boolean {
|
|
execute.CheckColType(cr.colMeta[j], flux.TBool)
|
|
return cr.cols[j].(*array.Boolean)
|
|
}
|
|
|
|
func (cr *colReader) Ints(j int) *array.Int {
|
|
execute.CheckColType(cr.colMeta[j], flux.TInt)
|
|
return cr.cols[j].(*array.Int)
|
|
}
|
|
|
|
func (cr *colReader) UInts(j int) *array.Uint {
|
|
execute.CheckColType(cr.colMeta[j], flux.TUInt)
|
|
return cr.cols[j].(*array.Uint)
|
|
}
|
|
|
|
func (cr *colReader) Floats(j int) *array.Float {
|
|
execute.CheckColType(cr.colMeta[j], flux.TFloat)
|
|
return cr.cols[j].(*array.Float)
|
|
}
|
|
|
|
func (cr *colReader) Strings(j int) *array.String {
|
|
execute.CheckColType(cr.colMeta[j], flux.TString)
|
|
return cr.cols[j].(*array.String)
|
|
}
|
|
|
|
func (cr *colReader) Times(j int) *array.Int {
|
|
execute.CheckColType(cr.colMeta[j], flux.TTime)
|
|
return cr.cols[j].(*array.Int)
|
|
}
|
|
|
|
// readTags populates b.tags with the provided tags
|
|
func (t *table) readTags(tags models.Tags) {
|
|
for j := range t.tags {
|
|
t.tags[j] = t.defs[j]
|
|
}
|
|
|
|
if len(tags) == 0 {
|
|
return
|
|
}
|
|
|
|
for _, tag := range tags {
|
|
j := execute.ColIdx(string(tag.Key), t.cols)
|
|
// In the case of group aggregate, tags that are not referenced in group() are not included in the result, but
|
|
// readTags () still get a complete tag list. Here is just to skip the tags that should not be present in the result.
|
|
if j < 0 {
|
|
continue
|
|
}
|
|
t.tags[j] = tag.Value
|
|
}
|
|
}
|
|
|
|
// appendTheseTags fills the colBufs for the tag columns with the given tag values.
|
|
func (t *table) appendTheseTags(cr *colReader, tags [][]byte) {
|
|
for j := range t.cols {
|
|
v := tags[j]
|
|
if v != nil {
|
|
cr.cols[j] = t.cache.GetTag(string(v), cr.l, t.alloc)
|
|
}
|
|
}
|
|
}
|
|
|
|
// appendTags fills the colBufs for the tag columns with the tag values from the table structure.
|
|
func (t *table) appendTags(cr *colReader) {
|
|
t.appendTheseTags(cr, t.tags)
|
|
}
|
|
|
|
// appendBounds fills the colBufs for the time bounds
|
|
func (t *table) appendBounds(cr *colReader) {
|
|
start, stop := t.cache.GetBounds(t.bounds, cr.l, t.alloc)
|
|
cr.cols[startColIdx], cr.cols[stopColIdx] = start, stop
|
|
}
|
|
|
|
func (t *table) closeDone() {
|
|
if t.done != nil {
|
|
close(t.done)
|
|
t.done = nil
|
|
}
|
|
}
|
|
|
|
func (t *floatTable) toArrowBuffer(vs []float64) *array.Float {
|
|
return arrow.NewFloat(vs, t.alloc)
|
|
}
|
|
func (t *floatGroupTable) toArrowBuffer(vs []float64) *array.Float {
|
|
return arrow.NewFloat(vs, t.alloc)
|
|
}
|
|
func (t *floatWindowSelectorTable) toArrowBuffer(vs []float64) *array.Float {
|
|
return arrow.NewFloat(vs, t.alloc)
|
|
}
|
|
func (t *floatWindowTable) mergeValues(intervals []int64) *array.Float {
|
|
b := arrow.NewFloatBuilder(t.alloc)
|
|
b.Resize(len(intervals))
|
|
t.appendValues(intervals, b.Append, b.AppendNull)
|
|
return b.NewFloatArray()
|
|
}
|
|
func (t *floatEmptyWindowSelectorTable) arrowBuilder() *array.FloatBuilder {
|
|
return arrow.NewFloatBuilder(t.alloc)
|
|
}
|
|
func (t *floatEmptyWindowSelectorTable) append(builder *array.FloatBuilder, v float64) {
|
|
builder.Append(v)
|
|
}
|
|
func (t *integerTable) toArrowBuffer(vs []int64) *array.Int {
|
|
return arrow.NewInt(vs, t.alloc)
|
|
}
|
|
func (t *integerWindowSelectorTable) toArrowBuffer(vs []int64) *array.Int {
|
|
return arrow.NewInt(vs, t.alloc)
|
|
}
|
|
func (t *integerGroupTable) toArrowBuffer(vs []int64) *array.Int {
|
|
return arrow.NewInt(vs, t.alloc)
|
|
}
|
|
func (t *integerWindowTable) mergeValues(intervals []int64) *array.Int {
|
|
b := arrow.NewIntBuilder(t.alloc)
|
|
b.Resize(len(intervals))
|
|
appendNull := b.AppendNull
|
|
if t.fillValue != nil {
|
|
appendNull = func() { b.Append(*t.fillValue) }
|
|
}
|
|
t.appendValues(intervals, b.Append, appendNull)
|
|
return b.NewIntArray()
|
|
}
|
|
func (t *integerEmptyWindowSelectorTable) arrowBuilder() *array.IntBuilder {
|
|
return arrow.NewIntBuilder(t.alloc)
|
|
}
|
|
func (t *integerEmptyWindowSelectorTable) append(builder *array.IntBuilder, v int64) {
|
|
builder.Append(v)
|
|
}
|
|
func (t *unsignedTable) toArrowBuffer(vs []uint64) *array.Uint {
|
|
return arrow.NewUint(vs, t.alloc)
|
|
}
|
|
func (t *unsignedGroupTable) toArrowBuffer(vs []uint64) *array.Uint {
|
|
return arrow.NewUint(vs, t.alloc)
|
|
}
|
|
func (t *unsignedWindowSelectorTable) toArrowBuffer(vs []uint64) *array.Uint {
|
|
return arrow.NewUint(vs, t.alloc)
|
|
}
|
|
func (t *unsignedWindowTable) mergeValues(intervals []int64) *array.Uint {
|
|
b := arrow.NewUintBuilder(t.alloc)
|
|
b.Resize(len(intervals))
|
|
t.appendValues(intervals, b.Append, b.AppendNull)
|
|
return b.NewUintArray()
|
|
}
|
|
func (t *unsignedEmptyWindowSelectorTable) arrowBuilder() *array.UintBuilder {
|
|
return arrow.NewUintBuilder(t.alloc)
|
|
}
|
|
func (t *unsignedEmptyWindowSelectorTable) append(builder *array.UintBuilder, v uint64) {
|
|
builder.Append(v)
|
|
}
|
|
func (t *stringTable) toArrowBuffer(vs []string) *array.String {
|
|
return arrow.NewString(vs, t.alloc)
|
|
}
|
|
func (t *stringGroupTable) toArrowBuffer(vs []string) *array.String {
|
|
return arrow.NewString(vs, t.alloc)
|
|
}
|
|
func (t *stringWindowSelectorTable) toArrowBuffer(vs []string) *array.String {
|
|
return arrow.NewString(vs, t.alloc)
|
|
}
|
|
func (t *stringWindowTable) mergeValues(intervals []int64) *array.String {
|
|
b := arrow.NewStringBuilder(t.alloc)
|
|
b.Resize(len(intervals))
|
|
t.appendValues(intervals, b.Append, b.AppendNull)
|
|
return b.NewStringArray()
|
|
}
|
|
func (t *stringEmptyWindowSelectorTable) arrowBuilder() *array.StringBuilder {
|
|
return arrow.NewStringBuilder(t.alloc)
|
|
}
|
|
func (t *stringEmptyWindowSelectorTable) append(builder *array.StringBuilder, v string) {
|
|
builder.Append(v)
|
|
}
|
|
func (t *booleanTable) toArrowBuffer(vs []bool) *array.Boolean {
|
|
return arrow.NewBool(vs, t.alloc)
|
|
}
|
|
func (t *booleanGroupTable) toArrowBuffer(vs []bool) *array.Boolean {
|
|
return arrow.NewBool(vs, t.alloc)
|
|
}
|
|
func (t *booleanWindowSelectorTable) toArrowBuffer(vs []bool) *array.Boolean {
|
|
return arrow.NewBool(vs, t.alloc)
|
|
}
|
|
func (t *booleanWindowTable) mergeValues(intervals []int64) *array.Boolean {
|
|
b := arrow.NewBoolBuilder(t.alloc)
|
|
b.Resize(len(intervals))
|
|
t.appendValues(intervals, b.Append, b.AppendNull)
|
|
return b.NewBooleanArray()
|
|
}
|
|
func (t *booleanEmptyWindowSelectorTable) arrowBuilder() *array.BooleanBuilder {
|
|
return arrow.NewBoolBuilder(t.alloc)
|
|
}
|
|
func (t *booleanEmptyWindowSelectorTable) append(builder *array.BooleanBuilder, v bool) {
|
|
builder.Append(v)
|
|
}
|