com.datatorrent.api.DefaultPartition Java Examples

@Test
public void partition1Test()
{
  DummyOperator dummyOperator = new DummyOperator(5);
  StatelessPartitioner<DummyOperator> statelessPartitioner = new StatelessPartitioner<DummyOperator>();

  Collection<Partition<DummyOperator>> partitions = Lists.newArrayList();
  DefaultPartition<DummyOperator> defaultPartition = new DefaultPartition<DummyOperator>(dummyOperator);
  partitions.add(defaultPartition);

  Collection<Partition<DummyOperator>> newPartitions = statelessPartitioner.definePartitions(partitions, new PartitioningContextImpl(null, 0));
  Assert.assertEquals("Incorrect number of partitions", 1, newPartitions.size());

  for (Partition<DummyOperator> partition : newPartitions) {
    Assert.assertEquals("Incorrect cloned value", 5, partition.getPartitionedInstance().getValue());
  }
}
 

protected Partitioner.Partition<AbstractKafkaInputOperator<K>> createPartition(Set<KafkaPartition> pIds,
    Map<KafkaPartition, Long> initOffsets)
{

  Partitioner.Partition<AbstractKafkaInputOperator<K>> p = new DefaultPartition<>(KryoCloneUtils.cloneObject(this));
  if (p.getPartitionedInstance().getConsumer() instanceof SimpleKafkaConsumer) {
    p.getPartitionedInstance().getConsumer().resetPartitionsAndOffset(pIds, initOffsets);
    if (initOffsets != null) {
      //Don't send all offsets to all partitions
      //p.getPartitionedInstance().offsetStats.putAll(initOffsets);
      p.getPartitionedInstance().offsetStats.putAll(p.getPartitionedInstance().getConsumer().getCurrentOffsets());
    }
  }

  PartitionInfo pif = new PartitionInfo();
  pif.kpids = pIds;
  currentPartitionInfo.add(pif);
  return p;
}
 

@Override
public Collection<Partition<AbstractCouchBaseInputOperator<T>>> definePartitions(Collection<Partition<AbstractCouchBaseInputOperator<T>>> partitions, PartitioningContext incrementalCapacity)
{
  conf = store.getConf();
  int numPartitions = conf.getServers().size();
  List<String> list = conf.getServers();
  Collection<Partition<AbstractCouchBaseInputOperator<T>>> newPartitions = Lists.newArrayListWithExpectedSize(numPartitions);
  KryoCloneUtils<AbstractCouchBaseInputOperator<T>> cloneUtils = KryoCloneUtils.createCloneUtils(this);
  for (int i = 0; i < numPartitions; i++) {
    AbstractCouchBaseInputOperator<T> oper = cloneUtils.getClone();
    oper.setServerIndex(i);
    oper.setServerURIString(list.get(i));
    logger.debug("oper {} urlstring is {}", i, oper.getServerURIString());
    newPartitions.add(new DefaultPartition<>(oper));
  }

  return newPartitions;

}
 

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;
}
 

/**
 * <b>Note:</b> This partitioner does not support parallel partitioning.<br/><br/>
 * {@inheritDoc}
 */
@Override
public Collection<Partition<BenchmarkPartitionableKafkaOutputOperator>> definePartitions(
    Collection<Partition<BenchmarkPartitionableKafkaOutputOperator>> partitions, PartitioningContext context)
{

  ArrayList<Partition<BenchmarkPartitionableKafkaOutputOperator>> newPartitions =
      new ArrayList<Partitioner.Partition<BenchmarkPartitionableKafkaOutputOperator>>(partitionCount);

  for (int i = 0; i < partitionCount; i++) {
    BenchmarkPartitionableKafkaOutputOperator bpkoo = new BenchmarkPartitionableKafkaOutputOperator();
    bpkoo.setPartitionCount(partitionCount);
    bpkoo.setTopic(topic);
    bpkoo.setBrokerList(brokerList);
    bpkoo.setStickyKey(i);
    Partition<BenchmarkPartitionableKafkaOutputOperator> p =
        new DefaultPartition<BenchmarkPartitionableKafkaOutputOperator>(bpkoo);
    newPartitions.add(p);
  }
  return newPartitions;
}
 

