birschick-bq commented on code in PR #2847:
URL: https://github.com/apache/arrow-adbc/pull/2847#discussion_r2130463729
##########
csharp/src/Drivers/Apache/Hive2/HiveServer2Reader.cs:
##########
@@ -76,50 +78,54 @@ public HiveServer2Reader(
HiveServer2Statement statement,
Schema schema,
DataTypeConversion dataTypeConversion,
- bool enableBatchSizeStopCondition = true)
+ bool enableBatchSizeStopCondition = true) : base(statement)
{
_statement = statement;
Schema = schema;
_dataTypeConversion = dataTypeConversion;
_enableBatchSizeStopCondition = enableBatchSizeStopCondition;
}
- public Schema Schema { get; }
+ public override Schema Schema { get; }
- public async ValueTask<RecordBatch?>
ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)
+ public override async ValueTask<RecordBatch?>
ReadNextRecordBatchAsync(CancellationToken cancellationToken = default)
{
- // All records have been exhausted
- if (_statement == null)
+ return await TraceActivity(async activity =>
{
- return null;
- }
+ // All records have been exhausted
+ if (_hasNoMoreData)
+ {
+ return null;
+ }
+ try
+ {
+ // Await the fetch response
+ TFetchResultsResp response = await FetchNext(_statement,
cancellationToken);
+ ApacheUtility.HandleThriftResponse(response.Status,
HiveServer2Connection.GetResponseHandlers(activity));
- try
- {
- // Await the fetch response
- TFetchResultsResp response = await FetchNext(_statement,
cancellationToken);
+ int columnCount = GetColumnCount(response.Results);
+ int rowCount = GetRowCount(response.Results, columnCount);
+ if ((_enableBatchSizeStopCondition && _statement.BatchSize
> 0 && rowCount < _statement.BatchSize) || rowCount == 0)
+ {
+ // This is the last batch
+ _hasNoMoreData = true;
+ }
- int columnCount = GetColumnCount(response.Results);
- int rowCount = GetRowCount(response.Results, columnCount);
- if ((_enableBatchSizeStopCondition && _statement.BatchSize > 0
&& rowCount < _statement.BatchSize) || rowCount == 0)
+ activity?.AddTag(TagOptions.Db.Response.ReturnedRows,
rowCount);
Review Comment:
Converted the AddTag to AddEvent for batch reading.
--
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]
