feat: add dataframe export

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, MedRecordError> {
+        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;

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<AnyValue<'a>> for MedRecordValue {
@@ -85,6 +89,37 @@ impl<'a> TryFrom<AnyValue<'a>> for MedRecordAttribute {
     }
 }
 
+impl<'a> From<MedRecordValue> 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<MedRecordAttribute> 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<Self, MedRecordError> {
+        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<dyn Iterator<Item = &NodeIndex>> = 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<MedRecordAttribute, Vec<AnyValue>> = 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<dyn Iterator<Item = &EdgeIndex>> = 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<MedRecordAttribute, Vec<AnyValue>> = 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<Group, DataFramesGroupExport>,
+}
+
+impl DataFramesExport {
+    pub fn new(medrecord: &MedRecord) -> Result<Self, MedRecordError> {
+        let ungrouped = DataFramesGroupExport::new(medrecord, None)?;
+
+        let groups = medrecord
+            .groups()
+            .map(|group| {
+                Ok::<_, MedRecordError>((
+                    group.clone(),
+                    DataFramesGroupExport::new(medrecord, Some(group))?,
+                ))
+            })
+            .collect::<Result<_, _>>()?;
+
+        Ok(DataFramesExport { ungrouped, groups })
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::{dataframe_to_edges, dataframe_to_nodes, MedRecordValue};

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: ...

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.
 

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.