• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

vortex-data / vortex / 16273663911

14 Jul 2025 05:31PM UTC coverage: 81.554% (+0.02%) from 81.537%
16273663911

Pull #3871

github

web-flow
Merge aeb0f7b4d into 9b0d852a1
Pull Request #3871: chore: Fewer ways to invoke display

83 of 86 new or added lines in 8 files covered. (96.51%)

6 existing lines in 2 files now uncovered.

46320 of 56797 relevant lines covered (81.55%)

146672.82 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

70.11
/vortex-array/src/compute/mod.rs
1
// SPDX-License-Identifier: Apache-2.0
2
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3

4
//! Compute kernels on top of Vortex Arrays.
5
//!
6
//! We aim to provide a basic set of compute kernels that can be used to efficiently index, slice,
7
//! and filter Vortex Arrays in their encoded forms.
8
//!
9
//! Every array encoding has the ability to implement their own efficient implementations of these
10
//! operators, else we will decode, and perform the equivalent operator from Arrow.
11

12
use std::any::{Any, type_name};
13
use std::fmt::{Debug, Formatter};
14

15
use arcref::ArcRef;
16
pub use between::*;
17
pub use boolean::*;
18
pub use cast::*;
19
pub use compare::*;
20
pub use fill_null::*;
21
pub use filter::*;
22
pub use invert::*;
23
pub use is_constant::*;
24
pub use is_sorted::*;
25
use itertools::Itertools;
26
pub use like::*;
27
pub use list_contains::*;
28
pub use mask::*;
29
pub use min_max::*;
30
pub use nan_count::*;
31
pub use numeric::*;
32
use parking_lot::RwLock;
33
pub use sum::*;
34
pub use take::*;
35
use vortex_dtype::DType;
36
use vortex_error::{VortexError, VortexResult, vortex_bail, vortex_err};
37
use vortex_mask::Mask;
38
use vortex_scalar::Scalar;
39

40
use crate::builders::ArrayBuilder;
41
use crate::display::DisplayOptions;
42
use crate::{Array, ArrayRef};
43

44
#[cfg(feature = "arbitrary")]
45
mod arbitrary;
46
mod between;
47
mod boolean;
48
mod cast;
49
mod compare;
50
#[cfg(feature = "test-harness")]
51
pub mod conformance;
52
mod fill_null;
53
mod filter;
54
mod invert;
55
mod is_constant;
56
mod is_sorted;
57
mod like;
58
mod list_contains;
59
mod mask;
60
mod min_max;
61
mod nan_count;
62
mod numeric;
63
mod sum;
64
mod take;
65

66
/// An instance of a compute function holding the implementation vtable and a set of registered
67
/// compute kernels.
68
pub struct ComputeFn {
69
    id: ArcRef<str>,
70
    vtable: ArcRef<dyn ComputeFnVTable>,
71
    kernels: RwLock<Vec<ArcRef<dyn Kernel>>>,
72
}
73

