Add StreamWriter tests

apache · Apr 15, 2024 · bed8aa8 · bed8aa8
1 parent 500101c
commit bed8aa8
Showing 1 changed file with 38 additions and 8 deletions.
diff --git a/csharp/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs b/csharp/test/Apache.Arrow.Tests/ArrowStreamWriterTests.cs
@@ -203,7 +203,37 @@ public async Task WriteBatchWithNullsAsync()
             await TestRoundTripRecordBatchAsync(originalBatch);
         }
 
-        private static void TestRoundTripRecordBatches(List<RecordBatch> originalBatches, IpcOptions options = null)
+        [Theory]
+        [InlineData(0, 45)]
+        [InlineData(3, 45)]
+        [InlineData(16, 45)]
+        public void WriteSlicedArrays(int sliceOffset, int sliceLength)
+        {
+            var originalBatch = TestData.CreateSampleRecordBatch(length: 100);
+            var slicedArrays = originalBatch.Arrays
+                .Select(array => ArrowArrayFactory.Slice(array, sliceOffset, sliceLength))
+                .ToList();
+            var slicedBatch = new RecordBatch(originalBatch.Schema, slicedArrays, sliceLength);
+
+            TestRoundTripRecordBatch(slicedBatch, strictCompare: false);
+        }
+
+        [Theory]
+        [InlineData(0, 45)]
+        [InlineData(3, 45)]
+        [InlineData(16, 45)]
+        public async Task WriteSlicedArraysAsync(int sliceOffset, int sliceLength)
+        {
+            var originalBatch = TestData.CreateSampleRecordBatch(length: 100);
+            var slicedArrays = originalBatch.Arrays
+                .Select(array => ArrowArrayFactory.Slice(array, sliceOffset, sliceLength))
+                .ToList();
+            var slicedBatch = new RecordBatch(originalBatch.Schema, slicedArrays, sliceLength);
+
+            await TestRoundTripRecordBatchAsync(slicedBatch, strictCompare: false);
+        }
+
+        private static void TestRoundTripRecordBatches(List<RecordBatch> originalBatches, IpcOptions options = null, bool strictCompare = true)
         {
             using (MemoryStream stream = new MemoryStream())
             {
@@ -223,13 +253,13 @@ private static void TestRoundTripRecordBatches(List<RecordBatch> originalBatches
                     foreach (RecordBatch originalBatch in originalBatches)
                     {
                         RecordBatch newBatch = reader.ReadNextRecordBatch();
-                        ArrowReaderVerifier.CompareBatches(originalBatch, newBatch);
+                        ArrowReaderVerifier.CompareBatches(originalBatch, newBatch, strictCompare: strictCompare);
                     }
                 }
             }
         }
 
-        private static async Task TestRoundTripRecordBatchesAsync(List<RecordBatch> originalBatches, IpcOptions options = null)
+        private static async Task TestRoundTripRecordBatchesAsync(List<RecordBatch> originalBatches, IpcOptions options = null, bool strictCompare = true)
         {
             using (MemoryStream stream = new MemoryStream())
             {
@@ -249,20 +279,20 @@ private static async Task TestRoundTripRecordBatchesAsync(List<RecordBatch> orig
                     foreach (RecordBatch originalBatch in originalBatches)
                     {
                         RecordBatch newBatch = reader.ReadNextRecordBatch();
-                        ArrowReaderVerifier.CompareBatches(originalBatch, newBatch);
+                        ArrowReaderVerifier.CompareBatches(originalBatch, newBatch, strictCompare: strictCompare);
                     }
                 }
             }
         }
 
-        private static void TestRoundTripRecordBatch(RecordBatch originalBatch, IpcOptions options = null)
+        private static void TestRoundTripRecordBatch(RecordBatch originalBatch, IpcOptions options = null, bool strictCompare = true)
         {
-            TestRoundTripRecordBatches(new List<RecordBatch> { originalBatch }, options);
+            TestRoundTripRecordBatches(new List<RecordBatch> { originalBatch }, options, strictCompare: strictCompare);
         }
 
-        private static async Task TestRoundTripRecordBatchAsync(RecordBatch originalBatch, IpcOptions options = null)
+        private static async Task TestRoundTripRecordBatchAsync(RecordBatch originalBatch, IpcOptions options = null, bool strictCompare = true)
         {
-            await TestRoundTripRecordBatchesAsync(new List<RecordBatch> { originalBatch }, options);
+            await TestRoundTripRecordBatchesAsync(new List<RecordBatch> { originalBatch }, options, strictCompare: strictCompare);
         }
 
         [Fact]