Java Code Examples for com.google.api.services.bigquery.model.TableRow#set()

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;
}
 

@Override
public TableRow apply(RowResult input) {
  TableRow outputTableRow = new TableRow();
  Schema kuduSchema = input.getSchema();

  for (int i = 0; i < kuduSchema.getColumnCount(); i++) {
    String columnName = kuduSchema.getColumnByIndex(i).getName();

    if (kuduSchema.getColumnByIndex(i).getType() == Type.UNIXTIME_MICROS) {
      outputTableRow.set(columnName, convertNanoSecondsToMilliSeconds(input, i));
    } else {
      // For other datatypes return the object value
      outputTableRow.set(columnName, input.getObject(i));
    }
  }
  return outputTableRow;
}
 

@ProcessElement
public void processElement(ProcessContext context) {
  FailsafeElement<KV<String, String>, String> failsafeElement = context.element();
  final KV<String, String> message = failsafeElement.getOriginalPayload();

  // Format the timestamp for insertion
  String timestamp =
      TIMESTAMP_FORMATTER.print(context.timestamp().toDateTime(DateTimeZone.UTC));

  // Build the table row
  final TableRow failedRow =
      new TableRow()
          .set("timestamp", timestamp)
          .set("errorMessage", failsafeElement.getErrorMessage())
          .set("stacktrace", failsafeElement.getStacktrace());

  // Only set the payload if it's populated on the message.
  failedRow.set(
      "payloadString",
      "key: "
          + (message.getKey() == null ? "" : message.getKey())
          + "value: "
          + (message.getValue() == null ? "" : message.getValue()));
  context.output(failedRow);
}
 

@Before
public void setUp() throws Exception {

    Mockito.when(resultSet.getObject(1)).thenReturn(NAME_VALUE);
    Mockito.when(resultSet.getObject(2)).thenReturn(AGE_VALUE);
    Mockito.when(resultSet.getMetaData()).thenReturn(resultSetMetaData);

    Mockito.when(resultSetMetaData.getColumnCount()).thenReturn(2);

    Mockito.when(resultSetMetaData.getColumnName(1)).thenReturn(NAME_KEY);

    Mockito.when(resultSetMetaData.getColumnName(2)).thenReturn(AGE_KEY);

    expectedTableRow = new TableRow();
    expectedTableRow.set(NAME_KEY, NAME_VALUE);
    expectedTableRow.set(AGE_KEY, AGE_VALUE);
}
 

/**
 * Tests {@link BigQueryConverters.JsonToTableRow} converts a valid Json TableRow to a TableRow.
 */
@Test
@Category(NeedsRunner.class)
public void testJsonToTableRowGood() throws Exception {
  TableRow expectedRow = new TableRow();
  expectedRow.set("location", "nyc");
  expectedRow.set("name", "Adam");
  expectedRow.set("age", 30);
  expectedRow.set("color", "blue");
  expectedRow.set("coffee", "black");

  ByteArrayOutputStream jsonStream = new ByteArrayOutputStream();
  TableRowJsonCoder.of().encode(expectedRow, jsonStream, Context.OUTER);
  String expectedJson = new String(jsonStream.toByteArray(), StandardCharsets.UTF_8.name());

  PCollection<TableRow> transformedJson =
      pipeline
          .apply("Create", Create.of(expectedJson))
          .apply(BigQueryConverters.jsonToTableRow());

  PAssert.that(transformedJson).containsInAnyOrder(expectedRow);

  pipeline.run();
}
 

@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 void processElement(ProcessContext c) {
  Double dollars = c.element();
  TableRow r = new TableRow();
  r.set("dollar_turnover", dollars);
  // the timing can be:
  // EARLY: the dollar amount is not yet final
  // ON_TIME: dataflow thinks the dollar amount is final but late data are still possible
  // LATE: late data has arrived
  r.set("dollar_timing", c.pane().getTiming()); // EARLY, ON_TIME or LATE
  r.set("dollar_window", ((IntervalWindow) c.window()).start().getMillis() / 1000.0 / 60.0); // timestamp in fractional minutes

  LOG.info("Outputting $ value {}} at {} with marker {} for window {}",
    dollars.toString(), new Date().getTime(), c.pane().getTiming().toString(), c.window().hashCode());
  c.output(r);
}
 