@Test
public void testParallelPartitionScaleDown()
{
  DummyOperator dummyOperator = new DummyOperator(5);
  StatelessPartitioner<DummyOperator> statelessPartitioner = new StatelessPartitioner<>();

  Collection<Partition<DummyOperator>> partitions = Lists.newArrayList();

  for (int i = 5; i-- > 0; ) {
    partitions.add(new DefaultPartition<>(dummyOperator));
  }

  Collection<Partition<DummyOperator>> newPartitions = statelessPartitioner.definePartitions(partitions,
      new PartitioningContextImpl(null, 1));
  Assert.assertEquals("after partition", 1, newPartitions.size());
}
 

@Test
public void partition5Test()
{
  DummyOperator dummyOperator = new DummyOperator(5);
  StatelessPartitioner<DummyOperator> statelessPartitioner = new StatelessPartitioner<>(5);

  Collection<Partition<DummyOperator>> partitions = Lists.newArrayList();
  DefaultPartition<DummyOperator> defaultPartition = new DefaultPartition<>(dummyOperator);
  partitions.add(defaultPartition);

  Collection<Partition<DummyOperator>> newPartitions = statelessPartitioner.definePartitions(partitions, new PartitioningContextImpl(null, 0));
  Assert.assertEquals("Incorrect number of partitions", 5, newPartitions.size());

  for (Partition<DummyOperator> partition : newPartitions) {
    Assert.assertEquals("Incorrect cloned value", 5, partition.getPartitionedInstance().getValue());
  }
}
 

@Test
public void partition1Test()
{
  DummyOperator dummyOperator = new DummyOperator(5);
  StatelessPartitioner<DummyOperator> statelessPartitioner = new StatelessPartitioner<>();

  Collection<Partition<DummyOperator>> partitions = Lists.newArrayList();
  DefaultPartition<DummyOperator> defaultPartition = new DefaultPartition<>(dummyOperator);
  partitions.add(defaultPartition);

  Collection<Partition<DummyOperator>> newPartitions = statelessPartitioner.definePartitions(partitions, new PartitioningContextImpl(null, 0));
  Assert.assertEquals("Incorrect number of partitions", 1, newPartitions.size());

  for (Partition<DummyOperator> partition : newPartitions) {
    Assert.assertEquals("Incorrect cloned value", 5, partition.getPartitionedInstance().getValue());
  }
}
 

/**
 * Adjust the partitions of an input operator (operator with no connected input stream).
 *
 * @param <T> The operator type
 * @param partitions
 * @return The new operators.
 */
public static <T extends Operator> Collection<Partition<T>> repartitionInputOperator(Collection<Partition<T>> partitions)
{
  List<Partition<T>> newPartitions = new ArrayList<>();
  List<Partition<T>> lowLoadPartitions = new ArrayList<>();
  for (Partition<T> p: partitions) {
    int load = p.getLoad();
    if (load < 0) {
      if (!lowLoadPartitions.isEmpty()) {
        newPartitions.add(lowLoadPartitions.remove(0));
      } else {
        lowLoadPartitions.add(p);
      }
    } else if (load > 0) {
      newPartitions.add(new DefaultPartition<>(p.getPartitionedInstance()));
      newPartitions.add(new DefaultPartition<>(p.getPartitionedInstance()));
    } else {
      newPartitions.add(p);
    }
  }
  newPartitions.addAll(lowLoadPartitions);
  return newPartitions;
}
 

