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
19/// A schema-resolved scalar column binding.
20///
21/// Resolve bindings once against the schema used by an
22/// [`EventBatch`](crate::data::EventBatch), then reuse them for every row.
23/// Bindings must not be used with a batch whose schema has a different column
24/// ordering.
25#[derive(Clone, Debug, PartialEq)]
26pub struct ScalarBinding {
27    schema: Arc<Schema>,
28    index: usize,
29}
30
31impl ScalarBinding {
32    /// Returns the physical scalar-column index.
33    pub fn index(&self) -> usize {
34        self.index
35    }
36
37    pub(crate) fn matches(&self, schema: &Schema) -> bool {
38        self.schema.as_ref() == schema
39    }
40}
41
42/// A schema-resolved four-momentum column binding.
43///
44/// Resolve bindings once against the schema used by an
45/// [`EventBatch`](crate::data::EventBatch), then reuse them for every row.
46/// Bindings must not be used with a batch whose schema has a different column
47/// ordering.
48#[derive(Clone, Debug, PartialEq)]
49pub struct P4Binding {
50    schema: Arc<Schema>,
51    index: usize,
52}
53
54impl P4Binding {
55    /// Returns the physical four-momentum-column index.
56    pub fn index(&self) -> usize {
57        self.index
58    }
59
60    pub(crate) fn matches(&self, schema: &Schema) -> bool {
61        self.schema.as_ref() == schema
62    }
63}
64
65impl PartialEq for Schema {
66    fn eq(&self, other: &Self) -> bool {
67        self.p4s == other.p4s
68            && self.scalars == other.scalars
69            && self.has_weight == other.has_weight
70    }
71}
72
73impl Schema {
74    /// Validates and constructs a logical schema.
75    ///
76    /// # Errors
77    ///
78    /// Returns [`LadduDataError`] when four-momentum or scalar column names are
79    /// duplicated.
80    pub fn new(
81        p4s: impl IntoIterator<Item = impl Into<Name>>,
82        scalars: impl IntoIterator<Item = impl Into<Name>>,
83        has_weight: bool,
84    ) -> LadduDataResult<Self> {
85        let p4s: Vec<Name> = p4s.into_iter().map(Into::into).collect();
86        let scalars: Vec<Name> = scalars.into_iter().map(Into::into).collect();
87        let p4_index = Arc::new(make_index(&p4s, "p4")?);
88        let scalar_index = Arc::new(make_index(&scalars, "scalar")?);
89        Ok(Self {
90            p4s,
91            scalars,
92            has_weight,
93            p4_index,
94            scalar_index,
95        })
96    }
97
98    /// Returns the four-momentum column index for `name`.
99    pub fn p4_index(&self, name: &str) -> Option<usize> {
100        self.p4_index.get(name).copied()
101    }
102
103    /// Resolves a four-momentum column name once for repeated row access.
104    pub fn bind_p4(&self, name: &str) -> Option<P4Binding> {
105        self.p4_index(name).map(|index| P4Binding {
106            schema: Arc::new(self.clone()),
107            index,
108        })
109    }
110
111    /// Returns the scalar column index for `name`.
112    pub fn scalar_index(&self, name: &str) -> Option<usize> {
113        self.scalar_index.get(name).copied()
114    }
115
116    /// Resolves a scalar column name once for repeated row access.
117    pub fn bind_scalar(&self, name: &str) -> Option<ScalarBinding> {
118        self.scalar_index(name).map(|index| ScalarBinding {
119            schema: Arc::new(self.clone()),
120            index,
121        })
122    }
123
124    /// Returns four-momentum names in column order.
125    pub fn p4s(&self) -> &[Name] {
126        &self.p4s
127    }
128
129    /// Returns scalar names in column order.
130    pub fn scalars(&self) -> &[Name] {
131        &self.scalars
132    }
133
134    /// Returns whether events carry explicit weights.
135    pub fn has_weight(&self) -> bool {
136        self.has_weight
137    }
138
139    /// Returns the number of four-momentum columns.
140    pub fn n_p4s(&self) -> usize {
141        self.p4s.len()
142    }
143
144    /// Returns the number of scalar columns.
145    pub fn n_scalars(&self) -> usize {
146        self.scalars.len()
147    }
148
149    /// Requires and returns a four-momentum column index.
150    ///
151    /// # Errors
152    ///
153    /// Returns [`LadduDataError::MissingColumn`] when `name` is not a
154    /// four-momentum column.
155    pub fn require_p4(&self, name: &str) -> LadduDataResult<usize> {
156        self.p4_index(name)
157            .ok_or_else(|| LadduDataError::MissingColumn(Name::from(name)))
158    }
159
160    /// Requires and returns a scalar column index.
161    ///
162    /// # Errors
163    ///
164    /// Returns [`LadduDataError::MissingColumn`] when `name` is not a scalar
165    /// column.
166    pub fn require_scalar(&self, name: &str) -> LadduDataResult<usize> {
167        self.scalar_index(name)
168            .ok_or_else(|| LadduDataError::MissingColumn(Name::from(name)))
169    }
170}
171
172fn make_index(names: &[Name], kind: &'static str) -> LadduDataResult<HashMap<Name, usize>> {
173    let mut out = HashMap::with_capacity(names.len());
174    for (i, name) in names.iter().cloned().enumerate() {
175        if out.insert(name.clone(), i).is_some() {
176            return Err(LadduDataError::Schema(format!(
177                "duplicate {kind} column: {name}"
178            )));
179        }
180    }
181    Ok(out)
182}
183
184/// Physical naming conventions used to map a logical schema to storage columns.
185#[derive(Clone, Debug)]
186pub struct SchemaColumnNames {
187    /// Physical weight-column name.
188    pub weight_column: Name,
189    /// Suffixes for four-momentum components.
190    pub p4_suffixes: P4Suffixes,
191}
192
193impl Default for SchemaColumnNames {
194    fn default() -> Self {
195        Self {
196            weight_column: Name::from("weight"),
197            p4_suffixes: P4Suffixes::default(),
198        }
199    }
200}
201
202/// Options controlling logical schema inference from physical columns.
203#[derive(Clone, Debug)]
204pub struct SchemaInferenceOptions {
205    /// Physical naming conventions.
206    pub column_names: SchemaColumnNames,
207    /// Whether inference fails if the weight column is absent.
208    pub require_weight: bool,
209    /// Whether incomplete four-momenta become independent scalar columns.
210    pub incomplete_p4_components_are_scalars: bool,
211}
212
213impl Default for SchemaInferenceOptions {
214    fn default() -> Self {
215        Self {
216            column_names: SchemaColumnNames::default(),
217            require_weight: false,
218            incomplete_p4_components_are_scalars: true,
219        }
220    }
221}
222
223/// Physical suffixes for `(E, px, py, pz)` columns.
224#[derive(Clone, Debug)]
225pub struct P4Suffixes {
226    /// Energy suffix.
227    pub e: &'static str,
228    /// X-momentum suffix.
229    pub px: &'static str,
230    /// Y-momentum suffix.
231    pub py: &'static str,
232    /// Z-momentum suffix.
233    pub pz: &'static str,
234}
235
236impl Default for P4Suffixes {
237    fn default() -> Self {
238        Self {
239            e: "_e",
240            px: "_px",
241            py: "_py",
242            pz: "_pz",
243        }
244    }
245}
246
247impl P4Suffixes {
248    /// Splits a matching physical name into logical prefix and component index.
249    pub fn component<'a>(&'a self, name: &'a str) -> Option<(&'a str, usize)> {
250        if let Some(prefix) = name.strip_suffix(self.e) {
251            Some((prefix, 0))
252        } else if let Some(prefix) = name.strip_suffix(self.px) {
253            Some((prefix, 1))
254        } else if let Some(prefix) = name.strip_suffix(self.py) {
255            Some((prefix, 2))
256        } else if let Some(prefix) = name.strip_suffix(self.pz) {
257            Some((prefix, 3))
258        } else {
259            None
260        }
261    }
262
263    /// Produces the four physical names for a logical four-momentum prefix.
264    pub fn physical_p4_names(&self, prefix: &str) -> [String; 4] {
265        [
266            format!("{prefix}{}", self.e),
267            format!("{prefix}{}", self.px),
268            format!("{prefix}{}", self.py),
269            format!("{prefix}{}", self.pz),
270        ]
271    }
272}
273
274/// Physical storage type relevant to schema inference.
275#[derive(Clone, Copy, Debug, PartialEq, Eq)]
276pub enum ColumnType {
277    /// 64-bit floating point.
278    F64,
279    /// 32-bit floating point.
280    F32,
281    /// Any unsupported type.
282    Other,
283}
284
285impl ColumnType {
286    /// Returns whether this type can populate an event-data column.
287    pub fn is_supported_float(self) -> bool {
288        matches!(self, Self::F64 | Self::F32)
289    }
290}
291
292/// Name and physical type of one available storage column.
293#[derive(Clone, Copy, Debug)]
294pub struct ColumnInfo<'a> {
295    /// Physical column name.
296    pub name: &'a str,
297    /// Physical column type.
298    pub dtype: ColumnType,
299}
300
301impl Schema {
302    /// Infers a logical schema from available physical columns.
303    ///
304    /// # Errors
305    ///
306    /// Returns [`LadduDataError`] when required momentum components or weights
307    /// are missing, names are ambiguous, or inferred logical names conflict.
308    pub fn infer_from_columns<'a>(
309        columns: impl IntoIterator<Item = ColumnInfo<'a>>,
310        options: &SchemaInferenceOptions,
311    ) -> LadduDataResult<Self> {
312        let mut p4_candidates = BTreeMap::<String, [bool; 4]>::new();
313        let mut scalar_names = Vec::<Name>::new();
314        let mut has_weight = false;
315
316        for col in columns {
317            if !col.dtype.is_supported_float() {
318                continue;
319            }
320
321            if col.name == options.column_names.weight_column.as_ref() {
322                has_weight = true;
323                continue;
324            }
325
326            if let Some((prefix, component)) = options.column_names.p4_suffixes.component(col.name)
327            {
328                p4_candidates.entry(prefix.to_owned()).or_default()[component] = true;
329            } else {
330                scalar_names.push(Name::from(col.name));
331            }
332        }
333
334        let mut p4s = Vec::<Name>::new();
335
336        for (prefix, seen) in p4_candidates {
337            if seen == [true, true, true, true] {
338                p4s.push(Name::from(prefix));
339            } else if options.incomplete_p4_components_are_scalars {
340                let names = options.column_names.p4_suffixes.physical_p4_names(&prefix);
341
342                for (i, name) in names.into_iter().enumerate() {
343                    if seen[i] {
344                        scalar_names.push(Name::from(name));
345                    }
346                }
347            }
348        }
349
350        if options.require_weight && !has_weight {
351            return Err(LadduDataError::MissingColumn(Arc::clone(
352                &options.column_names.weight_column,
353            )));
354        }
355
356        Schema::new(p4s, scalar_names, has_weight)
357    }
358
359    /// Returns all physical columns required to store this schema.
360    pub fn physical_columns(&self, column_names: &SchemaColumnNames) -> Vec<Name> {
361        PhysicalSchemaPlan::for_read(self, column_names)
362            .columns()
363            .iter()
364            .map(|column| Arc::clone(column.name()))
365            .collect()
366    }
367
368    /// Validates that all physical columns required by this schema are available.
369    ///
370    /// # Errors
371    ///
372    /// Returns [`LadduDataError::MissingColumn`] when a required physical
373    /// column is absent or has an unsupported type.
374    pub fn validate_required_columns<'a>(
375        &self,
376        available: impl IntoIterator<Item = ColumnInfo<'a>>,
377        options: &SchemaInferenceOptions,
378    ) -> LadduDataResult<()> {
379        let available: HashSet<&str> = available
380            .into_iter()
381            .filter(|c| c.dtype.is_supported_float())
382            .map(|c| c.name)
383            .collect();
384
385        for required in PhysicalSchemaPlan::for_read(self, &options.column_names).columns() {
386            if !available.contains(required.name().as_ref()) {
387                return Err(LadduDataError::MissingColumn(Arc::clone(required.name())));
388            }
389        }
390
391        Ok(())
392    }
393}
394
395/// Floating-point precision used when writing physical columns.
396#[derive(Copy, Clone, Debug, Default)]
397pub enum Precision {
398    /// Write 64-bit floating-point values.
399    #[default]
400    F64,
401    /// Write 32-bit floating-point values.
402    F32,
403}
404
405/// Policy controlling whether sinks emit a weight column.
406#[derive(Clone, Copy, Debug, Default)]
407pub enum WriteWeightColumn {
408    /// Always write weights, using unit weights when the schema has none.
409    #[default]
410    Always,
411    /// Write weights only when the logical schema contains them.
412    OnlyIfPresent,
413}
414
415/// Physical naming, precision, and weight policy for event sinks.
416#[derive(Clone, Debug, Default)]
417pub struct SchemaWriteOptions {
418    /// Physical column naming conventions.
419    pub column_names: SchemaColumnNames,
420    /// Floating-point output precision.
421    pub precision: Precision,
422    /// Weight-column emission policy.
423    pub write_weight_column: WriteWeightColumn,
424}
425
426/// The semantic role of one physical storage column.
427///
428/// This is deliberately crate-private: callers work with logical [`Schema`]
429/// values while the file-format adapters use this plan to keep physical
430/// column ordering and binding identical across backends.
431#[derive(Clone, Copy, Debug, PartialEq, Eq)]
432pub(crate) enum PhysicalColumnRole {
433    /// One component of a logical four-momentum column.
434    P4 {
435        /// Logical four-momentum column index.
436        index: usize,
437        /// Component index in `(E, px, py, pz)` order.
438        component: usize,
439    },
440    /// One logical scalar column.
441    Scalar {
442        /// Logical scalar column index.
443        index: usize,
444    },
445    /// The physical event-weight column.
446    Weight,
447}
448
449/// Ordered physical representation of a logical event schema.
450///
451/// A plan is built once at each storage boundary and then consumed by schema
452/// creation, projection, encoding, and decoding.  Keeping the role alongside
453/// the name avoids each backend reimplementing the canonical ordering.
454#[derive(Clone, Debug)]
455pub(crate) struct PhysicalSchemaPlan {
456    columns: Vec<PhysicalColumn>,
457}
458
459#[derive(Clone, Debug)]
460pub(crate) struct PhysicalColumn {
461    name: Name,
462    role: PhysicalColumnRole,
463}
464
465impl PhysicalSchemaPlan {
466    /// Builds the physical columns used when reading or validating a schema.
467    pub(crate) fn for_read(schema: &Schema, column_names: &SchemaColumnNames) -> Self {
468        Self::build(schema, column_names, schema.has_weight())
469    }
470
471    /// Builds the physical columns emitted by a sink.
472    pub(crate) fn for_write(
473        schema: &Schema,
474        options: &SchemaWriteOptions,
475        write_weight: WriteWeightColumn,
476    ) -> Self {
477        let should_write_weight =
478            matches!(write_weight, WriteWeightColumn::Always) || schema.has_weight();
479        Self::build(schema, &options.column_names, should_write_weight)
480    }
481
482    fn build(schema: &Schema, column_names: &SchemaColumnNames, include_weight: bool) -> Self {
483        let mut columns = Vec::with_capacity(
484            4 * schema.n_p4s() + schema.n_scalars() + usize::from(include_weight),
485        );
486
487        for (index, p4) in schema.p4s().iter().enumerate() {
488            for (component, name) in column_names
489                .p4_suffixes
490                .physical_p4_names(p4)
491                .into_iter()
492                .enumerate()
493            {
494                columns.push(PhysicalColumn {
495                    name: Name::from(name),
496                    role: PhysicalColumnRole::P4 { index, component },
497                });
498            }
499        }
500
501        for (index, name) in schema.scalars().iter().cloned().enumerate() {
502            columns.push(PhysicalColumn {
503                name,
504                role: PhysicalColumnRole::Scalar { index },
505            });
506        }
507
508        if include_weight {
509            columns.push(PhysicalColumn {
510                name: Arc::clone(&column_names.weight_column),
511                role: PhysicalColumnRole::Weight,
512            });
513        }
514
515        Self { columns }
516    }
517
518    /// Returns physical columns in canonical storage order.
519    pub(crate) fn columns(&self) -> &[PhysicalColumn] {
520        &self.columns
521    }
522}
523
524impl PhysicalColumn {
525    /// Returns the physical storage name.
526    pub(crate) fn name(&self) -> &Name {
527        &self.name
528    }
529
530    /// Returns the logical role of this physical column.
531    pub(crate) fn role(&self) -> PhysicalColumnRole {
532        self.role
533    }
534}
535
536#[cfg(test)]
537mod tests {
538    use super::*;
539
540    fn col(name: &'static str, dtype: ColumnType) -> ColumnInfo<'static> {
541        ColumnInfo { name, dtype }
542    }
543
544    #[test]
545    fn schema_new_rejects_duplicates_and_required_lookup_reports_missing_column() {
546        let duplicate_p4 = Schema::new(["p", "p"], ["mass"], false);
547        assert!(matches!(duplicate_p4, Err(LadduDataError::Schema(_))));
548
549        let duplicate_scalar = Schema::new(["p"], ["mass", "mass"], false);
550        assert!(matches!(duplicate_scalar, Err(LadduDataError::Schema(_))));
551
552        let schema = Schema::new(["beam", "recoil"], ["mass", "costheta"], true).unwrap();
553
554        assert_eq!(schema.require_p4("recoil").unwrap(), 1);
555        assert_eq!(schema.require_scalar("costheta").unwrap(), 1);
556
557        let err = schema.require_scalar("missing").unwrap_err();
558        assert!(matches!(err, LadduDataError::MissingColumn(name) if name.as_ref() == "missing"));
559    }
560
561    #[test]
562    fn schema_bindings_resolve_typed_indices_once() {
563        let schema = Schema::new(["beam", "recoil"], ["mass", "costheta"], true).unwrap();
564
565        assert_eq!(schema.bind_p4("recoil").unwrap().index(), 1);
566        assert_eq!(schema.bind_scalar("costheta").unwrap().index(), 1);
567        assert_eq!(schema.bind_p4("missing"), None);
568        assert_eq!(schema.bind_scalar("missing"), None);
569    }
570
571    #[test]
572    fn infer_from_columns_groups_complete_p4s_keeps_incomplete_components_as_scalars_and_ignores_nonfloats()
573     {
574        let options = SchemaInferenceOptions::default();
575
576        let schema = Schema::infer_from_columns(
577            [
578                col("gamma_px", ColumnType::F64),
579                col("gamma_py", ColumnType::F64),
580                col("gamma_pz", ColumnType::F32),
581                col("gamma_e", ColumnType::F64),
582                col("partial_px", ColumnType::F64),
583                col("partial_e", ColumnType::F64),
584                col("mass", ColumnType::F32),
585                col("ignored", ColumnType::Other),
586                col("weight", ColumnType::F64),
587            ],
588            &options,
589        )
590        .unwrap();
591
592        assert_eq!(
593            schema
594                .p4s()
595                .iter()
596                .map(|n| n.to_string())
597                .collect::<Vec<_>>(),
598            vec!["gamma"]
599        );
600
601        assert_eq!(
602            schema
603                .scalars()
604                .iter()
605                .map(|n| n.to_string())
606                .collect::<Vec<_>>(),
607            vec!["mass", "partial_e", "partial_px"]
608        );
609
610        assert!(schema.has_weight());
611    }
612
613    #[test]
614    fn infer_from_columns_can_discard_incomplete_p4_components_and_require_weight() {
615        let options = SchemaInferenceOptions {
616            incomplete_p4_components_are_scalars: false,
617            ..Default::default()
618        };
619
620        let schema = Schema::infer_from_columns(
621            [
622                col("partial_px", ColumnType::F64),
623                col("partial_e", ColumnType::F64),
624                col("mass", ColumnType::F64),
625                col("weight", ColumnType::F64),
626            ],
627            &options,
628        )
629        .unwrap();
630
631        assert!(schema.p4s().is_empty());
632        assert_eq!(
633            schema
634                .scalars()
635                .iter()
636                .map(|n| n.to_string())
637                .collect::<Vec<_>>(),
638            vec!["mass"]
639        );
640
641        let require_weight = SchemaInferenceOptions {
642            require_weight: true,
643            ..Default::default()
644        };
645
646        let err = Schema::infer_from_columns([col("mass", ColumnType::F64)], &require_weight)
647            .unwrap_err();
648
649        assert!(matches!(err, LadduDataError::MissingColumn(name) if name.as_ref() == "weight"));
650    }
651
652    #[test]
653    fn physical_columns_and_validation_respect_custom_names_and_float_types_only() {
654        let schema = Schema::new(["p"], ["mass"], true).unwrap();
655
656        let names = SchemaColumnNames {
657            weight_column: Name::from("event_weight"),
658            ..Default::default()
659        };
660
661        let physical = schema
662            .physical_columns(&names)
663            .into_iter()
664            .map(|n| n.to_string())
665            .collect::<Vec<_>>();
666
667        assert_eq!(
668            physical,
669            vec!["p_e", "p_px", "p_py", "p_pz", "mass", "event_weight"]
670        );
671
672        let options = SchemaInferenceOptions {
673            column_names: names,
674            ..Default::default()
675        };
676
677        let ok = schema.validate_required_columns(
678            [
679                col("p_e", ColumnType::F32),
680                col("p_px", ColumnType::F64),
681                col("p_py", ColumnType::F64),
682                col("p_pz", ColumnType::F32),
683                col("mass", ColumnType::F64),
684                col("event_weight", ColumnType::F64),
685            ],
686            &options,
687        );
688
689        assert!(ok.is_ok());
690
691        let missing_because_not_float = schema
692            .validate_required_columns(
693                [
694                    col("p_e", ColumnType::F32),
695                    col("p_px", ColumnType::F64),
696                    col("p_py", ColumnType::Other),
697                    col("p_pz", ColumnType::F32),
698                    col("mass", ColumnType::F64),
699                    col("event_weight", ColumnType::F64),
700                ],
701                &options,
702            )
703            .unwrap_err();
704
705        assert!(
706            matches!(missing_because_not_float, LadduDataError::MissingColumn(name) if name.as_ref() == "p_py")
707        );
708    }
709
710    #[test]
711    fn physical_schema_plan_preserves_order_roles_and_weight_policy() {
712        let schema = Schema::new(["p"], ["mass"], false).unwrap();
713        let options = SchemaWriteOptions {
714            column_names: SchemaColumnNames {
715                weight_column: Name::from("event_weight"),
716                ..Default::default()
717            },
718            ..Default::default()
719        };
720
721        let only_if_present =
722            PhysicalSchemaPlan::for_write(&schema, &options, WriteWeightColumn::OnlyIfPresent);
723        assert_eq!(
724            only_if_present
725                .columns()
726                .iter()
727                .map(|column| column.name().to_string())
728                .collect::<Vec<_>>(),
729            ["p_e", "p_px", "p_py", "p_pz", "mass"]
730        );
731        assert_eq!(
732            only_if_present
733                .columns()
734                .iter()
735                .map(PhysicalColumn::role)
736                .collect::<Vec<_>>(),
737            [
738                PhysicalColumnRole::P4 {
739                    index: 0,
740                    component: 0,
741                },
742                PhysicalColumnRole::P4 {
743                    index: 0,
744                    component: 1,
745                },
746                PhysicalColumnRole::P4 {
747                    index: 0,
748                    component: 2,
749                },
750                PhysicalColumnRole::P4 {
751                    index: 0,
752                    component: 3,
753                },
754                PhysicalColumnRole::Scalar { index: 0 },
755            ]
756        );
757
758        let always = PhysicalSchemaPlan::for_write(&schema, &options, WriteWeightColumn::Always);
759        assert_eq!(
760            always.columns().last().map(|column| column.name().as_ref()),
761            Some("event_weight")
762        );
763        assert_eq!(
764            always.columns().last().map(PhysicalColumn::role),
765            Some(PhysicalColumnRole::Weight)
766        );
767    }
768}
laddu_data/schema.rs

laddu_data/
schema.rs
