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  }