@Override
  public PTransform getTransform(FunctionSpec spec, PipelineOptions options) {
    SchemaTransformPayload payload;
    try {
      payload = SchemaTransformPayload.parseFrom(spec.getPayload());
      String identifier = payload.getIdentifier();
      if (!schemaTransformProviders.containsKey(identifier)) {
        throw new RuntimeException(
            "Did not find a SchemaTransformProvider with the identifier " + identifier);
      }

    } catch (InvalidProtocolBufferException e) {
      throw new IllegalArgumentException(
          "Invalid payload type for URN " + getUrn(ExpansionMethods.Enum.SCHEMA_TRANSFORM), e);
    }

    String identifier = payload.getIdentifier();
    org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider provider =
        schemaTransformProviders.get(identifier);
    if (provider == null) {
      throw new IllegalArgumentException(
          "Could not find a SchemaTransform with identifier " + identifier);
    }

    Schema configSchemaFromRequest =
        SchemaTranslation.schemaFromProto((payload.getConfigurationSchema()));
    Schema configSchemaFromProvider = provider.configurationSchema();

    if (!configSchemaFromRequest.assignableTo(configSchemaFromProvider)) {
      throw new IllegalArgumentException(
          String.format(
              "Config schema provided with the expansion request %s is not compatible with the "
                  + "config of the Schema transform %s.",
              configSchemaFromRequest, configSchemaFromProvider));
    }

    Row configRow;
    try {
      configRow =
          RowCoder.of(configSchemaFromRequest).decode(payload.getConfigurationRow().newInput());
    } catch (IOException e) {
      throw new RuntimeException("Error decoding payload", e);
    }

    return provider.from(configRow);
  }
@Test
  public void testSchemaTransformExpansionMultiInputMultiOutput() {
    Pipeline p = Pipeline.create();
    p.apply(Impulse.create());
    p.apply(Impulse.create());
    RunnerApi.Pipeline pipelineProto = PipelineTranslation.toProto(p);

    List<String> inputPcollIds = new ArrayList<>();
    for (RunnerApi.PTransform transform :
        pipelineProto.getComponents().getTransformsMap().values()) {
      inputPcollIds.add(Iterables.getOnlyElement(transform.getOutputsMap().values()));
    }
    assertEquals(2, inputPcollIds.size());

    Row configRow =
        Row.withSchema(TEST_SCHEMATRANSFORM_CONFIG_SCHEMA)
            .withFieldValue("int1", 111)
            .withFieldValue("int2", 222)
            .withFieldValue("str1", "aaa")
            .withFieldValue("str2", "bbb")
            .build();

    ExpansionApi.ExpansionRequest request =
        ExpansionApi.ExpansionRequest.newBuilder()
            .setComponents(pipelineProto.getComponents())
            .setTransform(
                RunnerApi.PTransform.newBuilder()
                    .setUniqueName(TEST_NAME)
                    .setSpec(createSpec("dummy_id_multi_input_multi_output", configRow))
                    .putInputs("input1", inputPcollIds.get(0))
                    .putInputs("input2", inputPcollIds.get(1)))
            .setNamespace(TEST_NAMESPACE)
            .build();

    ExpansionApi.ExpansionResponse response = expansionService.expand(request);
    RunnerApi.PTransform expandedTransform = response.getTransform();

    assertEquals(6, expandedTransform.getSubtransformsCount());
    assertEquals(2, expandedTransform.getInputsCount());
    assertEquals(2, expandedTransform.getOutputsCount());
    verifyLeafTransforms(response, 2);
  }