Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNOW-968741: Arrow performance optimizations #831

Merged
merged 13 commits into from
Jan 3, 2024
Merged

SNOW-968741: Arrow performance optimizations #831

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
1ea3f9a
SNOW-968741 PoC - arrow performance
sfc-gh-dstempniak Nov 19, 2023
a14d921
Merge branch 'master' into SNOW-968741_arrow_perf
sfc-gh-dstempniak Dec 5, 2023
493e4ff
SNOW-968741 disable stopwatch
sfc-gh-dstempniak Dec 5, 2023
012483c
Optimizations of ExtractCell, fix tests
sfc-gh-pbulawa Dec 12, 2023
bf0be28
Small optimization of extracting time
sfc-gh-pbulawa Dec 12, 2023
6e5f9ee
Cleanup
sfc-gh-pbulawa Dec 13, 2023
29d6338
Remove todo
sfc-gh-pbulawa Dec 13, 2023
d62389c
Cleanup
sfc-gh-pbulawa Dec 13, 2023
d2e075b
Further optimizations
sfc-gh-pbulawa Dec 18, 2023
0ca3dad
Upgrade Apache Arrow to 14.0.2
sfc-gh-pbulawa Dec 21, 2023
6f84c1b
Merge remote-tracking branch 'origin/master' into SNOW-968741_arrow_perf
sfc-gh-pbulawa Jan 2, 2024
0499a4e
Changed according to comments, added tests
sfc-gh-pbulawa Jan 2, 2024
34332d4
Fix test
sfc-gh-pbulawa Jan 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 4 additions & 5 deletions Snowflake.Data.Tests/IntegrationTests/SFDbDataReaderIT.cs
View file
Open in desktop
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
Open in desktop
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
Loading