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) }