Java Code Examples for org.apache.beam.sdk.values.PCollection#setSchema()

The following examples show how to use org.apache.beam.sdk.values.PCollection#setSchema() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: ParDo.java    From beam with Apache License 2.0 5 votes vote down vote up 
@Override
public PCollection<OutputT> expand(PCollection<? extends InputT> input) {
  SchemaRegistry schemaRegistry = input.getPipeline().getSchemaRegistry();
  CoderRegistry coderRegistry = input.getPipeline().getCoderRegistry();
  finishSpecifyingStateSpecs(fn, coderRegistry, schemaRegistry, input.getCoder());
  TupleTag<OutputT> mainOutput = new TupleTag<>(MAIN_OUTPUT_TAG);
  PCollection<OutputT> res =
      input.apply(withOutputTags(mainOutput, TupleTagList.empty())).get(mainOutput);

  TypeDescriptor<OutputT> outputTypeDescriptor = getFn().getOutputTypeDescriptor();
  try {
    res.setSchema(
        schemaRegistry.getSchema(outputTypeDescriptor),
        outputTypeDescriptor,
        schemaRegistry.getToRowFunction(outputTypeDescriptor),
        schemaRegistry.getFromRowFunction(outputTypeDescriptor));
  } catch (NoSuchSchemaException e) {
    try {
      res.setCoder(
          coderRegistry.getCoder(
              outputTypeDescriptor,
              getFn().getInputTypeDescriptor(),
              ((PCollection<InputT>) input).getCoder()));
    } catch (CannotProvideCoderException e2) {
      // Ignore and leave coder unset.
    }
  }

  return res;
}
Example 2
Source File: AvroIO.java    From beam with Apache License 2.0 5 votes vote down vote up 
@Experimental(Kind.SCHEMAS)
private static <T> PCollection<T> setBeamSchema(
    PCollection<T> pc, Class<T> clazz, @Nullable Schema schema) {
  org.apache.beam.sdk.schemas.Schema beamSchema =
      org.apache.beam.sdk.schemas.utils.AvroUtils.getSchema(clazz, schema);
  if (beamSchema != null) {
    pc.setSchema(
        beamSchema,
        TypeDescriptor.of(clazz),
        org.apache.beam.sdk.schemas.utils.AvroUtils.getToRowFunction(clazz, schema),
        org.apache.beam.sdk.schemas.utils.AvroUtils.getFromRowFunction(clazz));
  }
  return pc;
}
Example 3
Source File: JdbcIO.java    From beam with Apache License 2.0 5 votes vote down vote up 
@Override
public PCollection<OutputT> expand(PCollection<ParameterT> input) {
  PCollection<OutputT> output =
      input
          .apply(
              ParDo.of(
                  new ReadFn<>(
                      getDataSourceProviderFn(),
                      getQuery(),
                      getParameterSetter(),
                      getRowMapper(),
                      getFetchSize())))
          .setCoder(getCoder());

  if (getOutputParallelization()) {
    output = output.apply(new Reparallelize<>());
  }

  try {
    TypeDescriptor<OutputT> typeDesc = getCoder().getEncodedTypeDescriptor();
    SchemaRegistry registry = input.getPipeline().getSchemaRegistry();
    Schema schema = registry.getSchema(typeDesc);
    output.setSchema(
        schema,
        typeDesc,
        registry.getToRowFunction(typeDesc),
        registry.getFromRowFunction(typeDesc));
  } catch (NoSuchSchemaException e) {
    // ignore
  }

  return output;
}