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

Example 1 with Uncollect

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

the class DrillUnnestRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    final Uncollect uncollect = call.rel(0);
    final LogicalProject project = call.rel(1);
    RexNode projectedNode = project.getProjects().iterator().next();
    if (projectedNode.getKind() != SqlKind.FIELD_ACCESS) {
        return;
    }
    final RelTraitSet traits = uncollect.getTraitSet().plus(DrillRel.DRILL_LOGICAL);
    DrillUnnestRel unnest = new DrillUnnestRel(uncollect.getCluster(), traits, uncollect.getRowType(), projectedNode);
    call.transformTo(unnest);
}

Example 2 with Uncollect

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Uncollect in project beam by apache.

the class ArrayScanToJoinConverter method convert.

/**
 * Returns a LogicJoin.
 */
@Override
public RelNode convert(ResolvedArrayScan zetaNode, List<RelNode> inputs) {
    List<RexNode> projects = new ArrayList<>();
    RelNode leftInput = inputs.get(0);
    ResolvedColumnRef columnRef = (ResolvedColumnRef) zetaNode.getArrayExpr();
    CorrelationId correlationId = getCluster().createCorrel();
    getCluster().getQuery().mapCorrel(correlationId.getName(), leftInput);
    String columnName = String.format("%s%s", zetaNode.getElementColumn().getTableName(), zetaNode.getElementColumn().getName());
    projects.add(getCluster().getRexBuilder().makeFieldAccess(getCluster().getRexBuilder().makeCorrel(leftInput.getRowType(), correlationId), getExpressionConverter().indexOfProjectionColumnRef(columnRef.getColumn().getId(), zetaNode.getInputScan().getColumnList())));
    RelNode projectNode = LogicalProject.create(createOneRow(getCluster()), ImmutableList.of(), projects, ImmutableList.of(columnName));
    // Create an UnCollect
    boolean ordinality = (zetaNode.getArrayOffsetColumn() != null);
    // If they aren't true we need the Project to reorder columns.
    assert zetaNode.getElementColumn().getId() == 1;
    assert !ordinality || zetaNode.getArrayOffsetColumn().getColumn().getId() == 2;
    ZetaSqlUnnest uncollectNode = ZetaSqlUnnest.create(projectNode.getTraitSet(), projectNode, ordinality);
    List<RexInputRef> rightProjects = new ArrayList<>();
    List<String> rightNames = new ArrayList<>();
    rightProjects.add(getCluster().getRexBuilder().makeInputRef(uncollectNode, 0));
    rightNames.add(columnName);
    if (ordinality) {
        rightProjects.add(getCluster().getRexBuilder().makeInputRef(uncollectNode, 1));
        rightNames.add(String.format(zetaNode.getArrayOffsetColumn().getColumn().getTableName(), zetaNode.getArrayOffsetColumn().getColumn().getName()));
    }
    RelNode rightInput = LogicalProject.create(uncollectNode, ImmutableList.of(), rightProjects, rightNames);
    // Join condition should be a RexNode converted from join_expr.
    RexNode condition = getExpressionConverter().convertRexNodeFromResolvedExpr(zetaNode.getJoinExpr());
    JoinRelType joinRelType = zetaNode.getIsOuter() ? JoinRelType.LEFT : JoinRelType.INNER;
    return LogicalJoin.create(leftInput, rightInput, ImmutableList.of(), condition, ImmutableSet.of(), joinRelType, false, ImmutableList.of());
}

Example 3 with Uncollect

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Uncollect in project beam by apache.

the class BeamZetaSqlUnnestRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    LogicalCorrelate correlate = call.rel(0);
    RelNode outer = call.rel(1);
    RelNode uncollect = call.rel(2);
    if (correlate.getRequiredColumns().cardinality() != 1) {
        // can only unnest a single column
        return;
    }
    if (correlate.getJoinType() != JoinRelType.INNER) {
        return;
    }
    if (!(uncollect instanceof ZetaSqlUnnest)) {
        // Drop projection
        uncollect = ((SingleRel) uncollect).getInput();
        if (uncollect instanceof RelSubset) {
            uncollect = ((RelSubset) uncollect).getOriginal();
        }
        if (!(uncollect instanceof ZetaSqlUnnest)) {
            return;
        }
    }
    RelNode project = ((ZetaSqlUnnest) uncollect).getInput();
    if (project instanceof RelSubset) {
        project = ((RelSubset) project).getOriginal();
    }
    if (!(project instanceof LogicalProject)) {
        return;
    }
    if (((LogicalProject) project).getProjects().size() != 1) {
        // can only unnest a single column
        return;
    }
    RexNode exp = ((LogicalProject) project).getProjects().get(0);
    if (!(exp instanceof RexFieldAccess)) {
        return;
    }
    RexFieldAccess fieldAccess = (RexFieldAccess) exp;
    // Innermost field index comes first (e.g. struct.field1.field2 => [2, 1])
    ImmutableList.Builder<Integer> fieldAccessIndices = ImmutableList.builder();
    while (true) {
        fieldAccessIndices.add(fieldAccess.getField().getIndex());
        if (!(fieldAccess.getReferenceExpr() instanceof RexFieldAccess)) {
            break;
        }
        fieldAccess = (RexFieldAccess) fieldAccess.getReferenceExpr();
    }
    call.transformTo(new BeamZetaSqlUnnestRel(correlate.getCluster(), correlate.getTraitSet().replace(BeamLogicalConvention.INSTANCE), convert(outer, outer.getTraitSet().replace(BeamLogicalConvention.INSTANCE)), call.rel(2).getRowType(), fieldAccessIndices.build()));
}

