pabloem commented on code in PR #25155:
URL: https://github.com/apache/beam/pull/25155#discussion_r1090938376
##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/BigQueryStorageWriteApiSchemaTransformProvider.java:
##########
@@ -222,6 +226,32 @@ public void setBigQueryServices(BigQueryServices
testBigQueryServices) {
this.testBigQueryServices = testBigQueryServices;
}
+ // A counter of input BQ PCollection element count.
+ private static class ElementCounterFn extends DoFn<Row, Row> {
+ private static final Counter bqInputElementCounter =
+ Metrics.counter(
+ BigQueryStorageWriteApiPCollectionRowTupleTransform.class,
"element-counter");
+
+ @ProcessElement
+ public void process(ProcessContext c) {
+ bqInputElementCounter.inc();
+ c.output(c.element());
+ }
+ }
+
+ // A counter of output BQ insert failures.
+ private static class ErrorCounterFn extends DoFn<Row, Row> {
+ private static final Counter bqErrorCounter =
+ Metrics.counter(
+ BigQueryStorageWriteApiPCollectionRowTupleTransform.class,
"error-counter");
+
+ @ProcessElement
+ public void process(ProcessContext c) {
+ bqErrorCounter.inc();
+ c.output(c.element());
+ }
+ }
+
Review Comment:
I wonder if you can simplify these as a single DoFn that receives the
counter name as parameter?
Also, `Metrics.counter` is a little bit slow due to particulars of its
implementation. How about you do something like:
```
def process(self, element):
self._elements_in_bundle += 1
def finish_bundle(self):
self.bqErrorCounter.inc(self._elements_in_bundle)
self._elements_in_bundle = 0
```
WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
