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-unsubscr...@arrow.apache.org

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

Reply via email to