Examples with RexBuilder org.apache.calcite.rex.RexBuilder used on opensource projects

Example 56 with RexBuilder

use of org.apache.calcite.rex.RexBuilder in project drill by apache.

the class SqlHandlerUtil method qualifyPartitionCol.

/**
   *  Resolve the partition columns specified in "PARTITION BY" clause of CTAS statement.
   *
   *  A partition column is resolved, either (1) the same column appear in the select list of CTAS
   *  or (2) CTAS has a * in select list.
   *
   *  In the second case, a PROJECT with ITEM expression would be created and returned.
   *  Throw validation error if a partition column is not resolved correctly.
   *
   * @param input : the RelNode represents the select statement in CTAS.
   * @param partitionColumns : the list of partition columns.
   * @return : 1) the original RelNode input, if all partition columns are in select list of CTAS
   *           2) a New Project, if a partition column is resolved to * column in select list
   *           3) validation error, if partition column is not resolved.
   */
public static RelNode qualifyPartitionCol(RelNode input, List<String> partitionColumns) {
    final RelDataType inputRowType = input.getRowType();
    final List<RexNode> colRefStarExprs = Lists.newArrayList();
    final List<String> colRefStarNames = Lists.newArrayList();
    final RexBuilder builder = input.getCluster().getRexBuilder();
    final int originalFieldSize = inputRowType.getFieldCount();
    for (final String col : partitionColumns) {
        final RelDataTypeField field = inputRowType.getField(col, false, false);
        if (field == null) {
            throw UserException.validationError().message("Partition column %s is not in the SELECT list of CTAS!", col).build(logger);
        } else {
            if (field.getName().startsWith(StarColumnHelper.STAR_COLUMN)) {
                colRefStarNames.add(col);
                final List<RexNode> operands = Lists.newArrayList();
                operands.add(new RexInputRef(field.getIndex(), field.getType()));
                operands.add(builder.makeLiteral(col));
                final RexNode item = builder.makeCall(SqlStdOperatorTable.ITEM, operands);
                colRefStarExprs.add(item);
            }
        }
    }
    if (colRefStarExprs.isEmpty()) {
        return input;
    } else {
        final List<String> names = new AbstractList<String>() {

            @Override
            public String get(int index) {
                if (index < originalFieldSize) {
                    return inputRowType.getFieldNames().get(index);
                } else {
                    return colRefStarNames.get(index - originalFieldSize);
                }
            }

            @Override
            public int size() {
                return originalFieldSize + colRefStarExprs.size();
            }
        };
        final List<RexNode> refs = new AbstractList<RexNode>() {

            public int size() {
                return originalFieldSize + colRefStarExprs.size();
            }

            public RexNode get(int index) {
                if (index < originalFieldSize) {
                    return RexInputRef.of(index, inputRowType.getFieldList());
                } else {
                    return colRefStarExprs.get(index - originalFieldSize);
                }
            }
        };
        return RelOptUtil.createProject(input, refs, names, false);
    }
}

Example 57 with RexBuilder

use of org.apache.calcite.rex.RexBuilder in project drill by apache.

the class DefaultSqlHandler method addRenamedProject.

protected DrillRel addRenamedProject(DrillRel rel, RelDataType validatedRowType) {
    RelDataType t = rel.getRowType();
    RexBuilder b = rel.getCluster().getRexBuilder();
    List<RexNode> projections = Lists.newArrayList();
    int projectCount = t.getFieldList().size();
    for (int i = 0; i < projectCount; i++) {
        projections.add(b.makeInputRef(rel, i));
    }
    final List<String> fieldNames2 = SqlValidatorUtil.uniquify(validatedRowType.getFieldNames(), SqlValidatorUtil.F_SUGGESTER2, rel.getCluster().getTypeFactory().getTypeSystem().isSchemaCaseSensitive());
    RelDataType newRowType = RexUtil.createStructType(rel.getCluster().getTypeFactory(), projections, fieldNames2);
    DrillProjectRel topProj = DrillProjectRel.create(rel.getCluster(), rel.getTraitSet(), rel, projections, newRowType);
    // Add a final non-trivial Project to get the validatedRowType, if child is not project.
    if (rel instanceof Project && DrillRelOptUtil.isTrivialProject(topProj, true)) {
        return rel;
    } else {
        return topProj;
    }
}

