feat(csharp/src/Drivers/Databricks): Add option to enable using direct results for statements #2737

alexguo-db · 2025-04-23T00:35:20Z

Add option to set EnableDirectResults, which sends getDirectResults in the Thrift execute statement request
If getDirectResults is set in the request, then directResults is set on the response containing initial results (equivalent to the server calling GetOperationStatus, GetResultSetMetadata, FetchResults, and CloseOperation)
If directResults is set on the response, don't poll for the operation status
We already set getDirectResults on requests for metadata commands, just not the execute statement request

Tested E2E using dotnet test --filter CloudFetchE2ETest

[xUnit.net 00:00:00.11]   Starting:    Apache.Arrow.Adbc.Tests.Drivers.Databricks
[xUnit.net 00:01:27.27]   Finished:    Apache.Arrow.Adbc.Tests.Drivers.Databricks
  Apache.Arrow.Adbc.Tests.Drivers.Databricks test net8.0 succeeded (87.7s)

Test summary: total: 8, failed: 0, succeeded: 8, skipped: 0, duration: 87.7s
Build succeeded in 89.1s

… results for queries

CurtHagenlocher

Thanks! We need to mollify the nullability checker, and I've added a few suggestions for fairly superficial changes.

CurtHagenlocher · 2025-04-23T14:15:00Z

csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs

+            {
+                // The initial response has result data so we don't need to poll
+                metadata = DirectResults.ResultSetMetadata;
+            } else {


Suggested change

} else {

}

else

