Java Code Examples for org.apache.beam.sdk.coders.CoderRegistry#getCoder()

final Coder<DestinationT> getDestinationCoderWithDefault(CoderRegistry registry)
    throws CannotProvideCoderException {
  Coder<DestinationT> destinationCoder = getDestinationCoder();
  if (destinationCoder != null) {
    return destinationCoder;
  }
  // If dynamicDestinations doesn't provide a coder, try to find it in the coder registry.
  @Nullable
  TypeDescriptor<DestinationT> descriptor =
      extractFromTypeParameters(
          this,
          DynamicDestinations.class,
          new TypeVariableExtractor<
              DynamicDestinations<UserT, DestinationT, OutputT>, DestinationT>() {});
  try {
    return registry.getCoder(descriptor);
  } catch (CannotProvideCoderException e) {
    throw new CannotProvideCoderException(
        "Failed to infer coder for DestinationT from type "
            + descriptor
            + ", please provide it explicitly by overriding getDestinationCoder()",
        e);
  }
}
 

private static <T> Coder<T> inferCoder(
    @Nullable Coder<T> explicitCoder,
    SerializableFunction<GenericRecord, T> parseFn,
    CoderRegistry coderRegistry) {
  if (explicitCoder != null) {
    return explicitCoder;
  }
  // If a coder was not specified explicitly, infer it from parse fn.
  try {
    return coderRegistry.getCoder(TypeDescriptors.outputOf(parseFn));
  } catch (CannotProvideCoderException e) {
    throw new IllegalArgumentException(
        "Unable to infer coder for output of parseFn. Specify it explicitly using withCoder().",
        e);
  }
}
 

Coder<DestinationT> getDestinationCoderWithDefault(CoderRegistry registry)
    throws CannotProvideCoderException {
  Coder<DestinationT> destinationCoder = getDestinationCoder();
  if (destinationCoder != null) {
    return destinationCoder;
  }
  // If dynamicDestinations doesn't provide a coder, try to find it in the coder registry.
  TypeDescriptor<DestinationT> descriptor =
      extractFromTypeParameters(
          this,
          DynamicDestinations.class,
          new TypeDescriptors.TypeVariableExtractor<
              DynamicDestinations<T, DestinationT>, DestinationT>() {});
  try {
    return registry.getCoder(descriptor);
  } catch (CannotProvideCoderException e) {
    throw new CannotProvideCoderException(
        "Failed to infer coder for DestinationT from type "
            + descriptor
            + ", please provide it explicitly by overriding getDestinationCoder()",
        e);
  }
}
 

@Override
public PCollection<KV<K, V>> expand(PCollection<V> in) {
  PCollection<KV<K, V>> result =
      in.apply(
          "AddKeys",
          MapElements.via(
              new SimpleFunction<V, KV<K, V>>() {
                @Override
                public KV<K, V> apply(V element) {
                  return KV.of(fn.apply(element), element);
                }
              }));

  try {
    Coder<K> keyCoder;
    CoderRegistry coderRegistry = in.getPipeline().getCoderRegistry();
    if (keyType == null) {
      keyCoder = coderRegistry.getOutputCoder(fn, in.getCoder());
    } else {
      keyCoder = coderRegistry.getCoder(keyType);
    }
    // TODO: Remove when we can set the coder inference context.
    result.setCoder(KvCoder.of(keyCoder, in.getCoder()));
  } catch (CannotProvideCoderException exc) {
    // let lazy coder inference have a try
  }

  return result;
}
 

@Override
public Coder<AccumT> getAccumulatorCoder(CoderRegistry registry, Coder<InputT> inputCoder)
    throws CannotProvideCoderException {
  return registry.getCoder(
      getClass(),
      AbstractGlobalCombineFn.class,
      ImmutableMap.<Type, Coder<?>>of(getInputTVariable(), inputCoder),
      getAccumTVariable());
}
 

