fqaiser94 commented on a change in pull request #29322:
URL: https://github.com/apache/spark/pull/29322#discussion_r469534633
##########
File path: sql/core/src/main/scala/org/apache/spark/sql/Column.scala
##########
@@ -906,34 +906,84 @@ class Column(val expr: Expression) extends Logging {
*/
// scalastyle:on line.size.limit
def withField(fieldName: String, col: Column): Column = withExpr {
- require(fieldName != null, "fieldName cannot be null")
require(col != null, "col cannot be null")
+ updateFieldsHelper(expr, nameParts(fieldName), name => WithField(name,
col.expr))
+ }
+
+ // scalastyle:off line.size.limit
+ /**
+ * An expression that drops fields in `StructType` by name.
+ *
+ * {{{
+ * val df = sql("SELECT named_struct('a', 1, 'b', 2) struct_col")
+ * df.select($"struct_col".dropFields("b"))
+ * // result: {"a":1}
+ *
+ * val df = sql("SELECT named_struct('a', 1, 'b', 2) struct_col")
+ * df.select($"struct_col".dropFields("c"))
+ * // result: {"a":1,"b":2}
+ *
+ * val df = sql("SELECT named_struct('a', 1, 'b', 2, 'c', 3) struct_col")
+ * df.select($"struct_col".dropFields("b", "c"))
+ * // result: {"a":1}
+ *
+ * val df = sql("SELECT named_struct('a', 1, 'b', 2) struct_col")
+ * df.select($"struct_col".dropFields("a", "b"))
+ * // result: org.apache.spark.sql.AnalysisException: cannot resolve
'update_fields(update_fields(`struct_col`))' due to data type mismatch: cannot
drop all fields in struct
+ *
+ * val df = sql("SELECT CAST(NULL AS struct<a:int,b:int>) struct_col")
+ * df.select($"struct_col".dropFields("b"))
+ * // result: null of type struct<a:int>
+ *
+ * val df = sql("SELECT named_struct('a', 1, 'b', 2, 'b', 3) struct_col")
+ * df.select($"struct_col".dropFields("b"))
+ * // result: {"a":1}
+ *
+ * val df = sql("SELECT named_struct('a', named_struct('a', 1, 'b', 2))
struct_col")
+ * df.select($"struct_col".dropFields("a.b"))
+ * // result: {"a":{"a":1}}
+ *
+ * val df = sql("SELECT named_struct('a', named_struct('b', 1), 'a',
named_struct('c', 2)) struct_col")
+ * df.select($"struct_col".dropFields("a.c"))
+ * // result: org.apache.spark.sql.AnalysisException: Ambiguous reference
to fields
+ * }}}
+ *
+ * @group expr_ops
+ * @since 3.1.0
+ */
+ // scalastyle:on line.size.limit
+ def dropFields(fieldNames: String*): Column = withExpr {
+ def dropField(expr: Expression, fieldName: String): UpdateFields =
+ updateFieldsHelper(expr, nameParts(fieldName), name => DropField(name))
+
+ fieldNames.tail.foldLeft(dropField(expr, fieldNames.head)) {
+ (resExpr, fieldName) => dropField(resExpr, fieldName)
+ }
+ }
+
+ private def nameParts(fieldName: String): Seq[String] = {
+ require(fieldName != null, "fieldName cannot be null")
- val nameParts = if (fieldName.isEmpty) {
+ if (fieldName.isEmpty) {
fieldName :: Nil
} else {
CatalystSqlParser.parseMultipartIdentifier(fieldName)
}
- withFieldHelper(expr, nameParts, Nil, col.expr)
}
- private def withFieldHelper(
+ private def updateFieldsHelper(
struct: Expression,
namePartsRemaining: Seq[String],
- namePartsDone: Seq[String],
- value: Expression) : WithFields = {
- val name = namePartsRemaining.head
+ value: String => StructFieldsOperation): UpdateFields = {
Review comment:
Done
##########
File path:
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala
##########
@@ -541,57 +541,94 @@ case class StringToMap(text: Expression, pairDelim:
Expression, keyValueDelim: E
}
/**
- * Adds/replaces field in struct by name.
+ * Represents an operation to be applied to the fields of a struct.
*/
-case class WithFields(
- structExpr: Expression,
- names: Seq[String],
- valExprs: Seq[Expression]) extends Unevaluable {
+trait StructFieldsOperation {
- assert(names.length == valExprs.length)
+ val resolver: Resolver = SQLConf.get.resolver
+
+ /**
+ * Returns an updated list of expressions which will ultimately be used as
the children argument
+ * for [[CreateNamedStruct]].
+ */
+ def apply(exprs: Seq[(String, Expression)]): Seq[(String, Expression)]
+}
+
+/**
+ * Add or replace a field by name.
+ */
+case class WithField(name: String, valExpr: Expression)
+ extends Unevaluable with StructFieldsOperation {
Review comment:
Done.
----------------------------------------------------------------
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
