Github user mn-mikke commented on a diff in the pull request: https://github.com/apache/spark/pull/21208#discussion_r185540852 --- Diff: sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/collectionOperations.scala --- @@ -1229,3 +1229,98 @@ case class Flatten(child: Expression) extends UnaryExpression { override def prettyName: String = "flatten" } + +/** + * Returns the array containing the given input value (left) count (right) times. + */ +@ExpressionDescription( + usage = "_FUNC_(element, count) - Returns the array containing element count times.", + examples = """ + Examples: + > SELECT _FUNC_('123', 2); + ['123', '123'] + """) +case class ArrayRepeat(left: Expression, right: Expression) + extends BinaryExpression { + + override def dataType: ArrayType = ArrayType(left.dataType, left.nullable) + + override def checkInputDataTypes(): TypeCheckResult = { + val expected = IntegerType + if (!expected.acceptsType(right.dataType)) { + val mismatch = s"argument 2 requires ${expected.simpleString} type, " + + s"however, '${right.sql}' is of ${right.dataType.simpleString} type." + TypeCheckResult.TypeCheckFailure(mismatch) + } else { + TypeCheckResult.TypeCheckSuccess + } + } + + override def nullable: Boolean = false + + override def eval(input: InternalRow): Any = { + new GenericArrayData(List.fill(right.eval(input).asInstanceOf[Integer])(left.eval(input))) + } + + override def prettyName: String = "array_repeat" + + override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = { + + val leftGen = left.genCode(ctx) + val rightGen = right.genCode(ctx) + val element = leftGen.value + val count = rightGen.value + val et = dataType.elementType + val isPrimitive = CodeGenerator.isPrimitiveType(et) + + val arrayDataName = ctx.freshName("arrayData") + val arrayName = ctx.freshName("arrayObject") + val initialization = (numElements: String) => if (isPrimitive) { + val arrayName = ctx.freshName("array") + val baseOffset = Platform.BYTE_ARRAY_OFFSET + s""" + | int numBytes = ${et.defaultSize} * $numElements; + | int unsafeArraySizeInBytes = UnsafeArrayData.calculateHeaderPortionInBytes($numElements) + | + org.apache.spark.unsafe.array.ByteArrayMethods + | .roundNumberOfBytesToNearestWord(numBytes); + | byte[] $arrayName = new byte[unsafeArraySizeInBytes]; + | UnsafeArrayData $arrayDataName = new UnsafeArrayData(); + | Platform.putLong($arrayName, $baseOffset, $numElements); + | $arrayDataName.pointTo($arrayName, $baseOffset, unsafeArraySizeInBytes); + | ${ev.value} = $arrayDataName; + """.stripMargin + } else { + s"${ev.value} = new ${classOf[GenericArrayData].getName()}(new Object[$numElements]);" + } + + val primitiveValueTypeName = CodeGenerator.primitiveTypeName(et) + val assignments = { + val updateArray = if (isPrimitive) { + s"${ev.value}.set$primitiveValueTypeName(k, $element);" + } else { + s"${ev.value}.update(k, $element);" + } + s""" + | for (int k = 0; k < $count; k++) { --- End diff -- What about `count` is `null`?
--- --------------------------------------------------------------------- To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org For additional commands, e-mail: reviews-h...@spark.apache.org