Example 4 with Uncollect

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Uncollect in project beam by apache.

the class BeamUnnestRule method onMatch.

@Override
public void onMatch(RelOptRuleCall call) {
    LogicalCorrelate correlate = call.rel(0);
    RelNode outer = call.rel(1);
    RelNode uncollect = call.rel(2);
    if (correlate.getCorrelationId().getId() != 0) {
        // Only one level of correlation nesting is supported
        return;
    }
    if (correlate.getRequiredColumns().cardinality() != 1) {
        // can only unnest a single column
        return;
    }
    if (correlate.getJoinType() != JoinRelType.INNER) {
        return;
    }
    if (!(uncollect instanceof Uncollect)) {
        // Drop projection
        uncollect = ((SingleRel) uncollect).getInput();
        if (uncollect instanceof RelSubset) {
            uncollect = ((RelSubset) uncollect).getOriginal();
        }
        if (!(uncollect instanceof Uncollect)) {
            return;
        }
    }
    RelNode project = ((Uncollect) uncollect).getInput();
    if (project instanceof RelSubset) {
        project = ((RelSubset) project).getOriginal();
    }
    if (!(project instanceof LogicalProject)) {
        return;
    }
    if (((LogicalProject) project).getProjects().size() != 1) {
        // can only unnest a single column
        return;
    }
    RexNode exp = ((LogicalProject) project).getProjects().get(0);
    if (!(exp instanceof RexFieldAccess)) {
        return;
    }
    RexFieldAccess fieldAccess = (RexFieldAccess) exp;
    // Innermost field index comes first (e.g. struct.field1.field2 => [2, 1])
    ImmutableList.Builder<Integer> fieldAccessIndices = ImmutableList.builder();
    while (true) {
        fieldAccessIndices.add(fieldAccess.getField().getIndex());
        if (!(fieldAccess.getReferenceExpr() instanceof RexFieldAccess)) {
            break;
        }
        fieldAccess = (RexFieldAccess) fieldAccess.getReferenceExpr();
    }
    call.transformTo(new BeamUnnestRel(correlate.getCluster(), correlate.getTraitSet().replace(BeamLogicalConvention.INSTANCE), convert(outer, outer.getTraitSet().replace(BeamLogicalConvention.INSTANCE)), call.rel(2).getRowType(), fieldAccessIndices.build()));
}

Example 5 with Uncollect

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.rel.core.Uncollect in project beam by apache.

the class ArrayScanColumnRefToUncollect method convert.

@Override
public RelNode convert(ResolvedNodes.ResolvedArrayScan zetaNode, List<RelNode> inputs) {
    assert inputs.size() == 1;
    RelNode input = inputs.get(0);
    RexInputRef columnRef = (RexInputRef) getExpressionConverter().convertRexNodeFromResolvedExpr(getColumnRef(zetaNode.getArrayExpr()), zetaNode.getInputScan().getColumnList(), input.getRowType().getFieldList(), ImmutableMap.of());
    CorrelationId correlationId = new CorrelationId(0);
    RexNode convertedColumnRef = getCluster().getRexBuilder().makeFieldAccess(getCluster().getRexBuilder().makeCorrel(input.getRowType(), correlationId), columnRef.getIndex());
    String fieldName = String.format("%s%s", zetaNode.getElementColumn().getTableName(), zetaNode.getElementColumn().getName());
    RelNode projectNode = LogicalProject.create(createOneRow(getCluster()), ImmutableList.of(), Collections.singletonList(convertArrayExpr(zetaNode.getArrayExpr(), getCluster().getRexBuilder(), convertedColumnRef)), ImmutableList.of(fieldName));
    boolean ordinality = (zetaNode.getArrayOffsetColumn() != null);
    RelNode uncollect = ZetaSqlUnnest.create(projectNode.getTraitSet(), projectNode, ordinality);
    return LogicalCorrelate.create(input, uncollect, correlationId, ImmutableBitSet.of(columnRef.getIndex()), JoinRelType.INNER);
}