covering indexes

dolthub · bheni · Dec 16, 2020 · Dec 11, 2020 · Dec 11, 2020 · Dec 11, 2020
commit be2c221df9a14dcb58d710f74b4375e38cb6d4de
diff --git a/go/libraries/doltcore/sqle/index_lookup.go b/go/libraries/doltcore/sqle/index_lookup.go
@@ -130,13 +130,40 @@ func (il *doltIndexLookup) Union(indexLookups ...sql.IndexLookup) (sql.IndexLook
 
 // RowIter returns a row iterator for this index lookup. The iterator will return the single matching row for the index.
 func (il *doltIndexLookup) RowIter(ctx *sql.Context) (sql.RowIter, error) {
+	return il.RowIterWithProjection(ctx, nil)
+}
+
+func (il *doltIndexLookup) RowIterWithProjection(ctx *sql.Context, columns []string) (sql.RowIter, error) {
 	readRanges := make([]*noms.ReadRange, len(il.ranges))
 	for i, lookupRange := range il.ranges {
 		readRanges[i] = lookupRange.ToReadRange()
 	}
-	return NewIndexLookupRowIterAdapter(ctx, il.idx, &doltIndexKeyIter{
-		indexMapIter: noms.NewNomsRangeReader(il.idx.IndexSchema(), il.idx.IndexRowData(), readRanges),
-	}), nil
+
+	idxItr := &doltIndexKeyIter{indexMapIter: noms.NewNomsRangeReader(il.idx.IndexSchema(), il.idx.IndexRowData(), readRanges)}
+
+	covers := il.indexCoversCols(columns)
+	if covers {
+		return NewCoveringIndexRowIterAdapter(ctx, il.idx, idxItr), nil
+	} else {
+		return NewIndexLookupRowIterAdapter(ctx, il.idx, idxItr), nil
+	}
+}
+
+func (il *doltIndexLookup) indexCoversCols(cols []string) bool {
+	if cols == nil {
+		return false
+	}
+
+	idxCols := il.idx.IndexSchema().GetAllCols()
+	covers := true
+	for _, colName := range cols {
+		if _, ok := idxCols.GetByName(colName); !ok {
+			covers = false
+			break
+		}
+	}
+
+	return covers
 }
 
 type doltIndexKeyIter struct {

diff --git a/go/libraries/doltcore/sqle/index_row_iter.go b/go/libraries/doltcore/sqle/index_row_iter.go
@@ -16,6 +16,7 @@ package sqle
 
 import (
 	"context"
+	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
 	"io"
 	"runtime"
 
@@ -150,3 +151,57 @@ func (i *indexLookupRowIterAdapter) processKey(_ context.Context, valInt interfa
 	i.buffer[val.position] = sqlRow
 	return nil
 }
+
+type coveringIndexRowIterAdapter struct {
+	idx       DoltIndex
+	keyIter   IndexLookupKeyIterator
+	ctx       *sql.Context
+	pkCols	  *schema.ColCollection
+	nonPKCols *schema.ColCollection
+	nbf       *types.NomsBinFormat
+}
+
+func NewCoveringIndexRowIterAdapter(ctx *sql.Context, idx DoltIndex, keyIter IndexLookupKeyIterator) *coveringIndexRowIterAdapter {
+	sch := idx.Schema()
+	return &coveringIndexRowIterAdapter{
+		idx:     idx,
+		keyIter: keyIter,
+		ctx:     ctx,
+		pkCols:  sch.GetPKCols(),
+		nonPKCols: sch.GetNonPKCols(),
+		nbf:     idx.TableData().Format(),
+	}
+}
+
+// Next returns the next row from the iterator.
+func (ci *coveringIndexRowIterAdapter) Next() (sql.Row, error) {
+	taggedVals, err := ci.keyIter.NextKey(ci.ctx)
+
+	if err != nil {
+		return nil, err
+	}
+
+	pk, err := taggedVals.NomsTupleForPKCols(ci.nbf, ci.pkCols).Value(ci.ctx)
+
+	if err != nil {
+		return nil, err
+	}
+
+	nonPK, err := taggedVals.NomsTupleForNonPKCols(ci.nbf, ci.nonPKCols).Value(ci.ctx)
+
+	if err != nil {
+		return nil, err
+	}
+
+	r, err := row.FromNoms(ci.idx.Schema(), pk.(types.Tuple), nonPK.(types.Tuple))
+
+	if err != nil {
+		return nil, err
+	}
+
+	return sqlutil.DoltRowToSqlRow(r, ci.idx.Schema())
+}
+
+func (ci *coveringIndexRowIterAdapter) Close() error {
+	return nil
+}
diff --git a/go/libraries/doltcore/sqle/indexed_dolt_table.go b/go/libraries/doltcore/sqle/indexed_dolt_table.go
@@ -62,17 +62,28 @@ func (idt *IndexedDoltTable) PartitionRows(ctx *sql.Context, _ sql.Partition) (s
 type WritableIndexedDoltTable struct {
 	*WritableDoltTable
 	indexLookup *doltIndexLookup
+	projectedCols []string
 }
 
 var _ sql.IndexedTable = (*WritableIndexedDoltTable)(nil)
 var _ sql.UpdatableTable = (*WritableIndexedDoltTable)(nil)
 var _ sql.DeletableTable = (*WritableIndexedDoltTable)(nil)
 var _ sql.ReplaceableTable = (*WritableIndexedDoltTable)(nil)
+var _ sql.ProjectedTable = (*WritableIndexedDoltTable)(nil)
 
 func (t *WritableIndexedDoltTable) Partitions(ctx *sql.Context) (sql.PartitionIter, error) {
 	return sqlutil.NewSinglePartitionIter(), nil
 }
 
 func (t *WritableIndexedDoltTable) PartitionRows(ctx *sql.Context, _ sql.Partition) (sql.RowIter, error) {
-	return t.indexLookup.RowIter(ctx)
+	return t.indexLookup.RowIterWithProjection(ctx, t.projectedCols)
+}
+
+func (t *WritableIndexedDoltTable) WithProjection(colNames []string) sql.Table {
+	t.projectedCols = colNames
+	return t
+}
+
+func (t *WritableIndexedDoltTable) Projection() []string {
+	return t.projectedCols
 }
diff --git a/go/libraries/doltcore/sqle/sqlutil/convert.go b/go/libraries/doltcore/sqle/sqlutil/convert.go
@@ -183,3 +183,13 @@ func ToDoltCol(tag uint64, col *sql.Column) (schema.Column, error) {
 
 	return schema.NewColumnWithTypeInfo(col.Name, tag, typeInfo, col.PrimaryKey, col.Default.String(), col.AutoIncrement, col.Comment, constraints...)
 }
+
+func GetColNamesFromSqlSchema(sqlSch sql.Schema) []string {
+	colNames := make([]string, len(sqlSch))
+
+	for i, col := range sqlSch {
+		colNames[i] = col.Name
+	}
+
+	return colNames
+}
diff --git a/go/libraries/doltcore/sqle/tables.go b/go/libraries/doltcore/sqle/tables.go
@@ -270,6 +270,7 @@ func (t *WritableDoltTable) WithIndexLookup(lookup sql.IndexLookup) sql.Table {
 	return &WritableIndexedDoltTable{
 		WritableDoltTable: t,
 		indexLookup:       dil,
+		projectedCols:     sqlutil.GetColNamesFromSqlSchema(t.sqlSch),
 	}
 }