mihaibudiu commented on code in PR #4077:
URL: https://github.com/apache/calcite/pull/4077#discussion_r1870540984
##########
core/src/test/java/org/apache/calcite/sql2rel/RelFieldTrimmerTest.java:
##########
@@ -544,4 +548,36 @@ public static Frameworks.ConfigBuilder config() {
+ " LogicalTableScan(table=[[scott, EMP]])\n";
assertThat(trimmed, hasTree(expected));
}
+
+ @Test void testLogicalCorrelateFieldTrimmer() {
+ final RelBuilder builder = RelBuilder.create(config().build());
+ final Holder<@Nullable RexCorrelVariable> v = Holder.empty();
+ RelNode root = builder.scan("EMP")
+ .projectPlus(builder.call(SqlStdOperatorTable.PLUS, builder.field(0),
builder.field(0)))
+ .variable(v::set)
+ .scan("DEPT")
+ .filter(
+ builder.equals(builder.field(0),
+ builder.call(
+ SqlStdOperatorTable.PLUS,
+ builder.literal(10),
+ builder.field(v.get(), "DEPTNO"))))
+ .correlate(JoinRelType.LEFT, v.get().id, builder.field(2, 0, "DEPTNO"))
+ .aggregate(builder.groupKey("ENAME"),
builder.max(builder.field("EMPNO")))
+ .build();
+
+ final RelFieldTrimmer fieldTrimmer = new RelFieldTrimmer(null, builder);
+ final RelNode trimmed = fieldTrimmer.trim(root);
+ final String expected = ""
Review Comment:
a bit difficult to see what the trimmer has done.
perhaps if you have the expected tree before the trim it will be easier?
then you can also assert that they are different
##########
core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java:
##########
@@ -1336,6 +1338,132 @@ public TrimResult trimFields(
return result(newSnapshot, inputMapping, snapshot);
}
+ /**
+ * Trims {@link LogicalCorrelate} nodes.
+ */
+ public TrimResult trimFields(LogicalCorrelate correlate,
+ ImmutableBitSet fieldsUsed, Set<RelDataTypeField> extraFields) {
+ if (!extraFields.isEmpty()) {
+ // bail out with generic trim
+ return trimFields((RelNode) correlate, fieldsUsed, extraFields);
+ }
+
+ fieldsUsed = fieldsUsed.union(correlate.getRequiredColumns());
+
+ List<RelNode> newInputs = new ArrayList<>();
+ List<Mapping> inputMappings = new ArrayList<>();
+ int changeCount = 0;
+ int offset = 0;
+ for (RelNode input : correlate.getInputs()) {
+ final RelDataType inputRowType = input.getRowType();
+ final int inputFieldCount = inputRowType.getFieldCount();
+
+ ImmutableBitSet currentInputFieldsUsed = fieldsUsed
+ .intersect(ImmutableBitSet.range(offset, offset + inputFieldCount))
+ .shift(-offset);
+
+ TrimResult trimResult =
+ dispatchTrimFields(input, currentInputFieldsUsed, extraFields);
+
+ newInputs.add(trimResult.left);
+ inputMappings.add(trimResult.right);
+
+ offset += inputFieldCount;
+
+ if (trimResult.left != input) {
+ changeCount++;
+ }
+ }
+
+ if (changeCount == 0) {
+ return result(correlate,
+ Mappings.createIdentity(correlate.getRowType().getFieldCount()));
+ }
+
+ Mapping mapping = makeMapping(inputMappings);
+ RexBuilder rexBuilder = relBuilder.getRexBuilder();
+
+ RexCorrelVariableMapShuttle rexVisitor =
+ new RexCorrelVariableMapShuttle(correlate.getCorrelationId(),
+ newInputs.get(0).getRowType(), mapping, rexBuilder);
+ final LogicalCorrelate newCorrelate =
+ correlate
+ .copy(correlate.getTraitSet(),
+ newInputs.get(0),
+ newInputs.get(1).accept(new
RexRewritingRelShuttle(rexVisitor)),
+ correlate.getCorrelationId(),
+ correlate.getRequiredColumns().permute(mapping),
+ correlate.getJoinType());
+
+ return result(newCorrelate, mapping);
+ }
+
+ /**
+ * Updates correlate references in {@link RexNode} expressions.
+ */
+ static class RexCorrelVariableMapShuttle extends RexShuttle {
+ private final CorrelationId correlationId;
+ private final Mapping mapping;
+ private final RelDataType newCorrelRowType;
+ private final RexBuilder rexBuilder;
+
+ RexCorrelVariableMapShuttle(final CorrelationId correlationId,
+ RelDataType newCorrelRowType, Mapping mapping, RexBuilder rexBuilder) {
+ this.correlationId = correlationId;
+ this.newCorrelRowType = newCorrelRowType;
+ this.mapping = mapping;
+ this.rexBuilder = rexBuilder;
+ }
+
+ @Override public RexNode visitFieldAccess(final RexFieldAccess
fieldAccess) {
+ if (fieldAccess.getReferenceExpr() instanceof RexCorrelVariable) {
+ RexCorrelVariable referenceExpr =
+ (RexCorrelVariable) fieldAccess.getReferenceExpr();
+ if (referenceExpr.id.equals(correlationId)) {
+ int oldIndex = fieldAccess.getField().getIndex();
+ RexNode newCorrel =
+ rexBuilder.makeCorrel(newCorrelRowType, referenceExpr.id);
+ int newIndex = mapping.getTarget(oldIndex);
+ return rexBuilder.makeFieldAccess(newCorrel, newIndex);
+ }
+ }
+ return super.visitFieldAccess(fieldAccess);
+ }
+ }
+
+ /**
+ * Concatenates multiple mappings.
+ *
+ * <pre>
+ * [ 1:0, 2:1] // sourceCount:100
+ * [ 1:0, 2:1] // sourceCount:100
+ * output:
+ * [ 1:0, 2:1, 101:2, 102:3 ] ; sourceCount:200
+ * </pre>
+ */
+ private Mapping makeMapping(List<Mapping> inputMappings) {
Review Comment:
concatenateMappings is a more suggesting name.
##########
core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java:
##########
@@ -1336,6 +1338,132 @@ public TrimResult trimFields(
return result(newSnapshot, inputMapping, snapshot);
}
+ /**
+ * Trims {@link LogicalCorrelate} nodes.
+ */
+ public TrimResult trimFields(LogicalCorrelate correlate,
+ ImmutableBitSet fieldsUsed, Set<RelDataTypeField> extraFields) {
+ if (!extraFields.isEmpty()) {
+ // bail out with generic trim
+ return trimFields((RelNode) correlate, fieldsUsed, extraFields);
+ }
+
+ fieldsUsed = fieldsUsed.union(correlate.getRequiredColumns());
+
+ List<RelNode> newInputs = new ArrayList<>();
+ List<Mapping> inputMappings = new ArrayList<>();
+ int changeCount = 0;
+ int offset = 0;
+ for (RelNode input : correlate.getInputs()) {
+ final RelDataType inputRowType = input.getRowType();
+ final int inputFieldCount = inputRowType.getFieldCount();
+
+ ImmutableBitSet currentInputFieldsUsed = fieldsUsed
+ .intersect(ImmutableBitSet.range(offset, offset + inputFieldCount))
+ .shift(-offset);
+
+ TrimResult trimResult =
+ dispatchTrimFields(input, currentInputFieldsUsed, extraFields);
+
+ newInputs.add(trimResult.left);
+ inputMappings.add(trimResult.right);
+
+ offset += inputFieldCount;
+
+ if (trimResult.left != input) {
+ changeCount++;
+ }
+ }
+
+ if (changeCount == 0) {
+ return result(correlate,
+ Mappings.createIdentity(correlate.getRowType().getFieldCount()));
+ }
+
+ Mapping mapping = makeMapping(inputMappings);
+ RexBuilder rexBuilder = relBuilder.getRexBuilder();
+
+ RexCorrelVariableMapShuttle rexVisitor =
+ new RexCorrelVariableMapShuttle(correlate.getCorrelationId(),
+ newInputs.get(0).getRowType(), mapping, rexBuilder);
+ final LogicalCorrelate newCorrelate =
+ correlate
+ .copy(correlate.getTraitSet(),
+ newInputs.get(0),
+ newInputs.get(1).accept(new
RexRewritingRelShuttle(rexVisitor)),
Review Comment:
this also assumes implicitly that newInputs.size() == 2; perhaps this is
obvious, but then it would not harm to assert that this holds for the correlate
##########
core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java:
##########
@@ -1336,6 +1338,132 @@ public TrimResult trimFields(
return result(newSnapshot, inputMapping, snapshot);
}
+ /**
+ * Trims {@link LogicalCorrelate} nodes.
+ */
+ public TrimResult trimFields(LogicalCorrelate correlate,
+ ImmutableBitSet fieldsUsed, Set<RelDataTypeField> extraFields) {
+ if (!extraFields.isEmpty()) {
+ // bail out with generic trim
+ return trimFields((RelNode) correlate, fieldsUsed, extraFields);
+ }
+
+ fieldsUsed = fieldsUsed.union(correlate.getRequiredColumns());
+
+ List<RelNode> newInputs = new ArrayList<>();
+ List<Mapping> inputMappings = new ArrayList<>();
+ int changeCount = 0;
+ int offset = 0;
+ for (RelNode input : correlate.getInputs()) {
+ final RelDataType inputRowType = input.getRowType();
+ final int inputFieldCount = inputRowType.getFieldCount();
+
+ ImmutableBitSet currentInputFieldsUsed = fieldsUsed
+ .intersect(ImmutableBitSet.range(offset, offset + inputFieldCount))
+ .shift(-offset);
+
+ TrimResult trimResult =
+ dispatchTrimFields(input, currentInputFieldsUsed, extraFields);
+
+ newInputs.add(trimResult.left);
+ inputMappings.add(trimResult.right);
+
+ offset += inputFieldCount;
+
+ if (trimResult.left != input) {
+ changeCount++;
+ }
+ }
+
+ if (changeCount == 0) {
+ return result(correlate,
+ Mappings.createIdentity(correlate.getRowType().getFieldCount()));
+ }
+
+ Mapping mapping = makeMapping(inputMappings);
+ RexBuilder rexBuilder = relBuilder.getRexBuilder();
+
+ RexCorrelVariableMapShuttle rexVisitor =
+ new RexCorrelVariableMapShuttle(correlate.getCorrelationId(),
+ newInputs.get(0).getRowType(), mapping, rexBuilder);
+ final LogicalCorrelate newCorrelate =
+ correlate
+ .copy(correlate.getTraitSet(),
+ newInputs.get(0),
+ newInputs.get(1).accept(new
RexRewritingRelShuttle(rexVisitor)),
Review Comment:
I find it easier to debug if I avoid complex expressions - I would use a
temporary for the result of the visitor.
##########
core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java:
##########
@@ -1336,6 +1338,132 @@ public TrimResult trimFields(
return result(newSnapshot, inputMapping, snapshot);
}
+ /**
+ * Trims {@link LogicalCorrelate} nodes.
+ */
+ public TrimResult trimFields(LogicalCorrelate correlate,
+ ImmutableBitSet fieldsUsed, Set<RelDataTypeField> extraFields) {
+ if (!extraFields.isEmpty()) {
+ // bail out with generic trim
+ return trimFields((RelNode) correlate, fieldsUsed, extraFields);
+ }
+
+ fieldsUsed = fieldsUsed.union(correlate.getRequiredColumns());
+
+ List<RelNode> newInputs = new ArrayList<>();
+ List<Mapping> inputMappings = new ArrayList<>();
+ int changeCount = 0;
+ int offset = 0;
+ for (RelNode input : correlate.getInputs()) {
+ final RelDataType inputRowType = input.getRowType();
+ final int inputFieldCount = inputRowType.getFieldCount();
+
+ ImmutableBitSet currentInputFieldsUsed = fieldsUsed
+ .intersect(ImmutableBitSet.range(offset, offset + inputFieldCount))
+ .shift(-offset);
+
+ TrimResult trimResult =
+ dispatchTrimFields(input, currentInputFieldsUsed, extraFields);
+
+ newInputs.add(trimResult.left);
+ inputMappings.add(trimResult.right);
+
+ offset += inputFieldCount;
+
+ if (trimResult.left != input) {
+ changeCount++;
+ }
+ }
+
+ if (changeCount == 0) {
+ return result(correlate,
+ Mappings.createIdentity(correlate.getRowType().getFieldCount()));
+ }
+
+ Mapping mapping = makeMapping(inputMappings);
+ RexBuilder rexBuilder = relBuilder.getRexBuilder();
+
+ RexCorrelVariableMapShuttle rexVisitor =
+ new RexCorrelVariableMapShuttle(correlate.getCorrelationId(),
+ newInputs.get(0).getRowType(), mapping, rexBuilder);
+ final LogicalCorrelate newCorrelate =
+ correlate
+ .copy(correlate.getTraitSet(),
+ newInputs.get(0),
+ newInputs.get(1).accept(new
RexRewritingRelShuttle(rexVisitor)),
+ correlate.getCorrelationId(),
+ correlate.getRequiredColumns().permute(mapping),
+ correlate.getJoinType());
+
+ return result(newCorrelate, mapping);
+ }
+
+ /**
+ * Updates correlate references in {@link RexNode} expressions.
+ */
+ static class RexCorrelVariableMapShuttle extends RexShuttle {
+ private final CorrelationId correlationId;
+ private final Mapping mapping;
+ private final RelDataType newCorrelRowType;
+ private final RexBuilder rexBuilder;
+
+ RexCorrelVariableMapShuttle(final CorrelationId correlationId,
Review Comment:
Even though the class is not public, I think JavaDoc for the parameters
would make this easier to understand.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
