Add support for reading BitString

Giorgi · Nov 30, 2023 · eab8999 · eab8999
1 parent 294115b
commit eab8999
Show file tree

Hide file tree

Showing 2 changed files with 100 additions and 0 deletions.
diff --git a/DuckDB.NET.Data/Internal/Reader/StringVectorDataReader.cs b/DuckDB.NET.Data/Internal/Reader/StringVectorDataReader.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections;
 using System.IO;
 using System.Text;
 
@@ -14,6 +15,7 @@ internal override T GetValue<T>(ulong offset)
     {
         return DuckDBType switch
         {
+            DuckDBType.Bit => GetBitString<T>(offset),
             DuckDBType.Blob => (T)(object)GetStream(offset),
             DuckDBType.Varchar => (T)(object)GetString(offset),
             _ => base.GetValue<T>(offset)
@@ -24,12 +26,69 @@ internal override object GetValue(ulong offset, Type? targetType = null)
     {
         return DuckDBType switch
         {
+            DuckDBType.Bit => GetBitString(offset),
             DuckDBType.Blob => GetStream(offset),
             DuckDBType.Varchar => GetString(offset),
             _ => base.GetValue(offset, targetType)
         };
     }
 
+    private T GetBitString<T>(ulong offset)
+    {
+        if (typeof(T) == typeof(string))
+        {
+            return (T)(object)GetBitString(offset);
+        }
+
+        if (typeof(T) == typeof(BitArray))
+        {
+            return (T)(object)GetBitStringAsBitArray(offset);
+        }
+
+        return base.GetValue<T>(offset);
+    }
+
+    private string GetBitString(ulong offset)
+    {
+        var bitArray = GetBitStringAsBitArray(offset);
+
+        var output = new char[bitArray.Length];
+
+        for (var index = 0; index < bitArray.Count; index++)
+        {
+            output[index] = bitArray[index] ? '1' : '0';
+        }
+
+        return new string(output);
+    }
+
+    //Copied from https://github.com/duckdb/duckdb/blob/8a17511028d306561d88da9425f9e0e88dedd70c/src/common/types/bit.cpp#L63
+    private unsafe BitArray GetBitStringAsBitArray(ulong offset)
+    {
+        var bits = (DuckDBString*)DataPointer + offset;
+        var data = bits->Data;
+
+        var bitLength = (bits->Length - 1) * 8 - *data;
+
+        var output = new BitArray(bitLength);
+        var outputIndex = 0;
+
+        for (var bitIndex = *data; bitIndex < 8; bitIndex++)
+        {
+            output[outputIndex++] = (*(data + 1) & (1 << (7 - bitIndex))) > 0;
+        }
+
+        for (var byteIndex = 2; byteIndex < bits->Length; byteIndex++)
+        {
+            for (var bitIndex = 0; bitIndex < 8; bitIndex++)
+            {
+                output[outputIndex++] = (*(data + byteIndex) & (1 << (7 - bitIndex))) > 0;
+            }
+        }
+
+        return output;
+    }
+
     private unsafe string GetString(ulong offset)
     {
         var data = (DuckDBString*)DataPointer + offset;

diff --git a/DuckDB.NET.Test/DuckDBBitStringReaderTests.cs b/DuckDB.NET.Test/DuckDBBitStringReaderTests.cs
@@ -0,0 +1,41 @@
+using System.Collections;
+using System.Linq;
+using FluentAssertions;
+using Xunit;
+
+namespace DuckDB.NET.Test;
+
+public class DuckDBBitStringReaderTests : DuckDBTestBase
+{
+    public DuckDBBitStringReaderTests(DuckDBDatabaseFixture db) : base(db)
+    {
+    }
+
+    [Fact]
+    public void ReadBitString()
+    {
+        Command.CommandText = "SELECT bitstring('0101011', 12)";
+        var reader = Command.ExecuteReader();
+        reader.Read();
+
+        var value = reader.GetValue(0);
+        value.Should().Be("000000101011");
+
+        value = reader.GetFieldValue<string>(0);
+        value.Should().Be("000000101011");
+    }
+
+    [Fact]
+    public void ReadBitStringAsBitArray()
+    {
+        Command.CommandText = "SELECT bitstring('0101011', 12)";
+        var reader = Command.ExecuteReader();
+        reader.Read();
+
+        var expected = new BitArray(new bool[] { false, false, false, false, false, false, true, false, true, false, true, true });
+
+        var value = reader.GetFieldValue<BitArray>(0);
+
+        expected.Xor(value).OfType<bool>().All(b => !b).Should().BeTrue();
+    }
+}