chunkiterator.gen.go.tmpl

// Copyright 2019 Nick Poorman
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package iterator

import (
	"sync/atomic"

	"github.com/apache/arrow/go/arrow"
	"github.com/apache/arrow/go/arrow/array"
	"github.com/apache/arrow/go/arrow/decimal128"
	"github.com/apache/arrow/go/arrow/float16"
	"github.com/gomem/gomem/internal/debug"
)

{{range .In}}
{{if not (contains .ExcludeGenerate "chunkiterator") }}
// {{.Name}}ChunkIterator is an iterator for reading an Arrow Column value by value.
type {{.Name}}ChunkIterator struct {
	refCount int64
	col      *array.Column

	// Things Chunked maintains. We're going to maintain it ourselves.
	chunks []*array.{{.Name}} // cache the chunks on this iterator
	length int64              // this isn't set right on Chunked so we won't rely on it there. Instead we keep the correct value here.
	nulls  int64
	dtype  arrow.DataType

	// Things we need to maintain for the iterator
	currentIndex int              // current chunk
	currentChunk *array.{{.Name}} // current chunk
}

// New{{.Name}}ChunkIterator creates a new {{.Name}}ChunkIterator for reading an Arrow Column.
func New{{.Name}}ChunkIterator(col *array.Column) *{{.Name}}ChunkIterator {
	col.Retain()

	// Chunked is not using the correct type to keep track of length so we have to recalculate it.
	columnChunks := col.Data().Chunks()
	chunks := make([]*array.{{.Name}}, len(columnChunks))
	var length int64
	var nulls int64

	for i, chunk := range columnChunks {
		// Keep our own refs to chunks
		chunks[i] = chunk.(*array.{{.Name}})
		// Retain the chunk
		chunks[i].Retain()

		// Keep our own counters instead of Chunked's
		length += int64(chunk.Len())
		nulls += int64(chunk.NullN())
	}

	return &{{.Name}}ChunkIterator{
		refCount: 1,
		col:      col,

		chunks: chunks,
		length: length,
		nulls:  nulls,
		dtype:  col.DataType(),

		currentIndex: 0,
		currentChunk: nil,
	}
}

// Chunk will return the current chunk that the iterator is on.
func (cr *{{.Name}}ChunkIterator) Chunk() *array.{{.Name}} { return cr.currentChunk }

// ChunkValues returns the underlying []{{.Type}} chunk values.
// Keep in mind the []{{.Type}} type might not be able
// to account for nil values. You must check for those explicitly via the chunk.
func (cr *{{.Name}}ChunkIterator) ChunkValues() []{{.Type}} {
	{{if eq .Name "Null"}}
	return make([]interface{}, cr.Chunk().Len())
	{{else}}
	return cr.Chunk().{{- if .ValuesMethod -}}{{.ValuesMethod}}{{- else -}}{{.Name}}Values{{- end -}}()
	{{end}}
}

// Next moves the iterator to the next chunk. This will return false
// when there are no more chunks.
func (cr *{{.Name}}ChunkIterator) Next() bool {
	if cr.currentIndex >= len(cr.chunks) {
		return false
	}

	if cr.currentChunk != nil {
		cr.currentChunk.Release()
	}

	cr.currentChunk = cr.chunks[cr.currentIndex]
	cr.currentChunk.Retain()
	cr.currentIndex++

	return true
}

// Retain keeps a reference to the {{.Name}}ChunkIterator
func (cr *{{.Name}}ChunkIterator) Retain() {
	atomic.AddInt64(&cr.refCount, 1)
}

// Release removes a reference to the {{.Name}}ChunkIterator
func (cr *{{.Name}}ChunkIterator) Release() {
	debug.Assert(atomic.LoadInt64(&cr.refCount) > 0, "too many releases")
	ref := atomic.AddInt64(&cr.refCount, -1)
	if ref == 0 {
		cr.col.Release()
		for i := range cr.chunks {
			cr.chunks[i].Release()
		}
		if cr.currentChunk != nil {
			cr.currentChunk.Release()
			cr.currentChunk = nil
		}
		cr.col = nil
		cr.chunks = nil
		cr.dtype = nil
	}
}


{{end}}
{{end}}