Examples with RelCollation org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation used on opensource projects

Example 76 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project drill by apache.

the class NonCoveringIndexPlanGenerator method convertChild.

@Override
public RelNode convertChild(final RelNode topRel, final RelNode input) throws InvalidRelException {
    if (indexGroupScan == null) {
        logger.error("Null indexgroupScan in NonCoveringIndexPlanGenerator.convertChild");
        return null;
    }
    RelDataType dbscanRowType = convertRowType(origScan.getRowType(), origScan.getCluster().getTypeFactory());
    RelDataType indexScanRowType = FunctionalIndexHelper.convertRowTypeForIndexScan(origScan, indexContext.getOrigMarker(), indexGroupScan, functionInfo);
    DrillDistributionTrait partition = IndexPlanUtils.scanIsPartition(IndexPlanUtils.getGroupScan(origScan)) ? DrillDistributionTrait.RANDOM_DISTRIBUTED : DrillDistributionTrait.SINGLETON;
    ScanPrel indexScanPrel = new ScanPrel(origScan.getCluster(), origScan.getTraitSet().plus(Prel.DRILL_PHYSICAL).plus(partition), indexGroupScan, indexScanRowType, origScan.getTable());
    DbGroupScan origDbGroupScan = (DbGroupScan) IndexPlanUtils.getGroupScan(origScan);
    // right (build) side of the rowkey join: do a distribution of project-filter-indexscan subplan
    RexNode convertedIndexCondition = FunctionalIndexHelper.convertConditionForIndexScan(indexCondition, origScan, indexScanRowType, builder, functionInfo);
    FilterPrel rightIndexFilterPrel = new FilterPrel(indexScanPrel.getCluster(), indexScanPrel.getTraitSet(), indexScanPrel, convertedIndexCondition);
    double finalRowCount = indexGroupScan.getRowCount(indexContext.getOrigCondition(), origScan);
    // project the rowkey column from the index scan
    List<RexNode> rightProjectExprs = Lists.newArrayList();
    // indexGroupScan.getRowKeyOrdinal();
    int rightRowKeyIndex = getRowKeyIndex(indexScanPrel.getRowType(), origScan);
    assert rightRowKeyIndex >= 0;
    rightProjectExprs.add(RexInputRef.of(rightRowKeyIndex, indexScanPrel.getRowType()));
    final List<RelDataTypeField> indexScanFields = indexScanPrel.getRowType().getFieldList();
    final RelDataTypeFactory.FieldInfoBuilder rightFieldTypeBuilder = indexScanPrel.getCluster().getTypeFactory().builder();
    // build the row type for the right Project
    final RelDataTypeField rightRowKeyField = indexScanFields.get(rightRowKeyIndex);
    rightFieldTypeBuilder.add(rightRowKeyField);
    final RelDataType rightProjectRowType = rightFieldTypeBuilder.build();
    final ProjectPrel rightIndexProjectPrel = new ProjectPrel(indexScanPrel.getCluster(), indexScanPrel.getTraitSet(), rightIndexFilterPrel, rightProjectExprs, rightProjectRowType);
    // create a RANGE PARTITION on the right side (this could be removed later during ExcessiveExchangeIdentifier phase
    // if the estimated row count is smaller than slice_target
    final RelNode rangeDistRight = createRangeDistRight(rightIndexProjectPrel, rightRowKeyField, origDbGroupScan);
    // the range partitioning adds an extra column for the partition id but in the final plan we already have a
    // renaming Project for the _id field inserted as part of the JoinPrelRenameVisitor. Thus, we are not inserting
    // a separate Project here.
    final RelNode convertedRight = rangeDistRight;
    // left (probe) side of the rowkey join
    List<SchemaPath> cols = new ArrayList<SchemaPath>(origDbGroupScan.getColumns());
    if (!checkRowKey(cols)) {
        cols.add(origDbGroupScan.getRowKeyPath());
    }
    // Create a restricted groupscan from the primary table's groupscan
    DbGroupScan restrictedGroupScan = (DbGroupScan) origDbGroupScan.getRestrictedScan(cols);
    if (restrictedGroupScan == null) {
        logger.error("Null restricted groupscan in NonCoveringIndexPlanGenerator.convertChild");
        return null;
    }
    // Set left side (restricted scan) row count as rows returned from right side (index scan)
    DrillScanRel rightIdxRel = new DrillScanRel(origScan.getCluster(), origScan.getTraitSet(), origScan.getTable(), origScan.getRowType(), indexContext.getScanColumns());
    double rightIdxRowCount = indexGroupScan.getRowCount(indexCondition, rightIdxRel);
    restrictedGroupScan.setRowCount(null, rightIdxRowCount, rightIdxRowCount);
    RelTraitSet origScanTraitSet = origScan.getTraitSet();
    RelTraitSet restrictedScanTraitSet = origScanTraitSet.plus(Prel.DRILL_PHYSICAL);
    // Create the collation traits for restricted scan based on the index columns under the
    // conditions that (a) the index actually has collation property (e.g hash indexes don't)
    // and (b) if an explicit sort operation is not enforced
    RelCollation collation = null;
    if (indexDesc.getCollation() != null && !settings.isIndexForceSortNonCovering()) {
        collation = IndexPlanUtils.buildCollationNonCoveringIndexScan(indexDesc, indexScanRowType, dbscanRowType, indexContext);
        if (restrictedScanTraitSet.contains(RelCollationTraitDef.INSTANCE)) {
            // replace existing trait
            restrictedScanTraitSet = restrictedScanTraitSet.plus(partition).replace(collation);
        } else {
            // add new one
            restrictedScanTraitSet = restrictedScanTraitSet.plus(partition).plus(collation);
        }
    }
    ScanPrel dbScan = new ScanPrel(origScan.getCluster(), restrictedScanTraitSet, restrictedGroupScan, dbscanRowType, origScan.getTable());
    RelNode lastLeft = dbScan;
    // build the row type for the left Project
    List<RexNode> leftProjectExprs = Lists.newArrayList();
    int leftRowKeyIndex = getRowKeyIndex(dbScan.getRowType(), origScan);
    final RelDataTypeField leftRowKeyField = dbScan.getRowType().getFieldList().get(leftRowKeyIndex);
    final RelDataTypeFactory.FieldInfoBuilder leftFieldTypeBuilder = dbScan.getCluster().getTypeFactory().builder();
    // We are applying the same index condition to primary table's restricted scan. The reason is, the index may be an async
    // index .. i.e it is not synchronously updated along with the primary table update as part of a single transaction, so it
    // is possible that after or during index scan, the primary table rows may have been updated and no longer satisfy the index
    // condition. By re-applying the index condition here, we will ensure non-qualifying records are filtered out.
    // The remainder condition will be applied on top of RowKeyJoin.
    FilterPrel leftIndexFilterPrel = null;
    if (indexDesc.isAsyncIndex()) {
        leftIndexFilterPrel = new FilterPrel(dbScan.getCluster(), dbScan.getTraitSet(), dbScan, indexContext.getOrigCondition());
        lastLeft = leftIndexFilterPrel;
    }
    RelDataType origRowType = origProject == null ? origScan.getRowType() : origProject.getRowType();
    if (origProject != null) {
        // then we also  don't need a project
        // new Project's rowtype is original Project's rowtype [plus rowkey if rowkey is not in original rowtype]
        List<RelDataTypeField> origProjFields = origRowType.getFieldList();
        leftFieldTypeBuilder.addAll(origProjFields);
        // get the exprs from the original Project
        leftProjectExprs.addAll(IndexPlanUtils.getProjects(origProject));
        // add the rowkey IFF rowkey is not in orig scan
        if (getRowKeyIndex(origRowType, origScan) < 0) {
            leftFieldTypeBuilder.add(leftRowKeyField);
            leftProjectExprs.add(RexInputRef.of(leftRowKeyIndex, dbScan.getRowType()));
        }
        final RelDataType leftProjectRowType = leftFieldTypeBuilder.build();
        // build collation in project
        if (!settings.isIndexForceSortNonCovering()) {
            collation = IndexPlanUtils.buildCollationProject(leftProjectExprs, null, dbScan, functionInfo, indexContext);
        }
        final ProjectPrel leftIndexProjectPrel = new ProjectPrel(dbScan.getCluster(), collation != null ? dbScan.getTraitSet().plus(collation) : dbScan.getTraitSet(), leftIndexFilterPrel == null ? dbScan : leftIndexFilterPrel, leftProjectExprs, leftProjectRowType);
        lastLeft = leftIndexProjectPrel;
    }
    final RelTraitSet leftTraits = dbScan.getTraitSet().plus(Prel.DRILL_PHYSICAL);
    // final RelNode convertedLeft = convert(leftIndexProjectPrel, leftTraits);
    final RelNode convertedLeft = Prule.convert(lastLeft, leftTraits);
    // find the rowkey column on the left side of join
    final int leftRowKeyIdx = getRowKeyIndex(convertedLeft.getRowType(), origScan);
    // only rowkey field is being projected from right side
    final int rightRowKeyIdx = 0;
    assert leftRowKeyIdx >= 0;
    List<Integer> leftJoinKeys = ImmutableList.of(leftRowKeyIdx);
    List<Integer> rightJoinKeys = ImmutableList.of(rightRowKeyIdx);
    RexNode joinCondition = RelOptUtil.createEquiJoinCondition(convertedLeft, leftJoinKeys, convertedRight, rightJoinKeys, builder);
    RelNode newRel;
    if (settings.isIndexUseHashJoinNonCovering()) {
        // for hash join, collation will be cleared
        HashJoinPrel hjPrel = new HashJoinPrel(topRel.getCluster(), leftTraits, convertedLeft, convertedRight, joinCondition, JoinRelType.INNER, false, /* no swap */
        null, /* no runtime filter */
        true, /* useful for join-restricted scans */
        JoinControl.DEFAULT);
        newRel = hjPrel;
    } else {
        // if there is collation, add to rowkey join
        RowKeyJoinPrel rjPrel = new RowKeyJoinPrel(topRel.getCluster(), collation != null ? leftTraits.plus(collation) : leftTraits, convertedLeft, convertedRight, joinCondition, JoinRelType.INNER);
        rjPrel.setEstimatedRowCount(finalRowCount);
        newRel = rjPrel;
    }
    final RelDataTypeFactory.FieldInfoBuilder finalFieldTypeBuilder = origScan.getCluster().getTypeFactory().builder();
    List<RelDataTypeField> rjRowFields = newRel.getRowType().getFieldList();
    int toRemoveRowKeyCount = 1;
    if (getRowKeyIndex(origRowType, origScan) < 0) {
        toRemoveRowKeyCount = 2;
    }
    finalFieldTypeBuilder.addAll(rjRowFields.subList(0, rjRowFields.size() - toRemoveRowKeyCount));
    final RelDataType finalProjectRowType = finalFieldTypeBuilder.build();
    List<RexNode> resetExprs = Lists.newArrayList();
    for (int idx = 0; idx < rjRowFields.size() - toRemoveRowKeyCount; ++idx) {
        resetExprs.add(RexInputRef.of(idx, newRel.getRowType()));
    }
    // rewrite the collation for this projectPrel
    final ProjectPrel resetProjectPrel = new ProjectPrel(newRel.getCluster(), newRel.getTraitSet(), newRel, resetExprs, finalProjectRowType);
    newRel = resetProjectPrel;
    if (upperProject != null) {
        RelCollation newCollation = RelCollations.of(RelCollations.EMPTY.getFieldCollations());
        DrillDistributionTrait newDist = null;
        newDist = upperProject.getInput().getTraitSet().getTrait(DrillDistributionTraitDef.INSTANCE);
        if (!settings.isIndexForceSortNonCovering()) {
            newCollation = IndexPlanUtils.buildCollationProject(IndexPlanUtils.getProjects(upperProject), origProject, origScan, functionInfo, indexContext);
        }
        RelTraitSet newProjectTraits = newTraitSet(Prel.DRILL_PHYSICAL, newDist, newCollation);
        ProjectPrel cap = new ProjectPrel(upperProject.getCluster(), newProjectTraits, newRel, IndexPlanUtils.getProjects(upperProject), upperProject.getRowType());
        newRel = cap;
    }
    // whether to remove sort
    if (indexContext.getSort() != null) {
        // unsorted input because ordering is not guaranteed across different parallel inputs.
        if (toRemoveSort(indexContext.getCollation(), newRel.getTraitSet().getTrait(RelCollationTraitDef.INSTANCE))) {
            ((IndexGroupScan) indexScanPrel.getGroupScan()).setParallelizationWidth(1);
        }
        newRel = getSortNode(indexContext, newRel, false, true, true);
        Preconditions.checkArgument(newRel != null);
    }
    RelNode finalRel = Prule.convert(newRel, newRel.getTraitSet());
    logger.debug("NonCoveringIndexPlanGenerator got finalRel {} from origScan {}", finalRel.toString(), origScan.toString());
    return finalRel;
}

