Expose DataStreamWriter.Foreach API

src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/Streaming/DataStreamWriterTests.cs

@@ -2,10 +2,16 @@
 // 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.Linq;
+using Microsoft.Spark.E2ETest.Utils;
 using Microsoft.Spark.Sql;
 using Microsoft.Spark.Sql.Streaming;
+using Microsoft.Spark.Sql.Types;
 using Xunit;
+using static Microsoft.Spark.Sql.Functions;
 
 namespace Microsoft.Spark.E2ETest.IpcTests
 {
@@ -59,5 +65,213 @@ public void TestSignaturesV2_3_X()
 
             Assert.IsType<DataStreamWriter>(dsw.Trigger(Trigger.Once()));
         }
+
+        [SkipIfSparkVersionIsLessThan(Versions.V2_4_0)]
+        public void TestForeach()
+        {
+            // Temporary folder to put our test stream input.
+            using var srcTempDirectory = new TemporaryDirectory();
+            string streamInputPath = Path.Combine(srcTempDirectory.Path, "streamInput");
+
+            Func<Column, Column> intToStrUdf = Udf<int, string>(i => i.ToString());
+
+            // id column: [1, 2, ..., 99]
+            // idStr column: "id" column converted to string
+            // idAndIdStr column: Struct column composed from the "id" and "idStr" column.
+            _spark.Range(1, 100)
+                .WithColumn("idStr", intToStrUdf(Col("id")))
+                .WithColumn("idAndIdStr", Struct("id", "idStr"))
+                .Write()
+                .Json(streamInputPath);
+
+            // Test a scenario where IForeachWriter runs without issues.
+            // If everything is working as expected, then:
+            // - Triggering stream will not throw an exception
+            // - 3 CSV files will be created in the temporary directory.
+            // - 0 Exception files will be created in the temporary directory.
+            // - The CSV files will contain valid data to read, where the
+            //   expected entries will contain [1111, 2222, ..., 99999999]
+            TestAndValidateForeach(
+                streamInputPath,
+                new TestForeachWriter(),
+                3,
+                0,
+                Enumerable.Range(1, 99).Select(i => Convert.ToInt32($"{i}{i}{i}{i}")));
+
+            // Test scenario where IForeachWriter.Open returns false.
+            // When IForeachWriter.Open returns false, then IForeachWriter.Process
+            // is not called. Verify that:
+            // - Triggering stream will not throw an exception
+            // - 3 CSV files will be created in the temporary directory.
+            // - 0 Exception files will be created in the temporary directory.
+            // - The CSV files will not contain valid data to read.
+            TestAndValidateForeach(
+                streamInputPath,
+                new TestForeachWriterOpenFailure(),
+                3,
+                0,
+                Enumerable.Empty<int>());
+
+            // Test scenario where IForeachWriter.Process throws an Exception.
+            // When IForeachWriter.Process throws an Exception, then the exception
+            // is rethrown by ForeachWriterWrapper. We will limit the partitions
+            // to 1 to make validating this scenario simpler. Verify that:
+            // - Triggering stream throws an exception.
+            // - 1 CSV file will be created in the temporary directory.
+            // - 1 Exception will be created in the temporary directory. The
+            //   thrown exception from Process() will be sent to Close().
+            // - The CSV file will not contain valid data to read.
+            TestAndValidateForeach(
+                streamInputPath,
+                new TestForeachWriterProcessFailure(),
+                1,
+                1,
+                Enumerable.Empty<int>());
+        }
+
+        private void TestAndValidateForeach(
+            string streamInputPath,
+            TestForeachWriter foreachWriter,
+            int expectedCSVFiles,
+            int expectedExceptionFiles,
+            IEnumerable<int> expectedOutput)
+        {
+            // Temporary folder the TestForeachWriter will write to.
+            using var dstTempDirectory = new TemporaryDirectory();
+            foreachWriter.WritePath = dstTempDirectory.Path;
+
+            // Read streamInputPath, repartition data, then
+            // call TestForeachWriter on the data.
+            DataStreamWriter dsw = _spark
+                .ReadStream()
+                .Schema(new StructType(new[]
+                {
+                    new StructField("id", new IntegerType()),
+                    new StructField("idStr", new StringType()),
+                    new StructField("idAndIdStr", new StructType(new[]
+                    {
+                        new StructField("id", new IntegerType()),
+                        new StructField("idStr", new StringType())
+                    }))
+                }))
+                .Json(streamInputPath)
+                .Repartition(expectedCSVFiles)
+                .WriteStream()
+                .Foreach(foreachWriter);
+
+            // Trigger the stream batch once.
+            if (expectedExceptionFiles > 0)
+            {
+                Assert.Throws<Exception>(
+                    () => dsw.Trigger(Trigger.Once()).Start().AwaitTermination());
+            }
+            else
+            {
+                dsw.Trigger(Trigger.Once()).Start().AwaitTermination();
+            }
+
+            // Verify that TestForeachWriter created a unique .csv when
+            // ForeachWriter.Open was called on each partitionId.
+            Assert.Equal(
+                expectedCSVFiles,
+                Directory.GetFiles(dstTempDirectory.Path, "*.csv").Length);
+
+            // Only if ForeachWriter.Process(Row) throws an exception, will
+            // ForeachWriter.Close(Exception) create a file with the
+            // .exception extension.
+            Assert.Equal(
+                expectedExceptionFiles,
+                Directory.GetFiles(dstTempDirectory.Path, "*.exception").Length);
+
+            // Read in the *.csv file(s) generated by the TestForeachWriter.
+            // If there are multiple input files, sorting by "id" will make
+            // validation simpler. Contents of the *.csv will only be populated
+            // on successful calls to the ForeachWriter.Process method.
+            DataFrame foreachWriterOutputDF = _spark
+                .Read()
+                .Schema("id INT")
+                .Csv(dstTempDirectory.Path)
+                .Sort("id");
+
+            // Validate expected *.csv data.
+            Assert.Equal(
+                expectedOutput.Select(i => new object[] { i }),
+                foreachWriterOutputDF.Collect().Select(r => r.Values));
+        }
+
+        [Serializable]
+        private class TestForeachWriter : IForeachWriter
+        {
+            [NonSerialized]
+            private StreamWriter _streamWriter;
+
+            private long _partitionId;
+
+            private long _epochId;
+
+            internal string WritePath { get; set; }
+
+            public void Close(Exception errorOrNull)
+            {
+                if (errorOrNull != null)
+                {
+                    FileStream fs = File.Create(
+                        Path.Combine(
+                            WritePath,
+                            $"Close-{_partitionId}-{_epochId}.exception"));
+                    fs.Dispose();
+                }
+
+                _streamWriter?.Dispose();
+            }
+
+            public virtual bool Open(long partitionId, long epochId)
+            {
+                _partitionId = partitionId;
+                _epochId = epochId;
+                try
+                {
+                    _streamWriter = new StreamWriter(
+                        Path.Combine(
+                            WritePath,
+                            $"sink-foreachWriter-{_partitionId}-{_epochId}.csv"));
+                    return true;
+                }
+                catch
+                {
+                    return false;
+                }
+            }
+
+            public virtual void Process(Row value)
+            {
+                Row idAndIdStr = value.GetAs<Row>("idAndIdStr");
+                _streamWriter.WriteLine(
+                    string.Format("{0}{1}{2}{3}",
+                        value.GetAs<int>("id"),
+                        value.GetAs<string>("idStr"),
+                        idAndIdStr.GetAs<int>("id"),
+                        idAndIdStr.GetAs<string>("idStr")));
+            }
+        }
+
+        [Serializable]
+        private class TestForeachWriterOpenFailure : TestForeachWriter
+        {
+            public override bool Open(long partitionId, long epochId)
+            {
+                base.Open(partitionId, epochId);
+                return false;
+            }
+        }
+
+        [Serializable]
+        private class TestForeachWriterProcessFailure : TestForeachWriter
+        {
+            public override void Process(Row value)
+            {
+                throw new Exception("TestForeachWriterProcessFailure Process(Row) failure.");
+            }
+        }
     }
 }

