Changes to support Arrow 2.0 and Spark 3.0

src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/DataFrameTests.cs

@@ -20,6 +20,8 @@
 using DataFrame = Microsoft.Spark.Sql.DataFrame;
 using FxDataFrame = Microsoft.Data.Analysis.DataFrame;
 using Int32Type = Apache.Arrow.Types.Int32Type;
+using ArrowStructType = Apache.Arrow.Types.StructType;
+using System.Diagnostics;
 
 namespace Microsoft.Spark.E2ETest.IpcTests
 {
@@ -290,7 +292,7 @@ public void TestDataFrameVectorUdf()
             }
         }
 
-        [SkipIfSparkVersionIsGreaterOrEqualTo(Versions.V3_0_0)]
+        [Fact]
         public void TestGroupedMapUdf()
         {
             DataFrame df = _spark
@@ -355,9 +357,11 @@ private static RecordBatch ArrowBasedCountCharacters(RecordBatch records)
             // Return 1 record, if we were given any. 0, otherwise.
             int returnLength = records.Length > 0 ? 1 : 0;
 
+            ArrowStructType structType = new ArrowStructType(new List<Field> { ageField, new Field("name_CharCount", Int32Type.Default, true) });
+
             return new RecordBatch(
                 new Schema.Builder()
-                    .Field(ageField)
+                    .Field(ageField).Field(new Field("Ret Struct", structType, true))
                     .Field(f => f.Name("name_CharCount").DataType(Int32Type.Default))
                     .Build(),
                 new IArrowArray[]
@@ -368,7 +372,7 @@ private static RecordBatch ArrowBasedCountCharacters(RecordBatch records)
                 returnLength);
         }
 
-        [SkipIfSparkVersionIsGreaterOrEqualTo(Versions.V3_0_0)]
+        [Fact]
         public void TestDataFrameGroupedMapUdf()
         {
             DataFrame df = _spark

src/csharp/Microsoft.Spark.UnitTest/WorkerFunctionTests.cs

@@ -110,7 +110,7 @@ public void TestArrowWorkerFunctionForBool()
                 new ArrowUdfWrapper<StringArray, BooleanArray, BooleanArray>(
                     (strings, flags) => (BooleanArray)ToArrowArray(
                         Enumerable.Range(0, strings.Length)
-                            .Select(i => flags.GetBoolean(i) || strings.GetString(i).Contains("true"))
+                            .Select(i => flags.GetValue(i).Value || strings.GetString(i).Contains("true"))
                             .ToArray())).Execute);
 
             IArrowArray[] input = new[]
@@ -120,10 +120,10 @@ public void TestArrowWorkerFunctionForBool()
             };
             var results = (BooleanArray)func.Func(input, new[] { 0, 1 });
             Assert.Equal(4, results.Length);
-            Assert.True(results.GetBoolean(0));
-            Assert.True(results.GetBoolean(1));
-            Assert.True(results.GetBoolean(2));
-            Assert.False(results.GetBoolean(3));
+            Assert.True(results.GetValue(0).Value);
+            Assert.True(results.GetValue(1).Value);
+            Assert.True(results.GetValue(2).Value);
+            Assert.False(results.GetValue(3).Value);
         }
 
         /// <summary>

src/csharp/Microsoft.Spark.Worker.UnitTest/CommandExecutorTests.cs

