Copilot commented on code in PR #16399: URL: https://github.com/apache/pinot/pull/16399#discussion_r2224004946
########## pinot-common/src/test/java/org/apache/pinot/sql/parsers/rewriter/AggregationOptimizerTest.java: ########## @@ -0,0 +1,535 @@ +/** + * 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.pinot.sql.parsers.rewriter; + +import org.apache.pinot.common.request.Expression; +import org.apache.pinot.common.request.Function; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; + + +public class AggregationOptimizerTest { + + private final AggregationOptimizer _optimizer = new AggregationOptimizer(); + + @Test + public void testSumColumnPlusConstant() { + // Test: SELECT sum(met + 2) → SELECT sum(met) + 2 * count(1) + String query = "SELECT sum(met + 2) FROM mytable"; + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); + + // Apply optimizer + _optimizer.rewrite(pinotQuery); + + // Verify optimization + Expression selectExpression = pinotQuery.getSelectList().get(0); + verifyOptimizedAddition(selectExpression, "met", 2); + } + + @Test + public void testSumConstantPlusColumn() { + // Test: SELECT sum(2 + met) → SELECT sum(met) + 2 * count(1) + String query = "SELECT sum(2 + met) FROM mytable"; + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); + + // Apply optimizer + _optimizer.rewrite(pinotQuery); + + // Verify optimization + Expression selectExpression = pinotQuery.getSelectList().get(0); + verifyOptimizedAddition(selectExpression, "met", 2); + } + + @Test + public void testSumColumnMinusConstant() { + // Test: SELECT sum(met - 5) → SELECT sum(met) - 5 * count(1) + String query = "SELECT sum(met - 5) FROM mytable"; + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); + + // Apply optimizer + _optimizer.rewrite(pinotQuery); + + // Verify optimization + Expression selectExpression = pinotQuery.getSelectList().get(0); + verifyOptimizedSubtraction(selectExpression, "met", 5); + } + + @Test + public void testSumConstantMinusColumn() { + // Test: SELECT sum(10 - met) → SELECT 10 * count(1) - sum(met) + String query = "SELECT sum(10 - met) FROM mytable"; + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); + + // Apply optimizer + _optimizer.rewrite(pinotQuery); + + // Verify optimization + Expression selectExpression = pinotQuery.getSelectList().get(0); + verifyOptimizedSubtractionReversed(selectExpression, 10, "met"); + } + + @Test + public void testSumWithFloatConstant() { + // Test: SELECT sum(price + 2.5) → SELECT sum(price) + 2.5 * count(1) + String query = "SELECT sum(price + 2.5) FROM mytable"; + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); + + // Apply optimizer + _optimizer.rewrite(pinotQuery); + + // Verify optimization + Expression selectExpression = pinotQuery.getSelectList().get(0); + verifyOptimizedFloatAddition(selectExpression, "price", 2.5); + } + + @Test + public void testMultipleAggregations() { + // Test: SELECT sum(a + 1), sum(b - 2), avg(c) FROM mytable + String query = "SELECT sum(a + 1), sum(b - 2), avg(c) FROM mytable"; + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); + + // Apply optimizer + _optimizer.rewrite(pinotQuery); + + // Verify optimizations + assertEquals(pinotQuery.getSelectList().size(), 3); + + // First aggregation: sum(a + 1) → sum(a) + 1 * count(1) + Expression firstExpression = pinotQuery.getSelectList().get(0); + verifyOptimizedAddition(firstExpression, "a", 1); + + // Second aggregation: sum(b - 2) → sum(b) - 2 * count(1) + Expression secondExpression = pinotQuery.getSelectList().get(1); + verifyOptimizedSubtraction(secondExpression, "b", 2); + + // Third aggregation: avg(c) should remain unchanged + Expression thirdExpression = pinotQuery.getSelectList().get(2); + assertEquals(thirdExpression.getFunctionCall().getOperator(), "avg"); + } + + @Test + public void testNoOptimizationForUnsupportedPatterns() { + // Test patterns that should NOT be optimized + String[] queries = { + "SELECT sum(a * 2) FROM mytable", // multiplication not supported + "SELECT sum(a / 2) FROM mytable", // division not supported + "SELECT sum(a + b) FROM mytable", // both operands are columns + "SELECT sum(1 + 2) FROM mytable", // both operands are constants + "SELECT avg(a + 2) FROM mytable", // not a sum function + "SELECT sum(a) FROM mytable", // no arithmetic expression + "SELECT sum(a + b + c) FROM mytable" // more than 2 operands + }; + + for (String query : queries) { + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); + PinotQuery originalQuery = new PinotQuery(pinotQuery); // Deep copy for comparison Review Comment: The PinotQuery copy constructor performs a shallow copy, not a deep copy as the comment suggests. This means both objects share the same underlying expression objects, making the comparison unreliable for detecting optimization changes. ```suggestion PinotQuery originalQuery = deepCopyPinotQuery(pinotQuery); // Deep copy for comparison ``` ########## pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/AggregationOptimizer.java: ########## @@ -0,0 +1,390 @@ +/** + * 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.pinot.sql.parsers.rewriter; + +import java.util.List; +import org.apache.pinot.common.request.Expression; +import org.apache.pinot.common.request.ExpressionType; +import org.apache.pinot.common.request.Function; +import org.apache.pinot.common.request.Literal; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.utils.request.RequestUtils; + + +/** + * AggregationOptimizer optimizes aggregation functions by leveraging mathematical properties. + * Currently supports: + * - sum(column + constant) → sum(column) + constant * count(1) + * - sum(column - constant) → sum(column) - constant * count(1) + * - sum(constant + column) → sum(column) + constant * count(1) + * - sum(constant - column) → constant * count(1) - sum(column) + */ +public class AggregationOptimizer implements QueryRewriter { + + @Override + public PinotQuery rewrite(PinotQuery pinotQuery) { + List<Expression> selectList = pinotQuery.getSelectList(); + if (selectList != null) { + for (int i = 0; i < selectList.size(); i++) { + Expression expression = selectList.get(i); + Expression optimized = optimizeExpression(expression); + if (optimized != null) { + selectList.set(i, optimized); + } + } + } + return pinotQuery; + } + + /** + * Optimizes an expression if it matches supported patterns. + * Returns the optimized expression or null if no optimization is possible. + */ + private Expression optimizeExpression(Expression expression) { + if (expression.getType() != ExpressionType.FUNCTION) { + return null; + } + + Function function = expression.getFunctionCall(); + if (function == null) { + return null; + } + + String operator = function.getOperator().toLowerCase(); Review Comment: Calling toLowerCase() on every function operator creates unnecessary string objects. Consider using equalsIgnoreCase() or caching the lowercase versions, especially since this optimization runs on every query. ```suggestion String operator = function.getOperator(); ``` ########## pinot-common/src/main/java/org/apache/pinot/sql/parsers/rewriter/AggregationOptimizer.java: ########## @@ -0,0 +1,390 @@ +/** + * 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.pinot.sql.parsers.rewriter; + +import java.util.List; +import org.apache.pinot.common.request.Expression; +import org.apache.pinot.common.request.ExpressionType; +import org.apache.pinot.common.request.Function; +import org.apache.pinot.common.request.Literal; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.utils.request.RequestUtils; + + +/** + * AggregationOptimizer optimizes aggregation functions by leveraging mathematical properties. + * Currently supports: + * - sum(column + constant) → sum(column) + constant * count(1) + * - sum(column - constant) → sum(column) - constant * count(1) + * - sum(constant + column) → sum(column) + constant * count(1) + * - sum(constant - column) → constant * count(1) - sum(column) + */ +public class AggregationOptimizer implements QueryRewriter { + + @Override + public PinotQuery rewrite(PinotQuery pinotQuery) { + List<Expression> selectList = pinotQuery.getSelectList(); + if (selectList != null) { + for (int i = 0; i < selectList.size(); i++) { + Expression expression = selectList.get(i); + Expression optimized = optimizeExpression(expression); + if (optimized != null) { + selectList.set(i, optimized); + } + } + } + return pinotQuery; + } + + /** + * Optimizes an expression if it matches supported patterns. + * Returns the optimized expression or null if no optimization is possible. + */ + private Expression optimizeExpression(Expression expression) { + if (expression.getType() != ExpressionType.FUNCTION) { + return null; + } + + Function function = expression.getFunctionCall(); + if (function == null) { + return null; + } + + String operator = function.getOperator().toLowerCase(); + List<Expression> operands = function.getOperands(); + + if (operands == null || operands.size() != 1) { + return null; + } + + Expression operand = operands.get(0); + + switch (operator) { + case "sum": + return optimizeSumExpression(operand); + case "avg": + return optimizeAvgExpression(operand); + case "min": + return optimizeMinExpression(operand); + case "max": + return optimizeMaxExpression(operand); + default: + return null; + } + } + + /** + * Optimizes sum(expression) based on the expression type. + */ + private Expression optimizeSumExpression(Expression sumOperand) { + return optimizeArithmeticExpression(sumOperand, "sum"); + } + + /** + * Optimizes avg(expression) based on the expression type. + * AVG(column + constant) = AVG(column) + constant + * AVG(column - constant) = AVG(column) - constant + * AVG(constant - column) = constant - AVG(column) + * AVG(column * constant) = AVG(column) * constant + */ + private Expression optimizeAvgExpression(Expression avgOperand) { + return optimizeArithmeticExpression(avgOperand, "avg"); + } + + /** + * Optimizes min(expression) based on the expression type. + * MIN(column + constant) = MIN(column) + constant + * MIN(column - constant) = MIN(column) - constant + * MIN(constant - column) = constant - MAX(column) + * MIN(column * constant) = MIN(column) * constant (if constant > 0) + * = MAX(column) * constant (if constant < 0) + */ + private Expression optimizeMinExpression(Expression minOperand) { + return optimizeArithmeticExpression(minOperand, "min"); + } + + /** + * Optimizes max(expression) based on the expression type. + * MAX(column + constant) = MAX(column) + constant + * MAX(column - constant) = MAX(column) - constant + * MAX(constant - column) = constant - MIN(column) + * MAX(column * constant) = MAX(column) * constant (if constant > 0) + * = MIN(column) * constant (if constant < 0) + */ + private Expression optimizeMaxExpression(Expression maxOperand) { + return optimizeArithmeticExpression(maxOperand, "max"); + } + + /** + * Generic method to optimize arithmetic expressions for different aggregation functions. + */ + private Expression optimizeArithmeticExpression(Expression operand, String aggregationFunction) { + if (operand.getType() != ExpressionType.FUNCTION) { + return null; + } + + Function innerFunction = operand.getFunctionCall(); + if (innerFunction == null) { + return null; + } + + String operator = innerFunction.getOperator(); Review Comment: Inconsistent case handling for operators. Line 69 uses toLowerCase() but line 147 doesn't. This could lead to case-sensitivity bugs when comparing operators in switch statements. ```suggestion String operator = innerFunction.getOperator().toLowerCase(); ``` ########## pinot-core/src/test/java/org/apache/pinot/core/query/reduce/ReducerDataSchemaUtilsTest.java: ########## @@ -72,11 +72,11 @@ public void testCanonicalizeDataSchemaForGroupBy() { queryContext = QueryContextConverterUtils.getQueryContext( "SELECT SUM(col1 + 1), MIN(col2 + 2), col4 FROM testTable GROUP BY col3, col4"); // Intentionally make data schema not matching the string representation of the expression - dataSchema = new DataSchema(new String[]{"col3", "col4", "sum(col1+1)", "min(col2+2)"}, + dataSchema = new DataSchema(new String[]{"col3", "col4", "sum(col1)", "count(1)", "min(col2)"}, new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.DOUBLE, ColumnDataType.DOUBLE}); Review Comment: The test data schema now includes 5 columns but the ColumnDataType array still has only 4 elements (DOUBLE, DOUBLE on line 76). This mismatch could cause issues or make the test less realistic. ```suggestion new ColumnDataType[]{ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.DOUBLE, ColumnDataType.LONG, ColumnDataType.DOUBLE}); ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
