com.google.cloud.dataflow.sdk.util.ReduceFnContextFactory Maven / Gradle / Ivy

Go to download
/*
 * Copyright (C) 2015 Google Inc.
 *
 * Licensed 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 com.google.cloud.dataflow.sdk.util;

import com.google.cloud.dataflow.sdk.coders.Coder;
import com.google.cloud.dataflow.sdk.transforms.windowing.BoundedWindow;
import com.google.cloud.dataflow.sdk.transforms.windowing.PaneInfo;
import com.google.cloud.dataflow.sdk.util.ReduceFn.MergingStateContext;
import com.google.cloud.dataflow.sdk.util.ReduceFn.StateContext;
import com.google.cloud.dataflow.sdk.util.ReduceFn.Timers;
import com.google.cloud.dataflow.sdk.util.TimerInternals.TimerData;
import com.google.cloud.dataflow.sdk.util.state.MergeableState;
import com.google.cloud.dataflow.sdk.util.state.State;
import com.google.cloud.dataflow.sdk.util.state.StateContents;
import com.google.cloud.dataflow.sdk.util.state.StateInternals;
import com.google.cloud.dataflow.sdk.util.state.StateNamespace;
import com.google.cloud.dataflow.sdk.util.state.StateNamespaces;
import com.google.cloud.dataflow.sdk.util.state.StateNamespaces.WindowNamespace;
import com.google.cloud.dataflow.sdk.util.state.StateTag;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;

import org.joda.time.Instant;

import java.util.ArrayList;
import java.util.Collection;
import java.util.List;
import java.util.Map;

/**
 * Factory for creating instances of the various {@link ReduceFn} contexts.
 */
class ReduceFnContextFactory {

  public interface OnTriggerCallbacks {
    void output(OutputT toOutput);
  }

  private final K key;
  private final ReduceFn reduceFn;
  private final WindowingStrategy windowingStrategy;
  private StateInternals stateInternals;
  private ActiveWindowSet activeWindows;
  private TimerInternals timerInternals;

  ReduceFnContextFactory(
      K key, ReduceFn reduceFn, WindowingStrategy windowingStrategy,
      StateInternals stateInternals, ActiveWindowSet activeWindows,
      TimerInternals timerInternals) {
    this.key = key;
    this.reduceFn = reduceFn;
    this.windowingStrategy = windowingStrategy;
    this.stateInternals = stateInternals;
    this.activeWindows = activeWindows;
    this.timerInternals = timerInternals;
  }

  private StateContextImpl stateContext(W window) {
    return new StateContextImpl<>(
        activeWindows, windowingStrategy.getWindowFn().windowCoder(), stateInternals, window);
  }

  public ReduceFn.Context base(W window) {
    return new ContextImpl(stateContext(window));
  }

  public ReduceFn.ProcessValueContext forValue(
      W window, InputT value, Instant timestamp) {
    return new ProcessValueContextImpl(stateContext(window), value, timestamp);
  }

  public ReduceFn.OnTriggerContext forTrigger(
      W window, StateContents pane, OnTriggerCallbacks callbacks) {
    return new OnTriggerContextImpl(stateContext(window), pane, callbacks);
  }

  public ReduceFn.OnMergeContext forMerge(
      Collection mergingWindows, W resultWindow) {
    return new OnMergeContextImpl(
        new MergingStateContextImpl(stateContext(resultWindow), mergingWindows));
  }


  private class TimersImpl implements ReduceFn.Timers {

    private final StateNamespace namespace;

    public TimersImpl(StateNamespace namespace) {
      Preconditions.checkArgument(namespace instanceof WindowNamespace);
      this.namespace = namespace;
    }

    @Override
    public void setTimer(Instant timestamp, TimeDomain timeDomain) {
      timerInternals.setTimer(TimerData.of(namespace, timestamp, timeDomain));
    }

    @Override
    public void deleteTimer(Instant timestamp, TimeDomain timeDomain) {
      timerInternals.deleteTimer(TimerData.of(namespace, timestamp, timeDomain));
    }

    @Override
    public Instant currentProcessingTime() {
      return timerInternals.currentProcessingTime();
    }
  }

  static class StateContextImpl
      implements ReduceFn.StateContext {

    private final ActiveWindowSet activeWindows;
    private final W window;
    protected StateNamespace namespace;
    protected final Coder windowCoder;
    private final StateInternals stateInternals;

    public StateContextImpl(
        ActiveWindowSet activeWindows,
        Coder windowCoder,
        StateInternals stateInternals,
        W window) {
      this.activeWindows = activeWindows;
      this.windowCoder = windowCoder;
      this.stateInternals = stateInternals;
      this.window = window;
      this.namespace = namespaceFor(window);
    }

    protected StateNamespace namespaceFor(W window) {
      return StateNamespaces.window(windowCoder, window);
    }

    W window() {
      return window;
    }

    StateNamespace namespace() {
      return namespace;
    }

    @Override
    public  StorageT access(StateTag address) {
      return stateInternals.state(namespace, address);
    }

    @Override
    public > StorageT accessAcrossMergedWindows(
        StateTag address) {
      List sourceNamespaces = new ArrayList<>();
      for (W sourceWindow : activeWindows.sourceWindows(window)) {
        sourceNamespaces.add(namespaceFor(sourceWindow));
      }

      return stateInternals.mergedState(sourceNamespaces, namespace, address);
    }
  }

