Make NodeMap Thread-safe

partiql-common/Cargo.toml

@@ -25,11 +25,13 @@ pretty = "0.12"
 serde = { version = "1.*", features = ["derive"], optional = true }
 smallvec = { version = "1.*" }
 thiserror = "1.0"
+dashmap = "6.0"
 
 [features]
 default = []
 serde = [
     "dep:serde",
     "indexmap/serde",
-    "smallvec/serde"
+    "smallvec/serde",
+    "dashmap/serde"
 ]

partiql-common/src/node.rs

@@ -1,11 +1,52 @@
-use indexmap::IndexMap;
-use std::hash::Hash;
+use dashmap::DashMap;
 use std::sync::{Arc, RwLock};
 
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 
-pub type NodeMap<T> = IndexMap<NodeId, T>;
+#[derive(Debug, Clone)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+pub struct NodeMap<T> {
+    map: DashMap<NodeId, T>,
+    #[cfg_attr(feature = "serde", serde(skip))]
+    order: Arc<RwLock<Vec<NodeId>>>,
+}
+
+impl<T> Default for NodeMap<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl<T> NodeMap<T> {
+    pub fn new() -> Self {
+        NodeMap {
+            map: DashMap::new(),
+            order: Arc::new(RwLock::new(Vec::new())),
+        }
+    }
+
+    pub fn with_capacity(capacity: usize) -> Self {
+        NodeMap {
+            map: DashMap::with_capacity(capacity),
+            order: Arc::new(RwLock::new(Vec::with_capacity(capacity))),
+        }
+    }
+
+    pub fn insert(&self, node_id: NodeId, value: T) -> Option<T> {
+        let mut order = self.order.write().expect("NodeMap order write lock");
+        if self.map.contains_key(&node_id) {
+            self.map.insert(node_id, value)
+        } else {
+            order.push(node_id);
+            self.map.insert(node_id, value)
+        }
+    }
+
+    pub fn get(&self, node_id: &NodeId) -> Option<dashmap::mapref::one::Ref<'_, NodeId, T>> {
+        self.map.get(node_id)
+    }
+}
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]

partiql-eval/Cargo.toml

@@ -41,6 +41,9 @@ delegate = "0.12"
 [dev-dependencies]
 criterion = "0.4"
 
+[features]
+default = []
+
 [[bench]]
 name = "bench_eval"
 harness = false

partiql-eval/src/eval/mod.rs

@@ -146,7 +146,6 @@ pub type EvalResult = Result<Evaluated, EvalErr>;
 /// Represents result of evaluation as an evaluated entity.
 #[non_exhaustive]
 #[derive(Debug)]
-#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 pub struct Evaluated {
     pub result: Value,
 }

partiql-parser/Cargo.toml

@@ -36,6 +36,7 @@ bigdecimal = "~0.2.0"
 rust_decimal = { version = "1.25.0", default-features = false, features = ["std"] }
 
 bitflags = "2"
+dashmap = "6.0.1"
 
 lalrpop-util = "0.20"
 logos = "0.12"

partiql-parser/src/parse/partiql.lalrpop

@@ -287,8 +287,8 @@ FromClause: ast::AstNode<ast::FromClause> = {
                     ast::FromSource::Join(node) => node.id,
                 };
 
-	            let start = state.locations.get(&start_id).unwrap_or(&total).start.0.clone();
-                let end = state.locations.get(&end_id).unwrap_or(&total).end.0.clone();
+                let start = state.locations.get(&start_id).map_or(total.start.0.clone(), |v| v.start.0.clone());
+                let end = state.locations.get(&end_id).map_or(total.end.0.clone(), |v| v.end.0.clone());
                 let range = start..end;
                 let join = state.node(ast::Join {
                     kind: ast::JoinKind::Cross,