refactor: Deduplicate POLARS_FORCE_ASYNC env var parsing (#14909)

pola-rs · Mar 8, 2024 · c6f37cf · c6f37cf
1 parent 48b469d
commit c6f37cf
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/crates/polars-core/src/config.rs b/crates/polars-core/src/config.rs
@@ -54,3 +54,9 @@ pub fn get_rg_prefetch_size() -> usize {
         // Set it to something big, but not unlimited.
         .unwrap_or_else(|_| std::cmp::max(get_file_prefetch_size(), 128))
 }
+
+pub fn env_force_async() -> bool {
+    std::env::var("POLARS_FORCE_ASYNC")
+        .map(|value| value == "1")
+        .unwrap_or_default()
+}
diff --git a/crates/polars-lazy/src/physical_plan/executors/scan/parquet.rs b/crates/polars-lazy/src/physical_plan/executors/scan/parquet.rs
@@ -1,5 +1,6 @@
 use std::path::PathBuf;
 
+use polars_core::config::env_force_async;
 #[cfg(feature = "cloud")]
 use polars_core::config::{get_file_prefetch_size, verbose};
 use polars_core::utils::accumulate_dataframes_vertical;
@@ -334,7 +335,7 @@ impl ParquetExec {
                 ));
             },
         };
-        let force_async = std::env::var("POLARS_FORCE_ASYNC").as_deref().unwrap_or("") == "1";
+        let force_async = env_force_async();
 
         let out = if is_cloud || force_async {
             #[cfg(not(feature = "cloud"))]

diff --git a/crates/polars-pipe/src/executors/sources/parquet.rs b/crates/polars-pipe/src/executors/sources/parquet.rs
@@ -4,7 +4,7 @@ use std::path::PathBuf;
 use std::sync::Arc;
 
 use arrow::datatypes::ArrowSchemaRef;
-use polars_core::config::get_file_prefetch_size;
+use polars_core::config::{env_force_async, get_file_prefetch_size};
 use polars_core::error::*;
 use polars_core::prelude::Series;
 use polars_core::POOL;
@@ -204,8 +204,7 @@ impl ParquetSource {
         if verbose {
             eprintln!("POLARS PREFETCH_SIZE: {}", prefetch_size)
         }
-        let run_async = paths.first().map(is_cloud_url).unwrap_or(false)
-            || std::env::var("POLARS_FORCE_ASYNC").as_deref().unwrap_or("") == "1";
+        let run_async = paths.first().map(is_cloud_url).unwrap_or(false) || env_force_async();
 
         let mut source = ParquetSource {
             batched_readers: VecDeque::new(),