From 0499a4e600f13659be311a8f49e8b91b583a2463 Mon Sep 17 00:00:00 2001 From: Piotr Bulawa Date: Tue, 2 Jan 2024 16:24:32 +0100 Subject: [PATCH] Changed according to comments, added tests --- .../UnitTests/ArrowResultChunkTest.cs | 31 ++++++++++++++-- Snowflake.Data/Core/ArrowResultSet.cs | 37 +++++++++++-------- Snowflake.Data/Core/BaseResultChunk.cs | 2 - 3 files changed, 48 insertions(+), 22 deletions(-) diff --git a/Snowflake.Data.Tests/UnitTests/ArrowResultChunkTest.cs b/Snowflake.Data.Tests/UnitTests/ArrowResultChunkTest.cs index a51c03df8..4505feff5 100755 --- a/Snowflake.Data.Tests/UnitTests/ArrowResultChunkTest.cs +++ b/Snowflake.Data.Tests/UnitTests/ArrowResultChunkTest.cs @@ -4,6 +4,7 @@ using System; using System.Collections; +using System.Collections.Generic; using System.Linq; using Apache.Arrow; using Apache.Arrow.Types; @@ -156,10 +157,32 @@ public void TestUnusedExtractCellThrowsNotSupportedException() [Test] public void TestExtractCellReturnsNull() { - var chunk = new ArrowResultChunk(RecordBatchWithNullValue); - chunk.Next(); - - Assert.AreEqual(DBNull.Value, chunk.ExtractCell(0, SFDataType.FIXED, 0)); + var cases = new Dictionary + { + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Int8", false, col => col.Int8(array => array.AppendNull())).Build()), SFDataType.FIXED }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Int16", false, col => col.Int16(array => array.AppendNull())).Build()), SFDataType.FIXED }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Int32", false, col => col.Int32(array => array.AppendNull())).Build()), SFDataType.FIXED }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Int64", false, col => col.Int64(array => array.AppendNull())).Build()), SFDataType.FIXED }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Decimal128", false, col => col.Decimal128(new Decimal128Type(0, 0), array => array.AppendNull())).Build()), SFDataType.FIXED }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Boolean", false, col => col.Boolean(array => array.AppendNull())).Build()), SFDataType.BOOLEAN }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Real", false, col => col.Double(array => array.AppendNull())).Build()), SFDataType.REAL }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Text", false, col => col.String(array => array.AppendNull())).Build()), SFDataType.TEXT }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Array", false, col => col.String(array => array.AppendNull())).Build()), SFDataType.ARRAY }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Variant", false, col => col.String(array => array.AppendNull())).Build()), SFDataType.VARIANT }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Object", false, col => col.String(array => array.AppendNull())).Build()), SFDataType.OBJECT }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Binary", false, col => col.Binary(array => array.AppendNull())).Build()), SFDataType.BINARY }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Date", false, col => col.Date32(array => array.AppendNull())).Build()), SFDataType.DATE }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Time", false, col => col.Int32(array => array.AppendNull())).Build()), SFDataType.TIME }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Timestamp_TZ", false, col => col.Int32(array => array.AppendNull())).Build()), SFDataType.TIMESTAMP_TZ }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Timestamp_LTZ", false, col => col.Int32(array => array.AppendNull())).Build()), SFDataType.TIMESTAMP_LTZ }, + { new ArrowResultChunk(new RecordBatch.Builder().Append("Col_Timestamp_NTZ", false, col => col.Int32(array => array.AppendNull())).Build()), SFDataType.TIMESTAMP_NTZ }, + }; + + foreach (var (chunk, type) in cases) + { + chunk.Next(); + Assert.AreEqual(DBNull.Value, chunk.ExtractCell(0, type, 0), $"Expected DBNull.Value for SFDataType: {type}"); + } } [Test] diff --git a/Snowflake.Data/Core/ArrowResultSet.cs b/Snowflake.Data/Core/ArrowResultSet.cs index bcfead3f6..31e0eccca 100755 --- a/Snowflake.Data/Core/ArrowResultSet.cs +++ b/Snowflake.Data/Core/ArrowResultSet.cs @@ -45,30 +45,35 @@ public ArrowResultSet(QueryExecResponseData responseData, SFStatement sfStatemen queryId = responseData.queryId; - if (responseData.rowsetBase64.Length > 0) + ReadChunk(responseData); + } + catch(Exception ex) + { + s_logger.Error("Result set error queryId="+responseData.queryId, ex); + throw; + } + } + + private void ReadChunk(QueryExecResponseData responseData) + { + if (responseData.rowsetBase64.Length > 0) + { + using (var stream = new MemoryStream(Convert.FromBase64String(responseData.rowsetBase64))) { - using (var stream = new MemoryStream(Convert.FromBase64String(responseData.rowsetBase64))) + using (var reader = new ArrowStreamReader(stream)) { - using (var reader = new ArrowStreamReader(stream)) + var recordBatch = reader.ReadNextRecordBatch(); + _currentChunk = new ArrowResultChunk(recordBatch); + while ((recordBatch = reader.ReadNextRecordBatch()) != null) { - var recordBatch = reader.ReadNextRecordBatch(); - _currentChunk = new ArrowResultChunk(recordBatch); - while ((recordBatch = reader.ReadNextRecordBatch()) != null) - { - ((ArrowResultChunk)_currentChunk).AddRecordBatch(recordBatch); - } + ((ArrowResultChunk)_currentChunk).AddRecordBatch(recordBatch); } } } - else - { - _currentChunk = new ArrowResultChunk(columnCount); - } } - catch(Exception ex) + else { - s_logger.Error("Result set error queryId="+responseData.queryId, ex); - throw; + _currentChunk = new ArrowResultChunk(columnCount); } } diff --git a/Snowflake.Data/Core/BaseResultChunk.cs b/Snowflake.Data/Core/BaseResultChunk.cs index 40073446c..37e8fa114 100755 --- a/Snowflake.Data/Core/BaseResultChunk.cs +++ b/Snowflake.Data/Core/BaseResultChunk.cs @@ -42,10 +42,8 @@ internal virtual void Reset(ExecResponseChunk chunkInfo, int chunkIndex) RowCount = chunkInfo.rowCount; Url = chunkInfo.url; ChunkIndex = chunkIndex; - CompressedSize = chunkInfo.compressedSize; UncompressedSize = chunkInfo.uncompressedSize; - } internal virtual void ResetForRetry()