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