org.apache.pig.Expression Java Examples

private org.apache.iceberg.expressions.Expression convert(OpType op, Column col, Const constant) {
  String name = col.getName();
  Object value = constant.getValue();

  switch (op) {
    case OP_GE: return Expressions.greaterThanOrEqual(name, value);
    case OP_GT: return Expressions.greaterThan(name, value);
    case OP_LE: return Expressions.lessThanOrEqual(name, value);
    case OP_LT: return Expressions.lessThan(name, value);
    case OP_EQ: return Expressions.equal(name, value);
    case OP_NE: return Expressions.notEqual(name, value);
  }

  throw new RuntimeException(
      String.format("[%s]: Failed to pushdown expression: %s %s %s", signature, col, op, constant));
}
 

/**
 * Return cql where clauses for the corresponding partition filter. Make sure the data format matches
 * Only support the following Pig data types: int, long, float, double, boolean and chararray
 * */
private String partitionFilterToWhereClauseString(Expression expression) throws IOException
{
    Expression.BinaryExpression be = (Expression.BinaryExpression) expression;
    OpType op = expression.getOpType();
    String opString = op.toString();
    switch (op)
    {
        case OP_EQ:
            opString = " = ";
        case OP_GE:
        case OP_GT:
        case OP_LE:
        case OP_LT:
            String name = be.getLhs().toString();
            String value = be.getRhs().toString();
            return String.format("%s %s %s", name, opString, value);
        case OP_AND:
            return String.format("%s AND %s", partitionFilterToWhereClauseString(be.getLhs()), partitionFilterToWhereClauseString(be.getRhs()));
        default:
            throw new IOException("Unsupported expression type: " + opString);
    }
}
 

@Override
public List<Expression.OpType> getSupportedExpressionTypes() {
  OpType supportedTypes [] = {
      OpType.OP_EQ,
      OpType.OP_NE,
      OpType.OP_GT,
      OpType.OP_GE,
      OpType.OP_LT,
      OpType.OP_LE,
      OpType.OP_AND,
      OpType.OP_OR,
      //OpType.OP_BETWEEN, // not implemented in Pig yet
      //OpType.OP_IN,      // not implemented in Pig yet
      OpType.OP_NOT
  };

  return Arrays.asList(supportedTypes);
}
 

@VisibleForTesting
SearchArgument getSearchArgument(Expression expr) {
    if (expr == null) {
        return null;
    }
    Builder builder = SearchArgumentFactory.newBuilder();
    boolean beginWithAnd = !(expr.getOpType().equals(OpType.OP_AND) || expr.getOpType().equals(OpType.OP_OR) || expr.getOpType().equals(OpType.OP_NOT));
    if (beginWithAnd) {
        builder.startAnd();
    }
    buildSearchArgument(expr, builder);
    if (beginWithAnd) {
        builder.end();
    }
    SearchArgument sArg = builder.build();
    return sArg;
}
 

@Test
public void testNegativeMatchesExpr() throws Exception {
    // matches operator is not a supported op type
    String q = query + "b = filter a by name matches 'foo*';" + "store b into 'out';";
    Expression expr = getExpressionForTest(q, Arrays.asList("name"));
    Assert.assertNull(expr);
    SearchArgument sarg = orcStorage.getSearchArgument(expr);
    Assert.assertNull(sarg);

    // AND in LHS/RHS
    q = query + "b = filter a by name matches 'foo*' and srcid == 10;" + "store b into 'out';";
    expr = getExpressionForTest(q, Arrays.asList("srcid", "name"));
    sarg = orcStorage.getSearchArgument(expr);
    assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
            "expr = leaf-0", sarg.toString());

    q = query + "b = filter a by srcid == 10 and name matches 'foo*';" + "store b into 'out';";
    expr = getExpressionForTest(q, Arrays.asList("srcid", "name"));
    sarg = orcStorage.getSearchArgument(expr);
    assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
            "expr = leaf-0", sarg.toString());

    // OR - Nothing should be pushed
    q = query + "b = filter a by name matches 'foo*' or srcid == 10;" + "store b into 'out';";
    expr = getExpressionForTest(q, Arrays.asList("srcid", "name"));
    Assert.assertNull(expr);
    sarg = orcStorage.getSearchArgument(expr);
    Assert.assertNull(sarg);
}
 