{

CurtHagenlocher · 2025-04-23T14:16:39Z

csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs

+
+                if (!string.IsNullOrEmpty(DirectResults.OperationStatus?.DisplayMessage))
+                {
+                    throw new HiveServer2Exception(DirectResults.OperationStatus.DisplayMessage)


Suggested change

throw new HiveServer2Exception(DirectResults.OperationStatus.DisplayMessage)

throw new HiveServer2Exception(DirectResults.OperationStatus!.DisplayMessage)

I think the null-forgiveness operator is required here because the nullability analysis doesn't know about string.IsNullOrEmpty.

CurtHagenlocher · 2025-04-23T15:18:36Z

csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs

            statement.QueryTimeout = QueryTimeoutSeconds;
        }

+        public TSparkDirectResults? DirectResults { get; set; }


Suggested change

public TSparkDirectResults? DirectResults { get; set; }

protected TSparkDirectResults? DirectResults { get; set; }

CurtHagenlocher · 2025-04-23T15:19:37Z

csharp/src/Drivers/Databricks/DatabricksConnection.cs

        private bool _applySSPWithQueries = false;
+        private bool _enableDirectResults = true;
+
+        internal static TSparkGetDirectResults defaultGetDirectResults = new(){


Suggested change

internal static TSparkGetDirectResults defaultGetDirectResults = new(){

internal static TSparkGetDirectResults defaultGetDirectResults = new()

{

CurtHagenlocher · 2025-04-23T15:21:46Z

csharp/src/Drivers/Databricks/DatabricksStatement.cs

+        /// Checks if direct results are available.
+        /// </summary>
+        /// <returns>True if direct results are available and contain result data, false otherwise.</returns>
+        public bool HasDirectResults()


Suggested change

public bool HasDirectResults()

public bool HasDirectResults => DirectResults?.ResultSet != null && DirectResults?.ResultSetMetadata != null;

(This is a good candidate to be a property instead of a method.)

jadewang-db · 2025-04-23T17:37:26Z

csharp/src/Drivers/Databricks/DatabricksReader.cs

            this.statement = statement;
            this.schema = schema;
            this.isLz4Compressed = isLz4Compressed;
+


same here, maybe start a new direct result reader. make sure doing one thing in one class

maybe not, after a 2nd look at the code

alexguo-db · 2025-04-23T20:32:41Z

@CurtHagenlocher
I believe the feedback is addressed now, but please wait for #2678 before merging

CurtHagenlocher · 2025-04-24T15:57:59Z

@CurtHagenlocher I believe the feedback is addressed now, but please wait for #2678 before merging

Now that #2678 has been checked in there's a merge conflict that will need to be resolved.

alexguo-db · 2025-04-24T22:51:46Z

csharp/src/Drivers/Databricks/CloudFetch/CloudFetchResultFetcher.cs

+                if (_statement.HasDirectResults && _statement.DirectResults?.ResultSet?.Results?.ResultLinks?.Count > 0)
+                {
+                    // Yield execution so the download queue doesn't get blocked before downloader is started
+                    await Task.Yield();


Without this line, there's a deadlock because the _downloadQueue is full, and there's no queue consumer since the downloader hasn't started yet

CurtHagenlocher · 2025-04-28T18:44:21Z

csharp/src/Drivers/Databricks/DatabricksConnection.cs

            return baseHandler;
        }

+        protected internal override bool AreResultsAvailableDirectly() => _enableDirectResults;


Another good place for a property:

Suggested change

protected internal override bool AreResultsAvailableDirectly() => _enableDirectResults;

protected internal override bool AreResultsAvailableDirectly => _enableDirectResults;

Sorry; I missed this in my previous review.

This is overriding AreResultsAvailableDirectly in HiveServer2Connection/SparkConnection, where it is a function. Should I update the field in both parent classes to be a property as well?

Yes, please!

CurtHagenlocher

Thanks again! Sorry; I missed another method that should be a property in my previous review. Looks like we're good after that change.

CurtHagenlocher · 2025-04-28T18:45:00Z

csharp/src/Drivers/Databricks/DatabricksStatement.cs

+            get { return _directResults; }
        }

+


nit: extra blank line

CurtHagenlocher

Thanks!

…t results for statements (apache#2737) - Add option to set EnableDirectResults, which sends getDirectResults in the Thrift execute statement request - If getDirectResults is set in the request, then directResults is set on the response containing initial results (equivalent to the server calling GetOperationStatus, GetResultSetMetadata, FetchResults, and CloseOperation) - If directResults is set on the response, don't poll for the operation status - We already set getDirectResults on requests for metadata commands, just not the execute statement request Tested E2E using `dotnet test --filter CloudFetchE2ETest` ``` [xUnit.net 00:00:00.11] Starting: Apache.Arrow.Adbc.Tests.Drivers.Databricks [xUnit.net 00:01:27.27] Finished: Apache.Arrow.Adbc.Tests.Drivers.Databricks Apache.Arrow.Adbc.Tests.Drivers.Databricks test net8.0 succeeded (87.7s) Test summary: total: 8, failed: 0, succeeded: 8, skipped: 0, duration: 87.7s Build succeeded in 89.1s ```

alexguo-db requested a review from CurtHagenlocher as a code owner April 23, 2025 00:35

alexguo-db changed the title ~~feat(csharp/src/Drivers/Databricks) Add option to enable using direct results for queries~~ feat(csharp/src/Drivers/Databricks) Add option to enable using direct results for statements Apr 23, 2025

github-actions bot added this to the ADBC Libraries 18 milestone Apr 23, 2025

alexguo-db added 2 commits April 22, 2025 17:36

feat(csharp/src/Drivers/Databricks) Add option to enable using direct…

62ffee3

… results for queries

Add tests

56cd239

alexguo-db force-pushed the alex-guo_data/alexguo-support-direct-results branch from c076812 to 56cd239 Compare April 23, 2025 00:36

alexguo-db changed the title ~~feat(csharp/src/Drivers/Databricks) Add option to enable using direct results for statements~~ feat(csharp/src/Drivers/Databricks): Add option to enable using direct results for statements Apr 23, 2025

CurtHagenlocher requested changes Apr 23, 2025

View reviewed changes

alexguo-db added 2 commits April 23, 2025 10:09

Address feedback

ecf62cf

Reset direct results before query execution

d05fd34

jadewang-db reviewed Apr 23, 2025

View reviewed changes

alexguo-db requested a review from CurtHagenlocher April 23, 2025 20:32

Merge branch 'main' into alexguo-support-direct-results

3a99e25

alexguo-db commented Apr 24, 2025

View reviewed changes

CurtHagenlocher reviewed Apr 28, 2025

View reviewed changes

CurtHagenlocher requested changes Apr 28, 2025

View reviewed changes

Fix nit

6831798

alexguo-db requested a review from CurtHagenlocher April 28, 2025 20:50

Convert AreResultsAvailableDirectly to a property

b5ad54c

CurtHagenlocher approved these changes Apr 28, 2025

View reviewed changes

CurtHagenlocher merged commit 1038658 into apache:main Apr 28, 2025
6 checks passed

	throw new HiveServer2Exception(DirectResults.OperationStatus.DisplayMessage)
	throw new HiveServer2Exception(DirectResults.OperationStatus!.DisplayMessage)

	public TSparkDirectResults? DirectResults { get; set; }
	protected TSparkDirectResults? DirectResults { get; set; }

	internal static TSparkGetDirectResults defaultGetDirectResults = new(){
	internal static TSparkGetDirectResults defaultGetDirectResults = new()
	{

	public bool HasDirectResults()
	public bool HasDirectResults => DirectResults?.ResultSet != null && DirectResults?.ResultSetMetadata != null;

	protected internal override bool AreResultsAvailableDirectly() => _enableDirectResults;
	protected internal override bool AreResultsAvailableDirectly => _enableDirectResults;

feat(csharp/src/Drivers/Databricks): Add option to enable using direct results for statements #2737

feat(csharp/src/Drivers/Databricks): Add option to enable using direct results for statements #2737

Uh oh!

Conversation

alexguo-db commented Apr 23, 2025

Uh oh!

CurtHagenlocher left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alexguo-db commented Apr 23, 2025

Uh oh!

CurtHagenlocher commented Apr 24, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

CurtHagenlocher left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

CurtHagenlocher left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants