Skip to content

Commit d93dbaf

Browse files
fix: Resolve issue with unset thread-local options (#741)
1 parent b5a3928 commit d93dbaf

File tree

2 files changed

+64
-21
lines changed

2 files changed

+64
-21
lines changed

bigframes/_config/__init__.py

Lines changed: 23 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,12 @@
1717
DataFrames from this package.
1818
"""
1919

20+
from __future__ import annotations
21+
2022
import copy
23+
from dataclasses import dataclass, field
2124
import threading
25+
from typing import Optional
2226

2327
import bigframes_vendored.pandas._config.config as pandas_config
2428

@@ -28,18 +32,27 @@
2832
import bigframes._config.sampling_options as sampling_options
2933

3034

35+
@dataclass
36+
class ThreadLocalConfig(threading.local):
37+
# If unset, global settings will be used
38+
bigquery_options: Optional[bigquery_options.BigQueryOptions] = None
39+
# Note: use default factory instead of default instance so each thread initializes to default values
40+
display_options: display_options.DisplayOptions = field(
41+
default_factory=display_options.DisplayOptions
42+
)
43+
sampling_options: sampling_options.SamplingOptions = field(
44+
default_factory=sampling_options.SamplingOptions
45+
)
46+
compute_options: compute_options.ComputeOptions = field(
47+
default_factory=compute_options.ComputeOptions
48+
)
49+
50+
3151
class Options:
3252
"""Global options affecting BigQuery DataFrames behavior."""
3353

3454
def __init__(self):
35-
self._local = threading.local()
36-
37-
# Initialize these in the property getters to make sure we do have a
38-
# separate instance per thread.
39-
self._local.bigquery_options = None
40-
self._local.display_options = None
41-
self._local.sampling_options = None
42-
self._local.compute_options = None
55+
self._local = ThreadLocalConfig()
4356

4457
# BigQuery options are special because they can only be set once per
4558
# session, so we need an indicator as to whether we are using the
@@ -61,21 +74,16 @@ def _init_bigquery_thread_local(self):
6174
@property
6275
def bigquery(self) -> bigquery_options.BigQueryOptions:
6376
"""Options to use with the BigQuery engine."""
64-
if (
65-
bigquery_options := getattr(self._local, "bigquery_options", None)
66-
) is not None:
77+
if self._local.bigquery_options is not None:
6778
# The only way we can get here is if someone called
6879
# _init_bigquery_thread_local.
69-
return bigquery_options
80+
return self._local.bigquery_options
7081

7182
return self._bigquery_options
7283

7384
@property
7485
def display(self) -> display_options.DisplayOptions:
7586
"""Options controlling object representation."""
76-
if self._local.display_options is None:
77-
self._local.display_options = display_options.DisplayOptions()
78-
7987
return self._local.display_options
8088

8189
@property
@@ -88,17 +96,11 @@ def sampling(self) -> sampling_options.SamplingOptions:
8896
matplotlib plotting). This option can be overriden by
8997
parameters in specific functions.
9098
"""
91-
if self._local.sampling_options is None:
92-
self._local.sampling_options = sampling_options.SamplingOptions()
93-
9499
return self._local.sampling_options
95100

96101
@property
97102
def compute(self) -> compute_options.ComputeOptions:
98103
"""Thread-local options controlling object computation."""
99-
if self._local.compute_options is None:
100-
self._local.compute_options = compute_options.ComputeOptions()
101-
102104
return self._local.compute_options
103105

104106
@property
Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
# Copyright 2023 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import threading
16+
17+
import bigframes._config
18+
19+
20+
def test_mutate_options_threaded():
21+
options = bigframes._config.Options()
22+
options.display.max_rows = 50
23+
result_dict = {"this_before": options.display.max_rows}
24+
25+
def mutate_options_threaded(options, result_dict):
26+
result_dict["other_before"] = options.display.max_rows
27+
28+
options.display.max_rows = 100
29+
result_dict["other_after"] = options.display.max_rows
30+
31+
thread = threading.Thread(
32+
target=(lambda: mutate_options_threaded(options, result_dict))
33+
)
34+
thread.start()
35+
thread.join(1)
36+
result_dict["this_after"] = options.display.max_rows
37+
38+
assert result_dict["this_before"] == 50
39+
assert result_dict["this_after"] == 50
40+
assert result_dict["other_before"] == 25
41+
assert result_dict["other_after"] == 100

0 commit comments

Comments
 (0)