You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by xi...@apache.org on 2020/10/09 21:06:06 UTC
[incubator-pinot] 01/01: Rewrite possible array aggregation
functions to one level
This is an automated email from the ASF dual-hosted git repository.
xiangfu pushed a commit to branch rewrite_array_aggregation_functions
in repository https://gitbox.apache.org/repos/asf/incubator-pinot.git
commit 463d6348e3ed246483fc40b15c482bef5a3df1e3
Author: Xiang Fu <fx...@gmail.com>
AuthorDate: Fri Oct 9 14:05:34 2020 -0700
Rewrite possible array aggregation functions to one level
---
.../apache/pinot/sql/parsers/CalciteSqlParser.java | 66 ++++++++++++++++++++++
.../pinot/sql/parsers/CalciteSqlCompilerTest.java | 36 ++++++++++++
2 files changed, 102 insertions(+)
diff --git a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
index 5275a1a..3ac041b 100644
--- a/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
+++ b/pinot-common/src/main/java/org/apache/pinot/sql/parsers/CalciteSqlParser.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.sql.parsers;
+import com.google.common.base.Preconditions;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
@@ -47,11 +48,13 @@ import org.apache.calcite.sql.parser.SqlParseException;
import org.apache.calcite.sql.parser.SqlParser;
import org.apache.calcite.sql.parser.babel.SqlBabelParserImpl;
import org.apache.calcite.sql.validate.SqlConformanceEnum;
+import org.apache.commons.lang3.StringUtils;
import org.apache.pinot.common.function.AggregationFunctionType;
import org.apache.pinot.common.function.FunctionDefinitionRegistry;
import org.apache.pinot.common.function.FunctionInfo;
import org.apache.pinot.common.function.FunctionInvoker;
import org.apache.pinot.common.function.FunctionRegistry;
+import org.apache.pinot.common.function.TransformFunctionType;
import org.apache.pinot.common.request.DataSource;
import org.apache.pinot.common.request.Expression;
import org.apache.pinot.common.request.ExpressionType;
@@ -343,6 +346,9 @@ public class CalciteSqlParser {
// Invoke compilation time functions
invokeCompileTimeFunctions(pinotQuery);
+ // Rewrite Selection list
+ rewriteSelections(pinotQuery.getSelectList());
+
// Update Predicate Comparison
Expression filterExpression = pinotQuery.getFilterExpression();
if (filterExpression != null) {
@@ -364,6 +370,66 @@ public class CalciteSqlParser {
validate(aliasMap, pinotQuery);
}
+ private static void rewriteSelections(List<Expression> selectList) {
+ for (Expression expression : selectList) {
+ // Rewrite aggregation
+ tryToRewriteArrayFunction(expression);
+ }
+ }
+
+ private static void tryToRewriteArrayFunction(Expression expression) {
+ if (!expression.isSetFunctionCall()) {
+ return;
+ }
+ Function functionCall = expression.getFunctionCall();
+ switch (canonicalize(functionCall.getOperator())) {
+ case "sum":
+ Preconditions.checkState(functionCall.getOperands().size() == 1, "Expected 1 argument in SUM function");
+ if (functionCall.getOperands().get(0).isSetFunctionCall()) {
+ Function innerFunction = functionCall.getOperands().get(0).getFunctionCall();
+ if (isSameFunction(innerFunction.getOperator(), TransformFunctionType.ARRAYSUM.getName())) {
+ Function sumMvFunc = new Function(AggregationFunctionType.SUMMV.getName());
+ sumMvFunc.setOperands(innerFunction.getOperands());
+ expression.setFunctionCall(sumMvFunc);
+ }
+ }
+ return;
+ case "min":
+ Preconditions.checkState(functionCall.getOperands().size() == 1, "Expected 1 argument in MIN function");
+ if (functionCall.getOperands().get(0).isSetFunctionCall()) {
+ Function innerFunction = functionCall.getOperands().get(0).getFunctionCall();
+ if (isSameFunction(innerFunction.getOperator(), TransformFunctionType.ARRAYMIN.getName())) {
+ Function sumMvFunc = new Function(AggregationFunctionType.MINMV.getName());
+ sumMvFunc.setOperands(innerFunction.getOperands());
+ expression.setFunctionCall(sumMvFunc);
+ }
+ }
+ return;
+ case "max":
+ Preconditions.checkState(functionCall.getOperands().size() == 1, "Expected 1 argument in MAX function");
+ if (functionCall.getOperands().get(0).isSetFunctionCall()) {
+ Function innerFunction = functionCall.getOperands().get(0).getFunctionCall();
+ if (isSameFunction(innerFunction.getOperator(), TransformFunctionType.ARRAYMAX.getName())) {
+ Function sumMvFunc = new Function(AggregationFunctionType.MAXMV.getName());
+ sumMvFunc.setOperands(innerFunction.getOperands());
+ expression.setFunctionCall(sumMvFunc);
+ }
+ }
+ return;
+ }
+ for (Expression operand : functionCall.getOperands()) {
+ tryToRewriteArrayFunction(operand);
+ }
+ }
+
+ private static String canonicalize(String functionName) {
+ return StringUtils.remove(functionName, '_').toLowerCase();
+ }
+
+ private static boolean isSameFunction(String function1, String function2) {
+ return canonicalize(function1).equals(canonicalize(function2));
+ }
+
/**
* Rewrite non-aggregate group by query to distinct query.
* E.g.
diff --git a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
index c025d6b..7d2c08d 100644
--- a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
+++ b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlCompilerTest.java
@@ -2048,4 +2048,40 @@ public class CalciteSqlCompilerTest {
Assert.assertNull(brokerRequest.getHavingFilterSubQueryMap());
}
}
+
+ @Test
+ public void testArrayAggregationRewrite() {
+ String sql;
+ PinotQuery pinotQuery;
+ sql = "SELECT sum(array_sum(a)) FROM Foo";
+ pinotQuery = CalciteSqlParser.compileToPinotQuery(sql);
+ Assert.assertEquals(pinotQuery.getSelectListSize(), 1);
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "sumMV");
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().size(), 1);
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().get(0).getIdentifier().getName(), "a");
+
+ sql = "SELECT MIN(ARRAYMIN(a)) FROM Foo";
+ pinotQuery = CalciteSqlParser.compileToPinotQuery(sql);
+ Assert.assertEquals(pinotQuery.getSelectListSize(), 1);
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "minMV");
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().size(), 1);
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().get(0).getIdentifier().getName(), "a");
+
+ sql = "SELECT Max(ArrayMax(a)) FROM Foo";
+ pinotQuery = CalciteSqlParser.compileToPinotQuery(sql);
+ Assert.assertEquals(pinotQuery.getSelectListSize(), 1);
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "maxMV");
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().size(), 1);
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().get(0).getIdentifier().getName(), "a");
+
+ sql = "SELECT Max(ArrayMax(a)) + 1 FROM Foo";
+ pinotQuery = CalciteSqlParser.compileToPinotQuery(sql);
+ Assert.assertEquals(pinotQuery.getSelectListSize(), 1);
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperator(), "PLUS");
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().size(), 2);
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().get(0).getFunctionCall().getOperator(), "maxMV");
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().get(0).getFunctionCall().getOperands().size(), 1);
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().get(0).getFunctionCall().getOperands().get(0).getIdentifier().getName(), "a");
+ Assert.assertEquals(pinotQuery.getSelectList().get(0).getFunctionCall().getOperands().get(1).getLiteral().getLongValue(), 1L);
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org