src/csharp/Microsoft.Spark.Worker/Payload.cs

@@ -3,81 +3,10 @@
 // See the LICENSE file in the project root for more information.
 
 using System.Collections.Generic;
-using System.Linq;
 using Microsoft.Spark.Utils;
 
 namespace Microsoft.Spark.Worker
 {
-    /// <summary>
-    /// TaskContext stores information related to a task.
-    /// </summary>
-    internal class TaskContext
-    {
-        internal int StageId { get; set; }
-
-        internal int PartitionId { get; set; }
-
-        internal int AttemptNumber { get; set; }
-
-        internal long AttemptId { get; set; }
-
-        internal bool IsBarrier { get; set; }
-
-        internal int Port { get; set; }
-
-        internal string Secret { get; set; }
-
-        internal IEnumerable<Resource> Resources { get; set; } = new List<Resource>();
-
-        internal Dictionary<string, string> LocalProperties { get; set; } =
-            new Dictionary<string, string>();
-
-        public override bool Equals(object obj)
-        {
-            if (!(obj is TaskContext other))
-            {
-                return false;
-            }
-
-            return (StageId == other.StageId) &&
-                (PartitionId == other.PartitionId) &&
-                (AttemptNumber == other.AttemptNumber) &&
-                (AttemptId == other.AttemptId) &&
-                Resources.SequenceEqual(other.Resources) &&
-                (LocalProperties.Count == other.LocalProperties.Count) &&
-                !LocalProperties.Except(other.LocalProperties).Any();
-        }
-
-        public override int GetHashCode()
-        {
-            return StageId;
-        }
-
-        internal class Resource
-        {
-            internal string Key { get; set; }
-            internal string Value { get; set; }
-            internal IEnumerable<string> Addresses { get; set; } = new List<string>();
-
-            public override bool Equals(object obj)
-            {
-                if (!(obj is Resource other))
-                {
-                    return false;
-                }
-
-                return (Key == other.Key) &&
-                    (Value == other.Value) &&
-                    Addresses.SequenceEqual(Addresses);
-            }
-
-            public override int GetHashCode()
-            {
-                return Key.GetHashCode();
-            }
-        }
-    }
-
     /// <summary>
     /// BroadcastVariables stores information on broadcast variables.
     /// </summary>

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

@@ -54,7 +54,10 @@ internal Payload Process(Stream stream)
 
             payload.SplitIndex = BinaryPrimitives.ReadInt32BigEndian(splitIndexBytes);
             payload.Version = SerDe.ReadString(stream);
+
             payload.TaskContext = new TaskContextProcessor(_version).Process(stream);
+            TaskContextHolder.Set(payload.TaskContext);
+
             payload.SparkFilesDir = SerDe.ReadString(stream);
 
             if (Utils.SettingUtils.IsDatabricks)

src/csharp/Microsoft.Spark/Attributes.cs

@@ -73,4 +73,12 @@ public DeprecatedAttribute(string version)
         {
         }
     }
+
+    /// <summary>
+    /// Custom attribute to denote that a class is a Udf Wrapper.
+    /// </summary>
+    [AttributeUsage(AttributeTargets.Class)]
+    internal sealed class UdfWrapperAttribute : Attribute
+    {
+    }
 }