feat: add ICompressionCodecFactory support to Flight record batch readers

src/Apache.Arrow.Flight/Client/FlightClient.cs

@@ -27,15 +27,42 @@ public class FlightClient
         internal static readonly Empty EmptyInstance = new Empty();
 
         private readonly FlightService.FlightServiceClient _client;
+        private readonly ArrowContext _context;
 
         public FlightClient(ChannelBase grpcChannel)
+            : this(grpcChannel, null)
+        {
+        }
+
+        public FlightClient(ChannelBase grpcChannel, ArrowContext context)
         {
             _client = new FlightService.FlightServiceClient(grpcChannel);
+            _context = context;
         }
 
         public FlightClient(CallInvoker callInvoker)
+            : this(callInvoker, null)
+        {
+        }
+
+        private FlightClient(CallInvoker callInvoker, ArrowContext context)
         {
             _client = new FlightService.FlightServiceClient(callInvoker);
+            _context = context;
+        }
+
+        /// <summary>
+        /// Creates a <see cref="FlightClient"/> from a <see cref="CallInvoker"/> with an <see cref="ArrowContext"/>.
+        /// </summary>
+        /// <remarks>
+        /// A factory method is used instead of a constructor overload so that dependency injection
+        /// (e.g. Grpc.Net.ClientFactory) can still resolve the single-argument constructor unambiguously.
+        /// </remarks>
+        /// <param name="callInvoker">The <see cref="CallInvoker"/> to use for gRPC calls</param>
+        /// <param name="context">The <see cref="ArrowContext"/> carrying the compression codec factory and other reader configuration</param>
+        public static FlightClient Create(CallInvoker callInvoker, ArrowContext context)
+        {
+            return new FlightClient(callInvoker, context);
         }
 
         public AsyncServerStreamingCall<FlightInfo> ListFlights(FlightCriteria criteria = null, Metadata headers = null)