Example 77 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project drill by apache.

the class IndexPlanUtils method buildCollationNonCoveringIndexScan.

public static RelCollation buildCollationNonCoveringIndexScan(IndexDescriptor indexDesc, RelDataType indexScanRowType, RelDataType restrictedScanRowType, IndexCallContext context) {
    if (context.getSortExprs() == null) {
        return RelCollations.of(RelCollations.EMPTY.getFieldCollations());
    }
    final List<RelDataTypeField> indexFields = indexScanRowType.getFieldList();
    final List<RelDataTypeField> rsFields = restrictedScanRowType.getFieldList();
    final Map<LogicalExpression, RelFieldCollation> collationMap = indexDesc.getCollationMap();
    assert collationMap != null : "Invalid collation map for index";
    List<RelFieldCollation> fieldCollations = Lists.newArrayList();
    Map<Integer, RelFieldCollation> rsScanCollationMap = Maps.newTreeMap();
    // restricted scan's row type.
    for (int i = 0; i < indexScanRowType.getFieldCount(); i++) {
        RelDataTypeField f1 = indexFields.get(i);
        for (int j = 0; j < rsFields.size(); j++) {
            RelDataTypeField f2 = rsFields.get(j);
            if (f1.getName().equals(f2.getName())) {
                FieldReference ref = FieldReference.getWithQuotedRef(f1.getName());
                RelFieldCollation origCollation = collationMap.get(ref);
                if (origCollation != null) {
                    RelFieldCollation fc = new RelFieldCollation(j, origCollation.direction, origCollation.nullDirection);
                    rsScanCollationMap.put(origCollation.getFieldIndex(), fc);
                }
            }
        }
    }
    // should sort by the order of these fields in indexDesc
    for (Map.Entry<Integer, RelFieldCollation> entry : rsScanCollationMap.entrySet()) {
        RelFieldCollation fc = entry.getValue();
        if (fc != null) {
            fieldCollations.add(fc);
        }
    }
    final RelCollation collation = RelCollations.of(fieldCollations);
    return collation;
}

