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

vortex_array/
context.rs
