Add support for writing TimeTz to data chunk

Giorgi · May 10, 2024 · f87a95a · f87a95a
1 parent 668c65e
commit f87a95a
Show file tree

Hide file tree

Showing 6 changed files with 24 additions and 5 deletions.
diff --git a/DuckDB.NET.Bindings/NativeMethods/NativeMethods.DateTime.cs b/DuckDB.NET.Bindings/NativeMethods/NativeMethods.DateTime.cs
@@ -16,6 +16,9 @@ public static class DateTimeHelpers
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_from_time")]
         public static extern DuckDBTimeOnly DuckDBFromTime(DuckDBTime time);
 
+        [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_create_time_tz")]
+        public static extern DuckDBTimeTzStruct DuckDBCreateTimeTz(long micros, int offset);
+
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_from_time_tz")]
         public static extern DuckDBTimeTz DuckDBFromTimeTz(DuckDBTimeTzStruct micros);
 

diff --git a/DuckDB.NET.Data/DuckDBAppenderRow.cs b/DuckDB.NET.Data/DuckDBAppenderRow.cs
@@ -92,6 +92,8 @@ public void EndRow()
 
     public DuckDBAppenderRow AppendValue(DateTime? value) => AppendValueInternal(value);
 
+    public DuckDBAppenderRow AppendValue(DateTimeOffset? value) => AppendValueInternal(value);
+
     public DuckDBAppenderRow AppendValue(TimeSpan? value)
     {
         return AppendValueInternal(value);

diff --git a/DuckDB.NET.Data/Internal/Writer/DateTimeVectorDataWriter.cs b/DuckDB.NET.Data/Internal/Writer/DateTimeVectorDataWriter.cs
@@ -32,6 +32,14 @@ internal override bool AppendDateTime(DateTime value, int rowIndex)
         return AppendValueInternal(timestamp, rowIndex);
     }
 
+    internal override bool AppendDateTimeOffset(DateTimeOffset value, int rowIndex)
+    {
+        var time = NativeMethods.DateTimeHelpers.DuckDBToTime((DuckDBTimeOnly)value.DateTime);
+        var timeTz = NativeMethods.DateTimeHelpers.DuckDBCreateTimeTz(time.Micros, (int)value.Offset.TotalSeconds);
+
+        return AppendValueInternal(timeTz, rowIndex);
+    }
+
 #if NET6_0_OR_GREATER
     internal override bool AppendDateOnly(DateOnly value, int rowIndex) => AppendValueInternal(NativeMethods.DateTimeHelpers.DuckDBToDate(value), rowIndex);
 

diff --git a/DuckDB.NET.Data/Internal/Writer/VectorDataWriterBase.cs b/DuckDB.NET.Data/Internal/Writer/VectorDataWriterBase.cs
@@ -58,6 +58,7 @@ public unsafe void AppendValue<T>(T value, int rowIndex)
             DateOnly val => AppendDateOnly(val, rowIndex),
             TimeOnly val => AppendTimeOnly(val, rowIndex),
 #endif
+            DateTimeOffset val => AppendDateTimeOffset(val, rowIndex),
             _ => ThrowException<T>()
         };
     }
@@ -86,6 +87,8 @@ public unsafe void AppendValue<T>(T value, int rowIndex)
 
     internal virtual bool AppendTimeOnly(DuckDBTimeOnly value, int rowIndex) => ThrowException<DuckDBTimeOnly>();
 
+    internal virtual bool AppendDateTimeOffset(DateTimeOffset value, int rowIndex) => ThrowException<DateTimeOffset>();
+
     internal virtual bool AppendNumeric<T>(T value, int rowIndex) where T : unmanaged => ThrowException<T>();
 
     internal virtual bool AppendBigInteger(BigInteger value, int rowIndex) => ThrowException<BigInteger>();

