influxdb/storage/reads/array_cursor.gen.go.tmpl

package reads

import (
	"errors"
	"fmt"
	"math"

    "github.com/influxdata/flux/execute"
    "github.com/influxdata/flux/values"
    "github.com/influxdata/influxdb/v2"
	"github.com/influxdata/influxdb/v2/tsdb/cursors"
)

const (
	// MaxPointsPerBlock is the maximum number of points in an encoded
	// block in a TSM file. It should match the value in the tsm1
	// package, but we don't want to import it.
	MaxPointsPerBlock = 1000
)

func newWindowMeanArrayCursor(cur cursors.Cursor, window execute.Window) (cursors.Cursor, error) {
	switch cur := cur.(type) {
{{range .}}
{{$Type := .Name}}
{{range .Aggs}}
{{if eq .Name "Mean"}}
	case cursors.{{$Type}}ArrayCursor:
		return new{{$Type}}WindowMeanArrayCursor(cur, window), nil
{{end}}
{{end}}{{/* for each supported agg fn */}}
{{end}}{{/* for each field type */}}
	default:
		return nil, &influxdb.Error{
			Code: influxdb.EInvalid,
			Msg: fmt.Sprintf("unsupported input type for mean aggregate: %s", arrayCursorType(cur)),
		}
	}
}
{{range .}}
{{$arrayType := print "*cursors." .Name "Array"}}
{{$type := print .name "ArrayFilterCursor"}}
{{$Type := print .Name "ArrayFilterCursor"}}

// ********************
// {{.Name}} Array Cursor

type {{$type}} struct {
	cursors.{{.Name}}ArrayCursor
	cond expression
	m    *singleValue
	res  {{$arrayType}}
	tmp  {{$arrayType}}
}

func new{{.Name}}FilterArrayCursor(cond expression) *{{$type}} {
	return &{{$type}}{
		cond: cond,
		m:    &singleValue{},
		res:  cursors.New{{.Name}}ArrayLen(MaxPointsPerBlock),
		tmp:  &cursors.{{.Name}}Array{},
	}
}

func (c *{{$type}}) reset(cur cursors.{{.Name}}ArrayCursor) {
	c.{{.Name}}ArrayCursor = cur
	c.tmp.Timestamps, c.tmp.Values = nil, nil
}

func (c *{{$type}}) Stats() cursors.CursorStats { return c.{{.Name}}ArrayCursor.Stats() }

func (c *{{$type}}) Next() {{$arrayType}} {
	pos := 0
	c.res.Timestamps = c.res.Timestamps[:cap(c.res.Timestamps)]
	c.res.Values = c.res.Values[:cap(c.res.Values)]

	var a {{$arrayType}}

	if c.tmp.Len() > 0 {
		a = c.tmp
	} else {
		a = c.{{.Name}}ArrayCursor.Next()
	}

LOOP:
	for len(a.Timestamps) > 0 {
		for i, v := range a.Values {
			c.m.v = v
			if c.cond.EvalBool(c.m) {
				c.res.Timestamps[pos] = a.Timestamps[i]
				c.res.Values[pos] = v
				pos++
				if pos >= MaxPointsPerBlock {
					c.tmp.Timestamps = a.Timestamps[i+1:]
					c.tmp.Values = a.Values[i+1:]
					break LOOP
				}
			}
		}


		// Clear buffered timestamps & values if we make it through a cursor.
		// The break above will skip this if a cursor is partially read.
		c.tmp.Timestamps = nil
		c.tmp.Values = nil

		a = c.{{.Name}}ArrayCursor.Next()
	}

	c.res.Timestamps = c.res.Timestamps[:pos]
	c.res.Values = c.res.Values[:pos]

	return c.res
}

type {{.name}}MultiShardArrayCursor struct {
	cursors.{{.Name}}ArrayCursor
	cursorContext
	filter *{{$type}}
}

func (c *{{.name}}MultiShardArrayCursor) reset(cur cursors.{{.Name}}ArrayCursor, itrs cursors.CursorIterators, cond expression) {
	if cond != nil {
		if c.filter == nil {
			c.filter = new{{.Name}}FilterArrayCursor(cond)
		}
		c.filter.reset(cur)
		cur = c.filter
	}

	c.{{.Name}}ArrayCursor = cur
	c.itrs = itrs
	c.err = nil
	c.count = 0
}


