github.com/cayleygraph/cayley@v0.7.7/graph/kv/iterators.go (about)

     1  package kv
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	"github.com/hidal-go/hidalgo/kv"
     8  
     9  	"github.com/cayleygraph/cayley/graph"
    10  	"github.com/cayleygraph/cayley/graph/iterator"
    11  	"github.com/cayleygraph/cayley/graph/shape"
    12  	"github.com/cayleygraph/quad"
    13  )
    14  
    15  func (qs *QuadStore) NodesAllIterator() graph.Iterator {
    16  	return NewAllIterator(true, qs, nil)
    17  }
    18  
    19  func (qs *QuadStore) QuadsAllIterator() graph.Iterator {
    20  	return NewAllIterator(false, qs, nil)
    21  }
    22  
    23  func (qs *QuadStore) indexSize(ctx context.Context, ind QuadIndex, vals []uint64) (graph.Size, error) {
    24  	var sz int64
    25  	err := kv.View(qs.db, func(tx kv.Tx) error {
    26  		val, err := tx.Get(ctx, ind.Key(vals))
    27  		if err != nil {
    28  			return err
    29  		}
    30  		sz, err = countIndex(val)
    31  		if err != nil {
    32  			return err
    33  		}
    34  		return nil
    35  	})
    36  	if err != nil {
    37  		return graph.Size{}, err
    38  	}
    39  	if len(ind.Dirs) == len(vals) {
    40  		return graph.Size{
    41  			Size:  sz,
    42  			Exact: true,
    43  		}, nil
    44  	}
    45  	return graph.Size{
    46  		Size:  1 + sz/2,
    47  		Exact: false,
    48  	}, nil
    49  }
    50  
    51  func (qs *QuadStore) QuadIteratorSize(ctx context.Context, d quad.Direction, v graph.Ref) (graph.Size, error) {
    52  	vi, ok := v.(Int64Value)
    53  	if !ok {
    54  		return graph.Size{Size: 0, Exact: true}, nil
    55  	}
    56  	qs.indexes.RLock()
    57  	all := qs.indexes.all
    58  	qs.indexes.RUnlock()
    59  	for _, ind := range all {
    60  		if len(ind.Dirs) == 1 && ind.Dirs[0] == d {
    61  			return qs.indexSize(ctx, ind, []uint64{uint64(vi)})
    62  		}
    63  	}
    64  	st, err := qs.Stats(ctx, false)
    65  	if err != nil {
    66  		return graph.Size{}, err
    67  	}
    68  	return graph.Size{
    69  		Size:  1 + st.Quads.Size/2,
    70  		Exact: false,
    71  	}, nil
    72  }
    73  
    74  func (qs *QuadStore) QuadIterator(dir quad.Direction, v graph.Ref) graph.Iterator {
    75  	if v == nil {
    76  		return iterator.NewNull()
    77  	}
    78  	vi, ok := v.(Int64Value)
    79  	if !ok {
    80  		return iterator.NewError(fmt.Errorf("unexpected node type: %T", v))
    81  	}
    82  	// Find the best index for this direction.
    83  	if ind := qs.bestIndexes([]quad.Direction{dir}); len(ind) == 1 {
    84  		// this will scan the prefix automatically
    85  		return NewQuadIterator(qs, ind[0], []uint64{uint64(vi)})
    86  	}
    87  	// Fallback: iterate all quads and check the corresponding direction.
    88  	return NewAllIterator(false, qs, &constraint{
    89  		dir: dir,
    90  		val: vi,
    91  	})
    92  }
    93  
    94  func (qs *QuadStore) OptimizeShape(s shape.Shape) (shape.Shape, bool) {
    95  	switch s := s.(type) {
    96  	case shape.QuadsAction:
    97  		return qs.optimizeQuadsAction(s)
    98  	}
    99  	return s, false
   100  }
   101  
   102  func (qs *QuadStore) optimizeQuadsAction(s shape.QuadsAction) (shape.Shape, bool) {
   103  	if len(s.Filter) == 0 {
   104  		return s, false
   105  	}
   106  	dirs := make([]quad.Direction, 0, len(s.Filter))
   107  	for d := range s.Filter {
   108  		dirs = append(dirs, d)
   109  	}
   110  	ind := qs.bestIndexes(dirs)
   111  	if len(ind) != 1 {
   112  		return s, false // TODO(dennwc): allow intersecting indexes
   113  	}
   114  	quads := IndexScan{Index: ind[0]}
   115  	for _, d := range ind[0].Dirs {
   116  		v, ok := s.Filter[d].(Int64Value)
   117  		if !ok {
   118  			return s, false
   119  		}
   120  		quads.Values = append(quads.Values, uint64(v))
   121  	}
   122  	return s.SimplifyFrom(quads), true
   123  }
   124  
   125  type IndexScan struct {
   126  	Index  QuadIndex
   127  	Values []uint64
   128  }
   129  
   130  func (s IndexScan) BuildIterator(qs graph.QuadStore) graph.Iterator {
   131  	kqs, ok := qs.(*QuadStore)
   132  	if !ok {
   133  		return iterator.NewError(fmt.Errorf("expected KV quadstore, got: %T", qs))
   134  	}
   135  	return NewQuadIterator(kqs, s.Index, s.Values)
   136  }
   137  
   138  func (s IndexScan) Optimize(r shape.Optimizer) (shape.Shape, bool) {
   139  	return s, false
   140  }