Search in sources :

Example 1 with TimestampPrefixingWindowCoder

use of org.apache.beam.sdk.coders.TimestampPrefixingWindowCoder in project beam by apache.

the class CommonCoderTest method convertValue.

/**
 * Converts from JSON-auto-deserialized types into the proper Java types for the known coders.
 */
private static Object convertValue(Object value, CommonCoder coderSpec, Coder coder) {
    String s = coderSpec.getUrn();
    if (s.equals(getUrn(StandardCoders.Enum.BYTES))) {
        return ((String) value).getBytes(StandardCharsets.ISO_8859_1);
    } else if (s.equals(getUrn(StandardCoders.Enum.BOOL))) {
        return value;
    } else if (s.equals(getUrn(StandardCoders.Enum.STRING_UTF8))) {
        return value;
    } else if (s.equals(getUrn(StandardCoders.Enum.KV))) {
        Coder keyCoder = ((KvCoder) coder).getKeyCoder();
        Coder valueCoder = ((KvCoder) coder).getValueCoder();
        Map<String, Object> kvMap = (Map<String, Object>) value;
        Object k = convertValue(kvMap.get("key"), coderSpec.getComponents().get(0), keyCoder);
        Object v = convertValue(kvMap.get("value"), coderSpec.getComponents().get(1), valueCoder);
        return KV.of(k, v);
    } else if (s.equals(getUrn(StandardCoders.Enum.VARINT))) {
        return ((Number) value).longValue();
    } else if (s.equals(getUrn(StandardCoders.Enum.TIMER))) {
        Map<String, Object> kvMap = (Map<String, Object>) value;
        Coder<?> keyCoder = ((Timer.Coder) coder).getValueCoder();
        Coder<? extends BoundedWindow> windowCoder = ((Timer.Coder) coder).getWindowCoder();
        List<BoundedWindow> windows = new ArrayList<>();
        for (Object window : (List<Object>) kvMap.get("windows")) {
            windows.add((BoundedWindow) convertValue(window, coderSpec.getComponents().get(1), windowCoder));
        }
        if ((boolean) kvMap.get("clearBit")) {
            return Timer.cleared(convertValue(kvMap.get("userKey"), coderSpec.getComponents().get(0), keyCoder), (String) kvMap.get("dynamicTimerTag"), windows);
        }
        Map<String, Object> paneInfoMap = (Map<String, Object>) kvMap.get("pane");
        PaneInfo paneInfo = PaneInfo.createPane((boolean) paneInfoMap.get("is_first"), (boolean) paneInfoMap.get("is_last"), PaneInfo.Timing.valueOf((String) paneInfoMap.get("timing")), (int) paneInfoMap.get("index"), (int) paneInfoMap.get("on_time_index"));
        return Timer.of(convertValue(kvMap.get("userKey"), coderSpec.getComponents().get(0), keyCoder), (String) kvMap.get("dynamicTimerTag"), windows, new Instant(((Number) kvMap.get("fireTimestamp")).longValue()), new Instant(((Number) kvMap.get("holdTimestamp")).longValue()), paneInfo);
    } else if (s.equals(getUrn(StandardCoders.Enum.INTERVAL_WINDOW))) {
        Map<String, Object> kvMap = (Map<String, Object>) value;
        Instant end = new Instant(((Number) kvMap.get("end")).longValue());
        Duration span = Duration.millis(((Number) kvMap.get("span")).longValue());
        return new IntervalWindow(end.minus(span), span);
    } else if (s.equals(getUrn(StandardCoders.Enum.ITERABLE)) || s.equals(getUrn(StandardCoders.Enum.STATE_BACKED_ITERABLE))) {
        Coder elementCoder = ((IterableLikeCoder) coder).getElemCoder();
        List<Object> elements = (List<Object>) value;
        List<Object> convertedElements = new ArrayList<>();
        for (Object element : elements) {
            convertedElements.add(convertValue(element, coderSpec.getComponents().get(0), elementCoder));
        }
        return convertedElements;
    } else if (s.equals(getUrn(StandardCoders.Enum.GLOBAL_WINDOW))) {
        return GlobalWindow.INSTANCE;
    } else if (s.equals(getUrn(StandardCoders.Enum.WINDOWED_VALUE)) || s.equals(getUrn(StandardCoders.Enum.PARAM_WINDOWED_VALUE))) {
        Map<String, Object> kvMap = (Map<String, Object>) value;
        Coder valueCoder = ((WindowedValue.FullWindowedValueCoder) coder).getValueCoder();
        Coder windowCoder = ((WindowedValue.FullWindowedValueCoder) coder).getWindowCoder();
        Object windowValue = convertValue(kvMap.get("value"), coderSpec.getComponents().get(0), valueCoder);
        Instant timestamp = new Instant(((Number) kvMap.get("timestamp")).longValue());
        List<BoundedWindow> windows = new ArrayList<>();
        for (Object window : (List<Object>) kvMap.get("windows")) {
            windows.add((BoundedWindow) convertValue(window, coderSpec.getComponents().get(1), windowCoder));
        }
        Map<String, Object> paneInfoMap = (Map<String, Object>) kvMap.get("pane");
        PaneInfo paneInfo = PaneInfo.createPane((boolean) paneInfoMap.get("is_first"), (boolean) paneInfoMap.get("is_last"), PaneInfo.Timing.valueOf((String) paneInfoMap.get("timing")), (int) paneInfoMap.get("index"), (int) paneInfoMap.get("on_time_index"));
        return WindowedValue.of(windowValue, timestamp, windows, paneInfo);
    } else if (s.equals(getUrn(StandardCoders.Enum.DOUBLE))) {
        return Double.parseDouble((String) value);
    } else if (s.equals(getUrn(StandardCoders.Enum.ROW))) {
        Schema schema;
        try {
            schema = SchemaTranslation.schemaFromProto(SchemaApi.Schema.parseFrom(coderSpec.getPayload()));
        } catch (InvalidProtocolBufferException e) {
            throw new RuntimeException("Failed to parse schema payload for row coder", e);
        }
        return parseField(value, Schema.FieldType.row(schema));
    } else if (s.equals(getUrn(StandardCoders.Enum.SHARDED_KEY))) {
        Map<String, Object> kvMap = (Map<String, Object>) value;
        Coder<?> keyCoder = ((ShardedKey.Coder) coder).getKeyCoder();
        byte[] shardId = ((String) kvMap.get("shardId")).getBytes(StandardCharsets.ISO_8859_1);
        return ShardedKey.of(convertValue(kvMap.get("key"), coderSpec.getComponents().get(0), keyCoder), shardId);
    } else if (s.equals(getUrn(StandardCoders.Enum.CUSTOM_WINDOW))) {
        Map<String, Object> kvMap = (Map<String, Object>) value;
        Coder windowCoder = ((TimestampPrefixingWindowCoder) coder).getWindowCoder();
        return convertValue(kvMap.get("window"), coderSpec.getComponents().get(0), windowCoder);
    } else {
        throw new IllegalStateException("Unknown coder URN: " + coderSpec.getUrn());
    }
}
Also used : DoubleCoder(org.apache.beam.sdk.coders.DoubleCoder) IntervalWindowCoder(org.apache.beam.sdk.transforms.windowing.IntervalWindow.IntervalWindowCoder) ByteCoder(org.apache.beam.sdk.coders.ByteCoder) KvCoder(org.apache.beam.sdk.coders.KvCoder) VarLongCoder(org.apache.beam.sdk.coders.VarLongCoder) BooleanCoder(org.apache.beam.sdk.coders.BooleanCoder) TimestampPrefixingWindowCoder(org.apache.beam.sdk.coders.TimestampPrefixingWindowCoder) Coder(org.apache.beam.sdk.coders.Coder) RowCoder(org.apache.beam.sdk.coders.RowCoder) StringUtf8Coder(org.apache.beam.sdk.coders.StringUtf8Coder) IterableCoder(org.apache.beam.sdk.coders.IterableCoder) IterableLikeCoder(org.apache.beam.sdk.coders.IterableLikeCoder) Instant(org.joda.time.Instant) Schema(org.apache.beam.sdk.schemas.Schema) InvalidProtocolBufferException(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.InvalidProtocolBufferException) Duration(org.joda.time.Duration) ByteString(org.apache.beam.vendor.grpc.v1p43p2.com.google.protobuf.ByteString) Timer(org.apache.beam.runners.core.construction.Timer) WindowedValue(org.apache.beam.sdk.util.WindowedValue) PaneInfo(org.apache.beam.sdk.transforms.windowing.PaneInfo) BoundedWindow(org.apache.beam.sdk.transforms.windowing.BoundedWindow) List(java.util.List) ImmutableList.toImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList.toImmutableList) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) ArrayList(java.util.ArrayList) ImmutableMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap) ImmutableBiMap(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableBiMap) Map(java.util.Map) HashMap(java.util.HashMap) IntervalWindow(org.apache.beam.sdk.transforms.windowing.IntervalWindow)

