Skip to content

Commit

Permalink
SNOW-968741: Arrow performance optimizations (#831)
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-pbulawa
sfc-gh-pbulawa authored Jan 3, 2024
1 parent 1683d66 commit fa1b879
Show file tree
Hide file tree
Showing 9 changed files with 280 additions and 101 deletions.
9 changes: 4 additions & 5 deletions Snowflake.Data.Tests/IntegrationTests/SFDbDataReaderIT.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
{
Expand Down Expand Up @@ -571,7 +570,7 @@ public void TestGetTimestampLTZ()
}

[Test]
public void TestGetBoolean()
public void TestGetBoolean([Values]bool value)
{
using (var conn = CreateAndOpenConnection())
{
Expand All @@ -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();
Expand All @@ -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);
Expand Down
35 changes: 30 additions & 5 deletions Snowflake.Data.Tests/UnitTests/ArrowResultChunkTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using Apache.Arrow;
using Apache.Arrow.Types;
Expand Down Expand Up @@ -141,7 +142,7 @@ public void TestGetChunkIndexReturnsFirstChunk()
{
var chunk = new ArrowResultChunk(_recordBatchOne);

Assert.AreEqual(0, chunk.ChunkIndex);
Assert.AreEqual(-1, chunk.ChunkIndex);
}

[Test]
Expand All @@ -156,10 +157,34 @@ 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<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 pair in cases)
{
var chunk = pair.Key;
var type = pair.Value;
chunk.Next();
Assert.AreEqual(DBNull.Value, chunk.ExtractCell(0, type, 0), $"Expected DBNull.Value for SFDataType: {type}");
}
}

[Test]
Expand Down
Loading

0 comments on commit fa1b879

Please sign in to comment.