record_reader.go

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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 bigquery

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"log"
	"sync/atomic"
	"time"

	"cloud.google.com/go/bigquery"
	"cloud.google.com/go/civil"

	"github.com/apache/arrow-adbc/go/adbc"
	"github.com/apache/arrow-go/v18/arrow"
	"github.com/apache/arrow-go/v18/arrow/array"
	"github.com/apache/arrow-go/v18/arrow/ipc"
	"github.com/apache/arrow-go/v18/arrow/memory"
	"golang.org/x/sync/errgroup"
	"google.golang.org/api/iterator"
)

type reader struct {
	refCount   int64
	schema     *arrow.Schema
	chs        []chan arrow.RecordBatch
	curChIndex int
	rec        arrow.RecordBatch
	err        error

	cancelFn context.CancelFunc
}

func checkContext(ctx context.Context, maybeErr error) error {
	if maybeErr != nil {
		return maybeErr
	} else if errors.Is(ctx.Err(), context.Canceled) {
		return adbc.Error{Msg: ctx.Err().Error(), Code: adbc.StatusCancelled}
	} else if errors.Is(ctx.Err(), context.DeadlineExceeded) {
		return adbc.Error{Msg: ctx.Err().Error(), Code: adbc.StatusTimeout}
	}
	return ctx.Err()
}

func runUpdate(ctx context.Context, query *bigquery.Query) (int64, error) {
	if _, err := query.Run(ctx); err != nil {
		return -1, err
	}
	return 0, nil
}

func runQuery(ctx context.Context, query *bigquery.Query, linkFailedJob bool, alloc memory.Allocator, outStats *jobStats) (bigquery.ArrowIterator, int64, error) {
	job, err := query.Run(ctx)
	if err != nil {
		return nil, -1, err
	}

	// The project id, location, and job id are all URL-safe:
	// Project id and job id can only contain URL safe characters: https://cloud.google.com/bigquery/docs/reference/rest/v2/JobReference
	// Locations are also URL-safe, listed here: https://cloud.google.com/bigquery/docs/locations
	jobLink := fmt.Sprintf("https://console.cloud.google.com/bigquery?project=%s&j=bq:%s:%s&page=queryresults", job.ProjectID(), job.Location(), job.ID())
	iter, err := job.Read(ctx)
	if err != nil {
		if linkFailedJob {
			return nil, -1, fmt.Errorf("%w (Query: %s)", err, jobLink)
		}
		return nil, -1, err
	}

	var arrowIterator bigquery.ArrowIterator
	useLegacyAPI := ctx.Value(ContextKeyUseStorageApiDisabledClient).(bool)
	if iter.TotalRows > 0 {
		if !useLegacyAPI {
			if !iter.IsAccelerated() {
				return nil, -1, fmt.Errorf("Storage API is not available for query: %s", jobLink)
			}
			// Storage API is available, use it
			if arrowIterator, err = iter.ArrowIterator(); err != nil {
				if linkFailedJob {
					return nil, -1, fmt.Errorf("%w (Query: %s)", err, jobLink)
				}
				return nil, -1, err
			}
		} else {
			arrowIterator = newRowBasedArrowIterator(iter, alloc)
		}
	} else {
		arrowIterator = emptyArrowIterator{iter.Schema}
	}
	totalRows := int64(iter.TotalRows)

	// Store job ID in query for adding to metadata
	query.JobID = job.ID()
	if outStats != nil {
		*outStats = newJobStats(job)
		outStats.fetch(ctx, job)
	}
	return arrowIterator, totalRows, nil
}

func ipcReaderFromArrowIterator(arrowIterator bigquery.ArrowIterator, alloc memory.Allocator) (*ipc.Reader, error) {
	arrowItReader := bigquery.NewArrowIteratorReader(arrowIterator)
	return ipc.NewReader(arrowItReader, ipc.WithAllocator(alloc))
}

