Search in sources :

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));
    }
}
Also used : AbstractList(java.util.AbstractList) RelDataTypeField(org.apache.calcite.rel.type.RelDataTypeField) RexBuilder(org.apache.calcite.rex.RexBuilder) RexInputRef(org.apache.calcite.rex.RexInputRef) RelDataType(org.apache.calcite.rel.type.RelDataType) RexNode(org.apache.calcite.rex.RexNode)

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);
}
Also used : RelDataTypeField(org.apache.calcite.rel.type.RelDataTypeField) RexInputRef(org.apache.calcite.rex.RexInputRef) RelDataType(org.apache.calcite.rel.type.RelDataType)

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;
}
Also used : RelDataTypeField(org.apache.calcite.rel.type.RelDataTypeField) PlannerSettings(org.apache.drill.exec.planner.physical.PlannerSettings) RexInputRef(org.apache.calcite.rex.RexInputRef) RelDataType(org.apache.calcite.rel.type.RelDataType)

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);
}
Also used : RelDataType(org.apache.calcite.rel.type.RelDataType) RelTraitSet(org.apache.calcite.plan.RelTraitSet) DrillWindowRel(org.apache.drill.exec.planner.logical.DrillWindowRel) Window(org.apache.calcite.rel.core.Window) SqlAggFunction(org.apache.calcite.sql.SqlAggFunction) RelRecordType(org.apache.calcite.rel.type.RelRecordType) RelCollation(org.apache.calcite.rel.RelCollation) RelDataTypeField(org.apache.calcite.rel.type.RelDataTypeField) RelNode(org.apache.calcite.rel.RelNode) RexInputRef(org.apache.calcite.rex.RexInputRef) RexNode(org.apache.calcite.rex.RexNode)

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;
    }
}
Also used : ProjectPrel(org.apache.drill.exec.planner.physical.ProjectPrel) RexInputRef(org.apache.calcite.rex.RexInputRef) RelDataType(org.apache.calcite.rel.type.RelDataType) ScanPrel(org.apache.drill.exec.planner.physical.ScanPrel) Prel(org.apache.drill.exec.planner.physical.Prel) ProjectAllowDupPrel(org.apache.drill.exec.planner.physical.ProjectAllowDupPrel) ProjectPrel(org.apache.drill.exec.planner.physical.ProjectPrel) ScreenPrel(org.apache.drill.exec.planner.physical.ScreenPrel) WriterPrel(org.apache.drill.exec.planner.physical.WriterPrel) UnnestPrel(org.apache.drill.exec.planner.physical.UnnestPrel) MetadataControllerPrel(org.apache.drill.exec.planner.physical.MetadataControllerPrel) LeafPrel(org.apache.drill.exec.planner.physical.LeafPrel) RexNode(org.apache.calcite.rex.RexNode)

Aggregations

RexInputRef (org.apache.calcite.rex.RexInputRef)241 RexNode (org.apache.calcite.rex.RexNode)200 ArrayList (java.util.ArrayList)105 RelNode (org.apache.calcite.rel.RelNode)85 RelDataTypeField (org.apache.calcite.rel.type.RelDataTypeField)80 RexCall (org.apache.calcite.rex.RexCall)67 RelDataType (org.apache.calcite.rel.type.RelDataType)63 RexBuilder (org.apache.calcite.rex.RexBuilder)54 ImmutableBitSet (org.apache.calcite.util.ImmutableBitSet)52 HashMap (java.util.HashMap)47 AggregateCall (org.apache.calcite.rel.core.AggregateCall)36 List (java.util.List)35 HashSet (java.util.HashSet)32 Pair (org.apache.calcite.util.Pair)32 RexLiteral (org.apache.calcite.rex.RexLiteral)29 Map (java.util.Map)24 RelOptUtil (org.apache.calcite.plan.RelOptUtil)24 Set (java.util.Set)20 ImmutableList (com.google.common.collect.ImmutableList)19 LinkedHashMap (java.util.LinkedHashMap)19