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]