nodedb_array/query/
elementwise.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Pairwise binary ops between two coord-aligned sparse tiles.
4//!
5//! Both inputs must share the same schema (same dim arity, same attr
6//! columns, same dtypes). We outer-join on coordinates: cells present
7//! in only one operand contribute a Null on the other side, which
8//! propagates as Null through the op (SQL-style null semantics).
9//!
10//! Numeric attrs (Int64, Float64) participate in the op; non-numeric
11//! attrs (String, Bytes) are passed through from the left operand
12//! unchanged — the op is undefined for them. Division by zero yields
13//! Null rather than infinity so downstream aggregates stay finite.
14
15use std::collections::BTreeMap;
16
17use crate::error::{ArrayError, ArrayResult};
18use crate::schema::ArraySchema;
19use crate::tile::sparse_tile::{RowKind, SparseTile, SparseTileBuilder};
20use crate::types::cell_value::value::CellValue;
21use crate::types::coord::value::CoordValue;
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum BinaryOp {
25    Add,
26    Sub,
27    Mul,
28    Div,
29}
30
31/// Apply `op` to every aligned pair of cells. The result tile carries
32/// the union of both coord sets.
33pub fn elementwise(
34    schema: &ArraySchema,
35    a: &SparseTile,
36    b: &SparseTile,
37    op: BinaryOp,
38) -> ArrayResult<SparseTile> {
39    let n_attrs = schema.attrs.len();
40    let by_coord_a = index_rows(a)?;
41    let by_coord_b = index_rows(b)?;
42    let mut keys: BTreeMap<Vec<CoordKey>, ()> = BTreeMap::new();
43    for k in by_coord_a.keys().chain(by_coord_b.keys()) {
44        keys.insert(k.clone(), ());
45    }
46
47    let mut builder = SparseTileBuilder::new(schema);
48    for key in keys.keys() {
49        let coord = decode_key(key);
50        let lhs = by_coord_a.get(key);
51        let rhs = by_coord_b.get(key);
52        let mut out = Vec::with_capacity(n_attrs);
53        for i in 0..n_attrs {
54            let l = lhs.map(|v| v[i].clone()).unwrap_or(CellValue::Null);
55            let r = rhs.map(|v| v[i].clone()).unwrap_or(CellValue::Null);
56            out.push(apply(&l, &r, op));
57        }
58        builder.push(&coord, &out)?;
59    }
60    Ok(builder.build())
61}
62
63fn apply(l: &CellValue, r: &CellValue, op: BinaryOp) -> CellValue {
64    let (lf, rf) = match (to_f64(l), to_f64(r)) {
65        (Some(a), Some(b)) => (a, b),
66        _ => return passthrough(l, r),
67    };
68    let v = match op {
69        BinaryOp::Add => lf + rf,
70        BinaryOp::Sub => lf - rf,
71        BinaryOp::Mul => lf * rf,
72        BinaryOp::Div => {
73            if rf == 0.0 {
74                return CellValue::Null;
75            }
76            lf / rf
77        }
78    };
79    // Preserve Int64 when both operands were integral and result is
80    // exact — keeps integer attrs from drifting to floats unnecessarily.
81    if let (CellValue::Int64(_), CellValue::Int64(_)) = (l, r)
82        && v.fract() == 0.0
83        && v.is_finite()
84    {
85        return CellValue::Int64(v as i64);
86    }
87    CellValue::Float64(v)
88}
89
90fn passthrough(l: &CellValue, r: &CellValue) -> CellValue {
91    if !l.is_null() {
92        l.clone()
93    } else if !r.is_null() {
94        r.clone()
95    } else {
96        CellValue::Null
97    }
98}
99
100fn to_f64(v: &CellValue) -> Option<f64> {
101    match v {
102        CellValue::Int64(x) => Some(*x as f64),
103        CellValue::Float64(x) => Some(*x),
104        _ => None,
105    }
106}
107
108/// Hashable, totally-ordered key for a coordinate tuple.
109#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
110enum CoordKey {
111    I(i64),
112    Ts(i64),
113    F(u64), // f64 bits
114    S(String),
115}
116
117fn encode_key(c: &[CoordValue]) -> Vec<CoordKey> {
118    c.iter()
119        .map(|v| match v {
120            CoordValue::Int64(x) => CoordKey::I(*x),
121            CoordValue::TimestampMs(x) => CoordKey::Ts(*x),
122            CoordValue::Float64(x) => CoordKey::F(x.to_bits()),
123            CoordValue::String(s) => CoordKey::S(s.clone()),
124        })
125        .collect()
126}
127
128fn decode_key(k: &[CoordKey]) -> Vec<CoordValue> {
129    k.iter()
130        .map(|v| match v {
131            CoordKey::I(x) => CoordValue::Int64(*x),
132            CoordKey::Ts(x) => CoordValue::TimestampMs(*x),
133            CoordKey::F(b) => CoordValue::Float64(f64::from_bits(*b)),
134            CoordKey::S(s) => CoordValue::String(s.clone()),
135        })
136        .collect()
137}
138
139/// The live cells of `tile` by coordinate. Tombstone and erasure rows hold
140/// no attribute entries, so attributes are read at the live-row index.
141fn index_rows(tile: &SparseTile) -> ArrayResult<BTreeMap<Vec<CoordKey>, Vec<CellValue>>> {
142    let corrupt = |row: usize| ArrayError::SegmentCorruption {
143        detail: format!("elementwise: sparse tile row {row} is out of range"),
144    };
145    let mut out = BTreeMap::new();
146    let mut live_row = 0usize;
147    for row in 0..tile.row_count() {
148        if tile.row_kind(row)? != RowKind::Live {
149            continue;
150        }
151        let coord = tile
152            .dim_dicts
153            .iter()
154            .map(|d| {
155                d.indices
156                    .get(row)
157                    .and_then(|&i| d.values.get(i as usize))
158                    .cloned()
159                    .ok_or_else(|| corrupt(row))
160            })
161            .collect::<ArrayResult<Vec<CoordValue>>>()?;
162        let attrs = tile
163            .attr_cols
164            .iter()
165            .map(|col| col.get(live_row).cloned().ok_or_else(|| corrupt(row)))
166            .collect::<ArrayResult<Vec<CellValue>>>()?;
167        live_row += 1;
168        out.insert(encode_key(&coord), attrs);
169    }
170    Ok(out)
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176    use crate::schema::ArraySchemaBuilder;
177    use crate::schema::attr_spec::{AttrSpec, AttrType};
178    use crate::schema::dim_spec::{DimSpec, DimType};
179    use crate::types::domain::{Domain, DomainBound};
180
181    fn schema() -> ArraySchema {
182        ArraySchemaBuilder::new("g")
183            .dim(DimSpec::new(
184                "k",
185                DimType::Int64,
186                Domain::new(DomainBound::Int64(0), DomainBound::Int64(15)),
187            ))
188            .attr(AttrSpec::new("v", AttrType::Int64, true))
189            .tile_extents(vec![16])
190            .build()
191            .unwrap()
192    }
193
194    fn tile(rows: &[(i64, i64)]) -> SparseTile {
195        let s = schema();
196        let mut b = SparseTileBuilder::new(&s);
197        for (k, v) in rows {
198            b.push(&[CoordValue::Int64(*k)], &[CellValue::Int64(*v)])
199                .unwrap();
200        }
201        b.build()
202    }
203
204    #[test]
205    fn add_aligned_cells() {
206        let s = schema();
207        let a = tile(&[(0, 1), (1, 2)]);
208        let b = tile(&[(0, 10), (1, 20)]);
209        let out = elementwise(&s, &a, &b, BinaryOp::Add).unwrap();
210        assert_eq!(out.nnz(), 2);
211        assert_eq!(out.attr_cols[0][0], CellValue::Int64(11));
212        assert_eq!(out.attr_cols[0][1], CellValue::Int64(22));
213    }
214
215    #[test]
216    fn outer_join_propagates_null() {
217        let s = schema();
218        let a = tile(&[(0, 1)]);
219        let b = tile(&[(1, 2)]);
220        let out = elementwise(&s, &a, &b, BinaryOp::Add).unwrap();
221        assert_eq!(out.nnz(), 2);
222        // Both rows are missing one operand → both become Null+something
223        // → passthrough gives the present operand.
224        let v0 = &out.attr_cols[0][0];
225        let v1 = &out.attr_cols[0][1];
226        assert!(matches!(v0, CellValue::Int64(1) | CellValue::Int64(2)));
227        assert!(matches!(v1, CellValue::Int64(1) | CellValue::Int64(2)));
228    }
229
230    #[test]
231    fn div_by_zero_returns_null() {
232        let s = schema();
233        let a = tile(&[(0, 5)]);
234        let b = tile(&[(0, 0)]);
235        let out = elementwise(&s, &a, &b, BinaryOp::Div).unwrap();
236        assert_eq!(out.attr_cols[0][0], CellValue::Null);
237    }
238
239    /// A tombstone row before a live row neither joins the result nor
240    /// shifts the live row's attributes.
241    #[test]
242    fn a_tombstone_row_is_skipped() {
243        use crate::tile::sparse_tile::SparseRow;
244        use nodedb_types::{OPEN_UPPER, Surrogate};
245
246        let s = schema();
247        let mut builder = SparseTileBuilder::new(&s);
248        builder
249            .push_row(SparseRow {
250                coord: &[CoordValue::Int64(0)],
251                attrs: &[],
252                surrogate: None,
253                valid_from_ms: 0,
254                valid_until_ms: OPEN_UPPER,
255                kind: RowKind::Tombstone,
256            })
257            .unwrap();
258        builder
259            .push_row(SparseRow {
260                coord: &[CoordValue::Int64(1)],
261                attrs: &[CellValue::Int64(4)],
262                surrogate: Some(Surrogate::new(1)),
263                valid_from_ms: 0,
264                valid_until_ms: OPEN_UPPER,
265                kind: RowKind::Live,
266            })
267            .unwrap();
268        let a = builder.build();
269        let b = tile(&[(1, 6)]);
270        let out = elementwise(&s, &a, &b, BinaryOp::Add).unwrap();
271        assert_eq!(out.nnz(), 1);
272        assert_eq!(out.attr_cols[0][0], CellValue::Int64(10));
273    }
274
275    #[test]
276    fn sub_and_mul() {
277        let s = schema();
278        let a = tile(&[(0, 10)]);
279        let b = tile(&[(0, 3)]);
280        let s1 = elementwise(&s, &a, &b, BinaryOp::Sub).unwrap();
281        let s2 = elementwise(&s, &a, &b, BinaryOp::Mul).unwrap();
282        assert_eq!(s1.attr_cols[0][0], CellValue::Int64(7));
283        assert_eq!(s2.attr_cols[0][0], CellValue::Int64(30));
284    }
285}
nodedb_array/query/elementwise.rs

nodedb_array/query/
elementwise.rs
