Re: [PR] feat(csharp/src/Drivers/Databricks): Add option to enable using direct results for statements [arrow-adbc]

["CurtHagenlocher (via GitHub)" <gi...@apache.org>]

CurtHagenlocher commented on code in PR #2737:
URL: https://github.com/apache/arrow-adbc/pull/2737#discussion_r2056167393


##########
csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs:
##########
@@ -104,12 +106,19 @@ private async Task<QueryResult> ExecuteQueryAsyncInternal(CancellationToken canc
             // OR
             // take QueryTimeoutSeconds (but this could be restricting)
             await ExecuteStatementAsync(cancellationToken); // --> get QueryTimeout +
-            await HiveServer2Connection.PollForResponseAsync(OperationHandle!, Connection.Client, PollTimeMilliseconds, cancellationToken); // + poll, up to QueryTimeout
-            TGetResultSetMetadataResp response = await HiveServer2Connection.GetResultSetMetadataAsync(OperationHandle!, Connection.Client, cancellationToken);
-            Schema schema = Connection.SchemaParser.GetArrowSchema(response.Schema, Connection.DataTypeConversion);
 
+            TGetResultSetMetadataResp metadata;
+            if (DirectResults?.OperationStatus?.OperationState == TOperationState.FINISHED_STATE)
+            {
+                // The initial response has result data so we don't need to poll
+                metadata = DirectResults.ResultSetMetadata;
+            } else {

Review Comment:
   ```suggestion
               }
               else
               {
   ```



##########
csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs:
##########
@@ -54,6 +54,8 @@ protected virtual void SetStatementProperties(TExecuteStatementReq statement)
             statement.QueryTimeout = QueryTimeoutSeconds;
         }
 
+        public TSparkDirectResults? DirectResults { get; set; }

Review Comment:
   ```suggestion
           protected TSparkDirectResults? DirectResults { get; set; }
   ```



##########
csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs:
##########
@@ -257,6 +266,19 @@ protected async Task ExecuteStatementAsync(CancellationToken cancellationToken =
                     .SetNativeError(executeResponse.Status.ErrorCode);
             }
             OperationHandle = executeResponse.OperationHandle;
+
+            // Capture direct results if they're available
+            if (executeResponse.DirectResults != null)
+            {
+                DirectResults = executeResponse.DirectResults;
+
+                if (!string.IsNullOrEmpty(DirectResults.OperationStatus?.DisplayMessage))
+                {
+                    throw new HiveServer2Exception(DirectResults.OperationStatus.DisplayMessage)

Review Comment:
   ```suggestion
                       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`.



##########
csharp/src/Drivers/Databricks/DatabricksStatement.cs:
##########
@@ -48,6 +48,20 @@ protected override void SetStatementProperties(TExecuteStatementReq statement)
             statement.CanDownloadResult = useCloudFetch;
             statement.CanDecompressLZ4Result = canDecompressLz4;
             statement.MaxBytesPerFile = maxBytesPerFile;
+
+            if (Connection.AreResultsAvailableDirectly())
+            {
+                statement.GetDirectResults = DatabricksConnection.defaultGetDirectResults;
+            }
+        }
+
+        /// <summary>
+        /// Checks if direct results are available.
+        /// </summary>
+        /// <returns>True if direct results are available and contain result data, false otherwise.</returns>
+        public bool HasDirectResults()

Review Comment:
   ```suggestion
           public bool HasDirectResults => DirectResults?.ResultSet != null && DirectResults?.ResultSetMetadata != null;
   ```
   (This is a good candidate to be a property instead of a method.)



##########
csharp/src/Drivers/Databricks/DatabricksConnection.cs:
##########
@@ -33,6 +33,12 @@ namespace Apache.Arrow.Adbc.Drivers.Databricks
     internal class DatabricksConnection : SparkHttpConnection
     {
         private bool _applySSPWithQueries = false;
+        private bool _enableDirectResults = true;
+
+        internal static TSparkGetDirectResults defaultGetDirectResults = new(){

Review Comment:
   ```suggestion
           internal static TSparkGetDirectResults defaultGetDirectResults = new()
           {
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org