mdebray/tvl-depot
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor(users/edef/weave): simplify OwningRef handling

Browse source 
We switch to working with the underlying Arc<Bytes<u8>> type, since
Buffer<u8> is already essentially an OwningRef-esque subslice.

Because we're now working with an exposed Arc directly, we don't need
to have our own `unsafe impl StableAddress` any more.

Change-Id: I9ce2edc6899177145e15b72aa5380f708a62173c
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12668
Reviewed-by: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
This commit is contained in:
edef 2024-10-19 12:59:07 +00:00
parent 549c51f200
commit bdc2891053
2 changed files with 12 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
users/edef/weave/src/bytes.rs
View file

@ -1,27 +0,0 @@
use owning_ref::{OwningRef, StableAddress};
use polars::export::arrow::buffer::Buffer;
use std::ops::Deref;
/// An shared `[[u8; N]]` backed by a Polars [Buffer].
pub type FixedBytes<const N: usize> = OwningRef<'static, Bytes, [[u8; N]]>;
/// Wrapper struct to make [Buffer] implement [StableAddress].
/// TODO(edef): upstream the `impl`
pub struct Bytes(pub Buffer<u8>);
/// SAFETY: [Buffer] is always an Arc+Vec indirection.
unsafe impl StableAddress for Bytes {}
impl Bytes {
pub fn map<U: ?Sized>(self, f: impl FnOnce(&[u8]) -> &U) -> OwningRef<'static, Self, U> {
OwningRef::new(self).map(f)
}
}
impl Deref for Bytes {
type Target = [u8];
fn deref(&self) -> &Self::Target {
&*self.0
}
}

17
users/edef/weave/src/lib.rs
View file

@ -1,4 +1,5 @@
use anyhow::Result; use anyhow::Result;
use owning_ref::ArcRef;
use rayon::prelude::*; use rayon::prelude::*;
use std::{fs::File, ops::Range, slice}; use std::{fs::File, ops::Range, slice};
@ -8,8 +9,9 @@ use polars::{
prelude::{ParquetReader, SerReader}, prelude::{ParquetReader, SerReader},
}; };
pub use crate::bytes::*; /// An shared `[[u8; N]]` backed by a Polars [Buffer].
mod bytes; pub type FixedBytes<const N: usize> =
ArcRef<'static, polars::export::arrow::buffer::Bytes<u8>, [[u8; N]]>;
pub const INDEX_NULL: u32 = !0; pub const INDEX_NULL: u32 = !0;
pub const DONE: &str = "\u{2714}"; pub const DONE: &str = "\u{2714}";
@ -60,10 +62,15 @@ fn into_fixed_binary_rechunk<const N: usize>(chunked: &BinaryChunked) -> FixedBy
let chunked = chunked.rechunk(); let chunked = chunked.rechunk();
let mut iter = chunked.downcast_iter(); let mut iter = chunked.downcast_iter();
let array = iter.next().unwrap(); let array = iter.next().unwrap();
assert!(iter.next().is_none());
let range = assert_fixed_dense::<N>(array); let (buf, off, len) = {
Bytes(array.values().clone().sliced(range.start, range.len())) let range = assert_fixed_dense::<N>(array);
.map(|buf| exact_chunks(buf).unwrap()) array.values().clone().sliced(range.start, range.len())
}
.into_inner();
ArcRef::new(buf).map(|bytes| exact_chunks(&bytes[off..off + len]).unwrap())
} }
/// Ensures that the supplied Arrow array consists of densely packed bytestrings of length `N`. /// Ensures that the supplied Arrow array consists of densely packed bytestrings of length `N`.