github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/libraries/doltcore/sqle/rows.go (about)

     1  // Copyright 2019 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package sqle
    16  
    17  import (
    18  	"context"
    19  
    20  	"github.com/dolthub/go-mysql-server/sql"
    21  
    22  	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
    23  	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
    24  	"github.com/dolthub/dolt/go/libraries/doltcore/table"
    25  	"github.com/dolthub/dolt/go/libraries/utils/set"
    26  	"github.com/dolthub/dolt/go/store/types"
    27  )
    28  
    29  // An iterator over the rows of a table.
    30  type doltTableRowIter struct {
    31  	sql.RowIter
    32  	ctx    context.Context
    33  	reader table.SqlTableReader
    34  }
    35  
    36  // Returns a new row iterator for the table given
    37  func newRowIterator(ctx *sql.Context, tbl *doltdb.Table, projCols []string, partition *doltTablePartition) (sql.RowIter, error) {
    38  	sch, err := tbl.GetSchema(ctx)
    39  
    40  	if err != nil {
    41  		return nil, err
    42  	}
    43  
    44  	if schema.IsKeyless(sch) {
    45  		// would be more optimal to project columns into keyless tables also
    46  		return newKeylessRowIterator(ctx, tbl, partition)
    47  	} else {
    48  		return newKeyedRowIter(ctx, tbl, projCols, partition)
    49  	}
    50  }
    51  
    52  func newKeylessRowIterator(ctx *sql.Context, tbl *doltdb.Table, partition *doltTablePartition) (*doltTableRowIter, error) {
    53  	var iter table.SqlTableReader
    54  	var err error
    55  	if partition.end == NoUpperBound {
    56  		iter, err = table.NewBufferedTableReader(ctx, tbl)
    57  	} else {
    58  		iter, err = table.NewBufferedTableReaderForPartition(ctx, tbl, partition.start, partition.end)
    59  	}
    60  
    61  	if err != nil {
    62  		return nil, err
    63  	}
    64  
    65  	return &doltTableRowIter{
    66  		ctx:    ctx,
    67  		reader: iter,
    68  	}, nil
    69  }
    70  
    71  func newKeyedRowIter(ctx context.Context, tbl *doltdb.Table, projectedCols []string, partition *doltTablePartition) (sql.RowIter, error) {
    72  	var err error
    73  	var mapIter types.MapTupleIterator
    74  	rowData := partition.rowData
    75  	if partition.end == NoUpperBound {
    76  		mapIter, err = rowData.RangeIterator(ctx, 0, rowData.Len())
    77  	} else {
    78  		mapIter, err = partition.IteratorForPartition(ctx, rowData)
    79  	}
    80  
    81  	if err != nil {
    82  		return nil, err
    83  	}
    84  
    85  	sch, err := tbl.GetSchema(ctx)
    86  	if err != nil {
    87  		return nil, err
    88  	}
    89  
    90  	cols := sch.GetAllCols().GetColumns()
    91  	tagToSqlColIdx := make(map[uint64]int)
    92  
    93  	resultColSet := set.NewCaseInsensitiveStrSet(projectedCols)
    94  	for i, col := range cols {
    95  		if len(projectedCols) == 0 || resultColSet.Contains(col.Name) {
    96  			tagToSqlColIdx[col.Tag] = i
    97  		}
    98  	}
    99  
   100  	conv := NewKVToSqlRowConverter(tbl.Format(), tagToSqlColIdx, cols, len(cols))
   101  	return NewDoltMapIter(ctx, mapIter.NextTuple, nil, conv), nil
   102  }
   103  
   104  // Next returns the next row in this row iterator, or an io.EOF error if there aren't any more.
   105  func (itr *doltTableRowIter) Next() (sql.Row, error) {
   106  	return itr.reader.ReadSqlRow(itr.ctx)
   107  }
   108  
   109  // Close required by sql.RowIter interface
   110  func (itr *doltTableRowIter) Close(*sql.Context) error {
   111  	return nil
   112  }