feat: Implement ScyllaDB Online Store #138

msistla96 · 2024-09-11T21:13:46Z

What this PR does / why we need it:

(WIP)
This PR is to implement the ScyllaDB online store for customers seeking a low cost Online Store when compared to Redis.
More details in https://confluence.expedia.biz/display/~msistla/Scylla+DB+Online+Store+Implementation
This PR contains the following changes:

Changes to RepoConfig to support ScyllaDB as an online store.
Changes to setup.py, go.mod and go.sum, Dockerfile, Makefile to add the ScyllaDB specific driver dependency for Python and Go, in order to take advantage of the shard aware functionality it offers. See https://github.com/scylladb/gocql and https://github.com/scylladb/python-driver for more details.
Changes to CassandraOnlineStore in Python.
Addition of CassandraOnlineStore as a new Online Store in Go.

TODO

Add tests
Metrics?

Which issue(s) this PR fixes:

N/A

Fixes

N/A

…nup code a bit

EXPEbdodla · 2024-11-18T20:25:31Z

Makefile

@@ -486,7 +486,7 @@ build-java-docker-dev:
 build-go-docker-dev:
 	docker buildx build --build-arg VERSION=dev \
 		-t feastdev/feature-server-go:dev \
-		-f go/infra/docker/feature-server/Dockerfile --load .
+		-f SCYLLADB_SHARD_AWARE_ENABLED=$(SCYLLADB_SHARD_AWARE_ENABLED) go/infra/docker/feature-server/Dockerfile --load .


Why do we need this? Doesn't the drivers auto detect this?

Deleted, we should only be using the scylladb gocql fork

EXPEbdodla · 2024-11-18T20:25:50Z

go.mod

 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
+
+require github.com/aws/aws-sdk-go v1.34.28


Do we need this?

You're right, I guess this was part of Meenakshi's testing? Deleted it anyway along with the code that used it.

EXPEbdodla · 2024-11-18T20:26:53Z

go.mod

+
+require github.com/aws/aws-sdk-go v1.34.28
+
+require github.com/gocql/gocql v1.6.0


cassandra 1.6.0. Instructions for scylladb gocql: https://github.com/scylladb/gocql?tab=readme-ov-file#2-installation

Run go mod tidy after.

EXPEbdodla · 2024-11-18T20:29:52Z

go/infra/docker/feature-server/Dockerfile

@@ -1,5 +1,7 @@
 FROM golang:1.22.5

+ARG SCYLLADB_SHARD_AWARE_ENABLED=False
+ENV SCYLLADB_SHARD_AWARE_ENABLED=$SCYLLADB_SHARD_AWARE_ENABLED


What's the use of this ENV?

Was for optionally selecting the ScyllaDB forked driver, but we're just going to use it by default.

EXPEbdodla · 2024-11-18T20:29:59Z

go/infra/docker/feature-server/Dockerfile

@@ -23,6 +25,10 @@ RUN find ./protos -name "*.proto" \
 # Build the Go application
 RUN go build -o feast ./go/main.go

+RUN if ["$SCYLLADB_SHARD_AWARE_ENABLED"== True]; then  \


Move it to go.mod

EXPEbdodla · 2024-11-19T00:47:42Z

go/internal/feast/onlinestore/cassandraonlinestore.go

