1use std::collections::BTreeMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::adapter::{AdapterPath, AdapterRegistry, PlanningPolicy};
6use crate::alignment::{alignment_metadata, alignment_mode_from_fusion};
7use crate::error::{DataError, Result};
8use crate::ids::{RepresentationId, SourceId};
9use crate::model::{DatasetSchema, SourceDescriptor};
10use crate::plan::{DataPlan, DataPlanStep, DataPlanStepKind, FitScope};
11use crate::ModelInputSpec;
12
13#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
14pub struct DataPlanRequest {
15 pub id: String,
16 #[serde(default)]
17 pub source_ids: Option<Vec<SourceId>>,
18 #[serde(default)]
19 pub planning_policy: PlanningPolicy,
20}
21
22impl DataPlanRequest {
23 pub fn new(id: impl Into<String>) -> Self {
24 Self {
25 id: id.into(),
26 source_ids: None,
27 planning_policy: PlanningPolicy::default(),
28 }
29 }
30}
31
32struct ResolvedSource<'a> {
33 source: &'a SourceDescriptor,
34 target_representation: RepresentationId,
35 path: AdapterPath,
36}
37
38pub fn plan_model_input(
39 schema: &DatasetSchema,
40 model_input: &ModelInputSpec,
41 adapters: &AdapterRegistry,
42 request: &DataPlanRequest,
43) -> Result<DataPlan> {
44 schema.validate()?;
45 model_input.validate()?;
46 validate_request(request)?;
47
48 let candidate_sources = candidate_sources(schema, request)?;
49 let mut steps = Vec::new();
50 let mut output_representation = None;
51 let mut adapt_idx = 0usize;
52 let mut align_idx = 0usize;
53
54 for port in &model_input.ports {
55 let resolved = candidate_sources
56 .iter()
57 .filter_map(|source| resolve_source(source, port, adapters, &request.planning_policy))
58 .collect::<Vec<_>>();
59
60 if resolved.is_empty() {
61 if port.optional {
62 continue;
63 }
64 return Err(DataError::Validation(format!(
65 "no source can satisfy required model input port `{}`",
66 port.name
67 )));
68 }
69 if resolved.len() > 1 && !port.multi_source {
70 return Err(DataError::Validation(format!(
71 "model input port `{}` accepts one source but {} sources resolved",
72 port.name,
73 resolved.len()
74 )));
75 }
76
77 let port_output_representation = resolved[0].target_representation.clone();
78 let mut join_inputs = Vec::new();
79 for source_plan in resolved {
80 let source_output = source_output_id(&source_plan.source.id);
81 steps.push(DataPlanStep {
82 kind: DataPlanStepKind::Materialize,
83 source_id: Some(source_plan.source.id.clone()),
84 adapter_id: None,
85 input_representation: None,
86 output_representation: Some(source_plan.source.native_representation.id.clone()),
87 fit_scope: FitScope::Stateless,
88 requires_user_choice: false,
89 metadata: BTreeMap::from([(
90 "output".to_string(),
91 serde_json::Value::String(source_output.clone()),
92 )]),
93 });
94
95 let mut current_output = source_output;
96 let mut current_representation = source_plan.source.native_representation.id.clone();
97 for adapter in source_plan.path.adapters {
98 let step_output = format!("step:adapt:{adapt_idx}");
99 steps.push(DataPlanStep {
100 kind: DataPlanStepKind::Adapt,
101 source_id: Some(source_plan.source.id.clone()),
102 adapter_id: Some(adapter.id),
103 input_representation: Some(current_representation),
104 output_representation: Some(adapter.output_representation.clone()),
105 fit_scope: adapter.fit_scope,
106 requires_user_choice: false,
107 metadata: BTreeMap::from([
108 (
109 "input".to_string(),
110 serde_json::Value::String(current_output),
111 ),
112 (
113 "output".to_string(),
114 serde_json::Value::String(step_output.clone()),
115 ),
116 ]),
117 });
118 current_output = step_output;
119 current_representation = adapter.output_representation;
120 adapt_idx += 1;
121 }
122 join_inputs.push(current_output);
123 }
124
125 let join_inputs = if join_inputs.len() > 1 {
126 let align_output = format!("step:align:{align_idx}");
127 let align_input_representation = output_representation
128 .clone()
129 .unwrap_or_else(|| resolved_target_representation(&steps));
130 let alignment_mode = alignment_mode_from_fusion(model_input.default_fusion.as_ref())?;
131 steps.push(DataPlanStep {
132 kind: DataPlanStepKind::Align,
133 source_id: None,
134 adapter_id: None,
135 input_representation: Some(align_input_representation),
136 output_representation: Some(port_output_representation.clone()),
137 fit_scope: FitScope::Stateless,
138 requires_user_choice: false,
139 metadata: alignment_metadata(join_inputs, align_output.clone(), alignment_mode)?,
140 });
141 align_idx += 1;
142 vec![align_output]
143 } else {
144 join_inputs
145 };
146
147 steps.push(DataPlanStep {
148 kind: DataPlanStepKind::Join,
149 source_id: None,
150 adapter_id: None,
151 input_representation: Some(
152 output_representation
153 .clone()
154 .unwrap_or_else(|| resolved_target_representation(&steps)),
155 ),
156 output_representation: Some(port_output_representation.clone()),
157 fit_scope: FitScope::Stateless,
158 requires_user_choice: false,
159 metadata: BTreeMap::from([
160 (
161 "inputs".to_string(),
162 serde_json::Value::Array(
163 join_inputs
164 .into_iter()
165 .map(serde_json::Value::String)
166 .collect(),
167 ),
168 ),
169 (
170 "output".to_string(),
171 serde_json::Value::String(format!("port:{}", port.name)),
172 ),
173 ]),
174 });
175 output_representation = Some(port_output_representation);
176 }
177
178 let output_representation = output_representation
179 .ok_or_else(|| DataError::Validation("model input plan produced no outputs".to_string()))?;
180 let plan = DataPlan {
181 id: request.id.clone(),
182 steps,
183 output_representation,
184 issues: Vec::new(),
185 };
186 plan.validate()?;
187 Ok(plan)
188}
189
190fn validate_request(request: &DataPlanRequest) -> Result<()> {
191 if request.id.trim().is_empty() {
192 return Err(DataError::Validation(
193 "data plan request id is empty".to_string(),
194 ));
195 }
196 if let Some(source_ids) = &request.source_ids {
197 let mut seen = std::collections::BTreeSet::new();
198 for source_id in source_ids {
199 if !seen.insert(source_id) {
200 return Err(DataError::Validation(format!(
201 "data plan request contains duplicate source `{source_id}`"
202 )));
203 }
204 }
205 }
206 Ok(())
207}
208
209fn candidate_sources<'a>(
210 schema: &'a DatasetSchema,
211 request: &DataPlanRequest,
212) -> Result<Vec<&'a SourceDescriptor>> {
213 let mut sources = schema.sources.iter().collect::<Vec<_>>();
214 sources.sort_by(|left, right| left.id.cmp(&right.id));
215 if let Some(requested) = &request.source_ids {
216 let by_id = sources
217 .iter()
218 .map(|source| (&source.id, *source))
219 .collect::<BTreeMap<_, _>>();
220 requested
221 .iter()
222 .map(|source_id| {
223 by_id.get(source_id).copied().ok_or_else(|| {
224 DataError::Validation(format!(
225 "data plan request references unknown source `{source_id}`"
226 ))
227 })
228 })
229 .collect()
230 } else {
231 Ok(sources)
232 }
233}
234
235fn resolve_source<'a>(
236 source: &'a SourceDescriptor,
237 port: &crate::InputPortSpec,
238 adapters: &AdapterRegistry,
239 policy: &PlanningPolicy,
240) -> Option<ResolvedSource<'a>> {
241 for target_type in &port.accepted_types {
242 for target_representation in &port.accepted_representations {
243 let resolution = adapters.find_path(
244 &source.type_id,
245 &source.native_representation.id,
246 target_type,
247 target_representation,
248 policy,
249 );
250 if let Some(path) = resolution.path {
251 return Some(ResolvedSource {
252 source,
253 target_representation: target_representation.clone(),
254 path,
255 });
256 }
257 }
258 }
259 None
260}
261
262fn source_output_id(source_id: &SourceId) -> String {
263 format!("src:{source_id}")
264}
265
266fn resolved_target_representation(steps: &[DataPlanStep]) -> RepresentationId {
267 steps
268 .iter()
269 .rev()
270 .find_map(|step| step.output_representation.clone())
271 .expect("at least one source step was emitted")
272}
273
274#[cfg(test)]
275mod tests {
276 use super::*;
277 use crate::adapter::AdapterRegistrySpec;
278 use crate::ids::{RepresentationId, SourceId, TypeId};
279 use crate::plan::DataPlan;
280
281 fn load_schema() -> DatasetSchema {
282 serde_json::from_str(include_str!(
283 "../../../examples/fixtures/oof_campaign/schema_nir_6_samples.json"
284 ))
285 .unwrap()
286 }
287
288 fn load_model_input() -> ModelInputSpec {
289 serde_json::from_str(include_str!(
290 "../../../examples/fixtures/oof_campaign/model_input_tabular_numeric.json"
291 ))
292 .unwrap()
293 }
294
295 fn load_registry() -> AdapterRegistry {
296 let spec: AdapterRegistrySpec = serde_json::from_str(include_str!(
297 "../../../examples/fixtures/oof_campaign/adapter_registry_signal_to_tabular.json"
298 ))
299 .unwrap();
300 AdapterRegistry::from_spec(spec).unwrap()
301 }
302
303 #[test]
304 fn fixture_plan_matches_expected_json() {
305 let plan = plan_model_input(
306 &load_schema(),
307 &load_model_input(),
308 &load_registry(),
309 &DataPlanRequest {
310 id: "nir-to-tabular".to_string(),
311 source_ids: Some(vec![SourceId::new("nir").unwrap()]),
312 planning_policy: PlanningPolicy::default(),
313 },
314 )
315 .unwrap();
316 let expected: DataPlan = serde_json::from_str(include_str!(
317 "../../../examples/fixtures/oof_campaign/expected_data_plan_nir_to_tabular.json"
318 ))
319 .unwrap();
320
321 assert_eq!(plan, expected);
322 }
323
324 #[test]
325 fn planner_refuses_missing_required_port_source() {
326 let err = plan_model_input(
327 &load_schema(),
328 &load_model_input(),
329 &load_registry(),
330 &DataPlanRequest {
331 id: "missing-source".to_string(),
332 source_ids: Some(vec![SourceId::new("other").unwrap()]),
333 planning_policy: PlanningPolicy::default(),
334 },
335 )
336 .unwrap_err();
337
338 assert!(err.to_string().contains("unknown source"));
339 }
340
341 #[test]
342 fn planner_emits_alignment_step_for_multi_source_port() {
343 let mut schema = load_schema();
344 let mut chem = schema.sources[0].clone();
345 chem.id = SourceId::new("chem").unwrap();
346 chem.name = "Chemistry".to_string();
347 chem.type_id = TypeId::new("table").unwrap();
348 chem.native_representation.id = RepresentationId::new("tabular_numeric").unwrap();
349 chem.native_representation.type_id = TypeId::new("table").unwrap();
350 chem.native_representation.container = "dataframe".to_string();
351 schema.sources.push(chem);
352
353 let plan = plan_model_input(
354 &schema,
355 &load_model_input(),
356 &load_registry(),
357 &DataPlanRequest {
358 id: "nir-chem-to-tabular".to_string(),
359 source_ids: Some(vec![
360 SourceId::new("nir").unwrap(),
361 SourceId::new("chem").unwrap(),
362 ]),
363 planning_policy: PlanningPolicy::default(),
364 },
365 )
366 .unwrap();
367
368 let align = plan
369 .steps
370 .iter()
371 .find(|step| step.kind == DataPlanStepKind::Align)
372 .expect("multi-source plan should expose an Align step");
373 assert_eq!(align.metadata["alignment"], serde_json::json!("left"));
374 assert_eq!(align.metadata["output"], serde_json::json!("step:align:0"));
375 assert_eq!(
376 align.metadata["inputs"].as_array().unwrap().len(),
377 2,
378 "alignment should receive both source feature streams"
379 );
380
381 let join = plan
382 .steps
383 .iter()
384 .find(|step| step.kind == DataPlanStepKind::Join)
385 .unwrap();
386 assert_eq!(join.metadata["inputs"], serde_json::json!(["step:align:0"]));
387 }
388}