Skip to content

Commit 66bc832

Browse files
authored
DRILL-8013: Drill attempts to push "$SUM0" to JDBC storage plugin for AVG (#2521)
1 parent fd836a2 commit 66bc832

File tree

4 files changed

+197
-5
lines changed

4 files changed

+197
-5
lines changed

contrib/storage-jdbc/src/test/java/org/apache/drill/exec/store/jdbc/TestJdbcPluginWithPostgres.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,4 +313,16 @@ public void testLimitPushDownWithConvertFromJson() throws Exception {
313313
.exclude("Limit\\(")
314314
.match();
315315
}
316+
317+
@Test // DRILL-8013
318+
public void testAvgFunction() throws Exception {
319+
String query = "select avg(person_id) `avg` from pg.`public`.person";
320+
321+
testBuilder()
322+
.sqlQuery(query)
323+
.unOrdered()
324+
.baselineColumns("avg")
325+
.baselineValues(2.75)
326+
.go();
327+
}
316328
}

exec/java-exec/src/main/java/org/apache/drill/exec/planner/logical/DrillReduceAggregatesRule.java

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
*/
1818
package org.apache.drill.exec.planner.logical;
1919

20+
import org.apache.drill.exec.planner.sql.DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper;
2021
import org.apache.drill.shaded.guava.com.google.common.collect.ImmutableList;
2122
import org.apache.drill.shaded.guava.com.google.common.collect.Lists;
2223
import org.apache.drill.shaded.guava.com.google.common.collect.Maps;
@@ -340,7 +341,8 @@ private RexNode reduceAvg(
340341
sumType,
341342
sumType.isNullable() || nGroups == 0);
342343
SqlAggFunction sumAgg =
343-
new DrillCalciteSqlAggFunctionWrapper(new SqlSumEmptyIsZeroAggFunction(), sumType);
344+
new DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
345+
new SqlSumEmptyIsZeroAggFunction(), sumType);
344346
AggregateCall sumCall = AggregateCall.create(sumAgg, oldCall.isDistinct(),
345347
oldCall.isApproximate(), oldCall.getArgList(), -1, sumType, null);
346348
final SqlCountAggFunction countAgg = (SqlCountAggFunction) SqlStdOperatorTable.COUNT;
@@ -437,7 +439,7 @@ private RexNode reduceSum(
437439
typeFactory.createTypeWithNullability(
438440
oldCall.getType(), argType.isNullable());
439441
}
440-
sumZeroAgg = new DrillCalciteSqlAggFunctionWrapper(
442+
sumZeroAgg = new DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
441443
new SqlSumEmptyIsZeroAggFunction(), sumType);
442444
AggregateCall sumZeroCall = AggregateCall.create(sumZeroAgg, oldCall.isDistinct(),
443445
oldCall.isApproximate(), oldCall.getArgList(), -1, sumType, null);
@@ -713,7 +715,7 @@ public void onMatch(RelOptRuleCall call) {
713715
final RelDataType argType = oldAggregateCall.getType();
714716
final RelDataType sumType = oldAggRel.getCluster().getTypeFactory()
715717
.createTypeWithNullability(argType, argType.isNullable());
716-
final SqlAggFunction sumZeroAgg = new DrillCalciteSqlAggFunctionWrapper(
718+
final SqlAggFunction sumZeroAgg = new DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
717719
new SqlSumEmptyIsZeroAggFunction(), sumType);
718720
AggregateCall sumZeroCall =
719721
AggregateCall.create(
@@ -775,7 +777,7 @@ public void onMatch(RelOptRuleCall call) {
775777
final RelDataType argType = rexWinAggCall.getType();
776778
final RelDataType sumType = oldWinRel.getCluster().getTypeFactory()
777779
.createTypeWithNullability(argType, argType.isNullable());
778-
final SqlAggFunction sumZeroAgg = new DrillCalciteSqlAggFunctionWrapper(
780+
final SqlAggFunction sumZeroAgg = new DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
779781
new SqlSumEmptyIsZeroAggFunction(), sumType);
780782
final Window.RexWinAggCall sumZeroCall =
781783
new Window.RexWinAggCall(
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing, software
13+
* distributed under the License is distributed on an "AS IS" BASIS,
14+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15+
* See the License for the specific language governing permissions and
16+
* limitations under the License.
17+
*/
18+
package org.apache.drill.exec.planner.sql;
19+
20+
import org.apache.calcite.rel.type.RelDataType;
21+
import org.apache.calcite.sql.SqlAggFunction;
22+
import org.apache.calcite.sql.SqlCall;
23+
import org.apache.calcite.sql.SqlCallBinding;
24+
import org.apache.calcite.sql.SqlOperator;
25+
import org.apache.calcite.sql.SqlOperatorBinding;
26+
import org.apache.calcite.sql.SqlSyntax;
27+
import org.apache.calcite.sql.fun.SqlSumEmptyIsZeroAggFunction;
28+
import org.apache.calcite.sql.type.SqlReturnTypeInference;
29+
import org.apache.calcite.sql.validate.SqlMonotonicity;
30+
import org.apache.calcite.sql.validate.SqlValidator;
31+
import org.apache.calcite.sql.validate.SqlValidatorScope;
32+
import org.apache.calcite.util.Litmus;
33+
import org.apache.calcite.util.Util;
34+
import org.apache.drill.exec.expr.fn.DrillFuncHolder;
35+
36+
import java.util.List;
37+
38+
/**
39+
* This class serves as a wrapper class for {@link SqlSumEmptyIsZeroAggFunction}
40+
* with the same goal as {@link DrillCalciteSqlAggFunctionWrapper}
41+
* but extends {@link SqlSumEmptyIsZeroAggFunction} to allow using
42+
* additional Calcite functionality designated for {@link SqlSumEmptyIsZeroAggFunction}.
43+
*/
44+
public class DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper
45+
extends SqlSumEmptyIsZeroAggFunction implements DrillCalciteSqlWrapper {
46+
47+
private final SqlAggFunction operator;
48+
49+
private final SqlReturnTypeInference sqlReturnTypeInference;
50+
51+
private DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
52+
SqlSumEmptyIsZeroAggFunction sqlAggFunction,
53+
SqlReturnTypeInference sqlReturnTypeInference) {
54+
this.sqlReturnTypeInference = sqlReturnTypeInference;
55+
this.operator = sqlAggFunction;
56+
}
57+
58+
public DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
59+
SqlSumEmptyIsZeroAggFunction sqlAggFunction,
60+
List<DrillFuncHolder> functions) {
61+
this(sqlAggFunction,
62+
TypeInferenceUtils.getDrillSqlReturnTypeInference(
63+
sqlAggFunction.getName(),
64+
functions));
65+
}
66+
67+
public DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
68+
SqlSumEmptyIsZeroAggFunction sqlAggFunction,
69+
RelDataType relDataType) {
70+
this(sqlAggFunction, opBinding -> relDataType);
71+
}
72+
73+
@Override
74+
public SqlOperator getOperator() {
75+
return operator;
76+
}
77+
78+
@Override
79+
public SqlReturnTypeInference getReturnTypeInference() {
80+
return this.sqlReturnTypeInference;
81+
}
82+
83+
@Override
84+
public RelDataType inferReturnType(SqlOperatorBinding opBinding) {
85+
if (this.sqlReturnTypeInference != null) {
86+
RelDataType returnType = this.sqlReturnTypeInference.inferReturnType(opBinding);
87+
if (returnType == null) {
88+
throw new IllegalArgumentException(String.format(
89+
"Cannot infer return type for %s; operand types: %s",
90+
opBinding.getOperator(), opBinding.collectOperandTypes()));
91+
} else {
92+
return returnType;
93+
}
94+
} else {
95+
throw Util.needToImplement(this);
96+
}
97+
}
98+
99+
@Override
100+
public boolean validRexOperands(int count, Litmus litmus) {
101+
return true;
102+
}
103+
104+
@Override
105+
public String getAllowedSignatures(String opNameToUse) {
106+
return operator.getAllowedSignatures(opNameToUse);
107+
}
108+
109+
@Override
110+
public boolean isAggregator() {
111+
return operator.isAggregator();
112+
}
113+
114+
@Override
115+
public boolean allowsFraming() {
116+
return operator.allowsFraming();
117+
}
118+
119+
@Override
120+
public SqlMonotonicity getMonotonicity(SqlOperatorBinding call) {
121+
return operator.getMonotonicity(call);
122+
}
123+
124+
@Override
125+
public boolean isDeterministic() {
126+
return operator.isDeterministic();
127+
}
128+
129+
@Override
130+
public boolean isDynamicFunction() {
131+
return operator.isDynamicFunction();
132+
}
133+
134+
@Override
135+
public boolean requiresDecimalExpansion() {
136+
return operator.requiresDecimalExpansion();
137+
}
138+
139+
@Override
140+
public boolean argumentMustBeScalar(int ordinal) {
141+
return operator.argumentMustBeScalar(ordinal);
142+
}
143+
144+
@Override
145+
public boolean checkOperandTypes(
146+
SqlCallBinding callBinding,
147+
boolean throwOnFailure) {
148+
return true;
149+
}
150+
151+
@Override
152+
public SqlSyntax getSyntax() {
153+
return operator.getSyntax();
154+
}
155+
156+
@Override
157+
public List<String> getParamNames() {
158+
return operator.getParamNames();
159+
}
160+
161+
@Override
162+
public String getSignatureTemplate(final int operandsCount) {
163+
return operator.getSignatureTemplate(operandsCount);
164+
}
165+
166+
@Override
167+
public RelDataType deriveType(
168+
SqlValidator validator,
169+
SqlValidatorScope scope,
170+
SqlCall call) {
171+
return operator.deriveType(validator, scope, call);
172+
}
173+
}

exec/java-exec/src/main/java/org/apache/drill/exec/planner/sql/DrillOperatorTable.java

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
*/
1818
package org.apache.drill.exec.planner.sql;
1919

20+
import org.apache.calcite.sql.fun.SqlSumEmptyIsZeroAggFunction;
2021
import org.apache.calcite.sql.validate.SqlNameMatcher;
2122
import org.apache.drill.shaded.guava.com.google.common.collect.ArrayListMultimap;
2223
import org.apache.drill.shaded.guava.com.google.common.collect.Lists;
@@ -169,7 +170,11 @@ public List<SqlOperator> getSqlOperator(String name) {
169170
private void populateWrappedCalciteOperators() {
170171
for (SqlOperator calciteOperator : inner.getOperatorList()) {
171172
final SqlOperator wrapper;
172-
if (calciteOperator instanceof SqlAggFunction) {
173+
if (calciteOperator instanceof SqlSumEmptyIsZeroAggFunction) {
174+
wrapper = new DrillCalciteSqlSumEmptyIsZeroAggFunctionWrapper(
175+
(SqlSumEmptyIsZeroAggFunction) calciteOperator,
176+
getFunctionListWithInference(calciteOperator.getName()));
177+
} else if (calciteOperator instanceof SqlAggFunction) {
173178
wrapper = new DrillCalciteSqlAggFunctionWrapper((SqlAggFunction) calciteOperator,
174179
getFunctionListWithInference(calciteOperator.getName()));
175180
} else if (calciteOperator instanceof SqlFunction) {

0 commit comments

Comments
 (0)