StructReader::partition_expr doesn't hold lock if the expression is being newly partitioned

vortex-array/src/expr/mod.rs

@@ -106,7 +106,7 @@ fn split_inner(expr: &Expression, exprs: &mut Vec<Expression>) {
 }
 
 /// 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<H: Hasher>(&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");

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<Partitioned> {
         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<Partitioned> {