@Test
  public void testCombineValuesFnCoders() throws Exception {
    CoderRegistry registry = new CoderRegistry();
    registry.registerStandardCoders();

    MeanInts meanInts = new MeanInts();
    MeanInts.CountSum countSum = meanInts.new CountSum(6, 27);

    Coder<MeanInts.CountSum> coder =
        meanInts.getAccumulatorCoder(
            registry, registry.getDefaultCoder(TypeDescriptor.of(Integer.class)));

    assertEquals(
        countSum,
        CoderUtils.decodeFromByteArray(coder, CoderUtils.encodeToByteArray(coder, countSum)));
  }
  <T> void runTestRead(T[] expected, Coder<T> coder) throws Exception {
    File tmpFile = tmpFolder.newFile("file.txt");
    String filename = tmpFile.getPath();

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

    Pipeline p = TestPipeline.create();

    TextIO.Read.Bound<T> read;
    if (coder.equals(StringUtf8Coder.of())) {
      TextIO.Read.Bound<String> readStrings = TextIO.Read.from(filename);
      // T==String
      read = (TextIO.Read.Bound<T>) readStrings;
    } else {
      read = TextIO.Read.from(filename).withCoder(coder);
    }

    PCollection<T> output = p.apply(read);

    DataflowAssert.that(output).containsInAnyOrder(expected);
    p.run();
  }
示例#3
0
 @Override
 protected T nextImpl() throws IOException {
   String encodedElementString = encodedElements.get(nextIndex++);
   // TODO: Replace with the real encoding used by the
   // front end, when we know what it is.
   byte[] encodedElement = StringUtils.jsonStringToByteArray(encodedElementString);
   notifyElementRead(encodedElement.length);
   return CoderUtils.decodeFromByteArray(coder, encodedElement);
 }
  <T> void runTestWrite(T[] elems, Coder<T> coder) throws Exception {
    File tmpFile = tmpFolder.newFile("file.txt");
    String filename = tmpFile.getPath();

    Pipeline p = TestPipeline.create();

    PCollection<T> input = p.apply(Create.of(Arrays.asList(elems)).withCoder(coder));

    TextIO.Write.Bound<T> write;
    if (coder.equals(StringUtf8Coder.of())) {
      TextIO.Write.Bound<String> writeStrings = TextIO.Write.to(filename).withoutSharding();
      // T==String
      write = (TextIO.Write.Bound<T>) writeStrings;
    } else {
      write = TextIO.Write.to(filename).withCoder(coder).withoutSharding();
    }

    input.apply(write);

    p.run();

    List<String> actual = new ArrayList<>();
    try (BufferedReader reader = new BufferedReader(new FileReader(tmpFile))) {
      for (; ; ) {
        String line = reader.readLine();
        if (line == null) {
          break;
        }
        actual.add(line);
      }
    }

    String[] expected = new String[elems.length];
    for (int i = 0; i < elems.length; i++) {
      T elem = elems[i];
      byte[] encodedElem = CoderUtils.encodeToByteArray(coder, elem);
      String line = new String(encodedElem);
      expected[i] = line;
    }

    assertThat(actual, containsInAnyOrder(expected));
  }