Reduce allocations when materializing MAP values

DuckDB.NET.Data/DataChunk/Reader/MapVectorDataReader.cs

@@ -1,9 +1,15 @@
-namespace DuckDB.NET.Data.DataChunk.Reader;
+using System.Collections.Concurrent;
+
+namespace DuckDB.NET.Data.DataChunk.Reader;
 
 internal sealed class MapVectorDataReader : VectorDataReaderBase
 {
+    private static readonly ConcurrentDictionary<Type, IMapMaterializer> Materializers = new();
+
     private readonly VectorDataReaderBase keyReader;
     private readonly VectorDataReaderBase valueReader;
+    private Type? cachedTargetType;
+    private IMapMaterializer? cachedMaterializer;
 
     internal unsafe MapVectorDataReader(IntPtr vector, void* dataPointer, ulong* validityMaskPointer, DuckDBType columnType, DuckDBLogicalType logicalColumnType, string columnName) 
                     : base(dataPointer, validityMaskPointer, columnType, columnName)
@@ -37,17 +43,22 @@ internal override unsafe object GetValue(ulong offset, Type targetType)
             return base.GetValue(offset, targetType);
         }
 
+        var listData = (DuckDBListEntry*)DataPointer + offset;
+        var materializer = GetMaterializer(targetType);
+
+        if (materializer != null)
+        {
+            return materializer.Materialize(keyReader, valueReader, listData->Offset, listData->Length, ColumnName);
+        }
+
         if (Activator.CreateInstance(targetType) is not IDictionary instance)
         {
             throw new InvalidOperationException($"Cannot read Map column {ColumnName} in a non-dictionary type");
         }
 
         var arguments = targetType.GetGenericArguments();
-
         var allowsNullValues = arguments.Length == 2 && arguments[1].AllowsNullValue(out _, out _);
 
-        var listData = (DuckDBListEntry*)DataPointer + offset;
-
         for (ulong i = 0; i < listData->Length; i++)
         {
             var childOffset = i + listData->Offset;
@@ -68,6 +79,43 @@ internal override unsafe object GetValue(ulong offset, Type targetType)
         return instance;
     }
 
+    private IMapMaterializer? GetMaterializer(Type targetType)
+    {
+        if (targetType == cachedTargetType)
+        {
+            return cachedMaterializer;
+        }
+
+        cachedTargetType = targetType;
+        cachedMaterializer = null;
+
+        if (!targetType.IsGenericType || targetType.GetGenericTypeDefinition() != typeof(Dictionary<,>))
+        {
+            return null;
+        }
+
+        var arguments = targetType.GetGenericArguments();
+        if (!CanReadDirectly(keyReader, arguments[0]) || !CanReadDirectly(valueReader, arguments[1]))
+        {
+            return null;
+        }
+
+        cachedMaterializer = Materializers.GetOrAdd(targetType, static type => CreateMaterializer(type));
+        return cachedMaterializer;
+    }
+
+    private static bool CanReadDirectly(VectorDataReaderBase reader, Type targetType)
+    {
+        var valueType = Nullable.GetUnderlyingType(targetType) ?? targetType;
+        return valueType == reader.ClrType || valueType == reader.ProviderSpecificClrType;
+    }
+
+    private static IMapMaterializer CreateMaterializer(Type targetType)
+    {
+        var materializerType = typeof(MapMaterializer<,>).MakeGenericType(targetType.GetGenericArguments());
+        return (IMapMaterializer)Activator.CreateInstance(materializerType)!;
+    }
+
     internal override void Reset(IntPtr vector)
     {
         base.Reset(vector);
@@ -84,4 +132,49 @@ public override void Dispose()
         valueReader.Dispose();
         base.Dispose();
     }
-}
+
+    private interface IMapMaterializer
+    {
+        object Materialize(VectorDataReaderBase keys, VectorDataReaderBase values, ulong offset, ulong length, string columnName);
+    }
+
+    private sealed class MapMaterializer<TKey, TValue> : IMapMaterializer where TKey : notnull
+    {
+        private static readonly bool AllowsNullValues = typeof(TValue).AllowsNullValue(out _, out _);
+
+        public object Materialize(VectorDataReaderBase keys, VectorDataReaderBase values, ulong offset, ulong length, string columnName)
+        {
+            var result = new Dictionary<TKey, TValue>(checked((int)length));
+
+            for (ulong i = 0; i < length; i++)
+            {
+                var childOffset = offset + i;
+                var key = Read<TKey>(keys, childOffset);
+
+                if (values.IsValid(childOffset))
+                {
+                    result.Add(key, Read<TValue>(values, childOffset));
+                }
+                else if (AllowsNullValues)
+                {
+                    result.Add(key, default!);
+                }
+                else
+                {
+                    throw new InvalidCastException($"The Map in column {columnName} contains null value but dictionary does not allow null values");
+                }
+            }
+
+            return result;
+        }
+
+        private static T Read<T>(VectorDataReaderBase reader, ulong offset)
+        {
+            // Object-valued readers need their natural non-generic path. Concrete value types use
+            // the generic path so they can flow into Dictionary<TKey, TValue> without boxing.
+            return typeof(T) == typeof(object)
+                ? (T)reader.GetValue(offset)
+                : reader.GetValueStrict<T>(offset);
+        }
+    }
+}

DuckDB.NET.Test/DuckDBDataReaderMapTests.cs

@@ -53,6 +53,32 @@ public void ReadMapStronglyTyped()
         value.Should().BeEquivalentTo(expectation);
     }
 
+    [Fact]
+    public void ReadMapWithWiderValueTypeUsesCompatibleFallback()
+    {
+        Command.CommandText = "SELECT MAP { 'key1': 1, 'key2': 5, 'key3': 7 }";
+        var reader = Command.ExecuteReader();
+
+        reader.Read();
+        var value = reader.GetFieldValue<Dictionary<string, object>>(0);
+
+        var expectation = new Dictionary<string, object>() { { "key1", 1 }, { "key2", 5 }, { "key3", 7 } };
+        value.Should().BeEquivalentTo(expectation);
+    }
+
+    [Fact]
+    public void ReadMapIntoDerivedDictionaryUsesCompatibleFallback()
+    {
+        Command.CommandText = "SELECT MAP { 'key1': 1, 'key2': 5, 'key3': 7 }";
+        var reader = Command.ExecuteReader();
+
+        reader.Read();
+        var value = reader.GetFieldValue<CustomDictionary<string, int>>(0);
+
+        var expectation = new Dictionary<string, int>() { { "key1", 1 }, { "key2", 5 }, { "key3", 7 } };
+        value.Should().BeEquivalentTo(expectation);
+    }
+
     [Fact]
     public void ReadMapWithNullInNullableDictionary()
     {
@@ -147,4 +173,8 @@ public int GetHashCode(List<string> obj)
             return string.Join(",", obj).GetHashCode();
         }
     }
-}
+
+    private sealed class CustomDictionary<TKey, TValue> : Dictionary<TKey, TValue> where TKey : notnull
+    {
+    }
+}