Merge pull request #763 from niaow/sql-leak

Fix SQL memory leak

proto/interface.go

@@ -40,36 +40,6 @@ func (EmptyStream) Recv() (*RowResponse, error) {
 	return nil, io.EOF
 }
 
-// ReadIntoTable reads from a StreamClient and stores the result into a table response.
-func ReadIntoTable(cli StreamClient) (*TableResponse, error) {
-	var headers []*ColumnInfo
-	rows := []*Row{}
-
-rx:
-	for {
-		row, err := cli.Recv()
-		switch err {
-		case nil:
-		case io.EOF:
-			break rx
-		default:
-			return nil, err
-		}
-
-		if headers == nil {
-			headers = row.Headers
-		}
-		rows = append(rows, &Row{
-			Columns: row.Columns,
-		})
-	}
-
-	return &TableResponse{
-		Headers: headers,
-		Rows:    rows,
-	}, nil
-}
-
 // StreamServer is an interface for a stream
 // which can accept a RowResponse to be later
 // returned by the stream via Recv().
@@ -93,24 +63,20 @@ type ToRowser interface {
 
 // RowsToTable is a helper function which takes a ToRowser,
 // along with the number of rows, and returns a TableResponse.
-// Obviously passing the number of rows seems unnecessary,
-// and we could remove that requirement, but for now we
-// do it to allow for pre-allocation of the rows slice.
+// The number of rows is treated as a hint.
 func RowsToTable(tr ToRowser, n int) (*TableResponse, error) {
 	var headers []*ColumnInfo
-	rows := make([]*Row, n)
+	rows := make([]*Row, 0, n)
 
 	// This callback gets called for every "row" in r.
 	// Each row populates its position in the pre-allocated
 	// `rows`. The headers get set based on those received
 	// in the first row.
-	var idx int
 	cb := func(rr *RowResponse) error {
-		if idx == 0 {
+		if len(rows) == 0 {
 			headers = rr.GetHeaders()
 		}
-		rows[idx] = &Row{Columns: rr.GetColumns()}
-		idx++
+		rows = append(rows, &Row{Columns: rr.GetColumns()})
 		return nil
 	}
 
@@ -124,60 +90,6 @@ func RowsToTable(tr ToRowser, n int) (*TableResponse, error) {
 	}, nil
 }
 
-// RowBuffer acts as a Sender/Receiver of RowResponses.
-// Note that sending a nil value will cause the Recv
-// method to return an io.EOF error.
-type RowBuffer struct {
-	ch chan *RowResponse
-}
-
-// NewRowBuffer returns a new instance of RowBuffer.
-// sz is the size of the buffer.
-func NewRowBuffer(sz int) *RowBuffer {
-	var chSz int
-	if sz > 0 {
-		// Add one to allow for the EOF record.
-		chSz = sz + 1
-	}
-	return &RowBuffer{
-		ch: make(chan *RowResponse, chSz),
-	}
-}
-
-// Recv returns a RowResponse. When the buffer is empty,
-// calling Recv will return an io.EOF error.
-func (rb *RowBuffer) Recv() (*RowResponse, error) {
-	r := <-rb.ch
-
-	// If the StatusError contains a message then return
-	// with the approprate error.
-	se := r.GetStatusError()
-	code := codes.Code(se.GetCode())
-	msg := se.GetMessage()
-	if code != codes.OK {
-		return nil, status.Error(code, msg)
-	} else if msg == "EOF" {
-		return nil, io.EOF
-	} else if msg != "" {
-		return nil, errors.New(msg)
-	}
-
-	return r, nil
-}
-
-func (rb *RowBuffer) Send(rr *RowResponse) error {
-	rb.ch <- rr
-	return nil
-}
-
-// EOF acts as an io.EOF encoded into a RowResponse.
-var EOF *RowResponse = &RowResponse{
-	StatusError: &StatusError{
-		Code:    0,
-		Message: "EOF",
-	},
-}
-
 // Error is a helper function to create a RowResponse
 // based on an error message. If the error is a grpc
 // Status, then the status code is passed through.
@@ -398,3 +310,18 @@ func (r RowResponseSorter) Less(i, j int) bool {
 	}
 	return false
 }
+
+// ConstRowser implements ToRowser with a slice of row responses.
+type ConstRowser []RowResponse
+
+// ToRows calls a function with a pointer to each element of the slice.
+func (c ConstRowser) ToRows(fn func(*RowResponse) error) error {
+	for i := range c {
+		err := fn(&c[i])
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

server/grpc.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
-	"io"
 	"net"
 	"strings"
 	"sync"
@@ -167,7 +166,7 @@ func (h *GRPCHandler) DeleteVDS(ctx context.Context, req *pb.DeleteVDSRequest) (
 	return &pb.DeleteVDSResponse{}, nil
 }
 
-func (h *GRPCHandler) execSQL(ctx context.Context, queryStr string) (pb.StreamClient, error) {
+func (h *GRPCHandler) execSQL(ctx context.Context, queryStr string) (pb.ToRowser, error) {
 	return execSQL(ctx, h.api, h.logger, queryStr)
 }
 
@@ -178,21 +177,12 @@ func (h *GRPCHandler) QuerySQL(req *pb.QuerySQLRequest, stream pb.Pilosa_QuerySQ
 		return err
 	}
 
-	for {
-		row, err := results.Recv()
-		switch err {
-		case nil:
-		case io.EOF:
-			return nil
-		default:
-			return errors.Wrap(err, "failed to load next row")
-		}
-
-		err = stream.Send(row)
-		if err != nil {
-			return errors.Wrap(err, "failed to send row")
-		}
+	err = results.ToRows(stream.Send)
+	if err != nil {
+		return errors.Wrap(err, "streaming result")
 	}
+
+	return nil
 }
 
 // QuerySQLUnary is a unary-response (non-streaming) version of QuerySQL, returning a TableResponse.
@@ -212,7 +202,10 @@ func (h *GRPCHandler) QuerySQLUnary(ctx context.Context, req *pb.QuerySQLRequest
 	if err != nil {
 		return nil, err
 	}
-	return pb.ReadIntoTable(results)
+	if results, ok := results.(pb.ToTabler); ok {
+		return results.ToTable()
+	}
+	return pb.RowsToTable(results, 0)
 }
 
 // QueryPQL handles the PQL request and sends RowResponses to the stream.

server/pg.go

@@ -19,7 +19,6 @@ import (
 	"crypto/tls"
 	"encoding/json"
 	"fmt"
-	"io"
 	"net"
 	"strconv"
 	"strings"
@@ -358,28 +357,6 @@ func pgWriteRowser(w pg.QueryResultWriter, result pb.ToRowser) error {
 	})
 }
 
-type clientRowser struct {
-	pb.StreamClient
-}
-
-func (cr *clientRowser) ToRows(f func(*pb.RowResponse) error) error {
-	for {
-		resp, err := cr.StreamClient.Recv()
-		if err != nil {
-			if err == io.EOF {
-				return nil
-			}
-
-			return err
-		}
-
-		err = f(resp)
-		if err != nil {
-			return err
-		}
-	}
-}
-
 func pgWriteResult(w pg.QueryResultWriter, result interface{}) error {
 	switch result := result.(type) {
 	case *pilosa.Row:
@@ -392,8 +369,6 @@ func pgWriteResult(w pg.QueryResultWriter, result interface{}) error {
 		return pgWriteGroupCount(w, result)
 	case pb.ToRowser: // we should avoid protobuf where we can...
 		return pgWriteRowser(w, result)
-	case pb.StreamClient:
-		return pgWriteRowser(w, &clientRowser{result})
 	case uint64:
 		err := w.WriteHeader(pg.ColumnInfo{
 			Name: "count",

server/sql.go

@@ -26,14 +26,14 @@ import (
 	"google.golang.org/grpc/status"
 )
 
-func execSQL(ctx context.Context, api *pilosa.API, logger logger.Logger, queryStr string) (pb.StreamClient, error) {
+func execSQL(ctx context.Context, api *pilosa.API, logger logger.Logger, queryStr string) (pb.ToRowser, error) {
 	mapper := sql.NewMapper()
 	mapper.Logger = logger
 	query, err := mapper.MapSQL(queryStr)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to map SQL")
 	}
-	var results pb.StreamClient
+	var results pb.ToRowser
 	switch query.SQLType {
 	case sql.SQLTypeSelect:
 		handler := sql.NewSelectHandler(api)

sql/ddl.go

@@ -37,7 +37,7 @@ func NewDDLHandler(api *pilosa.API) *DDLHandler {
 }
 
 // Handle executes mapped SQL
-func (h *DDLHandler) Handle(ctx context.Context, mapped *MappedSQL) (pproto.StreamClient, error) {
+func (h *DDLHandler) Handle(ctx context.Context, mapped *MappedSQL) (pproto.ToRowser, error) {
 	stmt, ok := mapped.Statement.(*sqlparser.DDL)
 	if !ok {
 		return nil, fmt.Errorf("statement is not type DDL: %T", mapped.Statement)
@@ -52,7 +52,7 @@ func (h *DDLHandler) Handle(ctx context.Context, mapped *MappedSQL) (pproto.Stre
 	}
 }
 
-func (h *DDLHandler) execDropTable(ctx context.Context, stmt *sqlparser.DDL) (pproto.StreamClient, error) {
+func (h *DDLHandler) execDropTable(ctx context.Context, stmt *sqlparser.DDL) (pproto.ToRowser, error) {
 	if n := len(stmt.FromTables); n != 1 {
 		return nil, fmt.Errorf("statement can only contain a single drop table, but got: %d", n)
 	}
@@ -61,5 +61,5 @@ func (h *DDLHandler) execDropTable(ctx context.Context, stmt *sqlparser.DDL) (pp
 	if err := h.api.DeleteIndex(ctx, indexName); err != nil {
 		return nil, errors.Wrapf(err, "deleting index %s", indexName)
 	}
-	return pproto.EmptyStream{}, nil
+	return pproto.ConstRowser{}, nil
 }

sql/reduce.go

@@ -15,110 +15,105 @@
 package sql
 
 import (
-	"io"
 	"sort"
 
+	"github.com/pilosa/pilosa/v2"
 	"github.com/pilosa/pilosa/v2/pql"
 	pproto "github.com/pilosa/pilosa/v2/proto"
 	"github.com/pkg/errors"
-	"google.golang.org/grpc/codes"
 )
 
-// DataType contants describe the possible values
-// for the Datatype value in the RowResponse header.
-const (
-	DataTypeDecimal     = "decimal"
-	DataTypeFloat64     = "float64"
-	DataTypeInt64       = "int64"
-	DataTypeString      = "string"
-	DataTypeUint64Array = "[]uint64"
-)
-
-type Reducer interface {
-	Reduce(pproto.StreamClient, pproto.StreamServer) error
+type limitRowser struct {
+	rowser pproto.ToRowser
+	limit  uint
 }
 
-// LimitReducer limits the number of messages passed through.
-type LimitReducer struct {
-	limit  uint
+func (l *limitRowser) ToRows(fn func(*pproto.RowResponse) error) error {
+	limit := l.limit
+	return l.rowser.ToRows(func(row *pproto.RowResponse) error {
+		if limit == 0 {
+			return nil
+		}
+		limit--
+
+		return fn(row)
+	})
+}
+
+// LimitRows applies a limit to a ToRowser.
+func LimitRows(rowser pproto.ToRowser, limit uint) pproto.ToRowser {
+	switch rowser := rowser.(type) {
+	case pilosa.ExtractedTable:
+		if uint(len(rowser.Columns)) > limit {
+			rowser.Columns = rowser.Columns[:limit]
+		}
+		return rowser
+	default:
+		return &limitRowser{rowser, limit}
+	}
+}
+
+type offsetRowser struct {
+	rowser pproto.ToRowser
 	offset uint
 }
 
-// NewLimitReducer returns a new instance of LimitReducer.
-func NewLimitReducer(limit, offset uint) *LimitReducer {
-	return &LimitReducer{
-		limit:  limit,
-		offset: offset,
-	}
-}
-
-// Reduce applies the limit reducer to the client stream and sends the results
-// to the server stream.
-func (l *LimitReducer) Reduce(c pproto.StreamClient, s pproto.StreamServer) error {
-	offsetCountdown := l.offset
-
-	// in the case of an offset, since we'll be skipping the first record
-	// which contains the headers, we need to pull the headers, save them,
-	// and apply them to the first record that we actually send through.
+func (o *offsetRowser) ToRows(fn func(*pproto.RowResponse) error) error {
+	offset := o.offset
 	var headers []*pproto.ColumnInfo
+	return o.rowser.ToRows(func(row *pproto.RowResponse) error {
+		if headers == nil {
+			headers = row.Headers
+		}
+		if offset > 0 {
+			offset--
+			return nil
+		}
+		row.Headers = headers
 
-	for i := uint(0); i < l.limit+l.offset || l.limit == 0; i++ {
-		r, err := c.Recv()
-		if err == io.EOF {
-			break
-		} else if err != nil {
-			return s.Send(pproto.ErrorWrap(err, "receiving on client stream"))
-		}
-		if offsetCountdown > 0 {
-			if headers == nil {
-				headers = r.Headers
-			}
-			offsetCountdown--
-			continue
-		}
-		if headers != nil {
-			r.Headers = headers
-			headers = nil
-		}
-		if err := s.Send(r); err != nil {
-			return s.Send(pproto.ErrorWrap(err, "sending on server stream"))
-		}
-	}
-	return s.Send(pproto.EOF)
+		return fn(row)
+	})
 }
 
-// OrderByReducer orders the results based on the provide conditions.
-// It also takes limit and offset to reduce the amount of items
-// needing to be held in memory for sorting.
-type OrderByReducer struct {
+// OffsetRows applies an offset to a ToRowser.
+func OffsetRows(rowser pproto.ToRowser, offset uint) pproto.ToRowser {
+	if offset == 0 {
+		return rowser
+	}
+
+	switch rowser := rowser.(type) {
+	case pilosa.ExtractedTable:
+		if uint(len(rowser.Columns)) > offset {
+			rowser.Columns = rowser.Columns[:0]
+		} else {
+			rowser.Columns = rowser.Columns[offset:]
+		}
+		return rowser
+	default:
+		return &offsetRowser{rowser, offset}
+	}
+}
+
+type orderByRowser struct {
+	rowser       pproto.ToRowser
 	fields       []string
 	isDescending []bool // direction[asc: false, desc: true]
-	limit        uint
-	offset       uint
 }
 
-// NewOrderByReducer returns a new instance of OrderByReducer.
-func NewOrderByReducer(fields, dirs []string, limit, offset uint) *OrderByReducer {
-	descendings := make([]bool, len(fields))
-	for i := range dirs {
-		if dirs[i] == "desc" {
-			descendings[i] = true
-		}
-	}
-	return &OrderByReducer{
-		fields:       fields,
-		isDescending: descendings,
-		limit:        limit,
-		offset:       offset,
-	}
-}
-
-// Reduce applies the order by reducer to the client stream and sends the results
-// to the server stream.
-func (o *OrderByReducer) Reduce(c pproto.StreamClient, s pproto.StreamServer) error {
+func (o *orderByRowser) ToRows(fn func(*pproto.RowResponse) error) error {
 	// hold is a slice of row responses, to be sent to the output
 	// stream sorted by the sort conditions.
 	var hold []*pproto.RowResponse
+	err := o.rowser.ToRows(func(row *pproto.RowResponse) error {
+		hold = append(hold, row)
+		return nil
+	})
+	if err != nil {
+		return err
+	}
+	if len(hold) == 0 {
+		return nil
+	}
 
 	// sortColNames contains the names of the columns to
 	// sort on.
@@ -138,49 +133,16 @@ func (o *OrderByReducer) Reduce(c pproto.StreamClient, s pproto.StreamServer) er
 	// the first row) so they can be applied later
 	// to what will eventually be the first row after
 	// sorting has occurred.
-	var holdHeaders []*pproto.ColumnInfo
-
-	ii := 0
-	for {
-		rr, err := c.Recv()
-		if err != nil {
-			if err == io.EOF {
-				break
+	holdHeaders := hold[0].Headers
+	for i, hdr := range holdHeaders {
+		hdrName := hdr.GetName()
+		hdrType := hdr.GetDatatype()
+		for j := range sortColNames {
+			if sortColNames[j] == hdrName {
+				sortColIdxs[j] = i
+				sortColTypes[j] = hdrType
 			}
-			return s.Send(pproto.ErrorWrap(err, "receiving row response"))
 		}
-
-		// On the first row, get the sort column information
-		// from the headers. Also, stash the headers for
-		// later in the `holdHeaders` var.
-		if ii == 0 {
-			holdHeaders = rr.Headers
-			for i, rrHdr := range rr.Headers {
-				hdrName := rrHdr.GetName()
-				hdrType := rrHdr.GetDatatype()
-				for j := range sortColNames {
-					if sortColNames[j] == hdrName {
-						sortColIdxs[j] = i
-						sortColTypes[j] = hdrType
-					}
-				}
-			}
-			// Clear the headers in case this record is
-			// no longer first (we re-apply the headers
-			// to the first outgoing record later).
-			rr.Headers = nil
-		}
-
-		// Put each row in the hold.
-		hold = append(hold, rr)
-
-		ii++
-
-		// TODO: in the case where limit is provided and the number of possible
-		// rows is large, it might be more efficient to periodically sort/trim
-		// the hold so it doesn't become too large. For example, it could
-		// be constrained to size (limit + offset + buffer), where buffer is
-		// an amount that the hold can grow before being trimmed.
 	}
 
 	// Sort the hold.
@@ -191,61 +153,59 @@ func (o *OrderByReducer) Reduce(c pproto.StreamClient, s pproto.StreamServer) er
 		hold,
 	)
 	if err != nil {
-		return s.Send(pproto.ErrorWrap(err, "creating row response sorter"))
+		return errors.Wrap(err, "creating row response sorter")
 	}
 	sort.Sort(sorter)
 
-	var rowsToConsider uint = uint(len(hold))
-	var offsetCountdown uint
-	if o.limit > 0 {
-		offsetCountdown = o.offset
-		if o.limit+o.offset < rowsToConsider {
-			rowsToConsider = o.limit + o.offset
-		}
-	}
-
 	// Loop over hold and send each row response.
 	// Apply the header to the first row that is sent.
 	var headerApplied bool
-	for i := uint(0); i < rowsToConsider; i++ {
-		if offsetCountdown > 0 {
-			offsetCountdown--
-			continue
-		}
+	for i := range hold {
 		// Re-apply the headers to the first record.
 		if !headerApplied {
 			hold[i].Headers = holdHeaders
 			headerApplied = true
 		}
-		err := s.Send(hold[i])
+		err := fn(hold[i])
 		if err != nil {
-			return s.Send(pproto.ErrorWrap(err, "sending hold row"))
+			return errors.Wrap(err, "sending hold row")
 		}
 	}
-	return s.Send(pproto.EOF)
+
+	return nil
 }
 
-// ValCountFuncReducer converts a ValCount result to the proper
-// result for Func.
-type ValCountFuncReducer struct {
-	fn FuncName
-}
-
-// NewValCountFuncReducer returns a new instance of ValCountFuncReducer.
-func NewValCountFuncReducer(fn FuncName) *ValCountFuncReducer {
-	return &ValCountFuncReducer{
-		fn: fn,
+// OrderBy sorts a rowser.
+func OrderBy(rowser pproto.ToRowser, fields, dirs []string) pproto.ToRowser {
+	descendings := make([]bool, len(fields))
+	for i := range dirs {
+		if dirs[i] == "desc" {
+			descendings[i] = true
+		}
+	}
+	return &orderByRowser{
+		rowser:       rowser,
+		fields:       fields,
+		isDescending: descendings,
 	}
 }
 
-// Reduce modifies the stream according to the function.
-func (v *ValCountFuncReducer) Reduce(c pproto.StreamClient, s pproto.StreamServer) error {
-	r, err := c.Recv()
+type valCountRowser struct {
+	rowser pproto.ToRowser
+	fn     FuncName
+}
+
+func (v *valCountRowser) ToRows(fn func(row *pproto.RowResponse) error) error {
+	var r *pproto.RowResponse
+	err := v.rowser.ToRows(func(row *pproto.RowResponse) error {
+		if r != nil {
+			return errors.New("extra row in valcount")
+		}
+		r = row
+		return nil
+	})
 	if err != nil {
-		if err == io.EOF {
-			return s.Send(pproto.EOF)
-		}
-		return s.Send(pproto.Error(err))
+		return err
 	}
 
 	// Get the index of the column with header of "value".
@@ -268,7 +228,7 @@ func (v *ValCountFuncReducer) Reduce(c pproto.StreamClient, s pproto.StreamServe
 	returnDataType = sourceDataType
 	switch v.fn {
 	case FuncAvg:
-		returnDataType = DataTypeFloat64
+		returnDataType = "float64"
 	}
 
 	rr := pproto.RowResponse{
@@ -280,29 +240,20 @@ func (v *ValCountFuncReducer) Reduce(c pproto.StreamClient, s pproto.StreamServe
 
 	cols := r.GetColumns()
 	if len(cols) == 0 {
-		return s.Send(pproto.ErrorCode(
-			errors.New("empty column set"),
-			codes.NotFound,
-		))
+		return errors.New("empty column set")
 	}
 
 	if idxVal == -1 {
-		return s.Send(pproto.ErrorCode(
-			errors.New("result set has no column: value"),
-			codes.NotFound,
-		))
+		return errors.New("result set has no column: value")
 	}
 	if idxCnt == -1 {
-		return s.Send(pproto.ErrorCode(
-			errors.New("result set has no column: count"),
-			codes.NotFound,
-		))
+		return errors.New("result set has no column: count")
 	}
 
 	switch v.fn {
 	case FuncAvg:
 		var avg float64
-		if sourceDataType == DataTypeDecimal {
+		if sourceDataType == "decimal" {
 			val := cols[idxVal].GetDecimalVal()
 			dec := pql.NewDecimal(val.Value, val.Scale)
 			cnt := cols[idxCnt].GetInt64Val()
@@ -314,7 +265,7 @@ func (v *ValCountFuncReducer) Reduce(c pproto.StreamClient, s pproto.StreamServe
 		}
 		rr.Columns[0] = &pproto.ColumnResponse{ColumnVal: &pproto.ColumnResponse_Float64Val{Float64Val: avg}}
 	default:
-		if sourceDataType == DataTypeDecimal {
+		if sourceDataType == "decimal" {
 			val := cols[idxVal].GetDecimalVal()
 			rr.Columns[0] = &pproto.ColumnResponse{ColumnVal: &pproto.ColumnResponse_DecimalVal{DecimalVal: &pproto.Decimal{Value: val.Value, Scale: val.Scale}}}
 		} else {
@@ -323,127 +274,90 @@ func (v *ValCountFuncReducer) Reduce(c pproto.StreamClient, s pproto.StreamServe
 		}
 	}
 
-	if err := s.Send(&rr); err != nil {
-		return errors.Wrap(err, "sending row response")
-	}
-	return s.Send(pproto.EOF)
+	return fn(&rr)
 }
 
-// CountIDReducer returns a stream of _id's as a count.
-type CountIDReducer struct{}
+// ApplyValCountFunc converts a ValCount result to the proper
+// result for Func
+func ApplyValCountFunc(rowser pproto.ToRowser, fn FuncName) pproto.ToRowser {
+	return &valCountRowser{
+		rowser: rowser,
+		fn:     fn,
+	}
+}
 
-// Reduce counts the stream of IDs and returns a single record.
-func (r *CountIDReducer) Reduce(c pproto.StreamClient, s pproto.StreamServer) error {
-	var cnt uint64
+type countIDRowser struct {
+	rowser pproto.ToRowser
+}
 
-	for {
-		_, err := c.Recv()
-		if err != nil {
-			if err == io.EOF {
-				break
-			}
-			return s.Send(pproto.ErrorWrap(err, "receiving on client stream"))
-		}
-		cnt++
+func (c *countIDRowser) ToRows(fn func(*pproto.RowResponse) error) error {
+	var count uint64
+	err := c.rowser.ToRows(func(row *pproto.RowResponse) error {
+		count++
+		return nil
+	})
+	if err != nil {
+		return err
 	}
 
-	rr := pproto.RowResponse{
+	return fn(&pproto.RowResponse{
 		Headers: []*pproto.ColumnInfo{
 			{Name: string(FuncCount), Datatype: "uint64"},
 		},
 		Columns: []*pproto.ColumnResponse{
-			&pproto.ColumnResponse{ColumnVal: &pproto.ColumnResponse_Uint64Val{Uint64Val: cnt}},
+			{
+				ColumnVal: &pproto.ColumnResponse_Uint64Val{Uint64Val: count},
+			},
 		},
-	}
-
-	if err := s.Send(&rr); err != nil {
-		return errors.Wrap(err, "sending row response")
-	}
-	return s.Send(pproto.EOF)
+	})
 }
 
-// AssignHeadersReducer overwrites the headers on the first record
-// according to field names and aliases from sql. It also reorders
-// the columns in the result stream to match the sql select clause.
-type AssignHeadersReducer struct {
-	cols []Column
+// CountRows counts the rows from the input rowser.
+func CountRows(rowser pproto.ToRowser) pproto.ToRowser {
+	return &countIDRowser{rowser}
 }
 
-// NewAssignHeadersReducer returns a new instance of AssignHeadersReducer.
-func NewAssignHeadersReducer(cols []Column) *AssignHeadersReducer {
-	return &AssignHeadersReducer{
-		cols: cols,
-	}
+type assignHeadersRowser struct {
+	rowser pproto.ToRowser
+	cols   []Column
 }
 
-// Reduce modifies the stream.
-func (r *AssignHeadersReducer) Reduce(c pproto.StreamClient, s pproto.StreamServer) error {
+func (a *assignHeadersRowser) ToRows(fn func(*pproto.RowResponse) error) error {
 	var placement []uint
-	var labels []string
 
-	var cnt int
-	for {
-		rr, err := c.Recv()
-		if err != nil {
-			if err == io.EOF {
-				break
-			}
-			return s.Send(pproto.ErrorWrap(err, "receiving on client stream"))
-		}
-
-		// If the placement slice is [0-n] where n == len(Headers)
-		// then we don't need to alter rr on records after cnt == 0.
-		// If we don't apply aliases, we don't have to alter Headers
-		// either, but that may not be worth messing with.
-
-		if cnt == 0 {
-			placement, labels, err = headerAssignment(r.cols, rr.Headers)
+	return a.rowser.ToRows(func(row *pproto.RowResponse) error {
+		var out pproto.RowResponse
+		if placement == nil {
+			// Assign headers and generate placement.
+			var err error
+			var labels []string
+			placement, labels, err = headerAssignment(a.cols, row.Headers)
 			if err != nil {
-				return s.Send(pproto.ErrorWrap(err, "getting header assignment"))
+				return errors.Wrap(err, "getting header assignment")
 			}
-
-			// mod is the modified RowResponse object that gets populated
-			// according to placement and labels, then sent.
-			mod := &pproto.RowResponse{
-				Headers: make([]*pproto.ColumnInfo, len(placement)),
-				Columns: make([]*pproto.ColumnResponse, len(placement)),
-			}
-
-			// For now, we assume that the column count in each RowResponse
-			// is consistent (i.e. we can validate one time, here, on the
-			// first row, and not every time, in the `else` statement below).
-			if len(placement) > len(rr.Columns) {
-				return s.Send(pproto.ErrorCode(
-					errors.New("mismatched header placement and column count"),
-					codes.InvalidArgument,
-				))
-			}
-
-			for i := 0; i < len(placement); i++ {
-				mod.Headers[i] = rr.Headers[placement[i]]
-				mod.Headers[i].Name = labels[i]
-				mod.Columns[i] = rr.Columns[placement[i]]
-			}
-			if err := s.Send(mod); err != nil {
-				return errors.Wrap(err, "sending mod")
-			}
-		} else {
-			// mod is the modified RowResponse object that gets populated
-			// according to placement and labels, then sent.
-			mod := &pproto.RowResponse{
-				Columns: make([]*pproto.ColumnResponse, len(placement)),
-			}
-			for i := 0; i < len(placement); i++ {
-				mod.Columns[i] = rr.Columns[placement[i]]
-			}
-			if err := s.Send(mod); err != nil {
-				return errors.Wrap(err, "sending mod")
+			headers := make([]*pproto.ColumnInfo, len(placement))
+			for i, v := range placement {
+				header := row.Headers[v]
+				header.Name = labels[i]
+				headers[i] = header
 			}
+			out.Headers = headers
 		}
-		cnt++
-	}
 
-	return s.Send(pproto.EOF)
+		// Re-order the columns.
+		cols := make([]*pproto.ColumnResponse, len(placement))
+		for i, v := range placement {
+			cols[i] = row.Columns[v]
+		}
+		out.Columns = cols
+
+		return fn(&out)
+	})
+}
+
+// AssignHeaders assigns headers to a ToRowser.
+func AssignHeaders(rowser pproto.ToRowser, headers ...Column) pproto.ToRowser {
+	return &assignHeadersRowser{rowser, headers}
 }
 
 var (

sql/select.go

@@ -41,7 +41,7 @@ func NewSelectHandler(api *pilosa.API) *SelectHandler {
 }
 
 // Handle executes mapped SQL
-func (s *SelectHandler) Handle(ctx context.Context, mapped *MappedSQL) (pproto.StreamClient, error) {
+func (s *SelectHandler) Handle(ctx context.Context, mapped *MappedSQL) (pproto.ToRowser, error) {
 	stmt, ok := mapped.Statement.(*sqlparser.Select)
 	if !ok {
 		return nil, fmt.Errorf("statement is not type select: %T", mapped.Statement)
@@ -74,7 +74,7 @@ func (s *SelectHandler) mapSelect(ctx context.Context, selectStmt *sqlparser.Sel
 	return mr, nil
 }
 
-func (s *SelectHandler) execMappingResult(ctx context.Context, mr *MappingResult) (pproto.StreamClient, error) {
+func (s *SelectHandler) execMappingResult(ctx context.Context, mr *MappingResult) (pproto.ToRowser, error) {
 	if mr.Query == "" {
 		return nil, errors.New("no pql query created")
 	}
@@ -85,29 +85,15 @@ func (s *SelectHandler) execMappingResult(ctx context.Context, mr *MappingResult
 	}
 	res := resp.Results[0]
 
-	// TODO: synchronize this properly somehow.
-	// It would probbably help to get rid of the streaming too.
-	respRows := pproto.NewRowBuffer(0)
+	var result pproto.ToRowser
 	switch res := res.(type) {
 	case pproto.ToRowser:
-		go func() {
-			if err := res.ToRows(respRows.Send); err != nil {
-				respRows.Send(pproto.Error(err)) //nolint:errcheck
-			} else {
-				_ = respRows.Send(pproto.EOF) //nolint:errcheck
-			}
-		}()
+		result = res
 	case []pilosa.GroupCount:
-		go func() {
-			if err := pilosa.GroupCounts(res).ToRows(respRows.Send); err != nil {
-				respRows.Send(pproto.Error(err)) //nolint:errcheck
-			} else {
-				respRows.Send(pproto.EOF) //nolint:errcheck
-			}
-		}()
+		result = pilosa.GroupCounts(res)
 	case uint64:
-		go func() {
-			respRows.Send(&pproto.RowResponse{ //nolint:errcheck
+		result = pproto.ConstRowser{
+			{
 				Headers: []*pproto.ColumnInfo{
 					{
 						Name:     "count",
@@ -121,12 +107,11 @@ func (s *SelectHandler) execMappingResult(ctx context.Context, mr *MappingResult
 						},
 					},
 				},
-			})
-			respRows.Send(pproto.EOF) //nolint:errcheck
-		}()
+			},
+		}
 	case bool:
-		go func() {
-			respRows.Send(&pproto.RowResponse{ //nolint:errcheck
+		result = pproto.ConstRowser{
+			{
 				Headers: []*pproto.ColumnInfo{
 					{
 						Name:     "result",
@@ -140,24 +125,15 @@ func (s *SelectHandler) execMappingResult(ctx context.Context, mr *MappingResult
 						},
 					},
 				},
-			})
-			respRows.Send(pproto.EOF) //nolint:errcheck
-		}()
+			},
+		}
 	default:
 		return nil, fmt.Errorf("unsupported result type %T", res)
 	}
 
-	// Apply Reducers
-	result := respRows
-	for _, red := range mr.Reducers {
-		out := pproto.NewRowBuffer(0)
-
-		// Run Reducers asyncronously.
-		// TODO: stop swallowing this error.
-		// TODO: does this need an EOF as input?
-		go red.Reduce(result, out) //nolint:errcheck
-
-		result = out
+	// Apply reducers.
+	for _, reducer := range mr.Reducers {
+		result = reducer(result)
 	}
 
 	return result, nil
@@ -172,10 +148,10 @@ type MappingResult struct {
 	Offset       uint64
 	Query        string
 	Header       []Column
-	Reducers     []Reducer
+	Reducers     []func(pproto.ToRowser) pproto.ToRowser
 }
 
-func (mr *MappingResult) addReducer(r Reducer) {
+func (mr *MappingResult) addReducer(r func(pproto.ToRowser) pproto.ToRowser) {
 	mr.Reducers = append(mr.Reducers, r)
 }
 
@@ -278,21 +254,24 @@ func (h handlerSelectFieldsFromTableWhere) Apply(stmt *sqlparser.Select, qm Quer
 		Header: selectFields,
 	}
 
-	// TODO: assign headers
-	mr.addReducer(NewAssignHeadersReducer(selectFields))
+	// assign headers
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return AssignHeaders(result, selectFields...)
+	})
 
-	// TODO: If both order and limit/offset are required, then
-	// we can't supply limit/offset to the InspectRequest; we
-	// have to get all records, which we don't want to do on
-	// a large data set. We need to come up with a better
-	// way to handle that situation.
-	switch {
-	case qm.HasOrderBy():
-		mr.addReducer(NewOrderByReducer(orderByFlds, orderByDirs, limit, offset))
-	case limit != 0:
+	if qm.HasOrderBy() {
+		// Sort the results.
+		mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+			return OrderBy(result, orderByFlds, orderByDirs)
+		})
+
+		// Apply the limit and offset after sorting.
+		mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+			return LimitRows(OffsetRows(result, offset), limit)
+		})
+	} else {
+		// Apply the limit and offset inside the query.
 		whereQuery = Limit(whereQuery, limit, offset)
-	case offset != 0:
-		whereQuery = Offset(whereQuery, offset)
 	}
 
 	if len(fields) > 0 && fields[0] == "_id" {
@@ -359,13 +338,23 @@ func (h handlerSelectDistinctFromTable) Apply(stmt *sqlparser.Select, qm QueryMa
 		Query:     qo,
 	}
 
-	mr.addReducer(NewAssignHeadersReducer(selectFields))
+	// Assign headers to the result.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return AssignHeaders(result, selectFields...)
+	})
+
 	if qm.HasOrderBy() {
-		mr.addReducer(NewOrderByReducer(orderByFlds, orderByDirs, limit, offset))
-	} else {
-		mr.addReducer(NewLimitReducer(limit, offset))
+		// Sort the result.
+		mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+			return OrderBy(result, orderByFlds, orderByDirs)
+		})
 	}
 
+	// Apply a limit and offset to the result.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return LimitRows(OffsetRows(result, offset), limit)
+	})
+
 	return mr, nil
 }
 
@@ -374,7 +363,7 @@ type handlerSelectCountFromTableWhere struct{}
 
 func (h handlerSelectCountFromTableWhere) Apply(stmt *sqlparser.Select, qm QueryMask, indexFunc func(string) *pilosa.Index) (*MappingResult, error) {
 	var qo string
-	var reducers []Reducer
+	var reducers []func(pproto.ToRowser) pproto.ToRowser
 
 	indexName, err := extractIndexName(stmt)
 	if err != nil {
@@ -412,7 +401,9 @@ func (h handlerSelectCountFromTableWhere) Apply(stmt *sqlparser.Select, qm Query
 	} else {
 		// TODO: add the Distinct (for Int fields) here (like we do in handlerSelectDistinctFromTable)
 		qo = Rows(funcs[0].field.Name())
-		reducers = append(reducers, &CountIDReducer{})
+		reducers = append(reducers, func(result pproto.ToRowser) pproto.ToRowser {
+			return CountRows(result)
+		})
 	}
 	mr := &MappingResult{
 		IndexName: indexName,
@@ -421,9 +412,10 @@ func (h handlerSelectCountFromTableWhere) Apply(stmt *sqlparser.Select, qm Query
 		Reducers:  reducers,
 	}
 
-	mr.addReducer(NewAssignHeadersReducer(selectFields))
-	// NOTE: limit and order by don't make sense in this handler
-	// because it just returns a single row.
+	// Assign headers to the result.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return AssignHeaders(result, selectFields...)
+	})
 
 	return mr, nil
 }
@@ -493,14 +485,28 @@ func (h handlerSelectFuncFromTableWhere) Apply(stmt *sqlparser.Select, qm QueryM
 		Query:     qo,
 	}
 
-	mr.addReducer(NewValCountFuncReducer(funcs[0].funcName))
-	mr.addReducer(NewAssignHeadersReducer(selectFields))
+	// Apply the ValCount function.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return ApplyValCountFunc(result, funcs[0].funcName)
+	})
+
+	// Assign headers to the result.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return AssignHeaders(result, selectFields...)
+	})
+
 	if qm.HasOrderBy() {
-		mr.addReducer(NewOrderByReducer(orderByFlds, orderByDirs, limit, offset))
-	} else {
-		mr.addReducer(NewLimitReducer(limit, offset))
+		// Sort the result.
+		mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+			return OrderBy(result, orderByFlds, orderByDirs)
+		})
 	}
 
+	// Apply a limit and offset to the result.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return LimitRows(OffsetRows(result, offset), limit)
+	})
+
 	return mr, nil
 }
 
@@ -627,13 +633,23 @@ func (h handlerSelectGroupBy) Apply(stmt *sqlparser.Select, qm QueryMask, indexF
 		Query:     qo,
 	}
 
-	mr.addReducer(NewAssignHeadersReducer(selectFields))
+	// Assign headers to the result.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return AssignHeaders(result, selectFields...)
+	})
+
 	if qm.HasOrderBy() {
-		mr.addReducer(NewOrderByReducer(orderByFlds, orderByDirs, limit, offset))
-	} else {
-		mr.addReducer(NewLimitReducer(limit, offset))
+		// Sort the result.
+		mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+			return OrderBy(result, orderByFlds, orderByDirs)
+		})
 	}
 
+	// Apply a limit and offset to the result.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return LimitRows(OffsetRows(result, offset), limit)
+	})
+
 	return mr, nil
 }
 
@@ -689,13 +705,21 @@ func (f handlerSelectIDCountFromTable) Apply(stmt *sqlparser.Select, qm QueryMas
 		Query:     qo,
 	}
 
-	mr.addReducer(NewAssignHeadersReducer(selectFields))
-	mr.addReducer(NewLimitReducer(limit, offset))
+	// Assign headers to the result.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return AssignHeaders(result, selectFields...)
+	})
+
 	// TODO: order by is not implemented on this method because order desc
 	// is handled in pilosa TopN. In order to support asc here, we would
 	// have to return the entire TopN cache. Instead, we should consider
 	// supported something like this in Pilosa itself.
 
+	// Apply a limit and offset to the result.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return LimitRows(OffsetRows(result, offset), limit)
+	})
+
 	return mr, nil
 }
 
@@ -784,12 +808,22 @@ func (h handlerSelectJoin) Apply(stmt *sqlparser.Select, qm QueryMask, indexFunc
 		Query:     qo,
 	}
 
-	mr.addReducer(NewAssignHeadersReducer(selectFields))
+	// Assign headers to the result.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return AssignHeaders(result, selectFields...)
+	})
+
 	if qm.HasOrderBy() {
-		mr.addReducer(NewOrderByReducer(orderByFlds, orderByDirs, limit, offset))
-	} else {
-		mr.addReducer(NewLimitReducer(limit, offset))
+		// Sort the result.
+		mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+			return OrderBy(result, orderByFlds, orderByDirs)
+		})
 	}
 
+	// Apply a limit and offset to the result.
+	mr.addReducer(func(result pproto.ToRowser) pproto.ToRowser {
+		return LimitRows(OffsetRows(result, offset), limit)
+	})
+
 	return mr, nil
 }

sql/show.go

@@ -37,7 +37,7 @@ func NewShowHandler(api *pilosa.API) *ShowHandler {
 }
 
 // Handle executes mapped SQL
-func (s *ShowHandler) Handle(ctx context.Context, mapped *MappedSQL) (pproto.StreamClient, error) {
+func (s *ShowHandler) Handle(ctx context.Context, mapped *MappedSQL) (pproto.ToRowser, error) {
 	stmt, ok := mapped.Statement.(*sqlparser.Show)
 	if !ok {
 		return nil, fmt.Errorf("statement is not type show: %T", mapped.Statement)
@@ -53,20 +53,12 @@ func (s *ShowHandler) Handle(ctx context.Context, mapped *MappedSQL) (pproto.Str
 	}
 }
 
-func (s *ShowHandler) execShowTables(ctx context.Context, showStmt *sqlparser.Show) (pproto.StreamClient, error) {
+func (s *ShowHandler) execShowTables(ctx context.Context, showStmt *sqlparser.Show) (pproto.ToRowser, error) {
 	indexInfo := s.api.Schema(ctx)
-	sz := len(indexInfo)
-	// If there aren't any indexes, don't bother creating
-	// a result row buffer.
-	if sz == 0 {
-		return pproto.EmptyStream{}, nil
-	}
 
-	// Create a buffer large enough to hold the entire result
-	// set. This way we don't have to use a goroutine.
-	result := pproto.NewRowBuffer(sz)
-	for _, ii := range indexInfo {
-		rr := &pproto.RowResponse{
+	result := make(pproto.ConstRowser, len(indexInfo))
+	for i, ii := range indexInfo {
+		result[i] = pproto.RowResponse{
 			Headers: []*pproto.ColumnInfo{
 				{Name: "Table", Datatype: "string"},
 			},
@@ -74,24 +66,13 @@ func (s *ShowHandler) execShowTables(ctx context.Context, showStmt *sqlparser.Sh
 				{ColumnVal: &pproto.ColumnResponse_StringVal{StringVal: ii.Name}},
 			},
 		}
-		if err := result.Send(rr); err != nil {
-			return nil, errors.Wrap(err, "sending row response")
-		}
-	}
-	if err := result.Send(pproto.EOF); err != nil {
-		return nil, errors.Wrap(err, "sending EOF")
 	}
 
-	// Apply Sort Reducer
-	out := pproto.NewRowBuffer(0)
-	red := NewOrderByReducer([]string{"Table"}, []string{"asc"}, 0, 0)
-	go red.Reduce(result, out) //nolint:errcheck
-
-	result = out
-	return result, nil
+	// Sort the result.
+	return OrderBy(result, []string{"Table"}, []string{"asc"}), nil
 }
 
-func (s *ShowHandler) execShowFields(ctx context.Context, showStmt *sqlparser.Show) (pproto.StreamClient, error) {
+func (s *ShowHandler) execShowFields(ctx context.Context, showStmt *sqlparser.Show) (pproto.ToRowser, error) {
 	indexName := showStmt.OnTable.ToViewName().Name.String()
 	index, err := s.api.Index(ctx, indexName)
 	if err != nil {
@@ -101,16 +82,8 @@ func (s *ShowHandler) execShowFields(ctx context.Context, showStmt *sqlparser.Sh
 		return nil, errors.WithMessage(pilosa.ErrIndexNotFound, indexName)
 	}
 	fields := index.Fields()
-	sz := len(fields)
-	// If there aren't any fields, don't bother creating
-	// a result row buffer.
-	if sz == 0 {
-		return pproto.EmptyStream{}, nil
-	}
 
-	// Create a buffer large enough to hold the entire result
-	// set. This way we don't have to use a goroutine.
-	result := pproto.NewRowBuffer(sz)
+	result := make(pproto.ConstRowser, 0, len(fields))
 	for _, f := range fields {
 		if f.Name() == "_exists" {
 			continue
@@ -120,7 +93,7 @@ func (s *ShowHandler) execShowFields(ctx context.Context, showStmt *sqlparser.Sh
 		if err != nil {
 			return nil, errors.Wrapf(err, "field %s", f.Name())
 		}
-		rr := &pproto.RowResponse{
+		result = append(result, pproto.RowResponse{
 			Headers: []*pproto.ColumnInfo{
 				{Name: "Field", Datatype: "string"},
 				{Name: "Type", Datatype: "string"},
@@ -129,20 +102,9 @@ func (s *ShowHandler) execShowFields(ctx context.Context, showStmt *sqlparser.Sh
 				{ColumnVal: &pproto.ColumnResponse_StringVal{StringVal: f.Name()}},
 				{ColumnVal: &pproto.ColumnResponse_StringVal{StringVal: dt}},
 			},
-		}
-		if err := result.Send(rr); err != nil {
-			return nil, errors.Wrap(err, "sending row response")
-		}
-	}
-	if err := result.Send(pproto.EOF); err != nil {
-		return nil, errors.Wrap(err, "sending EOF")
+		})
 	}
 
-	// Apply Sort Reducer
-	out := pproto.NewRowBuffer(0)
-	red := NewOrderByReducer([]string{"Field"}, []string{"asc"}, 0, 0)
-	go red.Reduce(result, out) //nolint:errcheck
-
-	result = out
-	return result, nil
+	// Sort the result.
+	return OrderBy(result, []string{"Field"}, []string{"asc"}), nil
 }