content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.beam.sdk.testing.expansion;
  
  import com.google.auto.service.AutoService;
  import java.io.Serializable;
  import java.util.Map;
  import java.util.Set;
  import java.util.stream.StreamSupport;
  import org.apache.avro.Schema;
  import org.apache.avro.generic.GenericRecord;
  import org.apache.beam.model.pipeline.v1.RunnerApi;
  import org.apache.beam.sdk.Pipeline;
  import org.apache.beam.sdk.expansion.ExternalTransformRegistrar;
  import org.apache.beam.sdk.expansion.service.ExpansionService;
  import org.apache.beam.sdk.expansion.service.TransformProvider;
  import org.apache.beam.sdk.extensions.avro.coders.AvroCoder;
  import org.apache.beam.sdk.io.FileIO;
  import org.apache.beam.sdk.io.TextIO;
  import org.apache.beam.sdk.io.parquet.ParquetIO;
  import org.apache.beam.sdk.options.PipelineOptions;
  import org.apache.beam.sdk.transforms.Count;
  import org.apache.beam.sdk.transforms.DoFn;
  import org.apache.beam.sdk.transforms.ExternalTransformBuilder;
  import org.apache.beam.sdk.transforms.Filter;
  import org.apache.beam.sdk.transforms.Flatten;
  import org.apache.beam.sdk.transforms.GroupByKey;
  import org.apache.beam.sdk.transforms.MapElements;
  import org.apache.beam.sdk.transforms.PTransform;
  import org.apache.beam.sdk.transforms.ParDo;
  import org.apache.beam.sdk.transforms.Partition;
  import org.apache.beam.sdk.transforms.Sum;
  import org.apache.beam.sdk.transforms.Values;
  import org.apache.beam.sdk.transforms.View;
  import org.apache.beam.sdk.transforms.join.CoGbkResult;
  import org.apache.beam.sdk.transforms.join.CoGroupByKey;
  import org.apache.beam.sdk.transforms.join.KeyedPCollectionTuple;
  import org.apache.beam.sdk.values.KV;
  import org.apache.beam.sdk.values.PBegin;
  import org.apache.beam.sdk.values.PCollection;
  import org.apache.beam.sdk.values.PCollectionList;
  import org.apache.beam.sdk.values.PCollectionTuple;
  import org.apache.beam.sdk.values.PCollectionView;
  import org.apache.beam.sdk.values.TupleTag;
  import org.apache.beam.sdk.values.TypeDescriptors;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
  import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
  
  /** An {@link org.apache.beam.sdk.util.construction.expansion.ExpansionService} useful for tests. */
  @SuppressWarnings({
    "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
  })
  public class TestExpansionService {
  
    private static final String TEST_PREFIX_URN = "beam:transforms:xlang:test:prefix";
    private static final String TEST_MULTI_URN = "beam:transforms:xlang:test:multi";
    private static final String TEST_GBK_URN = "beam:transforms:xlang:test:gbk";
    private static final String TEST_CGBK_URN = "beam:transforms:xlang:test:cgbk";
    private static final String TEST_COMGL_URN = "beam:transforms:xlang:test:comgl";
    private static final String TEST_COMPK_URN = "beam:transforms:xlang:test:compk";
    private static final String TEST_FLATTEN_URN = "beam:transforms:xlang:test:flatten";
    private static final String TEST_PARTITION_URN = "beam:transforms:xlang:test:partition";
    private static final String TEST_PARQUET_WRITE_URN = "beam:transforms:xlang:test:parquet_write";
  
    private static final String TEST_COUNT_URN = "beam:transforms:xlang:count";
    private static final String TEST_FILTER_URN = "beam:transforms:xlang:filter_less_than_eq";
    private static final String TEST_PARQUET_READ_URN = "beam:transforms:xlang:parquet_read";
    private static final String TEST_TEXTIO_READ_URN = "beam:transforms:xlang:textio_read";
  
    @AutoService(ExpansionService.ExpansionServiceRegistrar.class)
    public static class TestServiceRegistrar implements ExpansionService.ExpansionServiceRegistrar {
      @Override
      public Map<String, TransformProvider> knownTransforms() {
        ImmutableMap.Builder<String, TransformProvider> builder = ImmutableMap.builder();
        builder.put(TEST_CGBK_URN, new TestCoGroupByKeyTransformProvider());
        builder.put(TEST_FLATTEN_URN, new TestFlattenTransformProvider());
        return builder.build();
      }
  
      public static class TestCoGroupByKeyTransformProvider
          implements TransformProvider<
              KeyedPCollectionTuple<Long>, PCollection<KV<Long, Iterable<String>>>> {
        public static class TestCoGroupByKeyTransform
            extends PTransform<KeyedPCollectionTuple<Long>, PCollection<KV<Long, Iterable<String>>>> {
          @Override
          public PCollection<KV<Long, Iterable<String>>> expand(KeyedPCollectionTuple<Long> input) {
            Set<String> tagSet = ImmutableSet.of("col1", "col2");
            return input
                .apply(CoGroupByKey.create())
                .apply(
                    ParDo.of(
                        new DoFn<KV<Long, CoGbkResult>, KV<Long, Iterable<String>>>() {
                          @ProcessElement
                          public void processElement(
                              @Element KV<Long, CoGbkResult> kv,
                              OutputReceiver<KV<Long, Iterable<String>>> out) {
                            Iterable<String> iter =
                                () ->
                                    tagSet.stream()
                                        .flatMap(
                                            (String t) ->
                                                StreamSupport.stream(
                                                    kv.getValue().<String>getAll(t).spliterator(),
                                                    false))
                                        .iterator();
                            out.output(KV.of(kv.getKey(), iter));
                          }
                        }));
          }
        }
  
        @Override
        public KeyedPCollectionTuple<Long> createInput(
            Pipeline p, Map<String, PCollection<?>> inputs) {
          KeyedPCollectionTuple inputTuple = KeyedPCollectionTuple.empty(p);
          for (Map.Entry<String, PCollection<?>> entry : inputs.entrySet()) {
            inputTuple = inputTuple.and(new TupleTag(entry.getKey()), entry.getValue());
          }
          return inputTuple;
        }
  
        @Override
        public PTransform<KeyedPCollectionTuple<Long>, PCollection<KV<Long, Iterable<String>>>>
            getTransform(RunnerApi.FunctionSpec spec, PipelineOptions options) {
          return new TestCoGroupByKeyTransform();
        }
      }
  
      public static class TestFlattenTransformProvider
          implements TransformProvider<PCollectionList<Long>, PCollection<Long>> {
        @Override
        public PCollectionList<Long> createInput(Pipeline p, Map<String, PCollection<?>> inputs) {
          PCollectionList<Long> inputList = PCollectionList.empty(p);
          for (PCollection<?> collection : inputs.values()) {
            inputList = inputList.and((PCollection<Long>) collection);
          }
          return inputList;
        }
  
        @Override
        public PTransform<PCollectionList<Long>, PCollection<Long>> getTransform(
            RunnerApi.FunctionSpec spec, PipelineOptions options) {
          return Flatten.pCollections();
        }
      }
    }
  
    @AutoService(ExternalTransformRegistrar.class)
    public static class TestTransformRegistrar implements ExternalTransformRegistrar {
      private static String rawSchema =
          "{ \"type\": \"record\", \"name\": \"testrecord\", \"fields\": "
              + "[ {\"name\": \"name\", \"type\": \"string\"} ]}";
      private static Schema schema = new Schema.Parser().parse(rawSchema);
  
      @Override
      public Map<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>> knownBuilders() {
        ImmutableMap.Builder<String, Class<? extends ExternalTransformBuilder<?, ?, ?>>> builder =
            ImmutableMap.builder();
        builder.put(TEST_PREFIX_URN, PrefixBuilder.class);
        builder.put(TEST_MULTI_URN, MultiBuilder.class);
        builder.put(TEST_GBK_URN, GBKBuilder.class);
        builder.put(TEST_COMGL_URN, CombineGloballyBuilder.class);
        builder.put(TEST_COMPK_URN, CombinePerKeyBuilder.class);
        builder.put(TEST_COUNT_URN, CountBuilder.class);
        builder.put(TEST_FILTER_URN, FilterBuilder.class);
        builder.put(TEST_PARTITION_URN, PartitionBuilder.class);
        builder.put(TEST_PARQUET_WRITE_URN, ParquetWriteBuilder.class);
        builder.put(TEST_PARQUET_READ_URN, ParquetReadBuilder.class);
        builder.put(TEST_TEXTIO_READ_URN, TextIOReadBuilder.class);
        return builder.build();
      }
  
      public static class EmptyConfiguration {}
  
      public static class StringConfiguration implements Serializable {
        private String data;
  
        public void setData(String data) {
          this.data = data;
        }
      }
  
      public static class PrefixBuilder
          implements ExternalTransformBuilder<
              StringConfiguration, PCollection<? extends String>, PCollection<String>> {
        @Override
        public PTransform<PCollection<? extends String>, PCollection<String>> buildExternal(
            StringConfiguration configuration) {
          return MapElements.into(TypeDescriptors.strings())
              .via((String x) -> configuration.data + x);
        }
      }
  
      public static class MultiBuilder
          implements ExternalTransformBuilder<
              EmptyConfiguration, PCollectionTuple, PCollectionTuple> {
        public static class TestMultiPTransform
            extends PTransform<PCollectionTuple, PCollectionTuple> {
          @Override
          public PCollectionTuple expand(PCollectionTuple input) {
            PCollectionView<String> sideView = input.<String>get("side").apply(View.asSingleton());
            PCollection<String> main =
                PCollectionList.<String>of(input.get("main1"))
                    .and(input.get("main2"))
                    .apply(Flatten.pCollections())
                    .apply(
                        ParDo.of(
                                new DoFn<String, String>() {
                                  @ProcessElement
                                  public void processElement(
                                      @Element String x,
                                      OutputReceiver<String> out,
                                      DoFn<String, String>.ProcessContext c) {
                                    out.output(x + c.sideInput(sideView));
                                  }
                                })
                            .withSideInputs(sideView));
            PCollection<String> side =
                input
                    .<String>get("side")
                    .apply(MapElements.into(TypeDescriptors.strings()).via((String x) -> x + x));
            return PCollectionTuple.of("main", main).and("side", side);
          }
        }
  
        @Override
        public PTransform<PCollectionTuple, PCollectionTuple> buildExternal(
            EmptyConfiguration configuration) {
          return new TestMultiPTransform();
        }
      }
  
      public static class GBKBuilder
          implements ExternalTransformBuilder<
              EmptyConfiguration,
              PCollection<KV<Long, String>>,
              PCollection<KV<Long, Iterable<String>>>> {
        @Override
        public PTransform<PCollection<KV<Long, String>>, PCollection<KV<Long, Iterable<String>>>>
            buildExternal(EmptyConfiguration configuration) {
          return GroupByKey.create();
        }
      }
  
      public static class CombineGloballyBuilder
          implements ExternalTransformBuilder<
              EmptyConfiguration, PCollection<Long>, PCollection<Long>> {
        @Override
        public PTransform<PCollection<Long>, PCollection<Long>> buildExternal(
            EmptyConfiguration configuration) {
          return Sum.longsGlobally();
        }
      }
  
      public static class CombinePerKeyBuilder
          implements ExternalTransformBuilder<
              EmptyConfiguration, PCollection<KV<String, Long>>, PCollection<KV<String, Long>>> {
        @Override
        public PTransform<PCollection<KV<String, Long>>, PCollection<KV<String, Long>>> buildExternal(
            EmptyConfiguration configuration) {
          return Sum.longsPerKey();
        }
      }
  
      public static class CountBuilder
          implements ExternalTransformBuilder<
              EmptyConfiguration, PCollection<Object>, PCollection<KV<Object, Long>>> {
        @Override
        public PTransform<PCollection<Object>, PCollection<KV<Object, Long>>> buildExternal(
            EmptyConfiguration configuration) {
          return Count.perElement();
        }
      }
  
      public static class FilterBuilder
          implements ExternalTransformBuilder<
              StringConfiguration, PCollection<String>, PCollection<String>> {
        @Override
        public PTransform<PCollection<String>, PCollection<String>> buildExternal(
            StringConfiguration configuration) {
          return Filter.lessThanEq(configuration.data);
        }
      }
  
      public static class PartitionBuilder
          implements ExternalTransformBuilder<
              EmptyConfiguration, PCollection<Long>, PCollectionList<Long>> {
        @Override
        public PTransform<PCollection<Long>, PCollectionList<Long>> buildExternal(
            EmptyConfiguration configuration) {
          return Partition.of(2, (Long elem, int numP) -> elem % 2 == 0 ? 0 : 1);
        }
      }
  
      public static class ParquetWriteBuilder
          implements ExternalTransformBuilder<
              StringConfiguration, PCollection<GenericRecord>, PCollection<String>> {
        @Override
        public PTransform<PCollection<GenericRecord>, PCollection<String>> buildExternal(
            StringConfiguration configuration) {
          return new PTransform<PCollection<GenericRecord>, PCollection<String>>() {
            @Override
            public PCollection<String> expand(PCollection<GenericRecord> input) {
              return input
                  .apply(
                      FileIO.<GenericRecord>write()
                          .via(ParquetIO.sink(schema))
                          .to(configuration.data))
                  .getPerDestinationOutputFilenames()
                  .apply(Values.create());
            }
          };
        }
      }
  
      public static class ParquetReadBuilder
          implements ExternalTransformBuilder<
              StringConfiguration, PBegin, PCollection<GenericRecord>> {
        @Override
        public PTransform<PBegin, PCollection<GenericRecord>> buildExternal(
            StringConfiguration configuration) {
          return new PTransform<PBegin, PCollection<GenericRecord>>() {
            @Override
            public PCollection<GenericRecord> expand(PBegin input) {
              return input
                  .apply(FileIO.match().filepattern(configuration.data))
                  .apply(FileIO.readMatches())
                  .apply(ParquetIO.readFiles(schema))
                  .setCoder(AvroCoder.of(schema));
            }
          };
        }
      }
  
      public static class TextIOReadBuilder
          implements ExternalTransformBuilder<StringConfiguration, PBegin, PCollection<String>> {
        @Override
        public PTransform<PBegin, PCollection<String>> buildExternal(
            StringConfiguration configuration) {
          return new PTransform<PBegin, PCollection<String>>() {
            @Override
            public PCollection<String> expand(PBegin input) {
              return input.apply(TextIO.read().from(configuration.data));
            }
          };
        }
      }
    }
  }
  