From d421558f5838128df60406c3f9fac74d16a55e64 Mon Sep 17 00:00:00 2001 From: Nia Weiss Date: Mon, 31 Aug 2020 11:49:50 -0400 Subject: [PATCH] fix SQL memory leak --- proto/interface.go | 111 ++--------- server/grpc.go | 27 +-- server/pg.go | 25 --- server/sql.go | 4 +- sql/ddl.go | 6 +- sql/reduce.go | 452 ++++++++++++++++++--------------------------- sql/select.go | 192 +++++++++++-------- sql/show.go | 64 ++----- 8 files changed, 343 insertions(+), 538 deletions(-) diff --git a/proto/interface.go b/proto/interface.go index 6a6df7213..af2599707 100644 --- a/proto/interface.go +++ b/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 +} diff --git a/server/grpc.go b/server/grpc.go index 6ed428772..2c55bc9e0 100644 --- a/server/grpc.go +++ b/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. diff --git a/server/pg.go b/server/pg.go index 964472ece..49f697606 100644 --- a/server/pg.go +++ b/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", diff --git a/server/sql.go b/server/sql.go index 7647e0377..26686c561 100644 --- a/server/sql.go +++ b/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) diff --git a/sql/ddl.go b/sql/ddl.go index 3b693850e..fd098f6c5 100644 --- a/sql/ddl.go +++ b/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 } diff --git a/sql/reduce.go b/sql/reduce.go index b349f8ceb..e670d1220 100644 --- a/sql/reduce.go +++ b/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 ( diff --git a/sql/select.go b/sql/select.go index a195a7c85..8d125a96c 100644 --- a/sql/select.go +++ b/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 } diff --git a/sql/show.go b/sql/show.go index d145344d2..062a80d4a 100644 --- a/sql/show.go +++ b/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 }