@Test
public void testAndOr() throws Exception {
    String q = query + "b = filter a by (srcid > 10 or dstid <= 5) and name == 'foo' and mrkt is null;" + "store b into 'out';";
    Expression expr = getExpressionForTest(q, Arrays.asList("srcid", "dstid", "name", "mrkt"));
    SearchArgument sarg = orcStorage.getSearchArgument(expr);
    assertEquals("leaf-0 = (LESS_THAN_EQUALS srcid 10)\n" +
            "leaf-1 = (LESS_THAN_EQUALS dstid 5)\n" +
            "leaf-2 = (EQUALS name foo)\n" +
            "leaf-3 = (IS_NULL mrkt)\n" +
            "expr = (and (or (not leaf-0) leaf-1) leaf-2 leaf-3)", sarg.toString());
}
 

@Override
   public void setPartitionFilter(Expression partitionFilter)
    throws IOException {
getUDFContext().setProperty(
	PathPartitionHelper.PARITITION_FILTER_EXPRESSION,
	partitionFilter.toString());
   }
 

private String getColumnName(Expression expr) {
    try {
        return ((Column) expr).getName();
    } catch (ClassCastException e) {
        throw new RuntimeException("Expected a Column but found " + expr.getClass().getName() +
                " in expression " + expr, e);
    }
}
 

/** get a list of Cassandra IndexExpression from Pig expression */
private List<IndexExpression> filterToIndexExpressions(Expression expression) throws IOException
{
    List<IndexExpression> indexExpressions = new ArrayList<IndexExpression>();
    Expression.BinaryExpression be = (Expression.BinaryExpression)expression;
    ByteBuffer name = ByteBuffer.wrap(be.getLhs().toString().getBytes());
    ByteBuffer value = ByteBuffer.wrap(be.getRhs().toString().getBytes());
    switch (expression.getOpType())
    {
        case OP_EQ:
            indexExpressions.add(new IndexExpression(name, IndexOperator.EQ, value));
            break;
        case OP_GE:
            indexExpressions.add(new IndexExpression(name, IndexOperator.GTE, value));
            break;
        case OP_GT:
            indexExpressions.add(new IndexExpression(name, IndexOperator.GT, value));
            break;
        case OP_LE:
            indexExpressions.add(new IndexExpression(name, IndexOperator.LTE, value));
            break;
        case OP_LT:
            indexExpressions.add(new IndexExpression(name, IndexOperator.LT, value));
            break;
        case OP_AND:
            indexExpressions.addAll(filterToIndexExpressions(be.getLhs()));
            indexExpressions.addAll(filterToIndexExpressions(be.getRhs()));
            break;
        default:
            throw new IOException("Unsupported expression type: " + expression.getOpType().name());
    }
    return indexExpressions;
}
 

/** set partition filter */
public void setPartitionFilter(Expression partitionFilter) throws IOException
{
    UDFContext context = UDFContext.getUDFContext();
    Properties property = context.getUDFProperties(AbstractCassandraStorage.class);
    property.setProperty(PARTITION_FILTER_SIGNATURE, indexExpressionsToString(filterToIndexExpressions(partitionFilter)));
}
 

private Object getExpressionValue(Expression expr) {
    switch(expr.getOpType()) {
    case TERM_COL:
        return ((Column) expr).getName();
    case TERM_CONST:
        return getSearchArgObjValue(((Const) expr).getValue());
    default:
        throw new RuntimeException("Unsupported expression type: " + expr.getOpType() + " in " + expr);
    }
}
 

@Override
public void setPushdownPredicate(Expression predicate) throws IOException {
  LOG.info(format("[%s]: setPushdownPredicate()", signature));
  LOG.info(format("[%s]: Pig predicate expression: %s", signature, predicate));

  com.netflix.iceberg.expressions.Expression icebergExpression = convert(predicate);

  LOG.info(format("[%s]: Iceberg predicate expression: %s", signature, icebergExpression));

  storeInUDFContext(ICEBERG_FILTER_EXPRESSION, icebergExpression);
}
 

@Test
public void testSimple() throws Exception {
    String q = query + "b = filter a by srcid == 10;" + "store b into 'out';";
    Expression expr = getExpressionForTest(q, Arrays.asList("srcid"));
    SearchArgument sarg = orcStorage.getSearchArgument(expr);
    assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
            "expr = leaf-0", sarg.toString());
}
 

