DataSourceTask extends AbstractInputTask<InputSplit> { @Override public void invoke() throws Exception { if (LOG.isInfoEnabled()) LOG.info(getLogString("Start PACT code")); try { RegularPactTask.openChainedTasks(this.chainedTasks, this); final Iterator<InputSplit> splitIterator = getInputSplits(); final OT record = this.serializer.createInstance(); while (!this.taskCanceled && splitIterator.hasNext()) { final InputSplit split = splitIterator.next(); if (LOG.isDebugEnabled()) LOG.debug(getLogString("Opening input split " + split.toString())); final InputFormat<OT, InputSplit> format = this.format; format.open(split); if (LOG.isDebugEnabled()) LOG.debug(getLogString("Starting to read input from split " + split.toString())); try { if (record.getClass() == PactRecord.class) { final PactRecord pactRecord = (PactRecord) record; @SuppressWarnings("unchecked") final InputFormat<PactRecord, InputSplit> inFormat = (InputFormat<PactRecord, InputSplit>) format; if (this.output instanceof PactRecordOutputCollector) { final PactRecordOutputCollector output = (PactRecordOutputCollector) this.output; while (!this.taskCanceled && !inFormat.reachedEnd()) { pactRecord.clear(); if (inFormat.nextRecord(pactRecord)) { output.collect(pactRecord); } } } else if (this.output instanceof ChainedMapDriver) { @SuppressWarnings("unchecked") final ChainedMapDriver<PactRecord, ?> output = (ChainedMapDriver<PactRecord, ?>) this.output; while (!this.taskCanceled && !inFormat.reachedEnd()) { pactRecord.clear(); if (inFormat.nextRecord(pactRecord)) { output.collect(pactRecord); } } } else { @SuppressWarnings("unchecked") final Collector<PactRecord> output = (Collector<PactRecord>) this.output; while (!this.taskCanceled && !inFormat.reachedEnd()) { pactRecord.clear(); if (inFormat.nextRecord(pactRecord)) { output.collect(pactRecord); } } } } else { if (this.output instanceof OutputCollector) { final OutputCollector<OT> output = (OutputCollector<OT>) this.output; while (!this.taskCanceled && !format.reachedEnd()) { if (format.nextRecord(record)) { output.collect(record); } } } else if (this.output instanceof ChainedMapDriver) { @SuppressWarnings("unchecked") final ChainedMapDriver<OT, ?> output = (ChainedMapDriver<OT, ?>) this.output; while (!this.taskCanceled && !format.reachedEnd()) { if (format.nextRecord(record)) { output.collect(record); } } } else { final Collector<OT> output = this.output; while (!this.taskCanceled && !format.reachedEnd()) { if (format.nextRecord(record)) { output.collect(record); } } } } if (LOG.isDebugEnabled() && !this.taskCanceled) { LOG.debug(getLogString("Closing input split " + split.toString())); } } finally { format.close(); } } this.output.close(); RegularPactTask.closeChainedTasks(this.chainedTasks, this); } catch (Exception ex) { try { this.format.close(); } catch (Throwable t) {} RegularPactTask.cancelChainedTasks(this.chainedTasks); if (!this.taskCanceled) { RegularPactTask.logAndThrowException(ex, this); } } if (!this.taskCanceled) { if (LOG.isInfoEnabled()) LOG.info(getLogString("Finished PACT code")); } else { if (LOG.isWarnEnabled()) LOG.warn(getLogString("PACT code cancelled")); } }  @Override void registerInputOutput(); @Override void invoke(); @Override void cancel(); void setUserCodeClassLoader(ClassLoader cl); @Override InputSplit[] computeInputSplits(int requestedMinNumber); @SuppressWarnings("unchecked") @Override Class<InputSplit> getInputSplitType(); @Override int getMinimumNumberOfSubtasks(); @Override int getMaximumNumberOfSubtasks();  }
@Test public void testDataSourceTask() { int keyCnt = 100; int valCnt = 20; this.outList = new ArrayList<PactRecord>(); try { InputFilePreparator.prepareInputFile(new UniformPactRecordGenerator(keyCnt, valCnt, false), this.tempTestPath, true); } catch (IOException e1) { Assert.fail("Unable to set-up test input file"); } super.initEnvironment(1024 * 1024); super.addOutput(this.outList); DataSourceTask<PactRecord> testTask = new DataSourceTask<PactRecord>(); super.registerFileInputTask(testTask, MockInputFormat.class, "file: try { testTask.invoke(); } catch (Exception e) { System.err.println(e); Assert.fail("Invoke method caused exception."); } Assert.assertTrue("Invalid output size. Expected: "+(keyCnt*valCnt)+" Actual: "+this.outList.size(), this.outList.size() == keyCnt * valCnt); HashMap<Integer,HashSet<Integer>> keyValueCountMap = new HashMap<Integer, HashSet<Integer>>(keyCnt); for (PactRecord kvp : this.outList) { int key = kvp.getField(0, PactInteger.class).getValue(); int val = kvp.getField(1, PactInteger.class).getValue(); if(!keyValueCountMap.containsKey(key)) { keyValueCountMap.put(key,new HashSet<Integer>()); } keyValueCountMap.get(key).add(val); } Assert.assertTrue("Invalid key count in out file. Expected: "+keyCnt+" Actual: "+keyValueCountMap.keySet().size(), keyValueCountMap.keySet().size() == keyCnt); for(Integer mapKey : keyValueCountMap.keySet()) { Assert.assertTrue("Invalid value count for key: "+mapKey+". Expected: "+valCnt+" Actual: "+keyValueCountMap.get(mapKey).size(), keyValueCountMap.get(mapKey).size() == valCnt); } }
@Test public void testFailingDataSourceTask() { int keyCnt = 20; int valCnt = 10; this.outList = new NirvanaOutputList(); try { InputFilePreparator.prepareInputFile(new UniformPactRecordGenerator(keyCnt, valCnt, false), this.tempTestPath, false); } catch (IOException e1) { Assert.fail("Unable to set-up test input file"); } super.initEnvironment(1024 * 1024); super.addOutput(this.outList); DataSourceTask<PactRecord> testTask = new DataSourceTask<PactRecord>(); super.registerFileInputTask(testTask, MockFailingInputFormat.class, "file: boolean stubFailed = false; try { testTask.invoke(); } catch (Exception e) { stubFailed = true; } Assert.assertTrue("Stub exception was not forwarded.", stubFailed); File tempTestFile = new File(this.tempTestPath); Assert.assertTrue("Temp output file does not exist",tempTestFile.exists()); }
@Test public void testCancelDataSourceTask() { int keyCnt = 20; int valCnt = 4; super.initEnvironment(1024 * 1024); super.addOutput(new NirvanaOutputList()); try { InputFilePreparator.prepareInputFile(new UniformPactRecordGenerator(keyCnt, valCnt, false), this.tempTestPath, false); } catch (IOException e1) { Assert.fail("Unable to set-up test input file"); } final DataSourceTask<PactRecord> testTask = new DataSourceTask<PactRecord>(); super.registerFileInputTask(testTask, MockDelayingInputFormat.class, "file: Thread taskRunner = new Thread() { @Override public void run() { try { testTask.invoke(); } catch (Exception ie) { ie.printStackTrace(); Assert.fail("Task threw exception although it was properly canceled"); } } }; taskRunner.start(); TaskCancelThread tct = new TaskCancelThread(1, taskRunner, testTask); tct.start(); try { tct.join(); taskRunner.join(); } catch(InterruptedException ie) { Assert.fail("Joining threads failed"); } File tempTestFile = new File(this.tempTestPath); Assert.assertTrue("Temp output file does not exist",tempTestFile.exists()); }