Skip to content

Commit 03b0ad6

Browse files
author
zhilingc
committed
Extract kafka record to feature row transform to separate dofn
1 parent e1e7815 commit 03b0ad6

4 files changed

Lines changed: 134 additions & 73 deletions

File tree

‎ingestion/src/main/java/feast/ingestion/transform/ReadFromSource.java‎

Lines changed: 10 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import com.google.protobuf.InvalidProtocolBufferException;
66
import feast.core.SourceProto.Source;
77
import feast.core.SourceProto.SourceType;
8+
import feast.ingestion.transform.fn.KafkaRecordToFeatureRowDoFn;
89
import feast.ingestion.values.FailedElement;
910
import feast.ingestion.values.Field;
1011
import feast.types.FeatureRowProto.FeatureRow;
@@ -26,6 +27,7 @@
2627

2728
@AutoValue
2829
public abstract class ReadFromSource extends PTransform<PBegin, PCollectionTuple> {
30+
2931
public abstract Source getSource();
3032

3133
public abstract Map<String, Field> getFieldByName();
@@ -44,6 +46,7 @@ public static Builder newBuilder() {
4446

4547
@AutoValue.Builder
4648
public abstract static class Builder {
49+
4750
public abstract Builder setSource(Source source);
4851

4952
public abstract Builder setFeatureSetName(String featureSetName);
@@ -89,72 +92,13 @@ public PCollectionTuple expand(PBegin input) {
8992
.withReadCommitted()
9093
.commitOffsetsInFinalize())
9194
.apply(
92-
"KafkaRecordToFeatureRow",
93-
ParDo.of(
94-
new DoFn<KafkaRecord<byte[], byte[]>, FeatureRow>() {
95-
@ProcessElement
96-
public void processElement(ProcessContext context) {
97-
byte[] value = context.element().getKV().getValue();
98-
FeatureRow featureRow;
99-
100-
try {
101-
featureRow = FeatureRow.parseFrom(value);
102-
} catch (InvalidProtocolBufferException e) {
103-
context.output(
104-
getFailureTag(),
105-
FailedElement.newBuilder()
106-
.setTransformName("KafkaRecordToFeatureRow")
107-
.setStackTrace(ExceptionUtils.getStackTrace(e))
108-
.setJobName(context.getPipelineOptions().getJobName())
109-
.setPayload(new String(Base64.getEncoder().encode(value)))
110-
.setErrorMessage(e.getMessage())
111-
.build());
112-
return;
113-
}
114-
115-
// If FeatureRow contains field names that do not exist as EntitySpec
116-
// or FeatureSpec in FeatureSetSpec, mark the FeatureRow as FailedElement.
117-
String error = null;
118-
for (FieldProto.Field field : featureRow.getFieldsList()) {
119-
if (!getFieldByName().containsKey(field.getName())) {
120-
error =
121-
String.format(
122-
"FeatureRow contains field '%s' which do not exists in FeatureSet '%s' version '%d'. Please check the FeatureRow data.",
123-
field.getName(), getFeatureSetName(), getFeatureSetVersion());
124-
break;
125-
}
126-
// If value is set in the FeatureRow, make sure the value type matches
127-
// that defined in FeatureSetSpec
128-
if (!field.getValue().getValCase().equals(ValCase.VAL_NOT_SET)) {
129-
int expectedTypeFieldNumber =
130-
getFieldByName().get(field.getName()).getType().getNumber();
131-
int actualTypeFieldNumber = field.getValue().getValCase().getNumber();
132-
if (expectedTypeFieldNumber != actualTypeFieldNumber) {
133-
error =
134-
String.format(
135-
"FeatureRow contains field '%s' with invalid type '%s'. Feast expects the field type to match that in FeatureSet '%s'. Please check the FeatureRow data.",
136-
field.getName(),
137-
field.getValue().getValCase(),
138-
getFieldByName().get(field.getName()).getType());
139-
break;
140-
}
141-
}
142-
}
143-
144-
if (error != null) {
145-
context.output(
146-
getFailureTag(),
147-
FailedElement.newBuilder()
148-
.setTransformName("KafkaRecordToFeatureRow")
149-
.setJobName(context.getPipelineOptions().getJobName())
150-
.setPayload(featureRow.toString())
151-
.setErrorMessage(error)
152-
.build());
153-
} else {
154-
context.output(featureRow);
155-
}
156-
}
157-
})
95+
"KafkaRecordToFeatureRow", ParDo.of(KafkaRecordToFeatureRowDoFn.newBuilder()
96+
.setFeatureSetName(getFeatureSetName())
97+
.setFeatureSetVersion(getFeatureSetVersion())
98+
.setFieldByName(getFieldByName())
99+
.setSuccessTag(getSuccessTag())
100+
.setFailureTag(getFailureTag())
101+
.build())
158102
.withOutputTags(getSuccessTag(), TupleTagList.of(getFailureTag())));
159103
}
160104

‎ingestion/src/main/java/feast/ingestion/transform/fn/KafkaRecordToFeatureRowDoFn.java‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
package feast.ingestion.transform.fn;
2+
3+
import com.google.auto.value.AutoValue;
4+
import com.google.protobuf.InvalidProtocolBufferException;
5+
import feast.ingestion.values.FailedElement;
6+
import feast.ingestion.values.Field;
7+
import feast.types.FeatureRowProto.FeatureRow;
8+
import feast.types.FieldProto;
9+
import feast.types.ValueProto.Value.ValCase;
10+
import java.util.Base64;
11+
import java.util.Map;
12+
import org.apache.beam.sdk.io.kafka.KafkaRecord;
13+
import org.apache.beam.sdk.transforms.DoFn;
14+
import org.apache.beam.sdk.values.TupleTag;
15+
import org.apache.commons.lang3.exception.ExceptionUtils;
16+
17+
@AutoValue
18+
public abstract class KafkaRecordToFeatureRowDoFn extends DoFn<KafkaRecord<byte[], byte[]>, FeatureRow> {
19+
20+
public abstract String getFeatureSetName();
21+
22+
public abstract int getFeatureSetVersion();
23+
24+
public abstract Map<String, Field> getFieldByName();
25+
26+
public abstract TupleTag<FeatureRow> getSuccessTag();
27+
28+
public abstract TupleTag<FailedElement> getFailureTag();
29+
30+
public static KafkaRecordToFeatureRowDoFn.Builder newBuilder() {
31+
return new AutoValue_KafkaRecordToFeatureRowDoFn.Builder();
32+
}
33+
34+
@AutoValue.Builder
35+
public abstract static class Builder {
36+
37+
public abstract Builder setFeatureSetName(String featureSetName);
38+
39+
public abstract Builder setFeatureSetVersion(int featureSetVersion);
40+
41+
public abstract Builder setFieldByName(Map<String, Field> fieldByName);
42+
43+
public abstract Builder setSuccessTag(TupleTag<FeatureRow> successTag);
44+
45+
public abstract Builder setFailureTag(TupleTag<FailedElement> failureTag);
46+
47+
public abstract KafkaRecordToFeatureRowDoFn build();
48+
}
49+
50+
@ProcessElement
51+
public void processElement(ProcessContext context) {
52+
byte[] value = context.element().getKV().getValue();
53+
FeatureRow featureRow;
54+
55+
try {
56+
featureRow = FeatureRow.parseFrom(value);
57+
} catch (InvalidProtocolBufferException e) {
58+
context.output(
59+
getFailureTag(),
60+
FailedElement.newBuilder()
61+
.setTransformName("KafkaRecordToFeatureRow")
62+
.setStackTrace(ExceptionUtils.getStackTrace(e))
63+
.setJobName(context.getPipelineOptions().getJobName())
64+
.setPayload(new String(Base64.getEncoder().encode(value)))
65+
.setErrorMessage(e.getMessage())
66+
.build());
67+
return;
68+
}
69+
70+
// If FeatureRow contains field names that do not exist as EntitySpec
71+
// or FeatureSpec in FeatureSetSpec, mark the FeatureRow as FailedElement.
72+
String error = null;
73+
for (FieldProto.Field field : featureRow.getFieldsList()) {
74+
if (!getFieldByName().containsKey(field.getName())) {
75+
error =
76+
String.format(
77+
"FeatureRow contains field '%s' which do not exists in FeatureSet '%s' version '%d'. Please check the FeatureRow data.",
78+
field.getName(), getFeatureSetName(), getFeatureSetVersion());
79+
break;
80+
}
81+
// If value is set in the FeatureRow, make sure the value type matches
82+
// that defined in FeatureSetSpec
83+
if (!field.getValue().getValCase().equals(ValCase.VAL_NOT_SET)) {
84+
int expectedTypeFieldNumber =
85+
getFieldByName().get(field.getName()).getType().getNumber();
86+
int actualTypeFieldNumber = field.getValue().getValCase().getNumber();
87+
if (expectedTypeFieldNumber != actualTypeFieldNumber) {
88+
error =
89+
String.format(
90+
"FeatureRow contains field '%s' with invalid type '%s'. Feast expects the field type to match that in FeatureSet '%s'. Please check the FeatureRow data.",
91+
field.getName(),
92+
field.getValue().getValCase(),
93+
getFieldByName().get(field.getName()).getType());
94+
break;
95+
}
96+
}
97+
}
98+
99+
if (error != null) {
100+
context.output(
101+
getFailureTag(),
102+
FailedElement.newBuilder()
103+
.setTransformName("KafkaRecordToFeatureRow")
104+
.setJobName(context.getPipelineOptions().getJobName())
105+
.setPayload(featureRow.toString())
106+
.setErrorMessage(error)
107+
.build());
108+
} else {
109+
context.output(featureRow);
110+
}
111+
}
112+
}

‎ingestion/src/main/java/feast/store/serving/redis/RedisCustomIO.java‎

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,11 @@
1818
// package io.suryawirawan.henry.beam.redis.io;
1919
package feast.store.serving.redis;
2020

21+
import lombok.AllArgsConstructor;
2122
import lombok.Builder;
2223
import lombok.Data;
24+
import lombok.Getter;
25+
import lombok.NoArgsConstructor;
2326
import org.apache.beam.sdk.coders.AvroCoder;
2427
import org.apache.beam.sdk.coders.DefaultCoder;
2528
import org.apache.beam.sdk.transforms.DoFn;
@@ -87,14 +90,16 @@ public enum Method {
8790
}
8891

8992
@Builder
90-
@Data
93+
@Getter
94+
@NoArgsConstructor
95+
@AllArgsConstructor
9196
@DefaultCoder(value = AvroCoder.class)
9297
public static class RedisMutation {
93-
private final Method method;
94-
private final byte[] key;
95-
private final byte[] value;
96-
private final long expiryMillis;
97-
private final long score; // Score is only utilized when method is ZSET
98+
private Method method;
99+
private byte[] key;
100+
private byte[] value;
101+
private long expiryMillis;
102+
private long score; // Score is only utilized when method is ZSET
98103
}
99104

100105
/** ServingStoreWrite data to a Redis server. */

‎ingestion/src/test/java/feast/store/serving/redis/RedisStoreOptionsTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ public void testParse() {
3434

3535
RedisStoreOptions opts = OptionsParser.parse(map, RedisStoreOptions.class);
3636
assertEquals("localhost", opts.host);
37-
assertEquals(1234, (int) opts.port);
37+
assertEquals(1234, opts.port);
3838
}
3939

4040
@Test(expected = IllegalArgumentException.class)

0 commit comments

Comments
 (0)