@@ -981,15 +981,18 @@ await arrowWriter.WriteRecordBatchAsync(
             RecordBatch outputBatch = await arrowReader.ReadNextRecordBatchAsync();
 
             Assert.Equal(numRows, outputBatch.Length);
-            Assert.Equal(2, outputBatch.ColumnCount);
+            Assert.Equal(1, outputBatch.ColumnCount);
 
-            var stringArray = (StringArray)outputBatch.Column(0);
+            var structArray = (StructArray)outputBatch.Column(0);
+            Assert.Equal(2, structArray.Fields.Count);
+
+            var stringArray = (StringArray)structArray.Fields[0];
             for (int i = 0; i < numRows; ++i)
             {
                 Assert.Equal($"udf: {i}", stringArray.GetString(i));
             }
 
-            var doubleArray = (DoubleArray)outputBatch.Column(1);
+            var doubleArray = (DoubleArray)structArray.Fields[1];
             for (int i = 0; i < numRows; ++i)
             {
                 Assert.Equal(100 + i, doubleArray.Values[i]);

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

@@ -12,6 +12,7 @@
 using Apache.Arrow.Ipc;
 using Apache.Arrow.Types;
 using Microsoft.Data.Analysis;
+using Microsoft.Spark.Interop;
 using Microsoft.Spark.Interop.Ipc;
 using Microsoft.Spark.Sql;
 using Microsoft.Spark.Utils;
@@ -737,6 +738,25 @@ protected internal override CommandExecutorStat ExecuteCore(
             return ExecuteArrowGroupedMapCommand(inputStream, outputStream, commands);
         }
 
+        private RecordBatch WrapArrowRecordBatchColumnsInAStruct(RecordBatch batch)
+        {
+            if (SparkEnvironment.SparkVersion >= new Version(Versions.V3_0_0))
+            {
+                ArrowBuffer.BitmapBuilder validityBitmapBuilder = new ArrowBuffer.BitmapBuilder();
+                for (int i = 0; i < batch.Length; i++)
+                {
+                    validityBitmapBuilder.Append(true);
+                }
+                ArrowBuffer validityBitmap = validityBitmapBuilder.Build();
+
+                StructType structType = new StructType(batch.Schema.Fields.Select((KeyValuePair<string, Field> pair) => pair.Value).ToList());
+                StructArray structArray = new StructArray(structType, batch.Length, batch.Arrays.Cast<Apache.Arrow.Array>(), validityBitmap);
+                Schema schema = new Schema.Builder().Field(new Field("Struct", structType, false)).Build();
+                return new RecordBatch(schema, new[] { structArray }, batch.Length);
+            }
+            return batch;
+        }
+
         private CommandExecutorStat ExecuteArrowGroupedMapCommand(
             Stream inputStream,
             Stream outputStream,
@@ -754,8 +774,9 @@ private CommandExecutorStat ExecuteArrowGroupedMapCommand(
             ArrowStreamWriter writer = null;
             foreach (RecordBatch input in GetInputIterator(inputStream))
             {
-                RecordBatch result = worker.Func(input);
+                RecordBatch batch = worker.Func(input);
 
+                RecordBatch result = WrapArrowRecordBatchColumnsInAStruct(batch);
                 int numEntries = result.Length;
                 stat.NumEntriesProcessed += numEntries;
 
@@ -794,20 +815,22 @@ private CommandExecutorStat ExecuteDataFrameGroupedMapCommand(
             {
                 FxDataFrame dataFrame = FxDataFrame.FromArrowRecordBatch(input);
                 FxDataFrame resultDataFrame = worker.Func(dataFrame);
+
                 IEnumerable<RecordBatch> recordBatches = resultDataFrame.ToArrowRecordBatches();
 
-                foreach (RecordBatch result in recordBatches)
+                foreach (RecordBatch batch in recordBatches)
                 {
-                    stat.NumEntriesProcessed += result.Length;
+                    RecordBatch final = WrapArrowRecordBatchColumnsInAStruct(batch);
+                    stat.NumEntriesProcessed += final.Length;
 
                     if (writer == null)
                     {
                         writer =
-                            new ArrowStreamWriter(outputStream, result.Schema, leaveOpen: true, ipcOptions);
+                            new ArrowStreamWriter(outputStream, final.Schema, leaveOpen: true, ipcOptions);
                     }
 
                     // TODO: Remove sync-over-async once WriteRecordBatch exists.
-                    writer.WriteRecordBatchAsync(result).GetAwaiter().GetResult();
+                    writer.WriteRecordBatchAsync(final).GetAwaiter().GetResult();
                 }
             }
 

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

@@ -28,7 +28,7 @@
   </ItemGroup>
 
   <ItemGroup>
-    <PackageReference Include="Apache.Arrow" Version="0.15.1" />
+    <PackageReference Include="Apache.Arrow" Version="2.0.0" />
     <PackageReference Include="Microsoft.CSharp" Version="4.5.0" />
     <PackageReference Include="Microsoft.Data.Analysis" Version="0.4.0" />
     <PackageReference Include="Newtonsoft.Json" Version="11.0.2" />
@@ -37,10 +37,7 @@
   </ItemGroup>
 
   <ItemGroup>
-    <Content Include="..\..\scala\microsoft-spark-*\target\microsoft-spark-*.jar"
-             Link="jars\%(Filename)%(Extension)"
-             Pack="true"
-             PackagePath="jars\%(Filename)%(Extension)" />
+    <Content Include="..\..\scala\microsoft-spark-*\target\microsoft-spark-*.jar" Link="jars\%(Filename)%(Extension)" Pack="true" PackagePath="jars\%(Filename)%(Extension)" />
     <Content Include="build\**" Pack="true" PackagePath="build" />
   </ItemGroup>