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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
