@Test
  public void testCreateNormalParDoFn() throws Exception {
    String stringState = "some state";
    long longState = 42L;

    TestDoFn fn = new TestDoFn(stringState, longState);

    String serializedFn =
        StringUtils.byteArrayToJsonString(
            SerializableUtils.serializeToByteArray(
                new DoFnInfo(fn, WindowingStrategy.globalDefault())));

    CloudObject cloudUserFn = CloudObject.forClassName("DoFn");
    addString(cloudUserFn, "serialized_fn", serializedFn);

    String tag = "output";
    MultiOutputInfo multiOutputInfo = new MultiOutputInfo();
    multiOutputInfo.setTag(tag);
    List<MultiOutputInfo> multiOutputInfos = Arrays.asList(multiOutputInfo);

    PipelineOptions options = PipelineOptionsFactory.create();
    DataflowExecutionContext context = BatchModeExecutionContext.fromOptions(options);
    CounterSet counters = new CounterSet();
    StateSampler stateSampler = new StateSampler("test", counters.getAddCounterMutator());
    ParDoFn parDoFn =
        factory.create(
            options,
            cloudUserFn,
            "name",
            "transformName",
            null,
            multiOutputInfos,
            1,
            context,
            counters.getAddCounterMutator(),
            stateSampler);

    // Test that the factory created the correct class
    assertThat(parDoFn, instanceOf(NormalParDoFn.class));

    // Test that the DoFnInfo reflects the one passed in
    NormalParDoFn normalParDoFn = (NormalParDoFn) parDoFn;
    DoFnInfo doFnInfo = normalParDoFn.getDoFnInfo();
    DoFn actualDoFn = doFnInfo.getDoFn();
    assertThat(actualDoFn, instanceOf(TestDoFn.class));
    assertThat(doFnInfo.getWindowingStrategy().getWindowFn(), instanceOf(GlobalWindows.class));
    assertThat(
        doFnInfo.getWindowingStrategy().getTrigger().getSpec(), instanceOf(DefaultTrigger.class));

    // Test that the deserialized user DoFn is as expected
    TestDoFn actualTestDoFn = (TestDoFn) actualDoFn;
    assertEquals(stringState, actualTestDoFn.stringState);
    assertEquals(longState, actualTestDoFn.longState);
    assertEquals(context, normalParDoFn.getExecutionContext());
  }
    @Override
    public PCollectionTuple apply(PCollection<Integer> input) {
      PCollection<Integer> sum = input.apply(Sum.integersGlobally());

      // Fails here when attempting to construct a tuple with an unbound object.
      return PCollectionTuple.of(sumTag, sum)
          .and(
              doneTag,
              PCollection.<Void>createPrimitiveOutputInternal(
                  input.getPipeline(), WindowingStrategy.globalDefault(), input.isBounded()));
    }
Пример #3
0
 @Override
 public PCollection<T> apply(PInput input) {
   if (topic == null && subscription == null) {
     throw new IllegalStateException(
         "need to set either the topic or the subscription for "
             + "a PubsubIO.Read transform");
   }
   if (topic != null && subscription != null) {
     throw new IllegalStateException(
         "Can't set both the topic and the subscription for a " + "PubsubIO.Read transform");
   }
   return PCollection.<T>createPrimitiveOutputInternal(
           input.getPipeline(), WindowingStrategy.globalDefault(), IsBounded.UNBOUNDED)
       .setCoder(coder);
 }
  protected <T> WindowedValue<T> makeWindowedValue(
      T output, Instant timestamp, Collection<? extends BoundedWindow> windows, PaneInfo pane) {
    final Instant inputTimestamp = timestamp;
    final WindowFn windowFn = windowingStrategy.getWindowFn();

    if (timestamp == null) {
      timestamp = BoundedWindow.TIMESTAMP_MIN_VALUE;
    }

    if (windows == null) {
      try {
        windows =
            windowFn.assignWindows(
                windowFn.new AssignContext() {
                  @Override
                  public Object element() {
                    throw new UnsupportedOperationException(
                        "WindowFn attempted to access input element when none was available"); // TODO: 12/16/15 aljoscha's comment in slack
                  }

                  @Override
                  public Instant timestamp() {
                    if (inputTimestamp == null) {
                      throw new UnsupportedOperationException(
                          "WindowFn attempted to access input timestamp when none was available");
                    }
                    return inputTimestamp;
                  }

                  @Override
                  public Collection<? extends BoundedWindow> windows() {
                    throw new UnsupportedOperationException(
                        "WindowFn attempted to access input windows when none were available");
                  }
                });
      } catch (Exception e) {
        Throwables.propagateIfInstanceOf(e, UserCodeException.class);
        throw new UserCodeException(e);
      }
    }

    return WindowedValue.of(output, timestamp, windows, pane);
  }
 @Override
 public PCollection<String> apply(PCollection<String> input) {
   return PCollection.createPrimitiveOutputInternal(
       input.getPipeline(), WindowingStrategy.globalDefault(), input.isBounded());
 }
  @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);
  }
Пример #7
0
 @Override
 public PCollection<KV<K, V>> apply(PInput input) {
   return PCollection.createPrimitiveOutputInternal(
       input.getPipeline(), WindowingStrategy.globalDefault(), PCollection.IsBounded.BOUNDED);
 }