Make GetOrdinal faster

Giorgi · Nov 22, 2024 · 1c96434 · 1c96434
1 parent a241693
commit 1c96434
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 9 deletions.
diff --git a/DuckDB.NET.Data/DuckDBDataReader.cs b/DuckDB.NET.Data/DuckDBDataReader.cs
@@ -1,11 +1,11 @@
-using System;
+using DuckDB.NET.Data.Internal.Reader;
+using DuckDB.NET.Native;
+using System;
 using System.Collections;
 using System.Collections.Generic;
 using System.Data;
 using System.Data.Common;
 using System.IO;
-using DuckDB.NET.Data.Internal.Reader;
-using DuckDB.NET.Native;
 
 namespace DuckDB.NET.Data;
 
@@ -29,6 +29,7 @@ public class DuckDBDataReader : DbDataReader
 
     private readonly IEnumerator<DuckDBResult> resultEnumerator;
     private VectorDataReaderBase[] vectorReaders = [];
+    private Dictionary<string, int> columnMapping = [];
 
     internal DuckDBDataReader(DuckDBCommand command, IEnumerable<DuckDBResult> queryResults, CommandBehavior behavior)
     {
@@ -48,6 +49,7 @@ private bool InitNextReader()
                 currentChunkIndex = 0;
                 currentResult = resultEnumerator.Current;
 
+                columnMapping = [];
                 fieldCount = (int)NativeMethods.Query.DuckDBColumnCount(ref currentResult);
                 streamingResult = NativeMethods.Types.DuckDBResultIsStreaming(currentResult) > 0;
 
@@ -80,7 +82,7 @@ private bool InitChunkData()
             {
                 vectorReaders = new VectorDataReaderBase[fieldCount];
             }
-            
+
             for (int index = 0; index < fieldCount; index++)
             {
                 var vector = NativeMethods.DataChunks.DuckDBDataChunkGetVector(currentChunk, index);
@@ -91,6 +93,17 @@ private bool InitChunkData()
                                                                             NativeMethods.Query.DuckDBColumnName(ref currentResult, index).ToManagedString(false));
             }
 
+            if (columnMapping.Count == 0)
+            {
+                for (var i = 0; i < vectorReaders.Length; i++)
+                {
+                    if (!columnMapping.ContainsKey(vectorReaders[i].ColumnName))
+                    {
+                        columnMapping.Add(vectorReaders[i].ColumnName, i);
+                    }
+                }
+            }
+
             return currentChunkRowCount > 0;
         }
     }
@@ -182,12 +195,9 @@ public override string GetName(int ordinal)
 
     public override int GetOrdinal(string name)
     {
-        for (var i = 0; i < fieldCount; i++)
+        if (columnMapping.TryGetValue(name, out var index))
         {
-            if (GetName(i) == name)
-            {
-                return i;
-            }
+            return index;
         }
 
         throw new DuckDBException($"Column with name {name} was not found.");

diff --git a/DuckDB.NET.Test/DuckDBDataReaderTests.cs b/DuckDB.NET.Test/DuckDBDataReaderTests.cs
@@ -29,6 +29,22 @@ public void GetOrdinalReturnsColumnIndex()
         reader.Invoking(dataReader => dataReader.GetOrdinal("Random")).Should().Throw<DuckDBException>();
     }
 
+    [Fact]
+    public void GetOrdinalRepeatedColumnReturnsFirstIndex()
+    {
+        Command.CommandText = "CREATE TABLE GetOrdinalTests (key INTEGER, value TEXT, State Boolean)";
+        Command.ExecuteNonQuery();
+
+        Command.CommandText = "select value, key, value from GetOrdinalTests";
+        Command.UseStreamingMode = true;
+        var reader = Command.ExecuteReader();
+
+        reader.GetOrdinal("key").Should().Be(1);
+        reader.GetOrdinal("value").Should().Be(0);
+
+        reader.Invoking(dataReader => dataReader.GetOrdinal("Random")).Should().Throw<DuckDBException>();
+    }
+
     [Fact]
     public void CloseConnectionClosesConnection()
     {