@Test
public void partition5Test()
{
  DummyOperator dummyOperator = new DummyOperator(5);
  StatelessPartitioner<DummyOperator> statelessPartitioner = new StatelessPartitioner<DummyOperator>(5);

  Collection<Partition<DummyOperator>> partitions = Lists.newArrayList();
  DefaultPartition<DummyOperator> defaultPartition = new DefaultPartition<DummyOperator>(dummyOperator);
  partitions.add(defaultPartition);

  Collection<Partition<DummyOperator>> newPartitions = statelessPartitioner.definePartitions(partitions, new PartitioningContextImpl(null, 0));
  Assert.assertEquals("Incorrect number of partitions", 5, newPartitions.size());

  for (Partition<DummyOperator> partition : newPartitions) {
    Assert.assertEquals("Incorrect cloned value", 5, partition.getPartitionedInstance().getValue());
  }
}
 

@Override
public Collection<Partition<PartitioningTestOperator>> definePartitions(Collection<Partition<PartitioningTestOperator>> partitions, PartitioningContext context)
{
  final int newPartitionCount = DefaultPartition.getRequiredPartitionCount(context, this.partitionCount);

  if (!fixedCapacity) {
    partitionKeys = new Integer[newPartitionCount];
    for (int i = 0; i < partitionKeys.length; i++) {
      partitionKeys[i] = i;
    }
  }

  List<Partition<PartitioningTestOperator>> newPartitions = new ArrayList<>(this.partitionKeys.length);
  for (Integer partitionKey : partitionKeys) {
    PartitioningTestOperator temp = new PartitioningTestOperator();
    temp.setPartitionCount(newPartitionCount);
    Partition<PartitioningTestOperator> p = new DefaultPartition<>(temp);
    PartitionKeys lpks = new PartitionKeys(2, Sets.newHashSet(partitionKey));
    p.getPartitionKeys().put(this.inport1, lpks);
    p.getPartitionKeys().put(this.inportWithCodec, lpks);
    p.getPartitionedInstance().pks = p.getPartitionKeys().values().toString();
    newPartitions.add(p);
  }

  return newPartitions;
}
 

@Override
    public Collection<com.datatorrent.api.Partitioner.Partition<T>> definePartitions(
        Collection<com.datatorrent.api.Partitioner.Partition<T>> partitions,
        com.datatorrent.api.Partitioner.PartitioningContext context) {
      //Get a partition
      DefaultPartition<T> partition = (DefaultPartition<T>) partitions.iterator().next();
      Collection<Partition<T>> newPartitions;

      // first call to define partitions
      newPartitions = Lists.newArrayList();

      for (int partitionCounter = 0; partitionCounter < newPartitionCount; partitionCounter++) {
        newPartitions.add(new DefaultPartition<T>(partition.getPartitionedInstance()));
      }
      List<InputPort<?>> inputPortList = context.getInputPorts();
//      if (inputPortList != null && !inputPortList.isEmpty()) {
//        DefaultPartition.assignPartitionKeys(newPartitions, inputPortList.iterator().next());
//      }

      return newPartitions;
    }
 

@Override
public Collection<Partition<PartitionableInputOperator>> definePartitions(Collection<Partition<PartitionableInputOperator>> partitions, PartitioningContext context)
{
  List<Partition<PartitionableInputOperator>> newPartitions = new ArrayList<>(3);
  Iterator<? extends Partition<PartitionableInputOperator>> iterator = partitions.iterator();
  Partition<PartitionableInputOperator> templatePartition;
  for (int i = 0; i < 3; i++) {
    PartitionableInputOperator op = new PartitionableInputOperator();
    if (iterator.hasNext()) {
      templatePartition = iterator.next();
      op.partitionProperty = templatePartition.getPartitionedInstance().partitionProperty;
    }
    op.partitionProperty += "_" + i;
    newPartitions.add(new DefaultPartition<>(op));
  }
  return newPartitions;
}
 

