You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by cu...@apache.org on 2026/07/06 19:54:46 UTC

(arrow-dotnet) branch main updated: feat: add ICompressionCodecFactory support to Flight record batch readers (#285)

This is an automated email from the ASF dual-hosted git repository.

CurtHagenlocher pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-dotnet.git


The following commit(s) were added to refs/heads/main by this push:
     new 2e54e4e  feat: add ICompressionCodecFactory support to Flight record batch readers (#285)
2e54e4e is described below

commit 2e54e4ee4d51eec0083714f700c281884ec10cd4
Author: David Linton <51...@users.noreply.github.com>
AuthorDate: Mon Jul 6 21:54:41 2026 +0200

    feat: add ICompressionCodecFactory support to Flight record batch readers (#285)
    
    ## Summary
    
    Resolves #184.
    
    ### The Problem
    
    When a Flight server sends IPC-compressed record batches (LZ4/Zstd), the
    C# Flight client throws:
    
    > "Body is compressed with codec LZ4_FRAME but no
    ICompressionCodecFactory has been configured to decompress buffers"
    
    The decompression infrastructure already exists — `ArrowStreamReader`
    accepts `ICompressionCodecFactory` and the `Apache.Arrow.Compression`
    package provides LZ4/Zstd codecs. But the Flight reader chain never
    passes the factory through:
    
    ```
    FlightClient.GetStream()
      → FlightClientRecordBatchStreamReader(stream)
        → FlightRecordBatchStreamReader(stream)
          → RecordBatchReaderImplementation(stream)
            → ArrowReaderImplementation()          ← parameterless, _compressionCodecFactory = null
    ```
    
    `RecordBatchReaderImplementation` calls the base constructor without a
    codec factory, so `GetBufferCreator()` always hits the "no factory
    configured" error path.
    
    ### The Fix
    
    Thread `ICompressionCodecFactory` through the reader hierarchy (5 files,
    all with `= null` defaults for backward compatibility):
    
    1. **`RecordBatchReaderImplementation`** — accept
    `ICompressionCodecFactory`, pass to `base(allocator: null,
    compressionCodecFactory)`
    2. **`FlightRecordBatchStreamReader`** — accept and forward to
    `RecordBatchReaderImplementation`
    3. **`FlightClientRecordBatchStreamReader`** — forward to base
    4. **`FlightServerRecordBatchStreamReader`** — forward to base (both
    constructors)
    5. **`FlightClient`** — store as field, add constructor overloads, pass
    to readers in `GetStream()` and `DoExchange()`
    
    After the fix:
    
    ```
    FlightClient(channel, new CompressionCodecFactory())
      → stores _compressionCodecFactory
    
    FlightClient.GetStream(ticket)
      → FlightClientRecordBatchStreamReader(stream, _compressionCodecFactory)
        → FlightRecordBatchStreamReader(stream, compressionCodecFactory)
          → RecordBatchReaderImplementation(stream, compressionCodecFactory)
            → ArrowReaderImplementation(null, compressionCodecFactory)  ← decompression works
    ```
    
    ### Usage
    
    ```csharp
    using Apache.Arrow.Compression;
    
    // Before (compressed batches cause error):
    var client = new FlightClient(channel);
    
    // After (LZ4/Zstd decompression handled transparently):
    var client = new FlightClient(channel, new CompressionCodecFactory());
    ```
    
    ### Impact
    
    - **25 lines added, 8 removed** across 5 files
    - **No breaking changes** — all new parameters default to `null`
    - **No new dependencies** — `ICompressionCodecFactory` is in
    `Apache.Arrow.Ipc`, already referenced by the Flight package
    
    ## Test plan
    
    - [ ] Existing Flight tests pass unchanged (backward compatibility via
    `null` defaults)
    - [ ] Test with a Flight server sending LZ4-compressed record batches
    - [ ] Test with a Flight server sending Zstd-compressed record batches
    - [ ] Confirm `GetStream()` and `DoExchange()` correctly decompress when
    `ICompressionCodecFactory` is provided
    - [ ] Confirm no regression when `ICompressionCodecFactory` is `null`
    (existing behavior)
---
 src/Apache.Arrow.Flight/Client/FlightClient.cs     |  31 +++++-
 .../Client/FlightClientRecordBatchStreamReader.cs  |   2 +-
 .../FlightRecordBatchStreamReader.cs               |   4 +-
 .../Internal/RecordBatchReaderImplementation.cs    |   3 +-
 .../Server/FlightServerRecordBatchStreamReader.cs  |   9 +-
 test/Apache.Arrow.Flight.Tests/FlightTests.cs      | 105 +++++++++++++++++++++
 6 files changed, 146 insertions(+), 8 deletions(-)

diff --git a/src/Apache.Arrow.Flight/Client/FlightClient.cs b/src/Apache.Arrow.Flight/Client/FlightClient.cs
index d094266..c8dc82f 100644
--- a/src/Apache.Arrow.Flight/Client/FlightClient.cs
+++ b/src/Apache.Arrow.Flight/Client/FlightClient.cs
@@ -27,15 +27,42 @@ namespace Apache.Arrow.Flight.Client
         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 @@ namespace Apache.Arrow.Flight.Client
         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 @@ namespace Apache.Arrow.Flight.Client
         {
             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,
diff --git a/src/Apache.Arrow.Flight/Client/FlightClientRecordBatchStreamReader.cs b/src/Apache.Arrow.Flight/Client/FlightClientRecordBatchStreamReader.cs
index 65eb6b6..fc0bd89 100644
--- a/src/Apache.Arrow.Flight/Client/FlightClientRecordBatchStreamReader.cs
+++ b/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)
         {
         }
     }
diff --git a/src/Apache.Arrow.Flight/FlightRecordBatchStreamReader.cs b/src/Apache.Arrow.Flight/FlightRecordBatchStreamReader.cs
index 7400ec1..7336b85 100644
--- a/src/Apache.Arrow.Flight/FlightRecordBatchStreamReader.cs
+++ b/src/Apache.Arrow.Flight/FlightRecordBatchStreamReader.cs
@@ -40,9 +40,9 @@ namespace Apache.Arrow.Flight
 
         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();
diff --git a/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs b/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs
index 16e080d..54fd8c2 100644
--- a/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs
+++ b/src/Apache.Arrow.Flight/Internal/RecordBatchReaderImplementation.cs
@@ -32,7 +32,8 @@ namespace Apache.Arrow.Flight.Internal
         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>();
diff --git a/src/Apache.Arrow.Flight/Server/FlightServerRecordBatchStreamReader.cs b/src/Apache.Arrow.Flight/Server/FlightServerRecordBatchStreamReader.cs
index c52b761..ca93610 100644
--- a/src/Apache.Arrow.Flight/Server/FlightServerRecordBatchStreamReader.cs
+++ b/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)
         {
         }
 
diff --git a/test/Apache.Arrow.Flight.Tests/FlightTests.cs b/test/Apache.Arrow.Flight.Tests/FlightTests.cs
index 8c69211..46662cc 100644
--- a/test/Apache.Arrow.Flight.Tests/FlightTests.cs
+++ b/test/Apache.Arrow.Flight.Tests/FlightTests.cs
@@ -577,5 +577,110 @@ namespace Apache.Arrow.Flight.Tests
 
             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);
+        }
     }
 }