/*
 *  Copyright (c) 2012 Malhar, Inc.
 *  All Rights Reserved.
 */
package com.malhartech.stream;

import java.util.HashSet;
import java.util.Set;

import com.malhartech.api.Sink;
import com.malhartech.api.StreamCodec;
import com.malhartech.engine.Tuple;

/**
 *
 * @author Chetan Narsude <chetan@malhar-inc.com>
 */
public class PartitionAwareSink<T> implements Sink<T>
{
  private final StreamCodec<T> serde;
  private final Set<Integer> partitions;
  private volatile Sink<T> output;

  /**
   *
   * @param serde
   * @param partitions
   * @param output
   */
  public PartitionAwareSink(StreamCodec<T> serde, Set<Integer> partitions, Sink<T> output)
  {
    this.serde = serde;
    this.partitions = partitions;
    this.output = output;
  }

  /**
   *
   * @param payload
   */
  @Override
  public void process(T payload)
  {
    if (payload instanceof Tuple) {
      output.process(payload);
    }
    else if (partitions.contains(serde.getPartition(payload))) {
      output.process(payload);
    }
  }
}

@@ -97,57 +97,8 @@
 m;%0A%0A
-import java.util.HashSet;%0Aimport java.util.Set;%0A%0A
 impo
@@ -204,15 +204,51 @@
 le;%0A
-%0A/**%0A *
+import java.util.Set;%0A%0A/**%0A *%0A * @param %3CT%3E
 %0A * 
@@ -433,16 +433,42 @@
 itions;%0A
+  private final int mask;%0A
   privat
@@ -547,16 +547,33 @@
 titions%0A
+   * @param mask%0A
    * @pa
@@ -663,16 +663,26 @@
 titions,
+ int mask,
  Sink%3CT%3E
@@ -778,16 +778,38 @@
 output;%0A
+    this.mask = mask;%0A
   %7D%0A%0A  /
@@ -1024,16 +1024,23 @@
 payload)
+ & mask
 )) %7B%0A   
@@ -1077,10 +1077,11 @@
   %7D%0A  %7D%0A
+%0A
 %7D%0A

Make PartitionAwareSink compatible with bitwise partition scheme.