Example 78 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project drill by apache.

the class IndexPlanUtils method buildCollationForExpressions.

/**
 * Given index, compute the collations for a list of projected expressions(from Scan's rowType or Project's )
 * in the context
 * @param projectExprs the output expression list of a RelNode
 * @param indexDesc  the index for which we are building index plan
 * @param context  the context of this index planning process
 * @return the collation provided by index that will be exposed by the expression list
 */
public static RelCollation buildCollationForExpressions(Map<LogicalExpression, Integer> projectExprs, IndexDescriptor indexDesc, IndexCallContext context) {
    assert projectExprs != null;
    final List<LogicalExpression> sortExpressions = context.getSortExprs();
    // if leading fields of index are here, add them to RelCollation
    List<RelFieldCollation> newFields = Lists.newArrayList();
    if (indexDesc.getCollation() == null) {
        return RelCollations.of(newFields);
    }
    // go through indexed fields to build collation
    // break out of the loop when found first indexed field [not projected && not _only_ in equality condition of filter]
    // or the leading field is not projected
    List<LogicalExpression> indexedCols = indexDesc.getIndexColumns();
    for (int idxFieldCount = 0; idxFieldCount < indexedCols.size(); ++idxFieldCount) {
        LogicalExpression expr = indexedCols.get(idxFieldCount);
        if (!projectExprs.containsKey(expr)) {
            // but it is only-in-equality field, -- we continue to next indexed field, but we don't generate collation for this field
            if (exprOnlyInEquality(expr, context)) {
                continue;
            }
            // else no more collation is needed to be generated, since we now have one leading field which is not in equality condition
            break;
        }
        // and we are okay to continue: generate collation for next indexed field.
        if (sortExpressions != null && !sortExpressions.contains(expr) && exprOnlyInEquality(expr, context)) {
            continue;
        }
        RelCollation idxCollation = indexDesc.getCollation();
        RelFieldCollation.NullDirection nullsDir = idxCollation == null ? RelFieldCollation.NullDirection.UNSPECIFIED : idxCollation.getFieldCollations().get(idxFieldCount).nullDirection;
        RelFieldCollation.Direction dir = (idxCollation == null) ? null : idxCollation.getFieldCollations().get(idxFieldCount).direction;
        if (dir == null) {
            break;
        }
        newFields.add(new RelFieldCollation(projectExprs.get(expr), dir, nullsDir));
    }
    return RelCollations.of(newFields);
}

