com.hazelcast.org.apache.calcite.rel.rules.AggregateProjectMergeRule Maven / Gradle / Ivy
/*
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to you under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.hazelcast.org.apache.calcite.rel.rules;
import com.hazelcast.org.apache.calcite.plan.RelOptRuleCall;
import com.hazelcast.org.apache.calcite.plan.RelOptUtil;
import com.hazelcast.org.apache.calcite.plan.RelRule;
import com.hazelcast.org.apache.calcite.rel.RelNode;
import com.hazelcast.org.apache.calcite.rel.core.Aggregate;
import com.hazelcast.org.apache.calcite.rel.core.Aggregate.Group;
import com.hazelcast.org.apache.calcite.rel.core.AggregateCall;
import com.hazelcast.org.apache.calcite.rel.core.Project;
import com.hazelcast.org.apache.calcite.rex.RexInputRef;
import com.hazelcast.org.apache.calcite.rex.RexNode;
import com.hazelcast.org.apache.calcite.tools.RelBuilder;
import com.hazelcast.org.apache.calcite.tools.RelBuilderFactory;
import com.hazelcast.org.apache.calcite.util.ImmutableBitSet;
import com.hazelcast.org.apache.calcite.util.Util;
import com.hazelcast.org.apache.calcite.util.mapping.Mappings;
import com.hazelcast.com.google.common.collect.ImmutableList;
import com.hazelcast.org.checkerframework.checker.nullness.qual.Nullable;
import org.immutables.value.Value;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import static java.util.Objects.requireNonNull;
/**
* Planner rule that recognizes a {@link com.hazelcast.org.apache.calcite.rel.core.Aggregate}
* on top of a {@link com.hazelcast.org.apache.calcite.rel.core.Project} and if possible
* aggregate through the project or removes the project.
*
* This is only possible when the grouping expressions and arguments to
* the aggregate functions are field references (i.e. not expressions).
*
*
In some cases, this rule has the effect of trimming: the aggregate will
* use fewer columns than the project did.
*
* @see CoreRules#AGGREGATE_PROJECT_MERGE
*/
@Value.Enclosing
public class AggregateProjectMergeRule
extends RelRule
implements TransformationRule {
/** Creates an AggregateProjectMergeRule. */
protected AggregateProjectMergeRule(Config config) {
super(config);
}
@Deprecated // to be removed before 2.0
public AggregateProjectMergeRule(
Class aggregateClass,
Class projectClass,
RelBuilderFactory relBuilderFactory) {
this(CoreRules.AGGREGATE_PROJECT_MERGE.config
.withRelBuilderFactory(relBuilderFactory)
.as(Config.class)
.withOperandFor(aggregateClass, projectClass));
}
@Override public void onMatch(RelOptRuleCall call) {
final Aggregate aggregate = call.rel(0);
final Project project = call.rel(1);
RelNode x = apply(call, aggregate, project);
if (x != null) {
call.transformTo(x);
}
}
public static @Nullable RelNode apply(RelOptRuleCall call, Aggregate aggregate,
Project project) {
// Find all fields which we need to be straightforward field projections.
final Set interestingFields = RelOptUtil.getAllFields(aggregate);
// Build the map from old to new; abort if any entry is not a
// straightforward field projection.
final Map map = new HashMap<>();
for (int source : interestingFields) {
final RexNode rex = project.getProjects().get(source);
if (!(rex instanceof RexInputRef)) {
return null;
}
map.put(source, ((RexInputRef) rex).getIndex());
}
final ImmutableBitSet newGroupSet = aggregate.getGroupSet().permute(map);
ImmutableList newGroupingSets = null;
if (aggregate.getGroupType() != Group.SIMPLE) {
newGroupingSets =
ImmutableBitSet.ORDERING.immutableSortedCopy(
ImmutableBitSet.permute(aggregate.getGroupSets(), map));
}
final ImmutableList.Builder aggCalls =
ImmutableList.builder();
final int sourceCount = aggregate.getInput().getRowType().getFieldCount();
final int targetCount = project.getInput().getRowType().getFieldCount();
final Mappings.TargetMapping targetMapping =
Mappings.target(map, sourceCount, targetCount);
for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
aggCalls.add(aggregateCall.transform(targetMapping));
}
final Aggregate newAggregate =
aggregate.copy(aggregate.getTraitSet(), project.getInput(),
newGroupSet, newGroupingSets, aggCalls.build());
// Add a project if the group set is not in the same order or
// contains duplicates.
final RelBuilder relBuilder = call.builder();
relBuilder.push(newAggregate);
final List newKeys =
Util.transform(aggregate.getGroupSet().asList(),
key -> requireNonNull(map.get(key),
() -> "no value found for key " + key + " in " + map));
if (!newKeys.equals(newGroupSet.asList())) {
final List posList = new ArrayList<>();
for (int newKey : newKeys) {
posList.add(newGroupSet.indexOf(newKey));
}
for (int i = newAggregate.getGroupCount();
i < newAggregate.getRowType().getFieldCount(); i++) {
posList.add(i);
}
relBuilder.project(relBuilder.fields(posList));
}
return relBuilder.build();
}
/** Rule configuration. */
@Value.Immutable
public interface Config extends RelRule.Config {
Config DEFAULT = ImmutableAggregateProjectMergeRule.Config.of()
.withOperandFor(Aggregate.class, Project.class);
@Override default AggregateProjectMergeRule toRule() {
return new AggregateProjectMergeRule(this);
}
/** Defines an operand tree for the given classes. */
default Config withOperandFor(Class aggregateClass,
Class projectClass) {
return withOperandSupplier(b0 ->
b0.operand(aggregateClass).oneInput(b1 ->
b1.operand(projectClass).anyInputs())).as(Config.class);
}
}
}