[arrow] branch main updated: GH-35267: [C#] Serialize TotalBytes and TotalRecords in FlightInfo (#35222)

[westonpace]

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

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


The following commit(s) were added to refs/heads/main by this push:
     new b43f4cd4ce GH-35267: [C#] Serialize TotalBytes and TotalRecords in 
FlightInfo (#35222)
b43f4cd4ce is described below

commit b43f4cd4ce28e2ddb7bad89081e141f381b09d5e
Author: jeremyosterhoudt <jeremyosterho...@yahoo.com>
AuthorDate: Fri May 5 11:07:27 2023 -0700

    GH-35267: [C#] Serialize TotalBytes and TotalRecords in FlightInfo (#35222)
    
    Currently the `TotalBytes` and `TotalRecords` properties are not serialized 
into the `FlightInfo` ProtoBuf for C#/.NET.  This breaks compatibility with 
Arrow Flight implementation done in other languages (c++ for example).
    
    This PR adds the `TotalBytes` and `TotalRecords` properties to the protoBuf 
`FlightInfo` for the C# implementation.
    * Closes: #35267
    
    Authored-by: Jeremy Osterhoudt <jeremy.osterho...@ge.com>
    Signed-off-by: Weston Pace <weston.p...@gmail.com>
---
 csharp/src/Apache.Arrow.Flight/FlightInfo.cs         |  6 ++++--
 .../test/Apache.Arrow.Flight.TestWeb/FlightHolder.cs |  4 +++-
 csharp/test/Apache.Arrow.Flight.Tests/FlightTests.cs | 20 ++++++++++++++++++++
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/csharp/src/Apache.Arrow.Flight/FlightInfo.cs 
b/csharp/src/Apache.Arrow.Flight/FlightInfo.cs
index 44a7965cce..16ddb6fbfb 100644
--- a/csharp/src/Apache.Arrow.Flight/FlightInfo.cs
+++ b/csharp/src/Apache.Arrow.Flight/FlightInfo.cs
@@ -39,7 +39,7 @@ namespace Apache.Arrow.Flight
             TotalRecords = flightInfo.TotalRecords;
         }
 
-        public FlightInfo(Schema schema, FlightDescriptor descriptor, 
IReadOnlyList<FlightEndpoint> endpoints, long totalRecords = 0, long totalBytes 
= 0)
+        public FlightInfo(Schema schema, FlightDescriptor descriptor, 
IReadOnlyList<FlightEndpoint> endpoints, long totalRecords = -1, long 
totalBytes = -1)
         {
             Schema = schema;
             Descriptor = descriptor;
@@ -64,7 +64,9 @@ namespace Apache.Arrow.Flight
             var response = new Protocol.FlightInfo()
             {
                 Schema = serializedSchema,
-                FlightDescriptor = Descriptor.ToProtocol()
+                FlightDescriptor = Descriptor.ToProtocol(),
+                TotalBytes = TotalBytes,
+                TotalRecords = TotalRecords
             };
 
             foreach(var endpoint in Endpoints)
diff --git a/csharp/test/Apache.Arrow.Flight.TestWeb/FlightHolder.cs 
b/csharp/test/Apache.Arrow.Flight.TestWeb/FlightHolder.cs
index 34a5270182..c6f7e66c6c 100644
--- a/csharp/test/Apache.Arrow.Flight.TestWeb/FlightHolder.cs
+++ b/csharp/test/Apache.Arrow.Flight.TestWeb/FlightHolder.cs
@@ -51,12 +51,14 @@ namespace Apache.Arrow.Flight.TestWeb
 
         public FlightInfo GetFlightInfo()
         {
+            int batchArrayLength = _recordBatches.Sum(rb => 
rb.RecordBatch.Length);
+            int batchBytes = _recordBatches.Sum(rb => 
rb.RecordBatch.Arrays.Sum(arr => arr.Data.Buffers.Sum(b=>b.Length)));
             return new FlightInfo(_schema, _flightDescriptor, new 
List<FlightEndpoint>()
             {
                 new FlightEndpoint(new 
FlightTicket(_flightDescriptor.Paths.FirstOrDefault()), new 
List<FlightLocation>(){
                     new FlightLocation(_location)
                 })
-            });
+            }, batchArrayLength, batchBytes);
         }
     }
 }
diff --git a/csharp/test/Apache.Arrow.Flight.Tests/FlightTests.cs 
b/csharp/test/Apache.Arrow.Flight.Tests/FlightTests.cs
index 79025a2178..3556c6a17f 100644
--- a/csharp/test/Apache.Arrow.Flight.Tests/FlightTests.cs
+++ b/csharp/test/Apache.Arrow.Flight.Tests/FlightTests.cs
@@ -312,5 +312,25 @@ namespace Apache.Arrow.Flight.Tests
             ArrowReaderVerifier.CompareBatches(expectedBatch1, resultList[0]);
             ArrowReaderVerifier.CompareBatches(expectedBatch2, resultList[1]);
         }
+
+        [Fact]
+        public async Task 
EnsureTheSerializedBatchContainsTheProperTotalRecordsAndTotalBytesProperties()
+        {
+            var flightDescriptor1 = 
FlightDescriptor.CreatePathDescriptor("test1");
+            var expectedBatch = CreateTestBatch(0, 100);
+            var expectedTotalBytes = expectedBatch.Arrays.Sum(arr => 
arr.Data.Buffers.Sum(b => b.Length));
+
+            List<FlightInfo> expectedFlightInfo = new List<FlightInfo>();
+
+            expectedFlightInfo.Add(GivenStoreBatches(flightDescriptor1, new 
RecordBatchWithMetadata(expectedBatch)));
+
+            var listFlightStream = _flightClient.ListFlights();
+
+            var actualFlights = await 
listFlightStream.ResponseStream.ToListAsync();
+            var result = actualFlights.First();
+
+            Assert.Equal(expectedBatch.Length, result.TotalRecords);
+            Assert.Equal(expectedTotalBytes, result.TotalBytes);
+        }
     }
 }