Skip to content

Commit 974a74d

Browse files
suhsteveeerhardt
authored andcommitted
ARROW-5034: [C#] ArrowStreamWriter and ArrowFileWriter implement sync WriteRecordBatch
`ArrowStreamWriter` and `ArrowFileWriter` currently only supports an `WriteRecordBatchAsync`. This PR implements a synced version of it. All the logic has been reused and all `async` calls have been replaced with `sync` counterparts. Closes apache#8146 from suhsteve/writerecordbatch Authored-by: Steve Suh <suhsteve@gmail.com> Signed-off-by: Eric Erhardt <eric.erhardt@microsoft.com>
1 parent cd94749 commit 974a74d

8 files changed

Lines changed: 592 additions & 37 deletions

File tree

‎csharp/src/Apache.Arrow/Apache.Arrow.csproj‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
<Project Sdk="Microsoft.NET.Sdk">
1+
<Project Sdk="Microsoft.NET.Sdk">
22

33
<PropertyGroup>
44
<TargetFrameworks>netstandard1.3;netcoreapp2.1</TargetFrameworks>
@@ -37,5 +37,6 @@
3737
</ItemGroup>
3838
<ItemGroup Condition="'$(TargetFramework)' == 'netcoreapp2.1'">
3939
<Compile Remove="Extensions\StreamExtensions.netstandard.cs" />
40+
<Compile Remove="Extensions\TupleExtensions.netstandard.cs" />
4041
</ItemGroup>
4142
</Project>

‎csharp/src/Apache.Arrow/Extensions/StreamExtensions.netcoreapp2.1.cs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,5 +25,10 @@ public static int Read(this Stream stream, Memory<byte> buffer)
2525
{
2626
return stream.Read(buffer.Span);
2727
}
28+
29+
public static void Write(this Stream stream, ReadOnlyMemory<byte> buffer)
30+
{
31+
stream.Write(buffer.Span);
32+
}
2833
}
2934
}

‎csharp/src/Apache.Arrow/Extensions/StreamExtensions.netstandard.cs‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,27 @@ async ValueTask<int> FinishReadAsync(Task<int> readTask, byte[] localBuffer, Mem
7474
}
7575
}
7676

