diff --git a/ql/src/java/org/apache/hadoop/hive/ql/optimizer/calcite/rules/HiveRelFieldTrimmer.java b/ql/src/java/org/apache/hadoop/hive/ql/optimizer/calcite/rules/HiveRelFieldTrimmer.java index 055da42e10..f3930a118c 100644 --- a/ql/src/java/org/apache/hadoop/hive/ql/optimizer/calcite/rules/HiveRelFieldTrimmer.java +++ b/ql/src/java/org/apache/hadoop/hive/ql/optimizer/calcite/rules/HiveRelFieldTrimmer.java @@ -24,20 +24,28 @@ import java.util.Map; import java.util.Set; +import com.google.common.collect.ImmutableList; import org.apache.calcite.adapter.druid.DruidQuery; import org.apache.calcite.linq4j.Ord; import org.apache.calcite.plan.RelOptTable; import org.apache.calcite.plan.RelOptUtil; +import org.apache.calcite.rel.RelCollation; +import org.apache.calcite.rel.RelFieldCollation; import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.CorrelationId; import org.apache.calcite.rel.core.Project; import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rel.metadata.RelMetadataQuery; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeField; import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexCorrelVariable; +import org.apache.calcite.rex.RexFieldAccess; import org.apache.calcite.rex.RexNode; import org.apache.calcite.rex.RexPermuteInputsShuttle; import org.apache.calcite.rex.RexVisitor; import org.apache.calcite.sql.validate.SqlValidator; +import org.apache.calcite.sql2rel.CorrelationReferenceFinder; import org.apache.calcite.sql2rel.RelFieldTrimmer; import org.apache.calcite.tools.RelBuilder; import org.apache.calcite.util.ImmutableBitSet; @@ -83,6 +91,40 @@ public HiveRelFieldTrimmer(SqlValidator validator, RelBuilder relBuilder, boolea this.fetchStats = fetchStats; } + /** + * Trims the fields of an input relational expression. + * + * @param rel Relational expression + * @param input Input relational expression, whose fields to trim + * @param fieldsUsed Bitmap of fields needed by the consumer + * @return New relational expression and its field mapping + */ + protected TrimResult trimChild( + RelNode rel, + RelNode input, + final ImmutableBitSet fieldsUsed, + Set extraFields) { + final ImmutableBitSet.Builder fieldsUsedBuilder = fieldsUsed.rebuild(); + + // Correlating variables are a means for other relational expressions to use + // fields. + for (final CorrelationId correlation : rel.getVariablesSet()) { + rel.accept( + new CorrelationReferenceFinder() { + protected RexNode handle(RexFieldAccess fieldAccess) { + final RexCorrelVariable v = + (RexCorrelVariable) fieldAccess.getReferenceExpr(); + if (v.id.equals(correlation)) { + fieldsUsedBuilder.set(fieldAccess.getField().getIndex()); + } + return fieldAccess; + } + }); + } + + return dispatchTrimFields(input, fieldsUsedBuilder.build(), extraFields); + } + /** * Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for * {@link org.apache.hadoop.hive.ql.optimizer.calcite.reloperators.HiveMultiJoin}.