package onlinestore

import (

"context"

"encoding/binary"

"fmt"

"sort"

"github.com/feast-dev/feast/go/internal/feast/registry"

"github.com/feast-dev/feast/go/protos/feast/serving"

"github.com/feast-dev/feast/go/protos/feast/types"

"github.com/golang/protobuf/ptypes/timestamp"

)

type FeatureData struct {

Reference serving.FeatureReferenceV2

Timestamp timestamp.Timestamp

Value types.Value

}

type OnlineStore interface {

// OnlineRead reads multiple features (specified in featureReferences) for multiple

// entity keys (specified in entityKeys) and returns an array of array of features,

// where each feature contains 3 fields:

// 1. feature Reference

// 2. feature event timestamp

// 3. feature value

// The inner array will have the same size as featureReferences,

// while the outer array will have the same size as entityKeys.

// TODO: Can we return [][]FeatureData, []timstamps, error

// instead and remove timestamp from FeatureData struct to mimic Python's code

// and reduces repeated memory storage for the same timstamp (which is stored as value and not as a pointer).

// Should each attribute in FeatureData be stored as a pointer instead since the current

// design forces value copied in OnlineRead + GetOnlineFeatures

// (array is destructed so we cannot use the same fields in each

// Feature object as pointers in GetOnlineFeaturesResponse)

// => allocate memory for each field once in OnlineRead

// and reuse them in GetOnlineFeaturesResponse?

OnlineRead(ctx context.Context, entityKeys []*types.EntityKey, featureViewNames []string, featureNames []string) ([][]FeatureData, error)

// Destruct must be call once user is done using OnlineStore

// This is to comply with the Connector since we have to close the plugin

Destruct()

}

func getOnlineStoreType(onlineStoreConfig map[string]interface{}) (string, bool) {

if onlineStoreType, ok := onlineStoreConfig["type"]; !ok {

// If online store type isn't specified, default to sqlite

return "sqlite", true

} else {

result, ok := onlineStoreType.(string)

return result, ok

}

}

func NewOnlineStore(config *registry.RepoConfig) (OnlineStore, error) {

onlineStoreType, ok := getOnlineStoreType(config.OnlineStore)

if !ok {

return nil, fmt.Errorf("could not get online store type from online store config: %+v", config.OnlineStore)

} else if onlineStoreType == "sqlite" {

onlineStore, err := NewSqliteOnlineStore(config.Project, config, config.OnlineStore)

return onlineStore, err

} else if onlineStoreType == "redis" {

onlineStore, err := NewRedisOnlineStore(config.Project, config, config.OnlineStore)

return onlineStore, err

} else if onlineStoreType == "dynamodb" {

onlineStore, err := NewDynamodbOnlineStore(config.Project, config, config.OnlineStore)

return onlineStore, err

} else {

return nil, fmt.Errorf("%s online store type is currently not supported; only redis and sqlite are supported", onlineStoreType)

}

}

func serializeEntityKey(entityKey *types.EntityKey, entityKeySerializationVersion int64) (*[]byte, error) {

// Serialize entity key to a bytestring so that it can be used as a lookup key in a hash table.

// Ensure that we have the right amount of join keys and entity values

if len(entityKey.JoinKeys) != len(entityKey.EntityValues) {

return nil, fmt.Errorf("the amount of join key names and entity values don't match: %s vs %s", entityKey.JoinKeys, entityKey.EntityValues)

}

// Make sure that join keys are sorted so that we have consistent key building

m := make(map[string]*types.Value)

for i := 0; i < len(entityKey.JoinKeys); i++ {

m[entityKey.JoinKeys[i]] = entityKey.EntityValues[i]

}

keys := make([]string, 0, len(m))

keys = append(keys, entityKey.JoinKeys...)

sort.Strings(keys) // Sort the keys

// Build the key

length := 7 * len(keys)

bufferList := make([][]byte, length)

offset := 0

// For entityKeySerializationVersion 3 and above, we add the number of join keys

// as the first 4 bytes of the serialized key.

if entityKeySerializationVersion >= 3 {

byteBuffer := make([]byte, 4)

binary.LittleEndian.PutUint32(byteBuffer, uint32(len(keys)))

bufferList[offset] = byteBuffer // First buffer is always the length of the keys

offset++

}

for i := 0; i < len(keys); i++ {

// Add the key type STRING info

byteBuffer := make([]byte, 4)

binary.LittleEndian.PutUint32(byteBuffer, uint32(types.ValueType_Enum_value["STRING"]))

bufferList[offset] = byteBuffer

offset++

// Add the size of current "key" string

keyLenByteBuffer := make([]byte, 4)

binary.LittleEndian.PutUint32(keyLenByteBuffer, uint32(len(keys[i])))

bufferList[offset] = keyLenByteBuffer

offset++

// Add value

bufferList[offset] = []byte(keys[i])

offset++

}

for i := 0; i < len(keys); i++ {

value := m[keys[i]].GetVal()

valueBytes, valueTypeBytes, err := serializeValue(value, entityKeySerializationVersion)

if err != nil {

return valueBytes, err

}

// Add value type info

typeBuffer := make([]byte, 4)

binary.LittleEndian.PutUint32(typeBuffer, uint32(valueTypeBytes))

bufferList[offset] = typeBuffer

offset++

// Add length info

lenBuffer := make([]byte, 4)

binary.LittleEndian.PutUint32(lenBuffer, uint32(len(*valueBytes)))

bufferList[offset] = lenBuffer

offset++

bufferList[offset] = *valueBytes

offset++

}

// Convert from an array of byte arrays to a single byte array

var entityKeyBuffer []byte

for i := 0; i < len(bufferList); i++ {

entityKeyBuffer = append(entityKeyBuffer, bufferList[i]...)

}

return &entityKeyBuffer, nil

}

func serializeValue(value interface{}, entityKeySerializationVersion int64) (*[]byte, types.ValueType_Enum, error) {

// TODO: Implement support for other types (at least the major types like ints, strings, bytes)

switch x := (value).(type) {

case *types.Value_StringVal:

valueString := []byte(x.StringVal)

return &valueString, types.ValueType_STRING, nil

case *types.Value_BytesVal:

return &x.BytesVal, types.ValueType_BYTES, nil

case *types.Value_Int32Val:

valueBuffer := make([]byte, 4)

binary.LittleEndian.PutUint32(valueBuffer, uint32(x.Int32Val))

return &valueBuffer, types.ValueType_INT32, nil

case *types.Value_Int64Val:

valueBuffer := make([]byte, 8)

binary.LittleEndian.PutUint64(valueBuffer, uint64(x.Int64Val))

return &valueBuffer, types.ValueType_INT64, nil

case nil:

return nil, types.ValueType_INVALID, fmt.Errorf("could not detect type for %v", x)

default:

return nil, types.ValueType_INVALID, fmt.Errorf("could not detect type for %v", x)

}

}