Java tutorial
/* * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.apache.pig.builtin; import java.io.IOException; import java.util.ArrayList; import java.util.Comparator; import java.util.Iterator; import java.util.List; import java.util.PriorityQueue; import java.util.Random; import org.apache.commons.logging.Log; import org.apache.commons.logging.LogFactory; import org.apache.pig.AccumulatorEvalFunc; import org.apache.pig.Algebraic; import org.apache.pig.EvalFunc; import org.apache.pig.FuncSpec; import org.apache.pig.backend.executionengine.ExecException; import org.apache.pig.data.BagFactory; import org.apache.pig.data.DataBag; import org.apache.pig.data.DataType; import org.apache.pig.data.Tuple; import org.apache.pig.data.TupleFactory; import org.apache.pig.impl.logicalLayer.FrontendException; import org.apache.pig.impl.logicalLayer.schema.Schema; import org.apache.pig.impl.logicalLayer.schema.Schema.FieldSchema; /** * Top UDF accepts a bag of tuples and returns top-n tuples depending upon the * tuple field value of type long. Both n and field number needs to be provided * to the UDF. The UDF iterates through the input bag and just retains top-n * tuples by storing them in a priority queue of size n+1 where priority is the * long field. This is efficient as priority queue provides constant time - O(1) * removal of the least element and O(log n) time for heap restructuring. The * UDF is especially helpful for turning the nested grouping operation inside * out and retaining top-n in a nested group. * * Assumes all tuples in the bag contain an element of the same type in the compared column. * * Sample usage: * DEFINE TOP_ASC TOP("ASC") * DEFINE TOP_DESC TOP("DESC") * * A = LOAD 'test.tsv' as (first: chararray, second: chararray); * B = GROUP A BY (first, second); * C = FOREACH B generate FLATTEN(group), COUNT(*) as count; * D = GROUP C BY first; // again group by first * topResults = FOREACH D { * result = TOP_ASC(10, 1, C); // and retain top 10 occurrences of 'second' in first * GENERATE FLATTEN(result); * topDescResults = FOREACH D { * result = TOP_DESC(10, 1, C); // and retain top 10 occurrences of 'second' in first * GENERATE FLATTEN(result); * * } */ public class TOP extends AccumulatorEvalFunc<DataBag> implements Algebraic { private static final Log log = LogFactory.getLog(TOP.class); private static BagFactory mBagFactory = BagFactory.getInstance(); private static TupleFactory mTupleFactory = TupleFactory.getInstance(); private final Random randomizer = new Random(); private boolean sortDesc; // By default, set the sorting order to descending. This is to make it // backward compatible. public TOP() { sortDesc = true; } // Allow to set the sorting order explicitly. public TOP(String arg) { sortDesc = decideSortOrder(arg); } private static boolean decideSortOrder(String arg) { if ("ASC".equalsIgnoreCase(arg)) { return false; } else if ("DESC".equalsIgnoreCase(arg)) { return true; } else { log.warn("Unknown order '" + arg + "' is given. Defaulting to descending order."); return true; } } private static class TupleComparator implements Comparator<Tuple> { private final int fieldNum; private byte datatype; private boolean typeFound = false; private boolean isDescOrder = true; public TupleComparator(int fieldNum, boolean isDescOrder) { this.fieldNum = fieldNum; this.isDescOrder = isDescOrder; } /* * (non-Javadoc) * @see java.util.Comparator#compare(java.lang.Object, java.lang.Object) */ @Override public int compare(Tuple o1, Tuple o2) { int ret = 0; if (o1 == null) { ret = -1; } else if (o2 == null) { ret = 1; } else { try { Object field1 = o1.get(fieldNum); Object field2 = o2.get(fieldNum); if (!typeFound) { datatype = DataType.findType(field1); if (datatype != DataType.NULL) { typeFound = true; } } ret = DataType.compare(field1, field2, datatype, datatype); } catch (ExecException e) { throw new RuntimeException("Error while comparing o1:" + o1 + " and o2:" + o2, e); } } return isDescOrder ? ret : ret * -1; } } // for Accumulator interface private PriorityQueue<Tuple> store = null; @Override public void accumulate(Tuple tuple) throws IOException { if (tuple == null || tuple.size() < 3) { return; } try { int n = (Integer) tuple.get(0); int fieldNum = (Integer) tuple.get(1); DataBag inputBag = (DataBag) tuple.get(2); if (inputBag == null) { return; } if (store == null) { store = new PriorityQueue<Tuple>(n + 1, new TupleComparator(fieldNum, sortDesc)); } updateTop(store, n, inputBag); } catch (ExecException e) { throw new RuntimeException("ExecException executing function: ", e); } catch (Exception e) { throw new RuntimeException("General Exception executing function: ", e); } } @Override public DataBag getValue() { if (store == null) { return null; } DataBag outputBag = mBagFactory.newDefaultBag(); for (Tuple t : store) { outputBag.add(t); } if (log.isDebugEnabled()) { if (randomizer.nextInt(1000) == 1) { log.debug("outputting a bag: "); try { for (Tuple t : outputBag) { log.debug("outputting " + t.toDelimitedString("\t")); } } catch (ExecException e) { throw new RuntimeException("ExecException executing function: ", e); } log.debug("=================="); } } return outputBag; } @Override public void cleanup() { store = null; } protected static void updateTop(PriorityQueue<Tuple> store, int limit, DataBag inputBag) { Iterator<Tuple> itr = inputBag.iterator(); while (itr.hasNext()) { Tuple t = itr.next(); store.add(t); if (store.size() > limit) store.poll(); } } /* * (non-Javadoc) * * @see org.apache.pig.EvalFunc#getArgToFuncMapping() */ @Override public List<FuncSpec> getArgToFuncMapping() throws FrontendException { List<FieldSchema> fields = new ArrayList<FieldSchema>(3); fields.add(new Schema.FieldSchema(null, DataType.INTEGER)); fields.add(new Schema.FieldSchema(null, DataType.INTEGER)); fields.add(new Schema.FieldSchema(null, DataType.BAG)); FuncSpec funcSpec = new FuncSpec(this.getClass().getName(), new Schema(fields)); List<FuncSpec> funcSpecs = new ArrayList<FuncSpec>(1); funcSpecs.add(funcSpec); return funcSpecs; } @Override public Schema outputSchema(Schema input) { try { if (input.size() < 3) { return null; } return new Schema(input.getField(2)); } catch (Exception e) { return null; } } @Override public String getInitial() { return Initial.class.getName(); } @Override public String getIntermed() { return Intermed.class.getName(); } @Override public String getFinal() { return Final.class.getName(); } /* * Same as normal code-path exec, but outputs a Tuple with the schema * <Int, Int, DataBag> -- same schema as expected input. */ public static class Initial extends EvalFunc<Tuple> { @SuppressWarnings("unused") private boolean sortDesc; public Initial() { this.sortDesc = true; } public Initial(String arg) { sortDesc = decideSortOrder(arg); } @Override public Tuple exec(Tuple tuple) throws IOException { if (tuple == null || tuple.size() < 3) { return null; } try { int n = (Integer) tuple.get(0); int fieldNum = (Integer) tuple.get(1); DataBag inputBag = (DataBag) tuple.get(2); if (inputBag == null) { return null; } Tuple retTuple = mTupleFactory.newTuple(3); DataBag outputBag = mBagFactory.newDefaultBag(); // initially, there should only be one, so not much point in doing the priority queue for (Tuple t : inputBag) { outputBag.add(t); } retTuple.set(0, n); retTuple.set(1, fieldNum); retTuple.set(2, outputBag); return retTuple; } catch (Exception e) { throw new RuntimeException("General Exception executing function: ", e); } } } public static class Intermed extends EvalFunc<Tuple> { private static final Log log = LogFactory.getLog(Intermed.class); private final Random randomizer = new Random(); private boolean sortDesc; public Intermed() { this.sortDesc = true; } public Intermed(String arg) { sortDesc = decideSortOrder(arg); } /* The input is a tuple that contains a single bag. * This bag contains outputs of the Initial step -- * tuples of the format (limit, index, { top_tuples }) * * We need to take the top of tops and return a similar tuple. * * (non-Javadoc) * @see org.apache.pig.EvalFunc#exec(org.apache.pig.data.Tuple) */ @Override public Tuple exec(Tuple input) throws IOException { if (input == null || input.size() < 1) { return null; } try { DataBag bagOfIntermediates = (DataBag) input.get(0); Iterator<Tuple> intermediateIterator = bagOfIntermediates.iterator(); if (!intermediateIterator.hasNext()) { return null; } Tuple peekTuple = intermediateIterator.next(); if (peekTuple == null || peekTuple.size() < 3) return null; int n = (Integer) peekTuple.get(0); int fieldNum = (Integer) peekTuple.get(1); DataBag inputBag = (DataBag) peekTuple.get(2); boolean allInputBagsNull = true; PriorityQueue<Tuple> store = new PriorityQueue<Tuple>(n + 1, new TupleComparator(fieldNum, sortDesc)); if (inputBag != null) { allInputBagsNull = false; updateTop(store, n, inputBag); } while (intermediateIterator.hasNext()) { Tuple t = intermediateIterator.next(); if (t == null || t.size() < 3) continue; inputBag = (DataBag) t.get(2); if (inputBag != null) { allInputBagsNull = false; updateTop(store, n, inputBag); } } Tuple retTuple = mTupleFactory.newTuple(3); retTuple.set(0, n); retTuple.set(1, fieldNum); DataBag outputBag = null; if (!allInputBagsNull) { outputBag = mBagFactory.newDefaultBag(); for (Tuple t : store) { outputBag.add(t); } } retTuple.set(2, outputBag); if (log.isDebugEnabled()) { if (randomizer.nextInt(1000) == 1) log.debug("outputting " + retTuple.toDelimitedString("\t")); } return retTuple; } catch (ExecException e) { throw new RuntimeException("ExecException executing function: ", e); } catch (Exception e) { throw new RuntimeException("General Exception executing function: ", e); } } } public static class Final extends EvalFunc<DataBag> { private static final Log log = LogFactory.getLog(Final.class); private final Random randomizer = new Random(); private boolean sortDesc; public Final() { this.sortDesc = true; } public Final(String arg) { sortDesc = decideSortOrder(arg); } /* * The input to this function is a tuple that contains a single bag. * This bag, in turn, contains outputs of the Intermediate step -- * tuples of the format (limit, index, { top_tuples } ) * * we want to return a bag of top tuples * * (non-Javadoc) * @see org.apache.pig.EvalFunc#exec(org.apache.pig.data.Tuple) */ @Override public DataBag exec(Tuple tuple) throws IOException { if (tuple == null || tuple.size() < 1) { return null; } try { DataBag bagOfIntermediates = (DataBag) tuple.get(0); Iterator<Tuple> intermediateIterator = bagOfIntermediates.iterator(); if (!intermediateIterator.hasNext()) { return null; } Tuple peekTuple = intermediateIterator.next(); if (peekTuple == null || peekTuple.size() < 3) return null; int n = (Integer) peekTuple.get(0); int fieldNum = (Integer) peekTuple.get(1); DataBag inputBag = (DataBag) peekTuple.get(2); boolean allInputBagsNull = true; PriorityQueue<Tuple> store = new PriorityQueue<Tuple>(n + 1, new TupleComparator(fieldNum, sortDesc)); if (inputBag != null) { allInputBagsNull = false; updateTop(store, n, inputBag); } while (intermediateIterator.hasNext()) { Tuple t = intermediateIterator.next(); if (t == null || t.size() < 3) continue; inputBag = (DataBag) t.get(2); if (inputBag != null) { allInputBagsNull = false; updateTop(store, n, inputBag); } } if (allInputBagsNull) { return null; } DataBag outputBag = mBagFactory.newDefaultBag(); for (Tuple t : store) { outputBag.add(t); } if (log.isDebugEnabled()) { if (randomizer.nextInt(1000) == 1) for (Tuple t : outputBag) log.debug("outputting " + t.toDelimitedString("\t")); } return outputBag; } catch (ExecException e) { throw new RuntimeException("ExecException executing function: ", e); } catch (Exception e) { throw new RuntimeException("General Exception executing function: ", e); } } } }