korlov42 commented on code in PR #4221: URL: https://github.com/apache/ignite-3/pull/4221#discussion_r1731326114
########## modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/PlannerHelper.java: ########## @@ -343,4 +355,164 @@ static boolean hasSubQuery(SqlNode node) { return super.visit(call); } } + + /** + * Tries to optimize a query that looks like {@code SELECT count(*)}. + * + * @param planner Planner. + * @param node Query node. + * @return Plan node with list of aliases, if the optimization is applicable. + */ + public static @Nullable Pair<IgniteRel, List<String>> tryOptimizeSelectCount( + IgnitePlanner planner, + SqlNode node + ) { + SqlSelect select = getSelectCountOptimizationNode(node); + if (select == null) { + return null; + } + + assert select.getFrom() != null : "FROM is missing"; + + IgniteSqlToRelConvertor converter = planner.sqlToRelConverter(); + + // Convert PUBLIC.T AS X (a,b) to PUBLIC.T + SqlNode from = SqlUtil.stripAs(select.getFrom()); + RelOptTable targetTable = converter.getTargetTable(from); + + IgniteDataSource dataSource = targetTable.unwrap(IgniteDataSource.class); + if (!(dataSource instanceof IgniteTable)) { + return null; + } + + IgniteTypeFactory typeFactory = planner.getTypeFactory(); + + // SELECT COUNT(*) ... row type + RelDataType countResultType = typeFactory.createSqlType(SqlTypeName.BIGINT); + + // Build projection + // Rewrites SELECT count(*) ... as Project(exprs = [lit, $0, ... ]), where $0 references a row that stores a count. + // So we can feed results of get count operation into a projection to compute final results. + + List<RexNode> expressions = new ArrayList<>(); + List<String> expressionNames = new ArrayList<>(); + boolean countAdded = false; + + for (SqlNode selectItem : select.getSelectList()) { + SqlNode expr = SqlUtil.stripAs(selectItem); + + if (isCountStar(expr)) { + RexBuilder rexBuilder = planner.cluster().getRexBuilder(); + RexSlot countValRef = rexBuilder.makeInputRef(countResultType, 0); + + expressions.add(countValRef); + + countAdded = true; + } else if (expr instanceof SqlLiteral || expr instanceof SqlDynamicParam) { + RexNode rexNode = converter.convertExpression(expr); + + expressions.add(rexNode); + } else { + return null; + } + + String alias = planner.validator().deriveAlias(selectItem, expressionNames.size()); + expressionNames.add(alias); + } + + if (!countAdded) { + return null; + } + + IgniteSelectCount rel = new IgniteSelectCount( + planner.cluster(), + planner.cluster().traitSetOf(IgniteConvention.INSTANCE), + targetTable, + expressions + ); + + return new Pair<>(rel, expressionNames); + } + + private static @Nullable SqlSelect getSelectCountOptimizationNode(SqlNode node) { + // Unwrap SELECT .. from SELECT x FROM t ORDER BY ... + if (node instanceof SqlOrderBy) { + SqlOrderBy orderBy = (SqlOrderBy) node; + + // Skip ORDER BY with OFFSET/FETCH + if (orderBy.fetch != null || orderBy.offset != null) { + return null; + } + + // Skip ORDER BY with non literals + for (SqlNode arg : orderBy.orderList) { + if (!SqlUtil.isLiteral(arg)) { + return null; + } + } + + assert orderBy.getOperandList().size() == 4 : "Expected 4 operands, but was " + orderBy.getOperandList().size(); + node = orderBy.query; + } + + if (!(node instanceof SqlSelect)) { + return null; + } + + SqlSelect select = (SqlSelect) node; + + if (select.getGroup() != null + || select.getFrom() == null + || select.getWhere() != null + || select.getHaving() != null + || select.getQualify() != null + || !select.getWindowList().isEmpty() + || select.getOffset() != null + || select.getFetch() != null) { + + return null; + } + + // make sure that the following IF statement does not leave out any operand of the SELECT node + assert select.getOperandList().size() == 12 : "Expected 12 operands, but was " + select.getOperandList().size(); + + // Convert PUBLIC.T AS X (a,b) to PUBLIC.T + SqlNode from = SqlUtil.stripAs(select.getFrom()); + // Skip non-references such as VALUES .. + if (from.getKind() == SqlKind.IDENTIFIER) { + return select; + } else { + return null; + } + } + + private static boolean isCountStar(SqlNode node) { + if (!SqlUtil.isCallTo(node, SqlStdOperatorTable.COUNT)) { + // The SQL node was checked by the validator and the call has correct number of arguments. + return false; + } else { + SqlCall call = (SqlCall) node; + // Reject COUNT(DISTINCT ...) + if (call.getFunctionQuantifier() != null) { + return false; + } + if (call.getOperandList().isEmpty()) { + return false; + } + SqlNode operand = call.getOperandList().get(0); + return !isCountStartOperand(operand); + } + } + + private static boolean isCountStartOperand(SqlNode node) { Review Comment: name of this method is quite confusing. Why `NULL` is countStart operand while any literal is not? -- 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: notifications-unsubscr...@ignite.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org