runners/core-construction-java/src/main/java/org/apache/beam/runners/core/construction/PTransformMatchers.java

/*
 * 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.runners.core.construction;

import com.google.common.base.MoreObjects;
import java.util.HashSet;
import java.util.Set;
import org.apache.beam.sdk.annotations.Experimental;
import org.apache.beam.sdk.annotations.Experimental.Kind;
import org.apache.beam.sdk.io.WriteFiles;
import org.apache.beam.sdk.runners.AppliedPTransform;
import org.apache.beam.sdk.runners.PTransformMatcher;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.transforms.PTransform;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.View.CreatePCollectionView;
import org.apache.beam.sdk.transforms.ViewFn;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature;
import org.apache.beam.sdk.transforms.reflect.DoFnSignature.ProcessElementMethod;
import org.apache.beam.sdk.transforms.reflect.DoFnSignatures;
import org.apache.beam.sdk.util.WindowedValue;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PValue;

/**
 * A {@link PTransformMatcher} that matches {@link PTransform PTransforms} based on the class of the
 * transform.
 *
 * <p>Once {@link PTransform PTransforms} have URNs, this will be removed and replaced with a
 * UrnPTransformMatcher.
 */
@Experimental(Kind.CORE_RUNNERS_ONLY)
public class PTransformMatchers {
  private PTransformMatchers() {}

  /**
   * Returns a {@link PTransformMatcher} that matches a {@link PTransform} if the class of the
   * {@link PTransform} is equal to the {@link Class} provided ot this matcher.
   */
  public static PTransformMatcher classEqualTo(Class<? extends PTransform> clazz) {
    return new EqualClassPTransformMatcher(clazz);
  }

  private static class EqualClassPTransformMatcher implements PTransformMatcher {
    private final Class<? extends PTransform> clazz;

    private EqualClassPTransformMatcher(Class<? extends PTransform> clazz) {
      this.clazz = clazz;
    }

    @Override
    public boolean matches(AppliedPTransform<?, ?, ?> application) {
      return application.getTransform().getClass().equals(clazz);
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper(EqualClassPTransformMatcher.class)
          .add("class", clazz)
          .toString();
    }
  }

  /**
   * A {@link PTransformMatcher} that matches a {@link ParDo.SingleOutput} containing a {@link DoFn}
   * that is splittable, as signified by {@link ProcessElementMethod#isSplittable()}.
   */
  public static PTransformMatcher splittableParDoSingle() {
    return new PTransformMatcher() {
      @Override
      public boolean matches(AppliedPTransform<?, ?, ?> application) {
        PTransform<?, ?> transform = application.getTransform();
        if (transform instanceof ParDo.SingleOutput) {
          DoFn<?, ?> fn = ((ParDo.SingleOutput<?, ?>) transform).getFn();
          DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
          return signature.processElement().isSplittable();
        }
        return false;
      }

      @Override
      public String toString() {
        return MoreObjects.toStringHelper("SplittableParDoSingleMatcher").toString();
      }
    };
  }

  /**
   * A {@link PTransformMatcher} that matches a {@link ParDo.SingleOutput} containing a {@link DoFn}
   * that uses state or timers, as specified by {@link DoFnSignature#usesState()} and {@link
   * DoFnSignature#usesTimers()}.
   */
  public static PTransformMatcher stateOrTimerParDoSingle() {
    return new PTransformMatcher() {
      @Override
      public boolean matches(AppliedPTransform<?, ?, ?> application) {
        PTransform<?, ?> transform = application.getTransform();
        if (transform instanceof ParDo.SingleOutput) {
          DoFn<?, ?> fn = ((ParDo.SingleOutput<?, ?>) transform).getFn();
          DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
          return signature.usesState() || signature.usesTimers();
        }
        return false;
      }

      @Override
      public String toString() {
        return MoreObjects.toStringHelper("StateOrTimerParDoSingleMatcher").toString();
      }
    };
  }

  /**
   * A {@link PTransformMatcher} that matches a {@link ParDo.MultiOutput} containing a {@link DoFn}
   * that is splittable, as signified by {@link ProcessElementMethod#isSplittable()}.
   */
  public static PTransformMatcher splittableParDoMulti() {
    return new PTransformMatcher() {
      @Override
      public boolean matches(AppliedPTransform<?, ?, ?> application) {
        PTransform<?, ?> transform = application.getTransform();
        if (transform instanceof ParDo.MultiOutput) {
          DoFn<?, ?> fn = ((ParDo.MultiOutput<?, ?>) transform).getFn();
          DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
          return signature.processElement().isSplittable();
        }
        return false;
      }

      @Override
      public String toString() {
        return MoreObjects.toStringHelper("SplittableParDoMultiMatcher").toString();
      }
    };
  }

