| /* |
| * 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.kylin.query.runtime.plan |
| |
| import org.apache.calcite.rel.core.AggregateCall |
| import org.apache.calcite.rex.RexLiteral |
| import org.apache.calcite.sql.SqlKind |
| import org.apache.kylin.common.KylinConfig |
| import org.apache.kylin.engine.spark.utils.LogEx |
| import org.apache.kylin.measure.percentile.PercentileCounter |
| import org.apache.kylin.metadata.model.FunctionDesc |
| import org.apache.kylin.query.relnode.{KapAggregateRel, KapProjectRel, KylinAggregateCall, OLAPAggregateRel} |
| import org.apache.kylin.query.util.RuntimeHelper |
| import org.apache.spark.sql.KapFunctions._ |
| import org.apache.spark.sql._ |
| import org.apache.spark.sql.catalyst.expressions.aggregate.ApproximatePercentile |
| import org.apache.spark.sql.catalyst.expressions.{CreateArray, In} |
| import org.apache.spark.sql.catalyst.plans.logical.Project |
| import org.apache.spark.sql.catalyst.util.ArrayData |
| import org.apache.spark.sql.execution.utils.SchemaProcessor |
| import org.apache.spark.sql.functions._ |
| import org.apache.spark.sql.types.{ArrayType, StructType} |
| import org.apache.spark.sql.udaf.SingleValueAgg |
| import org.apache.spark.sql.util.SparderTypeUtil |
| |
| import java.util.Locale |
| import scala.collection.JavaConverters._ |
| |
| // scalastyle:off |
| object AggregatePlan extends LogEx { |
| val binaryMeasureType = |
| List("PERCENTILE", "PERCENTILE_APPROX", "INTERSECT_COUNT", "COUNT_DISTINCT", "BITMAP_UUID", |
| FunctionDesc.FUNC_BITMAP_BUILD, FunctionDesc.FUNC_SUM_LC) |
| |
| def agg(inputs: java.util.List[DataFrame], |
| rel: KapAggregateRel): DataFrame = logTime("aggregate", debug = true) { |
| |
| var dataFrame = inputs.get(0) |
| val schemaNames = dataFrame.schema.fieldNames |
| val groupList = rel.getRewriteGroupKeys.asScala.map(groupId => col(schemaNames.apply(groupId))).toList |
| |
| if (rel.getContext != null && rel.getContext.isExactlyAggregate && !rel.getContext.isNeedToManyDerived) { |
| // exactly match, skip agg, direct project. |
| val aggCols = rel.getRewriteAggCalls.asScala.zipWithIndex.map { |
| case (call: KylinAggregateCall, index: Int) => |
| val funcName = OLAPAggregateRel.getAggrFuncName(call); |
| val dataType = call.getFunc.getReturnDataType |
| val argNames = call.getArgList.asScala.map(dataFrame.schema.names.apply(_)) |
| val columnName = argNames.map(col) |
| val hash = System.identityHashCode(rel).toString |
| funcName match { |
| case FunctionDesc.FUNC_COUNT_DISTINCT => |
| if (call.isHllCountDistinctFunc) { |
| val aggName = SchemaProcessor.replaceToAggravateSchemaName(index, "APPROX_COUNT_DISTINCT_DECODE", hash, argNames: _*) |
| KapFunctions.approx_count_distinct_decode(columnName.head, dataType.getPrecision).alias(aggName) |
| } else if (call.isBitmapCountDistinctFunc) { |
| if (rel.getContext.isExactlyFastBitmap) { |
| col(schemaNames.apply(call.getArgList.get(0))) |
| } else { |
| val aggName = SchemaProcessor.replaceToAggravateSchemaName(index, "PRECISE_COUNT_DISTINCT_DECODE", hash, argNames: _*) |
| KapFunctions.precise_count_distinct_decode(columnName.head).alias(aggName) |
| } |
| } else { |
| throw new IllegalArgumentException( |
| s"""Unsupported function name $funcName""") |
| } |
| case FunctionDesc.FUNC_PERCENTILE => |
| val aggName = SchemaProcessor.replaceToAggravateSchemaName(index, "PERCENTILE_DECODE", hash, argNames: _*) |
| KapFunctions.k_percentile_decode(columnName.head, columnName(1), dataType.getPrecision).alias(aggName) |
| case FunctionDesc.FUNC_BITMAP_BUILD => |
| val aggName = SchemaProcessor.replaceToAggravateSchemaName(index, "BITMAP_BUILD_DECODE", hash, argNames: _*) |
| KapFunctions.precise_bitmap_build_decode(columnName.head).alias(aggName) |
| case FunctionDesc.FUNC_SUM_LC => |
| val aggName = SchemaProcessor.replaceToAggravateSchemaName(index, "SUM_LC_DECODE", hash, argNames: _*) |
| val sparkDataType = SparderTypeUtil.toSparkType(dataType) |
| KapFunctions.k_sum_lc_decode(columnName.head, sparkDataType.json).alias(aggName) |
| case _ => |
| col(schemaNames.apply(call.getArgList.get(0))) |
| } |
| case (call: Any, _: Int) => |
| col(schemaNames.apply(call.getArgList.get(0))) |
| }.toList |
| |
| val prjList = groupList ++ aggCols |
| logInfo(s"Query exactly match index, skip agg, project $prjList.") |
| dataFrame.select(prjList: _*) |
| } else { |
| dataFrame = genFiltersWhenIntersectCount(rel, dataFrame) |
| val aggList = buildAgg(dataFrame.schema, rel) |
| val groupSets = rel.getRewriteGroupSets.asScala |
| .map(groupSet => groupSet.asScala.map(groupId => col(schemaNames.apply(groupId))).toList).toList |
| SparkOperation.agg(AggArgc(dataFrame, groupList, aggList, groupSets, rel.isSimpleGroupType)) |
| } |
| } |
| |
| private def genFiltersWhenIntersectCount(rel: KapAggregateRel, dataFrame: DataFrame): DataFrame = { |
| try { |
| val intersects = rel.getRewriteAggCalls.asScala.filter(_.isInstanceOf[KylinAggregateCall]) |
| .filter(!_.asInstanceOf[KylinAggregateCall].getFunc.isCount) |
| .map(_.asInstanceOf[KylinAggregateCall]) |
| .filter(call => !call.getFunc.isCount && OLAPAggregateRel.getAggrFuncName(call).equals(FunctionDesc.FUNC_INTERSECT_COUNT)) |
| val names = dataFrame.schema.names |
| val children = dataFrame.queryExecution.logical |
| if (intersects.nonEmpty && intersects.size == rel.getRewriteAggCalls.size() && children.isInstanceOf[Project]) { |
| // only exists intersect count function in agg |
| val list = children.asInstanceOf[Project].projectList |
| val supportGenFilter = intersects.forall { call => |
| val listIndex = call.getArgList.get(2) |
| call.getArgList.size() == 3 && list.apply(listIndex).eval().asInstanceOf[ArrayData].array.map(_.toString) |
| .forall(!_.contains(KylinConfig.getInstanceFromEnv.getIntersectFilterOrSeparator)) |
| } |
| if (supportGenFilter) { |
| val filters = intersects.map { call => |
| val filterColumnIndex = call.getArgList.get(1) |
| val litIndex = call.getArgList.get(2) |
| new Column(In(col(names(filterColumnIndex)).expr, list.apply(litIndex).children.head.asInstanceOf[CreateArray].children)) |
| } |
| val column = filters.reduceLeft(_.or(_)) |
| dataFrame.filter(column) |
| } else { |
| dataFrame |
| } |
| } else { |
| dataFrame |
| } |
| } catch { |
| case e: Throwable => logWarning("Error occurred when generate filters", e) |
| dataFrame |
| } |
| } |
| |
| def buildAgg(schema: StructType, |
| rel: KapAggregateRel): List[Column] = { |
| val hash = System.identityHashCode(rel).toString |
| |
| rel.getRewriteAggCalls.asScala.zipWithIndex.map { |
| case (call: KylinAggregateCall, index: Int) |
| if binaryMeasureType.contains(OLAPAggregateRel.getAggrFuncName(call)) => |
| val dataType = call.getFunc.getReturnDataType |
| val isCount = call.getFunc.isCount |
| val funcName = |
| if (isCount) FunctionDesc.FUNC_COUNT else OLAPAggregateRel.getAggrFuncName(call) |
| val argNames = call.getArgList.asScala.map(schema.names.apply(_)) |
| val columnName = argNames.map(col) |
| val registeredFuncName = RuntimeHelper.registerSingleByColName(funcName, dataType) |
| val aggName = SchemaProcessor.replaceToAggravateSchemaName(index, funcName, hash, argNames: _*) |
| if (funcName == FunctionDesc.FUNC_COUNT_DISTINCT) { |
| if (dataType.getName == "hllc") { |
| org.apache.spark.sql.KapFunctions |
| .approx_count_distinct(columnName.head, dataType.getPrecision) |
| .alias(aggName) |
| } else { |
| KapFunctions.precise_count_distinct(columnName.head).alias(aggName) |
| } |
| } else if (funcName.equalsIgnoreCase(FunctionDesc.FUNC_BITMAP_UUID)) { |
| KapFunctions.precise_bitmap_uuid(columnName.head).alias(aggName) |
| } else if (funcName.equalsIgnoreCase(FunctionDesc.FUNC_BITMAP_BUILD)) { |
| KapFunctions.precise_bitmap_build(columnName.head).alias(aggName) |
| } else if (funcName.equalsIgnoreCase(FunctionDesc.FUNC_INTERSECT_COUNT)) { |
| require(columnName.size >= 3, s"Input columns size ${columnName.size} don't greater than or equal to 3.") |
| val columns = columnName.slice(0, 3).zipWithIndex.map { |
| case (column: Column, 2) => column.cast(ArrayType.apply(schema.fields.apply(call.getArgList.get(1)).dataType)) |
| case (column: Column, _) => column |
| } |
| val separator = s"\\${KylinConfig.getInstanceFromEnv.getIntersectFilterOrSeparator}" |
| val upperBound = KylinConfig.getInstanceFromEnv.getBitmapValuesUpperBound |
| call.name.toUpperCase(Locale.ROOT) match { |
| case FunctionDesc.FUNC_INTERSECT_COUNT => KapFunctions.intersect_count(separator, upperBound, columns.toList: _*).alias(aggName) |
| case FunctionDesc.FUNC_INTERSECT_VALUE => KapFunctions.intersect_value(separator, upperBound, columns.toList: _*).alias(aggName) |
| case FunctionDesc.FUNC_INTERSECT_BITMAP_UUID => KapFunctions.intersect_bitmap(separator, upperBound, columns.toList: _*).alias(aggName) |
| case FunctionDesc.FUNC_INTERSECT_COUNT_V2 => KapFunctions.intersect_count_v2(columnName.last, separator, upperBound, columns.toList: _*).alias(aggName) |
| case FunctionDesc.FUNC_INTERSECT_VALUE_V2 => KapFunctions.intersect_value_v2(columnName.last, separator, upperBound, columns.toList: _*).alias(aggName) |
| case FunctionDesc.FUNC_INTERSECT_BITMAP_UUID_V2 => KapFunctions.intersect_bitmap_v2(columnName.last, separator, upperBound, columns.toList: _*).alias(aggName) |
| case func => throw new UnsupportedOperationException(s"Unsupported intersect count function: $func, please check the sql.") |
| } |
| } else if (funcName.equalsIgnoreCase(FunctionDesc.FUNC_PERCENTILE)) { |
| require(columnName.size == 2, s"Input columns size ${columnName.size} don't equal to 2.") |
| KapFunctions.k_percentile(columnName.head, columnName(1), dataType.getPrecision).alias(aggName) |
| } else if (funcName.equalsIgnoreCase(FunctionDesc.FUNC_SUM_LC)) { |
| KapFunctions.k_sum_lc(columnName.head, SparderTypeUtil.toSparkType(dataType)).alias(aggName) |
| } else { |
| callUDF(registeredFuncName, columnName.toList: _*).alias(aggName) |
| } |
| case (call: Any, index: Int) => |
| val funcName = OLAPAggregateRel.getAggrFuncName(call) |
| val schemaNames = schema.names |
| val argNames = call.getArgList.asScala.map(id => schemaNames.apply(id)) |
| val columnName = argNames.map(col) |
| val inputType = call.getType |
| val aggName = SchemaProcessor.replaceToAggravateSchemaName(index, |
| funcName, |
| hash, |
| argNames: _*) |
| funcName match { |
| case FunctionDesc.FUNC_PERCENTILE => |
| rel.getInput match { |
| case projectRel: KapProjectRel => |
| val percentageArg = projectRel.getChildExps.get(call.getArgList.get(1)) |
| val accuracyArg = if (call.getArgList.size() < 3) { None } else { Some(projectRel.getChildExps.get(call.getArgList.get(2))) } |
| (percentageArg, accuracyArg) match { |
| case (percentageLitRex: RexLiteral, accuracyArgLitRex: Option[RexLiteral]) => |
| if (KylinConfig.getInstanceFromEnv.getPercentileApproxAlgorithm.equalsIgnoreCase("t-digest")) { |
| KapFunctions.k_percentile(columnName.head, columnName(1), PercentileCounter.DEFAULT_PERCENTILE_ACCURACY).alias(aggName) |
| } else { |
| val percentage = percentageLitRex.getValue |
| val accuracy = accuracyArgLitRex.map(arg => arg.getValue).getOrElse(ApproximatePercentile.DEFAULT_PERCENTILE_ACCURACY) |
| percentile_approx(col(argNames.head), lit(percentage), lit(accuracy)).alias(aggName) |
| } |
| } |
| case _ => |
| throw new UnsupportedOperationException(s"Invalid percentile_approx parameters, " + |
| s"expecting approx_percentile(col, percentage [, accuracy]), percentage/accuracy must be of constant literal") |
| } |
| case FunctionDesc.FUNC_SUM => |
| if (isSum0(call)) { |
| sum0( |
| col(argNames.head).cast( |
| SparderTypeUtil.convertSqlTypeToSparkType(inputType))) |
| .alias(aggName) |
| } else { |
| sum( |
| col(argNames.head).cast( |
| SparderTypeUtil.convertSqlTypeToSparkType(inputType))) |
| .alias(aggName) |
| } |
| case FunctionDesc.FUNC_COUNT => |
| count(if (argNames.isEmpty) k_lit(1) else col(argNames.head)) |
| .alias(aggName) |
| case FunctionDesc.FUNC_MAX => |
| max( |
| col(argNames.head).cast( |
| SparderTypeUtil.convertSqlTypeToSparkType(inputType))) |
| .alias(aggName) |
| case FunctionDesc.FUNC_MIN => |
| min( |
| col(argNames.head).cast( |
| SparderTypeUtil.convertSqlTypeToSparkType(inputType))) |
| .alias(aggName) |
| case FunctionDesc.FUNC_COUNT_DISTINCT if call.getAggregation.getName == "BITMAP_COUNT" => |
| KapFunctions.precise_count_distinct(col(argNames.head)).alias(aggName) |
| case FunctionDesc.FUNC_COUNT_DISTINCT => |
| countDistinct(argNames.head, argNames.drop(1): _*) |
| .alias(aggName) |
| case FunctionDesc.FUNC_BITMAP_BUILD => |
| KapFunctions.precise_bitmap_build_pushdown(columnName.head).alias(aggName) |
| // Issue 4337: Supported select (select '2012-01-02') as data, xxx from table group by xxx |
| case SqlKind.SINGLE_VALUE.sql => |
| SingleValueAgg(schema.head).apply(col(argNames.head)).alias(aggName) |
| case FunctionDesc.FUNC_GROUPING => |
| if (!rel.isSimpleGroupType) { |
| grouping(argNames.head).alias(aggName) |
| } else { |
| if (rel.getRewriteGroupKeys.contains(call.getArgList.get(0))) { |
| k_lit(0).alias(aggName) |
| } else { |
| k_lit(1).alias(aggName) |
| } |
| } |
| case FunctionDesc.FUNC_COLLECT_SET => |
| call match { |
| case kac: KylinAggregateCall => |
| array_distinct(flatten(collect_set(col(argNames.head)))) |
| .alias(aggName) |
| case _ => |
| collect_set(col(argNames.head)).alias(aggName) |
| } |
| case _ => |
| throw new IllegalArgumentException( |
| s"""Unsupported function name $funcName""") |
| } |
| }.toList |
| } |
| |
| private def isSum0(call: AggregateCall) = { |
| call.isInstanceOf[KylinAggregateCall] && call |
| .asInstanceOf[KylinAggregateCall] |
| .isSum0 |
| } |
| } |
