Examples with TypeDescriptor org.apache.beam.sdk.values.TypeDescriptor used on opensource projects

Example 11 with TypeDescriptor

use of org.apache.beam.sdk.values.TypeDescriptor in project beam by apache.

the class ParDoTest method testTimerReceivedInOriginalWindow.

@Test
@Category({ ValidatesRunner.class, UsesTimersInParDo.class })
public void testTimerReceivedInOriginalWindow() throws Exception {
    final String timerId = "foo";
    DoFn<KV<String, Integer>, BoundedWindow> fn = new DoFn<KV<String, Integer>, BoundedWindow>() {

        @TimerId(timerId)
        private final TimerSpec spec = TimerSpecs.timer(TimeDomain.EVENT_TIME);

        @ProcessElement
        public void processElement(ProcessContext context, @TimerId(timerId) Timer timer) {
            timer.offset(Duration.standardSeconds(1)).setRelative();
        }

        @OnTimer(timerId)
        public void onTimer(OnTimerContext context, BoundedWindow window) {
            context.output(context.window());
        }

        public TypeDescriptor<BoundedWindow> getOutputTypeDescriptor() {
            return (TypeDescriptor) TypeDescriptor.of(IntervalWindow.class);
        }
    };
    SlidingWindows windowing = SlidingWindows.of(Duration.standardMinutes(3)).every(Duration.standardMinutes(1));
    PCollection<BoundedWindow> output = pipeline.apply(Create.timestamped(TimestampedValue.of(KV.of("hello", 24), new Instant(0L)))).apply(Window.<KV<String, Integer>>into(windowing)).apply(ParDo.of(fn));
    PAssert.that(output).containsInAnyOrder(new IntervalWindow(new Instant(0), Duration.standardMinutes(3)), new IntervalWindow(new Instant(0).minus(Duration.standardMinutes(1)), Duration.standardMinutes(3)), new IntervalWindow(new Instant(0).minus(Duration.standardMinutes(2)), Duration.standardMinutes(3)));
    pipeline.run();
}

Example 12 with TypeDescriptor

use of org.apache.beam.sdk.values.TypeDescriptor in project beam by apache.

the class SchemaCoderCloudObjectTranslator method fromCloudObject.

/**
 * Convert from a cloud object.
 */
@Override
public SchemaCoder fromCloudObject(CloudObject cloudObject) {
    try {
        TypeDescriptor typeDescriptor = (TypeDescriptor) SerializableUtils.deserializeFromByteArray(StringUtils.jsonStringToByteArray(Structs.getString(cloudObject, TYPE_DESCRIPTOR)), "typeDescriptor");
        SerializableFunction toRowFunction = (SerializableFunction) SerializableUtils.deserializeFromByteArray(StringUtils.jsonStringToByteArray(Structs.getString(cloudObject, TO_ROW_FUNCTION)), "toRowFunction");
        SerializableFunction fromRowFunction = (SerializableFunction) SerializableUtils.deserializeFromByteArray(StringUtils.jsonStringToByteArray(Structs.getString(cloudObject, FROM_ROW_FUNCTION)), "fromRowFunction");
        SchemaApi.Schema.Builder schemaBuilder = SchemaApi.Schema.newBuilder();
        JsonFormat.parser().merge(Structs.getString(cloudObject, SCHEMA), schemaBuilder);
        Schema schema = SchemaTranslation.schemaFromProto(schemaBuilder.build());
        @Nullable UUID uuid = schema.getUUID();
        if (schema.isEncodingPositionsOverridden() && uuid != null) {
            SchemaCoder.overrideEncodingPositions(uuid, schema.getEncodingPositions());
        }
        return SchemaCoder.of(schema, typeDescriptor, toRowFunction, fromRowFunction);
    } catch (IOException e) {
        throw new RuntimeException(e);
    }
}

Example 13 with TypeDescriptor

use of org.apache.beam.sdk.values.TypeDescriptor in project beam by apache.

the class DirectRunnerTest method byteArrayCountShouldSucceed.