@Override
public Coder<OutputT> getDefaultOutputCoder(CoderRegistry registry, Coder<InputT> inputCoder)
    throws CannotProvideCoderException {
  return registry.getCoder(
      getClass(),
      AbstractGlobalCombineFn.class,
      ImmutableMap.<Type, Coder<?>>of(
          getInputTVariable(),
          inputCoder,
          getAccumTVariable(),
          this.getAccumulatorCoder(registry, inputCoder)),
      getOutputTVariable());
}
 

/**
 * Try to provide coders for as many of the type arguments of given {@link
 * DoFnSignature.StateDeclaration} as possible.
 */
private static <InputT> Coder[] codersForStateSpecTypes(
    DoFnSignature.StateDeclaration stateDeclaration,
    CoderRegistry coderRegistry,
    Coder<InputT> inputCoder) {
  Type stateType = stateDeclaration.stateType().getType();

  if (!(stateType instanceof ParameterizedType)) {
    // No type arguments means no coders to infer.
    return new Coder[0];
  }

  Type[] typeArguments = ((ParameterizedType) stateType).getActualTypeArguments();
  Coder[] coders = new Coder[typeArguments.length];

  for (int i = 0; i < typeArguments.length; i++) {
    Type typeArgument = typeArguments[i];
    TypeDescriptor<?> typeDescriptor = TypeDescriptor.of(typeArgument);
    try {
      coders[i] = coderRegistry.getCoder(typeDescriptor);
    } catch (CannotProvideCoderException e) {
      try {
        coders[i] =
            coderRegistry.getCoder(
                typeDescriptor, inputCoder.getEncodedTypeDescriptor(), inputCoder);
      } catch (CannotProvideCoderException ignored) {
        // Since not all type arguments will have a registered coder we ignore this exception.
      }
    }
  }

  return coders;
}
 

@VisibleForTesting
Coder<T> inferCoder(CoderRegistry coderRegistry) {
  if (getCoder() != null) {
    return getCoder();
  }

  try {
    return coderRegistry.getCoder(TypeDescriptors.outputOf(getParseFn()));
  } catch (CannotProvideCoderException e) {
    throw new IllegalArgumentException(
        "Unable to infer coder for output of parseFn. Specify it explicitly using withCoder().",
        e);
  }
}
 

/**
 * Returns the default coder for a given type descriptor. Coder Registry is queried for correct
 * coder, if not found in Coder Registry, then check if the type descriptor provided is of type
 * Writable, then WritableCoder is returned, else exception is thrown "Cannot find coder".
 */
@SuppressWarnings({"unchecked", "WeakerAccess"})
public <T> Coder<T> getDefaultCoder(TypeDescriptor<?> typeDesc, CoderRegistry coderRegistry) {
  Class classType = typeDesc.getRawType();
  try {
    return (Coder<T>) coderRegistry.getCoder(typeDesc);
  } catch (CannotProvideCoderException e) {
    if (Writable.class.isAssignableFrom(classType)) {
      return (Coder<T>) WritableCoder.of(classType);
    }
    throw new IllegalStateException(
        String.format("Cannot find coder for %s  : ", typeDesc) + e.getMessage(), e);
  }
}
 

@VisibleForTesting
Coder<T> inferCoder(CoderRegistry coderRegistry) {
  try {
    return getCoder() != null
        ? getCoder()
        : coderRegistry.getCoder(TypeDescriptors.outputOf(getParseFn()));
  } catch (CannotProvideCoderException e) {
    throw new IllegalArgumentException(
        "Unable to infer coder for output of parseFn ("
            + TypeDescriptors.outputOf(getParseFn())
            + "). Specify it explicitly using withCoder().",
        e);
  }
}
 

