com.datatorrent.api.Partitioner Java Examples

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

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

@KuduClusterTestContext(kuduClusterBasedTest = false)
@Test
public void testAssignPartitions() throws Exception
{
  AbstractKuduInputOperator<UnitTestTablePojo,InputOperatorControlTuple> mockedInputOperator =
      PowerMockito.mock(AbstractKuduInputOperator.class);
  when(mockedInputOperator.getNumberOfPartitions()).thenReturn(9);
  PowerMockito.mockStatic(KryoCloneUtils.class);
  when(KryoCloneUtils.cloneObject(mockedInputOperator)).thenReturn(mockedInputOperator);
  KuduOneToOnePartitioner kuduOneToOnePartitioner = new KuduOneToOnePartitioner(mockedInputOperator);
  buildMockWiring(mockedInputOperator, KuduClientTestCommons.TOTAL_KUDU_TABLETS_FOR_UNITTEST_TABLE);
  kuduOneToOnePartitioner.setPrototypeKuduInputOperator(mockedInputOperator);
  Map<Integer,List<KuduPartitionScanAssignmentMeta>> assignedPartitions = kuduOneToOnePartitioner.assign(
      kuduOneToOnePartitioner.getListOfPartitionAssignments(
      new ArrayList<Partitioner.Partition<AbstractKuduInputOperator>>(),
      partitioningContext),partitioningContext);
  assertThat(assignedPartitions.size(), is(12));
  for (List<KuduPartitionScanAssignmentMeta> eachOperatorassignment: assignedPartitions.values()) {
    assertThat(eachOperatorassignment.size(), is(1));
  }
}
 

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

public static void initCommonConfigsForAllTypesOfTests() throws Exception
{
  KuduClientTestCommons.buildSchemaForUnitTestsTable();
  partitioningContext = new Partitioner.PartitioningContext()
  {
    @Override
    public int getParallelPartitionCount()
    {
      return numberOfKuduInputOperatorPartitions;
    }

    @Override
    public List<Operator.InputPort<?>> getInputPorts()
    {
      return null;
    }
  };
}
 

RepartitionContext(Partitioner<Operator> partitioner, PMapping currentMapping, int partitionCount)
{
  super(currentMapping, partitionCount);
  this.operators = currentMapping.partitions;
  this.currentPartitions = new ArrayList<>(operators.size());
  this.currentPartitionMap = Maps.newHashMapWithExpectedSize(operators.size());
  this.operatorIdToPartition = Maps.newHashMapWithExpectedSize(operators.size());

  // collect current partitions with committed operator state
  // those will be needed by the partitioner for split/merge
  for (PTOperator pOperator : operators) {
    Map<InputPort<?>, PartitionKeys> pks = pOperator.getPartitionKeys();
    if (pks == null) {
      throw new AssertionError("Null partition: " + pOperator);
    }

    // if partitions checkpoint at different windows, processing for new or modified
    // partitions will start from earliest checkpoint found (at least once semantics)
    if (minCheckpoint == null) {
      minCheckpoint = pOperator.recoveryCheckpoint;
    } else if (minCheckpoint.windowId > pOperator.recoveryCheckpoint.windowId) {
      minCheckpoint = pOperator.recoveryCheckpoint;
    }

    Operator partitionedOperator = loadOperator(pOperator);
    DefaultPartition<Operator> partition = new DefaultPartition<>(partitionedOperator, pks, pOperator.loadIndicator, pOperator.stats);
    currentPartitions.add(partition);
    currentPartitionMap.put(partition, pOperator);
    LOG.debug("partition load: {} {} {}", pOperator, partition.getPartitionKeys(), partition.getLoad());
    operatorIdToPartition.put(pOperator.getId(), partition);
  }

  newPartitions = partitioner.definePartitions(new ArrayList<Partition<Operator>>(currentPartitions), this);
}
 

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

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

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

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

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

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

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

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

@Override
public void partitioned(Map<Integer, com.datatorrent.api.Partitioner.Partition<RabbitMQLogsInputOperator>> partitions)
{
}
 

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

@Override
public void partitioned(
    Map<Integer, com.datatorrent.api.Partitioner.Partition<AbstractJdbcPollInputOperator<T>>> partitions)
{
  // Nothing to implement here
}
 

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

