github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/sliceio/reader.go

// Copyright 2018 GRAIL, Inc. All rights reserved.
// Use of this source code is governed by the Apache 2.0
// license that can be found in the LICENSE file.

// Package sliceio provides utilities for managing I/O for Bigslice
// operations.
package sliceio

import (
	"context"
	"io"
	"reflect"
	"runtime/pprof"

	"github.com/grailbio/base/errors"
	"github.com/grailbio/bigslice/frame"
	"github.com/grailbio/bigslice/internal/defaultsize"
	"github.com/grailbio/bigslice/slicetype"
)

// DefaultChunksize is the default size used for I/O vectors within the
// sliceio package.
var defaultChunksize = defaultsize.Chunk

// EOF is the error returned by Reader.Read when no more data is
// available. EOF is intended as a sentinel error: it signals a
// graceful end of output. If output terminates unexpectedly, a
// different error should be returned.
var EOF = errors.New("EOF")

// A Reader represents a stateful stream of records. Each call to
// Read reads the next set of available records.
type Reader interface {
	// Read reads a vector of records from the underlying Slice. Each
	// passed-in column should be a value containing a slice of column
	// values. The number of columns should match the number of columns
	// in the slice; their types should match the corresponding column
	// types of the slice. Each column should have the same slice
	// length.
	//
	// Read returns the total number of records read, or an error. When
	// no more records are available, Read returns EOF. Read may return
	// EOF when n > 0. In this case, n records were read, but no more
	// are available.
	//
	// Read should never reuse any allocated memory in the frame;
	// its callers should not mutate the data returned.
	//
	// Read should not be called concurrently.
	Read(ctx context.Context, frame frame.Frame) (int, error)
}

// ReadCloser groups the Read and Close methods.
type ReadCloser interface {
	Reader
	io.Closer
}

// nopCloser decorates a reader with a no-op Close method. Use it to adapt a
// Reader to a ReadCloser when the Reader has no resources to release on Close.
type nopCloser struct {
	Reader
}

func (nopCloser) Close() error {
	return nil
}

func NopCloser(r Reader) ReadCloser {
	return nopCloser{r}
}

type multiReader struct {
	q   []ReadCloser
	err error
}

// MultiReader returns a ReadCloser that's the logical concatenation of the
// provided input readers. Once every underlying ReadCloser has returned EOF,
// Read will return EOF, too. Non-EOF errors are returned immediately.
func MultiReader(readers ...ReadCloser) ReadCloser {
	return &multiReader{q: readers}
}

func (m *multiReader) Read(ctx context.Context, out frame.Frame) (n int, err error) {
	if m.err != nil {
		return 0, m.err
	}
	for len(m.q) > 0 {
		n, err := m.q[0].Read(ctx, out)
		switch {
		case err == EOF:
			// There's not much for us to do if the Close fails, so we just
			// ignore it.
			_ = m.q[0].Close()
			m.q[0] = nil
			m.q = m.q[1:]
		case err != nil:
			m.err = err
			return n, err
		case n > 0:
			return n, err
		}
	}
	return 0, EOF
}

func (m *multiReader) Close() error {
	var err error
	for i, r := range m.q {
		if r == nil {
			continue
		}
		cerr := r.Close()
		if err == nil {
			err = cerr
		}
		m.q[i] = nil
	}
	return err
}

// FrameReader implements a Reader for a single Frame.
type frameReader struct {
	frame.Frame
}

// FrameReader returns a Reader that reads the provided
// Frame to completion.
func FrameReader(frame frame.Frame) Reader {
	return &frameReader{frame}
}

func (f *frameReader) Read(ctx context.Context, out frame.Frame) (int, error) {
	n := out.Len()
	max := f.Frame.Len()
	if max < n {
		n = max
	}
	frame.Copy(out, f.Frame)
	f.Frame = f.Frame.Slice(n, max)
	if f.Frame.Len() == 0 {
		return n, EOF
	}
	return n, nil
}

// ReadAll copies all elements from reader r into the provided column
// pointers. ReadAll is not tuned for performance and is intended for
// testing purposes.
func ReadAll(ctx context.Context, r Reader, columns ...interface{}) error {
	columnsv := make([]reflect.Value, len(columns))
	types := make([]reflect.Type, len(columns))
	for i := range columns {
		columnsv[i] = reflect.ValueOf(columns[i])
		if columnsv[i].Type().Kind() != reflect.Ptr {
			return errors.E(errors.Invalid, "attempted to read into non-pointer")
		}
		types[i] = reflect.TypeOf(columns[i]).Elem().Elem()
	}
	buf := frame.Make(slicetype.New(types...), defaultChunksize, defaultChunksize)
	for {
		n, err := r.Read(ctx, buf)
		if err != nil && err != EOF {
			return err
		}
		buf = buf.Slice(0, n)
		for i := range columnsv {
			columnsv[i].Elem().Set(reflect.AppendSlice(columnsv[i].Elem(), buf.Value(i)))
		}
		if err == EOF {
			break
		}
		buf = buf.Slice(0, buf.Cap())
	}
	return nil
}

// ReadFull reads the full length of the frame. ReadFull reads short
// frames only on EOF.
func ReadFull(ctx context.Context, r Reader, f frame.Frame) (n int, err error) {
	len := f.Len()
	for n < len {
		m, err := r.Read(ctx, f.Slice(n, len))
		n += m
		if err != nil {
			return n, err
		}
	}
	return n, nil
}

// An errReader is a reader that only returns errors.
type errReader struct{ Err error }

// ErrReader returns a reader that returns the provided error
// on every call to read. ErrReader panics if err is nil.
func ErrReader(err error) Reader {
	if err == nil {
		panic("nil error")
	}
	return &errReader{err}
}

func (e errReader) Read(ctx context.Context, f frame.Frame) (int, error) {
	return 0, e.Err
}

// ReaderWithCloseFunc is a ReadCloser that wraps an existing Reader and uses a
// provided function for its Close.
type ReaderWithCloseFunc struct {
	Reader
	CloseFunc func() error
}

// Close implements io.Closer.
func (r ReaderWithCloseFunc) Close() error {
	return r.CloseFunc()
}

// TODO(jcharumilind): Get rid of ClosingReader, as it makes it too tempting to
// not properly handle errors. We use it in cases where we expect to read from
// many readers (e.g. mergeReader). On failure, we should close all of them, but
// ClosingReader obscures this a bit and makes it so that the only way to close
// is by reading until non-nil error.

// ClosingReader closes the wrapped ReadCloser when Read returns any error.
type ClosingReader struct {
	r      ReadCloser
	closed bool
}

// NewClosingReader returns a new ClosingReader for r.
func NewClosingReader(r ReadCloser) *ClosingReader {
	return &ClosingReader{r: r}
}

// Read implements sliceio.Reader.
func (c *ClosingReader) Read(ctx context.Context, out frame.Frame) (int, error) {
	n, err := c.r.Read(ctx, out)
	if err != nil && !c.closed {
		c.r.Close()
		c.closed = true
	}
	return n, err
}

// EmptyReader returns an EOF.
type EmptyReader struct{}

func (EmptyReader) Read(ctx context.Context, f frame.Frame) (int, error) {
	return 0, EOF
}

// PprofReader executes Read in a labeled Context.
type PprofReader struct {
	Reader
	Label string
}

func (r *PprofReader) Read(ctx context.Context, frame frame.Frame) (n int, err error) {
	labels := pprof.Labels("sliceName", r.Label)
	pprof.Do(ctx, labels, func(ctx context.Context) {
		n, err = r.Reader.Read(ctx, frame)
	})
	return
}