Example 79 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project drill by apache.

the class IndexSelector method buildAndCheckCollation.

private boolean buildAndCheckCollation(IndexProperties indexProps) {
    IndexDescriptor indexDesc = indexProps.getIndexDesc();
    FunctionalIndexInfo functionInfo = indexDesc.getFunctionalInfo();
    RelCollation inputCollation;
    // for the purpose of collation we can assume that a covering index scan would provide
    // the collation property that would be relevant for non-covering as well
    ScanPrel indexScanPrel = IndexPlanUtils.buildCoveringIndexScan(indexContext.getScan(), indexDesc.getIndexGroupScan(), indexContext, indexDesc);
    inputCollation = indexScanPrel.getTraitSet().getTrait(RelCollationTraitDef.INSTANCE);
    if (indexContext.hasLowerProject()) {
        inputCollation = IndexPlanUtils.buildCollationProject(indexContext.getLowerProject().getProjects(), null, indexContext.getScan(), functionInfo, indexContext);
    }
    if (indexContext.hasUpperProject()) {
        inputCollation = IndexPlanUtils.buildCollationProject(indexContext.getUpperProject().getProjects(), indexContext.getLowerProject(), indexContext.getScan(), functionInfo, indexContext);
    }
    if ((inputCollation != null) && (inputCollation.satisfies(indexContext.getCollation()))) {
        return true;
    }
    return false;
}