74
impl ComputeFn {
75
    /// Create a new compute function from the given [`ComputeFnVTable`].
76
    pub fn new(id: ArcRef<str>, vtable: ArcRef<dyn ComputeFnVTable>) -> Self {
35,418✔
77
        Self {
35,418✔
78
            id,
35,418✔
79
            vtable,
35,418✔
80
            kernels: Default::default(),
35,418✔
81
        }
35,418✔
82
    }
35,418✔
83

84
    /// Returns the string identifier of the compute function.
85
    pub fn id(&self) -> &ArcRef<str> {
×
86
        &self.id
×
87
    }
×
88

89
    /// Register a kernel for the compute function.
90
    pub fn register_kernel(&self, kernel: ArcRef<dyn Kernel>) {
308,607✔
91
        self.kernels.write().push(kernel);
308,607✔
92
    }
308,607✔
93

94
    /// Invokes the compute function with the given arguments.
95
    pub fn invoke(&self, args: &InvocationArgs) -> VortexResult<Output> {
377,120✔
96
        // Perform some pre-condition checks against the arguments and the function properties.
377,120✔
97
        if self.is_elementwise() {
377,120✔
98
            // For element-wise functions, all input arrays must be the same length.
99
            if !args
54,830✔
100
                .inputs
54,830✔
101
                .iter()
54,830✔
102
                .filter_map(|input| input.array())
99,513✔
103
                .map(|array| array.len())
73,141✔
104
                .all_equal()
54,830✔
105
            {
106
                vortex_bail!(
×
107
                    "Compute function {} is elementwise but input arrays have different lengths",
×
108
                    self.id
×
109
                );
×
110
            }
54,830✔
111
        }
322,290✔
112

113
        let expected_dtype = self.vtable.return_dtype(args)?;
377,120✔
114
        let expected_len = self.vtable.return_len(args)?;
377,119✔
115

116
        let output = self.vtable.invoke(args, &self.kernels.read())?;
377,119✔
117

118
        if output.dtype() != &expected_dtype {
375,628✔
119
            vortex_bail!(
×
120
                "Internal error: compute function {} returned a result of type {} but expected {}\n{}",
×
121
                self.id,
×
122
                output.dtype(),
×
123
                &expected_dtype,
×
124
                args.inputs
×
125
                    .iter()
×
126
                    .filter_map(|input| input.array())
×
NEW
127
                    .format_with(",", |array, f| f(
×
NEW
128
                        &array.display_as(DisplayOptions::TreeDisplay)
×
NEW
129
                    ))
×
UNCOV
130
            );
×
131
        }
375,628✔
132
        if output.len() != expected_len {
375,628✔
133
            vortex_bail!(
×
134
                "Internal error: compute function {} returned a result of length {} but expected {}",
×
135
                self.id,
×
136
                output.len(),
×
137
                expected_len
×
138
            );
×
139
        }
375,628✔
140

375,628✔
141
        Ok(output)
375,628✔
142
    }
377,120✔
143

144
    /// Compute the return type of the function given the input arguments.
145
    pub fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType> {
15✔
146
        self.vtable.return_dtype(args)
15✔
147
    }
15✔
148

149
    /// Compute the return length of the function given the input arguments.
150
    pub fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize> {
×
151
        self.vtable.return_len(args)
×
152
    }
×
153

154
    /// Returns whether the compute function is elementwise, i.e. the output is the same shape as
155
    pub fn is_elementwise(&self) -> bool {
377,150✔
156
        // TODO(ngates): should this just be a constant passed in the constructor?
377,150✔
157
        self.vtable.is_elementwise()
377,150✔
158
    }
377,150✔
159

160
    /// Returns the compute function's kernels.
161
    pub fn kernels(&self) -> Vec<ArcRef<dyn Kernel>> {
4,453✔
162
        self.kernels.read().to_vec()
4,453✔
163
    }
4,453✔
164
}
165

166
/// VTable for the implementation of a compute function.
167
pub trait ComputeFnVTable: 'static + Send + Sync {
168
    /// Invokes the compute function entry-point with the given input arguments and options.
169
    ///
170
    /// The entry-point logic can short-circuit compute using statistics, update result array
171
    /// statistics, search for relevant compute kernels, and canonicalize the inputs in order
172
    /// to successfully compute a result.
173
    fn invoke(&self, args: &InvocationArgs, kernels: &[ArcRef<dyn Kernel>])
174
    -> VortexResult<Output>;
175

176
    /// Computes the return type of the function given the input arguments.
177
    ///
178
    /// All kernel implementations will be validated to return the [`DType`] as computed here.
179
    fn return_dtype(&self, args: &InvocationArgs) -> VortexResult<DType>;
180

181
    /// Computes the return length of the function given the input arguments.
182
    ///
183
    /// All kernel implementations will be validated to return the len as computed here.
184
    /// Scalars are considered to have length 1.
185
    fn return_len(&self, args: &InvocationArgs) -> VortexResult<usize>;
186

187
    /// Returns whether the function operates elementwise, i.e. the output is the same shape as the
188
    /// input and no information is shared between elements.
189
    ///
190
    /// Examples include `add`, `subtract`, `and`, `cast`, `fill_null` etc.
191
    /// Examples that are not elementwise include `sum`, `count`, `min`, `fill_forward` etc.
192
    ///
193
    /// All input arrays to an elementwise function *must* have the same length.
194
    fn is_elementwise(&self) -> bool;
195
}
196

197
/// Arguments to a compute function invocation.
198
#[derive(Clone)]
199
pub struct InvocationArgs<'a> {
200
    pub inputs: &'a [Input<'a>],
201
    pub options: &'a dyn Options,
202
}
203

204
/// For unary compute functions, it's useful to just have this short-cut.
205
pub struct UnaryArgs<'a, O: Options> {
206
    pub array: &'a dyn Array,
207
    pub options: &'a O,
