dag_ml_data_core/
plan.rs

1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::error::{DataError, Result};
6use crate::ids::{RepresentationId, SourceId};
7
8#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
9#[serde(rename_all = "snake_case")]
10pub enum FitScope {
11    Stateless,
12    FoldTrain,
13    FullTrain,
14    InferenceOnly,
15}
16
17#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
18#[serde(rename_all = "snake_case")]
19pub enum DataPlanStepKind {
20    Materialize,
21    Adapt,
22    Align,
23    Join,
24    Collate,
25}
26
27#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
28pub struct DataPlanStep {
29    pub kind: DataPlanStepKind,
30    pub source_id: Option<SourceId>,
31    pub adapter_id: Option<String>,
32    pub input_representation: Option<RepresentationId>,
33    pub output_representation: Option<RepresentationId>,
34    pub fit_scope: FitScope,
35    #[serde(default)]
36    pub requires_user_choice: bool,
37    #[serde(default)]
38    pub metadata: BTreeMap<String, serde_json::Value>,
39}
40
41#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
42pub struct PlanIssue {
43    pub code: String,
44    pub message: String,
45    #[serde(default)]
46    pub choices: Vec<String>,
47}
48
49#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
50pub struct DataPlan {
51    pub id: String,
52    pub steps: Vec<DataPlanStep>,
53    pub output_representation: RepresentationId,
54    #[serde(default)]
55    pub issues: Vec<PlanIssue>,
56}
57
58impl DataPlan {
59    pub fn validate(&self) -> Result<()> {
60        if self.id.trim().is_empty() {
61            return Err(DataError::Validation("data plan id is empty".to_string()));
62        }
63        if self.steps.is_empty() {
64            return Err(DataError::Validation(format!(
65                "data plan `{}` contains no steps",
66                self.id
67            )));
68        }
69        for (idx, step) in self.steps.iter().enumerate() {
70            match step.kind {
71                DataPlanStepKind::Materialize if step.source_id.is_none() => {
72                    return Err(DataError::Validation(format!(
73                        "data plan `{}` step {} materializes without source_id",
74                        self.id, idx
75                    )));
76                }
77                DataPlanStepKind::Adapt if step.adapter_id.is_none() => {
78                    return Err(DataError::Validation(format!(
79                        "data plan `{}` step {} adapts without adapter_id",
80                        self.id, idx
81                    )));
82                }
83                _ => {}
84            }
85        }
86        Ok(())
87    }
88
89    pub fn requires_user_choice(&self) -> bool {
90        self.steps.iter().any(|step| step.requires_user_choice) || !self.issues.is_empty()
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use crate::ids::{RepresentationId, SourceId};
98
99    #[test]
100    fn rejects_empty_plan() {
101        let plan = DataPlan {
102            id: "p".to_string(),
103            steps: vec![],
104            output_representation: RepresentationId::new("tabular").unwrap(),
105            issues: vec![],
106        };
107
108        assert!(plan.validate().is_err());
109    }
110
111    #[test]
112    fn flags_user_choice() {
113        let plan = DataPlan {
114            id: "p".to_string(),
115            steps: vec![DataPlanStep {
116                kind: DataPlanStepKind::Materialize,
117                source_id: Some(SourceId::new("nir").unwrap()),
118                adapter_id: None,
119                input_representation: None,
120                output_representation: Some(RepresentationId::new("signal").unwrap()),
121                fit_scope: FitScope::Stateless,
122                requires_user_choice: true,
123                metadata: BTreeMap::new(),
124            }],
125            output_representation: RepresentationId::new("signal").unwrap(),
126            issues: vec![],
127        };
128
129        assert!(plan.validate().is_ok());
130        assert!(plan.requires_user_choice());
131    }
132}
dag_ml_data_core/plan.rs

dag_ml_data_core/
plan.rs
