Skip to content

Commit

Permalink
testing progress
Browse files Browse the repository at this point in the history
  • Loading branch information
@max-hoffman
max-hoffman committed Jun 26, 2024
1 parent 9ef0aee commit ef8cc8b
Show file tree
Hide file tree
Showing 4 changed files with 268 additions and 31 deletions.
164 changes: 164 additions & 0 deletions go/libraries/doltcore/sqle/index/index_lookup.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,11 @@ import (
"github.com/dolthub/dolt/go/store/val"
)

func ProllyRangesForIndex(ctx *sql.Context, index sql.Index, ranges sql.RangeCollection) ([]prolly.Range, error) {
idx := index.(*doltIndex)
return idx.prollyRanges(ctx, idx.ns, ranges...)
}

func RowIterForIndexLookup(ctx *sql.Context, t DoltTableable, lookup sql.IndexLookup, pkSch sql.PrimaryKeySchema, columns []uint64) (sql.RowIter, error) {
idx := lookup.Index.(*doltIndex)
durableState, err := idx.getDurableState(ctx, t)
Expand Down Expand Up @@ -226,6 +231,7 @@ func (rp rangePartition) Key() []byte {
type LookupBuilder interface {
// NewRowIter returns a new index iter for the given partition
NewRowIter(ctx *sql.Context, part sql.Partition) (sql.RowIter, error)
NewMapIter(ctx context.Context, r prolly.Range, reverse bool) (prolly.MapIter, error)
Key() doltdb.DataCacheKey
}

Expand Down Expand Up @@ -357,6 +363,11 @@ func (lb *baseLookupBuilder) NewRowIter(ctx *sql.Context, part sql.Partition) (s
panic("cannot call NewRowIter on baseLookupBuilder")
}

// NewMapIter implements IndexLookup
func (lb *baseLookupBuilder) NewMapIter(ctx context.Context, r prolly.Range, reverse bool) (prolly.MapIter, error) {
panic("cannot call NewMapIter on baseLookupBuilder")
}

// newPointLookup will create a cursor once, and then use the same cursor for
// every subsequent point lookup. Note that equality joins can have a mix of
// point lookups on concrete values, and range lookups for null matches.
Expand Down Expand Up @@ -398,6 +409,51 @@ type coveringLookupBuilder struct {
keyMap, valMap, ordMap val.OrdinalMapping
}

type sequenceMapIter struct {
cur prolly.MapIter
lb LookupBuilder
reverse bool
left []prolly.Range
}

func NewSequenceMapIter(ctx context.Context, lb LookupBuilder, ranges []prolly.Range, reverse bool) (prolly.MapIter, error) {
cur, err := lb.NewMapIter(ctx, ranges[0], reverse)
if err != nil || len(ranges) < 2 {
return cur, err
}
return &sequenceMapIter{
cur: cur,
lb: lb,
reverse: reverse,
left: ranges[1:],
}, nil

}
func (i *sequenceMapIter) Next(ctx context.Context) (val.Tuple, val.Tuple, error) {
k, v, err := i.cur.Next(ctx)
if err == io.EOF {
if len(i.left) == 0 {
return nil, nil, io.EOF
}
i.cur, err = i.lb.NewMapIter(ctx, i.left[0], i.reverse)
if err != nil {
return nil, nil, err
}
i.left = i.left[1:]
return i.Next(ctx)
}
return k, v, nil
}

// NewMapIter implements IndexLookup
func (lb *coveringLookupBuilder) NewMapIter(ctx context.Context, r prolly.Range, reverse bool) (prolly.MapIter, error) {
if reverse {
return lb.sec.IterRangeReverse(ctx, r)
} else {
return lb.sec.IterRange(ctx, r)
}
}

// NewRowIter implements IndexLookup
func (lb *coveringLookupBuilder) NewRowIter(ctx *sql.Context, part sql.Partition) (sql.RowIter, error) {
rangeIter, err := lb.rangeIter(ctx, part)
Expand Down Expand Up @@ -431,6 +487,55 @@ type nonCoveringLookupBuilder struct {
pkMap, keyMap, valMap, ordMap val.OrdinalMapping
}

type nonCoveringMapIter struct {
idx DoltIndex
indexIter prolly.MapIter
primary prolly.Map
pkMap val.OrdinalMapping
pkBld *val.TupleBuilder
}

func (i *nonCoveringMapIter) Next(ctx context.Context) (val.Tuple, val.Tuple, error) {
idxKey, _, err := i.indexIter.Next(ctx)
if err != nil {
return nil, nil, err
}
for to := range i.pkMap {
from := i.pkMap.MapOrdinal(to)
i.pkBld.PutRaw(to, idxKey.GetField(from))
}
pk := i.pkBld.Build(sharePool)

var value val.Tuple
err = i.primary.Get(ctx, pk, func(_, v val.Tuple) error {
value = v
return nil
})
return pk, value, nil
}

// NewMapIter implements IndexLookup
func (lb *nonCoveringLookupBuilder) NewMapIter(ctx context.Context, r prolly.Range, reverse bool) (prolly.MapIter, error) {
var secIter prolly.MapIter
var err error
if reverse {
secIter, err = lb.sec.IterRangeReverse(ctx, r)
} else {
secIter, err = lb.sec.IterRange(ctx, r)
}
if err != nil {
return nil, err
}

return &nonCoveringMapIter{
idx: lb.idx,
indexIter: secIter,
primary: lb.pri,
pkBld: lb.pkBld,
pkMap: lb.pkMap,
}, nil
}

// NewRowIter implements IndexLookup
func (lb *nonCoveringLookupBuilder) NewRowIter(ctx *sql.Context, part sql.Partition) (sql.RowIter, error) {
rangeIter, err := lb.rangeIter(ctx, part)
Expand Down Expand Up @@ -459,6 +564,60 @@ type keylessLookupBuilder struct {
s *durableIndexState
}

// NewMapIter implements IndexLookup
func (lb *keylessLookupBuilder) NewMapIter(ctx context.Context, r prolly.Range, reverse bool) (prolly.MapIter, error) {
rows := lb.s.Primary
dsecondary := lb.s.Secondary
secondary := durable.ProllyMapFromIndex(dsecondary)
indexIter, err := secondary.IterRange(ctx, r)
if err != nil {
return nil, err
}
clustered := durable.ProllyMapFromIndex(rows)
keyDesc := clustered.KeyDesc()
indexMap := ordinalMappingFromIndex(lb.idx)

keyBld := val.NewTupleBuilder(keyDesc)
return &keylessMapIter{
indexIter: indexIter,
clustered: clustered,
clusteredMap: indexMap,
clusteredBld: keyBld,
}, nil
}

type keylessMapIter struct {
indexIter prolly.MapIter
clustered prolly.Map
// clusteredMap transforms secondary index keys
// into clustered index keys
clusteredMap val.OrdinalMapping
clusteredBld *val.TupleBuilder
}

func (i *keylessMapIter) Next(ctx context.Context) (val.Tuple, val.Tuple, error) {
idxKey, _, err := i.indexIter.Next(ctx)
if err != nil {
return nil, nil, err
}

for to := range i.clusteredMap {
from := i.clusteredMap.MapOrdinal(to)
i.clusteredBld.PutRaw(to, idxKey.GetField(from))
}
pk := i.clusteredBld.Build(sharePool)

var value val.Tuple
err = i.clustered.Get(ctx, pk, func(k, v val.Tuple) error {
value = v
return nil
})
if err != nil {
return nil, nil, err
}
return pk, value, nil
}

// NewRowIter implements IndexLookup
func (lb *keylessLookupBuilder) NewRowIter(ctx *sql.Context, part sql.Partition) (sql.RowIter, error) {
var prollyRange prolly.Range
Expand All @@ -483,6 +642,11 @@ func (lb *nomsLookupBuilder) NewRowIter(ctx *sql.Context, part sql.Partition) (s
return RowIterForNomsRanges(ctx, lb.idx, ranges, lb.projections, lb.s)
}

// NewRowIter implements IndexLookup
func (lb *nomsLookupBuilder) NewMapIter(ctx context.Context, r prolly.Range, reverse bool) (prolly.MapIter, error) {
panic("cannot call NewMapIter on *]nomsLookupBuilder")
}

// boundsCase determines the case upon which the bounds are tested.
type boundsCase byte

Expand Down
17 changes: 17 additions & 0 deletions go/libraries/doltcore/sqle/indexed_dolt_table.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,23 @@ type WritableIndexedDoltTable struct {
mu *sync.Mutex
}

func (t *WritableIndexedDoltTable) Index() index.DoltIndex {
return t.idx
}

func (t *WritableIndexedDoltTable) LookupBuilder(ctx *sql.Context) (index.LookupBuilder, error) {
t.mu.Lock()
defer t.mu.Unlock()
key, canCache, err := t.DataCacheKey(ctx)
if err != nil {
return nil, err
}
if t.lb == nil || !canCache || t.lb.Key() != key {
return index.NewLookupBuilder(ctx, t.DoltTable, t.idx, key, t.DoltTable.projectedCols, t.DoltTable.sqlSch, t.isDoltFormat)
}
return t.lb, nil
}

func (t *WritableIndexedDoltTable) LookupPartitions(ctx *sql.Context, lookup sql.IndexLookup) (sql.PartitionIter, error) {
return index.NewRangePartitionIter(ctx, t.DoltTable, lookup, t.isDoltFormat)
}
Expand Down
108 changes: 84 additions & 24 deletions go/libraries/doltcore/sqle/rowexec/builder.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,12 @@ package rowexec
import (
"context"
"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
"github.com/dolthub/dolt/go/libraries/doltcore/sqle/index"
"github.com/dolthub/dolt/go/store/prolly/tree"
"github.com/dolthub/dolt/go/store/val"
"github.com/dolthub/go-mysql-server/sql"
"github.com/dolthub/go-mysql-server/sql/plan"
"strings"

"github.com/dolthub/dolt/go/libraries/doltcore/doltdb/durable"
"github.com/dolthub/dolt/go/libraries/doltcore/schema"
Expand All @@ -23,9 +25,9 @@ func (b Builder) Build(ctx *sql.Context, n sql.Node, r sql.Row) (sql.RowIter, er
case *plan.JoinNode:
if n.Op.IsLookup() {
if ita, ok := getIta(n.Right()); ok {
if dstIter, dstSchema, dstTags, dstFilter, err := getSourceKvIter(ctx, n.Right()); err == nil && dstSchema != nil {
if srcIter, srcSchema, srcTags, srcFilter, err := getSourceKvIter(ctx, n.Left()); err == nil && dstSchema != nil {
return rowIterTableLookupJoin(ctx, srcIter, dstIter, srcSchema, dstSchema, srcTags, dstTags, ita.Expressions(), srcFilter, dstFilter)
if dstMap, _, dstSchema, dstTags, dstFilter, err := getSourceKv(ctx, n.Right(), false); err == nil && dstSchema != nil {
if srcMap, srcIter, srcSchema, srcTags, srcFilter, err := getSourceKv(ctx, n.Left(), true); err == nil && srcSchema != nil {
return rowIterTableLookupJoin(ctx, srcIter, srcMap, dstMap, srcSchema, dstSchema, srcTags, dstTags, ita.Expressions(), srcFilter, dstFilter)
}
}
}
Expand Down Expand Up @@ -174,56 +176,114 @@ func getMap(ctx *sql.Context, dt *sqle.DoltTable) (prolly.Map, schema.Schema, er
return durable.ProllyMapFromIndex(priIndex), sch, nil
}

func getSourceKvIter(ctx *sql.Context, n sql.Node) (prolly.Map, schema.Schema, []uint64, sql.Expression, error) {
func getSourceKv(ctx *sql.Context, n sql.Node, iter bool) (prolly.Map, prolly.MapIter, schema.Schema, []uint64, sql.Expression, error) {
var table *doltdb.Table
var tags []uint64
var err error
var indexMap prolly.Map
var mapIter prolly.MapIter
switch n := n.(type) {
case *plan.TableAlias:
return getSourceKvIter(ctx, n.Child)
return getSourceKv(ctx, n.Child, iter)
case *plan.Filter:
m, s, t, _, err := getSourceKvIter(ctx, n.Child)
m, mIter, s, t, _, err := getSourceKv(ctx, n.Child, iter)
if err != nil {
return prolly.Map{}, nil, nil, nil, err
return prolly.Map{}, nil, nil, nil, nil, err
}
return m, s, t, n.Expression, nil
case *plan.ResolvedTable, *plan.IndexedTableAccess:
switch dt := n.(sql.TableNode).UnderlyingTable().(type) {
case *sqle.WritableDoltTable:
tags = dt.ProjectedTags()
table, err = dt.DoltTable.DoltTable(ctx)
return m, mIter, s, t, n.Expression, nil
case *plan.IndexedTableAccess:
switch dt := n.UnderlyingTable().(type) {
case *sqle.WritableIndexedDoltTable:
tags = dt.ProjectedTags()
table, err = dt.DoltTable.DoltTable(ctx)
if err != nil {
return prolly.Map{}, nil, nil, nil, nil, err
}

var rowData durable.Index
switch strings.ToLower(dt.Index().ID()) {
case "primary":
rowData, err = table.GetRowData(ctx)
default:
rowData, err = table.GetIndexRowData(ctx, dt.Index().ID())
}
if err != nil {
return prolly.Map{}, nil, nil, nil, nil, err
}
indexMap = durable.ProllyMapFromIndex(rowData)
if iter {
l, err := n.GetLookup(ctx, nil)
if err != nil {
return prolly.Map{}, nil, nil, nil, nil, err
}
prollyRanges, err := index.ProllyRangesForIndex(ctx, l.Index, l.Ranges)
if err != nil {
return prolly.Map{}, nil, nil, nil, nil, err
}

lb, err := dt.LookupBuilder(ctx)
if err != nil {
return prolly.Map{}, nil, nil, nil, nil, err
}

mapIter, err = index.NewSequenceMapIter(ctx, lb, prollyRanges, l.IsReverse)
if err != nil {
return prolly.Map{}, nil, nil, nil, nil, err
}
}

case *sqle.IndexedDoltTable:
tags = dt.ProjectedTags()
table, err = dt.DoltTable.DoltTable(ctx)
if err != nil {
return prolly.Map{}, nil, nil, nil, nil, err
}
rowData, err := table.GetIndexRowData(ctx, dt.Index().ID())
if err != nil {
return prolly.Map{}, nil, nil, nil, nil, err
}
indexMap = durable.ProllyMapFromIndex(rowData)
}
case *plan.ResolvedTable:
switch dt := n.UnderlyingTable().(type) {
case *sqle.WritableDoltTable:
tags = dt.ProjectedTags()
table, err = dt.DoltTable.DoltTable(ctx)
case *sqle.AlterableDoltTable:
tags = dt.ProjectedTags()
table, err = dt.DoltTable.DoltTable(ctx)
case *sqle.DoltTable:
tags = dt.ProjectedTags()
table, err = dt.DoltTable(ctx)
default:
return prolly.Map{}, nil, nil, nil, nil
return prolly.Map{}, nil, nil, nil, nil, nil
}
if err != nil {
return prolly.Map{}, nil, nil, nil, nil, err
}
priIndex, err := table.GetRowData(ctx)
if err != nil {
return prolly.Map{}, nil, nil, nil, nil, err
}
indexMap = durable.ProllyMapFromIndex(priIndex)

if iter {
mapIter, err = indexMap.IterAll(ctx)
if err != nil {
return prolly.Map{}, nil, nil, nil, nil, err
}
}
default:
return prolly.Map{}, nil, nil, nil, nil
return prolly.Map{}, nil, nil, nil, nil, nil
}
if err != nil {
return prolly.Map{}, nil, nil, nil, err
}

priIndex, err := table.GetRowData(ctx)
if err != nil {
return prolly.Map{}, nil, nil, nil, err
return prolly.Map{}, nil, nil, nil, nil, err
}

m := durable.ProllyMapFromIndex(priIndex)
sch, err := table.GetSchema(ctx)
if err != nil {
return prolly.Map{}, nil, nil, nil, err
return prolly.Map{}, nil, nil, nil, nil, err
}

return m, sch, tags, nil, nil
return indexMap, mapIter, sch, tags, nil, nil
}

0 comments on commit ef8cc8b

Please sign in to comment.