Examples with IsmShardKey org.apache.beam.runners.dataflow.worker.IsmReaderImpl.IsmShardKey used on opensource projects

Example 1 with IsmShardKey

use of org.apache.beam.runners.dataflow.worker.IsmReaderImpl.IsmShardKey in project beam by apache.

the class IsmReaderTest method testInitializationForSmallFilesIsCached.

@Test
public void testInitializationForSmallFilesIsCached() throws Exception {
    File tmpFile = tmpFolder.newFile();
    IsmShardKey expectedShardKey = new IsmShardKey(tmpFile.getAbsolutePath(), new RandomAccessData(0), 0, 13);
    List<IsmRecord<byte[]>> data = new ArrayList<>();
    data.add(IsmRecord.<byte[]>of(ImmutableList.of(EMPTY, new byte[] { 0x04 }), new byte[] { 0x04 }));
    data.add(IsmRecord.<byte[]>of(ImmutableList.of(EMPTY, new byte[] { 0x08 }), new byte[] { 0x08 }));
    writeElementsToFile(data, tmpFile);
    IsmReader<byte[]> reader = new IsmReaderImpl<byte[]>(FileSystems.matchSingleFileSpec(tmpFile.getAbsolutePath()).resourceId(), CODER, cache);
    // Validate that reader and cache are in initial state
    assertFalse(reader.isInitialized());
    assertEquals(0, cache.size());
    // Force initialization
    reader.overKeyComponents(ImmutableList.of());
    // Validate reader is initialized and expected entry is cached
    assertTrue(reader.isInitialized());
    WeightedValue<NavigableMap<RandomAccessData, WindowedValue<IsmRecord<byte[]>>>> block = cache.getIfPresent(expectedShardKey);
    assertNotNull(block);
    assertArrayEquals(new byte[] { 0x04 }, block.getValue().firstEntry().getValue().getValue().getValue());
    assertArrayEquals(new byte[] { 0x08 }, block.getValue().lastEntry().getValue().getValue().getValue());
}

Example 2 with IsmShardKey

use of org.apache.beam.runners.dataflow.worker.IsmReaderImpl.IsmShardKey in project beam by apache.

the class IsmReaderTest method testReadKeyThatEncodesToEmptyByteArray.

@Test
public void testReadKeyThatEncodesToEmptyByteArray() throws Exception {
    File tmpFile = tmpFolder.newFile();
    IsmRecordCoder<Void> coder = IsmRecordCoder.of(1, 0, ImmutableList.<Coder<?>>of(VoidCoder.of()), VoidCoder.of());
    IsmSink<Void> sink = new IsmSink<>(FileSystems.matchNewResource(tmpFile.getPath(), false), coder, BLOOM_FILTER_SIZE_LIMIT);
    IsmRecord<Void> element = IsmRecord.of(Arrays.asList((Void) null), (Void) null);
    try (SinkWriter<WindowedValue<IsmRecord<Void>>> writer = sink.writer()) {
        writer.add(new ValueInEmptyWindows<>(element));
    }
    Cache<IsmShardKey, WeightedValue<NavigableMap<RandomAccessData, WindowedValue<IsmRecord<Void>>>>> cache = CacheBuilder.newBuilder().weigher(Weighers.fixedWeightKeys(1)).maximumWeight(10_000).build();
    IsmReader<Void> reader = new IsmReaderImpl<>(FileSystems.matchSingleFileSpec(tmpFile.getAbsolutePath()).resourceId(), coder, cache);
    IsmReader<Void>.IsmPrefixReaderIterator iterator = reader.iterator();
    assertTrue(iterator.start());
    assertEquals(coder.structuralValue(element), coder.structuralValue(iterator.getCurrent().getValue()));
}