eric-wang-1990 commented on code in PR #2766:
URL: https://github.com/apache/arrow-adbc/pull/2766#discussion_r2078771822
##########
csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs:
##########
@@ -587,5 +592,370 @@ protected internal QueryResult
EnhanceGetColumnsResult(Schema originalSchema, IR
return new QueryResult(rowCount, new
HiveServer2Connection.HiveInfoArrowStream(enhancedSchema, enhancedData));
}
+
+ // Helper method to read all batches from a stream
+ private async Task<(List<RecordBatch> Batches, Schema Schema, int
TotalRows)> ReadAllBatchesAsync(
+ IArrowArrayStream stream, CancellationToken cancellationToken)
+ {
+ List<RecordBatch> batches = new List<RecordBatch>();
+ int totalRows = 0;
+ Schema schema = stream.Schema;
+
+ // Read all batches
+ while (true)
+ {
+ var batch = await
stream.ReadNextRecordBatchAsync(cancellationToken);
+ if (batch == null) break;
+
+ if (batch.Length > 0)
+ {
+ batches.Add(batch);
+ totalRows += batch.Length;
+ }
+ else
+ {
+ batch.Dispose();
+ }
+ }
+
+ return (batches, schema, totalRows);
+ }
+
+ // Helper method to create empty string columns filled with nulls
+ private List<IArrowArray> CreateEmptyStringColumns(string[]
fieldNames, int rowCount)
+ {
+ var result = new List<IArrowArray>();
+
+ foreach (var fieldName in fieldNames)
+ {
+ var builder = new StringArray.Builder();
+ for (int i = 0; i < rowCount; i++)
+ {
+ builder.AppendNull();
+ }
+ result.Add(builder.Build());
+ }
+
+ return result;
+ }
+
+ // Helper class to manage builder creation and value appending
+ // This is only for metadata query result so we only have int or
string types
+ private class TypedBuilder
+ {
+ private readonly IArrowArrayBuilder _builder;
+ private readonly ArrowTypeId _typeId;
+
+ public TypedBuilder(ArrowTypeId typeId)
+ {
+ _typeId = typeId;
+ _builder = CreateBuilder(typeId);
+ }
+
+ private static IArrowArrayBuilder CreateBuilder(ArrowTypeId
typeId) => typeId switch
+ {
+ ArrowTypeId.Int16 => new Int16Array.Builder(),
+ ArrowTypeId.Int32 => new Int32Array.Builder(),
+ _ => new StringArray.Builder() // Default to string for
unsupported types
+ };
+
+ public void AppendNull()
+ {
+ switch (_typeId)
+ {
+ case ArrowTypeId.Int16:
+ ((Int16Array.Builder)_builder).AppendNull();
+ break;
+ case ArrowTypeId.Int32:
+ ((Int32Array.Builder)_builder).AppendNull();
+ break;
+ default:
+ ((StringArray.Builder)_builder).AppendNull();
+ break;
+ }
+ }
+
+ public IArrowArray Build() => _typeId switch
+ {
+ ArrowTypeId.Int16 => ((Int16Array.Builder)_builder).Build(),
+ ArrowTypeId.Int32 => ((Int32Array.Builder)_builder).Build(),
+ _ => ((StringArray.Builder)_builder).Build()
+ };
+
+ public void AppendValue(IArrowArray columnArray, int rowIndex)
+ {
+ try
+ {
+ switch (_typeId)
+ {
+ case ArrowTypeId.Int16:
+
((Int16Array.Builder)_builder).Append(((Int16Array)columnArray).GetValue(rowIndex)!.Value);
+ break;
+ case ArrowTypeId.Int32:
+
((Int32Array.Builder)_builder).Append(((Int32Array)columnArray).GetValue(rowIndex)!.Value);
+ break;
+ default: // Handles String and other types that
default to StringArray.Builder
+ // Try to cast to StringArray and get string value
- if it fails,
+ // the outer try-catch will handle it and call
AppendNull()
+
((StringArray.Builder)_builder).Append(((StringArray)columnArray).GetString(rowIndex));
+ break;
+ }
+ }
+ catch
+ {
+ // If any conversion fails or columnArray is null, append
null as fallback
+ AppendNull();
+ }
+ }
+
+
+ }
+
+ private async Task<QueryResult>
GetColumnsExtendedAsync(CancellationToken cancellationToken = default)
+ {
+ // 1. Get all three results at once
+ var columnsResult = await GetColumnsAsync(cancellationToken);
+ if (columnsResult.Stream == null) return columnsResult;
+
+ var pkResult = await GetPrimaryKeysAsync(cancellationToken);
+
+ // For FK lookup, we need to pass in the current
catalog/schema/table as the foreign table
+ var resp = await Connection.GetCrossReferenceAsync(
+ null,
+ null,
+ null,
+ CatalogName,
+ SchemaName,
+ TableName,
+ cancellationToken);
+
+ var fkResult = await GetQueryResult(resp.DirectResults,
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)
+ {
+ colNameIndex = stream.Schema.GetFieldIndex("COLUMN_NAME");
+ if (colNameIndex < 0) return columnsResult; // Can't match
without column names
+
+ var batchResult = await ReadAllBatchesAsync(stream,
cancellationToken);
+ columnsBatches = batchResult.Batches;
+ columnsSchema = batchResult.Schema;
+ totalRows = batchResult.TotalRows;
+
+ if (columnsBatches.Count == 0) return columnsResult;
+
+ // Create column names array from all batches
+ var builder = new StringArray.Builder();
+ foreach (var batch in columnsBatches)
+ {
+ StringArray batchColNames =
(StringArray)batch.Column(colNameIndex);
+ for (int i = 0; i < batch.Length; i++)
+ {
+ builder.Append(batchColNames.GetString(i));
+ }
+ }
+ columnNames = builder.Build();
+ }
+
+ // 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
+ // Create a combined array for each column across all batches
+ for (int colIdx = 0; colIdx < columnsSchema.FieldsList.Count;
colIdx++)
Review Comment:
sure
--
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
