Java Code Examples for org.apache.pig.newplan.logical.relational.LogicalSchema#getField()

@Test
public void testSchemaPatcher() throws FrontendException {
    SchemaPatcher patcher = new SchemaPatcher();
    patcher.transformed(lp, changedPlan);

    // Check that the filter now has the proper schema.
    List<Operator> roots = changedPlan.getSources();
    assertEquals(1, roots.size());
    LOFilter D = (LOFilter)roots.get(0);
    assertNotNull(D);
    LogicalSchema dschema = D.getSchema();
    assertEquals(1, dschema.size());
    LogicalSchema.LogicalFieldSchema y = dschema.getField(0);
    assertEquals("y", y.alias);
    assertEquals(2, y.uid);
}
 

@Test    
// See PIG-767
public void testCogroupSchema1() throws Exception {
    String query = ("a = load '1.txt' as (a0, a1);" +
    "b = group a by a0;" +
    "store b into 'empty';");  
    
    LogicalPlan newLogicalPlan = buildPlan(query);
    Operator store = newLogicalPlan.getSinks().get(0);
    LOCogroup cogroup = (LOCogroup)newLogicalPlan.getPredecessors(store).get(0);
    
    LogicalSchema cogroupSchema = cogroup.getSchema();
    assertEquals(cogroupSchema.getField(1).type, DataType.BAG);
    assertTrue(cogroupSchema.getField(1).alias.equals("a"));
    LogicalSchema bagSchema = cogroupSchema.getField(1).schema;
    assertEquals(bagSchema.getField(0).type, DataType.TUPLE);
    assertEquals(bagSchema.getField(0).alias, null);
    LogicalSchema tupleSchema = bagSchema.getField(0).schema;
    assertEquals(tupleSchema.size(), 2);
}
 

@Test    
// See PIG-767
public void testCogroupSchema2() throws Exception {
    String query = "a = load '1.txt' as (a0, a1);" +
    "b = group a by a0;" +
    "c = foreach b generate a.a1;" +
    "store c into 'empty';";  
    
    LogicalPlan newLogicalPlan = buildPlan(query);
    Operator store = newLogicalPlan.getSinks().get(0);
    LOForEach foreach = (LOForEach)newLogicalPlan.getPredecessors(store).get(0);
    
    LogicalSchema foreachSchema = foreach.getSchema();
    assertEquals(foreachSchema.getField(0).type, DataType.BAG);
    LogicalSchema bagSchema = foreachSchema.getField(0).schema;
    assertEquals(bagSchema.getField(0).type, DataType.TUPLE);
    assertEquals(bagSchema.getField(0).alias, null);
    LogicalSchema tupleSchema = bagSchema.getField(0).schema;
    assertEquals(tupleSchema.size(), 1);
    assertTrue(tupleSchema.getField(0).alias.equals("a1"));
}
 

@Test
public void testNewNormalNestedMerge1() throws Exception {
    LogicalSchema a = org.apache.pig.newplan.logical.Util.translateSchema(Utils.getSchemaFromString(
        "a1:bytearray, b1:(b11:int, b12:float), c1:long"));
    LogicalSchema b = org.apache.pig.newplan.logical.Util.translateSchema(Utils.getSchemaFromString(
        "a2:bytearray, b2:(b21:double, b22:long), c2:int"));

    LogicalSchema mergedSchema = LogicalSchema.merge(a, b, LogicalSchema.MergeMode.Union);
    LogicalSchema expected = org.apache.pig.newplan.logical.Util.translateSchema(Utils.getSchemaFromString(
        "a1:bytearray, b1:(), c1:long"));
    expected.getField(1).schema = new LogicalSchema();
    assertTrue(LogicalSchema.equals(mergedSchema, expected, false, false));

    mergedSchema = LogicalSchema.merge(a, b, LogicalSchema.MergeMode.LoadForEach);
    expected = org.apache.pig.newplan.logical.Util.translateSchema(Utils.getSchemaFromString(
        "a1:bytearray, b1:(b11:int, b12:float), c1:long"));
    assertTrue(LogicalSchema.equals(mergedSchema, expected, false, false));

    mergedSchema = LogicalSchema.merge(b, a, LogicalSchema.MergeMode.LoadForEach);
    expected = org.apache.pig.newplan.logical.Util.translateSchema(Utils.getSchemaFromString(
        "a2:bytearray, b2:(b21:double, b22:long), c2:int"));
    assertTrue(LogicalSchema.equals(mergedSchema, expected, false, false));
}
 