func getQueryParameter(values arrow.RecordBatch, row int, parameterMode string) ([]bigquery.QueryParameter, error) {
	parameters := make([]bigquery.QueryParameter, values.NumCols())
	includeName := parameterMode == OptionValueQueryParameterModeNamed
	schema := values.Schema()
	for i, v := range values.Columns() {
		pi, err := arrowValueToQueryParameterValue(schema.Field(i), v, row)
		if err != nil {
			return nil, err
		}
		parameters[i] = pi
		if includeName {
			parameters[i].Name = values.ColumnName(i)
		}
	}
	return parameters, nil
}

func runPlainQuery(ctx context.Context, query *bigquery.Query, alloc memory.Allocator, resultRecordBufferSize int, linkFailedJob bool, fetchJobStats bool) (bigqueryRdr *reader, totalRows int64, err error) {
	// Only allocate a stats struct (which triggers job.Status in runQuery) if
	// the caller opted in. Otherwise runQuery skips the extra API call and the
	// schema is returned without BIGQUERY:* stats metadata.
	var statsPtr *jobStats
	if fetchJobStats {
		statsPtr = &jobStats{}
	}
	arrowIterator, totalRows, err := runQuery(ctx, query, linkFailedJob, alloc, statsPtr)
	if err != nil {
		return nil, -1, err
	}
	rdr, err := ipcReaderFromArrowIterator(arrowIterator, alloc)
	if err != nil {
		return nil, -1, err
	}

	chs := make([]chan arrow.RecordBatch, 1)
	ctx, cancelFn := context.WithCancel(ctx)
	ch := make(chan arrow.RecordBatch, resultRecordBufferSize)
	chs[0] = ch

	defer func() {
		if err != nil {
			close(ch)
			cancelFn()
		}
	}()

	schema := statsPtr.attachToSchema(rdr.Schema())

	bigqueryRdr = &reader{
		refCount:   1,
		chs:        chs,
		curChIndex: 0,
		err:        nil,
		cancelFn:   cancelFn,
		schema:     schema,
	}

	go func() {
		defer rdr.Release()
		for rdr.Next() && ctx.Err() == nil {
			rec := rdr.RecordBatch()
			rec.Retain()
			ch <- rec
		}

		err = checkContext(ctx, rdr.Err())
		defer close(ch)
	}()
	return bigqueryRdr, totalRows, nil
}

func queryRecordWithSchemaCallback(ctx context.Context, group *errgroup.Group, query *bigquery.Query, rec arrow.RecordBatch, ch chan arrow.RecordBatch, parameterMode string, alloc memory.Allocator, rdrSchema func(schema *arrow.Schema), linkFailedJob bool, fetchJobStats bool) (int64, error) {
	totalRows := int64(-1)
	for i := 0; i < int(rec.NumRows()); i++ {
		parameters, err := getQueryParameter(rec, i, parameterMode)
		if err != nil {
			return -1, err
		}
		if parameters != nil {
			query.Parameters = parameters
		}

		var statsPtr *jobStats
		if fetchJobStats {
			statsPtr = &jobStats{}
		}
		arrowIterator, rows, err := runQuery(ctx, query, linkFailedJob, alloc, statsPtr)
		if err != nil {
			return -1, err
		}
		totalRows = rows
		rdr, err := ipcReaderFromArrowIterator(arrowIterator, alloc)
		if err != nil {
			return -1, err
		}
		schema := statsPtr.attachToSchema(rdr.Schema())
		rdrSchema(schema)
		group.Go(func() error {
			defer rdr.Release()
			for rdr.Next() && ctx.Err() == nil {
				rec := rdr.RecordBatch()
				rec.Retain()
				ch <- rec
			}
			return checkContext(ctx, rdr.Err())
		})
	}
	return totalRows, nil
}

