Java BoundedReader.close方法代码示例

import org.apache.beam.sdk.io.BoundedSource.BoundedReader; //导入方法依赖的package包/类
/**
 * This test validates behavior of
 * {@link HadoopInputFormatBoundedSource.HadoopInputFormatReader#start() start()} method if
 * InputFormat's {@link InputFormat#getSplits() getSplits()} returns InputSplitList having zero
 * records.
 */
@Test
public void testReadersStartWhenZeroRecords() throws Exception {

  InputFormat mockInputFormat = Mockito.mock(EmployeeInputFormat.class);
  EmployeeRecordReader mockReader = Mockito.mock(EmployeeRecordReader.class);
  Mockito.when(
      mockInputFormat.createRecordReader(Mockito.any(InputSplit.class),
          Mockito.any(TaskAttemptContext.class))).thenReturn(mockReader);
  Mockito.when(mockReader.nextKeyValue()).thenReturn(false);
  InputSplit mockInputSplit = Mockito.mock(NewObjectsEmployeeInputSplit.class);
  HadoopInputFormatBoundedSource<Text, Employee> boundedSource =
      new HadoopInputFormatBoundedSource<Text, Employee>(
          serConf,
          WritableCoder.of(Text.class),
          AvroCoder.of(Employee.class),
          null, // No key translation required.
          null, // No value translation required.
          new SerializableSplit(mockInputSplit));
  boundedSource.setInputFormatObj(mockInputFormat);
  BoundedReader<KV<Text, Employee>> reader = boundedSource.createReader(p.getOptions());
  assertEquals(false, reader.start());
  assertEquals(Double.valueOf(1), reader.getFractionConsumed());
  reader.close();
}
 

import org.apache.beam.sdk.io.BoundedSource.BoundedReader; //导入方法依赖的package包/类
/**
 * This test validates the method getFractionConsumed()- when a bad progress value is returned by
 * the inputformat.
 */
@Test
public void testGetFractionConsumedForBadProgressValue() throws Exception {
  InputFormat<Text, Employee> mockInputFormat = Mockito.mock(EmployeeInputFormat.class);
  EmployeeRecordReader mockReader = Mockito.mock(EmployeeRecordReader.class);
  Mockito.when(
      mockInputFormat.createRecordReader(Mockito.any(InputSplit.class),
          Mockito.any(TaskAttemptContext.class))).thenReturn(mockReader);
  Mockito.when(mockReader.nextKeyValue()).thenReturn(true);
  // Set to a bad value , not in range of 0 to 1
  Mockito.when(mockReader.getProgress()).thenReturn(2.0F);
  InputSplit mockInputSplit = Mockito.mock(NewObjectsEmployeeInputSplit.class);
  HadoopInputFormatBoundedSource<Text, Employee> boundedSource =
      new HadoopInputFormatBoundedSource<Text, Employee>(
          serConf,
          WritableCoder.of(Text.class),
          AvroCoder.of(Employee.class),
          null, // No key translation required.
          null, // No value translation required.
          new SerializableSplit(mockInputSplit));
  boundedSource.setInputFormatObj(mockInputFormat);
  BoundedReader<KV<Text, Employee>> reader = boundedSource.createReader(p.getOptions());
  assertEquals(Double.valueOf(0), reader.getFractionConsumed());
  boolean start = reader.start();
  assertEquals(true, start);
  if (start) {
    boolean advance = reader.advance();
    assertEquals(null, reader.getFractionConsumed());
    assertEquals(true, advance);
    if (advance) {
      advance = reader.advance();
      assertEquals(null, reader.getFractionConsumed());
    }
  }
  // Validate if getFractionConsumed() returns null after few number of reads as getProgress
  // returns invalid value '2' which is not in the range of 0 to 1.
  assertEquals(null, reader.getFractionConsumed());
  reader.close();
}
 

import org.apache.beam.sdk.io.BoundedSource.BoundedReader; //导入方法依赖的package包/类
@Test
public void testGetSplitPointsConsumed() throws Exception {
  final String table = "TEST-TABLE";
  final int numRows = 100;
  int splitPointsConsumed = 0;

  makeTableData(table, numRows);

  BigtableSource source =
      new BigtableSource(serviceFactory, table, null, Arrays.asList(ByteKeyRange.ALL_KEYS), null);

  BoundedReader<Row> reader = source.createReader(TestPipeline.testingPipelineOptions());

  reader.start();
  // Started, 0 split points consumed
  assertEquals("splitPointsConsumed starting",
      splitPointsConsumed, reader.getSplitPointsConsumed());

  // Split points consumed increases for each row read
  while (reader.advance()) {
    assertEquals("splitPointsConsumed advancing",
        ++splitPointsConsumed, reader.getSplitPointsConsumed());
  }

  // Reader marked as done, 100 split points consumed
  assertEquals("splitPointsConsumed done", numRows, reader.getSplitPointsConsumed());

  reader.close();
}
 

import org.apache.beam.sdk.io.BoundedSource.BoundedReader; //导入方法依赖的package包/类
/**
 * This test validates the method getFractionConsumed()- which indicates the progress of the read
 * in range of 0 to 1.
 */
@Test
public void testReadersGetFractionConsumed() throws Exception {
  List<KV<Text, Employee>> referenceRecords = TestEmployeeDataSet.getEmployeeData();
  HadoopInputFormatBoundedSource<Text, Employee> hifSource = getTestHIFSource(
      EmployeeInputFormat.class,
      Text.class,
      Employee.class,
      WritableCoder.of(Text.class),
      AvroCoder.of(Employee.class));
  long estimatedSize = hifSource.getEstimatedSizeBytes(p.getOptions());
  // Validate if estimated size is equal to the size of records.
  assertEquals(referenceRecords.size(), estimatedSize);
  List<BoundedSource<KV<Text, Employee>>> boundedSourceList =
      hifSource.split(0, p.getOptions());
  // Validate if split() has split correctly.
  assertEquals(TestEmployeeDataSet.NUMBER_OF_SPLITS, boundedSourceList.size());
  List<KV<Text, Employee>> bundleRecords = new ArrayList<>();
  for (BoundedSource<KV<Text, Employee>> source : boundedSourceList) {
    List<KV<Text, Employee>> elements = new ArrayList<KV<Text, Employee>>();
    BoundedReader<KV<Text, Employee>> reader = source.createReader(p.getOptions());
    float recordsRead = 0;
    // When start is not called, getFractionConsumed() should return 0.
    assertEquals(Double.valueOf(0), reader.getFractionConsumed());
    boolean start = reader.start();
    assertEquals(true, start);
    if (start) {
      elements.add(reader.getCurrent());
      boolean advance = reader.advance();
      // Validate if getFractionConsumed() returns the correct fraction based on
      // the number of records read in the split.
      assertEquals(
          Double.valueOf(++recordsRead / TestEmployeeDataSet.NUMBER_OF_RECORDS_IN_EACH_SPLIT),
          reader.getFractionConsumed());
      assertEquals(true, advance);
      while (advance) {
        elements.add(reader.getCurrent());
        advance = reader.advance();
        assertEquals(
            Double.valueOf(++recordsRead / TestEmployeeDataSet.NUMBER_OF_RECORDS_IN_EACH_SPLIT),
            reader.getFractionConsumed());
      }
      bundleRecords.addAll(elements);
    }
    // Validate if getFractionConsumed() returns 1 after reading is complete.
    assertEquals(Double.valueOf(1), reader.getFractionConsumed());
    reader.close();
  }
  assertThat(bundleRecords, containsInAnyOrder(referenceRecords.toArray()));
}