Repository navigation
Expand file tree
/
Copy pathDuckDBArrowArrayStream.cs
More file actions
170 lines (142 loc) · 4.95 KB
/
Copy pathDuckDBArrowArrayStream.cs
File metadata and controls
170 lines (142 loc) · 4.95 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
using System.Runtime.InteropServices;
using System.Threading;
using System.Threading.Tasks;
using Apache.Arrow;
using Apache.Arrow.C;
using Apache.Arrow.Ipc;
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 DuckDBNativeConnection connection;
private readonly DuckDBArrowOptions arrowOptions;
private readonly bool streaming;
private bool disposed;
public Schema Schema { get; }
internal DuckDBArrowArrayStream(DuckDBResult result, DuckDBNativeConnection connection)
{
this.result = result;
this.connection = connection;
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();
if (streaming)
{
result.ThrowOnError(connection);
}
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();
}
}