// kicks off a goroutine for each endpoint and returns a reader which
// gathers all of the records as they come in.
func newRecordReader(ctx context.Context, query *bigquery.Query, boundParameters array.RecordReader, parameterMode string, alloc memory.Allocator, resultRecordBufferSize, prefetchConcurrency int, linkFailedJob bool, fetchJobStats bool) (bigqueryRdr *reader, totalRows int64, err error) {
	if boundParameters == nil {
		return runPlainQuery(ctx, query, alloc, resultRecordBufferSize, linkFailedJob, fetchJobStats)
	}
	defer boundParameters.Release()

	totalRows = 0
	// BigQuery can expose result sets as multiple streams when using certain APIs
	// for now lets keep this and set the number of channels to 1
	// when we need to adapt to multiple streams we can change the value here
	chs := make([]chan arrow.RecordBatch, 1)

	ch := make(chan arrow.RecordBatch, resultRecordBufferSize)
	group, ctx := errgroup.WithContext(ctx)
	group.SetLimit(prefetchConcurrency)
	ctx, cancelFn := context.WithCancel(ctx)
	chs[0] = ch

	defer func() {
		if err != nil {
			close(ch)
			cancelFn()
		}
	}()

	bigqueryRdr = &reader{
		refCount: 1,
		chs:      chs,
		err:      nil,
		cancelFn: cancelFn,
		schema:   nil,
	}

	for boundParameters.Next() {
		rec := boundParameters.RecordBatch()
		// Each call to Record() on the record reader is allowed to release the previous record
		// and since we're doing this sequentially
		// we don't need to call rec.Retain() here and call call rec.Release() in queryRecordWithSchemaCallback
		batchRows, err := queryRecordWithSchemaCallback(ctx, group, query, rec, ch, parameterMode, alloc, func(schema *arrow.Schema) {
			bigqueryRdr.schema = schema
		}, linkFailedJob, fetchJobStats)
		if err != nil {
			return nil, -1, err
		}
		totalRows += batchRows
	}
	bigqueryRdr.err = group.Wait()
	defer close(ch)
	return bigqueryRdr, totalRows, nil
}

func (r *reader) Retain() {
	atomic.AddInt64(&r.refCount, 1)
}

func (r *reader) Release() {
	if atomic.AddInt64(&r.refCount, -1) == 0 {
		if r.rec != nil {
			r.rec.Release()
		}
		r.cancelFn()
		for _, ch := range r.chs {
			for rec := range ch {
				rec.Release()
			}
		}
	}
}

func (r *reader) Err() error {
	return r.err
}

func (r *reader) Next() bool {
	if r.rec != nil {
		r.rec.Release()
		r.rec = nil
	}

	if r.curChIndex >= len(r.chs) {
		return false
	}
	var ok bool
	for r.curChIndex < len(r.chs) {
		if r.rec, ok = <-r.chs[r.curChIndex]; ok {
			break
		}
		r.curChIndex++
	}
	return r.rec != nil
}

func (r *reader) Schema() *arrow.Schema {
	return r.schema
}

func (r *reader) Record() arrow.RecordBatch {
	return r.rec
}

func (r *reader) RecordBatch() arrow.RecordBatch {
	return r.rec
}

type emptyArrowIterator struct {
	schema bigquery.Schema
}

func (e emptyArrowIterator) Next() (*bigquery.ArrowRecordBatch, error) {
	return nil, iterator.Done
}

func (e emptyArrowIterator) Schema() bigquery.Schema {
	return e.schema
}

func (e emptyArrowIterator) SerializedArrowSchema() []byte {

	fields := make([]arrow.Field, len(e.schema))
	for i, field := range e.schema {
		f, err := buildField(field, 0)
		if err != nil {
			log.Fatalf("Error building field %s: %v", field.Name, err)
		}
		fields[i] = f
	}

	arrowSchema := arrow.NewSchema(fields, nil)

	var buf bytes.Buffer
	writer := ipc.NewWriter(&buf, ipc.WithSchema(arrowSchema))

	err := writer.Close()
	if err != nil {
		log.Fatalf("Error serializing an empty schema: %v", err)
	}

	return buf.Bytes()
}

// RowBasedArrowIterator wraps a bigquery.RowIterator and implements the ArrowIterator interface
// This is used when the Storage Read API cannot be used (e.g. to read data for pseudo-columns like _PARTITIONTIME)
type RowBasedArrowIterator struct {
	iter   *bigquery.RowIterator
	schema bigquery.Schema
	alloc  memory.Allocator
	done   bool
}

func newRowBasedArrowIterator(iter *bigquery.RowIterator, alloc memory.Allocator) bigquery.ArrowIterator {
	return &RowBasedArrowIterator{
		iter:   iter,
		schema: iter.Schema,
		alloc:  alloc,
		done:   false,
	}
}

