vitess.io/vitess@v0.16.2/go/vt/vtgate/vindexes/numeric_static_map.go (about) 1 /* 2 Copyright 2019 The Vitess Authors. 3 4 Licensed under the Apache License, Version 2.0 (the "License"); 5 you may not use this file except in compliance with the License. 6 You may obtain a copy of the License at 7 8 http://www.apache.org/licenses/LICENSE-2.0 9 10 Unless required by applicable law or agreed to in writing, software 11 distributed under the License is distributed on an "AS IS" BASIS, 12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 See the License for the specific language governing permissions and 14 limitations under the License. 15 */ 16 17 package vindexes 18 19 import ( 20 "bytes" 21 "context" 22 "encoding/binary" 23 "encoding/json" 24 "errors" 25 "os" 26 "strconv" 27 28 "vitess.io/vitess/go/vt/vtgate/evalengine" 29 30 "vitess.io/vitess/go/sqltypes" 31 "vitess.io/vitess/go/vt/key" 32 ) 33 34 var ( 35 _ SingleColumn = (*NumericStaticMap)(nil) 36 _ Hashing = (*NumericStaticMap)(nil) 37 ) 38 39 // NumericLookupTable stores the mapping of keys. 40 type NumericLookupTable map[uint64]uint64 41 42 // NumericStaticMap is similar to vindex Numeric but first attempts a lookup via 43 // a JSON file. 44 type NumericStaticMap struct { 45 name string 46 lookup NumericLookupTable 47 } 48 49 func init() { 50 Register("numeric_static_map", NewNumericStaticMap) 51 } 52 53 // NewNumericStaticMap creates a NumericStaticMap vindex. 54 func NewNumericStaticMap(name string, params map[string]string) (Vindex, error) { 55 jsonPath, ok := params["json_path"] 56 if !ok { 57 return nil, errors.New("NumericStaticMap: Could not find `json_path` param in vschema") 58 } 59 60 lt, err := loadNumericLookupTable(jsonPath) 61 if err != nil { 62 return nil, err 63 } 64 65 return &NumericStaticMap{ 66 name: name, 67 lookup: lt, 68 }, nil 69 } 70 71 // String returns the name of the vindex. 72 func (vind *NumericStaticMap) String() string { 73 return vind.name 74 } 75 76 // Cost returns the cost of this vindex as 1. 77 func (*NumericStaticMap) Cost() int { 78 return 1 79 } 80 81 // IsUnique returns true since the Vindex is unique. 82 func (vind *NumericStaticMap) IsUnique() bool { 83 return true 84 } 85 86 // NeedsVCursor satisfies the Vindex interface. 87 func (vind *NumericStaticMap) NeedsVCursor() bool { 88 return false 89 } 90 91 // Verify returns true if ids and ksids match. 92 func (vind *NumericStaticMap) Verify(ctx context.Context, vcursor VCursor, ids []sqltypes.Value, ksids [][]byte) ([]bool, error) { 93 out := make([]bool, 0, len(ids)) 94 for i, id := range ids { 95 ksid, err := vind.Hash(id) 96 if err != nil { 97 return nil, err 98 } 99 out = append(out, bytes.Equal(ksid, ksids[i])) 100 } 101 return out, nil 102 } 103 104 // Map can map ids to key.Destination objects. 105 func (vind *NumericStaticMap) Map(ctx context.Context, vcursor VCursor, ids []sqltypes.Value) ([]key.Destination, error) { 106 out := make([]key.Destination, 0, len(ids)) 107 for _, id := range ids { 108 ksid, err := vind.Hash(id) 109 if err != nil { 110 out = append(out, key.DestinationNone{}) 111 continue 112 } 113 out = append(out, key.DestinationKeyspaceID(ksid)) 114 } 115 return out, nil 116 } 117 118 func (vind *NumericStaticMap) Hash(id sqltypes.Value) ([]byte, error) { 119 num, err := evalengine.ToUint64(id) 120 if err != nil { 121 return nil, err 122 } 123 lookupNum, ok := vind.lookup[num] 124 if ok { 125 num = lookupNum 126 } 127 var keybytes [8]byte 128 binary.BigEndian.PutUint64(keybytes[:], num) 129 return keybytes[:], nil 130 } 131 132 func loadNumericLookupTable(path string) (NumericLookupTable, error) { 133 var m map[string]uint64 134 lt := make(map[uint64]uint64) 135 data, err := os.ReadFile(path) 136 if err != nil { 137 return lt, err 138 } 139 err = json.Unmarshal(data, &m) 140 if err != nil { 141 return lt, err 142 } 143 for k, v := range m { 144 newK, err := strconv.ParseUint(k, 10, 64) 145 if err != nil { 146 return lt, err 147 } 148 lt[newK] = v 149 } 150 151 return lt, nil 152 }