Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

138 infer schema of medrecord #325

Draft
wants to merge 17 commits into
base: main
Choose a base branch
from
Draft

138 infer schema of medrecord #325

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
2f578e0
feat: initial trials on schema inference
JabobKrauskopf Feb 3, 2025
6b93b5b
feat: separated ProvidedSchema and InferredSchema
JabobKrauskopf Feb 5, 2025
05ed842
feat: finish rust implementation and started python bindings
JabobKrauskopf Feb 6, 2025
09bfe1d
feat: continue python bindings
JabobKrauskopf Feb 6, 2025
5be28a3
refactor: remove schema strictness
JabobKrauskopf Feb 7, 2025
4716085
refactor: merge inferred and provided schema
JabobKrauskopf Feb 7, 2025
9307824
feat: bindings for schema
JabobKrauskopf Feb 12, 2025
ae4b004
feat: bindings for schema and typing
JabobKrauskopf Feb 12, 2025
676ea4e
feat: bindings for schema and typing
JabobKrauskopf Feb 12, 2025
7aed988
fix: typing and tests
JabobKrauskopf Feb 12, 2025
3841ebb
feat: implement new functions
JabobKrauskopf Feb 12, 2025
417491c
feat: add overloads
JabobKrauskopf Feb 18, 2025
1d4cbed
feat: add data-type <-> attribute-type checking to rust schema
JabobKrauskopf Mar 4, 2025
07ccc7a
test: start writing tests
JabobKrauskopf Mar 4, 2025
6850570
refactor: continue tests and refactor AttributeSchema
JabobKrauskopf Mar 4, 2025
443679f
fix: linting
JabobKrauskopf Mar 5, 2025
d3cc258
refactor: improve edge case handling
JabobKrauskopf Mar 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,3 +13,4 @@ target
.python-version
docs/api/_autosummary/*
.coverage
supply-chain
10 changes: 5 additions & 5 deletions crates/medmodels-core/src/medrecord/datatypes/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,9 +104,9 @@ impl Display for DataType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
DataType::String => write!(f, "String"),
DataType::Int => write!(f, "Integer"),
DataType::Int => write!(f, "Int"),
DataType::Float => write!(f, "Float"),
DataType::Bool => write!(f, "Boolean"),
DataType::Bool => write!(f, "Bool"),
DataType::DateTime => write!(f, "DateTime"),
DataType::Duration => write!(f, "Duration"),
DataType::Null => write!(f, "Null"),
Expand Down Expand Up @@ -376,14 +376,14 @@ mod test {
#[test]
fn test_display() {
assert_eq!("String", format!("{}", DataType::String));
assert_eq!("Integer", format!("{}", DataType::Int));
assert_eq!("Int", format!("{}", DataType::Int));
assert_eq!("Float", format!("{}", DataType::Float));
assert_eq!("Boolean", format!("{}", DataType::Bool));
assert_eq!("Bool", format!("{}", DataType::Bool));
assert_eq!("DateTime", format!("{}", DataType::DateTime));
assert_eq!("Null", format!("{}", DataType::Null));
assert_eq!("Any", format!("{}", DataType::Any));
assert_eq!(
"Union[String, Integer]",
"Union[String, Int]",
format!(
"{}",
DataType::Union((Box::new(DataType::String), Box::new(DataType::Int)))
Expand Down
197 changes: 102 additions & 95 deletions crates/medmodels-core/src/medrecord/example_dataset/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
use super::{datatypes::DataType, AttributeType, GroupSchema, MedRecordAttribute, Schema};
use super::{
datatypes::DataType,
schema::{GroupSchema, Schema},
AttributeSchema, AttributeType, MedRecordAttribute,
};
use crate::MedRecord;
use polars::{
io::SerReader,
Expand All @@ -8,28 +12,32 @@ use std::{collections::HashMap, io::Cursor, sync::Arc};

macro_rules! simple_dataset_schema {
() => {
Schema {
groups: HashMap::from([
Schema::new_provided(
HashMap::from([
(
"diagnosis".into(),
GroupSchema {
nodes: HashMap::from([("description".into(), DataType::String.into())]),
edges: HashMap::new(),
strict: Some(true),
},
GroupSchema::new(
AttributeSchema::from([(
"description".into(),
(DataType::String, AttributeType::Unstructured).into(),
)]),
AttributeSchema::default(),
),
),
(
"drug".into(),
GroupSchema {
nodes: HashMap::from([("description".into(), DataType::String.into())]),
edges: HashMap::new(),
strict: Some(true),
},
GroupSchema::new(
AttributeSchema::from([(
"description".into(),
(DataType::String, AttributeType::Unstructured).into(),
)]),
AttributeSchema::default(),
),
),
(
"patient".into(),
GroupSchema {
nodes: HashMap::from([
GroupSchema::new(
AttributeSchema::from([
(
"gender".into(),
(DataType::String, AttributeType::Categorical).into(),
Expand All @@ -39,23 +47,24 @@ macro_rules! simple_dataset_schema {
(DataType::Int, AttributeType::Continuous).into(),
),
]),
edges: HashMap::new(),
strict: Some(true),
},
AttributeSchema::default(),
),
),
(
"procedure".into(),
GroupSchema {
nodes: HashMap::from([("description".into(), DataType::String.into())]),
edges: HashMap::new(),
strict: Some(true),
},
GroupSchema::new(
AttributeSchema::from([(
"description".into(),
(DataType::String, AttributeType::Unstructured).into(),
)]),
AttributeSchema::default(),
),
),
(
"patient_diagnosis".into(),
GroupSchema {
nodes: HashMap::new(),
edges: HashMap::from([
GroupSchema::new(
AttributeSchema::default(),
AttributeSchema::from([
(
"time".into(),
(DataType::DateTime, AttributeType::Temporal).into(),
Expand All @@ -69,14 +78,13 @@ macro_rules! simple_dataset_schema {
.into(),
),
]),
strict: Some(true),
},
),
),
(
"patient_drug".into(),
GroupSchema {
nodes: HashMap::new(),
edges: HashMap::from([
GroupSchema::new(
AttributeSchema::default(),
AttributeSchema::from([
(
"time".into(),
(DataType::DateTime, AttributeType::Temporal).into(),
Expand All @@ -90,14 +98,13 @@ macro_rules! simple_dataset_schema {
(DataType::Float, AttributeType::Continuous).into(),
),
]),
strict: Some(true),
},
),
),
(
"patient_procedure".into(),
GroupSchema {
nodes: HashMap::new(),
edges: HashMap::from([
GroupSchema::new(
AttributeSchema::default(),
AttributeSchema::from([
(
"time".into(),
(DataType::DateTime, AttributeType::Temporal).into(),
Expand All @@ -107,40 +114,42 @@ macro_rules! simple_dataset_schema {
(DataType::Float, AttributeType::Continuous).into(),
),
]),
strict: Some(true),
},
),
),
]),
default: None,
strict: Some(true),
}
GroupSchema::new(Default::default(), Default::default()),
)
};
}

macro_rules! advanced_dataset_schema {
() => {
Schema {
groups: HashMap::from([
Schema::new_provided(
HashMap::from([
(
"diagnosis".into(),
GroupSchema {
nodes: HashMap::from([("description".into(), DataType::String.into())]),
edges: HashMap::new(),
strict: Some(true),
},
GroupSchema::new(
AttributeSchema::from([(
"description".into(),
(DataType::String, AttributeType::Unstructured).into(),
)]),
AttributeSchema::default(),
),
),
(
"drug".into(),
GroupSchema {
nodes: HashMap::from([("description".into(), DataType::String.into())]),
edges: HashMap::new(),
strict: Some(true),
},
GroupSchema::new(
AttributeSchema::from([(
"description".into(),
(DataType::String, AttributeType::Unstructured).into(),
)]),
AttributeSchema::default(),
),
),
(
"patient".into(),
GroupSchema {
nodes: HashMap::from([
GroupSchema::new(
AttributeSchema::from([
(
"gender".into(),
(DataType::String, AttributeType::Categorical).into(),
Expand All @@ -150,31 +159,28 @@ macro_rules! advanced_dataset_schema {
(DataType::Int, AttributeType::Continuous).into(),
),
]),
edges: HashMap::new(),
strict: Some(true),
},
AttributeSchema::default(),
),
),
(
"procedure".into(),
GroupSchema {
nodes: HashMap::from([("description".into(), DataType::String.into())]),
edges: HashMap::new(),
strict: Some(true),
},
GroupSchema::new(
AttributeSchema::from([(
"description".into(),
(DataType::String, AttributeType::Unstructured).into(),
)]),
AttributeSchema::default(),
),
),
(
"event".into(),
GroupSchema {
nodes: HashMap::new(),
edges: HashMap::new(),
strict: Some(true),
},
GroupSchema::new(AttributeSchema::default(), AttributeSchema::default()),
),
(
"patient_diagnosis".into(),
GroupSchema {
nodes: HashMap::new(),
edges: HashMap::from([
GroupSchema::new(
AttributeSchema::default(),
AttributeSchema::from([
(
"time".into(),
(DataType::DateTime, AttributeType::Temporal).into(),
Expand All @@ -188,14 +194,13 @@ macro_rules! advanced_dataset_schema {
.into(),
),
]),
strict: Some(true),
},
),
),
(
"patient_drug".into(),
GroupSchema {
nodes: HashMap::new(),
edges: HashMap::from([
GroupSchema::new(
AttributeSchema::default(),
AttributeSchema::from([
(
"time".into(),
(DataType::DateTime, AttributeType::Temporal).into(),
Expand All @@ -209,14 +214,13 @@ macro_rules! advanced_dataset_schema {
(DataType::Float, AttributeType::Continuous).into(),
),
]),
strict: Some(true),
},
),
),
(
"patient_procedure".into(),
GroupSchema {
nodes: HashMap::new(),
edges: HashMap::from([
GroupSchema::new(
AttributeSchema::default(),
AttributeSchema::from([
(
"time".into(),
(DataType::DateTime, AttributeType::Temporal).into(),
Expand All @@ -226,24 +230,21 @@ macro_rules! advanced_dataset_schema {
(DataType::Float, AttributeType::Continuous).into(),
),
]),
strict: Some(true),
},
),
),
(
"patient_event".into(),
GroupSchema {
nodes: HashMap::new(),
edges: HashMap::from([(
GroupSchema::new(
AttributeSchema::default(),
AttributeSchema::from([(
"time".into(),
(DataType::DateTime, AttributeType::Temporal).into(),
)]),
strict: Some(true),
},
),
),
]),
default: None,
strict: Some(true),
}
GroupSchema::new(Default::default(), Default::default()),
)
};
}

Expand Down Expand Up @@ -460,7 +461,7 @@ impl MedRecord {
)
.expect("Group can be added");

medrecord.schema = simple_dataset_schema!();
unsafe { medrecord.update_schema_unchecked(&mut simple_dataset_schema!()) };

medrecord
}
Expand Down Expand Up @@ -695,16 +696,22 @@ impl MedRecord {
.add_group("patient_event".into(), None, Some(patient_event_ids))
.expect("Group can be added");

medrecord.schema = advanced_dataset_schema!();
unsafe { medrecord.update_schema_unchecked(&mut advanced_dataset_schema!()) };

medrecord
}
}

#[cfg(test)]
mod test {
use super::{AttributeType, DataType, GroupSchema, Schema};
use crate::MedRecord;
use super::{AttributeType, DataType};
use crate::{
medrecord::{
schema::{GroupSchema, Schema},
AttributeSchema,
},
MedRecord,
};
use std::collections::HashMap;

#[test]
Expand Down
8 changes: 4 additions & 4 deletions crates/medmodels-core/src/medrecord/group_mapping.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,10 +7,10 @@ pub type Group = MedRecordAttribute;

#[derive(Debug, Serialize, Deserialize, Clone)]
pub(super) struct GroupMapping {
nodes_in_group: MrHashMap<Group, MrHashSet<NodeIndex>>,
edges_in_group: MrHashMap<Group, MrHashSet<EdgeIndex>>,
groups_of_node: MrHashMap<NodeIndex, MrHashSet<Group>>,
groups_of_edge: MrHashMap<EdgeIndex, MrHashSet<Group>>,
pub(super) nodes_in_group: MrHashMap<Group, MrHashSet<NodeIndex>>,
pub(super) edges_in_group: MrHashMap<Group, MrHashSet<EdgeIndex>>,
pub(super) groups_of_node: MrHashMap<NodeIndex, MrHashSet<Group>>,
pub(super) groups_of_edge: MrHashMap<EdgeIndex, MrHashSet<Group>>,
}

impl GroupMapping {
Expand Down
Loading