Allow empty streams to return an empty schema asynchronously.

Author: CurtHagenlocher

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

@@ -67,7 +67,7 @@ public override void ReadSchema()
             ReadSchemaAsync(CancellationToken.None).AsTask().Wait();
         }
 
-        public override async ValueTask ReadSchemaAsync(CancellationToken cancellationToken)
+        public override async ValueTask<Schema> ReadSchemaAsync(CancellationToken cancellationToken)
         {
             while (!HasReadSchema)
             {
@@ -107,6 +107,7 @@ public override async ValueTask ReadSchemaAsync(CancellationToken cancellationTo
                         throw new Exception($"Expected schema as the first message, but got: {message.HeaderType.ToString()}");
                 }
             }
+            return _schema;
         }
 
         public override async ValueTask<RecordBatch> ReadNextRecordBatchAsync(CancellationToken cancellationToken)

src/Apache.Arrow/Ipc/ArrowFileReaderImplementation.cs

@@ -52,11 +52,11 @@ public async ValueTask<int> RecordBatchCountAsync(CancellationToken cancellation
             return _footer.RecordBatchCount;
         }
 
-        public override async ValueTask ReadSchemaAsync(CancellationToken cancellationToken = default)
+        public override async ValueTask<Schema> ReadSchemaAsync(CancellationToken cancellationToken = default)
         {
             if (HasReadSchema)
             {
-                return;
+                return _schema;
             }
 
             await ValidateFileAsync(cancellationToken).ConfigureAwait(false);
@@ -82,6 +82,8 @@ public override async ValueTask ReadSchemaAsync(CancellationToken cancellationTo
                 EnsureFullRead(buffer, bytesRead);
 
                 ReadSchema(buffer);
+
+                return _schema;
             }
         }
 

src/Apache.Arrow/Ipc/ArrowMemoryReaderImplementation.cs

@@ -33,11 +33,11 @@ public ArrowMemoryReaderImplementation(ReadOnlyMemory<byte> buffer, ICompression
             _buffer = buffer;
         }
 
-        public override ValueTask ReadSchemaAsync(CancellationToken cancellationToken)
+        public override ValueTask<Schema> ReadSchemaAsync(CancellationToken cancellationToken)
         {
             cancellationToken.ThrowIfCancellationRequested();
             ReadSchema();
-            return default;
+            return new ValueTask<Schema>(_schema);
         }
 
         public override ValueTask<RecordBatch> ReadNextRecordBatchAsync(CancellationToken cancellationToken)

src/Apache.Arrow/Ipc/ArrowReaderImplementation.cs

@@ -69,7 +69,7 @@ protected virtual void Dispose(bool disposing)
         {
         }
 
-        public abstract ValueTask ReadSchemaAsync(CancellationToken cancellationToken);
+        public abstract ValueTask<Schema> ReadSchemaAsync(CancellationToken cancellationToken);
         public abstract void ReadSchema();
 
         public abstract ValueTask<RecordBatch> ReadNextRecordBatchAsync(CancellationToken cancellationToken);

src/Apache.Arrow/Ipc/ArrowStreamReader.cs

@@ -102,11 +102,11 @@ protected virtual void Dispose(bool disposing)
 
         public async ValueTask<Schema> GetSchema(CancellationToken cancellationToken = default)
         {
-            if (!_implementation.HasReadSchema)
+            if (_implementation.HasReadSchema)
             {
-                await _implementation.ReadSchemaAsync(cancellationToken);
+                return _implementation.Schema;
             }
-            return _implementation.Schema;
+            return await _implementation.ReadSchemaAsync(cancellationToken);
         }
 
         public ValueTask<RecordBatch> ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)

src/Apache.Arrow/Ipc/ArrowStreamReaderImplementation.cs

@@ -151,19 +151,19 @@ protected ReadResult ReadMessage()
             return new ReadResult(messageLength, result);
         }
 
-        public override async ValueTask ReadSchemaAsync(CancellationToken cancellationToken = default)
+        public override async ValueTask<Schema> ReadSchemaAsync(CancellationToken cancellationToken = default)
         {
             if (HasReadSchema)
             {
-                return;
+                return _schema;
             }
 
             // Figure out length of schema
             int schemaMessageLength = await ReadMessageLengthAsync(throwOnFullRead: true, returnOnEmptyStream: true, cancellationToken)
                 .ConfigureAwait(false);
             if (schemaMessageLength == 0)
             {
-                return;
+                return null;
             }
 
             using (ArrayPool<byte>.Shared.RentReturn(schemaMessageLength, out Memory<byte> buff))
@@ -174,6 +174,7 @@ public override async ValueTask ReadSchemaAsync(CancellationToken cancellationTo
 
                 Google.FlatBuffers.ByteBuffer schemabb = CreateByteBuffer(buff);
                 _schema = MessageSerializer.GetSchema(ReadMessage<Flatbuf.Schema>(schemabb), ref _dictionaryMemo);
+                return _schema;
             }
         }
 

test/Apache.Arrow.Tests/ArrowStreamReaderTests.cs

@@ -240,6 +240,17 @@ private static async Task TestReaderFromPartialReadStream(Func<ArrowStreamReader
             }
         }
 
+        [Fact]
+        public async Task EmptyStreamNoSyncRead()
+        {
+            using (var stream = new EmptyAsyncOnlyStream())
+            {
+                var reader = new ArrowStreamReader(stream);
+                var schema = await reader.GetSchema();
+                Assert.Null(schema);
+            }
+        }
+
         /// <summary>
         /// A stream class that only returns a part of the data at a time.
         /// </summary>
@@ -280,6 +291,20 @@ public override Task<int> ReadAsync(byte[] buffer, int offset, int length, Cance
             }
 #endif
         }
+
+        private class EmptyAsyncOnlyStream : Stream
+        {
+            public override bool CanRead => true;
+            public override bool CanSeek => false;
+            public override bool CanWrite => false;
+            public override long Length => 0;
+            public override long Position { get => 0; set => throw new NotSupportedException(); }
+            public override void Flush() { }
+            public override int Read(byte[] buffer, int offset, int count) => throw new NotSupportedException();
+            public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
+            public override void SetLength(long value) => throw new NotSupportedException();
+            public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
+            public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) => Task.FromResult(0);
+        }
     }
 }
-