Support UDF that returns RowType (dotnet#376)

Author: elvaliuliuliu

src/csharp/Microsoft.Spark.E2ETest/UdfTests/UdfComplexTypesTests.cs

@@ -48,7 +48,7 @@ public void TestUdfWithArrayType()
 
             Row[] rows = _df.Select(workingUdf(_df["ids"])).Collect().ToArray();
             Assert.Equal(3, rows.Length);
-            
+
             var expected = new[] { "1", "3,5", "2,4" };
             string[] rowsToArray = rows.Select(x => x[0].ToString()).ToArray();
             Assert.Equal(expected, rowsToArray);
@@ -158,25 +158,76 @@ public void TestUdfWithRowType()
         [Fact]
         public void TestUdfWithReturnAsRowType()
         {
-            // UDF with return as RowType throws a following exception:
-            // Unhandled Exception: System.Reflection.TargetInvocationException: Exception has been thrown by the target of an invocation. 
-            // --->System.ArgumentException: System.Object is not supported.
-            // at Microsoft.Spark.Utils.UdfUtils.GetReturnType(Type type) in Microsoft.Spark\Utils\UdfUtils.cs:line 142
-            // at Microsoft.Spark.Utils.UdfUtils.GetReturnType(Type type) in Microsoft.Spark\Utils\UdfUtils.cs:line 136
-            // at Microsoft.Spark.Sql.Functions.CreateUdf[TResult](String name, Delegate execute, PythonEvalType evalType) in Microsoft.Spark\Sql\Functions.cs:line 4053
-            // at Microsoft.Spark.Sql.Functions.CreateUdf[TResult](String name, Delegate execute) in Microsoft.Spark\Sql\Functions.cs:line 4040
-            // at Microsoft.Spark.Sql.Functions.Udf[T, TResult](Func`2 udf) in Microsoft.Spark\Sql\Functions.cs:line 3607
-            Assert.Throws<ArgumentException>(() => Udf<string, Row>(
-                (str) =>
+            // Single GenericRow
+            {
+                var schema = new StructType(new[]
+                {
+                    new StructField("col1", new IntegerType()),
+                    new StructField("col2", new StringType())
+                });
+                Func<Column, Column> udf = Udf<string>(
+                    str => new GenericRow(new object[] { 1, "abc" }), schema);
+
+                Row[] rows = _df.Select(udf(_df["name"])).Collect().ToArray();
+                Assert.Equal(3, rows.Length);
+
+                foreach (Row row in rows)
+                {
+                    Assert.Equal(2, row.Size());
+                    Assert.Equal(1, row.GetAs<int>("col1"));
+                    Assert.Equal("abc", row.GetAs<string>("col2"));
+                }
+            }
+
+            // Nested GenericRow
+            {
+                var subSchema1 = new StructType(new[]
+                {
+                    new StructField("col1", new IntegerType()),
+                });
+                var subSchema2 = new StructType(new[]
+                {
+                    new StructField("col1", new StringType()),
+                    new StructField("col2", subSchema1),
+                });
+                var schema = new StructType(new[]
+                {
+                    new StructField("col1", new IntegerType()),
+                    new StructField("col2", subSchema1),
+                    new StructField("col3", subSchema2)
+                });
+
+                Func<Column, Column> udf = Udf<string>(
+                    str => new GenericRow(
+                        new object[]
+                        {
+                            1,
+                            new GenericRow(new object[] { 1 }),
+                            new GenericRow(new object[]
+                                {
+                                    "abc",
+                                    new GenericRow(new object[] { 10 })
+                                })
+                        }),
+                        schema);
+
+                Row[] rows = _df.Select(udf(_df["name"])).Collect().ToArray();
+                Assert.Equal(3, rows.Length);
+
+                foreach (Row row in rows)
                 {
-                    var structFields = new List<StructField>()
-                    {
-                        new StructField("name", new StringType()),
-                    };
-                    var schema = new StructType(structFields);
-                    var row = new Row(new object[] { str }, schema);
-                    return row;
-                }));
+                    Assert.Equal(3, row.Size());
+                    Assert.Equal(1, row.GetAs<int>("col1"));
+                    Assert.Equal(
+                        new Row(new object[] { 1 }, subSchema1),
+                        row.GetAs<Row>("col2"));
+                    Assert.Equal(
+                        new Row(
+                            new object[] { "abc", new Row(new object[] { 10 }, subSchema1) },
+                            subSchema2),
+                        row.GetAs<Row>("col3"));
+                }
+            }
         }
     }
 }

