Java Code Examples for com.datatorrent.api.Partitioner#Partition

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

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

/**
 * Create a new partition with the partition Ids and initial offset positions
 *
 * @deprecated use {@link #createPartition(Set, Map)}
 */
@Deprecated
protected Partitioner.Partition<AbstractKafkaInputOperator<K>> createPartition(Set<KafkaPartition> pIds,
    Map<KafkaPartition, Long> initOffsets,
    @SuppressWarnings("UnusedParameters") Collection<WindowDataManager> newManagers)
{
  return createPartition(pIds, initOffsets);
}
 

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

public void initOperatorState() throws Exception
{
  Attribute.AttributeMap.DefaultAttributeMap attributeMapForInputOperator =
      new Attribute.AttributeMap.DefaultAttributeMap();
  attributeMapForInputOperator.put(DAG.APPLICATION_ID, APP_ID);
  operatorContext = mockOperatorContext(OPERATOR_ID_FOR_ONE_TO_ONE_PARTITIONER,
    attributeMapForInputOperator);

  Attribute.AttributeMap.DefaultAttributeMap portAttributesForInputOperator =
      new Attribute.AttributeMap.DefaultAttributeMap();
  portAttributesForInputOperator.put(Context.PortContext.TUPLE_CLASS, UnitTestTablePojo.class);
  testPortContext = new TestPortContext(portAttributesForInputOperator);

  unitTestStepwiseScanInputOperator = new UnitTestStepwiseScanInputOperator(
      getConnectionConfigForTable(), UnitTestTablePojo.class);
  unitTestStepwiseScanInputOperator.setNumberOfPartitions(numberOfKuduInputOperatorPartitions);
  unitTestStepwiseScanInputOperator.setPartitionScanStrategy(partitonScanStrategy);
  unitTestStepwiseScanInputOperator.setScanOrderStrategy(scanOrderStrategy);
  initCommonConfigsForAllTypesOfTests();
  partitions = unitTestStepwiseScanInputOperator.definePartitions(
    new ArrayList(), partitioningContext);
  Iterator<Partitioner.Partition<AbstractKuduInputOperator>> iteratorForMeta = partitions.iterator();
  UnitTestStepwiseScanInputOperator actualOperator =
      (UnitTestStepwiseScanInputOperator)iteratorForMeta.next()
      .getPartitionedInstance();
  // Adjust the bindings as if apex has completed the partioning.The runtime of the framework does this in reality
  unitTestStepwiseScanInputOperator = actualOperator;
  unitTestStepwiseScanInputOperator.setup(operatorContext);
  unitTestStepwiseScanInputOperator.activate(operatorContext);
  //rewire parent operator to enable proper unit testing method calls
  unitTestStepwiseScanInputOperator.getPartitioner().setPrototypeKuduInputOperator(unitTestStepwiseScanInputOperator);
  unitTestStepwiseScanInputOperator.getScanner().setParentOperator(unitTestStepwiseScanInputOperator);
}
 

@KuduClusterTestContext(kuduClusterBasedTest = true)
@Test
public void testDefinePartitions() throws Exception
{
  // case of setting num partitions to 7 but setting one to one mapping.
  numberOfKuduInputOperatorPartitions = 7;
  initOperatorState();
  AbstractKuduInputPartitioner partitioner = unitTestStepwiseScanInputOperator.getPartitioner();
  assertEquals(KuduClientTestCommons.TOTAL_KUDU_TABLETS_FOR_UNITTEST_TABLE,partitioner.definePartitions(
      new ArrayList<Partitioner.Partition<AbstractKuduInputOperator>>(),partitioningContext).size());
  // case of setting num partition to 7 but go with many to one mapping
  partitonScanStrategy = KuduPartitionScanStrategy.MANY_TABLETS_PER_OPERATOR;
  initOperatorState();
  partitioner = unitTestStepwiseScanInputOperator.getPartitioner();
  Collection<Partitioner.Partition<AbstractKuduInputOperator>> partitionsCalculated = partitioner.definePartitions(
      new ArrayList<Partitioner.Partition<AbstractKuduInputOperator>>(),partitioningContext);
  assertEquals(7,partitionsCalculated.size());
  int maxPartitionPerOperator = -1;
  int minPartitionPerOperator = Integer.MAX_VALUE;
  Iterator<Partitioner.Partition<AbstractKuduInputOperator>> iteratorForPartitionScan =
      partitionsCalculated.iterator();
  while ( iteratorForPartitionScan.hasNext()) {
    Partitioner.Partition<AbstractKuduInputOperator> anOperatorPartition = iteratorForPartitionScan.next();
    AbstractKuduInputOperator<UnitTestTablePojo,InputOperatorControlTuple> anOperatorHandle =
        anOperatorPartition.getPartitionedInstance();
    List<KuduPartitionScanAssignmentMeta> partitionsAssigned = anOperatorHandle.getPartitionPieAssignment();
    if (  partitionsAssigned.size() > maxPartitionPerOperator ) {
      maxPartitionPerOperator = partitionsAssigned.size();
    }
    if ( partitionsAssigned.size() < minPartitionPerOperator) {
      minPartitionPerOperator = partitionsAssigned.size();
    }
  }
  assertEquals(2,maxPartitionPerOperator); // 7 kudu operator instances dealing with 12 kudu tablets
  assertEquals(1,minPartitionPerOperator); // 7 kudu operator instances dealing with 12 kudu tablets
  // revert all the changes to defaults for the next test
  numberOfKuduInputOperatorPartitions = 5;
  partitonScanStrategy = KuduPartitionScanStrategy.ONE_TABLET_PER_OPERATOR;
}
 

