Skip to content

Commit 8d7d49a

Browse files
committed
Add option to write metrics to InfluxDB in ingestion
1 parent 3d4e2a1 commit 8d7d49a

5 files changed

Lines changed: 193 additions & 13 deletions

File tree

ingestion/pom.xml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -349,5 +349,11 @@
349349
<version>${org.apache.beam.version}</version>
350350
<scope>runtime</scope>
351351
</dependency>
352+
353+
<dependency>
354+
<groupId>org.influxdb</groupId>
355+
<artifactId>influxdb-java</artifactId>
356+
<version>2.15</version>
357+
</dependency>
352358
</dependencies>
353359
</project>

ingestion/src/main/java/feast/ingestion/ImportJob.java

Lines changed: 38 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import feast.ingestion.transform.ToFeatureRowExtended;
3939
import feast.ingestion.transform.ValidateTransform;
4040
import feast.ingestion.transform.WarehouseStoreTransform;
41+
import feast.ingestion.transform.WriteFeatureMetricsToInfluxDB;
4142
import feast.ingestion.transform.fn.ConvertTypesDoFn;
4243
import feast.ingestion.transform.fn.LoggerDoFn;
4344
import feast.ingestion.values.PFeatureRows;
@@ -110,6 +111,7 @@ public static void main(String[] args) {
110111
mainWithResult(args);
111112
}
112113

114+
@SuppressWarnings({"UnusedReturnValue", "WeakerAccess"})
113115
public static PipelineResult mainWithResult(String[] args) {
114116
log.info("Arguments: " + Arrays.toString(args));
115117
ImportJobPipelineOptions options =
@@ -118,12 +120,30 @@ public static PipelineResult mainWithResult(String[] args) {
118120
options.setJobName(generateName());
119121
}
120122
log.info("options: " + options.toString());
121-
ImportJobSpecs importJobSpecs = new ImportJobSpecsSupplier(options.getWorkspace())
122-
.get();
123+
ImportJobSpecs importJobSpecs = new ImportJobSpecsSupplier(options.getWorkspace()).get();
123124
Injector injector =
124125
Guice.createInjector(new ImportJobModule(options, importJobSpecs), new PipelineModule());
125126
ImportJob job = injector.getInstance(ImportJob.class);
126127

128+
// Validate Influx DB configuration if options to write feature metrics to Influx DB is enabled
129+
if (options.isWriteFeatureMetricsToInfluxDb()) {
130+
if (options.getInfluxDbUrl() == null) {
131+
throw new IllegalArgumentException(
132+
"Influx DB url is required to write feature metrics to "
133+
+ "Influx DB. Please set this value like so '--influxDbUrl=http://localhost:8086'");
134+
}
135+
if (options.getInfluxDbDatabase() == null) {
136+
throw new IllegalArgumentException(
137+
"Influx DB database is required to write feature metrics to "
138+
+ "Influx DB. Please set this value like so '--influxDbDatabase=myinfluxdatabase'");
139+
}
140+
if (options.getInfluxDbMeasurement() == null) {
141+
throw new IllegalArgumentException(
142+
"Influx DB measurement name is required to write feature metrics to "
143+
+ "Influx DB. Please set this value like so '--influxDbMeasurement=mymeasurement'");
144+
}
145+
}
146+
127147
job.expand();
128148
return job.run();
129149
}
@@ -143,8 +163,8 @@ public void expand() {
143163
TypeDescriptor.of(FeatureRowExtended.class), ProtoCoder.of(FeatureRowExtended.class));
144164
coderRegistry.registerCoderForType(TypeDescriptor.of(TableRow.class), TableRowJsonCoder.of());
145165

146-
JobOptions jobOptions = OptionsParser
147-
.parse(importJobSpecs.getImportSpec().getJobOptionsMap(), JobOptions.class);
166+
JobOptions jobOptions =
167+
OptionsParser.parse(importJobSpecs.getImportSpec().getJobOptionsMap(), JobOptions.class);
148168

