/*
* 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 static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import java.io.IOException;
import org.apache.hadoop.mapreduce.Job;
import org.apache.pig.ExecType;
import org.apache.pig.Expression;
import org.apache.pig.LoadMetadata;
import org.apache.pig.PigServer;
import org.apache.pig.ResourceSchema;
import org.apache.pig.ResourceStatistics;
import org.apache.pig.backend.executionengine.ExecException;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.MapReduceOper;
import org.apache.pig.backend.hadoop.executionengine.mapReduceLayer.plans.MROperPlan;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.expressionOperators.POProject;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.plans.PhysicalPlan;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POSort;
import org.apache.pig.backend.hadoop.executionengine.physicalLayer.relationalOperators.POStore;
import org.apache.pig.builtin.PigStorage;
import org.apache.pig.data.DataType;
import org.apache.pig.impl.PigContext;
import org.apache.pig.impl.logicalLayer.schema.Schema;
import org.apache.pig.impl.util.Utils;
import org.apache.pig.newplan.Operator;
import org.apache.pig.newplan.logical.expression.LogicalExpression;
import org.apache.pig.newplan.logical.relational.LOCogroup;
import org.apache.pig.newplan.logical.relational.LOFilter;
import org.apache.pig.newplan.logical.relational.LOForEach;
import org.apache.pig.newplan.logical.relational.LOLoad;
import org.apache.pig.newplan.logical.relational.LOSort;
import org.apache.pig.newplan.logical.relational.LOStore;
import org.apache.pig.newplan.logical.relational.LogicalPlan;
import org.apache.pig.newplan.logical.relational.LogicalSchema;
import org.apache.pig.newplan.logical.relational.LogicalSchema.LogicalFieldSchema;
import org.apache.pig.parser.ParserException;
import org.junit.BeforeClass;
import org.junit.Test;
public class TestPlanGeneration {
private static PigContext pc;
private static PigServer ps;
@BeforeClass
public static void setUp() throws ExecException {
ps = new PigServer(ExecType.LOCAL);
pc = ps.getPigContext();
pc.connect();
}
@Test
public void testGenerateStar() throws Exception {
String query = "a = load 'x';" +
"b = foreach a generate *;" +
"store b into '111';";
LogicalPlan lp = Util.parseAndPreprocess(query, pc);
Util.optimizeNewLP(lp);
LOStore loStore = (LOStore)lp.getSinks().get(0);
LOForEach loForEach = (LOForEach)lp.getPredecessors(loStore).get(0);
assertNull(loForEach.getSchema());
}
@Test
public void testEmptyBagDereference() throws Exception {
String query = "A = load 'x' as ( u:bag{} );" +
"B = foreach A generate u.$100;" +
"store B into '111';";
LogicalPlan lp = Util.parseAndPreprocess(query, pc);
Util.optimizeNewLP(lp);
LOStore loStore = (LOStore)lp.getSinks().get(0);
LOForEach loForEach = (LOForEach)lp.getPredecessors(loStore).get(0);
LogicalSchema schema = loForEach.getSchema();
assertEquals(1, schema.size());
LogicalFieldSchema bagFieldSchema = schema.getField(0);
assertEquals(DataType.BAG, bagFieldSchema.type);
LogicalFieldSchema tupleFieldSchema = bagFieldSchema.schema.getField(0);
assertEquals(1, tupleFieldSchema.schema.size());
assertEquals(DataType.BYTEARRAY, tupleFieldSchema.schema.getField(0).type);
}
@Test
public void testEmptyTupleDereference() throws Exception {
String query = "A = load 'x' as ( u:tuple() );" +
"B = foreach A generate u.$100;" +
"store B into '111';";
LogicalPlan lp = Util.parseAndPreprocess(query, pc);
Util.optimizeNewLP(lp);
LOStore loStore = (LOStore)lp.getSinks().get(0);
LOForEach loForEach = (LOForEach)lp.getPredecessors(loStore).get(0);
LogicalSchema schema = loForEach.getSchema();
assertEquals(1, schema.size());
assertEquals(DataType.BYTEARRAY, schema.getField(0).type);
}
@Test
public void testEmptyBagInnerPlan() throws Exception {
String query = "A = load 'x' as ( u:bag{} );" +
"B = foreach A { B1 = filter u by $1==0; generate B1;};" +
"store B into '111';";
LogicalPlan lp = Util.parseAndPreprocess(query, pc);
Util.optimizeNewLP(lp);
LOStore loStore = (LOStore)lp.getSinks().get(0);
LOForEach loForEach = (LOForEach)lp.getPredecessors(loStore).get(0);
LogicalSchema schema = loForEach.getSchema();
assertEquals(1, schema.size());
LogicalFieldSchema bagFieldSchema = schema.getField(0);
assertEquals(DataType.BAG, bagFieldSchema.type);
LogicalFieldSchema tupleFieldSchema = bagFieldSchema.schema.getField(0);
assertNull(tupleFieldSchema.schema);
}
@Test
public void testOrderByNullFieldSchema() throws Exception {
String query = "A = load 'x';" +
"B = order A by *;" +
"store B into '111';";
LogicalPlan lp = Util.parseAndPreprocess(query, pc);
Util.optimizeNewLP(lp);
LOStore loStore = (LOStore)lp.getSinks().get(0);
LOSort loSort = (LOSort)lp.getPredecessors(loStore).get(0);
Operator sortPlanLeaf = loSort.getSortColPlans().get(0).getSources().get(0);
LogicalFieldSchema sortPlanFS = ((LogicalExpression)sortPlanLeaf).getFieldSchema();
assertNull(sortPlanFS);
PhysicalPlan pp = Util.buildPhysicalPlanFromNewLP(lp, pc);
POStore poStore = (POStore)pp.getLeaves().get(0);
POSort poSort = (POSort)pp.getPredecessors(poStore).get(0);
POProject poProject = (POProject)poSort.getSortPlans().get(0).getLeaves().get(0);
assertEquals(DataType.TUPLE, poProject.getResultType());
}
@Test
public void testGroupByNullFieldSchema() throws Exception {
String query = "A = load 'x';" +
"B = group A by *;" +
"store B into '111';";
LogicalPlan lp = Util.parseAndPreprocess(query, pc);
Util.optimizeNewLP(lp);
LOStore loStore = (LOStore)lp.getSinks().get(0);
LOCogroup loCoGroup = (LOCogroup)lp.getPredecessors(loStore).get(0);
LogicalFieldSchema groupFieldSchema = loCoGroup.getSchema().getField(0);
assertEquals(DataType.TUPLE, groupFieldSchema.type);
assertNull(groupFieldSchema.schema);
}
@Test
public void testStoreAlias() throws Exception {
String query = "A = load 'data' as (a0, a1);" +
"B = filter A by a0 > 1;" +
"store B into 'output';";
LogicalPlan lp = Util.parse(query, pc);
Util.optimizeNewLP(lp);
LOStore loStore = (LOStore)lp.getSinks().get(0);
assertEquals("B", loStore.getAlias());
PhysicalPlan pp = Util.buildPhysicalPlanFromNewLP(lp, pc);
POStore poStore = (POStore)pp.getLeaves().get(0);
assertEquals("B", poStore.getAlias());
MROperPlan mrp = Util.buildMRPlanWithOptimizer(pp, pc);
MapReduceOper mrOper = mrp.getLeaves().get(0);
poStore = (POStore)mrOper.mapPlan.getLeaves().get(0);
assertEquals("B", poStore.getAlias());
}
// See PIG-2119
@Test
public void testDanglingNestedNode() throws Exception {
String query = "a = load 'b.txt' AS (id:chararray, num:int); " +
"b = group a by id;" +
"c = foreach b {" +
" d = order a by num DESC;" +
" n = COUNT(a);" +
" e = limit d 1;" +
" generate n;" +
"};";
LogicalPlan lp = Util.parse(query, pc);
Util.optimizeNewLP(lp);
}
public static class SchemaLoader extends PigStorage implements LoadMetadata {
Schema schema;
public SchemaLoader(String schemaString) throws ParserException {
schema = Utils.getSchemaFromString(schemaString);
}
@Override
public ResourceSchema getSchema(String location, Job job)
throws IOException {
return new ResourceSchema(schema);
}
@Override
public ResourceStatistics getStatistics(String location, Job job)
throws IOException {
return null;
}
@Override
public String[] getPartitionKeys(String location, Job job)
throws IOException {
return null;
}
@Override
public void setPartitionFilter(Expression partitionFilter)
throws IOException {
}
}
@Test
public void testLoaderWithSchema() throws Exception {
String query = "a = load 'foo' using " + SchemaLoader.class.getName()
+ "('name,age,gpa');\n"
+ "b = filter a by age==20;"
+ "store b into 'output';";
LogicalPlan lp = Util.parse(query, pc);
Util.optimizeNewLP(lp);
LOLoad loLoad = (LOLoad)lp.getSources().get(0);
LOFilter loFilter = (LOFilter)lp.getSuccessors(loLoad).get(0);
LOStore loStore = (LOStore)lp.getSuccessors(loFilter).get(0);
assertNull(lp.getSuccessors(loStore));
}
public static class PartitionedLoader extends PigStorage implements LoadMetadata {
Schema schema;
String[] partCols;
static Expression partFilter = null;
public PartitionedLoader(String schemaString, String commaSepPartitionCols)
throws ParserException {
schema = Utils.getSchemaFromString(schemaString);
partCols = commaSepPartitionCols.split(",");
}
@Override
public ResourceSchema getSchema(String location, Job job)
throws IOException {
return new ResourceSchema(schema);
}
@Override
public ResourceStatistics getStatistics(String location,
Job job) throws IOException {
return null;
}
@Override
public void setPartitionFilter(Expression partitionFilter)
throws IOException {
partFilter = partitionFilter;
}
@Override
public String[] getPartitionKeys(String location, Job job)
throws IOException {
return partCols;
}
public Expression getPartFilter() {
return partFilter;
}
}
@Test
// See PIG-2339
public void testPartitionFilterOptimizer() throws Exception {
String query = "a = load 'foo' using " + PartitionedLoader.class.getName() +
"('name:chararray, dt:chararray', 'dt');\n" +
"b = filter a by dt=='2011';\n" +
"store b into 'output';";
LogicalPlan lp = Util.parse(query, pc);
Util.optimizeNewLP(lp);
LOLoad loLoad = (LOLoad)lp.getSources().get(0);
LOStore loStore = (LOStore)lp.getSuccessors(loLoad).get(0);
assertNotNull(((PartitionedLoader)loLoad.getLoadFunc()).getPartFilter());
assertEquals("b", loStore.getAlias());
}
}