feat(csharp/src/Drivers/BigQuery): support max stream count setting w…

…hen creating read session (#2289)
apache · Oct 30, 2024 · 8524b91 · 8524b91
1 parent f6a8861
commit 8524b91
Show file tree

Hide file tree

Showing 8 changed files with 29 additions and 3 deletions.
diff --git a/csharp/src/Drivers/BigQuery/BigQueryConnection.cs b/csharp/src/Drivers/BigQuery/BigQueryConnection.cs
@@ -1002,7 +1002,8 @@ private IReadOnlyDictionary<string, string> ParseOptions()
                 BigQueryParameters.UseLegacySQL,
                 BigQueryParameters.LargeDecimalsAsString,
                 BigQueryParameters.LargeResultsDestinationTable,
-                BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes
+                BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes,
+                BigQueryParameters.MaxFetchConcurrency
             };
 
             foreach (string key in statementOptions)

diff --git a/csharp/src/Drivers/BigQuery/BigQueryParameters.cs b/csharp/src/Drivers/BigQuery/BigQueryParameters.cs
@@ -35,6 +35,7 @@ public class BigQueryParameters
         public const string Scopes = "adbc.bigquery.scopes";
         public const string IncludeConstraintsWithGetObjects = "adbc.bigquery.include_constraints_getobjects";
         public const string GetQueryResultsOptionsTimeoutMinutes = "adbc.bigquery.get_query_results_options.timeout";
+        public const string MaxFetchConcurrency = "adbc.bigquery.max_fetch_concurrency";
     }
 
     /// <summary>

diff --git a/csharp/src/Drivers/BigQuery/BigQueryStatement.cs b/csharp/src/Drivers/BigQuery/BigQueryStatement.cs
@@ -103,8 +103,19 @@ public override QueryResult ExecuteQuery()
 
             string table = $"projects/{results.TableReference.ProjectId}/datasets/{results.TableReference.DatasetId}/tables/{results.TableReference.TableId}";
 
+            int maxStreamCount = 1;
+            if (this.Options?.TryGetValue(BigQueryParameters.MaxFetchConcurrency, out string? maxStreamCountString) == true)
+            {
+                if (int.TryParse(maxStreamCountString, out int count))
+                {
+                    if (count >= 0)
+                    {
+                        maxStreamCount = count;
+                    }
+                }
+            }
             ReadSession rs = new ReadSession { Table = table, DataFormat = DataFormat.Arrow };
-            ReadSession rrs = readClient.CreateReadSession("projects/" + results.TableReference.ProjectId, rs, 1);
+            ReadSession rrs = readClient.CreateReadSession("projects/" + results.TableReference.ProjectId, rs, maxStreamCount);
 
             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)));

diff --git a/csharp/src/Drivers/BigQuery/readme.md b/csharp/src/Drivers/BigQuery/readme.md
@@ -54,6 +54,9 @@ https://cloud.google.com/dotnet/docs/reference/Google.Cloud.BigQuery.V2/latest/G
 **adbc.bigquery.get_query_results_options.timeout**<br>
 &nbsp;&nbsp;&nbsp;&nbsp;Optional. Sets the timeout (in minutes) for the GetQueryResultsOptions value. If not set, defaults to 5 minutes.
 
+**adbc.bigquery.max_fetch_concurrency**<br>
+&nbsp;&nbsp;&nbsp;&nbsp;Optional. Sets the [maxStreamCount](https://cloud.google.com/dotnet/docs/reference/Google.Cloud.BigQuery.Storage.V1/latest/Google.Cloud.BigQuery.Storage.V1.BigQueryReadClient#Google_Cloud_BigQuery_Storage_V1_BigQueryReadClient_CreateReadSession_System_String_Google_Cloud_BigQuery_Storage_V1_ReadSession_System_Int32_Google_Api_Gax_Grpc_CallSettings_) for the CreateReadSession method. If not set, defaults to 1.
+
 **adbc.bigquery.include_constraints_getobjects**<br>
 &nbsp;&nbsp;&nbsp;&nbsp;Optional. Some callers do not need the constraint details when they get the table information and can improve the speed of obtaining the results. Setting this value to `"false"` will not include the constraint details. The default value is `"true"`.
 

diff --git a/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs b/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs
@@ -59,5 +59,8 @@ public BigQueryTestConfiguration()
 
         [JsonPropertyName("timeoutMinutes")]
         public int? TimeoutMinutes { get; set; }
+
+        [JsonPropertyName("maxStreamCount")]
+        public int? MaxStreamCount { get; set; }
     }
 }
diff --git a/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs b/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs
@@ -91,6 +91,11 @@ internal static Dictionary<string, string> GetBigQueryParameters(BigQueryTestCon
                 parameters.Add(BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes, testConfiguration.TimeoutMinutes.Value.ToString());
             }
 
+            if (testConfiguration.MaxStreamCount.HasValue)
+            {
+                parameters.Add(BigQueryParameters.MaxFetchConcurrency, testConfiguration.MaxStreamCount.Value.ToString());
+            }
+
             return parameters;
         }
 

diff --git a/csharp/test/Drivers/BigQuery/Resources/bigqueryconfig.json b/csharp/test/Drivers/BigQuery/Resources/bigqueryconfig.json
@@ -3,11 +3,12 @@
     "clientId": "",
     "clientSecret": "",
     "refreshToken": "",
+    "maxStreamCount": 1,
     "metadata": {
         "catalog": "",
         "schema": "",
         "table": "",
-        "expectedColumnCount":  0
+        "expectedColumnCount": 0
     },
     "query": "",
     "expectedResults": 0

diff --git a/csharp/test/Drivers/BigQuery/readme.md b/csharp/test/Drivers/BigQuery/readme.md
@@ -39,6 +39,7 @@ The following values can be setup in the configuration
 - **query** - The query to use.
 - **expectedResults** - The expected number of results from the query.
 - **timeoutMinutes** - The timeout (in minutes).
+- **maxStreamCount** - The max stream count.
 - **includeTableConstraints** - Whether to include table constraints in the GetObjects query.
 - **largeResultsDestinationTable** - 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.
 - **allowLargeResults** - Whether to allow large results .