use itertools::Itertools;
use num_traits::AsPrimitive;
use vortex_buffer::Buffer;
use vortex_buffer::BufferMut;
use vortex_buffer::ByteBuffer;
pub use crate::arrays::varbinview::BinaryView;
use crate::dtype::NativePType;
#[inline]
pub fn offsets_to_lengths<P: NativePType>(offsets: &[P]) -> Buffer<P> {
offsets
.iter()
.tuple_windows::<(_, _)>()
.map(|(&start, &end)| end - start)
.collect()
}
pub const MAX_BUFFER_LEN: usize = i32::MAX as usize;
pub fn build_views<P: NativePType + AsPrimitive<usize>>(
start_buf_index: u32,
max_buffer_len: usize,
bytes: ByteBuffer,
lens: &[P],
) -> (Vec<ByteBuffer>, Buffer<BinaryView>) {
let mut views = BufferMut::with_capacity(lens.len());
let buffers = extend_views(
&mut views,
start_buf_index,
max_buffer_len,
&bytes,
lens.len(),
|i| lens[i].as_(),
);
(buffers, views.freeze())
}
pub fn build_views_from_offsets<P: NativePType + AsPrimitive<usize>>(
start_buf_index: u32,
max_buffer_len: usize,
bytes: ByteBuffer,
offsets: &[P],
) -> (Vec<ByteBuffer>, Buffer<BinaryView>) {
assert!(!offsets.is_empty(), "offsets must hold at least one entry");
let first: usize = offsets[0].as_();
let last: usize = offsets[offsets.len() - 1].as_();
let bytes = bytes.slice(first..last);
let count = offsets.len() - 1;
let mut views = BufferMut::with_capacity(count);
let buffers = extend_views(
&mut views,
start_buf_index,
max_buffer_len,
&bytes,
count,
|i| {
AsPrimitive::<usize>::as_(offsets[i + 1])
.wrapping_sub(AsPrimitive::<usize>::as_(offsets[i]))
},
);
(buffers, views.freeze())
}
pub(crate) fn extend_views(
views: &mut BufferMut<BinaryView>,
start_buf_index: u32,
max_buffer_len: usize,
bytes: &ByteBuffer,
count: usize,
len_at: impl Fn(usize) -> usize,
) -> Vec<ByteBuffer> {
assert!(
max_buffer_len <= MAX_BUFFER_LEN,
"max_buffer_len cannot exceed MAX_BUFFER_LEN, offsets must fit in u32"
);
if bytes.len() <= max_buffer_len {
extend_views_single_buffer(views, start_buf_index, bytes, count, len_at);
if bytes.is_empty() {
Vec::new()
} else {
vec![bytes.clone()]
}
} else {
extend_views_rolling(views, start_buf_index, max_buffer_len, bytes, count, len_at)
}
}
fn extend_views_single_buffer(
views: &mut BufferMut<BinaryView>,
buf_index: u32,
bytes: &ByteBuffer,
count: usize,
len_at: impl Fn(usize) -> usize,
) {
views.reserve(count);
let base = views.len();
let data = bytes.as_slice();
let mut offset = 0usize;
let spare = &mut views.spare_capacity_mut()[..count];
for (i, slot) in spare.iter_mut().enumerate() {
let len = len_at(i);
let value = &data[offset..offset + len];
let view = if len > BinaryView::MAX_INLINED_SIZE {
let mut prefix = [0u8; 4];
prefix.copy_from_slice(&value[..4]);
BinaryView::new_ref(len.as_(), prefix, buf_index, offset.as_())
} else {
BinaryView::make_view(value, buf_index, offset.as_())
};
slot.write(view);
offset += len;
}
assert_eq!(
offset,
data.len(),
"value lengths must describe the byte heap exactly"
);
unsafe { views.set_len(base + count) };
}
fn extend_views_rolling(
views: &mut BufferMut<BinaryView>,
start_buf_index: u32,
max_buffer_len: usize,
bytes: &ByteBuffer,
count: usize,
len_at: impl Fn(usize) -> usize,
) -> Vec<ByteBuffer> {
views.reserve(count);
let mut buffers = Vec::new();
let mut buf_index = start_buf_index;
let data = bytes.as_slice();
let mut segment_start = 0usize;
let mut offset = 0usize;
for i in 0..count {
let len = len_at(i);
assert!(len <= max_buffer_len, "values cannot exceed max_buffer_len");
if offset + len > max_buffer_len {
buffers.push(bytes.slice(segment_start..segment_start + offset));
buf_index += 1;
segment_start += offset;
offset = 0;
}
let start = segment_start + offset;
let view = BinaryView::make_view(&data[start..start + len], buf_index, offset.as_());
views.push(view);
offset += len;
}
assert_eq!(
segment_start + offset,
data.len(),
"value lengths must describe the byte heap exactly"
);
if segment_start < data.len() {
buffers.push(bytes.slice(segment_start..data.len()));
}
buffers
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use vortex_buffer::ByteBuffer;
use vortex_buffer::ByteBufferMut;
use crate::arrays::varbinview::BinaryView;
use crate::arrays::varbinview::build_views::MAX_BUFFER_LEN;
use crate::arrays::varbinview::build_views::build_views;
use crate::arrays::varbinview::build_views::build_views_from_offsets;
fn flatten(values: &[&[u8]]) -> (ByteBuffer, Vec<u32>) {
let mut bytes = ByteBufferMut::empty();
let mut lens = Vec::with_capacity(values.len());
for v in values {
bytes.extend_from_slice(v);
lens.push(u32::try_from(v.len()).unwrap());
}
(bytes.freeze(), lens)
}
fn reconstruct(
buffers: &[ByteBuffer],
views: &[BinaryView],
start_buf_index: u32,
) -> Vec<Vec<u8>> {
views
.iter()
.map(|view| {
if view.is_inlined() {
view.as_inlined().value().to_vec()
} else {
let r = view.as_view();
let buf = &buffers[(r.buffer_index - start_buf_index) as usize];
buf[r.as_range()].to_vec()
}
})
.collect()
}
#[rstest]
#[case::mixed(&[b"a".as_slice(), b"this is a long reference value", b"short", b"another long value here!!"])]
#[case::inline_boundary(&[&[b'x'; 12] as &[u8], &[b'y'; 13], &[b'z'; 12], &[b'w'; 13]])]
#[case::all_inlined(&[b"".as_slice(), b"a", b"bb", b"ccc", b"dddddddddddd"])]
#[case::all_reference(&[&[b'a'; 100] as &[u8], &[b'b'; 50], &[b'c'; 4096]])]
#[case::empty_values_interleaved(&[b"".as_slice(), b"a long value that is referenced", b"", b"", b"trailing long reference value"])]
#[case::single_long(&[&[7u8; 1 << 16] as &[u8]])]
fn fast_path_roundtrip(#[case] values: &[&[u8]]) {
let (bytes, lens) = flatten(values);
let total = bytes.len();
let start_buf_index = 3;
let (buffers, views) = build_views(start_buf_index, total + 1, bytes, &lens);
assert_eq!(views.len(), values.len());
if total == 0 {
assert!(buffers.is_empty(), "empty heap must not allocate a buffer");
} else {
assert_eq!(buffers.len(), 1, "whole heap must stay in one buffer");
let concatenated: Vec<u8> = values.concat();
assert_eq!(buffers[0].as_slice(), concatenated.as_slice());
}
for view in views.iter() {
if !view.is_inlined() {
assert_eq!(view.as_view().buffer_index, start_buf_index);
}
}
let expected: Vec<Vec<u8>> = values.iter().map(|v| v.to_vec()).collect();
assert_eq!(reconstruct(&buffers, &views, start_buf_index), expected);
}
#[test]
fn output_buffers_are_zero_copy() {
let values: &[&[u8]] = &[
b"first long reference value",
b"tiny",
b"second long reference value!!",
b"third looooong reference value",
];
let (bytes, lens) = flatten(values);
let base = bytes.as_ptr();
let (buffers, _views) = build_views(0, bytes.len() + 1, bytes.clone(), &lens);
assert_eq!(buffers.len(), 1);
assert_eq!(buffers[0].as_ptr(), base, "fast path must not copy");
let longest = values.iter().map(|v| v.len()).max().unwrap();
let (buffers, _views) = build_views(0, longest, bytes, &lens);
assert!(buffers.len() > 1);
let mut expected_ptr = base;
for buffer in &buffers {
assert_eq!(buffer.as_ptr(), expected_ptr, "rolling path must not copy");
expected_ptr = unsafe { expected_ptr.add(buffer.len()) };
}
}
#[test]
fn fast_path_large_offsets() {
const N: usize = 9000;
const LEN: usize = 1000;
const { assert!((N - 1) * LEN > (1 << 23)) };
let values: Vec<Vec<u8>> = (0..N)
.map(|i| {
let mut v = vec![0u8; LEN];
v[..4].copy_from_slice(&u32::try_from(i).unwrap().to_le_bytes());
v
})
.collect();
let refs: Vec<&[u8]> = values.iter().map(|v| v.as_slice()).collect();
let (bytes, lens) = flatten(&refs);
let total = bytes.len();
let (buffers, views) = build_views(0, total + 1, bytes, &lens);
assert_eq!(buffers.len(), 1);
for (i, view) in views.iter().enumerate() {
let r = view.as_view();
assert_eq!(r.offset as usize, i * LEN, "wrong offset for view {i}");
assert_eq!(r.size as usize, LEN);
}
assert_eq!(reconstruct(&buffers, &views, 0), values);
}
#[test]
fn fast_path_taken_at_exact_boundary() {
let (bytes, lens) =
flatten(&[b"this value is definitely long", b"and so is this one here"]);
let total = bytes.len();
let (buffers, views) = build_views(0, total, bytes, &lens);
assert_eq!(
buffers.len(),
1,
"len == max_buffer_len must stay on fast path"
);
assert_eq!(views.len(), 2);
}
#[test]
fn fast_and_slow_paths_agree() {
let values: &[&[u8]] = &[
b"first long reference value",
b"tiny",
b"second long reference value!!",
b"third looooong reference value",
];
let expected: Vec<Vec<u8>> = values.iter().map(|v| v.to_vec()).collect();
let (fast_bytes, lens) = flatten(values);
let total = fast_bytes.len();
let (fast_buffers, fast_views) = build_views(0, total + 1, fast_bytes, &lens);
assert_eq!(fast_buffers.len(), 1);
assert_eq!(reconstruct(&fast_buffers, &fast_views, 0), expected);
let longest = values.iter().map(|v| v.len()).max().unwrap();
let (slow_bytes, _) = flatten(values);
let (slow_buffers, slow_views) = build_views(0, longest, slow_bytes, &lens);
assert!(
slow_buffers.len() > 1,
"small cap should split into many buffers"
);
assert_eq!(reconstruct(&slow_buffers, &slow_views, 0), expected);
assert_eq!(
reconstruct(&fast_buffers, &fast_views, 0),
reconstruct(&slow_buffers, &slow_views, 0)
);
}
#[test]
fn fast_path_empty_input() {
let lens: Vec<u32> = Vec::new();
let (buffers, views) = build_views(0, 1024, ByteBuffer::empty(), &lens);
assert!(buffers.is_empty());
assert!(views.is_empty());
}
#[test]
fn fast_path_matches_make_view() {
let values: &[&[u8]] = &[b"inline", b"this is a long reference value", b""];
let (bytes, lens) = flatten(values);
let total = bytes.len();
let (_buffers, views) = build_views(0, total + 1, bytes, &lens);
let expected = [
BinaryView::make_view(b"inline", 0, 0),
BinaryView::make_view(b"this is a long reference value", 0, 6),
BinaryView::make_view(b"", 0, 36),
];
assert_eq!(views.as_slice(), &expected);
}
#[test]
fn from_offsets_matches_lengths_and_is_zero_copy() {
let heap = ByteBuffer::copy_from(b"..a long value that is referenced!tiny..".as_slice());
let offsets: Vec<u32> = vec![2, 34, 38];
let (buffers, views) = build_views_from_offsets(5, MAX_BUFFER_LEN, heap.clone(), &offsets);
assert_eq!(buffers.len(), 1);
assert_eq!(buffers[0].as_ptr(), unsafe { heap.as_ptr().add(2) });
assert_eq!(buffers[0].len(), 36);
assert_eq!(
reconstruct(&buffers, &views, 5),
vec![
b"a long value that is referenced!".to_vec(),
b"tiny".to_vec()
]
);
}
#[test]
#[should_panic(expected = "value lengths must describe the byte heap exactly")]
fn short_lengths_panic() {
let (bytes, _) = flatten(&[b"a long value that is referenced", b"tiny"]);
build_views(0, MAX_BUFFER_LEN, bytes, &[31u32]);
}
#[test_with::env(CI)]
#[test_with::no_env(VORTEX_SKIP_SLOW_TESTS)]
fn build_views_offsets_overflow_i32() {
const STRING_LEN: usize = 64 * 1024;
const TOTAL_BYTES: usize = (1usize << 31) + (256 << 20); const N: usize = TOTAL_BYTES / STRING_LEN;
let nth_string = |i: usize| {
let mut s = vec![b'x'; STRING_LEN];
s[..8].copy_from_slice(&(i as u64).to_le_bytes());
s
};
let mut bytes = ByteBufferMut::with_capacity(N * STRING_LEN);
let mut value = vec![b'x'; STRING_LEN];
for i in 0..N {
value[..8].copy_from_slice(&(i as u64).to_le_bytes());
bytes.extend_from_slice(&value);
}
let lens = vec![u32::try_from(STRING_LEN).unwrap(); N];
let (buffers, views) = build_views(0, MAX_BUFFER_LEN, bytes.freeze(), &lens);
assert_eq!(views.len(), N);
assert!(
buffers.len() >= 2,
"heap exceeding MAX_BUFFER_LEN must roll over into multiple buffers, got {}",
buffers.len()
);
for (i, b) in buffers.iter().enumerate() {
assert!(
b.len() <= MAX_BUFFER_LEN,
"buffer {i} of {} bytes exceeds MAX_BUFFER_LEN",
b.len()
);
}
let boundary = MAX_BUFFER_LEN / STRING_LEN;
for i in [0, boundary - 1, boundary, boundary + 1, N / 2, N - 1] {
let view = &views[i];
let r = view.as_view();
let got = &buffers[r.buffer_index as usize][r.as_range()];
assert_eq!(got, nth_string(i).as_slice(), "value mismatch at row {i}");
assert_eq!(r.size as usize, STRING_LEN);
}
}
#[test]
fn test_to_canonical_large() {
let raw_data =
ByteBuffer::copy_from("aaaaaaaaaaaaabbbbbbbbbbbbbcccccccccccccddddddddddddd");
let lens = vec![13u8; 4];
let (buffers, views) = build_views(0, 26, raw_data, &lens);
assert_eq!(
buffers,
vec![
ByteBuffer::copy_from("aaaaaaaaaaaaabbbbbbbbbbbbb"),
ByteBuffer::copy_from("cccccccccccccddddddddddddd"),
]
);
assert_eq!(
views.as_slice(),
&[
BinaryView::make_view(b"aaaaaaaaaaaaa", 0, 0),
BinaryView::make_view(b"bbbbbbbbbbbbb", 0, 13),
BinaryView::make_view(b"ccccccccccccc", 1, 0),
BinaryView::make_view(b"ddddddddddddd", 1, 13),
]
)
}
#[test]
#[should_panic(expected = "max_buffer_len cannot exceed MAX_BUFFER_LEN")]
fn test_max_buffer_len_too_large_panics() {
build_views(0, MAX_BUFFER_LEN + 1, ByteBuffer::copy_from("abc"), &[3u32]);
}
}