Add sqlite online store for go for testing

Signed-off-by: Kevin Zhang <kzhang@tecton.ai>
feast-dev · feast-ci-bot · Mar 29, 2022 · Mar 23, 2022 · Mar 23, 2022 · Mar 24, 2022
commit 21c8900a4d62733c5c25eef1bfb0972e0c16cd4c
@@ -59,7 +59,8 @@ func NewOnlineStore(config *RepoConfig) (OnlineStore, error) {
 		return onlineStore, err
 	}
 	if onlineStoreType == "sqlite" {
-		onlineStore, err := NewSqliteOnlineStore(config.Project, config.OnlineStore)
+		print("ASFDSASF!!")
+		onlineStore, err := NewSqliteOnlineStore(config.Project, config, config.OnlineStore)
 		return onlineStore, err
 	} else {
 		return nil, fmt.Errorf("%s online store type is currently not supported; only Redis is supported", onlineStoreType)

@@ -3,32 +3,37 @@ package feast
 import (
 	"database/sql"
 	"errors"
+	"log"
+	"strings"
+	"time"
 
 	"context"
 	"fmt"
 
+	"github.com/feast-dev/feast/go/protos/feast/serving"
 	"github.com/feast-dev/feast/go/protos/feast/types"
 	_ "github.com/mattn/go-sqlite3"
+	"google.golang.org/protobuf/proto"
+	timestamppb "google.golang.org/protobuf/types/known/timestamppb"
 )
 
 type SqliteOnlineStore struct {
 	// Feast project name
-	// TODO (woop): Should we remove project as state that is tracked at the store level?
 	project string
-
-	path string
-	db   *sql.DB
+	path    string
+	db      *sql.DB
 }
 
-func NewSqliteOnlineStore(project string, onlineStoreConfig map[string]interface{}) (*SqliteOnlineStore, error) {
+// Creates a new sqlite online store object. onlineStoreConfig should have relative path of database file with respect to repoConfig.repoPath.
+func NewSqliteOnlineStore(project string, repoConfig *RepoConfig, onlineStoreConfig map[string]interface{}) (*SqliteOnlineStore, error) {
 	store := SqliteOnlineStore{project: project}
 	if db_path, ok := onlineStoreConfig["path"]; !ok {
 		return nil, fmt.Errorf("cannot find sqlite path %s", db_path)
 	} else if dbPathStr, ok := db_path.(string); !ok {
 		return nil, fmt.Errorf("cannot find convert sqlite path to string %s", db_path)
 	} else {
-		store.path = dbPathStr
-		db, err := initializeConnection(dbPathStr)
+		store.path = fmt.Sprintf("%s/%s", repoConfig.RepoPath, dbPathStr)
+		db, err := initializeConnection(store.path)
 		if err != nil {
 			return nil, err
 		}
@@ -41,17 +46,69 @@ func (s *SqliteOnlineStore) Destruct() {
 	s.db.Close()
 }
 
+// Returns FeatureData 2D array. Each row corresponds to one entity value and each column corresponds to a single feature where the number of columns should be
+// same length as the length of featureNames. Reads from every table in featureViewNames with the entity keys described.
 func (s *SqliteOnlineStore) OnlineRead(ctx context.Context, entityKeys []types.EntityKey, featureViewNames []string, featureNames []string) ([][]FeatureData, error) {
+	featureCount := len(featureNames)
 	_, err := s.getConnection()
 	if err != nil {
 		return nil, err
 	}
-	return nil, nil
-}
 
-// feature views, entities, dataframe?
-func (s *SqliteOnlineStore) WriteToOnlineStore(ctx context.Context, featureViewName string, data [][]FeatureData) error {
-	return nil
+	project := s.project
+	results := make([][]FeatureData, len(entityKeys))
+	entityNameToEntityIndex := make(map[string]int)
+	in_query := make([]string, len(entityKeys))
+	serialized_entities := make([]interface{}, len(entityKeys))
+	for i := 0; i < len(entityKeys); i++ {
+		serKey, err := serializeEntityKey(&entityKeys[i])
+		if err != nil {
+			return nil, err
+		}
+		// TODO: fix this, string conversion is not safe
+		entityNameToEntityIndex[string(*serKey)] = i
+		// for IN clause in read query
+		in_query[i] = "?"
+		serialized_entities[i] = *serKey
+	}
+	featureNamesToIdx := make(map[string]int)
+	for idx, name := range featureNames {
+		featureNamesToIdx[name] = idx
+	}
+
+	for idx := 0; idx < len(entityKeys); idx++ {
+		results[idx] = make([]FeatureData, featureCount)
+	}
+	for _, featureViewName := range featureViewNames {
+		query_string := fmt.Sprintf(`SELECT entity_key, feature_name, value, event_ts
+									FROM %s
+									WHERE entity_key IN (%s)
+									ORDER BY entity_key`, tableId(project, featureViewName), strings.Join(in_query, ","))
+		rows, err := s.db.Query(query_string, serialized_entities...)
+		if err != nil {
+			return nil, err
+		}
+		defer rows.Close()
+		for rows.Next() {
+			var entity_key string
+			var feature_name string
+			var valueString string
+			var event_ts time.Time
+			var value types.Value
+			err = rows.Scan(&entity_key, &feature_name, &valueString, &event_ts)
+			if err != nil {
+				log.Fatal(err)
+			}
+			if err := proto.Unmarshal([]byte(valueString), &value); err != nil {
+				return nil, errors.New("error converting parsed value to types.Value")
+			}
+			results[entityNameToEntityIndex[entity_key]][featureNamesToIdx[feature_name]] = FeatureData{reference: serving.FeatureReferenceV2{FeatureViewName: featureViewName, FeatureName: feature_name},
+				timestamp: *timestamppb.New(event_ts),
+				value:     types.Value{Val: value.Val},
+			}
+		}
+	}
+	return results, nil
 }
 
 func (s *SqliteOnlineStore) Update(ctx context.Context, config *RepoConfig, tables_to_delete []*FeatureView, tables_to_keep []*FeatureView) error {
@@ -62,16 +119,17 @@ func (s *SqliteOnlineStore) Update(ctx context.Context, config *RepoConfig, tabl
 	project := config.Project
 	for _, table := range tables_to_keep {
 		s.db.Exec(
-			fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (entity_key BLOB, feature_name TEXT, value BLOB, event_ts timestamp, created_ts timestamp,  PRIMARY KEY(entity_key, feature_name))", tableId(project, table)))
+			fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (entity_key BLOB, feature_name TEXT, value BLOB, event_ts timestamp, created_ts timestamp,  PRIMARY KEY(entity_key, feature_name))", tableId(project, table.base.name)))
 		s.db.Exec(
-			fmt.Sprintf("CREATE INDEX IF NOT EXISTS %s_ek ON %s (entity_key);", tableId(project, table), tableId(project, table)))
+			fmt.Sprintf("CREATE INDEX IF NOT EXISTS %s_ek ON %s (entity_key);", tableId(project, table.base.name), tableId(project, table.base.name)))
 	}
 	for _, table := range tables_to_delete {
-		s.db.Exec("DROP TABLE IF EXISTS %s", tableId(project, table))
+		s.db.Exec("DROP TABLE IF EXISTS %s", tableId(project, table.base.name))
 	}
 	return nil
 }
 
+// Gets a sqlite connection and sets it to the online store and also returns a pointer to the connection.
 func (s *SqliteOnlineStore) getConnection() (*sql.DB, error) {
 	if s.db == nil {
 		if s.path == "" {
@@ -86,12 +144,14 @@ func (s *SqliteOnlineStore) getConnection() (*sql.DB, error) {
 	return s.db, nil
 }
 
-func tableId(project string, table *FeatureView) string {
-	return fmt.Sprintf("%s_%s", project, table.base.name)
+// Constructs the table id from the project and table(featureViewName) string.
+func tableId(project string, featureViewName string) string {
+	return fmt.Sprintf("%s_%s", project, featureViewName)
 }
 
+// Creates a connection to the sqlite database and returns the connection.
 func initializeConnection(db_path string) (*sql.DB, error) {
-	db, err := sql.Open("sqlite3", "./foo.db")
+	db, err := sql.Open("sqlite3", db_path)
 	if err != nil {
 		return nil, err
 	}

@@ -1,18 +1,19 @@
 package feast
 
 import (
-	"fmt"
-	"log"
+	"context"
+	"reflect"
 	"testing"
 
+	"github.com/feast-dev/feast/go/protos/feast/types"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestSqliteSetup(t *testing.T) {
-	dir := "../test/selected_tetra"
+	dir := "../test/feature_repo"
 	config, err := NewRepoConfigFromFile(dir)
 	assert.Nil(t, err)
-	assert.Equal(t, "selected_tetra", config.Project)
+	assert.Equal(t, "feature_repo", config.Project)
 	assert.Equal(t, "data/registry.db", config.GetRegistryConfig().Path)
 	assert.Equal(t, "local", config.Provider)
 	assert.Equal(t, map[string]interface{}{
@@ -23,32 +24,50 @@ func TestSqliteSetup(t *testing.T) {
 	assert.Empty(t, config.Flags)
 }
 
-func TestSqliteUpdate(t *testing.T) {
-	dir := "../test/selected_tetra"
+func TestSqliteOnlineRead(t *testing.T) {
+	dir := "../test/feature_repo"
 	config, err := NewRepoConfigFromFile(dir)
 	assert.Nil(t, err)
-	store, err := NewSqliteOnlineStore("test", config.OnlineStore)
+	store, err := NewSqliteOnlineStore("feature_repo", config, config.OnlineStore)
 	defer store.Destruct()
 	assert.Nil(t, err)
-	show_tables := `SELECT
-    	*
-	FROM
-		selected_tetra_driver_hourly_stats;`
-	rows, err := store.db.Query(show_tables)
-	if err != nil {
-		log.Fatal(err)
+	entity_key1 := types.EntityKey{
+		JoinKeys:     []string{"driver_id"},
+		EntityValues: []*types.Value{{Val: &types.Value_Int64Val{1005}}},
 	}
-	defer rows.Close()
-	for rows.Next() {
-		fmt.Println("in")
-		var id int
-		var name string
-		err = rows.Scan(&id, &name)
-		if err != nil {
-			log.Fatal(err)
+	entity_key2 := types.EntityKey{
+		JoinKeys:     []string{"driver_id"},
+		EntityValues: []*types.Value{{Val: &types.Value_Int64Val{1001}}},
+	}
+	entity_key3 := types.EntityKey{
+		JoinKeys:     []string{"driver_id"},
+		EntityValues: []*types.Value{{Val: &types.Value_Int64Val{1003}}},
+	}
+	entityKeys := []types.EntityKey{entity_key1, entity_key2, entity_key3}
+	tableNames := []string{"driver_hourly_stats"}
+	featureNames := []string{"conv_rate", "acc_rate", "avg_daily_trips"}
+	featureData, err := store.OnlineRead(context.Background(), entityKeys, tableNames, featureNames)
+	assert.Nil(t, err)
+	returnedFeatureValues := make([]types.Value, 0)
+	returnedFeatureNames := make([]string, 0)
+	for _, featureVector := range featureData {
+		for _, feature := range featureVector {
+			returnedFeatureValues = append(returnedFeatureValues, feature.value)
+			returnedFeatureNames = append(returnedFeatureNames, feature.reference.FeatureName)
 		}
-		fmt.Println(id, name)
 	}
-	assert.False(t, true)
-
+	expectedFeatureValues := []types.Value{
+		{Val: &types.Value_FloatVal{0.78135854}},
+		{Val: &types.Value_FloatVal{0.38527268}},
+		{Val: &types.Value_Int64Val{755}},
+		{Val: &types.Value_FloatVal{0.49661186}},
+		{Val: &types.Value_FloatVal{0.9440974}},
+		{Val: &types.Value_Int64Val{169}},
+		{Val: &types.Value_FloatVal{0.80762655}},
+		{Val: &types.Value_FloatVal{0.71510273}},
+		{Val: &types.Value_Int64Val{545}},
+	}
+	expectedFeatureNames := []string{"conv_rate", "acc_rate", "avg_daily_trips", "conv_rate", "acc_rate", "avg_daily_trips", "conv_rate", "acc_rate", "avg_daily_trips"}
+	assert.True(t, reflect.DeepEqual(expectedFeatureValues, returnedFeatureValues))
+	assert.True(t, reflect.DeepEqual(expectedFeatureNames, returnedFeatureNames))
 }
diff --git a/go/internal/test/selected_tetra/__init__.py → go/internal/test/feature_repo/__init__.py b/go/internal/test/selected_tetra/__init__.py → go/internal/test/feature_repo/__init__.py
diff --git a/go/internal/test/feature_repo/data/driver_stats.parquet b/go/internal/test/feature_repo/data/driver_stats.parquet
diff --git a/.../test/selected_tetra/data/online_store.db → ...al/test/feature_repo/data/online_store.db b/.../test/selected_tetra/data/online_store.db → ...al/test/feature_repo/data/online_store.db
diff --git a/go/internal/test/selected_tetra/example.py → go/internal/test/feature_repo/example.py b/go/internal/test/selected_tetra/example.py → go/internal/test/feature_repo/example.py
@@ -8,7 +8,7 @@
 # production, you can use your favorite DWH, such as BigQuery. See Feast documentation
 # for more info.
 driver_hourly_stats = FileSource(
-    path="/Users/kevinzhang/tecton-ai/offline_store/feast/go/internal/test/selected_tetra/data/driver_stats.parquet",
+    path="/Users/kevinzhang/tecton-ai/offline_store/feast/go/internal/test/feature_repo/data/driver_stats.parquet",
     event_timestamp_column="event_timestamp",
     created_timestamp_column="created",
 )

diff --git a/...al/test/selected_tetra/feature_store.yaml → ...rnal/test/feature_repo/feature_store.yaml b/...al/test/selected_tetra/feature_store.yaml → ...rnal/test/feature_repo/feature_store.yaml
@@ -1,4 +1,4 @@
-project: selected_tetra
+project: feature_repo
 registry: data/registry.db
 provider: local
 online_store:

diff --git a/go/internal/test/selected_tetra/data/driver_stats.parquet b/go/internal/test/selected_tetra/data/driver_stats.parquet