Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(csharp/src/Drivers/BigQuery): support max stream count setting when creating read session #2289

Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion csharp/src/Drivers/BigQuery/BigQueryConnection.cs
Original file line number Diff line number Diff line change
Expand Up @@ -1002,7 +1002,8 @@ private IReadOnlyDictionary<string, string> ParseOptions()
BigQueryParameters.UseLegacySQL,
BigQueryParameters.LargeDecimalsAsString,
BigQueryParameters.LargeResultsDestinationTable,
BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes
BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes,
BigQueryParameters.CreateReadSessionMaxStreamCount
};

foreach (string key in statementOptions)
Expand Down
1 change: 1 addition & 0 deletions csharp/src/Drivers/BigQuery/BigQueryParameters.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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 CreateReadSessionMaxStreamCount = "adbc.bigquery.create_read_session.max_stream_count";
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we name this something a little more meaningful to the user like adbc.bigquery.max_fetch_concurrency?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

have changed

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you please change the constant name to match the property name?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OK, thanks for the correction.

}

/// <summary>
Expand Down
13 changes: 12 additions & 1 deletion csharp/src/Drivers/BigQuery/BigQueryStatement.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is 1 really the best default value here? (I don't know how we ended up with 1; the documentation clearly recommends 0.

Copy link
Contributor Author

@qifanzhang-ms qifanzhang-ms Oct 30, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The default value is set to 1, mainly due to two reasons:

1 The original driver setting is 1. The GBQ MSDI driver is now being used by a large number of CXs, and changing this setting may cause potential regression.

2 The doc does say: If unset or zero, the server will provide a value of streams so as to produce reasonable throughput. But I found in my test that this is not 100% accurate. At least in my test of reading the same table, it took 42.5 minutes when max_stream_count=1 and 46.4 minutes when max_stream_count=0. Although this may be due to other reasons such as the test table content or network fluctuations, it at least shows that potential regression on perf exists.

if (this.Options?.TryGetValue(BigQueryParameters.CreateReadSessionMaxStreamCount, 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)));
Expand Down
3 changes: 3 additions & 0 deletions csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs
Original file line number Diff line number Diff line change
Expand Up @@ -59,5 +59,8 @@ public BigQueryTestConfiguration()

[JsonPropertyName("timeoutMinutes")]
public int? TimeoutMinutes { get; set; }

[JsonPropertyName("maxStreamCount")]
public int? MaxStreamCount { get; set; }
}
}
5 changes: 5 additions & 0 deletions csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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.CreateReadSessionMaxStreamCount, testConfiguration.MaxStreamCount.Value.ToString());
}

return parameters;
}

Expand Down
Loading