Github user kiszk commented on a diff in the pull request:
https://github.com/apache/spark/pull/21121#discussion_r183214860
--- Diff:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala
---
@@ -883,3 +884,139 @@ case class Concat(children: Seq[Expression]) extends
Expression {
override def sql: String = s"concat(${children.map(_.sql).mkString(",
")})"
}
+
+/**
+ * Returns the maximum value in the array.
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = "_FUNC_(array[, indexFirst]) - Transforms the input array by
encapsulating elements into pairs with indexes indicating the order.",
+ examples = """
+ Examples:
+ > SELECT _FUNC_(array("d", "a", null, "b"));
+ [("d",0),("a",1),(null,2),("b",3)]
+ > SELECT _FUNC_(array("d", "a", null, "b"), true);
+ [(0,"d"),(1,"a"),(2,null),(3,"b")]
+ """,
+ since = "2.4.0")
+case class ZipWithIndex(child: Expression, indexFirst: Expression)
+ extends UnaryExpression with ExpectsInputTypes {
+
+ def this(e: Expression) = this(e, Literal.FalseLiteral)
+
+ val indexFirstValue: Boolean = indexFirst match {
+ case Literal(v: Boolean, BooleanType) => v
+ case _ => throw new AnalysisException("The second argument has to be a
boolean constant.")
+ }
+
+ private val MAX_ARRAY_LENGTH: Int =
ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH
+
+ override def inputTypes: Seq[AbstractDataType] = Seq(ArrayType)
+
+ lazy val childArrayType: ArrayType =
child.dataType.asInstanceOf[ArrayType]
+
+ override def dataType: DataType = {
+ val elementField = StructField("value", childArrayType.elementType,
childArrayType.containsNull)
+ val indexField = StructField("index", IntegerType, false)
+
+ val fields = if (indexFirstValue) Seq(indexField, elementField) else
Seq(elementField, indexField)
+
+ ArrayType(StructType(fields), false)
+ }
+
+ override protected def nullSafeEval(input: Any): Any = {
+ val array =
input.asInstanceOf[ArrayData].toObjectArray(childArrayType.elementType)
+
+ val makeStruct = (v: Any, i: Int) => if (indexFirstValue)
InternalRow(i, v) else InternalRow(v, i)
+ val resultData = array.zipWithIndex.map{case (v, i) => makeStruct(v,
i)}
+
+ new GenericArrayData(resultData)
+ }
+
+ override protected def doGenCode(ctx: CodegenContext, ev: ExprCode):
ExprCode = {
+ nullSafeCodeGen(ctx, ev, c => {
+ if (CodeGenerator.isPrimitiveType(childArrayType.elementType)) {
+ genCodeForPrimitiveElements(ctx, c, ev.value)
+ } else {
+ genCodeForNonPrimitiveElements(ctx, c, ev.value)
+ }
+ })
+ }
+
+ private def genCodeForPrimitiveElements(
+ ctx: CodegenContext,
+ childVariableName: String,
+ arrayData: String): String = {
+ val numElements = ctx.freshName("numElements")
+ val byteArraySize = ctx.freshName("byteArraySize")
+ val data = ctx.freshName("byteArray")
+ val unsafeRow = ctx.freshName("unsafeRow")
+ val structSize = ctx.freshName("structSize")
+ val unsafeArrayData = ctx.freshName("unsafeArrayData")
+ val structsOffset = ctx.freshName("structsOffset")
+ val calculateArraySize =
"UnsafeArrayData.calculateSizeOfUnderlyingByteArray"
+ val calculateHeader = "UnsafeArrayData.calculateHeaderPortionInBytes"
+
+ val baseOffset = Platform.BYTE_ARRAY_OFFSET
+ val longSize = LongType.defaultSize
+ val primitiveValueTypeName =
CodeGenerator.primitiveTypeName(childArrayType.elementType)
+ val valuePosition = if (indexFirstValue) "1" else "0"
+ val indexPosition = if (indexFirstValue) "0" else "1"
--- End diff --
nit: How about `val (valuePosition, indexPosition) = if (indexFirstValue)
("1", "0") else ("0", "1")`?
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