Example 58 with RexBuilder

use of org.apache.calcite.rex.RexBuilder in project drill by apache.

the class DrillExtractConvertlet method convertCall.

/*
   * Custom convertlet to handle extract functions. Optiq rewrites
   * extract functions as divide and modulo functions, based on the
   * data type. We cannot do that in Drill since we don't know the data type
   * till we start scanning. So we don't rewrite extract and treat it as
   * a regular function.
   */
@Override
public RexNode convertCall(SqlRexContext cx, SqlCall call) {
    final RexBuilder rexBuilder = cx.getRexBuilder();
    final List<SqlNode> operands = call.getOperandList();
    final List<RexNode> exprs = new LinkedList<>();
    String timeUnit = ((SqlIntervalQualifier) operands.get(0)).timeUnitRange.toString();
    RelDataTypeFactory typeFactory = cx.getTypeFactory();
    for (SqlNode node : operands) {
        exprs.add(cx.convertExpression(node));
    }
    final RelDataType returnType;
    if (call.getOperator() == SqlStdOperatorTable.EXTRACT) {
        // Legacy code:
        // The return type is wrong!
        // Legacy code choose SqlTypeName.BIGINT simply to avoid conflicting against Calcite's inference mechanism
        // (, which chose BIGINT in validation phase already)
        // Determine NULL-able using 2nd argument's Null-able.
        returnType = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), exprs.get(1).getType().isNullable());
    } else {
        // Determine NULL-able using 2nd argument's Null-able.
        returnType = typeFactory.createTypeWithNullability(typeFactory.createSqlType(TypeInferenceUtils.getSqlTypeNameForTimeUnit(timeUnit)), exprs.get(1).getType().isNullable());
    }
    return rexBuilder.makeCall(returnType, call.getOperator(), exprs);
}

Example 59 with RexBuilder

use of org.apache.calcite.rex.RexBuilder in project drill by apache.

the class SplitUpComplexExpressions method visitProject.

