/*
* 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.test;
import static org.junit.Assert.assertEquals;
import java.io.IOException;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Random;
import org.apache.pig.Algebraic;
import org.apache.pig.ComparisonFunc;
import org.apache.pig.EvalFunc;
import org.apache.pig.FuncSpec;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.POStatus;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.PhysicalOperator;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.Result;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POUserComparisonFunc;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POUserFunc;
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.schema.Schema;
import org.apache.pig.impl.plan.OperatorKey;
import org.apache.pig.test.utils.GenRandomData;
import org.junit.Test;
public class TestPOUserFunc {
Random r = new Random(42L);
int MAX_TUPLES = 10;
public static class ARITY extends EvalFunc<Integer> {
@Override
public Integer exec(Tuple input) throws IOException {
try {
return new Integer(((Tuple)input.get(0)).size());
} catch (ExecException e) {
// TODO Auto-generated catch block
e.printStackTrace();
}
return 0;
}
@Override
public Schema outputSchema(Schema input) {
return new Schema(new Schema.FieldSchema(null, DataType.INTEGER));
}
}
public static class WeirdComparator extends ComparisonFunc {
@Override
public int compare(Tuple t1, Tuple t2) {
// TODO Auto-generated method stub
Object o1 = null;
Object o2 = null;
try {
o1 = t1.get(2);
o2 = t2.get(2);
} catch (ExecException e) {
// TODO Auto-generated catch block
e.printStackTrace();
}
if ( o1==null || o2==null ){
return -1;
}
int i1 = (Integer) o1 - 2;
int i2 = (Integer) o2 - 2;
return (int) (i1 * i1 - i2 * i2);
}
}
/**
* Generates the average of the values of the first field of a tuple. This
* class is Algebraic in implemenation, so if possible the execution will be
* split into a local and global application
*/
public static class AVG extends EvalFunc<Double> implements Algebraic {
private static TupleFactory mTupleFactory = TupleFactory.getInstance();
@Override
public Double exec(Tuple input) throws IOException {
double sum = 0;
double count = 0;
try {
sum = sum(input);
count = count(input);
} catch (ExecException e) {
e.printStackTrace();
}
double avg = 0;
if (count > 0)
avg = sum / count;
return new Double(avg);
}
public String getInitial() {
return Initial.class.getName();
}
public String getIntermed() {
return Intermed.class.getName();
}
public String getFinal() {
return Final.class.getName();
}
static public class Initial extends EvalFunc<Tuple> {
@Override
public Tuple exec(Tuple input) throws IOException {
try {
Tuple t = mTupleFactory.newTuple(2);
t.set(0, new Double(sum(input)));
t.set(1, new Long(count(input)));
return t;
} catch (ExecException t) {
throw new RuntimeException(t.getMessage() + ": " + input);
}
}
}
static public class Intermed extends EvalFunc<Tuple> {
@Override
public Tuple exec(Tuple input) throws IOException {
DataBag b = null;
Tuple t = null;
try {
b = (DataBag) input.get(0);
t = combine(b);
} catch (ExecException e) {
// TODO Auto-generated catch block
e.printStackTrace();
}
return t;
}
}
static public class Final extends EvalFunc<Double> {
@Override
public Double exec(Tuple input) throws IOException {
double sum = 0;
double count = 0;
try {
DataBag b = (DataBag) input.get(0);
Tuple combined = combine(b);
sum = (Double) combined.get(0);
count = (Long) combined.get(1);
} catch (ExecException e) {
e.printStackTrace();
}
double avg = 0;
if (count > 0) {
avg = sum / count;
}
return new Double(avg);
}
}
static protected Tuple combine(DataBag values) throws ExecException {
double sum = 0;
long count = 0;
Tuple output = mTupleFactory.newTuple(2);
for (Iterator<Tuple> it = values.iterator(); it.hasNext();) {
Tuple t = it.next();
sum += (Double) t.get(0);
count += (Long) t.get(1);
}
output.set(0, new Double(sum));
output.set(1, new Long(count));
return output;
}
static protected long count(Tuple input) throws ExecException {
DataBag values = (DataBag) input.get(0);
return values.size();
}
static protected double sum(Tuple input) throws ExecException {
DataBag values = (DataBag) input.get(0);
double sum = 0;
for (Iterator<Tuple> it = values.iterator(); it.hasNext();) {
Tuple t = it.next();
Double d = DataType.toDouble(t.get(0));
if (d == null)
continue;
sum += d;
}
return sum;
}
@Override
public Schema outputSchema(Schema input) {
return new Schema(new Schema.FieldSchema(null, DataType.DOUBLE));
}
}
@Test
public void testUserFuncArity() throws ExecException {
DataBag input = (DataBag) GenRandomData.genRandSmallTupDataBag(r,
MAX_TUPLES, 100);
userFuncArity( input );
}
@Test
public void testUserFuncArityWithNulls() throws ExecException {
DataBag input = (DataBag) GenRandomData.genRandSmallTupDataBagWithNulls(r,
MAX_TUPLES, 100);
userFuncArity( input );
}
public void userFuncArity(DataBag input ) throws ExecException {
String funcSpec = ARITY.class.getName() + "()";
PORead read = new PORead(new OperatorKey("", r.nextLong()), input);
List<PhysicalOperator> inputs = new LinkedList<PhysicalOperator>();
inputs.add(read);
POUserFunc userFunc = new POUserFunc(new OperatorKey("", r.nextLong()),
-1, inputs, new FuncSpec(funcSpec));
Result res = new Result();
Integer i = null;
res = userFunc.getNextInteger();
while (res.returnStatus != POStatus.STATUS_EOP) {
// System.out.println(res.result);
int result = (Integer) res.result;
assertEquals(2, result);
res = userFunc.getNextInteger();
}
}
@Test
public void testUDFCompare() throws ExecException {
DataBag input = (DataBag) GenRandomData.genRandSmallTupDataBag(r, 2, 100);
udfCompare(input);
}
@Test
public void testUDFCompareWithNulls() throws ExecException {
DataBag input = (DataBag) GenRandomData.genRandSmallTupDataBagWithNulls(r, 2, 100);
udfCompare(input);
}
public void udfCompare(DataBag input) throws ExecException {
String funcSpec = WeirdComparator.class.getName() + "()";
POUserComparisonFunc userFunc = new POUserComparisonFunc(new OperatorKey("", r.nextLong()),
-1, null, new FuncSpec(funcSpec));
Iterator<Tuple> it = input.iterator();
Tuple t1 = it.next();
Tuple t2 = it.next();
t1.append(2);
t2.append(3);
userFunc.attachInput(t1, t2);
Integer i = null;
// System.out.println(t1 + " " + t2);
int result = (Integer) (userFunc.getNextInteger().result);
assertEquals(-1, result);
}
@Test
public void testAlgebraicAVG() throws IOException, ExecException {
Integer input[] = { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 };
algebraicAVG( input, 55.0, 10L, 110.0, 20L, 5.5 );
}
/* NOTE: for calculating the average
*
* A pig "count" will include data that had "null",and the sum will
* A pig "count" will include data that had "null",and the sum will
* treat the null as a 0, impacting the average
* A SQL "count" will exclude data that had "null"
*/
@Test
public void testAlgebraicAVGWithNulls() throws IOException, ExecException {
Integer input[] = { 1, 2, 3, 4, null, 6, 7, 8, 9, 10 };
algebraicAVG( input, 50.0, 10L, 100.0, 20L, 5.0 );
}
public void algebraicAVG(
Integer[] input
, Double initialExpectedSum, Long initialExpectedCount
, Double intermedExpectedSum, Long intermedExpectedCount
, Double expectedAvg
) throws IOException, ExecException {
// generate data
byte INIT = 0;
byte INTERMED = 1;
byte FINAL = 2;
Tuple tup1 = Util.loadNestTuple(TupleFactory.getInstance().newTuple(1),
input);
Tuple tup2 = Util.loadNestTuple(TupleFactory.getInstance().newTuple(1),
input);
// System.out.println("Input = " + tup1);
String funcSpec = AVG.class.getName() + "()";
POUserFunc po = new POUserFunc(new OperatorKey("", r.nextLong()), -1,
null, new FuncSpec(funcSpec));
//************ Initial Calculations ******************
TupleFactory tf = TupleFactory.getInstance();
po.setAlgebraicFunction(INIT);
po.attachInput(tup1);
Tuple t = null;
Result res = po.getNextTuple();
Tuple outputInitial1 = (res.returnStatus == POStatus.STATUS_OK) ? (Tuple) res.result
: null;
Tuple outputInitial2 = (res.returnStatus == POStatus.STATUS_OK) ? (Tuple) res.result
: null;
System.out.println(outputInitial1 + " " + outputInitial2);
assertEquals(outputInitial1, outputInitial2);
Double sum = (Double) outputInitial1.get(0);
Long count = (Long) outputInitial1.get(1);
assertEquals(initialExpectedSum, sum);
assertEquals(initialExpectedCount, count);
//************ Intermediate Data and Calculations ******************
DataBag bag = BagFactory.getInstance().newDefaultBag();
bag.add(outputInitial1);
bag.add(outputInitial2);
Tuple outputInitial = tf.newTuple();
outputInitial.append(bag);
// Tuple outputIntermed = intermed.exec(outputInitial);
po = new POUserFunc(new OperatorKey("", r.nextLong()), -1, null,
new FuncSpec(funcSpec));
po.setAlgebraicFunction(INTERMED);
po.attachInput(outputInitial);
res = po.getNextTuple();
Tuple outputIntermed = (res.returnStatus == POStatus.STATUS_OK) ? (Tuple) res.result
: null;
sum = (Double) outputIntermed.get(0);
count = (Long) outputIntermed.get(1);
assertEquals(intermedExpectedSum, sum);
assertEquals(intermedExpectedCount, count);
System.out.println(outputIntermed);
//************ Final Calculations ******************
po = new POUserFunc(new OperatorKey("", r.nextLong()), -1, null,
new FuncSpec(funcSpec));
po.setAlgebraicFunction(FINAL);
po.attachInput(outputInitial);
res = po.getNextTuple();
Double output = (res.returnStatus == POStatus.STATUS_OK) ? (Double) res.result
: null;
// Double output = fin.exec(outputInitial);
assertEquals((Double)expectedAvg, output);
// System.out.println("output = " + output);
}
}