in gluten-substrait/src/main/scala/org/apache/gluten/backendsapi/SparkPlanExecApi.scala [465:613]
def extraExpressionConverter(
substraitExprName: String,
expr: Expression,
attributeSeq: Seq[Attribute]): Option[ExpressionTransformer] =
None
/**
* Define whether the join operator is fallback because of the join operator is not supported by
* backend
*/
def joinFallback(
JoinType: JoinType,
leftOutputSet: AttributeSet,
right: AttributeSet,
condition: Option[Expression]): Boolean = false
/** default function to generate window function node */
def genWindowFunctionsNode(
windowExpression: Seq[NamedExpression],
windowExpressionNodes: JList[WindowFunctionNode],
originalInputAttributes: Seq[Attribute],
context: SubstraitContext): Unit = {
windowExpression.map {
windowExpr =>
val aliasExpr = windowExpr.asInstanceOf[Alias]
val columnName = s"${aliasExpr.name}_${aliasExpr.exprId.id}"
val wExpression = aliasExpr.child.asInstanceOf[WindowExpression]
wExpression.windowFunction match {
case wf @ (RowNumber() | Rank(_) | DenseRank(_) | CumeDist() | PercentRank(_)) =>
val aggWindowFunc = wf.asInstanceOf[AggregateWindowFunction]
val frame = aggWindowFunc.frame.asInstanceOf[SpecifiedWindowFrame]
val windowFunctionNode = ExpressionBuilder.makeWindowFunction(
WindowFunctionsBuilder.create(context, aggWindowFunc).toInt,
new JArrayList[ExpressionNode](),
columnName,
ConverterUtils.getTypeNode(aggWindowFunc.dataType, aggWindowFunc.nullable),
frame.upper,
frame.lower,
frame.frameType.sql,
originalInputAttributes.asJava
)
windowExpressionNodes.add(windowFunctionNode)
case aggExpression: AggregateExpression =>
val frame = wExpression.windowSpec.frameSpecification.asInstanceOf[SpecifiedWindowFrame]
val aggregateFunc = aggExpression.aggregateFunction
val substraitAggFuncName = ExpressionMappings.expressionsMap.get(aggregateFunc.getClass)
if (substraitAggFuncName.isEmpty) {
throw new GlutenNotSupportException(s"Not currently supported: $aggregateFunc.")
}
val childrenNodeList = aggregateFunc.children
.map(
ExpressionConverter
.replaceWithExpressionTransformer(_, originalInputAttributes)
.doTransform(context))
.asJava
val windowFunctionNode = ExpressionBuilder.makeWindowFunction(
AggregateFunctionsBuilder.create(context, aggExpression.aggregateFunction).toInt,
childrenNodeList,
columnName,
ConverterUtils.getTypeNode(aggExpression.dataType, aggExpression.nullable),
frame.upper,
frame.lower,
frame.frameType.sql,
originalInputAttributes.asJava
)
windowExpressionNodes.add(windowFunctionNode)
case wf @ (_: Lead | _: Lag) =>
val offsetWf = wf.asInstanceOf[FrameLessOffsetWindowFunction]
val frame = offsetWf.frame.asInstanceOf[SpecifiedWindowFrame]
val childrenNodeList = new JArrayList[ExpressionNode]()
childrenNodeList.add(
ExpressionConverter
.replaceWithExpressionTransformer(
offsetWf.input,
attributeSeq = originalInputAttributes)
.doTransform(context))
// Spark only accepts foldable offset. Converts it to LongType literal.
val offset = offsetWf.offset.eval(EmptyRow).asInstanceOf[Int]
// Velox only allows negative offset. WindowFunctionsBuilder#create converts
// lag/lead with negative offset to the function with positive offset. So just
// makes offsetNode store positive value.
val offsetNode = ExpressionBuilder.makeLiteral(Math.abs(offset.toLong), LongType, false)
childrenNodeList.add(offsetNode)
// NullType means Null is the default value. Don't pass it to native.
if (offsetWf.default.dataType != NullType) {
childrenNodeList.add(
ExpressionConverter
.replaceWithExpressionTransformer(
offsetWf.default,
attributeSeq = originalInputAttributes)
.doTransform(context))
}
val windowFunctionNode = ExpressionBuilder.makeWindowFunction(
WindowFunctionsBuilder.create(context, offsetWf).toInt,
childrenNodeList,
columnName,
ConverterUtils.getTypeNode(offsetWf.dataType, offsetWf.nullable),
frame.upper,
frame.lower,
frame.frameType.sql,
offsetWf.ignoreNulls,
originalInputAttributes.asJava
)
windowExpressionNodes.add(windowFunctionNode)
case wf @ NthValue(input, offset: Literal, ignoreNulls: Boolean) =>
val frame = wExpression.windowSpec.frameSpecification.asInstanceOf[SpecifiedWindowFrame]
val childrenNodeList = new JArrayList[ExpressionNode]()
childrenNodeList.add(
ExpressionConverter
.replaceWithExpressionTransformer(input, attributeSeq = originalInputAttributes)
.doTransform(context))
childrenNodeList.add(LiteralTransformer(offset).doTransform(context))
val windowFunctionNode = ExpressionBuilder.makeWindowFunction(
WindowFunctionsBuilder.create(context, wf).toInt,
childrenNodeList,
columnName,
ConverterUtils.getTypeNode(wf.dataType, wf.nullable),
frame.upper,
frame.lower,
frame.frameType.sql,
ignoreNulls,
originalInputAttributes.asJava
)
windowExpressionNodes.add(windowFunctionNode)
case wf @ NTile(buckets: Expression) =>
val frame = wExpression.windowSpec.frameSpecification.asInstanceOf[SpecifiedWindowFrame]
val childrenNodeList = new JArrayList[ExpressionNode]()
val literal = buckets.asInstanceOf[Literal]
childrenNodeList.add(LiteralTransformer(literal).doTransform(context))
val windowFunctionNode = ExpressionBuilder.makeWindowFunction(
WindowFunctionsBuilder.create(context, wf).toInt,
childrenNodeList,
columnName,
ConverterUtils.getTypeNode(wf.dataType, wf.nullable),
frame.upper,
frame.lower,
frame.frameType.sql,
originalInputAttributes.asJava
)
windowExpressionNodes.add(windowFunctionNode)
case _ =>
throw new GlutenNotSupportException(
"unsupported window function type: " +
wExpression.windowFunction)
}
}
}