src/csharp/Microsoft.Spark/Sql/Functions.cs

@@ -8,6 +8,7 @@
 using Microsoft.Spark.Interop;
 using Microsoft.Spark.Interop.Ipc;
 using Microsoft.Spark.Sql.Expressions;
+using Microsoft.Spark.Sql.Types;
 using Microsoft.Spark.Utils;
 
 namespace Microsoft.Spark.Sql
@@ -3797,6 +3798,189 @@ public static Func<Column, Column, Column, Column, Column, Column, Column, Colum
             return CreateUdf<TResult>(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf)).Apply10;
         }
 
+        /// <summary>Creates a UDF from the specified delegate.</summary>
+        /// <param name="udf">The UDF function implementation.</param>
+        /// <param name="returnType">Schema associated with this row</param>
+        /// <returns>
+        /// A delegate that returns a <see cref="Column"/> for the result of the UDF.
+        /// </returns>
+        public static Func<Column> Udf(Func<GenericRow> udf, StructType returnType)
+        {
+            return CreateUdf(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf), returnType).Apply0;
+        }
+
+        /// <summary>Creates a UDF from the specified delegate.</summary>
+        /// <typeparam name="T">Specifies the type of the first argument to the UDF.</typeparam>
+        /// <param name="udf">The UDF function implementation.</param>
+        /// <param name="returnType">Schema associated with this row</param>
+        /// <returns>
+        /// A delegate that returns a <see cref="Column"/> for the result of the UDF.
+        /// </returns>
+        public static Func<Column, Column> Udf<T>(Func<T, GenericRow> udf, StructType returnType)
+        {
+            return CreateUdf(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf), returnType).Apply1;
+        }
+
+        /// <summary>Creates a UDF from the specified delegate.</summary>
+        /// <typeparam name="T1">Specifies the type of the first argument to the UDF.</typeparam>
+        /// <typeparam name="T2">Specifies the type of the second argument to the UDF.</typeparam>
+        /// <param name="udf">The UDF function implementation.</param>
+        /// <param name="returnType">Schema associated with this row</param>
+        /// <returns>
+        /// A delegate that returns a <see cref="Column"/> for the result of the UDF.
+        /// </returns>
+        public static Func<Column, Column, Column> Udf<T1, T2>(
+            Func<T1, T2, GenericRow> udf, StructType returnType)
+        {
+            return CreateUdf(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf), returnType).Apply2;
+        }
+
+        /// <summary>Creates a UDF from the specified delegate.</summary>
+        /// <typeparam name="T1">Specifies the type of the first argument to the UDF.</typeparam>
+        /// <typeparam name="T2">Specifies the type of the second argument to the UDF.</typeparam>
+        /// <typeparam name="T3">Specifies the type of the third argument to the UDF.</typeparam>
+        /// <param name="udf">The UDF function implementation.</param>
+        /// <param name="returnType">Schema associated with this row</param>
+        /// <returns>
+        /// A delegate that returns a <see cref="Column"/> for the result of the UDF.
+        /// </returns>
+        public static Func<Column, Column, Column, Column> Udf<T1, T2, T3>(
+            Func<T1, T2, T3, GenericRow> udf, StructType returnType)
+        {
+            return CreateUdf(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf), returnType).Apply3;
+        }
+
+        /// <summary>Creates a UDF from the specified delegate.</summary>
+        /// <typeparam name="T1">Specifies the type of the first argument to the UDF.</typeparam>
+        /// <typeparam name="T2">Specifies the type of the second argument to the UDF.</typeparam>
+        /// <typeparam name="T3">Specifies the type of the third argument to the UDF.</typeparam>
+        /// <typeparam name="T4">Specifies the type of the fourth argument to the UDF.</typeparam>
+        /// <param name="udf">The UDF function implementation.</param>
+        /// <param name="returnType">Schema associated with this row</param>
+        /// <returns>
+        /// A delegate that returns a <see cref="Column"/> for the result of the UDF.
+        /// </returns>
+        public static Func<Column, Column, Column, Column, Column> Udf<T1, T2, T3, T4>(
+            Func<T1, T2, T3, T4, GenericRow> udf, StructType returnType)
+        {
+            return CreateUdf(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf), returnType).Apply4;
+        }
+
+        /// <summary>Creates a UDF from the specified delegate.</summary>
+        /// <typeparam name="T1">Specifies the type of the first argument to the UDF.</typeparam>
+        /// <typeparam name="T2">Specifies the type of the second argument to the UDF.</typeparam>
+        /// <typeparam name="T3">Specifies the type of the third argument to the UDF.</typeparam>
+        /// <typeparam name="T4">Specifies the type of the fourth argument to the UDF.</typeparam>
+        /// <typeparam name="T5">Specifies the type of the fifth argument to the UDF.</typeparam>
+        /// <param name="udf">The UDF function implementation.</param>
+        /// <param name="returnType">Schema associated with this row</param>
+        /// <returns>
+        /// A delegate that returns a <see cref="Column"/> for the result of the UDF.
+        /// </returns>
+        public static Func<Column, Column, Column, Column, Column, Column> Udf<T1, T2, T3, T4, T5>(
+            Func<T1, T2, T3, T4, T5, GenericRow> udf, StructType returnType)
+        {
+            return CreateUdf(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf), returnType).Apply5;
+        }
+
+        /// <summary>Creates a UDF from the specified delegate.</summary>
+        /// <typeparam name="T1">Specifies the type of the first argument to the UDF.</typeparam>
+        /// <typeparam name="T2">Specifies the type of the second argument to the UDF.</typeparam>
+        /// <typeparam name="T3">Specifies the type of the third argument to the UDF.</typeparam>
+        /// <typeparam name="T4">Specifies the type of the fourth argument to the UDF.</typeparam>
+        /// <typeparam name="T5">Specifies the type of the fifth argument to the UDF.</typeparam>
+        /// <typeparam name="T6">Specifies the type of the sixth argument to the UDF.</typeparam>
+        /// <param name="udf">The UDF function implementation.</param>
+        /// <param name="returnType">Schema associated with this row</param>
+        /// <returns>
+        /// A delegate that returns a <see cref="Column"/> for the result of the UDF.
+        /// </returns>
+        public static Func<Column, Column, Column, Column, Column, Column, Column> Udf<T1, T2, T3, T4, T5, T6>(
+            Func<T1, T2, T3, T4, T5, T6, GenericRow> udf, StructType returnType)
+        {
+            return CreateUdf(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf), returnType).Apply6;
+        }
+
+        /// <summary>Creates a UDF from the specified delegate.</summary>
+        /// <typeparam name="T1">Specifies the type of the first argument to the UDF.</typeparam>
+        /// <typeparam name="T2">Specifies the type of the second argument to the UDF.</typeparam>
+        /// <typeparam name="T3">Specifies the type of the third argument to the UDF.</typeparam>
+        /// <typeparam name="T4">Specifies the type of the fourth argument to the UDF.</typeparam>
+        /// <typeparam name="T5">Specifies the type of the fifth argument to the UDF.</typeparam>
+        /// <typeparam name="T6">Specifies the type of the sixth argument to the UDF.</typeparam>
+        /// <typeparam name="T7">Specifies the type of the seventh argument to the UDF.</typeparam>
+        /// <param name="udf">The UDF function implementation.</param>
+        /// <param name="returnType">Schema associated with this row</param>
+        /// <returns>
+        /// A delegate that returns a <see cref="Column"/> for the result of the UDF.
+        /// </returns>
+        public static Func<Column, Column, Column, Column, Column, Column, Column, Column> Udf<T1, T2, T3, T4, T5, T6, T7>(
+            Func<T1, T2, T3, T4, T5, T6, T7, GenericRow> udf, StructType returnType)
+        {
+            return CreateUdf(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf), returnType).Apply7;
+        }
+
+        /// <summary>Creates a UDF from the specified delegate.</summary>
+        /// <typeparam name="T1">Specifies the type of the first argument to the UDF.</typeparam>
+        /// <typeparam name="T2">Specifies the type of the second argument to the UDF.</typeparam>
+        /// <typeparam name="T3">Specifies the type of the third argument to the UDF.</typeparam>
+        /// <typeparam name="T4">Specifies the type of the fourth argument to the UDF.</typeparam>
+        /// <typeparam name="T5">Specifies the type of the fifth argument to the UDF.</typeparam>
+        /// <typeparam name="T6">Specifies the type of the sixth argument to the UDF.</typeparam>
+        /// <typeparam name="T7">Specifies the type of the seventh argument to the UDF.</typeparam>
+        /// <typeparam name="T8">Specifies the type of the eighth argument to the UDF.</typeparam>
+        /// <param name="udf">The UDF function implementation.</param>
+        /// <param name="returnType">Schema associated with this row</param>
+        /// <returns>
+        /// A delegate that returns a <see cref="Column"/> for the result of the UDF.
+        /// </returns>
+        public static Func<Column, Column, Column, Column, Column, Column, Column, Column, Column> Udf<T1, T2, T3, T4, T5, T6, T7, T8>(
+            Func<T1, T2, T3, T4, T5, T6, T7, T8, GenericRow> udf, StructType returnType)
+        {
+            return CreateUdf(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf), returnType).Apply8;
+        }
+
+        /// <summary>Creates a UDF from the specified delegate.</summary>
+        /// <typeparam name="T1">Specifies the type of the first argument to the UDF.</typeparam>
+        /// <typeparam name="T2">Specifies the type of the second argument to the UDF.</typeparam>
+        /// <typeparam name="T3">Specifies the type of the third argument to the UDF.</typeparam>
+        /// <typeparam name="T4">Specifies the type of the fourth argument to the UDF.</typeparam>
+        /// <typeparam name="T5">Specifies the type of the fifth argument to the UDF.</typeparam>
+        /// <typeparam name="T6">Specifies the type of the sixth argument to the UDF.</typeparam>
+        /// <typeparam name="T7">Specifies the type of the seventh argument to the UDF.</typeparam>
+        /// <typeparam name="T8">Specifies the type of the eighth argument to the UDF.</typeparam>
+        /// <typeparam name="T9">Specifies the type of the ninth argument to the UDF.</typeparam>
+        /// <param name="udf">The UDF function implementation.</param>
+        /// <param name="returnType">Schema associated with this row</param>
+        /// <returns>A delegate that when invoked will return a <see cref="Column"/> for the result of the UDF.</returns>
+        public static Func<Column, Column, Column, Column, Column, Column, Column, Column, Column, Column> Udf<T1, T2, T3, T4, T5, T6, T7, T8, T9>(
+            Func<T1, T2, T3, T4, T5, T6, T7, T8, T9, GenericRow> udf, StructType returnType)
+        {
+            return CreateUdf(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf), returnType).Apply9;
+        }
+
+        /// <summary>Creates a UDF from the specified delegate.</summary>
+        /// <typeparam name="T1">Specifies the type of the first argument to the UDF.</typeparam>
+        /// <typeparam name="T2">Specifies the type of the second argument to the UDF.</typeparam>
+        /// <typeparam name="T3">Specifies the type of the third argument to the UDF.</typeparam>
+        /// <typeparam name="T4">Specifies the type of the fourth argument to the UDF.</typeparam>
+        /// <typeparam name="T5">Specifies the type of the fifth argument to the UDF.</typeparam>
+        /// <typeparam name="T6">Specifies the type of the sixth argument to the UDF.</typeparam>
+        /// <typeparam name="T7">Specifies the type of the seventh argument to the UDF.</typeparam>
+        /// <typeparam name="T8">Specifies the type of the eighth argument to the UDF.</typeparam>
+        /// <typeparam name="T9">Specifies the type of the ninth argument to the UDF.</typeparam>
+        /// <typeparam name="T10">Specifies the type of the tenth argument to the UDF.</typeparam>
+        /// <param name="udf">The UDF function implementation.</param>
+        /// <param name="returnType">Schema associated with this row</param>
+        /// <returns>
+        /// A delegate that returns a <see cref="Column"/> for the result of the UDF.
+        /// </returns>
+        public static Func<Column, Column, Column, Column, Column, Column, Column, Column, Column, Column, Column> Udf<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10>(
+            Func<T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, GenericRow> udf, StructType returnType)
+        {
+            return CreateUdf(udf.Method.ToString(), UdfUtils.CreateUdfWrapper(udf), returnType).Apply10;
+        }
+
         /// <summary>Creates a Vector UDF from the specified delegate.</summary>
         /// <typeparam name="T">Specifies the type of the first argument to the UDF.</typeparam>
         /// <typeparam name="TResult">Specifies the return type of the UDF.</typeparam>
