toddmeng-db commented on code in PR #3241:
URL: https://github.com/apache/arrow-adbc/pull/3241#discussion_r2261888507
##########
csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs:
##########
@@ -739,81 +750,91 @@ protected virtual async Task<QueryResult>
GetColumnsExtendedAsync(CancellationTo
// For FK lookup, we need to pass in the current
catalog/schema/table as the foreign table
var fkResult = await
GetCrossReferenceAsForeignTableAsync(cancellationToken);
- // 2. Read all batches into memory
- List<RecordBatch> columnsBatches;
- int totalRows;
- Schema columnsSchema;
- StringArray? columnNames = null;
- int colNameIndex = -1;
-
- // Extract column data
- using (var stream = columnsResult.Stream)
+ try
{
- colNameIndex = stream.Schema.GetFieldIndex("COLUMN_NAME");
- if (colNameIndex < 0)
+ // 2. Read all batches into memory
+ List<RecordBatch> columnsBatches;
+ int totalRows;
+ Schema columnsSchema;
+ StringArray? columnNames = null;
+ int colNameIndex = -1;
+
+ // Extract column data
+ using (var stream = columnsResult.Stream)
{
- // TODO: Add log or throw
- return columnsResult; // Can't match without column names
- }
+ colNameIndex = stream.Schema.GetFieldIndex("COLUMN_NAME");
+ if (colNameIndex < 0)
+ {
+ // TODO: Add log or throw
+ return columnsResult; // Can't match without column
names
+ }
- var batchResult = await ReadAllBatchesAsync(stream,
cancellationToken);
- columnsBatches = batchResult.Batches;
- columnsSchema = batchResult.Schema;
- totalRows = batchResult.TotalRows;
+ var batchResult = await ReadAllBatchesAsync(stream,
cancellationToken);
+ columnsBatches = batchResult.Batches;
+ columnsSchema = batchResult.Schema;
+ totalRows = batchResult.TotalRows;
- if (columnsBatches.Count == 0)
- {
- // Return empty result with complete schema
- return CreateEmptyExtendedColumnsResult(columnsSchema);
+ if (columnsBatches.Count == 0)
+ {
+ // Return empty result with complete schema
+ return CreateEmptyExtendedColumnsResult(columnsSchema);
+ }
+
+ // Create column names array from all batches using
ArrayDataConcatenator.Concatenate
+ List<ArrayData> columnNameArrayDataList =
columnsBatches.Select(batch =>
+ batch.Column(colNameIndex).Data).ToList();
+ ArrayData? concatenatedColumnNames =
ArrayDataConcatenator.Concatenate(columnNameArrayDataList);
+ columnNames =
(StringArray)ArrowArrayFactory.BuildArray(concatenatedColumnNames!);
}
- // Create column names array from all batches using
ArrayDataConcatenator.Concatenate
- List<ArrayData> columnNameArrayDataList =
columnsBatches.Select(batch =>
- batch.Column(colNameIndex).Data).ToList();
- ArrayData? concatenatedColumnNames =
ArrayDataConcatenator.Concatenate(columnNameArrayDataList);
- columnNames =
(StringArray)ArrowArrayFactory.BuildArray(concatenatedColumnNames!);
- }
+ // 3. Create combined schema and prepare data
+ var allFields = new List<Field>(columnsSchema.FieldsList);
+ var combinedData = new List<IArrowArray>();
- // 3. Create combined schema and prepare data
- var allFields = new List<Field>(columnsSchema.FieldsList);
- var combinedData = new List<IArrowArray>();
+ // 4. Add all columns data by combining all batches
+ for (int colIdx = 0; colIdx < columnsSchema.FieldsList.Count;
colIdx++)
+ {
+ if (columnsBatches.Count == 0)
+ continue;
- // 4. Add all columns data by combining all batches
- for (int colIdx = 0; colIdx < columnsSchema.FieldsList.Count;
colIdx++)
- {
- if (columnsBatches.Count == 0)
- continue;
+ var field = columnsSchema.GetFieldByIndex(colIdx);
- var field = columnsSchema.GetFieldByIndex(colIdx);
+ // Collect arrays for this column from all batches
+ var columnArrays = new List<IArrowArray>();
+ foreach (var batch in columnsBatches)
+ {
+ columnArrays.Add(batch.Column(colIdx));
+ }
- // Collect arrays for this column from all batches
- var columnArrays = new List<IArrowArray>();
- foreach (var batch in columnsBatches)
- {
- columnArrays.Add(batch.Column(colIdx));
+ List<ArrayData> arrayDataList = columnArrays.Select(arr =>
arr.Data).ToList();
+ ArrayData? concatenatedData =
ArrayDataConcatenator.Concatenate(arrayDataList);
+ IArrowArray array =
ArrowArrayFactory.BuildArray(concatenatedData);
+ combinedData.Add(array);
}
- List<ArrayData> arrayDataList = columnArrays.Select(arr =>
arr.Data).ToList();
- ArrayData? concatenatedData =
ArrayDataConcatenator.Concatenate(arrayDataList);
- IArrowArray array =
ArrowArrayFactory.BuildArray(concatenatedData);
- combinedData.Add(array);
- }
-
- // 5. Process PK and FK data using helper methods with selected
fields
- await ProcessRelationshipDataSafe(pkResult, PrimaryKeyPrefix,
"COLUMN_NAME",
- PrimaryKeyFields, // Selected PK fields
- columnNames, totalRows,
- allFields, combinedData, cancellationToken);
+ // 5. Process PK and FK data using helper methods with
selected fields
+ await ProcessRelationshipDataSafe(pkResult, PrimaryKeyPrefix,
"COLUMN_NAME",
+ PrimaryKeyFields, // Selected PK fields
+ columnNames, totalRows,
+ allFields, combinedData, cancellationToken);
- await ProcessRelationshipDataSafe(fkResult, ForeignKeyPrefix,
"FKCOLUMN_NAME",
- ForeignKeyFields, // Selected FK fields
- columnNames, totalRows,
- allFields, combinedData, cancellationToken);
+ await ProcessRelationshipDataSafe(fkResult, ForeignKeyPrefix,
"FKCOLUMN_NAME",
+ ForeignKeyFields, // Selected FK fields
+ columnNames, totalRows,
+ allFields, combinedData, cancellationToken);
- // 6. Return the combined result
- var combinedSchema = new Schema(allFields, columnsSchema.Metadata);
+ // 6. Return the combined result
+ var combinedSchema = new Schema(allFields,
columnsSchema.Metadata);
- return new QueryResult(totalRows, new
HiveServer2Connection.HiveInfoArrowStream(combinedSchema, combinedData));
+ return new QueryResult(totalRows, new
HiveServer2Connection.HiveInfoArrowStream(combinedSchema, combinedData));
+ }
+ finally
+ {
+ // Dispose internal query results to ensure proper resource
cleanup
+ columnsResult.Stream?.Dispose();
Review Comment:
I wonder if we will consider awaiting on the results further down this
function, though?
--
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
