antecedent_data/
column.rs

1//! Columnar storage and typed column views.
2//!
3//! SPDX-License-Identifier: MIT OR Apache-2.0
4
5use std::sync::Arc;
6
7use antecedent_core::VariableId;
8use antecedent_kernels::{BitMaskView, F64VectorView};
9
10use crate::buffer::F64Buffer;
11use crate::categorical::CategoricalColumn;
12use crate::error::DataError;
13
14/// Packed validity bitmap (`1` = valid). Missingness is never a sentinel value.
15#[derive(Clone, Debug, Default, Eq, PartialEq)]
16pub struct ValidityBitmap {
17    bytes: Arc<[u8]>,
18    len: usize,
19}
20
21impl ValidityBitmap {
22    /// All-valid bitmap of `len` bits.
23    #[must_use]
24    pub fn all_valid(len: usize) -> Self {
25        let n = len.div_ceil(8);
26        Self { bytes: Arc::from(vec![0xFFu8; n].into_boxed_slice()), len }
27    }
28
29    /// Construct from raw bytes.
30    ///
31    /// # Errors
32    ///
33    /// When the buffer is shorter than `ceil(len/8)`.
34    pub fn from_bytes(bytes: impl Into<Arc<[u8]>>, len: usize) -> Result<Self, DataError> {
35        let bytes = bytes.into();
36        if bytes.len() < len.div_ceil(8) {
37            return Err(DataError::InvalidValidity { message: "validity buffer too short" });
38        }
39        Ok(Self { bytes, len })
40    }
41
42    /// Bit length.
43    #[must_use]
44    pub const fn len(&self) -> usize {
45        self.len
46    }
47
48    /// Whether empty.
49    #[must_use]
50    pub const fn is_empty(&self) -> bool {
51        self.len == 0
52    }
53
54    /// Borrow as a kernel mask view.
55    ///
56    /// # Errors
57    ///
58    /// Propagates view construction errors.
59    pub fn as_mask_view(&self) -> Result<BitMaskView<'_>, DataError> {
60        BitMaskView::new(&self.bytes, self.len)
61            .map_err(|_| DataError::InvalidValidity { message: "mask view rejected buffer" })
62    }
63
64    /// Whether row `i` is valid.
65    #[must_use]
66    pub fn is_valid(&self, i: usize) -> bool {
67        self.as_mask_view().is_ok_and(|m| m.get(i))
68    }
69
70    /// Whether every bit is valid.
71    #[must_use]
72    pub fn is_all_valid(&self) -> bool {
73        self.as_mask_view().is_ok_and(|m| (0..self.len).all(|i| m.get(i)))
74    }
75
76    /// Gather bits through a row map (`out[i] = self[row_map[i]]`).
77    ///
78    /// # Errors
79    ///
80    /// When a mapped row is out of range.
81    pub fn gather(&self, row_map: &[u32]) -> Result<Self, DataError> {
82        let mask = self.as_mask_view()?;
83        let n = row_map.len();
84        let mut bytes = vec![0u8; n.div_ceil(8)];
85        for (i, &r) in row_map.iter().enumerate() {
86            let r = r as usize;
87            if r >= self.len {
88                return Err(DataError::InvalidValidity { message: "row map exceeds bitmap" });
89            }
90            if mask.get(r) {
91                bytes[i / 8] |= 1 << (i % 8);
92            }
93        }
94        Self::from_bytes(bytes, n)
95    }
96
97    /// Gather bits through a `usize` row map.
98    ///
99    /// # Errors
100    ///
101    /// When a mapped row is out of range.
102    pub fn gather_rows(&self, row_map: &[usize]) -> Result<Self, DataError> {
103        let mask = self.as_mask_view()?;
104        let n = row_map.len();
105        let mut bytes = vec![0u8; n.div_ceil(8)];
106        for (i, &r) in row_map.iter().enumerate() {
107            if r >= self.len {
108                return Err(DataError::InvalidValidity { message: "row map exceeds bitmap" });
109            }
110            if mask.get(r) {
111                bytes[i / 8] |= 1 << (i % 8);
112            }
113        }
114        Self::from_bytes(bytes, n)
115    }
116
117    /// Compact to rows where `keep[i]` is true.
118    ///
119    /// # Errors
120    ///
121    /// Length mismatch.
122    pub fn compact(&self, keep: &[bool]) -> Result<Self, DataError> {
123        if keep.len() != self.len {
124            return Err(DataError::LengthMismatch {
125                expected: self.len,
126                actual: keep.len(),
127                context: "validity compact keep",
128            });
129        }
130        let n_new = keep.iter().filter(|&&k| k).count();
131        let mut bytes = vec![0u8; n_new.div_ceil(8)];
132        let mut j = 0usize;
133        for (i, &k) in keep.iter().enumerate() {
134            if k {
135                if self.is_valid(i) {
136                    bytes[j / 8] |= 1 << (j % 8);
137                }
138                j += 1;
139            }
140        }
141        Self::from_bytes(bytes, n_new)
142    }
143
144    /// Concatenate bitmaps end-to-end.
145    ///
146    /// # Errors
147    ///
148    /// Propagates bitmap construction errors.
149    pub fn concat(parts: &[&Self]) -> Result<Self, DataError> {
150        let n: usize = parts.iter().map(|p| p.len).sum();
151        let mut bytes = vec![0u8; n.div_ceil(8)];
152        let mut offset = 0usize;
153        for part in parts {
154            for i in 0..part.len {
155                if part.is_valid(i) {
156                    let j = offset + i;
157                    bytes[j / 8] |= 1 << (j % 8);
158                }
159            }
160            offset += part.len;
161        }
162        Self::from_bytes(bytes, n)
163    }
164}
165
166/// Float64 column (owned or foreign-backed values).
167#[derive(Clone, Debug, PartialEq)]
168pub struct Float64Column {
169    /// Variable id.
170    pub id: VariableId,
171    /// Values (sentinel-free; use validity for missing).
172    pub values: F64Buffer,
173    /// Validity bitmap.
174    pub validity: ValidityBitmap,
175}
176
177impl Float64Column {
178    /// Construct a column from owned values; lengths must match.
179    ///
180    /// # Errors
181    ///
182    /// [`DataError::LengthMismatch`] when validity length differs.
183    pub fn new(
184        id: VariableId,
185        values: impl Into<F64Buffer>,
186        validity: ValidityBitmap,
187    ) -> Result<Self, DataError> {
188        let values = values.into();
189        if validity.len() != values.len() {
190            return Err(DataError::LengthMismatch {
191                expected: values.len(),
192                actual: validity.len(),
193                context: "float64 validity",
194            });
195        }
196        Ok(Self { id, values, validity })
197    }
198
199    /// Row count.
200    #[must_use]
201    pub fn len(&self) -> usize {
202        self.values.len()
203    }
204
205    /// Whether empty.
206    #[must_use]
207    pub fn is_empty(&self) -> bool {
208        self.values.is_empty()
209    }
210
211    /// Borrowed contiguous view (no allocation).
212    #[must_use]
213    pub fn as_f64_view(&self) -> F64VectorView<'_> {
214        F64VectorView::contiguous(self.values.as_slice())
215    }
216}
217
218/// Owned int64 column.
219#[derive(Clone, Debug, Eq, PartialEq)]
220pub struct Int64Column {
221    /// Variable id.
222    pub id: VariableId,
223    /// Values.
224    pub values: Arc<[i64]>,
225    /// Validity.
226    pub validity: ValidityBitmap,
227}
228
229impl Int64Column {
230    /// Construct with matching lengths.
231    ///
232    /// # Errors
233    ///
234    /// Length mismatch.
235    pub fn new(
236        id: VariableId,
237        values: impl Into<Arc<[i64]>>,
238        validity: ValidityBitmap,
239    ) -> Result<Self, DataError> {
240        let values = values.into();
241        if validity.len() != values.len() {
242            return Err(DataError::LengthMismatch {
243                expected: values.len(),
244                actual: validity.len(),
245                context: "int64 validity",
246            });
247        }
248        Ok(Self { id, values, validity })
249    }
250
251    /// Row count.
252    #[must_use]
253    pub fn len(&self) -> usize {
254        self.values.len()
255    }
256
257    /// Whether empty.
258    #[must_use]
259    pub fn is_empty(&self) -> bool {
260        self.values.is_empty()
261    }
262}
263
264/// Owned boolean column (bytes: 0/1 per row for simplicity).
265#[derive(Clone, Debug, Eq, PartialEq)]
266pub struct BooleanColumn {
267    /// Variable id.
268    pub id: VariableId,
269    /// Values as 0/1 bytes.
270    pub values: Arc<[u8]>,
271    /// Validity.
272    pub validity: ValidityBitmap,
273}
274
275impl BooleanColumn {
276    /// Construct with matching lengths.
277    ///
278    /// # Errors
279    ///
280    /// Length mismatch.
281    pub fn new(
282        id: VariableId,
283        values: impl Into<Arc<[u8]>>,
284        validity: ValidityBitmap,
285    ) -> Result<Self, DataError> {
286        let values = values.into();
287        if validity.len() != values.len() {
288            return Err(DataError::LengthMismatch {
289                expected: values.len(),
290                actual: validity.len(),
291                context: "bool validity",
292            });
293        }
294        Ok(Self { id, values, validity })
295    }
296
297    /// Row count.
298    #[must_use]
299    pub fn len(&self) -> usize {
300        self.values.len()
301    }
302
303    /// Whether empty.
304    #[must_use]
305    pub fn is_empty(&self) -> bool {
306        self.values.is_empty()
307    }
308}
309
310/// Owned timestamp column (nanoseconds since epoch; timezone metadata lives in schema).
311#[derive(Clone, Debug, Eq, PartialEq)]
312pub struct TimestampColumn {
313    /// Variable id.
314    pub id: VariableId,
315    /// Values in nanoseconds.
316    pub values_ns: Arc<[i64]>,
317    /// Validity.
318    pub validity: ValidityBitmap,
319}
320
321impl TimestampColumn {
322    /// Construct with matching lengths.
323    ///
324    /// # Errors
325    ///
326    /// Length mismatch.
327    pub fn new(
328        id: VariableId,
329        values_ns: impl Into<Arc<[i64]>>,
330        validity: ValidityBitmap,
331    ) -> Result<Self, DataError> {
332        let values_ns = values_ns.into();
333        if validity.len() != values_ns.len() {
334            return Err(DataError::LengthMismatch {
335                expected: values_ns.len(),
336                actual: validity.len(),
337                context: "timestamp validity",
338            });
339        }
340        Ok(Self { id, values_ns, validity })
341    }
342
343    /// Row count.
344    #[must_use]
345    pub fn len(&self) -> usize {
346        self.values_ns.len()
347    }
348
349    /// Whether empty.
350    #[must_use]
351    pub fn is_empty(&self) -> bool {
352        self.values_ns.is_empty()
353    }
354}
355
356/// Owned fixed-size vector column (row-major: `values[row * dim + component]`).
357#[derive(Clone, Debug, PartialEq)]
358pub struct FixedVectorColumn {
359    /// Variable id.
360    pub id: VariableId,
361    /// Vector dimensionality.
362    pub dim: usize,
363    /// Flattened values.
364    pub values: Arc<[f64]>,
365    /// Per-row validity.
366    pub validity: ValidityBitmap,
367}
368
369impl FixedVectorColumn {
370    /// Construct a fixed-vector column.
371    ///
372    /// # Errors
373    ///
374    /// Length / shape mismatch.
375    pub fn new(
376        id: VariableId,
377        dim: usize,
378        values: impl Into<Arc<[f64]>>,
379        validity: ValidityBitmap,
380    ) -> Result<Self, DataError> {
381        if dim == 0 {
382            return Err(DataError::InvalidValidity { message: "fixed vector dim must be > 0" });
383        }
384        let values = values.into();
385        let expected = validity
386            .len()
387            .checked_mul(dim)
388            .ok_or(DataError::InvalidValidity { message: "fixed vector shape overflow" })?;
389        if values.len() != expected {
390            return Err(DataError::LengthMismatch {
391                expected,
392                actual: values.len(),
393                context: "fixed vector values",
394            });
395        }
396        Ok(Self { id, dim, values, validity })
397    }
398
399    /// Row count.
400    #[must_use]
401    pub fn len(&self) -> usize {
402        self.validity.len()
403    }
404
405    /// Whether empty.
406    #[must_use]
407    pub fn is_empty(&self) -> bool {
408        self.len() == 0
409    }
410}
411
412/// Borrowed typed column view (library-owned; not Arrow types).
413#[derive(Clone, Copy, Debug)]
414pub enum ColumnView<'a> {
415    /// Float64 column.
416    Float64(&'a Float64Column),
417    /// Int64 column.
418    Int64(&'a Int64Column),
419    /// Boolean column.
420    Boolean(&'a BooleanColumn),
421    /// Dictionary categorical column.
422    Categorical(&'a CategoricalColumn),
423    /// Timestamp column.
424    Timestamp(&'a TimestampColumn),
425    /// Fixed-size vector column.
426    FixedVector(&'a FixedVectorColumn),
427}
428
429impl<'a> ColumnView<'a> {
430    /// Variable id.
431    #[must_use]
432    pub fn id(self) -> VariableId {
433        match self {
434            Self::Float64(c) => c.id,
435            Self::Int64(c) => c.id,
436            Self::Boolean(c) => c.id,
437            Self::Categorical(c) => c.id,
438            Self::Timestamp(c) => c.id,
439            Self::FixedVector(c) => c.id,
440        }
441    }
442
443    /// Row count.
444    #[must_use]
445    pub fn len(self) -> usize {
446        match self {
447            Self::Float64(c) => c.len(),
448            Self::Int64(c) => c.len(),
449            Self::Boolean(c) => c.len(),
450            Self::Categorical(c) => c.len(),
451            Self::Timestamp(c) => c.len(),
452            Self::FixedVector(c) => c.len(),
453        }
454    }
455
456    /// Whether empty.
457    #[must_use]
458    pub fn is_empty(self) -> bool {
459        self.len() == 0
460    }
461
462    /// Borrow the column validity bitmap.
463    #[must_use]
464    pub fn validity(self) -> &'a ValidityBitmap {
465        match self {
466            Self::Float64(c) => &c.validity,
467            Self::Int64(c) => &c.validity,
468            Self::Boolean(c) => &c.validity,
469            Self::Categorical(c) => &c.validity,
470            Self::Timestamp(c) => &c.validity,
471            Self::FixedVector(c) => &c.validity,
472        }
473    }
474}
475
476/// Owned column enum stored in a table.
477#[derive(Clone, Debug)]
478pub enum OwnedColumn {
479    /// Float64.
480    Float64(Float64Column),
481    /// Int64.
482    Int64(Int64Column),
483    /// Boolean.
484    Boolean(BooleanColumn),
485    /// Categorical.
486    Categorical(CategoricalColumn),
487    /// Timestamp.
488    Timestamp(TimestampColumn),
489    /// Fixed-size vector.
490    FixedVector(FixedVectorColumn),
491}
492
493impl OwnedColumn {
494    /// Variable id.
495    #[must_use]
496    pub fn id(&self) -> VariableId {
497        match self {
498            Self::Float64(c) => c.id,
499            Self::Int64(c) => c.id,
500            Self::Boolean(c) => c.id,
501            Self::Categorical(c) => c.id,
502            Self::Timestamp(c) => c.id,
503            Self::FixedVector(c) => c.id,
504        }
505    }
506
507    /// Clone the column with a remapped dense id (value buffers stay shared).
508    #[must_use]
509    pub fn with_id(&self, id: VariableId) -> Self {
510        match self {
511            Self::Float64(c) => Self::Float64(Float64Column {
512                id,
513                values: c.values.clone(),
514                validity: c.validity.clone(),
515            }),
516            Self::Int64(c) => Self::Int64(Int64Column {
517                id,
518                values: Arc::clone(&c.values),
519                validity: c.validity.clone(),
520            }),
521            Self::Boolean(c) => Self::Boolean(BooleanColumn {
522                id,
523                values: Arc::clone(&c.values),
524                validity: c.validity.clone(),
525            }),
526            Self::Categorical(c) => Self::Categorical(CategoricalColumn {
527                id,
528                codes: Arc::clone(&c.codes),
529                validity: c.validity.clone(),
530                domain: Arc::clone(&c.domain),
531            }),
532            Self::Timestamp(c) => Self::Timestamp(TimestampColumn {
533                id,
534                values_ns: Arc::clone(&c.values_ns),
535                validity: c.validity.clone(),
536            }),
537            Self::FixedVector(c) => Self::FixedVector(FixedVectorColumn {
538                id,
539                values: Arc::clone(&c.values),
540                dim: c.dim,
541                validity: c.validity.clone(),
542            }),
543        }
544    }
545
546    /// Row count.
547    #[must_use]
548    pub fn len(&self) -> usize {
549        match self {
550            Self::Float64(c) => c.len(),
551            Self::Int64(c) => c.len(),
552            Self::Boolean(c) => c.len(),
553            Self::Categorical(c) => c.len(),
554            Self::Timestamp(c) => c.len(),
555            Self::FixedVector(c) => c.len(),
556        }
557    }
558
559    /// Whether empty.
560    #[must_use]
561    pub fn is_empty(&self) -> bool {
562        self.len() == 0
563    }
564
565    /// Borrow as a [`ColumnView`].
566    #[must_use]
567    pub fn as_view(&self) -> ColumnView<'_> {
568        match self {
569            Self::Float64(c) => ColumnView::Float64(c),
570            Self::Int64(c) => ColumnView::Int64(c),
571            Self::Boolean(c) => ColumnView::Boolean(c),
572            Self::Categorical(c) => ColumnView::Categorical(c),
573            Self::Timestamp(c) => ColumnView::Timestamp(c),
574            Self::FixedVector(c) => ColumnView::FixedVector(c),
575        }
576    }
577}
antecedent_data/column.rs

antecedent_data/
column.rs
