vortex_array/arrays/scalar_fn/vtable/
canonical.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_error::VortexExpect;
6use vortex_vector::Datum;
7
8use crate::Array;
9use crate::Canonical;
10use crate::arrays::scalar_fn::array::ScalarFnArray;
11use crate::arrays::scalar_fn::vtable::SCALAR_FN_SESSION;
12use crate::arrays::scalar_fn::vtable::ScalarFnVTable;
13use crate::expr::functions::ExecutionArgs;
14use crate::vectors::VectorIntoArray;
15use crate::vtable::CanonicalVTable;
16
17impl CanonicalVTable<ScalarFnVTable> for ScalarFnVTable {
18    fn canonicalize(array: &ScalarFnArray) -> Canonical {
19        let child_dtypes: Vec<_> = array.children.iter().map(|c| c.dtype().clone()).collect();
20        let child_datums: Vec<_> = array
21            .children()
22            .iter()
23            // TODO(ngates): we could make all execution operate over datums
24            .map(|child| child.execute(&SCALAR_FN_SESSION).map(Datum::Vector))
25            .try_collect()
26            // FIXME(ngates): canonicalizing really ought to be fallible
27            .vortex_expect(
28                "Failed to execute child array during canonicalization of ScalarFnArray",
29            );
30
31        let ctx = ExecutionArgs::new(array.len, array.dtype.clone(), child_dtypes, child_datums);
32
33        let result_vector = array
34            .scalar_fn
35            .execute(&ctx)
36            .vortex_expect("Canonicalize should be fallible")
37            .into_vector()
38            .vortex_expect("Canonicalize should return a vector");
39
40        result_vector.into_array(&array.dtype).to_canonical()
41    }
42}