SerDe referenced assemblies (#180)

Author: suhsteve

src/csharp/Microsoft.Spark.E2ETest.ExternalLibrary/ExternalClass.cs

@@ -0,0 +1,29 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+
+namespace Microsoft.Spark.E2ETest.ExternalLibrary
+{
+    [Serializable]
+    public class ExternalClass
+    {
+        private string _s;
+
+        public ExternalClass(string s)
+        {
+            _s = s;
+        }
+
+        public static string HelloWorld()
+        {
+            return "Hello World";
+        }
+
+        public string Concat(string s)
+        {
+            return _s + s;
+        }
+    }
+}

src/csharp/Microsoft.Spark.E2ETest.ExternalLibrary/Microsoft.Spark.E2ETest.ExternalLibrary.csproj

@@ -0,0 +1,8 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFramework>netstandard2.0</TargetFramework>
+    <PublicSign>false</PublicSign>
+  </PropertyGroup>
+
+</Project>

src/csharp/Microsoft.Spark.E2ETest/Microsoft.Spark.E2ETest.csproj

@@ -11,6 +11,7 @@
   </ItemGroup>
 
   <ItemGroup>
+    <ProjectReference Include="..\Microsoft.Spark.E2ETest.ExternalLibrary\Microsoft.Spark.E2ETest.ExternalLibrary.csproj" />
     <ProjectReference Include="..\Microsoft.Spark.Experimental\Microsoft.Spark.Experimental.csproj" />
     <ProjectReference Include="..\Microsoft.Spark.Worker\Microsoft.Spark.Worker.csproj" />
     <ProjectReference Include="..\Microsoft.Spark\Microsoft.Spark.csproj" />

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

@@ -0,0 +1,94 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using Microsoft.Spark.E2ETest.ExternalLibrary;
+using Microsoft.Spark.Sql;
+using Xunit;
+using static Microsoft.Spark.Sql.Functions;
+
+namespace Microsoft.Spark.E2ETest.UdfTests
+{
+    [Collection("Spark E2E Tests")]
+    public class UdfSerDeTests
+    {
+        private readonly SparkSession _spark;
+        private readonly DataFrame _df;
+
+        public UdfSerDeTests(SparkFixture fixture)
+        {
+            _spark = fixture.Spark;
+            _df = _spark
+                .Read()
+                .Schema("age INT, name STRING")
+                .Json($"{TestEnvironment.ResourceDirectory}people.json");
+        }
+
+        [Fact]
+        public void TestUdfClosure()
+        {
+            var ec = new ExternalClass("Hello");
+            Func<Column, Column> udf = Udf<string, string>(
+                (str) =>
+                {
+                    return ec.Concat(str);
+                });
+
+            Row[] rows = _df.Select(udf(_df["name"])).Collect().ToArray();
+            Assert.Equal(3, rows.Length);
+
+            var expected = new[] { "HelloMichael", "HelloAndy", "HelloJustin" };
+            for (int i = 0; i < rows.Length; ++i)
+            {
+                Row row = rows[i];
+                Assert.Equal(1, row.Size());
+                Assert.Equal(expected[i], row.GetAs<string>(0));
+            }
+        }
+
+        [Fact]
+        public void TestExternalStaticMethodCall()
+        {
+            Func<Column, Column> udf = Udf<string, string>(str =>
+            {
+                return ExternalClass.HelloWorld();
+            });
+
+            Row[] rows = _df.Select(udf(_df["name"])).Collect().ToArray();
+            Assert.Equal(3, rows.Length);
+
+            for (int i = 0; i < rows.Length; ++i)
+            {
+                Row row = rows[i];
+                Assert.Equal(1, row.Size());
+                Assert.Equal("Hello World", row.GetAs<string>(0));
+            }
+        }
+
+        [Fact]
+        public void TestInitExternalClassInUdf()
+        {
+            // Instantiate external assembly class within body of Udf.
+            Func<Column, Column> udf = Udf<string, string>(
+                    (str) =>
+                    {
+                        var ec = new ExternalClass("Hello");
+                        return ec.Concat(str);
+                    });
+
+            Row[] rows = _df.Select(udf(_df["name"])).Collect().ToArray();
+            Assert.Equal(3, rows.Length);
+
+            var expected = new[] { "HelloMichael", "HelloAndy", "HelloJustin" };
+            for (int i = 0; i < rows.Length; ++i)
+            {
+                Row row = rows[i];
+                Assert.Equal(1, row.Size());
+                Assert.Equal(expected[i], row.GetAs<string>(0));
+            }
+        }
+    }
+}

src/csharp/Microsoft.Spark.Worker/Processor/CommandProcessor.cs

@@ -10,6 +10,7 @@
 using static Microsoft.Spark.Utils.UdfUtils;
 
 #if NETCOREAPP
+using System.Reflection;
 using System.Runtime.Loader;
 #endif
 
@@ -19,12 +20,19 @@ internal sealed class CommandProcessor
     {
         private readonly Version _version;
 
-#if NETCOREAPP
         static CommandProcessor()
         {
-            UdfSerDe.AssemblyLoader = AssemblyLoadContext.Default.LoadFromAssemblyPath;
-        }
+#if NETCOREAPP
+            AssemblyLoader.LoadFromFile = AssemblyLoadContext.Default.LoadFromAssemblyPath;
+            AssemblyLoader.LoadFromName = (asmName) =>
+                AssemblyLoadContext.Default.LoadFromAssemblyName(new AssemblyName(asmName));
+            AssemblyLoadContext.Default.Resolving += (assemblyLoadContext, assemblyName) =>
+                AssemblyLoader.ResolveAssembly(assemblyName.FullName);
+#else
+            AppDomain.CurrentDomain.AssemblyResolve += (object sender, ResolveEventArgs args) =>
+                AssemblyLoader.ResolveAssembly(args.Name);
 #endif
+        }
 
         internal CommandProcessor(Version version)
         {

src/csharp/Microsoft.Spark.sln

@@ -1,6 +1,6 @@
 Microsoft Visual Studio Solution File, Format Version 12.00
-# Visual Studio 15
-VisualStudioVersion = 15.0.28010.2046
+# Visual Studio Version 16
+VisualStudioVersion = 16.0.29009.5
 MinimumVisualStudioVersion = 10.0.40219.1
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.Spark", "Microsoft.Spark\Microsoft.Spark.csproj", "{2B4236DD-00A9-4B24-9041-6F9727F3B025}"
 EndProject
@@ -25,6 +25,8 @@ Project("{6EC3EE1D-3C4E-46DD-8F32-0CC8E7565705}") = "Microsoft.Spark.FSharp.Exam
 EndProject
 Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.Spark.Experimental", "Microsoft.Spark.Experimental\Microsoft.Spark.Experimental.csproj", "{7F276D07-6D94-49E3-A305-4B965DE7A700}"
 EndProject
+Project("{9A19103F-16F7-4668-BE54-9A1E7A4F7556}") = "Microsoft.Spark.E2ETest.ExternalLibrary", "Microsoft.Spark.E2ETest.ExternalLibrary\Microsoft.Spark.E2ETest.ExternalLibrary.csproj", "{920A2CC8-4075-41D1-BC30-C496430BD9E4}"
+EndProject
 Global
 	GlobalSection(SolutionConfigurationPlatforms) = preSolution
 		Debug|Any CPU = Debug|Any CPU
@@ -63,6 +65,10 @@ Global
 		{7F276D07-6D94-49E3-A305-4B965DE7A700}.Debug|Any CPU.Build.0 = Debug|Any CPU
 		{7F276D07-6D94-49E3-A305-4B965DE7A700}.Release|Any CPU.ActiveCfg = Release|Any CPU
 		{7F276D07-6D94-49E3-A305-4B965DE7A700}.Release|Any CPU.Build.0 = Release|Any CPU
+		{920A2CC8-4075-41D1-BC30-C496430BD9E4}.Debug|Any CPU.ActiveCfg = Debug|Any CPU
+		{920A2CC8-4075-41D1-BC30-C496430BD9E4}.Debug|Any CPU.Build.0 = Debug|Any CPU
+		{920A2CC8-4075-41D1-BC30-C496430BD9E4}.Release|Any CPU.ActiveCfg = Release|Any CPU
+		{920A2CC8-4075-41D1-BC30-C496430BD9E4}.Release|Any CPU.Build.0 = Release|Any CPU
 	EndGlobalSection
 	GlobalSection(SolutionProperties) = preSolution
 		HideSolutionNode = FALSE

src/csharp/Microsoft.Spark/Utils/AssemblyLoader.cs

@@ -0,0 +1,132 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Reflection;
+using System.Runtime.InteropServices;
+
+namespace Microsoft.Spark.Utils
+{
+    internal static class AssemblyLoader
+    {
+        internal static Func<string, Assembly> LoadFromFile { get; set; } = Assembly.LoadFrom;
+
+        internal static Func<string, Assembly> LoadFromName { get; set; } = Assembly.Load;
+
+        private static readonly Dictionary<string, Assembly> s_assemblyCache =
+            new Dictionary<string, Assembly>();
+
+        private static readonly string[] s_searchPaths =
+            new[] { Directory.GetCurrentDirectory(), AppDomain.CurrentDomain.BaseDirectory };
+
+        private static readonly string[] s_extensions =
+            RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ?
+                    new[] { ".dll", ".exe", ".ni.dll", ".ni.exe" } :
+                    new[] { ".dll", ".ni.dll" };
+
+        private static readonly object s_cacheLock = new object();
+
+        /// <summary>
+        /// Return the cached assembly, otherwise attempt to load and cache the assembly
+        /// by searching for the assembly filename in the search paths.
+        /// </summary>
+        /// <param name="assemblyName">The full name of the assembly</param>
+        /// <param name="assemblyFileName">Name of the file that contains the assembly</param>
+        /// <returns>Cached or Loaded Assembly</returns>
+        /// <exception cref="FileNotFoundException">Thrown if the assembly is not
+        /// found.</exception>
+        internal static Assembly LoadAssembly(string assemblyName, string assemblyFileName)
+        {
+            lock (s_cacheLock)
+            {
+                if (s_assemblyCache.TryGetValue(assemblyName, out Assembly assembly))
+                {
+                    return assembly;
+                }
+
+                if (TryLoadAssembly(assemblyFileName, ref assembly))
+                {
+                    s_assemblyCache[assemblyName] = assembly;
+                    return assembly;
+                }
+
+                throw new FileNotFoundException($"Assembly '{assemblyName}' file not found: '{assemblyFileName}'");
+            }
+        }
+
+        /// <summary>
+        /// Return the cached assembly, otherwise look in the following probing paths,
+        /// searching for the simple assembly name and s_extension combination.
+        /// 1) The working directory
+        /// 2) The directory of the application
+        /// </summary>
+        /// <param name="assemblyName">The fullname of the assembly to load</param>
+        /// <returns>The loaded assembly</returns>
+        /// <exception cref="FileNotFoundException">Thrown if the assembly is not
+        /// found.</exception>
+        internal static Assembly ResolveAssembly(string assemblyName)
+        {
+            lock (s_cacheLock)
+            {
+                if (s_assemblyCache.TryGetValue(assemblyName, out Assembly assembly))
+                {
+                    return assembly;
+                }
+
+                string simpleAsmName = new AssemblyName(assemblyName).Name;
+                foreach (string extension in s_extensions)
+                {
+                    string assemblyFileName = $"{simpleAsmName}{extension}";
+                    if (TryLoadAssembly(assemblyFileName, ref assembly))
+                    {
+                        s_assemblyCache[assemblyName] = assembly;
+                        return assembly;
+                    }
+                }
+
+                throw new FileNotFoundException($"Assembly file not found: '{assemblyName}'");
+            }
+        }
+
+        /// <summary>
+        /// Returns the loaded assembly by probing the following locations in order:
+        /// 1) The working directory
+        /// 2) The directory of the application
+        /// </summary>
+        /// <remarks>
+        /// The probing order is important in cases when spark is launched on
+        /// YARN. The executors are run inside 'containers' and files that are passed
+        /// via 'spark-submit --files' will be pushed to these 'containers'. This path
+        /// is the working directory and the 1st probing path that will be checked.
+        /// </remarks>
+        /// <param name="assemblyFileName">Name of the file that contains the assembly</param>
+        /// <param name="assembly">The loaded assembly.</param>
+        /// <returns>True if assembly is loaded, false otherwise.</returns>
+        private static bool TryLoadAssembly(string assemblyFileName, ref Assembly assembly)
+        {
+            foreach (string searchPath in s_searchPaths)
+            {
+                string assemblyPath = Path.Combine(searchPath, assemblyFileName);
+                if (File.Exists(assemblyPath))
+                {
+                    try
+                    {
+                        assembly = LoadFromFile(assemblyPath);
+                        return true;
+                    }
+                    catch (Exception ex) when (
+                        ex is FileLoadException ||
+                        ex is BadImageFormatException)
+                    {
+                        // Ignore invalid assemblies.
+                    }
+                }
+            }
+
+            return false;
+        }
+    }
+}

src/csharp/Microsoft.Spark/Utils/CommandSerDe.cs

@@ -212,7 +212,7 @@ private static void SerializeUdfs(
 
             foreach (UdfSerDe.FieldData field in fields)
             {
-                SerializeUdfs((Delegate)field.ValueData.Value, curNode, udfWrapperNodes, udfs);
+                SerializeUdfs((Delegate)field.Value, curNode, udfWrapperNodes, udfs);
             }
         }