@Test
public void testPollWithOffsetRebase() throws Exception
{
  insertEvents(0, true, 0); // clear table

  JdbcStore store = new JdbcStore();
  store.setDatabaseDriver(DB_DRIVER);
  store.setDatabaseUrl(URL);

  List<FieldInfo> fieldInfos = getFieldInfos();

  Attribute.AttributeMap.DefaultAttributeMap portAttributes = new Attribute.AttributeMap.DefaultAttributeMap();
  portAttributes.put(Context.PortContext.TUPLE_CLASS, TestPOJOEvent.class);
  TestPortContext tpc = new TestPortContext(portAttributes);

  JdbcPOJOPollInputOperator inputOperator = new JdbcPOJOPollInputOperator();
  inputOperator.setStore(store);
  inputOperator.setTableName(TABLE_POJO_NAME);
  inputOperator.setColumnsExpression("ID,STARTDATE,STARTTIME,STARTTIMESTAMP");
  inputOperator.setKey("id");
  inputOperator.setFieldInfos(fieldInfos);
  inputOperator.setFetchSize(100);
  inputOperator.setBatchSize(100);
  inputOperator.setPartitionCount(1);
  inputOperator.setRebaseOffset(true);

  Collection<com.datatorrent.api.Partitioner.Partition<AbstractJdbcPollInputOperator<Object>>> newPartitions = inputOperator
      .definePartitions(new ArrayList<Partitioner.Partition<AbstractJdbcPollInputOperator<Object>>>(), null);

  int operatorId = 0;
  for (com.datatorrent.api.Partitioner.Partition<AbstractJdbcPollInputOperator<Object>> partition : newPartitions) {

    Attribute.AttributeMap.DefaultAttributeMap partitionAttributeMap = new Attribute.AttributeMap.DefaultAttributeMap();
    partitionAttributeMap.put(DAG.APPLICATION_ID, APP_ID);
    partitionAttributeMap.put(Context.DAGContext.APPLICATION_PATH, dir);

    OperatorContext partitioningContext = mockOperatorContext(operatorId++, partitionAttributeMap);

    JdbcPOJOPollInputOperator parition = (JdbcPOJOPollInputOperator)partition.getPartitionedInstance();
    parition.outputPort.setup(tpc);
    parition.setScheduledExecutorService(mockscheduler);
    parition.setup(partitioningContext);
    parition.activate(partitioningContext);
  }

  Iterator<com.datatorrent.api.Partitioner.Partition<AbstractJdbcPollInputOperator<Object>>> itr = newPartitions
      .iterator();
  // First partition is for range queries,last is for polling queries
  JdbcPOJOPollInputOperator firstInstance = (JdbcPOJOPollInputOperator)itr.next().getPartitionedInstance();

  int rows = 0;
  int windowId = 0;
  insertEvents(4, false, rows);
  rows += 4;
  JdbcPOJOPollInputOperator poller = (JdbcPOJOPollInputOperator)itr.next().getPartitionedInstance();
  CollectorTestSink<Object> sink3 = new CollectorTestSink<>();
  poller.outputPort.setSink(sink3);
  poller.beginWindow(windowId++);
  poller.pollRecords();
  poller.emitTuples();
  Assert.assertEquals("emitted", rows, sink3.collectedTuples.size());
  poller.endWindow();

  insertEvents(1, false, rows);
  rows += 1;
  poller.beginWindow(windowId++);
  poller.pollRecords(); // offset rebase, fetch 1 record
  poller.emitTuples();
  Assert.assertEquals("emitted", rows, sink3.collectedTuples.size());
  poller.endWindow();

}
 

@Override
public void partitioned(Map<Integer, com.datatorrent.api.Partitioner.Partition<UniqueValueCountAppender<V>>> partitions)
{
}
 

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

private void initPartitioning(PMapping m, int partitionCnt)
{
  Operator operator = m.logicalOperator.getOperator();
  Collection<Partition<Operator>> partitions;

  @SuppressWarnings("unchecked")
  Partitioner<Operator> partitioner = m.logicalOperator.getAttributes().contains(OperatorContext.PARTITIONER)
      ? (Partitioner<Operator>)m.logicalOperator.getValue(OperatorContext.PARTITIONER)
      : operator instanceof Partitioner ? (Partitioner<Operator>)operator : null;

  Collection<Partition<Operator>> collection = new ArrayList<>(1);
  DefaultPartition<Operator> firstPartition = new DefaultPartition<>(operator);
  collection.add(firstPartition);

  if (partitioner != null) {
    partitions = partitioner.definePartitions(collection, new PartitioningContextImpl(m, partitionCnt));

    if (partitions == null || partitions.isEmpty()) {
      throw new IllegalStateException("Partitioner returns null or empty.");
    }
  } else {
    //This handles the case when parallel partitioning is occurring. Partition count will be
    //Non zero in the case of parallel partitioning.
    for (int partitionCounter = 0; partitionCounter < partitionCnt - 1; partitionCounter++) {
      collection.add(firstPartition);
    }
    partitions = collection;
  }

  Collection<StatsListener> statsListeners = m.logicalOperator.getValue(OperatorContext.STATS_LISTENERS);
  if (statsListeners != null && !statsListeners.isEmpty()) {
    if (m.statsHandlers == null) {
      m.statsHandlers = new ArrayList<>(statsListeners.size());
    }
    for (StatsListener sl : statsListeners) {
      m.statsHandlers.add(getStatsListenerAdapter(sl));
    }
  }

  if (m.logicalOperator.getOperator() instanceof StatsListener) {
    if (m.statsHandlers == null) {
      m.statsHandlers = new ArrayList<>(1);
    }
    m.statsHandlers.add(new StatsListenerProxy(m.logicalOperator));
  }

  // create operator instance per partition
  Map<Integer, Partition<Operator>> operatorIdToPartition = Maps.newHashMapWithExpectedSize(partitions.size());
  for (Partition<Operator> partition : partitions) {
    PTOperator p = addPTOperator(m, partition, null);
    operatorIdToPartition.put(p.getId(), partition);
  }

  if (partitioner != null) {
    partitioner.partitioned(operatorIdToPartition);
  }
}
 

@Override
public void partitioned(Map<Integer, com.datatorrent.api.Partitioner.Partition<T>> partitions) {
}