private static <InputT, KeyT, ValueT, AccT, OutputT>
    Combine.CombineFn<ValueT, AccT, OutputT> asCombineFn(
        ReduceByKey<InputT, KeyT, ValueT, AccT, OutputT> operator) {

  @SuppressWarnings("unchecked")
  ReduceByKey<InputT, KeyT, ValueT, AccT, OutputT> cast = (ReduceByKey) operator;

  VoidFunction<AccT> accumulatorFactory = cast.getAccumulatorFactory();
  BinaryFunction<AccT, ValueT, AccT> accumulate = cast.getAccumulate();
  CombinableBinaryFunction<AccT> mergeAccumulators = cast.getMergeAccumulators();
  UnaryFunction<AccT, OutputT> outputFn = cast.getOutputFn();
  TypeDescriptor<AccT> accumulatorType = cast.getAccumulatorType();

  return new Combine.CombineFn<ValueT, AccT, OutputT>() {

    @Override
    public AccT createAccumulator() {
      return accumulatorFactory.apply();
    }

    @Override
    public Coder<AccT> getAccumulatorCoder(CoderRegistry registry, Coder<ValueT> inputCoder)
        throws CannotProvideCoderException {
      return registry.getCoder(accumulatorType);
    }

    @Override
    public AccT addInput(AccT mutableAccumulator, ValueT input) {
      return accumulate.apply(mutableAccumulator, input);
    }

    @Override
    public AccT mergeAccumulators(Iterable<AccT> accumulators) {
      AccT accumulated = null;
      for (AccT o : accumulators) {
        if (accumulated == null) {
          accumulated = o;
        } else {
          accumulated = mergeAccumulators.apply(accumulated, o);
        }
      }
      return accumulated;
    }

    @Override
    public OutputT extractOutput(AccT accumulator) {
      return outputFn.apply(accumulator);
    }
  };
}
 

/**
 * Attempt to infer the type for some very common Apache Beam parameterized types.
 *
 * <p>TODO: Instead, build a TypeDescriptor so that the {@link CoderRegistry} is invoked for the
 * type instead of hard coding the coders for common types.
 */
private static Coder<?> inferCoderFromObject(
    CoderRegistry coderRegistry, SchemaRegistry schemaRegistry, Object o)
    throws CannotProvideCoderException {

  if (o == null) {
    return VoidCoder.of();
  }

  try {
    return SchemaCoder.of(
        schemaRegistry.getSchema(o.getClass()),
        TypeDescriptor.of(o.getClass()),
        (SerializableFunction) schemaRegistry.getToRowFunction(o.getClass()),
        (SerializableFunction) schemaRegistry.getFromRowFunction(o.getClass()));
  } catch (NoSuchSchemaException e) {
    // No schema.
  }

  if (o instanceof TimestampedValue) {
    return TimestampedValueCoder.of(
        inferCoderFromObject(coderRegistry, schemaRegistry, ((TimestampedValue) o).getValue()));
  } else if (o instanceof List) {
    return ListCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Set) {
    return SetCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Collection) {
    return CollectionCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Iterable) {
    return IterableCoder.of(inferCoderFromObjects(coderRegistry, schemaRegistry, (Iterable) o));
  } else if (o instanceof Map) {
    return MapCoder.of(
        inferCoderFromObjects(coderRegistry, schemaRegistry, ((Map) o).keySet()),
        inferCoderFromObjects(coderRegistry, schemaRegistry, ((Map) o).entrySet()));
  } else if (o instanceof KV) {
    return KvCoder.of(
        inferCoderFromObject(coderRegistry, schemaRegistry, ((KV) o).getKey()),
        inferCoderFromObject(coderRegistry, schemaRegistry, ((KV) o).getValue()));
  } else {
    return coderRegistry.getCoder(o.getClass());
  }
}
 

/** Doesn't split the restriction. */
@SuppressWarnings({"unused", "unchecked"})
public <RestrictionT> Coder<RestrictionT> invokeGetRestrictionCoder(CoderRegistry registry)
    throws CannotProvideCoderException {
  return (Coder) registry.getCoder(restrictionType);
}
 

@SuppressWarnings({"unused", "unchecked"})
public <WatermarkEstimatorStateT>
    Coder<WatermarkEstimatorStateT> invokeGetWatermarkEstimatorStateCoder(
        CoderRegistry registry) throws CannotProvideCoderException {
  return (Coder) registry.getCoder(watermarkEstimatorStateType);
}