Github user cloud-fan commented on a diff in the pull request:
https://github.com/apache/spark/pull/21912#discussion_r213540868
--- Diff:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala
---
@@ -1490,6 +1423,63 @@ object CodeGenerator extends Logging {
}
}
+ /**
+ * 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 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]]
+ */
+ def createArrayData(
+ arrayName: String,
+ elementType: DataType,
+ numElements: String,
+ additionalErrorMessage: String,
+ elementSize: Option[Int] = None): String = {
+ val (isPrimitiveType, elemSize) = if (elementSize.isDefined) {
+ (false, elementSize.get)
+ } else {
+ (CodeGenerator.isPrimitiveType(elementType), elementType.defaultSize)
+ }
+
+ s"""
+ |ArrayData $arrayName = ArrayData.allocateArrayData(
+ | $elemSize, $numElements, $isPrimitiveType,
"$additionalErrorMessage");
+ """.stripMargin
+ }
+
+ /**
+ * Generates assignment code for an [[ArrayData]]
--- End diff --
shall we mention that the returned code should be put inside a loop?
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