@ProcessElement
public void processElement(ProcessContext c) throws Exception {
    if (c.element().equalsIgnoreCase(HEADER)) return;
    String[] split = c.element().split(",");
    if (split.length > 7) return;
    TableRow row = new TableRow();
    for (int i = 0; i < split.length; i++) {
        TableFieldSchema col = getTableSchema().getFields().get(i);
        row.set(col.getName(), split[i]);
    }
    c.output(row);
}
 

@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);
}
 

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, BoundedWindow window) {

  TableRow row = new TableRow();
  for (Map.Entry<String, FieldInfo<T>> entry : fieldInfo.entrySet()) {
    String key = entry.getKey();
    FieldInfo<T> fcnInfo = entry.getValue();
    row.set(key, fcnInfo.getFieldFn().apply(c, window));
  }
  c.output(row);
}
 

@ProcessElement
public void processElement(ProcessContext c) {
  String day = (String) c.element().get("day");
  String month = (String) c.element().get("month");
  String year = (String) c.element().get("year");

  TableRow row = new TableRow();
  row.set("station_number", c.element().get("station_number"));
  row.set("date", String.format("%s-%s-%s", year, month, day));
  c.output(row);
}
 

public TableRow toTableRow() {
  TableRow result = new TableRow();
  result.set("latitude", lat);
  result.set("longitude", lon);
  result.set("ride_id", rideId);
  result.set("timestamp", Instant.ofEpochMilli(timestamp).toString());
  result.set("ride_status", status);
  return result;
}
 

@Override
public TableRow apply(KV<LatLon, Long> ridegrp) {
  TableRow result = new TableRow();
  result.set("latitude", ridegrp.getKey().lat);
  result.set("longitude", ridegrp.getKey().lon);
  result.set("ntaxis", ridegrp.getValue());

  return result;
}
 

@Override
public TableRow apply(Row row) {

  TableRow bqRow = new TableRow();
  String[] headers = row.getHeader();
  String[] values = row.getValue();
  for (int i = 0; i < values.length; i++) {
    bqRow.set(headers[i], values[i]);
  }
  return bqRow;
}
 

@ProcessElement
public void processElement(ProcessContext c, BoundedWindow window) {

  TableRow row = new TableRow();
  for (Map.Entry<String, FieldInfo<T>> entry : fieldInfo.entrySet()) {
    String key = entry.getKey();
    FieldInfo<T> fcnInfo = entry.getValue();
    row.set(key, fcnInfo.getFieldFn().apply(c, window));
  }
  c.output(row);
}
 

@ProcessElement
public void processElement(ProcessContext c, BoundedWindow window) {

  TableRow row = new TableRow();
  for (Map.Entry<String, FieldInfo<InputT>> entry : fieldInfo.entrySet()) {
    String key = entry.getKey();
    FieldInfo<InputT> fcnInfo = entry.getValue();
    FieldFn<InputT> fcn = fcnInfo.getFieldFn();
    row.set(key, fcn.apply(c, window));
  }
  c.output(row);
}
 

