RedisFeatureSink implements FeatureSink { @Override public PTransform<PCollection<FeatureRow>, WriteResult> writer() { int flushFrequencySeconds = DEFAULT_FREQUENCY_SECONDS; if (getRedisClusterConfig() != null) { if (getRedisClusterConfig().getFlushFrequencySeconds() > 0) { flushFrequencySeconds = getRedisClusterConfig().getFlushFrequencySeconds(); } RedisKeySerializer serializer = new RedisKeyPrefixSerializer(getRedisClusterConfig().getKeyPrefix()); return new RedisCustomIO.Write( new RedisClusterIngestionClient(getRedisClusterConfig()), getSpecsView(), serializer) .withFlushFrequency(Duration.standardSeconds(flushFrequencySeconds)) .withBatchSize(DEFAULT_BATCH_SIZE); } else if (getRedisConfig() != null) { if (getRedisConfig().getFlushFrequencySeconds() > 0) { flushFrequencySeconds = getRedisConfig().getFlushFrequencySeconds(); } RedisKeySerializer serializer = new RedisKeyProtoSerializer(); return new RedisCustomIO.Write( new RedisStandaloneIngestionClient(getRedisConfig()), getSpecsView(), serializer) .withFlushFrequency(Duration.standardSeconds(flushFrequencySeconds)) .withBatchSize(DEFAULT_BATCH_SIZE); } else { throw new RuntimeException( "At least one RedisConfig or RedisClusterConfig must be provided to Redis Sink"); } }  static FeatureSink fromConfig(RedisConfig redisConfig); static FeatureSink fromConfig(RedisClusterConfig redisConfig); @Nullable abstract RedisConfig getRedisConfig(); @Nullable abstract RedisClusterConfig getRedisClusterConfig(); abstract Builder toBuilder(); static Builder builder(); RedisFeatureSink withSpecsView(
      PCollectionView<Map<String, Iterable<FeatureSetSpec>>> specsView); @Override PCollection<FeatureSetReference> prepareWrite(
      PCollection<KV<FeatureSetReference, FeatureSetProto.FeatureSetSpec>> featureSetSpecs); @Override PTransform<PCollection<FeatureRow>, WriteResult> writer();  }
