@Test public void featureRowShouldBeAllocatedToStoreTagsAccordingToSubscription() { StoreProto.Store bqOnlyStore = newStore("bq*"); StoreProto.Store redisOnlyStore = newStore("redis*"); StoreProto.Store anyStore = newStore("*"); Map<StoreProto.Store, TupleTag<FeatureRowProto.FeatureRow>> storeTags = ImmutableMap.of( bqOnlyStore, new TupleTag<>(), redisOnlyStore, new TupleTag<>(), anyStore, new TupleTag<>()); PCollectionTuple allocatedRows = p.apply( Create.of( FeatureRowProto.FeatureRow.newBuilder().setFeatureSet("project/bq_1").build(), FeatureRowProto.FeatureRow.newBuilder().setFeatureSet("project/bq_2").build(), FeatureRowProto.FeatureRow.newBuilder() .setFeatureSet("project/redis_1") .build(), FeatureRowProto.FeatureRow.newBuilder() .setFeatureSet("project/redis_2") .build(), FeatureRowProto.FeatureRow.newBuilder() .setFeatureSet("project/redis_3") .build())) .apply( FeatureRowToStoreAllocator.newBuilder() .setStoreTags(storeTags) .setStores(ImmutableList.of(bqOnlyStore, redisOnlyStore, anyStore)) .build()); PAssert.that( allocatedRows .get(storeTags.get(bqOnlyStore)) .setCoder(ProtoCoder.of(FeatureRowProto.FeatureRow.class)) .apply("CountBq", Count.globally())) .containsInAnyOrder(2L); PAssert.that( allocatedRows .get(storeTags.get(redisOnlyStore)) .setCoder(ProtoCoder.of(FeatureRowProto.FeatureRow.class)) .apply("CountRedis", Count.globally())) .containsInAnyOrder(3L); PAssert.that( allocatedRows .get(storeTags.get(anyStore)) .setCoder(ProtoCoder.of(FeatureRowProto.FeatureRow.class)) .apply("CountAll", Count.globally())) .containsInAnyOrder(5L); p.run(); }
public static Builder newBuilder() { return new AutoValue_FeatureRowToStoreAllocator.Builder(); }
FeatureRowToStoreAllocator extends PTransform<PCollection<FeatureRowProto.FeatureRow>, PCollectionTuple> { public static Builder newBuilder() { return new AutoValue_FeatureRowToStoreAllocator.Builder(); } }
FeatureRowToStoreAllocator extends PTransform<PCollection<FeatureRowProto.FeatureRow>, PCollectionTuple> { public static Builder newBuilder() { return new AutoValue_FeatureRowToStoreAllocator.Builder(); }  }
FeatureRowToStoreAllocator extends PTransform<PCollection<FeatureRowProto.FeatureRow>, PCollectionTuple> { public static Builder newBuilder() { return new AutoValue_FeatureRowToStoreAllocator.Builder(); }  abstract List<StoreProto.Store> getStores(); abstract Map<StoreProto.Store, TupleTag<FeatureRowProto.FeatureRow>> getStoreTags(); static Builder newBuilder(); @Override PCollectionTuple expand(PCollection<FeatureRowProto.FeatureRow> input); }
FeatureRowToStoreAllocator extends PTransform<PCollection<FeatureRowProto.FeatureRow>, PCollectionTuple> { public static Builder newBuilder() { return new AutoValue_FeatureRowToStoreAllocator.Builder(); }  abstract List<StoreProto.Store> getStores(); abstract Map<StoreProto.Store, TupleTag<FeatureRowProto.FeatureRow>> getStoreTags(); static Builder newBuilder(); @Override PCollectionTuple expand(PCollection<FeatureRowProto.FeatureRow> input);  }