Jackie-Jiang commented on a change in pull request #6998: URL: https://github.com/apache/incubator-pinot/pull/6998#discussion_r656506126
########## File path: pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/JsonStatementOptimizer.java ########## @@ -0,0 +1,516 @@ +/** + * 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.core.query.optimizer.statement; + +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.pinot.common.function.scalar.ArithmeticFunctions; +import org.apache.pinot.common.function.scalar.DateTimeFunctions; +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.Identifier; +import org.apache.pinot.common.request.Literal; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.request.RequestUtils; +import org.apache.pinot.pql.parsers.pql2.ast.FilterKind; +import org.apache.pinot.pql.parsers.pql2.ast.FloatingPointLiteralAstNode; +import org.apache.pinot.pql.parsers.pql2.ast.IntegerLiteralAstNode; +import org.apache.pinot.pql.parsers.pql2.ast.LiteralAstNode; +import org.apache.pinot.pql.parsers.pql2.ast.StringLiteralAstNode; +import org.apache.pinot.segment.spi.AggregationFunctionType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.Pair; + + +/** + * This class will rewrite a query that has json path expressions into a query that uses JSON_EXTRACT_SCALAR and + * JSON_MATCH functions. + * + * Example 1: + * From : SELECT jsonColumn.name.first + * FROM testTable + * WHERE jsonColumn.name.first IS NOT NULL + * TO : SELECT JSON_EXTRACT_SCALAR(jsonColumn, '$.name.first', 'STRING', 'null') AS jsonColum.name.first + * FROM testTable + * WHERE JSON_MATCH('"$.name.first" IS NOT NULL') + * + * Output datatype of any json path expression is 'STRING'. However, if json path expression appears as an argument to + * a numerical function, then output of json path expression is set to 'DOUBLE' as shown in the example below. + * + * Example 2: + * From: SELECT MIN(jsonColumn.id - 5) + * FROM testTable + * WHERE jsonColumn.id IS NOT NULL + * To: SELECT MIN(MINUS(JSON_EXTRACT_SCALAR(jsonColumn, '$.id', 'DOUBLE', Double.NEGATIVE_INFINITY),5)) AS min(minus(jsonColum.id, '5')) + * FROM testTable + * WHERE JSON_MATCH('"$.id" IS NOT NULL') + * + * Example 3: + * From: SELECT jsonColumn.id, count(*) + * FROM testTable + * WHERE jsonColumn.name.first = 'Daffy' OR jsonColumn.id = 101 + * GROUP BY jsonColumn.id + * To: SELECT JSON_EXTRACT_SCALAR(jsonColumn, '$.id', 'STRING', 'null') AS jsonColumn.id, count(*) + * FROM testTable + * WHERE JSON_MATCH('"$.name.first" = ''Daffy''') OR JSON_MATCH('"$.id" = 101') + * GROUP BY JSON_EXTRACT_SCALAR(jsonColumn, '$.id', 'STRING', 'null'); + * + * Example 4: + * From: SELECT jsonColumn.name.last, count(*) + * FROM testTable + * GROUP BY jsonColumn.name.last + * HAVING jsonColumn.name.last = 'mouse' + * To: SELECT JSON_EXTRACT_SCALAR(jsonColumn, '$.name.last', 'STRING', 'null') AS jsonColumn.name.last, count(*) + * FROM testTable + * GROUP BY JSON_EXTRACT_SCALAR(jsonColumn, '$.name.last', 'STRING', 'null') + * HAVING JSON_EXTRACT_SCALAR(jsonColumn, '$.name.last', 'STRING', 'null') = 'mouse' + * + * Notes: + * 1) In a filter expression, if json path appears on the left-hand side, the right-hand side must be a literal. In + * future this can be changed to have any expression on the right-hand side by implementing a function that would + * convert any {@link Expression} into SQL fragment that can be used in JSON_MATCH. Currently only literals are + * converted into SQL fragments {see @link #getLiteralSQL} function. + * 2) In WHERE clause each json path expression will be replaced with a JSON_MATCH function. If there are multiple + * json path expressions, they will be replaced by multiple JSON_MATCH functions. We currently don't fold multiple + * JSON_MATCH functions into a single JSON_MATCH_FUNCTION. + */ +public class JsonStatementOptimizer implements StatementOptimizer { + + /** + * Maintain a list of numerical functions that requiring json path expression to output numerical values. This allows + * us to implicitly convert the output of json path expression to DOUBLE. TODO: There are better ways of doing this + * if we were to move to a new storage (currently STRING) for JSON or functions were to pre-declare their input + * data types. + */ + private static Set<String> numericalFunctions = getNumericalFunctionList(); + + /** + * A list of functions that require json path expression to output LONG value. This allows us to implicitly convert + * the output of json path expression to LONG. + */ + private static Set<String> datetimeFunctions = getDateTimeFunctionList(); + + /** + * Null value constants for different column types. Used while rewriting json path expression to JSON_EXTRACT_SCALAR function. + */ + private static LiteralAstNode DEFAULT_DIMENSION_NULL_VALUE_OF_INT_AST = new IntegerLiteralAstNode(FieldSpec.DEFAULT_DIMENSION_NULL_VALUE_OF_INT); + private static LiteralAstNode DEFAULT_DIMENSION_NULL_VALUE_OF_LONG_AST = new IntegerLiteralAstNode(FieldSpec.DEFAULT_DIMENSION_NULL_VALUE_OF_LONG); + private static LiteralAstNode DEFAULT_DIMENSION_NULL_VALUE_OF_FLOAT_AST = new FloatingPointLiteralAstNode(FieldSpec.DEFAULT_DIMENSION_NULL_VALUE_OF_FLOAT); + private static LiteralAstNode DEFAULT_DIMENSION_NULL_VALUE_OF_DOUBLE_AST = new FloatingPointLiteralAstNode(FieldSpec.DEFAULT_DIMENSION_NULL_VALUE_OF_DOUBLE); + private static LiteralAstNode DEFAULT_DIMENSION_NULL_VALUE_OF_STRING_AST = new StringLiteralAstNode(FieldSpec.DEFAULT_DIMENSION_NULL_VALUE_OF_STRING); + + @Override + public void optimize(PinotQuery query, @Nullable Schema schema) { + // In SELECT clause, replace JSON path expressions with JSON_EXTRACT_SCALAR function with an alias. + List<Expression> expressions = query.getSelectList(); + for (Expression expression : expressions) { + Pair<String, Boolean> result = optimizeJsonIdentifier(expression, schema, DataSchema.ColumnDataType.STRING, true); + if (expression.getType() == ExpressionType.FUNCTION && !expression.getFunctionCall().getOperator().equals("AS") + && result.getSecond()) { + // Since this is not an AS function (user-specified alias) and the function or its arguments contain json path + // expression, set an alias for the expression after replacing json path expression with JSON_EXTRACT_SCALAR + // function. + Function aliasFunction = getAliasFunction(result.getFirst(), expression.getFunctionCall()); + expression.setFunctionCall(aliasFunction); + } + } + + // In WHERE clause, replace JSON path expressions with JSON_MATCH function. + Expression filter = query.getFilterExpression(); + if (filter != null) { + optimizeJsonPredicate(filter, schema); + } + + // In GROUP BY clause, replace JSON path expressions with JSON_EXTRACT_SCALAR function without an alias. + expressions = query.getGroupByList(); + if (expressions != null) { + for (Expression expression : expressions) { + optimizeJsonIdentifier(expression, schema, DataSchema.ColumnDataType.STRING, false); + } + } + + // In ORDER BY clause, replace JSON path expression with JSON_EXTRACT_SCALAR. This expression must match the + // corresponding SELECT list expression except for the alias. + expressions = query.getOrderByList(); + if (expressions != null) { + for (Expression expression : expressions) { + optimizeJsonIdentifier(expression, schema, DataSchema.ColumnDataType.STRING, false); + } + } + + // In HAVING clause, replace JSON path expressions with JSON_EXTRACT_SCALAR. This expression must match the + // corresponding SELECT list expression except for the alias. + Expression expression = query.getHavingExpression(); + if (expression != null) { + optimizeJsonIdentifier(expression, schema, DataSchema.ColumnDataType.STRING, false); + } + } + + /** + * Replace an json path expression with an aliased JSON_EXTRACT_SCALAR function. + * @param expression input expression to rewrite into JSON_EXTRACT_SCALAR function if the expression is json path. + * @param outputDataType to keep track of output datatype of JSON_EXTRACT_SCALAR function which depends upon the outer + * function that json path expression appears in. + * @return A {@link Pair} of values where the first value is alias for the input expression and second + * value indicates whether json path expression was found (true) or not (false) in the expression. + */ + private static Pair<String, Boolean> optimizeJsonIdentifier(Expression expression, Schema schema, DataSchema.ColumnDataType outputDataType, + boolean hasColumnAlias) { Review comment: `hasColumnAlias` is unused? ########## File path: pinot-core/src/main/java/org/apache/pinot/core/query/optimizer/statement/JsonStatementOptimizer.java ########## @@ -0,0 +1,516 @@ +/** + * 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.core.query.optimizer.statement; + +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.pinot.common.function.scalar.ArithmeticFunctions; +import org.apache.pinot.common.function.scalar.DateTimeFunctions; +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.Identifier; +import org.apache.pinot.common.request.Literal; +import org.apache.pinot.common.request.PinotQuery; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.request.RequestUtils; +import org.apache.pinot.pql.parsers.pql2.ast.FilterKind; +import org.apache.pinot.pql.parsers.pql2.ast.FloatingPointLiteralAstNode; +import org.apache.pinot.pql.parsers.pql2.ast.IntegerLiteralAstNode; +import org.apache.pinot.pql.parsers.pql2.ast.LiteralAstNode; +import org.apache.pinot.pql.parsers.pql2.ast.StringLiteralAstNode; +import org.apache.pinot.segment.spi.AggregationFunctionType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.Pair; + + +/** + * This class will rewrite a query that has json path expressions into a query that uses JSON_EXTRACT_SCALAR and + * JSON_MATCH functions. Review comment: We can only rewrite the query to `JSON_MATCH` when json index is available. One option is to read that info from the table config, but segments on server might not be always in-sync with the table config ########## File path: pinot-core/src/main/java/org/apache/pinot/core/operator/combine/BaseCombineOperator.java ########## @@ -163,7 +163,7 @@ protected void processSegments(int taskIndex) { LOGGER.error("Caught exception while executing operator of index: {} (query: {})", operatorIndex, _queryContext, e); _blockingQueue.offer(new IntermediateResultsBlock(e)); - return; + throw e; Review comment: (Critical) Why changing this? The exception is already embedded in the results block, and we don't want to kill the working thread -- 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. For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@pinot.apache.org For additional commands, e-mail: commits-h...@pinot.apache.org
