Databricks client - Add support for array type (#274)

* Support array types * Increment package version * Make test more readable * Update documention
Energinet-DataHub · Nov 17, 2023 · 55e04ad · 55e04ad
1 parent e5d0229
commit 55e04ad
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 2 deletions.
diff --git a/source/Databricks/documents/documentation.md b/source/Databricks/documents/documentation.md
@@ -79,6 +79,24 @@ await foreach (var record in records)
     allSheldons.Add(new Person(record.name, record.date));
 ```
 
+A query can contain arrays. When used with Apache Arrow the array is encoded as object[]. If the format is JsonArray the array is encoded as Json string array.
+
+```c#
+var statement = DatabricksStatement.FromRawSql(
+            @"SELECT a, b FROM VALUES
+                ('one', array(0, 1)),
+                ('two', array(2, 3)) AS data(a, b);").Build();
+
+var result = client.ExecuteStatementAsync(statement, Format.ApacheArrow);
+var row = await result.FirstAsync();
+
+// Apache arrow
+var values = ((object[])row.b).OfType<int>();
+
+// JsonArray
+var values = JsonConvert.DeserializeObject<string[]>((string)row.b).Select(int.Parse);
+```
+
 #### Adhoc queries
 
 It's possible to create adhoc queries from `DatabricksStatement` class.

diff --git a/source/Databricks/documents/release-notes/release-notes.md b/source/Databricks/documents/release-notes/release-notes.md
@@ -1,5 +1,10 @@
 # Databricks Release Notes
 
+## Version 8.2.0
+
+- Add support for ListArray in Apache Arrow format
+- Update documentation with usage guide for Apache Arrow format and JsonArray format for arrays
+
 ## Version 8.1.0
 
 - Create objects from statements as an alternative to dynamic

diff --git a/source/Databricks/source/Jobs/Jobs.csproj b/source/Databricks/source/Jobs/Jobs.csproj
@@ -31,7 +31,7 @@ limitations under the License.
 
     <PropertyGroup>
         <PackageId>Energinet.DataHub.Core.Databricks.Jobs</PackageId>
-        <PackageVersion>8.1.0$(VersionSuffix)</PackageVersion>
+        <PackageVersion>8.2.0$(VersionSuffix)</PackageVersion>
         <Title>Databricks Jobs</Title>
         <Company>Energinet-DataHub</Company>
         <Authors>Energinet-DataHub</Authors>

diff --git a/...abricks/source/SqlStatementExecution.IntegrationTests/Client/DatabricksStatementsTests.cs b/...abricks/source/SqlStatementExecution.IntegrationTests/Client/DatabricksStatementsTests.cs
@@ -28,6 +28,25 @@ public DatabricksStatementsTests(DatabricksSqlWarehouseFixture sqlWarehouseFixtu
         _sqlWarehouseFixture = sqlWarehouseFixture;
     }
 
+    [Fact]
+    public async Task CanHandleArrayTypes()
+    {
+        // Arrange
+        var client = _sqlWarehouseFixture.CreateSqlStatementClient();
+        var statement = DatabricksStatement.FromRawSql(
+            @"SELECT a, b FROM VALUES
+                ('one', array(0, 1)),
+                ('two', array(2, 3)) AS data(a, b);").Build();
+
+        // Act
+        var result = client.ExecuteStatementAsync(statement, Format.ApacheArrow);
+        var row = await result.FirstAsync();
+        var values = ((object[])row.b).OfType<int>();
+
+        // Assert
+        values.Should().BeEquivalentTo(new[] { 0, 1 });
+    }
+
     [Fact]
     public async Task ExecuteStatement_FromRawSqlWithParameters_ShouldReturnRows()
     {

diff --git a/source/Databricks/source/SqlStatementExecution/Formats/IArrowArrayExtensions.cs b/source/Databricks/source/SqlStatementExecution/Formats/IArrowArrayExtensions.cs
@@ -38,6 +38,19 @@ internal static class IArrowArrayExtensions
             TimestampArray timestampArray => timestampArray.GetTimestamp(i),
             Decimal128Array decimal128Array => decimal128Array.GetValue(i),
             StringArray stringArray => stringArray.GetString(i),
+            ListArray listArray => ReadArray(listArray, i),
             _ => throw new NotSupportedException($"Unsupported data type {arrowArray}"),
         };
+
+    private static object? ReadArray(ListArray array, int i)
+    {
+        var objectArray = new object?[array.Length];
+        var offset = array.ValueOffsets[i];
+        for (var j = 0; j < array.Length; j++)
+        {
+            objectArray[j] = array.Values.GetValue(j + offset);
+        }
+
+        return objectArray;
+    }
 }
diff --git a/source/Databricks/source/SqlStatementExecution/SqlStatementExecution.csproj b/source/Databricks/source/SqlStatementExecution/SqlStatementExecution.csproj
@@ -30,7 +30,7 @@ limitations under the License.
 
   <PropertyGroup>
     <PackageId>Energinet.DataHub.Core.Databricks.SqlStatementExecution</PackageId>
-    <PackageVersion>8.1.0$(VersionSuffix)</PackageVersion>
+    <PackageVersion>8.2.0$(VersionSuffix)</PackageVersion>
     <Title>Databricks SQL Statement Execution</Title>
     <Company>Energinet-DataHub</Company>
     <Authors>Energinet-DataHub</Authors>