  /**
   * A {@link PTransformMatcher} that matches a {@link ParDo.MultiOutput} containing a {@link DoFn}
   * that uses state or timers, as specified by {@link DoFnSignature#usesState()} and
   * {@link DoFnSignature#usesTimers()}.
   */
  public static PTransformMatcher stateOrTimerParDoMulti() {
    return new PTransformMatcher() {
      @Override
      public boolean matches(AppliedPTransform<?, ?, ?> application) {
        PTransform<?, ?> transform = application.getTransform();
        if (transform instanceof ParDo.MultiOutput) {
          DoFn<?, ?> fn = ((ParDo.MultiOutput<?, ?>) transform).getFn();
          DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
          return signature.usesState() || signature.usesTimers();
        }
        return false;
      }

      @Override
      public String toString() {
        return MoreObjects.toStringHelper("StateOrTimerParDoMultiMatcher").toString();
      }
    };
  }

  /**
   * A {@link PTransformMatcher} which matches a {@link ParDo.SingleOutput} or {@link
   * ParDo.MultiOutput} where the {@link DoFn} is of the provided type.
   */
  public static PTransformMatcher parDoWithFnType(final Class<? extends DoFn> fnType) {
    return new PTransformMatcher() {
      @Override
      public boolean matches(AppliedPTransform<?, ?, ?> application) {
        DoFn<?, ?> fn;
        if (application.getTransform() instanceof ParDo.SingleOutput) {
          fn = ((ParDo.SingleOutput) application.getTransform()).getFn();
        } else if (application.getTransform() instanceof ParDo.MultiOutput) {
          fn = ((ParDo.MultiOutput) application.getTransform()).getFn();
        } else {
          return false;
        }
        return fnType.equals(fn.getClass());
      }

      @Override
      public String toString() {
        return MoreObjects.toStringHelper("ParDoWithFnTypeMatcher")
            .add("fnType", fnType)
            .toString();
      }
    };
  }

  public static PTransformMatcher createViewWithViewFn(final Class<? extends ViewFn> viewFnType) {
    return new PTransformMatcher() {
      @Override
      public boolean matches(AppliedPTransform<?, ?, ?> application) {
        if (!(application.getTransform() instanceof CreatePCollectionView)) {
          return false;
        }
        CreatePCollectionView<?, ?> createView =
            (CreatePCollectionView<?, ?>) application.getTransform();
        ViewFn<Iterable<WindowedValue<?>>, ?> viewFn = createView.getView().getViewFn();
        return viewFn.getClass().equals(viewFnType);
      }
    };
  }

  /**
   * A {@link PTransformMatcher} which matches a {@link Flatten.PCollections} which
   * consumes no input {@link PCollection PCollections}.
   */
  public static PTransformMatcher emptyFlatten() {
    return new PTransformMatcher() {
      @Override
      public boolean matches(AppliedPTransform<?, ?, ?> application) {
        return (application.getTransform() instanceof Flatten.PCollections)
            && application.getInputs().isEmpty();
      }

      @Override
      public String toString() {
        return MoreObjects.toStringHelper("EmptyFlattenMatcher").toString();
      }
    };
  }

  /**
   * A {@link PTransformMatcher} which matches a {@link Flatten.PCollections} which
   * consumes a single input {@link PCollection} multiple times.
   */
  public static PTransformMatcher flattenWithDuplicateInputs() {
    return new PTransformMatcher() {
      @Override
      public boolean matches(AppliedPTransform<?, ?, ?> application) {
        if (application.getTransform() instanceof Flatten.PCollections) {
          Set<PValue> observed = new HashSet<>();
          for (PValue pvalue : application.getInputs().values()) {
            boolean firstInstance = observed.add(pvalue);
            if (!firstInstance) {
              return true;
            }
          }
        }
        return false;
      }

      @Override
      public String toString() {
        return MoreObjects.toStringHelper("FlattenWithDuplicateInputsMatcher").toString();
      }
    };
  }

  public static PTransformMatcher writeWithRunnerDeterminedSharding() {
    return new PTransformMatcher() {
      @Override
      public boolean matches(AppliedPTransform<?, ?, ?> application) {
        if (application.getTransform() instanceof WriteFiles) {
          WriteFiles write = (WriteFiles) application.getTransform();
          return write.getSharding() == null && write.getNumShards() == null;
        }
        return false;
      }
    };
  }
}