@Override
public Collection definePartitions(Collection partitions, PartitioningContext context)
{
  Collection<Partition> newPartitions = new ArrayList<>();

  int partitionMask = 0x03;

  // No partitioning done so far..
  // Single partition again, but with only even numbers ok?
  // First partition
  PassThruOperatorWithCodec newInstance = new PassThruOperatorWithCodec();
  Partition partition = new DefaultPartition<>(newInstance);
  PartitionKeys value = new PartitionKeys(partitionMask, Sets.newHashSet(0));
  partition.getPartitionKeys().put(input, value);
  newPartitions.add(partition);

  // Second partition
  newInstance = new PassThruOperatorWithCodec();
  partition = new DefaultPartition<>(newInstance);
  value = new PartitionKeys(partitionMask, Sets.newHashSet(1));
  partition.getPartitionKeys().put(input, value);

  newPartitions.add(partition);

  return newPartitions;
}
 

@Test
public void testParallelPartitionScaleDown()
{
  DummyOperator dummyOperator = new DummyOperator(5);
  StatelessPartitioner<DummyOperator> statelessPartitioner = new StatelessPartitioner<DummyOperator>();

  TestUtils.MockBatchedOperatorStats mockStats = new TestUtils.MockBatchedOperatorStats(2);
  mockStats.operatorStats = Lists.newArrayList();

  Collection<Partition<DummyOperator>> partitions = Lists.newArrayList();

  for (int i = 5; i-- > 0; ) {
    DefaultPartition<DummyOperator> defaultPartition = new DefaultPartition<DummyOperator>(dummyOperator);
    TestUtils.MockPartition<DummyOperator> mockPartition = new TestUtils.MockPartition<DummyOperator>(defaultPartition, mockStats);

    partitions.add(mockPartition);
  }

  Collection<Partition<DummyOperator>> newPartitions = statelessPartitioner.definePartitions(partitions,
      new PartitioningContextImpl(null, 1));
  Assert.assertEquals("after partition", 1, newPartitions.size());
}
 

@Override
public Collection<Partition<T>> definePartitions(Collection<Partition<T>> partitions, PartitioningContext context)
{
  Collection<Partition<T>> newPartitions = Lists.newArrayList(partitions);
  int numTotal = partitions.size();
  Partition<T> first = partitions.iterator().next();
  if (first.getStats() == null) {
    // Initial partition
    numTotal = initalPartitionCount;
  } else {
    for (Partition<T> p : partitions) {
      // Assumption load is non-negative
      numTotal += p.getLoad();
    }
  }
  T paritionable = first.getPartitionedInstance();
  for (int i = partitions.size(); i < numTotal; ++i) {
    newPartitions.add(new DefaultPartition<>(paritionable));
  }
  return newPartitions;
}
 

@Override
public Collection<Partition<T>> definePartitions(Collection<Partition<T>> partitions, PartitioningContext context)
{
  if (!extraPartitions.isEmpty()) {
    partitions.addAll(extraPartitions);
    return partitions;
  }

  Collection<Partition<T>> newPartitions = super.definePartitions(partitions, context);
  if (context.getParallelPartitionCount() > 0 && newPartitions.size() < context.getParallelPartitionCount()) {
    // parallel partitioned, fill to requested count
    for (int i = newPartitions.size(); i < context.getParallelPartitionCount(); i++) {
      newPartitions.add(new DefaultPartition<>(partitions.iterator().next().getPartitionedInstance()));
    }
  }
  return newPartitions;
}
 

@Override
public Collection definePartitions(Collection partitions, PartitioningContext context)
{
  Collection<Partition> newPartitions = new ArrayList<>();

  int partitionMask = 0x03;

  // No partitioning done so far..
  // Single partition with mask 0x03 and set {0}
  // First partition
  PartitionedTestPersistanceOperator newInstance = new PartitionedTestPersistanceOperator();
  Partition partition = new DefaultPartition<>(newInstance);
  PartitionKeys value = new PartitionKeys(partitionMask, Sets.newHashSet(0));
  partition.getPartitionKeys().put(inport, value);
  newPartitions.add(partition);

  return newPartitions;
}
 

