Skip to content

Commit dafbfa4

Browse files
mhk197mhk197
committed
Read the permitted serialized IDs from the CompressorContext
The compressor seeds each root CompressorContext with its permitted serialized IDs and every descent inherits them, so a scheme asks compress_ctx.allows_serialized_id both while estimating and while compressing and picks the same mode in both. The per-compressor accessor goes; allowed_serialized_ids remains for inspection. Signed-off-by: Matt Katz <mhkatz97@gmail.com>
1 parent 98adaee commit dafbfa4

5 files changed

Lines changed: 127 additions & 58 deletions

File tree

‎vortex-btrblocks/src/builder.rs‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -371,13 +371,15 @@ mod tests {
371371
#[test]
372372
fn allowed_serialized_ids_reach_the_compressor() {
373373
let default = BtrBlocksCompressorBuilder::default().build();
374-
assert!(default.0.allows_serialized_id(Bool.id()));
374+
assert!(default.0.allowed_serialized_ids().is_none());
375375

376376
let narrowed = BtrBlocksCompressorBuilder::default()
377377
.allow_serialized_ids(&HashSet::from([FoR.id()]))
378378
.build();
379-
assert!(narrowed.0.allows_serialized_id(FoR.id()));
380-
assert!(!narrowed.0.allows_serialized_id(Bool.id()));
379+
assert_eq!(
380+
narrowed.0.allowed_serialized_ids(),
381+
Some(&HashSet::from([FoR.id()]))
382+
);
381383
}
382384

383385
#[test]

‎vortex-compressor/src/compressor/cascade.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ impl CascadingCompressor {
5959

6060
let canonical = array.clone().execute::<CanonicalValidity>(exec_ctx)?.0;
6161
let compact = canonical.compact(exec_ctx)?;
62-
let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?;
62+
let compressed = self.compress_canonical(compact, self.root_context(), exec_ctx)?;
6363

6464
trace::record_compress_outcome(&span, before_nbytes, compressed.nbytes());
6565

‎vortex-compressor/src/compressor/mod.rs‎

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,14 @@ mod sample;
99
mod select;
1010
mod structural;
1111

12+
use std::sync::Arc;
13+
1214
use vortex_array::ArrayId;
1315
use vortex_utils::aliases::hash_set::HashSet;
1416

1517
use crate::builtins::IntDictScheme;
1618
use crate::scheme::ChildSelection;
19+
use crate::scheme::CompressorContext;
1720
use crate::scheme::DescendantExclusion;
1821
use crate::scheme::Scheme;
1922
use crate::scheme::SchemeExt;
@@ -50,9 +53,9 @@ pub struct CascadingCompressor {
5053
/// list offsets).
5154
root_exclusions: Vec<DescendantExclusion>,
5255

53-
/// The serialized IDs the output may use, or `None` for no restriction. See
54-
/// [`allows_serialized_id`](Self::allows_serialized_id).
55-
allowed_serialized_ids: Option<HashSet<ArrayId>>,
56+
/// The serialized IDs the writer may emit, or `None` for no restriction. Seeds every root
57+
/// [`CompressorContext`], where schemes read it.
58+
allowed_serialized_ids: Option<Arc<HashSet<ArrayId>>>,
5659
}
5760

5861
impl CascadingCompressor {
@@ -77,27 +80,26 @@ impl CascadingCompressor {
7780
/// Hands the compressor the serialized IDs the writer may emit, intersecting with any earlier
7881
/// call.
7982
///
80-
/// The file writer passes the serialized IDs its enabled editions permit. A scheme whose
81-
/// encoding has several wire formats picks its compression mode from this set, the newest
82-
/// permitted one, before estimating or compressing.
83+
/// The file writer passes the serialized IDs its enabled editions permit. Schemes read the
84+
/// set through [`CompressorContext::allows_serialized_id`], so a scheme whose encoding has
85+
/// several wire formats picks the newest permitted one as its mode, while estimating and
86+
/// while compressing alike.
8387
pub fn with_allowed_serialized_ids(mut self, allowed: HashSet<ArrayId>) -> Self {
84-
self.allowed_serialized_ids = Some(match self.allowed_serialized_ids.take() {
88+
self.allowed_serialized_ids = Some(Arc::new(match self.allowed_serialized_ids.take() {
8589
Some(existing) => existing.intersection(&allowed).copied().collect(),
8690
None => allowed,
87-
});
91+
}));
8892
self
8993
}
9094

91-
/// Returns whether the writer may emit the serialized ID `id`.
92-
///
93-
/// Schemes whose encoding has several wire formats consult this to pick their compression
94-
/// mode. Without a restriction every ID is allowed, so the newest mode is chosen. The
95-
/// serializer still emits the oldest wire form the resulting array fits, and the
96-
/// serialization context validates that ID.
97-
pub fn allows_serialized_id(&self, id: ArrayId) -> bool {
98-
self.allowed_serialized_ids
99-
.as_ref()
100-
.is_none_or(|allowed| allowed.contains(&id))
95+
/// The serialized IDs the writer may emit, or `None` when unrestricted.
96+
pub fn allowed_serialized_ids(&self) -> Option<&HashSet<ArrayId>> {
97+
self.allowed_serialized_ids.as_deref()
98+
}
99+
100+
/// The context a compress call starts from.
101+
pub(crate) fn root_context(&self) -> CompressorContext {
102+
CompressorContext::new(self.allowed_serialized_ids.clone())
101103
}
102104

103105
/// Returns whether the compressor was configured with `scheme`.

‎vortex-compressor/src/compressor/tests.rs‎

Lines changed: 77 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -132,12 +132,12 @@ impl Scheme for FormatRecordingScheme {
132132

133133
fn compress(
134134
&self,
135-
compressor: &CascadingCompressor,
135+
_compressor: &CascadingCompressor,
136136
data: &ArrayAndStats,
137-
_compress_ctx: CompressorContext,
137+
compress_ctx: CompressorContext,
138138
_exec_ctx: &mut ExecutionCtx,
139139
) -> VortexResult<ArrayRef> {
140-
*SEEN_FORMAT.lock() = Some(compressor.allows_serialized_id(Constant.id()));
140+
*SEEN_FORMAT.lock() = Some(compress_ctx.allows_serialized_id(Constant.id()));
141141
Ok(data.array().clone())
142142
}
143143
}
@@ -420,8 +420,12 @@ fn immediate_always_use_wins_immediately() -> VortexResult<()> {
420420
let data = estimate_test_data();
421421
let mut exec_ctx = SESSION.create_execution_ctx();
422422

423-
let winner =
424-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
423+
let winner = compressor.choose_best_scheme(
424+
&schemes,
425+
&data,
426+
CompressorContext::new(None),
427+
&mut exec_ctx,
428+
)?;
425429

426430
assert!(matches!(
427431
winner,
@@ -438,8 +442,12 @@ fn callback_always_use_wins_immediately() -> VortexResult<()> {
438442
let data = estimate_test_data();
439443
let mut exec_ctx = SESSION.create_execution_ctx();
440444

441-
let winner =
442-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
445+
let winner = compressor.choose_best_scheme(
446+
&schemes,
447+
&data,
448+
CompressorContext::new(None),
449+
&mut exec_ctx,
450+
)?;
443451

444452
assert!(matches!(
445453
winner,
@@ -456,8 +464,12 @@ fn callback_skip_is_ignored() -> VortexResult<()> {
456464
let data = estimate_test_data();
457465
let mut exec_ctx = SESSION.create_execution_ctx();
458466

459-
let winner =
460-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
467+
let winner = compressor.choose_best_scheme(
468+
&schemes,
469+
&data,
470+
CompressorContext::new(None),
471+
&mut exec_ctx,
472+
)?;
461473

462474
assert!(matches!(
463475
winner,
@@ -474,8 +486,12 @@ fn callback_ratio_competes_numerically() -> VortexResult<()> {
474486
let data = estimate_test_data();
475487
let mut exec_ctx = SESSION.create_execution_ctx();
476488

477-
let winner =
478-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
489+
let winner = compressor.choose_best_scheme(
490+
&schemes,
491+
&data,
492+
CompressorContext::new(None),
493+
&mut exec_ctx,
494+
)?;
479495

480496
assert!(matches!(
481497
winner,
@@ -492,8 +508,12 @@ fn zero_byte_sample_loses_to_finite_ratio() -> VortexResult<()> {
492508
let data = estimate_test_data();
493509
let mut exec_ctx = SESSION.create_execution_ctx();
494510

495-
let winner =
496-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
511+
let winner = compressor.choose_best_scheme(
512+
&schemes,
513+
&data,
514+
CompressorContext::new(None),
515+
&mut exec_ctx,
516+
)?;
497517

498518
assert!(matches!(
499519
winner,
@@ -510,8 +530,12 @@ fn finite_ratio_displaces_zero_byte_sample() -> VortexResult<()> {
510530
let data = estimate_test_data();
511531
let mut exec_ctx = SESSION.create_execution_ctx();
512532

513-
let winner =
514-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
533+
let winner = compressor.choose_best_scheme(
534+
&schemes,
535+
&data,
536+
CompressorContext::new(None),
537+
&mut exec_ctx,
538+
)?;
515539

516540
assert!(matches!(
517541
winner,
@@ -528,8 +552,12 @@ fn zero_byte_sample_alone_selects_no_scheme() -> VortexResult<()> {
528552
let data = estimate_test_data();
529553
let mut exec_ctx = SESSION.create_execution_ctx();
530554

531-
let winner =
532-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
555+
let winner = compressor.choose_best_scheme(
556+
&schemes,
557+
&data,
558+
CompressorContext::new(None),
559+
&mut exec_ctx,
560+
)?;
533561

534562
assert!(winner.is_none());
535563
Ok(())
@@ -630,8 +658,12 @@ fn callback_always_use_overrides_pass_one_best() -> VortexResult<()> {
630658
let data = estimate_test_data();
631659
let mut exec_ctx = SESSION.create_execution_ctx();
632660

633-
let winner =
634-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
661+
let winner = compressor.choose_best_scheme(
662+
&schemes,
663+
&data,
664+
CompressorContext::new(None),
665+
&mut exec_ctx,
666+
)?;
635667

636668
assert!(matches!(
637669
winner,
@@ -651,7 +683,7 @@ fn threshold_reflects_pass_one_best() -> VortexResult<()> {
651683
let data = estimate_test_data();
652684
let mut exec_ctx = SESSION.create_execution_ctx();
653685

654-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
686+
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?;
655687

656688
let observed = *OBSERVED_THRESHOLD.lock();
657689
assert!(matches!(
@@ -672,7 +704,7 @@ fn threshold_is_none_when_only_prior_is_zero_bytes() -> VortexResult<()> {
672704
let data = estimate_test_data();
673705
let mut exec_ctx = SESSION.create_execution_ctx();
674706

675-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
707+
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?;
676708

677709
// The observing callback was invoked (outer `Some`) and `best_so_far` was `None` (inner
678710
// `None`) because the zero-byte sample is never stored as the best.
@@ -691,7 +723,7 @@ fn threshold_is_none_when_no_prior_scheme() -> VortexResult<()> {
691723
let data = estimate_test_data();
692724
let mut exec_ctx = SESSION.create_execution_ctx();
693725

694-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
726+
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?;
695727

696728
let observed = *OBSERVED_THRESHOLD.lock();
697729
assert_eq!(observed, Some(None));
@@ -711,7 +743,7 @@ fn threshold_updates_from_earlier_deferred_callback() -> VortexResult<()> {
711743
let data = estimate_test_data();
712744
let mut exec_ctx = SESSION.create_execution_ctx();
713745

714-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
746+
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(None), &mut exec_ctx)?;
715747

716748
let observed = *OBSERVED_THRESHOLD.lock();
717749
assert!(matches!(
@@ -732,8 +764,12 @@ fn ratio_tie_between_immediate_and_deferred_favors_immediate() -> VortexResult<(
732764
let data = estimate_test_data();
733765
let mut exec_ctx = SESSION.create_execution_ctx();
734766

735-
let winner =
736-
compressor.choose_best_scheme(&schemes, &data, CompressorContext::new(), &mut exec_ctx)?;
767+
let winner = compressor.choose_best_scheme(
768+
&schemes,
769+
&data,
770+
CompressorContext::new(None),
771+
&mut exec_ctx,
772+
)?;
737773

738774
assert!(matches!(
739775
winner,
@@ -780,7 +816,7 @@ fn sampling_uses_scheme_stats_options() -> VortexResult<()> {
780816
// A context with default stats_options (count_distinct_values = false) and
781817
// marked as a sample so the function skips the sampling step and compresses
782818
// the array directly.
783-
let ctx = CompressorContext::new().with_sampling();
819+
let ctx = CompressorContext::new(None).with_sampling();
784820

785821
// Before the fix this panicked with:
786822
// "this must be present since `DictScheme` declared that we need distinct values"
@@ -891,19 +927,27 @@ fn map_compression_preserves_repeated_entry_children() -> VortexResult<()> {
891927
#[test]
892928
fn allowed_serialized_ids_default_to_everything_and_intersect() {
893929
let compressor = compressor();
894-
assert!(compressor.allows_serialized_id(Constant.id()));
895-
assert!(compressor.allows_serialized_id(Bool.id()));
930+
let root = compressor.root_context();
931+
assert!(root.allows_serialized_id(Constant.id()));
932+
assert!(root.allows_serialized_id(Bool.id()));
896933

897934
let restricted =
898935
compressor.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Constant.id()]));
899-
assert!(restricted.allows_serialized_id(Constant.id()));
900-
assert!(!restricted.allows_serialized_id(Bool.id()));
936+
let root = restricted.root_context();
937+
assert!(root.allows_serialized_id(Constant.id()));
938+
assert!(!root.allows_serialized_id(Bool.id()));
901939

902940
let narrowed =
903941
restricted.with_allowed_serialized_ids(HashSet::from([Primitive.id(), Bool.id()]));
904-
assert!(narrowed.allows_serialized_id(Primitive.id()));
905-
assert!(!narrowed.allows_serialized_id(Constant.id()));
906-
assert!(!narrowed.allows_serialized_id(Bool.id()));
942+
let root = narrowed.root_context();
943+
assert!(root.allows_serialized_id(Primitive.id()));
944+
assert!(!root.allows_serialized_id(Constant.id()));
945+
assert!(!root.allows_serialized_id(Bool.id()));
946+
947+
// Descending keeps the set.
948+
let child = root.descend_with_scheme(IntDictScheme.id(), 0);
949+
assert!(child.allows_serialized_id(Primitive.id()));
950+
assert!(!child.allows_serialized_id(Constant.id()));
907951
}
908952

909953
/// A scheme sees the restriction through the compressor it is handed: everything is allowed until

‎vortex-compressor/src/scheme/ctx.rs‎

Lines changed: 24 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,11 @@
44
//! Compression context for recursive compression.
55
66
use std::fmt;
7+
use std::sync::Arc;
78

9+
use vortex_array::ArrayId;
810
use vortex_error::VortexExpect;
11+
use vortex_utils::aliases::hash_set::HashSet;
912

1013
use crate::compressor::ROOT_SCHEME_ID;
1114
use crate::scheme::SchemeId;
@@ -38,26 +41,32 @@ pub struct CompressorContext {
3841
/// [`descendant_exclusions`]: crate::scheme::Scheme::descendant_exclusions
3942
/// [`ancestor_exclusions`]: crate::scheme::Scheme::ancestor_exclusions
4043
cascade_history: Vec<(SchemeId, usize)>,
44+
45+
/// The serialized IDs the writer may emit, or `None` for no restriction. Shared by every
46+
/// context of one compress call, so cloning at each descent is a pointer bump.
47+
allowed_serialized_ids: Option<Arc<HashSet<ArrayId>>>,
4148
}
4249

4350
impl CompressorContext {
44-
/// Creates a new `CompressorContext`.
51+
/// Creates a new root `CompressorContext` for a compressor that may emit the given serialized
52+
/// IDs, or any ID when `None`.
4553
///
4654
/// This should **only** be created by the compressor.
47-
pub(crate) fn new() -> Self {
55+
pub(crate) fn new(allowed_serialized_ids: Option<Arc<HashSet<ArrayId>>>) -> Self {
4856
Self {
4957
is_sample: false,
5058
allowed_cascading: MAX_CASCADE,
5159
merged_stats_options: GenerateStatsOptions::default(),
5260
cascade_history: Vec::new(),
61+
allowed_serialized_ids,
5362
}
5463
}
5564
}
5665

5766
#[cfg(test)]
5867
impl Default for CompressorContext {
5968
fn default() -> Self {
60-
Self::new()
69+
Self::new(None)
6170
}
6271
}
6372

@@ -67,6 +76,18 @@ impl CompressorContext {
6776
self.is_sample
6877
}
6978

79+
/// Returns whether the writer may emit the serialized ID `id`.
80+
///
81+
/// A scheme whose encoding has several wire formats picks its compression mode from this,
82+
/// the newest permitted one, and the same answer is available while estimating and while
83+
/// compressing. Without a restriction every ID is allowed. The serializer still emits the
84+
/// oldest wire form the resulting array fits, and the serialization context validates it.
85+
pub fn allows_serialized_id(&self, id: ArrayId) -> bool {
86+
self.allowed_serialized_ids
87+
.as_ref()
88+
.is_none_or(|allowed| allowed.contains(&id))
89+
}
90+
7091
/// Returns the merged stats generation options for this compression site.
7192
pub fn merged_stats_options(&self) -> GenerateStatsOptions {
7293
self.merged_stats_options

0 commit comments

Comments
 (0)