@Test
public void testNewNormalNestedMerge2() throws Exception {
    LogicalSchema a = org.apache.pig.newplan.logical.Util.translateSchema(Utils.getSchemaFromString(
        "a1:(a11:chararray, a12:float), b1:(b11:chararray, b12:float), c1:long"));
    LogicalSchema b = org.apache.pig.newplan.logical.Util.translateSchema(Utils.getSchemaFromString(
        "a2:bytearray, b2:(b21:double, b22:long), c2:chararray"));

    LogicalSchema mergedSchema = LogicalSchema.merge(a, b, LogicalSchema.MergeMode.Union);
    LogicalSchema expected = org.apache.pig.newplan.logical.Util.translateSchema(Utils.getSchemaFromString(
        "a1:(a11:chararray, a12:float), b1:(), c1:bytearray"));
    expected.getField(1).schema = new LogicalSchema();
    assertTrue(LogicalSchema.equals(mergedSchema, expected, false, false));

    mergedSchema = LogicalSchema.merge(a, b, LogicalSchema.MergeMode.LoadForEach);
    expected = org.apache.pig.newplan.logical.Util.translateSchema(Utils.getSchemaFromString(
        "a1:(a11:chararray, a12:float), b1:(b11:chararray, b12:float), c1:long"));
    assertTrue(LogicalSchema.equals(mergedSchema, expected, false, false));

    mergedSchema = LogicalSchema.merge(b, a, LogicalSchema.MergeMode.LoadForEach);
    expected = org.apache.pig.newplan.logical.Util.translateSchema(Utils.getSchemaFromString(
            "a2:(a11:chararray, a12:float), b2:(b21:double, b22:long), c2:chararray"));
    assertTrue(LogicalSchema.equals(mergedSchema, expected, false, false));
}
 

@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 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);
}
 

/**
 * {@inheritDoc}.
 */
@Override
public Type toMetacatType(@Nonnull @NonNull final String pigType) {
    try {
        final LogicalSchema schema = new QueryParserDriver(PIG_CONTEXT,
            "util", new HashMap<>()).parseSchema(pigType);
        final LogicalSchema.LogicalFieldSchema field = schema.getField(0);
        return toCanonicalType(field);
    } catch (Exception e) {
        throw new IllegalArgumentException(String.format("Invalid type signature: '%s'", pigType));
    }
}
 

/**
 * Converts a Pig tuple schema to a SQL row type.
 *
 * @param pigSchema Pig tuple schema
 * @param nullable true if the type is nullable
 * @return a SQL row type
 */
static RelDataType convertSchema(LogicalSchema pigSchema, boolean nullable) {
  if (pigSchema != null && pigSchema.size() > 0) {
    List<String> fieldNameList = new ArrayList<>();
    List<RelDataType> typeList = new ArrayList<>();
    for (int i = 0; i < pigSchema.size(); i++) {
      final LogicalSchema.LogicalFieldSchema subPigField = pigSchema.getField(i);
      fieldNameList.add(subPigField.alias != null ? subPigField.alias : "$" + i);
      typeList.add(convertSchemaField(subPigField, nullable));
    }
    return TYPE_FACTORY.createStructType(typeList, fieldNameList, nullable);
  }
  return new DynamicTupleRecordType(TYPE_FACTORY);
}
 

private boolean atLeastOneCastNeeded(LogicalSchema determinedSchema, LogicalSchema s) {
    for (int i = 0; i < s.size(); i++) {
        LogicalSchema.LogicalFieldSchema fs = s.getField(i);
        if (fs.type != DataType.BYTEARRAY && (determinedSchema == null || (!fs.isEqual(determinedSchema.getField(i))))) {
            // we have to cast this field from the default BYTEARRAY type to
            // whatever the user specified in the 'AS' clause of the LOAD
            // statement (the fs.type).
            return true;
        }
    }
    return false;
}
 

@Override
public void visit(LOInnerLoad innerLoad) throws FrontendException{
    LOForEach foreach = innerLoad.getLOForEach();
    LogicalRelationalOperator pred = 
        ((LogicalRelationalOperator)foreach.getPlan().getPredecessors(foreach).get(0));

    LogicalSchema predSchema = pred.getSchema();
    
    //if this has a schema, the lineage can be tracked using the uid in input schema
    if(innerLoad.getSchema() != null){
        if(innerLoad.getSchema().size() == 1 
                && innerLoad.getSchema().getField(0).type == DataType.BYTEARRAY
                && uid2LoadFuncMap.get(innerLoad.getSchema().getField(0).uid) == null
                && predSchema != null
        ){
            long inpUid = predSchema.getField(innerLoad.getProjection().getColNum()).uid;
            if(uid2LoadFuncMap.get(inpUid) != null){
                addUidLoadFuncToMap(innerLoad.getSchema().getField(0).uid, uid2LoadFuncMap.get(inpUid));
            }
            return;
        }

    }

    // associated load func could not be found using uid, use
    // the single load func associated with input relation (if any)
    if(getAssociatedLoadFunc(pred) != null){
        mapRelToPredLoadFunc(innerLoad, pred);
    }
}
 