@Override
public void setPushdownPredicate(Expression predicate) throws IOException {
  LOG.info("[{}]: setPushdownPredicate()", signature);
  LOG.info("[{}]: Pig predicate expression: {}", signature, predicate);

  org.apache.iceberg.expressions.Expression icebergExpression = convert(predicate);

  LOG.info("[{}]: Iceberg predicate expression: {}", signature, icebergExpression);

  storeInUDFContext(IcebergPigInputFormat.ICEBERG_FILTER_EXPRESSION, icebergExpression);
}
 

private com.netflix.iceberg.expressions.Expression convert(OpType op, Column col, Const constant) {
  String name = col.getName();
  Object value = constant.getValue();

  switch (op) {
    case OP_GE: return greaterThanOrEqual(name, value);
    case OP_GT: return greaterThan(name, value);
    case OP_LE: return lessThanOrEqual(name, value);
    case OP_LT: return lessThan(name, value);
    case OP_EQ: return equal(name, value);
    case OP_NE: return notEqual(name, value);
  }

  throw new RuntimeException(format("[%s]: Failed to pushdown expression: %s %s %s", signature, col, op, constant));
}
 

@Override
public void setPushdownPredicate(Expression e) throws IOException {
  LOG.info("Pig pushdown expression: {}", e);

  FilterPredicate pred = buildFilter(e);
  LOG.info("Parquet filter predicate expression: {}", pred);

  storeInUDFContext(ParquetInputFormat.FILTER_PREDICATE, pred);
}
 

@Override
public void setPushdownPredicate(Expression expr) throws IOException {
    SearchArgument sArg = getSearchArgument(expr);
    if (sArg != null) {
        log.info("Pushdown predicate expression is " + expr);
        log.info("Pushdown predicate SearchArgument is:\n" + sArg);
        Properties p = UDFContext.getUDFContext().getUDFProperties(this.getClass());
        try {
            p.setProperty(signature + SearchArgsSuffix, sArg.toKryo());
        } catch (Exception e) {
            throw new IOException("Cannot serialize SearchArgument: " + sArg);
        }
    }
}
 

@Test
public void testNot() throws Exception {
    String q = query + "b = filter a by srcid != 10 and mrkt is not null;" + "store b into 'out';";
    Expression expr = getExpressionForTest(q, Arrays.asList("srcid", "dstid", "name", "mrkt"));
    SearchArgument sarg = orcStorage.getSearchArgument(expr);
    assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
            "leaf-1 = (IS_NULL mrkt)\n" +
            "expr = (and (not leaf-0) (not leaf-1))", sarg.toString());
}
 

@Test
public void testBetweenExpression() throws Exception {
    // TODO: Add support for OP_BETWEEN expression type
    String q = query + "b = filter a by srcid > 10 or srcid < 20;" + "store b into 'out';";
    Expression expr = getExpressionForTest(q, Arrays.asList("srcid"));
    SearchArgument sarg = orcStorage.getSearchArgument(expr);
    assertEquals("leaf-0 = (LESS_THAN_EQUALS srcid 10)\n" +
            "leaf-1 = (LESS_THAN srcid 20)\n" +
            "expr = (or (not leaf-0) leaf-1)", sarg.toString());
}
 

@Test
public void testInExpression() throws Exception {
    // TODO: Add support for OP_IN expression type
    String q = query + "b = filter a by srcid == 10 or srcid == 11;" + "store b into 'out';";
    Expression expr = getExpressionForTest(q, Arrays.asList("srcid"));
    SearchArgument sarg = orcStorage.getSearchArgument(expr);
    assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
            "leaf-1 = (EQUALS srcid 11)\n" +
            "expr = (or leaf-0 leaf-1)", sarg.toString());
}
 

@Override
public void transform(OperatorPlan matched) throws FrontendException {
	subPlan = new OperatorSubPlan( currentPlan );

	setupColNameMaps();

    FilterExtractor filterFinder = new PartitionFilterExtractor(loFilter.getFilterPlan(),
            getMappedKeys(partitionKeys));
    filterFinder.visit();
    Expression partitionFilter = filterFinder.getPushDownExpression();

    if(partitionFilter != null) {
        // the column names in the filter may be the ones provided by
        // the user in the schema in the load statement - we may need
        // to replace them with partition column names as given by
        // LoadFunc.getSchema()
        updateMappedColNames(partitionFilter);
        try {
            loadMetadata.setPartitionFilter(partitionFilter);
        } catch (IOException e) {
            throw new FrontendException( e );
        }
        if(filterFinder.isFilterRemovable()) {
            currentPlan.removeAndReconnect( loFilter );
        } else {
            loFilter.setFilterPlan(filterFinder.getFilteredPlan());
        }
    }
}
 

