Skip to main content

pond_rs/
column.rs

1//! Value columns: numbers, exact decimals, text, true/false values and
2//! arrays.
3//!
4//! In a number column a gap is stored as `NaN`, as in pond-ts
5//! ([PND-NANREP]): any non-finite value in the buffer means "missing".
6//! Decimal, text and true/false columns mark gaps with an Arrow null; an
7//! array column marks them with `None`.
8//!
9//! The values live in an Arrow buffer (shared, immutable, reference
10//! counted). So a column read from Arrow can point at Arrow's memory instead
11//! of copying it, a column written to Arrow is handed over without a copy,
12//! and cloning a column or a series does not copy values.
13
14use std::sync::Arc;
15
16use arrow_array::builder::StringBuilder;
17use arrow_array::cast::AsArray;
18use arrow_array::types::{Decimal128Type, Int32Type};
19use arrow_array::{
20    Array, ArrayRef, BooleanArray, Decimal128Array, DictionaryArray, StringArray, UInt32Array,
21};
22use arrow_buffer::{NullBuffer, ScalarBuffer};
23use arrow_schema::DataType;
24
25use crate::series::ColumnKind;
26
27/// A column of `f64` values where any non-finite value (`NaN`, `±inf`) is a
28/// gap. Reading a gap gives `None`.
29#[derive(Debug, Clone)]
30pub struct Float64Column {
31    values: ScalarBuffer<f64>,
32    /// True when every value is finite (no gaps). Only a speed hint: the
33    /// results are the same either way.
34    all_finite: bool,
35}
36
37impl Float64Column {
38    /// Build from raw values; non-finite values become gaps.
39    pub fn from_values(values: Vec<f64>) -> Self {
40        Self::from_buffer(values.into())
41    }
42
43    /// Wrap an Arrow buffer as it is, without copying: non-finite values are
44    /// gaps.
45    pub(crate) fn from_buffer(values: ScalarBuffer<f64>) -> Self {
46        let all_finite = values.iter().all(|v| v.is_finite());
47        Self { values, all_finite }
48    }
49
50    /// The Arrow buffer holding the values (cheap to clone: shared).
51    pub(crate) fn buffer(&self) -> &ScalarBuffer<f64> {
52        &self.values
53    }
54
55    /// Build from optional values; `None` is a gap.
56    pub fn from_options(values: impl IntoIterator<Item = Option<f64>>) -> Self {
57        Self::from_values(values.into_iter().map(|v| v.unwrap_or(f64::NAN)).collect())
58    }
59
60    /// Number of values.
61    pub fn len(&self) -> usize {
62        self.values.len()
63    }
64
65    /// `true` when there are no values.
66    pub fn is_empty(&self) -> bool {
67        self.values.is_empty()
68    }
69
70    /// The value at `i`, or `None` for a gap. Panics if `i` is out of range.
71    pub fn get(&self, i: usize) -> Option<f64> {
72        let v = self.values[i];
73        v.is_finite().then_some(v)
74    }
75
76    /// The raw buffer: gaps are non-finite values.
77    pub fn values(&self) -> &[f64] {
78        &self.values
79    }
80
81    /// True when the column has no gaps.
82    pub fn all_finite(&self) -> bool {
83        self.all_finite
84    }
85
86    /// Number of gaps.
87    pub fn missing_count(&self) -> usize {
88        if self.all_finite {
89            0
90        } else {
91            self.values.iter().filter(|v| !v.is_finite()).count()
92        }
93    }
94
95    /// Iterate as `Option<f64>`.
96    pub fn iter(&self) -> impl Iterator<Item = Option<f64>> + '_ {
97        self.values.iter().map(|&v| v.is_finite().then_some(v))
98    }
99}
100
101/// Two number columns are equal when they hold the same values, bit for bit
102/// (so `-0.0` and `0.0` differ), and gaps in the same rows, however each gap
103/// is stored (`NaN` or an infinity).
104impl PartialEq for Float64Column {
105    fn eq(&self, other: &Self) -> bool {
106        let bits = |c: &Self| c.iter().map(|v| v.map(f64::to_bits)).collect::<Vec<_>>();
107        self.len() == other.len() && bits(self) == bits(other)
108    }
109}
110
111/// A column of text values (pond-ts `kind: 'string'`). A gap is an Arrow
112/// null; an empty string is a value, not a gap.
113///
114/// Held as an Arrow `Utf8` array, or as a dictionary (`Dictionary<Int32,
115/// Utf8>`: each row an index into a list of distinct strings) when it was
116/// read that way from Arrow. Reading gives the same values either way;
117/// [`StringColumn::is_dictionary`] says which, and `to_arrow` writes it back
118/// the same way. Two columns are equal when they hold the same values,
119/// however each is held.
120#[derive(Debug, Clone)]
121pub struct StringColumn {
122    repr: StringRepr,
123}
124
125#[derive(Debug, Clone)]
126enum StringRepr {
127    Plain(StringArray),
128    Dictionary(DictionaryArray<Int32Type>),
129}
130
131impl StringColumn {
132    /// Build from optional values; `None` is a gap.
133    pub fn from_options<S: AsRef<str>>(values: impl IntoIterator<Item = Option<S>>) -> Self {
134        Self::plain(values.into_iter().collect::<StringArray>())
135    }
136
137    /// Wrap an Arrow `Utf8` array without copying.
138    pub(crate) fn plain(array: StringArray) -> Self {
139        Self {
140            repr: StringRepr::Plain(array),
141        }
142    }
143
144    /// Wrap an Arrow dictionary array without copying. Its values must be
145    /// `Utf8` with no nulls.
146    pub(crate) fn dictionary(array: DictionaryArray<Int32Type>) -> Option<Self> {
147        let values = array.values();
148        (values.data_type() == &DataType::Utf8 && values.null_count() == 0).then_some(Self {
149            repr: StringRepr::Dictionary(array),
150        })
151    }
152
153    /// Number of values.
154    pub fn len(&self) -> usize {
155        match &self.repr {
156            StringRepr::Plain(a) => a.len(),
157            StringRepr::Dictionary(a) => a.len(),
158        }
159    }
160
161    /// `true` when there are no values.
162    pub fn is_empty(&self) -> bool {
163        self.len() == 0
164    }
165
166    /// The value at `i`, or `None` for a gap. Panics if `i` is out of range.
167    pub fn get(&self, i: usize) -> Option<&str> {
168        assert!(i < self.len(), "row {i} out of range ({} rows)", self.len());
169        match &self.repr {
170            StringRepr::Plain(a) => a.is_valid(i).then(|| a.value(i)),
171            StringRepr::Dictionary(a) => a.is_valid(i).then(|| {
172                let key = a.keys().value(i) as usize;
173                a.values().as_string::<i32>().value(key)
174            }),
175        }
176    }
177
178    /// Iterate as `Option<&str>`.
179    pub fn iter(&self) -> impl Iterator<Item = Option<&str>> + '_ {
180        (0..self.len()).map(|i| self.get(i))
181    }
182
183    /// Number of gaps.
184    pub fn missing_count(&self) -> usize {
185        self.arrow().logical_null_count()
186    }
187
188    /// True when held as a dictionary (see the type's notes).
189    pub fn is_dictionary(&self) -> bool {
190        matches!(self.repr, StringRepr::Dictionary(_))
191    }
192
193    /// The Arrow array behind the column (shared, not copied).
194    pub(crate) fn arrow(&self) -> ArrayRef {
195        match &self.repr {
196            StringRepr::Plain(a) => Arc::new(a.clone()),
197            StringRepr::Dictionary(a) => Arc::new(a.clone()),
198        }
199    }
200
201    /// The rows named by `rows`, in that order; `None` gives a gap. A
202    /// dictionary column stays one, sharing its whole list of strings, so a
203    /// short result (an `aggregate` of a long column) still carries, and
204    /// writes to Arrow, every distinct string of the source.
205    pub(crate) fn take(&self, rows: &[Option<usize>]) -> Self {
206        self.take_indices(&indices_of(rows))
207    }
208}
209
210impl StringColumn {
211    /// The rows at `indices` (a null index gives a gap), with Arrow's
212    /// gather kernel: offsets first, then one copy per row, about three
213    /// times faster than appending to a builder. A dictionary column
214    /// gathers its keys and keeps its values.
215    pub(crate) fn take_indices(&self, indices: &UInt32Array) -> Self {
216        let array: &dyn Array = match &self.repr {
217            StringRepr::Plain(a) => a,
218            StringRepr::Dictionary(a) => a,
219        };
220        let out = arrow_select::take::take(array, indices, None).expect("row numbers are in range");
221        match &self.repr {
222            StringRepr::Plain(_) => Self::plain(out.as_string::<i32>().clone()),
223            StringRepr::Dictionary(_) => Self {
224                repr: StringRepr::Dictionary(out.as_dictionary::<Int32Type>().clone()),
225            },
226        }
227    }
228}
229
230/// `rows` as Arrow row numbers, `None` as a null.
231pub(crate) fn indices_of(rows: &[Option<usize>]) -> UInt32Array {
232    rows.iter()
233        .map(|r| r.map(|i| u32::try_from(i).expect("a text column has under 2^32 rows")))
234        .collect()
235}
236
237impl PartialEq for StringColumn {
238    fn eq(&self, other: &Self) -> bool {
239        self.len() == other.len() && self.iter().eq(other.iter())
240    }
241}
242
243/// A column of true/false values (pond-ts `kind: 'boolean'`). A gap is an
244/// Arrow null. Held as an Arrow `Boolean` array (one bit per row).
245#[derive(Debug, Clone, PartialEq)]
246pub struct BooleanColumn {
247    values: BooleanArray,
248}
249
250impl BooleanColumn {
251    /// Build from optional values; `None` is a gap.
252    pub fn from_options(values: impl IntoIterator<Item = Option<bool>>) -> Self {
253        Self::from_arrow(values.into_iter().collect())
254    }
255
256    /// Wrap an Arrow array without copying.
257    pub(crate) fn from_arrow(values: BooleanArray) -> Self {
258        Self { values }
259    }
260
261    /// Number of values.
262    pub fn len(&self) -> usize {
263        self.values.len()
264    }
265
266    /// `true` when there are no values.
267    pub fn is_empty(&self) -> bool {
268        self.values.is_empty()
269    }
270
271    /// The value at `i`, or `None` for a gap. Panics if `i` is out of range.
272    pub fn get(&self, i: usize) -> Option<bool> {
273        assert!(i < self.len(), "row {i} out of range ({} rows)", self.len());
274        self.values.is_valid(i).then(|| self.values.value(i))
275    }
276
277    /// Iterate as `Option<bool>`.
278    pub fn iter(&self) -> impl Iterator<Item = Option<bool>> + '_ {
279        self.values.iter()
280    }
281
282    /// Number of gaps.
283    pub fn missing_count(&self) -> usize {
284        self.values.null_count()
285    }
286
287    /// The Arrow array behind the column (shared, not copied).
288    pub(crate) fn arrow(&self) -> ArrayRef {
289        Arc::new(self.values.clone())
290    }
291
292    /// The rows named by `rows`, in that order; `None` gives a gap.
293    pub(crate) fn take(&self, rows: &[Option<usize>]) -> Self {
294        Self::from_options(rows.iter().map(|r| r.and_then(|i| self.get(i))))
295    }
296}
297
298/// One value of any column kind: a [`Fill::Value`](crate::Fill::Value)
299/// literal, or one part of a partition key
300/// ([`PartitionedSeries`](crate::PartitionedSeries)), where `None` stands
301/// for a gap.
302#[derive(Debug, Clone, PartialEq)]
303pub enum Value {
304    /// A number.
305    Number(f64),
306    /// As a fill value, must have the column's scale.
307    Decimal(Decimal),
308    /// Text.
309    String(String),
310    /// A true/false value.
311    Boolean(bool),
312    /// The elements of one array cell (see [`ArrayColumn`]).
313    Array(Vec<Scalar>),
314}
315
316impl From<f64> for Value {
317    fn from(v: f64) -> Self {
318        Value::Number(v)
319    }
320}
321
322impl From<&str> for Value {
323    fn from(v: &str) -> Self {
324        Value::String(v.to_string())
325    }
326}
327
328impl From<String> for Value {
329    fn from(v: String) -> Self {
330        Value::String(v)
331    }
332}
333
334impl From<bool> for Value {
335    fn from(v: bool) -> Self {
336        Value::Boolean(v)
337    }
338}
339
340impl From<Vec<Scalar>> for Value {
341    fn from(v: Vec<Scalar>) -> Self {
342        Value::Array(v)
343    }
344}
345
346impl From<Decimal> for Value {
347    fn from(v: Decimal) -> Self {
348        Value::Decimal(v)
349    }
350}
351
352impl Value {
353    /// The number, if this is a [`Value::Number`].
354    pub fn as_f64(&self) -> Option<f64> {
355        match self {
356            Value::Number(v) => Some(*v),
357            _ => None,
358        }
359    }
360
361    /// The decimal, if this is a [`Value::Decimal`].
362    pub fn as_decimal(&self) -> Option<Decimal> {
363        match self {
364            Value::Decimal(v) => Some(*v),
365            _ => None,
366        }
367    }
368
369    /// The text, if this is a [`Value::String`]. For a partition key part
370    /// (`Option<Value>`), `part.as_ref().and_then(Value::as_str)`:
371    ///
372    /// ```
373    /// # use pond_rs::{Column, TimeSeries, Value};
374    /// let s = TimeSeries::new("s", vec![0, 1], vec![
375    ///     Column::required("symbol", vec!["AAPL", "MSFT"]),
376    /// ])?;
377    /// let parts = s.partition_by(&["symbol"])?;
378    /// let names: Vec<&str> = parts
379    ///     .iter()
380    ///     .filter_map(|(key, _)| key[0].as_ref().and_then(Value::as_str))
381    ///     .collect();
382    /// assert_eq!(names, ["AAPL", "MSFT"]);
383    /// # Ok::<(), pond_rs::PondError>(())
384    /// ```
385    pub fn as_str(&self) -> Option<&str> {
386        match self {
387            Value::String(v) => Some(v),
388            _ => None,
389        }
390    }
391
392    /// The true/false value, if this is a [`Value::Boolean`].
393    pub fn as_bool(&self) -> Option<bool> {
394        match self {
395            Value::Boolean(v) => Some(*v),
396            _ => None,
397        }
398    }
399
400    /// The elements, if this is a [`Value::Array`].
401    pub fn as_array(&self) -> Option<&[Scalar]> {
402        match self {
403            Value::Array(v) => Some(v),
404            _ => None,
405        }
406    }
407}
408
409/// One exact decimal: `value × 10^-scale`, so `Decimal { value: 125, scale:
410/// 2 }` is `1.25`. Prints as a decimal (`1.25`). Equality compares `value`
411/// and `scale` as stored, so `1.25` at scale 2 is not equal to `1.250` at
412/// scale 3 (a column has one scale, so within a column this never arises).
413#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
414pub struct Decimal {
415    /// The whole number stored (Arrow's raw `i128`).
416    pub value: i128,
417    /// Digits after the decimal point (negative: zeros before it).
418    pub scale: i8,
419}
420
421impl Decimal {
422    /// The nearest `f64` (exact when the decimal is one; correctly rounded
423    /// otherwise).
424    pub fn to_f64(self) -> f64 {
425        // Both parts are exact and one division rounds once.
426        const EXACT: u128 = 1 << 53;
427        if (-22..=22).contains(&self.scale) && self.value.unsigned_abs() <= EXACT {
428            let power = 10f64.powi(i32::from(self.scale.abs()));
429            return if self.scale >= 0 {
430                self.value as f64 / power
431            } else {
432                self.value as f64 * power
433            };
434        }
435        // Rust's float parsing rounds correctly.
436        self.to_string().parse().expect("a decimal parses as f64")
437    }
438}
439
440impl std::fmt::Display for Decimal {
441    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
442        let digits = self.value.unsigned_abs().to_string();
443        let sign = if self.value < 0 { "-" } else { "" };
444        let scale = i32::from(self.scale);
445        if scale <= 0 {
446            let zeros = if self.value == 0 {
447                0
448            } else {
449                (-scale) as usize
450            };
451            return write!(f, "{sign}{digits}{}", "0".repeat(zeros));
452        }
453        let scale = scale as usize;
454        let padded = format!("{digits:0>width$}", width = scale + 1);
455        let (whole, frac) = padded.split_at(padded.len() - scale);
456        write!(f, "{sign}{whole}.{frac}")
457    }
458}
459
460/// A column of exact decimals: Arrow `Decimal128(precision, scale)`. Each
461/// value is a whole number (`i128`) and the column has one scale, so `1.25`
462/// at scale 2 is stored as `125`. A gap is an Arrow null.
463///
464/// pond-ts has no decimal kind. pond-rs keeps decimals exact through
465/// `count`, `first`, `last`, `min`, `max`, `keep`, `sum` and `difference`;
466/// every other reducer works on the nearest `f64` and gives a number
467/// column (or, for `unique`, `top<N>` and `samples`, a list of numbers).
468#[derive(Debug, Clone, PartialEq)]
469pub struct DecimalColumn {
470    values: Decimal128Array,
471}
472
473impl DecimalColumn {
474    /// Build from raw whole numbers at `scale` (`None` is a gap), with the
475    /// widest precision, 38 digits. An error if a value has more than 38
476    /// digits or `scale` is outside `-38..=38`.
477    pub fn from_options(
478        values: impl IntoIterator<Item = Option<i128>>,
479        scale: i8,
480    ) -> crate::error::Result<Self> {
481        let array: Decimal128Array = values.into_iter().collect();
482        Self::with_type(array, 38, scale)
483    }
484
485    /// [`DecimalColumn::from_options`] with no gaps.
486    pub fn from_values(values: Vec<i128>, scale: i8) -> crate::error::Result<Self> {
487        Self::from_options(values.into_iter().map(Some), scale)
488    }
489
490    pub(crate) fn with_type(
491        array: Decimal128Array,
492        precision: u8,
493        scale: i8,
494    ) -> crate::error::Result<Self> {
495        let invalid = |e: arrow_schema::ArrowError| crate::error::PondError::InvalidData {
496            call: "DecimalColumn",
497            message: e.to_string(),
498        };
499        let array = array
500            .with_precision_and_scale(precision, scale)
501            .map_err(invalid)?;
502        // Not arrow's validate_decimal_precision: it casts the scale to u8,
503        // so it refuses every negative scale.
504        let max = 10u128.pow(u32::from(precision)) - 1;
505        if let Some(i) =
506            (0..array.len()).find(|&i| array.is_valid(i) && array.value(i).unsigned_abs() > max)
507        {
508            return Err(crate::error::PondError::InvalidData {
509                call: "DecimalColumn",
510                message: format!(
511                    "row {i}: {} has more than {precision} digits (the precision)",
512                    Decimal {
513                        value: array.value(i),
514                        scale
515                    }
516                ),
517            });
518        }
519        Ok(Self { values: array })
520    }
521
522    /// Number of values.
523    pub fn len(&self) -> usize {
524        self.values.len()
525    }
526
527    /// `true` when there are no values.
528    pub fn is_empty(&self) -> bool {
529        self.values.is_empty()
530    }
531
532    /// Digits after the decimal point, the same for every row.
533    pub fn scale(&self) -> i8 {
534        self.values.scale()
535    }
536
537    /// The most digits a value may have (Arrow's precision, at most 38).
538    pub fn precision(&self) -> u8 {
539        self.values.precision()
540    }
541
542    /// The value at `i`, or `None` for a gap. Panics if `i` is out of range.
543    pub fn get(&self, i: usize) -> Option<Decimal> {
544        assert!(i < self.len(), "row {i} out of range ({} rows)", self.len());
545        self.values.is_valid(i).then(|| Decimal {
546            value: self.values.value(i),
547            scale: self.scale(),
548        })
549    }
550
551    /// Iterate as `Option<Decimal>`.
552    pub fn iter(&self) -> impl Iterator<Item = Option<Decimal>> + '_ {
553        (0..self.len()).map(|i| self.get(i))
554    }
555
556    /// The raw whole numbers (gaps hold an unspecified value; see
557    /// [`DecimalColumn::get`]).
558    pub(crate) fn raw(&self) -> &[i128] {
559        self.values.values()
560    }
561
562    /// Number of gaps.
563    pub fn missing_count(&self) -> usize {
564        self.values.null_count()
565    }
566
567    /// The Arrow array behind the column (shared, not copied).
568    pub(crate) fn arrow(&self) -> ArrayRef {
569        Arc::new(self.values.clone())
570    }
571
572    pub(crate) fn nulls(&self) -> Option<&NullBuffer> {
573        self.values.nulls()
574    }
575
576    /// Each row as the nearest `f64`, `NaN` for a gap.
577    pub(crate) fn to_f64s(&self) -> Vec<f64> {
578        let scale = self.scale();
579        // Decimal::to_f64's fast path with the divisor worked out once.
580        let divisor = (0..=22)
581            .contains(&scale)
582            .then(|| 10f64.powi(i32::from(scale)));
583        const EXACT: i128 = 1 << 53;
584        let raw = self.raw();
585        (0..self.len())
586            .map(|i| {
587                let value = raw[i];
588                match divisor {
589                    _ if self.values.is_null(i) => f64::NAN,
590                    Some(d) if (-EXACT..=EXACT).contains(&value) => value as i64 as f64 / d,
591                    _ => Decimal { value, scale }.to_f64(),
592                }
593            })
594            .collect()
595    }
596
597    /// Raw values at this column's scale and the given precision.
598    pub(crate) fn rebuild(
599        &self,
600        values: impl IntoIterator<Item = Option<i128>>,
601        precision: u8,
602    ) -> Self {
603        let array: Decimal128Array = values.into_iter().collect();
604        Self {
605            values: array
606                .with_precision_and_scale(precision, self.scale())
607                .expect("this column's own scale"),
608        }
609    }
610
611    /// The rows named by `rows`, in that order; `None` gives a gap.
612    pub(crate) fn take(&self, rows: &[Option<usize>]) -> Self {
613        self.rebuild(
614            rows.iter()
615                .map(|r| r.and_then(|i| self.values.is_valid(i).then(|| self.values.value(i)))),
616            self.precision(),
617        )
618    }
619}
620
621/// One element of an array cell: a number, text or a true/false value, as
622/// pond-ts's `ScalarValue`. Numbers must be finite (checked when a series
623/// is built).
624///
625/// `==` compares as JavaScript's `includes` does on finite elements
626/// (SameValueZero): `0.0` equals `-0.0`, and values of different kinds never
627/// match (`1` is not `"1"`).
628#[derive(Debug, Clone, PartialEq)]
629pub enum Scalar {
630    /// A number.
631    Number(f64),
632    /// Text.
633    String(String),
634    /// A true/false value.
635    Boolean(bool),
636}
637
638impl Scalar {
639    /// The kind of column this element would sit in: number, string or
640    /// boolean.
641    pub fn kind(&self) -> ColumnKind {
642        match self {
643            Scalar::Number(_) => ColumnKind::Number,
644            Scalar::String(_) => ColumnKind::String,
645            Scalar::Boolean(_) => ColumnKind::Boolean,
646        }
647    }
648
649    /// The number, if this is a [`Scalar::Number`].
650    pub fn as_f64(&self) -> Option<f64> {
651        match self {
652            Scalar::Number(v) => Some(*v),
653            _ => None,
654        }
655    }
656
657    /// The text, if this is a [`Scalar::String`].
658    pub fn as_str(&self) -> Option<&str> {
659        match self {
660            Scalar::String(v) => Some(v),
661            _ => None,
662        }
663    }
664
665    /// The true/false value, if this is a [`Scalar::Boolean`].
666    pub fn as_bool(&self) -> Option<bool> {
667        match self {
668            Scalar::Boolean(v) => Some(*v),
669            _ => None,
670        }
671    }
672
673    /// The same element as a [`Value`].
674    pub fn to_value(&self) -> Value {
675        match self {
676            Scalar::Number(v) => Value::Number(*v),
677            Scalar::String(v) => Value::String(v.clone()),
678            Scalar::Boolean(v) => Value::Boolean(*v),
679        }
680    }
681
682    /// Equal bit for bit: as `==`, except `0.0` and `-0.0` differ.
683    pub(crate) fn same(&self, other: &Scalar) -> bool {
684        match (self, other) {
685            (Scalar::Number(a), Scalar::Number(b)) => a.to_bits() == b.to_bits(),
686            _ => self == other,
687        }
688    }
689}
690
691impl From<f64> for Scalar {
692    fn from(v: f64) -> Self {
693        Scalar::Number(v)
694    }
695}
696
697impl From<i32> for Scalar {
698    fn from(v: i32) -> Self {
699        Scalar::Number(f64::from(v))
700    }
701}
702
703impl From<&str> for Scalar {
704    fn from(v: &str) -> Self {
705        Scalar::String(v.to_string())
706    }
707}
708
709impl From<String> for Scalar {
710    fn from(v: String) -> Self {
711        Scalar::String(v)
712    }
713}
714
715impl From<bool> for Scalar {
716    fn from(v: bool) -> Self {
717        Scalar::Boolean(v)
718    }
719}
720
721/// One array cell: its elements, shared (cloning a cell or a column does
722/// not copy them).
723pub(crate) type Cell = Arc<[Scalar]>;
724
725/// A column of arrays (pond-ts `kind: 'array'`): each row holds a list of
726/// numbers, text and true/false values, mixed freely, or is a gap. An empty
727/// list is a value, not a gap. Numbers must be finite.
728///
729/// ```
730/// # use pond_rs::{ArrayColumn, Column, Scalar, TimeSeries};
731/// let tags = ArrayColumn::from_options([
732///     Some(vec!["a".into(), "b".into()]),
733///     Some(vec![]),
734///     None,
735///     Some(vec![Scalar::from(1.0), true.into()]),
736/// ]);
737/// let s = TimeSeries::new("s", vec![0, 1, 2, 3], vec![Column::optional("tags", tags)])?;
738/// let tags = s.column("tags")?.arrays()?;
739/// assert_eq!(tags.get(0), Some(&[Scalar::from("a"), "b".into()][..]));
740/// assert_eq!(tags.get(1), Some(&[][..])); // empty, not a gap
741/// assert_eq!(tags.get(2), None);
742/// # Ok::<(), pond_rs::PondError>(())
743/// ```
744///
745/// In Arrow an array column is a list: [`TimeSeries::to_arrow`] writes one
746/// when every element in the column is the same kind, and
747/// [`TimeSeries::from_arrow`] reads lists of numbers, text or true/false
748/// values. pond-ts has no Arrow form for arrays.
749///
750/// [`TimeSeries::to_arrow`]: crate::TimeSeries::to_arrow
751/// [`TimeSeries::from_arrow`]: crate::TimeSeries::from_arrow
752#[derive(Debug, Clone)]
753pub struct ArrayColumn {
754    cells: Arc<[Option<Cell>]>,
755    offset: usize,
756    len: usize,
757}
758
759impl ArrayColumn {
760    /// Build from optional lists; `None` is a gap.
761    pub fn from_options(values: impl IntoIterator<Item = Option<Vec<Scalar>>>) -> Self {
762        Self::from_cells(values.into_iter().map(|v| v.map(Cell::from)).collect())
763    }
764
765    pub(crate) fn from_cells(cells: Vec<Option<Cell>>) -> Self {
766        let len = cells.len();
767        Self {
768            cells: cells.into(),
769            offset: 0,
770            len,
771        }
772    }
773
774    /// Number of values.
775    pub fn len(&self) -> usize {
776        self.len
777    }
778
779    /// `true` when there are no values.
780    pub fn is_empty(&self) -> bool {
781        self.len == 0
782    }
783
784    fn cells(&self) -> &[Option<Cell>] {
785        &self.cells[self.offset..self.offset + self.len]
786    }
787
788    /// The list at `i`, or `None` for a gap. Panics if `i` is out of range.
789    pub fn get(&self, i: usize) -> Option<&[Scalar]> {
790        self.cells()[i].as_deref()
791    }
792
793    /// The shared cell at `i`.
794    pub(crate) fn cell(&self, i: usize) -> Option<&Cell> {
795        self.cells()[i].as_ref()
796    }
797
798    /// Iterate as `Option<&[Scalar]>`.
799    pub fn iter(&self) -> impl Iterator<Item = Option<&[Scalar]>> + '_ {
800        self.cells().iter().map(|c| c.as_deref())
801    }
802
803    /// Number of gaps.
804    pub fn missing_count(&self) -> usize {
805        self.cells().iter().filter(|c| c.is_none()).count()
806    }
807
808    /// The first element that is not a finite number, string or boolean
809    /// (a non-finite number): `(row, element)`.
810    pub(crate) fn first_bad_element(&self) -> Option<(usize, usize)> {
811        self.iter().enumerate().find_map(|(r, cell)| {
812            cell?
813                .iter()
814                .position(|e| matches!(e, Scalar::Number(v) if !v.is_finite()))
815                .map(|e| (r, e))
816        })
817    }
818
819    /// The rows named by `rows`, in that order; `None` gives a gap. Cells
820    /// are shared, not copied.
821    pub(crate) fn take(&self, rows: &[Option<usize>]) -> Self {
822        Self::from_cells(
823            rows.iter()
824                .map(|r| r.and_then(|i| self.cell(i).cloned()))
825                .collect(),
826        )
827    }
828}
829
830/// Equal when the same rows are gaps and the lists hold the same elements
831/// bit for bit (so `-0.0` and `0.0` differ, as in number columns).
832impl PartialEq for ArrayColumn {
833    fn eq(&self, other: &Self) -> bool {
834        self.len == other.len
835            && self.iter().zip(other.iter()).all(|(a, b)| match (a, b) {
836                (None, None) => true,
837                (Some(a), Some(b)) => a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.same(y)),
838                _ => false,
839            })
840    }
841}
842
843/// What a value column holds: numbers, exact decimals, text, true/false
844/// values, or arrays.
845#[derive(Debug, Clone, PartialEq)]
846pub enum ColumnData {
847    /// Numbers (`f64`, `NaN` marks a gap).
848    Number(Float64Column),
849    /// Exact decimals.
850    Decimal(DecimalColumn),
851    /// Text.
852    String(StringColumn),
853    /// True/false values.
854    Boolean(BooleanColumn),
855    /// Arrays of scalars.
856    Array(ArrayColumn),
857}
858
859impl ColumnData {
860    /// Number of values.
861    pub fn len(&self) -> usize {
862        match self {
863            ColumnData::Number(c) => c.len(),
864            ColumnData::Decimal(c) => c.len(),
865            ColumnData::String(c) => c.len(),
866            ColumnData::Boolean(c) => c.len(),
867            ColumnData::Array(c) => c.len(),
868        }
869    }
870
871    /// `true` when there are no values.
872    pub fn is_empty(&self) -> bool {
873        self.len() == 0
874    }
875
876    /// Number of gaps.
877    pub fn missing_count(&self) -> usize {
878        match self {
879            ColumnData::Number(c) => c.missing_count(),
880            ColumnData::Decimal(c) => c.missing_count(),
881            ColumnData::String(c) => c.missing_count(),
882            ColumnData::Boolean(c) => c.missing_count(),
883            ColumnData::Array(c) => c.missing_count(),
884        }
885    }
886
887    /// The validity bits of a decimal, text, true/false or array column
888    /// (`None` when it has no gaps, and always for numbers, whose gaps are
889    /// non-finite values instead).
890    pub(crate) fn arrow_nulls(&self) -> Option<NullBuffer> {
891        match self {
892            ColumnData::Number(_) => None,
893            ColumnData::Decimal(c) => c.nulls().cloned(),
894            ColumnData::String(c) => c.arrow().logical_nulls(),
895            ColumnData::Boolean(c) => c.values.nulls().cloned(),
896            ColumnData::Array(c) => (c.missing_count() > 0)
897                .then(|| NullBuffer::from(c.iter().map(|v| v.is_some()).collect::<Vec<bool>>())),
898        }
899    }
900
901    /// The rows in `range`, sharing this column's buffers (no copy).
902    pub(crate) fn slice(&self, range: std::ops::Range<usize>) -> ColumnData {
903        let (offset, len) = (range.start, range.len());
904        match self {
905            ColumnData::Number(c) => ColumnData::Number(Float64Column {
906                values: c.values.slice(offset, len),
907                // A slice of a column with no gaps has none either.
908                all_finite: c.all_finite || c.values[range].iter().all(|v| v.is_finite()),
909            }),
910            ColumnData::Decimal(c) => ColumnData::Decimal(DecimalColumn {
911                values: c.values.slice(offset, len),
912            }),
913            ColumnData::String(c) => ColumnData::String(StringColumn {
914                repr: match &c.repr {
915                    StringRepr::Plain(a) => StringRepr::Plain(a.slice(offset, len)),
916                    StringRepr::Dictionary(a) => StringRepr::Dictionary(a.slice(offset, len)),
917                },
918            }),
919            ColumnData::Boolean(c) => ColumnData::Boolean(BooleanColumn {
920                values: c.values.slice(offset, len),
921            }),
922            ColumnData::Array(c) => ColumnData::Array(ArrayColumn {
923                cells: c.cells.clone(),
924                offset: c.offset + offset,
925                len,
926            }),
927        }
928    }
929
930    /// Rows picked from several columns of this kind: `(source, row)`
931    /// pairs, in order. Every source must be the kind of the first.
932    pub(crate) fn gather(sources: &[&ColumnData], picks: &[(u32, u32)]) -> ColumnData {
933        macro_rules! each {
934            ($variant:ident) => {
935                sources
936                    .iter()
937                    .map(|s| match s {
938                        ColumnData::$variant(c) => c,
939                        _ => unreachable!("sources checked to share a kind"),
940                    })
941                    .collect::<Vec<_>>()
942            };
943        }
944        match sources[0] {
945            ColumnData::Number(_) => {
946                let cols = each!(Number);
947                let values: Vec<&[f64]> = cols.iter().map(|c| c.values()).collect();
948                ColumnData::from(
949                    picks
950                        .iter()
951                        .map(|&(s, r)| values[s as usize][r as usize])
952                        .collect::<Vec<f64>>(),
953                )
954            }
955            ColumnData::Decimal(first) => {
956                let cols = each!(Decimal);
957                let precision = cols.iter().map(|c| c.precision()).max().unwrap_or(38);
958                first
959                    .rebuild(
960                        picks
961                            .iter()
962                            .map(|&(s, r)| cols[s as usize].get(r as usize).map(|d| d.value)),
963                        precision,
964                    )
965                    .into()
966            }
967            ColumnData::String(_) => {
968                let cols = each!(String);
969                let picked: Vec<Option<&str>> = picks
970                    .iter()
971                    .map(|&(s, r)| cols[s as usize].get(r as usize))
972                    .collect();
973                // Sized up front, as in `StringColumn::take`.
974                let bytes = picked.iter().flatten().map(|s| s.len()).sum();
975                let mut out = StringBuilder::with_capacity(picked.len(), bytes);
976                for v in picked {
977                    out.append_option(v);
978                }
979                StringColumn::plain(out.finish()).into()
980            }
981            ColumnData::Boolean(_) => {
982                let cols = each!(Boolean);
983                BooleanColumn::from_options(
984                    picks.iter().map(|&(s, r)| cols[s as usize].get(r as usize)),
985                )
986                .into()
987            }
988            ColumnData::Array(_) => {
989                let cols = each!(Array);
990                ArrayColumn::from_cells(
991                    picks
992                        .iter()
993                        .map(|&(s, r)| cols[s as usize].cell(r as usize).cloned())
994                        .collect(),
995                )
996                .into()
997            }
998        }
999    }
1000
1001    /// The rows named by `rows`, in that order; `None` gives a gap.
1002    pub(crate) fn take(&self, rows: &[Option<usize>]) -> ColumnData {
1003        match self {
1004            ColumnData::Number(c) => {
1005                let v = c.values();
1006                ColumnData::from(
1007                    rows.iter()
1008                        .map(|r| r.map_or(f64::NAN, |i| v[i]))
1009                        .collect::<Vec<f64>>(),
1010                )
1011            }
1012            ColumnData::Decimal(c) => c.take(rows).into(),
1013            ColumnData::String(c) => c.take(rows).into(),
1014            ColumnData::Boolean(c) => c.take(rows).into(),
1015            ColumnData::Array(c) => c.take(rows).into(),
1016        }
1017    }
1018
1019    /// The rows at `indices`, in that order; a null index gives a gap.
1020    /// Text, true/false and decimal columns use Arrow's gather kernel.
1021    pub(crate) fn take_indices(&self, indices: &UInt32Array) -> ColumnData {
1022        let gather = |a: &dyn Array| {
1023            arrow_select::take::take(a, indices, None).expect("row numbers are in range")
1024        };
1025        match self {
1026            ColumnData::Number(c) => {
1027                let v = c.values();
1028                ColumnData::from(
1029                    indices
1030                        .iter()
1031                        .map(|r| r.map_or(f64::NAN, |i| v[i as usize]))
1032                        .collect::<Vec<f64>>(),
1033                )
1034            }
1035            ColumnData::Decimal(c) => DecimalColumn {
1036                values: gather(&c.values).as_primitive::<Decimal128Type>().clone(),
1037            }
1038            .into(),
1039            ColumnData::String(c) => c.take_indices(indices).into(),
1040            ColumnData::Boolean(c) => {
1041                BooleanColumn::from_arrow(gather(&c.values).as_boolean().clone()).into()
1042            }
1043            ColumnData::Array(c) => ArrayColumn::from_cells(
1044                indices
1045                    .iter()
1046                    .map(|r| r.and_then(|i| c.cell(i as usize).cloned()))
1047                    .collect(),
1048            )
1049            .into(),
1050        }
1051    }
1052
1053    /// True when row `i` holds a value. Panics if `i` is out of range.
1054    pub fn is_present(&self, i: usize) -> bool {
1055        match self {
1056            ColumnData::Number(c) => c.get(i).is_some(),
1057            ColumnData::Decimal(c) => c.get(i).is_some(),
1058            ColumnData::String(c) => c.get(i).is_some(),
1059            ColumnData::Boolean(c) => c.get(i).is_some(),
1060            ColumnData::Array(c) => c.get(i).is_some(),
1061        }
1062    }
1063
1064    /// What the column holds: [`ColumnKind::Number`], `Decimal`, `String`,
1065    /// `Boolean` or `Array`.
1066    pub fn kind(&self) -> ColumnKind {
1067        match self {
1068            ColumnData::Number(_) => ColumnKind::Number,
1069            ColumnData::Decimal(_) => ColumnKind::Decimal,
1070            ColumnData::String(_) => ColumnKind::String,
1071            ColumnData::Boolean(_) => ColumnKind::Boolean,
1072            ColumnData::Array(_) => ColumnKind::Array,
1073        }
1074    }
1075}
1076
1077impl From<Float64Column> for ColumnData {
1078    fn from(c: Float64Column) -> Self {
1079        ColumnData::Number(c)
1080    }
1081}
1082
1083impl From<DecimalColumn> for ColumnData {
1084    fn from(c: DecimalColumn) -> Self {
1085        ColumnData::Decimal(c)
1086    }
1087}
1088
1089impl From<StringColumn> for ColumnData {
1090    fn from(c: StringColumn) -> Self {
1091        ColumnData::String(c)
1092    }
1093}
1094
1095impl From<BooleanColumn> for ColumnData {
1096    fn from(c: BooleanColumn) -> Self {
1097        ColumnData::Boolean(c)
1098    }
1099}
1100
1101impl From<ArrayColumn> for ColumnData {
1102    fn from(c: ArrayColumn) -> Self {
1103        ColumnData::Array(c)
1104    }
1105}
1106
1107/// `From` for plain values (`Vec<f64>`, `Vec<&str>`, `Vec<bool>`,
1108/// `Vec<Vec<Scalar>>`, …), their
1109/// `Option` forms (`None` is a gap), and arrays of either.
1110macro_rules! column_data_from {
1111    ($($item:ty => |$v:ident| $build:expr;)*) => {$(
1112        impl From<Vec<$item>> for ColumnData {
1113            fn from($v: Vec<$item>) -> Self {
1114                $build
1115            }
1116        }
1117        impl<const N: usize> From<[$item; N]> for ColumnData {
1118            fn from(values: [$item; N]) -> Self {
1119                Vec::from(values).into()
1120            }
1121        }
1122    )*};
1123}
1124
1125column_data_from! {
1126    f64 => |v| ColumnData::Number(Float64Column::from_values(v));
1127    Option<f64> => |v| ColumnData::Number(Float64Column::from_options(v));
1128    &str => |v| ColumnData::String(StringColumn::from_options(v.into_iter().map(Some)));
1129    Option<&str> => |v| ColumnData::String(StringColumn::from_options(v));
1130    String => |v| ColumnData::String(StringColumn::from_options(v.into_iter().map(Some)));
1131    Option<String> => |v| ColumnData::String(StringColumn::from_options(v));
1132    bool => |v| ColumnData::Boolean(BooleanColumn::from_options(v.into_iter().map(Some)));
1133    Option<bool> => |v| ColumnData::Boolean(BooleanColumn::from_options(v));
1134    Vec<Scalar> => |v| ColumnData::Array(ArrayColumn::from_options(v.into_iter().map(Some)));
1135    Option<Vec<Scalar>> => |v| ColumnData::Array(ArrayColumn::from_options(v));
1136}