208
}
209

210
impl<'a, O: Options> TryFrom<&InvocationArgs<'a>> for UnaryArgs<'a, O> {
211
    type Error = VortexError;
212

213
    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
426,348✔
214
        if value.inputs.len() != 1 {
426,348✔
215
            vortex_bail!("Expected 1 input, found {}", value.inputs.len());
×
216
        }
426,348✔
217
        let array = value.inputs[0]
426,348✔
218
            .array()
426,348✔
219
            .ok_or_else(|| vortex_err!("Expected input 0 to be an array"))?;
426,348✔
220
        let options =
426,348✔
221
            value.options.as_any().downcast_ref::<O>().ok_or_else(|| {
426,348✔
222
                vortex_err!("Expected options to be of type {}", type_name::<O>())
×
223
            })?;
426,348✔
224
        Ok(UnaryArgs { array, options })
426,348✔
225
    }
426,348✔
226
}
227

228
/// For binary compute functions, it's useful to just have this short-cut.
229
pub struct BinaryArgs<'a, O: Options> {
230
    pub lhs: &'a dyn Array,
231
    pub rhs: &'a dyn Array,
232
    pub options: &'a O,
233
}
234

235
impl<'a, O: Options> TryFrom<&InvocationArgs<'a>> for BinaryArgs<'a, O> {
236
    type Error = VortexError;
237

238
    fn try_from(value: &InvocationArgs<'a>) -> Result<Self, Self::Error> {
1,101✔
239
        if value.inputs.len() != 2 {
1,101✔
240
            vortex_bail!("Expected 2 input, found {}", value.inputs.len());
×
241
        }
1,101✔
242
        let lhs = value.inputs[0]
1,101✔
243
            .array()
1,101✔
244
            .ok_or_else(|| vortex_err!("Expected input 0 to be an array"))?;
1,101✔
245
        let rhs = value.inputs[1]
1,101✔
246
            .array()
1,101✔
247
            .ok_or_else(|| vortex_err!("Expected input 1 to be an array"))?;
1,101✔
248
        let options =
1,101✔
249
            value.options.as_any().downcast_ref::<O>().ok_or_else(|| {
1,101✔
250
                vortex_err!("Expected options to be of type {}", type_name::<O>())
×
251
            })?;
1,101✔
252
        Ok(BinaryArgs { lhs, rhs, options })
1,101✔
253
    }
1,101✔
254
}
255

256
/// Input to a compute function.
257
pub enum Input<'a> {
258
    Scalar(&'a Scalar),
259
    Array(&'a dyn Array),
260
    Mask(&'a Mask),
261
    Builder(&'a mut dyn ArrayBuilder),
262
    DType(&'a DType),
263
}
264

265
impl Debug for Input<'_> {
266
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
×
267
        let mut f = f.debug_struct("Input");
×
268
        match self {
×
269
            Input::Scalar(scalar) => f.field("Scalar", scalar),
×
270
            Input::Array(array) => f.field("Array", array),
×
271
            Input::Mask(mask) => f.field("Mask", mask),
×
272
            Input::Builder(builder) => f.field("Builder", &builder.len()),
×
273
            Input::DType(dtype) => f.field("DType", dtype),
×
274
        };
275
        f.finish()
×
276
    }
×
277
}
278

279
impl<'a> From<&'a dyn Array> for Input<'a> {
280
    fn from(value: &'a dyn Array) -> Self {
573,929✔
281
        Input::Array(value)
573,929✔
282
    }
573,929✔
283
}
284

285
impl<'a> From<&'a Scalar> for Input<'a> {
286
    fn from(value: &'a Scalar) -> Self {
4,302✔
287
        Input::Scalar(value)
4,302✔
288
    }
4,302✔
289
}
290

291
impl<'a> From<&'a Mask> for Input<'a> {
292
    fn from(value: &'a Mask) -> Self {
14,632✔
293
        Input::Mask(value)
14,632✔
294
    }
14,632✔
295
}
296

297
impl<'a> From<&'a DType> for Input<'a> {
298
    fn from(value: &'a DType) -> Self {
19,552✔
299
        Input::DType(value)
19,552✔
300
    }
19,552✔
301
}
302

303
impl<'a> Input<'a> {
304
    pub fn scalar(&self) -> Option<&'a Scalar> {
14,272✔
305
        match self {
14,272✔
306
            Input::Scalar(scalar) => Some(*scalar),
14,272✔
307
            _ => None,
×
308
        }
309
    }
