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

     1  // Copyright 2020 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  	"errors"
    20  	"io"
    21  
    22  	"github.com/dolthub/go-mysql-server/sql"
    23  
    24  	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
    25  	"github.com/dolthub/dolt/go/store/types"
    26  )
    27  
    28  func maxU64(x, y uint64) uint64 {
    29  	if x > y {
    30  		return x
    31  	}
    32  
    33  	return y
    34  }
    35  
    36  // KVToSqlRowConverter takes noms types.Value key value pairs and converts them directly to a sql.Row.  It
    37  // can be configured to only process a portion of the columns and map columns to desired output columns.
    38  type KVToSqlRowConverter struct {
    39  	nbf            *types.NomsBinFormat
    40  	cols           []schema.Column
    41  	tagToSqlColIdx map[uint64]int
    42  	// rowSize is the number of columns in the output row.  This may be bigger than the number of columns being converted,
    43  	// but not less.  When rowSize is bigger than the number of columns being processed that means that some of the columns
    44  	// in the output row will be filled with nils
    45  	rowSize     int
    46  	valsFromKey int
    47  	valsFromVal int
    48  	maxValTag   uint64
    49  }
    50  
    51  func NewKVToSqlRowConverter(nbf *types.NomsBinFormat, tagToSqlColIdx map[uint64]int, cols []schema.Column, rowSize int) *KVToSqlRowConverter {
    52  	valsFromKey, valsFromVal, maxValTag := getValLocations(tagToSqlColIdx, cols)
    53  
    54  	return &KVToSqlRowConverter{
    55  		nbf:            nbf,
    56  		cols:           cols,
    57  		tagToSqlColIdx: tagToSqlColIdx,
    58  		rowSize:        rowSize,
    59  		valsFromKey:    valsFromKey,
    60  		valsFromVal:    valsFromVal,
    61  		maxValTag:      maxValTag,
    62  	}
    63  }
    64  
    65  // get counts of where the values we want converted come from so we can skip entire tuples at times.
    66  func getValLocations(tagToSqlColIdx map[uint64]int, cols []schema.Column) (int, int, uint64) {
    67  	var fromKey int
    68  	var fromVal int
    69  	var maxValTag uint64
    70  	for _, col := range cols {
    71  		if _, ok := tagToSqlColIdx[col.Tag]; ok {
    72  			if col.IsPartOfPK {
    73  				fromKey++
    74  			} else {
    75  				fromVal++
    76  				maxValTag = maxU64(maxValTag, col.Tag)
    77  			}
    78  		}
    79  	}
    80  
    81  	return fromKey, fromVal, maxValTag
    82  }
    83  
    84  // NewKVToSqlRowConverterForCols returns a KVToSqlConverter instance based on the list of columns passed in
    85  func NewKVToSqlRowConverterForCols(nbf *types.NomsBinFormat, cols []schema.Column) *KVToSqlRowConverter {
    86  	tagToSqlColIdx := make(map[uint64]int)
    87  	for i, col := range cols {
    88  		tagToSqlColIdx[col.Tag] = i
    89  	}
    90  
    91  	return NewKVToSqlRowConverter(nbf, tagToSqlColIdx, cols, len(cols))
    92  }
    93  
    94  // ConvertKVToSqlRow returns a sql.Row generated from the key and value provided.
    95  func (conv *KVToSqlRowConverter) ConvertKVToSqlRow(k, v types.Value) (sql.Row, error) {
    96  	keyTup, ok := k.(types.Tuple)
    97  
    98  	if !ok {
    99  		return nil, errors.New("invalid key is not a tuple")
   100  	}
   101  
   102  	var valTup types.Tuple
   103  	if !types.IsNull(v) {
   104  		valTup, ok = v.(types.Tuple)
   105  
   106  		if !ok {
   107  			return nil, errors.New("invalid value is not a tuple")
   108  		}
   109  	} else {
   110  		valTup = types.EmptyTuple(conv.nbf)
   111  	}
   112  
   113  	return conv.ConvertKVTuplesToSqlRow(keyTup, valTup)
   114  }
   115  
   116  // ConvertKVToSqlRow returns a sql.Row generated from the key and value provided.
   117  func (conv *KVToSqlRowConverter) ConvertKVTuplesToSqlRow(k, v types.Tuple) (sql.Row, error) {
   118  	tupItr := types.TupleItrPool.Get().(*types.TupleIterator)
   119  	defer types.TupleItrPool.Put(tupItr)
   120  
   121  	cols := make([]interface{}, conv.rowSize)
   122  	if conv.valsFromKey > 0 {
   123  		// keys are not in sorted order so cannot use max tag to early exit
   124  		err := conv.processTuple(cols, conv.valsFromKey, 0xFFFFFFFFFFFFFFFF, k, tupItr)
   125  
   126  		if err != nil {
   127  			return nil, err
   128  		}
   129  	}
   130  
   131  	if conv.valsFromVal > 0 {
   132  		err := conv.processTuple(cols, conv.valsFromVal, conv.maxValTag, v, tupItr)
   133  
   134  		if err != nil {
   135  			return nil, err
   136  		}
   137  	}
   138  
   139  	return cols, nil
   140  }
   141  
   142  func (conv *KVToSqlRowConverter) processTuple(cols []interface{}, valsToFill int, maxTag uint64, tup types.Tuple, tupItr *types.TupleIterator) error {
   143  	err := tupItr.InitForTuple(tup)
   144  
   145  	if err != nil {
   146  		return err
   147  	}
   148  
   149  	nbf := tup.Format()
   150  	primReader, numPrimitives := tupItr.CodecReader()
   151  
   152  	filled := 0
   153  	for pos := uint64(0); pos+1 < numPrimitives; pos += 2 {
   154  		if filled >= valsToFill {
   155  			break
   156  		}
   157  
   158  		tagKind := primReader.ReadKind()
   159  
   160  		if tagKind != types.UintKind {
   161  			return errors.New("Encountered unexpected kind while attempting to read tag")
   162  		}
   163  
   164  		tag64 := primReader.ReadUint()
   165  		if tag64 > maxTag {
   166  			break
   167  		}
   168  
   169  		if sqlColIdx, ok := conv.tagToSqlColIdx[tag64]; !ok {
   170  			err = primReader.SkipValue(nbf)
   171  
   172  			if err != nil {
   173  				return err
   174  			}
   175  		} else {
   176  			cols[sqlColIdx], err = conv.cols[sqlColIdx].TypeInfo.ReadFrom(nbf, primReader)
   177  
   178  			if err != nil {
   179  				return err
   180  			}
   181  
   182  			filled++
   183  		}
   184  	}
   185  
   186  	return nil
   187  }
   188  
   189  // KVGetFunc defines a function that returns a Key Value pair
   190  type KVGetFunc func(ctx context.Context) (types.Tuple, types.Tuple, error)
   191  
   192  func GetGetFuncForMapIter(nbf *types.NomsBinFormat, mapItr types.MapIterator) func(ctx context.Context) (types.Tuple, types.Tuple, error) {
   193  	return func(ctx context.Context) (types.Tuple, types.Tuple, error) {
   194  		k, v, err := mapItr.Next(ctx)
   195  
   196  		if err != nil {
   197  			return types.Tuple{}, types.Tuple{}, err
   198  		} else if k == nil {
   199  			return types.Tuple{}, types.Tuple{}, io.EOF
   200  		}
   201  
   202  		valTup, ok := v.(types.Tuple)
   203  		if !ok {
   204  			valTup = types.EmptyTuple(nbf)
   205  		}
   206  
   207  		return k.(types.Tuple), valTup, nil
   208  	}
   209  }
   210  
   211  // DoltMapIter uses a types.MapIterator to iterate over a types.Map and returns sql.Row instances that it reads and
   212  // converts
   213  type DoltMapIter struct {
   214  	ctx           context.Context
   215  	kvGet         KVGetFunc
   216  	closeKVGetter func() error
   217  	conv          *KVToSqlRowConverter
   218  }
   219  
   220  // NewDoltMapIter returns a new DoltMapIter
   221  func NewDoltMapIter(ctx context.Context, keyValGet KVGetFunc, closeKVGetter func() error, conv *KVToSqlRowConverter) *DoltMapIter {
   222  	return &DoltMapIter{
   223  		ctx:           ctx,
   224  		kvGet:         keyValGet,
   225  		closeKVGetter: closeKVGetter,
   226  		conv:          conv,
   227  	}
   228  }
   229  
   230  // Next returns the next sql.Row until all rows are returned at which point (nil, io.EOF) is returned.
   231  func (dmi *DoltMapIter) Next() (sql.Row, error) {
   232  	k, v, err := dmi.kvGet(dmi.ctx)
   233  
   234  	if err != nil {
   235  		return nil, err
   236  	}
   237  
   238  	return dmi.conv.ConvertKVTuplesToSqlRow(k, v)
   239  }
   240  
   241  func (dmi *DoltMapIter) Close(*sql.Context) error {
   242  	if dmi.closeKVGetter != nil {
   243  		return dmi.closeKVGetter()
   244  	}
   245  
   246  	return nil
   247  }