Examples with RexInputRef org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef used on opensource projects

Example 61 with RexInputRef

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef in project drill by axbaretto.

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 (SchemaPath.DYNAMIC_STAR.equals(field.getName())) {
                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>() {

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

            @Override
            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, DrillRelFactories.LOGICAL_BUILDER.create(input.getCluster(), null));
    }
}

Example 62 with RexInputRef

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef in project drill by apache.

the class ConvertHiveParquetScanToDrillParquetScan method createPartitionColumnCast.

/**
 * Create a cast for partition column. Partition column is output as "VARCHAR" in native parquet reader. Cast it
 * appropriate type according the partition type in HiveScan.
 */
private RexNode createPartitionColumnCast(final DrillScanRel hiveScanRel, final DrillScanRel nativeScanRel, final String outputColName, final String dirColName, final RexBuilder rb) {
    final RelDataType outputType = hiveScanRel.getRowType().getField(outputColName, false, false).getType();
    final RelDataTypeField inputField = nativeScanRel.getRowType().getField(dirColName, false, false);
    final RexInputRef inputRef = rb.makeInputRef(rb.getTypeFactory().createSqlType(SqlTypeName.VARCHAR), inputField.getIndex());
    if (outputType.getSqlTypeName() == SqlTypeName.CHAR) {
        return rb.makeCall(RTRIM, inputRef);
    }
    return rb.makeCast(outputType, inputRef);
}

Example 63 with RexInputRef

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef in project drill by apache.

the class ConvertHiveParquetScanToDrillParquetScan method createColumnFormatConversion.

/**
 * Apply any data format conversion expressions.
 */
private RexNode createColumnFormatConversion(DrillScanRel hiveScanRel, DrillScanRel nativeScanRel, String colName, RexBuilder rb) {
    RelDataType outputType = hiveScanRel.getRowType().getField(colName, false, false).getType();
    RelDataTypeField inputField = nativeScanRel.getRowType().getField(colName, false, false);
    RexInputRef inputRef = rb.makeInputRef(inputField.getType(), inputField.getIndex());
    PlannerSettings settings = PrelUtil.getPlannerSettings(hiveScanRel.getCluster().getPlanner());
    boolean conversionToTimestampEnabled = settings.getOptions().getBoolean(ExecConstants.PARQUET_READER_INT96_AS_TIMESTAMP);
    if (outputType.getSqlTypeName() == SqlTypeName.TIMESTAMP && !conversionToTimestampEnabled) {
        // disabled to avoid double conversion after reading value from parquet and here.
        return rb.makeCall(INT96_TO_TIMESTAMP, inputRef);
    }
    return inputRef;
}

Example 64 with RexInputRef

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef in project drill by apache.