@Test
public void testParallelPartitionScaleUP()
{
  DummyOperator dummyOperator = new DummyOperator(5);
  StatelessPartitioner<DummyOperator> statelessPartitioner = new StatelessPartitioner<DummyOperator>();

  TestUtils.MockBatchedOperatorStats mockStats = new TestUtils.MockBatchedOperatorStats(2);
  mockStats.operatorStats = Lists.newArrayList();

  Collection<Partition<DummyOperator>> partitions = Lists.newArrayList();

  DefaultPartition<DummyOperator> defaultPartition = new DefaultPartition<DummyOperator>(dummyOperator);
  TestUtils.MockPartition<DummyOperator> mockPartition = new TestUtils.MockPartition<DummyOperator>(defaultPartition, mockStats);

  partitions.add(mockPartition);

  Collection<Partition<DummyOperator>> newPartitions = statelessPartitioner.definePartitions(partitions,
      new PartitioningContextImpl(null, 5));
  Assert.assertEquals("after partition", 5, newPartitions.size());
}
 

@Test
public void testBandwidthOnReducedPartitions()
{
  when(partitionContextMock.getParallelPartitionCount()).thenReturn(2);
  when(bandwidthManagerMock.getBandwidth()).thenReturn(2L);
  Collection<Partition<BandwidthLimitingOperator>> partitions = Lists.newArrayList();

  for (int i = 5; i-- > 0;) {
    partitions.add(new DefaultPartition<BandwidthLimitingOperator>(operatorMock));
  }

  underTest.definePartitions(partitions, partitionContextMock);
  verify(bandwidthManagerMock, times(2)).setBandwidth(5L);
}
 

@Test
public void testPartitioning() throws Exception
{
  LineByLineFileInputOperator oper = new LineByLineFileInputOperator();
  oper.getScanner().setFilePatternRegexp(".*partition([\\d]*)");
  oper.setDirectory(new File(testMeta.dir).getAbsolutePath());

  Path path = new Path(new File(testMeta.dir).getAbsolutePath());
  FileContext.getLocalFSFileContext().delete(path, true);
  for (int file = 0; file < 4; file++) {
    FileUtils.write(new File(testMeta.dir, "partition00" + file), "");
  }

  List<Partition<AbstractFileInputOperator<String>>> partitions = Lists.newArrayList();
  partitions.add(new DefaultPartition<AbstractFileInputOperator<String>>(oper));
  Collection<Partition<AbstractFileInputOperator<String>>> newPartitions = oper.definePartitions(partitions,
      new PartitioningContextImpl(null, 2));
  Assert.assertEquals(2, newPartitions.size());
  Assert.assertEquals(1, oper.getCurrentPartitions()); // partitioned() wasn't called

  for (Partition<AbstractFileInputOperator<String>> p : newPartitions) {
    Assert.assertNotSame(oper, p.getPartitionedInstance());
    Assert.assertNotSame(oper.getScanner(), p.getPartitionedInstance().getScanner());
    Set<String> consumed = Sets.newHashSet();
    LinkedHashSet<Path> files = p.getPartitionedInstance().getScanner().scan(FileSystem.getLocal(new Configuration(false)), path, consumed);
    Assert.assertEquals("partition " + files, 3, files.size());
  }
}
 

@Test
public void testBandwidthOnIncresedPartitions()
{
  when(partitionContextMock.getParallelPartitionCount()).thenReturn(5);
  Collection<Partition<BandwidthLimitingOperator>> partitions = Lists.newArrayList();
  DefaultPartition<BandwidthLimitingOperator> defaultPartition = new DefaultPartition<BandwidthLimitingOperator>(operatorMock);
  partitions.add(defaultPartition);

  underTest.definePartitions(partitions, partitionContextMock);
  verify(bandwidthManagerMock, times(5)).setBandwidth(2L);
}
 

