SNOW-968741: Arrow performance optimizations

Snowflake.Data.Tests/IntegrationTests/SFDbDataReaderIT.cs

@@ -15,8 +15,7 @@
 
 namespace Snowflake.Data.Tests.IntegrationTests
 {
-    // TODO: enable tests for Arrow
-    //[TestFixture(ResultFormat.ARROW)]
+    [TestFixture(ResultFormat.ARROW)]
     [TestFixture(ResultFormat.JSON)]
     class SFDbDataReaderIT : SFBaseTest
     {
@@ -571,7 +570,7 @@ public void TestGetTimestampLTZ()
         }
 
         [Test]
-        public void TestGetBoolean()
+        public void TestGetBoolean([Values]bool value)
         {
             using (var conn = CreateAndOpenConnection())
             {
@@ -585,7 +584,7 @@ public void TestGetBoolean()
                 var p1 = cmd.CreateParameter();
                 p1.ParameterName = "1";
                 p1.DbType = DbType.Boolean;
-                p1.Value = true;
+                p1.Value = value;
                 cmd.Parameters.Add(p1);
 
                 var count = cmd.ExecuteNonQuery();
@@ -597,7 +596,7 @@ public void TestGetBoolean()
                 ValidateResultFormat(reader);
 
                 Assert.IsTrue(reader.Read());
-                Assert.IsTrue(reader.GetBoolean(0));
+                Assert.AreEqual(value, reader.GetBoolean(0));
                 reader.Close();
 
                 CloseConnection(conn);

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;
@@ -141,7 +142,7 @@
         {
             var chunk = new ArrowResultChunk(_recordBatchOne);
 
-            Assert.AreEqual(0, chunk.ChunkIndex);
+            Assert.AreEqual(-1, chunk.ChunkIndex);
         }
 
         [Test]
@@ -156,10 +157,32 @@
         [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<ArrowResultChunk, SFDataType>
+            {
+                { 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]