Skip to main content

columnar/
string.rs

1use alloc::{vec::Vec, string::String, string::ToString, boxed::Box};
2use super::{Clear, Columnar, Container, Len, Index, IndexAs, Push, Borrow};
3
4/// A stand-in for `Vec<String>`.
5///
6/// The reference type for `Strings` is `&[u8]` rather than `&str` to remove utf8 validation
7/// from the critical path of reads. You get to make the call about whether and how you'd like
8/// to manage this validation. The `copy_from` and `into_owned` methods panic on invalid data.
9#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
10#[derive(Copy, Clone, Debug, Default, PartialEq)]
11pub struct Strings<BC = Vec<u64>, VC = Vec<u8>> {
12    /// Bounds container; provides indexed access to offsets.
13    pub bounds: BC,
14    /// Values container; provides slice access to bytes.
15    pub values: VC,
16}
17
18impl Columnar for String {
19    #[inline(always)]
20    fn copy_from<'a>(&mut self, other: crate::Ref<'a, Self>) {
21        self.clear();
22        self.push_str(core::str::from_utf8(other).expect("invalid utf8 in Strings column"));
23    }
24    #[inline(always)]
25    fn into_owned<'a>(other: crate::Ref<'a, Self>) -> Self {
26        core::str::from_utf8(other).expect("invalid utf8 in Strings column").to_string()
27    }
28    type Container = Strings;
29}
30
31impl Columnar for Box<str> {
32    #[inline(always)]
33    fn copy_from<'a>(&mut self, other: crate::Ref<'a, Self>) {
34        let mut s = String::from(core::mem::take(self));
35        s.clear();
36        s.push_str(core::str::from_utf8(other).expect("invalid utf8 in Strings column"));
37        *self = s.into_boxed_str();
38    }
39    #[inline(always)]
40    fn into_owned<'a>(other: crate::Ref<'a, Self>) -> Self {
41        Self::from(core::str::from_utf8(other).expect("invalid utf8 in Strings column"))
42    }
43    type Container = Strings;
44}
45
46impl<BC: crate::common::BorrowIndexAs<u64>> Borrow for Strings<BC, Vec<u8>> {
47    type Ref<'a> = &'a [u8];
48    type Borrowed<'a> = Strings<BC::Borrowed<'a>, &'a [u8]> where BC: 'a;
49    #[inline(always)]
50    fn borrow<'a>(&'a self) -> Self::Borrowed<'a> {
51        Strings {
52            bounds: self.bounds.borrow(),
53            values: self.values.borrow(),
54        }
55    }
56    #[inline(always)]
57    fn reborrow<'c, 'a: 'c>(thing: Self::Borrowed<'a>) -> Self::Borrowed<'c> where BC: 'a {
58        Strings {
59            bounds: BC::reborrow(thing.bounds),
60            values: thing.values,
61        }
62    }
63    #[inline(always)]
64    fn reborrow_ref<'b, 'a: 'b>(thing: Self::Ref<'a>) -> Self::Ref<'b> where Self: 'a { thing }
65}
66
67impl<BC: crate::common::PushIndexAs<u64>> Container for Strings<BC, Vec<u8>> {
68    #[inline(always)]
69    fn extend_from_self(&mut self, other: Self::Borrowed<'_>, range: core::ops::Range<usize>) {
70        if !range.is_empty() {
71            // Imported bounds will be relative to this starting offset.
72            let values_len = self.values.len() as u64;
73
74            // Push all bytes that we can, all at once.
75            let other_lower = if range.start == 0 { 0 } else { other.bounds.index_as(range.start-1) };
76            let other_upper = other.bounds.index_as(range.end-1);
77            self.values.extend_from_self(other.values, other_lower as usize .. other_upper as usize);
78
79            // Each bound needs to be shifted by `values_len - other_lower`.
80            if values_len == other_lower {
81                self.bounds.extend_from_self(other.bounds, range);
82            }
83            else {
84                for index in range {
85                    let shifted = other.bounds.index_as(index) - other_lower + values_len;
86                    self.bounds.push(&shifted)
87                }
88            }
89        }
90    }
91
92    fn reserve_for<'a, I>(&mut self, selves: I) where Self: 'a, I: Iterator<Item = Self::Borrowed<'a>> + Clone {
93        self.bounds.reserve_for(selves.clone().map(|x| x.bounds));
94        self.values.reserve_for(selves.map(|x| x.values));
95    }
96
97}
98
99impl<'a, BC: crate::AsBytes<'a>, VC: crate::AsBytes<'a>> crate::AsBytes<'a> for Strings<BC, VC> {
100    #[inline(always)]
101    fn as_bytes(&self) -> impl Iterator<Item=(u64, &'a [u8])> {
102        crate::chain(self.bounds.as_bytes(), self.values.as_bytes())
103    }
104}
105impl<'a, BC: crate::FromBytes<'a>, VC: crate::FromBytes<'a>> crate::FromBytes<'a> for Strings<BC, VC> {
106    const SLICE_COUNT: usize = BC::SLICE_COUNT + VC::SLICE_COUNT;
107    #[inline(always)]
108    fn from_bytes(bytes: &mut impl Iterator<Item=&'a [u8]>) -> Self {
109        Self {
110            bounds: crate::FromBytes::from_bytes(bytes),
111            values: crate::FromBytes::from_bytes(bytes),
112        }
113    }
114    #[inline(always)]
115    fn from_store(store: &crate::bytes::indexed::DecodedStore<'a>, offset: &mut usize) -> Self {
116        Self {
117            bounds: BC::from_store(store, offset),
118            values: VC::from_store(store, offset),
119        }
120    }
121    fn element_sizes(sizes: &mut Vec<usize>) -> Result<(), String> {
122        BC::element_sizes(sizes)?;
123        VC::element_sizes(sizes)?;
124        Ok(())
125    }
126}
127
128impl<BC: Len, VC> Len for Strings<BC, VC> {
129    #[inline(always)] fn len(&self) -> usize { self.bounds.len() }
130}
131
132impl<'a, BC: Len+IndexAs<u64>> Strings<BC, &'a [u8]> {
133    /// Returns the `index`-th string as `&str`, validating UTF-8.
134    ///
135    /// This is a convenience wrapper around [`Index::get`] (which returns `&[u8]`)
136    /// for callers who need `&str`. The UTF-8 validation has a measurable cost;
137    /// use `get` directly if you can work with `&[u8]`.
138    #[inline(always)]
139    pub fn get_str(&self, index: usize) -> &'a str {
140        core::str::from_utf8(self.get(index)).expect("invalid utf8 in Strings column")
141    }
142}
143
144impl<'a, BC: Len+IndexAs<u64>> Index for Strings<BC, &'a [u8]> {
145    type Ref = &'a [u8];
146    #[inline(always)] fn get(&self, index: usize) -> Self::Ref {
147        let lower = if index == 0 { 0 } else { self.bounds.index_as(index - 1) };
148        let upper = self.bounds.index_as(index);
149        let lower: usize = lower.try_into().expect("bounds must fit in `usize`");
150        let upper: usize = upper.try_into().expect("bounds must fit in `usize`");
151        &self.values[lower .. upper]
152    }
153}
154impl<'a, BC: Len+IndexAs<u64>> Index for &'a Strings<BC, Vec<u8>> {
155    type Ref = &'a [u8];
156    #[inline(always)] fn get(&self, index: usize) -> Self::Ref {
157        let lower = if index == 0 { 0 } else { self.bounds.index_as(index - 1) };
158        let upper = self.bounds.index_as(index);
159        let lower: usize = lower.try_into().expect("bounds must fit in `usize`");
160        let upper: usize = upper.try_into().expect("bounds must fit in `usize`");
161        &self.values[lower .. upper]
162    }
163}
164
165// This is a simpler implementation, but it leads to a performance regression
166// for Strings and str because it loses access to `Vec::extend_from_slice`.
167//
168// impl<BC: Push<u64>, D: std::fmt::Display> Push<D> for Strings<BC> {
169//     #[inline(always)]
170//     fn push(&mut self, item: D) {
171//         use std::io::Write;
172//         write!(self.values, "{}", item).unwrap();
173//         self.bounds.push(self.values.len() as u64);
174//     }
175// }
176
177impl<BC: for<'a> Push<&'a u64>> Push<&[u8]> for Strings<BC> {
178    #[inline(always)] fn push(&mut self, item: &[u8]) {
179        self.values.extend_from_slice(item);
180        self.bounds.push(&(self.values.len() as u64));
181    }
182}
183impl<BC: for<'a> Push<&'a u64>> Push<&String> for Strings<BC> {
184    #[inline(always)] fn push(&mut self, item: &String) {
185        self.values.extend_from_slice(item.as_bytes());
186        self.bounds.push(&(self.values.len() as u64));
187    }
188}
189impl<BC: for<'a> Push<&'a u64>> Push<&str> for Strings<BC> {
190    #[inline]
191    fn push(&mut self, item: &str) {
192        self.values.extend_from_slice(item.as_bytes());
193        self.bounds.push(&(self.values.len() as u64));
194    }
195}
196impl<BC: for<'a> Push<&'a u64>> Push<&Box<str>> for Strings<BC> {
197    #[inline]
198    fn push(&mut self, item: &Box<str>) {
199        self.values.extend_from_slice(item.as_bytes());
200        self.bounds.push(&(self.values.len() as u64));
201    }
202}
203impl<'a, BC: for<'b> Push<&'b u64>> Push<core::fmt::Arguments<'a>> for Strings<BC> {
204    #[inline]
205    fn push(&mut self, item: core::fmt::Arguments<'a>) {
206        // Use core::fmt::Write via a wrapper to avoid requiring std::io::Write.
207        struct VecWriter<'a>(&'a mut alloc::vec::Vec<u8>);
208        impl core::fmt::Write for VecWriter<'_> {
209            fn write_str(&mut self, s: &str) -> core::fmt::Result {
210                self.0.extend_from_slice(s.as_bytes());
211                Ok(())
212            }
213        }
214        core::fmt::Write::write_fmt(&mut VecWriter(&mut self.values), item).expect("write_fmt failed");
215        self.bounds.push(&(self.values.len() as u64));
216    }
217}
218impl<'a, 'b, BC: for<'c> Push<&'c u64>> Push<&'b core::fmt::Arguments<'a>> for Strings<BC> {
219    #[inline]
220    fn push(&mut self, item: &'b core::fmt::Arguments<'a>) {
221        self.push(*item);
222    }
223}
224impl<BC: Clear, VC: Clear> Clear for Strings<BC, VC> {
225    #[inline(always)]
226    fn clear(&mut self) {
227        self.bounds.clear();
228        self.values.clear();
229    }
230}
231