feat: support reading dictionary-encoded columns over Arrow Flight

src/Apache.Arrow.Flight/Internal/FlightDataStream.cs

@@ -36,6 +36,8 @@ internal class FlightDataStream : ArrowStreamWriter
         private readonly FlightDescriptor _flightDescriptor;
         private readonly IAsyncStreamWriter<Protocol.FlightData> _clientStreamWriter;
         private Protocol.FlightData _currentFlightData;
+        private bool _hasPendingMessage;
+        private ByteString _recordBatchAppMetadata;
 
         public FlightDataStream(IAsyncStreamWriter<Protocol.FlightData> clientStreamWriter, FlightDescriptor flightDescriptor, Schema schema)
             : base(new MemoryStream(), schema)
@@ -46,17 +48,9 @@ public FlightDataStream(IAsyncStreamWriter<Protocol.FlightData> 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<long> WriteMessageAsync<T>(MessageHeader headerType, Offset<T> headerOffset, int bodyLength, VectorOffset customMetadataOffset, CancellationToken cancellationToken)
+        private protected override async ValueTask<long> WriteMessageAsync<T>(MessageHeader headerType, Offset<T> 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<Flatbuf.Message> messageOffset = Flatbuf.Message.CreateMessage(
                 Builder, CurrentMetadataVersion, headerType, headerOffset.Value,
                 bodyLength, customMetadataOffset);
@@ -101,9 +101,32 @@ private protected override ValueTask<long> WriteMessageAsync<T>(MessageHeader he
 
             ReadOnlyMemory<byte> 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<long>(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;
         }
     }
 }

src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs

@@ -103,7 +103,9 @@ public override async ValueTask<Schema> 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<RecordBatch> 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<RecordBatch> 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;

src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs

@@ -178,6 +178,15 @@ private static IReadOnlyDictionary<string, string> ReadMessageCustomMetadata(Fla
             return metadata;
         }
 
+        /// <summary>
+        /// Decode a schema message, registering any dictionary-encoded fields in the
+        /// reader's DictionaryMemo so subsequent dictionary batches can be resolved.
+        /// </summary>
+        protected Schema ReadSchemaFromMessage(ByteBuffer schemaBuffer)
+        {
+            return MessageSerializer.GetSchema(ReadMessage<Flatbuf.Schema>(schemaBuffer), ref _dictionaryMemo, _extensionRegistry);
+        }
+
         internal static ByteBuffer CreateByteBuffer(ReadOnlyMemory<byte> buffer)
         {
             return new ByteBuffer(new ReadOnlyMemoryBufferAllocator(buffer), 0);

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; }
 

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<int32, string>.
+            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<RpcException>(
+                async () => await getStream.ResponseStream.MoveNext(cts.Token));
+            Assert.Equal(StatusCode.Cancelled, exception.StatusCode);
+        }
+
+
         private Schema GetStoreSchema(FlightDescriptor flightDescriptor)
         {
             Assert.Contains(flightDescriptor, (IReadOnlyDictionary<FlightDescriptor, FlightHolder>)_flightStore.Flights);