Skip to content

Commit af2e74b

Browse files
terryyylimterryyylim
committed
Fix IT
Signed-off-by: Terence Lim <terencelimxp@gmail.com>
1 parent b309886 commit af2e74b

File tree

3 files changed

+37
-20
lines changed

3 files changed

+37
-20
lines changed

‎serving/src/test/java/feast/serving/it/BaseAuthIT.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ public class BaseAuthIT {
5959
static final String CASSANDRA_DATACENTER = "datacenter1";
6060
static final String CASSANDRA_KEYSPACE = "feast";
6161
static final String CASSANDRA_SCHEMA_TABLE = "feast_schema_reference";
62+
static final String CASSANDRA_ENTITY_KEY = "key";
6263

6364
static final int FEAST_CORE_PORT = 6565;
6465

‎serving/src/test/java/feast/serving/it/ServingServiceCassandraIT.java‎

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import feast.proto.serving.ServingAPIProto;
3232
import feast.proto.serving.ServingServiceGrpc;
3333
import feast.proto.types.ValueProto;
34+
import io.grpc.ManagedChannel;
3435
import java.io.ByteArrayOutputStream;
3536
import java.io.File;
3637
import java.io.IOException;
@@ -48,6 +49,7 @@
4849
import org.apache.avro.io.Encoder;
4950
import org.apache.avro.io.EncoderFactory;
5051
import org.junit.ClassRule;
52+
import org.junit.jupiter.api.AfterAll;
5153
import org.junit.jupiter.api.BeforeAll;
5254
import org.junit.jupiter.api.Test;
5355
import org.springframework.boot.test.context.SpringBootTest;
@@ -187,12 +189,13 @@ static void globalSetup() throws IOException {
187189
// Single Entity Cassandra Table
188190
cqlSession.execute(
189191
String.format(
190-
"CREATE TABLE %s.%s (key BLOB PRIMARY KEY, schema_ref BLOB);",
191-
CASSANDRA_KEYSPACE, cassandraTableName));
192+
"CREATE TABLE %s.%s (key BLOB PRIMARY KEY);", CASSANDRA_KEYSPACE, cassandraTableName));
192193

193194
// Add column families
194195
cqlSession.execute(
195-
String.format("ALTER TABLE %s.%s ADD rides BLOB;", CASSANDRA_KEYSPACE, cassandraTableName));
196+
String.format(
197+
"ALTER TABLE %s.%s ADD (%s BLOB, %s__schema_ref BLOB);",
198+
CASSANDRA_KEYSPACE, cassandraTableName, ridesFeatureTableName, ridesFeatureTableName));
196199

