]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
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 | } |