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    const SLICE_COUNT: usize = BC::SLICE_COUNT + VC::SLICE_COUNT;
101    #[inline]
102    fn get_byte_slice(&self, index: usize) -> (u64, &'a [u8]) {
103        debug_assert!(index < Self::SLICE_COUNT);
104        if index < BC::SLICE_COUNT {
105            self.bounds.get_byte_slice(index)
106        } else {
107            self.values.get_byte_slice(index - BC::SLICE_COUNT)
108        }
109    }
110}
111impl<'a, BC: crate::FromBytes<'a>, VC: crate::FromBytes<'a>> crate::FromBytes<'a> for Strings<BC, VC> {
112    const SLICE_COUNT: usize = BC::SLICE_COUNT + VC::SLICE_COUNT;
113    #[inline(always)]
114    fn from_bytes(bytes: &mut impl Iterator<Item=&'a [u8]>) -> Self {
115        Self {
116            bounds: crate::FromBytes::from_bytes(bytes),
117            values: crate::FromBytes::from_bytes(bytes),
118        }
119    }
120    #[inline(always)]
121    fn from_store(store: &crate::bytes::indexed::DecodedStore<'a>, offset: &mut usize) -> Self {
122        Self {
123            bounds: BC::from_store(store, offset),
124            values: VC::from_store(store, offset),
125        }
126    }
127    fn element_sizes(sizes: &mut Vec<usize>) -> Result<(), String> {
128        BC::element_sizes(sizes)?;
129        VC::element_sizes(sizes)?;
130        Ok(())
131    }
132}
133
134impl<BC: Len, VC> Len for Strings<BC, VC> {
135    #[inline(always)] fn len(&self) -> usize { self.bounds.len() }
136}
137
138impl<'a, BC: Len+IndexAs<u64>> Strings<BC, &'a [u8]> {
139    /// Returns the `index`-th string as `&str`, validating UTF-8.
140    ///
141    /// This is a convenience wrapper around [`Index::get`] (which returns `&[u8]`)
142    /// for callers who need `&str`. The UTF-8 validation has a measurable cost;
143    /// use `get` directly if you can work with `&[u8]`.
144    #[inline(always)]
145    pub fn get_str(&self, index: usize) -> &'a str {
146        core::str::from_utf8(self.get(index)).expect("invalid utf8 in Strings column")
147    }
148}
149
150impl<'a, BC: Len+IndexAs<u64>> Index for Strings<BC, &'a [u8]> {
151    type Ref = &'a [u8];
152    #[inline(always)] fn get(&self, index: usize) -> Self::Ref {
153        let lower = if index == 0 { 0 } else { self.bounds.index_as(index - 1) };
154        let upper = self.bounds.index_as(index);
155        let lower: usize = lower.try_into().expect("bounds must fit in `usize`");
156        let upper: usize = upper.try_into().expect("bounds must fit in `usize`");
157        &self.values[lower .. upper]
158    }
159}
160impl<'a, BC: Len+IndexAs<u64>> Index for &'a Strings<BC, Vec<u8>> {
161    type Ref = &'a [u8];
162    #[inline(always)] fn get(&self, index: usize) -> Self::Ref {
163        let lower = if index == 0 { 0 } else { self.bounds.index_as(index - 1) };
164        let upper = self.bounds.index_as(index);
165        let lower: usize = lower.try_into().expect("bounds must fit in `usize`");
166        let upper: usize = upper.try_into().expect("bounds must fit in `usize`");
167        &self.values[lower .. upper]
168    }
169}
170
171// This is a simpler implementation, but it leads to a performance regression
172// for Strings and str because it loses access to `Vec::extend_from_slice`.
173//
174// impl<BC: Push<u64>, D: std::fmt::Display> Push<D> for Strings<BC> {
175//     #[inline(always)]
176//     fn push(&mut self, item: D) {
177//         use std::io::Write;
178//         write!(self.values, "{}", item).unwrap();
179//         self.bounds.push(self.values.len() as u64);
180//     }
181// }
182
183impl<BC: for<'a> Push<&'a u64>> Push<&[u8]> for Strings<BC> {
184    #[inline(always)] fn push(&mut self, item: &[u8]) {
185        self.values.extend_from_slice(item);
186        self.bounds.push(&(self.values.len() as u64));
187    }
188}
189impl<BC: for<'a> Push<&'a u64>> Push<&String> for Strings<BC> {
190    #[inline(always)] fn push(&mut self, item: &String) {
191        self.values.extend_from_slice(item.as_bytes());
192        self.bounds.push(&(self.values.len() as u64));
193    }
194}
195impl<BC: for<'a> Push<&'a u64>> Push<&str> for Strings<BC> {
196    #[inline]
197    fn push(&mut self, item: &str) {
198        self.values.extend_from_slice(item.as_bytes());
199        self.bounds.push(&(self.values.len() as u64));
200    }
201}
202impl<BC: for<'a> Push<&'a u64>> Push<&Box<str>> for Strings<BC> {
203    #[inline]
204    fn push(&mut self, item: &Box<str>) {
205        self.values.extend_from_slice(item.as_bytes());
206        self.bounds.push(&(self.values.len() as u64));
207    }
208}
209impl<'a, BC: for<'b> Push<&'b u64>> Push<core::fmt::Arguments<'a>> for Strings<BC> {
210    #[inline]
211    fn push(&mut self, item: core::fmt::Arguments<'a>) {
212        // Use core::fmt::Write via a wrapper to avoid requiring std::io::Write.
213        struct VecWriter<'a>(&'a mut alloc::vec::Vec<u8>);
214        impl core::fmt::Write for VecWriter<'_> {
215            fn write_str(&mut self, s: &str) -> core::fmt::Result {
216                self.0.extend_from_slice(s.as_bytes());
217                Ok(())
218            }
219        }
220        core::fmt::Write::write_fmt(&mut VecWriter(&mut self.values), item).expect("write_fmt failed");
221        self.bounds.push(&(self.values.len() as u64));
222    }
223}
224impl<'a, 'b, BC: for<'c> Push<&'c u64>> Push<&'b core::fmt::Arguments<'a>> for Strings<BC> {
225    #[inline]
226    fn push(&mut self, item: &'b core::fmt::Arguments<'a>) {
227        self.push(*item);
228    }
229}
230impl<BC: Clear, VC: Clear> Clear for Strings<BC, VC> {
231    #[inline(always)]
232    fn clear(&mut self) {
233        self.bounds.clear();
234        self.values.clear();
235    }
236}
237