refactor: query engine (#231)

limebit · Oct 11, 2024 · 1bd90e8 · 1bd90e8
2 parents 10ebca6 + b8f363c
commit 1bd90e8
Show file tree

Hide file tree

Showing 62 changed files with 12,406 additions and 7,382 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -13,11 +13,8 @@ description = "Limebit MedModels Crate"
 [workspace.dependencies]
 hashbrown = { version = "0.14.5", features = ["serde"] }
 serde = { version = "1.0.203", features = ["derive"] }
-ron = "0.8.1"
-chrono = { version = "0.4.38", features = ["serde"] }
-pyo3 = { version = "0.21.2", features = ["chrono"] }
 polars = { version = "0.40.0", features = ["polars-io"] }
-pyo3-polars = "0.14.0"
+chrono = { version = "0.4.38", features = ["serde"] }
 
 medmodels = { version = "0.1.2", path = "crates/medmodels" }
 medmodels-core = { version = "0.1.2", path = "crates/medmodels-core" }

diff --git a/crates/medmodels-core/Cargo.toml b/crates/medmodels-core/Cargo.toml
@@ -12,5 +12,7 @@ medmodels-utils = { workspace = true }
 
 polars = { workspace = true }
 serde = { workspace = true }
-ron = { workspace = true }
 chrono = { workspace = true }
+ron = "0.8.1"
+roaring = "0.10.6"
+itertools = "0.13.0"
diff --git a/crates/medmodels-core/src/errors/medrecord.rs b/crates/medmodels-core/src/errors/medrecord.rs
@@ -10,6 +10,7 @@ pub enum MedRecordError {
     ConversionError(String),
     AssertionError(String),
     SchemaError(String),
+    QueryError(String),
 }
 
 impl Error for MedRecordError {
@@ -20,6 +21,7 @@ impl Error for MedRecordError {
             MedRecordError::ConversionError(message) => message,
             MedRecordError::AssertionError(message) => message,
             MedRecordError::SchemaError(message) => message,
+            MedRecordError::QueryError(message) => message,
         }
     }
 }
@@ -32,6 +34,7 @@ impl Display for MedRecordError {
             Self::ConversionError(message) => write!(f, "ConversionError: {}", message),
             Self::AssertionError(message) => write!(f, "AssertionError: {}", message),
             Self::SchemaError(message) => write!(f, "SchemaError: {}", message),
+            Self::QueryError(message) => write!(f, "QueryError: {}", message),
         }
     }
 }

diff --git a/crates/medmodels-core/src/errors/mod.rs b/crates/medmodels-core/src/errors/mod.rs
@@ -14,6 +14,8 @@ impl From<GraphError> for MedRecordError {
     }
 }
 
