diff --git a/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs b/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs index 46a7ac20..ecec016d 100644 --- a/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs +++ b/src/Apache.Arrow.Flight/Internal/FlightDataStream.cs @@ -36,6 +36,8 @@ internal class FlightDataStream : ArrowStreamWriter private readonly FlightDescriptor _flightDescriptor; private readonly IAsyncStreamWriter _clientStreamWriter; private Protocol.FlightData _currentFlightData; + private bool _hasPendingMessage; + private ByteString _recordBatchAppMetadata; public FlightDataStream(IAsyncStreamWriter clientStreamWriter, FlightDescriptor flightDescriptor, Schema schema) : base(new MemoryStream(), schema) @@ -46,17 +48,9 @@ public FlightDataStream(IAsyncStreamWriter clientStreamWrit public async Task SendSchema() { - _currentFlightData = new Protocol.FlightData(); - - if (_flightDescriptor != null) - { - _currentFlightData.FlightDescriptor = _flightDescriptor.ToProtocol(); - } - var offset = SerializeSchema(Schema); - CancellationTokenSource cancellationTokenSource = new CancellationTokenSource(); - await WriteMessageAsync(MessageHeader.Schema, offset, 0, default, cancellationTokenSource.Token).ConfigureAwait(false); - await _clientStreamWriter.WriteAsync(_currentFlightData).ConfigureAwait(false); + await WriteMessageAsync(MessageHeader.Schema, offset, 0, default, CancellationToken.None).ConfigureAwait(false); + await FlushCurrentMessageAsync().ConfigureAwait(false); HasWrittenSchema = true; } @@ -74,25 +68,31 @@ public async Task Write(RecordBatch recordBatch, ByteString applicationMetadata) } ResetStream(); - _currentFlightData = new Protocol.FlightData(); + // Attached to the record-batch message, not to any preceding dictionary messages. + _recordBatchAppMetadata = applicationMetadata; - if (applicationMetadata != null) - { - _currentFlightData.AppMetadata = applicationMetadata; - } + // Resend the full dictionary before every record batch rather than just the first (#180). + HasWrittenDictionaryBatch = false; + // Writes any dictionary-batch messages followed by the record-batch message. Each is + // flushed as its own FlightData frame (see WriteMessageAsync) so that dictionary batches + // are delivered before the record batch that references them. await WriteRecordBatchInternalAsync(recordBatch, customMetadata: null).ConfigureAwait(false); - //Reset stream position - this.BaseStream.Position = 0; - var bodyData = await ByteString.FromStreamAsync(this.BaseStream).ConfigureAwait(false); - - _currentFlightData.DataBody = bodyData; - await _clientStreamWriter.WriteAsync(_currentFlightData).ConfigureAwait(false); + // Flush the final (record-batch) message. + await FlushCurrentMessageAsync().ConfigureAwait(false); + _recordBatchAppMetadata = null; } - private protected override ValueTask WriteMessageAsync(MessageHeader headerType, Offset headerOffset, int bodyLength, VectorOffset customMetadataOffset, CancellationToken cancellationToken) + private protected override async ValueTask WriteMessageAsync(MessageHeader headerType, Offset headerOffset, int bodyLength, VectorOffset customMetadataOffset, CancellationToken cancellationToken) { + // A new message is beginning; the previous message's body is now fully buffered, so flush + // it as its own FlightData frame before starting the next one. + if (_hasPendingMessage) + { + await FlushCurrentMessageAsync().ConfigureAwait(false); + } + Offset messageOffset = Flatbuf.Message.CreateMessage( Builder, CurrentMetadataVersion, headerType, headerOffset.Value, bodyLength, customMetadataOffset); @@ -101,9 +101,32 @@ private protected override ValueTask WriteMessageAsync(MessageHeader he ReadOnlyMemory messageData = Builder.DataBuffer.ToReadOnlyMemory(Builder.DataBuffer.Position, Builder.Offset); - _currentFlightData.DataHeader = ByteString.CopyFrom(messageData.Span); + _currentFlightData = new Protocol.FlightData + { + DataHeader = ByteString.CopyFrom(messageData.Span) + }; + + if (headerType == MessageHeader.Schema && _flightDescriptor != null) + { + _currentFlightData.FlightDescriptor = _flightDescriptor.ToProtocol(); + } + + if (headerType == MessageHeader.RecordBatch && _recordBatchAppMetadata != null) + { + _currentFlightData.AppMetadata = _recordBatchAppMetadata; + } + + _hasPendingMessage = true; + return 0; + } - return new ValueTask(0); + private async Task FlushCurrentMessageAsync() + { + this.BaseStream.Position = 0; + _currentFlightData.DataBody = await ByteString.FromStreamAsync(this.BaseStream).ConfigureAwait(false); + await _clientStreamWriter.WriteAsync(_currentFlightData).ConfigureAwait(false); + ResetStream(); + _hasPendingMessage = false; } } } diff --git a/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs b/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs index bb4aba7e..d5be69bd 100644 --- a/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs +++ b/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs @@ -103,7 +103,9 @@ public override async ValueTask ReadSchemaAsync(CancellationToken cancel switch (message.HeaderType) { case MessageHeader.Schema: - _schema = FlightMessageSerializer.DecodeSchema(message.ByteBuffer); + // Use the base reader's DictionaryMemo so dictionary-encoded + // fields are registered (FlightMessageSerializer discards them). + _schema = ReadSchemaFromMessage(message.ByteBuffer); break; default: throw new Exception($"Expected schema as the first message, but got: {message.HeaderType.ToString()}"); @@ -120,8 +122,9 @@ public override async ValueTask ReadNextRecordBatchAsync(Cancellati { await ReadSchemaAsync(cancellationToken).ConfigureAwait(false); } - var moveNextResult = await _flightDataStream.MoveNext().ConfigureAwait(false); - if (moveNextResult) + // Dictionary batches precede the record batch that references them; keep + // reading until CreateArrowObjectFromMessage yields a record batch. + while (await _flightDataStream.MoveNext(cancellationToken).ConfigureAwait(false)) { //AppMetadata will never be null, but length 0 if empty //Those are skipped @@ -131,21 +134,27 @@ public override async ValueTask ReadNextRecordBatchAsync(Cancellati } var header = _flightDataStream.Current.DataHeader.Memory; + if (header.IsEmpty) + { + continue; + } Message message = Message.GetRootAsMessage(CreateByteBuffer(header)); - switch (message.HeaderType) + if (message.BodyLength < 0 || message.BodyLength > int.MaxValue) { - case MessageHeader.RecordBatch: - if (message.BodyLength < 0 || message.BodyLength > int.MaxValue) - { - throw new InvalidDataException( - $"Cannot read batch. Message body of {message.BodyLength} bytes is out of range"); - } - - var body = _flightDataStream.Current.DataBody.Memory; - return CreateArrowObjectFromMessage(message, CreateByteBuffer(body.Slice(0, checked((int)message.BodyLength))), null); - default: - throw new NotImplementedException(); + throw new InvalidDataException( + $"Cannot read batch. Message body of {message.BodyLength} bytes is out of range"); + } + + var body = _flightDataStream.Current.DataBody.Memory; + var arrowObject = CreateArrowObjectFromMessage( + message, + CreateByteBuffer(body.Slice(0, checked((int)message.BodyLength))), + null); + + if (arrowObject != null) + { + return arrowObject; } } return null; diff --git a/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs b/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs index 45fd7920..c13be1ad 100644 --- a/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs +++ b/src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs @@ -178,6 +178,15 @@ private static IReadOnlyDictionary ReadMessageCustomMetadata(Fla return metadata; } + /// + /// Decode a schema message, registering any dictionary-encoded fields in the + /// reader's DictionaryMemo so subsequent dictionary batches can be resolved. + /// + protected Schema ReadSchemaFromMessage(ByteBuffer schemaBuffer) + { + return MessageSerializer.GetSchema(ReadMessage(schemaBuffer), ref _dictionaryMemo, _extensionRegistry); + } + internal static ByteBuffer CreateByteBuffer(ReadOnlyMemory buffer) { return new ByteBuffer(new ReadOnlyMemoryBufferAllocator(buffer), 0); diff --git a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs index aa63985e..e19ed8f3 100644 --- a/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs +++ b/src/Apache.Arrow/Ipc/ArrowStreamWriter.cs @@ -728,7 +728,7 @@ public void Visit(IArrowArray array) protected bool HasWrittenSchema { get; set; } - private bool HasWrittenDictionaryBatch { get; set; } + protected bool HasWrittenDictionaryBatch { get; set; } private bool HasWrittenStart { get; set; } diff --git a/test/Apache.Arrow.Flight.Tests/FlightTests.cs b/test/Apache.Arrow.Flight.Tests/FlightTests.cs index 46662cce..867b64b0 100644 --- a/test/Apache.Arrow.Flight.Tests/FlightTests.cs +++ b/test/Apache.Arrow.Flight.Tests/FlightTests.cs @@ -21,6 +21,7 @@ using Apache.Arrow.Flight.Client; using Apache.Arrow.Flight.TestWeb; using Apache.Arrow.Tests; +using Apache.Arrow.Types; using Google.Protobuf; using Grpc.Core; using Grpc.Core.Utils; @@ -58,6 +59,100 @@ private RecordBatch CreateTestBatch(int startValue, int length) return batchBuilder.Build(); } + private static Schema CreateDictionaryTestSchema() + { + return new Schema.Builder() + .Field(f => f.Name("value").DataType(Int32Type.Default).Nullable(true)) + .Field(f => f.Name("symbol").DataType(new DictionaryType(Int32Type.Default, StringType.Default, false)).Nullable(true)) + .Build(); + } + + // Batches in the same stream must share the same Schema/Field instances. + private RecordBatch CreateTestBatchWithDictionaryColumn(Schema schema, int startValue, int length, string[] dictionaryValues) + { + Int32Array.Builder valueBuilder = new Int32Array.Builder(); + for (int i = 0; i < length; i++) + { + valueBuilder.Append(startValue + i); + } + + // Dictionary-encoded (categorical) string column: dictionary. + StringArray dictionary = new StringArray.Builder().AppendRange(dictionaryValues).Build(); + Int32Array.Builder indicesBuilder = new Int32Array.Builder(); + for (int i = 0; i < length; i++) + { + indicesBuilder.Append(i % dictionaryValues.Length); + } + var dictionaryType = (DictionaryType)schema.GetFieldByIndex(1).DataType; + var dictionaryArray = new DictionaryArray(dictionaryType, indicesBuilder.Build(), dictionary); + + return new RecordBatch(schema, new IArrowArray[] { valueBuilder.Build(), dictionaryArray }, length); + } + + private RecordBatch CreateTestBatchWithDictionaryColumn(int startValue, int length) + { + return CreateTestBatchWithDictionaryColumn(CreateDictionaryTestSchema(), startValue, length, new[] { "a", "b", "c" }); + } + + [Fact] + public async Task TestGetRecordBatchWithDictionaryColumn() + { + var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test"); + var expectedBatch = CreateTestBatchWithDictionaryColumn(0, 100); + GivenStoreBatches(flightDescriptor, new RecordBatchWithMetadata(expectedBatch)); + + var flightInfo = await _flightClient.GetInfo(flightDescriptor); + var endpoint = flightInfo.Endpoints.First(); + var getStream = _flightClient.GetStream(endpoint.Ticket); + var batches = await getStream.ResponseStream.ToListAsync(); + + Assert.Single(batches); + ArrowReaderVerifier.CompareBatches(expectedBatch, batches[0]); + } + + [Fact] + public async Task TestGetRecordBatchesWithReplacementDictionary() + { + // Batches carry different dictionary vocabularies, exercising the per-batch resend (#180). + var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test"); + var schema = CreateDictionaryTestSchema(); + var expectedBatch1 = CreateTestBatchWithDictionaryColumn(schema, 0, 50, new[] { "a", "b", "c" }); + var expectedBatch2 = CreateTestBatchWithDictionaryColumn(schema, 50, 50, new[] { "w", "x", "y", "z" }); + GivenStoreBatches(flightDescriptor, new RecordBatchWithMetadata(expectedBatch1), new RecordBatchWithMetadata(expectedBatch2)); + + var flightInfo = await _flightClient.GetInfo(flightDescriptor); + var endpoint = flightInfo.Endpoints.First(); + var getStream = _flightClient.GetStream(endpoint.Ticket); + var batches = await getStream.ResponseStream.ToListAsync(); + + Assert.Equal(2, batches.Count); + ArrowReaderVerifier.CompareBatches(expectedBatch1, batches[0]); + ArrowReaderVerifier.CompareBatches(expectedBatch2, batches[1]); + } + + [Fact] + public async Task TestGetStreamReadHonoursCancellation() + { + var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test"); + var expectedBatch = CreateTestBatch(0, 100); + GivenStoreBatches(flightDescriptor, new RecordBatchWithMetadata(expectedBatch)); + + var flightInfo = await _flightClient.GetInfo(flightDescriptor); + var endpoint = flightInfo.Endpoints.First(); + var getStream = _flightClient.GetStream(endpoint.Ticket); + + // Read the schema first so cancellation is exercised on the record-batch read loop. + await getStream.ResponseStream.Schema; + + var cts = new CancellationTokenSource(); + cts.Cancel(); + + var exception = await Assert.ThrowsAsync( + async () => await getStream.ResponseStream.MoveNext(cts.Token)); + Assert.Equal(StatusCode.Cancelled, exception.StatusCode); + } + + private Schema GetStoreSchema(FlightDescriptor flightDescriptor) { Assert.Contains(flightDescriptor, (IReadOnlyDictionary)_flightStore.Flights);