Skip to content

Commit 57962f4

Browse files
authored
StructReader::partition_expr doesn't hold lock if the expression is being newly partitioned (#8892)
We also simplify ExactExpr hash behaviour to match the equality semantics and avoid deep tree hashing --------- Signed-off-by: Robert Kruszewski <github@robertk.io>
1 parent 759562a commit 57962f4

2 files changed

Lines changed: 39 additions & 10 deletions

File tree

vortex-array/src/expr/mod.rs

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,7 @@ fn split_inner(expr: &Expression, exprs: &mut Vec<Expression>) {
106106
}
107107

108108
/// An expression wrapper that performs pointer equality on child expressions.
109-
#[derive(Clone)]
109+
#[derive(Clone, Debug)]
110110
pub struct ExactExpr(pub Expression);
111111
impl PartialEq for ExactExpr {
112112
fn eq(&self, other: &Self) -> bool {
@@ -118,7 +118,8 @@ impl Eq for ExactExpr {}
118118

119119
impl Hash for ExactExpr {
120120
fn hash<H: Hasher>(&self, state: &mut H) {
121-
self.0.hash(state);
121+
self.0.scalar_fn().hash(state);
122+
Arc::as_ptr(self.0.children()).hash(state);
122123
}
123124
}
124125

@@ -148,6 +149,9 @@ pub mod test_harness {
148149

149150
#[cfg(test)]
150151
mod tests {
152+
use std::collections::hash_map::RandomState;
153+
use std::hash::BuildHasher;
154+
151155
use super::*;
152156
use crate::dtype::DType;
153157
use crate::dtype::FieldNames;
@@ -189,6 +193,22 @@ mod tests {
189193
assert_eq!(conjunction.len(), 2, "Conjunction is {conjunction:?}");
190194
}
191195

196+
#[test]
197+
fn exact_expr_hash_consistent_with_eq() {
198+
let state = RandomState::new();
199+
let expr = eq(get_item("col1", root()), lit(1));
200+
201+
// Clones share the children Arc, so they are equal and must hash equally.
202+
let a = ExactExpr(expr.clone());
203+
let b = ExactExpr(expr);
204+
assert_eq!(a, b);
205+
assert_eq!(state.hash_one(&a), state.hash_one(&b));
206+
207+
// Structurally identical expressions built separately are distinct keys.
208+
let rebuilt = ExactExpr(eq(get_item("col1", root()), lit(1)));
209+
assert_ne!(a, rebuilt);
210+
}
211+
192212
#[test]
193213
fn expr_display() {
194214
assert_eq!(col("a").to_string(), "$.a");

vortex-layout/src/layouts/struct_/reader.rs

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -157,17 +157,26 @@ impl StructReader {
157157
/// Utility for partitioning an expression over the fields of a struct.
158158
fn partition_expr(&self, expr: Expression) -> VortexResult<Partitioned> {
159159
let key = ExactExpr(expr.clone());
160-
let binding = self
161-
.partitioned_expr_cache
162-
.entry(key)
163-
.or_insert_with(|| Arc::new(OnceLock::new()));
164-
let entry = binding.value();
165-
if let Some(value) = entry.get() {
160+
161+
// Look up the cell under a shared shard lock; only a miss takes the write lock, and
162+
// only for as long as it takes to insert an empty cell.
163+
let cell = match self.partitioned_expr_cache.get(&key) {
164+
Some(entry) => Arc::clone(entry.value()),
165+
None => Arc::clone(
166+
self.partitioned_expr_cache
167+
.entry(key)
168+
.or_insert_with(|| Arc::new(OnceLock::new()))
169+
.value(),
170+
),
171+
};
172+
// All map guards are dropped here, so partitioning runs outside any shard lock.
173+
// Concurrent misses may compute redundantly; `get_or_init` keeps a single winner.
174+
175+
if let Some(value) = cell.get() {
166176
return Ok(value.clone());
167177
}
168178
let result = self.compute_partitioned_expr(expr)?;
169-
let result = entry.get_or_init(|| result);
170-
Ok(result.clone())
179+
Ok(cell.get_or_init(|| result).clone())
171180
}
172181

173182
fn compute_partitioned_expr(&self, expr: Expression) -> VortexResult<Partitioned> {

0 commit comments

Comments
 (0)