Github user cloud-fan commented on a diff in the pull request:
https://github.com/apache/spark/pull/21912#discussion_r209573854
--- Diff:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
---
@@ -735,70 +735,100 @@ class CodegenContext {
}
/**
- * Generates code creating a [[UnsafeArrayData]].
+ * Generates code creating a [[UnsafeArrayData]] or [[GenericArrayData]]
based on
+ * given parameters.
*
* @param arrayName name of the array to create
+ * @param elementType data type of the elements in source array
* @param numElements code representing the number of elements the array
should contain
- * @param elementType data type of the elements in the array
* @param additionalErrorMessage string to include in the error message
+ * @param elementSize optional value which shows the size of an element
of the allocated
+ * [[UnsafeArrayData]] or [[GenericArrayData]]
+ *
+ * @return code representing the allocation of [[ArrayData]]
+ * code representing a setter of an assignment for the generated
array
*/
- def createUnsafeArray(
+ def createArrayData(
arrayName: String,
- numElements: String,
elementType: DataType,
- additionalErrorMessage: String): String = {
- val arraySize = freshName("size")
- val arrayBytes = freshName("arrayBytes")
+ numElements: String,
+ additionalErrorMessage: String,
+ elementSize: Option[Int] = None): (String, String) = {
+ val isPrimitiveType = if (elementSize.isDefined) {
+ false
+ } else {
+ CodeGenerator.isPrimitiveType(elementType)
+ }
- s"""
- |long $arraySize =
UnsafeArrayData.calculateSizeOfUnderlyingByteArray(
- | $numElements,
- | ${elementType.defaultSize});
- |if ($arraySize > ${ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH}) {
- | throw new RuntimeException("Unsuccessful try create array with "
+ $arraySize +
- | " bytes of data due to exceeding the limit " +
- | "${ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH} bytes for
UnsafeArrayData." +
- | "$additionalErrorMessage");
- |}
- |byte[] $arrayBytes = new byte[(int)$arraySize];
- |UnsafeArrayData $arrayName = new UnsafeArrayData();
- |Platform.putLong($arrayBytes, ${Platform.BYTE_ARRAY_OFFSET},
$numElements);
- |$arrayName.pointTo($arrayBytes, ${Platform.BYTE_ARRAY_OFFSET},
(int)$arraySize);
- """.stripMargin
+ val setFunc = if (isPrimitiveType) {
+ s"set${CodeGenerator.primitiveTypeName(elementType)}"
+ } else {
+ "update"
+ }
+
+ val elemSize = if (elementSize.isDefined) {
+ elementSize.get
+ } else {
+ elementType.defaultSize
+ }
+ val arrayData = classOf[ArrayData].getName
+ val allocation =
+ s"""
+ |ArrayData $arrayName = $arrayData$$.MODULE$$.allocateArrayData(
--- End diff --
IIRC java can call methods of scala object directly, can you verify if
`ArrayData.allocateArrayData` works?
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
