Skip to content

Commit f924a34

Browse files
AdamGSAdamGS
authored
Unify all runtime env var access to functions wrapping LazyLock (#7468)
## Summary Env variable access isn't free, this is a fairly common pattern we already used in a couple of places, this PR just makes it consistent across the codebase. ## API Changes Replaced the experimental patches from a public `LazyLock` to a function. I think this should've been a feature an not a runtime thing, but that's a different change. I've also removed the `EnvVarGuard` from `vortex-utils`, it was only used in one place and I've just replaced it with `temp-env` that we already use elsewhere. ## Testing For the few tests that toggled env variables, I've changed them to only run on nextest, relying on it process isolation. --------- Signed-off-by: Adam Gutglick <adam@spiraldb.com>
1 parent 9f9f8f1 commit f924a34

20 files changed

Lines changed: 260 additions & 442 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
@@ -236,6 +236,7 @@ taffy = "0.9.0"
236236
take_mut = "0.2.2"
237237
tar = "0.4"
238238
target-lexicon = "0.13"
239+
temp-env = "0.3"
239240
tempfile = "3"
240241
termtree = { version = "1.0" }
241242
test-with = "0.15"

‎encodings/alp/src/lib.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ use vortex_array::ArrayVTable;
2222
use vortex_array::aggregate_fn::AggregateFnVTable;
2323
use vortex_array::aggregate_fn::fns::nan_count::NanCount;
2424
use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
25-
use vortex_array::arrays::patched::USE_EXPERIMENTAL_PATCHES;
25+
use vortex_array::arrays::patched::use_experimental_patches;
2626
use vortex_array::session::ArraySessionExt;
2727
use vortex_session::VortexSession;
2828

@@ -33,7 +33,7 @@ mod alp_rd;
3333
pub fn initialize(session: &VortexSession) {
3434
// If we're using the experimental Patched encoding, register a shim
3535
// for ALP with interior patches to decode as Patched array.
36-
if *USE_EXPERIMENTAL_PATCHES {
36+
if use_experimental_patches() {
3737
session.arrays().register(ALPPatchedPlugin);
3838
} else {
3939
session.arrays().register(ALP);

‎encodings/fastlanes/src/lib.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,15 +29,15 @@ use vortex_array::aggregate_fn::AggregateFnVTable;
2929
use vortex_array::aggregate_fn::fns::is_constant::IsConstant;
3030
use vortex_array::aggregate_fn::fns::is_sorted::IsSorted;
3131
use vortex_array::aggregate_fn::session::AggregateFnSessionExt;
32-
use vortex_array::arrays::patched::USE_EXPERIMENTAL_PATCHES;
32+
use vortex_array::arrays::patched::use_experimental_patches;
3333
use vortex_array::session::ArraySessionExt;
3434
use vortex_session::VortexSession;
3535

3636
/// Initialize fastlanes encodings in the given session.
3737
pub fn initialize(session: &VortexSession) {
3838
// If we're using the experimental Patched encoding, register a shim
3939
// for BitPacked with interior patches decode as Patched array.
40-
if *USE_EXPERIMENTAL_PATCHES {
40+
if use_experimental_patches() {
4141
session.arrays().register(BitPackedPatchedPlugin);
4242
} else {
4343
session.arrays().register(BitPacked);

‎vortex-array/public-api.lock‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3466,8 +3466,6 @@ pub fn vortex_array::arrays::patched::PatchedSlotsView<'a>::fmt(&self, f: &mut c
34663466

34673467
impl<'a> core::marker::Copy for vortex_array::arrays::patched::PatchedSlotsView<'a>
34683468

3469-
pub static vortex_array::arrays::patched::USE_EXPERIMENTAL_PATCHES: std::sync::lazy_lock::LazyLock<bool>
3470-
34713469
pub trait vortex_array::arrays::patched::PatchedArrayExt: vortex_array::arrays::patched::PatchedArraySlotsExt
34723470

34733471
pub fn vortex_array::arrays::patched::PatchedArrayExt::lane_range(&self, chunk: usize, lane: usize) -> vortex_error::VortexResult<core::ops::range::Range<usize>>
@@ -3512,6 +3510,8 @@ pub fn T::patch_values(&self) -> &vortex_array::ArrayRef
35123510

35133511
pub fn T::slots_view(&self) -> vortex_array::arrays::patched::PatchedSlotsView<'_>
35143512

3513+
pub fn vortex_array::arrays::patched::use_experimental_patches() -> bool
3514+
35153515
pub type vortex_array::arrays::patched::PatchedArray = vortex_array::Array<vortex_array::arrays::patched::Patched>
35163516

35173517
pub mod vortex_array::arrays::primitive

‎vortex-array/src/aggregate_fn/accumulator.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ use crate::aggregate_fn::AggregateFnRef;
1414
use crate::aggregate_fn::AggregateFnVTable;
1515
use crate::aggregate_fn::session::AggregateFnSessionExt;
1616
use crate::dtype::DType;
17-
use crate::executor::MAX_ITERATIONS;
17+
use crate::executor::max_iterations;
1818
use crate::scalar::Scalar;
1919

2020
/// Reference-counted type-erased accumulator.
@@ -108,7 +108,7 @@ impl<V: AggregateFnVTable> DynAccumulator for Accumulator<V> {
108108
let kernels = &session.aggregate_fns().kernels;
109109

110110
let mut batch = batch.clone();
111-
for _ in 0..*MAX_ITERATIONS {
111+
for _ in 0..max_iterations() {
112112
if batch.is::<AnyCanonical>() {
113113
break;
114114
}

‎vortex-array/src/aggregate_fn/accumulator_grouped.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ use crate::builders::builder_with_capacity;
3232
use crate::builtins::ArrayBuiltins;
3333
use crate::dtype::DType;
3434
use crate::dtype::IntegerPType;
35-
use crate::executor::MAX_ITERATIONS;
35+
use crate::executor::max_iterations;
3636
use crate::match_each_integer_ptype;
3737

3838
/// Reference-counted type-erased grouped accumulator.
@@ -165,7 +165,7 @@ impl<V: AggregateFnVTable> GroupedAccumulator<V> {
165165
let session = ctx.session().clone();
166166
let kernels = &session.aggregate_fns().grouped_kernels;
167167

168-
for _ in 0..*MAX_ITERATIONS {
168+
for _ in 0..max_iterations() {
169169
if elements.is::<AnyCanonical>() {
170170
break;
171171
}

‎vortex-array/src/arrays/patched/mod.rs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -108,5 +108,8 @@ const fn patch_lanes<V: Sized>() -> usize {
108108
/// array, and eliminate the interior patches.
109109
///
110110
/// The builtin compressor will also generate Patched arrays.
111-
pub static USE_EXPERIMENTAL_PATCHES: LazyLock<bool> =
112-
LazyLock::new(|| env::var("VORTEX_EXPERIMENTAL_PATCHED_ARRAY").is_ok());
111+
pub fn use_experimental_patches() -> bool {
112+
static USE_EXPERIMENTAL_PATCHES: LazyLock<bool> =
113+
LazyLock::new(|| env::var("VORTEX_EXPERIMENTAL_PATCHED_ARRAY").is_ok_and(|v| v == "1"));
114+
*USE_EXPERIMENTAL_PATCHES
115+
}

‎vortex-array/src/executor.rs‎

Lines changed: 17 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -38,18 +38,21 @@ use crate::memory::HostAllocatorRef;
3838
use crate::memory::MemorySessionExt;
3939
use crate::optimizer::ArrayOptimizer;
4040

41-
/// Maximum number of iterations to attempt when executing an array before giving up and returning
42-
/// an error.
43-
pub(crate) static MAX_ITERATIONS: LazyLock<usize> =
44-
LazyLock::new(|| match std::env::var("VORTEX_MAX_ITERATIONS") {
45-
Ok(val) => val
46-
.parse::<usize>()
47-
.unwrap_or_else(|e| vortex_panic!("VORTEX_MAX_ITERATIONS is not a valid usize: {e}")),
48-
Err(VarError::NotPresent) => 128,
49-
Err(VarError::NotUnicode(_)) => {
50-
vortex_panic!("VORTEX_MAX_ITERATIONS is not a valid unicode string")
51-
}
52-
});
41+
/// Returns the maximum number of iterations to attempt when executing an array before giving up and returning
42+
/// an error, can be by the `VORTEX_MAX_ITERATIONS` env variables, otherwise defaults to 128.
43+
pub(crate) fn max_iterations() -> usize {
44+
static MAX_ITERATIONS: LazyLock<usize> =
45+
LazyLock::new(|| match std::env::var("VORTEX_MAX_ITERATIONS") {
46+
Ok(val) => val.parse::<usize>().unwrap_or_else(|e| {
47+
vortex_panic!("VORTEX_MAX_ITERATIONS is not a valid usize: {e}")
48+
}),
49+
Err(VarError::NotPresent) => 128,
50+
Err(VarError::NotUnicode(_)) => {
51+
vortex_panic!("VORTEX_MAX_ITERATIONS is not a valid unicode string")
52+
}
53+
});
54+
*MAX_ITERATIONS
55+
}
5356

5457
/// Marker trait for types that an [`ArrayRef`] can be executed into.
5558
///
@@ -95,22 +98,11 @@ impl ArrayRef {
9598
/// For safety, we will error when the number of execution iterations reaches a configurable
9699
/// maximum (default 128, override with `VORTEX_MAX_ITERATIONS`).
97100
pub fn execute_until<M: Matcher>(self, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
98-
static MAX_ITERATIONS: LazyLock<usize> =
99-
LazyLock::new(|| match std::env::var("VORTEX_MAX_ITERATIONS") {
100-
Ok(val) => val.parse::<usize>().unwrap_or_else(|e| {
101-
vortex_panic!("VORTEX_MAX_ITERATIONS is not a valid usize: {e}")
102-
}),
103-
Err(VarError::NotPresent) => 128,
104-
Err(VarError::NotUnicode(_)) => {
105-
vortex_panic!("VORTEX_MAX_ITERATIONS is not a valid unicode string")
106-
}
107-
});
108-
109101
let mut current = self.optimize()?;
110102
// Stack frames: (parent, slot_idx, done_predicate_for_slot)
111103
let mut stack: Vec<(ArrayRef, usize, DonePredicate)> = Vec::new();
112104

113-
for _ in 0..*MAX_ITERATIONS {
105+
for _ in 0..max_iterations() {
114106
// Check for termination: use the stack frame's done predicate, or the root matcher.
115107
let is_done = stack
116108
.last()
@@ -179,7 +171,7 @@ impl ArrayRef {
179171

180172
vortex_bail!(
181173
"Exceeded maximum execution iterations ({}) while executing array",
182-
*MAX_ITERATIONS,
174+
max_iterations(),
183175
)
184176
}
185177
}

‎vortex-btrblocks/src/schemes/float.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ use vortex_array::LEGACY_SESSION;
1717
use vortex_array::ToCanonical;
1818
use vortex_array::VortexSessionExecute;
1919
use vortex_array::arrays::Patched;
20-
use vortex_array::arrays::patched::USE_EXPERIMENTAL_PATCHES;
20+
use vortex_array::arrays::patched::use_experimental_patches;
2121
use vortex_array::arrays::primitive::PrimitiveArrayExt;
2222
use vortex_array::dtype::PType;
2323
use vortex_compressor::estimate::CompressionEstimate;
@@ -120,7 +120,7 @@ impl Scheme for ALPScheme {
120120
let alp_stats = alp_encoded.as_array().statistics().to_owned();
121121
let exponents = alp_encoded.exponents();
122122

123-
if *USE_EXPERIMENTAL_PATCHES {
123+
if use_experimental_patches() {
124124
let patches = alp_encoded.patches();
125125

126126
// Create ALP array without interior patches.

0 commit comments

Comments
 (0)