@Override public Coder<KV<K, V>> getOutputCoder() { return KvCoder.of(keyCoder, valueCoder); }
private KvCoder<K, InputT> getKvCoder( Coder<? extends KV<K, ? extends Iterable<InputT>>> inputCoder) { if (!(inputCoder instanceof KvCoder)) { throw new IllegalStateException("Combine.GroupedValues requires its input to use KvCoder"); } @SuppressWarnings({"unchecked", "rawtypes"}) KvCoder<K, ? extends Iterable<InputT>> kvCoder = (KvCoder) inputCoder; Coder<K> keyCoder = kvCoder.getKeyCoder(); Coder<? extends Iterable<InputT>> kvValueCoder = kvCoder.getValueCoder(); if (!(kvValueCoder instanceof IterableCoder)) { throw new IllegalStateException( "Combine.GroupedValues requires its input values to use " + "IterableCoder"); } @SuppressWarnings("unchecked") IterableCoder<InputT> inputValuesCoder = (IterableCoder<InputT>) kvValueCoder; Coder<InputT> inputValueCoder = inputValuesCoder.getElemCoder(); return KvCoder.of(keyCoder, inputValueCoder); }
@Override public Object structuralValue(KV<K, V> kv) { if (consistentWithEquals()) { return kv; } else { return KV.of( getKeyCoder().structuralValue(kv.getKey()), getValueCoder().structuralValue(kv.getValue())); } }
@Override public void verifyDeterministic() throws NonDeterministicException { verifyDeterministic(this, "Key coder must be deterministic", getKeyCoder()); verifyDeterministic(this, "Value coder must be deterministic", getValueCoder()); }
@Override public List<? extends Coder<?>> getComponents(KvCoder<?, ?> from) { return ImmutableList.of(from.getKeyCoder(), from.getValueCoder()); }
private static <K, V> PCollection<KV<K, V>> setValueCoder( PCollection<KV<K, V>> kvs, Coder<V> valueCoder) { // safe case because PCollection of KV always has KvCoder KvCoder<K, V> coder = (KvCoder<K, V>) kvs.getCoder(); return kvs.setCoder(KvCoder.of(coder.getKeyCoder(), valueCoder)); }
/** * Returns the {@code Coder} of the keys of the input to this transform, which is also used as the * {@code Coder} of the keys of the output of this transform. */ public static <K, V> Coder<K> getKeyCoder(Coder<KV<K, V>> inputCoder) { return getInputKvCoder(inputCoder).getKeyCoder(); }
@Override public Coder<OutputT> getDefaultOutputCoder( CoderRegistry registry, Coder<InputOrAccum<InputT, AccumT>> accumulatorCoder) throws CannotProvideCoderException { return fn.getDefaultOutputCoder(registry, inputCoder.getValueCoder()); }
@Override public boolean consistentWithEquals() { return kvCoder.consistentWithEquals(); } }
@Override public CloudObject toCloudObject(KvCoder target, SdkComponents sdkComponents) { CloudObject result = CloudObject.forClassName(CloudObjectKinds.KIND_PAIR); Structs.addBoolean(result, PropertyNames.IS_PAIR_LIKE, true); return addComponents( result, ImmutableList.<Coder<?>>of(target.getKeyCoder(), target.getValueCoder()), sdkComponents); }
private static <K, V> Coder<K> getKeyCoder(PCollection<KV<K, V>> pc) { // TODO: This should already have run coder inference for output, but may not have been consumed // as input yet (and won't be fully specified); This is fine // Assumes that the PCollection uses a KvCoder. Coder<?> entryCoder = pc.getCoder(); if (!(entryCoder instanceof KvCoder<?, ?>)) { throw new IllegalArgumentException("PCollection does not use a KvCoder"); } @SuppressWarnings("unchecked") KvCoder<K, V> coder = (KvCoder<K, V>) entryCoder; return coder.getKeyCoder(); }
@Override public Coder<OutputT> getDefaultOutputCoder( CoderRegistry registry, Coder<InputOrAccum<InputT, AccumT>> accumulatorCoder) throws CannotProvideCoderException { return fnWithContext.getDefaultOutputCoder(registry, inputCoder.getValueCoder()); }
@Override public boolean consistentWithEquals() { return kvCoder.consistentWithEquals(); } }
@Override public KvCoder<?, ?> fromComponents(List<Coder<?>> components) { return KvCoder.of(components.get(0), components.get(1)); } };
@Override public PCollection<KV<K1, Iterable<KV<K2, V>>>> expand(PCollection<KV<K1, KV<K2, V>>> input) { @SuppressWarnings("unchecked") KvCoder<K1, KV<K2, V>> inputCoder = (KvCoder<K1, KV<K2, V>>) input.getCoder(); return PCollection.createPrimitiveOutputInternal( input.getPipeline(), WindowingStrategy.globalDefault(), IsBounded.BOUNDED, KvCoder.of(inputCoder.getKeyCoder(), IterableCoder.of(inputCoder.getValueCoder()))); } }
/** * Constructs a {@link SingletonAssert} for the value of the provided {@link PCollection} with the * specified reason. The {@link PCollection} must have at most one value per key. * * <p>Note that the actual value must be coded by a {@link KvCoder}, not just any {@code Coder<K, * V>}. */ public static <K, V> SingletonAssert<Map<K, V>> thatMap( String reason, PCollection<KV<K, V>> actual) { @SuppressWarnings("unchecked") KvCoder<K, V> kvCoder = (KvCoder<K, V>) actual.getCoder(); return new PCollectionViewAssert<>( actual, View.asMap(), MapCoder.of(kvCoder.getKeyCoder(), kvCoder.getValueCoder()), PAssertionSite.capture(reason)); }
private Coder<K> getKeyCoder(Coder<KV<K, V>> coder) { checkState( coder instanceof KvCoder, "%s requires a coder of class %s." + " This is an internal error; this is checked during pipeline construction" + " but became corrupted.", getClass().getSimpleName(), KvCoder.class.getSimpleName()); @SuppressWarnings("unchecked") Coder<K> keyCoder = ((KvCoder<K, V>) coder).getKeyCoder(); return keyCoder; }
/** Returns the {@code Coder} of the values associated with the secondary keys. */ private static <PrimaryKeyT, SecondaryKeyT, ValueT> Coder<ValueT> getValueCoder( Coder<KV<PrimaryKeyT, Iterable<KV<SecondaryKeyT, ValueT>>>> inputCoder) { return getSecondaryKeyValueCoder(inputCoder).getValueCoder(); }
@Override public KvCoder fromCloudObject(CloudObject object) { List<Coder<?>> components = getComponents(object); checkArgument(components.size() == 2, "Expecting 2 components, got %s", components.size()); return KvCoder.of(components.get(0), components.get(1)); }
@Override public PCollection<KV<K, V>> expand(PCollection<KV<K, TimestampedValue<V>>> input) { KvCoder<K, TimestampedValue<V>> kvCoder = (KvCoder<K, TimestampedValue<V>>) input.getCoder(); TimestampedValueCoder<V> tvCoder = (TimestampedValueCoder<V>) kvCoder.getValueCoder(); return input .apply( ParDo.of( new DoFn<KV<K, TimestampedValue<V>>, KV<K, V>>() { @Override public Duration getAllowedTimestampSkew() { return Duration.millis(Long.MAX_VALUE); } @ProcessElement public void processElement( @Element KV<K, TimestampedValue<V>> kv, OutputReceiver<KV<K, V>> r) { r.outputWithTimestamp( KV.of(kv.getKey(), kv.getValue().getValue()), kv.getValue().getTimestamp()); } })) .setCoder(KvCoder.of(kvCoder.getKeyCoder(), tvCoder.getValueCoder())); } }