Initial implementation for table-valued functions

Giorgi · Giorgi · commit f98ae60734e1 · 2024-11-04T00:07:18.000+04:00
diff --git a/DuckDB.NET.Bindings/DuckDBWrapperObjects.cs b/DuckDB.NET.Bindings/DuckDBWrapperObjects.cs
@@ -1,5 +1,5 @@
-﻿using System;
-using Microsoft.Win32.SafeHandles;
+﻿using Microsoft.Win32.SafeHandles;
+using System;
 
 namespace DuckDB.NET.Native;
 
@@ -89,7 +89,7 @@ protected override bool ReleaseHandle()
     }
 }
 
-public class DuckDBValue() : SafeHandleZeroOrMinusOneIsInvalid(true)
+public class DuckDBValue() : SafeHandleZeroOrMinusOneIsInvalid(true), IDuckDBValueReader
 {
     private DuckDBValue[] childValues = [];
 
@@ -108,4 +108,11 @@ internal void SetChildValues(DuckDBValue[] values)
     { 
         childValues = values;
     }
+
+    public T GetValue<T>()
+    {
+        var value = NativeMethods.Value.DuckDBGetInt32(this);
+        return (T)(object)value;
+        //return Unsafe.As<int, T>(ref value);
+    }
 }
diff --git a/DuckDB.NET.Bindings/IDuckDBValueReader.cs b/DuckDB.NET.Bindings/IDuckDBValueReader.cs
@@ -0,0 +1,6 @@
+﻿namespace DuckDB.NET.Native;
+
+public interface IDuckDBValueReader
+{
+    T GetValue<T>();
+}
diff --git a/DuckDB.NET.Bindings/NativeMethods/NativeMethods.TableFunction.cs b/DuckDB.NET.Bindings/NativeMethods/NativeMethods.TableFunction.cs
@@ -7,10 +7,10 @@ public partial class NativeMethods
 {
     public static class TableFunction
     {
-        [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_create_tableFunction")]
+        [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_create_table_function")]
         public static extern IntPtr DuckDBCreateTableFunction();
 
-        [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_destroy_tableFunction")]
+        [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_destroy_table_function")]
         public static extern void DuckDBDestroyTableFunction(out IntPtr tableFunction);
 
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_table_function_set_name")]
@@ -25,17 +25,20 @@ public static class TableFunction
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_table_function_set_bind")]
         public static extern unsafe void DuckDBTableFunctionSetBind(IntPtr tableFunction, delegate* unmanaged[Cdecl]<IntPtr, void> bind);
 
-        //[DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_table_function_set_init")]
-        //public static extern unsafe void DuckDBTableFunctionSetInit(IntPtr tableFunction, duckdb_tableFunction_init_t init);
+        [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_table_function_set_init")]
+        public static extern unsafe void DuckDBTableFunctionSetInit(IntPtr tableFunction, delegate* unmanaged[Cdecl]<IntPtr, void> init);
 
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_table_function_set_function")]
-        public static extern unsafe void DuckDBTableFunctionSetFunction(IntPtr tableFunction, delegate* unmanaged[Cdecl]<IntPtr, IntPtr, IntPtr, void> callback);
+        public static extern unsafe void DuckDBTableFunctionSetFunction(IntPtr tableFunction, delegate* unmanaged[Cdecl]<IntPtr, IntPtr, void> callback);
 
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_register_table_function")]
         public static extern DuckDBState DuckDBRegisterTableFunction(DuckDBNativeConnection con, IntPtr tableFunction);
 
         #region TableFunctionBind
 
+        [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_bind_get_extra_info")]
+        public static extern unsafe IntPtr DuckDBBindGetExtraInfo(IntPtr info);
+
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_bind_add_result_column")]
         public static extern unsafe void DuckDBBindAddResultColumn(IntPtr info, SafeUnmanagedMemoryHandle name, DuckDBLogicalType type);
 
@@ -52,8 +55,11 @@ public static class TableFunction
 
         #region TableFunction
 
+        [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_function_get_extra_info")]
+        public static extern unsafe IntPtr DuckDBFunctionGetExtraInfo(IntPtr info);
+
         [DllImport(DuckDbLibrary, CallingConvention = CallingConvention.Cdecl, EntryPoint = "duckdb_function_get_bind_data")]
-        public static extern unsafe void* duckdb_function_get_bind_data(IntPtr info);
+        public static extern unsafe IntPtr DuckDBFunctionGetBindData(IntPtr info);
 
         #endregion
     }
diff --git a/DuckDB.NET.Data/DuckDBConnection.ScalarFunction.cs b/DuckDB.NET.Data/DuckDBConnection.ScalarFunction.cs
@@ -97,7 +97,7 @@ private unsafe void RegisterScalarMethod(string name, Action<IDuckDBDataReader[]
 
         if (!state.IsSuccess())
         {
-            throw new InvalidOperationException("Error registering user defined scalar function");
+            throw new InvalidOperationException($"Error registering user defined scalar function: {name}");
         }
     }
 
diff --git a/DuckDB.NET.Data/DuckDBConnection.TableFunction.cs b/DuckDB.NET.Data/DuckDBConnection.TableFunction.cs
@@ -1,15 +1,31 @@
-﻿using System;
-using System.Runtime.CompilerServices;
-using System.Runtime.InteropServices;
-using DuckDB.NET.Data.Extensions;
+﻿using DuckDB.NET.Data.Extensions;
 using DuckDB.NET.Data.Internal;
 using DuckDB.NET.Native;
+using System;
+using System.Collections;
+using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
+using System.Linq;
+using System.Runtime.CompilerServices;
+using System.Runtime.InteropServices;
+using DuckDB.NET.Data.Internal.Writer;
+using DuckDB.NET.Data.Writer;
 
 namespace DuckDB.NET.Data;
 
+public record ColumnInfo(string Name, Type Type)
+{
+}
+
+public record TableFunction(IReadOnlyList<ColumnInfo> Columns, IEnumerable Data)
+{
+}
+
 partial class DuckDBConnection
 {
-    public unsafe void RegisterTableFunction<T, TResult>(string name)
+#if NET8_0_OR_GREATER
+    [Experimental("DuckDBNET001")]
+    public unsafe void RegisterTableFunction<T>(string name, Func<IEnumerable<IDuckDBValueReader>, TableFunction> resultCallback, Action<object?, IDuckDBDataWriter[]> mapperCallback)
     {
         var function = NativeMethods.TableFunction.DuckDBCreateTableFunction();
         NativeMethods.TableFunction.DuckDBTableFunctionSetName(function, name.ToUnmanagedString());
@@ -19,20 +35,105 @@ public unsafe void RegisterTableFunction<T, TResult>(string name)
             NativeMethods.TableFunction.DuckDBTableFunctionAddParameter(function, logicalType);
         }
 
+        var tableFunctionInfo = new TableFunctionInfo(resultCallback, mapperCallback);
+
         NativeMethods.TableFunction.DuckDBTableFunctionSetBind(function, &Bind);
+        NativeMethods.TableFunction.DuckDBTableFunctionSetInit(function, &Init);
+        NativeMethods.TableFunction.DuckDBTableFunctionSetFunction(function, &TableFunction);
+        NativeMethods.TableFunction.DuckDBTableFunctionSetExtraInfo(function, tableFunctionInfo.ToHandle(), &DestroyExtraInfo);
+
+        var state = NativeMethods.TableFunction.DuckDBRegisterTableFunction(NativeConnection, function);
+
+        if (!state.IsSuccess())
+        {
+            throw new InvalidOperationException($"Error registering user defined table function: {name}");
+        }
+
+        NativeMethods.TableFunction.DuckDBDestroyTableFunction(out function);
     }
 
     [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])]
     public static unsafe void Bind(IntPtr info)
     {
-        var parameters = new object[NativeMethods.TableFunction.DuckDBBindGetParameterCount(info)];
+        var handle = GCHandle.FromIntPtr(NativeMethods.TableFunction.DuckDBBindGetExtraInfo(info));
+
+        if (handle.Target is not TableFunctionInfo functionInfo)
+        {
+            throw new InvalidOperationException("User defined table function bind failed. Bind extra info is null");
+        }
+
+        var parameters = new IDuckDBValueReader[NativeMethods.TableFunction.DuckDBBindGetParameterCount(info)];
 
         for (var i = 0; i < parameters.Length; i++)
         {
-            using var value = NativeMethods.TableFunction.DuckDBBindGetParameter(info, (ulong)i);
-            parameters[i] = NativeMethods.Value.DuckDBGetInt32(value);
+            var value = NativeMethods.TableFunction.DuckDBBindGetParameter(info, (ulong)i);
+            parameters[i] = value;
+        }
+
+        var tableFunctionData = functionInfo.Bind(parameters);
+
+        foreach (var parameter in parameters)
+        {
+            (parameter as IDisposable)?.Dispose();
+        }
+
+        foreach (var columnInfo in tableFunctionData.Columns)
+        {
+            using var logicalType = DuckDBTypeMap.GetLogicalType(columnInfo.Type);
+            NativeMethods.TableFunction.DuckDBBindAddResultColumn(info, columnInfo.Name.ToUnmanagedString(), logicalType);
+        }
+
+        var bindData = new TableFunctionBindData(tableFunctionData.Columns, tableFunctionData.Data.GetEnumerator());
+
+        NativeMethods.TableFunction.DuckDBBindSetBindData(info, bindData.ToHandle(), &DestroyExtraInfo);
+    }
+
+    [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])]
+    public static unsafe void Init(IntPtr info) { }
+
+    [UnmanagedCallersOnly(CallConvs = [typeof(CallConvCdecl)])]
+    public static unsafe void TableFunction(IntPtr info, IntPtr chunk)
+    {
+        var bindData = GCHandle.FromIntPtr(NativeMethods.TableFunction.DuckDBFunctionGetBindData(info));
+        var extraInfo = GCHandle.FromIntPtr(NativeMethods.TableFunction.DuckDBFunctionGetExtraInfo(info));
+
+        if (bindData.Target is not TableFunctionBindData tableFunctionBindData)
+        {
+            throw new InvalidOperationException("User defined table function failed. Function bind data is null");
+        }
+
+        if (extraInfo.Target is not TableFunctionInfo tableFunctionInfo)
+        {
+            throw new InvalidOperationException("User defined table function failed. Function extra info is null");
+        }
+
+        var dataChunk = new DuckDBDataChunk(chunk);
+
+        var writers = new VectorDataWriterBase[tableFunctionBindData.Columns.Count];
+        for (var columnIndex = 0; columnIndex < tableFunctionBindData.Columns.Count; columnIndex++)
+        {
+            var column = tableFunctionBindData.Columns[columnIndex];
+            var vector = NativeMethods.DataChunks.DuckDBDataChunkGetVector(dataChunk, columnIndex);
+
+            using var logicalType = DuckDBTypeMap.GetLogicalType(column.Type);
+            writers[columnIndex] = VectorDataWriterFactory.CreateWriter(vector, logicalType);
+        }
+
+        ulong size = 0;
+
+        for (; size < DuckDBGlobalData.VectorSize; size++)
+        {
+            if (tableFunctionBindData.DataEnumerator.MoveNext())
+            {
+                tableFunctionInfo.Mapper(tableFunctionBindData.DataEnumerator.Current, writers);
+            }
+            else
+            {
+                break;
+            }
         }
 
-        NativeMethods.TableFunction.DuckDBBindSetBindData(info, parameters.ToHandle(), &DestroyExtraInfo);
+        NativeMethods.DataChunks.DuckDBDataChunkSetSize(dataChunk, size);
     }
+#endif
 }
diff --git a/DuckDB.NET.Data/Internal/DuckDBTypeMap.cs b/DuckDB.NET.Data/Internal/DuckDBTypeMap.cs
@@ -79,13 +79,15 @@ public static DbType GetDbTypeForValue(object? value)
         return DbType.Object;
     }
 
-    public static DuckDBLogicalType GetLogicalType<T>()
+    public static DuckDBLogicalType GetLogicalType<T>() => GetLogicalType(typeof(T));
+
+    public static DuckDBLogicalType GetLogicalType(Type type)
     {
-        if (ClrToDuckDBTypeMap.TryGetValue(typeof(T), out var duckDBType))
+        if (ClrToDuckDBTypeMap.TryGetValue(type, out var duckDBType))
         {
             return NativeMethods.LogicalType.DuckDBCreateLogicalType(duckDBType);
         }
 
-        throw new InvalidOperationException($"Cannot map type {typeof(T).FullName} to DuckDBType.");
+        throw new InvalidOperationException($"Cannot map type {type.FullName} to DuckDBType.");
     }
 }
diff --git a/DuckDB.NET.Data/Internal/ScalarFunctionInfo.cs b/DuckDB.NET.Data/Internal/ScalarFunctionInfo.cs
@@ -1,8 +1,7 @@
-﻿using System;
-using System.Collections.Generic;
-using DuckDB.NET.Data.Internal.Reader;
+﻿using DuckDB.NET.Data.Internal.Reader;
 using DuckDB.NET.Data.Internal.Writer;
 using DuckDB.NET.Native;
+using System;
 
 namespace DuckDB.NET.Data.Internal;
 
diff --git a/DuckDB.NET.Data/Internal/TableFunctionInfo.cs b/DuckDB.NET.Data/Internal/TableFunctionInfo.cs
@@ -0,0 +1,24 @@
+﻿using DuckDB.NET.Data.Internal.Writer;
+using DuckDB.NET.Native;
+using System;
+using System.Collections;
+using System.Collections.Generic;
+
+namespace DuckDB.NET.Data.Internal;
+
+class TableFunctionInfo(Func<IEnumerable<IDuckDBValueReader>, TableFunction> bind, Action<object?, VectorDataWriterBase[]> mapper)
+{
+    public Func<IEnumerable<IDuckDBValueReader>, TableFunction> Bind { get; private set; } = bind;
+    public Action<object?, VectorDataWriterBase[]> Mapper { get; private set; } = mapper;
+}
+
+class TableFunctionBindData(IReadOnlyList<ColumnInfo> columns, IEnumerator dataEnumerator) : IDisposable
+{
+    public IReadOnlyList<ColumnInfo> Columns { get; } = columns;
+    public IEnumerator DataEnumerator { get; private set; } = dataEnumerator;
+
+    public void Dispose()
+    {
+        (DataEnumerator as IDisposable)?.Dispose();
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`		`-using System;`
`2`		`-using Microsoft.Win32.SafeHandles;`
	`1`	`+using Microsoft.Win32.SafeHandles;`
	`2`	`+using System;`
`3`	`3`
`4`	`4`	`namespace DuckDB.NET.Native;`
`5`	`5`
`@@ -89,7 +89,7 @@ protected override bool ReleaseHandle()`
`89`	`89`	`}`
`90`	`90`	`}`
`91`	`91`
`92`		`-public class DuckDBValue() : SafeHandleZeroOrMinusOneIsInvalid(true)`
	`92`	`+public class DuckDBValue() : SafeHandleZeroOrMinusOneIsInvalid(true), IDuckDBValueReader`
`93`	`93`	`{`
`94`	`94`	`private DuckDBValue[] childValues = [];`
`95`	`95`
`@@ -108,4 +108,11 @@ internal void SetChildValues(DuckDBValue[] values)`
`108`	`108`	`{`
`109`	`109`	`childValues = values;`
`110`	`110`	`}`
	`111`	`+`
	`112`	`+ public T GetValue<T>()`
	`113`	`+ {`
	`114`	`+ var value = NativeMethods.Value.DuckDBGetInt32(this);`
	`115`	`+ return (T)(object)value;`
	`116`	`+ //return Unsafe.As<int, T>(ref value);`
	`117`	`+ }`
`111`	`118`	`}`
Original file line number	Diff line number	Diff line change
`@@ -97,7 +97,7 @@ private unsafe void RegisterScalarMethod(string name, Action<IDuckDBDataReader[]`
`97`	`97`
`98`	`98`	`if (!state.IsSuccess())`
`99`	`99`	`{`
`100`		`- throw new InvalidOperationException("Error registering user defined scalar function");`
	`100`	`+ throw new InvalidOperationException($"Error registering user defined scalar function: {name}");`
`101`	`101`	`}`
`102`	`102`	`}`
`103`	`103`
Original file line number	Diff line number	Diff line change
`@@ -79,13 +79,15 @@ public static DbType GetDbTypeForValue(object? value)`
`79`	`79`	`return DbType.Object;`
`80`	`80`	`}`
`81`	`81`
`82`		`- public static DuckDBLogicalType GetLogicalType<T>()`
	`82`	`+ public static DuckDBLogicalType GetLogicalType<T>() => GetLogicalType(typeof(T));`
	`83`	`+`
	`84`	`+ public static DuckDBLogicalType GetLogicalType(Type type)`
`83`	`85`	`{`
`84`		`- if (ClrToDuckDBTypeMap.TryGetValue(typeof(T), out var duckDBType))`
	`86`	`+ if (ClrToDuckDBTypeMap.TryGetValue(type, out var duckDBType))`
`85`	`87`	`{`
`86`	`88`	`return NativeMethods.LogicalType.DuckDBCreateLogicalType(duckDBType);`
`87`	`89`	`}`
`88`	`90`
`89`		`- throw new InvalidOperationException($"Cannot map type {typeof(T).FullName} to DuckDBType.");`
	`91`	`+ throw new InvalidOperationException($"Cannot map type {type.FullName} to DuckDBType.");`
`90`	`92`	`}`
`91`	`93`	`}`