14,272✔
310

311
    pub fn array(&self) -> Option<&'a dyn Array> {
2,390,157✔
312
        match self {
2,390,157✔
313
            Input::Array(array) => Some(*array),
2,363,785✔
314
            _ => None,
26,372✔
315
        }
316
    }
2,390,157✔
317

318
    pub fn mask(&self) -> Option<&'a Mask> {
188,543✔
319
        match self {
188,543✔
320
            Input::Mask(mask) => Some(*mask),
188,543✔
321
            _ => None,
×
322
        }
323
    }
188,543✔
324

325
    pub fn builder(&'a mut self) -> Option<&'a mut dyn ArrayBuilder> {
×
326
        match self {
×
327
            Input::Builder(builder) => Some(*builder),
×
328
            _ => None,
×
329
        }
330
    }
×
331

332
    pub fn dtype(&self) -> Option<&'a DType> {
128,626✔
333
        match self {
128,626✔
334
            Input::DType(dtype) => Some(*dtype),
128,626✔
335
            _ => None,
×
336
        }
337
    }
128,626✔
338
}
339

340
/// Output from a compute function.
341
#[derive(Debug)]
342
pub enum Output {
343
    Scalar(Scalar),
344
    Array(ArrayRef),
345
}
346

347
#[allow(clippy::len_without_is_empty)]
348
impl Output {
349
    pub fn dtype(&self) -> &DType {
375,628✔
350
        match self {
375,628✔
351
            Output::Scalar(scalar) => scalar.dtype(),
280,205✔
352
            Output::Array(array) => array.dtype(),
95,423✔
353
        }
354
    }
375,628✔
355

356
    pub fn len(&self) -> usize {
375,628✔
357
        match self {
375,628✔
358
            Output::Scalar(_) => 1,
280,205✔
359
            Output::Array(array) => array.len(),
95,423✔
360
        }
361
    }
375,628✔
362

363
    pub fn unwrap_scalar(self) -> VortexResult<Scalar> {
418,615✔
364
        match self {
418,615✔
365
            Output::Array(_) => vortex_bail!("Expected array output, got Array"),
×
366
            Output::Scalar(scalar) => Ok(scalar),
418,615✔
367
        }
368
    }
418,615✔
369

370
    pub fn unwrap_array(self) -> VortexResult<ArrayRef> {
99,828✔
371
        match self {
99,828✔
372
            Output::Array(array) => Ok(array),
99,828✔
373
            Output::Scalar(_) => vortex_bail!("Expected array output, got Scalar"),
×
374
        }
375
    }
99,828✔
376
}
377

378
impl From<ArrayRef> for Output {
379
    fn from(value: ArrayRef) -> Self {
99,523✔
380
        Output::Array(value)
99,523✔
381
    }
99,523✔
382
}
383

384
impl From<Scalar> for Output {
385
    fn from(value: Scalar) -> Self {
418,615✔
386
        Output::Scalar(value)
418,615✔
387
    }
418,615✔
388
}
389

390
/// Options for a compute function invocation.
391
pub trait Options: 'static {
392
    fn as_any(&self) -> &dyn Any;
393
}
394

395
impl Options for () {
396
    fn as_any(&self) -> &dyn Any {
427,449✔
397
        self
427,449✔
398
    }
427,449✔
399
}
400

401
/// Compute functions can ask arrays for compute kernels for a given invocation.
402
///
403
/// The kernel is invoked with the input arguments and options, and can return `None` if it is
404
/// unable to compute the result for the given inputs due to missing implementation logic.
405
/// For example, if kernel doesn't support the `LTE` operator. By returning `None`, the kernel
406
/// is indicating that it cannot compute the result for the given inputs, and another kernel should
407
/// be tried. *Not* that the given inputs are invalid for the compute function.
408
///
409
/// If the kernel fails to compute a result, it should return a `Some` with the error.
410
pub trait Kernel: 'static + Send + Sync + Debug {
411
    /// Invokes the kernel with the given input arguments and options.
412
    fn invoke(&self, args: &InvocationArgs) -> VortexResult<Option<Output>>;
413
}
414

415
/// Register a kernel for a compute function.
416
/// See each compute function for the correct type of kernel to register.
417
#[macro_export]
418
macro_rules! register_kernel {
419
    ($T:expr) => {
420
        $crate::aliases::inventory::submit!($T);
421
    };
422
}
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc