diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 4331b2817f0..15fc88a4c23 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -106,7 +106,7 @@ fn split_inner(expr: &Expression, exprs: &mut Vec) { } /// An expression wrapper that performs pointer equality on child expressions. -#[derive(Clone)] +#[derive(Clone, Debug)] pub struct ExactExpr(pub Expression); impl PartialEq for ExactExpr { fn eq(&self, other: &Self) -> bool { @@ -118,7 +118,8 @@ impl Eq for ExactExpr {} impl Hash for ExactExpr { fn hash(&self, state: &mut H) { - self.0.hash(state); + self.0.scalar_fn().hash(state); + Arc::as_ptr(self.0.children()).hash(state); } } @@ -148,6 +149,9 @@ pub mod test_harness { #[cfg(test)] mod tests { + use std::collections::hash_map::RandomState; + use std::hash::BuildHasher; + use super::*; use crate::dtype::DType; use crate::dtype::FieldNames; @@ -189,6 +193,22 @@ mod tests { assert_eq!(conjunction.len(), 2, "Conjunction is {conjunction:?}"); } + #[test] + fn exact_expr_hash_consistent_with_eq() { + let state = RandomState::new(); + let expr = eq(get_item("col1", root()), lit(1)); + + // Clones share the children Arc, so they are equal and must hash equally. + let a = ExactExpr(expr.clone()); + let b = ExactExpr(expr); + assert_eq!(a, b); + assert_eq!(state.hash_one(&a), state.hash_one(&b)); + + // Structurally identical expressions built separately are distinct keys. + let rebuilt = ExactExpr(eq(get_item("col1", root()), lit(1))); + assert_ne!(a, rebuilt); + } + #[test] fn expr_display() { assert_eq!(col("a").to_string(), "$.a"); diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 18a14a2c97b..0dae97f0de0 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -157,17 +157,26 @@ impl StructReader { /// Utility for partitioning an expression over the fields of a struct. fn partition_expr(&self, expr: Expression) -> VortexResult { let key = ExactExpr(expr.clone()); - let binding = self - .partitioned_expr_cache - .entry(key) - .or_insert_with(|| Arc::new(OnceLock::new())); - let entry = binding.value(); - if let Some(value) = entry.get() { + + // Look up the cell under a shared shard lock; only a miss takes the write lock, and + // only for as long as it takes to insert an empty cell. + let cell = match self.partitioned_expr_cache.get(&key) { + Some(entry) => Arc::clone(entry.value()), + None => Arc::clone( + self.partitioned_expr_cache + .entry(key) + .or_insert_with(|| Arc::new(OnceLock::new())) + .value(), + ), + }; + // All map guards are dropped here, so partitioning runs outside any shard lock. + // Concurrent misses may compute redundantly; `get_or_init` keeps a single winner. + + if let Some(value) = cell.get() { return Ok(value.clone()); } let result = self.compute_partitioned_expr(expr)?; - let result = entry.get_or_init(|| result); - Ok(result.clone()) + Ok(cell.get_or_init(|| result).clone()) } fn compute_partitioned_expr(&self, expr: Expression) -> VortexResult {