diff --git a/src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/DataFrameTests.cs b/src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/DataFrameTests.cs index 378f3b86c..77891945b 100644 --- a/src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/DataFrameTests.cs +++ b/src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/DataFrameTests.cs @@ -290,7 +290,7 @@ public void TestDataFrameVectorUdf() } } - [SkipIfSparkVersionIsGreaterOrEqualTo(Versions.V3_0_0)] + [Fact] public void TestGroupedMapUdf() { DataFrame df = _spark @@ -368,7 +368,7 @@ private static RecordBatch ArrowBasedCountCharacters(RecordBatch records) returnLength); } - [SkipIfSparkVersionIsGreaterOrEqualTo(Versions.V3_0_0)] + [Fact] public void TestDataFrameGroupedMapUdf() { DataFrame df = _spark diff --git a/src/csharp/Microsoft.Spark.UnitTest/WorkerFunctionTests.cs b/src/csharp/Microsoft.Spark.UnitTest/WorkerFunctionTests.cs index 3f0dfc8e3..bc1f3989e 100644 --- a/src/csharp/Microsoft.Spark.UnitTest/WorkerFunctionTests.cs +++ b/src/csharp/Microsoft.Spark.UnitTest/WorkerFunctionTests.cs @@ -110,7 +110,7 @@ public void TestArrowWorkerFunctionForBool() new ArrowUdfWrapper( (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); } /// diff --git a/src/csharp/Microsoft.Spark.Worker.UnitTest/CommandExecutorTests.cs b/src/csharp/Microsoft.Spark.Worker.UnitTest/CommandExecutorTests.cs index 6de026a83..d3ef6175d 100644 --- a/src/csharp/Microsoft.Spark.Worker.UnitTest/CommandExecutorTests.cs +++ b/src/csharp/Microsoft.Spark.Worker.UnitTest/CommandExecutorTests.cs @@ -702,7 +702,7 @@ public void TestDataFrameSqlCommandExecutorWithEmptyInput( IpcOptions ipcOptions) { var udfWrapper = new Sql.DataFrameUdfWrapper( - (strings) => strings.Apply(cur=> $"udf: {cur}")); + (strings) => strings.Apply(cur => $"udf: {cur}")); var command = new SqlCommand() { @@ -874,15 +874,28 @@ await arrowWriter.WriteRecordBatchAsync( RecordBatch outputBatch = await arrowReader.ReadNextRecordBatchAsync(); Assert.Equal(numRows, outputBatch.Length); - Assert.Equal(2, outputBatch.ColumnCount); + StringArray stringArray; + Int64Array longArray; + if (sparkVersion < new Version(Versions.V3_0_0)) + { + Assert.Equal(2, outputBatch.ColumnCount); + stringArray = (StringArray)outputBatch.Column(0); + longArray = (Int64Array)outputBatch.Column(1); + } + else + { + Assert.Equal(1, outputBatch.ColumnCount); + var structArray = (StructArray)outputBatch.Column(0); + Assert.Equal(2, structArray.Fields.Count); + stringArray = (StringArray)structArray.Fields[0]; + longArray = (Int64Array)structArray.Fields[1]; + } - var stringArray = (StringArray)outputBatch.Column(0); for (int i = 0; i < numRows; ++i) { Assert.Equal($"udf: {i}", stringArray.GetString(i)); } - var longArray = (Int64Array)outputBatch.Column(1); for (int i = 0; i < numRows; ++i) { Assert.Equal(100 + i, longArray.Values[i]); @@ -981,15 +994,28 @@ await arrowWriter.WriteRecordBatchAsync( RecordBatch outputBatch = await arrowReader.ReadNextRecordBatchAsync(); Assert.Equal(numRows, outputBatch.Length); - Assert.Equal(2, outputBatch.ColumnCount); + StringArray stringArray; + DoubleArray doubleArray; + if (sparkVersion < new Version(Versions.V3_0_0)) + { + Assert.Equal(2, outputBatch.ColumnCount); + stringArray = (StringArray)outputBatch.Column(0); + doubleArray = (DoubleArray)outputBatch.Column(1); + } + else + { + Assert.Equal(1, outputBatch.ColumnCount); + var structArray = (StructArray)outputBatch.Column(0); + Assert.Equal(2, structArray.Fields.Count); + stringArray = (StringArray)structArray.Fields[0]; + doubleArray = (DoubleArray)structArray.Fields[1]; + } - var stringArray = (StringArray)outputBatch.Column(0); for (int i = 0; i < numRows; ++i) { Assert.Equal($"udf: {i}", stringArray.GetString(i)); } - var doubleArray = (DoubleArray)outputBatch.Column(1); for (int i = 0; i < numRows; ++i) { Assert.Equal(100 + i, doubleArray.Values[i]); diff --git a/src/csharp/Microsoft.Spark.Worker/Command/SqlCommandExecutor.cs b/src/csharp/Microsoft.Spark.Worker/Command/SqlCommandExecutor.cs index ed66cf021..54cc0577b 100644 --- a/src/csharp/Microsoft.Spark.Worker/Command/SqlCommandExecutor.cs +++ b/src/csharp/Microsoft.Spark.Worker/Command/SqlCommandExecutor.cs @@ -422,8 +422,7 @@ private CommandExecutorStat ExecuteArrowSqlCommand( var recordBatch = new RecordBatch(resultSchema, results, numEntries); - // TODO: Remove sync-over-async once WriteRecordBatch exists. - writer.WriteRecordBatchAsync(recordBatch).GetAwaiter().GetResult(); + writer.WriteRecordBatch(recordBatch); } WriteEnd(outputStream, ipcOptions); @@ -468,8 +467,7 @@ private CommandExecutorStat ExecuteDataFrameSqlCommand( new ArrowStreamWriter(outputStream, result.Schema, leaveOpen: true, ipcOptions); } - // TODO: Remove sync-over-async once WriteRecordBatch exists. - writer.WriteRecordBatchAsync(result).GetAwaiter().GetResult(); + writer.WriteRecordBatch(result); } } @@ -737,6 +735,29 @@ protected internal override CommandExecutorStat ExecuteCore( return ExecuteArrowGroupedMapCommand(inputStream, outputStream, commands); } + private RecordBatch WrapColumnsInStructIfApplicable(RecordBatch batch) + { + if (_version >= new Version(Versions.V3_0_0)) + { + var fields = new Field[batch.Schema.Fields.Count]; + for (int i = 0; i < batch.Schema.Fields.Count; ++i) + { + fields[i] = batch.Schema.GetFieldByIndex(i); + } + + var structType = new StructType(fields); + var structArray = new StructArray( + structType, + batch.Length, + batch.Arrays.Cast(), + ArrowBuffer.Empty); + 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,19 +775,19 @@ private CommandExecutorStat ExecuteArrowGroupedMapCommand( ArrowStreamWriter writer = null; foreach (RecordBatch input in GetInputIterator(inputStream)) { - RecordBatch result = worker.Func(input); + RecordBatch batch = worker.Func(input); - int numEntries = result.Length; + RecordBatch final = WrapColumnsInStructIfApplicable(batch); + int numEntries = final.Length; stat.NumEntriesProcessed += numEntries; 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.WriteRecordBatch(final); } WriteEnd(outputStream, ipcOptions); @@ -794,20 +815,21 @@ private CommandExecutorStat ExecuteDataFrameGroupedMapCommand( { FxDataFrame dataFrame = FxDataFrame.FromArrowRecordBatch(input); FxDataFrame resultDataFrame = worker.Func(dataFrame); + IEnumerable recordBatches = resultDataFrame.ToArrowRecordBatches(); - foreach (RecordBatch result in recordBatches) + foreach (RecordBatch batch in recordBatches) { - stat.NumEntriesProcessed += result.Length; + RecordBatch final = WrapColumnsInStructIfApplicable(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.WriteRecordBatch(final); } } diff --git a/src/csharp/Microsoft.Spark/Microsoft.Spark.csproj b/src/csharp/Microsoft.Spark/Microsoft.Spark.csproj index 30fdd0432..bc6a0c239 100644 --- a/src/csharp/Microsoft.Spark/Microsoft.Spark.csproj +++ b/src/csharp/Microsoft.Spark/Microsoft.Spark.csproj @@ -28,7 +28,7 @@ - + @@ -37,10 +37,7 @@ - +