Skip to content

Commit

Permalink
Fix null check, read rows from all chunks
Browse files Browse the repository at this point in the history
  • Loading branch information
@Giorgi
Giorgi committed Sep 13, 2023
1 parent c7cca70 commit 4ff6ae2
Show file tree
Hide file tree
Showing 3 changed files with 148 additions and 69 deletions.
147 changes: 85 additions & 62 deletions DuckDB.NET.Data/DuckDBDataReader.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,11 @@ public class DuckDBDataReader : DbDataReader
private readonly Dictionary<int, IntPtr> vectors = new();
private readonly Dictionary<int, IntPtr> vectorValidityMask = new();

private long chunkCount;
private int currentChunkIndex;
private int rowsReadFromCurrentChunk;
private long currentChunkRowCount;

internal DuckDBDataReader(DuckDbCommand command, List<DuckDBResult> queryResults, CommandBehavior behavior)
{
this.command = command;
Expand All @@ -40,41 +45,51 @@ internal DuckDBDataReader(DuckDbCommand command, List<DuckDBResult> queryResults

private void InitReaderData()
{
var chunkCount = NativeMethods.Types.DuckDBResultChunkCount(currentResult);
var chunk = NativeMethods.Types.DuckDBResultGetChunk(currentResult, 0);
currentRow = -1;

rowCount = NativeMethods.Query.DuckDBRowCount(currentResult);
fieldCount = (int)NativeMethods.Query.DuckDBColumnCount(currentResult);
chunkCount = NativeMethods.Types.DuckDBResultChunkCount(currentResult);

currentChunkIndex = 0;
rowsReadFromCurrentChunk = 0;

InitChunkData();

var columnCount = NativeMethods.DataChunks.DuckDBDataChunkGetColumnCount(chunk);
// recordsAffected = (int)NativeMethods.Query.DuckDBRowsChanged(currentResult);
}

private void InitChunkData()
{
var chunk = NativeMethods.Types.DuckDBResultGetChunk(currentResult, currentChunkIndex);

var size = NativeMethods.DataChunks.DuckDBDataChunkGetSize(chunk);
currentChunkRowCount = NativeMethods.DataChunks.DuckDBDataChunkGetSize(chunk);

for (int i = 0; i < columnCount; i++)
for (int i = 0; i < fieldCount; i++)
{
vectors[i] = NativeMethods.DataChunks.DuckDBDataChunkGetVector(chunk, i);
vectorValidityMask[i] = (NativeMethods.DataChunks.DuckDBVectorGetValidity(vectors[i]));
}
var vector = NativeMethods.DataChunks.DuckDBDataChunkGetVector(chunk, i);

currentRow = -1;
rowCount = NativeMethods.Query.DuckDBRowCount(currentResult);
fieldCount = (int)NativeMethods.Query.DuckDBColumnCount(currentResult);
recordsAffected = (int)NativeMethods.Query.DuckDBRowsChanged(currentResult);
vectors[i] = NativeMethods.DataChunks.DuckDBVectorGetData(vector);
vectorValidityMask[i] = NativeMethods.DataChunks.DuckDBVectorGetValidity(vector);
}
}

public override bool GetBoolean(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
return Marshal.ReadByte(data, currentRow * Marshal.SizeOf<byte>()) != 0;
var data = vectors[ordinal];
return Marshal.ReadByte(data, (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<byte>()) != 0;
}

public override byte GetByte(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
return Marshal.ReadByte(data, currentRow * Marshal.SizeOf<byte>());
var data = vectors[ordinal];
return Marshal.ReadByte(data, (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<byte>());
}

private sbyte GetSByte(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
return (sbyte)Marshal.ReadByte(data, currentRow * Marshal.SizeOf<byte>());
var data = vectors[ordinal];
return (sbyte)Marshal.ReadByte(data, (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<byte>());
}

public override long GetBytes(int ordinal, long dataOffset, byte[] buffer, int bufferOffset, int length)
Expand All @@ -99,23 +114,23 @@ public override string GetDataTypeName(int ordinal)

public override DateTime GetDateTime(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
var timestampStruct = Marshal.PtrToStructure<DuckDBTimestampStruct>(data + currentRow * Marshal.SizeOf<DuckDBTimestampStruct>());
var data = vectors[ordinal];
var timestampStruct = Marshal.PtrToStructure<DuckDBTimestampStruct>(data + (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<DuckDBTimestampStruct>());

return timestampStruct.ToDateTime();
}

private DuckDBDateOnly GetDateOnly(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
var date = Marshal.PtrToStructure<DuckDBDate>(data + currentRow * Marshal.SizeOf<DuckDBDate>());
var data = vectors[ordinal];
var date = Marshal.PtrToStructure<DuckDBDate>(data + (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<DuckDBDate>());
return NativeMethods.DateTime.DuckDBFromDate(date);
}

private DuckDBTimeOnly GetTimeOnly(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
var time = Marshal.PtrToStructure<DuckDBTime>(data + currentRow * Marshal.SizeOf<DuckDBTime>());
var data = vectors[ordinal];
var time = Marshal.PtrToStructure<DuckDBTime>(data + (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<DuckDBTime>());
return NativeMethods.DateTime.DuckDBFromTime(time);
}

Expand All @@ -124,7 +139,6 @@ public override decimal GetDecimal(int ordinal)
using (var logicalType = NativeMethods.Query.DuckDBColumnLogicalType(currentResult, ordinal))
{
var scale = NativeMethods.LogicalType.DuckDBDecimalScale(logicalType);
var width = NativeMethods.LogicalType.DuckDBDecimalWidth(logicalType);
var internalType = NativeMethods.LogicalType.DuckDBDecimalInternalType(logicalType);

decimal result = 0;
Expand All @@ -143,14 +157,14 @@ public override decimal GetDecimal(int ordinal)
result = decimal.Divide(GetInt64(ordinal), pow);
break;
case DuckDBType.DuckdbTypeHugeInt:
{
var hugeInt = GetBigInteger(ordinal);
{
var hugeInt = GetBigInteger(ordinal);

result = (decimal)BigInteger.DivRem(hugeInt, (BigInteger)pow, out var remainder);
result = (decimal)BigInteger.DivRem(hugeInt, (BigInteger)pow, out var remainder);

result += decimal.Divide((decimal)remainder, pow);
break;
}
result += decimal.Divide((decimal)remainder, pow);
break;
}
}

return result;
Expand All @@ -159,8 +173,8 @@ public override decimal GetDecimal(int ordinal)

public override double GetDouble(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
return Marshal.PtrToStructure<double>(data + currentRow * Marshal.SizeOf<double>());
var data = vectors[ordinal];
return Marshal.PtrToStructure<double>(data + (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<double>());
}

public override Type GetFieldType(int ordinal)
Expand Down Expand Up @@ -193,8 +207,8 @@ public override Type GetFieldType(int ordinal)

public override float GetFloat(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
return Marshal.PtrToStructure<float>(data + currentRow * Marshal.SizeOf<float>());
var data = vectors[ordinal];
return Marshal.PtrToStructure<float>(data + (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<float>());
}

public override Guid GetGuid(int ordinal)
Expand All @@ -204,44 +218,44 @@ public override Guid GetGuid(int ordinal)

public override short GetInt16(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
return Marshal.ReadInt16(data, currentRow * Marshal.SizeOf<short>());
var data = vectors[ordinal];
return Marshal.ReadInt16(data, (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<short>());
}

public override int GetInt32(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
return Marshal.ReadInt32(data, currentRow * Marshal.SizeOf<int>());
var data = vectors[ordinal];
return Marshal.ReadInt32(data, (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<int>());
}

public override long GetInt64(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
return Marshal.ReadInt64(data, currentRow * Marshal.SizeOf<long>());
var data = vectors[ordinal];
return Marshal.ReadInt64(data, (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<long>());
}

private ushort GetUInt16(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
return (ushort)Marshal.ReadInt32(data, currentRow * Marshal.SizeOf<ushort>());
var data = vectors[ordinal];
return (ushort)Marshal.ReadInt32(data, (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<ushort>());
}

private uint GetUInt32(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
return (uint)Marshal.ReadInt32(data, currentRow * Marshal.SizeOf<uint>());
var data = vectors[ordinal];
return (uint)Marshal.ReadInt32(data, (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<uint>());
}

private ulong GetUInt64(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
return (ulong)Marshal.ReadInt32(data, currentRow * Marshal.SizeOf<ulong>());
var data = vectors[ordinal];
return (ulong)Marshal.ReadInt32(data, (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<ulong>());
}

private BigInteger GetBigInteger(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
var hugeInt = Marshal.PtrToStructure<DuckDBHugeInt>(data + currentRow * Marshal.SizeOf<DuckDBHugeInt>());
var data = vectors[ordinal];
var hugeInt = Marshal.PtrToStructure<DuckDBHugeInt>(data + (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<DuckDBHugeInt>());

return hugeInt.ToBigInteger();
}
Expand All @@ -268,12 +282,7 @@ public override int GetOrdinal(string name)

public override string GetString(int ordinal)
{
if (IsDBNull(ordinal))
{
return null;
}

var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]) + currentRow * Marshal.SizeOf<DuckDBString>();
var data = vectors[ordinal] + (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<DuckDBString>();

var length = Marshal.ReadInt32(data);

Expand Down Expand Up @@ -323,8 +332,8 @@ public override object GetValue(int ordinal)

private DuckDBInterval GetDuckDBInterval(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]);
var interval = Marshal.PtrToStructure<DuckDBInterval>(data + currentRow * Marshal.SizeOf<DuckDBInterval>());
var data = vectors[ordinal];
var interval = Marshal.PtrToStructure<DuckDBInterval>(data + (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<DuckDBInterval>());
return interval;
}

Expand All @@ -340,7 +349,7 @@ public override int GetValues(object[] values)

public override Stream GetStream(int ordinal)
{
var data = NativeMethods.DataChunks.DuckDBVectorGetData(vectors[ordinal]) + currentRow * Marshal.SizeOf<DuckDBString>();
var data = vectors[ordinal] + (rowsReadFromCurrentChunk - 1) * Marshal.SizeOf<DuckDBString>();

var length = Marshal.ReadInt32(data);

Expand All @@ -353,10 +362,10 @@ public override Stream GetStream(int ordinal)

public override bool IsDBNull(int ordinal)
{
var validityMaskEntryIndex = currentRow / 64;
var validityBitIndex = currentRow % 64;
var validityMaskEntryIndex = (rowsReadFromCurrentChunk - 1) / 64;
var validityBitIndex = (rowsReadFromCurrentChunk - 1) % 64;

var validityMaskEntryPtr = vectorValidityMask[ordinal] + validityMaskEntryIndex;
var validityMaskEntryPtr = vectorValidityMask[ordinal] + validityMaskEntryIndex * Marshal.SizeOf<ulong>();
var validityBit = 1ul << validityBitIndex;

var isValid = (Marshal.PtrToStructure<ulong>(validityMaskEntryPtr) & validityBit) != 0;
Expand Down Expand Up @@ -392,7 +401,21 @@ public override bool NextResult()

public override bool Read()
{
return ++currentRow < rowCount;
var hasMoreRows = ++currentRow < rowCount;

if (!hasMoreRows) return false;

if (rowsReadFromCurrentChunk == currentChunkRowCount)
{
currentChunkIndex++;
rowsReadFromCurrentChunk = 0;

InitChunkData();
}

rowsReadFromCurrentChunk++;

return true;
}

public override int Depth { get; }
Expand Down
Loading

0 comments on commit 4ff6ae2

Please sign in to comment.