From eb3f7e7143f8a1897935c753dd1abd1c6bc21a38 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 16 Sep 2026 11:51:18 -0400 Subject: [PATCH 1/2] Filter compression schemes by the serialized IDs they write Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 15 ++++--- vortex-btrblocks/src/schemes/binary/varbin.rs | 2 +- vortex-btrblocks/src/schemes/binary/zstd.rs | 2 +- .../src/schemes/binary/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/decimal.rs | 2 +- vortex-btrblocks/src/schemes/float/alp.rs | 2 +- vortex-btrblocks/src/schemes/float/alprd.rs | 2 +- vortex-btrblocks/src/schemes/float/pco.rs | 2 +- vortex-btrblocks/src/schemes/float/rle.rs | 2 +- vortex-btrblocks/src/schemes/float/sparse.rs | 2 +- .../src/schemes/integer/bitpacking.rs | 2 +- vortex-btrblocks/src/schemes/integer/delta.rs | 2 +- vortex-btrblocks/src/schemes/integer/for_.rs | 2 +- vortex-btrblocks/src/schemes/integer/pco.rs | 2 +- vortex-btrblocks/src/schemes/integer/rle.rs | 2 +- .../src/schemes/integer/runend.rs | 2 +- .../src/schemes/integer/sequence.rs | 2 +- .../src/schemes/integer/sparse.rs | 2 +- .../src/schemes/integer/zigzag.rs | 2 +- vortex-btrblocks/src/schemes/string/fsst.rs | 2 +- vortex-btrblocks/src/schemes/string/onpair.rs | 2 +- vortex-btrblocks/src/schemes/string/sparse.rs | 2 +- vortex-btrblocks/src/schemes/string/zstd.rs | 2 +- .../src/schemes/string/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/temporal.rs | 2 +- vortex-compressor/src/builtins/dict/binary.rs | 2 +- vortex-compressor/src/builtins/dict/float.rs | 2 +- .../src/builtins/dict/integer.rs | 2 +- vortex-compressor/src/builtins/dict/string.rs | 2 +- vortex-compressor/src/compressor/tests.rs | 18 ++++----- vortex-compressor/src/scheme/mod.rs | 14 ++++--- vortex-file/src/writer.rs | 39 ++++++++----------- 32 files changed, 71 insertions(+), 71 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index c6faaba93d8..7870e8517fe 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -200,13 +200,16 @@ impl BtrBlocksCompressorBuilder { self } - /// Retains only schemes whose produced encodings all belong to `allowed`. + /// Retains only schemes whose produced serialized IDs all belong to `allowed`. /// - /// The file writer uses this to restrict compression to the encodings of its configured - /// editions. + /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions + /// permit. pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { - self.schemes - .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); + self.schemes.retain(|s| { + s.produced_serialized_ids() + .iter() + .all(|id| allowed.contains(id)) + }); self } @@ -250,7 +253,7 @@ mod tests { fn retaining_all_declared_outputs_keeps_every_scheme() { let allowed: HashSet = ALL_SCHEMES .iter() - .flat_map(|scheme| scheme.produced_encodings()) + .flat_map(|scheme| scheme.produced_serialized_ids()) .collect(); let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); diff --git a/vortex-btrblocks/src/schemes/binary/varbin.rs b/vortex-btrblocks/src/schemes/binary/varbin.rs index 849402a6493..bf8237ed85f 100644 --- a/vortex-btrblocks/src/schemes/binary/varbin.rs +++ b/vortex-btrblocks/src/schemes/binary/varbin.rs @@ -44,7 +44,7 @@ impl Scheme for VarBinScheme { canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![VarBin.id()] } diff --git a/vortex-btrblocks/src/schemes/binary/zstd.rs b/vortex-btrblocks/src/schemes/binary/zstd.rs index d652e344db2..3f304860c9c 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdScheme { canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![vortex_zstd::Zstd.id()] } diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index 3f06d65b061..c240fd006a6 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index 1dff2171f60..af8f7f1998a 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -40,7 +40,7 @@ impl Scheme for DecimalScheme { matches!(canonical, Canonical::Decimal(_)) } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![DecimalByteParts.id()] } diff --git a/vortex-btrblocks/src/schemes/float/alp.rs b/vortex-btrblocks/src/schemes/float/alp.rs index f9fc7066bf4..bc39f434013 100644 --- a/vortex-btrblocks/src/schemes/float/alp.rs +++ b/vortex-btrblocks/src/schemes/float/alp.rs @@ -42,7 +42,7 @@ impl Scheme for ALPScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { let mut encodings = vec![ALP.id()]; if use_experimental_patches() { encodings.push(Patched.id()); diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/vortex-btrblocks/src/schemes/float/alprd.rs index 09c0ee85b0d..9efecaf83c6 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -40,7 +40,7 @@ impl Scheme for ALPRDScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![vortex_alp::ALPRD.id()] } diff --git a/vortex-btrblocks/src/schemes/float/pco.rs b/vortex-btrblocks/src/schemes/float/pco.rs index 416668c2fd0..bf517a21dfe 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/vortex-btrblocks/src/schemes/float/pco.rs @@ -31,7 +31,7 @@ impl Scheme for PcoScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![vortex_pco::Pco.id()] } diff --git a/vortex-btrblocks/src/schemes/float/rle.rs b/vortex-btrblocks/src/schemes/float/rle.rs index 71158b9dc3b..6b19eb577ee 100644 --- a/vortex-btrblocks/src/schemes/float/rle.rs +++ b/vortex-btrblocks/src/schemes/float/rle.rs @@ -38,7 +38,7 @@ impl Scheme for FloatRLEScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![RLE.id()] } diff --git a/vortex-btrblocks/src/schemes/float/sparse.rs b/vortex-btrblocks/src/schemes/float/sparse.rs index 3d9c25b18e4..d4e18683baf 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/vortex-btrblocks/src/schemes/float/sparse.rs @@ -41,7 +41,7 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![Sparse.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index 5ac7d0e4078..a76826b5ce6 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -40,7 +40,7 @@ impl Scheme for BitPackingScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { let mut encodings = vec![BitPacked.id()]; if use_experimental_patches() { encodings.push(Patched.id()); diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 46b2f1e302e..8fc0a06a3ab 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -97,7 +97,7 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![Delta.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index 476a0dec282..38299c2a425 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -44,7 +44,7 @@ impl Scheme for FoRScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![FoR.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index 675a112d44f..13cf16544da 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -32,7 +32,7 @@ impl Scheme for PcoScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![vortex_pco::Pco.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/vortex-btrblocks/src/schemes/integer/rle.rs index 0909e1428d4..3b617b4f212 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/vortex-btrblocks/src/schemes/integer/rle.rs @@ -144,7 +144,7 @@ impl Scheme for IntRLEScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![RLE.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/runend.rs b/vortex-btrblocks/src/schemes/integer/runend.rs index 6a97f7ec37d..ce42712786e 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/vortex-btrblocks/src/schemes/integer/runend.rs @@ -48,7 +48,7 @@ impl Scheme for RunEndScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![RunEnd.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/sequence.rs b/vortex-btrblocks/src/schemes/integer/sequence.rs index edcefb99fc2..35139c7e62d 100644 --- a/vortex-btrblocks/src/schemes/integer/sequence.rs +++ b/vortex-btrblocks/src/schemes/integer/sequence.rs @@ -43,7 +43,7 @@ impl Scheme for SequenceScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![Sequence.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/sparse.rs b/vortex-btrblocks/src/schemes/integer/sparse.rs index 429ff5c1a31..44ffbd59ef2 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/vortex-btrblocks/src/schemes/integer/sparse.rs @@ -46,7 +46,7 @@ impl Scheme for SparseScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![Sparse.id(), Constant.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/zigzag.rs b/vortex-btrblocks/src/schemes/integer/zigzag.rs index 0e4be01845a..af6ca3df7a3 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/vortex-btrblocks/src/schemes/integer/zigzag.rs @@ -46,7 +46,7 @@ impl Scheme for ZigZagScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![ZigZag.id()] } diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index fd3fd28696a..dd720f38de1 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -50,7 +50,7 @@ impl Scheme for FSSTScheme { canonical.dtype().is_utf8() || canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![FSST.id(), VarBin.id()] } diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index a1bc8643775..e5358cfddbd 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -50,7 +50,7 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![OnPair.id()] } diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index 8620c366f77..f6392cddaab 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -42,7 +42,7 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![Sparse.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd.rs b/vortex-btrblocks/src/schemes/string/zstd.rs index 84e8860d626..faa1c98c358 100644 --- a/vortex-btrblocks/src/schemes/string/zstd.rs +++ b/vortex-btrblocks/src/schemes/string/zstd.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![vortex_zstd::Zstd.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index cf691c70fcb..3f524b65cce 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index 79748b69450..9afa52b206e 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -55,7 +55,7 @@ impl Scheme for TemporalScheme { ) } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![DateTimeParts.id()] } diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index c407d0251c6..abf62f914c5 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -48,7 +48,7 @@ impl Scheme for BinaryDictScheme { canonical.dtype().is_binary() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index f962c3ff967..4d9ed43319a 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -54,7 +54,7 @@ impl Scheme for FloatDictScheme { canonical.dtype().is_float() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 27a17ef94ad..1c4e5d5ba0e 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -49,7 +49,7 @@ impl Scheme for IntDictScheme { canonical.dtype().is_int() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index f5cbcd54d89..05abec9f946 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -48,7 +48,7 @@ impl Scheme for StringDictScheme { canonical.dtype().is_utf8() } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 3a2a6281047..b91ac594e00 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -72,7 +72,7 @@ impl Scheme for DirectRatioScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { Vec::new() } @@ -108,7 +108,7 @@ impl Scheme for ImmediateAlwaysUseScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { Vec::new() } @@ -144,7 +144,7 @@ impl Scheme for CallbackAlwaysUseScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { Vec::new() } @@ -182,7 +182,7 @@ impl Scheme for CallbackSkipScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { Vec::new() } @@ -220,7 +220,7 @@ impl Scheme for CallbackRatioScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { Vec::new() } @@ -258,7 +258,7 @@ impl Scheme for HugeRatioScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { Vec::new() } @@ -294,7 +294,7 @@ impl Scheme for ZeroBytesSamplingScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { Vec::new() } @@ -507,7 +507,7 @@ impl Scheme for ThresholdObservingScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { Vec::new() } @@ -548,7 +548,7 @@ impl Scheme for CallbackMatchingRatioScheme { matches_integer_primitive(canonical) } - fn produced_encodings(&self) -> Vec { + fn produced_serialized_ids(&self) -> Vec { Vec::new() } diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index de9e67690d4..1b5bd801252 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -124,12 +124,16 @@ pub trait Scheme: Debug + Send + Sync { /// Whether this scheme can compress the given canonical array. fn matches(&self, canonical: &Canonical) -> bool; - /// The array encodings this scheme itself may introduce into its compressed output. + /// The serialized IDs this scheme itself may write into its compressed output. /// - /// Cascaded children are compressed by other schemes, which declare their own encodings, - /// so only encodings constructed directly by [`compress`](Scheme::compress) belong here. - /// Canonical arrays the scheme merely rearranges do not need to be declared. - fn produced_encodings(&self) -> Vec; + /// Every declared ID must be permitted for the scheme to be used. Cascaded children are + /// compressed by other schemes, which declare their own IDs, so only arrays constructed + /// directly by [`compress`](Scheme::compress) belong here. Canonical arrays the scheme + /// merely rearranges do not need to be declared. + /// + /// For most encodings this is the in-memory encoding ID. An encoding with several wire + /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID. + fn produced_serialized_ids(&self) -> Vec; /// Returns the stats generation options this scheme requires. The compressor merges all /// eligible schemes' options before generating stats so that a single stats pass satisfies diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 8167d31ede2..7b411074635 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -239,8 +239,7 @@ impl VortexWriteOptions { let enforce_editions = !self.disable_editions; // The array context is built here, rather than when the options were constructed, so that // encodings registered on the session in between are still eligible for the file. - let (array_ctx, allowed_array_encodings) = - new_array_context(&self.session, enforce_editions); + let (array_ctx, permitted_ids) = new_array_context(&self.session, enforce_editions); let ctx = LayoutWriterContext::new(array_ctx) .with_buffered_bytes_tracker(self.buffered_bytes.clone()); let ctx = if enforce_editions { @@ -252,8 +251,7 @@ impl VortexWriteOptions { Some(strategy) => strategy, None => WriteStrategyBuilder::default() .with_btrblocks_builder( - BtrBlocksCompressorBuilder::default() - .retain_allowed_encodings(&allowed_array_encodings), + BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&permitted_ids), ) .build(), }; @@ -401,11 +399,8 @@ fn new_array_context( .registry() .read(|registry| registry.keys().copied().collect()) }; - let allowed_array_encodings = serialized_ids - .iter() - .filter_map(|serialized_id| arrays.registry().get(serialized_id)) - .map(|plugin| plugin.id()) - .collect(); + // Compression schemes declare the serialized IDs they write, so the same set restricts them. + let permitted_ids = serialized_ids.iter().copied().collect(); let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect()); let array_ctx = if enforce_editions { // Only permit serialized IDs in the enabled editions. @@ -413,7 +408,7 @@ fn new_array_context( } else { array_ctx }; - (array_ctx, allowed_array_encodings) + (array_ctx, permitted_ids) } /// The ids of `kind` the enabled editions permit. @@ -787,29 +782,27 @@ mod tests { session.register_edition(&DECLARATION)?; session.enable_edition(EDITION)?; - let (ctx, allowed_array_encodings) = new_array_context(&session, true); + let (ctx, permitted_ids) = new_array_context(&session, true); assert_eq!(ctx.to_ids(), [Primitive.id()]); assert!(ctx.intern(&Bool.id()).is_none()); - assert_eq!(allowed_array_encodings, HashSet::from([Primitive.id()])); + assert_eq!(permitted_ids, HashSet::from([Primitive.id()])); Ok(()) } #[test] fn disabling_editions_allows_all_registered_array_ids() { let session = array_session(); - let (registered_ids, registered_encodings) = session.arrays().registry().read(|registry| { - ( - registry.keys().copied().sorted().collect::>(), - registry - .values() - .map(|plugin| plugin.id()) - .collect::>(), - ) - }); + let registered_ids = session + .arrays() + .registry() + .read(|registry| registry.keys().copied().sorted().collect::>()); - let (ctx, allowed_array_encodings) = new_array_context(&session, false); + let (ctx, permitted_ids) = new_array_context(&session, false); assert_eq!(ctx.to_ids(), registered_ids); - assert_eq!(allowed_array_encodings, registered_encodings); + assert_eq!( + permitted_ids, + registered_ids.iter().copied().collect::>() + ); assert!(ctx.intern(&Bool.id()).is_some()); } From d308072d7c911adc6ac7c8e8dc29ae11bdbdad14 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Wed, 16 Sep 2026 12:12:34 -0400 Subject: [PATCH 2/2] Keep produced_encodings and read the writer's permitted IDs from the array context Signed-off-by: Matt Katz --- vortex-btrblocks/src/builder.rs | 9 ++--- vortex-btrblocks/src/schemes/binary/varbin.rs | 2 +- vortex-btrblocks/src/schemes/binary/zstd.rs | 2 +- .../src/schemes/binary/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/decimal.rs | 2 +- vortex-btrblocks/src/schemes/float/alp.rs | 2 +- vortex-btrblocks/src/schemes/float/alprd.rs | 2 +- vortex-btrblocks/src/schemes/float/pco.rs | 2 +- vortex-btrblocks/src/schemes/float/rle.rs | 2 +- vortex-btrblocks/src/schemes/float/sparse.rs | 2 +- .../src/schemes/integer/bitpacking.rs | 2 +- vortex-btrblocks/src/schemes/integer/delta.rs | 2 +- vortex-btrblocks/src/schemes/integer/for_.rs | 2 +- vortex-btrblocks/src/schemes/integer/pco.rs | 2 +- vortex-btrblocks/src/schemes/integer/rle.rs | 2 +- .../src/schemes/integer/runend.rs | 2 +- .../src/schemes/integer/sequence.rs | 2 +- .../src/schemes/integer/sparse.rs | 2 +- .../src/schemes/integer/zigzag.rs | 2 +- vortex-btrblocks/src/schemes/string/fsst.rs | 2 +- vortex-btrblocks/src/schemes/string/onpair.rs | 2 +- vortex-btrblocks/src/schemes/string/sparse.rs | 2 +- vortex-btrblocks/src/schemes/string/zstd.rs | 2 +- .../src/schemes/string/zstd_buffers.rs | 2 +- vortex-btrblocks/src/schemes/temporal.rs | 2 +- vortex-compressor/src/builtins/dict/binary.rs | 2 +- vortex-compressor/src/builtins/dict/float.rs | 2 +- .../src/builtins/dict/integer.rs | 2 +- vortex-compressor/src/builtins/dict/string.rs | 2 +- vortex-compressor/src/compressor/tests.rs | 18 ++++----- vortex-compressor/src/scheme/mod.rs | 2 +- vortex-file/src/writer.rs | 37 ++++++++----------- 32 files changed, 57 insertions(+), 65 deletions(-) diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 7870e8517fe..341569f408c 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -205,11 +205,8 @@ impl BtrBlocksCompressorBuilder { /// `allowed` holds serialized IDs. The file writer passes the array IDs its enabled editions /// permit. pub fn retain_allowed_encodings(mut self, allowed: &HashSet) -> Self { - self.schemes.retain(|s| { - s.produced_serialized_ids() - .iter() - .all(|id| allowed.contains(id)) - }); + self.schemes + .retain(|s| s.produced_encodings().iter().all(|id| allowed.contains(id))); self } @@ -253,7 +250,7 @@ mod tests { fn retaining_all_declared_outputs_keeps_every_scheme() { let allowed: HashSet = ALL_SCHEMES .iter() - .flat_map(|scheme| scheme.produced_serialized_ids()) + .flat_map(|scheme| scheme.produced_encodings()) .collect(); let builder = BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&allowed); assert_eq!(builder.schemes.len(), ALL_SCHEMES.len()); diff --git a/vortex-btrblocks/src/schemes/binary/varbin.rs b/vortex-btrblocks/src/schemes/binary/varbin.rs index bf8237ed85f..849402a6493 100644 --- a/vortex-btrblocks/src/schemes/binary/varbin.rs +++ b/vortex-btrblocks/src/schemes/binary/varbin.rs @@ -44,7 +44,7 @@ impl Scheme for VarBinScheme { canonical.dtype().is_binary() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![VarBin.id()] } diff --git a/vortex-btrblocks/src/schemes/binary/zstd.rs b/vortex-btrblocks/src/schemes/binary/zstd.rs index 3f304860c9c..d652e344db2 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdScheme { canonical.dtype().is_binary() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_zstd::Zstd.id()] } diff --git a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs index c240fd006a6..3f06d65b061 100644 --- a/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/binary/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_binary() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/decimal.rs b/vortex-btrblocks/src/schemes/decimal.rs index af8f7f1998a..1dff2171f60 100644 --- a/vortex-btrblocks/src/schemes/decimal.rs +++ b/vortex-btrblocks/src/schemes/decimal.rs @@ -40,7 +40,7 @@ impl Scheme for DecimalScheme { matches!(canonical, Canonical::Decimal(_)) } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![DecimalByteParts.id()] } diff --git a/vortex-btrblocks/src/schemes/float/alp.rs b/vortex-btrblocks/src/schemes/float/alp.rs index bc39f434013..f9fc7066bf4 100644 --- a/vortex-btrblocks/src/schemes/float/alp.rs +++ b/vortex-btrblocks/src/schemes/float/alp.rs @@ -42,7 +42,7 @@ impl Scheme for ALPScheme { canonical.dtype().is_float() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { let mut encodings = vec![ALP.id()]; if use_experimental_patches() { encodings.push(Patched.id()); diff --git a/vortex-btrblocks/src/schemes/float/alprd.rs b/vortex-btrblocks/src/schemes/float/alprd.rs index 9efecaf83c6..09c0ee85b0d 100644 --- a/vortex-btrblocks/src/schemes/float/alprd.rs +++ b/vortex-btrblocks/src/schemes/float/alprd.rs @@ -40,7 +40,7 @@ impl Scheme for ALPRDScheme { canonical.dtype().is_float() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_alp::ALPRD.id()] } diff --git a/vortex-btrblocks/src/schemes/float/pco.rs b/vortex-btrblocks/src/schemes/float/pco.rs index bf517a21dfe..416668c2fd0 100644 --- a/vortex-btrblocks/src/schemes/float/pco.rs +++ b/vortex-btrblocks/src/schemes/float/pco.rs @@ -31,7 +31,7 @@ impl Scheme for PcoScheme { canonical.dtype().is_float() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_pco::Pco.id()] } diff --git a/vortex-btrblocks/src/schemes/float/rle.rs b/vortex-btrblocks/src/schemes/float/rle.rs index 6b19eb577ee..71158b9dc3b 100644 --- a/vortex-btrblocks/src/schemes/float/rle.rs +++ b/vortex-btrblocks/src/schemes/float/rle.rs @@ -38,7 +38,7 @@ impl Scheme for FloatRLEScheme { canonical.dtype().is_float() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![RLE.id()] } diff --git a/vortex-btrblocks/src/schemes/float/sparse.rs b/vortex-btrblocks/src/schemes/float/sparse.rs index d4e18683baf..3d9c25b18e4 100644 --- a/vortex-btrblocks/src/schemes/float/sparse.rs +++ b/vortex-btrblocks/src/schemes/float/sparse.rs @@ -41,7 +41,7 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_float() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Sparse.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/bitpacking.rs b/vortex-btrblocks/src/schemes/integer/bitpacking.rs index a76826b5ce6..5ac7d0e4078 100644 --- a/vortex-btrblocks/src/schemes/integer/bitpacking.rs +++ b/vortex-btrblocks/src/schemes/integer/bitpacking.rs @@ -40,7 +40,7 @@ impl Scheme for BitPackingScheme { canonical.dtype().is_int() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { let mut encodings = vec![BitPacked.id()]; if use_experimental_patches() { encodings.push(Patched.id()); diff --git a/vortex-btrblocks/src/schemes/integer/delta.rs b/vortex-btrblocks/src/schemes/integer/delta.rs index 8fc0a06a3ab..46b2f1e302e 100644 --- a/vortex-btrblocks/src/schemes/integer/delta.rs +++ b/vortex-btrblocks/src/schemes/integer/delta.rs @@ -97,7 +97,7 @@ impl Scheme for DeltaScheme { canonical.dtype().is_int() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Delta.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/for_.rs b/vortex-btrblocks/src/schemes/integer/for_.rs index 38299c2a425..476a0dec282 100644 --- a/vortex-btrblocks/src/schemes/integer/for_.rs +++ b/vortex-btrblocks/src/schemes/integer/for_.rs @@ -44,7 +44,7 @@ impl Scheme for FoRScheme { canonical.dtype().is_int() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![FoR.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/pco.rs b/vortex-btrblocks/src/schemes/integer/pco.rs index 13cf16544da..675a112d44f 100644 --- a/vortex-btrblocks/src/schemes/integer/pco.rs +++ b/vortex-btrblocks/src/schemes/integer/pco.rs @@ -32,7 +32,7 @@ impl Scheme for PcoScheme { canonical.dtype().is_int() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_pco::Pco.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/rle.rs b/vortex-btrblocks/src/schemes/integer/rle.rs index 3b617b4f212..0909e1428d4 100644 --- a/vortex-btrblocks/src/schemes/integer/rle.rs +++ b/vortex-btrblocks/src/schemes/integer/rle.rs @@ -144,7 +144,7 @@ impl Scheme for IntRLEScheme { canonical.dtype().is_int() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![RLE.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/runend.rs b/vortex-btrblocks/src/schemes/integer/runend.rs index ce42712786e..6a97f7ec37d 100644 --- a/vortex-btrblocks/src/schemes/integer/runend.rs +++ b/vortex-btrblocks/src/schemes/integer/runend.rs @@ -48,7 +48,7 @@ impl Scheme for RunEndScheme { canonical.dtype().is_int() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![RunEnd.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/sequence.rs b/vortex-btrblocks/src/schemes/integer/sequence.rs index 35139c7e62d..edcefb99fc2 100644 --- a/vortex-btrblocks/src/schemes/integer/sequence.rs +++ b/vortex-btrblocks/src/schemes/integer/sequence.rs @@ -43,7 +43,7 @@ impl Scheme for SequenceScheme { canonical.dtype().is_int() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Sequence.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/sparse.rs b/vortex-btrblocks/src/schemes/integer/sparse.rs index 44ffbd59ef2..429ff5c1a31 100644 --- a/vortex-btrblocks/src/schemes/integer/sparse.rs +++ b/vortex-btrblocks/src/schemes/integer/sparse.rs @@ -46,7 +46,7 @@ impl Scheme for SparseScheme { canonical.dtype().is_int() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Sparse.id(), Constant.id()] } diff --git a/vortex-btrblocks/src/schemes/integer/zigzag.rs b/vortex-btrblocks/src/schemes/integer/zigzag.rs index af6ca3df7a3..0e4be01845a 100644 --- a/vortex-btrblocks/src/schemes/integer/zigzag.rs +++ b/vortex-btrblocks/src/schemes/integer/zigzag.rs @@ -46,7 +46,7 @@ impl Scheme for ZigZagScheme { canonical.dtype().is_int() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![ZigZag.id()] } diff --git a/vortex-btrblocks/src/schemes/string/fsst.rs b/vortex-btrblocks/src/schemes/string/fsst.rs index dd720f38de1..fd3fd28696a 100644 --- a/vortex-btrblocks/src/schemes/string/fsst.rs +++ b/vortex-btrblocks/src/schemes/string/fsst.rs @@ -50,7 +50,7 @@ impl Scheme for FSSTScheme { canonical.dtype().is_utf8() || canonical.dtype().is_binary() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![FSST.id(), VarBin.id()] } diff --git a/vortex-btrblocks/src/schemes/string/onpair.rs b/vortex-btrblocks/src/schemes/string/onpair.rs index e5358cfddbd..a1bc8643775 100644 --- a/vortex-btrblocks/src/schemes/string/onpair.rs +++ b/vortex-btrblocks/src/schemes/string/onpair.rs @@ -50,7 +50,7 @@ impl Scheme for OnPairScheme { canonical.dtype().is_utf8() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![OnPair.id()] } diff --git a/vortex-btrblocks/src/schemes/string/sparse.rs b/vortex-btrblocks/src/schemes/string/sparse.rs index f6392cddaab..8620c366f77 100644 --- a/vortex-btrblocks/src/schemes/string/sparse.rs +++ b/vortex-btrblocks/src/schemes/string/sparse.rs @@ -42,7 +42,7 @@ impl Scheme for NullDominatedSparseScheme { canonical.dtype().is_utf8() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Sparse.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd.rs b/vortex-btrblocks/src/schemes/string/zstd.rs index faa1c98c358..84e8860d626 100644 --- a/vortex-btrblocks/src/schemes/string/zstd.rs +++ b/vortex-btrblocks/src/schemes/string/zstd.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdScheme { canonical.dtype().is_utf8() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_zstd::Zstd.id()] } diff --git a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs index 3f524b65cce..cf691c70fcb 100644 --- a/vortex-btrblocks/src/schemes/string/zstd_buffers.rs +++ b/vortex-btrblocks/src/schemes/string/zstd_buffers.rs @@ -31,7 +31,7 @@ impl Scheme for ZstdBuffersScheme { canonical.dtype().is_utf8() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![vortex_zstd::ZstdBuffers.id()] } diff --git a/vortex-btrblocks/src/schemes/temporal.rs b/vortex-btrblocks/src/schemes/temporal.rs index 9afa52b206e..79748b69450 100644 --- a/vortex-btrblocks/src/schemes/temporal.rs +++ b/vortex-btrblocks/src/schemes/temporal.rs @@ -55,7 +55,7 @@ impl Scheme for TemporalScheme { ) } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![DateTimeParts.id()] } diff --git a/vortex-compressor/src/builtins/dict/binary.rs b/vortex-compressor/src/builtins/dict/binary.rs index abf62f914c5..c407d0251c6 100644 --- a/vortex-compressor/src/builtins/dict/binary.rs +++ b/vortex-compressor/src/builtins/dict/binary.rs @@ -48,7 +48,7 @@ impl Scheme for BinaryDictScheme { canonical.dtype().is_binary() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/builtins/dict/float.rs b/vortex-compressor/src/builtins/dict/float.rs index 4d9ed43319a..f962c3ff967 100644 --- a/vortex-compressor/src/builtins/dict/float.rs +++ b/vortex-compressor/src/builtins/dict/float.rs @@ -54,7 +54,7 @@ impl Scheme for FloatDictScheme { canonical.dtype().is_float() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/builtins/dict/integer.rs b/vortex-compressor/src/builtins/dict/integer.rs index 1c4e5d5ba0e..27a17ef94ad 100644 --- a/vortex-compressor/src/builtins/dict/integer.rs +++ b/vortex-compressor/src/builtins/dict/integer.rs @@ -49,7 +49,7 @@ impl Scheme for IntDictScheme { canonical.dtype().is_int() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/builtins/dict/string.rs b/vortex-compressor/src/builtins/dict/string.rs index 05abec9f946..f5cbcd54d89 100644 --- a/vortex-compressor/src/builtins/dict/string.rs +++ b/vortex-compressor/src/builtins/dict/string.rs @@ -48,7 +48,7 @@ impl Scheme for StringDictScheme { canonical.dtype().is_utf8() } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { vec![Dict.id()] } diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index b91ac594e00..3a2a6281047 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -72,7 +72,7 @@ impl Scheme for DirectRatioScheme { matches_integer_primitive(canonical) } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { Vec::new() } @@ -108,7 +108,7 @@ impl Scheme for ImmediateAlwaysUseScheme { matches_integer_primitive(canonical) } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { Vec::new() } @@ -144,7 +144,7 @@ impl Scheme for CallbackAlwaysUseScheme { matches_integer_primitive(canonical) } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { Vec::new() } @@ -182,7 +182,7 @@ impl Scheme for CallbackSkipScheme { matches_integer_primitive(canonical) } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { Vec::new() } @@ -220,7 +220,7 @@ impl Scheme for CallbackRatioScheme { matches_integer_primitive(canonical) } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { Vec::new() } @@ -258,7 +258,7 @@ impl Scheme for HugeRatioScheme { matches_integer_primitive(canonical) } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { Vec::new() } @@ -294,7 +294,7 @@ impl Scheme for ZeroBytesSamplingScheme { matches_integer_primitive(canonical) } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { Vec::new() } @@ -507,7 +507,7 @@ impl Scheme for ThresholdObservingScheme { matches_integer_primitive(canonical) } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { Vec::new() } @@ -548,7 +548,7 @@ impl Scheme for CallbackMatchingRatioScheme { matches_integer_primitive(canonical) } - fn produced_serialized_ids(&self) -> Vec { + fn produced_encodings(&self) -> Vec { Vec::new() } diff --git a/vortex-compressor/src/scheme/mod.rs b/vortex-compressor/src/scheme/mod.rs index 1b5bd801252..0ba1c90202a 100644 --- a/vortex-compressor/src/scheme/mod.rs +++ b/vortex-compressor/src/scheme/mod.rs @@ -133,7 +133,7 @@ pub trait Scheme: Debug + Send + Sync { /// /// For most encodings this is the in-memory encoding ID. An encoding with several wire /// formats declares the wire IDs the scheme writes, which may differ from its in-memory ID. - fn produced_serialized_ids(&self) -> Vec; + fn produced_encodings(&self) -> Vec; /// Returns the stats generation options this scheme requires. The compressor merges all /// eligible schemes' options before generating stats so that a single stats pass satisfies diff --git a/vortex-file/src/writer.rs b/vortex-file/src/writer.rs index 7b411074635..6c0d01dfa05 100644 --- a/vortex-file/src/writer.rs +++ b/vortex-file/src/writer.rs @@ -239,7 +239,7 @@ impl VortexWriteOptions { let enforce_editions = !self.disable_editions; // The array context is built here, rather than when the options were constructed, so that // encodings registered on the session in between are still eligible for the file. - let (array_ctx, permitted_ids) = new_array_context(&self.session, enforce_editions); + let array_ctx = new_array_context(&self.session, enforce_editions); let ctx = LayoutWriterContext::new(array_ctx) .with_buffered_bytes_tracker(self.buffered_bytes.clone()); let ctx = if enforce_editions { @@ -247,11 +247,14 @@ impl VortexWriteOptions { } else { ctx }; + let allowed_serialized_ids: HashSet = + ctx.array_ctx().to_ids().into_iter().collect(); let strategy = match self.strategy { Some(strategy) => strategy, None => WriteStrategyBuilder::default() .with_btrblocks_builder( - BtrBlocksCompressorBuilder::default().retain_allowed_encodings(&permitted_ids), + BtrBlocksCompressorBuilder::default() + .retain_allowed_encodings(&allowed_serialized_ids), ) .build(), }; @@ -382,33 +385,30 @@ impl VortexWriteOptions { } } -fn new_array_context( - session: &VortexSession, - enforce_editions: bool, -) -> (ArrayContext, HashSet) { +fn new_array_context(session: &VortexSession, enforce_editions: bool) -> ArrayContext { // NOTE(os): Set up an array context with all eligible serialized IDs pre-populated. // This is preferred for now over having an empty context here, because only the // serialised array order is deterministic. The serialisation of arrays are done // parallel and with an empty context they can register their encodings to the context // in different order, changing the written bytes from run to run. - let arrays = session.arrays(); - let serialized_ids = if enforce_editions { + // + // The seeded IDs are also what the writer may emit: callers read them back with + // `ArrayContext::to_ids` to restrict compression to the same set. + let serialized_ids: Vec = if enforce_editions { session.enabled_component_ids(ComponentKind::Array) } else { - arrays + session + .arrays() .registry() .read(|registry| registry.keys().copied().collect()) }; - // Compression schemes declare the serialized IDs they write, so the same set restricts them. - let permitted_ids = serialized_ids.iter().copied().collect(); let array_ctx = ArrayContext::new(serialized_ids.iter().copied().sorted().collect()); - let array_ctx = if enforce_editions { + if enforce_editions { // Only permit serialized IDs in the enabled editions. array_ctx.with_allowed_ids(serialized_ids.into_iter().collect()) } else { array_ctx - }; - (array_ctx, permitted_ids) + } } /// The ids of `kind` the enabled editions permit. @@ -782,10 +782,9 @@ mod tests { session.register_edition(&DECLARATION)?; session.enable_edition(EDITION)?; - let (ctx, permitted_ids) = new_array_context(&session, true); + let ctx = new_array_context(&session, true); assert_eq!(ctx.to_ids(), [Primitive.id()]); assert!(ctx.intern(&Bool.id()).is_none()); - assert_eq!(permitted_ids, HashSet::from([Primitive.id()])); Ok(()) } @@ -797,12 +796,8 @@ mod tests { .registry() .read(|registry| registry.keys().copied().sorted().collect::>()); - let (ctx, permitted_ids) = new_array_context(&session, false); + let ctx = new_array_context(&session, false); assert_eq!(ctx.to_ids(), registered_ids); - assert_eq!( - permitted_ids, - registered_ids.iter().copied().collect::>() - ); assert!(ctx.intern(&Bool.id()).is_some()); }