197200
/** Single Entity Ingestion Workflow */
198201
Schema ftSchema =
@@ -216,14 +219,18 @@ static void globalSetup() throws IOException {
216219
.build();
217220
byte[] entityFeatureKey =
218221
String.valueOf(DataGenerator.createInt64Value(1).getInt64Val()).getBytes();
219-
byte[] entityFeatureValue = createEntityValue(ftSchema, schemaReference, record);
222+
byte[] entityFeatureValue = createEntityValue(ftSchema, record);
220223
byte[] schemaKey = createSchemaKey(schemaReference);
221224

222225
PreparedStatement statement =
223226
cqlSession.prepare(
224227
String.format(
225-
"INSERT INTO %s.%s (key, schema_ref, rides) VALUES (?, ?, ?)",
226-
CASSANDRA_KEYSPACE, cassandraTableName));
228+
"INSERT INTO %s.%s (%s, %s__schema_ref, %s) VALUES (?, ?, ?)",
229+
CASSANDRA_KEYSPACE,
230+
cassandraTableName,
231+
CASSANDRA_ENTITY_KEY,
232+
ridesFeatureTableName,
233+
ridesFeatureTableName));
227234
cqlSession.execute(
228235
statement.bind(
229236
ByteBuffer.wrap(entityFeatureKey),
@@ -262,14 +269,11 @@ private static byte[] createSchemaKey(byte[] schemaReference) throws IOException
262269
return schemaKey;
263270
}
264271

265-
private static byte[] createEntityValue(
266-
Schema schema, byte[] schemaReference, GenericRecord record) throws IOException {
272+
private static byte[] createEntityValue(Schema schema, GenericRecord record) throws IOException {
267273
// Entity-Feature Row
268274
byte[] avroSerializedFeatures = recordToAvro(record, schema);
269275

270276
ByteArrayOutputStream concatOutputStream = new ByteArrayOutputStream();
271-
concatOutputStream.write(schemaReference);
272-
concatOutputStream.write("".getBytes());
273277
concatOutputStream.write(avroSerializedFeatures);
274278
byte[] entityFeatureValue = concatOutputStream.toByteArray();
275279

@@ -286,6 +290,11 @@ private static byte[] recordToAvro(GenericRecord datum, Schema schema) throws IO
286290
return output.toByteArray();
287291
}
288292

293+
@AfterAll
294+
static void tearDown() {
295+
((ManagedChannel) servingStub.getChannel()).shutdown();
296+
}
297+
289298
@Test
290299
public void shouldRegisterSingleEntityAndGetOnlineFeatures() {
291300
String projectName = "default";

‎storage/connectors/cassandra/src/main/java/feast/storage/connectors/cassandra/retriever/CassandraOnlineRetriever.java‎

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,7 @@
2929
import feast.storage.api.retriever.OnlineRetrieverV2;
3030
import java.io.IOException;
3131
import java.nio.ByteBuffer;
32-
import java.util.Collections;
33-
import java.util.List;
34-
import java.util.Map;
35-
import java.util.Objects;
32+
import java.util.*;
3633
import java.util.function.Function;
3734
import java.util.stream.Collectors;
3835
import java.util.stream.StreamSupport;
@@ -50,7 +47,6 @@ public class CassandraOnlineRetriever implements OnlineRetrieverV2 {
5047

5148
private static String ENTITY_KEY = "key";
5249
private static String SCHEMA_REF_KEY = "schema_ref";
53-
private static String TIMESTAMP_COLUMN = String.format("writetime(%s)", SCHEMA_REF_KEY);
5450

5551
public CassandraOnlineRetriever(CqlSession session) {
5652
this.session = session;
@@ -140,9 +136,12 @@ private List<Feature> decodeFeatures(
140136
CassandraSchemaRegistry.SchemaReference schemaReference =
141137
new CassandraSchemaRegistry.SchemaReference(schemaRefKey);
142138

139+
// Convert ByteBuffer to ByteArray
140+
byte[] bytesArray = new byte[value.remaining()];
141+
value.get(bytesArray, 0, bytesArray.length);
143142
Schema schema = schemaRegistry.getSchema(schemaReference);
144143
GenericDatumReader<GenericRecord> reader = new GenericDatumReader<>(schema);
145-
Decoder decoder = DecoderFactory.get().binaryDecoder(value.array(), null);
144+
Decoder decoder = DecoderFactory.get().binaryDecoder(bytesArray, null);
146145
GenericRecord record = reader.read(null, decoder);
147146

148147
return featureReferences.stream()
@@ -205,15 +204,19 @@ public List<List<Feature>> getOnlineFeatures(
205204
private Map<ByteBuffer, Row> getFeaturesFromCassandra(
206205
String tableName, List<ByteBuffer> rowKeys, List<String> columnFamilies) {
207206
SelectFrom query = QueryBuilder.selectFrom(String.format("\"%s\"", tableName));
207+
List<String> schemaRefKeyColumns =
208+
columnFamilies.stream()
209+
.map(cf -> String.format("%s__%s", cf, SCHEMA_REF_KEY))
210+
.collect(Collectors.toList());
208211

209212
BoundStatement statement =
210213
session
211214
.prepare(
212215
query
213216
.columns(columnFamilies)
214-
.column(SCHEMA_REF_KEY)
217+
.columns(schemaRefKeyColumns)
215218
.column(ENTITY_KEY)
216-
.writeTime(SCHEMA_REF_KEY)
219+
.writeTime(schemaRefKeyColumns.get(0))
217220
.whereColumn(ENTITY_KEY)
218221
.in(QueryBuilder.bindMarker())
219222
.build())
@@ -246,7 +249,11 @@ private List<List<Feature>> convertRowToFeature(
246249
Row row = rows.get(rowKey);
247250

248251
String featureTableColumn = columnFamilies.get(0);
249-
ByteBuffer schemaRefKey = row.getByteBuffer(SCHEMA_REF_KEY);
252+
String schemaRefKeyColumn =
253+
String.format("%s__%s", featureTableColumn, SCHEMA_REF_KEY);
254+
String timestampColumn = String.format("writetime(%s)", schemaRefKeyColumn);
255+
256+
ByteBuffer schemaRefKey = row.getByteBuffer(schemaRefKeyColumn);
250257
ByteBuffer featureValues = row.getByteBuffer(featureTableColumn);
251258

252259
List<Feature> features;
@@ -256,7 +263,7 @@ private List<List<Feature>> convertRowToFeature(
256263
schemaRefKey,
257264
featureValues,
258265
featureReferences,
259-
row.getLong(TIMESTAMP_COLUMN));
266+
row.getLong(timestampColumn));
260267
} catch (IOException e) {
261268
throw new RuntimeException("Failed to decode features from Cassandra");
262269
}

0 commit comments

Comments
 (0)