@Override
public Prel visitProject(ProjectPrel project, Object unused) throws RelConversionException {
    // Apply the rule to the child
    RelNode originalInput = ((Prel) project.getInput(0)).accept(this, null);
    project = (ProjectPrel) project.copy(project.getTraitSet(), Lists.newArrayList(originalInput));
    List<RexNode> exprList = new ArrayList<>();
    List<RelDataTypeField> relDataTypes = new ArrayList<>();
    List<RelDataTypeField> origRelDataTypes = new ArrayList<>();
    int i = 0;
    final int lastColumnReferenced = PrelUtil.getLastUsedColumnReference(project.getProjects());
    if (lastColumnReferenced == -1) {
        return project;
    }
    final int lastRexInput = lastColumnReferenced + 1;
    RexVisitorComplexExprSplitter exprSplitter = new RexVisitorComplexExprSplitter(factory, funcReg, lastRexInput);
    for (RexNode rex : project.getChildExps()) {
        origRelDataTypes.add(project.getRowType().getFieldList().get(i));
        i++;
        exprList.add(rex.accept(exprSplitter));
    }
    List<RexNode> complexExprs = exprSplitter.getComplexExprs();
    if (complexExprs.size() == 1 && findTopComplexFunc(project.getChildExps()).size() == 1) {
        return project;
    }
    ProjectPrel childProject;
    List<RexNode> allExprs = new ArrayList<>();
    int exprIndex = 0;
    List<String> fieldNames = originalInput.getRowType().getFieldNames();
    for (int index = 0; index < lastRexInput; index++) {
        RexBuilder builder = new RexBuilder(factory);
        allExprs.add(builder.makeInputRef(new RelDataTypeDrillImpl(new RelDataTypeHolder(), factory), index));
        if (fieldNames.get(index).contains(StarColumnHelper.STAR_COLUMN)) {
            relDataTypes.add(new RelDataTypeFieldImpl(fieldNames.get(index), allExprs.size(), factory.createSqlType(SqlTypeName.ANY)));
        } else {
            relDataTypes.add(new RelDataTypeFieldImpl("EXPR$" + exprIndex, allExprs.size(), factory.createSqlType(SqlTypeName.ANY)));
            exprIndex++;
        }
    }
    RexNode currRexNode;
    int index = lastRexInput - 1;
    // if the projection expressions contained complex outputs, split them into their own individual projects
    if (complexExprs.size() > 0) {
        while (complexExprs.size() > 0) {
            if (index >= lastRexInput) {
                allExprs.remove(allExprs.size() - 1);
                RexBuilder builder = new RexBuilder(factory);
                allExprs.add(builder.makeInputRef(new RelDataTypeDrillImpl(new RelDataTypeHolder(), factory), index));
            }
            index++;
            exprIndex++;
            currRexNode = complexExprs.remove(0);
            allExprs.add(currRexNode);
            relDataTypes.add(new RelDataTypeFieldImpl("EXPR$" + exprIndex, allExprs.size(), factory.createSqlType(SqlTypeName.ANY)));
            childProject = new ProjectPrel(project.getCluster(), project.getTraitSet(), originalInput, ImmutableList.copyOf(allExprs), new RelRecordType(relDataTypes));
            originalInput = childProject;
        }
        // copied from above, find a better way to do this
        allExprs.remove(allExprs.size() - 1);
        RexBuilder builder = new RexBuilder(factory);
        allExprs.add(builder.makeInputRef(new RelDataTypeDrillImpl(new RelDataTypeHolder(), factory), index));
        relDataTypes.add(new RelDataTypeFieldImpl("EXPR$" + index, allExprs.size(), factory.createSqlType(SqlTypeName.ANY)));
    }
    return (Prel) project.copy(project.getTraitSet(), originalInput, exprList, new RelRecordType(origRelDataTypes));
}

Example 60 with RexBuilder

use of org.apache.calcite.rex.RexBuilder in project drill by apache.

the class RewriteProjectToFlatten method visitProject.

@Override
public Prel visitProject(ProjectPrel node, Object unused) throws RelConversionException {
    ProjectPrel project = node;
    List<RexNode> exprList = new ArrayList<>();
    boolean rewrite = false;
    List<RelDataTypeField> relDataTypes = new ArrayList<>();
    int i = 0;
    RexNode flatttenExpr = null;
    for (RexNode rex : project.getChildExps()) {
        RexNode newExpr = rex;
        if (rex instanceof RexCall) {
            RexCall function = (RexCall) rex;
            String functionName = function.getOperator().getName();
            if (functionName.equalsIgnoreCase("flatten")) {
                rewrite = true;
                if (function.getOperands().size() != 1) {
                    throw new RelConversionException("Flatten expression expects a single input.");
                }
                newExpr = function.getOperands().get(0);
                RexBuilder builder = new RexBuilder(factory);
                flatttenExpr = builder.makeInputRef(new RelDataTypeDrillImpl(new RelDataTypeHolder(), factory), i);
            }
        }
        relDataTypes.add(project.getRowType().getFieldList().get(i));
        i++;
        exprList.add(newExpr);
    }
    if (rewrite == true) {
        // TODO - figure out what is the right setting for the traits
        Prel newChild = ((Prel) project.getInput(0)).accept(this, null);
        ProjectPrel newProject = new ProjectPrel(node.getCluster(), project.getTraitSet(), newChild, exprList, new RelRecordType(relDataTypes));
        FlattenPrel flatten = new FlattenPrel(project.getCluster(), project.getTraitSet(), newProject, flatttenExpr);
        return flatten;
    }
    Prel child = ((Prel) project.getInput()).accept(this, null);
    return (Prel) project.copy(project.getTraitSet(), child, exprList, new RelRecordType(relDataTypes));
}