From b4c7985ad66cbc22c4dceeac5ea461d5d254986e Mon Sep 17 00:00:00 2001 From: Jakob Kraus <52459467+JabobKrauskopf@users.noreply.github.com> Date: Fri, 12 Dec 2025 15:28:40 +0100 Subject: [PATCH] feat: add dataframe export --- crates/medmodels-core/src/medrecord/mod.rs | 9 +- crates/medmodels-core/src/medrecord/polars.rs | 232 +++++++++++++++++- medmodels/_medmodels.pyi | 2 + medmodels/medrecord/medrecord.py | 38 +++ medmodels/medrecord/types.py | 28 +++ rustmodels/src/medrecord/mod.rs | 52 +++- rustmodels/src/medrecord/overview.rs | 22 +- tests/medrecord/test_medrecord.py | 36 +++ 8 files changed, 399 insertions(+), 20 deletions(-) diff --git a/crates/medmodels-core/src/medrecord/mod.rs b/crates/medmodels-core/src/medrecord/mod.rs index 2da4072c..dcfedfef 100644 --- a/crates/medmodels-core/src/medrecord/mod.rs +++ b/crates/medmodels-core/src/medrecord/mod.rs @@ -14,7 +14,10 @@ pub use self::{ }; use crate::{ errors::MedRecordError, - medrecord::overview::{GroupOverview, Overview, DEFAULT_TRUNCATE_DETAILS}, + medrecord::{ + overview::{GroupOverview, Overview, DEFAULT_TRUNCATE_DETAILS}, + polars::DataFramesExport, + }, }; use ::polars::frame::DataFrame; use graph::Graph; @@ -243,6 +246,10 @@ impl MedRecord { }) } + pub fn to_dataframes(&self) -> Result { + DataFramesExport::new(self) + } + pub fn set_schema(&mut self, mut schema: Schema) -> Result<(), MedRecordError> { let mut nodes_group_cache = HashMap::<&Group, usize>::new(); let mut nodes_ungrouped_visited = false; diff --git a/crates/medmodels-core/src/medrecord/polars.rs b/crates/medmodels-core/src/medrecord/polars.rs index dc66a2f8..3c0d3e5d 100644 --- a/crates/medmodels-core/src/medrecord/polars.rs +++ b/crates/medmodels-core/src/medrecord/polars.rs @@ -1,9 +1,13 @@ use crate::{ errors::MedRecordError, medrecord::{Attributes, MedRecordAttribute, MedRecordValue, NodeIndex}, + prelude::{EdgeIndex, Group}, + MedRecord, }; use chrono::{DateTime, TimeDelta}; -use polars::{datatypes::AnyValue, frame::DataFrame}; +use medmodels_utils::aliases::MrHashMap; +use polars::{datatypes::AnyValue, frame::DataFrame, prelude::Column}; +use std::collections::HashMap; // TODO: Add tests for Duration impl<'a> TryFrom> for MedRecordValue { @@ -85,6 +89,37 @@ impl<'a> TryFrom> for MedRecordAttribute { } } +impl<'a> From for AnyValue<'a> { + fn from(value: MedRecordValue) -> Self { + match value { + MedRecordValue::String(value) => AnyValue::StringOwned(value.into()), + MedRecordValue::Int(value) => AnyValue::Int64(value), + MedRecordValue::Float(value) => AnyValue::Float64(value), + MedRecordValue::Bool(value) => AnyValue::Boolean(value), + MedRecordValue::DateTime(value) => { + let timestamp = value.and_utc().timestamp_millis(); + + AnyValue::Datetime(timestamp, polars::prelude::TimeUnit::Milliseconds, None) + } + MedRecordValue::Duration(value) => { + let duration_ms = value.num_milliseconds(); + + AnyValue::Duration(duration_ms, polars::prelude::TimeUnit::Milliseconds) + } + MedRecordValue::Null => AnyValue::Null, + } + } +} + +impl<'a> From for AnyValue<'a> { + fn from(value: MedRecordAttribute) -> Self { + match value { + MedRecordAttribute::String(value) => AnyValue::StringOwned(value.into()), + MedRecordAttribute::Int(value) => AnyValue::Int64(value), + } + } +} + pub(crate) fn dataframe_to_nodes( mut nodes: DataFrame, index_column_name: &str, @@ -200,6 +235,201 @@ pub(crate) fn dataframe_to_edges( .collect() } +pub struct DataFramesGroupExport { + pub nodes: DataFrame, + pub edges: DataFrame, +} + +impl DataFramesGroupExport { + fn new(medrecord: &MedRecord, group: Option<&Group>) -> Result { + let group_schema = match group { + Some(group) => medrecord.get_schema().group(group)?, + None => medrecord.get_schema().ungrouped(), + }; + let group_string = match group { + Some(group) => format!("{group}"), + None => "ungrouped".to_string(), + }; + + let node_indices: Box> = match group { + Some(group) => Box::new(medrecord.nodes_in_group(group)?), + None => Box::new(medrecord.ungrouped_nodes()), + }; + + let group_node_attributes = node_indices.map(|node_index| { + ( + node_index, + medrecord + .node_attributes(node_index) + .expect("Node index must exist"), + ) + }); + + let node_attributes: Vec<_> = group_schema.nodes().keys().collect(); + + let mut node_columns: MrHashMap> = node_attributes + .iter() + .map(|attribute_name| ((*attribute_name).clone(), Vec::new())) + .collect(); + + let node_index_attribute = MedRecordAttribute::String("node_index".into()); + + if node_columns.contains_key(&node_index_attribute) { + return Err(MedRecordError::ConversionError( + "Node attribute name 'node_index' is reserved".into(), + )); + } + + node_columns.insert(node_index_attribute.clone(), Vec::new()); + + for (node_index, attributes) in group_node_attributes { + node_columns + .get_mut(&node_index_attribute) + .expect("Attribute must exist in columns") + .push(node_index.clone().into()); + + for attribute_name in node_attributes.iter() { + let attribute_value = attributes + .get(attribute_name) + .cloned() + .unwrap_or(MedRecordValue::Null); + + node_columns + .get_mut(*attribute_name) + .expect("Attribute must exist in columns") + .push(attribute_value.into()); + } + } + + let node_columns: Vec<_> = node_columns + .into_iter() + .map(|(attribute_name, values)| Column::new(attribute_name.to_string().into(), values)) + .collect(); + + let node_dataframe = DataFrame::new(node_columns).map_err(|_| { + MedRecordError::ConversionError(format!( + "Failed to create node DataFrame for group {group_string}" + )) + })?; + + let edge_indices: Box> = match group { + Some(group) => Box::new(medrecord.edges_in_group(group)?), + None => Box::new(medrecord.ungrouped_edges()), + }; + + let group_edge_attributes = edge_indices.map(|edge_index| { + let edge_endpoints = medrecord + .edge_endpoints(edge_index) + .expect("Edge index must exist"); + + ( + edge_index, + edge_endpoints, + medrecord + .edge_attributes(edge_index) + .expect("Edge index must exist"), + ) + }); + + let edge_attributes: Vec<_> = group_schema.edges().keys().collect(); + + let mut edge_columns: MrHashMap> = edge_attributes + .iter() + .map(|attribute_name| ((*attribute_name).clone(), Vec::new())) + .collect(); + + let edge_index_attribute = MedRecordAttribute::String("edge_index".into()); + let source_node_index_attribute = MedRecordAttribute::String("source_node_index".into()); + let target_node_index_attribute = MedRecordAttribute::String("target_node_index".into()); + + if edge_columns.contains_key(&edge_index_attribute) { + return Err(MedRecordError::ConversionError( + "Edge attribute name 'edge_index' is reserved".into(), + )); + } + if edge_columns.contains_key(&source_node_index_attribute) { + return Err(MedRecordError::ConversionError( + "Edge attribute name 'source_node_index' is reserved".into(), + )); + } + if edge_columns.contains_key(&target_node_index_attribute) { + return Err(MedRecordError::ConversionError( + "Edge attribute name 'target_node_index' is reserved".into(), + )); + } + + edge_columns.insert(edge_index_attribute.clone(), Vec::new()); + edge_columns.insert(source_node_index_attribute.clone(), Vec::new()); + edge_columns.insert(target_node_index_attribute.clone(), Vec::new()); + + for (edge_index, edge_endpoints, attributes) in group_edge_attributes { + edge_columns + .get_mut(&edge_index_attribute) + .expect("Attribute must exist in columns") + .push((*edge_index).into()); + edge_columns + .get_mut(&source_node_index_attribute) + .expect("Attribute must exist in columns") + .push(edge_endpoints.0.clone().into()); + edge_columns + .get_mut(&target_node_index_attribute) + .expect("Attribute must exist in columns") + .push(edge_endpoints.1.clone().into()); + + for attribute_name in edge_attributes.iter() { + let attribute_value = attributes + .get(attribute_name) + .cloned() + .unwrap_or(MedRecordValue::Null); + + edge_columns + .get_mut(*attribute_name) + .expect("Attribute must exist in columns") + .push(attribute_value.into()); + } + } + + let edge_columns: Vec<_> = edge_columns + .into_iter() + .map(|(attribute_name, values)| Column::new(attribute_name.to_string().into(), values)) + .collect(); + + let edge_dataframe = DataFrame::new(edge_columns).map_err(|_| { + MedRecordError::ConversionError(format!( + "Failed to create edge DataFrame for group {group_string}" + )) + })?; + + Ok(DataFramesGroupExport { + nodes: node_dataframe, + edges: edge_dataframe, + }) + } +} + +pub struct DataFramesExport { + pub ungrouped: DataFramesGroupExport, + pub groups: HashMap, +} + +impl DataFramesExport { + pub fn new(medrecord: &MedRecord) -> Result { + let ungrouped = DataFramesGroupExport::new(medrecord, None)?; + + let groups = medrecord + .groups() + .map(|group| { + Ok::<_, MedRecordError>(( + group.clone(), + DataFramesGroupExport::new(medrecord, Some(group))?, + )) + }) + .collect::>()?; + + Ok(DataFramesExport { ungrouped, groups }) + } +} + #[cfg(test)] mod test { use super::{dataframe_to_edges, dataframe_to_nodes, MedRecordValue}; diff --git a/medmodels/_medmodels.pyi b/medmodels/_medmodels.pyi index f380ee73..750939da 100644 --- a/medmodels/_medmodels.pyi +++ b/medmodels/_medmodels.pyi @@ -26,6 +26,7 @@ from medmodels.medrecord.types import ( NodeIndex, NodeIndexInputList, NodeTuple, + PolarsDataFramesExport, PolarsEdgeDataFrameInput, PolarsNodeDataFrameInput, PyCategoricalAttributeOverview, @@ -192,6 +193,7 @@ class PyMedRecord: @staticmethod def from_ron(path: str) -> PyMedRecord: ... def to_ron(self, path: str) -> None: ... + def to_dataframes(self) -> PolarsDataFramesExport: ... def get_schema(self) -> PySchema: ... def set_schema(self, schema: PySchema) -> None: ... def freeze_schema(self) -> None: ... diff --git a/medmodels/medrecord/medrecord.py b/medmodels/medrecord/medrecord.py index 9d667c94..dadfc505 100644 --- a/medmodels/medrecord/medrecord.py +++ b/medmodels/medrecord/medrecord.py @@ -132,8 +132,12 @@ NodeIndexInputList, NodeInput, NodeTuple, + PandasDataFramesExport, + PandasDataFramesGroupExport, PandasEdgeDataFrameInput, PandasNodeDataFrameInput, + PolarsDataFramesExport, + PolarsDataFramesGroupExport, PolarsEdgeDataFrameInput, PolarsNodeDataFrameInput, is_edge_tuple, @@ -486,6 +490,40 @@ def to_ron(self, path: str) -> None: """ self._medrecord.to_ron(path) + def to_pandas(self) -> PandasDataFramesExport: + """Exports the MedRecord instance to Pandas DataFrames. + + Returns: + PandasDataFramesExport: A dictionary containing 'ungrouped' and + 'groups' DataFrames. + """ + export = self._medrecord.to_dataframes() + + def _convert_group_export( + group_export: PolarsDataFramesGroupExport, + ) -> PandasDataFramesGroupExport: + return { + "nodes": group_export["nodes"].to_pandas(), + "edges": group_export["edges"].to_pandas(), + } + + return { + "ungrouped": _convert_group_export(export["ungrouped"]), + "groups": { + group: _convert_group_export(group_export) + for group, group_export in export["groups"].items() + }, + } + + def to_polars(self) -> PolarsDataFramesExport: + """Exports the MedRecord instance to Polars DataFrames. + + Returns: + PolarsDataFramesExport: A dictionary containing 'ungrouped' and + 'groups' DataFrames. + """ + return self._medrecord.to_dataframes() + def get_schema(self) -> Schema: """Returns a copy of the MedRecords schema. diff --git a/medmodels/medrecord/types.py b/medmodels/medrecord/types.py index 5a40b5b1..f639c591 100644 --- a/medmodels/medrecord/types.py +++ b/medmodels/medrecord/types.py @@ -187,6 +187,34 @@ class UnstructuredAttributeOverview(TypedDict): distinct_count: int +class PolarsDataFramesGroupExport(TypedDict): + """Dictionary for Polars DataFrames export for a group.""" + + nodes: pl.DataFrame + edges: pl.DataFrame + + +class PolarsDataFramesExport(TypedDict): + """Dictionary for Polars DataFrame export.""" + + ungrouped: PolarsDataFramesGroupExport + groups: Dict[Group, PolarsDataFramesGroupExport] + + +class PandasDataFramesGroupExport(TypedDict): + """Dictionary for Pandas DataFrames export for a group.""" + + nodes: pd.DataFrame + edges: pd.DataFrame + + +class PandasDataFramesExport(TypedDict): + """Dictionary for Pandas DataFrame export.""" + + ungrouped: PandasDataFramesGroupExport + groups: Dict[Group, PandasDataFramesGroupExport] + + def is_medrecord_attribute(value: object) -> TypeIs[MedRecordAttribute]: """Check if a value is a MedRecord attribute. diff --git a/rustmodels/src/medrecord/mod.rs b/rustmodels/src/medrecord/mod.rs index fd912217..d7f39ab9 100644 --- a/rustmodels/src/medrecord/mod.rs +++ b/rustmodels/src/medrecord/mod.rs @@ -19,7 +19,10 @@ use medmodels::core::{ errors::MedRecordError, medrecord::{Attributes, EdgeIndex, MedRecord, MedRecordAttribute, MedRecordValue}, }; -use pyo3::{prelude::*, types::PyFunction}; +use pyo3::{ + prelude::*, + types::{PyDict, PyFunction}, +}; use pyo3_polars::PyDataFrame; use querying::{edges::PyEdgeOperand, nodes::PyNodeOperand, PyReturnOperand, PyReturnValue}; use schema::PySchema; @@ -124,6 +127,53 @@ impl PyMedRecord { Ok(self.0.to_ron(path).map_err(PyMedRecordError::from)?) } + pub fn to_dataframes(&self, py: Python<'_>) -> PyResult { + let export = self.0.to_dataframes().map_err(PyMedRecordError::from)?; + + let outer_dict = PyDict::new(py); + let inner_dict = PyDict::new(py); + + for (group, group_export) in export.groups { + let group_dict = PyDict::new(py); + + let nodes_df = PyDataFrame(group_export.nodes); + group_dict + .set_item("nodes", nodes_df) + .expect("Setting item must succeed"); + + let edges_df = PyDataFrame(group_export.edges); + group_dict + .set_item("edges", edges_df) + .expect("Setting item must succeed"); + + inner_dict + .set_item(PyMedRecordAttribute::from(group), group_dict) + .expect("Setting item must succeed"); + } + + outer_dict + .set_item("groups", inner_dict) + .expect("Setting item must succeed"); + + let ungrouped_dict = PyDict::new(py); + + let nodes_df = PyDataFrame(export.ungrouped.nodes); + ungrouped_dict + .set_item("nodes", nodes_df) + .expect("Setting item must succeed"); + + let edges_df = PyDataFrame(export.ungrouped.edges); + ungrouped_dict + .set_item("edges", edges_df) + .expect("Setting item must succeed"); + + outer_dict + .set_item("ungrouped", ungrouped_dict) + .expect("Setting item must succeed"); + + Ok(outer_dict.into()) + } + pub fn get_schema(&self) -> PySchema { self.0.get_schema().clone().into() } diff --git a/rustmodels/src/medrecord/overview.rs b/rustmodels/src/medrecord/overview.rs index f44c2e0d..fd979e10 100644 --- a/rustmodels/src/medrecord/overview.rs +++ b/rustmodels/src/medrecord/overview.rs @@ -57,10 +57,10 @@ impl PyAttributeOverview { #[getter] pub fn data(&self, py: Python<'_>) -> PyObject { + let dict = PyDict::new(py); + match &self.0.data { AttributeOverviewData::Categorical { distinct_values } => { - let dict = PyDict::new(py); - let distinct_values: Vec = distinct_values.clone().deep_into(); dict.set_item("distinct_values", distinct_values) @@ -73,8 +73,6 @@ impl PyAttributeOverview { .into() } AttributeOverviewData::Continuous { min, mean, max } => { - let dict = PyDict::new(py); - dict.set_item("min", PyMedRecordValue::from(min.clone())) .expect("Setting item must succeed"); dict.set_item("mean", PyMedRecordValue::from(mean.clone())) @@ -84,13 +82,9 @@ impl PyAttributeOverview { dict.set_item("attribute_type", PyAttributeType::Continuous) .expect("Setting item must succeed"); - dict.into_pyobject(py) - .expect("Conversion must succeed") - .into() + dict.into() } AttributeOverviewData::Temporal { min, max } => { - let dict = PyDict::new(py); - dict.set_item("min", PyMedRecordValue::from(min.clone())) .expect("Setting item must succeed"); dict.set_item("max", PyMedRecordValue::from(max.clone())) @@ -98,21 +92,15 @@ impl PyAttributeOverview { dict.set_item("attribute_type", PyAttributeType::Temporal) .expect("Setting item must succeed"); - dict.into_pyobject(py) - .expect("Conversion must succeed") - .into() + dict.into() } AttributeOverviewData::Unstructured { distinct_count } => { - let dict = PyDict::new(py); - dict.set_item("distinct_count", *distinct_count) .expect("Setting item must succeed"); dict.set_item("attribute_type", PyAttributeType::Unstructured) .expect("Setting item must succeed"); - dict.into_pyobject(py) - .expect("Conversion must succeed") - .into() + dict.into() } } } diff --git a/tests/medrecord/test_medrecord.py b/tests/medrecord/test_medrecord.py index dc0de3aa..4be6db6a 100644 --- a/tests/medrecord/test_medrecord.py +++ b/tests/medrecord/test_medrecord.py @@ -343,6 +343,42 @@ def test_ron(self) -> None: assert medrecord.node_count() == loaded_medrecord.node_count() assert medrecord.edge_count() == loaded_medrecord.edge_count() + def test_to_polars(self) -> None: + medrecord = create_medrecord() + + export = medrecord.to_polars() + + assert "ungrouped" in export + assert "nodes" in export["ungrouped"] + assert "edges" in export["ungrouped"] + + nodes_df = export["ungrouped"]["nodes"] + edges_df = export["ungrouped"]["edges"] + + assert isinstance(nodes_df, pl.DataFrame) + assert isinstance(edges_df, pl.DataFrame) + + assert nodes_df.shape[0] == medrecord.node_count() + assert edges_df.shape[0] == medrecord.edge_count() + + def test_to_pandas(self) -> None: + medrecord = create_medrecord() + + export = medrecord.to_pandas() + + assert "ungrouped" in export + assert "nodes" in export["ungrouped"] + assert "edges" in export["ungrouped"] + + nodes_df = export["ungrouped"]["nodes"] + edges_df = export["ungrouped"]["edges"] + + assert isinstance(nodes_df, pd.DataFrame) + assert isinstance(edges_df, pd.DataFrame) + + assert nodes_df.shape[0] == medrecord.node_count() + assert edges_df.shape[0] == medrecord.edge_count() + def test_schema(self) -> None: medrecord = MedRecord()