ImportJob { @SuppressWarnings("UnusedReturnValue") public static PipelineResult runPipeline(ImportOptions options) throws IOException { PipelineOptionsValidator.validate(ImportOptions.class, options); Pipeline pipeline = Pipeline.create(options); log.info("Starting import job with settings: \n{}", options.toString()); List<Store> stores = SpecUtil.parseStoreJsonList(options.getStoresJson()); Source source = SpecUtil.parseSourceJson(options.getSourceJson()); SpecsStreamingUpdateConfig specsStreamingUpdateConfig = SpecUtil.parseSpecsStreamingUpdateConfig(options.getSpecsStreamingUpdateConfigJson()); PCollection<KV<FeatureSetReference, FeatureSetSpec>> featureSetSpecs = pipeline.apply( "ReadFeatureSetSpecs", ReadFeatureSetSpecs.newBuilder() .setSource(source) .setStores(stores) .setSpecsStreamingUpdateConfig(specsStreamingUpdateConfig) .build()); PCollectionView<Map<String, Iterable<FeatureSetSpec>>> globalSpecView = featureSetSpecs .apply(MapElements.via(new ReferenceToString())) .apply("GlobalSpecView", View.asMultimap()); PCollectionTuple convertedFeatureRows = pipeline.apply( "ReadFeatureRowFromSource", ReadFromSource.newBuilder() .setSource(source) .setSuccessTag(FEATURE_ROW_OUT) .setFailureTag(DEADLETTER_OUT) .setKafkaConsumerProperties( options.getKafkaConsumerProperties() == null ? new HashMap<>() : options.getKafkaConsumerProperties()) .build()); PCollectionTuple validatedRows = convertedFeatureRows .get(FEATURE_ROW_OUT) .apply( ProcessAndValidateFeatureRows.newBuilder() .setDefaultProject(options.getDefaultFeastProject()) .setFeatureSetSpecs(globalSpecView) .setSuccessTag(FEATURE_ROW_OUT) .setFailureTag(DEADLETTER_OUT) .build()); Map<Store, TupleTag<FeatureRow>> storeTags = stores.stream() .map(s -> Pair.of(s, new TupleTag<FeatureRow>())) .collect(Collectors.toMap(Pair::getLeft, Pair::getRight)); PCollectionTuple storeAllocatedRows = validatedRows .get(FEATURE_ROW_OUT) .apply( FeatureRowToStoreAllocator.newBuilder() .setStores(stores) .setStoreTags(storeTags) .build()); PCollectionList<FeatureSetReference> sinkReadiness = PCollectionList.empty(pipeline); for (Store store : stores) { FeatureSink featureSink = getFeatureSink(store); sinkReadiness = sinkReadiness.and(featureSink.prepareWrite(featureSetSpecs)); PCollection<FeatureRow> rowsForStore = storeAllocatedRows.get(storeTags.get(store)).setCoder(ProtoCoder.of(FeatureRow.class)); rowsForStore.apply( "WriteInflightMetrics", WriteInflightMetricsTransform.create(store.getName())); WriteResult writeFeatureRows = rowsForStore.apply("WriteFeatureRowToStore", featureSink.writer()); if (options.getDeadLetterTableSpec() != null) { DeadletterSink deadletterSink = new BigQueryDeadletterSink(options.getDeadLetterTableSpec()); writeFeatureRows .getFailedInserts() .apply("WriteFailedElements_WriteFeatureRowToStore", deadletterSink.write()); } writeFeatureRows .getSuccessfulInserts() .apply("WriteSuccessMetrics", WriteSuccessMetricsTransform.create(store.getName())); writeFeatureRows .getFailedInserts() .apply("WriteFailureMetrics", WriteFailureMetricsTransform.create(store.getName())); } if (options.getDeadLetterTableSpec() != null) { DeadletterSink deadletterSink = new BigQueryDeadletterSink(options.getDeadLetterTableSpec()); convertedFeatureRows .get(DEADLETTER_OUT) .apply("WriteFailedElements_ReadFromSource", deadletterSink.write()); validatedRows .get(DEADLETTER_OUT) .apply("WriteFailedElements_ValidateRows", deadletterSink.write()); } sinkReadiness .apply(Flatten.pCollections()) .apply( "WriteAck", WriteFeatureSetSpecAck.newBuilder() .setSinksCount(stores.size()) .setSpecsStreamingUpdateConfig(specsStreamingUpdateConfig) .build()); return pipeline.run(); }  static void main(String[] args); @SuppressWarnings("UnusedReturnValue") static PipelineResult runPipeline(ImportOptions options);  }
@Test public void runPipeline_ShouldWriteToRedisCorrectlyGivenValidSpecAndFeatureRow() throws IOException, InterruptedException { Source featureSource = Source.newBuilder() .setType(SourceType.KAFKA) .setKafkaSourceConfig( KafkaSourceConfig.newBuilder() .setBootstrapServers(kafkaContainer.getBootstrapServers()) .setTopic(KAFKA_TOPIC) .build()) .build(); IngestionJobProto.SpecsStreamingUpdateConfig specsStreamingUpdateConfig = IngestionJobProto.SpecsStreamingUpdateConfig.newBuilder() .setSource( KafkaSourceConfig.newBuilder() .setBootstrapServers(kafkaContainer.getBootstrapServers()) .setTopic(KAFKA_SPECS_TOPIC) .build()) .setAck( KafkaSourceConfig.newBuilder() .setBootstrapServers(kafkaContainer.getBootstrapServers()) .setTopic(KAFKA_SPECS_ACK_TOPIC) .build()) .build(); FeatureSetSpec spec = FeatureSetSpec.newBuilder() .setName("feature_set") .setProject("myproject") .addEntities( EntitySpec.newBuilder() .setName("entity_id_primary") .setValueType(Enum.INT32) .build()) .addEntities( EntitySpec.newBuilder() .setName("entity_id_secondary") .setValueType(Enum.STRING) .build()) .addFeatures( FeatureSpec.newBuilder() .setName("feature_1") .setValueType(Enum.STRING_LIST) .build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_2").setValueType(Enum.STRING).build()) .addFeatures( FeatureSpec.newBuilder().setName("feature_3").setValueType(Enum.INT64).build()) .setSource(featureSource) .build(); FeatureSet featureSet = FeatureSet.newBuilder().setSpec(spec).build(); Store redis = Store.newBuilder() .setName(StoreType.REDIS.toString()) .setType(StoreType.REDIS) .setRedisConfig( RedisConfig.newBuilder() .setHost(redisContainer.getHost()) .setPort(redisContainer.getFirstMappedPort()) .build()) .addSubscriptions( Subscription.newBuilder() .setProject(spec.getProject()) .setName(spec.getName()) .build()) .build(); ImportOptions options = PipelineOptionsFactory.create().as(ImportOptions.class); options.setSpecsStreamingUpdateConfigJson( JsonFormat.printer().print(specsStreamingUpdateConfig)); options.setSourceJson(JsonFormat.printer().print(featureSource)); options.setStoresJson(Collections.singletonList(JsonFormat.printer().print(redis))); options.setDefaultFeastProject("myproject"); options.setProject(""); options.setBlockOnRun(false); List<Pair<String, FeatureRow>> input = new ArrayList<>(); Map<RedisKey, FeatureRow> expected = new HashMap<>(); LOGGER.info("Generating test data ..."); IntStream.range(0, IMPORT_JOB_SAMPLE_FEATURE_ROW_SIZE) .forEach( i -> { FeatureRow randomRow = TestUtil.createRandomFeatureRow(featureSet.getSpec()); RedisKey redisKey = TestUtil.createRedisKey(featureSet.getSpec(), randomRow); input.add(Pair.of("", randomRow)); List<FieldProto.Field> fields = randomRow.getFieldsList().stream() .filter( field -> spec.getFeaturesList().stream() .map(FeatureSpec::getName) .collect(Collectors.toList()) .contains(field.getName())) .map( field -> field.toBuilder().setName(TestUtil.hash(field.getName())).build()) .collect(Collectors.toList()); randomRow = randomRow .toBuilder() .clearFields() .addAllFields(fields) .clearFeatureSet() .build(); expected.put(redisKey, randomRow); }); LOGGER.info("Starting Import Job with the following options: {}", options.toString()); PipelineResult pipelineResult = ImportJob.runPipeline(options); Thread.sleep(Duration.standardSeconds(IMPORT_JOB_READY_DURATION_SEC).getMillis()); assertThat(pipelineResult.getState(), equalTo(State.RUNNING)); LOGGER.info("Publishing {} Feature Row messages to Kafka ...", input.size()); TestUtil.publishToKafka( kafkaContainer.getBootstrapServers(), KAFKA_SPECS_TOPIC, ImmutableList.of(Pair.of(getFeatureSetStringRef(spec), spec)), ByteArraySerializer.class, KAFKA_PUBLISH_TIMEOUT_SEC); TestUtil.publishToKafka( kafkaContainer.getBootstrapServers(), KAFKA_TOPIC, input, ByteArraySerializer.class, KAFKA_PUBLISH_TIMEOUT_SEC); TestUtil.waitUntilAllElementsAreWrittenToStore( pipelineResult, Duration.standardSeconds(IMPORT_JOB_MAX_RUN_DURATION_SEC), Duration.standardSeconds(IMPORT_JOB_CHECK_INTERVAL_DURATION_SEC)); LOGGER.info("Validating the actual values written to Redis ..."); RedisClient redisClient = RedisClient.create( new RedisURI( redisContainer.getHost(), redisContainer.getFirstMappedPort(), java.time.Duration.ofMillis(2000))); StatefulRedisConnection<byte[], byte[]> connection = redisClient.connect(new ByteArrayCodec()); RedisCommands<byte[], byte[]> sync = connection.sync(); for (Map.Entry<RedisKey, FeatureRow> entry : expected.entrySet()) { RedisKey key = entry.getKey(); FeatureRow expectedValue = entry.getValue(); byte[] actualByteValue = sync.get(key.toByteArray()); assertThat("Key not found in Redis: " + key, actualByteValue, notNullValue()); FeatureRow actualValue = null; actualValue = FeatureRow.parseFrom(actualByteValue); assertThat(actualValue, equalTo(expectedValue)); } redisClient.shutdown(); }