Add Apache Arrow result streaming to DuckDBCommand

DuckDB.NET.Bindings/DuckDBWrapperObjects.cs

@@ -68,6 +68,28 @@ protected override bool ReleaseHandle()
     }
 }
 
+public class DuckDBArrowOptions() : SafeHandleZeroOrMinusOneIsInvalid(true)
+{
+    protected override bool ReleaseHandle()
+    {
+        NativeMethods.Arrow.DuckDBDestroyArrowOptions(ref handle);
+        return true;
+    }
+}
+
+public class DuckDBErrorData() : SafeHandleZeroOrMinusOneIsInvalid(true)
+{
+    public bool HasError => !IsInvalid && NativeMethods.ErrorData.DuckDBErrorDataHasError(this);
+
+    public string? Message => IsInvalid ? null : NativeMethods.ErrorData.DuckDBErrorDataMessage(this);
+
+    protected override bool ReleaseHandle()
+    {
+        NativeMethods.ErrorData.DuckDBDestroyErrorData(ref handle);
+        return true;
+    }
+}
+
 public class DuckDBDataChunk : SafeHandleZeroOrMinusOneIsInvalid
 {
     public DuckDBDataChunk() : base(true)

DuckDB.NET.Bindings/NativeMethods/NativeMethods.Arrow.cs

@@ -0,0 +1,27 @@
+namespace DuckDB.NET.Native;
+
+public partial class NativeMethods
+{
+    //https://duckdb.org/docs/stable/clients/c/api#arrow-interface
+    public static partial class Arrow
+    {
+        [LibraryImport(DuckDbLibrary, EntryPoint = "duckdb_result_get_arrow_options")]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
+        public static partial DuckDBArrowOptions DuckDBResultGetArrowOptions(ref DuckDBResult result);
+
+        [LibraryImport(DuckDbLibrary, EntryPoint = "duckdb_destroy_arrow_options")]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
+        public static partial void DuckDBDestroyArrowOptions(ref IntPtr arrowOptions);
+
+        // duckdb_error_data duckdb_to_arrow_schema(duckdb_arrow_options, duckdb_logical_type *types,
+        //                                          const char **names, idx_t column_count, ArrowSchema *out_schema)
+        [LibraryImport(DuckDbLibrary, EntryPoint = "duckdb_to_arrow_schema")]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
+        public static partial DuckDBErrorData DuckDBToArrowSchema(DuckDBArrowOptions arrowOptions, IntPtr types, IntPtr names, ulong columnCount, IntPtr outSchema);
+
+        // duckdb_error_data duckdb_data_chunk_to_arrow(duckdb_arrow_options, duckdb_data_chunk, ArrowArray *out_arrow_array)
+        [LibraryImport(DuckDbLibrary, EntryPoint = "duckdb_data_chunk_to_arrow")]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
+        public static partial DuckDBErrorData DuckDBDataChunkToArrow(DuckDBArrowOptions arrowOptions, DuckDBDataChunk chunk, IntPtr outArray);
+    }
+}

DuckDB.NET.Bindings/NativeMethods/NativeMethods.ErrorData.cs

@@ -0,0 +1,23 @@
+namespace DuckDB.NET.Native;
+
+public partial class NativeMethods
+{
+    //https://duckdb.org/docs/stable/clients/c/api#error-data
+    public static partial class ErrorData
+    {
+        [SuppressGCTransition]
+        [LibraryImport(DuckDbLibrary, EntryPoint = "duckdb_error_data_has_error")]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
+        [return: MarshalAs(UnmanagedType.I1)]
+        public static partial bool DuckDBErrorDataHasError(DuckDBErrorData errorData);
+
+        [LibraryImport(DuckDbLibrary, EntryPoint = "duckdb_error_data_message")]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
+        [return: MarshalUsing(typeof(DuckDBOwnedStringMarshaller))]
+        public static partial string DuckDBErrorDataMessage(DuckDBErrorData errorData);
+
+        [LibraryImport(DuckDbLibrary, EntryPoint = "duckdb_destroy_error_data")]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
+        public static partial void DuckDBDestroyErrorData(ref IntPtr errorData);
+    }
+}

DuckDB.NET.Data/Arrow/DuckDBArrowArrayStream.cs

@@ -0,0 +1,163 @@
+using System.Runtime.InteropServices;
+using System.Threading;
+using System.Threading.Tasks;
+using Apache.Arrow;
+using Apache.Arrow.C;
+using Apache.Arrow.Ipc;
+using DuckDB.NET.Native;
+
+namespace DuckDB.NET.Data.Arrow;
+
+/// <summary>
+/// Streams the rows of a DuckDB query result as Apache Arrow <see cref="RecordBatch"/> values using
+/// DuckDB's Arrow C Data Interface (<c>duckdb_to_arrow_schema</c> / <c>duckdb_data_chunk_to_arrow</c>).
+/// Each DuckDB data chunk is converted into one Arrow record batch and imported with no row-by-row marshaling.
+/// </summary>
+internal sealed class DuckDBArrowArrayStream : IArrowArrayStream
+{
+    private DuckDBResult result;
+    private readonly DuckDBArrowOptions arrowOptions;
+    private readonly bool streaming;
+    private bool disposed;
+
+    public Schema Schema { get; }
+
+    internal DuckDBArrowArrayStream(DuckDBResult result)
+    {
+        this.result = result;
+
+        arrowOptions = NativeMethods.Arrow.DuckDBResultGetArrowOptions(ref this.result);
+        if (arrowOptions.IsInvalid)
+        {
+            this.result.Close();
+            throw new InvalidOperationException("Failed to obtain Arrow options from the DuckDB result.");
+        }
+
+        streaming = NativeMethods.Types.DuckDBResultIsStreaming(this.result) > 0;
+
+        try
+        {
+            Schema = BuildSchema();
+        }
+        catch
+        {
+            arrowOptions.Dispose();
+            this.result.Close();
+            throw;
+        }
+    }
+
+    private unsafe Schema BuildSchema()
+    {
+        var columnCount = NativeMethods.Query.DuckDBColumnCount(ref result);
+
+        var logicalTypes = new DuckDBLogicalType[columnCount];
+        var typeHandles = new IntPtr[columnCount];
+        var namePointers = new IntPtr[columnCount];
+
+        try
+        {
+            for (var index = 0UL; index < columnCount; index++)
+            {
+                var logicalType = NativeMethods.Query.DuckDBColumnLogicalType(ref result, (long)index);
+                logicalTypes[index] = logicalType;
+                typeHandles[index] = logicalType.DangerousGetHandle();
+
+                var name = NativeMethods.Query.DuckDBColumnName(ref result, (long)index);
+                namePointers[index] = Marshal.StringToCoTaskMemUTF8(name);
+            }
+
+            var cSchema = CArrowSchema.Create();
+
+            try
+            {
+                fixed (IntPtr* typesPointer = typeHandles)
+                fixed (IntPtr* namesPointer = namePointers)
+                {
+                    var error = NativeMethods.Arrow.DuckDBToArrowSchema(arrowOptions, (IntPtr)typesPointer, (IntPtr)namesPointer, columnCount, (IntPtr)cSchema);
+                    error.ThrowOnError("Failed to convert the DuckDB result schema to an Arrow schema.");
+                }
+
+                return CArrowSchemaImporter.ImportSchema(cSchema);
+            }
+            finally
+            {
+                CArrowSchema.Free(cSchema);
+            }
+        }
+        finally
+        {
+            foreach (var pointer in namePointers)
+            {
+                if (pointer != IntPtr.Zero)
+                {
+                    Marshal.FreeCoTaskMem(pointer);
+                }
+            }
+
+            foreach (var logicalType in logicalTypes)
+            {
+                logicalType?.Dispose();
+            }
+        }
+    }
+
+    public ValueTask<RecordBatch?> ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)
+    {
+        ObjectDisposedException.ThrowIf(disposed, this);
+
+        if (cancellationToken.IsCancellationRequested)
+        {
+            return new ValueTask<RecordBatch?>(Task.FromCanceled<RecordBatch?>(cancellationToken));
+        }
+
+        var chunk = streaming
+            ? NativeMethods.StreamingResult.DuckDBStreamFetchChunk(result)
+            : NativeMethods.Query.DuckDBFetchChunk(result);
+
+        if (chunk.IsInvalid)
+        {
+            chunk.Dispose();
+            return new ValueTask<RecordBatch?>((RecordBatch?)null);
+        }
+
+        try
+        {
+            return new ValueTask<RecordBatch?>(ConvertChunk(chunk));
+        }
+        finally
+        {
+            chunk.Dispose();
+        }
+    }
+
+    private unsafe RecordBatch ConvertChunk(DuckDBDataChunk chunk)
+    {
+        var cArray = CArrowArray.Create();
+
+        try
+        {
+            var error = NativeMethods.Arrow.DuckDBDataChunkToArrow(arrowOptions, chunk, (IntPtr)cArray);
+            error.ThrowOnError("Failed to convert a DuckDB data chunk to an Arrow array.");
+
+            return CArrowArrayImporter.ImportRecordBatch(cArray, Schema);
+        }
+        finally
+        {
+            CArrowArray.Free(cArray);
+        }
+    }
+
+    public void Dispose()
+    {
+        if (disposed)
+        {
+            return;
+        }
+
+        disposed = true;
+
+        arrowOptions.Dispose();
+        result.Close();
+    }
+}

