1use std::collections::BTreeSet;
2
3use serde::{Deserialize, Serialize};
4
5use crate::error::{DataError, Result};
6use crate::handle::CoordinatorFeatureBlock;
7use crate::ids::{ObservationId, RepresentationId, SampleId};
8
9#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
10#[serde(rename_all = "snake_case")]
11pub enum CollationPadding {
12 #[default]
13 None,
14 Right,
15 Left,
16 Center,
17}
18
19#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
20pub struct CollationPolicy {
21 #[serde(default)]
22 pub padding: CollationPadding,
23 #[serde(default)]
24 pub truncate: bool,
25 #[serde(default)]
26 pub batch_container: Option<String>,
27 #[serde(default = "default_true")]
28 pub emit_mask: bool,
29 #[serde(default)]
30 pub max_length: Option<usize>,
31 #[serde(default)]
32 pub pad_value: f64,
33}
34
35impl Default for CollationPolicy {
36 fn default() -> Self {
37 Self {
38 padding: CollationPadding::None,
39 truncate: false,
40 batch_container: None,
41 emit_mask: true,
42 max_length: None,
43 pad_value: 0.0,
44 }
45 }
46}
47
48fn default_true() -> bool {
49 true
50}
51
52#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
53pub struct NumericCollationInputBlock {
54 pub block_id: String,
55 pub representation_id: RepresentationId,
56 pub observation_ids: Vec<ObservationId>,
57 pub sample_ids: Vec<SampleId>,
58 pub rows: Vec<Vec<Option<f64>>>,
59 #[serde(default)]
60 pub feature_names: Option<Vec<String>>,
61}
62
63#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
64pub struct NumericTensorBlock {
65 pub block_id: String,
66 pub representation_id: RepresentationId,
67 pub batch_container: String,
68 pub observation_ids: Vec<ObservationId>,
69 pub sample_ids: Vec<SampleId>,
70 pub shape: Vec<usize>,
71 pub values: Vec<f64>,
72 #[serde(default)]
73 pub presence_mask: Option<Vec<bool>>,
74 #[serde(default)]
75 pub validity_mask: Option<Vec<bool>>,
76 #[serde(default)]
77 pub feature_names: Option<Vec<String>>,
78}
79
80pub fn numeric_input_from_feature_block(
81 block: &CoordinatorFeatureBlock,
82) -> Result<NumericCollationInputBlock> {
83 validate_feature_block_shape(block)?;
84 let rows = block
85 .values
86 .iter()
87 .enumerate()
88 .map(|(row_idx, row)| {
89 row.iter()
90 .enumerate()
91 .map(|(feature_idx, value)| match value {
92 serde_json::Value::Null => Ok(None),
93 serde_json::Value::Number(number) => number.as_f64().map(Some).ok_or_else(|| {
94 DataError::Validation(format!(
95 "feature block `{}` row `{}` feature `{}` contains a non-f64 numeric value",
96 block.feature_set_id,
97 block.observation_ids[row_idx],
98 block.feature_names[feature_idx]
99 ))
100 }),
101 _ => Err(DataError::Validation(format!(
102 "feature block `{}` row `{}` feature `{}` must be numeric or null for collation",
103 block.feature_set_id,
104 block.observation_ids[row_idx],
105 block.feature_names[feature_idx]
106 ))),
107 })
108 .collect::<Result<Vec<_>>>()
109 })
110 .collect::<Result<Vec<_>>>()?;
111 Ok(NumericCollationInputBlock {
112 block_id: block.feature_set_id.clone(),
113 representation_id: block.representation_id.clone(),
114 observation_ids: block.observation_ids.clone(),
115 sample_ids: block.sample_ids.clone(),
116 rows,
117 feature_names: Some(block.feature_names.clone()),
118 })
119}
120
121pub fn collate_feature_block(
122 block: &CoordinatorFeatureBlock,
123 policy: &CollationPolicy,
124) -> Result<NumericTensorBlock> {
125 let input = numeric_input_from_feature_block(block)?;
126 collate_numeric_block(&input, policy)
127}
128
129pub fn collate_numeric_block(
130 block: &NumericCollationInputBlock,
131 policy: &CollationPolicy,
132) -> Result<NumericTensorBlock> {
133 validate_numeric_input_block(block)?;
134 validate_collation_policy(policy)?;
135 let target_len = target_length(block, policy)?;
136 validate_feature_names_for_collation(block, policy, target_len)?;
137
138 let batch = block.rows.len();
139 let mut values = Vec::with_capacity(batch * target_len);
140 let mut presence = Vec::with_capacity(batch * target_len);
141 let mut validity = Vec::with_capacity(batch * target_len);
142 let mut has_invalid = false;
143 for row in &block.rows {
144 let projected = project_row(row, target_len, policy)?;
145 values.extend(
146 projected
147 .values
148 .iter()
149 .map(|value| value.unwrap_or(policy.pad_value)),
150 );
151 presence.extend(projected.presence.iter().copied());
152 for (value, present) in projected.values.iter().zip(projected.presence.iter()) {
153 let valid = *present && value.is_some();
154 has_invalid |= !valid;
155 validity.push(valid);
156 }
157 }
158
159 Ok(NumericTensorBlock {
160 block_id: block.block_id.clone(),
161 representation_id: block.representation_id.clone(),
162 batch_container: policy
163 .batch_container
164 .clone()
165 .unwrap_or_else(|| "ndarray".to_string()),
166 observation_ids: block.observation_ids.clone(),
167 sample_ids: block.sample_ids.clone(),
168 shape: vec![batch, target_len],
169 values,
170 presence_mask: policy.emit_mask.then_some(presence),
171 validity_mask: has_invalid.then_some(validity),
172 feature_names: projected_feature_names(block.feature_names.as_deref(), target_len, policy)?,
173 })
174}
175
176struct ProjectedRow {
177 values: Vec<Option<f64>>,
178 presence: Vec<bool>,
179}
180
181fn validate_collation_policy(policy: &CollationPolicy) -> Result<()> {
182 if policy.max_length == Some(0) {
183 return Err(DataError::Validation(
184 "collation max_length must be greater than zero".to_string(),
185 ));
186 }
187 if let Some(container) = &policy.batch_container {
188 if container.trim().is_empty() {
189 return Err(DataError::Validation(
190 "collation batch_container must not be empty".to_string(),
191 ));
192 }
193 }
194 if !policy.pad_value.is_finite() {
195 return Err(DataError::Validation(
196 "collation pad_value must be finite".to_string(),
197 ));
198 }
199 Ok(())
200}
201
202fn validate_numeric_input_block(block: &NumericCollationInputBlock) -> Result<()> {
203 if block.block_id.trim().is_empty() {
204 return Err(DataError::Validation(
205 "collation input block_id is empty".to_string(),
206 ));
207 }
208 if block.observation_ids.is_empty() {
209 return Err(DataError::Validation(format!(
210 "collation input block `{}` contains no rows",
211 block.block_id
212 )));
213 }
214 if block.observation_ids.len() != block.sample_ids.len()
215 || block.sample_ids.len() != block.rows.len()
216 {
217 return Err(DataError::Validation(format!(
218 "collation input block `{}` row identity/value lengths differ",
219 block.block_id
220 )));
221 }
222 let mut observations = BTreeSet::new();
223 for (idx, row) in block.rows.iter().enumerate() {
224 if !observations.insert(&block.observation_ids[idx]) {
225 return Err(DataError::Validation(format!(
226 "collation input block `{}` contains duplicate observation `{}`",
227 block.block_id, block.observation_ids[idx]
228 )));
229 }
230 if row.is_empty() {
231 return Err(DataError::Validation(format!(
232 "collation input block `{}` row `{}` is empty",
233 block.block_id, block.observation_ids[idx]
234 )));
235 }
236 for value in row.iter().flatten() {
237 if !value.is_finite() {
238 return Err(DataError::Validation(format!(
239 "collation input block `{}` row `{}` contains a non-finite value",
240 block.block_id, block.observation_ids[idx]
241 )));
242 }
243 }
244 }
245 if let Some(feature_names) = &block.feature_names {
246 if feature_names.is_empty() {
247 return Err(DataError::Validation(format!(
248 "collation input block `{}` has empty feature_names",
249 block.block_id
250 )));
251 }
252 if feature_names.iter().any(|name| name.trim().is_empty()) {
253 return Err(DataError::Validation(format!(
254 "collation input block `{}` has an empty feature name",
255 block.block_id
256 )));
257 }
258 let mut names = BTreeSet::new();
259 for name in feature_names {
260 if !names.insert(name) {
261 return Err(DataError::Validation(format!(
262 "collation input block `{}` has duplicate feature `{name}`",
263 block.block_id
264 )));
265 }
266 }
267 }
268 Ok(())
269}
270
271fn validate_feature_block_shape(block: &CoordinatorFeatureBlock) -> Result<()> {
272 if block.feature_set_id.trim().is_empty() {
273 return Err(DataError::Validation(
274 "feature block feature_set_id is empty".to_string(),
275 ));
276 }
277 if block.feature_names.is_empty() {
278 return Err(DataError::Validation(format!(
279 "feature block `{}` contains no features",
280 block.feature_set_id
281 )));
282 }
283 if block.observation_ids.len() != block.sample_ids.len()
284 || block.sample_ids.len() != block.values.len()
285 {
286 return Err(DataError::Validation(format!(
287 "feature block `{}` row identity/value lengths differ",
288 block.feature_set_id
289 )));
290 }
291 for (idx, row) in block.values.iter().enumerate() {
292 if row.len() != block.feature_names.len() {
293 return Err(DataError::Validation(format!(
294 "feature block `{}` row `{}` has {} values for {} features",
295 block.feature_set_id,
296 block.observation_ids[idx],
297 row.len(),
298 block.feature_names.len()
299 )));
300 }
301 }
302 Ok(())
303}
304
305fn validate_feature_names_for_collation(
306 block: &NumericCollationInputBlock,
307 policy: &CollationPolicy,
308 target_len: usize,
309) -> Result<()> {
310 let Some(feature_names) = &block.feature_names else {
311 return Ok(());
312 };
313 for row in &block.rows {
314 if row.len() != feature_names.len() {
315 return Err(DataError::Validation(format!(
316 "collation input block `{}` named features require rectangular rows",
317 block.block_id
318 )));
319 }
320 }
321 if target_len > feature_names.len() {
322 return Err(DataError::Validation(format!(
323 "collation input block `{}` cannot pad named feature rows",
324 block.block_id
325 )));
326 }
327 if target_len < feature_names.len() && !policy.truncate {
328 return Err(DataError::Validation(format!(
329 "collation input block `{}` feature names require truncation for max_length",
330 block.block_id
331 )));
332 }
333 Ok(())
334}
335
336fn target_length(block: &NumericCollationInputBlock, policy: &CollationPolicy) -> Result<usize> {
337 let max_observed = block.rows.iter().map(Vec::len).max().unwrap_or(0);
338 let target = policy.max_length.unwrap_or(max_observed);
339 if target == 0 {
340 return Err(DataError::Validation(format!(
341 "collation input block `{}` produced an empty target length",
342 block.block_id
343 )));
344 }
345 if !policy.truncate && block.rows.iter().any(|row| row.len() > target) {
346 return Err(DataError::Validation(format!(
347 "collation input block `{}` has rows longer than max_length without truncate",
348 block.block_id
349 )));
350 }
351 if policy.padding == CollationPadding::None && block.rows.iter().any(|row| row.len() < target) {
352 return Err(DataError::Validation(format!(
353 "collation input block `{}` has ragged rows but padding is none",
354 block.block_id
355 )));
356 }
357 Ok(target)
358}
359
360fn project_row(
361 row: &[Option<f64>],
362 target_len: usize,
363 policy: &CollationPolicy,
364) -> Result<ProjectedRow> {
365 let truncated = if row.len() > target_len {
366 if !policy.truncate {
367 return Err(DataError::Validation(
368 "collation row is longer than target length without truncate".to_string(),
369 ));
370 }
371 let start = truncate_start(row.len(), target_len, policy.padding);
372 row[start..start + target_len].to_vec()
373 } else {
374 row.to_vec()
375 };
376
377 if truncated.len() == target_len {
378 return Ok(ProjectedRow {
379 presence: vec![true; target_len],
380 values: truncated,
381 });
382 }
383 if policy.padding == CollationPadding::None {
384 return Err(DataError::Validation(
385 "collation row is shorter than target length but padding is none".to_string(),
386 ));
387 }
388
389 let missing = target_len - truncated.len();
390 let (left_pad, right_pad) = match policy.padding {
391 CollationPadding::None => unreachable!("padding none handled above"),
392 CollationPadding::Right => (0, missing),
393 CollationPadding::Left => (missing, 0),
394 CollationPadding::Center => (missing / 2, missing - (missing / 2)),
395 };
396 let mut values = Vec::with_capacity(target_len);
397 let mut presence = Vec::with_capacity(target_len);
398 values.extend(std::iter::repeat_n(None, left_pad));
399 presence.extend(std::iter::repeat_n(false, left_pad));
400 values.extend(truncated);
401 presence.extend(std::iter::repeat_n(true, target_len - left_pad - right_pad));
402 values.extend(std::iter::repeat_n(None, right_pad));
403 presence.extend(std::iter::repeat_n(false, right_pad));
404 Ok(ProjectedRow { values, presence })
405}
406
407fn truncate_start(row_len: usize, target_len: usize, padding: CollationPadding) -> usize {
408 match padding {
409 CollationPadding::Left => row_len - target_len,
410 CollationPadding::Center => (row_len - target_len) / 2,
411 CollationPadding::None | CollationPadding::Right => 0,
412 }
413}
414
415fn projected_feature_names(
416 feature_names: Option<&[String]>,
417 target_len: usize,
418 policy: &CollationPolicy,
419) -> Result<Option<Vec<String>>> {
420 let Some(feature_names) = feature_names else {
421 return Ok(None);
422 };
423 if target_len == feature_names.len() {
424 return Ok(Some(feature_names.to_vec()));
425 }
426 if target_len > feature_names.len() {
427 return Err(DataError::Validation(
428 "named feature collation cannot add padded feature names".to_string(),
429 ));
430 }
431 let start = truncate_start(feature_names.len(), target_len, policy.padding);
432 Ok(Some(feature_names[start..start + target_len].to_vec()))
433}
434
435#[cfg(test)]
436mod tests {
437 use super::*;
438 use crate::ids::RepresentationId;
439 use serde_json::json;
440
441 fn obs(value: &str) -> ObservationId {
442 ObservationId::new(value).unwrap()
443 }
444
445 fn sample(value: &str) -> SampleId {
446 SampleId::new(value).unwrap()
447 }
448
449 fn feature_block() -> CoordinatorFeatureBlock {
450 CoordinatorFeatureBlock {
451 feature_set_id: "x".to_string(),
452 representation_id: RepresentationId::new("tabular_numeric").unwrap(),
453 feature_names: vec!["f0".to_string(), "f1".to_string()],
454 observation_ids: vec![obs("obs.S001"), obs("obs.S002")],
455 sample_ids: vec![sample("S001"), sample("S002")],
456 values: vec![vec![json!(1.0), json!(2.0)], vec![json!(3.0), json!(4.0)]],
457 }
458 }
459
460 #[test]
461 fn collates_rectangular_feature_block_to_row_major_tensor() {
462 let tensor = collate_feature_block(
463 &feature_block(),
464 &CollationPolicy {
465 emit_mask: false,
466 ..Default::default()
467 },
468 )
469 .unwrap();
470
471 assert_eq!(tensor.shape, vec![2, 2]);
472 assert_eq!(tensor.values, vec![1.0, 2.0, 3.0, 4.0]);
473 assert_eq!(tensor.presence_mask, None);
474 assert_eq!(tensor.validity_mask, None);
475 assert_eq!(
476 tensor.feature_names,
477 Some(vec!["f0".to_string(), "f1".to_string()])
478 );
479 }
480
481 #[test]
482 fn right_padding_emits_presence_and_validity_masks() {
483 let block = NumericCollationInputBlock {
484 block_id: "seq".to_string(),
485 representation_id: RepresentationId::new("sequence_tensor").unwrap(),
486 observation_ids: vec![obs("obs.S001"), obs("obs.S002")],
487 sample_ids: vec![sample("S001"), sample("S002")],
488 rows: vec![vec![Some(1.0), Some(2.0)], vec![Some(3.0), None]],
489 feature_names: None,
490 };
491 let tensor = collate_numeric_block(
492 &block,
493 &CollationPolicy {
494 padding: CollationPadding::Right,
495 max_length: Some(3),
496 pad_value: -1.0,
497 ..Default::default()
498 },
499 )
500 .unwrap();
501
502 assert_eq!(tensor.shape, vec![2, 3]);
503 assert_eq!(tensor.values, vec![1.0, 2.0, -1.0, 3.0, -1.0, -1.0]);
504 assert_eq!(
505 tensor.presence_mask,
506 Some(vec![true, true, false, true, true, false])
507 );
508 assert_eq!(
509 tensor.validity_mask,
510 Some(vec![true, true, false, true, false, false])
511 );
512 }
513
514 #[test]
515 fn no_padding_refuses_ragged_rows() {
516 let block = NumericCollationInputBlock {
517 block_id: "seq".to_string(),
518 representation_id: RepresentationId::new("sequence_tensor").unwrap(),
519 observation_ids: vec![obs("obs.S001"), obs("obs.S002")],
520 sample_ids: vec![sample("S001"), sample("S002")],
521 rows: vec![vec![Some(1.0), Some(2.0)], vec![Some(3.0)]],
522 feature_names: None,
523 };
524
525 let err = collate_numeric_block(&block, &CollationPolicy::default()).unwrap_err();
526
527 assert!(err.to_string().contains("ragged rows"));
528 }
529
530 #[test]
531 fn left_truncation_keeps_suffix_and_projects_feature_names() {
532 let tensor = collate_feature_block(
533 &feature_block(),
534 &CollationPolicy {
535 padding: CollationPadding::Left,
536 truncate: true,
537 max_length: Some(1),
538 ..Default::default()
539 },
540 )
541 .unwrap();
542
543 assert_eq!(tensor.shape, vec![2, 1]);
544 assert_eq!(tensor.values, vec![2.0, 4.0]);
545 assert_eq!(tensor.feature_names, Some(vec!["f1".to_string()]));
546 }
547
548 #[test]
549 fn collation_refuses_non_numeric_feature_values() {
550 let mut block = feature_block();
551 block.values[0][0] = json!("bad");
552
553 let err = collate_feature_block(&block, &CollationPolicy::default()).unwrap_err();
554
555 assert!(err.to_string().contains("must be numeric or null"));
556 }
557}