From 8a0d844bab9096aec616965751cfc4690c9a3e89 Mon Sep 17 00:00:00 2001
From: Sammy Sidhu <samster25@users.noreply.github.com>
Date: Tue, 28 May 2024 22:48:33 -0500
Subject: [PATCH] [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">
---
 Cargo.lock       |  1 +
 Cargo.toml       |  1 +
 daft/__init__.py |  7 +++++++
 daft/daft.pyi    |  1 +
 src/lib.rs       | 14 ++++++++++++--
 5 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 298f61fc5d..fb2a9ed0c7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1179,6 +1179,7 @@ dependencies = [
  "daft-scan",
  "daft-stats",
  "daft-table",
+ "lazy_static",
  "libc",
  "lzma-sys",
  "pyo3",
diff --git a/Cargo.toml b/Cargo.toml
index 22c0ec458c..b77e10d980 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -13,6 +13,7 @@ daft-plan = {path = "src/daft-plan", default-features = false}
 daft-scan = {path = "src/daft-scan", default-features = false}
 daft-stats = {path = "src/daft-stats", default-features = false}
 daft-table = {path = "src/daft-table", default-features = false}
+lazy_static = {workspace = true}
 lzma-sys = {version = "*", features = ["static"]}
 pyo3 = {workspace = true, optional = true}
 pyo3-log = {workspace = true, optional = true}
diff --git a/daft/__init__.py b/daft/__init__.py
index 18063fc996..e6c614f262 100644
--- a/daft/__init__.py
+++ b/daft/__init__.py
@@ -25,6 +25,7 @@
 
 from daft.daft import build_type as _build_type
 from daft.daft import version as _version
+from daft.daft import refresh_logger as _refresh_logger
 
 
 def get_version() -> str:
@@ -35,6 +36,11 @@ def get_build_type() -> str:
     return _build_type()
 
 
+def refresh_logger() -> None:
+    """Refreshes Daft's internal rust logging to the current python log level"""
+    _refresh_logger()
+
+
 __version__ = get_version()
 
 
@@ -112,6 +118,7 @@ def get_build_type() -> str:
     "Series",
     "TimeUnit",
     "register_viz_hook",
+    "refresh_logger",
     "udf",
     "ResourceRequest",
     "set_planning_config",
diff --git a/daft/daft.pyi b/daft/daft.pyi
index 6aefd6fe82..270a4cf700 100644
--- a/daft/daft.pyi
+++ b/daft/daft.pyi
@@ -1498,6 +1498,7 @@ class PyDaftPlanningConfig:
 
 def build_type() -> str: ...
 def version() -> str: ...
+def refresh_logger() -> None: ...
 def __getattr__(name) -> Any: ...
 def io_glob(
     path: str,
diff --git a/src/lib.rs b/src/lib.rs
index 6ca4d5a317..dd77645316 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -33,8 +33,11 @@ pub static malloc_conf: Option<&'static libc::c_char> = Some(unsafe {
 
 #[cfg(feature = "python")]
 pub mod pylib {
+    use lazy_static::lazy_static;
     use pyo3::prelude::*;
-
+    lazy_static! {
+        static ref LOG_RESET_HANDLE: pyo3_log::ResetHandle = pyo3_log::init();
+    }
     #[pyfunction]
     pub fn version() -> &'static str {
         daft_core::VERSION
@@ -45,9 +48,15 @@ pub mod pylib {
         daft_core::DAFT_BUILD_TYPE
     }
 
+    #[pyfunction]
+    pub fn refresh_logger() {
+        LOG_RESET_HANDLE.reset();
+    }
+
     #[pymodule]
     fn daft(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
-        pyo3_log::init();
+        refresh_logger();
+
         common_daft_config::register_modules(_py, m)?;
         common_system_info::register_modules(_py, m)?;
         daft_core::register_modules(_py, m)?;
@@ -63,6 +72,7 @@ pub mod pylib {
         daft_scan::register_modules(_py, m)?;
         m.add_wrapped(wrap_pyfunction!(version))?;
         m.add_wrapped(wrap_pyfunction!(build_type))?;
+        m.add_wrapped(wrap_pyfunction!(refresh_logger))?;
         Ok(())
     }
 }