77+
public static void Write(this Stream stream, ReadOnlyMemory<byte> buffer)
78+
{
79+
if (MemoryMarshal.TryGetArray(buffer, out ArraySegment<byte> array))
80+
{
81+
stream.Write(array.Array, array.Offset, array.Count);
82+
}
83+
else
84+
{
85+
byte[] sharedBuffer = ArrayPool<byte>.Shared.Rent(buffer.Length);
86+
try
87+
{
88+
buffer.Span.CopyTo(sharedBuffer);
89+
stream.Write(sharedBuffer, 0, buffer.Length);
90+
}
91+
finally
92+
{
93+
ArrayPool<byte>.Shared.Return(sharedBuffer);
94+
}
95+
}
96+
}
97+
7798
public static ValueTask WriteAsync(this Stream stream, ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)
7899
{
79100
if (MemoryMarshal.TryGetArray(buffer, out ArraySegment<byte> array))
Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
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 System;
17+
18+
namespace Apache.Arrow
19+
{
20+
// Helpers to Deconstruct Tuples on netstandard
21+
internal static partial class TupleExtensions
22+
{
23+
public static void Deconstruct<T1, T2>(this Tuple<T1, T2> value, out T1 item1, out T2 item2)
24+
{
25+
item1 = value.Item1;
26+
item2 = value.Item2;
27+
}
28+
}
29+
}

‎csharp/src/Apache.Arrow/Ipc/ArrowFileWriter.cs‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,19 @@ public ArrowFileWriter(Stream stream, Schema schema, bool leaveOpen, IpcOptions
6161
RecordBatchBlocks = new List<Block>();
6262
}
6363

64+
public override void WriteRecordBatch(RecordBatch recordBatch)
65+
{
66+
// TODO: Compare record batch schema
67+
68+
if (!HasWrittenHeader)
69+
{
70+
WriteHeader();
71+
HasWrittenHeader = true;
72+
}
73+
74+
WriteRecordBatchInternal(recordBatch);
75+
}
76+
6477
public override async Task WriteRecordBatchAsync(RecordBatch recordBatch, CancellationToken cancellationToken = default)
6578
{
6679
// TODO: Compare record batch schema
@@ -104,13 +117,28 @@ private protected override void FinishedWritingRecordBatch(long bodyLength, long
104117
_currentRecordBatchOffset = -1;
105118
}
106119

120+
private protected override void WriteEndInternal()
121+
{
122+
base.WriteEndInternal();
123+
124+
WriteFooter(Schema);
125+
}
126+
107127
private protected override async ValueTask WriteEndInternalAsync(CancellationToken cancellationToken)
108128
{
109129
await base.WriteEndInternalAsync(cancellationToken);
110130

111131
await WriteFooterAsync(Schema, cancellationToken);
112132
}
113133

134+
private void WriteHeader()
135+
{
136+
// Write magic number and empty padding up to the 8-byte boundary
137+
138+
WriteMagic();
139+
WritePadding(CalculatePadding(ArrowFileConstants.Magic.Length));
140+
}
141+
114142
private async Task WriteHeaderAsync(CancellationToken cancellationToken)
115143
{
116144
// Write magic number and empty padding up to the 8-byte boundary
@@ -120,6 +148,64 @@ await WritePaddingAsync(CalculatePadding(ArrowFileConstants.Magic.Length))
120148
.ConfigureAwait(false);
121149
}
122150

151+
private void WriteFooter(Schema schema)
152+
{
153+
Builder.Clear();
154+
155+
long offset = BaseStream.Position;
156+
157+
// Serialize the schema
158+
159+
FlatBuffers.Offset<Flatbuf.Schema> schemaOffset = SerializeSchema(schema);
160+
161+
// Serialize all record batches
162+
163+
Flatbuf.Footer.StartRecordBatchesVector(Builder, RecordBatchBlocks.Count);
164+
165+
foreach (Block recordBatch in RecordBatchBlocks)
166+
{
167+
Flatbuf.Block.CreateBlock(
168+
Builder, recordBatch.Offset, recordBatch.MetadataLength, recordBatch.BodyLength);
169+
}
170+
171+
FlatBuffers.VectorOffset recordBatchesVectorOffset = Builder.EndVector();
172+
173+
// Serialize all dictionaries
174+
// NOTE: Currently unsupported.
175+
176+
Flatbuf.Footer.StartDictionariesVector(Builder, 0);
177+
178+
FlatBuffers.VectorOffset dictionaryBatchesOffset = Builder.EndVector();
179+
180+
// Serialize and write the footer flatbuffer
181+
182+
FlatBuffers.Offset<Flatbuf.Footer> footerOffset = Flatbuf.Footer.CreateFooter(Builder, CurrentMetadataVersion,
183+
schemaOffset, dictionaryBatchesOffset, recordBatchesVectorOffset);
184+
185+
Builder.Finish(footerOffset.Value);
186+
187+
WriteFlatBuffer();
188+
189+
// Write footer length
190+
191+
Buffers.RentReturn(4, (buffer) =>
192+
{
193+
int footerLength;
194+
checked
195+
{
196+
footerLength = (int)(BaseStream.Position - offset);
197+
}
198+
199+
BinaryPrimitives.WriteInt32LittleEndian(buffer.Span, footerLength);
200+
201+
BaseStream.Write(buffer);
202+
});
203+
204+
// Write magic
205+
206+
WriteMagic();
207+
}
208+
123209
private async Task WriteFooterAsync(Schema schema, CancellationToken cancellationToken)
124210
{
125211
Builder.Clear();
@@ -182,6 +268,11 @@ await Buffers.RentReturnAsync(4, async (buffer) =>
182268
await WriteMagicAsync(cancellationToken).ConfigureAwait(false);
183269
}
184270

271+
private void WriteMagic()
272+
{
273+
BaseStream.Write(ArrowFileConstants.Magic);
274+
}
275+
185276
private ValueTask WriteMagicAsync(CancellationToken cancellationToken)
186277
{
187278
return BaseStream.WriteAsync(ArrowFileConstants.Magic, cancellationToken);

0 commit comments

Comments
 (0)