vortex_array/expr/functions/
execution.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! This module defines the API for vectorized execution within Vortex.
5// TODO(ngates): these definitions should be lifted out of the functions module.
6
7use vortex_dtype::DType;
8use vortex_vector::Datum;
9use vortex_vector::VectorOps;
10
11/// Arguments provided when executing an object using vectorized execution.
12pub struct ExecutionArgs {
13    /// The number of rows to be processed.
14    row_count: usize,
15    /// The expected return dtype of the execution.
16    return_dtype: DType,
17    /// The input data types.
18    input_dtypes: Vec<DType>,
19    /// The input datums.
20    input_datums: Vec<Datum>,
21}
22
23impl ExecutionArgs {
24    pub fn new(
25        row_count: usize,
26        return_dtype: DType,
27        input_types: Vec<DType>,
28        input_datums: Vec<impl Into<Datum>>,
29    ) -> Self {
30        let input_datums: Vec<Datum> = input_datums.into_iter().map(|d| d.into()).collect();
31        assert!(
32            input_datums
33                .iter()
34                .all(|d| d.as_vector().is_none_or(|v| v.len() == row_count)),
35            "All input vectors must have the same length as row_count"
36        );
37        Self {
38            row_count,
39            return_dtype,
40            input_dtypes: input_types,
41            input_datums,
42        }
43    }
44
45    pub fn row_count(&self) -> usize {
46        self.row_count
47    }
48
49    pub fn return_dtype(&self) -> &DType {
50        &self.return_dtype
51    }
52
53    pub fn input_type(&self, idx: usize) -> &DType {
54        &self.input_dtypes[idx]
55    }
56
57    pub fn input_datums(&self, idx: usize) -> &Datum {
58        &self.input_datums[idx]
59    }
60}