DataSinkTask extends AbstractOutputTask { @Override public void invoke() throws Exception { if (LOG.isInfoEnabled()) LOG.info(getLogString("Start PACT code")); final UnilateralSortMerger<IT> sorter; if (this.config.getLocalStrategy() == LocalStrategy.SORT) { final Class<? extends TypeComparatorFactory<IT>> comparatorFactoryClass = this.config.getComparatorFactoryForInput(0, this.userCodeClassLoader); final TypeComparatorFactory<IT> comparatorFactory; if (comparatorFactoryClass == null) { @SuppressWarnings("unchecked") TypeComparatorFactory<IT> cf = (TypeComparatorFactory<IT>) PactRecordComparatorFactory.get(); comparatorFactory = cf; } else { comparatorFactory = InstantiationUtil.instantiate(comparatorFactoryClass, TypeComparatorFactory.class); } TypeComparator<IT> comparator; try { comparator = comparatorFactory.createComparator(this.config.getConfigForInputParameters(0), this.userCodeClassLoader); } catch (ClassNotFoundException cnfex) { throw new Exception("The instantiation of the type comparator from factory '" + comparatorFactory.getClass().getName() + "' failed. A referenced class from the user code could not be loaded."); } final long availableMemory = this.config.getMemorySize(); final int maxFileHandles = this.config.getNumFilehandles(); final float spillThreshold = this.config.getSortSpillingTreshold(); sorter = new UnilateralSortMerger<IT>(getEnvironment().getMemoryManager(), getEnvironment().getIOManager(), this.reader, this, this.inputTypeSerializer, comparator, availableMemory, maxFileHandles, spillThreshold); this.reader = sorter.getIterator(); } else { sorter = null; } try { final MutableObjectIterator<IT> reader = this.reader; final OutputFormat<IT> format = this.format; final IT record = this.inputTypeSerializer.createInstance(); if (this.taskCanceled) { return; } if (LOG.isDebugEnabled()) { LOG.debug(getLogString("Starting to produce output")); } format.open(this.getEnvironment().getIndexInSubtaskGroup() + 1); if (record.getClass() == PactRecord.class && format instanceof FileOutputFormat) { @SuppressWarnings("unchecked") final MutableObjectIterator<PactRecord> pi = (MutableObjectIterator<PactRecord>) reader; final PactRecord pr = (PactRecord) record; final FileOutputFormat pf = (FileOutputFormat) format; while (!this.taskCanceled && pi.next(pr)) { pf.writeRecord(pr); } } else { while (!this.taskCanceled && reader.next(record)) { format.writeRecord(record); } } if (!this.taskCanceled) { this.format.close(); this.format = null; } } catch (Exception ex) { if (!this.taskCanceled) { if (LOG.isErrorEnabled()) LOG.error(getLogString("Error in Pact user code: " + ex.getMessage()), ex); throw ex; } } finally { if (this.format != null) { try { this.format.close(); } catch (Throwable t) { if (LOG.isWarnEnabled()) LOG.warn(getLogString("Error closing the ouput format."), t); } } if (sorter != null) { sorter.close(); } } if (!this.taskCanceled) { if (LOG.isDebugEnabled()) LOG.debug(getLogString("Finished producing output")); 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 int getMaximumNumberOfSubtasks(); static final String DEGREE_OF_PARALLELISM_KEY; static final String SORT_ORDER; }
@Test public void testDataSinkTask() { int keyCnt = 100; int valCnt = 20; super.initEnvironment(1024 * 1024); super.addInput(new UniformPactRecordGenerator(keyCnt, valCnt, false), 1); DataSinkTask<PactRecord> testTask = new DataSinkTask<PactRecord>(); super.getTaskConfig().setLocalStrategy(LocalStrategy.NONE); super.getConfiguration().setString(DataSinkTask.SORT_ORDER, Order.NONE.name()); super.registerFileOutputTask(testTask, MockOutputFormat.class, "file: try { testTask.invoke(); } catch (Exception e) { LOG.debug(e); Assert.fail("Invoke method caused exception."); } File tempTestFile = new File(this.tempTestPath); Assert.assertTrue("Temp output file does not exist",tempTestFile.exists()); try { FileReader fr = new FileReader(tempTestFile); BufferedReader br = new BufferedReader(fr); HashMap<Integer,HashSet<Integer>> keyValueCountMap = new HashMap<Integer, HashSet<Integer>>(keyCnt); while(br.ready()) { String line = br.readLine(); Integer key = Integer.parseInt(line.substring(0,line.indexOf("_"))); Integer val = Integer.parseInt(line.substring(line.indexOf("_")+1,line.length())); 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 key : keyValueCountMap.keySet()) { Assert.assertTrue("Invalid value count for key: "+key+". Expected: "+valCnt+" Actual: "+keyValueCountMap.get(key).size(), keyValueCountMap.get(key).size() == valCnt); } } catch (FileNotFoundException e) { Assert.fail("Out file got lost..."); } catch (IOException ioe) { Assert.fail("Caught IOE while reading out file"); } }
@Test public void testFailingDataSinkTask() { int keyCnt = 100; int valCnt = 20; super.initEnvironment(1024 * 1024); super.addInput(new UniformPactRecordGenerator(keyCnt, valCnt, false), 1); DataSinkTask<PactRecord> testTask = new DataSinkTask<PactRecord>(); super.getTaskConfig().setLocalStrategy(LocalStrategy.NONE); Configuration stubParams = new Configuration(); stubParams.setString(DataSinkTask.SORT_ORDER, Order.NONE.name()); super.getTaskConfig().setStubParameters(stubParams); super.registerFileOutputTask(testTask, MockFailingOutputFormat.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 testCancelDataSinkTask() { super.initEnvironment(1024 * 1024); super.addInput(new InfiniteInputIterator(), 1); final DataSinkTask<PactRecord> testTask = new DataSinkTask<PactRecord>(); super.getTaskConfig().setLocalStrategy(LocalStrategy.NONE); Configuration stubParams = new Configuration(); stubParams.setString(DataSinkTask.SORT_ORDER, Order.NONE.name()); super.getTaskConfig().setStubParameters(stubParams); super.registerFileOutputTask(testTask, MockOutputFormat.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()); }