Example 80 with RelCollation

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.RelCollation in project drill by apache.

the class CoveringIndexPlanGenerator method convertChild.

@Override
public RelNode convertChild(final RelNode filter, final RelNode input) throws InvalidRelException {
    if (indexGroupScan == null) {
        logger.error("Null indexgroupScan in CoveringIndexPlanGenerator.convertChild");
        return null;
    }
    RexNode coveringCondition;
    ScanPrel indexScanPrel = IndexPlanUtils.buildCoveringIndexScan(origScan, indexGroupScan, indexContext, indexDesc);
    // If remainder condition, then combine the index and remainder conditions. This is a covering plan so we can
    // pushed the entire condition into the index.
    coveringCondition = IndexPlanUtils.getTotalFilter(indexCondition, remainderCondition, indexScanPrel.getCluster().getRexBuilder());
    RexNode newIndexCondition = rewriteFunctionalCondition(coveringCondition, indexScanPrel.getRowType(), functionInfo);
    // build collation for filter
    RelTraitSet indexFilterTraitSet = indexScanPrel.getTraitSet();
    FilterPrel indexFilterPrel = new FilterPrel(indexScanPrel.getCluster(), indexFilterTraitSet, indexScanPrel, newIndexCondition);
    ProjectPrel indexProjectPrel = null;
    if (origProject != null) {
        RelCollation collation = IndexPlanUtils.buildCollationProject(IndexPlanUtils.getProjects(origProject), null, origScan, functionInfo, indexContext);
        indexProjectPrel = new ProjectPrel(origScan.getCluster(), indexFilterTraitSet.plus(collation), indexFilterPrel, IndexPlanUtils.getProjects(origProject), origProject.getRowType());
    }
    RelNode finalRel;
    if (indexProjectPrel != null) {
        finalRel = indexProjectPrel;
    } else {
        finalRel = indexFilterPrel;
    }
    if (upperProject != null) {
        RelCollation newCollation = IndexPlanUtils.buildCollationProject(IndexPlanUtils.getProjects(upperProject), origProject, origScan, functionInfo, indexContext);
        ProjectPrel cap = new ProjectPrel(upperProject.getCluster(), newCollation == null ? finalRel.getTraitSet() : finalRel.getTraitSet().plus(newCollation), finalRel, IndexPlanUtils.getProjects(upperProject), upperProject.getRowType());
        if (functionInfo.hasFunctional()) {
            // if there is functional index field, then a rewrite may be needed in upperProject/indexProject
            // merge upperProject with indexProjectPrel(from origProject) if both exist,
            ProjectPrel newProject = cap;
            if (indexProjectPrel != null) {
                newProject = (ProjectPrel) DrillMergeProjectRule.replace(newProject, indexProjectPrel);
            }
            // then rewrite functional expressions in new project.
            List<RexNode> newProjects = Lists.newArrayList();
            DrillParseContext parseContxt = new DrillParseContext(PrelUtil.getPlannerSettings(newProject.getCluster()));
            for (RexNode projectRex : newProject.getProjects()) {
                RexNode newRex = IndexPlanUtils.rewriteFunctionalRex(indexContext, parseContxt, null, origScan, projectRex, indexScanPrel.getRowType(), functionInfo);
                newProjects.add(newRex);
            }
            ProjectPrel rewrittenProject = new ProjectPrel(newProject.getCluster(), newCollation == null ? newProject.getTraitSet() : newProject.getTraitSet().plus(newCollation), indexFilterPrel, newProjects, newProject.getRowType());
            cap = rewrittenProject;
        }
        finalRel = cap;
    }
    if (indexContext.getSort() != null) {
        finalRel = getSortNode(indexContext, finalRel, false, true, true);
        Preconditions.checkArgument(finalRel != null);
    }
    finalRel = Prule.convert(finalRel, finalRel.getTraitSet().plus(Prel.DRILL_PHYSICAL));
    logger.debug("CoveringIndexPlanGenerator got finalRel {} from origScan {}, original digest {}, new digest {}.", finalRel.toString(), origScan.toString(), upperProject == null ? indexContext.getFilter().getDigest() : upperProject.getDigest(), finalRel.getDigest());
    return finalRel;
}