@@ -77,7 +104,7 @@ public FlightRecordBatchStreamingCall GetStream(FlightTicket ticket, Metadata he
         public FlightRecordBatchStreamingCall GetStream(FlightTicket ticket, Metadata headers, System.DateTime? deadline, CancellationToken cancellationToken = default)
         {
             var stream = _client.DoGet(ticket.ToProtocol(), headers, deadline, cancellationToken);
-            var responseStream = new FlightClientRecordBatchStreamReader(stream.ResponseStream);
+            var responseStream = new FlightClientRecordBatchStreamReader(stream.ResponseStream, _context);
             return new FlightRecordBatchStreamingCall(responseStream, stream.ResponseHeadersAsync, stream.GetStatus, stream.GetTrailers, stream.Dispose);
         }
 
@@ -207,7 +234,7 @@ public FlightRecordBatchExchangeCall DoExchange(FlightDescriptor flightDescripto
         {
             var channel = _client.DoExchange(headers, deadline, cancellationToken);
             var requestStream = new FlightClientRecordBatchStreamWriter(channel.RequestStream, flightDescriptor);
-            var responseStream = new FlightClientRecordBatchStreamReader(channel.ResponseStream);
+            var responseStream = new FlightClientRecordBatchStreamReader(channel.ResponseStream, _context);
             var call = new FlightRecordBatchExchangeCall(
                 requestStream,
                 responseStream,

src/Apache.Arrow.Flight/Client/FlightClientRecordBatchStreamReader.cs

@@ -21,7 +21,7 @@ namespace Apache.Arrow.Flight.Client
 {
     public class FlightClientRecordBatchStreamReader : FlightRecordBatchStreamReader
     {
-        internal FlightClientRecordBatchStreamReader(IAsyncStreamReader<Protocol.FlightData> flightDataStream) : base(flightDataStream)
+        internal FlightClientRecordBatchStreamReader(IAsyncStreamReader<Protocol.FlightData> flightDataStream, ArrowContext context = null) : base(flightDataStream, context)
         {
         }
     }

src/Apache.Arrow.Flight/FlightRecordBatchStreamReader.cs

@@ -40,9 +40,9 @@ public abstract class FlightRecordBatchStreamReader : IAsyncStreamReader<RecordB
 
         private readonly RecordBatchReaderImplementation _arrowReaderImplementation;
 
-        private protected FlightRecordBatchStreamReader(IAsyncStreamReader<Protocol.FlightData> flightDataStream)
+        private protected FlightRecordBatchStreamReader(IAsyncStreamReader<Protocol.FlightData> flightDataStream, ArrowContext context = null)
         {
-            _arrowReaderImplementation = new RecordBatchReaderImplementation(flightDataStream);
+            _arrowReaderImplementation = new RecordBatchReaderImplementation(flightDataStream, context);
         }
 
         public ValueTask<Schema> Schema => _arrowReaderImplementation.GetSchemaAsync();

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

@@ -31,7 +31,8 @@ internal class RecordBatchReaderImplementation : ArrowReaderImplementation
         private FlightDescriptor _flightDescriptor;
         private readonly List<ByteString> _applicationMetadatas;
 
-        public RecordBatchReaderImplementation(IAsyncStreamReader<Protocol.FlightData> streamReader)
+        public RecordBatchReaderImplementation(IAsyncStreamReader<Protocol.FlightData> streamReader, ArrowContext context = null)
+            : base(context?.Allocator, context?.CompressionCodecFactory, context?.ExtensionRegistry)
         {
             _flightDataStream = streamReader;
             _applicationMetadatas = new List<ByteString>();

src/Apache.Arrow.Flight/Server/FlightServerRecordBatchStreamReader.cs

@@ -21,11 +21,16 @@ namespace Apache.Arrow.Flight.Server
 {
     public class FlightServerRecordBatchStreamReader : FlightRecordBatchStreamReader
     {
-        public FlightServerRecordBatchStreamReader(IAsyncStreamReader<FlightData> flightDataStream) : base(new StreamReader<FlightData, Protocol.FlightData>(flightDataStream, data => data.ToProtocol()))
+        public FlightServerRecordBatchStreamReader(IAsyncStreamReader<FlightData> flightDataStream)
+            : this(flightDataStream, null)
         {
         }
 
-        internal FlightServerRecordBatchStreamReader(IAsyncStreamReader<Protocol.FlightData> flightDataStream) : base(flightDataStream)
+        public FlightServerRecordBatchStreamReader(IAsyncStreamReader<FlightData> flightDataStream, ArrowContext context) : base(new StreamReader<FlightData, Protocol.FlightData>(flightDataStream, data => data.ToProtocol()), context)
+        {
+        }
+
+        internal FlightServerRecordBatchStreamReader(IAsyncStreamReader<Protocol.FlightData> flightDataStream, ArrowContext context = null) : base(flightDataStream, context)
         {
         }
 

test/Apache.Arrow.Flight.Tests/FlightTests.cs

@@ -577,5 +577,110 @@ public async Task TestIntegrationWithGrpcNetClientFactory()
 
             SchemaComparer.Compare(expectedSchema, actualSchema);
         }
+
+        [Fact]
+        public async Task TestIntegrationWithGrpcNetClientFactoryAndArrowContext()
+        {
+            IServiceCollection services = new ServiceCollection();
+
+            services.AddGrpcClient<FlightClient>(grpc => grpc.Address = new Uri(_testWebFactory.GetAddress()))
+                .ConfigureGrpcClientCreator(invoker =>
+                {
+                    return FlightClient.Create(invoker, new ArrowContext());
+                });
+
+            IServiceProvider provider = services.BuildServiceProvider();
+
+            // Test that a FlightClient carrying an ArrowContext can be created via the
+            // Grpc.Net.ClientFactory library using the FlightClient.Create factory method.
+            FlightClient flightClient = provider.GetRequiredService<FlightClient>();
+
+            // Test that the resolved client is functional.
+            var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test-factory-context");
+            var expectedBatch = CreateTestBatch(0, 100);
+            var expectedSchema = expectedBatch.Schema;
+
+            GivenStoreBatches(flightDescriptor, new RecordBatchWithMetadata(expectedBatch));
+
+            var actualSchema = await flightClient.GetSchema(flightDescriptor);
+
+            SchemaComparer.Compare(expectedSchema, actualSchema);
+        }
+
+        [Fact]
+        public async Task TestGetWithArrowContext()
+        {
+            // Verify that FlightClient works when constructed with an ArrowContext
+            var context = new ArrowContext();
+            var flightClient = new FlightClient(_testWebFactory.GetChannel(), context);
+
+            var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test-context");
+            var expectedBatch = CreateTestBatch(0, 100);
+
+            GivenStoreBatches(flightDescriptor, new RecordBatchWithMetadata(expectedBatch));
+
+            var flightInfo = await flightClient.GetInfo(flightDescriptor);
+            var endpoint = flightInfo.Endpoints.FirstOrDefault();
+
+            var getStream = flightClient.GetStream(endpoint.Ticket);
+            var resultList = await getStream.ResponseStream.ToListAsync();
+
+            Assert.Single(resultList);
+            ArrowReaderVerifier.CompareBatches(expectedBatch, resultList[0]);
+        }
+
+        [Fact]
+        public async Task TestGetWithNullArrowContext()
+        {
+            // Verify that FlightClient works with null ArrowContext (backward compat)
+            var flightClient = new FlightClient(_testWebFactory.GetChannel(), null);
+
+            var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test-null-context");
+            var expectedBatch = CreateTestBatch(0, 100);
+
+            GivenStoreBatches(flightDescriptor, new RecordBatchWithMetadata(expectedBatch));
+
+            var flightInfo = await flightClient.GetInfo(flightDescriptor);
+            var endpoint = flightInfo.Endpoints.FirstOrDefault();
+
+            var getStream = flightClient.GetStream(endpoint.Ticket);
+            var resultList = await getStream.ResponseStream.ToListAsync();
+
+            Assert.Single(resultList);
+            ArrowReaderVerifier.CompareBatches(expectedBatch, resultList[0]);
+        }
+
+        [Fact]
+        public async Task TestPutAndGetWithArrowContext()
+        {
+            // Verify put + get round-trip works with ArrowContext
+            var context = new ArrowContext();
+            var flightClient = new FlightClient(_testWebFactory.GetChannel(), context);
+
+            var flightDescriptor = FlightDescriptor.CreatePathDescriptor("test-context-roundtrip");
+            var expectedBatch = CreateTestBatch(0, 100);
+
+            var putStream = await flightClient.StartPut(flightDescriptor, expectedBatch.Schema);
+            await putStream.RequestStream.WriteAsync(expectedBatch);
+            await putStream.RequestStream.CompleteAsync();
+            await putStream.ResponseStream.ToListAsync();
+
+            var flightInfo = await flightClient.GetInfo(flightDescriptor);
+            var endpoint = flightInfo.Endpoints.FirstOrDefault();
+
+            var getStream = flightClient.GetStream(endpoint.Ticket);
+            var resultList = await getStream.ResponseStream.ToListAsync();
+
+            Assert.Single(resultList);
+            ArrowReaderVerifier.CompareBatches(expectedBatch, resultList[0]);
+        }
+
+        [Fact]
+        public void TestFlightClientDefaultConstructorStillWorks()
+        {
+            // Verify the original constructor without ArrowContext still works
+            var client = new FlightClient(_testWebFactory.GetChannel());
+            Assert.NotNull(client);
+        }
     }
 }