@Test
public void kuduRowResultShouldConvertedToBigQueryTableRowAsExpected() {
  // BigQuery TableRows expected result
  TableRow expectedTableRow = new TableRow();
  expectedTableRow.set("INT32", 1);
  expectedTableRow.set("BOOL", true);
  expectedTableRow.set("DOUBLE", 1.2);
  expectedTableRow.set("FLOAT", 1.3);
  expectedTableRow.set("INT8", 8);
  expectedTableRow.set("INT16", 16);
  expectedTableRow.set("INT64", 64);
  expectedTableRow.set("STRING", "string");
  expectedTableRow.set("UNIXTIME_MICROS", 1586422251000L);
  // Using Mockito to mock Kudu RowResults. The objects needs to be in column order
  RowResult kuduResult = Mockito.mock(RowResult.class);
  Mockito.doReturn(createTestKuduSchema()).when(kuduResult).getSchema();
  Mockito.doReturn(1).when(kuduResult).getObject(0);
  Mockito.doReturn(true).when(kuduResult).getObject(1);
  Mockito.doReturn(1.2).when(kuduResult).getObject(2);
  Mockito.doReturn(1.3).when(kuduResult).getObject(3);
  Mockito.doReturn(8).when(kuduResult).getObject(4);
  Mockito.doReturn(16).when(kuduResult).getObject(5);
  Mockito.doReturn(64).when(kuduResult).getObject(6);
  Mockito.doReturn("string").when(kuduResult).getObject(7);
  Mockito.doReturn(1586422251000000000L).when(kuduResult).getLong(8);

  // Convert using KuduConterters Class
  KuduToBigQuery.KuduConverters kuduConverters = new KuduToBigQuery.KuduConverters();
  TableRow convertedTableRow = kuduConverters.apply(kuduResult);
  List<TableRow> rows = ImmutableList.of(convertedTableRow);
  PCollection<TableRow> resultRows = pipeline.apply("CreateInput", Create.of(rows));

  // Check if TableRow expected equal to convertion result
  assertThat(convertedTableRow, equalTo(expectedTableRow));
  PAssert.that(resultRows)
      .satisfies(
          collection -> {
            TableRow result = collection.iterator().next();
            assertThat(result.get("INT32"), equalTo(1));
            assertThat(result.get("BOOL"), equalTo(true));
            assertThat(result.get("DOUBLE"), equalTo(1.2));
            assertThat(result.get("FLOAT"), equalTo(1.3));
            assertThat(result.get("INT8"), equalTo(8));
            assertThat(result.get("INT16"), equalTo(16));
            assertThat(result.get("INT64"), equalTo(64));
            assertThat(result.get("STRING"), equalTo("string"));
            assertThat(result.get("UNIXTIME_MICROS"), equalTo(1586422251000L));
            return null;
          });

  // Execute pipeline
  pipeline.run();
}
 

public TableRow apply(FeatureRow featureRow) {

    TableRow tableRow = new TableRow();
    tableRow.set(EVENT_TIMESTAMP_COLUMN, Timestamps.toString(featureRow.getEventTimestamp()));
    tableRow.set(CREATED_TIMESTAMP_COLUMN, Instant.now().toString());
    tableRow.set(INGESTION_ID_COLUMN, featureRow.getIngestionId());
    tableRow.set(JOB_ID_COLUMN, jobId);

    for (Field field : featureRow.getFieldsList()) {
      switch (field.getValue().getValCase()) {
        case BYTES_VAL:
          tableRow.set(
              field.getName(),
              Base64.getEncoder().encodeToString(field.getValue().getBytesVal().toByteArray()));
          break;
        case STRING_VAL:
          tableRow.set(field.getName(), field.getValue().getStringVal());
          break;
        case INT32_VAL:
          tableRow.set(field.getName(), field.getValue().getInt32Val());
          break;
        case INT64_VAL:
          tableRow.set(field.getName(), field.getValue().getInt64Val());
          break;
        case DOUBLE_VAL:
          tableRow.set(field.getName(), field.getValue().getDoubleVal());
          break;
        case FLOAT_VAL:
          tableRow.set(field.getName(), field.getValue().getFloatVal());
          break;
        case BOOL_VAL:
          tableRow.set(field.getName(), field.getValue().getBoolVal());
          break;
        case BYTES_LIST_VAL:
          tableRow.set(
              field.getName(),
              field.getValue().getBytesListVal().getValList().stream()
                  .map(x -> Base64.getEncoder().encodeToString(x.toByteArray()))
                  .collect(Collectors.toList()));
          break;
        case STRING_LIST_VAL:
          tableRow.set(field.getName(), field.getValue().getStringListVal().getValList());
          break;
        case INT32_LIST_VAL:
          tableRow.set(field.getName(), field.getValue().getInt32ListVal().getValList());
          break;
        case INT64_LIST_VAL:
          tableRow.set(field.getName(), field.getValue().getInt64ListVal().getValList());
          break;
        case DOUBLE_LIST_VAL:
          tableRow.set(field.getName(), field.getValue().getDoubleListVal().getValList());
          break;
        case FLOAT_LIST_VAL:
          tableRow.set(field.getName(), field.getValue().getFloatListVal().getValList());
          break;
        case BOOL_LIST_VAL:
          tableRow.set(field.getName(), field.getValue().getBytesListVal().getValList());
          break;
        case VAL_NOT_SET:
          break;
      }
    }

    return tableRow;
  }
 

@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;
}