CurtHagenlocher commented on code in PR #1192: URL: https://github.com/apache/arrow-adbc/pull/1192#discussion_r1368845052
########## csharp/src/Drivers/BigQuery/BigQueryStatement.cs: ########## @@ -0,0 +1,394 @@ +/* +* Licensed to the Apache Software Foundation (ASF) under one or more +* contributor license agreements. See the NOTICE file distributed with +* this work for additional information regarding copyright ownership. +* The ASF licenses this file to You under the Apache License, Version 2.0 +* (the "License"); you may not use this file except in compliance with +* the License. You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + +using System; +using System.Collections.Generic; +using System.IO; +using System.Linq; +using System.Text.Json; +using System.Text.RegularExpressions; +using System.Threading; +using System.Threading.Tasks; +using Apache.Arrow.Ipc; +using Apache.Arrow.Types; +using Google.Apis.Auth.OAuth2; +using Google.Cloud.BigQuery.Storage.V1; +using Google.Cloud.BigQuery.V2; +using TableFieldSchema = Google.Apis.Bigquery.v2.Data.TableFieldSchema; +using TableSchema = Google.Apis.Bigquery.v2.Data.TableSchema; + +namespace Apache.Arrow.Adbc.Drivers.BigQuery +{ + /// <summary> + /// BigQuery-specific implementation of <see cref="AdbcStatement"/> + /// </summary> + public class BigQueryStatement : AdbcStatement + { + readonly BigQueryClient client; + readonly GoogleCredential credential; + + public BigQueryStatement(BigQueryClient client, GoogleCredential credential) + { + this.client = client; + this.credential = credential; + } + + public IReadOnlyDictionary<string, string>? Options { get; set; } + + public override QueryResult ExecuteQuery() + { + QueryOptions? queryOptions = ValidateOptions(); + BigQueryJob job = this.client.CreateQueryJob(SqlQuery, null, queryOptions); + BigQueryResults results = job.GetQueryResults(); + + BigQueryReadClientBuilder readClientBuilder = new BigQueryReadClientBuilder(); + readClientBuilder.Credential = this.credential; + BigQueryReadClient readClient = readClientBuilder.Build(); + + // TODO: translate the schema + + string table = $"projects/{results.TableReference.ProjectId}/datasets/{results.TableReference.DatasetId}/tables/{results.TableReference.TableId}"; + + ReadSession rs = new ReadSession { Table = table, DataFormat = DataFormat.Arrow }; + ReadSession rrs = readClient.CreateReadSession("projects/" + results.TableReference.ProjectId, rs, 1); + + long totalRows = results.TotalRows == null ? -1L : (long)results.TotalRows.Value; + IArrowArrayStream stream = new MultiArrowReader(TranslateSchema(results.Schema), rrs.Streams.Select(s => ReadChunk(readClient, s.Name))); + + return new QueryResult(totalRows, stream); + } + + public override UpdateResult ExecuteUpdate() + { + BigQueryResults result = this.client.ExecuteQuery(SqlQuery, parameters: null); + long updatedRows = result.NumDmlAffectedRows == null ? -1L : result.NumDmlAffectedRows.Value; + + return new UpdateResult(updatedRows); + } + + public override object GetValue(IArrowArray arrowArray, Field field, int index) + { + if(arrowArray == null) throw new ArgumentNullException(nameof(arrowArray)); + if (field == null) throw new ArgumentNullException(nameof(field)); + if (index < 0) throw new ArgumentOutOfRangeException(nameof(index)); + + try + { + switch (arrowArray) + { + case Int64Array int64Array: + return int64Array.GetValue(index); + case DoubleArray doubleArray: + return doubleArray.GetValue(index); + case Decimal128Array decimal128Array: + return decimal128Array.GetValue(index); + case Decimal256Array decimal256Array: + return decimal256Array.GetValue(index); + case BooleanArray booleanArray: + return booleanArray.GetValue(index); + case StringArray stringArray: + return stringArray.GetString(index); + case BinaryArray binaryArray: + + ReadOnlySpan<byte> bytes = binaryArray.GetBytes(index); + + if (bytes != null) Review Comment: Okay, then only the GetBytes would need to be fixed. -- 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]
