Skip to content

Commit 8a0d844

Browse files
authored
[FEAT] Add function to refresh logger state for rust (#2323)
When daft is imported it cached the current python log level to avoid grabbing the GIL for every log statement. See [this page](https://docs.rs/pyo3-log/latest/pyo3_log/#performance-filtering-and-caching) for more details. This PR exposes a method to "refresh" the cache. This is useful when the user updates the log level later in the code, like when they configure logging. Here is a example showing this. Note prior to refreshing the logger, no logs were coming out! <img width="521" alt="image" src="https://github.com/Eventual-Inc/Daft/assets/2550285/de7ad2d8-8335-4dd0-b7d5-8b36ef28d7ad">
1 parent 51451f9 commit 8a0d844

5 files changed

Lines changed: 22 additions & 2 deletions

File tree

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ daft-plan = {path = "src/daft-plan", default-features = false}
1313
daft-scan = {path = "src/daft-scan", default-features = false}
1414
daft-stats = {path = "src/daft-stats", default-features = false}
1515
daft-table = {path = "src/daft-table", default-features = false}
16+
lazy_static = {workspace = true}
1617
lzma-sys = {version = "*", features = ["static"]}
1718
pyo3 = {workspace = true, optional = true}
1819
pyo3-log = {workspace = true, optional = true}

daft/__init__.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525

2626
from daft.daft import build_type as _build_type
2727
from daft.daft import version as _version
28+
from daft.daft import refresh_logger as _refresh_logger
2829

2930

3031
def get_version() -> str:
@@ -35,6 +36,11 @@ def get_build_type() -> str:
3536
return _build_type()
3637

3738

39+
def refresh_logger() -> None:
40+
"""Refreshes Daft's internal rust logging to the current python log level"""
41+
_refresh_logger()
42+
43+
3844
__version__ = get_version()
3945

4046

@@ -112,6 +118,7 @@ def get_build_type() -> str:
112118
"Series",
113119
"TimeUnit",
114120
"register_viz_hook",
121+
"refresh_logger",
115122
"udf",
116123
"ResourceRequest",
117124
"set_planning_config",

daft/daft.pyi

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1498,6 +1498,7 @@ class PyDaftPlanningConfig:
14981498

14991499
def build_type() -> str: ...
15001500
def version() -> str: ...
1501+
def refresh_logger() -> None: ...
15011502
def __getattr__(name) -> Any: ...
15021503
def io_glob(
15031504
path: str,

src/lib.rs

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,11 @@ pub static malloc_conf: Option<&'static libc::c_char> = Some(unsafe {
3333

3434
#[cfg(feature = "python")]
3535
pub mod pylib {
36+
use lazy_static::lazy_static;
3637
use pyo3::prelude::*;
37-
38+
lazy_static! {
39+
static ref LOG_RESET_HANDLE: pyo3_log::ResetHandle = pyo3_log::init();
40+
}
3841
#[pyfunction]
3942
pub fn version() -> &'static str {
4043
daft_core::VERSION
@@ -45,9 +48,15 @@ pub mod pylib {
4548
daft_core::DAFT_BUILD_TYPE
4649
}
4750

51+
#[pyfunction]
52+
pub fn refresh_logger() {
53+
LOG_RESET_HANDLE.reset();
54+
}
55+
4856
#[pymodule]
4957
fn daft(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
50-
pyo3_log::init();
58+
refresh_logger();
59+
5160
common_daft_config::register_modules(_py, m)?;
5261
common_system_info::register_modules(_py, m)?;
5362
daft_core::register_modules(_py, m)?;
@@ -63,6 +72,7 @@ pub mod pylib {
6372
daft_scan::register_modules(_py, m)?;
6473
m.add_wrapped(wrap_pyfunction!(version))?;
6574
m.add_wrapped(wrap_pyfunction!(build_type))?;
75+
m.add_wrapped(wrap_pyfunction!(refresh_logger))?;
6676
Ok(())
6777
}
6878
}

0 commit comments

Comments
 (0)