@@ -4071,6 +4255,11 @@ private static UserDefinedFunction CreateUdf<TResult>(string name, Delegate exec
             return CreateUdf<TResult>(name, execute, UdfUtils.PythonEvalType.SQL_BATCHED_UDF);
         }
 
+        private static UserDefinedFunction CreateUdf(string name, Delegate execute, StructType returnType)
+        {
+            return CreateUdf(name, execute, UdfUtils.PythonEvalType.SQL_BATCHED_UDF, returnType);
+        }
+
         private static UserDefinedFunction CreateVectorUdf<TResult>(string name, Delegate execute)
         {
             return CreateUdf<TResult>(name, execute, UdfUtils.PythonEvalType.SQL_SCALAR_PANDAS_UDF);
@@ -4079,7 +4268,21 @@ private static UserDefinedFunction CreateVectorUdf<TResult>(string name, Delegat
         private static UserDefinedFunction CreateUdf<TResult>(
             string name,
             Delegate execute,
-            UdfUtils.PythonEvalType evalType)
+            UdfUtils.PythonEvalType evalType) =>
+            CreateUdf(name, execute, evalType, UdfUtils.GetReturnType(typeof(TResult)));
+
+        private static UserDefinedFunction CreateUdf(
+            string name,
+            Delegate execute,
+            UdfUtils.PythonEvalType evalType,
+            StructType returnType) =>
+            CreateUdf(name, execute, evalType, returnType.Json);
+
+        private static UserDefinedFunction CreateUdf(
+            string name,
+            Delegate execute,
+            UdfUtils.PythonEvalType evalType,
+            string returnType)
         {
             return UserDefinedFunction.Create(
                 name,
@@ -4088,7 +4291,7 @@ private static UserDefinedFunction CreateUdf<TResult>(
                     CommandSerDe.SerializedMode.Row,
                     CommandSerDe.SerializedMode.Row),
                 evalType,
-                UdfUtils.GetReturnType(typeof(TResult)));
+                returnType);
         }
 
         private static Column ApplyFunction(string funcName)