vortex_array/
normalize.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use itertools::Itertools;
5use vortex_error::VortexResult;
6use vortex_error::vortex_bail;
7use vortex_session::registry::Id;
8
9use crate::Array;
10use crate::ArrayRef;
11use crate::ArrayVisitor;
12use crate::session::ArrayRegistry;
13
14/// Options for normalizing an array.
15pub struct NormalizeOptions<'a> {
16    /// The set of allowed array encodings (in addition to the canonical ones) that are permitted
17    /// in the normalized array.
18    pub allowed: &'a ArrayRegistry,
19    /// The operation to perform when a non-allowed encoding is encountered.
20    pub operation: Operation,
21}
22
23/// The operation to perform when a non-allowed encoding is encountered.
24pub enum Operation {
25    Error,
26    // TODO(joe): add into canonical variant
27}
28
29impl dyn Array + '_ {
30    /// Normalize the array according to given options.
31    ///
32    /// This operation performs a recursive traversal of the array. Any non-allowed encoding is
33    /// normalized per the configured operation.
34    pub fn normalize(self: ArrayRef, options: &mut NormalizeOptions) -> VortexResult<ArrayRef> {
35        let array_ids = options.allowed.ids().collect_vec();
36        self.normalize_with_error(&array_ids)?;
37        // Note this takes ownership so we can at a later date remove non-allowed encodings.
38        Ok(self)
39    }
40
41    fn normalize_with_error(self: &ArrayRef, allowed: &[Id]) -> VortexResult<()> {
42        if !allowed.contains(&self.encoding_id()) {
43            vortex_bail!(AssertionFailed: "normalize forbids encoding ({})", self.encoding_id())
44        }
45
46        for child in ArrayVisitor::children(self) {
47            let child: ArrayRef = child;
48            child.normalize_with_error(allowed)?
49        }
50        Ok(())
51    }
52}
vortex_array/normalize.rs

vortex_array/
normalize.rs
