rubenada commented on code in PR #4077:
URL: https://github.com/apache/calcite/pull/4077#discussion_r1871617934
##########
core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java:
##########
@@ -1336,6 +1338,144 @@ 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() || correlate.getInputs().size() != 2) {
+ // 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 = concatenateMappings(inputMappings);
+ RexBuilder rexBuilder = relBuilder.getRexBuilder();
+
+ RelNode newLeft = newInputs.get(0);
+ RexCorrelVariableMapShuttle rexVisitor =
+ new RexCorrelVariableMapShuttle(correlate.getCorrelationId(),
+ newLeft.getRowType(), mapping, rexBuilder);
+ RelNode newRight =
+ newInputs.get(1).accept(new RexRewritingRelShuttle(rexVisitor));
+ final LogicalCorrelate newCorrelate =
+ correlate
+ .copy(correlate.getTraitSet(),
+ newLeft,
+ newRight,
+ 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;
+
+
+ /**
+ * Constructs a RexCorrelVariableMapShuttle.
+ *
+ * @param correlationId The ID of the correlation variable to update.
+ * @param newCorrelRowType The new row type for the correlate reference.
+ * @param mapping Mapping to transform field indices.
+ * @param rexBuilder A builder for constructing new RexNodes.
+ */
+ 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 concatenateMappings(List<Mapping> inputMappings) {
+ int fieldCount = 0;
+ int newFieldCount = 0;
+ for (Mapping mapping : inputMappings) {
+ fieldCount += mapping.getSourceCount();
+ newFieldCount += mapping.getTargetCount();
+ }
+
+ Mapping mapping =
+ Mappings.create(MappingType.INVERSE_SURJECTION, fieldCount,
newFieldCount);
+ int offset = 0;
+ int newOffset = 0;
+ for (int i = 0; i < inputMappings.size(); i++) {
Review Comment:
nit: we could have this for loop simplified as `for (Mapping inputMapping :
inputMappings) {...}`
--
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]
