@Override public TableRow mapRow(ResultSet resultSet) throws Exception { ResultSetMetaData metaData = resultSet.getMetaData(); TableRow outputTableRow = new TableRow(); for (int i = 1; i <= metaData.getColumnCount(); i++) { outputTableRow.set(metaData.getColumnName(i), resultSet.getObject(i)); } return outputTableRow; } }
public TableRowBuilder set(String fieldName, Object value) { row.set(fieldName, value); return this; }
@ProcessElement public void processElement(ProcessContext c) { List<TableRow> completions = new ArrayList<>(); for (CompletionCandidate cc : c.element().getValue()) { completions.add(new TableRow().set("count", cc.getCount()).set("tag", cc.getValue())); } TableRow row = new TableRow().set("prefix", c.element().getKey()).set("tags", completions); c.output(row); }
@ProcessElement public void processElement(ProcessContext c) { LaneInfo laneInfo = c.element().getValue(); TableRow row = new TableRow() .set("station_id", c.element().getKey()) .set("direction", laneInfo.getDirection()) .set("freeway", laneInfo.getFreeway()) .set("lane_max_flow", laneInfo.getLaneFlow()) .set("lane", laneInfo.getLane()) .set("avg_occ", laneInfo.getLaneAO()) .set("avg_speed", laneInfo.getLaneAS()) .set("total_flow", laneInfo.getTotalFlow()) .set("recorded_timestamp", laneInfo.getRecordedTimestamp()) .set("window_timestamp", c.timestamp().toString()); c.output(row); }
@ProcessElement public void processElement(ProcessContext c) { RouteInfo routeInfo = c.element().getValue(); TableRow row = new TableRow() .set("avg_speed", routeInfo.getAvgSpeed()) .set("slowdown_event", routeInfo.getSlowdownEvent()) .set("route", c.element().getKey()) .set("window_timestamp", c.timestamp().toString()); c.output(row); }
private static TableRow convertGenericRecordToTableRow( GenericRecord record, List<TableFieldSchema> fields) { TableRow row = new TableRow(); for (TableFieldSchema subSchema : fields) { // Per https://cloud.google.com/bigquery/docs/reference/v2/tables#schema, the name field // is required, so it may not be null. Field field = record.getSchema().getField(subSchema.getName()); Object convertedValue = getTypedCellValue(field.schema(), subSchema, record.get(field.name())); if (convertedValue != null) { // To match the JSON files exported by BigQuery, do not include null values in the output. row.set(field.name(), convertedValue); } } return row; }
/** In this example, put the whole string into single BigQuery field. */ @ProcessElement public void processElement(ProcessContext c) { c.output(new TableRow().set("string_field", c.element())); }
@Override public TableRow convertToDatum(IndexedRecord indexedRecord) { // When BigQueryOutput do not specify schema, so read it from the incoming data if (schema == null) { schema = indexedRecord.getSchema(); initFieldConverters(); } TableRow row = new TableRow(); for (Schema.Field field : schema.getFields()) { Object v = indexedRecord.get(field.pos()); if (v != null) { row.set(field.name(), fieldConverters.get(field.name()).convertToDatum(v)); } } return row; }
@Override public TableRow convertToDatum(IndexedRecord indexedRecord) { // When BigQueryOutput do not specify schema, so read it from the incoming data if (schema == null) { schema = indexedRecord.getSchema(); initFieldConverters(); } TableRow row = new TableRow(); for (Schema.Field field : schema.getFields()) { Object v = indexedRecord.get(field.pos()); if (v != null) { row.set(field.name(), fieldConverters.get(field.name()).convertToDatum(v)); } } return row; }
List<TableRow> convertStringsToLong(List<TableRow> toConvert) { // The numbers come back as String after JSON serialization. Change them back to // longs so that we can assert the output. List<TableRow> converted = Lists.newArrayList(); for (TableRow entry : toConvert) { TableRow convertedEntry = entry.clone(); convertedEntry.set("number", Long.parseLong((String) convertedEntry.get("number"))); converted.add(convertedEntry); } return converted; }
@ProcessElement public void processElement(ProcessContext c) throws Exception { TableRow element = c.element(); TableRow row = new TableRow() .set("trigger_type", element.get("trigger_type")) .set("freeway", element.get("freeway")) .set("total_flow", element.get("total_flow")) .set("number_of_records", element.get("number_of_records")) .set("isFirst", element.get("isFirst")) .set("isLast", element.get("isLast")) .set("timing", element.get("timing")) .set("window", element.get("window")); c.output(canonicalFormat(row)); } }
@Test @Category(ValidatesRunner.class) public void testFormatCounts() { PCollection<KV<Integer, Long>> inputs = p.apply(Create.of(KV.of(3, 0L), KV.of(4, Long.MAX_VALUE), KV.of(5, Long.MIN_VALUE))); PCollection<TableRow> result = inputs.apply(ParDo.of(new FormatCountsFn())); PAssert.that(result) .containsInAnyOrder( new TableRow().set("month", 3).set("tornado_count", 0), new TableRow().set("month", 4).set("tornado_count", Long.MAX_VALUE), new TableRow().set("month", 5).set("tornado_count", Long.MIN_VALUE)); p.run().waitUntilFinish(); } }
@Test public void testWriteToTableDecorator() throws Exception { TableRow row1 = new TableRow().set("name", "a").set("number", "1"); TableRow row2 = new TableRow().set("name", "b").set("number", "2"); TableSchema schema = new TableSchema() .setFields( ImmutableList.of(new TableFieldSchema().setName("number").setType("INTEGER"))); p.apply(Create.of(row1, row2)) .apply( BigQueryIO.writeTableRows() .to("project-id:dataset-id.table-id$20171127") .withTestServices(fakeBqServices) .withMethod(BigQueryIO.Write.Method.STREAMING_INSERTS) .withSchema(schema) .withoutValidation()); p.run(); }
@Test public void testCoder_nullCell() throws CoderException { TableRow row = new TableRow(); row.set("temperature", Data.nullOf(Object.class)); row.set("max_temperature", Data.nullOf(Object.class)); byte[] bytes = CoderUtils.encodeToByteArray(TableRowJsonCoder.of(), row); TableRow newRow = CoderUtils.decodeFromByteArray(TableRowJsonCoder.of(), bytes); byte[] newBytes = CoderUtils.encodeToByteArray(TableRowJsonCoder.of(), newRow); Assert.assertArrayEquals(bytes, newBytes); }
@Test @Category(ValidatesRunner.class) public void testNoTornadoes() { TableRow row = new TableRow().set("month", 6).set("tornado", false); PCollection<TableRow> inputs = p.apply(Create.of(ImmutableList.of(row))); PCollection<Integer> result = inputs.apply(ParDo.of(new ExtractTornadoesFn())); PAssert.that(result).empty(); p.run().waitUntilFinish(); }
@Test @Category(ValidatesRunner.class) public void testExtractTornadoes() { TableRow row = new TableRow().set("month", "6").set("tornado", true); PCollection<TableRow> input = p.apply(Create.of(ImmutableList.of(row))); PCollection<Integer> result = input.apply(ParDo.of(new ExtractTornadoesFn())); PAssert.that(result).containsInAnyOrder(6); p.run().waitUntilFinish(); }
@Test public void testWriteWithBrokenGetTable() throws Exception { p.apply(Create.<TableRow>of(new TableRow().set("foo", "bar"))) .apply( BigQueryIO.writeTableRows() .to(input -> null) .withCreateDisposition(BigQueryIO.Write.CreateDisposition.CREATE_NEVER) .withTestServices(fakeBqServices) .withoutValidation()); thrown.expectMessage("result of tableFunction can not be null"); thrown.expectMessage("foo"); p.run(); }