diff --git a/benchmark/csharp/Tpch/TpchFunctionalQueries.cs b/benchmark/csharp/Tpch/TpchFunctionalQueries.cs index 1c228f22e..cb53fdabd 100644 --- a/benchmark/csharp/Tpch/TpchFunctionalQueries.cs +++ b/benchmark/csharp/Tpch/TpchFunctionalQueries.cs @@ -10,6 +10,7 @@ using Microsoft.Spark.Sql; using static Microsoft.Spark.Sql.ExperimentalFunctions; using static Microsoft.Spark.Sql.Functions; +using Column = Microsoft.Spark.Sql.Column; namespace Tpch { diff --git a/benchmark/csharp/Tpch/VectorFunctions.cs b/benchmark/csharp/Tpch/VectorFunctions.cs index 0338ad1e2..590793c10 100644 --- a/benchmark/csharp/Tpch/VectorFunctions.cs +++ b/benchmark/csharp/Tpch/VectorFunctions.cs @@ -17,7 +17,7 @@ internal static DoubleArray ComputeTotal(DoubleArray price, DoubleArray discount } int length = price.Length; - var builder = new ArrowBuffer.Builder(length); + var builder = new DoubleArray.Builder().Reserve(length); ReadOnlySpan prices = price.Values; ReadOnlySpan discounts = discount.Values; ReadOnlySpan taxes = tax.Values; @@ -26,12 +26,7 @@ internal static DoubleArray ComputeTotal(DoubleArray price, DoubleArray discount builder.Append(prices[i] * (1 - discounts[i]) * (1 + taxes[i])); } - return new DoubleArray( - builder.Build(), - nullBitmapBuffer: ArrowBuffer.Empty, - length: length, - nullCount: 0, - offset: 0); + return builder.Build(); } internal static DoubleArray ComputeDiscountPrice(DoubleArray price, DoubleArray discount) @@ -42,7 +37,7 @@ internal static DoubleArray ComputeDiscountPrice(DoubleArray price, DoubleArray } int length = price.Length; - var builder = new ArrowBuffer.Builder(length); + var builder = new DoubleArray.Builder().Reserve(length); ReadOnlySpan prices = price.Values; ReadOnlySpan discounts = discount.Values; for (int i = 0; i < length; ++i) @@ -50,12 +45,7 @@ internal static DoubleArray ComputeDiscountPrice(DoubleArray price, DoubleArray builder.Append(prices[i] * (1 - discounts[i])); } - return new DoubleArray( - builder.Build(), - nullBitmapBuffer: ArrowBuffer.Empty, - length: length, - nullCount: 0, - offset: 0); + return builder.Build(); } } } diff --git a/examples/Microsoft.Spark.CSharp.Examples/Sql/VectorUdfs.cs b/examples/Microsoft.Spark.CSharp.Examples/Sql/VectorUdfs.cs index 10800752c..2bcac95ed 100644 --- a/examples/Microsoft.Spark.CSharp.Examples/Sql/VectorUdfs.cs +++ b/examples/Microsoft.Spark.CSharp.Examples/Sql/VectorUdfs.cs @@ -80,18 +80,13 @@ private static RecordBatch CountCharacters( return new RecordBatch( new Schema.Builder() - .Field(f => f.Name(groupField.Name).DataType(groupField.DataType)) + .Field(groupField) .Field(f => f.Name(stringFieldName + "_CharCount").DataType(Int32Type.Default)) .Build(), new IArrowArray[] { records.Column(groupFieldIndex), - new Int32Array( - new ArrowBuffer.Builder().Append(characterCount).Build(), - ArrowBuffer.Empty, - length: 1, - nullCount: 0, - offset: 0) + new Int32Array.Builder().Append(characterCount).Build() }, returnLength); } diff --git a/src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/DataFrameTests.cs b/src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/DataFrameTests.cs index cb4ab5225..94f1d1790 100644 --- a/src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/DataFrameTests.cs +++ b/src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/DataFrameTests.cs @@ -11,6 +11,7 @@ using Xunit; using static Microsoft.Spark.Sql.Functions; using static Microsoft.Spark.UnitTest.TestUtils.ArrowTestUtils; +using Column = Microsoft.Spark.Sql.Column; using Int32Type = Apache.Arrow.Types.Int32Type; namespace Microsoft.Spark.E2ETest.IpcTests @@ -190,16 +191,12 @@ public void TestGroupedMapUdf() Assert.Equal(3, rows.Length); foreach (Row row in rows) { - int age = row.GetAs("age"); - int? charCount = row.GetAs("nameCharCount"); + int? age = row.GetAs("age"); + int charCount = row.GetAs("nameCharCount"); switch (age) { - case 0: - // The results here are incorrect for the {name: "Michael" age: null} - // record because of https://issues.apache.org/jira/browse/ARROW-5887. - // When an updated Apache.Arrow library is available with the fix, - // this should change to check for age: null, charCount: 7. - Assert.Null(charCount); + case null: + Assert.Equal(7, charCount); break; case 19: Assert.Equal(11, charCount); @@ -234,18 +231,13 @@ private static RecordBatch CountCharacters(RecordBatch records) return new RecordBatch( new Schema.Builder() - .Field(f => f.Name(groupField.Name).DataType(groupField.DataType)) + .Field(groupField) .Field(f => f.Name("name_CharCount").DataType(Int32Type.Default)) .Build(), new IArrowArray[] { records.Column(groupFieldIndex), - new Int32Array( - new ArrowBuffer.Builder().Append(characterCount).Build(), - ArrowBuffer.Empty, - length: 1, - nullCount: 0, - offset: 0) + new Int32Array.Builder().Append(characterCount).Build() }, returnLength); } diff --git a/src/csharp/Microsoft.Spark.UnitTest/WorkerFunctionTests.cs b/src/csharp/Microsoft.Spark.UnitTest/WorkerFunctionTests.cs index f717d7201..7a98ea8b7 100644 --- a/src/csharp/Microsoft.Spark.UnitTest/WorkerFunctionTests.cs +++ b/src/csharp/Microsoft.Spark.UnitTest/WorkerFunctionTests.cs @@ -93,7 +93,7 @@ public void TestArrowWorkerFunctionForBool() new ArrowUdfWrapper( (strings, flags) => (BooleanArray)ToArrowArray( Enumerable.Range(0, strings.Length) - .Select(i => flags.GetBoolean(i).Value || strings.GetString(i).Contains("true")) + .Select(i => flags.GetBoolean(i) || strings.GetString(i).Contains("true")) .ToArray())).Execute); IArrowArray[] input = new[] diff --git a/src/csharp/Microsoft.Spark.Worker/Command/SqlCommandExecutor.cs b/src/csharp/Microsoft.Spark.Worker/Command/SqlCommandExecutor.cs index e86f74c12..82fcee8fe 100644 --- a/src/csharp/Microsoft.Spark.Worker/Command/SqlCommandExecutor.cs +++ b/src/csharp/Microsoft.Spark.Worker/Command/SqlCommandExecutor.cs @@ -291,9 +291,6 @@ public object Run(int splitId, object input) /// internal class ArrowSqlCommandExecutor : SqlCommandExecutor { - [ThreadStatic] - private static MemoryStream s_writeOutputStream; - protected override CommandExecutorStat ExecuteCore( Stream inputStream, Stream outputStream, @@ -304,11 +301,6 @@ protected override CommandExecutorStat ExecuteCore( SerDe.Write(outputStream, (int)SpecialLengths.START_ARROW_STREAM); - // TODO: Remove this MemoryStream once the arrow writer supports non-seekable streams. - // For now, we write to a temporary seekable MemoryStream which we then copy to - // the actual destination stream. - MemoryStream tmp = s_writeOutputStream ?? (s_writeOutputStream = new MemoryStream()); - ArrowStreamWriter writer = null; Schema resultSchema = null; foreach (ReadOnlyMemory input in GetInputIterator(inputStream)) @@ -319,24 +311,18 @@ protected override CommandExecutorStat ExecuteCore( int numEntries = results[0].Length; stat.NumEntriesProcessed += numEntries; - tmp.SetLength(0); - if (writer == null) { Debug.Assert(resultSchema == null); resultSchema = BuildSchema(results); - writer = new ArrowStreamWriter(tmp, resultSchema, leaveOpen: true); + writer = new ArrowStreamWriter(outputStream, resultSchema, leaveOpen: true); } var recordBatch = new RecordBatch(resultSchema, results, numEntries); // TODO: Remove sync-over-async once WriteRecordBatch exists. writer.WriteRecordBatchAsync(recordBatch).GetAwaiter().GetResult(); - - tmp.Position = 0; - tmp.CopyTo(outputStream); - outputStream.Flush(); } SerDe.Write(outputStream, 0); @@ -532,9 +518,6 @@ public IArrowArray[] Run(ReadOnlyMemory input) internal class ArrowGroupedMapCommandExecutor : SqlCommandExecutor { - [ThreadStatic] - private static MemoryStream s_writeOutputStream; - protected override CommandExecutorStat ExecuteCore( Stream inputStream, Stream outputStream, @@ -548,11 +531,6 @@ protected override CommandExecutorStat ExecuteCore( SerDe.Write(outputStream, (int)SpecialLengths.START_ARROW_STREAM); - // TODO: Remove this MemoryStream once the arrow writer supports non-seekable streams. - // For now, we write to a temporary seekable MemoryStream which we then copy to - // the actual destination stream. - MemoryStream tmp = s_writeOutputStream ?? (s_writeOutputStream = new MemoryStream()); - ArrowStreamWriter writer = null; foreach (RecordBatch input in GetInputIterator(inputStream)) { @@ -561,19 +539,13 @@ protected override CommandExecutorStat ExecuteCore( int numEntries = result.Length; stat.NumEntriesProcessed += numEntries; - tmp.SetLength(0); - if (writer == null) { - writer = new ArrowStreamWriter(tmp, result.Schema, leaveOpen: true); + writer = new ArrowStreamWriter(outputStream, result.Schema, leaveOpen: true); } // TODO: Remove sync-over-async once WriteRecordBatch exists. writer.WriteRecordBatchAsync(result).GetAwaiter().GetResult(); - - tmp.Position = 0; - tmp.CopyTo(outputStream); - outputStream.Flush(); } SerDe.Write(outputStream, 0); diff --git a/src/csharp/Microsoft.Spark/Microsoft.Spark.csproj b/src/csharp/Microsoft.Spark/Microsoft.Spark.csproj index d16ac4329..beba60fe4 100644 --- a/src/csharp/Microsoft.Spark/Microsoft.Spark.csproj +++ b/src/csharp/Microsoft.Spark/Microsoft.Spark.csproj @@ -22,7 +22,7 @@ - +