Skip to content

Commit 0dade7d

Browse files
ClSlaidClSlaid
committed
perf(arrow): export offset byte arrays directly
Signed-off-by: cl <cailue@apache.org>
1 parent 037b1a8 commit 0dade7d

23 files changed

Lines changed: 890 additions & 94 deletions

File tree

‎encodings/fsst/src/array.rs‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ use std::sync::OnceLock;
1111
use fsst::Compressor;
1212
use fsst::Decompressor;
1313
use fsst::Symbol;
14+
use num_traits::AsPrimitive;
1415
use prost::Message as _;
1516
use vortex_array::Array;
1617
use vortex_array::ArrayEq;
@@ -32,11 +33,13 @@ use vortex_array::arrays::VarBinArray;
3233
use vortex_array::arrays::varbin::VarBinArrayExt;
3334
use vortex_array::buffer::BufferHandle;
3435
use vortex_array::builders::ArrayBuilder;
36+
use vortex_array::builders::OffsetVarBinBuilder;
3537
use vortex_array::builders::VarBinViewBuilder;
3638
use vortex_array::dtype::DType;
3739
use vortex_array::dtype::Nullability;
3840
use vortex_array::dtype::PType;
3941
use vortex_array::legacy_session;
42+
use vortex_array::match_each_integer_ptype;
4043
use vortex_array::serde::ArrayChildren;
4144
use vortex_array::smallvec::smallvec;
4245
use vortex_array::validity::Validity;
@@ -56,6 +59,7 @@ use vortex_session::VortexSession;
5659
use vortex_session::registry::CachedId;
5760

5861
use crate::canonical::canonicalize_fsst;
62+
use crate::canonical::fsst_decode_bytes;
5963
use crate::canonical::fsst_decode_views;
6064
use crate::rules::RULES;
6165

@@ -304,6 +308,25 @@ impl VTable for FSST {
304308
builder: &mut dyn ArrayBuilder,
305309
ctx: &mut ExecutionCtx,
306310
) -> VortexResult<()> {
311+
if let Some(builder) = builder.as_any_mut().downcast_mut::<OffsetVarBinBuilder>() {
312+
let (bytes, lengths) = fsst_decode_bytes(array, ctx)?;
313+
let validity = array
314+
.array()
315+
.validity()?
316+
.execute_mask(array.array().len(), ctx)?;
317+
match_each_integer_ptype!(lengths.ptype(), |P| {
318+
builder.append_values(
319+
bytes.as_slice(),
320+
lengths
321+
.as_slice::<P>()
322+
.iter()
323+
.map(|length| AsPrimitive::<usize>::as_(*length)),
324+
&validity,
325+
);
326+
});
327+
return Ok(());
328+
}
329+
307330
let Some(builder) = builder.as_any_mut().downcast_mut::<VarBinViewBuilder>() else {
308331
return array
309332
.array()

‎encodings/fsst/src/canonical.rs‎

Lines changed: 26 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -40,20 +40,11 @@ pub(super) fn canonicalize_fsst(
4040
})
4141
}
4242

43-
pub(crate) fn fsst_decode_views(
43+
pub(crate) fn fsst_decode_bytes(
4444
fsst_array: ArrayView<'_, FSST>,
45-
start_buf_index: u32,
4645
ctx: &mut ExecutionCtx,
47-
) -> VortexResult<(Vec<ByteBuffer>, Buffer<BinaryView>)> {
48-
// FSSTArray has two child arrays:
49-
// 1. A VarBinArray, which holds the string heap of the compressed codes.
50-
// 2. An uncompressed_lengths primitive array, storing the length of each original
51-
// string element.
52-
// To speed up canonicalization, we can decompress the entire string-heap in a single
53-
// call. We then turn our uncompressed_lengths into an offsets buffer
54-
// necessary for a VarBinViewArray and construct the canonical array.
46+
) -> VortexResult<(ByteBufferMut, PrimitiveArray)> {
5547
let bytes = fsst_array.codes().sliced_bytes();
56-
5748
let uncompressed_lens_array = fsst_array
5849
.uncompressed_lengths()
5950
.clone()
@@ -68,14 +59,20 @@ pub(crate) fn fsst_decode_views(
6859
.sum()
6960
});
7061

71-
// Bulk-decompress the entire array.
7262
let decompressor = fsst_array.decompressor();
7363
let mut uncompressed_bytes = ByteBufferMut::with_capacity(total_size + 7);
7464
let len =
7565
decompressor.decompress_into(bytes.as_slice(), uncompressed_bytes.spare_capacity_mut());
7666
unsafe { uncompressed_bytes.set_len(len) };
67+
Ok((uncompressed_bytes, uncompressed_lens_array))
68+
}
7769

