Skip to content

Commit ed026ae

Browse files
eschgieschgi
committed
DuckDBAppender: Add enum support
1 parent 21940c3 commit ed026ae

File tree

6 files changed

+144
-1
lines changed

6 files changed

+144
-1
lines changed

‎DuckDB.NET.Bindings/NativeMethods/NativeMethods.LogicalType.cs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@ public static class LogicalType
2626
[DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_enum_internal_type")]
2727
public static extern DuckDBType DuckDBEnumInternalType(DuckDBLogicalType type);
2828

29+
[DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_enum_dictionary_size")]
30+
public static extern uint DuckDBEnumDictionarySize(DuckDBLogicalType type);
31+
2932
[DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_enum_dictionary_value")]
3033
public static extern IntPtr DuckDBEnumDictionaryValue(DuckDBLogicalType type, long index);
3134

‎DuckDB.NET.Data/DuckDBAppenderRow.cs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,12 @@ public void EndRow()
6868

6969
#endregion
7070

71+
#region Append Enum
72+
73+
public DuckDBAppenderRow AppendValue<TEnum>(TEnum value) where TEnum : Enum => AppendValueInternal(value);
74+
75+
#endregion
76+
7177
#region Append Float
7278

7379
public DuckDBAppenderRow AppendValue(float? value) => AppendValueInternal(value);

‎DuckDB.NET.Data/Internal/Writer/EnumVectorDataWriter.cs‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,103 @@
1+
using System;
2+
using System.Collections.Generic;
3+
using System.Runtime.CompilerServices;
4+
using DuckDB.NET.Native;
5+
6+
namespace DuckDB.NET.Data.Internal.Writer;
7+
8+
internal sealed unsafe class EnumVectorDataWriter : VectorDataWriterBase
9+
{
10+
private readonly DuckDBType enumType;
11+
12+
private readonly Dictionary<string, long> enumValueIndexDictionary;
13+
14+
public EnumVectorDataWriter(IntPtr vector, void* vectorData, DuckDBLogicalType logicalType, DuckDBType columnType) : base(vector, vectorData, columnType)
15+
{
16+
enumType = NativeMethods.LogicalType.DuckDBEnumInternalType(logicalType);
17+
uint size = NativeMethods.LogicalType.DuckDBEnumDictionarySize(logicalType);
18+
enumValueIndexDictionary = [];
19+
for (long index = 0; index < size; index++)
20+
{
21+
string enumValue = NativeMethods.LogicalType.DuckDBEnumDictionaryValue(logicalType, index).ToManagedString();
22+
enumValueIndexDictionary.Add(enumValue, index);
23+
}
24+
}
25+
26+
internal override bool AppendString(string value, int rowIndex)
27+
{
28+
switch (enumType)
29+
{
30+
case DuckDBType.UnsignedTinyInt:
31+
{
32+
if (enumValueIndexDictionary.TryGetValue(value, out long enumValueIndex) &&
33+
enumValueIndex >= byte.MinValue && enumValueIndex <= byte.MaxValue)
34+
{
35+
return AppendValueInternal((byte)enumValueIndex, rowIndex);
36+
}
37+
38+
return false;
39+
}
40+
case DuckDBType.UnsignedSmallInt:
41+
{
42+
if (enumValueIndexDictionary.TryGetValue(value, out long enumValueIndex) &&
43+
enumValueIndex >= ushort.MinValue && enumValueIndex <= ushort.MaxValue)
44+
{
45+
return AppendValueInternal((ushort)enumValueIndex, rowIndex);
46+
}
47+
48+
return false;
49+
}
50+
case DuckDBType.UnsignedInteger:
51+
{
52+
if (enumValueIndexDictionary.TryGetValue(value, out long enumValueIndex) &&
53+
enumValueIndex >= uint.MinValue && enumValueIndex <= uint.MaxValue)
54+
{
55+
return AppendValueInternal((uint)enumValueIndex, rowIndex);
56+
}
57+
58+
return false;
59+
}
60+
default:
61+
return false;
62+
}
63+
}
64+
65+
internal override bool AppendEnum<TEnum>(TEnum value, int rowIndex)
66+
{
67+
switch (enumType)
68+
{
69+
case DuckDBType.UnsignedTinyInt:
70+
{
71+
long enumValueIndex = Convert.ToInt64(value);
72+
if (enumValueIndex >= byte.MinValue && enumValueIndex <= byte.MaxValue)
73+
{
74+
return AppendValueInternal((byte)enumValueIndex, rowIndex);
75+
}
76+
77+
return false;
78+
}
79+
case DuckDBType.UnsignedSmallInt:
80+
{
81+
long enumValueIndex = Convert.ToInt64(value);
82+
if (enumValueIndex >= ushort.MinValue && enumValueIndex <= ushort.MaxValue)
83+
{
84+
return AppendValueInternal((ushort)enumValueIndex, rowIndex);
85+
}
86+
87+
return false;
88+
}
89+
case DuckDBType.UnsignedInteger:
90+
{
91+
long enumValueIndex = Convert.ToInt64(value);
92+
if (enumValueIndex >= uint.MinValue && enumValueIndex <= uint.MaxValue)
93+
{
94+
return AppendValueInternal((uint)enumValueIndex, rowIndex);
95+
}
96+
97+
return false;
98+
}
99+
default:
100+
return false;
101+
}
102+
}
103+
}

‎DuckDB.NET.Data/Internal/Writer/VectorDataWriterBase.cs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,8 @@ public void AppendValue<T>(T value, int rowIndex)
5050
decimal val => AppendDecimal(val, rowIndex),
5151
BigInteger val => AppendBigInteger(val, rowIndex),
5252

53+
Enum val => AppendEnum(val, rowIndex),
54+
5355
string val => AppendString(val, rowIndex),
5456
Guid val => AppendGuid(val, rowIndex),
5557
DateTime val => AppendDateTime(val, rowIndex),
@@ -96,6 +98,8 @@ public void AppendValue<T>(T value, int rowIndex)
9698

9799
internal virtual bool AppendBigInteger(BigInteger value, int rowIndex) => ThrowException<BigInteger>();
98100

101+
internal virtual bool AppendEnum<TEnum>(TEnum value, int rowIndex) where TEnum : Enum => ThrowException<TEnum>();
102+
99103
internal virtual bool AppendCollection(ICollection value, int rowIndex) => ThrowException<bool>();
100104

101105
private bool ThrowException<T>()

‎DuckDB.NET.Data/Internal/Writer/VectorDataWriterFactory.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ public static unsafe VectorDataWriterBase CreateWriter(IntPtr vector, DuckDBLogi
2727
DuckDBType.Blob => new StringVectorDataWriter(vector, dataPointer, columnType),
2828
DuckDBType.Varchar => new StringVectorDataWriter(vector, dataPointer, columnType),
2929
DuckDBType.Bit => throw new NotImplementedException($"Writing {columnType} to data chunk is not yet supported"),
30-
DuckDBType.Enum => throw new NotImplementedException($"Writing {columnType} to data chunk is not yet supported"),
30+
DuckDBType.Enum => new EnumVectorDataWriter(vector, dataPointer, logicalType, columnType),
3131
DuckDBType.Struct => throw new NotImplementedException($"Writing {columnType} to data chunk is not yet supported"),
3232
DuckDBType.Decimal => new DecimalVectorDataWriter(vector, dataPointer, logicalType, columnType),
3333
DuckDBType.TimestampS => new DateTimeVectorDataWriter(vector, dataPointer, columnType),

‎DuckDB.NET.Test/DuckDBManagedAppenderTests.cs‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,26 @@ public void TemporalValues()
275275
result.Select(tuple => tuple.Item8).Should().BeEquivalentTo(dates.Select(TimeOnly.FromDateTime));
276276
}
277277

278+
[Fact]
279+
public void EnumValues()
280+
{
281+
Command.CommandText = "CREATE TYPE test_enum AS ENUM ('test1', 'test2', 'test3')";
282+
Command.ExecuteNonQuery();
283+
284+
Command.CommandText = "CREATE TABLE managedAppenderEnum(a test_enum, b test_enum, c test_enum);";
285+
Command.ExecuteNonQuery();
286+
287+
using (var appender = Connection.CreateAppender("managedAppenderEnum"))
288+
{
289+
appender.CreateRow().AppendValue("test1").AppendValue("test2").AppendValue(TestEnum.Test3).EndRow();
290+
}
291+
292+
var result = Connection.Query<(string, string, TestEnum)>("SELECT a, b, c FROM managedAppenderEnum").ToList();
293+
result.Select(tuple => tuple.Item1).Should().BeEquivalentTo("test1");
294+
result.Select(tuple => tuple.Item2).Should().BeEquivalentTo("test2");
295+
result.Select(tuple => tuple.Item3).Should().Equal(TestEnum.Test3);
296+
}
297+
278298
[Fact]
279299
public void IncompleteRowThrowsException()
280300
{
@@ -509,4 +529,11 @@ private static string GetQualifiedObjectName(params string[] parts) =>
509529
Where(p => !string.IsNullOrWhiteSpace(p)).
510530
Select(p => '"' + p + '"')
511531
);
532+
533+
private enum TestEnum : long
534+
{
535+
Test1 = 0,
536+
Test2 = 1,
537+
Test3 = 2,
538+
}
512539
}

0 commit comments

Comments
 (0)