1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
use vortex_dtype::{NativePType, PType};
use vortex_error::VortexResult;

use crate::iter::Batch;
use crate::{Array, ArrayDType};

pub trait BinaryFn {
    fn binary<I: NativePType, U: NativePType, O: NativePType, F: Fn(I, U) -> O>(
        &self,
        rhs: Array,
        binary_fn: F,
    ) -> VortexResult<Array>;
}

pub trait UnaryFn {
    fn unary<I: NativePType, O: NativePType, F: Fn(I) -> O>(
        &self,
        unary_fn: F,
    ) -> VortexResult<Array>;
}

#[allow(clippy::unwrap_used)]
pub fn dyn_cast_array_iter<N: NativePType>(array: &Array) -> Box<dyn Iterator<Item = Batch<N>>> {
    match PType::try_from(array.dtype()).unwrap() {
        PType::U8 => Box::new(
            array
                .with_dyn(|a| a.as_primitive_array_unchecked().u8_iter())
                .unwrap()
                .map(|b| b.as_::<N>()),
        ),
        PType::U16 => Box::new(
            array
                .with_dyn(|a| a.as_primitive_array_unchecked().u16_iter())
                .unwrap()
                .map(|b| b.as_::<N>()),
        ),
        PType::U32 => Box::new(
            array
                .with_dyn(|a| a.as_primitive_array_unchecked().u32_iter())
                .unwrap()
                .map(|b| b.as_::<N>()),
        ),
        PType::U64 => Box::new(
            array
                .with_dyn(|a| a.as_primitive_array_unchecked().u64_iter())
                .unwrap()
                .map(|b| b.as_::<N>()),
        ),
        PType::I8 => Box::new(
            array
                .with_dyn(|a| a.as_primitive_array_unchecked().i8_iter())
                .unwrap()
                .map(|b| b.as_::<N>()),
        ),
        PType::I16 => Box::new(
            array
                .with_dyn(|a| a.as_primitive_array_unchecked().i16_iter())
                .unwrap()
                .map(|b| b.as_::<N>()),
        ),
        PType::I32 => Box::new(
            array
                .with_dyn(|a| a.as_primitive_array_unchecked().i32_iter())
                .unwrap()
                .map(|b| b.as_::<N>()),
        ),
        PType::I64 => Box::new(
            array
                .with_dyn(|a| a.as_primitive_array_unchecked().i64_iter())
                .unwrap()
                .map(|b| b.as_::<N>()),
        ),
        PType::F16 => Box::new(
            array
                .with_dyn(|a| a.as_primitive_array_unchecked().u64_iter())
                .unwrap()
                .map(|b| b.as_::<N>()),
        ),
        PType::F32 => Box::new(
            array
                .with_dyn(|a| a.as_primitive_array_unchecked().f32_iter())
                .unwrap()
                .map(|b| b.as_::<N>()),
        ),
        PType::F64 => Box::new(
            array
                .with_dyn(|a| a.as_primitive_array_unchecked().f64_iter())
                .unwrap()
                .map(|b| b.as_::<N>()),
        ),
    }
}