+	// Parse host_name and Ips
+	cassandraHosts, ok := onlineStoreConfig["hosts"]
+	if !ok {
+		cassandraHosts = []interface{}{"127.0.0.1"}
+		log.Warn().Msg("host not provided: Using 127.0.0.1 instead")
+	}
+
+	var rawCassandraHosts []interface{}
+	if rawCassandraHosts, ok = cassandraHosts.([]interface{}); !ok {
+		return nil, fmt.Errorf("didn't pass a list of hosts in the 'hosts' field")
+	}
+
+	var cassandraHostsStr = make([]string, len(rawCassandraHosts))
+	for i, rawHost := range rawCassandraHosts {
+		hostStr, ok := rawHost.(string)
+		if !ok {
+			return nil, fmt.Errorf("failed to convert a host to a string: %+v", rawHost)
+		}
+		cassandraHostsStr[i] = hostStr
+	}
+
+	username, ok := onlineStoreConfig["username"]
+	if !ok {
+		username = "cassandra"
+		log.Warn().Msg("username not defined: Using default username instead")
+	}
+
+	password, ok := onlineStoreConfig["password"]
+	if !ok {
+		password = "cassandra"
+		log.Warn().Msg("password not defined: Using default password instead")
+	}
+
+	var usernameStr string
+	if usernameStr, ok = username.(string); !ok {
+		return nil, fmt.Errorf("failed to convert username to string: %+v", usernameStr)
+	}
+
+	var passwordStr string
+	if passwordStr, ok = password.(string); !ok {
+		return nil, fmt.Errorf("failed to convert password to string: %+v", passwordStr)
+	}
+
+	keyspace, ok := onlineStoreConfig["keyspace"]
+	if !ok {
+		keyspace = "scylladb"
+		log.Warn().Msg("Keyspace not defined: Using 'scylladb' as keyspace instead")
+	}
+	store.keyspace = keyspace.(string)
+
+	var keyspaceStr string
+	if keyspaceStr, ok = keyspace.(string); !ok {
+		return nil, fmt.Errorf("failed to convert keyspace to string: %+v", keyspaceStr)
+	}
+
+	protocolVersion, ok := onlineStoreConfig["protocol_version"]
+	if !ok {
+		protocolVersion = 4.0
+		log.Warn().Msg("protocol_version not specified: Using 4 instead")
+	}
+	protocolVersionInt := int(protocolVersion.(float64))
+
+	redisTraceServiceName := os.Getenv("DD_SERVICE") + "-cassandra"
+	if redisTraceServiceName == "" {
+		redisTraceServiceName = "cassandra.client" // default service name if DD_SERVICE is not set
+	}
+	store.clusterConfigs = gocqltrace.NewCluster(cassandraHostsStr, gocqltrace.WithServiceName(redisTraceServiceName))
+	// TODO: Figure out if we need to offer users the ability to tune the timeouts
+	//store.clusterConfigs.ConnectTimeout = 1
+	//store.clusterConfigs.Timeout = 1
+	store.clusterConfigs.ProtoVersion = protocolVersionInt
+	//store.clusterConfigs.Consistency = gocql.Quorum
+	store.clusterConfigs.Keyspace = keyspaceStr
+	loadBalancingPolicy, ok := onlineStoreConfig["load_balancing"]
+	if !ok {
+		loadBalancingPolicy = gocql.RoundRobinHostPolicy()
+		log.Warn().Msg("No load balancing policy selected; setting Round Robin Host Policy")
+	}
+	loadBalancingPolicyStr, _ := loadBalancingPolicy.(string)
+	if loadBalancingPolicyStr == "DCAwareRoundRobinPolicy" {
+		store.clusterConfigs.PoolConfig.HostSelectionPolicy = gocql.RoundRobinHostPolicy()
+	} else if loadBalancingPolicyStr == "TokenAwarePolicy(DCAwareRoundRobinPolicy)" {
+		// Configure fallback policy if unable to reach the shard
+		fallback := gocql.RoundRobinHostPolicy()
+		// If using ScyllaDB and setting this policy, this makes the driver shard aware to improve performance
+		store.clusterConfigs.PoolConfig.HostSelectionPolicy = gocql.TokenAwareHostPolicy(fallback)
+		if config.OnlineStore["type"] == "scylladb" {
+			store.clusterConfigs.Port = 19042
+		} else {
+			store.clusterConfigs.Port = 9042
+		}
+	}


We can use type assertions by defining a struct type for online store config. If it doesn't work, move to a different function to do all the parsing.

Extracted to different function.

EXPEbdodla · 2024-11-19T00:57:38Z

go/internal/feast/onlinestore/cassandraonlinestore.go

+	store.clusterConfigs.ProtoVersion = protocolVersionInt
+	//store.clusterConfigs.Consistency = gocql.Quorum
+	store.clusterConfigs.Keyspace = keyspaceStr
+	loadBalancingPolicy, ok := onlineStoreConfig["load_balancing"]


load_balancing object has load_balancing_policy and local_dc fields. Seems they are not parsed correctly.

Yeah, that's right, addressed this.

EXPEbdodla · 2024-11-19T01:08:58Z

go/internal/feast/onlinestore/cassandraonlinestore.go

+	if redisTraceServiceName == "" {
+		redisTraceServiceName = "cassandra.client" // default service name if DD_SERVICE is not set
+	}
+	store.clusterConfigs = gocqltrace.NewCluster(cassandraHostsStr, gocqltrace.WithServiceName(redisTraceServiceName))


Seems its deprecated way of using it. https://github.com/DataDog/dd-trace-go/blob/cde886e584560f55306cbe0d302eb54d66b88cfc/contrib/gocql/gocql/gocql.go#L33-L36

EXPEbdodla · 2024-11-19T01:11:27Z

go/internal/feast/onlinestore/cassandraonlinestore.go

