[GitHub] spark pull request #21021: [SPARK-23921][SQL] Add array_sort function

[ueshin <gi...@git.apache.org>]

Github user ueshin commented on a diff in the pull request:

    https://github.com/apache/spark/pull/21021#discussion_r185166848
  
    --- Diff: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala ---
    @@ -191,28 +161,205 @@ case class SortArray(base: Expression, ascendingOrder: Expression)
             if (o1 == null && o2 == null) {
               0
             } else if (o1 == null) {
    -          1
    +          -nullOrder
             } else if (o2 == null) {
    -          -1
    +          nullOrder
             } else {
               -ordering.compare(o1, o2)
             }
           }
         }
       }
     
    -  override def nullSafeEval(array: Any, ascending: Any): Any = {
    -    val elementType = base.dataType.asInstanceOf[ArrayType].elementType
    +  def elementType: DataType = arrayExpression.dataType.asInstanceOf[ArrayType].elementType
    +
    +  def sortEval(array: Any, ascending: Boolean): Any = {
         val data = array.asInstanceOf[ArrayData].toArray[AnyRef](elementType)
         if (elementType != NullType) {
    -      java.util.Arrays.sort(data, if (ascending.asInstanceOf[Boolean]) lt else gt)
    +      java.util.Arrays.sort(data, if (ascending) lt else gt)
         }
         new GenericArrayData(data.asInstanceOf[Array[Any]])
       }
     
    +  def sortCodegen(ctx: CodegenContext, ev: ExprCode, base: String, order: String): String = {
    +    val arrayData = classOf[ArrayData].getName
    +    val genericArrayData = classOf[GenericArrayData].getName
    +    val array = ctx.freshName("array")
    +    val c = ctx.freshName("c")
    +    val sort = if (elementType == NullType) "" else {
    +      val sortOrder = ctx.freshName("sortOrder")
    +      val o1 = ctx.freshName("o1")
    +      val o2 = ctx.freshName("o2")
    +      val jt = CodeGenerator.javaType(elementType)
    +      val comp = if (CodeGenerator.isPrimitiveType(elementType)) {
    +        val bt = CodeGenerator.boxedType(elementType)
    +        val v1 = ctx.freshName("v1")
    +        val v2 = ctx.freshName("v2")
    +        s"""
    +           |$jt $v1 = (($bt) $o1).${jt}Value();
    +           |$jt $v2 = (($bt) $o2).${jt}Value();
    +           |int $c = ${ctx.genComp(elementType, v1, v2)};
    +         """.stripMargin
    +      } else {
    +        s"int $c = ${ctx.genComp(elementType, s"(($jt) $o1)", s"(($jt) $o2)")};"
    +      }
    +      s"""
    +         |final int $sortOrder = $order ? 1 : -1;
    +         |java.util.Arrays.sort($array, new java.util.Comparator() {
    +         |  @Override public int compare(Object $o1, Object $o2) {
    +         |    if ($o1 == null && $o2 == null) {
    +         |      return 0;
    +         |    } else if ($o1 == null) {
    +         |      return $sortOrder * $nullOrder;
    +         |    } else if ($o2 == null) {
    +         |      return -$sortOrder * $nullOrder;
    +         |    }
    +         |    $comp
    +         |    return $sortOrder * $c;
    +         |  }
    +         |});
    +       """.stripMargin
    +    }
    +    val dataTypes = elementType match {
    +      case DecimalType.Fixed(p, s) =>
    +        s"org.apache.spark.sql.types.DataTypes.createDecimalType($p, $s)"
    +      case ArrayType(et, cn) =>
    +        s"org.apache.spark.sql.types.DataTypes.createArrayType($et, $cn)"
    +      case MapType(kt, vt, cn) =>
    +        s"org.apache.spark.sql.types.DataTypes.createMapType($kt, $vt, $cn)"
    +      case StructType(f) =>
    +        "org.apache.spark.sql.types.StructType$.MODULE$." +
    +          s"apply(new java.util.ArrayList(${f.length}))"
    +      case _ =>
    +        s"org.apache.spark.sql.types.DataTypes.$elementType"
    +    }
    +    s"""
    +       |Object[] $array = (Object[]) (($arrayData) $base).toArray(
    --- End diff --
    
    How about using `toObjectArray` which doesn't need ClassTag?


---

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org