@Test public void shouldWriteToRedis() { HashMap<RedisKey, FeatureRow> kvs = new LinkedHashMap<>(); kvs.put( createRedisKey("myproject/fs", field("entity", 1, Enum.INT64)), createFeatureRow( null, Timestamp.getDefaultInstance(), field(hash("feature"), "one", Enum.STRING))); kvs.put( createRedisKey("myproject/fs", field("entity", 2, Enum.INT64)), createFeatureRow( null, Timestamp.getDefaultInstance(), field(hash("feature"), "two", Enum.STRING))); List<FeatureRow> featureRows = ImmutableList.of( createFeatureRow( "myproject/fs", null, field("entity", 1, Enum.INT64), field("feature", "one", Enum.STRING)), createFeatureRow( "myproject/fs", null, field("entity", 2, Enum.INT64), field("feature", "two", Enum.STRING))); p.apply(Create.of(featureRows)).apply(redisFeatureSink.writer()); p.run(); kvs.forEach( (key, value) -> { byte[] actual = sync.get(key.toByteArray()); assertThat(actual, equalTo(value.toByteArray())); }); }
@Test public void shouldConvertRowWithDuplicateEntitiesToValidKey() { FeatureRow offendingRow = createFeatureRow( "myproject/feature_set", Timestamp.newBuilder().setSeconds(10).build(), field("entity_id_primary", 1, Enum.INT32), field("entity_id_primary", 2, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING), field("feature_1", "strValue1", Enum.STRING), field("feature_2", 1001, Enum.INT64)); RedisKey expectedKey = createRedisKey( "myproject/feature_set", field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING)); FeatureRow expectedValue = createFeatureRow( "", Timestamp.newBuilder().setSeconds(10).build(), field(hash("feature_1"), "strValue1", Enum.STRING), field(hash("feature_2"), 1001, Enum.INT64)); p.apply(Create.of(offendingRow)).apply(redisFeatureSink.writer()); p.run(); byte[] actual = sync.get(expectedKey.toByteArray()); assertThat(actual, equalTo(expectedValue.toByteArray())); }
@Test public void shouldConvertRowWithOutOfOrderFieldsToValidKey() { FeatureRow offendingRow = createFeatureRow( "myproject/feature_set", Timestamp.newBuilder().setSeconds(10).build(), field("entity_id_secondary", "a", Enum.STRING), field("entity_id_primary", 1, Enum.INT32), field("feature_2", 1001, Enum.INT64), field("feature_1", "strValue1", Enum.STRING)); RedisKey expectedKey = createRedisKey( "myproject/feature_set", field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING)); FeatureRow expectedValue = createFeatureRow( "", Timestamp.newBuilder().setSeconds(10).build(), field(hash("feature_1"), "strValue1", Enum.STRING), field(hash("feature_2"), 1001, Enum.INT64)); p.apply(Create.of(offendingRow)).apply(redisFeatureSink.writer()); p.run(); byte[] actual = sync.get(expectedKey.toByteArray()); assertThat(actual, equalTo(expectedValue.toByteArray())); }
@Test public void shouldMergeDuplicateFeatureFields() { FeatureRow featureRowWithDuplicatedFeatureFields = createFeatureRow( "myproject/feature_set", Timestamp.newBuilder().setSeconds(10).build(), field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING), field("feature_2", 1001, Enum.INT64), field("feature_1", "strValue1", Enum.STRING), field("feature_1", "strValue1", Enum.STRING)); RedisKey expectedKey = createRedisKey( "myproject/feature_set", field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING)); FeatureRow expectedValue = createFeatureRow( "", Timestamp.newBuilder().setSeconds(10).build(), field(hash("feature_1"), "strValue1", Enum.STRING), field(hash("feature_2"), 1001, Enum.INT64)); p.apply(Create.of(featureRowWithDuplicatedFeatureFields)).apply(redisFeatureSink.writer()); p.run(); byte[] actual = sync.get(expectedKey.toByteArray()); assertThat(actual, equalTo(expectedValue.toByteArray())); }
@Test public void shouldPopulateMissingFeatureValuesWithDefaultInstance() { FeatureRow featureRowWithDuplicatedFeatureFields = createFeatureRow( "myproject/feature_set", Timestamp.newBuilder().setSeconds(10).build(), field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING), field("feature_1", "strValue1", Enum.STRING)); RedisKey expectedKey = createRedisKey( "myproject/feature_set", field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING)); FeatureRow expectedValue = createFeatureRow( "", Timestamp.newBuilder().setSeconds(10).build(), field(hash("feature_1"), "strValue1", Enum.STRING), Field.newBuilder() .setName(hash("feature_2")) .setValue(Value.getDefaultInstance()) .build()); p.apply(Create.of(featureRowWithDuplicatedFeatureFields)).apply(redisFeatureSink.writer()); p.run(); byte[] actual = sync.get(expectedKey.toByteArray()); assertThat(actual, equalTo(expectedValue.toByteArray())); }
@Test public void shouldDeduplicateRowsWithinBatch() { TestStream<FeatureRow> featureRowTestStream = TestStream.create(ProtoCoder.of(FeatureRow.class)) .addElements( createFeatureRow( "myproject/feature_set", Timestamp.newBuilder().setSeconds(20).build(), field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING), field("feature_2", 111, Enum.INT32))) .addElements( createFeatureRow( "myproject/feature_set", Timestamp.newBuilder().setSeconds(10).build(), field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING), field("feature_2", 222, Enum.INT32))) .addElements( createFeatureRow( "myproject/feature_set", Timestamp.getDefaultInstance(), field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING), field("feature_2", 333, Enum.INT32))) .advanceWatermarkToInfinity(); p.apply(featureRowTestStream).apply(redisFeatureSink.writer()); p.run(); RedisKey expectedKey = createRedisKey( "myproject/feature_set", field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING)); FeatureRow expectedValue = createFeatureRow( "", Timestamp.newBuilder().setSeconds(20).build(), Field.newBuilder() .setName(hash("feature_1")) .setValue(Value.getDefaultInstance()) .build(), field(hash("feature_2"), 111, Enum.INT32)); byte[] actual = sync.get(expectedKey.toByteArray()); assertThat(actual, equalTo(expectedValue.toByteArray())); }
@Test public void shouldWriteWithLatterTimestamp() { TestStream<FeatureRow> featureRowTestStream = TestStream.create(ProtoCoder.of(FeatureRow.class)) .addElements( createFeatureRow( "myproject/feature_set", Timestamp.newBuilder().setSeconds(20).build(), field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING), field("feature_2", 111, Enum.INT32))) .addElements( createFeatureRow( "myproject/feature_set", Timestamp.newBuilder().setSeconds(20).build(), field("entity_id_primary", 2, Enum.INT32), field("entity_id_secondary", "b", Enum.STRING), field("feature_2", 222, Enum.INT32))) .addElements( createFeatureRow( "myproject/feature_set", Timestamp.newBuilder().setSeconds(10).build(), field("entity_id_primary", 3, Enum.INT32), field("entity_id_secondary", "c", Enum.STRING), field("feature_2", 333, Enum.INT32))) .advanceWatermarkToInfinity(); RedisKey keyA = createRedisKey( "myproject/feature_set", field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING)); RedisKey keyB = createRedisKey( "myproject/feature_set", field("entity_id_primary", 2, Enum.INT32), field("entity_id_secondary", "b", Enum.STRING)); RedisKey keyC = createRedisKey( "myproject/feature_set", field("entity_id_primary", 3, Enum.INT32), field("entity_id_secondary", "c", Enum.STRING)); sync.set( keyA.toByteArray(), createFeatureRow("", Timestamp.newBuilder().setSeconds(30).build()).toByteArray()); sync.set( keyB.toByteArray(), createFeatureRow("", Timestamp.newBuilder().setSeconds(10).build()).toByteArray()); sync.set( keyC.toByteArray(), createFeatureRow("", Timestamp.newBuilder().setSeconds(10).build()).toByteArray()); p.apply(featureRowTestStream).apply(redisFeatureSink.writer()); p.run(); assertThat( sync.get(keyA.toByteArray()), equalTo(createFeatureRow("", Timestamp.newBuilder().setSeconds(30).build()).toByteArray())); assertThat( sync.get(keyB.toByteArray()), equalTo( createFeatureRow( "", Timestamp.newBuilder().setSeconds(20).build(), Field.newBuilder() .setName(hash("feature_1")) .setValue(Value.getDefaultInstance()) .build(), field(hash("feature_2"), 222, Enum.INT32)) .toByteArray())); assertThat( sync.get(keyC.toByteArray()), equalTo(createFeatureRow("", Timestamp.newBuilder().setSeconds(10).build()).toByteArray())); }
@Test public void shouldOverwriteInvalidRows() { TestStream<FeatureRow> featureRowTestStream = TestStream.create(ProtoCoder.of(FeatureRow.class)) .addElements( createFeatureRow( "myproject/feature_set", Timestamp.newBuilder().setSeconds(20).build(), field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING), field("feature_1", "text", Enum.STRING), field("feature_2", 111, Enum.INT32))) .advanceWatermarkToInfinity(); RedisKey expectedKey = createRedisKey( "myproject/feature_set", field("entity_id_primary", 1, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING)); sync.set(expectedKey.toByteArray(), "some-invalid-data".getBytes()); p.apply(featureRowTestStream).apply(redisFeatureSink.writer()); p.run(); FeatureRow expectedValue = createFeatureRow( "", Timestamp.newBuilder().setSeconds(20).build(), field(hash("feature_1"), "text", Enum.STRING), field(hash("feature_2"), 111, Enum.INT32)); byte[] actual = sync.get(expectedKey.toByteArray()); assertThat(actual, equalTo(expectedValue.toByteArray())); }
@Test public void loadTest() { List<FeatureRow> rows = IntStream.range(0, 10000) .mapToObj( i -> createFeatureRow( "myproject/feature_set", Timestamp.newBuilder().setSeconds(20).build(), field("entity_id_primary", i, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING), field("feature_1", "text", Enum.STRING), field("feature_2", 111, Enum.INT32))) .collect(Collectors.toList()); p.apply(Create.of(rows)).apply(redisFeatureSink.writer()); p.run(); List<byte[]> outcome = IntStream.range(0, 10000) .mapToObj( i -> createRedisKey( "myproject/feature_set", field("entity_id_primary", i, Enum.INT32), field("entity_id_secondary", "a", Enum.STRING)) .toByteArray()) .map(sync::get) .collect(Collectors.toList()); assertThat(outcome, hasSize(10000)); assertThat("All rows were stored", outcome.stream().allMatch(Objects::nonNull)); }