+	}
+
+	return fmt.Sprintf(
+		`SELECT "entity_key", "feature_name", "event_ts", "value" FROM %s WHERE "entity_key" IN (%s) AND "feature_name" IN (%s)`,


entity_key IN may not be efficient when working with partitions across different nodes. We may not be utilizing the shard aware functionality efficiently.

We can start testing with this first and try running concurrent queries as we start testing.

EXPEbdodla · 2024-11-19T01:12:35Z

go/internal/feast/onlinestore/cassandraonlinestore.go

+func (c *CassandraOnlineStore) buildCassandraEntityKeys(entityKeys []*types.EntityKey) ([]interface{}, map[string]int, error) {
+	cassandraKeys := make([]interface{}, len(entityKeys))
+	cassandraKeyToEntityIndex := make(map[string]int)
+	for i := 0; i < len(entityKeys); i++ {
+		var key, err = utils.SerializeEntityKey(entityKeys[i], c.config.EntityKeySerializationVersion)
+		if err != nil {
+			return nil, nil, err
+		}
+		// encoding to hex
+		encodedKey := hex.EncodeToString(*key)
+		cassandraKeys[i] = encodedKey
+		cassandraKeyToEntityIndex[encodedKey] = i
+	}
+	return cassandraKeys, cassandraKeyToEntityIndex, nil


Is it same as buildRedisEntityKeys? can we move to key_utils if its the same?

Almost, but I chose to use the same hex key in the array and map here for consistency, so they're slightly different.

EXPEbdodla · 2024-11-19T02:04:36Z

go/internal/feast/onlinestore/cassandraonlinestore_test.go

+	store, err := NewRedisOnlineStore("test", rc, config)
+	assert.Nil(t, err)
+	var opts = store.client.Options()
+	assert.Equal(t, opts.Addr, "redis://localhost:6379")
+	assert.Equal(t, opts.Password, "secret")


check this.

EXPEbdodla · 2024-11-19T02:05:25Z

go/internal/feast/utils/aws_utils.go

@@ -0,0 +1,117 @@
+package utils


Delete this if not necessary.

setup.py

EXPEbdodla · 2024-11-19T02:10:05Z

go/internal/feast/onlinestore/cassandraonlinestore.go

+	redisTraceServiceName := os.Getenv("DD_SERVICE") + "-cassandra"
+	if redisTraceServiceName == "" {
+		redisTraceServiceName = "cassandra.client" // default service name if DD_SERVICE is not set
+	}


variable names.

go/internal/feast/onlinestore/cassandraonlinestore.go

EXPEbdodla · 2024-11-19T02:12:26Z

go/internal/feast/onlinestore/cassandraonlinestore.go

+	// Session object that holds information about the connection to the cluster
+	session *gocqltrace.Session
+
+	// keyspace of the table. Defaulted to using the project name


if not specified, set to scylladb. It's set to feast_keyspace in CassandraOnlineStoreConfig in python. Bring the consistency to the keyspace name.

go/internal/feast/onlinestore/cassandraonlinestore.go

EXPEbdodla · 2024-11-19T02:29:40Z

go/internal/feast/onlinestore/cassandraonlinestore.go

+	tableName := c.getFqTableName(featureViewName)
+	cqlStatement := c.getCQLStatement(tableName, featureNames, len(serializedEntityKeys))
+	// Bundle the entity keys in one statement (gocql handles this as concurrent queries)
+	scanner := c.session.Query(cqlStatement, serializedEntityKeys...).Iter().Scanner()


We will easily be impacted the default page size set by client. https://github.com/scylladb/gocql/blob/30a838e39b1f0e4d559989b9f30404c5399b3d5e/cluster.go#L142-L144C2
We have a clients requesting 200-300 features and 100 keys (if they do batching at orchestrator. If not we get more keys.

Even if we partition queries, at one point the results of all those smaller queries will be in memory. I think such clients will have to request larger pods when deploying their feature servers.

EXPEbdodla · 2024-11-19T02:31:25Z

go/internal/feast/onlinestore/cassandraonlinestore.go

+		redisTraceServiceName = "cassandra.client" // default service name if DD_SERVICE is not set
+	}
+	store.clusterConfigs = gocqltrace.NewCluster(cassandraHostsStr, gocqltrace.WithServiceName(redisTraceServiceName))
+	// TODO: Figure out if we need to offer users the ability to tune the timeouts


Considering the NumConns set to 2 as default, I would like to add a config for that. Along with that add connection and request timeouts to play when testing

https://github.com/scylladb/gocql/blob/30a838e39b1f0e4d559989b9f30404c5399b3d5e/cluster.go#L40

Also add RetryPolicy

Added config for numConns, request and connection timeouts.

According to the README here,

When working with a Scylla cluster the driver always opens one connection per shard, so NumConns is ignored.

Not sure if setting NumConns will work, maybe this also needs to be tested.

Make sense. we can get rid of that.

go/internal/feast/onlinestore/cassandraonlinestore.go

EXPEbdodla · 2024-11-21T02:02:22Z

go/internal/feast/onlinestore/cassandraonlinestore.go

+		quotedFeatureNames[i] = fmt.Sprintf(`'%s'`, featureName)
+	}
+
+	return fmt.Sprintf(


IN has limit on 100 keys. We should see how we can handle keys with more than 100 features.

Yeah, the way to go once we reach that point is to partition requests and execute them in parallel, adds complexity but would work.

EXPEbdodla · 2024-11-21T02:25:25Z

go/internal/feast/onlinestore/cassandraonlinestore.go

+	var eventTs time.Time
+	var valueStr []byte
+	var deserializedValue types.Value
+	for scanner.Next() {


We can switch to iterator SliceMap(). It's the same code. May reduce the number of lines.

Looks nice, but I think it's more geared towards dynamic columns as it creates a map for every row, but in our case our table schema is fixed, so this surely has to come at a cost.

# Conflicts: # go.mod # go.sum # sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py

first set of changes for ScyllaDB Online Store

6807ffd

msistla96 changed the title ~~feat: Implement scylla db online store~~ feat: Implement ScyllaDB Online Store Sep 13, 2024

msistla96 added 2 commits September 17, 2024 17:49

Modify TTL clause for CREATE TABLE

5154914

Remove using AWS utils for querying EC2 instances for host name; Clea…

5d3bf5f

…nup code a bit

zabarn mentioned this pull request Oct 15, 2024

feat: Lazy Infra Creation with online_store Option Set to True #144

Merged

Fix config parsing, and multi key retrieval

95e0544

acevedosharp force-pushed the implementScyllaDBOnlineStore branch from ed81825 to 95e0544 Compare November 15, 2024 23:02

acevedosharp added 3 commits November 15, 2024 18:07

Add small comment

254074c

Fix build

aa1b41b

Added DD tracing to Go Cassandra store

69d46fb

EXPEbdodla reviewed Nov 18, 2024

View reviewed changes

Refactored common code across go servers

4b560b7

EXPEbdodla reviewed Nov 19, 2024

View reviewed changes

setup.py Outdated Show resolved Hide resolved

EXPEbdodla reviewed Nov 19, 2024

View reviewed changes

go/internal/feast/onlinestore/cassandraonlinestore.go Outdated Show resolved Hide resolved

EXPEbdodla reviewed Nov 19, 2024

View reviewed changes

go/internal/feast/onlinestore/cassandraonlinestore.go Outdated Show resolved Hide resolved

EXPEbdodla reviewed Nov 19, 2024

View reviewed changes

go/internal/feast/onlinestore/cassandraonlinestore.go Outdated Show resolved Hide resolved

EXPEbdodla reviewed Nov 21, 2024

View reviewed changes

go/internal/feast/onlinestore/cassandraonlinestore.go Outdated Show resolved Hide resolved

EXPEbdodla reviewed Nov 21, 2024

View reviewed changes

acevedosharp added 6 commits November 21, 2024 15:03

Address most of the PR comments

3ded707

Merge branch 'master' into implementScyllaDBOnlineStore

484a0d8

# Conflicts: # go.mod # go.sum # sdk/python/feast/infra/online_stores/contrib/cassandra_online_store/cassandra_online_store.py

Run go mod tidy

c6eccd4

Implement multi-key queries using parallel single-key queries

f89a56c

Small details

41e451d

Config changes, optimize data writing and use observer dd trace

2a07d1e


		require github.com/aws/aws-sdk-go v1.34.28

		require github.com/gocql/gocql v1.6.0

feat: Implement ScyllaDB Online Store #138

Are you sure you want to change the base?

feat: Implement ScyllaDB Online Store #138

Conversation

msistla96 commented Sep 11, 2024

What this PR does / why we need it:

Which issue(s) this PR fixes:

Fixes

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

EXPEbdodla Nov 18, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

EXPEbdodla Nov 19, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

acevedosharp Nov 21, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

EXPEbdodla Nov 18, 2024 •

edited

Loading

EXPEbdodla Nov 19, 2024 •

edited

Loading

acevedosharp Nov 21, 2024 •

edited

Loading