antecedent_data/
dataset.rs

1//! Concrete tabular and time-series dataset types.
2//!
3//! SPDX-License-Identifier: MIT OR Apache-2.0
4
5use std::sync::Arc;
6
7use antecedent_core::{
8    CausalSchema, CausalSchemaBuilder, MeasurementSpec, SmallRoleSet, ValueType, VariableId,
9};
10
11use crate::column::{ColumnView, Float64Column, OwnedColumn, ValidityBitmap};
12use crate::error::DataError;
13use crate::storage::OwnedColumnarStorage;
14use crate::table::TableView;
15use crate::temporal::{SamplingRegularity, TimeIndex};
16
17/// Tabular IID dataset.
18#[derive(Clone, Debug)]
19pub struct TabularData {
20    storage: OwnedColumnarStorage,
21}
22
23impl TabularData {
24    /// Wrap owned storage.
25    #[must_use]
26    pub fn new(storage: OwnedColumnarStorage) -> Self {
27        Self { storage }
28    }
29
30    /// Build continuous `f64` columns from named slices (equal length).
31    ///
32    /// Schema variables are continuous with empty role hints, in iterator order.
33    /// Dense [`VariableId`]s align with column order (`0..n`).
34    ///
35    /// # Errors
36    ///
37    /// Empty input, length mismatch across columns, or schema construction failure.
38    pub fn from_f64_columns<'a, S>(
39        columns: impl IntoIterator<Item = (S, &'a [f64])>,
40    ) -> Result<Self, DataError>
41    where
42        S: Into<Arc<str>>,
43    {
44        let collected: Vec<(Arc<str>, &'a [f64])> =
45            columns.into_iter().map(|(n, v)| (n.into(), v)).collect();
46        if collected.is_empty() {
47            return Err(DataError::InvalidArgument {
48                message: "from_f64_columns requires at least one column".into(),
49            });
50        }
51        let n = collected[0].1.len();
52        let mut b = CausalSchemaBuilder::new();
53        for (name, values) in &collected {
54            if values.len() != n {
55                return Err(DataError::LengthMismatch {
56                    expected: n,
57                    actual: values.len(),
58                    context: "from_f64_columns",
59                });
60            }
61            b.add_variable(
62                Arc::clone(name),
63                ValueType::Continuous,
64                SmallRoleSet::empty(),
65                None,
66                None,
67                MeasurementSpec::default(),
68            )
69            .map_err(|e| DataError::Schema(e.to_string()))?;
70        }
71        let schema = b.build().map_err(|e| DataError::Schema(e.to_string()))?;
72        Self::try_from_schema_f64(
73            schema,
74            collected.iter().map(|(name, values)| (name.as_ref(), *values)),
75        )
76    }
77
78    /// Bind named `f64` slices to an existing schema (names must match; order may differ).
79    ///
80    /// # Errors
81    ///
82    /// Missing/extra names, length mismatch, or storage construction failure.
83    pub fn try_from_schema_f64<'a>(
84        schema: CausalSchema,
85        columns: impl IntoIterator<Item = (&'a str, &'a [f64])>,
86    ) -> Result<Self, DataError> {
87        let n_vars = schema.len();
88        let mut by_name: std::collections::HashMap<&str, &[f64]> = columns.into_iter().collect();
89        if by_name.len() != n_vars {
90            return Err(DataError::InvalidArgument {
91                message: format!("expected {} columns for schema, got {}", n_vars, by_name.len()),
92            });
93        }
94        let mut row_count = None;
95        let mut owned = Vec::with_capacity(n_vars);
96        for var in schema.variables() {
97            let Some(values) = by_name.remove(var.name.as_ref()) else {
98                return Err(DataError::InvalidArgument {
99                    message: format!("missing column for schema variable '{}'", var.name),
100                });
101            };
102            let n = *row_count.get_or_insert(values.len());
103            if values.len() != n {
104                return Err(DataError::LengthMismatch {
105                    expected: n,
106                    actual: values.len(),
107                    context: "try_from_schema_f64",
108                });
109            }
110            let col = Float64Column::new(
111                var.id,
112                Arc::<[f64]>::from(values.to_vec()),
113                ValidityBitmap::all_valid(n),
114            )?;
115            owned.push(OwnedColumn::Float64(col));
116        }
117        if !by_name.is_empty() {
118            let extra: Vec<_> = by_name.keys().copied().collect();
119            return Err(DataError::InvalidArgument {
120                message: format!("columns not in schema: {extra:?}"),
121            });
122        }
123        let storage = OwnedColumnarStorage::try_new(schema, owned, None, None)?;
124        Ok(Self::new(storage))
125    }
126
127    /// Borrow underlying storage.
128    #[must_use]
129    pub fn storage(&self) -> &OwnedColumnarStorage {
130        &self.storage
131    }
132}
133
134impl TableView for TabularData {
135    fn schema(&self) -> &CausalSchema {
136        self.storage.schema()
137    }
138
139    fn row_count(&self) -> usize {
140        self.storage.row_count()
141    }
142
143    fn column(&self, id: VariableId) -> Result<ColumnView<'_>, DataError> {
144        self.storage.column(id)
145    }
146}
147
148/// Temporal / time-series dataset with time index metadata.
149#[derive(Clone, Debug)]
150pub struct TimeSeriesData {
151    storage: OwnedColumnarStorage,
152    time_index: TimeIndex,
153}
154
155impl TimeSeriesData {
156    /// Construct ensuring time index length matches row count.
157    ///
158    /// # Errors
159    ///
160    /// Length mismatch.
161    pub fn try_new(
162        storage: OwnedColumnarStorage,
163        time_index: TimeIndex,
164    ) -> Result<Self, DataError> {
165        if time_index.length != storage.row_count() {
166            return Err(DataError::LengthMismatch {
167                expected: storage.row_count(),
168                actual: time_index.length,
169                context: "time index",
170            });
171        }
172        Ok(Self { storage, time_index })
173    }
174
175    /// Build a regularly sampled series from named `f64` columns.
176    ///
177    /// # Errors
178    ///
179    /// Propagates [`TabularData::from_f64_columns`] errors.
180    pub fn from_f64_columns<'a, S>(
181        columns: impl IntoIterator<Item = (S, &'a [f64])>,
182        interval_ns: u64,
183    ) -> Result<Self, DataError>
184    where
185        S: Into<Arc<str>>,
186    {
187        let tabular = TabularData::from_f64_columns(columns)?;
188        let length = tabular.row_count();
189        Self::try_new(
190            tabular.storage().clone(),
191            TimeIndex { regularity: SamplingRegularity::Regular { interval_ns }, length },
192        )
193    }
194
195    /// Time index metadata.
196    #[must_use]
197    pub fn time_index(&self) -> &TimeIndex {
198        &self.time_index
199    }
200
201    /// Borrow storage.
202    #[must_use]
203    pub fn storage(&self) -> &OwnedColumnarStorage {
204        &self.storage
205    }
206
207    /// Pointer identity of the columnar Arc (tests: planning must not clone payloads).
208    #[cfg(test)]
209    #[must_use]
210    pub(crate) fn columnar_ptr(&self) -> *const [OwnedColumn] {
211        Arc::as_ptr(self.storage.columns_arc())
212    }
213
214    /// Restrict analysis to rows where `mask` is valid, intersected (AND) with any existing
215    /// analysis mask; preserves columns, validity, weights, and time index.
216    ///
217    /// # Errors
218    ///
219    /// Mask length mismatch.
220    pub fn with_analysis_mask(&self, mask: ValidityBitmap) -> Result<Self, DataError> {
221        let storage = &self.storage;
222        let n = storage.row_count();
223        if mask.len() != n {
224            return Err(DataError::LengthMismatch {
225                expected: n,
226                actual: mask.len(),
227                context: "analysis mask",
228            });
229        }
230        let combined = match storage.analysis_mask() {
231            Some(existing) => {
232                let mut bytes = vec![0u8; n.div_ceil(8)];
233                for i in 0..n {
234                    if existing.is_valid(i) && mask.is_valid(i) {
235                        bytes[i / 8] |= 1 << (i % 8);
236                    }
237                }
238                ValidityBitmap::from_bytes(bytes, n)?
239            }
240            None => mask,
241        };
242        let new_storage = OwnedColumnarStorage::try_new(
243            storage.schema().clone(),
244            storage.columns().to_vec(),
245            Some(combined),
246            storage.weights().map(Arc::from),
247        )?;
248        Self::try_new(new_storage, self.time_index.clone())
249    }
250}
251
252impl TableView for TimeSeriesData {
253    fn schema(&self) -> &CausalSchema {
254        self.storage.schema()
255    }
256
257    fn row_count(&self) -> usize {
258        self.storage.row_count()
259    }
260
261    fn column(&self, id: VariableId) -> Result<ColumnView<'_>, DataError> {
262        self.storage.column(id)
263    }
264}
antecedent_data/dataset.rs

antecedent_data/
dataset.rs