78-
// Directly create the binary views.
70+
pub(crate) fn fsst_decode_views(
71+
fsst_array: ArrayView<'_, FSST>,
72+
start_buf_index: u32,
73+
ctx: &mut ExecutionCtx,
74+
) -> VortexResult<(Vec<ByteBuffer>, Buffer<BinaryView>)> {
75+
let (uncompressed_bytes, uncompressed_lens_array) = fsst_decode_bytes(fsst_array, ctx)?;
7976
match_each_integer_ptype!(uncompressed_lens_array.ptype(), |P| {
8077
Ok(build_views(
8178
start_buf_index,
@@ -99,7 +96,9 @@ mod tests {
9996
use vortex_array::arrays::ChunkedArray;
10097
use vortex_array::arrays::VarBinArray;
10198
use vortex_array::arrays::VarBinViewArray;
99+
use vortex_array::arrays::varbin::VarBinArrayExt;
102100
use vortex_array::builders::ArrayBuilder;
101+
use vortex_array::builders::OffsetVarBinBuilder;
103102
use vortex_array::builders::VarBinViewBuilder;
104103
use vortex_array::dtype::DType;
105104
use vortex_array::dtype::Nullability;
@@ -199,6 +198,20 @@ mod tests {
199198
.collect::<Vec<_>>();
200199
assert_eq!(data, res2)
201200
};
201+
202+
{
203+
let mut builder =
204+
OffsetVarBinBuilder::with_capacity(chunked_arr.dtype().clone(), false, data.len());
205+
chunked_arr
206+
.into_array()
207+
.append_to_builder(&mut builder, &mut ctx)?;
208+
let arr = builder.finish_into_varbin();
209+
let mask = arr.validity()?.execute_mask(arr.len(), &mut ctx)?;
210+
let actual = (0..arr.len())
211+
.map(|i| mask.value(i).then(|| arr.bytes_at(i).to_vec()))
212+
.collect::<Vec<_>>();
213+
assert_eq!(data, actual);
214+
}
202215
Ok(())
203216
}
204217

‎encodings/onpair/src/array.rs‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ use std::hash::Hasher;
88
use std::sync::Arc;
99
use std::sync::OnceLock;
1010

11+
use num_traits::AsPrimitive;
1112
use onpair::CompactDictionaryView;
1213
use prost::Message as _;
1314
use vortex_array::Array;
@@ -25,10 +26,12 @@ use vortex_array::IntoArray;
2526
use vortex_array::array_slots;
2627
use vortex_array::buffer::BufferHandle;
2728
use vortex_array::builders::ArrayBuilder;
29+
use vortex_array::builders::OffsetVarBinBuilder;
2830
use vortex_array::builders::VarBinViewBuilder;
2931
use vortex_array::dtype::DType;
3032
use vortex_array::dtype::Nullability;
3133
use vortex_array::dtype::PType;
34+
use vortex_array::match_each_integer_ptype;
3235
use vortex_array::serde::ArrayChildren;
3336
use vortex_array::validity::Validity;
3437
use vortex_array::vtable::VTable;
@@ -46,6 +49,7 @@ use vortex_session::VortexSession;
4649
use vortex_session::registry::CachedId;
4750

4851
use crate::canonical::canonicalize_onpair;
52+
use crate::canonical::onpair_decode_bytes;
4953
use crate::canonical::onpair_decode_views;
5054
use crate::decode::collect_widened;
5155
use crate::rules::RULES;
@@ -544,6 +548,25 @@ impl VTable for OnPair {
544548
builder: &mut dyn ArrayBuilder,
545549
ctx: &mut ExecutionCtx,
546550
) -> VortexResult<()> {
551+
if let Some(builder) = builder.as_any_mut().downcast_mut::<OffsetVarBinBuilder>() {
552+
let (bytes, lengths) = onpair_decode_bytes(array, ctx)?;
553+
let validity = array
554+
.array()
555+
.validity()?
556+
.execute_mask(array.array().len(), ctx)?;
557+
match_each_integer_ptype!(lengths.ptype(), |P| {
558+
builder.append_values(
559+
bytes.as_slice(),
560+
lengths
561+
.as_slice::<P>()
562+
.iter()
563+
.map(|length| AsPrimitive::<usize>::as_(*length)),
564+
&validity,
565+
);
566+
});
567+
return Ok(());
568+
}
569+
547570
let Some(builder) = builder.as_any_mut().downcast_mut::<VarBinViewBuilder>() else {
548571
return array
549572
.array()

‎encodings/onpair/src/canonical.rs‎

Lines changed: 10 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -44,11 +44,10 @@ pub(super) fn canonicalize_onpair(
4444
})
4545
}
4646

47-
pub(crate) fn onpair_decode_views(
47+
pub(crate) fn onpair_decode_bytes(
4848
array: ArrayView<'_, OnPair>,
49-
start_buf_index: u32,
5049
ctx: &mut ExecutionCtx,
51-
) -> VortexResult<(Vec<ByteBuffer>, Buffer<BinaryView>)> {
50+
) -> VortexResult<(ByteBufferMut, PrimitiveArray)> {
5251
let lengths = array
5352
.uncompressed_lengths()
5453
.clone()
@@ -62,15 +61,6 @@ pub(crate) fn onpair_decode_views(
6261
.sum()
6362
});
6463

65-
// `codes_offsets` holds the per-row code boundaries and may itself be a
66-
// sliced or filtered view of the original. Its first and last entries
67-
// bound the contiguous run of `codes` belonging to the rows present in
68-
// this array: `slice` keeps the full `codes` child and only narrows
69-
// `codes_offsets` (so `code_start > 0` and/or `code_end < codes.len()`),
70-
// while `filter` rebuilds both children so the window is the whole stream.
71-
// OnPair has no `TakeExecute`, so a reordering take is served from the
72-
// canonical `VarBinView` and never reaches this path. We only need those
73-
// two boundaries, so point-look them up rather than decoding every offset.
7464
let codes_offsets = array.codes_offsets();
7565
let code_start = code_boundary_at(codes_offsets, 0, ctx)?;
7666
let code_end = code_boundary_at(codes_offsets, array.len(), ctx)?;
@@ -85,13 +75,8 @@ pub(crate) fn onpair_decode_views(
8575
array.codes().len()
8676
);
8777

88-
// Slice the `codes` child to that window *before* unpacking it, so a sliced
89-
// array materialises only its own codes rather than the whole column's. The
90-
// contiguous decoder walks `codes` in order and never reads the per-row
91-
// boundaries, so an empty boundary slice is sound.
9278
let codes = collect_widened::<u16>(&array.codes().slice(code_start..code_end)?, ctx)?;
9379
let dict = dict_view(array, ctx)?;
94-
9580
let mut out_bytes = ByteBufferMut::with_capacity(total_size);
9681
let written =
9782
match onpair::try_decode_into(codes.as_slice(), dict, out_bytes.spare_capacity_mut()) {
@@ -105,9 +90,16 @@ pub(crate) fn onpair_decode_views(
10590
"OnPair codes decoded to {written} bytes but uncompressed_lengths records {total_size}"
10691
);
10792
}
108-
// SAFETY: `try_decode_into` initialised exactly `written` bytes.
10993
unsafe { out_bytes.set_len(written) };
94+
Ok((out_bytes, lengths))
95+
}
11096

97+
pub(crate) fn onpair_decode_views(
98+
array: ArrayView<'_, OnPair>,
99+
start_buf_index: u32,
100+
ctx: &mut ExecutionCtx,
101+
) -> VortexResult<(Vec<ByteBuffer>, Buffer<BinaryView>)> {
102+
let (out_bytes, lengths) = onpair_decode_bytes(array, ctx)?;
111103
match_each_integer_ptype!(lengths.ptype(), |P| {
112104
Ok(build_views(
113105
start_buf_index,

‎encodings/onpair/src/tests.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ use vortex_array::arrays::VarBinViewArray;
1616
use vortex_array::arrays::filter::FilterKernel;
1717
use vortex_array::assert_arrays_eq;
1818
use vortex_array::buffer::BufferHandle;
19+
use vortex_array::builders::OffsetVarBinBuilder;
1920
use vortex_array::builtins::ArrayBuiltins;
2021
use vortex_array::dtype::DType;
2122
use vortex_array::dtype::Nullability;
@@ -60,6 +61,19 @@ fn sample_input() -> VarBinArray {
6061
)
6162
}
6263

64+
#[test]
65+
fn direct_offset_builder() -> vortex_error::VortexResult<()> {
66+
let mut ctx = SESSION.create_execution_ctx();
67+
let input = sample_input();
68+
let encoded = compress_onpair(input.as_ref(), &mut ctx)?;
69+
let mut builder = OffsetVarBinBuilder::with_capacity(input.dtype().clone(), false, input.len());
70+
encoded
71+
.into_array()
72+
.append_to_builder(&mut builder, &mut ctx)?;
73+
assert_arrays_eq!(builder.finish_into_varbin(), input, &mut ctx);
74+
Ok(())
75+
}
76+
6377
#[test]
6478
fn test_onpair_rejects_100k_token_dictionary() -> vortex_error::VortexResult<()> {
6579
let num_tokens = 100_000usize;

0 commit comments

Comments
 (0)