DuckDB.NET.Data/Data.csproj

@@ -31,6 +31,10 @@ New features:
     <InternalsVisibleTo Include="DuckDB.NET.Test, PublicKey=0024000004800000940000000602000000240000525341310004000001000100852ebcffb69a0dfb906bc0377ec8608f4a00ba9af2e29300d1ed77dcb2583f08116b1a1006d202d53a48ec0561c1816738368378f7c5d335fec3daa63ba1b5413298153f886aafc75304e7653715f2395ad370fe3b2f4bc44a36a2f6b958fd500a2f7eea9a69c6ab5819e0933db962630c56c1610c7c87ed6a3c2b36e1ca4ed2" />
   </ItemGroup>
 
+  <ItemGroup>
+    <PackageReference Include="Apache.Arrow" Version="23.0.0" />
+  </ItemGroup>
+
   <ItemGroup>
     <ProjectReference Include="..\DuckDB.NET.Bindings\Bindings.csproj" />
   </ItemGroup>

DuckDB.NET.Data/DuckDBCommand.cs

@@ -1,6 +1,11 @@
 using System.ComponentModel;
 using System.Diagnostics.CodeAnalysis;
 using System.Runtime.CompilerServices;
+using System.Threading;
+using System.Threading.Tasks;
+using Apache.Arrow;
+using Apache.Arrow.Ipc;
+using DuckDB.NET.Data.Arrow;
 
 namespace DuckDB.NET.Data;
 