@Test
public void testBandwidthOnPartitions()
{
  when(partitionContextMock.getParallelPartitionCount()).thenReturn(0); // no partitions
  Collection<Partition<BandwidthLimitingOperator>> partitions = Lists.newArrayList();
  DefaultPartition<BandwidthLimitingOperator> defaultPartition = new DefaultPartition<BandwidthLimitingOperator>(operatorMock);
  partitions.add(defaultPartition);

  underTest.definePartitions(partitions, partitionContextMock);
  verify(bandwidthManagerMock).setBandwidth(10L);
}
 

/**
 * Partition the operator in 2
 * create ten files with index of the file at the start, i.e 1_file, 2_file .. etc.
 * The scanner returns this index from getPartition method.
 * each partition should read 5 files as file index are from 0 to 9 (including 0 and 9).
 * @throws Exception
 */
@Test
public void testWithCustomScanner() throws Exception
{
  LineByLineFileInputOperator oper = new LineByLineFileInputOperator();
  oper.setScanner(new MyScanner());
  oper.getScanner().setFilePatternRegexp(".*partition_([\\d]*)");
  oper.setDirectory(new File(testMeta.dir).getAbsolutePath());

  Random rand = new Random();
  Path path = new Path(new File(testMeta.dir).getAbsolutePath());
  FileContext.getLocalFSFileContext().delete(path, true);
  for (int file = 0; file < 10; file++) {
    FileUtils.write(new File(testMeta.dir, file + "_partition_00" + rand.nextInt(100)), "");
  }

  List<Partition<AbstractFileInputOperator<String>>> partitions = Lists.newArrayList();
  partitions.add(new DefaultPartition<AbstractFileInputOperator<String>>(oper));
  Collection<Partition<AbstractFileInputOperator<String>>> newPartitions = oper.definePartitions(partitions,
      new PartitioningContextImpl(null, 2));
  Assert.assertEquals(2, newPartitions.size());
  Assert.assertEquals(1, oper.getCurrentPartitions()); // partitioned() wasn't called

  for (Partition<AbstractFileInputOperator<String>> p : newPartitions) {
    Assert.assertNotSame(oper, p.getPartitionedInstance());
    Assert.assertNotSame(oper.getScanner(), p.getPartitionedInstance().getScanner());
    Set<String> consumed = Sets.newHashSet();
    LinkedHashSet<Path> files = p.getPartitionedInstance().getScanner().scan(FileSystem.getLocal(new Configuration(false)), path, consumed);
    Assert.assertEquals("partition " + files, 6, files.size());
  }
}
 

/***
 * Clones the original operator and sets the partition pie assignments for this operator. Kryo is used for cloning
 * @param scanAssignmentsForThisPartition The partition pie that is assigned to the operator according to the
 *                                        configured partitioner
 * @return The partition that is used by the runtime to launch the new physical operator instance
 */
public Partitioner.Partition<AbstractKuduInputOperator> clonePartitionAndAssignScanMeta(
    List<KuduPartitionScanAssignmentMeta> scanAssignmentsForThisPartition)
{
  Partitioner.Partition<AbstractKuduInputOperator> clonedKuduInputOperator =
      new DefaultPartition<>(KryoCloneUtils.cloneObject(prototypeKuduInputOperator));
  clonedKuduInputOperator.getPartitionedInstance().setPartitionPieAssignment(scanAssignmentsForThisPartition);
  return clonedKuduInputOperator;
}
 

/**
 * <b>Note:</b> This partitioner does not support parallel partitioning.<br/><br/>
 * Partitions count will be the number of input routing keys.
 * Each partition receives tuples from its routing key.
 *
 * @param partitions
 * @param incrementalCapacity
 * @return
 */
