Support UDF that returns RowType

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

@@ -158,25 +158,58 @@ 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) =>
-                {
-                    var structFields = new List<StructField>()
+            // Single GenericRow
+            var schema1 = new StructType(new[]
+            {
+                new StructField("col1", new IntegerType()),
+                new StructField("col2", new StringType())
+            });
+            Func<Column, Column> udf1 = Udf<string>(
+                str => new GenericRow(new object[] { 1, "abc" }), schema1);
+
+            Row[] rows1 = _df.Select(udf1(_df["name"])).Collect().ToArray();
+            Assert.Equal(3, rows1.Length);
+
+            foreach (Row row in rows1)
+            {
+                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("subCol1", new IntegerType())
+            });
+            var subSchema2 = new StructType(new[]
+            {
+                new StructField("subCol2", new StringType())
+            });
+            var schema2 = new StructType(new[]
+            {
+                new StructField("col1", subSchema1),
+                new StructField("col2", subSchema2)
+            });
+            Func<Column, Column> udf2 = Udf<string>(
+                str => new GenericRow(
+                    new object[]
                     {
-                        new StructField("name", new StringType()),
-                    };
-                    var schema = new StructType(structFields);
-                    var row = new Row(new object[] { str }, schema);
-                    return row;
-                }));
+                        new GenericRow(new object[] { 1 }),
+                        new GenericRow(new object[] { "abc" })
+                    }), schema2);
+
+            Row[] rows2 = _df.Select(udf2(_df["name"])).Collect().ToArray();
+            Assert.Equal(3, rows2.Length);
+
+            foreach (Row row in rows2)
+            {
+                Assert.Equal(2, row.Size());
+                Assert.IsType<Row>(row.Get("col1"));
+                Assert.IsType<Row>(row.Get("col2"));
+                Assert.Equal(new Row(new object[] { 1 }, subSchema1), row.GetAs<Row>("col1"));
+                Assert.Equal(new Row(new object[] { "abc" }, subSchema2), row.GetAs<Row>("col2"));
+            }
         }
     }
 }

src/csharp/Microsoft.Spark.Worker/Command/SqlCommandExecutor.cs

@@ -165,7 +165,7 @@ private void WriteOutput(Stream stream, IEnumerable<object> rows, int sizeHint)
             if (s_outputBuffer == null)
                 s_outputBuffer = new byte[sizeHint];
 
-            Pickler pickler = s_pickler ?? (s_pickler = new Pickler(false));
+            Pickler pickler = s_pickler ?? (s_pickler = PythonSerDe.CreatePickler());
             pickler.dumps(rows, ref s_outputBuffer, out int bytesWritten);
 
             if (bytesWritten <= 0)

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)