Skip to main content

antecedent_data/
lib.rs

1//! Library-owned causal data views and storage.
2//!
3//! Arrow adapters are optional (`arrow` feature) and never leak Arrow types
4//! into the public causal API (ADR 0004).
5//!
6//! # Building tabular data
7//!
8//! ```
9//! use antecedent_data::{TableView, TabularData};
10//!
11//! let data = TabularData::from_f64_columns([
12//!     ("x", &[1.0_f64, 2.0, 3.0][..]),
13//! ]).unwrap();
14//! assert_eq!(data.row_count(), 3);
15//! ```
16//!
17//! SPDX-License-Identifier: MIT OR Apache-2.0
18
19#![deny(unsafe_code)]
20#![deny(missing_docs)]
21#![warn(clippy::missing_errors_doc, clippy::missing_panics_doc)]
22
23pub mod aligned_buffer;
24#[cfg(feature = "arrow")]
25pub mod arrow_adapter;
26#[cfg(feature = "arrow")]
27pub mod arrow_ffi;
28pub mod buffer;
29pub mod categorical;
30pub mod column;
31pub mod dataset;
32pub mod error;
33pub mod event;
34pub mod lagged_frame;
35pub mod materialize;
36pub mod multi_env;
37pub mod multi_env_plan;
38pub mod network;
39pub mod panel;
40pub mod pooled_frame;
41pub mod project;
42pub mod reference;
43pub mod resample;
44pub mod sample;
45pub mod sample_policy;
46pub mod sample_request;
47pub mod selection;
48pub mod sim;
49pub mod split;
50pub mod storage;
51pub mod surrogate;
52pub mod table;
53pub mod temporal;
54pub mod transforms;
55pub mod vector_vars;
56
57#[cfg(test)]
58mod testing;
59
60pub use aligned_buffer::AlignedBuffer;
61#[cfg(feature = "arrow")]
62pub use arrow_adapter::{ArrowLoadResult, tabular_from_arrow_c_columns, tabular_from_record_batch};
63#[cfg(feature = "arrow")]
64pub use arrow_ffi::{ArrowCColumn, FfiArrowArray, FfiArrowSchema};
65pub use buffer::{F64Buffer, ForeignBufferOwner, ForeignF64Buffer};
66pub use categorical::{
67    CategoricalColumn, CategoricalView, CategoryCode, CategoryDomain, CategoryLevel, Contrast,
68    ContrastMatrix, UnknownCategoryPolicy, compile_contrast_matrix,
69};
70pub use column::{
71    BooleanColumn, ColumnView, FixedVectorColumn, Float64Column, Int64Column, OwnedColumn,
72    TimestampColumn, ValidityBitmap,
73};
74pub use dataset::{TabularData, TimeSeriesData};
75pub use error::DataError;
76pub use event::EventData;
77pub use lagged_frame::{LaggedFrame, LaggedFrameOptions};
78pub use materialize::{MaterializationReason, materialization_diagnostic};
79pub use multi_env::MultiEnvironmentData;
80pub use multi_env_plan::{MultiEnvSamplePlan, PanelSamplePlan, plans_for_series_lengths};
81pub use network::{NetworkData, NetworkEdge};
82pub use panel::{PanelData, PanelUnit, PanelUnitView};
83pub use pooled_frame::{
84    DEFAULT_MAX_TIME_ONE_HOT_LEVELS, DummyOptions, PooledLaggedFrame, TimeDummyEncoding,
85    pool_multi_env_lagged_frame,
86};
87pub use project::{IdRemap, dedupe_variable_ids};
88pub use reference::ReferencePointPolicy;
89pub use resample::{
90    PermutationScheme, ResamplingPlan, fill_resample_index_batch, fill_resample_indexes,
91    fill_resample_indexes_grouped, fill_resample_weight_batch, fill_resample_weights,
92    resample_timeseries, resample_timeseries_grouped,
93};
94pub use sample::{
95    DropSummary, LagMap, LaggedColumn, LaggedPreparedSample, LaggedSamplePlan,
96    LaggedSampleWorkspace,
97};
98pub use sample_policy::{MaskPolicy, MissingPolicy, WeightPolicy};
99pub use sample_request::{
100    MatrixRef, PreparedColumn, PreparedRowSelector, PreparedSample, RowSelectionRef,
101    SampleCacheKey, SampleLayout, SamplePartitions, SamplePlan, SampleRequest, SampleWorkspace,
102};
103pub use sim::{KnownLaggedParent, LaggedLinearPair};
104pub use split::{
105    BlockedTemporalSplit, ClusterSplit, DiscoveryEstimationSplit, EnvHoldoutSplit, GroupedSplit,
106    RandomIidSplit, RegimeHoldoutSplit, RollingOriginSplit, RowSplit, TemporalFold,
107    TemporalRandomPolicy, TimeRange, ensure_random_allowed_on_temporal,
108};
109pub use storage::OwnedColumnarStorage;
110pub use surrogate::{surrogate_permute_columns, surrogate_phase_randomize};
111pub use table::TableView;
112pub use temporal::{SamplingRegularity, TemporalIndexer, TemporalNodeKey, TimeIndex};
113pub use transforms::{equal_width_bin, moving_average, ordinal_patterns};
114pub use vector_vars::{VectorVariableGroups, column_blocks_for_frame, expand_fixed_vector_columns};
115
116#[cfg(test)]
117#[allow(clippy::cast_precision_loss)]
118mod tests {
119    use std::sync::Arc;
120
121    use antecedent_core::{
122        CausalSchemaBuilder, MeasurementSpec, RoleHint, SmallRoleSet, ValueType, VariableId,
123    };
124
125    use super::*;
126
127    fn two_col_table() -> OwnedColumnarStorage {
128        let mut b = CausalSchemaBuilder::new();
129        b.add_variable(
130            "x",
131            ValueType::Continuous,
132            SmallRoleSet::from_hint(RoleHint::TreatmentCandidate),
133            None,
134            None,
135            MeasurementSpec::default(),
136        )
137        .unwrap();
138        b.add_variable(
139            "y",
140            ValueType::Continuous,
141            SmallRoleSet::from_hint(RoleHint::OutcomeCandidate),
142            None,
143            None,
144            MeasurementSpec::default(),
145        )
146        .unwrap();
147        let schema = b.build().unwrap();
148        let n = 1_000usize;
149        let x = Float64Column::new(
150            VariableId::from_raw(0),
151            Arc::<[f64]>::from((0..n).map(|i| i as f64).collect::<Vec<_>>()),
152            ValidityBitmap::all_valid(n),
153        )
154        .unwrap();
155        let y = Float64Column::new(
156            VariableId::from_raw(1),
157            Arc::<[f64]>::from((0..n).map(|i| (i * 2) as f64).collect::<Vec<_>>()),
158            ValidityBitmap::all_valid(n),
159        )
160        .unwrap();
161        OwnedColumnarStorage::try_new(
162            schema,
163            vec![OwnedColumn::Float64(x), OwnedColumn::Float64(y)],
164            None,
165            None,
166        )
167        .unwrap()
168    }
169
170    #[test]
171    fn table_view_returns_columns_by_id() {
172        let table = two_col_table();
173        assert_eq!(table.row_count(), 1000);
174        let col = table.column(VariableId::from_raw(0)).unwrap();
175        assert_eq!(col.len(), 1000);
176        match col {
177            ColumnView::Float64(c) => {
178                assert!((c.values[10] - 10.0).abs() < f64::EPSILON);
179            }
180            _ => panic!("expected float64"),
181        }
182    }
183
184    #[test]
185    fn prepared_column_view_does_not_reallocate() {
186        let table = two_col_table();
187        let col = table.column(VariableId::from_raw(0)).unwrap();
188        let ColumnView::Float64(c) = col else {
189            panic!("expected float");
190        };
191        let ptr = c.values.as_ptr();
192        for _ in 0..100 {
193            let again = table.column(VariableId::from_raw(0)).unwrap();
194            let ColumnView::Float64(c2) = again else {
195                panic!("expected float");
196            };
197            assert_eq!(c2.values.as_ptr(), ptr);
198            let view = c2.as_f64_view();
199            assert_eq!(view.len(), 1000);
200        }
201    }
202
203    #[test]
204    fn timeseries_wraps_storage() {
205        let storage = two_col_table();
206        let ts = TimeSeriesData::try_new(
207            storage,
208            TimeIndex {
209                regularity: SamplingRegularity::Regular { interval_ns: 1_000 },
210                length: 1000,
211            },
212        )
213        .unwrap();
214        assert_eq!(ts.row_count(), 1000);
215    }
216}