@Override
public Collection<Partition<RabbitMQLogsInputOperator>> definePartitions(Collection<Partition<RabbitMQLogsInputOperator>> partitions, PartitioningContext context)
{
  if (routingKeys == null || routingKeys.length == 0) {
    return partitions;
  }

  ArrayList<Partition<RabbitMQLogsInputOperator>> newPartitions = new ArrayList<Partition<RabbitMQLogsInputOperator>>();
  for (String rKey : routingKeys) {
    try {
      RabbitMQLogsInputOperator oper = new RabbitMQLogsInputOperator();
      oper.host = host;
      oper.port = port;
      oper.exchange = exchange;
      oper.exchangeType = exchangeType;
      oper.registry = registry;
      oper.routingKeys = routingKeys;
      oper.routingKey = rKey;
      oper.queueName = rKey;

      Partition<RabbitMQLogsInputOperator> partition = new DefaultPartition<RabbitMQLogsInputOperator>(oper);
      newPartitions.add(partition);
    }
    catch (Throwable ex) {
      DTThrowable.rethrow(ex);
    }
  }
  return newPartitions;
}
 

@Test
public void testPartitioner() throws Exception
{
  DummyOperator dummyOperator = new DummyOperator(5);

  StatelessThroughputBasedPartitioner<DummyOperator> statelessLatencyBasedPartitioner = new StatelessThroughputBasedPartitioner<DummyOperator>();
  statelessLatencyBasedPartitioner.setMaximumEvents(10);
  statelessLatencyBasedPartitioner.setMinimumEvents(1);
  statelessLatencyBasedPartitioner.setCooldownMillis(10);

  TestStats mockStats = new TestStats(2);
  mockStats.operatorStats = Lists.newArrayList();
  mockStats.tuplesProcessedPSMA = 11;

  List<Operator.InputPort<?>> ports = Lists.newArrayList();
  ports.add(dummyOperator.input);

  DefaultPartition<DummyOperator> defaultPartition = new DefaultPartition<DummyOperator>(dummyOperator);
  Collection<Partitioner.Partition<DummyOperator>> partitions = Lists.newArrayList();
  partitions.add(defaultPartition);
  partitions = statelessLatencyBasedPartitioner.definePartitions(partitions, new StatelessPartitionerTest.PartitioningContextImpl(ports, 1));
  Assert.assertTrue(1 == partitions.size());
  defaultPartition = (DefaultPartition<DummyOperator>)partitions.iterator().next();
  Map<Integer, Partitioner.Partition<DummyOperator>> partitionerMap = Maps.newHashMap();
  partitionerMap.put(2, defaultPartition);
  statelessLatencyBasedPartitioner.partitioned(partitionerMap);
  StatsListener.Response response = statelessLatencyBasedPartitioner.processStats(mockStats);
  Assert.assertEquals("repartition is false", false, response.repartitionRequired);
  Thread.sleep(100);
  response = statelessLatencyBasedPartitioner.processStats(mockStats);
  Assert.assertEquals("repartition is true", true, response.repartitionRequired);

  TestUtils.MockPartition<DummyOperator> mockPartition = new TestUtils.MockPartition<DummyOperator>(defaultPartition, mockStats);
  partitions.clear();
  partitions.add(mockPartition);

  Collection<Partitioner.Partition<DummyOperator>> newPartitions = statelessLatencyBasedPartitioner.definePartitions(
      partitions, new StatelessPartitionerTest.PartitioningContextImpl(ports, 5));
  Assert.assertEquals("after partition", 2, newPartitions.size());
}
 

private Partition<AbstractKinesisInputOperator> createPartition(Set<String> shardIds, Map<String, String> initShardPos)
{
  Partition<AbstractKinesisInputOperator> p = new DefaultPartition<AbstractKinesisInputOperator>(KryoCloneUtils.cloneObject(this));
  p.getPartitionedInstance().getConsumer().setShardIds(shardIds);
  p.getPartitionedInstance().getConsumer().resetShardPositions(initShardPos);

  PartitionInfo pif = new PartitionInfo();
  pif.kpids = shardIds;
  currentPartitionInfo.add(pif);
  return p;
}
 