@@ -109,6 +114,57 @@ protected override DbDataReader ExecuteDbDataReader(CommandBehavior behavior)
         return reader;
     }
 
+    /// <summary>
+    /// Executes the command and returns the first result set as an Apache Arrow
+    /// <see cref="IArrowArrayStream"/>. Each DuckDB data chunk is converted to an Arrow record batch
+    /// using DuckDB's Arrow C Data Interface, with no row-by-row marshaling.
+    /// When <see cref="UseStreamingMode"/> is enabled, record batches are produced lazily from a
+    /// streaming result (bounded memory), otherwise the result is materialized first.
+    /// The caller owns the returned stream and must dispose it.
+    /// </summary>
+    public IArrowArrayStream ExecuteArrowStream()
+    {
+        EnsureConnectionOpen();
+
+        var results = PreparedStatement.PreparedStatement.PrepareMultiple(connection!.NativeConnection, CommandText, parameters, UseStreamingMode);
+
+        foreach (var result in results)
+        {
+            var current = result;
+
+            if (NativeMethods.Query.DuckDBResultReturnType(current) == DuckDBResultType.QueryResult)
+            {
+                return new DuckDBArrowArrayStream(current);
+            }
+
+            current.Close();
+        }
+
+        throw new InvalidOperationException("The command did not return a result set.");
+    }
+
+    /// <summary>
+    /// Executes the command and asynchronously streams the first result set as Apache Arrow
+    /// <see cref="RecordBatch"/> values. The batches are produced lazily, one per DuckDB data chunk.
+    /// Set <see cref="UseStreamingMode"/> to stream from a streaming result with bounded memory.
+    /// </summary>
+    public async IAsyncEnumerable<RecordBatch> ExecuteArrowBatchesAsync([EnumeratorCancellation] CancellationToken cancellationToken = default)
+    {
+        var stream = ExecuteArrowStream();
+
+        try
+        {
+            while (await stream.ReadNextRecordBatchAsync(cancellationToken).ConfigureAwait(false) is { } batch)
+            {
+                yield return batch;
+            }
+        }
+        finally
+        {
+            stream.Dispose();
+        }
+    }
+
     public override void Prepare() { }
 
     protected override DbParameter CreateDbParameter() => new DuckDBParameter();

DuckDB.NET.Data/Extensions/DuckDBErrorDataExtensions.cs

@@ -0,0 +1,15 @@
+namespace DuckDB.NET.Data.Extensions;
+
+internal static class DuckDBErrorDataExtensions
+{
+    public static void ThrowOnError(this DuckDBErrorData errorData, string message)
+    {
+        using (errorData)
+        {
+            if (errorData.HasError)
+            {
+                throw new DuckDBException($"{message} {errorData.Message}".TrimEnd());
+            }
+        }
+    }
+}