diff --git a/DuckDB.NET.Data/Internal/Writer/VectorDataWriterFactory.cs b/DuckDB.NET.Data/Internal/Writer/VectorDataWriterFactory.cs
@@ -15,7 +15,7 @@ public static unsafe VectorDataWriterBase CreateWriter(IntPtr vector, DuckDBLogi
             DuckDBType.Uuid => new GuidVectorDataWriter(vector, dataPointer, columnType),
             DuckDBType.Date => new DateTimeVectorDataWriter(vector, dataPointer, columnType),
             DuckDBType.Time => new DateTimeVectorDataWriter(vector, dataPointer, columnType),
-            DuckDBType.TimeTz => throw new NotImplementedException($"Writing {columnType} to data chunk is not yet supported"),
+            DuckDBType.TimeTz => new DateTimeVectorDataWriter(vector, dataPointer, columnType),
             DuckDBType.Interval => new IntervalVectorDataWriter(vector, dataPointer, columnType),
             DuckDBType.Timestamp => new DateTimeVectorDataWriter(vector, dataPointer, columnType),
 

diff --git a/DuckDB.NET.Test/ManagedAppenderTests.cs b/DuckDB.NET.Test/ManagedAppenderTests.cs
@@ -9,6 +9,7 @@
 using System.Linq;
 using System.Numerics;
 using Dapper;
+using FluentAssertions.Common;
 
 namespace DuckDB.NET.Test;
 
@@ -235,10 +236,10 @@ public void IntervalValues()
     [Fact]
     public void TemporalValues()
     {
-        Command.CommandText = "CREATE TABLE managedAppenderTemporal(a Date, b TimeStamp, c TIMESTAMP_NS, d TIMESTAMP_MS, e TIMESTAMP_S, f TIMESTAMPTZ);";
+        Command.CommandText = "CREATE TABLE managedAppenderTemporal(a Date, b TimeStamp, c TIMESTAMP_NS, d TIMESTAMP_MS, e TIMESTAMP_S, f TIMESTAMPTZ, g TIMETZ);";
         Command.ExecuteNonQuery();
 
-        var dates = Enumerable.Range(0, 20).Select(i => new DateTime(1900, 1, 1).AddDays(Random.Shared.Next(1, 50000))).ToList();
+        var dates = Enumerable.Range(0, 20).Select(i => new DateTime(1900, 1, 1).AddDays(Random.Shared.Next(1, 50000)).AddSeconds(Random.Shared.Next(3600 * 2, 3600 * 24))).ToList();
 
         using (var appender = Connection.CreateAppender("managedAppenderTemporal"))
         {
@@ -247,19 +248,21 @@ public void TemporalValues()
                 appender.CreateRow()
                                 .AppendValue(value).AppendValue(value)
                                 .AppendValue(value).AppendValue(value)
-                                .AppendValue(value).AppendValue(value)
+                                .AppendValue(value).AppendValue(value).AppendValue(value.ToDateTimeOffset(TimeSpan.FromHours(1)))
                         .EndRow();
             }
         }
 
-        var result = Connection.Query<(DateOnly, DateTime, DateTime, DateTime, DateTime, DateTime)>("SELECT a, b, c, d, e, f FROM managedAppenderTemporal").ToList();
+        var result = Connection.Query<(DateOnly, DateTime, DateTime, DateTime, DateTime, DateTime, DateTimeOffset)>("SELECT a, b, c, d, e, f, g FROM managedAppenderTemporal").ToList();
 
         result.Select(tuple => tuple.Item1).Should().BeEquivalentTo(dates.Select(DateOnly.FromDateTime));
         result.Select(tuple => tuple.Item2).Should().BeEquivalentTo(dates);
         result.Select(tuple => tuple.Item3).Should().BeEquivalentTo(dates);
         result.Select(tuple => tuple.Item4).Should().BeEquivalentTo(dates);
         result.Select(tuple => tuple.Item5).Should().BeEquivalentTo(dates);
         result.Select(tuple => tuple.Item6).Should().BeEquivalentTo(dates);
+        result.Select(tuple => tuple.Item7).Should().BeEquivalentTo(dates.Select(time => time.ToDateTimeOffset(TimeSpan.FromHours(1))), 
+                                                              options => options.ComparingByMembers<DateTimeOffset>().Including(offset => offset.Offset).Including(offset => offset.TimeOfDay));
     }
 
     [Fact]