func (c *{{.name}}MultiShardArrayCursor) Err() error { return c.err }

func (c *{{.name}}MultiShardArrayCursor) Stats() cursors.CursorStats {
	return c.{{.Name}}ArrayCursor.Stats()
}

func (c *{{.name}}MultiShardArrayCursor) Next() {{$arrayType}} {
	for {
		a := c.{{.Name}}ArrayCursor.Next()
		if a.Len() == 0 {
			if c.nextArrayCursor() {
				continue
			}
		}
		c.count += int64(a.Len())
		if c.count > c.limit {
			diff := c.count - c.limit
			c.count -= diff
			rem := int64(a.Len()) - diff
			a.Timestamps = a.Timestamps[:rem]
			a.Values = a.Values[:rem]
		}
		return a
	}
}

func (c *{{.name}}MultiShardArrayCursor) nextArrayCursor() bool {
	if len(c.itrs) == 0 {
		return false
	}

	c.{{.Name}}ArrayCursor.Close()

	var itr cursors.CursorIterator
	var cur cursors.Cursor
	for cur == nil && len(c.itrs) > 0 {
		itr, c.itrs = c.itrs[0], c.itrs[1:]
		cur, _ = itr.Next(c.ctx, c.req)
	}

	var ok bool
	if cur != nil {
		var next cursors.{{.Name}}ArrayCursor
		next, ok = cur.(cursors.{{.Name}}ArrayCursor)
		if !ok {
			cur.Close()
			next = {{.Name}}EmptyArrayCursor
			c.itrs = nil
			c.err = errors.New("expected {{.name}} cursor")
		} else {
			if c.filter != nil {
				c.filter.reset(next)
				next = c.filter
			}
		}
		c.{{.Name}}ArrayCursor = next
	} else {
		c.{{.Name}}ArrayCursor = {{.Name}}EmptyArrayCursor
	}

	return ok
}

{{if .Agg}}
{{$type := print .name "ArraySumCursor"}}
{{$Type := print .Name "ArraySumCursor"}}


type {{$type}} struct {
	cursors.{{.Name}}ArrayCursor
	ts [1]int64
	vs [1]{{.Type}}
	res {{$arrayType}}
}

func new{{$Type}}(cur cursors.{{.Name}}ArrayCursor) *{{$type}} {
	return &{{$type}}{
		{{.Name}}ArrayCursor: cur,
		res:                  &cursors.{{.Name}}Array{},
	}
}

func (c {{$type}}) Stats() cursors.CursorStats { return c.{{.Name}}ArrayCursor.Stats() }

func (c {{$type}}) Next() {{$arrayType}} {
	a := c.{{.Name}}ArrayCursor.Next()
	if len(a.Timestamps) == 0 {
		return a
	}

	ts := a.Timestamps[0]
	var acc {{.Type}}

	for {
		for _, v := range a.Values {
			acc += v
		}
		a = c.{{.Name}}ArrayCursor.Next()
		if len(a.Timestamps) == 0 {
			c.ts[0] = ts
			c.vs[0] = acc
			c.res.Timestamps = c.ts[:]
			c.res.Values = c.vs[:]
			return c.res
		}
	}
}

{{end}}

type integer{{.Name}}CountArrayCursor struct {
	cursors.{{.Name}}ArrayCursor
}

func (c *integer{{.Name}}CountArrayCursor) Stats() cursors.CursorStats {
	return c.{{.Name}}ArrayCursor.Stats()
}

func (c *integer{{.Name}}CountArrayCursor) Next() *cursors.IntegerArray {
	a := c.{{.Name}}ArrayCursor.Next()
	if len(a.Timestamps) == 0 {
		return &cursors.IntegerArray{}
	}

	ts := a.Timestamps[0]
	var acc int64
	for {
		acc += int64(len(a.Timestamps))
		a = c.{{.Name}}ArrayCursor.Next()
		if len(a.Timestamps) == 0 {
			res := cursors.NewIntegerArrayLen(1)
			res.Timestamps[0] = ts
			res.Values[0] = acc
			return res
		}
	}
}

{{/* create an aggregate cursor for each aggregate function supported by the type */}}
{{$Name := .Name}}
{{$name := .name}}
{{range .Aggs}}
{{$aggName := .Name}}

type {{$name}}Window{{$aggName}}ArrayCursor struct {
	cursors.{{$Name}}ArrayCursor
	res   *cursors.{{.OutputTypeName}}Array
	tmp   {{$arrayType}}
	window execute.Window
}

