fix(csharp/src/Drivers/BigQuery): fix support for large results (#1507)

- Adds value for a DestinationTable, which is needed for large results - Removes old Arrow references in test projects --------- Co-authored-by: David Coe <coedavid@umich.edu>
apache · Feb 4, 2024 · b1947bf · b1947bf
1 parent 232ceeb
commit b1947bf
Show file tree

Hide file tree

Showing 8 changed files with 84 additions and 8 deletions.
diff --git a/csharp/src/Drivers/BigQuery/BigQueryConnection.cs b/csharp/src/Drivers/BigQuery/BigQueryConnection.cs
@@ -16,6 +16,7 @@
 */
 
 using System;
+using System.Collections;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
 using System.Diagnostics;
@@ -270,6 +271,23 @@ public override IArrowArrayStream GetObjects(
  return new BigQueryInfoArrowStream(StandardSchemas.GetObjectsSchema, dataArrays);
  }
 
+ /// <summary>
+ /// Executes the query using the BigQueryClient.
+ /// </summary>
+ /// <param name="sql">The query to execute.</param>
+ /// <param name="parameters">Parameters to include.</param>
+ /// <param name="queryOptions">Additional query options.</param>
+ /// <param name="resultsOptions">Additional result options.</param>
+ /// <returns></returns>
+ /// <remarks>
+ /// Can later add logging or metrics around query calls.
+ /// </remarks>
+ private BigQueryResults? ExecuteQuery(string sql, IEnumerable<BigQueryParameter>? parameters, QueryOptions? queryOptions = null, GetQueryResultsOptions? resultsOptions = null)
+ {
+ BigQueryResults? result = this.client?.ExecuteQuery(sql, parameters, queryOptions, resultsOptions);
+ return result;
+ }
+
  private List<IArrowArray> GetCatalogs(
  GetObjectsDepth depth,
  string catalogPattern,
@@ -405,7 +423,7 @@ private StructArray GetTableSchemas(
  }
  }
 
- BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null);
+ BigQueryResults? result = ExecuteQuery(query, parameters: null);
 
  if (result != null)
  {
@@ -481,7 +499,7 @@ private StructArray GetColumnSchema(
  query = string.Concat(query, string.Format("AND column_name LIKE '{0}'", Sanitize(columnNamePattern)));
  }
 
- BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null);
+ BigQueryResults? result = ExecuteQuery(query, parameters: null);
 
  if (result != null)
  {
@@ -570,7 +588,7 @@ private StructArray GetConstraintSchema(
  string query = string.Format("SELECT * FROM `{0}`.`{1}`.INFORMATION_SCHEMA.TABLE_CONSTRAINTS WHERE table_name = '{2}'",
  Sanitize(catalog), Sanitize(dbSchema), Sanitize(table));
 
- BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null);
+ BigQueryResults? result = ExecuteQuery(query, parameters: null);
 
  if (result != null)
  {
@@ -631,7 +649,7 @@ private StringArray GetConstraintColumnNames(
 
  StringArray.Builder constraintColumnNamesBuilder = new StringArray.Builder();
 
- BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null);
+ BigQueryResults? result = ExecuteQuery(query, parameters: null);
 
  if (result != null)
  {
@@ -661,7 +679,7 @@ private StructArray GetConstraintsUsage(
  string query = string.Format("SELECT * FROM `{0}`.`{1}`.INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE WHERE constraint_name = '{2}'",
  Sanitize(catalog), Sanitize(dbSchema), Sanitize(constraintName));
 
- BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null);
+ BigQueryResults? result = ExecuteQuery(query, parameters: null);
 
  if (result != null)
  {
@@ -788,7 +806,7 @@ public override Schema GetTableSchema(string catalog, string dbSchema, string ta
  string query = string.Format("SELECT * FROM `{0}`.`{1}`.INFORMATION_SCHEMA.COLUMNS WHERE table_name = '{2}'",
  Sanitize(catalog), Sanitize(dbSchema), Sanitize(tableName));
 
- BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null);
+ BigQueryResults? result = ExecuteQuery(query, parameters: null);
 
  List<Field> fields = new List<Field>();
 
@@ -1014,6 +1032,10 @@ private IReadOnlyDictionary<string, string> ParseOptions()
  {
  options[keyValuePair.Key] = keyValuePair.Value;
  }
+ if (keyValuePair.Key == BigQueryParameters.LargeResultsDestinationTable)
+ {
+ options[keyValuePair.Key] = keyValuePair.Value;
+ }
  }
 
  return new ReadOnlyDictionary<string, string>(options);

diff --git a/csharp/src/Drivers/BigQuery/BigQueryParameters.cs b/csharp/src/Drivers/BigQuery/BigQueryParameters.cs
@@ -29,6 +29,7 @@ public class BigQueryParameters
  public const string AuthenticationType = "adbc.bigquery.auth_type";
  public const string JsonCredential = "adbc.bigquery.auth_json_credential";
  public const string AllowLargeResults = "adbc.bigquery.allow_large_results";
+ public const string LargeResultsDestinationTable = "adbc.bigquery.large_results_destination_table";
  public const string UseLegacySQL = "adbc.bigquery.use_legacy_sql";
  public const string LargeDecimalsAsString = "adbc.bigquery.large_decimals_as_string";
  public const string Scopes = "adbc.bigquery.scopes";

diff --git a/csharp/src/Drivers/BigQuery/BigQueryStatement.cs b/csharp/src/Drivers/BigQuery/BigQueryStatement.cs
@@ -27,6 +27,7 @@
 using Apache.Arrow.Ipc;
 using Apache.Arrow.Types;
 using Google.Apis.Auth.OAuth2;
+using Google.Apis.Bigquery.v2.Data;
 using Google.Cloud.BigQuery.Storage.V1;
 using Google.Cloud.BigQuery.V2;
 using TableFieldSchema = Google.Apis.Bigquery.v2.Data.TableFieldSchema;
@@ -185,6 +186,36 @@ static IArrowReader ReadChunk(BigQueryReadClient readClient, string streamName)
  {
  options.AllowLargeResults = true ? keyValuePair.Value.ToLower().Equals("true") : false;
  }
+ if (keyValuePair.Key == BigQueryParameters.LargeResultsDestinationTable)
+ {
+ string destinationTable = keyValuePair.Value;
+
+ if (!destinationTable.Contains("."))
+ throw new InvalidOperationException($"{BigQueryParameters.LargeResultsDestinationTable} is invalid");
+
+ string projectId = string.Empty;
+ string datasetId = string.Empty;
+ string tableId = string.Empty;
+
+ string[] segments = destinationTable.Split('.');
+
+ if(segments.Length != 3)
+ throw new InvalidOperationException($"{BigQueryParameters.LargeResultsDestinationTable} cannot be parsed");
+
+ projectId = segments[0];
+ datasetId = segments[1];
+ tableId = segments[2];
+
+ if(string.IsNullOrEmpty(projectId.Trim()) || string.IsNullOrEmpty(datasetId.Trim()) || string.IsNullOrEmpty(tableId.Trim()))
+ throw new InvalidOperationException($"{BigQueryParameters.LargeResultsDestinationTable} contains invalid values");
+
+ options.DestinationTable = new TableReference()
+ {
+ ProjectId = projectId,
+ DatasetId = datasetId,
+ TableId = tableId
+ };
+ }
  if (keyValuePair.Key == BigQueryParameters.UseLegacySQL)
  {
  options.UseLegacySql = true ? keyValuePair.Value.ToLower().Equals("true") : false;

diff --git a/csharp/src/Drivers/BigQuery/readme.md b/csharp/src/Drivers/BigQuery/readme.md
@@ -49,6 +49,10 @@ https://cloud.google.com/dotnet/docs/reference/Google.Cloud.BigQuery.V2/latest/G
 **adbc.bigquery.auth_json_credential**<br>
 &nbsp;&nbsp;&nbsp;&nbsp;Required if using `service` authentication. This value is passed to the [GoogleCredential.FromJson](https://cloud.google.com/dotnet/docs/reference/Google.Apis/latest/Google.Apis.Auth.OAuth2.GoogleCredential#Google_Apis_Auth_OAuth2_GoogleCredential_FromJson_System_String) method.
 
+**adbc.bigquery.large_results_destination_table**<br>
+&nbsp;&nbsp;&nbsp;&nbsp;Sets the [DestinationTable](https://cloud.google.com/dotnet/docs/reference/Google.Cloud.BigQuery.V2/latest/Google.Cloud.BigQuery.V2.QueryOptions#Google_Cloud_BigQuery_V2_QueryOptions_DestinationTable) value of the QueryOptions if configured. Expects the format to be `{projectId}.{datasetId}.{tableId}` to set the corresponding values in the [TableReference](https://github.com/googleapis/google-api-dotnet-client/blob/6c415c73788b848711e47c6dd33c2f93c76faf97/Src/Generated/Google.Apis.Bigquery.v2/Google.Apis.Bigquery.v2.cs#L9348) class.
+
+
 **adbc.bigquery.project_id**<br>
 &nbsp;&nbsp;&nbsp;&nbsp;The [Project ID](https://cloud.google.com/resource-manager/docs/creating-managing-projects) used for accessing BigQuery.
 

diff --git a/csharp/test/Drivers/BigQuery/Apache.Arrow.Adbc.Tests.Drivers.BigQuery.csproj b/csharp/test/Drivers/BigQuery/Apache.Arrow.Adbc.Tests.Drivers.BigQuery.csproj
@@ -13,7 +13,6 @@
  </ItemGroup>
  <ItemGroup>
  <ProjectReference Include="..\..\..\src\Apache.Arrow.Adbc\Apache.Arrow.Adbc.csproj" />
- <ProjectReference Include="..\..\..\src\arrow\csharp\src\Apache.Arrow\Apache.Arrow.csproj" />
  <ProjectReference Include="..\..\..\src\Client\Apache.Arrow.Adbc.Client.csproj" />
  <ProjectReference Include="..\..\..\src\Drivers\BigQuery\Apache.Arrow.Adbc.Drivers.BigQuery.csproj" />
  <ProjectReference Include="..\..\Apache.Arrow.Adbc.Tests\Apache.Arrow.Adbc.Tests.csproj" />

diff --git a/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs b/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs
@@ -24,6 +24,11 @@ namespace Apache.Arrow.Adbc.Tests.Drivers.BigQuery
  /// </summary>
  internal class BigQueryTestConfiguration : TestConfiguration
  {
+ public BigQueryTestConfiguration()
+ {
+ AllowLargeResults = false;
+ }
+
  [JsonPropertyName("projectId")]
  public string ProjectId { get; set; }
 
@@ -42,5 +47,10 @@ internal class BigQueryTestConfiguration : TestConfiguration
  [JsonPropertyName("jsonCredential")]
  public string JsonCredential { get; set; }
 
+ [JsonPropertyName("allowLargeResults")]
+ public bool AllowLargeResults { get; set; }
+
+ [JsonPropertyName("largeResultsDestinationTable")]
+ public string LargeResultsDestinationTable { get; set; }
  }
 }
diff --git a/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs b/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs
@@ -74,6 +74,16 @@ internal static Dictionary<string,string> GetBigQueryParameters(BigQueryTestConf
  parameters.Add(BigQueryParameters.Scopes, testConfiguration.Scopes);
  }
 
+ if(testConfiguration.AllowLargeResults)
+ {
+ parameters.Add(BigQueryParameters.AllowLargeResults, testConfiguration.AllowLargeResults.ToString());
+ }
+
+ if(!string.IsNullOrEmpty(testConfiguration.LargeResultsDestinationTable))
+ {
+ parameters.Add(BigQueryParameters.LargeResultsDestinationTable, testConfiguration.LargeResultsDestinationTable);
+ }
+
  return parameters;
  }
 

diff --git a/csharp/test/Drivers/FlightSql/Apache.Arrow.Adbc.Tests.Drivers.FlightSql.csproj b/csharp/test/Drivers/FlightSql/Apache.Arrow.Adbc.Tests.Drivers.FlightSql.csproj
@@ -15,7 +15,6 @@
  </ItemGroup>
  <ItemGroup>
  <ProjectReference Include="..\..\..\src\Apache.Arrow.Adbc\Apache.Arrow.Adbc.csproj" />
- <ProjectReference Include="..\..\..\src\arrow\csharp\src\Apache.Arrow\Apache.Arrow.csproj" />
  <ProjectReference Include="..\..\..\src\Client\Apache.Arrow.Adbc.Client.csproj" />
  <ProjectReference Include="..\..\..\src\Drivers\FlightSql\Apache.Arrow.Adbc.Drivers.FlightSql.csproj" />
  <ProjectReference Include="..\..\Apache.Arrow.Adbc.Tests\Apache.Arrow.Adbc.Tests.csproj" />