Java Code Examples for org.apache.beam.sdk.util.CoderUtils#encodeToByteArray()

@Test
public void testTestStreamCoder() throws Exception {
  TestStream<String> testStream =
      TestStream.create(StringUtf8Coder.of())
          .addElements("hey")
          .advanceWatermarkTo(Instant.ofEpochMilli(22521600))
          .advanceProcessingTime(Duration.millis(42))
          .addElements("hey", "joe")
          .advanceWatermarkToInfinity();

  TestStream.TestStreamCoder<String> coder = TestStream.TestStreamCoder.of(StringUtf8Coder.of());

  byte[] bytes = CoderUtils.encodeToByteArray(coder, testStream);
  TestStream<String> recoveredStream = CoderUtils.decodeFromByteArray(coder, bytes);

  assertThat(recoveredStream, is(testStream));
}
 

/**
 * This is a change detector test for the sizes of encoded windows. Since these are present for
 * every element of every windowed PCollection, the size matters.
 *
 * <p>This test documents the expectation that encoding as a (endpoint, duration) pair using big
 * endian for the endpoint and variable length long for the duration should be about 25% smaller
 * than encoding two big endian Long values.
 */
@Test
public void testLengthsOfEncodingChoices() throws Exception {
  Instant start = Instant.parse("2015-04-01T00:00:00Z");
  Instant minuteEnd = Instant.parse("2015-04-01T00:01:00Z");
  Instant hourEnd = Instant.parse("2015-04-01T01:00:00Z");
  Instant dayEnd = Instant.parse("2015-04-02T00:00:00Z");

  Coder<Instant> instantCoder = InstantCoder.of();
  byte[] encodedStart = CoderUtils.encodeToByteArray(instantCoder, start);
  byte[] encodedMinuteEnd = CoderUtils.encodeToByteArray(instantCoder, minuteEnd);
  byte[] encodedHourEnd = CoderUtils.encodeToByteArray(instantCoder, hourEnd);
  byte[] encodedDayEnd = CoderUtils.encodeToByteArray(instantCoder, dayEnd);

  byte[] encodedMinuteWindow =
      CoderUtils.encodeToByteArray(TEST_CODER, new IntervalWindow(start, minuteEnd));
  byte[] encodedHourWindow =
      CoderUtils.encodeToByteArray(TEST_CODER, new IntervalWindow(start, hourEnd));
  byte[] encodedDayWindow =
      CoderUtils.encodeToByteArray(TEST_CODER, new IntervalWindow(start, dayEnd));

  assertThat(
      encodedMinuteWindow.length, equalTo(encodedStart.length + encodedMinuteEnd.length - 5));
  assertThat(encodedHourWindow.length, equalTo(encodedStart.length + encodedHourEnd.length - 4));
  assertThat(encodedDayWindow.length, equalTo(encodedStart.length + encodedDayEnd.length - 4));
}
 

private void runTestRead(String[] expected) throws Exception {
  File tmpFile = tempFolder.newFile();
  String filename = tmpFile.getPath();

  try (PrintStream writer = new PrintStream(new FileOutputStream(tmpFile))) {
    for (String elem : expected) {
      byte[] encodedElem = CoderUtils.encodeToByteArray(StringUtf8Coder.of(), elem);
      String line = new String(encodedElem, Charsets.UTF_8);
      writer.println(line);
    }
  }

  TextIO.Read read = TextIO.read().from(filename);
  PCollection<String> output = p.apply(read);

  PAssert.that(output).containsInAnyOrder(expected);
  p.run();
}
 

private SingletonViewFn(
    boolean hasDefault,
    T defaultValue,
    Coder<T> valueCoder,
    TypeDescriptorSupplier<T> typeDescriptorSupplier) {
  this.hasDefault = hasDefault;
  this.defaultValue = defaultValue;
  this.valueCoder = valueCoder;
  this.typeDescriptorSupplier = typeDescriptorSupplier;
  if (hasDefault) {
    try {
      this.encodedDefaultValue = CoderUtils.encodeToByteArray(valueCoder, defaultValue);
    } catch (IOException e) {
      throw new RuntimeException("Unexpected IOException: ", e);
    }
  }
}
 