func new{{$Name}}Window{{$aggName}}ArrayCursor(cur cursors.{{$Name}}ArrayCursor, window execute.Window) *{{$name}}Window{{$aggName}}ArrayCursor {
	resLen := MaxPointsPerBlock
	if window.Every.IsZero() {
		resLen = 1
	}
	return &{{$name}}Window{{$aggName}}ArrayCursor{
		{{$Name}}ArrayCursor: cur,
		res: cursors.New{{.OutputTypeName}}ArrayLen(resLen),
		tmp: &cursors.{{$Name}}Array{},
		window: window,
	}
}

func (c *{{$name}}Window{{$aggName}}ArrayCursor) Stats() cursors.CursorStats {
	return c.{{$Name}}ArrayCursor.Stats()
}

func (c *{{$name}}Window{{$aggName}}ArrayCursor) Next() *cursors.{{.OutputTypeName}}Array {
	pos := 0
	c.res.Timestamps = c.res.Timestamps[:cap(c.res.Timestamps)]
	c.res.Values = c.res.Values[:cap(c.res.Values)]

	var a *cursors.{{$Name}}Array
	if c.tmp.Len() > 0 {
		a = c.tmp
	} else {
		a = c.{{$Name}}ArrayCursor.Next()
	}

	if a.Len() == 0 {
		return &cursors.{{.OutputTypeName}}Array{}
	}

	rowIdx := 0
	{{.AccDecls}}

	var windowEnd int64
	if !c.window.Every.IsZero() {
		windowEnd = int64(c.window.GetEarliestBounds(values.Time(a.Timestamps[rowIdx])).Stop)
	} else {
		windowEnd = math.MaxInt64
	}
	windowHasPoints := false

	// enumerate windows
WINDOWS:
	for {
		for ; rowIdx < a.Len(); rowIdx++ {
			ts := a.Timestamps[rowIdx]
			if !c.window.Every.IsZero() && ts >= windowEnd {
				// new window detected, close the current window
				// do not generate a point for empty windows
				if windowHasPoints {
					{{.AccEmit}}
					pos++
					if pos >= MaxPointsPerBlock {
						// the output array is full,
						// save the remaining points in the input array in tmp.
						// they will be processed in the next call to Next()
						c.tmp.Timestamps = a.Timestamps[rowIdx:]
						c.tmp.Values = a.Values[rowIdx:]
						break WINDOWS
					}
				}

				// start the new window
				{{.AccReset}}
				windowEnd = int64(c.window.GetEarliestBounds(values.Time(ts)).Stop)
				windowHasPoints = false

				continue WINDOWS
			} else {
				{{.Accumulate}}
				windowHasPoints = true
			}
		}

		// Clear buffered timestamps & values if we make it through a cursor.
		// The break above will skip this if a cursor is partially read.
		c.tmp.Timestamps = nil
		c.tmp.Values = nil

		// get the next chunk
		a = c.{{$Name}}ArrayCursor.Next()
		if a.Len() == 0 {
			// write the final point
			// do not generate a point for empty windows
			if windowHasPoints {
				{{.AccEmit}}
				pos++
			}
			break WINDOWS
		}
		rowIdx = 0
	}

	c.res.Timestamps = c.res.Timestamps[:pos]
	c.res.Values = c.res.Values[:pos]

	return c.res
}

{{end}}{{/* range .Aggs */}}

type {{.name}}EmptyArrayCursor struct {
	res cursors.{{.Name}}Array
}

var {{.Name}}EmptyArrayCursor cursors.{{.Name}}ArrayCursor = &{{.name}}EmptyArrayCursor{}

func (c *{{.name}}EmptyArrayCursor) Err() error { return nil }
func (c *{{.name}}EmptyArrayCursor) Close() {}
func (c *{{.name}}EmptyArrayCursor) Stats() cursors.CursorStats { return cursors.CursorStats{} }
func (c *{{.name}}EmptyArrayCursor) Next() {{$arrayType}} { return &c.res }

{{end}}{{/* range . */}}

func arrayCursorType(cur cursors.Cursor) string {
	switch cur.(type) {
	{{range .}}
	case cursors.{{.Name}}ArrayCursor:
		return "{{.name}}"
	{{end}}{{/* range . */}}
	default:
		return "unknown"
	}
}