vortex_array/
context.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::sync::Arc;
6
7use itertools::Itertools;
8use parking_lot::RwLock;
9use vortex_error::VortexExpect;
10use vortex_error::VortexResult;
11use vortex_error::vortex_bail;
12use vortex_error::vortex_err;
13use vortex_session::registry::Registry;
14
15use crate::vtable::ArrayVTable;
16
17pub type ArrayContext = VTableContext<ArrayVTable>;
18
19/// A collection of encodings that can be addressed by a u16 positional index.
20/// This is used to map array encodings and layout encodings when reading from a file.
21#[derive(Debug, Clone)]
22pub struct VTableContext<T>(Arc<RwLock<Vec<T>>>);
23
24impl<T: Clone + Eq> VTableContext<T> {
25    pub fn new(encodings: Vec<T>) -> Self {
26        Self(Arc::new(RwLock::new(encodings)))
27    }
28
29    pub fn try_from_registry<'a>(
30        registry: &Registry<T>,
31        ids: impl IntoIterator<Item = &'a str>,
32    ) -> VortexResult<Self>
33    where
34        T: Display,
35    {
36        let items: Vec<T> = ids
37            .into_iter()
38            .map(|id| {
39                registry
40                    .find(id)
41                    .ok_or_else(|| vortex_err!("Registry missing encoding with id {}", id))
42            })
43            .try_collect()?;
44        if items.len() > u16::MAX as usize {
45            vortex_bail!(
46                "Cannot create VTableContext: registry has more than u16::MAX ({}) items",
47                u16::MAX
48            );
49        }
50        Ok(Self::new(items))
51    }
52
53    pub fn empty() -> Self {
54        Self(Arc::new(RwLock::new(Vec::new())))
55    }
56
57    pub fn with(self, encoding: T) -> Self {
58        {
59            let mut write = self.0.write();
60            if write.iter().all(|e| e != &encoding) {
61                write.push(encoding);
62            }
63        }
64        self
65    }
66
67    pub fn with_many<E: IntoIterator<Item = T>>(self, items: E) -> Self {
68        items.into_iter().fold(self, |ctx, e| ctx.with(e))
69    }
70
71    pub fn encodings(&self) -> Vec<T> {
72        self.0.read().clone()
73    }
74
75    /// Returns the index of the encoding in the context, or adds it if it doesn't exist.
76    ///
77    /// At write time the order encodings are registered by this method can change.
78    /// See [File Format specification](https://docs.vortex.rs/specs/file-format#file-determinism-and-reproducibility)
79    /// for more details.
80    pub fn encoding_idx(&self, encoding: &T) -> u16 {
81        let mut write = self.0.write();
82        if let Some(idx) = write.iter().position(|e| e == encoding) {
83            return u16::try_from(idx).vortex_expect("Cannot have more than u16::MAX encodings");
84        }
85        assert!(
86            write.len() < u16::MAX as usize,
87            "Cannot have more than u16::MAX encodings"
88        );
89        write.push(encoding.clone());
90        u16::try_from(write.len() - 1).vortex_expect("checked already")
91    }
92
93    /// Find an encoding by its position.
94    pub fn lookup_encoding(&self, idx: u16) -> Option<T> {
95        self.0.read().get(idx as usize).cloned()
96    }
97}
vortex_array/context.rs

vortex_array/
context.rs