/**
 * Generates a {@link PubsubMessage} with the {@link GenericRecord} payload and the user provided
 * error message as an attribute.
 *
 * @param record payload for the message
 * @param errorValue errors to be added as an attribute
 */
private static PubsubMessage getPubsubMessage(GenericRecord record, String errorValue)
    throws IOException {
  AvroCoder<GenericRecord> coder = AvroCoder.of(record.getSchema());
  String errorKey = "error";
  Map<String, String> attributeMap = ImmutableMap.<String, String>builder()
      .put(errorKey, errorValue)
      .build();
  return new PubsubMessage(CoderUtils.encodeToByteArray(coder, record), attributeMap);
}
 

@Override
public byte[] getCurrentRecordId() {
  try {
    return CoderUtils.encodeToByteArray(coder, getCurrent());
  } catch (CoderException e) {
    throw new RuntimeException(e);
  }
}
 

@Test
public void testEncodeDecode() throws Exception {
  MyCustomCoder coder = new MyCustomCoder("key");
  CoderProperties.coderDecodeEncodeEqual(coder, KV.of("key", 3L));

  byte[] encoded2 = CoderUtils.encodeToByteArray(coder, KV.of("ignored", 3L));
  Assert.assertEquals(KV.of("key", 3L), CoderUtils.decodeFromByteArray(coder, encoded2));
}
 

/**
 * Utility for adding new entries to the common coder spec -- prints the serialized bytes of the
 * given value in the given context using JSON-escaped strings.
 */
private static <T> String jsonByteString(Coder<T> coder, T value, Context context)
    throws CoderException {
  byte[] bytes = CoderUtils.encodeToByteArray(coder, value, context);
  ObjectMapper mapper = new ObjectMapper();
  mapper.configure(JsonGenerator.Feature.ESCAPE_NON_ASCII, true);
  try {
    return mapper.writeValueAsString(new String(bytes, StandardCharsets.ISO_8859_1));
  } catch (JsonProcessingException e) {
    throw new CoderException(String.format("Unable to encode %s with coder %s", value, coder), e);
  }
}
 

@Override
public byte[] getKey(WindowedValue<KV<K, InputT>> value) throws Exception {
  final byte[] encodedKey = CoderUtils.encodeToByteArray(keyCoder, value.getValue().getKey());
  @SuppressWarnings("unchecked")
  final byte[] encodedWindow =
      CoderUtils.encodeToByteArray(
          (Coder) windowCoder, Iterables.getOnlyElement(value.getWindows()));
  return Bytes.concat(encodedKey, encodedWindow);
}
 

@Test
public void testEncodeDecodeOuter() throws Exception {
  JAXBCoder<TestType> coder = JAXBCoder.of(TestType.class);

  byte[] encoded = CoderUtils.encodeToByteArray(coder, new TestType("abc", 9999));
  assertEquals(new TestType("abc", 9999), CoderUtils.decodeFromByteArray(coder, encoded));
}
 

public static <T> CreateSource<T> fromIterable(Iterable<T> elements, Coder<T> elemCoder)
    throws CoderException, IOException {
  ImmutableList.Builder<byte[]> allElementsBytes = ImmutableList.builder();
  long totalSize = 0L;
  for (T element : elements) {
    byte[] bytes = CoderUtils.encodeToByteArray(elemCoder, element);
    allElementsBytes.add(bytes);
    totalSize += bytes.length;
  }
  return new CreateSource<>(allElementsBytes.build(), totalSize, elemCoder);
}
 

@Test
public void testEncodedElementByteSize() throws Throwable {
  for (ByteString value : TEST_VALUES) {
    byte[] encoded = CoderUtils.encodeToByteArray(TEST_CODER, value, Context.NESTED);
    assertEquals(encoded.length, TEST_CODER.getEncodedElementByteSize(value));
  }
}
 

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

