ceph/src/arrow/csharp/test/Apache.Arrow.Tests/ArrowFileWriterTests.cs

   1 // Licensed to the Apache Software Foundation (ASF) under one or more
   2 // contributor license agreements. See the NOTICE file distributed with
   3 // this work for additional information regarding copyright ownership.
   4 // The ASF licenses this file to You under the Apache License, Version 2.0
   5 // (the "License"); you may not use this file except in compliance with
   6 // the License.  You may obtain a copy of the License at
   7 //
   8 //     http://www.apache.org/licenses/LICENSE-2.0
   9 //
  10 // Unless required by applicable law or agreed to in writing, software
  11 // distributed under the License is distributed on an "AS IS" BASIS,
  12 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 // See the License for the specific language governing permissions and
  14 // limitations under the License.
  15
  16 using Apache.Arrow.Ipc;
  17 using System;
  18 using System.IO;
  19 using System.Threading.Tasks;
  20 using Xunit;
  21
  22 namespace Apache.Arrow.Tests
  23 {
  24     public class ArrowFileWriterTests
  25     {
  26         [Fact]
  27         public void Ctor_LeaveOpenDefault_StreamClosedOnDispose()
  28         {
  29             RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100);
  30             var stream = new MemoryStream();
  31             new ArrowFileWriter(stream, originalBatch.Schema).Dispose();
  32             Assert.Throws<ObjectDisposedException>(() => stream.Position);
  33         }
  34
  35         [Fact]
  36         public void Ctor_LeaveOpenFalse_StreamClosedOnDispose()
  37         {
  38             RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100);
  39             var stream = new MemoryStream();
  40             new ArrowFileWriter(stream, originalBatch.Schema, leaveOpen: false).Dispose();
  41             Assert.Throws<ObjectDisposedException>(() => stream.Position);
  42         }
  43
  44         [Fact]
  45         public void Ctor_LeaveOpenTrue_StreamValidOnDispose()
  46         {
  47             RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100);
  48             var stream = new MemoryStream();
  49             new ArrowFileWriter(stream, originalBatch.Schema, leaveOpen: true).Dispose();
  50             Assert.Equal(0, stream.Position);
  51         }
  52
  53         /// <summary>
  54         /// Tests that writing an arrow file will always align the Block lengths
  55         /// to 8 bytes. There are asserts in both the reader and writer which will fail
  56         /// if this isn't the case.
  57         /// </summary>
  58         /// <returns></returns>
  59         [Fact]
  60         public async Task WritesFooterAlignedMulitpleOf8()
  61         {
  62             RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100);
  63
  64             var stream = new MemoryStream();
  65             var writer = new ArrowFileWriter(
  66                 stream,
  67                 originalBatch.Schema,
  68                 leaveOpen: true,
  69                 // use WriteLegacyIpcFormat, which only uses a 4-byte length prefix
  70                 // which causes the length prefix to not be 8-byte aligned by default
  71                 new IpcOptions() { WriteLegacyIpcFormat = true });
  72
  73             writer.WriteRecordBatch(originalBatch);
  74             writer.WriteEnd();
  75
  76             stream.Position = 0;
  77
  78             await ValidateRecordBatchFile(stream, originalBatch);
  79         }
  80
  81         /// <summary>
  82         /// Tests that writing an arrow file will always align the Block lengths
  83         /// to 8 bytes. There are asserts in both the reader and writer which will fail
  84         /// if this isn't the case.
  85         /// </summary>
  86         /// <returns></returns>
  87         [Fact]
  88         public async Task WritesFooterAlignedMulitpleOf8Async()
  89         {
  90             RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 100);
  91
  92             var stream = new MemoryStream();
  93             var writer = new ArrowFileWriter(
  94                 stream,
  95                 originalBatch.Schema,
  96                 leaveOpen: true,
  97                 // use WriteLegacyIpcFormat, which only uses a 4-byte length prefix
  98                 // which causes the length prefix to not be 8-byte aligned by default
  99                 new IpcOptions() { WriteLegacyIpcFormat = true });
 100
 101             await writer.WriteRecordBatchAsync(originalBatch);
 102             await writer.WriteEndAsync();
 103
 104             stream.Position = 0;
 105
 106             await ValidateRecordBatchFile(stream, originalBatch);
 107         }
 108
 109         private async Task ValidateRecordBatchFile(Stream stream, RecordBatch recordBatch)
 110         {
 111             var reader = new ArrowFileReader(stream);
 112             int count = await reader.RecordBatchCountAsync();
 113             Assert.Equal(1, count);
 114             RecordBatch readBatch = await reader.ReadRecordBatchAsync(0);
 115             ArrowReaderVerifier.CompareBatches(recordBatch, readBatch);
 116         }
 117
 118         /// <summary>
 119         /// Tests that writing an arrow file with no RecordBatches produces the correct
 120         /// file.
 121         /// </summary>
 122         [Fact]
 123         public async Task WritesEmptyFile()
 124         {
 125             RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 1);
 126
 127             var stream = new MemoryStream();
 128             var writer = new ArrowFileWriter(stream, originalBatch.Schema);
 129
 130             writer.WriteStart();
 131             writer.WriteEnd();
 132
 133             stream.Position = 0;
 134
 135             var reader = new ArrowFileReader(stream);
 136             int count = await reader.RecordBatchCountAsync();
 137             Assert.Equal(0, count);
 138             RecordBatch readBatch = reader.ReadNextRecordBatch();
 139             Assert.Null(readBatch);
 140             SchemaComparer.Compare(originalBatch.Schema, reader.Schema);
 141         }
 142
 143         /// <summary>
 144         /// Tests that writing an arrow file with no RecordBatches produces the correct
 145         /// file when using WriteStartAsync and WriteEndAsync.
 146         /// </summary>
 147         [Fact]
 148         public async Task WritesEmptyFileAsync()
 149         {
 150             RecordBatch originalBatch = TestData.CreateSampleRecordBatch(length: 1);
 151
 152             var stream = new MemoryStream();
 153             var writer = new ArrowFileWriter(stream, originalBatch.Schema);
 154
 155             await writer.WriteStartAsync();
 156             await writer.WriteEndAsync();
 157
 158             stream.Position = 0;
 159
 160             var reader = new ArrowFileReader(stream);
 161             int count = await reader.RecordBatchCountAsync();
 162             Assert.Equal(0, count);
 163             RecordBatch readBatch = reader.ReadNextRecordBatch();
 164             Assert.Null(readBatch);
 165             SchemaComparer.Compare(originalBatch.Schema, reader.Schema);
 166         }
 167     }
 168 }