the class WindowPrule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final DrillWindowRel window = call.rel(0);
    RelNode input = call.rel(1);
    // TODO: Order window based on existing partition by
    // input.getTraitSet().subsumes()
    boolean partitionby = false;
    boolean addMerge = false;
    // The start index of the constant fields of DrillWindowRel
    final int startConstantsIndex = window.getInput().getRowType().getFieldCount();
    int constantShiftIndex = 0;
    for (final Ord<Window.Group> w : Ord.zip(window.groups)) {
        Window.Group windowBase = w.getValue();
        RelTraitSet traits = call.getPlanner().emptyTraitSet().plus(Prel.DRILL_PHYSICAL);
        // For empty Over-Clause
        if (windowBase.keys.isEmpty() && windowBase.orderKeys.getFieldCollations().isEmpty()) {
            DrillDistributionTrait distEmptyKeys = new DrillDistributionTrait(DrillDistributionTrait.DistributionType.SINGLETON);
            traits = traits.plus(distEmptyKeys);
        } else if (windowBase.keys.size() > 0) {
            DrillDistributionTrait distOnAllKeys = new DrillDistributionTrait(DrillDistributionTrait.DistributionType.HASH_DISTRIBUTED, ImmutableList.copyOf(getDistributionFields(windowBase)));
            partitionby = true;
            traits = traits.plus(distOnAllKeys);
        } else if (windowBase.orderKeys.getFieldCollations().size() > 0) {
            // if only the order-by clause is specified, there is a single partition
            // consisting of all the rows, so we do a distributed sort followed by a
            // single merge as the input of the window operator
            DrillDistributionTrait distKeys = new DrillDistributionTrait(DrillDistributionTrait.DistributionType.HASH_DISTRIBUTED, ImmutableList.copyOf(getDistributionFieldsFromCollation(windowBase)));
            traits = traits.plus(distKeys);
            if (!isSingleMode(call)) {
                addMerge = true;
            }
        }
        // Add collation trait if either partition-by or order-by is specified.
        if (partitionby || windowBase.orderKeys.getFieldCollations().size() > 0) {
            RelCollation collation = getCollation(windowBase);
            traits = traits.plus(collation);
        }
        RelNode convertedInput = convert(input, traits);
        if (addMerge) {
            traits = traits.plus(DrillDistributionTrait.SINGLETON);
            convertedInput = new SingleMergeExchangePrel(window.getCluster(), traits, convertedInput, windowBase.collation());
        }
        List<RelDataTypeField> newRowFields = Lists.newArrayList();
        newRowFields.addAll(convertedInput.getRowType().getFieldList());
        Iterable<RelDataTypeField> newWindowFields = Iterables.filter(window.getRowType().getFieldList(), new Predicate<RelDataTypeField>() {

            @Override
            public boolean apply(RelDataTypeField relDataTypeField) {
                return relDataTypeField.getName().startsWith("w" + w.i + "$");
            }
        });
        for (RelDataTypeField newField : newWindowFields) {
            newRowFields.add(newField);
        }
        RelDataType rowType = new RelRecordType(newRowFields);
        List<Window.RexWinAggCall> newWinAggCalls = Lists.newArrayList();
        for (Ord<Window.RexWinAggCall> aggOrd : Ord.zip(windowBase.aggCalls)) {
            Window.RexWinAggCall aggCall = aggOrd.getValue();
            // If the argument points at the constant and
            // additional fields have been generated by the Window below,
            // the index of constants will be shifted
            final List<RexNode> newOperandsOfWindowFunction = Lists.newArrayList();
            for (RexNode operand : aggCall.getOperands()) {
                if (operand instanceof RexInputRef) {
                    final RexInputRef rexInputRef = (RexInputRef) operand;
                    final int refIndex = rexInputRef.getIndex();
                    // Check if this RexInputRef points at the constants
                    if (rexInputRef.getIndex() >= startConstantsIndex) {
                        operand = new RexInputRef(refIndex + constantShiftIndex, window.constants.get(refIndex - startConstantsIndex).getType());
                    }
                }
                newOperandsOfWindowFunction.add(operand);
            }
            aggCall = new Window.RexWinAggCall((SqlAggFunction) aggCall.getOperator(), aggCall.getType(), newOperandsOfWindowFunction, aggCall.ordinal, aggCall.distinct);
            newWinAggCalls.add(new Window.RexWinAggCall((SqlAggFunction) aggCall.getOperator(), aggCall.getType(), aggCall.getOperands(), aggOrd.i, aggCall.distinct));
        }
        windowBase = new Window.Group(windowBase.keys, windowBase.isRows, windowBase.lowerBound, windowBase.upperBound, windowBase.orderKeys, newWinAggCalls);
        input = new WindowPrel(window.getCluster(), window.getTraitSet().merge(traits), convertedInput, window.getConstants(), rowType, windowBase);
        constantShiftIndex += windowBase.aggCalls.size();
    }
    call.transformTo(input);
}

Example 65 with RexInputRef

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rex.RexInputRef in project drill by apache.

the class StarColumnConverter method visitProject.

@Override
public Prel visitProject(ProjectPrel prel, Void value) throws RuntimeException {
    // Require prefix rename : there exists other expression, in addition to a star column.
    if (// not set yet.
    !prefixedForStar && StarColumnHelper.containsStarColumnInProject(prel.getInput().getRowType(), prel.getProjects()) && prel.getRowType().getFieldNames().size() > 1) {
        prefixedForStar = true;
    }
    // For project, we need make sure that the project's field name is same as the input,
    // when the project expression is RexInPutRef, since we may insert a PAS which will
    // rename the projected fields.
    Prel child = ((Prel) prel.getInput(0)).accept(this, null);
    List<String> fieldNames = Lists.newArrayList();
    for (Pair<String, RexNode> pair : Pair.zip(prel.getRowType().getFieldNames(), prel.getProjects())) {
        if (pair.right instanceof RexInputRef) {
            String name = child.getRowType().getFieldNames().get(((RexInputRef) pair.right).getIndex());
            fieldNames.add(name);
        } else {
            fieldNames.add(pair.left);
        }
    }
    // Make sure the field names are unique : no allow of duplicate field names in a rowType.
    fieldNames = makeUniqueNames(fieldNames);
    RelDataType rowType = RexUtil.createStructType(prel.getCluster().getTypeFactory(), prel.getProjects(), fieldNames, null);
    ProjectPrel newProj = (ProjectPrel) prel.copy(prel.getTraitSet(), child, prel.getProjects(), rowType);
    if (ProjectRemoveRule.isTrivial(newProj)) {
        return child;
    } else {
        return newProj;
    }
}