Skip to content

Commit 5672977

Browse files
cfmcgradycomphead
authored andcommitted
feat: Support reverse function with ArrayType input (apache#2481)
* Support reverse function with ArrayType input * nit * refactor ut * assert * fix ci * Update spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala Co-authored-by: Oleks V <comphead@users.noreply.github.com> * Update spark/src/test/scala/org/apache/comet/CometArrayExpressionSuite.scala Co-authored-by: Oleks V <comphead@users.noreply.github.com> --------- Co-authored-by: Oleks V <comphead@users.noreply.github.com>
1 parent 41f31a4 commit 5672977

3 files changed

Lines changed: 482 additions & 371 deletions

File tree

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -916,6 +916,8 @@ object QueryPlanSerde extends Logging with CometExprShim {
916916
case l @ Length(child) if child.dataType == BinaryType =>
917917
withInfo(l, "Length on BinaryType is not supported")
918918
None
919+
case r @ Reverse(child) if child.dataType.isInstanceOf[ArrayType] =>
920+
convert(r, CometArrayReverse)
919921
case expr =>
920922
QueryPlanSerde.exprSerdeMap.get(expr.getClass) match {
921923
case Some(handler) =>

spark/src/main/scala/org/apache/comet/serde/arrays.scala

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ package org.apache.comet.serde
2121

2222
import scala.annotation.tailrec
2323

24-
import org.apache.spark.sql.catalyst.expressions.{ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayRemove, ArrayRepeat, ArraysOverlap, ArrayUnion, Attribute, CreateArray, ElementAt, Expression, Flatten, GetArrayItem, Literal}
24+
import org.apache.spark.sql.catalyst.expressions.{ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayRemove, ArrayRepeat, ArraysOverlap, ArrayUnion, Attribute, CreateArray, ElementAt, Expression, Flatten, GetArrayItem, Literal, Reverse}
2525
import org.apache.spark.sql.internal.SQLConf
2626
import org.apache.spark.sql.types._
2727

@@ -432,6 +432,22 @@ object CometGetArrayItem extends CometExpressionSerde[GetArrayItem] {
432432
}
433433
}
434434

435+
object CometArrayReverse extends CometExpressionSerde[Reverse] with ArraysBase {
436+
override def convert(
437+
expr: Reverse,
438+
inputs: Seq[Attribute],
439+
binding: Boolean): Option[ExprOuterClass.Expr] = {
440+
if (!isTypeSupported(expr.child.dataType)) {
441+
withInfo(expr, s"child data type not supported: ${expr.child.dataType}")
442+
return None
443+
}
444+
val reverseExprProto = exprToProto(expr.child, inputs, binding)
445+
val reverseScalarExpr = scalarFunctionExprToProto("array_reverse", reverseExprProto)
446+
optExprWithInfo(reverseScalarExpr, expr, expr.children: _*)
447+
}
448+
449+
}
450+
435451
object CometElementAt extends CometExpressionSerde[ElementAt] {
436452

437453
override def convert(

0 commit comments

Comments
 (0)