laddu_data/
schema.rs

1use std::{
2    collections::{BTreeMap, HashMap, HashSet},
3    sync::Arc,
4};
5
6use crate::{LadduDataError, LadduDataResult, Name};
7
8/// Logical names and lookup tables for four-momentum, scalar, and weight columns.
9#[derive(Clone, Debug)]
10pub struct Schema {
11    p4s: Vec<Name>,
12    scalars: Vec<Name>,
13    has_weight: bool,
14
15    p4_index: Arc<HashMap<Name, usize>>,
16    scalar_index: Arc<HashMap<Name, usize>>,
17}
18
19impl PartialEq for Schema {
20    fn eq(&self, other: &Self) -> bool {
21        self.p4s == other.p4s
22            && self.scalars == other.scalars
23            && self.has_weight == other.has_weight
24    }
25}
26
27impl Schema {
28    /// Validates and constructs a logical schema.
29    ///
30    /// # Errors
31    ///
32    /// Returns [`LadduDataError`] when four-momentum or scalar column names are
33    /// duplicated.
34    pub fn new(
35        p4s: impl IntoIterator<Item = impl Into<Name>>,
36        scalars: impl IntoIterator<Item = impl Into<Name>>,
37        has_weight: bool,
38    ) -> LadduDataResult<Self> {
39        let p4s: Vec<Name> = p4s.into_iter().map(Into::into).collect();
40        let scalars: Vec<Name> = scalars.into_iter().map(Into::into).collect();
41        let p4_index = Arc::new(make_index(&p4s, "p4")?);
42        let scalar_index = Arc::new(make_index(&scalars, "scalar")?);
43        Ok(Self {
44            p4s,
45            scalars,
46            has_weight,
47            p4_index,
48            scalar_index,
49        })
50    }
51
52    /// Returns the four-momentum column index for `name`.
53    pub fn p4_index(&self, name: &str) -> Option<usize> {
54        self.p4_index.get(name).copied()
55    }
56
57    /// Returns the scalar column index for `name`.
58    pub fn scalar_index(&self, name: &str) -> Option<usize> {
59        self.scalar_index.get(name).copied()
60    }
61
62    /// Returns four-momentum names in column order.
63    pub fn p4s(&self) -> &[Name] {
64        &self.p4s
65    }
66
67    /// Returns scalar names in column order.
68    pub fn scalars(&self) -> &[Name] {
69        &self.scalars
70    }
71
72    /// Returns whether events carry explicit weights.
73    pub fn has_weight(&self) -> bool {
74        self.has_weight
75    }
76
77    /// Returns the number of four-momentum columns.
78    pub fn n_p4s(&self) -> usize {
79        self.p4s.len()
80    }
81
82    /// Returns the number of scalar columns.
83    pub fn n_scalars(&self) -> usize {
84        self.scalars.len()
85    }
86
87    /// Requires and returns a four-momentum column index.
88    ///
89    /// # Errors
90    ///
91    /// Returns [`LadduDataError::MissingColumn`] when `name` is not a
92    /// four-momentum column.
93    pub fn require_p4(&self, name: &str) -> LadduDataResult<usize> {
94        self.p4_index(name)
95            .ok_or_else(|| LadduDataError::MissingColumn(Name::from(name)))
96    }
97
98    /// Requires and returns a scalar column index.
99    ///
100    /// # Errors
101    ///
102    /// Returns [`LadduDataError::MissingColumn`] when `name` is not a scalar
103    /// column.
104    pub fn require_scalar(&self, name: &str) -> LadduDataResult<usize> {
105        self.scalar_index(name)
106            .ok_or_else(|| LadduDataError::MissingColumn(Name::from(name)))
107    }
108}
109
110fn make_index(names: &[Name], kind: &'static str) -> LadduDataResult<HashMap<Name, usize>> {
111    let mut out = HashMap::with_capacity(names.len());
112    for (i, name) in names.iter().cloned().enumerate() {
113        if out.insert(name.clone(), i).is_some() {
114            return Err(LadduDataError::Schema(format!(
115                "duplicate {kind} column: {name}"
116            )));
117        }
118    }
119    Ok(out)
120}
121
122/// Physical naming conventions used to map a logical schema to storage columns.
123#[derive(Clone, Debug)]
124pub struct SchemaColumnNames {
125    /// Physical weight-column name.
126    pub weight_column: Name,
127    /// Suffixes for four-momentum components.
128    pub p4_suffixes: P4Suffixes,
129}
130
131impl Default for SchemaColumnNames {
132    fn default() -> Self {
133        Self {
134            weight_column: Name::from("weight"),
135            p4_suffixes: P4Suffixes::default(),
136        }
137    }
138}
139
140/// Options controlling logical schema inference from physical columns.
141#[derive(Clone, Debug)]
142pub struct SchemaInferenceOptions {
143    /// Physical naming conventions.
144    pub column_names: SchemaColumnNames,
145    /// Whether inference fails if the weight column is absent.
146    pub require_weight: bool,
147    /// Whether incomplete four-momenta become independent scalar columns.
148    pub incomplete_p4_components_are_scalars: bool,
149}
150
151impl Default for SchemaInferenceOptions {
152    fn default() -> Self {
153        Self {
154            column_names: SchemaColumnNames::default(),
155            require_weight: false,
156            incomplete_p4_components_are_scalars: true,
157        }
158    }
159}
160
161/// Physical suffixes for `(E, px, py, pz)` columns.
162#[derive(Clone, Debug)]
163pub struct P4Suffixes {
164    /// Energy suffix.
165    pub e: &'static str,
166    /// X-momentum suffix.
167    pub px: &'static str,
168    /// Y-momentum suffix.
169    pub py: &'static str,
170    /// Z-momentum suffix.
171    pub pz: &'static str,
172}
173
174impl Default for P4Suffixes {
175    fn default() -> Self {
176        Self {
177            e: "_e",
178            px: "_px",
179            py: "_py",
180            pz: "_pz",
181        }
182    }
183}
184
185impl P4Suffixes {
186    /// Splits a matching physical name into logical prefix and component index.
187    pub fn component<'a>(&'a self, name: &'a str) -> Option<(&'a str, usize)> {
188        if let Some(prefix) = name.strip_suffix(self.e) {
189            Some((prefix, 0))
190        } else if let Some(prefix) = name.strip_suffix(self.px) {
191            Some((prefix, 1))
192        } else if let Some(prefix) = name.strip_suffix(self.py) {
193            Some((prefix, 2))
194        } else if let Some(prefix) = name.strip_suffix(self.pz) {
195            Some((prefix, 3))
196        } else {
197            None
198        }
199    }
200
201    /// Produces the four physical names for a logical four-momentum prefix.
202    pub fn physical_p4_names(&self, prefix: &str) -> [String; 4] {
203        [
204            format!("{prefix}{}", self.e),
205            format!("{prefix}{}", self.px),
206            format!("{prefix}{}", self.py),
207            format!("{prefix}{}", self.pz),
208        ]
209    }
210}
211
212/// Physical storage type relevant to schema inference.
213#[derive(Clone, Copy, Debug, PartialEq, Eq)]
214pub enum ColumnType {
215    /// 64-bit floating point.
216    F64,
217    /// 32-bit floating point.
218    F32,
219    /// Any unsupported type.
220    Other,
221}
222
223impl ColumnType {
224    /// Returns whether this type can populate an event-data column.
225    pub fn is_supported_float(self) -> bool {
226        matches!(self, Self::F64 | Self::F32)
227    }
228}
229
230/// Name and physical type of one available storage column.
231#[derive(Clone, Copy, Debug)]
232pub struct ColumnInfo<'a> {
233    /// Physical column name.
234    pub name: &'a str,
235    /// Physical column type.
236    pub dtype: ColumnType,
237}
238
239impl Schema {
240    /// Infers a logical schema from available physical columns.
241    ///
242    /// # Errors
243    ///
244    /// Returns [`LadduDataError`] when required momentum components or weights
245    /// are missing, names are ambiguous, or inferred logical names conflict.
246    pub fn infer_from_columns<'a>(
247        columns: impl IntoIterator<Item = ColumnInfo<'a>>,
248        options: &SchemaInferenceOptions,
249    ) -> LadduDataResult<Self> {
250        let mut p4_candidates = BTreeMap::<String, [bool; 4]>::new();
251        let mut scalar_names = Vec::<Name>::new();
252        let mut has_weight = false;
253
254        for col in columns {
255            if !col.dtype.is_supported_float() {
256                continue;
257            }
258
259            if col.name == options.column_names.weight_column.as_ref() {
260                has_weight = true;
261                continue;
262            }
263
264            if let Some((prefix, component)) = options.column_names.p4_suffixes.component(col.name)
265            {
266                p4_candidates.entry(prefix.to_owned()).or_default()[component] = true;
267            } else {
268                scalar_names.push(Name::from(col.name));
269            }
270        }
271
272        let mut p4s = Vec::<Name>::new();
273
274        for (prefix, seen) in p4_candidates {
275            if seen == [true, true, true, true] {
276                p4s.push(Name::from(prefix));
277            } else if options.incomplete_p4_components_are_scalars {
278                let names = options.column_names.p4_suffixes.physical_p4_names(&prefix);
279
280                for (i, name) in names.into_iter().enumerate() {
281                    if seen[i] {
282                        scalar_names.push(Name::from(name));
283                    }
284                }
285            }
286        }
287
288        if options.require_weight && !has_weight {
289            return Err(LadduDataError::MissingColumn(Arc::clone(
290                &options.column_names.weight_column,
291            )));
292        }
293
294        Schema::new(p4s, scalar_names, has_weight)
295    }
296
297    /// Returns all physical columns required to store this schema.
298    pub fn physical_columns(&self, column_names: &SchemaColumnNames) -> Vec<Name> {
299        let mut names = Vec::with_capacity(4 * self.n_p4s() + self.n_scalars() + 1);
300
301        for p4 in self.p4s() {
302            for name in column_names.p4_suffixes.physical_p4_names(p4) {
303                names.push(Name::from(name));
304            }
305        }
306
307        names.extend(self.scalars().iter().cloned());
308
309        if self.has_weight() {
310            names.push(Arc::clone(&column_names.weight_column));
311        }
312
313        names
314    }
315
316    /// Validates that all physical columns required by this schema are available.
317    ///
318    /// # Errors
319    ///
320    /// Returns [`LadduDataError::MissingColumn`] when a required physical
321    /// column is absent or has an unsupported type.
322    pub fn validate_required_columns<'a>(
323        &self,
324        available: impl IntoIterator<Item = ColumnInfo<'a>>,
325        options: &SchemaInferenceOptions,
326    ) -> LadduDataResult<()> {
327        let available: HashSet<&str> = available
328            .into_iter()
329            .filter(|c| c.dtype.is_supported_float())
330            .map(|c| c.name)
331            .collect();
332
333        for required in self.physical_columns(&options.column_names) {
334            if !available.contains(required.as_ref()) {
335                return Err(LadduDataError::MissingColumn(required));
336            }
337        }
338
339        Ok(())
340    }
341}
342
343/// Floating-point precision used when writing physical columns.
344#[derive(Copy, Clone, Debug, Default)]
345pub enum Precision {
346    /// Write 64-bit floating-point values.
347    #[default]
348    F64,
349    /// Write 32-bit floating-point values.
350    F32,
351}
352
353/// Policy controlling whether sinks emit a weight column.
354#[derive(Clone, Copy, Debug, Default)]
355pub enum WriteWeightColumn {
356    /// Always write weights, using unit weights when the schema has none.
357    #[default]
358    Always,
359    /// Write weights only when the logical schema contains them.
360    OnlyIfPresent,
361}
362
363/// Physical naming, precision, and weight policy for event sinks.
364#[derive(Clone, Debug, Default)]
365pub struct SchemaWriteOptions {
366    /// Physical column naming conventions.
367    pub column_names: SchemaColumnNames,
368    /// Floating-point output precision.
369    pub precision: Precision,
370    /// Weight-column emission policy.
371    pub write_weight_column: WriteWeightColumn,
372}
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377
378    fn col(name: &'static str, dtype: ColumnType) -> ColumnInfo<'static> {
379        ColumnInfo { name, dtype }
380    }
381
382    #[test]
383    fn schema_new_rejects_duplicates_and_required_lookup_reports_missing_column() {
384        let duplicate_p4 = Schema::new(["p", "p"], ["mass"], false);
385        assert!(matches!(duplicate_p4, Err(LadduDataError::Schema(_))));
386
387        let duplicate_scalar = Schema::new(["p"], ["mass", "mass"], false);
388        assert!(matches!(duplicate_scalar, Err(LadduDataError::Schema(_))));
389
390        let schema = Schema::new(["beam", "recoil"], ["mass", "costheta"], true).unwrap();
391
392        assert_eq!(schema.require_p4("recoil").unwrap(), 1);
393        assert_eq!(schema.require_scalar("costheta").unwrap(), 1);
394
395        let err = schema.require_scalar("missing").unwrap_err();
396        assert!(matches!(err, LadduDataError::MissingColumn(name) if name.as_ref() == "missing"));
397    }
398
399    #[test]
400    fn infer_from_columns_groups_complete_p4s_keeps_incomplete_components_as_scalars_and_ignores_nonfloats()
401     {
402        let options = SchemaInferenceOptions::default();
403
404        let schema = Schema::infer_from_columns(
405            [
406                col("gamma_px", ColumnType::F64),
407                col("gamma_py", ColumnType::F64),
408                col("gamma_pz", ColumnType::F32),
409                col("gamma_e", ColumnType::F64),
410                col("partial_px", ColumnType::F64),
411                col("partial_e", ColumnType::F64),
412                col("mass", ColumnType::F32),
413                col("ignored", ColumnType::Other),
414                col("weight", ColumnType::F64),
415            ],
416            &options,
417        )
418        .unwrap();
419
420        assert_eq!(
421            schema
422                .p4s()
423                .iter()
424                .map(|n| n.to_string())
425                .collect::<Vec<_>>(),
426            vec!["gamma"]
427        );
428
429        assert_eq!(
430            schema
431                .scalars()
432                .iter()
433                .map(|n| n.to_string())
434                .collect::<Vec<_>>(),
435            vec!["mass", "partial_e", "partial_px"]
436        );
437
438        assert!(schema.has_weight());
439    }
440
441    #[test]
442    fn infer_from_columns_can_discard_incomplete_p4_components_and_require_weight() {
443        let options = SchemaInferenceOptions {
444            incomplete_p4_components_are_scalars: false,
445            ..Default::default()
446        };
447
448        let schema = Schema::infer_from_columns(
449            [
450                col("partial_px", ColumnType::F64),
451                col("partial_e", ColumnType::F64),
452                col("mass", ColumnType::F64),
453                col("weight", ColumnType::F64),
454            ],
455            &options,
456        )
457        .unwrap();
458
459        assert!(schema.p4s().is_empty());
460        assert_eq!(
461            schema
462                .scalars()
463                .iter()
464                .map(|n| n.to_string())
465                .collect::<Vec<_>>(),
466            vec!["mass"]
467        );
468
469        let require_weight = SchemaInferenceOptions {
470            require_weight: true,
471            ..Default::default()
472        };
473
474        let err = Schema::infer_from_columns([col("mass", ColumnType::F64)], &require_weight)
475            .unwrap_err();
476
477        assert!(matches!(err, LadduDataError::MissingColumn(name) if name.as_ref() == "weight"));
478    }
479
480    #[test]
481    fn physical_columns_and_validation_respect_custom_names_and_float_types_only() {
482        let schema = Schema::new(["p"], ["mass"], true).unwrap();
483
484        let names = SchemaColumnNames {
485            weight_column: Name::from("event_weight"),
486            ..Default::default()
487        };
488
489        let physical = schema
490            .physical_columns(&names)
491            .into_iter()
492            .map(|n| n.to_string())
493            .collect::<Vec<_>>();
494
495        assert_eq!(
496            physical,
497            vec!["p_e", "p_px", "p_py", "p_pz", "mass", "event_weight"]
498        );
499
500        let options = SchemaInferenceOptions {
501            column_names: names,
502            ..Default::default()
503        };
504
505        let ok = schema.validate_required_columns(
506            [
507                col("p_e", ColumnType::F32),
508                col("p_px", ColumnType::F64),
509                col("p_py", ColumnType::F64),
510                col("p_pz", ColumnType::F32),
511                col("mass", ColumnType::F64),
512                col("event_weight", ColumnType::F64),
513            ],
514            &options,
515        );
516
517        assert!(ok.is_ok());
518
519        let missing_because_not_float = schema
520            .validate_required_columns(
521                [
522                    col("p_e", ColumnType::F32),
523                    col("p_px", ColumnType::F64),
524                    col("p_py", ColumnType::Other),
525                    col("p_pz", ColumnType::F32),
526                    col("mass", ColumnType::F64),
527                    col("event_weight", ColumnType::F64),
528                ],
529                &options,
530            )
531            .unwrap_err();
532
533        assert!(
534            matches!(missing_because_not_float, LadduDataError::MissingColumn(name) if name.as_ref() == "p_py")
535        );
536    }
537}
laddu_data/schema.rs

laddu_data/
schema.rs
