github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/libraries/doltcore/doltdb/key_itr.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 doltdb
    16  
    17  import (
    18  	"context"
    19  
    20  	"github.com/dolthub/dolt/go/libraries/doltcore/row"
    21  	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
    22  	"github.com/dolthub/dolt/go/store/types"
    23  )
    24  
    25  // PKItr defines a function that iterates over a collection of noms values.  The PKItr will return a valid value
    26  // and true until all the values in the collection are exhausted.  At that time nil and false will be returned.
    27  type PKItr func() (val types.Tuple, ok bool, err error)
    28  
    29  func SingleColPKItr(nbf *types.NomsBinFormat, pkTag uint64, vals []types.Value) func() (types.Tuple, bool, error) {
    30  	next := 0
    31  	size := len(vals)
    32  	return func() (types.Tuple, bool, error) {
    33  		current := next
    34  		next++
    35  
    36  		if current < size {
    37  			tpl, err := types.NewTuple(nbf, types.Uint(pkTag), vals[current])
    38  
    39  			if err != nil {
    40  				return types.EmptyTuple(nbf), false, err
    41  			}
    42  
    43  			return tpl, true, nil
    44  		}
    45  
    46  		return types.EmptyTuple(nbf), false, nil
    47  	}
    48  }
    49  
    50  func TaggedValueSliceItr(nbf *types.NomsBinFormat, sch schema.Schema, vals []row.TaggedValues) func() (types.Tuple, bool, error) {
    51  	next := 0
    52  	size := len(vals)
    53  	return func() (types.Tuple, bool, error) {
    54  		current := next
    55  		next++
    56  
    57  		if current < size {
    58  			tpl := vals[current].NomsTupleForPKCols(nbf, sch.GetPKCols())
    59  			v, err := tpl.Value(context.TODO())
    60  
    61  			if err != nil {
    62  				return types.EmptyTuple(nbf), false, err
    63  			}
    64  
    65  			return v.(types.Tuple), true, nil
    66  		}
    67  
    68  		return types.EmptyTuple(nbf), false, nil
    69  	}
    70  }
    71  
    72  // TupleSliceItr returns a closure that has the signature of a PKItr and can be used to iterate over a slice of values
    73  func TupleSliceItr(nbf *types.NomsBinFormat, vals []types.Tuple) func() (types.Tuple, bool, error) {
    74  	next := 0
    75  	size := len(vals)
    76  	return func() (types.Tuple, bool, error) {
    77  		current := next
    78  		next++
    79  
    80  		if current < size {
    81  			return vals[current], true, nil
    82  		}
    83  
    84  		return types.EmptyTuple(nbf), false, nil
    85  	}
    86  }
    87  
    88  // SetItr returns a closure that has the signature of a PKItr and can be used to iterate over a noms Set of vaules
    89  func SetItr(ctx context.Context, valSet types.Set) (func() (types.Tuple, bool, error), error) {
    90  	itr, err := valSet.Iterator(ctx)
    91  
    92  	if err != nil {
    93  		return nil, err
    94  	}
    95  
    96  	return func() (types.Tuple, bool, error) {
    97  		// TODO: Should this be a `ctx` from the iter call?
    98  		v, err := itr.Next(ctx)
    99  
   100  		if err != nil {
   101  			return types.EmptyTuple(valSet.Format()), false, err
   102  		}
   103  
   104  		return v.(types.Tuple), v != nil, nil
   105  	}, nil
   106  }