com.google.cloud.dataflow.sdk.util.MergingActiveWindowSet 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.MapCoder;
import com.google.cloud.dataflow.sdk.coders.SetCoder;
import com.google.cloud.dataflow.sdk.transforms.windowing.BoundedWindow;
import com.google.cloud.dataflow.sdk.transforms.windowing.WindowFn;
import com.google.cloud.dataflow.sdk.util.state.StateInternals;
import com.google.cloud.dataflow.sdk.util.state.StateNamespaces;
import com.google.cloud.dataflow.sdk.util.state.StateTag;
import com.google.cloud.dataflow.sdk.util.state.StateTags;
import com.google.cloud.dataflow.sdk.util.state.ValueState;

import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * Implementation of {@link ActiveWindowSet} used with {@link WindowFn WindowFns} that support
 * merging.
 *
 * @param  the types of windows being managed
 */
public class MergingActiveWindowSet
    implements ActiveWindowSet {

  private final WindowFn windowFn;

  /**
   * A map of live windows to windows that were merged into them.
   *
   *  The keys of the map correspond to the set of (merged) windows and the values
   * are the no-longer-present windows that were merged into the keys.  A given
   * window can appear in both the key and value of a single entry, but other at
   * most once across all keys and values.
   */
  private final Map> mergeTree;

  /**
   * Used to determine if writing the mergeTree (which is relatively stable)
   * is necessary.
   */
  private final Map> originalMergeTree;

  private final ValueState>> mergeTreeValue;

  public MergingActiveWindowSet(WindowFn windowFn, StateInternals state) {
    this.windowFn = windowFn;

    StateTag>>> mergeTreeAddr = StateTags.makeSystemTagInternal(
        StateTags.value("tree",
            MapCoder.of(windowFn.windowCoder(), SetCoder.of(windowFn.windowCoder()))));
    this.mergeTreeValue = state.state(StateNamespaces.global(), mergeTreeAddr);
    this.mergeTree = emptyIfNull(mergeTreeValue.get().read());

    originalMergeTree = deepCopy(mergeTree);
  }

  @Override
  public void persist() {
    if (!mergeTree.equals(originalMergeTree)) {
      mergeTreeValue.set(mergeTree);
    }
  }

  @Override
  public boolean add(W window) {
    if (mergeTree.containsKey(window)) {
      return false;
    }

    mergeTree.put(window, new HashSet());
    return true;
  }

  @Override
  public void remove(W window) {
    mergeTree.remove(window);
  }

  private class MergeContextImpl extends WindowFn.MergeContext {

    private MergeCallback mergeCallback;

    public MergeContextImpl(MergeCallback mergeCallback) {
      windowFn.super();
      this.mergeCallback = mergeCallback;
    }

    @Override
    public Collection windows() {
      return mergeTree.keySet();
    }

    @Override
    public void merge(Collection toBeMerged, W mergeResult) throws Exception {
      boolean isResultNew = !mergeTree.containsKey(mergeResult);
      recordMerge(toBeMerged, mergeResult);
      mergeCallback.onMerge(toBeMerged, mergeResult, isResultNew);
    }
  }

  @Override
  public boolean mergeIfAppropriate(W window, MergeCallback mergeCallback) throws Exception {
    windowFn.mergeWindows(new MergeContextImpl(mergeCallback));
    return window == null || mergeTree.containsKey(window);
  }

  @Override
  public Iterable sourceWindows(W window) {
    Set curWindows = new HashSet<>();
    curWindows.add(window);

    Set sourceWindows = mergeTree.get(window);
    if (sourceWindows != null) {
      curWindows.addAll(sourceWindows);
    }
    return curWindows;
  }

  private void recordMerge(Collection otherWindows, W newWindow) throws Exception {
    Set subWindows = mergeTree.get(newWindow);
    if (subWindows == null) {
      subWindows = new HashSet<>();
    }

    for (W other : otherWindows) {
      if (!mergeTree.containsKey(other)) {
        throw new IllegalArgumentException("Tried to merge a non-existent window: " + other);
      }
      subWindows.addAll(mergeTree.get(other));
      subWindows.add(other);
      mergeTree.remove(other);
    }
    mergeTree.put(newWindow, subWindows);
  }


  private static  Map> emptyIfNull(Map> input) {
    if (input == null) {
      return new HashMap<>();
    } else {
      for (Map.Entry> entry : input.entrySet()) {
        if (entry.getValue() == null) {
          entry.setValue(new HashSet());
        }
      }
      return input;
    }
  }

  private Map> deepCopy(Map> mergeTree) {
    Map> newMergeTree = new HashMap<>();
    for (Map.Entry> entry : mergeTree.entrySet()) {
      newMergeTree.put(entry.getKey(), new HashSet(entry.getValue()));
    }
    return newMergeTree;
  }
}