nd_arrow_array/compute/
elementwise.rs

1//! Elementwise kernels with name-based broadcasting and NumPy type promotion.
2
3use std::sync::Arc;
4
5use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, Datum, Scalar};
6use arrow::compute::kernels::zip::zip;
7use arrow::compute::kernels::{boolean, cmp, numeric};
8use arrow::compute::{cast as arrow_cast, is_not_null, is_null};
9use arrow::datatypes::{
10    DataType, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type,
11};
12
13use crate::NdArrowArray;
14use crate::compute::shape::{broadcast_arrays, broadcast_to};
15use crate::dimensions::Dimensions;
16use crate::error::{Result, nd_err};
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19enum Kind {
20    Bool,
21    Int(u8),
22    UInt(u8),
23    Float(u8),
24}
25
26fn kind(data_type: &DataType) -> Option<Kind> {
27    Some(match data_type {
28        DataType::Boolean => Kind::Bool,
29        DataType::Int8 => Kind::Int(8),
30        DataType::Int16 => Kind::Int(16),
31        DataType::Int32 => Kind::Int(32),
32        DataType::Int64 => Kind::Int(64),
33        DataType::UInt8 => Kind::UInt(8),
34        DataType::UInt16 => Kind::UInt(16),
35        DataType::UInt32 => Kind::UInt(32),
36        DataType::UInt64 => Kind::UInt(64),
37        DataType::Float16 => Kind::Float(16),
38        DataType::Float32 => Kind::Float(32),
39        DataType::Float64 => Kind::Float(64),
40        _ => return None,
41    })
42}
43
44fn data_type(kind: Kind) -> DataType {
45    match kind {
46        Kind::Bool => DataType::Boolean,
47        Kind::Int(8) => DataType::Int8,
48        Kind::Int(16) => DataType::Int16,
49        Kind::Int(32) => DataType::Int32,
50        Kind::Int(_) => DataType::Int64,
51        Kind::UInt(8) => DataType::UInt8,
52        Kind::UInt(16) => DataType::UInt16,
53        Kind::UInt(32) => DataType::UInt32,
54        Kind::UInt(_) => DataType::UInt64,
55        Kind::Float(16) => DataType::Float16,
56        Kind::Float(32) => DataType::Float32,
57        Kind::Float(_) => DataType::Float64,
58    }
59}
60
61/// The NumPy result type of two numeric types, like `numpy.result_type`.
62///
63/// Returns `None` for pairs that are not both numeric or boolean. Kernels pass
64/// those types to Arrow unchanged, so temporal arithmetic keeps working.
65pub fn result_type(a: &DataType, b: &DataType) -> Option<DataType> {
66    if a == b {
67        return Some(a.clone());
68    }
69    let float_bits_for = |bits: u8| match bits {
70        8 => 16,
71        16 => 32,
72        _ => 64,
73    };
74    let kind = match (kind(a)?, kind(b)?) {
75        (Kind::Bool, other) | (other, Kind::Bool) => other,
76        (Kind::Int(x), Kind::Int(y)) => Kind::Int(x.max(y)),
77        (Kind::UInt(x), Kind::UInt(y)) => Kind::UInt(x.max(y)),
78        (Kind::Int(s), Kind::UInt(u)) | (Kind::UInt(u), Kind::Int(s)) => {
79            if s > u {
80                Kind::Int(s)
81            } else if u < 64 {
82                Kind::Int(u * 2)
83            } else {
84                Kind::Float(64)
85            }
86        }
87        (Kind::Float(x), Kind::Float(y)) => Kind::Float(x.max(y)),
88        (Kind::Float(f), Kind::Int(n) | Kind::UInt(n))
89        | (Kind::Int(n) | Kind::UInt(n), Kind::Float(f)) => Kind::Float(f.max(float_bits_for(n))),
90    };
91    Some(data_type(kind))
92}
93
94/// The floating point type that math functions compute in, like NumPy.
95fn float_type(data_type: &DataType) -> Result<DataType> {
96    match kind(data_type) {
97        Some(Kind::Float(64)) => Ok(DataType::Float64),
98        Some(Kind::Float(_)) => Ok(DataType::Float32),
99        Some(_) => Ok(DataType::Float64),
100        None => Err(nd_err!(
101            NotSupported,
102            "math function on non-numeric type {data_type}"
103        )),
104    }
105}
106
107fn cast_values(values: &ArrayRef, to: &DataType) -> Result<ArrayRef> {
108    if values.data_type() == to {
109        Ok(values.clone())
110    } else {
111        Ok(arrow_cast(values.as_ref(), to)?)
112    }
113}
114
115enum Operand {
116    Array(ArrayRef),
117    Scalar(Scalar<ArrayRef>),
118}
119
120impl Operand {
121    fn datum(&self) -> &dyn Datum {
122        match self {
123            Operand::Array(array) => array,
124            Operand::Scalar(scalar) => scalar,
125        }
126    }
127}
128
129/// Cast both sides to `to` and align them. Rank-0 sides stay Arrow scalars.
130fn prepare(
131    a: &NdArrowArray,
132    b: &NdArrowArray,
133    to: Option<&DataType>,
134) -> Result<(Operand, Operand, Dimensions)> {
135    let cast = |x: &NdArrowArray| match to {
136        Some(t) => cast_values(x.values(), t),
137        None => Ok(x.values().clone()),
138    };
139    if a.dims() == b.dims() {
140        return Ok((
141            Operand::Array(cast(a)?),
142            Operand::Array(cast(b)?),
143            a.dims().clone(),
144        ));
145    }
146    if b.is_scalar() {
147        return Ok((
148            Operand::Array(cast(a)?),
149            Operand::Scalar(Scalar::new(cast(b)?)),
150            a.dims().clone(),
151        ));
152    }
153    if a.is_scalar() {
154        return Ok((
155            Operand::Scalar(Scalar::new(cast(a)?)),
156            Operand::Array(cast(b)?),
157            b.dims().clone(),
158        ));
159    }
160    let aligned = broadcast_arrays(&[a, b])?;
161    Ok((
162        Operand::Array(cast(&aligned[0])?),
163        Operand::Array(cast(&aligned[1])?),
164        aligned[0].dims().clone(),
165    ))
166}
167
168fn wrap(values: ArrayRef, dims: Dimensions) -> Result<NdArrowArray> {
169    NdArrowArray::try_new(values, dims)
170}
171
172fn arithmetic(
173    a: &NdArrowArray,
174    b: &NdArrowArray,
175    to: Option<DataType>,
176    op: impl Fn(&dyn Datum, &dyn Datum) -> std::result::Result<ArrayRef, arrow::error::ArrowError>,
177) -> Result<NdArrowArray> {
178    let (x, y, dims) = prepare(a, b, to.as_ref())?;
179    wrap(op(x.datum(), y.datum())?, dims)
180}
181
182/// `a + b`. Integer overflow wraps, like NumPy.
183pub fn add(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
184    arithmetic(
185        a,
186        b,
187        result_type(a.data_type(), b.data_type()),
188        numeric::add_wrapping,
189    )
190}
191
192/// `a - b`. Integer overflow wraps, like NumPy.
193pub fn sub(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
194    arithmetic(
195        a,
196        b,
197        result_type(a.data_type(), b.data_type()),
198        numeric::sub_wrapping,
199    )
200}
201
202/// `a * b`. Integer overflow wraps, like NumPy.
203pub fn mul(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
204    arithmetic(
205        a,
206        b,
207        result_type(a.data_type(), b.data_type()),
208        numeric::mul_wrapping,
209    )
210}
211
212/// True division `a / b`, like `numpy.true_divide`: integers divide as `Float64`.
213pub fn div(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
214    let to = match result_type(a.data_type(), b.data_type()) {
215        Some(t) if matches!(kind(&t), Some(Kind::Float(_))) => Some(t),
216        Some(_) => Some(DataType::Float64),
217        None => None,
218    };
219    arithmetic(a, b, to, numeric::div)
220}
221
222/// Remainder with the sign of `a`, like `numpy.fmod`. Integer division by zero is an error.
223pub fn rem(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
224    arithmetic(
225        a,
226        b,
227        result_type(a.data_type(), b.data_type()),
228        numeric::rem,
229    )
230}
231
232/// Floor division `a // b`, like `numpy.floor_divide`.
233/// Integer division by zero gives null.
234pub fn floor_divide(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
235    floor_div_mod(a, b, false)
236}
237
238/// Remainder with the sign of `b`, like `numpy.mod` (Python `%`).
239/// Integer division by zero gives null.
240pub fn mod_(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
241    floor_div_mod(a, b, true)
242}
243
244fn floor_div_mod(a: &NdArrowArray, b: &NdArrowArray, modulo: bool) -> Result<NdArrowArray> {
245    let to = result_type(a.data_type(), b.data_type()).ok_or_else(|| {
246        nd_err!(
247            NotSupported,
248            "floor division on {} and {}",
249            a.data_type(),
250            b.data_type()
251        )
252    })?;
253    let aligned = broadcast_arrays(&[a, b])?;
254    let dims = aligned[0].dims().clone();
255    let values: ArrayRef = match kind(&to) {
256        Some(Kind::Float(bits)) => {
257            let x = cast_values(aligned[0].values(), &DataType::Float64)?;
258            let y = cast_values(aligned[1].values(), &DataType::Float64)?;
259            let out = arrow::compute::binary::<_, _, _, Float64Type>(
260                x.as_primitive::<Float64Type>(),
261                y.as_primitive::<Float64Type>(),
262                |x, y| {
263                    if modulo {
264                        let r = x % y;
265                        if r != 0.0 && (r < 0.0) != (y < 0.0) {
266                            r + y
267                        } else {
268                            r
269                        }
270                    } else {
271                        (x / y).floor()
272                    }
273                },
274            )?;
275            let to = if bits == 64 {
276                DataType::Float64
277            } else {
278                to.clone()
279            };
280            cast_values(&(Arc::new(out) as ArrayRef), &to)?
281        }
282        Some(Kind::Int(_) | Kind::UInt(_) | Kind::Bool) => {
283            let x = cast_values(aligned[0].values(), &DataType::Int64)?;
284            let y = cast_values(aligned[1].values(), &DataType::Int64)?;
285            let (x, y) = (x.as_primitive::<Int64Type>(), y.as_primitive::<Int64Type>());
286            let out: arrow::array::Int64Array = x
287                .iter()
288                .zip(y.iter())
289                .map(|(x, y)| match (x, y) {
290                    (Some(_), Some(0)) | (None, _) | (_, None) => None,
291                    (Some(x), Some(y)) => {
292                        let q = x.wrapping_div(y) - i64::from(x % y != 0 && ((x < 0) != (y < 0)));
293                        Some(if modulo {
294                            x.wrapping_sub(q.wrapping_mul(y))
295                        } else {
296                            q
297                        })
298                    }
299                })
300                .collect();
301            cast_values(&(Arc::new(out) as ArrayRef), &to)?
302        }
303        None => unreachable!("result_type returned a numeric type"),
304    };
305    wrap(values, dims)
306}
307
308/// `a ** b` in floating point, like `numpy.float_power`.
309pub fn pow(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
310    let to = result_type(a.data_type(), b.data_type()).ok_or_else(|| {
311        nd_err!(
312            NotSupported,
313            "pow on {} and {}",
314            a.data_type(),
315            b.data_type()
316        )
317    })?;
318    let to = float_type(&to)?;
319    let aligned = broadcast_arrays(&[a, b])?;
320    let x = cast_values(aligned[0].values(), &to)?;
321    let y = cast_values(aligned[1].values(), &to)?;
322    let values: ArrayRef = match to {
323        DataType::Float32 => Arc::new(arrow::compute::binary::<_, _, _, Float32Type>(
324            x.as_primitive::<Float32Type>(),
325            y.as_primitive::<Float32Type>(),
326            f32::powf,
327        )?),
328        _ => Arc::new(arrow::compute::binary::<_, _, _, Float64Type>(
329            x.as_primitive::<Float64Type>(),
330            y.as_primitive::<Float64Type>(),
331            f64::powf,
332        )?),
333    };
334    wrap(values, aligned[0].dims().clone())
335}
336
337fn comparison(
338    a: &NdArrowArray,
339    b: &NdArrowArray,
340    op: impl Fn(&dyn Datum, &dyn Datum) -> std::result::Result<BooleanArray, arrow::error::ArrowError>,
341) -> Result<NdArrowArray> {
342    let to = result_type(a.data_type(), b.data_type());
343    let (x, y, dims) = prepare(a, b, to.as_ref())?;
344    wrap(Arc::new(op(x.datum(), y.datum())?), dims)
345}
346
347/// `a == b`.
348pub fn eq(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
349    comparison(a, b, cmp::eq)
350}
351
352/// `a != b`.
353pub fn ne(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
354    comparison(a, b, cmp::neq)
355}
356
357/// `a < b`.
358pub fn lt(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
359    comparison(a, b, cmp::lt)
360}
361
362/// `a <= b`.
363pub fn le(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
364    comparison(a, b, cmp::lt_eq)
365}
366
367/// `a > b`.
368pub fn gt(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
369    comparison(a, b, cmp::gt)
370}
371
372/// `a >= b`.
373pub fn ge(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
374    comparison(a, b, cmp::gt_eq)
375}
376
377fn boolean_pair(
378    a: &NdArrowArray,
379    b: &NdArrowArray,
380    op: impl Fn(
381        &BooleanArray,
382        &BooleanArray,
383    ) -> std::result::Result<BooleanArray, arrow::error::ArrowError>,
384) -> Result<NdArrowArray> {
385    for x in [a, b] {
386        if x.data_type() != &DataType::Boolean {
387            return Err(nd_err!(
388                InvalidArgument,
389                "logical operation needs Boolean values, got {}",
390                x.data_type()
391            ));
392        }
393    }
394    let aligned = broadcast_arrays(&[a, b])?;
395    let values = op(
396        aligned[0].values().as_boolean(),
397        aligned[1].values().as_boolean(),
398    )?;
399    wrap(Arc::new(values), aligned[0].dims().clone())
400}
401
402/// Logical AND with Kleene null logic (`false AND null` is `false`).
403pub fn logical_and(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
404    boolean_pair(a, b, boolean::and_kleene)
405}
406
407/// Logical OR with Kleene null logic (`true OR null` is `true`).
408pub fn logical_or(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
409    boolean_pair(a, b, boolean::or_kleene)
410}
411
412/// Logical XOR. A null input gives a null.
413pub fn logical_xor(a: &NdArrowArray, b: &NdArrowArray) -> Result<NdArrowArray> {
414    boolean_pair(a, b, |x, y| cmp::neq(x, y))
415}
416
417/// Logical NOT.
418pub fn logical_not(a: &NdArrowArray) -> Result<NdArrowArray> {
419    let values = a
420        .values()
421        .as_boolean_opt()
422        .ok_or_else(|| nd_err!(InvalidArgument, "logical_not needs Boolean values"))?;
423    wrap(Arc::new(boolean::not(values)?), a.dims().clone())
424}
425
426/// `true` where the value is null.
427pub fn is_null_mask(a: &NdArrowArray) -> Result<NdArrowArray> {
428    wrap(Arc::new(is_null(a.values().as_ref())?), a.dims().clone())
429}
430
431/// `true` where the value is not null.
432pub fn is_not_null_mask(a: &NdArrowArray) -> Result<NdArrowArray> {
433    wrap(
434        Arc::new(is_not_null(a.values().as_ref())?),
435        a.dims().clone(),
436    )
437}
438
439/// `-a`. Integer overflow wraps, like NumPy.
440pub fn neg(a: &NdArrowArray) -> Result<NdArrowArray> {
441    wrap(
442        numeric::neg_wrapping(a.values().as_ref())?,
443        a.dims().clone(),
444    )
445}
446
447/// `|a|`, in the input type. `abs(i8::MIN)` wraps to `i8::MIN`, like NumPy.
448pub fn abs(a: &NdArrowArray) -> Result<NdArrowArray> {
449    let v = a.values();
450    let values: ArrayRef = match v.data_type() {
451        DataType::Int8 => Arc::new(
452            v.as_primitive::<Int8Type>()
453                .unary::<_, Int8Type>(i8::wrapping_abs),
454        ),
455        DataType::Int16 => Arc::new(
456            v.as_primitive::<Int16Type>()
457                .unary::<_, Int16Type>(i16::wrapping_abs),
458        ),
459        DataType::Int32 => Arc::new(
460            v.as_primitive::<Int32Type>()
461                .unary::<_, Int32Type>(i32::wrapping_abs),
462        ),
463        DataType::Int64 => Arc::new(
464            v.as_primitive::<Int64Type>()
465                .unary::<_, Int64Type>(i64::wrapping_abs),
466        ),
467        DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => v.clone(),
468        DataType::Float32 => Arc::new(
469            v.as_primitive::<Float32Type>()
470                .unary::<_, Float32Type>(f32::abs),
471        ),
472        DataType::Float64 => Arc::new(
473            v.as_primitive::<Float64Type>()
474                .unary::<_, Float64Type>(f64::abs),
475        ),
476        other => return Err(nd_err!(NotSupported, "abs on {other}")),
477    };
478    wrap(values, a.dims().clone())
479}
480
481/// Apply a floating point function. Integers compute as `Float64`, like NumPy.
482pub fn map_float(
483    a: &NdArrowArray,
484    f32_op: impl Fn(f32) -> f32,
485    f64_op: impl Fn(f64) -> f64,
486) -> Result<NdArrowArray> {
487    let to = float_type(a.data_type())?;
488    let values = cast_values(a.values(), &to)?;
489    let out: ArrayRef = match to {
490        DataType::Float32 => Arc::new(
491            values
492                .as_primitive::<Float32Type>()
493                .unary::<_, Float32Type>(f32_op),
494        ),
495        _ => Arc::new(
496            values
497                .as_primitive::<Float64Type>()
498                .unary::<_, Float64Type>(f64_op),
499        ),
500    };
501    wrap(out, a.dims().clone())
502}
503
504fn float_predicate(
505    a: &NdArrowArray,
506    f32_op: impl Fn(f32) -> bool,
507    f64_op: impl Fn(f64) -> bool,
508) -> Result<NdArrowArray> {
509    let to = float_type(a.data_type())?;
510    let values = cast_values(a.values(), &to)?;
511    let out = match to {
512        DataType::Float32 => BooleanArray::from_unary(values.as_primitive::<Float32Type>(), f32_op),
513        _ => BooleanArray::from_unary(values.as_primitive::<Float64Type>(), f64_op),
514    };
515    wrap(Arc::new(out), a.dims().clone())
516}
517
518macro_rules! float_functions {
519    ($($(#[$doc:meta])* $name:ident => $method:ident),* $(,)?) => {$(
520        $(#[$doc])*
521        pub fn $name(a: &NdArrowArray) -> Result<NdArrowArray> {
522            map_float(a, f32::$method, f64::$method)
523        }
524    )*};
525}
526
527float_functions!(
528    /// Square root, like `numpy.sqrt`.
529    sqrt => sqrt,
530    /// Cube root, like `numpy.cbrt`.
531    cbrt => cbrt,
532    /// `e^a`, like `numpy.exp`.
533    exp => exp,
534    /// Natural logarithm, like `numpy.log`.
535    log => ln,
536    /// Base-2 logarithm, like `numpy.log2`.
537    log2 => log2,
538    /// Base-10 logarithm, like `numpy.log10`.
539    log10 => log10,
540    /// `ln(1 + a)`, like `numpy.log1p`.
541    log1p => ln_1p,
542    /// `e^a - 1`, like `numpy.expm1`.
543    expm1 => exp_m1,
544    /// Sine, like `numpy.sin`.
545    sin => sin,
546    /// Cosine, like `numpy.cos`.
547    cos => cos,
548    /// Tangent, like `numpy.tan`.
549    tan => tan,
550    /// Inverse sine, like `numpy.arcsin`.
551    arcsin => asin,
552    /// Inverse cosine, like `numpy.arccos`.
553    arccos => acos,
554    /// Inverse tangent, like `numpy.arctan`.
555    arctan => atan,
556    /// Hyperbolic sine, like `numpy.sinh`.
557    sinh => sinh,
558    /// Hyperbolic cosine, like `numpy.cosh`.
559    cosh => cosh,
560    /// Hyperbolic tangent, like `numpy.tanh`.
561    tanh => tanh,
562    /// Round down, like `numpy.floor`.
563    floor => floor,
564    /// Round up, like `numpy.ceil`.
565    ceil => ceil,
566    /// Round toward zero, like `numpy.trunc`.
567    trunc => trunc,
568    /// Round half to even, like `numpy.round` with zero decimals.
569    round => round_ties_even,
570);
571
572/// Degrees to radians, like `numpy.deg2rad`.
573pub fn deg2rad(a: &NdArrowArray) -> Result<NdArrowArray> {
574    map_float(a, f32::to_radians, f64::to_radians)
575}
576
577/// Radians to degrees, like `numpy.rad2deg`.
578pub fn rad2deg(a: &NdArrowArray) -> Result<NdArrowArray> {
579    map_float(a, f32::to_degrees, f64::to_degrees)
580}
581
582/// `true` where the value is NaN, like `numpy.isnan`. Nulls stay null.
583pub fn isnan(a: &NdArrowArray) -> Result<NdArrowArray> {
584    float_predicate(a, f32::is_nan, f64::is_nan)
585}
586
587/// `true` where the value is infinite, like `numpy.isinf`.
588pub fn isinf(a: &NdArrowArray) -> Result<NdArrowArray> {
589    float_predicate(a, f32::is_infinite, f64::is_infinite)
590}
591
592/// `true` where the value is finite, like `numpy.isfinite`.
593pub fn isfinite(a: &NdArrowArray) -> Result<NdArrowArray> {
594    float_predicate(a, f32::is_finite, f64::is_finite)
595}
596
597/// Replace NaN with null, so that NaN counts as missing in all kernels.
598pub fn nan_to_null(a: &NdArrowArray) -> Result<NdArrowArray> {
599    if !matches!(kind(a.data_type()), Some(Kind::Float(_))) {
600        return Ok(a.clone());
601    }
602    let mask = isnan(a)?;
603    let values = arrow::compute::nullif(a.values().as_ref(), mask.values().as_boolean())?;
604    Ok(wrap(values, a.dims().clone())?.with_attributes(a.attributes().clone()))
605}
606
607/// Cast the values, like `ndarray.astype`. Keeps the attributes.
608pub fn cast(a: &NdArrowArray, to: &DataType) -> Result<NdArrowArray> {
609    Ok(wrap(cast_values(a.values(), to)?, a.dims().clone())?
610        .with_attributes(a.attributes().clone()))
611}
612
613fn scalar_operand(x: &NdArrowArray, target: &Dimensions, to: &DataType) -> Result<Operand> {
614    if x.is_scalar() {
615        return Ok(Operand::Scalar(Scalar::new(cast_values(x.values(), to)?)));
616    }
617    let aligned = broadcast_to(x, target)?;
618    Ok(Operand::Array(cast_values(aligned.values(), to)?))
619}
620
621/// `cond ? x : y`, like `numpy.where`. A null condition picks `y`.
622/// Keeps the attributes of `x`.
623pub fn where_(cond: &NdArrowArray, x: &NdArrowArray, y: &NdArrowArray) -> Result<NdArrowArray> {
624    if cond.data_type() != &DataType::Boolean {
625        return Err(nd_err!(
626            InvalidArgument,
627            "where condition must be Boolean, got {}",
628            cond.data_type()
629        ));
630    }
631    let target = Dimensions::broadcast(&[cond.dims(), x.dims(), y.dims()])?;
632    let to = result_type(x.data_type(), y.data_type()).unwrap_or_else(|| x.data_type().clone());
633    let mask = broadcast_to(cond, &target)?;
634    let truthy = scalar_operand(x, &target, &to)?;
635    let falsy = scalar_operand(y, &target, &to)?;
636    let values = if target.is_scalar() {
637        let pick = if mask.values().as_boolean().is_valid(0) && mask.values().as_boolean().value(0)
638        {
639            &truthy
640        } else {
641            &falsy
642        };
643        match pick {
644            Operand::Array(a) => a.clone(),
645            Operand::Scalar(s) => s.clone().into_inner(),
646        }
647    } else {
648        zip(mask.values().as_boolean(), truthy.datum(), falsy.datum())?
649    };
650    Ok(wrap(values, target)?.with_attributes(x.attributes().clone()))
651}
652
653/// Replace nulls with `value`, a length-1 array. Keeps the attributes.
654pub fn fill_null(a: &NdArrowArray, value: &ArrayRef) -> Result<NdArrowArray> {
655    if value.len() != 1 {
656        return Err(nd_err!(InvalidArgument, "fill value must have length 1"));
657    }
658    if a.values().null_count() == 0 {
659        return Ok(a.clone());
660    }
661    let fill = Scalar::new(cast_values(value, a.data_type())?);
662    let valid = is_not_null(a.values().as_ref())?;
663    Ok(wrap(zip(&valid, a.values(), &fill)?, a.dims().clone())?
664        .with_attributes(a.attributes().clone()))
665}
666
667/// Limit values to `[min, max]`, like `numpy.clip`. Keeps the attributes.
668pub fn clip(
669    a: &NdArrowArray,
670    min: Option<&ArrayRef>,
671    max: Option<&ArrayRef>,
672) -> Result<NdArrowArray> {
673    let mut values = a.values().clone();
674    for (bound, below) in [(min, true), (max, false)] {
675        let Some(bound) = bound else { continue };
676        let bound = Scalar::new(cast_values(bound, a.data_type())?);
677        let outside = if below {
678            cmp::lt(&values, &bound)?
679        } else {
680            cmp::gt(&values, &bound)?
681        };
682        values = zip(&outside, &bound, &values)?;
683    }
684    Ok(wrap(values, a.dims().clone())?.with_attributes(a.attributes().clone()))
685}
686
687#[cfg(test)]
688mod tests {
689    use arrow::array::{Float32Array, Float64Array, Int8Array, Int32Array, UInt8Array};
690    use arrow::datatypes::{Float64Type, Int32Type};
691
692    use super::*;
693    use crate::{Attributes, IntoScalar};
694
695    fn nd(values: ArrayRef, dims: &[(&str, usize)]) -> NdArrowArray {
696        NdArrowArray::try_new(values, Dimensions::try_new(dims.iter().copied()).unwrap()).unwrap()
697    }
698
699    fn scalar(value: impl IntoScalar) -> NdArrowArray {
700        NdArrowArray::try_new_scalar(value.into_scalar()).unwrap()
701    }
702
703    fn f64s(a: &NdArrowArray) -> Vec<Option<f64>> {
704        a.values().as_primitive::<Float64Type>().iter().collect()
705    }
706
707    #[test]
708    fn result_type_follows_numpy() {
709        use DataType::*;
710        let cases = [
711            (Int8, Int16, Int16),
712            (Int8, UInt8, Int16),
713            (Int64, UInt64, Float64),
714            (UInt8, Int32, Int32),
715            (Float32, Int16, Float32),
716            (Float32, Int32, Float64),
717            (Float16, Int8, Float16),
718            (Boolean, Int8, Int8),
719            (Float32, Float64, Float64),
720        ];
721        for (a, b, expected) in cases {
722            assert_eq!(result_type(&a, &b), Some(expected.clone()), "{a} + {b}");
723            assert_eq!(result_type(&b, &a), Some(expected), "{b} + {a}");
724        }
725        assert_eq!(result_type(&Utf8, &Int8), None);
726    }
727
728    #[test]
729    fn add_broadcasts_by_name() {
730        let a = nd(Arc::new(Int32Array::from(vec![1, 2, 3])), &[("lat", 3)]);
731        let b = nd(Arc::new(Int32Array::from(vec![10, 20])), &[("time", 2)]);
732        let c = add(&a, &b).unwrap();
733        assert_eq!(c.dims().names(), vec!["lat", "time"]);
734        assert_eq!(
735            c.values().as_primitive::<Int32Type>().values(),
736            &[11, 21, 12, 22, 13, 23]
737        );
738    }
739
740    #[test]
741    fn scalar_operands_and_promotion() {
742        let a = nd(Arc::new(Int8Array::from(vec![1, 2])), &[("x", 2)]);
743        let c = mul(&a, &scalar(1.5f32)).unwrap();
744        assert_eq!(c.data_type(), &DataType::Float32);
745        let d = div(&a, &scalar(2i8)).unwrap();
746        assert_eq!(f64s(&d), vec![Some(0.5), Some(1.0)]);
747        let wrapped = add(&nd(Arc::new(Int8Array::from(vec![127])), &[]), &scalar(1i8)).unwrap();
748        assert_eq!(wrapped.values().as_primitive::<Int8Type>().value(0), -128);
749        let u = sub(&nd(Arc::new(UInt8Array::from(vec![1])), &[]), &scalar(2i8)).unwrap();
750        assert_eq!(u.data_type(), &DataType::Int16);
751    }
752
753    #[test]
754    fn nulls_propagate_and_fill() {
755        let a = nd(
756            Arc::new(Float64Array::from(vec![Some(1.0), None])),
757            &[("x", 2)],
758        );
759        let b = add(&a, &scalar(1.0f64)).unwrap();
760        assert_eq!(f64s(&b), vec![Some(2.0), None]);
761        let filled = fill_null(&b, &0.0f64.into_scalar()).unwrap();
762        assert_eq!(f64s(&filled), vec![Some(2.0), Some(0.0)]);
763    }
764
765    #[test]
766    fn comparison_and_logic() {
767        let a = nd(Arc::new(Int32Array::from(vec![1, 5, 3])), &[("x", 3)]);
768        let gt3 = gt(&a, &scalar(2i32)).unwrap();
769        let lt5 = lt(&a, &scalar(5i64)).unwrap();
770        let both = logical_and(&gt3, &lt5).unwrap();
771        let v: Vec<bool> = both
772            .values()
773            .as_boolean()
774            .iter()
775            .map(Option::unwrap)
776            .collect();
777        assert_eq!(v, vec![false, false, true]);
778        let not = logical_not(&both).unwrap();
779        assert!(not.values().as_boolean().value(0));
780        assert!(logical_and(&a, &a).is_err());
781    }
782
783    #[test]
784    fn math_functions_promote_to_float() {
785        let a = nd(Arc::new(Int32Array::from(vec![1, 4, 9])), &[("x", 3)]);
786        assert_eq!(
787            f64s(&sqrt(&a).unwrap()),
788            vec![Some(1.0), Some(2.0), Some(3.0)]
789        );
790        let f = nd(
791            Arc::new(Float32Array::from(vec![2.5, -0.5, 1.5])),
792            &[("x", 3)],
793        );
794        let r = round(&f).unwrap();
795        assert_eq!(
796            r.values().as_primitive::<Float32Type>().values(),
797            &[2.0, -0.0, 2.0]
798        );
799        let ab = abs(&nd(Arc::new(Int32Array::from(vec![-3, 2])), &[("x", 2)])).unwrap();
800        assert_eq!(ab.values().as_primitive::<Int32Type>().values(), &[3, 2]);
801    }
802
803    #[test]
804    fn where_clip_and_nan() {
805        let a = nd(
806            Arc::new(Float64Array::from(vec![1.0, f64::NAN, 7.0])),
807            &[("x", 3)],
808        );
809        let a = a.with_attributes(Attributes::new().with("units", "K"));
810        let cond = gt(&a, &scalar(2.0f64)).unwrap();
811        let w = where_(&cond, &a, &scalar(0.0f64)).unwrap();
812        assert_eq!(w.attributes().get_str("units"), Some("K"));
813        let v = f64s(&w);
814        assert_eq!(v[0], Some(0.0));
815        assert_eq!(v[2], Some(7.0));
816
817        let c = clip(&a, Some(&2.0f64.into_scalar()), Some(&5.0f64.into_scalar())).unwrap();
818        assert_eq!(f64s(&c)[0], Some(2.0));
819        assert_eq!(f64s(&c)[2], Some(5.0));
820
821        let n = nan_to_null(&a).unwrap();
822        assert_eq!(n.values().null_count(), 1);
823    }
824
825    #[test]
826    fn pow_broadcasts() {
827        let a = nd(Arc::new(Int32Array::from(vec![2, 3])), &[("x", 2)]);
828        let p = pow(&a, &scalar(2i32)).unwrap();
829        assert_eq!(f64s(&p), vec![Some(4.0), Some(9.0)]);
830    }
831}
nd_arrow_array/compute/elementwise.rs

nd_arrow_array/compute/
elementwise.rs