149169
try {
150170
log.info(JsonFormat.printer().print(importJobSpecs));
@@ -174,12 +194,21 @@ public void expand() {
174194
PCollection<FeatureRowExtended> errorRows = pFeatureRows.getErrors();
175195
if (jobOptions.isCoalesceRowsEnabled()) {
176196
// Should we merge and dedupe rows before writing to the serving store?
177-
servingRows = servingRows.apply("Coalesce Rows", new CoalesceFeatureRowExtended(
178-
jobOptions.getCoalesceRowsDelaySeconds(),
179-
jobOptions.getCoalesceRowsTimeoutSeconds()));
197+
servingRows =
198+
servingRows.apply(
199+
"Coalesce Rows",
200+
new CoalesceFeatureRowExtended(
201+
jobOptions.getCoalesceRowsDelaySeconds(),
202+
jobOptions.getCoalesceRowsTimeoutSeconds()));
180203
}
181204

182205
if (!dryRun) {
206+
servingRows.apply(
207+
new WriteFeatureMetricsToInfluxDB(
208+
options.getInfluxDbUrl(),
209+
options.getInfluxDbDatabase(),
210+
options.getInfluxDbMeasurement()));
211+
183212
servingRows.apply("Write to Serving Stores", servingStoreTransform);
184213
if (!Strings.isNullOrEmpty(importJobSpecs.getWarehouseStorageSpec().getId())) {
185214
warehouseRows.apply("Write to Warehouse Stores", warehouseStoreTransform);
@@ -194,7 +223,8 @@ public PipelineResult run() {
194223
return result;
195224
}
196225

197-
public void logNRows(PFeatureRows pFeatureRows, String name, long limit, Duration period) {
226+
@SuppressWarnings("SameParameterValue")
227+
private void logNRows(PFeatureRows pFeatureRows, String name, long limit, Duration period) {
198228
PCollection<FeatureRowExtended> main = pFeatureRows.getMain();
199229
PCollection<FeatureRowExtended> errors = pFeatureRows.getErrors();
200230

ingestion/src/main/java/feast/ingestion/options/ImportJobPipelineOptions.java

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,7 @@
2525
import org.apache.beam.sdk.options.PipelineOptionsRegistrar;
2626
import org.apache.beam.sdk.options.Validation.Required;
2727

28-
/**
29-
* Options passed to Beam to influence the job's execution environment
30-
*/
28+
/** Options passed to Beam to influence the job's execution environment */
3129
public interface ImportJobPipelineOptions extends PipelineOptions {
3230

3331
@Description("Path to a workspace directory containing importJobSpecs.yaml")
@@ -36,11 +34,33 @@ public interface ImportJobPipelineOptions extends PipelineOptions {
3634

3735
void setWorkspace(String value);
3836

37+
@Description(
38+
"If set, Feast will write feature metrics (such as lag and value summaries) "
39+
+ "into Influx DB. If this options is set, influxDbURL, influxDbDatabase "
40+
+ "and influxDbMeasurement should be set as well.")
41+
@Default.Boolean(false)
42+
boolean isWriteFeatureMetricsToInfluxDb();
43+
44+
void setWriteFeatureMetricsToInfluxDb(boolean shouldWrite);
45+
46+
@Description("e.g. http://localhost:8086")
47+
String getInfluxDbUrl();
48+
49+
void setInfluxDbUrl(String influxDbUrl);
50+
51+
String getInfluxDbDatabase();
52+
53+
void setInfluxDbDatabase(String influxDbDatabase);
54+
55+
String getInfluxDbMeasurement();
56+
57+
void setInfluxDbMeasurement(String influxDbMeasurement);
58+
3959
@Description("If dry run is set, execute up to feature row validation")
4060
@Default.Boolean(false)
41-
Boolean isDryRun();
61+
boolean isDryRun();
4262

43-
void setDryRun(Boolean value);
63+
void setDryRun(boolean value);
4464

4565
@AutoService(PipelineOptionsRegistrar.class)
4666
class ImportJobPipelineOptionsRegistrar implements PipelineOptionsRegistrar {
Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
package feast.ingestion.transform;
2+
3+
import feast.types.FeatureProto.Feature;
4+
import feast.types.FeatureRowExtendedProto.FeatureRowExtended;
5+
import feast.types.FeatureRowProto.FeatureRow;
6+
import feast.types.ValueProto.Value;
7+
import java.util.concurrent.TimeUnit;
8+
import org.apache.beam.sdk.transforms.DoFn;
9+
import org.apache.beam.sdk.transforms.PTransform;
10+
import org.apache.beam.sdk.transforms.ParDo;
11+
import org.apache.beam.sdk.values.PCollection;
12+
import org.apache.beam.sdk.values.PDone;
13+
import org.influxdb.BatchOptions;
14+
import org.influxdb.InfluxDB;
15+
import org.influxdb.InfluxDBFactory;
16+
import org.influxdb.dto.Point;
17+
18+
public class WriteFeatureMetricsToInfluxDB
19+
extends PTransform<PCollection<FeatureRowExtended>, PDone> {
20+
21+
private String influxDbUrl;
22+
private String influxDbDatabase;
23+
private String influxDbMeasurement;
24+
25+
public WriteFeatureMetricsToInfluxDB(
26+
String influxDbUrl, String influxDbDatabase, String influxDbMeasurement) {
27+
this.influxDbUrl = influxDbUrl;
28+
this.influxDbDatabase = influxDbDatabase;
29+
this.influxDbMeasurement = influxDbMeasurement;
30+
}
31+
32+
@Override
33+
public PDone expand(PCollection<FeatureRowExtended> input) {
34+
input.apply(
35+
ParDo.of(
36+
new DoFn<FeatureRowExtended, Void>() {
37+
InfluxDB influxDB;
38+
39+
@Setup
40+
public void setup() {
41+
influxDB = InfluxDBFactory.connect(influxDbUrl);
42+
influxDB.setDatabase(influxDbDatabase);
43+
influxDB.enableBatch(BatchOptions.DEFAULTS);
44+
}
45+
46+
@FinishBundle
47+
public void finishBundle() {
48+
if (influxDB != null) {
49+
influxDB.close();
50+
}
51+
}
52+
53+
@ProcessElement
54+
public void processElement(
55+
ProcessContext c, @Element FeatureRowExtended featureRowExtended) {
56+
FeatureRow featureRow = featureRowExtended.getRow();
57+
for (Feature feature : featureRow.getFeaturesList()) {
58+
String featureId = feature.getId();
59+
long lagInSeconds =
60+
System.currentTimeMillis() / 1000L
61+
- featureRow.getEventTimestamp().getSeconds();
62+
double value = getValue(feature);
63+
influxDB.write(
64+
Point.measurement(influxDbMeasurement)
65+
.time(System.currentTimeMillis(), TimeUnit.MILLISECONDS)
66+
.addField("lag_in_seconds", lagInSeconds)
67+
.addField("value", value)
68+
.tag("feature_id", featureId)
69+
.build());
70+
}
71+
}
72+
}));
73+
74+
return PDone.in(input.getPipeline());
75+
}
76+
77+
/**
78+
* This method returns the numeric value of the feature in double type
79+
*
80+
* <ul>
81+
* <li>if the value is a numeric type, it will be cast to double type
82+
* <li>f the value is a timestamp type, value corresponds to epoch seconds
83+
* <li>if the value is a boolean type, value of 1 corresponds to "true" and value of 0
84+
* corresponds to "false"
85+
* <li>if the value type is other non number format, value will be set to 0
86+
* </ul>
87+
*
88+
* @param feature Feast feature instance of a feature row
89+
* @return numeric value of the feature in double type
90+
*/
91+
private double getValue(Feature feature) {
92+
double value;
93+
Value featureValue = feature.getValue();
94+
95+
switch (featureValue.getValCase()) {
96+
case INT32VAL:
97+
value = featureValue.getInt32Val();
98+
break;
99+
case INT64VAL:
100+
value = featureValue.getInt64Val();
101+
break;
102+
case DOUBLEVAL:
103+
value = featureValue.getDoubleVal();
104+
break;
105+
case FLOATVAL:
106+
value = featureValue.getFloatVal();
107+
break;
108+
case BOOLVAL:
109+
value = featureValue.getBoolVal() ? 1.0 : 0.0;
110+
break;
111+
case TIMESTAMPVAL:
112+
value = featureValue.getTimestampVal().getSeconds();
113+
break;
114+
default:
115+
value = 0.0;
116+
break;
117+
}
118+
119+
return value;
120+
}
121+
}

ingestion/src/main/resources/logback.xml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,4 +24,7 @@
2424
<root level="INFO">
2525
<appender-ref ref="STDOUT"/>
2626
</root>
27+
<logger name="org.apache.kafka" level="WARN">
28+
<appender-ref ref="STDOUT"/>
29+
</logger>
2730
</configuration>

0 commit comments

Comments
 (0)