vortex_session/
registry.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::cmp::Ordering;
5use std::fmt;
6use std::fmt::Debug;
7use std::fmt::Display;
8use std::fmt::Formatter;
9use std::hash::Hash;
10use std::ops::Deref;
11use std::sync::Arc;
12use std::sync::LazyLock;
13use std::sync::OnceLock;
14
15use lasso::Spur;
16use lasso::ThreadedRodeo;
17use parking_lot::RwLock;
18use vortex_error::VortexExpect;
19use vortex_utils::aliases::DefaultHashBuilder;
20use vortex_utils::aliases::hash_set::HashSet;
21
22/// Global string interner for [`Id`] values.
23static INTERNER: LazyLock<ThreadedRodeo<Spur, DefaultHashBuilder>> =
24    LazyLock::new(|| ThreadedRodeo::with_hasher(DefaultHashBuilder::default()));
25
26/// A lightweight, copyable identifier backed by a global string interner.
27///
28/// Used for array encoding IDs, scalar function IDs, layout IDs, and similar
29/// globally-unique string identifiers throughout Vortex. Equality and hashing
30/// are O(1) symbol comparisons.
31#[derive(Clone, Copy, PartialEq, Eq, Hash)]
32pub struct Id(Spur);
33
34impl Id {
35    /// Intern a string and return its `Id`.
36    pub fn new(s: &str) -> Self {
37        Self(INTERNER.get_or_intern(s))
38    }
39
40    /// Intern a string and return its `Id`.
41    pub fn new_static(s: &'static str) -> Self {
42        Self(INTERNER.get_or_intern_static(s))
43    }
44
45    /// Returns the interned string.
46    pub fn as_str(&self) -> &str {
47        let s = INTERNER.resolve(&self.0);
48        // SAFETY: INTERNER is 'static and its arena is append-only, so resolved string
49        // pointers are stable for the lifetime of the program.
50        unsafe { &*(s as *const str) }
51    }
52}
53
54impl From<&str> for Id {
55    #[expect(clippy::disallowed_methods, reason = "interning a dynamic id")]
56    fn from(s: &str) -> Self {
57        Self::new(s)
58    }
59}
60
61impl Display for Id {
62    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
63        f.write_str(self.as_str())
64    }
65}
66
67impl Debug for Id {
68    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
69        write!(f, "Id(\"{}\")", self.as_str())
70    }
71}
72
73impl PartialOrd for Id {
74    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
75        Some(self.cmp(other))
76    }
77}
78
79impl Ord for Id {
80    fn cmp(&self, other: &Self) -> Ordering {
81        self.as_str().cmp(other.as_str())
82    }
83}
84
85impl AsRef<str> for Id {
86    fn as_ref(&self) -> &str {
87        self.as_str()
88    }
89}
90
91impl PartialEq<&Id> for Id {
92    fn eq(&self, other: &&Id) -> bool {
93        self == *other
94    }
95}
96
97impl PartialEq<Id> for &Id {
98    fn eq(&self, other: &Id) -> bool {
99        *self == other
100    }
101}
102
103/// A lazily-initialized, cached [`Id`] for use as a `static`.
104///
105/// Avoids repeated interner write-lock acquisition by storing the interned [`Id`]
106/// on first access and returning the cached copy on all subsequent calls.
107///
108/// # Example
109///
110/// ```
111/// use vortex_session::registry::{CachedId, Id};
112///
113/// static MY_ID: CachedId = CachedId::new("my.encoding");
114///
115/// fn get_id() -> Id {
116///     *MY_ID
117/// }
118/// ```
119pub struct CachedId {
120    s: &'static str,
121    cached: OnceLock<Id>,
122}
123
124impl CachedId {
125    /// Create a new `CachedId` that will intern `s` on first access.
126    pub const fn new(s: &'static str) -> Self {
127        Self {
128            s,
129            cached: OnceLock::new(),
130        }
131    }
132}
133
134impl Deref for CachedId {
135    type Target = Id;
136
137    #[expect(
138        clippy::disallowed_methods,
139        reason = "CachedId interns its static id once here"
140    )]
141    fn deref(&self) -> &Id {
142        self.cached.get_or_init(|| Id::new_static(self.s))
143    }
144}
145
146/// A [`ReadContext`] holds a set of interned IDs for use during deserialization, mapping
147/// u16 indices to IDs.
148#[derive(Clone, Debug)]
149pub struct ReadContext {
150    ids: Arc<[Id]>,
151}
152
153impl ReadContext {
154    /// Create a context with the given initial IDs.
155    pub fn new(ids: impl Into<Arc<[Id]>>) -> Self {
156        Self { ids: ids.into() }
157    }
158
159    /// Resolve an interned ID by its index.
160    pub fn resolve(&self, idx: u16) -> Option<Id> {
161        self.ids.get(idx as usize).cloned()
162    }
163
164    pub fn ids(&self) -> &[Id] {
165        &self.ids
166    }
167}
168
169/// An [`Interner`] holds a set of interned IDs for use during serialization/deserialization,
170/// mapping IDs to u16 indices.
171///
172/// ## Upcoming Changes
173///
174/// This object holds an Arc of RwLock internally because we need concurrent access from the
175/// layout writer code path. We should update SegmentSink to take an Array rather than
176/// ByteBuffer such that serializing arrays is done sequentially.
177#[derive(Clone, Debug, Default)]
178pub struct Interner {
179    // TODO(ngates): it's a long story, but if we make SegmentSink and SegmentSource take an
180    //  enum of Segment { Array, DType, Buffer } then we don't actually need a mutable context
181    //  in the LayoutWriter, therefore we don't need a RwLock here and everyone is happier.
182    ids: Arc<RwLock<Vec<Id>>>,
183    // Optional set of permissible IDs; when present, only these may be interned.
184    allowed: Option<Arc<HashSet<Id>>>,
185}
186
187impl Interner {
188    /// Create an interner with the given initial IDs.
189    pub fn new(ids: Vec<Id>) -> Self {
190        Self {
191            ids: Arc::new(RwLock::new(ids)),
192            allowed: None,
193        }
194    }
195
196    /// Create an empty interner.
197    pub fn empty() -> Self {
198        Self::default()
199    }
200
201    /// Restrict the permissible set of interned IDs to `allowed`.
202    ///
203    /// The set is snapshotted at this call: IDs registered elsewhere afterwards are not
204    /// permitted.
205    pub fn with_allowed_ids(mut self, allowed: HashSet<Id>) -> Self {
206        self.allowed = Some(Arc::new(allowed));
207        self
208    }
209
210    /// Intern an ID, returning its index.
211    pub fn intern(&self, id: &Id) -> Option<u16> {
212        if let Some(allowed) = &self.allowed
213            && !allowed.contains(id)
214        {
215            // ID not permitted, cannot intern.
216            return None;
217        }
218
219        let mut ids = self.ids.write();
220        if let Some(idx) = ids.iter().position(|e| e == id) {
221            return Some(u16::try_from(idx).vortex_expect("Cannot have more than u16::MAX items"));
222        }
223
224        let idx = ids.len();
225        assert!(
226            idx < u16::MAX as usize,
227            "Cannot have more than u16::MAX items"
228        );
229        ids.push(*id);
230        Some(u16::try_from(idx).vortex_expect("checked already"))
231    }
232
233    /// Get the list of interned IDs.
234    pub fn to_ids(&self) -> Vec<Id> {
235        self.ids.read().clone()
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use vortex_utils::aliases::hash_set::HashSet;
242
243    use super::CachedId;
244    use super::Interner;
245
246    static VALID: CachedId = CachedId::new("vortex.test.valid");
247    static INVALID: CachedId = CachedId::new("vortex.test.invalid");
248
249    #[test]
250    fn context_filters_interned_ids() {
251        let valid = *VALID;
252        let invalid = *INVALID;
253        let context = Interner::empty().with_allowed_ids(HashSet::from([valid]));
254
255        assert_eq!(context.intern(&valid), Some(0));
256        assert_eq!(context.intern(&valid), Some(0));
257        assert_eq!(context.intern(&invalid), None);
258        assert_eq!(context.to_ids(), [valid]);
259    }
260}
vortex_session/registry.rs

vortex_session/
registry.rs
