-
Notifications
You must be signed in to change notification settings - Fork 52
feat: bigframes.options
and bigframes.option_context
now uses thread-local variables to prevent context managers in separate threads from affecting each other
#652
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
Changes from 10 commits
9d67c30
4fa1a92
5f3720c
d12c64e
a7dbb8d
3e3baf7
6a400aa
d65efba
5aa68a4
64e3700
8cd7931
a7f3c96
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -17,6 +17,9 @@ | |
DataFrames from this package. | ||
""" | ||
|
||
import copy | ||
import threading | ||
|
||
import bigframes_vendored.pandas._config.config as pandas_config | ||
|
||
import bigframes._config.bigquery_options as bigquery_options | ||
|
@@ -29,44 +32,93 @@ class Options: | |
"""Global options affecting BigQuery DataFrames behavior.""" | ||
|
||
def __init__(self): | ||
self._local = threading.local() | ||
|
||
# Initialize these in the property getters to make sure we do have a | ||
# separate instance per thread. | ||
self._local.bigquery_options = None | ||
self._local.display_options = None | ||
self._local.sampling_options = None | ||
self._local.compute_options = None | ||
|
||
# BigQuery options are special because they can only be set once per | ||
# session, so we need an indicator as to whether we are using the | ||
# thread-local session or the global session. | ||
self._local.is_bigquery_thread_local = False | ||
self._bigquery_options = bigquery_options.BigQueryOptions() | ||
self._display_options = display_options.DisplayOptions() | ||
self._sampling_options = sampling_options.SamplingOptions() | ||
self._compute_options = compute_options.ComputeOptions() | ||
|
||
def _init_bigquery_thread_local(self): | ||
"""Initialize thread-local options, based on current global options.""" | ||
|
||
# Already thread-local, so don't reset any options that have been set | ||
# already. No locks needed since this only modifies thread-local | ||
# variables. | ||
if self._local.is_bigquery_thread_local: | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. i.e., this could be "if self._local.bigquery_options is not None" |
||
return | ||
|
||
self._local.is_bigquery_thread_local = True | ||
self._local.bigquery_options = copy.deepcopy(self._bigquery_options) | ||
self._local.bigquery_options._session_started = False | ||
|
||
@property | ||
def bigquery(self) -> bigquery_options.BigQueryOptions: | ||
"""Options to use with the BigQuery engine.""" | ||
if self._local.is_bigquery_thread_local: | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Here too |
||
# The only way we can get here is if someone called | ||
# _init_bigquery_thread_local, which initializes | ||
# self._local.bigquery_options. | ||
return self._local.bigquery_options | ||
|
||
return self._bigquery_options | ||
|
||
@property | ||
def display(self) -> display_options.DisplayOptions: | ||
"""Options controlling object representation.""" | ||
return self._display_options | ||
if self._local.display_options is None: | ||
self._local.display_options = display_options.DisplayOptions() | ||
|
||
return self._local.display_options | ||
|
||
@property | ||
def sampling(self) -> sampling_options.SamplingOptions: | ||
"""Options controlling downsampling when downloading data | ||
to memory. The data will be downloaded into memory explicitly | ||
to memory. | ||
|
||
The data can be downloaded into memory explicitly | ||
(e.g., to_pandas, to_numpy, values) or implicitly (e.g., | ||
matplotlib plotting). This option can be overriden by | ||
parameters in specific functions.""" | ||
return self._sampling_options | ||
parameters in specific functions. | ||
""" | ||
if self._local.sampling_options is None: | ||
self._local.sampling_options = sampling_options.SamplingOptions() | ||
|
||
return self._local.sampling_options | ||
|
||
@property | ||
def compute(self) -> compute_options.ComputeOptions: | ||
"""Options controlling object computation.""" | ||
return self._compute_options | ||
"""Thread-local options controlling object computation.""" | ||
if self._local.compute_options is None: | ||
self._local.compute_options = compute_options.ComputeOptions() | ||
|
||
return self._local.compute_options | ||
|
||
@property | ||
def is_bigquery_thread_local(self) -> bool: | ||
"""Indicator that we're using a thread-local session. | ||
|
||
This is set to True by using `option_context` with a "bigquery" option. | ||
""" | ||
return self._local.is_bigquery_thread_local | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This could also check none (and the comment would need to be tweaked.) |
||
|
||
|
||
options = Options() | ||
"""Global options for default session.""" | ||
|
||
option_context = pandas_config.option_context | ||
|
||
|
||
__all__ = ( | ||
"Options", | ||
"options", | ||
"option_context", | ||
) | ||
|
||
|
||
option_context = pandas_config.option_context |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If I'm reading this correctly, would checking for self._local.bigquery_options is None be able to replace this boolean state? It seems a little more DRY to do it that way, if so.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I like it! Done in a7f3c96