/**
 * Find single load func spec associated with this relation.
 * If the relation has schema, all uids in schema should be associated
 * with same load func spec. if it does not have schema check the existing
 * mapping
 * @param relOp
 * @return
 * @throws FrontendException
 */
private FuncSpec getAssociatedLoadFunc(LogicalRelationalOperator relOp) throws FrontendException {
    LogicalSchema schema = relOp.getSchema();
    FuncSpec funcSpec = null;
    if(schema != null){
        if(schema.size() == 0)
            return null;
        funcSpec = uid2LoadFuncMap.get(schema.getField(0).uid);
        if(funcSpec != null) {
            for(int i=1; i<schema.size(); i++){
                LogicalFieldSchema fs = schema.getField(i);
                if(! haveIdenticalCasters(funcSpec,
                        uid2LoadFuncMap.get(fs.uid))){
                    //all uid are not associated with same func spec, there is no
                    // single func spec that represents all the fields
                    funcSpec = null;
                    break;
                }
            }
        }
    }
    
    if(funcSpec == null){
        // If relOp is LOForEach and contains UDF, byte field could come from UDF.
        // We don't assume it share the LoadCaster with predecessor
        if (relOp instanceof LOForEach) {
            UDFFinder udfFinder = new UDFFinder(((LOForEach) relOp).getInnerPlan());
            udfFinder.visit();
            if (udfFinder.getUDFList().size()!=0)
                return null;
        }
        
        funcSpec = rel2InputFuncMap.get(relOp);
    }

    return funcSpec;
}
 

@Before
public void setUp() throws Exception {
    // Build a plan that looks like it has just been transformed
    // It is roughly the logical plan for
    // A = load 'bla' as (x);
    // B = load 'morebla' as (y);
    // C = join A on x, B on y;
    // D = filter C by y > 0;
    // The plan is built with the filter having been pushed above the join
    // but the listners not yet having been called.
    // A = load
    lp = new LogicalPlan();
    LogicalSchema aschema = new LogicalSchema();
    aschema.addField(new LogicalSchema.LogicalFieldSchema(
        "x", null, DataType.INTEGER));
    aschema.getField(0).uid = 1;
    LOLoad A = new LOLoad(null, lp);
    A.neverUseForRealSetSchema(aschema);
    lp.add(A);

    // B = load
    LogicalSchema bschema = new LogicalSchema();
    bschema.addField(new LogicalSchema.LogicalFieldSchema(
        "y", null, DataType.INTEGER));
    bschema.getField(0).uid = 2;
    LOLoad B = new LOLoad(null, lp);
    B.neverUseForRealSetSchema(bschema);
    lp.add(B);

    // C = join
    LogicalSchema cschema = new LogicalSchema();
    cschema.addField(new LogicalSchema.LogicalFieldSchema(
        "x", null, DataType.INTEGER));
    cschema.addField(new LogicalSchema.LogicalFieldSchema(
        "y", null, DataType.INTEGER));
    cschema.getField(0).uid = 1;
    cschema.getField(1).uid = 2;

    MultiMap<Integer, LogicalExpressionPlan> mm =
        new MultiMap<Integer, LogicalExpressionPlan>();
    LOJoin C = new LOJoin(lp, mm, JOINTYPE.HASH, new boolean[] {true, true});

    LogicalExpressionPlan aprojplan = new LogicalExpressionPlan();
    ProjectExpression x = new ProjectExpression(aprojplan, 0, 0, C);
    x.neverUseForRealSetFieldSchema(new LogicalFieldSchema(null, null,
            DataType.INTEGER, 1));
    LogicalExpressionPlan bprojplan = new LogicalExpressionPlan();
    ProjectExpression y = new ProjectExpression(bprojplan, 1, 0, C);
    y.neverUseForRealSetFieldSchema(new LogicalFieldSchema(null, null,
            DataType.INTEGER, 2));
    mm.put(0, aprojplan);
    mm.put(1, bprojplan);

    C.neverUseForRealSetSchema(cschema);
    // Don't add it to the plan quite yet

    // D = filter
    LogicalExpressionPlan filterPlan = new LogicalExpressionPlan();
    LOFilter D = new LOFilter(lp, filterPlan);
    ProjectExpression fy = new ProjectExpression(filterPlan, 0, 1, D);
    fy.neverUseForRealSetFieldSchema(new LogicalFieldSchema(null, null,
            DataType.INTEGER, 2));
    ConstantExpression fc = new ConstantExpression(filterPlan, new Integer(0));
    new EqualExpression(filterPlan, fy, fc);

    D.neverUseForRealSetSchema(cschema);
    // Connect D to B, since the transform has happened.
    lp.add(D);
    lp.connect(B, D);

    // Now add in C, connected to A and D.
    lp.add(C);
    lp.connect(A, C);
    lp.connect(D, C);

    changedPlan = new LogicalPlan();
    changedPlan.add(D);
    changedPlan.add(C);
    changedPlan.connect(D, C);
}