nd_arrow_array/compute/
column.rs

1//! Kernels on a whole [`NdColumn`]: one kernel call for all rows.
2//!
3//! The `values` child of a column holds the values of all rows, one after
4//! another. An elementwise kernel can run on that child once and keep the
5//! offsets and dimensions. This replaces one call per row with one call per
6//! column.
7
8use std::sync::Arc;
9
10use arrow::array::{Array, ArrayRef, ListArray, StructArray, UInt64Array, new_null_array};
11use arrow::compute::concat;
12use arrow::datatypes::{DataType, Field, Int64Type};
13
14use crate::column::{NdColumn, NdColumnBuilder};
15use crate::compute::cf::decode_cf;
16use crate::compute::reduce::{Reduction, reduce};
17use crate::dimensions::Dimensions;
18use crate::error::{Result, nd_err};
19use crate::extension::{ITEM_FIELD, NdArrayType};
20use crate::{Attributes, NdArrowArray};
21
22/// The values child of all rows as one 1-D array along dimension `n`.
23fn flat_values(column: &NdColumn) -> Result<NdArrowArray> {
24    let values = column.values_list();
25    let first = values.value_offsets()[0] as usize;
26    let last = *values.value_offsets().last().unwrap() as usize;
27    let child = values.values().slice(first, last - first);
28    NdArrowArray::try_new(child.clone(), Dimensions::try_new([("n", child.len())])?)
29}
30
31/// Replace the values child, keep offsets and dimensions. Without
32/// `keep_attributes`, the `attrs` field goes away.
33fn with_values(column: &NdColumn, values: ArrayRef, keep_attributes: bool) -> Result<NdColumn> {
34    let list = column.values_list();
35    let first = list.value_offsets()[0];
36    let offsets = if first == 0 {
37        list.offsets().clone()
38    } else {
39        arrow::buffer::OffsetBuffer::new(list.offsets().iter().map(|o| o - first).collect())
40    };
41    let item = Arc::new(Field::new(ITEM_FIELD, values.data_type().clone(), true));
42    let new_list = ListArray::try_new(item.clone(), offsets, values, list.nulls().cloned())?;
43    let attrs = if keep_attributes {
44        column.extension_type().attribute_fields().clone()
45    } else {
46        Default::default()
47    };
48    let extension = NdArrayType::with_value_field(item, attrs);
49    let storage = column.storage();
50    let mut children: Vec<ArrayRef> = storage.columns()[..3].to_vec();
51    children[0] = Arc::new(new_list);
52    if extension.has_attributes() {
53        children.push(storage.column(3).clone());
54    }
55    let storage = StructArray::try_new(
56        extension.storage_fields(),
57        children,
58        storage.nulls().cloned(),
59    )?;
60    Ok(NdColumn::from_parts_unchecked(storage, extension))
61}
62
63/// Apply an elementwise unary kernel to all rows in one call, for example
64/// `column_unary(&c, compute::sqrt, false)`.
65///
66/// The kernel must keep the number of values. `keep_attributes` follows the
67/// attribute rule of the kernel: `true` for `cast` and `fill_null`, `false`
68/// for math functions.
69pub fn column_unary(
70    column: &NdColumn,
71    kernel: impl Fn(&NdArrowArray) -> Result<NdArrowArray>,
72    keep_attributes: bool,
73) -> Result<NdColumn> {
74    let flat = flat_values(column)?;
75    let out = kernel(&flat)?;
76    if out.len() != flat.len() {
77        return Err(nd_err!(
78            InvalidArgument,
79            "column kernel changed the number of values from {} to {}",
80            flat.len(),
81            out.len()
82        ));
83    }
84    with_values(column, out.values().clone(), keep_attributes)
85}
86
87/// Cast the values of all rows. Keeps the attributes.
88pub fn column_cast(column: &NdColumn, to: &DataType) -> Result<NdColumn> {
89    column_unary(column, |a| crate::compute::cast(a, to), true)
90}
91
92/// Apply an elementwise binary kernel to two columns, for example `compute::add`.
93///
94/// When every row has the same dimensions in both columns, the kernel runs
95/// once on the two values children. Otherwise it runs per row, with
96/// broadcasting. A null row in either column gives a null row.
97pub fn column_binary(
98    a: &NdColumn,
99    b: &NdColumn,
100    kernel: impl Fn(&NdArrowArray, &NdArrowArray) -> Result<NdArrowArray>,
101) -> Result<NdColumn> {
102    if a.len() != b.len() {
103        return Err(nd_err!(
104            ShapeMismatch,
105            "columns have {} and {} rows",
106            a.len(),
107            b.len()
108        ));
109    }
110    let same_layout = (0..a.len()).all(|row| {
111        a.is_null(row) == b.is_null(row) && a.dimensions_at(row) == b.dimensions_at(row)
112    }) && a.values_list().value_offsets() == b.values_list().value_offsets();
113    if same_layout {
114        let out = kernel(&flat_values(a)?, &flat_values(b)?)?;
115        return with_values(a, out.values().clone(), false);
116    }
117    let mut builder: Option<NdColumnBuilder> = None;
118    let mut pending_nulls = 0usize;
119    for row in 0..a.len() {
120        match (a.get(row), b.get(row)) {
121            (Some(x), Some(y)) => {
122                let out = kernel(&x, &y)?;
123                let builder = builder.get_or_insert_with(|| {
124                    let mut b = NdColumnBuilder::with_value_type(out.data_type().clone());
125                    (0..pending_nulls).for_each(|_| b.append_null());
126                    b
127                });
128                builder.append(&out)?;
129            }
130            _ => match &mut builder {
131                Some(builder) => builder.append_null(),
132                None => pending_nulls += 1,
133            },
134        }
135    }
136    match builder {
137        Some(builder) => builder.finish(),
138        None => Ok(NdColumn::new_null(
139            NdArrayType::new(a.value_type().clone(), Default::default()),
140            a.len(),
141        )),
142    }
143}
144
145/// Decode CF packing and fill values of every row, see [`decode_cf`].
146///
147/// When all rows share the same CF attributes, the decode runs once on the
148/// values of all rows. The CF keys leave the attributes.
149pub fn column_decode_cf(column: &NdColumn) -> Result<NdColumn> {
150    let rows: Vec<Option<NdArrowArray>> = column.iter().collect();
151    let cf_of = |a: &NdArrowArray| -> Attributes {
152        let mut kept = Attributes::new();
153        for (key, value) in a.attributes().iter() {
154            if crate::compute::cf::is_cf_key(key) {
155                kept.try_insert_array(key, value.clone()).expect("length 1");
156            }
157        }
158        kept
159    };
160    let present: Vec<&NdArrowArray> = rows.iter().flatten().collect();
161    let shared = present.first().map(|a| cf_of(a));
162    let uniform = shared
163        .as_ref()
164        .is_some_and(|first| present.iter().all(|a| cf_of(a) == *first));
165    if uniform {
166        let template = present[0];
167        let flat = flat_values(column)?.with_attributes(template.attributes().clone());
168        let (decoded, _) = decode_cf(&flat)?;
169        let mut builder = NdColumnBuilder::with_value_type(decoded.data_type().clone());
170        let offsets = column.values_list().value_offsets();
171        let base = offsets[0] as usize;
172        for (row, array) in rows.iter().enumerate() {
173            match array {
174                None => builder.append_null(),
175                Some(array) => {
176                    let start = offsets[row] as usize - base;
177                    let len = offsets[row + 1] as usize - offsets[row] as usize;
178                    let mut attrs = array.attributes().clone();
179                    for (key, _) in cf_of(array).iter() {
180                        attrs.remove(key);
181                    }
182                    let values = decoded.values().slice(start, len);
183                    builder.append(
184                        &NdArrowArray::try_new(values, array.dims().clone())?
185                            .with_attributes(attrs),
186                    )?;
187                }
188            }
189        }
190        return builder.finish();
191    }
192    let mut builder: Option<NdColumnBuilder> = None;
193    let decoded: Vec<Option<NdArrowArray>> = rows
194        .iter()
195        .map(|r| r.as_ref().map(|a| decode_cf(a).map(|(d, _)| d)).transpose())
196        .collect::<Result<_>>()?;
197    for row in &decoded {
198        let b = builder.get_or_insert_with(|| {
199            let value_type = decoded
200                .iter()
201                .flatten()
202                .next()
203                .map_or(column.value_type().clone(), |a| a.data_type().clone());
204            NdColumnBuilder::with_value_type(value_type)
205        });
206        b.append_option(row.as_ref())?;
207    }
208    builder
209        .unwrap_or_else(|| NdColumnBuilder::with_value_type(column.value_type().clone()))
210        .finish()
211}
212
213/// Reduce every row over all its axes: one value per row, like a SQL aggregate
214/// per row. A null row gives null.
215pub fn row_reduce(column: &NdColumn, op: Reduction) -> Result<ArrayRef> {
216    if op == Reduction::Count {
217        return Ok(Arc::new(row_counts(column)));
218    }
219    let parts: Vec<Option<ArrayRef>> = column
220        .iter()
221        .map(|row| {
222            row.map(|a| reduce(&a, None, op).map(|r| r.values().clone()))
223                .transpose()
224        })
225        .collect::<Result<_>>()?;
226    let out_type = parts
227        .iter()
228        .flatten()
229        .next()
230        .map_or(DataType::Null, |p| p.data_type().clone());
231    let null = new_null_array(&out_type, 1);
232    let refs: Vec<&dyn Array> = parts
233        .iter()
234        .map(|p| p.as_ref().map_or(null.as_ref(), |p| p.as_ref()))
235        .collect();
236    if refs.is_empty() {
237        return Ok(new_null_array(&out_type, 0));
238    }
239    Ok(concat(&refs)?)
240}
241
242/// Non-null values per row, computed from offsets and the null buffer only.
243fn row_counts(column: &NdColumn) -> arrow::array::PrimitiveArray<Int64Type> {
244    let list = column.values_list();
245    let values = list.values();
246    let offsets = list.value_offsets();
247    (0..column.len())
248        .map(|row| {
249            (!column.is_null(row)).then(|| {
250                let (start, end) = (offsets[row] as usize, offsets[row + 1] as usize);
251                let nulls = values
252                    .logical_nulls()
253                    .map_or(0, |n| (start..end).filter(|&i| n.is_null(i)).count());
254                (end - start - nulls) as i64
255            })
256        })
257        .collect()
258}
259
260/// Reduce every row over the named axes. Rows without an axis keep it out of
261/// the reduction. Returns a new column.
262pub fn row_reduce_axes(column: &NdColumn, axes: &[&str], op: Reduction) -> Result<NdColumn> {
263    let mut reduced: Vec<Option<NdArrowArray>> = Vec::with_capacity(column.len());
264    for array in column.iter() {
265        reduced.push(
266            array
267                .map(|a| {
268                    let own: Vec<&str> = axes
269                        .iter()
270                        .copied()
271                        .filter(|name| a.dims().contains(name))
272                        .collect();
273                    reduce(&a, Some(&own), op)
274                })
275                .transpose()?,
276        );
277    }
278    let value_type = reduced
279        .iter()
280        .flatten()
281        .next()
282        .map(|a| a.data_type().clone())
283        .unwrap_or(DataType::Float64);
284    let mut builder = NdColumnBuilder::with_value_type(value_type);
285    for row in &reduced {
286        builder.append_option(row.as_ref())?;
287    }
288    builder.finish()
289}
290
291/// Indices of non-null rows. Helper for callers that filter columns.
292pub fn valid_rows(column: &NdColumn) -> UInt64Array {
293    (0..column.len() as u64)
294        .filter(|&row| !column.is_null(row as usize))
295        .collect()
296}
297
298#[cfg(test)]
299mod tests {
300    use arrow::array::{AsArray, Int16Array, Int32Array};
301    use arrow::datatypes::{Float32Type, Float64Type, Int32Type};
302
303    use super::*;
304    use crate::compute::cf::{ADD_OFFSET, SCALE_FACTOR};
305    use crate::compute::{add, sqrt};
306
307    fn nd(values: Vec<i32>, dims: &[(&str, usize)]) -> NdArrowArray {
308        NdArrowArray::try_new(
309            Arc::new(Int32Array::from(values)),
310            Dimensions::try_new(dims.iter().copied()).unwrap(),
311        )
312        .unwrap()
313    }
314
315    fn column(rows: &[Option<NdArrowArray>]) -> NdColumn {
316        let mut b = NdColumnBuilder::with_value_type(DataType::Int32);
317        for r in rows {
318            b.append_option(r.as_ref()).unwrap();
319        }
320        b.finish().unwrap()
321    }
322
323    #[test]
324    fn unary_runs_once_and_keeps_layout() {
325        let c = column(&[
326            Some(nd(vec![1, 4], &[("x", 2)])),
327            None,
328            Some(nd(vec![9, 16, 25, 36], &[("y", 2), ("x", 2)])),
329        ])
330        .slice(0, 3);
331        let out = column_unary(&c, sqrt, false).unwrap();
332        assert_eq!(out.value_type(), &DataType::Float64);
333        assert!(out.is_null(1));
334        let row2 = out.value(2);
335        assert_eq!(row2.dims().names(), vec!["y", "x"]);
336        assert_eq!(
337            row2.values().as_primitive::<Float64Type>().values(),
338            &[3.0, 4.0, 5.0, 6.0]
339        );
340    }
341
342    #[test]
343    fn unary_on_a_sliced_column() {
344        let c = column(&[Some(nd(vec![1], &[])), Some(nd(vec![4, 9], &[("x", 2)]))]).slice(1, 1);
345        let out = column_unary(&c, sqrt, false).unwrap();
346        assert_eq!(out.len(), 1);
347        assert_eq!(
348            out.value(0).values().as_primitive::<Float64Type>().values(),
349            &[2.0, 3.0]
350        );
351    }
352
353    #[test]
354    fn binary_fast_and_per_row_paths() {
355        let a = column(&[Some(nd(vec![1, 2], &[("x", 2)])), Some(nd(vec![3], &[]))]);
356        let same = column_binary(&a, &a, add).unwrap();
357        assert_eq!(
358            same.value(0).values().as_primitive::<Int32Type>().values(),
359            &[2, 4]
360        );
361        let b = column(&[Some(nd(vec![10], &[])), None]);
362        let mixed = column_binary(&a, &b, add).unwrap();
363        assert_eq!(
364            mixed.value(0).values().as_primitive::<Int32Type>().values(),
365            &[11, 12]
366        );
367        assert!(mixed.is_null(1));
368    }
369
370    #[test]
371    fn row_reductions() {
372        let c = column(&[
373            Some(nd(vec![1, 2, 3, 4], &[("y", 2), ("x", 2)])),
374            None,
375            Some(nd(vec![5], &[])),
376        ]);
377        let sums = row_reduce(&c, Reduction::Sum).unwrap();
378        let sums: Vec<Option<i64>> = sums.as_primitive::<Int64Type>().iter().collect();
379        assert_eq!(sums, vec![Some(10), None, Some(5)]);
380        let counts = row_reduce(&c, Reduction::Count).unwrap();
381        assert_eq!(counts.as_primitive::<Int64Type>().value(0), 4);
382        assert!(counts.is_null(1));
383
384        let per_y = row_reduce_axes(&c, &["x"], Reduction::Max).unwrap();
385        assert_eq!(per_y.value(0).dims().names(), vec!["y"]);
386        assert_eq!(
387            per_y.value(0).values().as_primitive::<Int32Type>().values(),
388            &[2, 4]
389        );
390        assert!(per_y.value(2).is_scalar());
391        assert_eq!(valid_rows(&c).values(), &[0, 2]);
392    }
393
394    #[test]
395    fn decode_cf_uniform_and_mixed() {
396        let packed = |values: Vec<i16>, scale: f32| {
397            NdArrowArray::try_new(
398                Arc::new(Int16Array::from(values)),
399                Dimensions::try_new([("x", 2)]).unwrap(),
400            )
401            .unwrap()
402            .with_attributes(
403                Attributes::new()
404                    .with(SCALE_FACTOR, scale)
405                    .with(ADD_OFFSET, 1.0f32)
406                    .with("units", "K"),
407            )
408        };
409        let uniform =
410            NdColumn::try_from_arrays([&packed(vec![1, 2], 0.5), &packed(vec![3, 4], 0.5)])
411                .unwrap();
412        let out = column_decode_cf(&uniform).unwrap();
413        assert_eq!(out.value_type(), &DataType::Float32);
414        assert_eq!(
415            out.value(1).values().as_primitive::<Float32Type>().values(),
416            &[2.5, 3.0]
417        );
418        assert_eq!(out.value(1).attributes().get_str("units"), Some("K"));
419        assert!(!out.value(1).attributes().contains(SCALE_FACTOR));
420
421        let mixed = NdColumn::try_from_arrays([&packed(vec![1, 2], 0.5), &packed(vec![1, 2], 2.0)])
422            .unwrap();
423        let out = column_decode_cf(&mixed).unwrap();
424        assert_eq!(
425            out.value(1).values().as_primitive::<Float32Type>().values(),
426            &[3.0, 5.0]
427        );
428    }
429}
nd_arrow_array/compute/column.rs

nd_arrow_array/compute/
column.rs
