examples/partition/src/main/java/org/apache/apex/examples/partition/TestPartition.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.apex.examples.partition;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashSet;
import java.util.Map;

import javax.validation.constraints.Max;
import javax.validation.constraints.Min;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.datatorrent.api.Context;
import com.datatorrent.api.DefaultInputPort;
import com.datatorrent.api.DefaultPartition;
import com.datatorrent.api.Partitioner;

import com.datatorrent.common.util.BaseOperator;

/**
 * Simple operator to test partitioning
 *
 * @since 3.7.0
 */
public class TestPartition extends BaseOperator implements Partitioner<TestPartition>
{
  private static final Logger LOG = LoggerFactory.getLogger(TestPartition.class);

  private transient int id;             // operator/partition id
  private transient long curWindowId;   // current window id
  private transient long cnt;           // per-window tuple count

  @Min(1)
  @Max(20)
  private int nPartitions = 3;

  public final transient DefaultInputPort<Integer> in = new DefaultInputPort<Integer>()
  {
    @Override
    public void process(Integer tuple)
    {
      LOG.debug("{}: tuple = {}, operator id = {}", cnt, tuple, id);
      ++cnt;
    }
  };

  //public final transient DefaultOutputPort<Integer> out = new DefaultOutputPort<Integer>();

  @Override
  public void setup(Context.OperatorContext context)
  {
    super.setup(context);

    long appWindowId = context.getValue(context.ACTIVATION_WINDOW_ID);
    id = context.getId();
    LOG.debug("Started setup, appWindowId = {}, operator id = {}", appWindowId, id);
  }

  @Override
  public void beginWindow(long windowId)
  {
    cnt = 0;
    curWindowId = windowId;
    LOG.debug("window id = {}, operator id = {}", curWindowId, id);
  }

  @Override
  public void endWindow()
  {
    LOG.debug("window id = {}, operator id = {}, cnt = {}", curWindowId, id, cnt);
  }

  @Override
  public void partitioned(Map<Integer, Partition<TestPartition>> partitions)
  {
    //Do nothing
  }

  @Override
  public Collection<Partition<TestPartition>> definePartitions(Collection<Partition<TestPartition>> partitions, PartitioningContext context)
  {
    int oldSize = partitions.size();
    LOG.debug("partitionCount: current = {} requested = {}", oldSize, nPartitions);

    // each partition i in 0...nPartitions receives tuples divisible by i but not by any other
    // j in that range; all other tuples ignored
    //
    if (3 != nPartitions) {
      return getPartitions(partitions, context);
    }

    // special case of 3 partitions: All odd numbers to partition 0; even numbers divisible
    // by 4 to partition 1, those divisible by 2 but not 4 to partition 2.

    // mask used to extract discriminant from tuple hashcode
    int mask = 0x03;

    Partition<TestPartition>[] newPartitions = new Partition[] {
      new DefaultPartition<TestPartition>(new TestPartition()),
      new DefaultPartition<TestPartition>(new TestPartition()),
      new DefaultPartition<TestPartition>(new TestPartition()) };

    HashSet<Integer>[] set
      = new HashSet[] {new HashSet<>(), new HashSet<>(), new HashSet<>()};
    set[0].add(0);
    set[1].add(1);
    set[2].add(2);

    PartitionKeys[] keys = {
      new PartitionKeys(mask, set[0]),
      new PartitionKeys(mask, set[1]),
      new PartitionKeys(mask, set[2]) };

    for (int i = 0; i < 3; ++i ) {
      Partition<TestPartition> partition = newPartitions[i];
      partition.getPartitionKeys().put(in, keys[i]);
    }

    return new ArrayList<Partition<TestPartition>>(Arrays.asList(newPartitions));
  }  // definePartitions

  private Collection<Partition<TestPartition>> getPartitions(Collection<Partition<TestPartition>> partitions, PartitioningContext context)
  {
    // create array of partitions to return
    Collection<Partition<TestPartition>> result = new ArrayList<Partition<TestPartition>>(nPartitions);

    int mask = getMask(nPartitions);
    for (int i = 0; i < nPartitions; ++i) {
      HashSet<Integer> set = new HashSet<>();
      set.add(i);
      PartitionKeys keys = new PartitionKeys(mask, set);
      Partition partition = new DefaultPartition<TestPartition>(new TestPartition());
      partition.getPartitionKeys().put(in, keys);
    }

    return result;
  }  // getPartitions

  // return mask with bits 0..N set where N is the highest set bit of argument
  private int getMask(final int n)
  {
    return -1 >>> Integer.numberOfLeadingZeros(n);
  }  // getMask

  // accessors
  public int getNPartitions()
  {
    return nPartitions;
  }

  public void setNPartitions(int v)
  {
    nPartitions = v;
  }
}