@Test
public void byteArrayCountShouldSucceed() {
    Pipeline p = getPipeline();
    SerializableFunction<Integer, byte[]> getBytes = input -> {
        try {
            return CoderUtils.encodeToByteArray(VarIntCoder.of(), input);
        } catch (CoderException e) {
            fail("Unexpected Coder Exception " + e);
            throw new AssertionError("Unreachable");
        }
    };
    TypeDescriptor<byte[]> td = new TypeDescriptor<byte[]>() {
    };
    PCollection<byte[]> foos = p.apply(Create.of(1, 1, 1, 2, 2, 3)).apply(MapElements.into(td).via(getBytes));
    PCollection<byte[]> msync = p.apply(Create.of(1, -2, -8, -16)).apply(MapElements.into(td).via(getBytes));
    PCollection<byte[]> bytes = PCollectionList.of(foos).and(msync).apply(Flatten.pCollections());
    PCollection<KV<byte[], Long>> counts = bytes.apply(Count.perElement());
    PCollection<KV<Integer, Long>> countsBackToString = counts.apply(MapElements.via(new SimpleFunction<KV<byte[], Long>, KV<Integer, Long>>() {

        @Override
        public KV<Integer, Long> apply(KV<byte[], Long> input) {
            try {
                return KV.of(CoderUtils.decodeFromByteArray(VarIntCoder.of(), input.getKey()), input.getValue());
            } catch (CoderException e) {
                fail("Unexpected Coder Exception " + e);
                throw new AssertionError("Unreachable");
            }
        }
    }));
    Map<Integer, Long> expected = ImmutableMap.<Integer, Long>builder().put(1, 4L).put(2, 2L).put(3, 1L).put(-2, 1L).put(-8, 1L).put(-16, 1L).build();
    PAssert.thatMap(countsBackToString).isEqualTo(expected);
}

Example 14 with TypeDescriptor

use of org.apache.beam.sdk.values.TypeDescriptor in project beam by apache.

the class FlinkStreamingPortablePipelineTranslator method getSideInputIdToPCollectionViewMap.

private static LinkedHashMap<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>> getSideInputIdToPCollectionViewMap(RunnerApi.ExecutableStagePayload stagePayload, RunnerApi.Components components) {
    RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components);
    LinkedHashMap<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>> sideInputs = new LinkedHashMap<>();
    // for PCollectionView compatibility, not used to transform materialization
    ViewFn<Iterable<WindowedValue<?>>, ?> viewFn = (ViewFn) new PCollectionViews.MultimapViewFn<>((PCollectionViews.TypeDescriptorSupplier<Iterable<WindowedValue<Void>>>) () -> TypeDescriptors.iterables(new TypeDescriptor<WindowedValue<Void>>() {
    }), (PCollectionViews.TypeDescriptorSupplier<Void>) TypeDescriptors::voids);
    for (RunnerApi.ExecutableStagePayload.SideInputId sideInputId : stagePayload.getSideInputsList()) {
        // TODO: local name is unique as long as only one transform with side input can be within a
        // stage
        String sideInputTag = sideInputId.getLocalName();
        String collectionId = components.getTransformsOrThrow(sideInputId.getTransformId()).getInputsOrThrow(sideInputId.getLocalName());
        RunnerApi.WindowingStrategy windowingStrategyProto = components.getWindowingStrategiesOrThrow(components.getPcollectionsOrThrow(collectionId).getWindowingStrategyId());
        final WindowingStrategy<?, ?> windowingStrategy;
        try {
            windowingStrategy = WindowingStrategyTranslation.fromProto(windowingStrategyProto, rehydratedComponents);
        } catch (InvalidProtocolBufferException e) {
            throw new IllegalStateException(String.format("Unable to hydrate side input windowing strategy %s.", windowingStrategyProto), e);
        }
        Coder<WindowedValue<Object>> coder = instantiateCoder(collectionId, components);
        // side input materialization via GBK (T -> Iterable<T>)
        WindowedValueCoder wvCoder = (WindowedValueCoder) coder;
        coder = wvCoder.withValueCoder(IterableCoder.of(wvCoder.getValueCoder()));
        sideInputs.put(sideInputId, new RunnerPCollectionView<>(null, new TupleTag<>(sideInputTag), viewFn, // TODO: support custom mapping fn
        windowingStrategy.getWindowFn().getDefaultWindowMappingFn(), windowingStrategy, coder));
    }
    return sideInputs;
}

Example 15 with TypeDescriptor

use of org.apache.beam.sdk.values.TypeDescriptor in project beam by apache.

the class FieldValueTypeInformation method forSetter.

public static FieldValueTypeInformation forSetter(Method method, String setterPrefix) {
    String name;
    if (method.getName().startsWith(setterPrefix)) {
        name = ReflectUtils.stripPrefix(method.getName(), setterPrefix);
    } else {
        throw new RuntimeException("Setter has wrong prefix " + method.getName());
    }
    TypeDescriptor type = TypeDescriptor.of(method.getGenericParameterTypes()[0]);
    boolean nullable = hasSingleNullableParameter(method);
    return new AutoValue_FieldValueTypeInformation.Builder().setName(name).setNullable(nullable).setType(type).setRawType(type.getRawType()).setMethod(method).setElementType(getIterableComponentType(type)).setMapKeyType(getMapKeyType(type)).setMapValueType(getMapValueType(type)).setOneOfTypes(Collections.emptyMap()).build();
}