Add estimated cardinality support for table functions (#324)

Author: jdaigle

DuckDB.NET.Bindings/NativeMethods/NativeMethods.TableFunction.cs

@@ -89,6 +89,11 @@ public static partial class TableFunction
         [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
         public static unsafe partial void DuckDBBindSetBindData(IntPtr info, IntPtr bindData, delegate* unmanaged[Cdecl]<IntPtr, void> destroy);
 
+        [SuppressGCTransition]
+        [LibraryImport(DuckDbLibrary, EntryPoint = "duckdb_bind_set_cardinality")]
+        [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]
+        public static unsafe partial void DuckDBBindSetCardinality(IntPtr info, ulong cardinality, [MarshalAs(UnmanagedType.I1)] bool isExact);
+
         // Maybe [SuppressGCTransition]: strdup error string — one small allocation
         [LibraryImport(DuckDbLibrary, EntryPoint = "duckdb_bind_set_error", StringMarshalling = StringMarshalling.Utf8)]
         [UnmanagedCallConv(CallConvs = [typeof(CallConvCdecl)])]

DuckDB.NET.Data/DuckDBConnection.TableFunction.cs

@@ -8,7 +8,9 @@ namespace DuckDB.NET.Data;
 
 public record ColumnInfo(string Name, Type Type);
 
-public record TableFunction(IReadOnlyList<ColumnInfo> Columns, IEnumerable Data);
+public record CardinalityHint(ulong Value, bool IsExact = false);
+
+public record TableFunction(IReadOnlyList<ColumnInfo> Columns, IEnumerable Data, CardinalityHint? Cardinality = null);
 
 partial class DuckDBConnection
 {
@@ -142,6 +144,11 @@ private static unsafe void Bind(IntPtr info)
                 NativeMethods.TableFunction.DuckDBBindAddResultColumn(info, columnInfo.Name, logicalType);
             }
 
+            if (tableFunctionData.Cardinality is { } cardinality)
+            {
+                NativeMethods.TableFunction.DuckDBBindSetCardinality(info, cardinality.Value, isExact: cardinality.IsExact);
+            }
+
             var connectionId = UdfExceptionStore.GetTableFunctionBindConnectionId(info);
             var bindData = new TableFunctionBindData(tableFunctionData.Columns, tableFunctionData.Data.GetEnumerator(), connectionId);
 

DuckDB.NET.Test/TableFunctionTests.cs

@@ -1,4 +1,5 @@
 using System.Globalization;
+using System.Text.Json;
 using System.Threading;
 
 namespace DuckDB.NET.Test;
@@ -422,4 +423,52 @@ public void TableFunctionMapError_PreservesInnerException()
         ex.InnerException!.Message.Should().Be("custom map error");
         ex.InnerException.Should().BeSameAs(originalException);
     }
+
+    [Theory]
+    [InlineData(100, false, "card_estimated")]
+    [InlineData(50, true, "card_exact")]
+    public void RegisterTableFunctionWithCardinality(int count, bool isExact, string funcName)
+    {
+        var expectedData = Enumerable.Range(0, count).ToList();
+
+        Connection.RegisterTableFunction(funcName, parameters =>
+        {
+            return new TableFunction(new List<ColumnInfo>
+            {
+                new("value", typeof(int)),
+            }, expectedData, Cardinality: new CardinalityHint((ulong)count, IsExact: isExact));
+        }, (item, writers, rowIndex) =>
+        {
+            writers[0].WriteValue((int)item, rowIndex);
+        });
+
+        var data = Connection.Query<int>($"SELECT * FROM {funcName}();").ToList();
+        data.Should().BeEquivalentTo(expectedData);
+    }
+
+    [Theory]
+    [InlineData(42, false, "plan_estimated")]
+    [InlineData(50, true, "plan_exact")]
+    public void RegisterTableFunctionWithCardinality_AppearsInQueryPlan(int cardinality, bool isExact, string funcName)
+    {
+        Connection.RegisterTableFunction(funcName, _ =>
+        {
+            return new TableFunction(new List<ColumnInfo>
+            {
+                new("value", typeof(int)),
+            }, Enumerable.Range(0, 50), Cardinality: new CardinalityHint((ulong)cardinality, IsExact: isExact));
+        }, (item, writers, rowIndex) =>
+        {
+            writers[0].WriteValue((int)item, rowIndex);
+        });
+
+        var plan = Connection.Query<(string, string)>($"EXPLAIN (FORMAT JSON) SELECT * FROM {funcName}();").First();
+        var json = JsonDocument.Parse(plan.Item2);
+        var reported = json.RootElement[0]
+            .GetProperty("extra_info")
+            .GetProperty("Estimated Cardinality")
+            .GetString();
+
+        reported.Should().Be(cardinality.ToString());
+    }
 }