feat: support sort_array expression (#3706)

Author: grorge123

docs/source/user-guide/latest/compatibility.md

@@ -66,6 +66,7 @@ the [Comet Supported Expressions Guide](expressions.md) for more information on
 - **ArrayUnion**: Sorts input arrays before performing the union, while Spark preserves the order of the first array
   and appends unique elements from the second.
   [#3644](https://github.com/apache/datafusion-comet/issues/3644)
+- **SortArray**: Nested arrays with `Struct` or `Null` child values are not supported natively and will fall back to Spark.
 
 ### Date/Time Expressions
 

docs/spark_expressions_support.md

@@ -105,7 +105,7 @@
 - [ ] sequence
 - [ ] shuffle
 - [ ] slice
-- [ ] sort_array
+- [x] sort_array
 
 ### bitwise_funcs
 

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

@@ -60,6 +60,7 @@ object QueryPlanSerde extends Logging with CometExprShim {
     classOf[ArrayMin] -> CometArrayMin,
     classOf[ArrayRemove] -> CometArrayRemove,
     classOf[ArrayRepeat] -> CometArrayRepeat,
+    classOf[SortArray] -> CometSortArray,
     classOf[ArraysOverlap] -> CometArraysOverlap,
     classOf[ArrayUnion] -> CometArrayUnion,
     classOf[CreateArray] -> CometCreateArray,
@@ -796,30 +797,23 @@ object QueryPlanSerde extends Logging with CometExprShim {
    * TODO: Include SparkSQL's [[YearMonthIntervalType]] and [[DayTimeIntervalType]]
    */
   // scalastyle:on
-  def supportedSortType(op: SparkPlan, sortOrder: Seq[SortOrder]): Boolean = {
-    def canRank(dt: DataType): Boolean = {
-      dt match {
-        case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType |
-            _: DoubleType | _: DecimalType =>
-          true
-        case _: DateType | _: TimestampType | _: TimestampNTZType =>
-          true
-        case _: BooleanType | _: BinaryType | _: StringType => true
-        case _ => false
-      }
+  def supportedScalarSortElementType(dt: DataType): Boolean = {
+    dt match {
+      case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType |
+          _: DoubleType | _: DecimalType | _: DateType | _: TimestampType | _: TimestampNTZType |
+          _: BooleanType | _: BinaryType | _: StringType =>
+        true
+      case _ =>
+        false
     }
+  }
 
+  def supportedSortType(op: SparkPlan, sortOrder: Seq[SortOrder]): Boolean = {
     if (sortOrder.length == 1) {
       val canSort = sortOrder.head.dataType match {
-        case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType |
-            _: DoubleType | _: DecimalType =>
-          true
-        case _: DateType | _: TimestampType | _: TimestampNTZType =>
-          true
-        case _: BooleanType | _: BinaryType | _: StringType => true
-        case ArrayType(elementType, _) => canRank(elementType)
-        case MapType(_, valueType, _) => canRank(valueType)
-        case _ => false
+        case ArrayType(elementType, _) => supportedScalarSortElementType(elementType)
+        case MapType(_, valueType, _) => supportedScalarSortElementType(valueType)
+        case _ => supportedScalarSortElementType(sortOrder.head.dataType)
       }
       if (!canSort) {
         withInfo(op, s"Sort on single column of type ${sortOrder.head.dataType} is not supported")

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

@@ -21,11 +21,12 @@ package org.apache.comet.serde
 
 import scala.annotation.tailrec
 
-import org.apache.spark.sql.catalyst.expressions.{ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayFilter, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayRemove, ArrayRepeat, ArraysOverlap, ArrayUnion, Attribute, CreateArray, ElementAt, Expression, Flatten, GetArrayItem, IsNotNull, Literal, Reverse, Size}
+import org.apache.spark.sql.catalyst.expressions.{ArrayAppend, ArrayContains, ArrayDistinct, ArrayExcept, ArrayFilter, ArrayInsert, ArrayIntersect, ArrayJoin, ArrayMax, ArrayMin, ArrayRemove, ArrayRepeat, ArraysOverlap, ArrayUnion, Attribute, CreateArray, ElementAt, Expression, Flatten, GetArrayItem, IsNotNull, Literal, Reverse, Size, SortArray}
 import org.apache.spark.sql.catalyst.util.GenericArrayData
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 
+import org.apache.comet.CometConf
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.serde.QueryPlanSerde._
 import org.apache.comet.shims.CometExprShim
@@ -136,6 +137,80 @@ object CometArrayDistinct extends CometExpressionSerde[ArrayDistinct] {
   }
 }
 
+object CometSortArray extends CometExpressionSerde[SortArray] {
+  private def containsFloatingPoint(dt: DataType): Boolean = {
+    dt match {
+      case FloatType | DoubleType => true
+      case ArrayType(elementType, _) => containsFloatingPoint(elementType)
+      case StructType(fields) => fields.exists(f => containsFloatingPoint(f.dataType))
+      case MapType(keyType, valueType, _) =>
+        containsFloatingPoint(keyType) || containsFloatingPoint(valueType)
+      case _ => false
+    }
+  }
+
+  private def supportedSortArrayElementType(
+      dt: DataType,
+      nestedInArray: Boolean = false): Boolean = {
+    dt match {
+      // DataFusion's array_sort compares nested arrays through Arrow's rank kernel.
+      // That kernel does not support Struct or Null child values,
+      // so array<array<struct<...>>> and array<array<null>> would fail at runtime.
+      case _: NullType if !nestedInArray =>
+        true
+      case ArrayType(elementType, _) =>
+        supportedSortArrayElementType(elementType, nestedInArray = true)
+      case StructType(fields) if !nestedInArray =>
+        fields.forall(f => supportedSortArrayElementType(f.dataType))
+      case _ =>
+        supportedScalarSortElementType(dt)
+    }
+  }
+
+  override def getSupportLevel(expr: SortArray): SupportLevel = {
+    val elementType = expr.base.dataType.asInstanceOf[ArrayType].elementType
+
+    if (!supportedSortArrayElementType(elementType)) {
+      Unsupported(Some(s"Sort on array element type $elementType is not supported"))
+    } else if (CometConf.COMET_EXEC_STRICT_FLOATING_POINT.get() &&
+      containsFloatingPoint(elementType)) {
+      Incompatible(
+        Some(
+          "Sorting on floating-point is not 100% compatible with Spark, and Comet is running " +
+            s"with ${CometConf.COMET_EXEC_STRICT_FLOATING_POINT.key}=true. " +
+            s"${CometConf.COMPAT_GUIDE}"))
+    } else {
+      Compatible()
+    }
+  }
+
+  override def convert(
+      expr: SortArray,
+      inputs: Seq[Attribute],
+      binding: Boolean): Option[ExprOuterClass.Expr] = {
+    val arrayExprProto = exprToProtoInternal(expr.base, inputs, binding)
+    val (sortDirectionExprProto, nullOrderingExprProto) = expr.ascendingOrder match {
+      case Literal(value: Boolean, BooleanType) =>
+        val direction = if (value) "ASC" else "DESC"
+        val nullOrdering = if (value) "NULLS FIRST" else "NULLS LAST"
+        (
+          exprToProtoInternal(Literal(direction), inputs, binding),
+          exprToProtoInternal(Literal(nullOrdering), inputs, binding))
+      case other =>
+        withInfo(expr, s"ascendingOrder must be a boolean literal: $other")
+        (None, None)
+    }
+
+    val sortArrayScalarExpr =
+      scalarFunctionExprToProto(
+        "array_sort",
+        arrayExprProto,
+        sortDirectionExprProto,
+        nullOrderingExprProto)
+    optExprWithInfo(sortArrayScalarExpr, expr, expr.children: _*)
+  }
+}
+
 object CometArrayIntersect extends CometExpressionSerde[ArrayIntersect] {
 
   override def getSupportLevel(expr: ArrayIntersect): SupportLevel = Incompatible(None)