@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());
}
 

@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());
}
 

protected Partitioner.Partition<AbstractKafkaInputOperator> createPartition(
    Set<AbstractKafkaPartitioner.PartitionMeta> partitionAssignment)
{
  Partitioner.Partition<AbstractKafkaInputOperator> p =
      new DefaultPartition<AbstractKafkaInputOperator>(KryoCloneUtils.cloneObject(prototypeOperator));
  p.getPartitionedInstance().assign(partitionAssignment);
  return p;
}
 

public Partitioner.Partition<BasePythonExecutionOperator> clonePartition()
{
  Partitioner.Partition<BasePythonExecutionOperator> clonedKuduInputOperator =
      new DefaultPartition<>(KryoCloneUtils.cloneObject(prototypePythonOperator));
  return clonedKuduInputOperator;
}
 

@Override
public void partitioned(Map<Integer, Partitioner.Partition<AbstractKafkaInputOperator<K>>> partitions)
{
  // update the last repartition time
  lastRepartitionTime = System.currentTimeMillis();
}
 

@KuduClusterTestContext(kuduClusterBasedTest = true)
@Test
public void testInit() throws Exception
{
  Attribute.AttributeMap.DefaultAttributeMap attributeMapForInputOperator =
      new Attribute.AttributeMap.DefaultAttributeMap();
  attributeMapForInputOperator.put(DAG.APPLICATION_ID, APP_ID);
  operatorContext = mockOperatorContext(OPERATOR_ID_FOR_ONE_TO_ONE_PARTITIONER,
    attributeMapForInputOperator);

  Attribute.AttributeMap.DefaultAttributeMap portAttributesForInputOperator =
      new Attribute.AttributeMap.DefaultAttributeMap();
  portAttributesForInputOperator.put(Context.PortContext.TUPLE_CLASS, UnitTestTablePojo.class);
  testPortContext = new TestPortContext(portAttributesForInputOperator);

  incrementalStepScanInputOperator = new IncrementalStepScanInputOperator(UnitTestTablePojo.class,
      "kuduincrementalstepscaninputoperator.properties");
  incrementalStepScanInputOperator.setNumberOfPartitions(numberOfKuduInputOperatorPartitions);
  incrementalStepScanInputOperator.setPartitionScanStrategy(partitonScanStrategy);
  incrementalStepScanInputOperator.setScanOrderStrategy(scanOrderStrategy);
  partitioningContext = new Partitioner.PartitioningContext()
  {
    @Override
    public int getParallelPartitionCount()
    {
      return numberOfKuduInputOperatorPartitions;
    }

    @Override
    public List<Operator.InputPort<?>> getInputPorts()
    {
      return null;
    }
  };
  partitions = incrementalStepScanInputOperator.definePartitions(
    new ArrayList(), partitioningContext);
  Iterator<Partitioner.Partition<AbstractKuduInputOperator>> iteratorForMeta = partitions.iterator();
  IncrementalStepScanInputOperator actualOperator =
      (IncrementalStepScanInputOperator)iteratorForMeta.next().getPartitionedInstance();
  // Adjust the bindings as if apex has completed the partioning.The runtime of the framework does this in reality
  incrementalStepScanInputOperator = actualOperator;
  incrementalStepScanInputOperator.setup(operatorContext);
  incrementalStepScanInputOperator.activate(operatorContext);
  //rewire parent operator to enable proper unit testing method calls
  incrementalStepScanInputOperator.getPartitioner().setPrototypeKuduInputOperator(incrementalStepScanInputOperator);
  incrementalStepScanInputOperator.getScanner().setParentOperator(incrementalStepScanInputOperator);
}
 

@Test
public void testCountersTransfer() 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();
  sliceReader.processStats(readerStats);

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

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

  for (Partitioner.Partition<AbstractBlockReader<Slice, BlockMetadata.FileBlockMetadata, FSDataInputStream>> partition
      : newPartitions) {

    partition.getPartitionedInstance().counters
        .setCounter(AbstractBlockReader.ReaderCounterKeys.BLOCKS, new MutableLong(1));

    newMocks.add(new TestUtils.MockPartition<>(
        (DefaultPartition<AbstractBlockReader<Slice, BlockMetadata.FileBlockMetadata, FSDataInputStream>>)partition,
        readerStats));
  }
  sliceReader.partitionCount = 1;
  newPartitions = sliceReader.definePartitions(newMocks, null);
  Assert.assertEquals(1, newPartitions.size());

  AbstractBlockReader<Slice, BlockMetadata.FileBlockMetadata, FSDataInputStream> last = newPartitions.iterator()
      .next().getPartitionedInstance();
  Assert.assertEquals("num blocks", 8,
      last.counters.getCounter(AbstractBlockReader.ReaderCounterKeys.BLOCKS).longValue());
}