org.apache.druid.sql.calcite.rel.Projection 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 org.apache.druid.sql.calcite.rel;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
import org.apache.calcite.rel.core.Project;
import org.apache.calcite.rel.type.RelDataType;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.sql.SqlKind;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.math.expr.ExpressionType;
import org.apache.druid.query.aggregation.PostAggregator;
import org.apache.druid.query.aggregation.post.ExpressionPostAggregator;
import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.column.RowSignature;
import org.apache.druid.segment.column.ValueType;
import org.apache.druid.sql.calcite.expression.DruidExpression;
import org.apache.druid.sql.calcite.expression.Expressions;
import org.apache.druid.sql.calcite.expression.OperatorConversions;
import org.apache.druid.sql.calcite.expression.PostAggregatorVisitor;
import org.apache.druid.sql.calcite.planner.Calcites;
import org.apache.druid.sql.calcite.planner.PlannerContext;
import org.apache.druid.sql.calcite.table.RowSignatures;
import javax.annotation.Nullable;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;
/**
* Used to represent projections (Calcite "Project"). These are embedded in {@link Sorting} and {@link Grouping} to
* store post-sorting and post-grouping projections, as well as directly in {@link DruidQuery} to store potential
* post-selection projections. They may be built using either virtual columns (pre-aggregation) or post-aggregators.
*
* It is expected that callers will create and use Projection instances in the same context (pre- or post-aggregation).
* If this isn't done properly (i.e. a caller creates a pre-aggregation Projection but then calls
* {@link #getPostAggregators()} then an exception will be thrown.
*/
public class Projection
{
@Nullable
private final List postAggregators;
@Nullable
private final List virtualColumns;
private final RowSignature outputRowSignature;
private Projection(
@Nullable final List postAggregators,
@Nullable final List virtualColumns,
final RowSignature outputRowSignature
)
{
if (postAggregators == null && virtualColumns == null) {
throw new IAE("postAggregators and virtualColumns cannot both be null");
} else if (postAggregators != null && virtualColumns != null) {
throw new IAE("postAggregators and virtualColumns cannot both be nonnull");
}
this.postAggregators = postAggregators;
this.virtualColumns = virtualColumns;
this.outputRowSignature = outputRowSignature;
}
private static void postAggregationHandleInputRefOrLiteral(
final Project project,
final PlannerContext plannerContext,
final RowSignature inputRowSignature,
final RexNode postAggregatorRexNode,
final List rowOrder,
final PostAggregatorVisitor postAggregatorVisitor
)
{
// Attempt to convert to PostAggregator.
final DruidExpression postAggregatorExpression = Expressions.toDruidExpression(
plannerContext,
inputRowSignature,
postAggregatorRexNode
);
if (postAggregatorExpression == null) {
throw new CannotBuildQueryException(project, postAggregatorRexNode);
}
handlePostAggregatorExpression(
plannerContext,
inputRowSignature,
postAggregatorRexNode,
rowOrder,
postAggregatorVisitor,
postAggregatorExpression
);
}
private static void postAggregationHandleOtherKinds(
final Project project,
final PlannerContext plannerContext,
final RowSignature inputRowSignature,
final RexNode postAggregatorRexNode,
final List rowOrder,
final PostAggregatorVisitor postAggregatorVisitor
)
{
PostAggregator pagg = OperatorConversions.toPostAggregator(
plannerContext,
inputRowSignature,
postAggregatorRexNode,
postAggregatorVisitor,
false
);
if (pagg != null) {
postAggregatorVisitor.addPostAgg(pagg);
rowOrder.add(pagg.getName());
} else {
final DruidExpression postAggregatorExpression = Expressions.toDruidExpressionWithPostAggOperands(
plannerContext,
inputRowSignature,
postAggregatorRexNode,
postAggregatorVisitor
);
if (postAggregatorExpression == null) {
throw new CannotBuildQueryException(project, postAggregatorRexNode);
}
handlePostAggregatorExpression(
plannerContext,
inputRowSignature,
postAggregatorRexNode,
rowOrder,
postAggregatorVisitor,
postAggregatorExpression
);
}
}
private static void handlePostAggregatorExpression(
final PlannerContext plannerContext,
final RowSignature inputRowSignature,
final RexNode postAggregatorRexNode,
final List rowOrder,
final PostAggregatorVisitor postAggregatorVisitor,
final DruidExpression postAggregatorExpression
)
{
if (postAggregatorDirectColumnIsOk(inputRowSignature, postAggregatorExpression, postAggregatorRexNode)) {
// Direct column access, without any type cast as far as Druid's runtime is concerned.
// (There might be a SQL-level type cast that we don't care about)
rowOrder.add(postAggregatorExpression.getDirectColumn());
} else {
final PostAggregator postAggregator = new ExpressionPostAggregator(
postAggregatorVisitor.getOutputNamePrefix() + postAggregatorVisitor.getAndIncrementCounter(),
postAggregatorExpression.getExpression(),
null,
postAggregatorExpression.getDruidType(),
plannerContext.parseExpression(postAggregatorExpression.getExpression())
);
postAggregatorVisitor.addPostAgg(postAggregator);
rowOrder.add(postAggregator.getName());
}
}
public static Projection postAggregation(
final Project project,
final PlannerContext plannerContext,
final RowSignature inputRowSignature,
final String basePrefix
)
{
final List rowOrder = new ArrayList<>();
final String outputNamePrefix = Calcites.findUnusedPrefixForDigits(basePrefix, inputRowSignature.getColumnNames());
final PostAggregatorVisitor postAggVisitor = new PostAggregatorVisitor(outputNamePrefix);
for (final RexNode postAggregatorRexNode : project.getProjects()) {
if (postAggregatorRexNode.getKind() == SqlKind.INPUT_REF || postAggregatorRexNode.getKind() == SqlKind.LITERAL) {
postAggregationHandleInputRefOrLiteral(
project,
plannerContext,
inputRowSignature,
postAggregatorRexNode,
rowOrder,
postAggVisitor
);
} else {
postAggregationHandleOtherKinds(
project,
plannerContext,
inputRowSignature,
postAggregatorRexNode,
rowOrder,
postAggVisitor
);
}
}
return new Projection(
postAggVisitor.getPostAggs(),
null,
RowSignatures.fromRelDataType(rowOrder, project.getRowType())
);
}
public static Projection preAggregation(
final Project project,
final PlannerContext plannerContext,
final RowSignature inputRowSignature,
final VirtualColumnRegistry virtualColumnRegistry
)
{
final List expressions = new ArrayList<>();
for (final RexNode rexNode : project.getProjects()) {
final DruidExpression expression = Expressions.toDruidExpression(
plannerContext,
inputRowSignature,
rexNode
);
if (expression == null) {
throw new CannotBuildQueryException(project, rexNode);
} else {
expressions.add(expression);
}
}
final Set virtualColumns = new HashSet<>();
final List rowOrder = new ArrayList<>();
for (int i = 0; i < expressions.size(); i++) {
final DruidExpression expression = expressions.get(i);
final RelDataType dataType = project.getRowType().getFieldList().get(i).getType();
if (expression.isDirectColumnAccess() &&
Objects.equals(
inputRowSignature.getColumnType(expression.getDirectColumn()).orElse(null),
Calcites.getColumnTypeForRelDataType(dataType)
)
) {
// Refer to column directly when it's a direct access with matching type.
rowOrder.add(expression.getDirectColumn());
} else {
String virtualColumnName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(
expression,
project.getProjects().get(i).getType()
);
virtualColumns.add(virtualColumnName);
rowOrder.add(virtualColumnName);
}
}
return new Projection(
null,
ImmutableList.copyOf(virtualColumns),
RowSignatures.fromRelDataType(rowOrder, project.getRowType())
);
}
/**
* Returns true if a post-aggregation "expression" can be realized as a direct field access. This is true if it's
* a direct column access that doesn't require an implicit cast.
*
* @param aggregateRowSignature signature of the aggregation
* @param expression post-aggregation expression
* @param rexNode RexNode for the post-aggregation expression
*
* @return yes or no
*/
private static boolean postAggregatorDirectColumnIsOk(
final RowSignature aggregateRowSignature,
final DruidExpression expression,
final RexNode rexNode
)
{
if (!expression.isDirectColumnAccess()) {
return false;
}
// We don't really have a way to cast complex type. So might as well not do anything and return.
final ColumnType columnValueType =
aggregateRowSignature.getColumnType(expression.getDirectColumn())
.orElseThrow(
() -> new ISE("Encountered null type for column[%s]", expression.getDirectColumn())
);
if (columnValueType.is(ValueType.COMPLEX)) {
return true;
}
// Check if a cast is necessary.
final ExpressionType toExprType = ExpressionType.fromColumnTypeStrict(columnValueType);
final ExpressionType fromExprType = ExpressionType.fromColumnTypeStrict(
Calcites.getColumnTypeForRelDataType(rexNode.getType())
);
return toExprType.equals(fromExprType);
}
public List getPostAggregators()
{
// If you ever see this error, it probably means a Projection was created in pre-aggregation mode, but then
// used in a post-aggregation context. This is likely a bug somewhere in DruidQuery. See class-level Javadocs.
return Preconditions.checkNotNull(postAggregators, "postAggregators");
}
public List getVirtualColumns()
{
// If you ever see this error, it probably means a Projection was created in post-aggregation mode, but then
// used in a pre-aggregation context. This is likely a bug somewhere in DruidQuery. See class-level Javadocs.
return Preconditions.checkNotNull(virtualColumns, "virtualColumns");
}
public RowSignature getOutputRowSignature()
{
return outputRowSignature;
}
@Override
public boolean equals(Object o)
{
if (this == o) {
return true;
}
if (o == null || getClass() != o.getClass()) {
return false;
}
Projection that = (Projection) o;
return Objects.equals(postAggregators, that.postAggregators) &&
Objects.equals(virtualColumns, that.virtualColumns) &&
Objects.equals(outputRowSignature, that.outputRowSignature);
}
@Override
public int hashCode()
{
return Objects.hash(postAggregators, virtualColumns, outputRowSignature);
}
@Override
public String toString()
{
return "PostSortingExpressions{" +
"postAggregators=" + postAggregators +
", virtualColumns=" + virtualColumns +
", outputRowSignature=" + outputRowSignature +
'}';
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy