Ejemplo n.º 1
0
  public static CombineValuesFn create(
      PipelineOptions options,
      CloudObject cloudUserFn,
      String stepName,
      @Nullable List<SideInputInfo> sideInputInfos,
      @Nullable List<MultiOutputInfo> multiOutputInfos,
      Integer numOutputs,
      ExecutionContext executionContext,
      CounterSet.AddCounterMutator addCounterMutator,
      StateSampler stateSampler /* unused */)
      throws Exception {
    Object deserializedFn =
        SerializableUtils.deserializeFromByteArray(
            getBytes(cloudUserFn, PropertyNames.SERIALIZED_FN), "serialized user fn");
    Preconditions.checkArgument(deserializedFn instanceof Combine.KeyedCombineFn);
    final Combine.KeyedCombineFn combineFn = (Combine.KeyedCombineFn) deserializedFn;

    // Get the combine phase, default to ALL. (The implementation
    // doesn't have to split the combiner).
    final String phase = getString(cloudUserFn, PropertyNames.PHASE, CombinePhase.ALL);

    Preconditions.checkArgument(
        sideInputInfos == null || sideInputInfos.size() == 0,
        "unexpected side inputs for CombineValuesFn");
    Preconditions.checkArgument(numOutputs == 1, "expected exactly one output for CombineValuesFn");

    DoFnInfoFactory fnFactory =
        new DoFnInfoFactory() {
          @Override
          public DoFnInfo createDoFnInfo() {
            DoFn doFn = null;
            switch (phase) {
              case CombinePhase.ALL:
                doFn = new CombineValuesDoFn(combineFn);
                break;
              case CombinePhase.ADD:
                doFn = new AddInputsDoFn(combineFn);
                break;
              case CombinePhase.MERGE:
                doFn = new MergeAccumulatorsDoFn(combineFn);
                break;
              case CombinePhase.EXTRACT:
                doFn = new ExtractOutputDoFn(combineFn);
                break;
              default:
                throw new IllegalArgumentException(
                    "phase must be one of 'all', 'add', 'merge', 'extract'");
            }
            return new DoFnInfo(doFn, null);
          }
        };
    return new CombineValuesFn(options, fnFactory, stepName, executionContext, addCounterMutator);
  }
Ejemplo n.º 2
0
 @Override
 public void setConf(Configuration conf) {
   this.conf = conf;
   if (conf != null) {
     String headerString = conf.get(SAM_HEADER_PROPERTY_NAME);
     if (headerString == null) {
       throw new IllegalStateException("SAM file header has not been set");
     }
     byte[] headerBytes = Base64.getDecoder().decode(headerString);
     setSAMHeader(
         (SAMFileHeader)
             SerializableUtils.deserializeFromByteArray(headerBytes, "SAMFileHeader"));
   }
 }
  @SuppressWarnings({"rawtypes", "unchecked"})
  public static GroupAlsoByWindowsParDoFn create(
      PipelineOptions options,
      CloudObject cloudUserFn,
      String stepName,
      @Nullable List<SideInputInfo> sideInputInfos,
      @Nullable List<MultiOutputInfo> multiOutputInfos,
      Integer numOutputs,
      ExecutionContext executionContext,
      CounterSet.AddCounterMutator addCounterMutator,
      StateSampler sampler /* unused */)
      throws Exception {
    Object windowingStrategyObj;
    byte[] encodedWindowingStrategy = getBytes(cloudUserFn, PropertyNames.SERIALIZED_FN);
    if (encodedWindowingStrategy.length == 0) {
      windowingStrategyObj = WindowingStrategy.globalDefault();
    } else {
      windowingStrategyObj =
          SerializableUtils.deserializeFromByteArray(
              encodedWindowingStrategy, "serialized windowing strategy");
      if (!(windowingStrategyObj instanceof WindowingStrategy)) {
        throw new Exception(
            "unexpected kind of WindowingStrategy: " + windowingStrategyObj.getClass().getName());
      }
    }
    WindowingStrategy windowingStrategy = (WindowingStrategy) windowingStrategyObj;

    byte[] serializedCombineFn = getBytes(cloudUserFn, PropertyNames.COMBINE_FN, null);
    KeyedCombineFn combineFn;
    if (serializedCombineFn != null) {
      Object combineFnObj =
          SerializableUtils.deserializeFromByteArray(serializedCombineFn, "serialized combine fn");
      if (!(combineFnObj instanceof KeyedCombineFn)) {
        throw new Exception(
            "unexpected kind of KeyedCombineFn: " + combineFnObj.getClass().getName());
      }
      combineFn = (KeyedCombineFn) combineFnObj;
    } else {
      combineFn = null;
    }

    Map<String, Object> inputCoderObject = getObject(cloudUserFn, PropertyNames.INPUT_CODER);

    Coder inputCoder = Serializer.deserialize(inputCoderObject, Coder.class);
    if (!(inputCoder instanceof WindowedValueCoder)) {
      throw new Exception(
          "Expected WindowedValueCoder for inputCoder, got: " + inputCoder.getClass().getName());
    }
    Coder elemCoder = ((WindowedValueCoder) inputCoder).getValueCoder();
    if (!(elemCoder instanceof KvCoder)) {
      throw new Exception(
          "Expected KvCoder for inputCoder, got: " + elemCoder.getClass().getName());
    }
    KvCoder kvCoder = (KvCoder) elemCoder;

    boolean isStreamingPipeline = false;
    if (options instanceof StreamingOptions) {
      isStreamingPipeline = ((StreamingOptions) options).isStreaming();
    }

    KeyedCombineFn maybeMergingCombineFn = null;
    if (combineFn != null) {
      String phase = getString(cloudUserFn, PropertyNames.PHASE, CombinePhase.ALL);
      Preconditions.checkArgument(
          phase.equals(CombinePhase.ALL) || phase.equals(CombinePhase.MERGE),
          "Unexpected phase: " + phase);
      if (phase.equals(CombinePhase.MERGE)) {
        maybeMergingCombineFn = new MergingKeyedCombineFn(combineFn);
      } else {
        maybeMergingCombineFn = combineFn;
      }
    }

    DoFnInfoFactory fnFactory;
    final DoFn groupAlsoByWindowsDoFn =
        getGroupAlsoByWindowsDoFn(
            isStreamingPipeline, windowingStrategy, kvCoder, maybeMergingCombineFn);

    fnFactory =
        new DoFnInfoFactory() {
          @Override
          public DoFnInfo createDoFnInfo() {
            return new DoFnInfo(groupAlsoByWindowsDoFn, null);
          }
        };
    return new GroupAlsoByWindowsParDoFn(
        options, fnFactory, stepName, executionContext, addCounterMutator);
  }