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 }