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