@Test
public void testPartitioning() throws Exception
{
  TestUtils.MockBatchedOperatorStats readerStats = new TestUtils.MockBatchedOperatorStats(2);
  readerStats.operatorStats = Lists.newArrayList();
  readerStats.operatorStats.add(new ReaderStats(10, 1, 100, 1));

  TestReader sliceReader = new TestReader();

  StatsListener.Response response = sliceReader.processStats(readerStats);

  Assert.assertTrue("partition needed", response.repartitionRequired);
  Assert.assertEquals("partition count changed", 8, sliceReader.getPartitionCount());

  List<Partitioner.Partition<AbstractBlockReader<Slice, BlockMetadata.FileBlockMetadata, FSDataInputStream>>>
      partitions = Lists
      .newArrayList();

  DefaultPartition<AbstractBlockReader<Slice, BlockMetadata.FileBlockMetadata, FSDataInputStream>> apartition = new
      DefaultPartition<AbstractBlockReader<Slice, BlockMetadata.FileBlockMetadata, FSDataInputStream>>(
      sliceReader);

  TestUtils.MockPartition<AbstractBlockReader<Slice, BlockMetadata.FileBlockMetadata, FSDataInputStream>>
      pseudoParttion = new TestUtils.MockPartition<>(
      apartition, readerStats);

  partitions.add(pseudoParttion);

  Collection<Partitioner.Partition<AbstractBlockReader<Slice, BlockMetadata.FileBlockMetadata, FSDataInputStream>>>
      newPartitions = sliceReader
      .definePartitions(partitions, null);
  Assert.assertEquals(8, newPartitions.size());
}
 

/**
 * Assigns the partitions according to certain key values and keeps track of the
 * keys that each partition will be processing so that in the case of a
 * rollback, each partition will only clear the data that it is responsible for.
 */
@Override
public Collection<com.datatorrent.api.Partitioner.Partition<UniqueValueCountAppender<V>>> definePartitions(Collection<com.datatorrent.api.Partitioner.Partition<UniqueValueCountAppender<V>>> partitions, PartitioningContext context)
{
  final int finalCapacity = DefaultPartition.getRequiredPartitionCount(context, this.partitionCount);
  UniqueValueCountAppender<V> anOldOperator = partitions.iterator().next().getPartitionedInstance();
  partitions.clear();

  Collection<Partition<UniqueValueCountAppender<V>>> newPartitions = Lists.newArrayListWithCapacity(finalCapacity);

  for (int i = 0; i < finalCapacity; i++) {
    try {
      @SuppressWarnings("unchecked")
      UniqueValueCountAppender<V> statefulUniqueCount = this.getClass().newInstance();
      DefaultPartition<UniqueValueCountAppender<V>> partition = new DefaultPartition<UniqueValueCountAppender<V>>(statefulUniqueCount);
      newPartitions.add(partition);
    } catch (Throwable cause) {
      DTThrowable.rethrow(cause);
    }
  }

  DefaultPartition.assignPartitionKeys(Collections.unmodifiableCollection(newPartitions), input);
  int lPartitionMask = newPartitions.iterator().next().getPartitionKeys().get(input).mask;

  for (Partition<UniqueValueCountAppender<V>> statefulUniqueCountPartition : newPartitions) {
    UniqueValueCountAppender<V> statefulUniqueCountInstance = statefulUniqueCountPartition.getPartitionedInstance();

    statefulUniqueCountInstance.partitionKeys = statefulUniqueCountPartition.getPartitionKeys().get(input).partitions;
    statefulUniqueCountInstance.partitionMask = lPartitionMask;
    statefulUniqueCountInstance.store = anOldOperator.store;
    statefulUniqueCountInstance.tableName = anOldOperator.tableName;
    statefulUniqueCountInstance.cacheManager = anOldOperator.cacheManager;
  }
  return newPartitions;
}