/** Write input elements to a file and return information about the Avro-encoded file. */
private <T> AvroFileInfo<T> initInputFile(List<List<T>> elemsList, Coder<T> coder)
    throws Exception {
  File tmpFile = tmpFolder.newFile("file.avro");
  AvroFileInfo<T> fileInfo = new AvroFileInfo<>();
  fileInfo.filename = tmpFile.getPath();

  // Write the data.
  OutputStream outStream =
      Channels.newOutputStream(
          FileSystems.create(
              FileSystems.matchNewResource(fileInfo.filename, false), MimeTypes.BINARY));
  Schema schema = Schema.create(Schema.Type.BYTES);
  DatumWriter<ByteBuffer> datumWriter = new GenericDatumWriter<>(schema);
  try (DataFileWriter<ByteBuffer> fileWriter = new DataFileWriter<>(datumWriter)) {
    fileWriter.create(schema, outStream);
    boolean first = true;
    for (List<T> elems : elemsList) {
      if (first) {
        first = false;
      } else {
        // Ensure a block boundary here.
        long syncPoint = fileWriter.sync();
        fileInfo.syncPoints.add(syncPoint);
      }
      for (T elem : elems) {
        byte[] encodedElement = CoderUtils.encodeToByteArray(coder, elem);
        fileWriter.append(ByteBuffer.wrap(encodedElement));
        fileInfo.elementSizes.add(encodedElement.length);
        fileInfo.totalElementEncodedSize += encodedElement.length;
      }
    }
  }

  return fileInfo;
}
 

private static <T> SerializableFunction<T, PubsubMessage> formatPayloadUsingCoder(
    Coder<T> coder) {
  return input -> {
    try {
      return new PubsubMessage(CoderUtils.encodeToByteArray(coder, input), ImmutableMap.of());
    } catch (CoderException e) {
      throw new RuntimeException("Could not encode Pubsub message", e);
    }
  };
}
 

@Override
public byte[] getKey(WindowedValue<KV<K, InputT>> value) throws Exception {
  return CoderUtils.encodeToByteArray(keyCoder, value.getValue().getKey());
}
 

@ProcessElement
public void processElement(ProcessContext c) throws IOException {
  byte[] payload = CoderUtils.encodeToByteArray(Event.CODER, c.element());
  c.output(new PubsubMessage(payload, Collections.emptyMap()));
}
 

@ProcessElement
public void processElement(ProcessContext c) throws IOException {
  byte[] encodedEvent = CoderUtils.encodeToByteArray(Event.CODER, c.element());
  c.output(encodedEvent);
}
 

private byte[] intervalWindowBytes(IntervalWindow window) throws Exception {
  return CoderUtils.encodeToByteArray(DEFAULT_WINDOW_COLLECTION_CODER, Arrays.asList(window));
}
 

private static void assertOutputFiles(
    String[] elems,
    final String header,
    final String footer,
    int numShards,
    ResourceId outputPrefix,
    String shardNameTemplate)
    throws Exception {
  List<File> expectedFiles = new ArrayList<>();
  if (numShards == 0) {
    String pattern = outputPrefix.toString() + "*";
    List<MatchResult> matches = FileSystems.match(Collections.singletonList(pattern));
    for (Metadata expectedFile : Iterables.getOnlyElement(matches).metadata()) {
      expectedFiles.add(new File(expectedFile.resourceId().toString()));
    }
  } else {
    for (int i = 0; i < numShards; i++) {
      expectedFiles.add(
          new File(
              DefaultFilenamePolicy.constructName(
                      outputPrefix, shardNameTemplate, "", i, numShards, null, null)
                  .toString()));
    }
  }

  List<List<String>> actual = new ArrayList<>();

  for (File tmpFile : expectedFiles) {
    List<String> currentFile = readLinesFromFile(tmpFile);
    actual.add(currentFile);
  }

  List<String> expectedElements = new ArrayList<>(elems.length);
  for (String elem : elems) {
    byte[] encodedElem = CoderUtils.encodeToByteArray(StringUtf8Coder.of(), elem);
    String line = new String(encodedElem, Charsets.UTF_8);
    expectedElements.add(line);
  }

  List<String> actualElements =
      Lists.newArrayList(
          Iterables.concat(
              FluentIterable.from(actual)
                  .transform(removeHeaderAndFooter(header, footer))
                  .toList()));

  assertThat(actualElements, containsInAnyOrder(expectedElements.toArray()));
  assertTrue(actual.stream().allMatch(haveProperHeaderAndFooter(header, footer)::apply));
}