1919
2020import com .google .common .base .Preconditions ;
2121import com .google .common .collect .Lists ;
22+ import feast .ingestion .exceptions .ErrorsHandler ;
23+ import feast .ingestion .model .Specs ;
24+ import feast .ingestion .transform .FeatureIO .Write ;
25+ import feast .ingestion .transform .SplitFeatures .MultiOutputSplit ;
26+ import feast .ingestion .values .PFeatureRows ;
27+ import feast .specs .FeatureSpecProto .FeatureSpec ;
28+ import feast .specs .StorageSpecProto .StorageSpec ;
29+ import feast .storage .FeatureStore ;
30+ import feast .storage .noop .NoOpIO ;
31+ import feast .types .FeatureRowExtendedProto .Attempt ;
32+ import feast .types .FeatureRowExtendedProto .Error ;
33+ import feast .types .FeatureRowExtendedProto .FeatureRowExtended ;
2234import java .util .Collection ;
2335import java .util .HashMap ;
2436import java .util .List ;
3547import org .apache .beam .sdk .values .PCollectionList ;
3648import org .apache .beam .sdk .values .PCollectionTuple ;
3749import org .apache .beam .sdk .values .TupleTag ;
38- import feast .ingestion .exceptions .ErrorsHandler ;
39- import feast .ingestion .model .Specs ;
40- import feast .ingestion .transform .FeatureIO .Write ;
41- import feast .ingestion .transform .SplitFeatures .MultiOutputSplit ;
42- import feast .ingestion .values .PFeatureRows ;
43- import feast .specs .FeatureSpecProto .FeatureSpec ;
44- import feast .specs .StorageSpecProto .StorageSpec ;
45- import feast .storage .FeatureStore ;
46- import feast .storage .noop .NoOpIO ;
47- import feast .types .FeatureRowExtendedProto .Attempt ;
48- import feast .types .FeatureRowExtendedProto .Error ;
49- import feast .types .FeatureRowExtendedProto .FeatureRowExtended ;
5050
5151@ AllArgsConstructor
5252@ Slf4j
5353public class SplitOutputByStore extends PTransform <PFeatureRows , PFeatureRows > {
54+
5455 private Collection <? extends FeatureStore > stores ;
5556 private SerializableFunction <FeatureSpec , String > selector ;
5657 private Specs specs ;
5758
5859 @ Override
5960 public PFeatureRows expand (PFeatureRows input ) {
6061 Map <String , Write > transforms = getFeatureStoreTransforms ();
62+ transforms .put ("" , new NoOpIO .Write ());
6163 Set <String > keys = transforms .keySet ();
6264 Preconditions .checkArgument (transforms .size () > 0 , "no write transforms found" );
6365
@@ -102,6 +104,7 @@ private Map<String, Write> getFeatureStoreTransforms() {
102104
103105 @ AllArgsConstructor
104106 public static class WriteTags extends PTransform <PCollectionTuple , PFeatureRows > {
107+
105108 private Map <TupleTag <FeatureRowExtended >, Write > transforms ;
106109 private TupleTag <FeatureRowExtended > mainTag ;
107110
@@ -120,9 +123,7 @@ public PFeatureRows expand(PCollectionTuple input) {
120123 }
121124
122125 String message =
123- "FeatureRow with output tag.no matching storage, these feature's "
124- + "specs may be specifying a store which was unknown when "
125- + "ingestion started as they somehow passed validation. " ;
126+ "FeatureRows have no matching write transform, these rows should not have passed validation." ;
126127 PCollection <FeatureRowExtended > errors =
127128 input .get (mainTag ).apply (ParDo .of (new WithErrors (getName (), message )));
128129
@@ -131,8 +132,11 @@ public PFeatureRows expand(PCollectionTuple input) {
131132 }
132133 }
133134
134- /** Sets the last attempt error for all rows with a given exception */
135+ /**
136+ * Sets the last attempt error for all rows with a given exception
137+ */
135138 public static class WithErrors extends DoFn <FeatureRowExtended , FeatureRowExtended > {
139+
136140 private Error error ;
137141
138142 public WithErrors (Error error ) {
0 commit comments