protected void updateMappedColNames(Expression expr) {
    if(expr instanceof BinaryExpression) {
        updateMappedColNames(((BinaryExpression) expr).getLhs());
        updateMappedColNames(((BinaryExpression) expr).getRhs());
    } else if (expr instanceof Column) {
        Column col = (Column) expr;
        col.setName(reverseColNameMap.get(col.getName()));
    }
}
 

@Override
public void transform(OperatorPlan matched) throws FrontendException {
    subPlan = new OperatorSubPlan( currentPlan );

    setupColNameMaps();

    PredicatePushDownFilterExtractor filterFinder = new PredicatePushDownFilterExtractor(
            loFilter.getFilterPlan(), getMappedKeys( predicateFields ), loadPredPushdown.getSupportedExpressionTypes() );
    filterFinder.visit();
    Expression pushDownPredicate = filterFinder.getPushDownExpression();

    if(pushDownPredicate != null) {
        // the column names in the filter may be the ones provided by
        // the user in the schema in the load statement - we may need
        // to replace them with partition column names as given by
        // LoadFunc.getSchema()
        updateMappedColNames(pushDownPredicate);
        try {
            loadPredPushdown.setPushdownPredicate(pushDownPredicate);
        } catch (IOException e) {
            throw new FrontendException( e );
        }

        //TODO: PIG-4093
        /*
        if (loadPredPushdown.removeFilterPredicateFromPlan()) {
            planChanged = true;
            if(filterFinder.isFilterRemovable()) {
                currentPlan.removeAndReconnect( loFilter );
            } else {
                loFilter.setFilterPlan(filterFinder.getFilteredPlan());
            }
        }
        */
    }
}
 

protected void updateMappedColNames(Expression expr) {
    if(expr instanceof BinaryExpression) {
        updateMappedColNames(((BinaryExpression) expr).getLhs());
        updateMappedColNames(((BinaryExpression) expr).getRhs());
    } else if (expr instanceof Column) {
        Column col = (Column) expr;
        col.setName(reverseColNameMap.get(col.getName()));
    }
}
 

private Expression getExpressionForTest(String query, List<String> predicateCols) throws Exception {
    LogicalPlan newLogicalPlan = Util.buildLp(pigServer, query);
    Operator op = newLogicalPlan.getSinks().get(0);
    LOFilter filter = (LOFilter) newLogicalPlan.getPredecessors(op).get(0);
    PredicatePushDownFilterExtractor filterExtractor = new PredicatePushDownFilterExtractor(filter.getFilterPlan(), predicateCols, supportedOpTypes);
    filterExtractor.visit();
    return filterExtractor.getPushDownExpression();
}
 

@Test
public void testUnSupportedFields() throws Exception {
    //Struct, Map and Bag are not supported
    // TODO: Change the test to use ORCStorage to test OrcStorage.getPredicateFields()
    String q = query + "b = filter a by srcid == 10 and browser#'type' == 'IE';" +
            "store b into 'out';";
    Expression expr = getExpressionForTest(q, Arrays.asList("srcid"));
    SearchArgument sarg = orcStorage.getSearchArgument(expr);
    assertEquals("leaf-0 = (EQUALS srcid 10)\n" +
            "expr = leaf-0", sarg.toString());
}
 

@Override
public void setPartitionFilter(Expression partitionFilter)
        throws IOException {
    LOG.debug("PartitionFilter: " + partitionFilter.toString());

    pathPartitionerHelper.setPartitionFilterExpression(
            partitionFilter.toString(), AllLoader.class, signature);

}
 

@Override
public void setPartitionFilter(Expression plan) throws IOException {
    throw new UnsupportedOperationException();
}
 

@Override
public void setPartitionFilter(Expression plan) throws IOException {
  throw new UnsupportedOperationException("setPartitionFilter() not yet supported");
}
 

@Override
public void setPartitionFilter(Expression plan) throws IOException {
    throw new UnsupportedOperationException();
}