Skip to content

Commit dd5b0c5

Browse files
committed
build PatchedArray with env var VORTEX_EXPERIMENTAL_PATCHED_ARRAY.
Signed-off-by: Andrew Duffy <andrew@a10y.dev>
1 parent abb955a commit dd5b0c5

1 file changed

Lines changed: 54 additions & 16 deletions

File tree

vortex-btrblocks/src/schemes/integer.rs

Lines changed: 54 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,17 @@
33

44
//! Integer compression schemes.
55
6+
use std::env;
7+
use std::sync::LazyLock;
8+
69
use vortex_array::ArrayRef;
710
use vortex_array::Canonical;
811
use vortex_array::IntoArray;
12+
use vortex_array::LEGACY_SESSION;
913
use vortex_array::ToCanonical;
14+
use vortex_array::VortexSessionExecute;
1015
use vortex_array::arrays::ConstantArray;
16+
use vortex_array::arrays::Patched;
1117
use vortex_array::arrays::primitive::PrimitiveArrayExt;
1218
use vortex_array::scalar::Scalar;
1319
use vortex_compressor::builtins::FloatDictScheme;
@@ -20,7 +26,7 @@ use vortex_error::VortexExpect;
2026
use vortex_error::VortexResult;
2127
use vortex_error::vortex_bail;
2228
use vortex_error::vortex_err;
23-
use vortex_fastlanes::BitPackedArrayExt;
29+
use vortex_fastlanes::BitPacked;
2430
use vortex_fastlanes::FoR;
2531
use vortex_fastlanes::FoRArrayExt;
2632
use vortex_fastlanes::bitpack_compress::bit_width_histogram;
@@ -279,6 +285,10 @@ impl Scheme for ZigZagScheme {
279285
}
280286
}
281287

288+
// replicated from vortex-file
289+
static USE_EXPERIMENTAL_PATCHES: LazyLock<bool> =
290+
LazyLock::new(|| env::var("VORTEX_EXPERIMENTAL_PATCHED_ARRAY").is_ok());
291+
282292
impl Scheme for BitPackingScheme {
283293
fn scheme_name(&self) -> &'static str {
284294
"vortex.int.bitpacking"
@@ -324,21 +334,49 @@ impl Scheme for BitPackingScheme {
324334

325335
let packed_stats = packed.statistics().to_owned();
326336
let ptype = packed.dtype().as_ptype();
327-
let patches = packed.patches().map(compress_patches).transpose()?;
328-
let mut parts = vortex_fastlanes::BitPacked::into_parts(packed);
329-
parts.patches = patches;
330-
331-
Ok(vortex_fastlanes::BitPacked::try_new(
332-
parts.packed,
333-
ptype,
334-
parts.validity,
335-
parts.patches,
336-
parts.bit_width,
337-
parts.len,
338-
parts.offset,
339-
)?
340-
.with_stats_set(packed_stats)
341-
.into_array())
337+
let mut parts = BitPacked::into_parts(packed);
338+
339+
let array = if *USE_EXPERIMENTAL_PATCHES {
340+
let patches = parts.patches.take();
341+
// Transpose patches into G-ALP style PatchedArray, wrapping an inner BitPackedArray.
342+
let array = BitPacked::try_new(
343+
parts.packed,
344+
ptype,
345+
parts.validity,
346+
None,
347+
parts.bit_width,
348+
parts.len,
349+
parts.offset,
350+
)?
351+
.into_array();
352+
353+
match patches {
354+
None => array,
355+
Some(p) => Patched::from_array_and_patches(
356+
array,
357+
&p,
358+
&mut LEGACY_SESSION.create_execution_ctx(),
359+
)?
360+
.into_array(),
361+
}
362+
} else {
363+
// Compress patches and place back into BitPackedArray.
364+
let patches = parts.patches.take().map(compress_patches).transpose()?;
365+
parts.patches = patches;
366+
BitPacked::try_new(
367+
parts.packed,
368+
ptype,
369+
parts.validity,
370+
parts.patches,
371+
parts.bit_width,
372+
parts.len,
373+
parts.offset,
374+
)?
375+
.with_stats_set(packed_stats)
376+
.into_array()
377+
};
378+
379+
Ok(array)
342380
}
343381
}
344382

0 commit comments

Comments
 (0)