use of org.apache.beam.sdk.coders.Coder in project beam by apache.
the class DoFnOperatorTest method keyedParDoSideInputCheckpointing.
@Test
public void keyedParDoSideInputCheckpointing() throws Exception {
sideInputCheckpointing(() -> {
StringUtf8Coder keyCoder = StringUtf8Coder.of();
Coder<WindowedValue<String>> coder = WindowedValue.getFullCoder(keyCoder, IntervalWindow.getCoder());
TupleTag<String> outputTag = new TupleTag<>("main-output");
KeySelector<WindowedValue<String>, ByteBuffer> keySelector = e -> FlinkKeyUtils.encodeKey(e.getValue(), keyCoder);
ImmutableMap<Integer, PCollectionView<?>> sideInputMapping = ImmutableMap.<Integer, PCollectionView<?>>builder().put(1, view1).put(2, view2).build();
DoFnOperator<String, String> doFnOperator = new DoFnOperator<>(new IdentityDoFn<>(), "stepName", coder, Collections.emptyMap(), outputTag, Collections.emptyList(), new DoFnOperator.MultiOutputOutputManagerFactory<>(outputTag, coder, new SerializablePipelineOptions(FlinkPipelineOptions.defaults())), WindowingStrategy.of(FixedWindows.of(Duration.millis(100))), sideInputMapping, /* side-input mapping */
ImmutableList.of(view1, view2), /* side inputs */
FlinkPipelineOptions.defaults(), keyCoder, keySelector, DoFnSchemaInformation.create(), Collections.emptyMap());
return new KeyedTwoInputStreamOperatorTestHarness<>(doFnOperator, keySelector, // we use a dummy key for the second input since it is considered to be broadcast
null, new CoderTypeInformation<>(FlinkKeyUtils.ByteBufferCoder.of(), FlinkPipelineOptions.defaults()));
});
}
use of org.apache.beam.sdk.coders.Coder in project beam by apache.
the class CommonCoderTest method executeSingleTest.
@Test
public void executeSingleTest() throws IOException {
assertCoderIsKnown(testSpec.getCoder());
Coder coder = instantiateCoder(testSpec.getCoder());
Object testValue = convertValue(testSpec.getValue(), testSpec.getCoder(), coder);
Context context = testSpec.getNested() ? Context.NESTED : Context.OUTER;
byte[] encoded = CoderUtils.encodeToByteArray(coder, testValue, context);
Object decodedValue = CoderUtils.decodeFromByteArray(coder, testSpec.getSerialized(), context);
if (!testSpec.getCoder().getNonDeterministic()) {
assertThat(testSpec.toString(), encoded, equalTo(testSpec.getSerialized()));
}
verifyDecodedValue(testSpec.getCoder(), decodedValue, testValue);
}
use of org.apache.beam.sdk.coders.Coder 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());
}
}
use of org.apache.beam.sdk.coders.Coder in project beam by apache.
the class FlattenP method tryProcess.
@Override
protected boolean tryProcess(int ordinal, @Nonnull Object item) {
Coder inputCoder = inputOrdinalCoders.get(ordinal);
WindowedValue<Object> windowedValue = Utils.decodeWindowedValue((byte[]) item, inputCoder);
return tryEmit(Utils.encode(windowedValue, outputCoder));
}
use of org.apache.beam.sdk.coders.Coder in project beam by apache.
the class CloudObjectTranslators method addComponents.
private static CloudObject addComponents(CloudObject base, List<? extends Coder<?>> components) {
if (!components.isEmpty()) {
List<CloudObject> cloudComponents = new ArrayList<>(components.size());
for (Coder component : components) {
cloudComponents.add(CloudObjects.asCloudObject(component));
}
Structs.addList(base, PropertyNames.COMPONENT_ENCODINGS, cloudComponents);
}
return base;
}
Aggregations