dbtsai commented on a change in pull request #27066: [SPARK-31317][SQL] Add
withField method to Column class
URL: https://github.com/apache/spark/pull/27066#discussion_r401072238
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
##########
@@ -514,3 +514,181 @@ case class StringToMap(text: Expression, pairDelim:
Expression, keyValueDelim: E
override def prettyName: String = "str_to_map"
}
+
+/**
+ * Adds/replaces fields in a struct.
+ * Returns null if struct is null.
+ * If multiple fields already exist with the one of the given fieldNames, they
will all be replaced.
+ */
+// scalastyle:off line.size.limit
+@ExpressionDescription(
+ usage = "_FUNC_(struct, name1, val1, name2, val2, ...) - Adds/replaces
fields in struct by name.",
+ examples = """
+ Examples:
+ > SELECT _FUNC_(NAMED_STRUCT("a", 1), "b", 2, "c", 3);
+ {"a":1,"b":2,"c":3}
+ """)
+// scalastyle:on line.size.limit
+case class AddFields(children: Seq[Expression]) extends Expression {
+
+ private lazy val struct: Expression = children.head
+ private lazy val (nameExprs, valExprs) = children.drop(1).grouped(2).map {
+ case Seq(name, value) => (name, value)
+ }.toList.unzip
+ private lazy val fieldNames =
nameExprs.map(_.eval().asInstanceOf[UTF8String].toString)
+ private lazy val pairs = fieldNames.zip(valExprs)
+
+ override def nullable: Boolean = struct.nullable
+
+ private lazy val ogStructType: StructType =
+ struct.dataType.asInstanceOf[StructType]
+
+ override lazy val dataType: StructType = {
+ val existingFields = ogStructType.fields.map { x => (x.name, x) }
+ val addOrReplaceFields = pairs.map { case (fieldName, field) =>
+ (fieldName, StructField(fieldName, field.dataType, field.nullable))
+ }
+ val newFields = loop(existingFields, addOrReplaceFields).map(_._2)
+ StructType(newFields)
+ }
+
+ override def checkInputDataTypes(): TypeCheckResult = {
+ if (children.size % 2 == 0) {
+ return TypeCheckResult.TypeCheckFailure(s"$prettyName expects an odd
number of arguments.")
+ }
+
Review comment:
Can you put it in big if else block? Having return in the middle of function
is fragile.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