Aggregations

ArrayList (java.util.ArrayList)1 HashMap (java.util.HashMap)1 List (java.util.List)1 Map (java.util.Map)1 Timer (org.apache.beam.runners.core.construction.Timer)1 BooleanCoder (org.apache.beam.sdk.coders.BooleanCoder)1 ByteCoder (org.apache.beam.sdk.coders.ByteCoder)1 Coder (org.apache.beam.sdk.coders.Coder)1 DoubleCoder (org.apache.beam.sdk.coders.DoubleCoder)1 IterableCoder (org.apache.beam.sdk.coders.IterableCoder)1 IterableLikeCoder (org.apache.beam.sdk.coders.IterableLikeCoder)1 KvCoder (org.apache.beam.sdk.coders.KvCoder)1 RowCoder (org.apache.beam.sdk.coders.RowCoder)1 StringUtf8Coder (org.apache.beam.sdk.coders.StringUtf8Coder)1 TimestampPrefixingWindowCoder (org.apache.beam.sdk.coders.TimestampPrefixingWindowCoder)1 VarLongCoder (org.apache.beam.sdk.coders.VarLongCoder)1 Schema (org.apache.beam.sdk.schemas.Schema)1 BoundedWindow (org.apache.beam.sdk.transforms.windowing.BoundedWindow)1 IntervalWindow (org.apache.beam.sdk.transforms.windowing.IntervalWindow)1 IntervalWindowCoder (org.apache.beam.sdk.transforms.windowing.IntervalWindow.IntervalWindowCoder)1