Github user cloud-fan commented on a diff in the pull request:

    https://github.com/apache/spark/pull/22544#discussion_r220772034
  
    --- Diff: 
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
 ---
    @@ -2140,21 +2140,34 @@ case class ElementAt(left: Expression, right: 
Expression) extends GetMapValueUti
       }
     
       override def inputTypes: Seq[AbstractDataType] = {
    -    Seq(TypeCollection(ArrayType, MapType),
    -      left.dataType match {
    -        case _: ArrayType => IntegerType
    -        case _: MapType => mapKeyType
    -        case _ => AnyDataType // no match for a wrong 'left' expression 
type
    -      }
    -    )
    +    (left.dataType, right.dataType) match {
    +      case (arr: ArrayType, e2: IntegralType) if (e2 != LongType) =>
    +        Seq(arr, IntegerType)
    +      case (MapType(keyType, valueType, hasNull), e2) =>
    +        TypeCoercion.findTightestCommonType(keyType, e2) match {
    +          case Some(dt) => Seq(MapType(dt, valueType, hasNull), dt)
    +          case _ => Seq.empty
    +        }
    +      case (l, r) => Seq.empty
    +
    +    }
       }
     
       override def checkInputDataTypes(): TypeCheckResult = {
    -    super.checkInputDataTypes() match {
    -      case f: TypeCheckResult.TypeCheckFailure => f
    -      case TypeCheckResult.TypeCheckSuccess if 
left.dataType.isInstanceOf[MapType] =>
    -        TypeUtils.checkForOrderingExpr(mapKeyType, s"function $prettyName")
    -      case TypeCheckResult.TypeCheckSuccess => 
TypeCheckResult.TypeCheckSuccess
    +    (left.dataType, right.dataType) match {
    +      case (ArrayType(e1, _), e2) if e2 != IntegerType =>
    --- End diff --
    
    nit: `e1` is not used, this can be `case (_: ArrayType, e2) if ...`


---

---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to