  static class MergingStateContextImpl
      implements ReduceFn.MergingStateContext {

    private final StateContextImpl delegate;
    private final Collection mergingWindows;

    public MergingStateContextImpl(StateContextImpl delegate, Collection mergingWindows) {
      this.delegate = delegate;
      this.mergingWindows = mergingWindows;
    }

    StateNamespace namespace() {
      return delegate.namespace;
    }

    W window() {
      return delegate.window();
    }

    Collection mergingWindows() {
      return mergingWindows;
    }

    @Override
    public  StorageT access(StateTag address) {
      return delegate.access(address);
    }

    @Override
    public > StorageT accessAcrossMergedWindows(
        StateTag address) {
      return delegate.accessAcrossMergedWindows(address);
    }

    @Override
    public > StateT accessAcrossMergingWindows(
        StateTag address) {
      List mergingNamespaces = new ArrayList<>();
      for (W mergingWindow : mergingWindows) {
        mergingNamespaces.add(delegate.namespaceFor(mergingWindow));
      }

      return delegate.stateInternals.mergedState(mergingNamespaces, delegate.namespace, address);
    }

    @Override
    public  Map accessInEachMergingWindow(
        StateTag address) {
      ImmutableMap.Builder builder = ImmutableMap.builder();
      for (W mergingWindow : mergingWindows) {
        StateT stateForWindow = delegate.stateInternals.state(
            delegate.namespaceFor(mergingWindow), address);
        builder.put(mergingWindow, stateForWindow);
      }
      return builder.build();
    }
  }

  private class ContextImpl extends ReduceFn.Context {

    private final StateContextImpl state;
    private final TimersImpl timers;

    private ContextImpl(StateContextImpl state) {
      reduceFn.super();
      this.state = state;
      this.timers = new TimersImpl(state.namespace);
    }

    @Override
    public K key() {
      return key;
    }

    @Override
    public W window() {
      return state.window;
    }

    @Override
    public WindowingStrategy windowingStrategy() {
      return windowingStrategy;
    }

    @Override
    public StateContext state() {
      return state;
    }

    @Override
    public Timers timers() {
      return timers;
    }
  }

  private class ProcessValueContextImpl
      extends ReduceFn.ProcessValueContext {

    private final InputT value;
    private final Instant timestamp;
    private final StateContextImpl state;
    private final TimersImpl timers;

    private ProcessValueContextImpl(StateContextImpl state, InputT value, Instant timestamp) {
      reduceFn.super();
      this.state = state;
      this.value = value;
      this.timestamp = timestamp;
      this.timers = new TimersImpl(state.namespace);
    }

    @Override
    public K key() {
      return key;
    }

    @Override
    public W window() {
      return state.window;
    }

    @Override
    public WindowingStrategy windowingStrategy() {
      return windowingStrategy;
    }

    @Override
    public StateContext state() {
      return state;
    }

    @Override
    public InputT value() {
      return value;
    }

    @Override
    public Instant timestamp() {
      return timestamp;
    }

    @Override
    public Timers timers() {
      return timers;
    }
  }

  private class OnTriggerContextImpl
      extends ReduceFn.OnTriggerContext {

    private final StateContextImpl state;
    private final StateContents pane;
    private final OnTriggerCallbacks callbacks;
    private final TimersImpl timers;

    private OnTriggerContextImpl(StateContextImpl state,
        StateContents pane, OnTriggerCallbacks callbacks) {
      reduceFn.super();
      this.state = state;
      this.pane = pane;
      this.callbacks = callbacks;
      this.timers = new TimersImpl(state.namespace);
    }

    @Override
    public K key() {
      return key;
    }

    @Override
    public W window() {
      return state.window;
    }

    @Override
    public WindowingStrategy windowingStrategy() {
      return windowingStrategy;
    }

    @Override
    public StateContext state() {
      return state;
    }

    @Override
    public PaneInfo paneInfo() {
      return pane.read();
    }

    @Override
    public void output(OutputT value) {
      callbacks.output(value);
    }

    @Override
    public Timers timers() {
      return timers;
    }
  }

  private class OnMergeContextImpl
      extends ReduceFn.OnMergeContext {

    private final MergingStateContextImpl state;
    private final TimersImpl timers;

    private OnMergeContextImpl(MergingStateContextImpl state) {
      reduceFn.super();
      this.state = state;
      this.timers = new TimersImpl(state.delegate.namespace);
    }

    @Override
    public K key() {
      return key;
    }

    @Override
    public WindowingStrategy windowingStrategy() {
      return windowingStrategy;
    }

    @Override
    public MergingStateContext state() {
      return state;
    }

    @Override
    public Collection mergingWindows() {
      return state.mergingWindows;
    }

    @Override
    public W window() {
      return state.delegate.window;
    }

    @Override
    public Timers timers() {
      return timers;
    }
  }
}