CurtHagenlocher commented on code in PR #1192: URL: https://github.com/apache/arrow-adbc/pull/1192#discussion_r1367731304
########## csharp/src/Drivers/BigQuery/BigQueryStatement.cs: ########## @@ -0,0 +1,422 @@ +/* +* 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 is Int64Array) + { + return ((Int64Array)arrowArray).Values[index]; + } + else if (arrowArray is DoubleArray) + { + return ((DoubleArray)arrowArray).Values[index]; + } + else if (arrowArray is Decimal128Array) + { + try + { + // the value may be <decimal.min or >decimal.max + // then Arrow throws an exception + // no good way to check prior to + return ((Decimal128Array)arrowArray).GetValue(index); + } + catch (OverflowException oex) + { + return ParseDecimalValueFromOverflowException(oex); + } + } + else if (arrowArray is Decimal256Array) + { + try + { + return ((Decimal256Array)arrowArray).GetValue(index); + } + catch (OverflowException oex) + { + return ParseDecimalValueFromOverflowException(oex); + } + } + else if (arrowArray is BooleanArray) + { + return ((BooleanArray)arrowArray).GetValue(index); + } + else if (arrowArray is StringArray) + { + return ((StringArray)arrowArray).GetString(index); + } + else if (arrowArray is BinaryArray) + { + ReadOnlySpan<byte> bytes = ((BinaryArray)arrowArray).GetBytes(index); + + if (bytes != null) + return bytes.ToArray(); + } + else if (arrowArray is Date32Array) + { + Date32Array date32Array = (Date32Array)arrowArray; + + return date32Array.GetDateTime(index); + } + else if (arrowArray is Date64Array) + { + Date64Array date64Array = (Date64Array)arrowArray; + + return date64Array.GetDateTime(index); + } + else if (arrowArray is Time64Array) + { + return ((Time64Array)arrowArray).GetValue(index); + } + else if (arrowArray is TimestampArray) + { + TimestampArray timestampArray = (TimestampArray)arrowArray; + DateTimeOffset dateTimeOffset = timestampArray.GetTimestamp(index).Value; + return dateTimeOffset; + } + else if (arrowArray is StructArray) + { + StructArray structArray = (StructArray)arrowArray; + return SerializeToJson(structArray, index); + } + // maybe not be needed? + else if (arrowArray is ListArray) + { + return ((ListArray)arrowArray).GetSlicedValues(index); + } + + return null; + } + + static Schema TranslateSchema(TableSchema schema) + { + return new Schema(schema.Fields.Select(TranslateField), null); + } + + static Field TranslateField(TableFieldSchema field) + { + return new Field(field.Name, TranslateType(field), field.Mode == "NULLABLE"); + } + + static IArrowType TranslateType(TableFieldSchema field) + { + // per https://developers.google.com/resources/api-libraries/documentation/bigquery/v2/java/latest/com/google/api/services/bigquery/model/TableFieldSchema.html#getType-- + + switch (field.Type) + { + case "INTEGER" or "INT64": + return Int64Type.Default; + case "FLOAT" or "FLOAT64": + return DoubleType.Default; + case "BOOL" or "BOOLEAN": + return BooleanType.Default; + case "STRING": + return StringType.Default; + case "BYTES": + return BinaryType.Default; + case "DATETIME": + return TimestampType.Default; + case "TIMESTAMP": + return TimestampType.Default; + case "TIME": + return Time64Type.Default; + case "DATE": + return Date64Type.Default; + case "RECORD" or "STRUCT": + // its a json string + return StringType.Default; + // get schema cannot get precison and scale + case "NUMERIC" or "DECIMAL": + return new Decimal128Type(38, 9); Review Comment: Okay, so is this like the "Snowflake problem" where each chunk of Arrow data which comes back might have a precision and scale that's different than we returned as the result schema? -- 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]
