Java Code Examples for org.apache.beam.sdk.transforms.SerializableFunctions#identity()

The following examples show how to use org.apache.beam.sdk.transforms.SerializableFunctions#identity() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: BigQueryIOWriteTest.java    From beam with Apache License 2.0 6 votes vote down vote up 
@Test
public void testRemoveTemporaryFiles() throws Exception {
  int numFiles = 10;
  List<String> fileNames = Lists.newArrayList();
  String tempFilePrefix = options.getTempLocation() + "/";
  for (int i = 0; i < numFiles; ++i) {
    TableRowWriter<TableRow> writer =
        new TableRowWriter<>(tempFilePrefix, SerializableFunctions.identity());
    writer.close();
    fileNames.add(writer.getResult().resourceId.toString());
  }
  fileNames.add(tempFilePrefix + String.format("files%05d", numFiles));

  File tempDir = new File(options.getTempLocation());
  testNumFiles(tempDir, 10);

  WriteTables.removeTemporaryFiles(fileNames);
}
Example 2
Source File: RowCoder.java    From beam with Apache License 2.0 5 votes vote down vote up 
private RowCoder(Schema schema) {
  super(
      schema,
      TypeDescriptors.rows(),
      SerializableFunctions.identity(),
      SerializableFunctions.identity());
}
Example 3
Source File: DynamicFileDestinations.java    From beam with Apache License 2.0 4 votes vote down vote up 
/**
 * A specialization of {@link #constant(FilenamePolicy, SerializableFunction)} for the case where
 * UserT and OutputT are the same type and the format function is the identity.
 */
public static <UserT> DynamicDestinations<UserT, Void, UserT> constant(
    FilenamePolicy filenamePolicy) {
  return new ConstantFilenamePolicy<>(filenamePolicy, SerializableFunctions.<UserT>identity());
}
Example 4
Source File: BigQueryIOWriteTest.java    From beam with Apache License 2.0 4 votes vote down vote up 
@Test
public void testWriteTables() throws Exception {
  long numTables = 3;
  long numPartitions = 3;
  long numFilesPerPartition = 10;
  String jobIdToken = "jobId";
  final Multimap<TableDestination, String> expectedTempTables = ArrayListMultimap.create();

  List<KV<ShardedKey<String>, List<String>>> partitions = Lists.newArrayList();
  for (int i = 0; i < numTables; ++i) {
    String tableName = String.format("project-id:dataset-id.table%05d", i);
    TableDestination tableDestination = new TableDestination(tableName, tableName);
    for (int j = 0; j < numPartitions; ++j) {
      String tempTableId = BigQueryHelpers.createJobId(jobIdToken, tableDestination, j, 0);
      List<String> filesPerPartition = Lists.newArrayList();
      for (int k = 0; k < numFilesPerPartition; ++k) {
        String filename =
            Paths.get(
                    testFolder.getRoot().getAbsolutePath(),
                    String.format("files0x%08x_%05d", tempTableId.hashCode(), k))
                .toString();
        TableRowWriter<TableRow> writer =
            new TableRowWriter<>(filename, SerializableFunctions.identity());
        try (TableRowWriter ignored = writer) {
          TableRow tableRow = new TableRow().set("name", tableName);
          writer.write(tableRow);
        }
        filesPerPartition.add(writer.getResult().resourceId.toString());
      }
      partitions.add(KV.of(ShardedKey.of(tableDestination.getTableSpec(), j), filesPerPartition));

      String json =
          String.format(
              "{\"datasetId\":\"dataset-id\",\"projectId\":\"project-id\",\"tableId\":\"%s\"}",
              tempTableId);
      expectedTempTables.put(tableDestination, json);
    }
  }

  PCollection<KV<ShardedKey<String>, List<String>>> writeTablesInput =
      p.apply(Create.of(partitions));
  PCollectionView<String> jobIdTokenView =
      p.apply("CreateJobId", Create.of("jobId")).apply(View.asSingleton());
  List<PCollectionView<?>> sideInputs = ImmutableList.of(jobIdTokenView);

  fakeJobService.setNumFailuresExpected(3);
  WriteTables<String> writeTables =
      new WriteTables<>(
          true,
          fakeBqServices,
          jobIdTokenView,
          BigQueryIO.Write.WriteDisposition.WRITE_EMPTY,
          BigQueryIO.Write.CreateDisposition.CREATE_IF_NEEDED,
          sideInputs,
          new IdentityDynamicTables(),
          null,
          4,
          false,
          null,
          "NEWLINE_DELIMITED_JSON",
          false,
          Collections.emptySet());

  PCollection<KV<TableDestination, String>> writeTablesOutput =
      writeTablesInput.apply(writeTables);

  PAssert.thatMultimap(writeTablesOutput)
      .satisfies(
          input -> {
            assertEquals(input.keySet(), expectedTempTables.keySet());
            for (Map.Entry<TableDestination, Iterable<String>> entry : input.entrySet()) {
              @SuppressWarnings("unchecked")
              String[] expectedValues =
                  Iterables.toArray(expectedTempTables.get(entry.getKey()), String.class);
              assertThat(entry.getValue(), containsInAnyOrder(expectedValues));
            }
            return null;
          });
  p.run();
}