fix: Update the serialization of EntityKey to version 3 in Go Feature server.

Signed-off-by: Shuchu Han <[email protected]>

Author: shuchu

go/internal/feast/onlinestore/onlinestore.go

@@ -87,41 +87,63 @@ func serializeEntityKey(entityKey *types.EntityKey, entityKeySerializationVersio
 	}
 
 	keys := make([]string, 0, len(m))
-	for k := range entityKey.JoinKeys {
-		keys = append(keys, entityKey.JoinKeys[k])
-	}
-	sort.Strings(keys)
+	keys = append(keys, entityKey.JoinKeys...)
+	sort.Strings(keys)  // Sort the keysgiut
 
 	// Build the key
-	length := 5 * len(keys)
+	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++ {
-		offset := i * 2
+		// Add the key type STRING info
 		byteBuffer := make([]byte, 4)
 		binary.LittleEndian.PutUint32(byteBuffer, uint32(types.ValueType_Enum_value["STRING"]))
 		bufferList[offset] = byteBuffer
-		bufferList[offset+1] = []byte(keys[i])
+		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++ {
-		offset := (2 * len(keys)) + (i * 3)
 		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+0] = typeBuffer
-		bufferList[offset+1] = lenBuffer
-		bufferList[offset+2] = *valueBytes
+		bufferList[offset] = *valueBytes
+		offset++
 	}
 
 	// Convert from an array of byte arrays to a single byte array
@@ -132,3 +154,27 @@ func serializeEntityKey(entityKey *types.EntityKey, entityKeySerializationVersio
 
 	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)
+	}
+}

go/internal/feast/onlinestore/onlinestore_test.go

@@ -0,0 +1,46 @@
+package onlinestore
+
+import (
+	"testing"
+	"reflect"
+	"github.com/stretchr/testify/assert"
+	"github.com/feast-dev/feast/go/protos/feast/types"
+)
+
+func Test_serializeEntityKey(t *testing.T) {
+	expect_res := []byte{1, 0, 0, 0, 2, 0, 0, 0, 9, 0, 0, 0, 100, 114, 105, 118, 101, 114, 95, 105, 100, 4, 0, 0, 0, 8, 0, 0, 0, 233, 3, 0, 0, 0, 0, 0, 0}
+	tests := []struct {
+		name string // description of this test case
+		// Named input parameters for target function.
+		entityKey                     *types.EntityKey
+		entityKeySerializationVersion int64
+		want                          []byte
+		wantErr                       bool
+	}{
+		{
+			"test a specific key",
+			&types.EntityKey{
+				JoinKeys:     []string{"driver_id"},
+				EntityValues: []*types.Value{{Val: &types.Value_Int64Val{Int64Val: 1001}}},
+			},
+			3,
+			expect_res,
+			false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, gotErr := serializeEntityKey(tt.entityKey, tt.entityKeySerializationVersion)
+			if gotErr != nil {
+				if !tt.wantErr {
+					t.Errorf("serializeEntityKey() failed: %v", gotErr)
+				}
+				return
+			}
+			if tt.wantErr {
+				t.Fatal("serializeEntityKey() succeeded unexpectedly")
+			}
+			assert.True(t, reflect.DeepEqual(*got, tt.want))
+		})
+	}
+}

go/internal/feast/onlinestore/redisonlinestore.go

@@ -338,26 +338,3 @@ func buildRedisKey(project string, entityKey *types.EntityKey, entityKeySerializ
 	fullKey := append(*serKey, []byte(project)...)
 	return &fullKey, 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)
-	}
-}