Examples with WeightedValue org.apache.beam.sdk.util.WeightedValue used on opensource projects

Example 1 with WeightedValue

use of org.apache.beam.sdk.util.WeightedValue in project beam by apache.

the class IsmReaderFactory method createImpl.

<V> NativeReader<?> createImpl(CloudObject spec, Coder<?> coder, PipelineOptions options, DataflowExecutionContext executionContext, DataflowOperationContext operationContext) throws Exception {
    final ResourceId resourceId = FileSystems.matchNewResource(getString(spec, WorkerPropertyNames.FILENAME), false);
    checkArgument(coder instanceof WindowedValueCoder, "%s only supports using %s but got %s.", IsmReader.class, WindowedValueCoder.class, coder);
    @SuppressWarnings("unchecked") WindowedValueCoder<IsmRecord<V>> windowedCoder = (WindowedValueCoder<IsmRecord<V>>) coder;
    checkArgument(windowedCoder.getValueCoder() instanceof IsmRecordCoder, "%s only supports using %s but got %s.", IsmReader.class, IsmRecordCoder.class, windowedCoder.getValueCoder());
    @SuppressWarnings("unchecked") final IsmRecordCoder<V> ismCoder = (IsmRecordCoder<V>) windowedCoder.getValueCoder();
    checkArgument(executionContext instanceof BatchModeExecutionContext, "%s only supports using %s but got %s.", IsmReader.class, BatchModeExecutionContext.class, executionContext);
    final BatchModeExecutionContext execContext = (BatchModeExecutionContext) executionContext;
    // the same file.
    return execContext.<IsmReaderKey, NativeReader<?>>getLogicalReferenceCache().get(new IsmReaderKey(resourceId.toString()), () -> new IsmReaderImpl<V>(resourceId, ismCoder, execContext.<IsmReaderImpl.IsmShardKey, WeightedValue<NavigableMap<RandomAccessData, WindowedValue<IsmRecord<V>>>>>getDataCache()));
}

Example 2 with WeightedValue

use of org.apache.beam.sdk.util.WeightedValue 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()));
}