+pub type MedRecordResult<T> = Result<T, MedRecordError>;
+
 #[cfg(test)]
 mod test {
     use super::{GraphError, MedRecordError};

diff --git a/crates/medmodels-core/src/medrecord/datatypes/attribute.rs b/crates/medmodels-core/src/medrecord/datatypes/attribute.rs
@@ -1,8 +1,16 @@
-use super::{Contains, EndsWith, MedRecordValue, StartsWith};
-use crate::errors::MedRecordError;
+use super::{
+    Abs, Contains, EndsWith, Lowercase, MedRecordValue, Mod, Pow, Slice, StartsWith, Trim, TrimEnd,
+    TrimStart, Uppercase,
+};
+use crate::errors::{MedRecordError, MedRecordResult};
 use medmodels_utils::implement_from_for_wrapper;
 use serde::{Deserialize, Serialize};
-use std::{cmp::Ordering, fmt::Display, hash::Hash};
+use std::{
+    cmp::Ordering,
+    fmt::Display,
+    hash::Hash,
+    ops::{Add, Mul, Sub},
+};
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum MedRecordAttribute {
@@ -43,15 +51,6 @@ impl TryFrom<MedRecordValue> for MedRecordAttribute {
     }
 }
 
-impl Display for MedRecordAttribute {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            Self::String(value) => write!(f, "{}", value),
-            Self::Int(value) => write!(f, "{}", value),
-        }
-    }
-}
-
 impl PartialEq for MedRecordAttribute {
     fn eq(&self, other: &Self) -> bool {
         match (self, other) {
@@ -80,6 +79,140 @@ impl PartialOrd for MedRecordAttribute {
     }
 }
 
+impl Display for MedRecordAttribute {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::String(value) => write!(f, "{}", value),
+            Self::Int(value) => write!(f, "{}", value),
+        }
+    }
+}
+
+// TODO: Add tests
+impl Add for MedRecordAttribute {
+    type Output = MedRecordResult<Self>;
+
+    fn add(self, rhs: Self) -> Self::Output {
+        match (self, rhs) {
+            (MedRecordAttribute::String(value), MedRecordAttribute::String(rhs)) => {
+                Ok(MedRecordAttribute::String(value + rhs.as_str()))
+            }
+            (MedRecordAttribute::String(value), MedRecordAttribute::Int(rhs)) => Err(
+                MedRecordError::AssertionError(format!("Cannot add {} to {}", rhs, value)),
+            ),
+            (MedRecordAttribute::Int(value), MedRecordAttribute::String(rhs)) => Err(
+                MedRecordError::AssertionError(format!("Cannot add {} to {}", rhs, value)),
+            ),
+            (MedRecordAttribute::Int(value), MedRecordAttribute::Int(rhs)) => {
+                Ok(MedRecordAttribute::Int(value + rhs))
+            }
+        }
+    }
+}
+
+// TODO: Add tests
+impl Sub for MedRecordAttribute {
+    type Output = MedRecordResult<Self>;
+
+    fn sub(self, rhs: Self) -> Self::Output {
+        match (self, rhs) {
+            (MedRecordAttribute::String(value), MedRecordAttribute::String(rhs)) => Err(
+                MedRecordError::AssertionError(format!("Cannot subtract {} from {}", rhs, value)),
+            ),
+            (MedRecordAttribute::String(value), MedRecordAttribute::Int(rhs)) => Err(
+                MedRecordError::AssertionError(format!("Cannot subtract {} from {}", rhs, value)),
+            ),
+            (MedRecordAttribute::Int(value), MedRecordAttribute::String(rhs)) => Err(
+                MedRecordError::AssertionError(format!("Cannot subtract {} from {}", rhs, value)),
+            ),
+            (MedRecordAttribute::Int(value), MedRecordAttribute::Int(rhs)) => {
+                Ok(MedRecordAttribute::Int(value - rhs))
+            }
+        }
+    }
+}
+
+// TODO: Add tests
+impl Mul for MedRecordAttribute {
+    type Output = MedRecordResult<Self>;
+
+    fn mul(self, rhs: Self) -> Self::Output {
+        match (self, rhs) {
+            (MedRecordAttribute::String(value), MedRecordAttribute::String(rhs)) => Err(
+                MedRecordError::AssertionError(format!("Cannot multiply {} by {}", value, rhs)),
+            ),
+            (MedRecordAttribute::String(value), MedRecordAttribute::Int(rhs)) => Err(
+                MedRecordError::AssertionError(format!("Cannot multiply {} by {}", value, rhs)),
+            ),
+            (MedRecordAttribute::Int(value), MedRecordAttribute::String(rhs)) => Err(
+                MedRecordError::AssertionError(format!("Cannot multiply {} by {}", value, rhs)),
+            ),
+            (MedRecordAttribute::Int(value), MedRecordAttribute::Int(rhs)) => {
+                Ok(MedRecordAttribute::Int(value * rhs))
+            }
+        }
+    }
+}
+
+// TODO: Add tests
+impl Pow for MedRecordAttribute {
+    fn pow(self, rhs: Self) -> MedRecordResult<Self> {
+        match (self, rhs) {
+            (MedRecordAttribute::String(value), MedRecordAttribute::String(rhs)) => {
+                Err(MedRecordError::AssertionError(format!(
+                    "Cannot raise {} to the power of {}",
+                    value, rhs
+                )))
+            }
+            (MedRecordAttribute::String(value), MedRecordAttribute::Int(rhs)) => {
+                Err(MedRecordError::AssertionError(format!(
+                    "Cannot raise {} to the power of {}",
+                    value, rhs
+                )))
+            }
+            (MedRecordAttribute::Int(value), MedRecordAttribute::String(rhs)) => {
+                Err(MedRecordError::AssertionError(format!(
+                    "Cannot raise {} to the power of {}",
+                    value, rhs
+                )))
+            }
+            (MedRecordAttribute::Int(value), MedRecordAttribute::Int(rhs)) => {
+                Ok(MedRecordAttribute::Int(value.pow(rhs as u32)))
+            }
+        }
+    }
+}
+
+// TODO: Add tests
+impl Mod for MedRecordAttribute {
+    fn r#mod(self, rhs: Self) -> MedRecordResult<Self> {
+        match (self, rhs) {
+            (MedRecordAttribute::String(value), MedRecordAttribute::String(rhs)) => Err(
+                MedRecordError::AssertionError(format!("Cannot mod {} by {}", value, rhs)),
+            ),
+            (MedRecordAttribute::String(value), MedRecordAttribute::Int(rhs)) => Err(
+                MedRecordError::AssertionError(format!("Cannot mod {} by {}", value, rhs)),
+            ),
+            (MedRecordAttribute::Int(value), MedRecordAttribute::String(rhs)) => Err(
+                MedRecordError::AssertionError(format!("Cannot mod {} by {}", value, rhs)),
+            ),
+            (MedRecordAttribute::Int(value), MedRecordAttribute::Int(rhs)) => {
+                Ok(MedRecordAttribute::Int(value % rhs))
+            }
+        }
+    }
+}
+
+// TODO: Add tests
+impl Abs for MedRecordAttribute {
+    fn abs(self) -> Self {
+        match self {
+            MedRecordAttribute::Int(value) => MedRecordAttribute::Int(value.abs()),
+            _ => self,
+        }
+    }
+}
+
 impl StartsWith for MedRecordAttribute {
     fn starts_with(&self, other: &Self) -> bool {
         match (self, other) {
@@ -137,6 +270,72 @@ impl Contains for MedRecordAttribute {
     }
 }
 
+// TODO: Add tests
+impl Slice for MedRecordAttribute {
+    fn slice(self, range: std::ops::Range<usize>) -> Self {
+        match self {
+            MedRecordAttribute::String(value) => value[range].into(),
+            MedRecordAttribute::Int(value) => value.to_string()[range].into(),
+        }
+    }
+}
+
+// TODO: Add tests
+impl Trim for MedRecordAttribute {
+    fn trim(self) -> Self {
+        match self {
+            MedRecordAttribute::String(value) => {
+                MedRecordAttribute::String(value.trim().to_string())
+            }
+            _ => self,
+        }
+    }
+}
+
+// TODO: Add tests
+impl TrimStart for MedRecordAttribute {
+    fn trim_start(self) -> Self {
+        match self {
+            MedRecordAttribute::String(value) => {
+                MedRecordAttribute::String(value.trim_start().to_string())
+            }
+            _ => self,
+        }
+    }
+}
+
+// TODO: Add tests
+impl TrimEnd for MedRecordAttribute {
+    fn trim_end(self) -> Self {
+        match self {
+            MedRecordAttribute::String(value) => {
+                MedRecordAttribute::String(value.trim_end().to_string())
+            }
+            _ => self,
+        }
+    }
+}
+
+// TODO: Add tests
+impl Lowercase for MedRecordAttribute {
+    fn lowercase(self) -> Self {
+        match self {
+            MedRecordAttribute::String(value) => MedRecordAttribute::String(value.to_lowercase()),
+            _ => self,
+        }
+    }
+}
+
+// TODO: Add tests
+impl Uppercase for MedRecordAttribute {
+    fn uppercase(self) -> Self {
+        match self {
+            MedRecordAttribute::String(value) => MedRecordAttribute::String(value.to_uppercase()),
+            _ => self,
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::MedRecordAttribute;