weixiangsun commented on a change in pull request #8029: URL: https://github.com/apache/pinot/pull/8029#discussion_r831473442
########## File path: pinot-core/src/main/java/org/apache/pinot/core/util/GapfillUtils.java ########## @@ -116,7 +132,265 @@ static public Serializable getDefaultValue(DataSchema.ColumnDataType dataType) { } } - private static String canonicalizeFunctionName(String functionName) { - return StringUtils.remove(functionName, '_').toLowerCase(); + public static boolean isGapfill(ExpressionContext expressionContext) { + if (expressionContext.getType() != ExpressionContext.Type.FUNCTION) { + return false; + } + + return GAP_FILL.equals(expressionContext.getFunction().getFunctionName()); + } + + private static boolean isGapfill(QueryContext queryContext) { + for (ExpressionContext expressionContext : queryContext.getSelectExpressions()) { + if (isGapfill(expressionContext)) { + return true; + } + } + return false; + } + + /** + * Get the gapfill type for queryContext. Also do the validation for gapfill request. + * @param queryContext + */ + public static GapfillType getGapfillType(QueryContext queryContext) { + GapfillType gapfillType = null; + if (queryContext.getSubquery() == null) { + if (isGapfill(queryContext)) { + Preconditions.checkArgument(queryContext.getAggregationFunctions() == null, + "Aggregation and Gapfill can not be in the same sql statement."); + gapfillType = GapfillType.GAP_FILL; + } + } else if (isGapfill(queryContext)) { + Preconditions.checkArgument(queryContext.getSubquery().getAggregationFunctions() != null, + "Select and Gapfill should be in the same sql statement."); + Preconditions.checkArgument(queryContext.getSubquery().getSubquery() == null, + "There is no three levels nesting sql when the outer query is gapfill."); + gapfillType = GapfillType.AGGREGATE_GAP_FILL; + } else if (isGapfill(queryContext.getSubquery())) { + if (queryContext.getAggregationFunctions() == null) { + gapfillType = GapfillType.GAP_FILL_SELECT; + } else if (queryContext.getSubquery().getSubquery() == null) { + gapfillType = GapfillType.GAP_FILL_AGGREGATE; + } else { + Preconditions + .checkArgument(queryContext.getSubquery().getSubquery().getAggregationFunctions() != null, + "Select cannot happen before gapfill."); + gapfillType = GapfillType.AGGREGATE_GAP_FILL_AGGREGATE; + } + } + + if (gapfillType == null) { + return gapfillType; + } + + ExpressionContext gapFillSelection = GapfillUtils.getGapfillExpressionContext(queryContext, gapfillType); + + Preconditions.checkArgument(gapFillSelection != null && gapFillSelection.getFunction() != null, + "Gapfill Expression should be function."); + List<ExpressionContext> args = gapFillSelection.getFunction().getArguments(); + Preconditions.checkArgument(args.size() > 5, "PreAggregateGapFill does not have correct number of arguments."); + Preconditions.checkArgument(args.get(1).getLiteral() != null, + "The second argument of PostAggregateGapFill should be TimeFormatter."); + Preconditions.checkArgument(args.get(2).getLiteral() != null, + "The third argument of PostAggregateGapFill should be start time."); + Preconditions.checkArgument(args.get(3).getLiteral() != null, + "The fourth argument of PostAggregateGapFill should be end time."); + Preconditions.checkArgument(args.get(4).getLiteral() != null, + "The fifth argument of PostAggregateGapFill should be time bucket size."); + + ExpressionContext timeseriesOn = GapfillUtils.getTimeSeriesOnExpressionContext(gapFillSelection); + Preconditions.checkArgument(timeseriesOn != null, "The TimeSeriesOn expressions should be specified."); + + if (queryContext.getAggregationFunctions() == null) { + return gapfillType; + } + + List<ExpressionContext> groupbyExpressions = queryContext.getGroupByExpressions(); + Preconditions.checkArgument(groupbyExpressions != null, "No GroupBy Clause."); + List<ExpressionContext> innerSelections = queryContext.getSubquery().getSelectExpressions(); + String timeBucketCol = null; + List<String> strAlias = queryContext.getSubquery().getAliasList(); + for (int i = 0; i < innerSelections.size(); i++) { + ExpressionContext innerSelection = innerSelections.get(i); + if (GapfillUtils.isGapfill(innerSelection)) { + if (strAlias.get(i) != null) { + timeBucketCol = strAlias.get(i); + } else { + timeBucketCol = innerSelection.getFunction().getArguments().get(0).toString(); + } + break; + } + } + + Preconditions.checkArgument(timeBucketCol != null, "No Group By timebucket."); + + boolean findTimeBucket = false; + for (ExpressionContext groupbyExp : groupbyExpressions) { + if (timeBucketCol.equals(groupbyExp.toString())) { + findTimeBucket = true; + break; + } + } + + Preconditions.checkArgument(findTimeBucket, "No Group By timebucket."); + return gapfillType; + } + + private static ExpressionContext findGapfillExpressionContext(QueryContext queryContext) { + for (ExpressionContext expressionContext : queryContext.getSelectExpressions()) { + if (isGapfill(expressionContext)) { + return expressionContext; + } + } + return null; + } + + public static ExpressionContext getGapfillExpressionContext(QueryContext queryContext, GapfillType gapfillType) { + if (gapfillType == GapfillType.AGGREGATE_GAP_FILL || gapfillType == GapfillType.GAP_FILL) { + return findGapfillExpressionContext(queryContext); + } else if (gapfillType == GapfillType.GAP_FILL_AGGREGATE || gapfillType == GapfillType.AGGREGATE_GAP_FILL_AGGREGATE + || gapfillType == GapfillType.GAP_FILL_SELECT) { + return findGapfillExpressionContext(queryContext.getSubquery()); + } else { + return null; + } + } + + public static int findTimeBucketColumnIndex(QueryContext queryContext, GapfillType gapfillType) { + if (gapfillType == GapfillType.GAP_FILL_AGGREGATE + || gapfillType == GapfillType.GAP_FILL_SELECT + || gapfillType == GapfillType.AGGREGATE_GAP_FILL_AGGREGATE) { + queryContext = queryContext.getSubquery(); + } + List<ExpressionContext> expressionContexts = queryContext.getSelectExpressions(); + for (int i = 0; i < expressionContexts.size(); i++) { + if (isGapfill(expressionContexts.get(i))) { + return i; + } + } + return -1; + } + + public static ExpressionContext getTimeSeriesOnExpressionContext(ExpressionContext gapFillSelection) { + List<ExpressionContext> args = gapFillSelection.getFunction().getArguments(); + for (int i = STARTING_INDEX_OF_OPTIONAL_ARGS_FOR_PRE_AGGREGATE_GAP_FILL; i < args.size(); i++) { + if (GapfillUtils.isTimeSeriesOn(args.get(i))) { + return args.get(i); + } + } + return null; + } + + public static Map<String, ExpressionContext> getFillExpressions(ExpressionContext gapFillSelection) { + Map<String, ExpressionContext> fillExpressions = new HashMap<>(); + List<ExpressionContext> args = gapFillSelection.getFunction().getArguments(); + for (int i = STARTING_INDEX_OF_OPTIONAL_ARGS_FOR_PRE_AGGREGATE_GAP_FILL; i < args.size(); i++) { + if (GapfillUtils.isFill(args.get(i))) { + ExpressionContext fillExpression = args.get(i); + fillExpressions.put(fillExpression.getFunction().getArguments().get(0).getIdentifier(), fillExpression); + } + } + return fillExpressions; + } + + public static BrokerRequest stripGapfill(BrokerRequest brokerRequest) { + if (brokerRequest.getPinotQuery().getDataSource() == null) { + return brokerRequest; + } + PinotQuery pinotQuery = brokerRequest.getPinotQuery(); + if (pinotQuery.getDataSource().getSubquery() == null && !hasGapfill(pinotQuery)) { + return brokerRequest; + } + + // carry over the query options from original query to server query. + Map<String, String> queryOptions = pinotQuery.getQueryOptions(); Review comment: Fixed -- 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: commits-unsubscr...@pinot.apache.org 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