Add CreateNamedStruct in both DataFrame function API and FunctionRegistery

Author: yjshen

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/FunctionRegistry.scala

@@ -96,6 +96,7 @@ object FunctionRegistry {
     expression[Rand]("rand"),
     expression[Randn]("randn"),
     expression[CreateStruct]("struct"),
+    expression[CreateNamedStruct]("named_struct"),
     expression[Sqrt]("sqrt"),
 
     // math functions

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/complexTypeCreator.scala

@@ -17,9 +17,11 @@
 
 package org.apache.spark.sql.catalyst.expressions
 
+import org.apache.spark.sql.catalyst
 import org.apache.spark.sql.catalyst.analysis.TypeCheckResult
 import org.apache.spark.sql.catalyst.util.TypeUtils
 import org.apache.spark.sql.types._
+import org.apache.spark.unsafe.types.UTF8String
 
 /**
  * Returns an Array containing the evaluation of all children expressions.
@@ -74,3 +76,44 @@ case class CreateStruct(children: Seq[Expression]) extends Expression {
 
   override def prettyName: String = "struct"
 }
+
+/**
+ * Creates a struct with the given field names and values
+ *
+ * @param children Seq(name1, val1, name2, val2, ...)
+ */
+case class CreateNamedStruct(children: Seq[Expression]) extends Expression {
+  assert(children.size % 2 == 0, "NamedStruct expects an even number of arguments.")
+
+  private val nameExprs = children.zipWithIndex.filter(_._2 % 2 == 0).map(_._1)
+  private val valExprs = children.zipWithIndex.filter(_._2 % 2 == 1).map(_._1)
+
+  private lazy val names = nameExprs.map { case name =>
+    name match {
+      case NonNullLiteral(str, StringType) =>
+        str.asInstanceOf[UTF8String].toString
+      case _ =>
+        throw new IllegalArgumentException("Expressions of odd index should be" +
+          s" Literal(_, StringType), get ${name.dataType} instead")
+    }
+  }
+
+  override def foldable: Boolean = children.forall(_.foldable)
+
+  override lazy val resolved: Boolean = childrenResolved
+
+  override lazy val dataType: StructType = {
+    assert(resolved,
+      s"CreateStruct contains unresolvable children: ${children.filterNot(_.resolved)}.")
+    val fields = names.zip(valExprs).map { case (name, valExpr) =>
+      StructField(name, valExpr.dataType, valExpr.nullable, Metadata.empty)
+    }
+    StructType(fields)
+  }
+
+  override def nullable: Boolean = false
+
+  override def eval(input: InternalRow): Any = {
+    InternalRow(valExprs.map(_.eval(input)): _*)
+  }
+}

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/ComplexTypeSuite.scala

@@ -100,6 +100,38 @@ class ComplexTypeSuite extends SparkFunSuite with ExpressionEvalHelper {
     assert(getStructField(nullStruct, "a").nullable === true)
   }
 
+  test("complex type") {
+    val row = create_row(
+      "^Ba*n", // 0
+      null.asInstanceOf[UTF8String], // 1
+      create_row("aa", "bb"), // 2
+      Map("aa" -> "bb"), // 3
+      Seq("aa", "bb") // 4
+    )
+
+    val typeS = StructType(
+      StructField("a", StringType, true) :: StructField("b", StringType, true) :: Nil
+    )
+    val typeMap = MapType(StringType, StringType)
+    val typeArray = ArrayType(StringType)
+
+    checkEvaluation(GetMapValue(BoundReference(3, typeMap, true),
+      Literal("aa")), "bb", row)
+    checkEvaluation(GetMapValue(Literal.create(null, typeMap), Literal("aa")), null, row)
+    checkEvaluation(
+      GetMapValue(Literal.create(null, typeMap), Literal.create(null, StringType)), null, row)
+    checkEvaluation(GetMapValue(BoundReference(3, typeMap, true),
+      Literal.create(null, StringType)), null, row)
+
+    checkEvaluation(GetArrayItem(BoundReference(4, typeArray, true),
+      Literal(1)), "bb", row)
+    checkEvaluation(GetArrayItem(Literal.create(null, typeArray), Literal(1)), null, row)
+    checkEvaluation(
+      GetArrayItem(Literal.create(null, typeArray), Literal.create(null, IntegerType)), null, row)
+    checkEvaluation(GetArrayItem(BoundReference(4, typeArray, true),
+      Literal.create(null, IntegerType)), null, row)
+  }
+
   test("GetArrayStructFields") {
     val typeAS = ArrayType(StructType(StructField("a", IntegerType) :: Nil))
     val arrayStruct = Literal.create(Seq(create_row(1)), typeAS)
@@ -119,11 +151,29 @@ class ComplexTypeSuite extends SparkFunSuite with ExpressionEvalHelper {
 
   test("CreateStruct") {
     val row = create_row(1, 2, 3)
-    val c1 = 'a.int.at(0).as("a")
-    val c3 = 'c.int.at(2).as("c")
+    val c1 = 'a.int.at(0)
+    val c3 = 'c.int.at(2)
     checkEvaluation(CreateStruct(Seq(c1, c3)), create_row(1, 3), row)
   }
 
+  test("CreateNamedStruct") {
+    val row = InternalRow(1, 2, 3)
+    val c1 = 'a.int.at(0)
+    val c3 = 'c.int.at(2)
+    checkEvaluation(CreateNamedStruct(Seq("a", c1, "b", c3)), InternalRow(1, 3), row)
+  }
+
+  test("CreateNamedStruct with literal field") {
+    val row = InternalRow(1, 2, 3)
+    val c1 = 'a.int.at(0)
+    checkEvaluation(CreateNamedStruct(Seq("a", c1, "b", "y")), InternalRow(1, "y"), row)
+  }
+
+  test("CreateNamedStruct from all literal fields") {
+    checkEvaluation(
+      CreateNamedStruct(Seq("a", "x", "b", 2.0)), InternalRow("x", 2.0), InternalRow.empty)
+  }
+
   test("test dsl for complex type") {
     def quickResolve(u: UnresolvedExtractValue): Expression = {
       ExtractValue(u.child, u.extraction, _ == _)

sql/core/src/main/scala/org/apache/spark/sql/functions.scala

@@ -762,6 +762,21 @@ object functions {
     struct((colName +: colNames).map(col) : _*)
   }
 
+  /**
+   * Creates a new struct column with given field names and columns.
+   * The input columns should be of length 2*n and follow (name1, col1, name2, col2),
+   * name* should be String Literal
+   *
+   * @group normal_funcs
+   * @since 1.5.0
+   */
+  @scala.annotation.varargs
+  def named_struct(cols: Column*): Column = {
+    require(cols.length % 2 == 0,
+      s"named_struct expects an even number of arguments.")
+    CreateNamedStruct(cols.map(_.expr))
+  }
+
   /**
    * Converts a string expression to upper case.
    *

sql/core/src/test/scala/org/apache/spark/sql/DataFrameFunctionsSuite.scala

@@ -85,6 +85,64 @@ class DataFrameFunctionsSuite extends QueryTest {
     }
   }
 
+  test("named_struct with column expression") {
+    val df = Seq((1, "str1"), (2, "str2")).toDF("a", "b")
+    val row = df.select(
+      named_struct(lit("x"), (col("a") * 2), lit("y"), col("b"))).take(2)
+
+    val expectedType = StructType(Seq(
+      StructField("x", IntegerType, nullable = false),
+      StructField("y", StringType)
+    ))
+
+    assert(row(0).schema(0).dataType === expectedType)
+    assert(row(0).getAs[Row](0) === Row(2, "str1"))
+    assert(row(1).getAs[Row](0) === Row(4, "str2"))
+  }
+
+  test("named_struct with literal columns") {
+    val df = Seq((1, "str1"), (2, "str2")).toDF("a", "b")
+    val row = df.select(
+      named_struct(lit("x"), (col("a") * 2), lit("y"), lit(5.0))).take(2)
+
+    val expectedType = StructType(Seq(
+      StructField("x", IntegerType, nullable = false),
+      StructField("y", DoubleType, nullable = false)
+    ))
+
+    assert(row(0).schema(0).dataType === expectedType)
+    assert(row(0).getAs[Row](0) === Row(2, 5.0))
+    assert(row(1).getAs[Row](0) === Row(4, 5.0))
+  }
+
+  test("named_struct with all literal columns") {
+    val df = Seq((1, "str1"), (2, "str2")).toDF("a", "b")
+    val row = df.select(
+      named_struct(lit("x"), lit("v"), lit("y"), lit(5.0))).take(2)
+
+    val expectedType = StructType(Seq(
+      StructField("x", StringType, nullable = false),
+      StructField("y", DoubleType, nullable = false)
+    ))
+
+    assert(row(0).schema(0).dataType === expectedType)
+    assert(row(0).getAs[Row](0) === Row("v", 5.0))
+    assert(row(1).getAs[Row](0) === Row("v", 5.0))
+  }
+
+  test("named_struct with odd arguments") {
+    intercept[IllegalArgumentException] {
+      named_struct(col("x"))
+    }
+  }
+
+  test("named_struct with non string literal names") {
+    val df = Seq((1, "str")).toDF("a", "b")
+    intercept[IllegalArgumentException] {
+      df.select(named_struct(lit(1), (col("a") * 2), lit("y"), lit(5.0)))
+    }
+  }
+
   test("constant functions") {
     checkAnswer(
       ctx.sql("SELECT E()"),