func (l *RowBasedArrowIterator) Next() (*bigquery.ArrowRecordBatch, error) {
	if l.done {
		return nil, iterator.Done
	}

	const batchSize = 1000
	rows := make([][]bigquery.Value, 0, batchSize)

	for i := 0; i < batchSize; i++ {
		var row []bigquery.Value
		err := l.iter.Next(&row)
		if err == iterator.Done {
			l.done = true
			break
		}
		if err != nil {
			return nil, err
		}
		rows = append(rows, row)
	}

	if len(rows) == 0 {
		return nil, iterator.Done
	}

	batch, err := rowsToArrowRecordBatch(l.schema, rows, l.alloc)
	if err != nil {
		log.Fatalf("Error converting rows to arrow record batch: %v", err)
		return nil, err
	}
	defer batch.Release()

	var buf bytes.Buffer
	writer := ipc.NewWriter(&buf, ipc.WithSchema(batch.Schema()), ipc.WithAllocator(l.alloc))
	if err := writer.Write(batch); err != nil {
		return nil, err
	}
	if err := writer.Close(); err != nil {
		return nil, err
	}

	return &bigquery.ArrowRecordBatch{
		Data: buf.Bytes(),
	}, nil
}

func (l *RowBasedArrowIterator) Schema() bigquery.Schema {
	return l.schema
}

func (l *RowBasedArrowIterator) SerializedArrowSchema() []byte {
	fields := make([]arrow.Field, len(l.schema))
	for i, field := range l.schema {
		f, err := buildField(field, 0)
		if err != nil {
			log.Fatalf("Error building field %s: %v", field.Name, err)
		}
		fields[i] = f
	}

	arrowSchema := arrow.NewSchema(fields, nil)

	var buf bytes.Buffer
	_ = ipc.NewWriter(&buf, ipc.WithSchema(arrowSchema))

	return buf.Bytes()
}

func rowsToArrowRecordBatch(schema bigquery.Schema, rows [][]bigquery.Value, alloc memory.Allocator) (arrow.Record, error) {
	if len(rows) == 0 {
		return nil, fmt.Errorf("no rows to convert")
	}

	// Build a schema
	fields := make([]arrow.Field, len(schema))
	for i, field := range schema {
		f, err := buildField(field, 0)
		if err != nil {
			return nil, err
		}
		fields[i] = f
	}
	arrowSchema := arrow.NewSchema(fields, nil)

	// Build arrays for columns
	builders := make([]array.Builder, len(schema))
	for i, field := range fields {
		builders[i] = array.NewBuilder(alloc, field.Type)
	}
	defer func() {
		for _, b := range builders {
			b.Release()
		}
	}()

	// Populate data
	for _, row := range rows {
		for colIdx, val := range row {
			if val == nil {
				builders[colIdx].AppendNull()
				continue
			}

			switch builder := builders[colIdx].(type) {
			case *array.Date32Builder:
				// BigQuery returns civil.Date for DATE columns
				if d, ok := val.(civil.Date); ok {
					t := time.Date(d.Year, time.Month(d.Month), d.Day, 0, 0, 0, 0, time.UTC)
					builder.Append(arrow.Date32FromTime(t))
				} else if t, ok := val.(time.Time); ok {
					builder.Append(arrow.Date32FromTime(t))
				} else {
					builder.AppendNull()
				}
			case *array.TimestampBuilder:
				if ts, ok := val.(time.Time); ok {
					builder.Append(arrow.Timestamp(ts.UnixMicro()))
				} else {
					builder.AppendNull()
				}
			// TODO: Add support for other types as needed
			default:
				return nil, fmt.Errorf("USE_STORAGE_API_DISABLED_CLIENT is enabled, unsupported type conversion for column type %s of value %v", builder.Type().String(), val)
			}
		}
	}

	arrays := make([]arrow.Array, len(builders))
	for i, b := range builders {
		arrays[i] = b.NewArray()
	}

	return array.NewRecordBatch(arrowSchema, arrays, int64(len(rows))), nil
}

var _ array.RecordReader = (*reader)(nil)