From 42eb78b92170f42fb63f60c9b55de16d131b17da Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 29 Oct 2026 16:48:50 +0000 Subject: [PATCH 1/8] Preserve dictionary encoding through the write pipeline Add a way to hand the compressor and writer a set of encodings to keep. The compressor compresses only a preserved array's child slots, and repartitioning merges dictionary blocks, including dictionaries nested in structs and lists, into one compacted dictionary instead of canonicalizing them. `DictStrategy` passes input that is already dictionary-encoded to its fallback rather than re-encoding it. The benchmark conversion keeps Arrow dictionary columns encoded when the Parquet schema has them (PolarSignals), so scans export the stored dictionaries to Arrow instead of hashing every value into a new one. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AUHM7GB9hiFe56JK863CKS Signed-off-by: Joe Isaacs --- vortex-bench/src/conversions.rs | 80 +++- vortex-btrblocks/src/builder.rs | 16 +- vortex-compressor/src/compressor/cascade.rs | 15 +- vortex-compressor/src/compressor/mod.rs | 16 + vortex-compressor/src/compressor/tests.rs | 35 ++ vortex-file/src/strategy.rs | 52 ++- vortex-layout/src/layouts/dict/writer.rs | 13 + vortex-layout/src/layouts/repartition.rs | 417 +++++++++++++++++++- 8 files changed, 594 insertions(+), 50 deletions(-) diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 2fa88de37da..228604f3c6d 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -7,6 +7,8 @@ use std::path::PathBuf; use std::sync::Arc; use arrow_array::RecordBatch; +use arrow_schema::DataType; +use arrow_schema::Fields; use arrow_select::concat::concat_batches; use futures::StreamExt; use futures::TryStreamExt; @@ -29,8 +31,10 @@ use vortex::VortexSessionDefault; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; +use vortex::array::VTable; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ChunkedArray; +use vortex::array::arrays::Dict; use vortex::array::arrays::ExtensionArray; use vortex::array::arrays::Struct; use vortex::array::arrays::StructArray; @@ -123,7 +127,7 @@ pub async fn parquet_to_vortex_chunks_with_batch_size( let builder = ParquetRecordBatchStreamBuilder::new(file).await?; let Some(batch_size) = batch_size.filter(|size| *size > 0) else { - let chunks: Vec = parquet_to_vortex_stream(builder.build()?) + let chunks: Vec = parquet_to_vortex_stream(builder.build()?, false) .map(|r| r.map_err(anyhow::Error::from)) .try_collect() .await?; @@ -146,16 +150,22 @@ pub async fn parquet_to_vortex_chunks_with_batch_size( let mut chunks = Vec::with_capacity(combined.num_rows().div_ceil(batch_size)); for start in (0..combined.num_rows()).step_by(batch_size) { let len = batch_size.min(combined.num_rows() - start); - chunks.push(record_batch_to_vortex(combined.slice(start, len))?); + chunks.push(record_batch_to_vortex(combined.slice(start, len), false)?); } Ok(ChunkedArray::from_iter(chunks)) } -/// Convert one Arrow [`RecordBatch`] into a canonical Vortex array. -fn record_batch_to_vortex(batch: RecordBatch) -> VortexResult { +/// Convert one Arrow [`RecordBatch`] into a Vortex array. +/// +/// The array is canonical unless `preserve_dicts` is set, in which case Arrow dictionary columns +/// stay dictionary-encoded for a writer that preserves them. +fn record_batch_to_vortex(batch: RecordBatch, preserve_dicts: bool) -> VortexResult { let schema = batch.schema(); let chunk = SESSION.arrow().from_arrow_record_batch(batch, &schema)?; + if preserve_dicts { + return Ok(chunk); + } let mut ctx = VortexSession::default().create_execution_ctx(); let mut builder = builder_with_capacity_in(chunk.dtype(), chunk.len(), ctx.allocator()); @@ -168,14 +178,16 @@ fn record_batch_to_vortex(batch: RecordBatch) -> VortexResult { /// Create a streaming Vortex array from a Parquet reader. /// /// Streams record batches and converts them to Vortex arrays on-the-fly, avoiding loading the -/// entire file into memory. +/// entire file into memory. With `preserve_dicts`, Arrow dictionary columns stay +/// dictionary-encoded instead of being canonicalized. pub fn parquet_to_vortex_stream( reader: ParquetRecordBatchStream, + preserve_dicts: bool, ) -> impl futures::Stream> { reader.map(move |result| { result .map_err(|e| vortex_err!(External: e)) - .and_then(record_batch_to_vortex) + .and_then(|batch| record_batch_to_vortex(batch, preserve_dicts)) }) } @@ -210,6 +222,7 @@ fn parquet_row_groups_to_vortex_stream( parquet_path: PathBuf, metadata: ArrowReaderMetadata, concurrency: usize, + preserve_dicts: bool, ) -> impl futures::Stream> { let num_row_groups = metadata.metadata().num_row_groups(); futures::stream::iter(0..num_row_groups) @@ -224,7 +237,7 @@ fn parquet_row_groups_to_vortex_stream( .with_row_groups(vec![row_group]) .build() .map_err(|e| vortex_err!(External: e))?; - parquet_to_vortex_stream(reader) + parquet_to_vortex_stream(reader, preserve_dicts) .try_collect::>() .await }) @@ -255,16 +268,20 @@ pub async fn convert_parquet_file_to_vortex( .from_arrow_schema(metadata.schema().as_ref())?, &spatial_columns, )?; + // Dictionary columns are kept dictionary-encoded on disk, so that scans exporting them to + // Arrow dictionaries hand back the stored dictionaries instead of re-encoding the values. + let preserve_dicts = has_dictionary_field(metadata.schema().fields()); let concurrency = row_group_decode_concurrency(metadata.metadata()); let stream = if concurrency > 1 { StreamExt::boxed(parquet_row_groups_to_vortex_stream( parquet_path.to_path_buf(), metadata, concurrency, + preserve_dicts, )) } else { let reader = ParquetRecordBatchStreamBuilder::new_with_metadata(file, metadata).build()?; - StreamExt::boxed(parquet_to_vortex_stream(reader)) + StreamExt::boxed(parquet_to_vortex_stream(reader, preserve_dicts)) } .map(move |chunk| chunk.and_then(|chunk| tag_spatial_array(chunk, &spatial_columns))); @@ -275,12 +292,17 @@ pub async fn convert_parquet_file_to_vortex( .open(output_path) .await?; - write_options_for(compaction, &dtype, is_spatialbench(parquet_path)) - .write( - &mut output_file, - ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype, stream)), - ) - .await?; + write_options_for( + compaction, + &dtype, + is_spatialbench(parquet_path), + preserve_dicts, + ) + .write( + &mut output_file, + ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype, stream)), + ) + .await?; Ok(()) } @@ -291,14 +313,39 @@ fn is_spatialbench(path: &Path) -> bool { .any(|component| component.as_os_str() == "spatialbench") } +/// Whether any field, at any depth, is an Arrow dictionary. +fn has_dictionary_field(fields: &Fields) -> bool { + fields + .iter() + .any(|field| data_type_has_dictionary(field.data_type())) +} + +fn data_type_has_dictionary(data_type: &DataType) -> bool { + match data_type { + DataType::Dictionary(..) => true, + DataType::Struct(fields) => has_dictionary_field(fields), + DataType::List(field) + | DataType::LargeList(field) + | DataType::ListView(field) + | DataType::LargeListView(field) + | DataType::FixedSizeList(field, _) + | DataType::Map(field, _) => data_type_has_dictionary(field.data_type()), + _ => false, + } +} + /// Vortex write options for converting `dtype`-shaped data. /// /// For SpatialBench (`skip_binary_dict`), the geometry blobs are large and /// unique, so the dictionary builder balloons memory (tens of GB) for zero gain. +/// +/// With `preserve_dicts`, dictionary-encoded input keeps its encoding and only the codes and +/// values are compressed. fn write_options_for( compaction: CompactionStrategy, dtype: &DType, skip_binary_dict: bool, + preserve_dicts: bool, ) -> VortexWriteOptions { let binary_fields: Vec<_> = match dtype { DType::Struct(fields, _) if skip_binary_dict => fields @@ -310,11 +357,14 @@ fn write_options_for( .collect(), _ => Vec::new(), }; - if binary_fields.is_empty() { + if binary_fields.is_empty() && !preserve_dicts { return compaction.apply_options(SESSION.write_options()); } let mut builder = WriteStrategyBuilder::from_session(&SESSION); + if preserve_dicts { + builder = builder.with_preserved_encodings([Dict.id()]); + } if matches!(compaction, CompactionStrategy::Compact) { builder = builder.with_btrblocks_builder( BtrBlocksCompressorBuilder::from_session(&SESSION).with_compact(), diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 9b6e9aaf7c4..15e6c3ee3cf 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -137,6 +137,7 @@ pub struct BtrBlocksCompressorBuilder { schemes: Vec<&'static dyn Scheme>, allowed: AllowedIds, mode: CompressionMode, + preserved_encodings: Vec, } impl BtrBlocksCompressorBuilder { @@ -148,6 +149,7 @@ impl BtrBlocksCompressorBuilder { schemes: session.compression().schemes().to_vec(), allowed: AllowedIds::from_session(session), mode: CompressionMode::Default, + preserved_encodings: Vec::new(), } } @@ -160,6 +162,7 @@ impl BtrBlocksCompressorBuilder { schemes: Vec::new(), allowed: AllowedIds::all(), mode: CompressionMode::All, + preserved_encodings: Vec::new(), } } @@ -202,6 +205,14 @@ impl BtrBlocksCompressorBuilder { self } + /// Keeps arrays of the given encodings, compressing only their child slots. + /// + /// See [`CascadingCompressor::with_preserved_encodings`]. + pub fn preserve_encodings(mut self, ids: impl IntoIterator) -> Self { + self.preserved_encodings.extend(ids); + self + } + /// Removes the specified compression schemes by their [`SchemeId`]. pub fn exclude_schemes(mut self, ids: impl IntoIterator) -> Self { let ids: HashSet<_> = ids.into_iter().collect(); @@ -228,7 +239,10 @@ impl BtrBlocksCompressorBuilder { /// Builds the configured [`BtrBlocksCompressor`] from the schemes that its preset does not /// exclude and whose produced serialized IDs are all allowed. pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor(CascadingCompressor::new(self.allowed_schemes())) + BtrBlocksCompressor( + CascadingCompressor::new(self.allowed_schemes()) + .with_preserved_encodings(self.preserved_encodings), + ) } fn allowed_schemes(&self) -> Vec<&'static dyn Scheme> { diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index f86ab6b136b..ec4b97b4fb1 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -63,7 +63,10 @@ impl Selection { impl CascadingCompressor { /// Compresses an array using cascading adaptive compression. /// - /// First canonicalizes and compacts the array, then applies optimal compression schemes. + /// First canonicalizes and compacts the array, then applies optimal compression schemes. An + /// array whose encoding is [preserved] keeps its encoding and has its child slots compressed. + /// + /// [preserved]: CascadingCompressor::with_preserved_encodings /// /// # Errors /// @@ -77,9 +80,13 @@ impl CascadingCompressor { let span = trace::compress_span(array.len(), array.dtype(), before_nbytes); let _enter = span.enter(); - let canonical = array.clone().execute::(exec_ctx)?.0; - let compact = canonical.compact(exec_ctx)?; - let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?; + let compressed = if self.preserved_encodings.contains(&array.encoding_id()) { + self.compress_physical_slots(array, exec_ctx)? + } else { + let canonical = array.clone().execute::(exec_ctx)?.0; + let compact = canonical.compact(exec_ctx)?; + self.compress_canonical(compact, CompressorContext::new(), exec_ctx)? + }; trace::record_compress_outcome(&span, before_nbytes, compressed.nbytes()); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index a661970950c..06e313512b4 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,6 +9,8 @@ mod sample; mod select; mod structural; +use vortex_array::ArrayId; + use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; @@ -46,6 +48,9 @@ pub struct CascadingCompressor { /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from /// list offsets). root_exclusions: Vec, + + /// Encodings kept as they are: only their child slots are compressed. + preserved_encodings: Vec, } impl CascadingCompressor { @@ -63,8 +68,19 @@ impl CascadingCompressor { Self { schemes, root_exclusions, + preserved_encodings: Vec::new(), } } + + /// Keeps arrays of the given encodings instead of canonicalizing them. + /// + /// The compressor rebuilds a preserved array from its compressed child slots. This keeps an + /// encoding a reader wants back as-is, such as a dictionary that is exported to an Arrow + /// dictionary, at the cost of the better compression decoding it might allow. + pub fn with_preserved_encodings(mut self, ids: impl IntoIterator) -> Self { + self.preserved_encodings.extend(ids); + self + } } // NB: Cascading compression logic is located in `vortex-compressor/src/compressor/cascade.rs`. diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index 29e32b433c7..d7695d3cded 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -5,16 +5,19 @@ use std::sync::Arc; use std::sync::LazyLock; use parking_lot::Mutex; +use rstest::rstest; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::Constant; +use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; @@ -23,8 +26,11 @@ use vortex_array::arrays::NullArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::Slice; use vortex_array::arrays::SliceArray; +use vortex_array::arrays::Struct; +use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinView; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::varbinview::BinaryView; use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; @@ -961,3 +967,32 @@ fn extension_with_lazy_storage_is_not_passed_through() -> VortexResult<()> { assert_arrays_eq!(&compressed, &array, &mut exec_ctx); Ok(()) } + +#[rstest] +#[case::canonicalized(false)] +#[case::preserved(true)] +fn preserved_dict_keeps_encoding_inside_struct(#[case] preserve: bool) -> VortexResult<()> { + let codes = PrimitiveArray::from_iter((0..1024u32).map(|i| i % 4)).into_array(); + let values = VarBinViewArray::from_iter_str(["alpha", "beta", "gamma", "delta"]).into_array(); + let dict = DictArray::try_new(codes, values)?.into_array(); + let array = StructArray::from_fields(&[("labels", dict)])?.into_array(); + let mut exec_ctx = SESSION.create_execution_ctx(); + + let compressor = CascadingCompressor::new(vec![&IntDictScheme]); + let compressor = if preserve { + compressor.with_preserved_encodings([VTable::id(&Dict)]) + } else { + compressor + }; + let compressed = compressor.compress(&array, &mut exec_ctx)?; + + let field = compressed.as_::().unmasked_field(0).clone(); + assert_eq!( + field.is::(), + preserve, + "{}", + compressed.display_tree() + ); + assert_arrays_eq!(&compressed, &array, &mut exec_ctx); + Ok(()) +} diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 8cf7bc5bc78..381995df9d1 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -6,6 +6,7 @@ use std::num::NonZeroUsize; use std::sync::Arc; +use vortex_array::ArrayId; use vortex_array::dtype::FieldPath; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_btrblocks::SchemeExt; @@ -67,6 +68,8 @@ pub struct WriteStrategyBuilder { /// /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy use_list_layout: bool, + /// Encodings kept through repartitioning and compression. + preserved_encodings: Vec, } impl WriteStrategyBuilder { @@ -85,6 +88,7 @@ impl WriteStrategyBuilder { use_list_layout: session .enabled_component_ids(ComponentKind::Layout) .contains(&List.id()), + preserved_encodings: Vec::new(), } } @@ -146,6 +150,18 @@ impl WriteStrategyBuilder { self } + /// Keep arrays of the given encodings as they are written, compressing only their children. + /// + /// Repartitioning passes preserved chunks on instead of canonicalizing them, merging blocks of + /// dictionary chunks into one dictionary, and the BtrBlocks compressor compresses their child + /// slots. Columns whose chunks arrive already dictionary-encoded skip the dictionary layout. + /// Use this when readers want an encoding back as written, for example an Arrow dictionary + /// column. A compressor set with [`Self::with_compressor`] must preserve them itself. + pub fn with_preserved_encodings(mut self, ids: impl IntoIterator) -> Self { + self.preserved_encodings.extend(ids); + self + } + /// Override the compressor used to probe whether a column is dict-eligible. pub fn with_probe_compressor(mut self, compressor: C) -> Self { self.probe_compressor = Some(Arc::new(compressor)); @@ -162,6 +178,7 @@ impl WriteStrategyBuilder { }; let compressor = self.compressor; + let preserved = self.preserved_encodings; // 7. for each chunk create a flat layout let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat)); @@ -177,6 +194,7 @@ impl WriteStrategyBuilder { builder .clone() .exclude_schemes([IntDictScheme.id()]) + .preserve_encodings(preserved.iter().copied()) .build(), ), CompressorConfig::Opaque(compressor) => Arc::clone(compressor), @@ -198,11 +216,16 @@ impl WriteStrategyBuilder { block_size_target: self.data_block_target_bytes, canonicalize: true, }, - ); + ) + .with_preserved_encodings(preserved.iter().copied()); // 2.1. | 3.1. compress stats tables and dict values. let stats_compressor: Arc = match compressor { - CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()), + CompressorConfig::BtrBlocks(builder) => Arc::new( + builder + .preserve_encodings(preserved.iter().copied()) + .build(), + ), CompressorConfig::Opaque(compressor) => compressor, }; let compress_then_flat = CompressingStrategy::new(flat, Arc::clone(&stats_compressor)); @@ -247,7 +270,8 @@ impl WriteStrategyBuilder { block_size_target: None, canonicalize: false, }, - ); + ) + .with_preserved_encodings(preserved.iter().copied()); // 0. start with splitting columns let validity_strategy = CollectStrategy::new(compress_then_flat.clone()); @@ -258,6 +282,7 @@ impl WriteStrategyBuilder { .with_field_writers(self.field_writers); if self.use_list_layout { + let preserved: Arc<[ArrayId]> = preserved.into(); // We need a closure here to enable recursive application of list layout. table_strategy = table_strategy.with_list_layout_factory( move |list_layout: ListLayoutStrategy| -> Arc { @@ -269,15 +294,18 @@ impl WriteStrategyBuilder { ..Default::default() }, ); - Arc::new(RepartitionStrategy::new( - zoned, - RepartitionWriterOptions { - block_size_minimum: 0, - block_len_multiple: row_block_size.get(), - block_size_target: None, - canonicalize: false, - }, - )) + Arc::new( + RepartitionStrategy::new( + zoned, + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: row_block_size.get(), + block_size_target: None, + canonicalize: false, + }, + ) + .with_preserved_encodings(preserved.iter().copied()), + ) }, ); } diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index 53d13e2a222..6c93751dec6 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -24,7 +24,9 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Chunked; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::arrays::Dict; use vortex_array::builders::dict::DictConstraints; use vortex_array::builders::dict::DictEncoder; @@ -169,6 +171,9 @@ impl LayoutStrategy for DictStrategy { let should_fallback = match sample { None => true, // empty stream + // Chunks that arrive dictionary-encoded were preserved upstream, so they keep their + // own dictionaries rather than being decoded into a shared one. + Some(sample) if is_dict_encoded(&sample) => true, Some(sample) => { let mut exec_ctx = session.create_execution_ctx(); let compressed = self @@ -527,6 +532,14 @@ impl Stream for DictionaryTransformer { /// Peeks chunks until they hold `probe_len` rows, returning the stream with those chunks back in /// front and its first `probe_len` rows, or `None` if the stream is empty. +/// Whether `sample` is a dictionary, or a chunked array of dictionaries. +fn is_dict_encoded(sample: &ArrayRef) -> bool { + sample.is::() + || sample + .as_opt::() + .is_some_and(|chunked| chunked.iter_chunks().all(|chunk| chunk.is::())) +} + async fn peek_probe_sample( mut stream: BoxStream<'static, SequencedChunk>, probe_len: usize, diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 7098b0d547c..06fd225bd70 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -8,18 +8,45 @@ use async_stream::try_stream; use async_trait::async_trait; use futures::StreamExt as _; use futures::pin_mut; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::Dict; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::List; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::Primitive; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinView; +use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::listview::list_from_list_view; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::builders::dict::dict_encode; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; use vortex_session::VortexSession; use crate::LayoutRef; @@ -84,6 +111,7 @@ impl RepartitionWriterOptions { pub struct RepartitionStrategy { child: Arc, options: RepartitionWriterOptions, + preserved_encodings: Arc<[ArrayId]>, } impl RepartitionStrategy { @@ -91,8 +119,19 @@ impl RepartitionStrategy { Self { child: Arc::new(child), options, + preserved_encodings: Arc::from([]), } } + + /// Keeps chunks of the given encodings instead of canonicalizing them. + /// + /// When [`Dict`] is preserved, a block made of dictionary chunks becomes a single dictionary + /// holding only the distinct values its codes reference. Any other block made of one preserved + /// chunk is passed on as it is, and every remaining block is canonicalized. + pub fn with_preserved_encodings(mut self, ids: impl IntoIterator) -> Self { + self.preserved_encodings = ids.into_iter().collect(); + self + } } #[async_trait] @@ -108,12 +147,17 @@ impl LayoutStrategy for RepartitionStrategy { // TODO(os): spawn stream below like: // canon_stream = stream.map(async {to_canonical}).map(spawn).buffered(parallelism) let dtype = stream.dtype().clone(); + let preserved = Arc::clone(&self.preserved_encodings); let stream = if self.options.canonicalize { let canonicalize_session = session.clone(); + let preserved = Arc::clone(&preserved); SequentialStreamAdapter::new( dtype.clone(), stream.map(move |chunk| { let (sequence_id, chunk) = chunk?; + if preserved.contains(&chunk.encoding_id()) { + return VortexResult::Ok((sequence_id, chunk)); + } let mut ctx = canonicalize_session.create_execution_ctx(); let canonical = chunk.execute::(&mut ctx)?.into_array(); VortexResult::Ok((sequence_id, canonical)) @@ -153,28 +197,19 @@ impl LayoutStrategy for RepartitionStrategy { if chunks.have_enough() { let output_chunks = chunks.collect_exact_blocks(&mut ctx)?; assert!(!output_chunks.is_empty()); - let chunked = - ChunkedArray::try_new(output_chunks, dtype_clone.clone())?; - if !chunked.is_empty() { - let canonical = chunked.into_array().execute::(&mut ctx)?.into_array(); - yield ( - sequence_pointer.advance(), - canonical, - ) + if let Some(block) = + finish_block(output_chunks, &dtype_clone, &preserved, &mut ctx)? + { + yield (sequence_pointer.advance(), block) } } } if canonical_stream.as_mut().peek().await.is_none() { - let to_flush = ChunkedArray::try_new( - chunks.data.drain(..).map(|(arr, _)| arr), - dtype_clone.clone(), - )?; - if !to_flush.is_empty() { - let canonical = to_flush.into_array().execute::(&mut ctx)?.into_array(); - yield ( - sequence_pointer.advance(), - canonical, - ) + let to_flush = chunks.data.drain(..).map(|(arr, _)| arr).collect(); + if let Some(block) = + finish_block(to_flush, &dtype_clone, &preserved, &mut ctx)? + { + yield (sequence_pointer.advance(), block) } } } @@ -272,6 +307,291 @@ impl ChunksBuffer { } } +/// Assembles the chunks of one block, or returns `None` for an empty block. +/// +/// See [`RepartitionStrategy::with_preserved_encodings`] for which blocks keep their encoding. +fn finish_block( + chunks: Vec, + dtype: &DType, + preserved: &[ArrayId], + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let chunks: Vec = chunks.into_iter().filter(|c| !c.is_empty()).collect(); + if chunks.is_empty() { + return Ok(None); + } + if preserved.contains(&Dict.id()) { + return concat_preserving_dicts(chunks, dtype, ctx).map(Some); + } + if let [chunk] = chunks.as_slice() + && preserved.contains(&chunk.encoding_id()) + { + return Ok(Some(chunk.clone())); + } + let chunked = ChunkedArray::try_new(chunks, dtype.clone())?; + Ok(Some( + chunked.into_array().execute::(ctx)?.into_array(), + )) +} + +/// Concatenates chunks, keeping the dictionaries found at any depth of structs and lists. +fn concat_preserving_dicts( + chunks: Vec, + dtype: &DType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let chunks: Vec = chunks.into_iter().filter(|c| !c.is_empty()).collect(); + if chunks.is_empty() { + return Ok(Canonical::empty(dtype).into_array()); + } + if let Some(dict) = concat_dicts(&chunks, dtype, ctx)? { + return Ok(dict); + } + + let len = chunks.iter().map(|c| c.len()).sum(); + match dtype { + DType::Struct(fields, nullability) => { + let structs = chunks + .into_iter() + .map(|chunk| chunk.execute::(ctx)) + .collect::>>()?; + let validity = concat_validity(&structs, *nullability, ctx)?; + let field_arrays = fields + .fields() + .enumerate() + .map(|(idx, field_dtype)| { + let field_chunks = structs + .iter() + .map(|chunk| chunk.unmasked_field(idx).clone()) + .collect(); + concat_preserving_dicts(field_chunks, &field_dtype, ctx) + }) + .collect::>>()?; + Ok( + StructArray::try_new(fields.names().clone(), field_arrays, len, validity)? + .into_array(), + ) + } + DType::List(element_dtype, nullability) => { + let lists = chunks + .into_iter() + .map(|chunk| match chunk.as_opt::() { + Some(list) => Ok(list.into_owned()), + None => list_from_list_view(chunk.execute::(ctx)?, ctx), + }) + .collect::>>()?; + let validity = concat_validity(&lists, *nullability, ctx)?; + + let mut offsets = BufferMut::::with_capacity(len + 1); + offsets.push(0); + let mut elements = Vec::with_capacity(lists.len()); + for list in &lists { + let list_offsets = list + .offsets() + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::(ctx)?; + let list_offsets = list_offsets.as_slice::(); + let (Some(&first), Some(&last)) = (list_offsets.first(), list_offsets.last()) + else { + continue; + }; + let base = offsets.last().copied().unwrap_or_default(); + offsets.extend(list_offsets[1..].iter().map(|o| o - first + base)); + elements.push( + list.elements() + .slice(usize::try_from(first)?..usize::try_from(last)?)?, + ); + } + let elements = concat_preserving_dicts(elements, element_dtype, ctx)?; + Ok( + ListArray::try_new(elements, offsets.freeze().into_array(), validity, ctx)? + .into_array(), + ) + } + _ => Ok(ChunkedArray::try_new(chunks, dtype.clone())? + .into_array() + .execute::(ctx)? + .into_array()), + } +} + +/// Concatenates the validity of `arrays`. +fn concat_validity>( + arrays: &[A], + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if nullability == Nullability::NonNullable { + return Ok(Validity::NonNullable); + } + let len = arrays.iter().map(|a| a.as_ref().len()).sum(); + let mut bits = BitBufferMut::with_capacity(len); + for array in arrays { + let array = array.as_ref(); + let mask = array.validity()?.execute_mask(array.len(), ctx)?; + bits.append_buffer(&mask.into_bit_buffer()); + } + Ok(Validity::from_mask(bits.freeze().into(), nullability)) +} + +/// Concatenates dictionary chunks into one compact dictionary. +/// +/// Each chunk's codes are rebased onto its values, then the values no code references are dropped +/// and duplicate values are merged. Only the values are hashed, never the rows, so chunks that +/// each carry a whole Parquet column chunk dictionary end up sharing one small dictionary. +/// +/// A constant chunk, as slicing a dictionary with constant codes produces, contributes a single +/// value. Returns `None` unless every chunk is a dictionary or a constant and at least one is a +/// dictionary. +fn concat_dicts( + chunks: &[ArrayRef], + dtype: &DType, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + if !chunks.iter().any(|c| c.is::()) + || !chunks.iter().all(|c| c.is::() || c.is::()) + { + return Ok(None); + } + + let len = chunks.iter().map(|c| c.len()).sum(); + let codes_nullability = Nullability::from(chunks.iter().any(|c| { + c.as_opt::() + .is_some_and(|dict| dict.codes().dtype().is_nullable()) + })); + let codes_dtype = DType::Primitive(PType::U32, codes_nullability); + let mut codes = BufferMut::::with_capacity(len); + let mut validity = BitBufferMut::with_capacity(len); + let mut values = Vec::with_capacity(chunks.len()); + let mut offset = 0u32; + + for chunk in chunks { + let (chunk_values, chunk_codes) = match chunk.as_opt::() { + Some(dict) => (dict.values().clone(), Some(dict.codes().clone())), + None => { + let constant = chunk.as_::(); + let value = ConstantArray::new(constant.scalar().clone(), 1).into_array(); + (value, None) + } + }; + let Some(next_offset) = u32::try_from(chunk_values.len()) + .ok() + .and_then(|values_len| offset.checked_add(values_len)) + else { + return Ok(None); + }; + + match chunk_codes { + Some(chunk_codes) => { + let chunk_codes = chunk_codes + .cast(codes_dtype.clone())? + .execute::(ctx)?; + let mask = chunk_codes + .validity()? + .execute_mask(chunk_codes.len(), ctx)?; + // Null slots may hold any code, so they are pointed at value 0 of the chunk. + match mask.bit_buffer() { + AllOr::All => codes.extend( + chunk_codes + .as_slice::() + .iter() + .map(|code| code + offset), + ), + AllOr::None => codes.push_n(offset, chunk_codes.len()), + AllOr::Some(bits) => codes.extend( + chunk_codes + .as_slice::() + .iter() + .zip(bits.iter()) + .map(|(code, valid)| if valid { code + offset } else { offset }), + ), + } + validity.append_buffer(&mask.into_bit_buffer()); + } + None => { + codes.push_n(offset, chunk.len()); + validity.append_n(true, chunk.len()); + } + } + + values.push(chunk_values.cast(dtype.clone())?); + offset = next_offset; + } + + let values = ChunkedArray::try_new(values, dtype.clone())? + .into_array() + .execute::(ctx)? + .into_array(); + let validity = validity.freeze(); + let (codes, values) = compact_dict_values(codes, &validity, values, ctx)?; + + let codes = PrimitiveArray::new( + codes.freeze(), + Validity::from_mask(validity.into(), codes_nullability), + ); + // SAFETY: every code was rebased onto the values its chunk contributed, then remapped onto + // the compacted values. + let dict = unsafe { DictArray::new_unchecked(codes.into_array(), values) }; + Ok(Some(dict.into_array())) +} + +/// Drops the values no valid code references and merges duplicate values, remapping the codes. +/// +/// Values that cannot be hashed are only filtered, and values holding nulls are left as they are. +fn compact_dict_values( + mut codes: BufferMut, + validity: &BitBuffer, + values: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult<(BufferMut, ArrayRef)> { + let mut referenced = vec![false; values.len()]; + for (code, valid) in codes.iter().zip(validity.iter()) { + if valid { + referenced[*code as usize] = true; + } + } + // Null codes point at value 0, so it is kept to stay in bounds. + if let Some(first) = referenced.first_mut() { + *first = true; + } + if values.all_invalid(ctx)? || !values.all_valid(ctx)? { + return Ok((codes, values)); + } + + let mut remap = vec![0u32; values.len()]; + let mut kept = 0u32; + for (slot, &is_referenced) in remap.iter_mut().zip(&referenced) { + if is_referenced { + *slot = kept; + kept += 1; + } + } + let values = values + .filter(Mask::from(BitBuffer::from(referenced.clone())))? + .execute::(ctx)? + .into_array(); + + let values = if values.is::() || values.is::() { + let unique = dict_encode(&values, ctx)?; + let unique_codes = unique + .codes() + .cast(DType::Primitive(PType::U32, Nullability::NonNullable))? + .execute::(ctx)?; + let unique_codes = unique_codes.as_slice::(); + for slot in &mut remap { + *slot = unique_codes[*slot as usize]; + } + unique.values().clone() + } else { + values + }; + + for code in codes.iter_mut() { + *code = remap[*code as usize]; + } + Ok((codes, values)) +} + /// The bytes a block's rows reference, or the bytes it holds for a dtype that cannot be /// measured, such as variant. fn exact_size(chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { @@ -294,8 +614,10 @@ mod tests { use vortex_array::arrays::ConstantArray; use vortex_array::arrays::FixedSizeListArray; use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::Struct; use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::varbinview::BinaryView; + use vortex_array::assert_arrays_eq; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::PType; @@ -594,4 +916,63 @@ mod tests { Ok(()) } + + fn str_dict(codes: &[u32], values: &[&str]) -> VortexResult { + let codes = PrimitiveArray::from_iter(codes.iter().copied()).into_array(); + let values = VarBinViewArray::from_iter_str(values.iter().copied()).into_array(); + Ok(DictArray::try_new(codes, values)?.into_array()) + } + + #[test] + fn concat_dicts_merges_and_compacts_values() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let dtype = DType::Utf8(NonNullable); + let chunks = vec![ + str_dict(&[0, 1, 1], &["x", "y", "unused"])?, + str_dict(&[1, 0], &["y", "x"])?, + ConstantArray::new("x", 2).into_array(), + ]; + let expected = ChunkedArray::try_new(chunks.clone(), dtype.clone())?; + + let block = finish_block(chunks, &dtype, &[Dict.id()], &mut ctx)? + .vortex_expect("block is not empty"); + + let dict = block + .as_opt::() + .vortex_expect("block stays a dictionary"); + assert_eq!(dict.values().len(), 2); + assert_arrays_eq!(&block, &expected, &mut ctx); + Ok(()) + } + + #[test] + fn concat_preserves_dicts_nested_in_lists() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut list = + |codes: &[u32], values: &[&str], offsets: Vec| -> VortexResult { + let elements = StructArray::from_fields(&[("name", str_dict(codes, values)?)])?; + Ok(ListArray::try_new( + elements.into_array(), + PrimitiveArray::from_iter(offsets).into_array(), + Validity::NonNullable, + &mut ctx, + )? + .into_array()) + }; + let first = list(&[0, 1, 2, 0], &["a", "b", "c"], vec![0, 1, 3, 4])?; + // A slice whose offsets do not start at zero. + let second = list(&[1, 1, 0], &["b", "d"], vec![0, 1, 3])?.slice(1..2)?; + let dtype = first.dtype().clone(); + let chunks = vec![first, second]; + let expected = ChunkedArray::try_new(chunks.clone(), dtype.clone())?; + + let block = finish_block(chunks, &dtype, &[Dict.id()], &mut ctx)? + .vortex_expect("block is not empty"); + + let elements = block.as_::().elements().clone(); + assert_eq!(elements.len(), 6); + assert!(elements.as_::().unmasked_field(0).is::()); + assert_arrays_eq!(&block, &expected, &mut ctx); + Ok(()) + } } From 0d8f6cf81de3948c72383a099a8f63fb6caa4f90 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 29 Oct 2026 16:50:38 +0000 Subject: [PATCH 2/8] Sort imports in dict writer Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AUHM7GB9hiFe56JK863CKS Signed-off-by: Joe Isaacs --- vortex-layout/src/layouts/dict/writer.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index 6c93751dec6..6836d1f3b5d 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -26,8 +26,8 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::Chunked; use vortex_array::arrays::ChunkedArray; -use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::arrays::Dict; +use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::builders::dict::DictConstraints; use vortex_array::builders::dict::DictEncoder; use vortex_array::builders::dict::dict_encoder_in; From f311cead97d7d31dc6f206d227ce7939449d6b92 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 29 Oct 2026 17:50:56 +0000 Subject: [PATCH 3/8] Keep dictionaries by field path instead of by encoding Replace the compressor-wide `preserved_encodings` with a per-column writer selected by field path. `WriteStrategyBuilder::with_dict_field(path)` writes that column without dictionary encoding: repartitioning merges its dictionary chunks into one compacted dictionary and `DictChildrenCompressor` compresses only the codes and values. The compressor and `DictStrategy` are unchanged from develop. `TableStrategy` now routes overrides into list elements through `Field::ElementType`, so a path like `locations.[].mapping_file` reaches its column. The repartition above a list keeps nested dictionaries so they survive until their field writer. The benchmark conversion collects the paths of the Arrow dictionary fields and passes them as dict fields. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AUHM7GB9hiFe56JK863CKS --- vortex-bench/src/conversions.rs | 58 +++++++----- vortex-btrblocks/src/builder.rs | 16 +--- vortex-compressor/src/compressor/cascade.rs | 15 +--- vortex-compressor/src/compressor/mod.rs | 16 ---- vortex-compressor/src/compressor/tests.rs | 35 -------- vortex-file/src/strategy.rs | 99 ++++++++++++++------- vortex-file/src/tests.rs | 83 +++++++++++++++++ vortex-layout/src/layouts/compressed.rs | 33 +++++++ vortex-layout/src/layouts/dict/writer.rs | 13 --- vortex-layout/src/layouts/repartition.rs | 52 ++++++----- vortex-layout/src/layouts/table.rs | 9 +- 11 files changed, 254 insertions(+), 175 deletions(-) diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 228604f3c6d..5c5df81a4d6 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -31,10 +31,8 @@ use vortex::VortexSessionDefault; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; -use vortex::array::VTable; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ChunkedArray; -use vortex::array::arrays::Dict; use vortex::array::arrays::ExtensionArray; use vortex::array::arrays::Struct; use vortex::array::arrays::StructArray; @@ -45,6 +43,7 @@ use vortex::array::stream::ArrayStreamAdapter; use vortex::array::stream::ArrayStreamExt; use vortex::compressor::BtrBlocksCompressorBuilder; use vortex::dtype::DType; +use vortex::dtype::Field; use vortex::dtype::FieldPath; use vortex::dtype::StructFields; use vortex::dtype::extension::ExtDType; @@ -270,7 +269,8 @@ pub async fn convert_parquet_file_to_vortex( )?; // Dictionary columns are kept dictionary-encoded on disk, so that scans exporting them to // Arrow dictionaries hand back the stored dictionaries instead of re-encoding the values. - let preserve_dicts = has_dictionary_field(metadata.schema().fields()); + let dict_fields = dictionary_field_paths(metadata.schema().fields()); + let preserve_dicts = !dict_fields.is_empty(); let concurrency = row_group_decode_concurrency(metadata.metadata()); let stream = if concurrency > 1 { StreamExt::boxed(parquet_row_groups_to_vortex_stream( @@ -296,7 +296,7 @@ pub async fn convert_parquet_file_to_vortex( compaction, &dtype, is_spatialbench(parquet_path), - preserve_dicts, + dict_fields, ) .write( &mut output_file, @@ -313,24 +313,38 @@ fn is_spatialbench(path: &Path) -> bool { .any(|component| component.as_os_str() == "spatialbench") } -/// Whether any field, at any depth, is an Arrow dictionary. -fn has_dictionary_field(fields: &Fields) -> bool { - fields - .iter() - .any(|field| data_type_has_dictionary(field.data_type())) +/// The paths of the Arrow dictionary fields at any depth of structs and lists. +fn dictionary_field_paths(fields: &Fields) -> Vec { + let mut paths = Vec::new(); + for field in fields { + collect_dictionary_paths( + field.data_type(), + FieldPath::from_name(field.name().as_str()), + &mut paths, + ); + } + paths } -fn data_type_has_dictionary(data_type: &DataType) -> bool { +fn collect_dictionary_paths(data_type: &DataType, path: FieldPath, paths: &mut Vec) { match data_type { - DataType::Dictionary(..) => true, - DataType::Struct(fields) => has_dictionary_field(fields), + DataType::Dictionary(..) => paths.push(path), + DataType::Struct(fields) => { + for field in fields { + collect_dictionary_paths( + field.data_type(), + path.clone().push(field.name().as_str()), + paths, + ); + } + } DataType::List(field) | DataType::LargeList(field) | DataType::ListView(field) - | DataType::LargeListView(field) - | DataType::FixedSizeList(field, _) - | DataType::Map(field, _) => data_type_has_dictionary(field.data_type()), - _ => false, + | DataType::LargeListView(field) => { + collect_dictionary_paths(field.data_type(), path.push(Field::ElementType), paths) + } + _ => {} } } @@ -339,13 +353,13 @@ fn data_type_has_dictionary(data_type: &DataType) -> bool { /// For SpatialBench (`skip_binary_dict`), the geometry blobs are large and /// unique, so the dictionary builder balloons memory (tens of GB) for zero gain. /// -/// With `preserve_dicts`, dictionary-encoded input keeps its encoding and only the codes and -/// values are compressed. +/// The `dict_fields` keep the dictionary encoding they arrive with, compressing only their codes +/// and values. fn write_options_for( compaction: CompactionStrategy, dtype: &DType, skip_binary_dict: bool, - preserve_dicts: bool, + dict_fields: Vec, ) -> VortexWriteOptions { let binary_fields: Vec<_> = match dtype { DType::Struct(fields, _) if skip_binary_dict => fields @@ -357,13 +371,13 @@ fn write_options_for( .collect(), _ => Vec::new(), }; - if binary_fields.is_empty() && !preserve_dicts { + if binary_fields.is_empty() && dict_fields.is_empty() { return compaction.apply_options(SESSION.write_options()); } let mut builder = WriteStrategyBuilder::from_session(&SESSION); - if preserve_dicts { - builder = builder.with_preserved_encodings([Dict.id()]); + for field in dict_fields { + builder = builder.with_dict_field(field); } if matches!(compaction, CompactionStrategy::Compact) { builder = builder.with_btrblocks_builder( diff --git a/vortex-btrblocks/src/builder.rs b/vortex-btrblocks/src/builder.rs index 15e6c3ee3cf..9b6e9aaf7c4 100644 --- a/vortex-btrblocks/src/builder.rs +++ b/vortex-btrblocks/src/builder.rs @@ -137,7 +137,6 @@ pub struct BtrBlocksCompressorBuilder { schemes: Vec<&'static dyn Scheme>, allowed: AllowedIds, mode: CompressionMode, - preserved_encodings: Vec, } impl BtrBlocksCompressorBuilder { @@ -149,7 +148,6 @@ impl BtrBlocksCompressorBuilder { schemes: session.compression().schemes().to_vec(), allowed: AllowedIds::from_session(session), mode: CompressionMode::Default, - preserved_encodings: Vec::new(), } } @@ -162,7 +160,6 @@ impl BtrBlocksCompressorBuilder { schemes: Vec::new(), allowed: AllowedIds::all(), mode: CompressionMode::All, - preserved_encodings: Vec::new(), } } @@ -205,14 +202,6 @@ impl BtrBlocksCompressorBuilder { self } - /// Keeps arrays of the given encodings, compressing only their child slots. - /// - /// See [`CascadingCompressor::with_preserved_encodings`]. - pub fn preserve_encodings(mut self, ids: impl IntoIterator) -> Self { - self.preserved_encodings.extend(ids); - self - } - /// Removes the specified compression schemes by their [`SchemeId`]. pub fn exclude_schemes(mut self, ids: impl IntoIterator) -> Self { let ids: HashSet<_> = ids.into_iter().collect(); @@ -239,10 +228,7 @@ impl BtrBlocksCompressorBuilder { /// Builds the configured [`BtrBlocksCompressor`] from the schemes that its preset does not /// exclude and whose produced serialized IDs are all allowed. pub fn build(self) -> BtrBlocksCompressor { - BtrBlocksCompressor( - CascadingCompressor::new(self.allowed_schemes()) - .with_preserved_encodings(self.preserved_encodings), - ) + BtrBlocksCompressor(CascadingCompressor::new(self.allowed_schemes())) } fn allowed_schemes(&self) -> Vec<&'static dyn Scheme> { diff --git a/vortex-compressor/src/compressor/cascade.rs b/vortex-compressor/src/compressor/cascade.rs index ec4b97b4fb1..f86ab6b136b 100644 --- a/vortex-compressor/src/compressor/cascade.rs +++ b/vortex-compressor/src/compressor/cascade.rs @@ -63,10 +63,7 @@ impl Selection { impl CascadingCompressor { /// Compresses an array using cascading adaptive compression. /// - /// First canonicalizes and compacts the array, then applies optimal compression schemes. An - /// array whose encoding is [preserved] keeps its encoding and has its child slots compressed. - /// - /// [preserved]: CascadingCompressor::with_preserved_encodings + /// First canonicalizes and compacts the array, then applies optimal compression schemes. /// /// # Errors /// @@ -80,13 +77,9 @@ impl CascadingCompressor { let span = trace::compress_span(array.len(), array.dtype(), before_nbytes); let _enter = span.enter(); - let compressed = if self.preserved_encodings.contains(&array.encoding_id()) { - self.compress_physical_slots(array, exec_ctx)? - } else { - let canonical = array.clone().execute::(exec_ctx)?.0; - let compact = canonical.compact(exec_ctx)?; - self.compress_canonical(compact, CompressorContext::new(), exec_ctx)? - }; + let canonical = array.clone().execute::(exec_ctx)?.0; + let compact = canonical.compact(exec_ctx)?; + let compressed = self.compress_canonical(compact, CompressorContext::new(), exec_ctx)?; trace::record_compress_outcome(&span, before_nbytes, compressed.nbytes()); diff --git a/vortex-compressor/src/compressor/mod.rs b/vortex-compressor/src/compressor/mod.rs index 06e313512b4..a661970950c 100644 --- a/vortex-compressor/src/compressor/mod.rs +++ b/vortex-compressor/src/compressor/mod.rs @@ -9,8 +9,6 @@ mod sample; mod select; mod structural; -use vortex_array::ArrayId; - use crate::builtins::IntDictScheme; use crate::scheme::ChildSelection; use crate::scheme::DescendantExclusion; @@ -48,9 +46,6 @@ pub struct CascadingCompressor { /// Descendant exclusion rules for the compressor's own cascading (e.g. excluding Dict from /// list offsets). root_exclusions: Vec, - - /// Encodings kept as they are: only their child slots are compressed. - preserved_encodings: Vec, } impl CascadingCompressor { @@ -68,19 +63,8 @@ impl CascadingCompressor { Self { schemes, root_exclusions, - preserved_encodings: Vec::new(), } } - - /// Keeps arrays of the given encodings instead of canonicalizing them. - /// - /// The compressor rebuilds a preserved array from its compressed child slots. This keeps an - /// encoding a reader wants back as-is, such as a dictionary that is exported to an Arrow - /// dictionary, at the cost of the better compression decoding it might allow. - pub fn with_preserved_encodings(mut self, ids: impl IntoIterator) -> Self { - self.preserved_encodings.extend(ids); - self - } } // NB: Cascading compression logic is located in `vortex-compressor/src/compressor/cascade.rs`. diff --git a/vortex-compressor/src/compressor/tests.rs b/vortex-compressor/src/compressor/tests.rs index d7695d3cded..29e32b433c7 100644 --- a/vortex-compressor/src/compressor/tests.rs +++ b/vortex-compressor/src/compressor/tests.rs @@ -5,19 +5,16 @@ use std::sync::Arc; use std::sync::LazyLock; use parking_lot::Mutex; -use rstest::rstest; use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; use vortex_array::arrays::BoolArray; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::Constant; -use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; use vortex_array::arrays::ExtensionArray; use vortex_array::arrays::FixedSizeListArray; @@ -26,11 +23,8 @@ use vortex_array::arrays::NullArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::Slice; use vortex_array::arrays::SliceArray; -use vortex_array::arrays::Struct; -use vortex_array::arrays::StructArray; use vortex_array::arrays::VarBinView; use vortex_array::arrays::VarBinViewArray; -use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::arrays::varbinview::BinaryView; use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; @@ -967,32 +961,3 @@ fn extension_with_lazy_storage_is_not_passed_through() -> VortexResult<()> { assert_arrays_eq!(&compressed, &array, &mut exec_ctx); Ok(()) } - -#[rstest] -#[case::canonicalized(false)] -#[case::preserved(true)] -fn preserved_dict_keeps_encoding_inside_struct(#[case] preserve: bool) -> VortexResult<()> { - let codes = PrimitiveArray::from_iter((0..1024u32).map(|i| i % 4)).into_array(); - let values = VarBinViewArray::from_iter_str(["alpha", "beta", "gamma", "delta"]).into_array(); - let dict = DictArray::try_new(codes, values)?.into_array(); - let array = StructArray::from_fields(&[("labels", dict)])?.into_array(); - let mut exec_ctx = SESSION.create_execution_ctx(); - - let compressor = CascadingCompressor::new(vec![&IntDictScheme]); - let compressor = if preserve { - compressor.with_preserved_encodings([VTable::id(&Dict)]) - } else { - compressor - }; - let compressed = compressor.compress(&array, &mut exec_ctx)?; - - let field = compressed.as_::().unmasked_field(0).clone(); - assert_eq!( - field.is::(), - preserve, - "{}", - compressed.display_tree() - ); - assert_arrays_eq!(&compressed, &array, &mut exec_ctx); - Ok(()) -} diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 381995df9d1..4f509fa6909 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -6,7 +6,6 @@ use std::num::NonZeroUsize; use std::sync::Arc; -use vortex_array::ArrayId; use vortex_array::dtype::FieldPath; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_btrblocks::SchemeExt; @@ -21,6 +20,7 @@ use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; use vortex_layout::layouts::collect::CollectStrategy; use vortex_layout::layouts::compressed::CompressingStrategy; use vortex_layout::layouts::compressed::CompressorPlugin; +use vortex_layout::layouts::compressed::DictChildrenCompressor; use vortex_layout::layouts::dict::writer::DictLayoutOptions; use vortex_layout::layouts::dict::writer::DictStrategy; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; @@ -62,14 +62,13 @@ pub struct WriteStrategyBuilder { row_block_size: usize, data_block_target_bytes: Option, field_writers: HashMap>, + dict_fields: Vec, flat_strategy: Option>, probe_compressor: Option>, /// Whether to write list fields using [`ListLayoutStrategy`]. /// /// [`ListLayoutStrategy`]: vortex_layout::layouts::list::writer::ListLayoutStrategy use_list_layout: bool, - /// Encodings kept through repartitioning and compression. - preserved_encodings: Vec, } impl WriteStrategyBuilder { @@ -83,12 +82,12 @@ impl WriteStrategyBuilder { row_block_size: 8192, data_block_target_bytes: Some(ONE_MEG), field_writers: HashMap::new(), + dict_fields: Vec::new(), flat_strategy: None, probe_compressor: None, use_list_layout: session .enabled_component_ids(ComponentKind::Layout) .contains(&List.id()), - preserved_encodings: Vec::new(), } } @@ -123,6 +122,17 @@ impl WriteStrategyBuilder { self } + /// Keep the dictionary encoding of the field at `field`, compressing only its codes and values. + /// + /// Use this for a column whose chunks arrive dictionary-encoded and whose readers want the + /// dictionary back, such as an Arrow dictionary column: the column is written without being + /// decoded, and a scan hands back the stored dictionary. Blocks of dictionary chunks are merged + /// into one dictionary of the values they reference. The path may step into list elements. + pub fn with_dict_field(mut self, field: impl Into) -> Self { + self.dict_fields.push(field.into()); + self + } + /// Override the flat layout strategy used for leaf chunks. /// /// By default, this uses [`FlatLayoutStrategy`]. This can be used to substitute a custom @@ -150,18 +160,6 @@ impl WriteStrategyBuilder { self } - /// Keep arrays of the given encodings as they are written, compressing only their children. - /// - /// Repartitioning passes preserved chunks on instead of canonicalizing them, merging blocks of - /// dictionary chunks into one dictionary, and the BtrBlocks compressor compresses their child - /// slots. Columns whose chunks arrive already dictionary-encoded skip the dictionary layout. - /// Use this when readers want an encoding back as written, for example an Arrow dictionary - /// column. A compressor set with [`Self::with_compressor`] must preserve them itself. - pub fn with_preserved_encodings(mut self, ids: impl IntoIterator) -> Self { - self.preserved_encodings.extend(ids); - self - } - /// Override the compressor used to probe whether a column is dict-eligible. pub fn with_probe_compressor(mut self, compressor: C) -> Self { self.probe_compressor = Some(Arc::new(compressor)); @@ -178,7 +176,6 @@ impl WriteStrategyBuilder { }; let compressor = self.compressor; - let preserved = self.preserved_encodings; // 7. for each chunk create a flat layout let chunked = ChunkedLayoutStrategy::new(Arc::clone(&flat)); @@ -194,12 +191,11 @@ impl WriteStrategyBuilder { builder .clone() .exclude_schemes([IntDictScheme.id()]) - .preserve_encodings(preserved.iter().copied()) .build(), ), CompressorConfig::Opaque(compressor) => Arc::clone(compressor), }; - let compressing = CompressingStrategy::new(buffered, data_compressor); + let compressing = CompressingStrategy::new(buffered, Arc::clone(&data_compressor)); // 4. prior to compression, coalesce up to a minimum size let coalescing = RepartitionStrategy::new( @@ -216,19 +212,15 @@ impl WriteStrategyBuilder { block_size_target: self.data_block_target_bytes, canonicalize: true, }, - ) - .with_preserved_encodings(preserved.iter().copied()); + ); // 2.1. | 3.1. compress stats tables and dict values. let stats_compressor: Arc = match compressor { - CompressorConfig::BtrBlocks(builder) => Arc::new( - builder - .preserve_encodings(preserved.iter().copied()) - .build(), - ), + CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()), CompressorConfig::Opaque(compressor) => compressor, }; - let compress_then_flat = CompressingStrategy::new(flat, Arc::clone(&stats_compressor)); + let compress_then_flat = + CompressingStrategy::new(Arc::clone(&flat), Arc::clone(&stats_compressor)); // 3. apply dict encoding or fallback let probe_compressor = if let Some(probe_compressor) = self.probe_compressor { @@ -270,8 +262,50 @@ impl WriteStrategyBuilder { block_size_target: None, canonicalize: false, }, - ) - .with_preserved_encodings(preserved.iter().copied()); + ); + + // Dictionary fields follow the same steps without dictionary encoding, keeping the + // dictionaries they arrive with and compressing only their codes and values. + let mut field_writers = self.field_writers; + if !self.dict_fields.is_empty() { + let dict_compressing = CompressingStrategy::new( + BufferedStrategy::new(ChunkedLayoutStrategy::new(Arc::clone(&flat)), 2 * ONE_MEG), + DictChildrenCompressor::new(data_compressor), + ); + let dict_coalescing = RepartitionStrategy::new( + dict_compressing, + RepartitionWriterOptions { + block_size_minimum: self.data_block_target_bytes.unwrap_or(0), + block_len_multiple: self.row_block_size, + block_size_target: self.data_block_target_bytes, + canonicalize: true, + }, + ) + .with_keep_dicts(); + let dict_stats = ZonedStrategy::new( + dict_coalescing, + compress_then_flat.clone(), + ZonedLayoutOptions { + block_size: row_block_size, + ..Default::default() + }, + ); + let dict_writer: Arc = Arc::new( + RepartitionStrategy::new( + dict_stats, + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: self.row_block_size, + block_size_target: None, + canonicalize: false, + }, + ) + .with_keep_dicts(), + ); + for field in self.dict_fields { + field_writers.insert(field, Arc::clone(&dict_writer)); + } + } // 0. start with splitting columns let validity_strategy = CollectStrategy::new(compress_then_flat.clone()); @@ -279,10 +313,9 @@ impl WriteStrategyBuilder { // Take any field overrides from the builder and apply them to the final strategy. let mut table_strategy = TableStrategy::new(Arc::new(validity_strategy), Arc::new(repartition)) - .with_field_writers(self.field_writers); + .with_field_writers(field_writers); if self.use_list_layout { - let preserved: Arc<[ArrayId]> = preserved.into(); // We need a closure here to enable recursive application of list layout. table_strategy = table_strategy.with_list_layout_factory( move |list_layout: ListLayoutStrategy| -> Arc { @@ -294,6 +327,8 @@ impl WriteStrategyBuilder { ..Default::default() }, ); + // Dictionaries nested in the list's elements stay encoded until the elements + // reach their own field writers. Arc::new( RepartitionStrategy::new( zoned, @@ -304,7 +339,7 @@ impl WriteStrategyBuilder { canonicalize: false, }, ) - .with_preserved_encodings(preserved.iter().copied()), + .with_keep_dicts(), ) }, ); diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 51386aca2d4..f7d64257afa 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -25,22 +25,26 @@ use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::NumericalAggregateOpts; use vortex_array::aggregate_fn::fns::sum_v2::SumV2; use vortex_array::arrays::BoolArray; +use vortex_array::arrays::Chunked; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::DecimalArray; use vortex_array::arrays::Dict; +use vortex_array::arrays::DictArray; use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::TemporalArray; use vortex_array::arrays::VarBinArray; use vortex_array::arrays::VarBinViewArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::assert_arrays_eq; use vortex_array::builders::MapBuilder; use vortex_array::dtype::DType; use vortex_array::dtype::DecimalDType; +use vortex_array::dtype::Field; use vortex_array::dtype::FieldPath; use vortex_array::dtype::MapDType; use vortex_array::dtype::Nullability; @@ -3388,3 +3392,82 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { assert_arrays_eq!(result, mk_struct(vec![big])?, &mut ctx); Ok(()) } + +/// Columns named with `with_dict_field`, including one inside a list's elements, keep the +/// dictionaries they were written with. The values are all distinct, so the default path would +/// not dictionary-encode them. +#[rstest] +#[case::default(false)] +#[case::dict_fields(true)] +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn dict_fields_keep_their_dictionaries(#[case] dict_fields: bool) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let distinct = |len: u32, prefix: &str| -> VortexResult { + let values: Vec = (0..len).map(|i| format!("{prefix}_{i:08}")).collect(); + Ok(DictArray::try_new( + PrimitiveArray::from_iter(0..len).into_array(), + VarBinViewArray::from_iter_str(values).into_array(), + )? + .into_array()) + }; + let chunk = |prefix: &str| -> VortexResult { + let elements = StructArray::from_fields(&[("tag", distinct(400, prefix)?)])?.into_array(); + let items = ListArray::try_new( + elements, + PrimitiveArray::from_iter((0..=200u32).map(|i| i * 2)).into_array(), + Validity::NonNullable, + &mut SESSION.create_execution_ctx(), + )? + .into_array(); + Ok( + StructArray::from_fields(&[("name", distinct(200, prefix)?), ("items", items)])? + .into_array(), + ) + }; + let input = ChunkedArray::from_iter([chunk("a")?, chunk("b")?]).into_array(); + + let mut strategy = crate::strategy::WriteStrategyBuilder::from_session(&SESSION); + if dict_fields { + strategy = strategy.with_dict_field(field_path!(name)).with_dict_field( + FieldPath::from_name("items") + .push(Field::ElementType) + .push("tag"), + ); + } + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .with_strategy(strategy.build()) + .write(&mut buf, input.to_array_stream()) + .await?; + + let result = SESSION + .open_options() + .open_buffer(buf)? + .scan()? + .into_array_stream()? + .read_all() + .await?; + assert_arrays_eq!(&result, &input, &mut ctx); + + let first = result + .as_opt::() + .map_or_else(|| result.clone(), |chunked| chunked.chunk(0).clone()) + .execute::(&mut ctx)?; + let tree = result.display_tree(); + assert_eq!( + first.unmasked_field_by_name("name")?.is::(), + dict_fields, + "{tree}" + ); + assert_eq!( + first + .unmasked_field_by_name("items")? + .depth_first_traversal() + .any(|array| array.is::()), + dict_fields, + "{tree}" + ); + Ok(()) +} diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index 07d83bf41c7..e290936ca58 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -7,9 +7,13 @@ use async_trait::async_trait; use futures::StreamExt as _; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; +use vortex_array::arrays::Dict; +use vortex_array::arrays::DictArray; +use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::expr::stats::Stat; use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexResult; @@ -54,6 +58,35 @@ impl CompressorPlugin for BtrBlocksCompressor { } } +/// Compresses a dictionary's codes and values separately, keeping the dictionary itself. +/// +/// Readers that export the column to an Arrow dictionary then get the stored dictionary back +/// instead of re-encoding the decoded values. Chunks that are not dictionaries, such as a +/// constant slice of one, go to the inner compressor unchanged. +pub struct DictChildrenCompressor { + inner: Arc, +} + +impl DictChildrenCompressor { + /// Creates a compressor that compresses dictionary children with `inner`. + pub fn new(inner: C) -> Self { + Self { + inner: Arc::new(inner), + } + } +} + +impl CompressorPlugin for DictChildrenCompressor { + fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let Some(dict) = chunk.as_opt::() else { + return self.inner.compress_chunk(chunk, ctx); + }; + let codes = self.inner.compress_chunk(dict.codes(), ctx)?; + let values = self.inner.compress_chunk(dict.values(), ctx)?; + Ok(DictArray::try_new(codes, values)?.into_array()) + } +} + /// A layout writer that compresses chunks. #[derive(Clone)] pub struct CompressingStrategy { diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index 6836d1f3b5d..53d13e2a222 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -24,10 +24,8 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::arrays::Chunked; use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::Dict; -use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::builders::dict::DictConstraints; use vortex_array::builders::dict::DictEncoder; use vortex_array::builders::dict::dict_encoder_in; @@ -171,9 +169,6 @@ impl LayoutStrategy for DictStrategy { let should_fallback = match sample { None => true, // empty stream - // Chunks that arrive dictionary-encoded were preserved upstream, so they keep their - // own dictionaries rather than being decoded into a shared one. - Some(sample) if is_dict_encoded(&sample) => true, Some(sample) => { let mut exec_ctx = session.create_execution_ctx(); let compressed = self @@ -532,14 +527,6 @@ impl Stream for DictionaryTransformer { /// Peeks chunks until they hold `probe_len` rows, returning the stream with those chunks back in /// front and its first `probe_len` rows, or `None` if the stream is empty. -/// Whether `sample` is a dictionary, or a chunked array of dictionaries. -fn is_dict_encoded(sample: &ArrayRef) -> bool { - sample.is::() - || sample - .as_opt::() - .is_some_and(|chunked| chunked.iter_chunks().all(|chunk| chunk.is::())) -} - async fn peek_probe_sample( mut stream: BoxStream<'static, SequencedChunk>, probe_len: usize, diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 06fd225bd70..8e7b309229a 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -8,12 +8,10 @@ use async_stream::try_stream; use async_trait::async_trait; use futures::StreamExt as _; use futures::pin_mut; -use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; -use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; @@ -111,7 +109,7 @@ impl RepartitionWriterOptions { pub struct RepartitionStrategy { child: Arc, options: RepartitionWriterOptions, - preserved_encodings: Arc<[ArrayId]>, + keep_dicts: bool, } impl RepartitionStrategy { @@ -119,17 +117,16 @@ impl RepartitionStrategy { Self { child: Arc::new(child), options, - preserved_encodings: Arc::from([]), + keep_dicts: false, } } - /// Keeps chunks of the given encodings instead of canonicalizing them. + /// Keeps dictionaries, at any depth of structs and lists, instead of canonicalizing them. /// - /// When [`Dict`] is preserved, a block made of dictionary chunks becomes a single dictionary - /// holding only the distinct values its codes reference. Any other block made of one preserved - /// chunk is passed on as it is, and every remaining block is canonicalized. - pub fn with_preserved_encodings(mut self, ids: impl IntoIterator) -> Self { - self.preserved_encodings = ids.into_iter().collect(); + /// A block made of dictionary chunks becomes a single dictionary holding only the distinct + /// values its codes reference. Blocks without dictionaries are canonicalized as usual. + pub fn with_keep_dicts(mut self) -> Self { + self.keep_dicts = true; self } } @@ -147,15 +144,14 @@ impl LayoutStrategy for RepartitionStrategy { // TODO(os): spawn stream below like: // canon_stream = stream.map(async {to_canonical}).map(spawn).buffered(parallelism) let dtype = stream.dtype().clone(); - let preserved = Arc::clone(&self.preserved_encodings); + let keep_dicts = self.keep_dicts; let stream = if self.options.canonicalize { let canonicalize_session = session.clone(); - let preserved = Arc::clone(&preserved); SequentialStreamAdapter::new( dtype.clone(), stream.map(move |chunk| { let (sequence_id, chunk) = chunk?; - if preserved.contains(&chunk.encoding_id()) { + if keep_dicts && contains_dict(&chunk) { return VortexResult::Ok((sequence_id, chunk)); } let mut ctx = canonicalize_session.create_execution_ctx(); @@ -198,7 +194,7 @@ impl LayoutStrategy for RepartitionStrategy { let output_chunks = chunks.collect_exact_blocks(&mut ctx)?; assert!(!output_chunks.is_empty()); if let Some(block) = - finish_block(output_chunks, &dtype_clone, &preserved, &mut ctx)? + finish_block(output_chunks, &dtype_clone, keep_dicts, &mut ctx)? { yield (sequence_pointer.advance(), block) } @@ -207,7 +203,7 @@ impl LayoutStrategy for RepartitionStrategy { if canonical_stream.as_mut().peek().await.is_none() { let to_flush = chunks.data.drain(..).map(|(arr, _)| arr).collect(); if let Some(block) = - finish_block(to_flush, &dtype_clone, &preserved, &mut ctx)? + finish_block(to_flush, &dtype_clone, keep_dicts, &mut ctx)? { yield (sequence_pointer.advance(), block) } @@ -309,31 +305,33 @@ impl ChunksBuffer { /// Assembles the chunks of one block, or returns `None` for an empty block. /// -/// See [`RepartitionStrategy::with_preserved_encodings`] for which blocks keep their encoding. +/// See [`RepartitionStrategy::with_keep_dicts`] for which blocks keep their dictionaries. fn finish_block( chunks: Vec, dtype: &DType, - preserved: &[ArrayId], + keep_dicts: bool, ctx: &mut ExecutionCtx, ) -> VortexResult> { let chunks: Vec = chunks.into_iter().filter(|c| !c.is_empty()).collect(); if chunks.is_empty() { return Ok(None); } - if preserved.contains(&Dict.id()) { + if keep_dicts && chunks.iter().any(contains_dict) { return concat_preserving_dicts(chunks, dtype, ctx).map(Some); } - if let [chunk] = chunks.as_slice() - && preserved.contains(&chunk.encoding_id()) - { - return Ok(Some(chunk.clone())); - } let chunked = ChunkedArray::try_new(chunks, dtype.clone())?; Ok(Some( chunked.into_array().execute::(ctx)?.into_array(), )) } +/// Whether a dictionary appears anywhere in `array`. +fn contains_dict(array: &ArrayRef) -> bool { + array + .depth_first_traversal() + .any(|child| child.is::()) +} + /// Concatenates chunks, keeping the dictionaries found at any depth of structs and lists. fn concat_preserving_dicts( chunks: Vec, @@ -934,8 +932,8 @@ mod tests { ]; let expected = ChunkedArray::try_new(chunks.clone(), dtype.clone())?; - let block = finish_block(chunks, &dtype, &[Dict.id()], &mut ctx)? - .vortex_expect("block is not empty"); + let block = + finish_block(chunks, &dtype, true, &mut ctx)?.vortex_expect("block is not empty"); let dict = block .as_opt::() @@ -966,8 +964,8 @@ mod tests { let chunks = vec![first, second]; let expected = ChunkedArray::try_new(chunks.clone(), dtype.clone())?; - let block = finish_block(chunks, &dtype, &[Dict.id()], &mut ctx)? - .vortex_expect("block is not empty"); + let block = + finish_block(chunks, &dtype, true, &mut ctx)?.vortex_expect("block is not empty"); let elements = block.as_::().elements().clone(); assert_eq!(elements.len(), 6); diff --git a/vortex-layout/src/layouts/table.rs b/vortex-layout/src/layouts/table.rs index 00853638e44..3dc78ae1eb3 100644 --- a/vortex-layout/src/layouts/table.rs +++ b/vortex-layout/src/layouts/table.rs @@ -212,13 +212,14 @@ impl TableStrategy { /// Build the [`ListLayoutStrategy`] used to write a list field stream at this level. /// - /// The `elements` sub-column is routed back through a clean descended dispatcher so nested - /// structs/lists recurse; `offsets` go straight to the leaf (they are always a primitive - /// column); and `validity` uses the shared validity strategy. + /// The `elements` sub-column is routed back through a dispatcher descended into + /// [`Field::ElementType`], so nested structs/lists recurse and overrides below the list still + /// apply; `offsets` go straight to the leaf (they are always a primitive column); and + /// `validity` uses the shared validity strategy. fn list_strategy(&self) -> Option> { let factory = self.list_layout_factory.as_ref()?; let list_layout = ListLayoutStrategy::default() - .with_elements(Arc::new(self.descend_clean())) + .with_elements(Arc::new(self.descend(&Field::ElementType))) .with_offsets(Arc::clone(&self.leaf)) .with_validity(Arc::clone(&self.validity)) .with_fallback(Arc::clone(&self.leaf)); From b0f0ae3ae34a83f7676162dfc468db04222b7f12 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 29 Oct 2026 18:20:20 +0000 Subject: [PATCH 4/8] Allow encodings per field path with `with_allow_encodings` Replace `with_dict_field(path)` with `with_allow_encodings(path, encodings)`: the field keeps arrays of the given encodings it arrives with, and only their child slots are compressed. `AllowedEncodingsCompressor` walks struct fields and list elements to the allowed paths, and repartitioning keeps allowed encodings, merging dictionary blocks. Without the list layout, which the default editions leave out, list elements are not written as columns of their own, so a path into a list is now kept by the writer of the list column that holds it. This is what the PolarSignals `locations` columns need in the benchmark session. The benchmark conversion allows `Dict` at each Arrow dictionary field path. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AUHM7GB9hiFe56JK863CKS --- vortex-bench/src/conversions.rs | 4 +- vortex-file/src/strategy.rs | 113 +++++++++----- vortex-file/src/tests.rs | 20 ++- vortex-layout/src/layouts/compressed.rs | 178 ++++++++++++++++++++--- vortex-layout/src/layouts/repartition.rs | 73 ++++++---- 5 files changed, 295 insertions(+), 93 deletions(-) diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 5c5df81a4d6..e1463563ea7 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -31,8 +31,10 @@ use vortex::VortexSessionDefault; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::IntoArray; +use vortex::array::VTable; use vortex::array::VortexSessionExecute; use vortex::array::arrays::ChunkedArray; +use vortex::array::arrays::Dict; use vortex::array::arrays::ExtensionArray; use vortex::array::arrays::Struct; use vortex::array::arrays::StructArray; @@ -377,7 +379,7 @@ fn write_options_for( let mut builder = WriteStrategyBuilder::from_session(&SESSION); for field in dict_fields { - builder = builder.with_dict_field(field); + builder = builder.with_allow_encodings(field, [Dict.id()]); } if matches!(compaction, CompactionStrategy::Compact) { builder = builder.with_btrblocks_builder( diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 4f509fa6909..d5dcf5cb0c6 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -6,6 +6,8 @@ use std::num::NonZeroUsize; use std::sync::Arc; +use vortex_array::ArrayId; +use vortex_array::dtype::Field; use vortex_array::dtype::FieldPath; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_btrblocks::SchemeExt; @@ -18,9 +20,9 @@ use vortex_layout::VTable; use vortex_layout::layouts::buffered::BufferedStrategy; use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; use vortex_layout::layouts::collect::CollectStrategy; +use vortex_layout::layouts::compressed::AllowedEncodingsCompressor; use vortex_layout::layouts::compressed::CompressingStrategy; use vortex_layout::layouts::compressed::CompressorPlugin; -use vortex_layout::layouts::compressed::DictChildrenCompressor; use vortex_layout::layouts::dict::writer::DictLayoutOptions; use vortex_layout::layouts::dict::writer::DictStrategy; use vortex_layout::layouts::flat::writer::FlatLayoutStrategy; @@ -62,7 +64,7 @@ pub struct WriteStrategyBuilder { row_block_size: usize, data_block_target_bytes: Option, field_writers: HashMap>, - dict_fields: Vec, + allowed_encodings: Vec<(FieldPath, Vec)>, flat_strategy: Option>, probe_compressor: Option>, /// Whether to write list fields using [`ListLayoutStrategy`]. @@ -82,7 +84,7 @@ impl WriteStrategyBuilder { row_block_size: 8192, data_block_target_bytes: Some(ONE_MEG), field_writers: HashMap::new(), - dict_fields: Vec::new(), + allowed_encodings: Vec::new(), flat_strategy: None, probe_compressor: None, use_list_layout: session @@ -122,14 +124,25 @@ impl WriteStrategyBuilder { self } - /// Keep the dictionary encoding of the field at `field`, compressing only its codes and values. + /// Allow the field at `field` to keep the given encodings it arrives with, compressing only + /// their children. /// - /// Use this for a column whose chunks arrive dictionary-encoded and whose readers want the - /// dictionary back, such as an Arrow dictionary column: the column is written without being - /// decoded, and a scan hands back the stored dictionary. Blocks of dictionary chunks are merged - /// into one dictionary of the values they reference. The path may step into list elements. - pub fn with_dict_field(mut self, field: impl Into) -> Self { - self.dict_fields.push(field.into()); + /// The field is written without being canonicalized where it holds an allowed encoding: such + /// an array keeps its encoding and only its child slots are compressed. Use this for a column + /// whose readers want the encoding back, such as an Arrow dictionary column written with + /// [`Dict`]: a scan hands back the stored dictionary instead of re-encoding the values. Blocks + /// of dictionary chunks are merged into one dictionary of the values they reference. The path + /// may step into list elements; without the list layout, the column holding the list keeps the + /// encodings inside it. + /// + /// [`Dict`]: vortex_array::arrays::Dict + pub fn with_allow_encodings( + mut self, + field: impl Into, + encodings: impl IntoIterator, + ) -> Self { + self.allowed_encodings + .push((field.into(), encodings.into_iter().collect())); self } @@ -264,16 +277,46 @@ impl WriteStrategyBuilder { }, ); - // Dictionary fields follow the same steps without dictionary encoding, keeping the - // dictionaries they arrive with and compressing only their codes and values. + // Fields with allowed encodings follow the same steps without dictionary encoding, keeping + // the allowed encodings they arrive with and compressing only their children. Without the + // list layout, list elements are not written as columns of their own, so a path into a list + // is kept by the writer of the list column that holds it. + let all_allowed: Vec = self + .allowed_encodings + .iter() + .flat_map(|(_, ids)| ids.iter().copied()) + .collect(); let mut field_writers = self.field_writers; - if !self.dict_fields.is_empty() { - let dict_compressing = CompressingStrategy::new( + let mut allowed_columns: HashMap)>> = + HashMap::new(); + for (field, ids) in self.allowed_encodings { + let split = if self.use_list_layout { + field.parts().len() + } else { + field + .parts() + .iter() + .position(|part| matches!(part, Field::ElementType)) + .unwrap_or_else(|| field.parts().len()) + }; + let column = FieldPath::from(field.parts()[..split].to_vec()); + let within = FieldPath::from(field.parts()[split..].to_vec()); + allowed_columns + .entry(column) + .or_default() + .push((within, ids)); + } + for (column, allowed) in allowed_columns { + let column_ids: Vec = allowed + .iter() + .flat_map(|(_, ids)| ids.iter().copied()) + .collect(); + let allowed_compressing = CompressingStrategy::new( BufferedStrategy::new(ChunkedLayoutStrategy::new(Arc::clone(&flat)), 2 * ONE_MEG), - DictChildrenCompressor::new(data_compressor), + AllowedEncodingsCompressor::new(Arc::clone(&data_compressor), allowed), ); - let dict_coalescing = RepartitionStrategy::new( - dict_compressing, + let allowed_coalescing = RepartitionStrategy::new( + allowed_compressing, RepartitionWriterOptions { block_size_minimum: self.data_block_target_bytes.unwrap_or(0), block_len_multiple: self.row_block_size, @@ -281,30 +324,26 @@ impl WriteStrategyBuilder { canonicalize: true, }, ) - .with_keep_dicts(); - let dict_stats = ZonedStrategy::new( - dict_coalescing, + .with_allowed_encodings(column_ids.iter().copied()); + let allowed_stats = ZonedStrategy::new( + allowed_coalescing, compress_then_flat.clone(), ZonedLayoutOptions { block_size: row_block_size, ..Default::default() }, ); - let dict_writer: Arc = Arc::new( - RepartitionStrategy::new( - dict_stats, - RepartitionWriterOptions { - block_size_minimum: 0, - block_len_multiple: self.row_block_size, - block_size_target: None, - canonicalize: false, - }, - ) - .with_keep_dicts(), - ); - for field in self.dict_fields { - field_writers.insert(field, Arc::clone(&dict_writer)); - } + let allowed_writer = RepartitionStrategy::new( + allowed_stats, + RepartitionWriterOptions { + block_size_minimum: 0, + block_len_multiple: self.row_block_size, + block_size_target: None, + canonicalize: false, + }, + ) + .with_allowed_encodings(column_ids); + field_writers.insert(column, Arc::new(allowed_writer)); } // 0. start with splitting columns @@ -327,7 +366,7 @@ impl WriteStrategyBuilder { ..Default::default() }, ); - // Dictionaries nested in the list's elements stay encoded until the elements + // Allowed encodings nested in the list's elements are kept until the elements // reach their own field writers. Arc::new( RepartitionStrategy::new( @@ -339,7 +378,7 @@ impl WriteStrategyBuilder { canonicalize: false, }, ) - .with_keep_dicts(), + .with_allowed_encodings(all_allowed.iter().copied()), ) }, ); diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f7d64257afa..a00b64a8e81 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -19,6 +19,7 @@ use rstest::rstest; use tempfile::tempdir; use vortex_array::ArrayRef; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnVTableExt; @@ -3393,15 +3394,15 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { Ok(()) } -/// Columns named with `with_dict_field`, including one inside a list's elements, keep the -/// dictionaries they were written with. The values are all distinct, so the default path would +/// Columns allowed to keep `Dict` with `with_allow_encodings`, including one inside a list's +/// elements, keep the dictionaries they were written with. The values are all distinct, so the default path would /// not dictionary-encode them. #[rstest] #[case::default(false)] #[case::dict_fields(true)] #[tokio::test] #[cfg_attr(miri, ignore)] -async fn dict_fields_keep_their_dictionaries(#[case] dict_fields: bool) -> VortexResult<()> { +async fn allowed_dicts_keep_their_dictionaries(#[case] dict_fields: bool) -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let distinct = |len: u32, prefix: &str| -> VortexResult { let values: Vec = (0..len).map(|i| format!("{prefix}_{i:08}")).collect(); @@ -3429,11 +3430,14 @@ async fn dict_fields_keep_their_dictionaries(#[case] dict_fields: bool) -> Vorte let mut strategy = crate::strategy::WriteStrategyBuilder::from_session(&SESSION); if dict_fields { - strategy = strategy.with_dict_field(field_path!(name)).with_dict_field( - FieldPath::from_name("items") - .push(Field::ElementType) - .push("tag"), - ); + strategy = strategy + .with_allow_encodings(field_path!(name), [Dict.id()]) + .with_allow_encodings( + FieldPath::from_name("items") + .push(Field::ElementType) + .push("tag"), + [Dict.id()], + ); } let mut buf = ByteBufferMut::empty(); SESSION diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index e290936ca58..158f4bb41a4 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -5,15 +5,24 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt as _; +use vortex_array::ArrayId; use vortex_array::ArrayRef; +use vortex_array::ArraySlots; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; -use vortex_array::arrays::Dict; -use vortex_array::arrays::DictArray; -use vortex_array::arrays::dict::DictArraySlotsExt; +use vortex_array::arrays::List; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::listview::list_from_list_view; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Field; +use vortex_array::dtype::FieldPath; use vortex_array::expr::stats::Stat; use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexResult; @@ -58,32 +67,110 @@ impl CompressorPlugin for BtrBlocksCompressor { } } -/// Compresses a dictionary's codes and values separately, keeping the dictionary itself. +/// Keeps arrays of allowed encodings at given paths of a column, compressing only their children. /// -/// Readers that export the column to an Arrow dictionary then get the stored dictionary back -/// instead of re-encoding the decoded values. Chunks that are not dictionaries, such as a -/// constant slice of one, go to the inner compressor unchanged. -pub struct DictChildrenCompressor { +/// Each path is relative to the column being compressed: the root path names the column itself, +/// and deeper paths step through struct fields and list elements. An array at such a path whose +/// encoding is allowed there keeps that encoding, and each of its child slots is compressed with +/// the inner compressor. For example, a dictionary keeps its codes and values as a dictionary, so +/// a reader exporting the field to an Arrow dictionary gets the stored dictionary back instead of +/// re-encoding the decoded values. Everything else goes to the inner compressor. +pub struct AllowedEncodingsCompressor { inner: Arc, + allowed: Arc<[(FieldPath, Vec)]>, } -impl DictChildrenCompressor { - /// Creates a compressor that compresses dictionary children with `inner`. - pub fn new(inner: C) -> Self { +impl AllowedEncodingsCompressor { + /// Creates a compressor that keeps the `allowed` encodings at their paths, compressing + /// everything with `inner`. + pub fn new( + inner: C, + allowed: impl IntoIterator)>, + ) -> Self { Self { inner: Arc::new(inner), + allowed: allowed.into_iter().collect(), + } + } + + fn compress_at( + &self, + array: &ArrayRef, + path: &FieldPath, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if self + .allowed + .iter() + .any(|(at, ids)| at == path && ids.contains(&array.encoding_id())) + { + let slots = array + .slots() + .iter() + .map(|slot| { + slot.as_ref() + .map(|child| self.inner.compress_chunk(child, ctx)) + .transpose() + }) + .collect::>()?; + // SAFETY: compression rewrites each child slot to an equivalent physical + // representation, preserving the parent array's logical values. + return unsafe { array.clone().with_slots(slots) }; + } + let below = |field: Field| { + let child = path.clone().push(field); + self.allowed + .iter() + .any(|(at, _)| at.parts().starts_with(child.parts())) + }; + + match array.dtype() { + DType::Struct(fields, _) + if fields + .names() + .iter() + .any(|name| below(Field::Name(name.clone()))) => + { + let array = array.clone().execute::(ctx)?; + let compressed = fields + .names() + .iter() + .enumerate() + .map(|(idx, name)| { + let field = array.unmasked_field(idx); + if below(Field::Name(name.clone())) { + self.compress_at(field, &path.clone().push(name.clone()), ctx) + } else { + self.inner.compress_chunk(field, ctx) + } + }) + .collect::>>()?; + Ok(StructArray::try_new( + fields.names().clone(), + compressed, + array.len(), + array.validity()?, + )? + .into_array()) + } + DType::List(..) if below(Field::ElementType) => { + let list = match array.as_opt::() { + Some(list) => list.into_owned(), + None => list_from_list_view(array.clone().execute::(ctx)?, ctx)?, + }; + let elements = + self.compress_at(list.elements(), &path.clone().push(Field::ElementType), ctx)?; + let offsets = self.inner.compress_chunk(list.offsets(), ctx)?; + Ok(ListArray::try_new(elements, offsets, list.validity()?, ctx)?.into_array()) + } + _ => self.inner.compress_chunk(array, ctx), } } } -impl CompressorPlugin for DictChildrenCompressor { +impl CompressorPlugin for AllowedEncodingsCompressor { fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let Some(dict) = chunk.as_opt::() else { - return self.inner.compress_chunk(chunk, ctx); - }; - let codes = self.inner.compress_chunk(dict.codes(), ctx)?; - let values = self.inner.compress_chunk(dict.values(), ctx)?; - Ok(DictArray::try_new(codes, values)?.into_array()) + self.compress_at(chunk, &FieldPath::root(), ctx) } } @@ -167,3 +254,58 @@ impl LayoutStrategy for CompressingStrategy { .await } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::VTable; + use vortex_array::arrays::Dict; + use vortex_array::arrays::DictArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::assert_arrays_eq; + use vortex_array::validity::Validity; + use vortex_btrblocks::BtrBlocksCompressorBuilder; + + use super::*; + use crate::test::new_session; + + /// A dictionary inside a list's elements keeps its encoding only where it is allowed. + #[rstest] + #[case::not_allowed(false)] + #[case::allowed(true)] + fn allowed_dict_in_list_elements_is_kept(#[case] allow: bool) -> VortexResult<()> { + let mut ctx = new_session().create_execution_ctx(); + let values: Vec = (0..6).map(|i| format!("value_{i}")).collect(); + let tags = DictArray::try_new( + PrimitiveArray::from_iter(0..6u32).into_array(), + VarBinViewArray::from_iter_str(values).into_array(), + )? + .into_array(); + let elements = StructArray::from_fields(&[("tag", tags)])?.into_array(); + let list = ListArray::try_new( + elements, + PrimitiveArray::from_iter([0u32, 2, 6]).into_array(), + Validity::NonNullable, + &mut ctx, + )? + .into_array(); + + let tag_path = FieldPath::root().push(Field::ElementType).push("tag"); + let allowed = allow.then(|| (tag_path, vec![Dict.id()])); + let compressor = + AllowedEncodingsCompressor::new(BtrBlocksCompressorBuilder::empty().build(), allowed); + let compressed = compressor.compress_chunk(&list, &mut ctx)?; + + let tag = compressed + .as_::() + .elements() + .clone() + .execute::(&mut ctx)? + .unmasked_field_by_name("tag")? + .clone(); + assert_eq!(tag.is::(), allow, "{}", compressed.display_tree()); + assert_arrays_eq!(&compressed, &list, &mut ctx); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index 8e7b309229a..aa960f49958 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -8,10 +8,12 @@ use async_stream::try_stream; use async_trait::async_trait; use futures::StreamExt as _; use futures::pin_mut; +use vortex_array::ArrayId; use vortex_array::ArrayRef; use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::VTable; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; @@ -109,7 +111,7 @@ impl RepartitionWriterOptions { pub struct RepartitionStrategy { child: Arc, options: RepartitionWriterOptions, - keep_dicts: bool, + allowed_encodings: Arc<[ArrayId]>, } impl RepartitionStrategy { @@ -117,16 +119,18 @@ impl RepartitionStrategy { Self { child: Arc::new(child), options, - keep_dicts: false, + allowed_encodings: Arc::from([]), } } - /// Keeps dictionaries, at any depth of structs and lists, instead of canonicalizing them. + /// Keeps arrays of the given encodings, at any depth of structs and lists, instead of + /// canonicalizing them. /// - /// A block made of dictionary chunks becomes a single dictionary holding only the distinct - /// values its codes reference. Blocks without dictionaries are canonicalized as usual. - pub fn with_keep_dicts(mut self) -> Self { - self.keep_dicts = true; + /// When [`Dict`] is allowed, a block made of dictionary chunks becomes a single dictionary + /// holding only the distinct values its codes reference. A block made of one chunk of an + /// allowed encoding is passed on as it is. Other blocks are canonicalized as usual. + pub fn with_allowed_encodings(mut self, ids: impl IntoIterator) -> Self { + self.allowed_encodings = ids.into_iter().collect(); self } } @@ -144,14 +148,15 @@ impl LayoutStrategy for RepartitionStrategy { // TODO(os): spawn stream below like: // canon_stream = stream.map(async {to_canonical}).map(spawn).buffered(parallelism) let dtype = stream.dtype().clone(); - let keep_dicts = self.keep_dicts; + let allowed = Arc::clone(&self.allowed_encodings); let stream = if self.options.canonicalize { let canonicalize_session = session.clone(); + let allowed = Arc::clone(&allowed); SequentialStreamAdapter::new( dtype.clone(), stream.map(move |chunk| { let (sequence_id, chunk) = chunk?; - if keep_dicts && contains_dict(&chunk) { + if contains_allowed(&chunk, &allowed) { return VortexResult::Ok((sequence_id, chunk)); } let mut ctx = canonicalize_session.create_execution_ctx(); @@ -194,7 +199,7 @@ impl LayoutStrategy for RepartitionStrategy { let output_chunks = chunks.collect_exact_blocks(&mut ctx)?; assert!(!output_chunks.is_empty()); if let Some(block) = - finish_block(output_chunks, &dtype_clone, keep_dicts, &mut ctx)? + finish_block(output_chunks, &dtype_clone, &allowed, &mut ctx)? { yield (sequence_pointer.advance(), block) } @@ -203,7 +208,7 @@ impl LayoutStrategy for RepartitionStrategy { if canonical_stream.as_mut().peek().await.is_none() { let to_flush = chunks.data.drain(..).map(|(arr, _)| arr).collect(); if let Some(block) = - finish_block(to_flush, &dtype_clone, keep_dicts, &mut ctx)? + finish_block(to_flush, &dtype_clone, &allowed, &mut ctx)? { yield (sequence_pointer.advance(), block) } @@ -305,19 +310,19 @@ impl ChunksBuffer { /// Assembles the chunks of one block, or returns `None` for an empty block. /// -/// See [`RepartitionStrategy::with_keep_dicts`] for which blocks keep their dictionaries. +/// See [`RepartitionStrategy::with_allowed_encodings`] for which blocks keep their encodings. fn finish_block( chunks: Vec, dtype: &DType, - keep_dicts: bool, + allowed: &[ArrayId], ctx: &mut ExecutionCtx, ) -> VortexResult> { let chunks: Vec = chunks.into_iter().filter(|c| !c.is_empty()).collect(); if chunks.is_empty() { return Ok(None); } - if keep_dicts && chunks.iter().any(contains_dict) { - return concat_preserving_dicts(chunks, dtype, ctx).map(Some); + if chunks.iter().any(|chunk| contains_allowed(chunk, allowed)) { + return concat_allowed(chunks, dtype, allowed, ctx).map(Some); } let chunked = ChunkedArray::try_new(chunks, dtype.clone())?; Ok(Some( @@ -325,26 +330,36 @@ fn finish_block( )) } -/// Whether a dictionary appears anywhere in `array`. -fn contains_dict(array: &ArrayRef) -> bool { - array - .depth_first_traversal() - .any(|child| child.is::()) +/// Whether an array of an allowed encoding appears anywhere in `array`. +fn contains_allowed(array: &ArrayRef, allowed: &[ArrayId]) -> bool { + !allowed.is_empty() + && array + .depth_first_traversal() + .any(|child| allowed.contains(&child.encoding_id())) } -/// Concatenates chunks, keeping the dictionaries found at any depth of structs and lists. -fn concat_preserving_dicts( +/// Concatenates chunks, keeping the arrays of allowed encodings found at any depth of structs and +/// lists. +fn concat_allowed( chunks: Vec, dtype: &DType, + allowed: &[ArrayId], ctx: &mut ExecutionCtx, ) -> VortexResult { let chunks: Vec = chunks.into_iter().filter(|c| !c.is_empty()).collect(); if chunks.is_empty() { return Ok(Canonical::empty(dtype).into_array()); } - if let Some(dict) = concat_dicts(&chunks, dtype, ctx)? { + if allowed.contains(&Dict.id()) + && let Some(dict) = concat_dicts(&chunks, dtype, ctx)? + { return Ok(dict); } + if let [chunk] = chunks.as_slice() + && allowed.contains(&chunk.encoding_id()) + { + return Ok(chunk.clone()); + } let len = chunks.iter().map(|c| c.len()).sum(); match dtype { @@ -362,7 +377,7 @@ fn concat_preserving_dicts( .iter() .map(|chunk| chunk.unmasked_field(idx).clone()) .collect(); - concat_preserving_dicts(field_chunks, &field_dtype, ctx) + concat_allowed(field_chunks, &field_dtype, allowed, ctx) }) .collect::>>()?; Ok( @@ -400,7 +415,7 @@ fn concat_preserving_dicts( .slice(usize::try_from(first)?..usize::try_from(last)?)?, ); } - let elements = concat_preserving_dicts(elements, element_dtype, ctx)?; + let elements = concat_allowed(elements, element_dtype, allowed, ctx)?; Ok( ListArray::try_new(elements, offsets.freeze().into_array(), validity, ctx)? .into_array(), @@ -932,8 +947,8 @@ mod tests { ]; let expected = ChunkedArray::try_new(chunks.clone(), dtype.clone())?; - let block = - finish_block(chunks, &dtype, true, &mut ctx)?.vortex_expect("block is not empty"); + let block = finish_block(chunks, &dtype, &[Dict.id()], &mut ctx)? + .vortex_expect("block is not empty"); let dict = block .as_opt::() @@ -964,8 +979,8 @@ mod tests { let chunks = vec![first, second]; let expected = ChunkedArray::try_new(chunks.clone(), dtype.clone())?; - let block = - finish_block(chunks, &dtype, true, &mut ctx)?.vortex_expect("block is not empty"); + let block = finish_block(chunks, &dtype, &[Dict.id()], &mut ctx)? + .vortex_expect("block is not empty"); let elements = block.as_::().elements().clone(); assert_eq!(elements.len(), 6); From 263a0b112a8b7b4ae3d96ff68c701787f6ca2f7f Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 29 Oct 2026 18:32:01 +0000 Subject: [PATCH 5/8] Let the dictionary layout write dictionary input without decoding it `DictStrategy` now takes chunks that are already dictionaries: it adds their values to the running dictionary and remaps their codes, so only values are hashed and rows are never decoded. Such input skips the probe and is always written as a dictionary layout, which a scan reads back as a dictionary. With that, `with_allow_encodings(path, encodings)` only needs repartitioning to keep the allowed encodings: the field takes the default steps and its layouts see the encodings it arrived with. `AllowedEncodingsCompressor` is removed and the compressor is unchanged from develop. Repartitioning no longer compacts merged dictionaries, since the dictionary layout dedupes their values. A path into list elements needs the list layout, which writes the elements as columns of their own; `with_list_layout` enables it, and the benchmark conversion does so when a dictionary field sits inside a list. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AUHM7GB9hiFe56JK863CKS --- vortex-bench/src/conversions.rs | 8 ++ vortex-file/src/strategy.rs | 98 ++++--------- vortex-layout/src/layouts/compressed.rs | 175 ----------------------- vortex-layout/src/layouts/dict/writer.rs | 155 ++++++++++++++++++++ vortex-layout/src/layouts/repartition.rs | 86 ++--------- 5 files changed, 199 insertions(+), 323 deletions(-) diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index e1463563ea7..76721debfe0 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -378,6 +378,14 @@ fn write_options_for( } let mut builder = WriteStrategyBuilder::from_session(&SESSION); + // The list layout writes list elements as columns of their own, so dictionary fields inside + // lists reach the dictionary layout too. + if dict_fields + .iter() + .any(|field| field.parts().contains(&Field::ElementType)) + { + builder = builder.with_list_layout(true); + } for field in dict_fields { builder = builder.with_allow_encodings(field, [Dict.id()]); } diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index d5dcf5cb0c6..64b6825a438 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -7,7 +7,6 @@ use std::num::NonZeroUsize; use std::sync::Arc; use vortex_array::ArrayId; -use vortex_array::dtype::Field; use vortex_array::dtype::FieldPath; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_btrblocks::SchemeExt; @@ -20,7 +19,6 @@ use vortex_layout::VTable; use vortex_layout::layouts::buffered::BufferedStrategy; use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; use vortex_layout::layouts::collect::CollectStrategy; -use vortex_layout::layouts::compressed::AllowedEncodingsCompressor; use vortex_layout::layouts::compressed::CompressingStrategy; use vortex_layout::layouts::compressed::CompressorPlugin; use vortex_layout::layouts::dict::writer::DictLayoutOptions; @@ -124,16 +122,14 @@ impl WriteStrategyBuilder { self } - /// Allow the field at `field` to keep the given encodings it arrives with, compressing only - /// their children. + /// Allow the field at `field` to reach its layouts in the given encodings it arrives with. /// - /// The field is written without being canonicalized where it holds an allowed encoding: such - /// an array keeps its encoding and only its child slots are compressed. Use this for a column - /// whose readers want the encoding back, such as an Arrow dictionary column written with - /// [`Dict`]: a scan hands back the stored dictionary instead of re-encoding the values. Blocks - /// of dictionary chunks are merged into one dictionary of the values they reference. The path - /// may step into list elements; without the list layout, the column holding the list keeps the - /// encodings inside it. + /// Repartitioning normally canonicalizes every chunk. For this field it keeps arrays of the + /// allowed encodings, so the layouts below can store them as they are. With [`Dict`], the + /// dictionary layout writes an Arrow-style dictionary column without decoding its rows, only + /// merging each chunk's values into its running dictionary, and a scan hands the dictionary + /// back instead of re-encoding the values. A path into list elements needs the list layout + /// (see [`Self::with_list_layout`]), which writes the elements as columns of their own. /// /// [`Dict`]: vortex_array::arrays::Dict pub fn with_allow_encodings( @@ -146,6 +142,15 @@ impl WriteStrategyBuilder { self } + /// Write list fields with the list layout, which splits a list into columns for its elements, + /// offsets and validity. + /// + /// The session's enabled editions decide the default. + pub fn with_list_layout(mut self, enabled: bool) -> Self { + self.use_list_layout = enabled; + self + } + /// Override the flat layout strategy used for leaf chunks. /// /// By default, this uses [`FlatLayoutStrategy`]. This can be used to substitute a custom @@ -208,7 +213,7 @@ impl WriteStrategyBuilder { ), CompressorConfig::Opaque(compressor) => Arc::clone(compressor), }; - let compressing = CompressingStrategy::new(buffered, Arc::clone(&data_compressor)); + let compressing = CompressingStrategy::new(buffered, data_compressor); // 4. prior to compression, coalesce up to a minimum size let coalescing = RepartitionStrategy::new( @@ -232,8 +237,7 @@ impl WriteStrategyBuilder { CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()), CompressorConfig::Opaque(compressor) => compressor, }; - let compress_then_flat = - CompressingStrategy::new(Arc::clone(&flat), Arc::clone(&stats_compressor)); + let compress_then_flat = CompressingStrategy::new(flat, Arc::clone(&stats_compressor)); // 3. apply dict encoding or fallback let probe_compressor = if let Some(probe_compressor) = self.probe_compressor { @@ -255,18 +259,18 @@ impl WriteStrategyBuilder { let row_block_size = NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0"); // 2. calculate stats for each row group - let stats = ZonedStrategy::new( + let stats: Arc = Arc::new(ZonedStrategy::new( dict, compress_then_flat.clone(), ZonedLayoutOptions { block_size: row_block_size, ..Default::default() }, - ); + )); // 1. repartition each column to fixed row counts let repartition = RepartitionStrategy::new( - stats, + Arc::clone(&stats), RepartitionWriterOptions { // No minimum block size in bytes block_size_minimum: 0, @@ -277,64 +281,18 @@ impl WriteStrategyBuilder { }, ); - // Fields with allowed encodings follow the same steps without dictionary encoding, keeping - // the allowed encodings they arrive with and compressing only their children. Without the - // list layout, list elements are not written as columns of their own, so a path into a list - // is kept by the writer of the list column that holds it. + // Fields with allowed encodings take the same steps, but repartitioning keeps the allowed + // encodings they arrive with, so the layouts below see them: `DictStrategy` writes a + // dictionary as a dictionary layout without decoding its rows. let all_allowed: Vec = self .allowed_encodings .iter() .flat_map(|(_, ids)| ids.iter().copied()) .collect(); let mut field_writers = self.field_writers; - let mut allowed_columns: HashMap)>> = - HashMap::new(); for (field, ids) in self.allowed_encodings { - let split = if self.use_list_layout { - field.parts().len() - } else { - field - .parts() - .iter() - .position(|part| matches!(part, Field::ElementType)) - .unwrap_or_else(|| field.parts().len()) - }; - let column = FieldPath::from(field.parts()[..split].to_vec()); - let within = FieldPath::from(field.parts()[split..].to_vec()); - allowed_columns - .entry(column) - .or_default() - .push((within, ids)); - } - for (column, allowed) in allowed_columns { - let column_ids: Vec = allowed - .iter() - .flat_map(|(_, ids)| ids.iter().copied()) - .collect(); - let allowed_compressing = CompressingStrategy::new( - BufferedStrategy::new(ChunkedLayoutStrategy::new(Arc::clone(&flat)), 2 * ONE_MEG), - AllowedEncodingsCompressor::new(Arc::clone(&data_compressor), allowed), - ); - let allowed_coalescing = RepartitionStrategy::new( - allowed_compressing, - RepartitionWriterOptions { - block_size_minimum: self.data_block_target_bytes.unwrap_or(0), - block_len_multiple: self.row_block_size, - block_size_target: self.data_block_target_bytes, - canonicalize: true, - }, - ) - .with_allowed_encodings(column_ids.iter().copied()); - let allowed_stats = ZonedStrategy::new( - allowed_coalescing, - compress_then_flat.clone(), - ZonedLayoutOptions { - block_size: row_block_size, - ..Default::default() - }, - ); - let allowed_writer = RepartitionStrategy::new( - allowed_stats, + let writer = RepartitionStrategy::new( + Arc::clone(&stats), RepartitionWriterOptions { block_size_minimum: 0, block_len_multiple: self.row_block_size, @@ -342,8 +300,8 @@ impl WriteStrategyBuilder { canonicalize: false, }, ) - .with_allowed_encodings(column_ids); - field_writers.insert(column, Arc::new(allowed_writer)); + .with_allowed_encodings(ids); + field_writers.insert(field, Arc::new(writer)); } // 0. start with splitting columns diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index 158f4bb41a4..07d83bf41c7 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -5,24 +5,11 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt as _; -use vortex_array::ArrayId; use vortex_array::ArrayRef; -use vortex_array::ArraySlots; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; -use vortex_array::arrays::List; -use vortex_array::arrays::ListArray; -use vortex_array::arrays::ListViewArray; -use vortex_array::arrays::StructArray; -use vortex_array::arrays::list::ListArraySlotsExt; -use vortex_array::arrays::listview::list_from_list_view; -use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::dtype::DType; -use vortex_array::dtype::Field; -use vortex_array::dtype::FieldPath; use vortex_array::expr::stats::Stat; use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexResult; @@ -67,113 +54,6 @@ impl CompressorPlugin for BtrBlocksCompressor { } } -/// Keeps arrays of allowed encodings at given paths of a column, compressing only their children. -/// -/// Each path is relative to the column being compressed: the root path names the column itself, -/// and deeper paths step through struct fields and list elements. An array at such a path whose -/// encoding is allowed there keeps that encoding, and each of its child slots is compressed with -/// the inner compressor. For example, a dictionary keeps its codes and values as a dictionary, so -/// a reader exporting the field to an Arrow dictionary gets the stored dictionary back instead of -/// re-encoding the decoded values. Everything else goes to the inner compressor. -pub struct AllowedEncodingsCompressor { - inner: Arc, - allowed: Arc<[(FieldPath, Vec)]>, -} - -impl AllowedEncodingsCompressor { - /// Creates a compressor that keeps the `allowed` encodings at their paths, compressing - /// everything with `inner`. - pub fn new( - inner: C, - allowed: impl IntoIterator)>, - ) -> Self { - Self { - inner: Arc::new(inner), - allowed: allowed.into_iter().collect(), - } - } - - fn compress_at( - &self, - array: &ArrayRef, - path: &FieldPath, - ctx: &mut ExecutionCtx, - ) -> VortexResult { - if self - .allowed - .iter() - .any(|(at, ids)| at == path && ids.contains(&array.encoding_id())) - { - let slots = array - .slots() - .iter() - .map(|slot| { - slot.as_ref() - .map(|child| self.inner.compress_chunk(child, ctx)) - .transpose() - }) - .collect::>()?; - // SAFETY: compression rewrites each child slot to an equivalent physical - // representation, preserving the parent array's logical values. - return unsafe { array.clone().with_slots(slots) }; - } - let below = |field: Field| { - let child = path.clone().push(field); - self.allowed - .iter() - .any(|(at, _)| at.parts().starts_with(child.parts())) - }; - - match array.dtype() { - DType::Struct(fields, _) - if fields - .names() - .iter() - .any(|name| below(Field::Name(name.clone()))) => - { - let array = array.clone().execute::(ctx)?; - let compressed = fields - .names() - .iter() - .enumerate() - .map(|(idx, name)| { - let field = array.unmasked_field(idx); - if below(Field::Name(name.clone())) { - self.compress_at(field, &path.clone().push(name.clone()), ctx) - } else { - self.inner.compress_chunk(field, ctx) - } - }) - .collect::>>()?; - Ok(StructArray::try_new( - fields.names().clone(), - compressed, - array.len(), - array.validity()?, - )? - .into_array()) - } - DType::List(..) if below(Field::ElementType) => { - let list = match array.as_opt::() { - Some(list) => list.into_owned(), - None => list_from_list_view(array.clone().execute::(ctx)?, ctx)?, - }; - let elements = - self.compress_at(list.elements(), &path.clone().push(Field::ElementType), ctx)?; - let offsets = self.inner.compress_chunk(list.offsets(), ctx)?; - Ok(ListArray::try_new(elements, offsets, list.validity()?, ctx)?.into_array()) - } - _ => self.inner.compress_chunk(array, ctx), - } - } -} - -impl CompressorPlugin for AllowedEncodingsCompressor { - fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - self.compress_at(chunk, &FieldPath::root(), ctx) - } -} - /// A layout writer that compresses chunks. #[derive(Clone)] pub struct CompressingStrategy { @@ -254,58 +134,3 @@ impl LayoutStrategy for CompressingStrategy { .await } } - -#[cfg(test)] -mod tests { - use rstest::rstest; - use vortex_array::VTable; - use vortex_array::arrays::Dict; - use vortex_array::arrays::DictArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::arrays::VarBinViewArray; - use vortex_array::assert_arrays_eq; - use vortex_array::validity::Validity; - use vortex_btrblocks::BtrBlocksCompressorBuilder; - - use super::*; - use crate::test::new_session; - - /// A dictionary inside a list's elements keeps its encoding only where it is allowed. - #[rstest] - #[case::not_allowed(false)] - #[case::allowed(true)] - fn allowed_dict_in_list_elements_is_kept(#[case] allow: bool) -> VortexResult<()> { - let mut ctx = new_session().create_execution_ctx(); - let values: Vec = (0..6).map(|i| format!("value_{i}")).collect(); - let tags = DictArray::try_new( - PrimitiveArray::from_iter(0..6u32).into_array(), - VarBinViewArray::from_iter_str(values).into_array(), - )? - .into_array(); - let elements = StructArray::from_fields(&[("tag", tags)])?.into_array(); - let list = ListArray::try_new( - elements, - PrimitiveArray::from_iter([0u32, 2, 6]).into_array(), - Validity::NonNullable, - &mut ctx, - )? - .into_array(); - - let tag_path = FieldPath::root().push(Field::ElementType).push("tag"); - let allowed = allow.then(|| (tag_path, vec![Dict.id()])); - let compressor = - AllowedEncodingsCompressor::new(BtrBlocksCompressorBuilder::empty().build(), allowed); - let compressed = compressor.compress_chunk(&list, &mut ctx)?; - - let tag = compressed - .as_::() - .elements() - .clone() - .execute::(&mut ctx)? - .unmasked_field_by_name("tag")? - .clone(); - assert_eq!(tag.is::(), allow, "{}", compressed.display_tree()); - assert_arrays_eq!(&compressed, &list, &mut ctx); - Ok(()) - } -} diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index 53d13e2a222..f2d19169850 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -21,17 +21,26 @@ use futures::stream::iter; use futures::stream::once; use futures::try_join; use vortex_array::ArrayRef; +use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; +use vortex_array::arrays::Chunked; use vortex_array::arrays::ChunkedArray; +use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Dict; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; +use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::builders::dict::DictConstraints; use vortex_array::builders::dict::DictEncoder; use vortex_array::builders::dict::dict_encoder_in; +use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_buffer::BufferMut; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -169,6 +178,9 @@ impl LayoutStrategy for DictStrategy { let should_fallback = match sample { None => true, // empty stream + // Chunks that arrive dictionary-encoded are written as a dictionary layout without + // being decoded. + Some(sample) if is_dict_encoded(&sample) => false, Some(sample) => { let mut exec_ctx = session.create_execution_ctx(); let compressed = self @@ -332,6 +344,9 @@ impl DictStreamState { chunk: ArrayRef, exec_ctx: &mut ExecutionCtx, ) -> VortexResult> { + if chunk.is::() { + return self.encode_dict(labeler, chunk, exec_ctx); + } let mut res = Vec::new(); let mut to_be_encoded = Some(chunk); while let Some(remaining) = to_be_encoded.take() { @@ -370,6 +385,76 @@ impl DictStreamState { Ok(res) } + /// Encodes a chunk that is already a dictionary by adding its values to the running dictionary + /// and remapping its codes, so only the values are hashed, never the rows. + /// + /// When the running dictionary fills up, it is closed and the chunk starts a new one. A chunk + /// whose values alone exceed the constraints is decoded and encoded row by row. + fn encode_dict( + &mut self, + labeler: &mut DictChunkLabeler, + chunk: ArrayRef, + exec_ctx: &mut ExecutionCtx, + ) -> VortexResult> { + let dict = chunk.as_::(); + let codes = dict.codes().clone().execute::(exec_ctx)?; + let mask = codes.validity()?.execute_mask(codes.len(), exec_ctx)?; + // Null rows point at an extra null value, since the encoder gives nulls a code of their own. + let mut values = dict.values().cast(chunk.dtype().clone())?; + let null_value = if mask.all_true() { + None + } else { + let null = ConstantArray::new(Scalar::null(chunk.dtype().clone()), 1).into_array(); + let value = u32::try_from(values.len())?; + values = ChunkedArray::try_new(vec![values, null], chunk.dtype().clone())?.into_array(); + Some(value) + }; + let values = values.execute::(exec_ctx)?.into_array(); + + let mut res = Vec::new(); + let fresh = self.encoder.is_none(); + let mut encoder = match self.encoder.take() { + Some(encoder) => encoder, + None => dict_encoder_in(&values, &self.constraints, exec_ctx.allocator().clone()), + }; + let value_codes = encoder.encode(&values, exec_ctx)?; + if value_codes.len() < values.len() { + if fresh { + // The chunk's values alone do not fit, so the values the encoder took are dropped + // with it and the rows are encoded one by one. + let canonical = chunk.execute::(exec_ctx)?.into_array(); + return self.encode(labeler, canonical, exec_ctx); + } + // Close the running dictionary; values it took from this chunk stay unreferenced. + res.push(labeler.values(encoder.reset())); + res.extend(self.encode_dict(labeler, chunk, exec_ctx)?); + return Ok(res); + } + + let row_codes = codes + .into_array() + .cast(DType::Primitive(PType::U32, Nullability::Nullable))? + .execute::(exec_ctx)?; + let indices: BufferMut = row_codes + .as_slice::() + .iter() + .enumerate() + .map(|(row, &code)| match null_value { + Some(null_value) if !mask.value(row) => null_value, + _ => code, + }) + .collect(); + let ptype = encoder.codes_ptype(); + let encoded = value_codes + .into_array() + .take(indices.freeze().into_array())? + .execute::(exec_ctx)? + .into_array(); + res.push(labeler.codes(encoded, ptype)); + self.encoder = Some(encoder); + Ok(res) + } + fn drain_values(&mut self, labeler: &mut DictChunkLabeler) -> Vec { match self.encoder.as_mut() { None => Vec::new(), @@ -527,6 +612,14 @@ impl Stream for DictionaryTransformer { /// Peeks chunks until they hold `probe_len` rows, returning the stream with those chunks back in /// front and its first `probe_len` rows, or `None` if the stream is empty. +/// Whether `sample` is a dictionary, or a chunked array of dictionaries. +fn is_dict_encoded(sample: &ArrayRef) -> bool { + sample.is::() + || sample + .as_opt::() + .is_some_and(|chunked| chunked.iter_chunks().all(|chunk| chunk.is::())) +} + async fn peek_probe_sample( mut stream: BoxStream<'static, SequencedChunk>, probe_len: usize, @@ -614,8 +707,11 @@ mod tests { use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ChunkedArray; + use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; + use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::builders::dict::DictConstraints; use vortex_array::builders::dict::dict_encode; @@ -747,6 +843,65 @@ mod tests { .sum() } + /// Rebuilds the rows that dictionary chunks encode: each run of codes refers to the values + /// that follow it. + fn decode_runs(chunks: Vec) -> VortexResult> { + let mut decoded = Vec::new(); + let mut pending = Vec::new(); + for chunk in chunks { + match chunk { + DictionaryChunk::Codes { codes, .. } => pending.push(codes), + DictionaryChunk::Values((_, values)) => { + for codes in pending.drain(..) { + decoded.push(DictArray::try_new(codes, values.clone())?.into_array()); + } + } + } + } + assert!(pending.is_empty(), "codes without values"); + Ok(decoded) + } + + /// Dictionary chunks, including null codes, are encoded from their values without decoding + /// their rows, and a full dictionary starts a new run. + #[rstest] + #[case::one_run(100)] + #[case::split_runs(4)] + fn encode_dictionary_chunks(#[case] max_len: usize) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let dict = |codes: Vec>, values: &[&str]| -> VortexResult { + Ok(DictArray::try_new( + PrimitiveArray::from_option_iter(codes).into_array(), + VarBinViewArray::from_iter_nullable_str(values.iter().map(Some)).into_array(), + )? + .into_array()) + }; + let chunks = vec![ + dict(vec![Some(0), None, Some(1), Some(0)], &["a", "b", "unused"])?, + dict(vec![Some(2), Some(0), None], &["b", "c", "d"])?, + ]; + let mut state = DictStreamState { + encoder: None, + constraints: DictConstraints { + max_bytes: 1 << 20, + max_len, + }, + }; + + let mut labeler = DictChunkLabeler::new(SequenceId::root().advance()); + let mut encoded = Vec::new(); + for chunk in &chunks { + encoded.extend(state.encode(&mut labeler, chunk.clone(), &mut ctx)?); + } + encoded.extend(state.drain_values(&mut labeler)); + + let dtype = chunks[0].dtype().clone(); + let actual = ChunkedArray::try_new(decode_runs(encoded)?, dtype.clone())?; + let expected = ChunkedArray::try_new(chunks, dtype)?; + assert_arrays_eq!(&actual, &expected, &mut ctx); + Ok(()) + } + /// A 49-byte value plus its 16-byte `BinaryView` exceeds the 64-byte budget of an empty /// dictionary. Encoding must still make progress instead of looping forever. #[test] diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index aa960f49958..e49e81a11d5 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -25,28 +25,23 @@ use vortex_array::arrays::DictArray; use vortex_array::arrays::List; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; -use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; -use vortex_array::arrays::VarBinView; use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::arrays::listview::list_from_list_view; use vortex_array::arrays::struct_::StructArrayExt; -use vortex_array::builders::dict::dict_encode; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::validity::Validity; -use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_mask::AllOr; -use vortex_mask::Mask; use vortex_session::VortexSession; use crate::LayoutRef; @@ -447,15 +442,11 @@ fn concat_validity>( Ok(Validity::from_mask(bits.freeze().into(), nullability)) } -/// Concatenates dictionary chunks into one compact dictionary. +/// Concatenates dictionary chunks into one dictionary over all of their values. /// -/// Each chunk's codes are rebased onto its values, then the values no code references are dropped -/// and duplicate values are merged. Only the values are hashed, never the rows, so chunks that -/// each carry a whole Parquet column chunk dictionary end up sharing one small dictionary. -/// -/// A constant chunk, as slicing a dictionary with constant codes produces, contributes a single -/// value. Returns `None` unless every chunk is a dictionary or a constant and at least one is a -/// dictionary. +/// Each chunk's codes are rebased onto its values. A constant chunk, as slicing a dictionary with +/// constant codes produces, contributes a single value. Returns `None` unless every chunk is a +/// dictionary or a constant and at least one is a dictionary. fn concat_dicts( chunks: &[ArrayRef], dtype: &DType, @@ -535,76 +526,15 @@ fn concat_dicts( .into_array() .execute::(ctx)? .into_array(); - let validity = validity.freeze(); - let (codes, values) = compact_dict_values(codes, &validity, values, ctx)?; - let codes = PrimitiveArray::new( codes.freeze(), - Validity::from_mask(validity.into(), codes_nullability), + Validity::from_mask(validity.freeze().into(), codes_nullability), ); - // SAFETY: every code was rebased onto the values its chunk contributed, then remapped onto - // the compacted values. + // SAFETY: every code was rebased onto the values its chunk contributed. let dict = unsafe { DictArray::new_unchecked(codes.into_array(), values) }; Ok(Some(dict.into_array())) } -/// Drops the values no valid code references and merges duplicate values, remapping the codes. -/// -/// Values that cannot be hashed are only filtered, and values holding nulls are left as they are. -fn compact_dict_values( - mut codes: BufferMut, - validity: &BitBuffer, - values: ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult<(BufferMut, ArrayRef)> { - let mut referenced = vec![false; values.len()]; - for (code, valid) in codes.iter().zip(validity.iter()) { - if valid { - referenced[*code as usize] = true; - } - } - // Null codes point at value 0, so it is kept to stay in bounds. - if let Some(first) = referenced.first_mut() { - *first = true; - } - if values.all_invalid(ctx)? || !values.all_valid(ctx)? { - return Ok((codes, values)); - } - - let mut remap = vec![0u32; values.len()]; - let mut kept = 0u32; - for (slot, &is_referenced) in remap.iter_mut().zip(&referenced) { - if is_referenced { - *slot = kept; - kept += 1; - } - } - let values = values - .filter(Mask::from(BitBuffer::from(referenced.clone())))? - .execute::(ctx)? - .into_array(); - - let values = if values.is::() || values.is::() { - let unique = dict_encode(&values, ctx)?; - let unique_codes = unique - .codes() - .cast(DType::Primitive(PType::U32, Nullability::NonNullable))? - .execute::(ctx)?; - let unique_codes = unique_codes.as_slice::(); - for slot in &mut remap { - *slot = unique_codes[*slot as usize]; - } - unique.values().clone() - } else { - values - }; - - for code in codes.iter_mut() { - *code = remap[*code as usize]; - } - Ok((codes, values)) -} - /// The bytes a block's rows reference, or the bytes it holds for a dtype that cannot be /// measured, such as variant. fn exact_size(chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { @@ -937,7 +867,7 @@ mod tests { } #[test] - fn concat_dicts_merges_and_compacts_values() -> VortexResult<()> { + fn concat_dicts_merges_chunks() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let dtype = DType::Utf8(NonNullable); let chunks = vec![ @@ -953,7 +883,7 @@ mod tests { let dict = block .as_opt::() .vortex_expect("block stays a dictionary"); - assert_eq!(dict.values().len(), 2); + assert_eq!(dict.values().len(), 6); assert_arrays_eq!(&block, &expected, &mut ctx); Ok(()) } From 5afcdf85db8a8328e80da0917ed36a227430e9d2 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Thu, 29 Oct 2026 19:48:23 +0000 Subject: [PATCH 6/8] Revert "Let the dictionary layout write dictionary input without decoding it" This reverts commit 263a0b112a8b7b4ae3d96ff68c701787f6ca2f7f. Keeping dictionaries nested in list elements that way needs the list layout, which only the unstable edition permits, so the PolarSignals conversion failed with "Layout encoding vortex.list not permitted by ctx". The previous version keeps them in the compressor of the list column instead, and needs no layout outside the default editions. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AUHM7GB9hiFe56JK863CKS --- vortex-bench/src/conversions.rs | 8 -- vortex-file/src/strategy.rs | 98 +++++++++---- vortex-layout/src/layouts/compressed.rs | 175 +++++++++++++++++++++++ vortex-layout/src/layouts/dict/writer.rs | 155 -------------------- vortex-layout/src/layouts/repartition.rs | 86 +++++++++-- 5 files changed, 323 insertions(+), 199 deletions(-) diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 76721debfe0..e1463563ea7 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -378,14 +378,6 @@ fn write_options_for( } let mut builder = WriteStrategyBuilder::from_session(&SESSION); - // The list layout writes list elements as columns of their own, so dictionary fields inside - // lists reach the dictionary layout too. - if dict_fields - .iter() - .any(|field| field.parts().contains(&Field::ElementType)) - { - builder = builder.with_list_layout(true); - } for field in dict_fields { builder = builder.with_allow_encodings(field, [Dict.id()]); } diff --git a/vortex-file/src/strategy.rs b/vortex-file/src/strategy.rs index 64b6825a438..d5dcf5cb0c6 100644 --- a/vortex-file/src/strategy.rs +++ b/vortex-file/src/strategy.rs @@ -7,6 +7,7 @@ use std::num::NonZeroUsize; use std::sync::Arc; use vortex_array::ArrayId; +use vortex_array::dtype::Field; use vortex_array::dtype::FieldPath; use vortex_btrblocks::BtrBlocksCompressorBuilder; use vortex_btrblocks::SchemeExt; @@ -19,6 +20,7 @@ use vortex_layout::VTable; use vortex_layout::layouts::buffered::BufferedStrategy; use vortex_layout::layouts::chunked::writer::ChunkedLayoutStrategy; use vortex_layout::layouts::collect::CollectStrategy; +use vortex_layout::layouts::compressed::AllowedEncodingsCompressor; use vortex_layout::layouts::compressed::CompressingStrategy; use vortex_layout::layouts::compressed::CompressorPlugin; use vortex_layout::layouts::dict::writer::DictLayoutOptions; @@ -122,14 +124,16 @@ impl WriteStrategyBuilder { self } - /// Allow the field at `field` to reach its layouts in the given encodings it arrives with. + /// Allow the field at `field` to keep the given encodings it arrives with, compressing only + /// their children. /// - /// Repartitioning normally canonicalizes every chunk. For this field it keeps arrays of the - /// allowed encodings, so the layouts below can store them as they are. With [`Dict`], the - /// dictionary layout writes an Arrow-style dictionary column without decoding its rows, only - /// merging each chunk's values into its running dictionary, and a scan hands the dictionary - /// back instead of re-encoding the values. A path into list elements needs the list layout - /// (see [`Self::with_list_layout`]), which writes the elements as columns of their own. + /// The field is written without being canonicalized where it holds an allowed encoding: such + /// an array keeps its encoding and only its child slots are compressed. Use this for a column + /// whose readers want the encoding back, such as an Arrow dictionary column written with + /// [`Dict`]: a scan hands back the stored dictionary instead of re-encoding the values. Blocks + /// of dictionary chunks are merged into one dictionary of the values they reference. The path + /// may step into list elements; without the list layout, the column holding the list keeps the + /// encodings inside it. /// /// [`Dict`]: vortex_array::arrays::Dict pub fn with_allow_encodings( @@ -142,15 +146,6 @@ impl WriteStrategyBuilder { self } - /// Write list fields with the list layout, which splits a list into columns for its elements, - /// offsets and validity. - /// - /// The session's enabled editions decide the default. - pub fn with_list_layout(mut self, enabled: bool) -> Self { - self.use_list_layout = enabled; - self - } - /// Override the flat layout strategy used for leaf chunks. /// /// By default, this uses [`FlatLayoutStrategy`]. This can be used to substitute a custom @@ -213,7 +208,7 @@ impl WriteStrategyBuilder { ), CompressorConfig::Opaque(compressor) => Arc::clone(compressor), }; - let compressing = CompressingStrategy::new(buffered, data_compressor); + let compressing = CompressingStrategy::new(buffered, Arc::clone(&data_compressor)); // 4. prior to compression, coalesce up to a minimum size let coalescing = RepartitionStrategy::new( @@ -237,7 +232,8 @@ impl WriteStrategyBuilder { CompressorConfig::BtrBlocks(builder) => Arc::new(builder.build()), CompressorConfig::Opaque(compressor) => compressor, }; - let compress_then_flat = CompressingStrategy::new(flat, Arc::clone(&stats_compressor)); + let compress_then_flat = + CompressingStrategy::new(Arc::clone(&flat), Arc::clone(&stats_compressor)); // 3. apply dict encoding or fallback let probe_compressor = if let Some(probe_compressor) = self.probe_compressor { @@ -259,18 +255,18 @@ impl WriteStrategyBuilder { let row_block_size = NonZeroUsize::new(self.row_block_size).vortex_expect("must be non 0"); // 2. calculate stats for each row group - let stats: Arc = Arc::new(ZonedStrategy::new( + let stats = ZonedStrategy::new( dict, compress_then_flat.clone(), ZonedLayoutOptions { block_size: row_block_size, ..Default::default() }, - )); + ); // 1. repartition each column to fixed row counts let repartition = RepartitionStrategy::new( - Arc::clone(&stats), + stats, RepartitionWriterOptions { // No minimum block size in bytes block_size_minimum: 0, @@ -281,18 +277,64 @@ impl WriteStrategyBuilder { }, ); - // Fields with allowed encodings take the same steps, but repartitioning keeps the allowed - // encodings they arrive with, so the layouts below see them: `DictStrategy` writes a - // dictionary as a dictionary layout without decoding its rows. + // Fields with allowed encodings follow the same steps without dictionary encoding, keeping + // the allowed encodings they arrive with and compressing only their children. Without the + // list layout, list elements are not written as columns of their own, so a path into a list + // is kept by the writer of the list column that holds it. let all_allowed: Vec = self .allowed_encodings .iter() .flat_map(|(_, ids)| ids.iter().copied()) .collect(); let mut field_writers = self.field_writers; + let mut allowed_columns: HashMap)>> = + HashMap::new(); for (field, ids) in self.allowed_encodings { - let writer = RepartitionStrategy::new( - Arc::clone(&stats), + let split = if self.use_list_layout { + field.parts().len() + } else { + field + .parts() + .iter() + .position(|part| matches!(part, Field::ElementType)) + .unwrap_or_else(|| field.parts().len()) + }; + let column = FieldPath::from(field.parts()[..split].to_vec()); + let within = FieldPath::from(field.parts()[split..].to_vec()); + allowed_columns + .entry(column) + .or_default() + .push((within, ids)); + } + for (column, allowed) in allowed_columns { + let column_ids: Vec = allowed + .iter() + .flat_map(|(_, ids)| ids.iter().copied()) + .collect(); + let allowed_compressing = CompressingStrategy::new( + BufferedStrategy::new(ChunkedLayoutStrategy::new(Arc::clone(&flat)), 2 * ONE_MEG), + AllowedEncodingsCompressor::new(Arc::clone(&data_compressor), allowed), + ); + let allowed_coalescing = RepartitionStrategy::new( + allowed_compressing, + RepartitionWriterOptions { + block_size_minimum: self.data_block_target_bytes.unwrap_or(0), + block_len_multiple: self.row_block_size, + block_size_target: self.data_block_target_bytes, + canonicalize: true, + }, + ) + .with_allowed_encodings(column_ids.iter().copied()); + let allowed_stats = ZonedStrategy::new( + allowed_coalescing, + compress_then_flat.clone(), + ZonedLayoutOptions { + block_size: row_block_size, + ..Default::default() + }, + ); + let allowed_writer = RepartitionStrategy::new( + allowed_stats, RepartitionWriterOptions { block_size_minimum: 0, block_len_multiple: self.row_block_size, @@ -300,8 +342,8 @@ impl WriteStrategyBuilder { canonicalize: false, }, ) - .with_allowed_encodings(ids); - field_writers.insert(field, Arc::new(writer)); + .with_allowed_encodings(column_ids); + field_writers.insert(column, Arc::new(allowed_writer)); } // 0. start with splitting columns diff --git a/vortex-layout/src/layouts/compressed.rs b/vortex-layout/src/layouts/compressed.rs index 07d83bf41c7..158f4bb41a4 100644 --- a/vortex-layout/src/layouts/compressed.rs +++ b/vortex-layout/src/layouts/compressed.rs @@ -5,11 +5,24 @@ use std::sync::Arc; use async_trait::async_trait; use futures::StreamExt as _; +use vortex_array::ArrayId; use vortex_array::ArrayRef; +use vortex_array::ArraySlots; use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::UncompressedSizeInBytes; use vortex_array::aggregate_fn::fns::uncompressed_size_in_bytes::uncompressed_size_in_bytes; +use vortex_array::arrays::List; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::listview::list_from_list_view; +use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Field; +use vortex_array::dtype::FieldPath; use vortex_array::expr::stats::Stat; use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexResult; @@ -54,6 +67,113 @@ impl CompressorPlugin for BtrBlocksCompressor { } } +/// Keeps arrays of allowed encodings at given paths of a column, compressing only their children. +/// +/// Each path is relative to the column being compressed: the root path names the column itself, +/// and deeper paths step through struct fields and list elements. An array at such a path whose +/// encoding is allowed there keeps that encoding, and each of its child slots is compressed with +/// the inner compressor. For example, a dictionary keeps its codes and values as a dictionary, so +/// a reader exporting the field to an Arrow dictionary gets the stored dictionary back instead of +/// re-encoding the decoded values. Everything else goes to the inner compressor. +pub struct AllowedEncodingsCompressor { + inner: Arc, + allowed: Arc<[(FieldPath, Vec)]>, +} + +impl AllowedEncodingsCompressor { + /// Creates a compressor that keeps the `allowed` encodings at their paths, compressing + /// everything with `inner`. + pub fn new( + inner: C, + allowed: impl IntoIterator)>, + ) -> Self { + Self { + inner: Arc::new(inner), + allowed: allowed.into_iter().collect(), + } + } + + fn compress_at( + &self, + array: &ArrayRef, + path: &FieldPath, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + if self + .allowed + .iter() + .any(|(at, ids)| at == path && ids.contains(&array.encoding_id())) + { + let slots = array + .slots() + .iter() + .map(|slot| { + slot.as_ref() + .map(|child| self.inner.compress_chunk(child, ctx)) + .transpose() + }) + .collect::>()?; + // SAFETY: compression rewrites each child slot to an equivalent physical + // representation, preserving the parent array's logical values. + return unsafe { array.clone().with_slots(slots) }; + } + let below = |field: Field| { + let child = path.clone().push(field); + self.allowed + .iter() + .any(|(at, _)| at.parts().starts_with(child.parts())) + }; + + match array.dtype() { + DType::Struct(fields, _) + if fields + .names() + .iter() + .any(|name| below(Field::Name(name.clone()))) => + { + let array = array.clone().execute::(ctx)?; + let compressed = fields + .names() + .iter() + .enumerate() + .map(|(idx, name)| { + let field = array.unmasked_field(idx); + if below(Field::Name(name.clone())) { + self.compress_at(field, &path.clone().push(name.clone()), ctx) + } else { + self.inner.compress_chunk(field, ctx) + } + }) + .collect::>>()?; + Ok(StructArray::try_new( + fields.names().clone(), + compressed, + array.len(), + array.validity()?, + )? + .into_array()) + } + DType::List(..) if below(Field::ElementType) => { + let list = match array.as_opt::() { + Some(list) => list.into_owned(), + None => list_from_list_view(array.clone().execute::(ctx)?, ctx)?, + }; + let elements = + self.compress_at(list.elements(), &path.clone().push(Field::ElementType), ctx)?; + let offsets = self.inner.compress_chunk(list.offsets(), ctx)?; + Ok(ListArray::try_new(elements, offsets, list.validity()?, ctx)?.into_array()) + } + _ => self.inner.compress_chunk(array, ctx), + } + } +} + +impl CompressorPlugin for AllowedEncodingsCompressor { + fn compress_chunk(&self, chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + self.compress_at(chunk, &FieldPath::root(), ctx) + } +} + /// A layout writer that compresses chunks. #[derive(Clone)] pub struct CompressingStrategy { @@ -134,3 +254,58 @@ impl LayoutStrategy for CompressingStrategy { .await } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::VTable; + use vortex_array::arrays::Dict; + use vortex_array::arrays::DictArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::VarBinViewArray; + use vortex_array::assert_arrays_eq; + use vortex_array::validity::Validity; + use vortex_btrblocks::BtrBlocksCompressorBuilder; + + use super::*; + use crate::test::new_session; + + /// A dictionary inside a list's elements keeps its encoding only where it is allowed. + #[rstest] + #[case::not_allowed(false)] + #[case::allowed(true)] + fn allowed_dict_in_list_elements_is_kept(#[case] allow: bool) -> VortexResult<()> { + let mut ctx = new_session().create_execution_ctx(); + let values: Vec = (0..6).map(|i| format!("value_{i}")).collect(); + let tags = DictArray::try_new( + PrimitiveArray::from_iter(0..6u32).into_array(), + VarBinViewArray::from_iter_str(values).into_array(), + )? + .into_array(); + let elements = StructArray::from_fields(&[("tag", tags)])?.into_array(); + let list = ListArray::try_new( + elements, + PrimitiveArray::from_iter([0u32, 2, 6]).into_array(), + Validity::NonNullable, + &mut ctx, + )? + .into_array(); + + let tag_path = FieldPath::root().push(Field::ElementType).push("tag"); + let allowed = allow.then(|| (tag_path, vec![Dict.id()])); + let compressor = + AllowedEncodingsCompressor::new(BtrBlocksCompressorBuilder::empty().build(), allowed); + let compressed = compressor.compress_chunk(&list, &mut ctx)?; + + let tag = compressed + .as_::() + .elements() + .clone() + .execute::(&mut ctx)? + .unmasked_field_by_name("tag")? + .clone(); + assert_eq!(tag.is::(), allow, "{}", compressed.display_tree()); + assert_arrays_eq!(&compressed, &list, &mut ctx); + Ok(()) + } +} diff --git a/vortex-layout/src/layouts/dict/writer.rs b/vortex-layout/src/layouts/dict/writer.rs index f2d19169850..53d13e2a222 100644 --- a/vortex-layout/src/layouts/dict/writer.rs +++ b/vortex-layout/src/layouts/dict/writer.rs @@ -21,26 +21,17 @@ use futures::stream::iter; use futures::stream::once; use futures::try_join; use vortex_array::ArrayRef; -use vortex_array::Canonical; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; -use vortex_array::arrays::Chunked; use vortex_array::arrays::ChunkedArray; -use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Dict; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::arrays::chunked::ChunkedArrayExt; -use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::builders::dict::DictConstraints; use vortex_array::builders::dict::DictEncoder; use vortex_array::builders::dict::dict_encoder_in; -use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::scalar::Scalar; -use vortex_buffer::BufferMut; use vortex_error::VortexError; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -178,9 +169,6 @@ impl LayoutStrategy for DictStrategy { let should_fallback = match sample { None => true, // empty stream - // Chunks that arrive dictionary-encoded are written as a dictionary layout without - // being decoded. - Some(sample) if is_dict_encoded(&sample) => false, Some(sample) => { let mut exec_ctx = session.create_execution_ctx(); let compressed = self @@ -344,9 +332,6 @@ impl DictStreamState { chunk: ArrayRef, exec_ctx: &mut ExecutionCtx, ) -> VortexResult> { - if chunk.is::() { - return self.encode_dict(labeler, chunk, exec_ctx); - } let mut res = Vec::new(); let mut to_be_encoded = Some(chunk); while let Some(remaining) = to_be_encoded.take() { @@ -385,76 +370,6 @@ impl DictStreamState { Ok(res) } - /// Encodes a chunk that is already a dictionary by adding its values to the running dictionary - /// and remapping its codes, so only the values are hashed, never the rows. - /// - /// When the running dictionary fills up, it is closed and the chunk starts a new one. A chunk - /// whose values alone exceed the constraints is decoded and encoded row by row. - fn encode_dict( - &mut self, - labeler: &mut DictChunkLabeler, - chunk: ArrayRef, - exec_ctx: &mut ExecutionCtx, - ) -> VortexResult> { - let dict = chunk.as_::(); - let codes = dict.codes().clone().execute::(exec_ctx)?; - let mask = codes.validity()?.execute_mask(codes.len(), exec_ctx)?; - // Null rows point at an extra null value, since the encoder gives nulls a code of their own. - let mut values = dict.values().cast(chunk.dtype().clone())?; - let null_value = if mask.all_true() { - None - } else { - let null = ConstantArray::new(Scalar::null(chunk.dtype().clone()), 1).into_array(); - let value = u32::try_from(values.len())?; - values = ChunkedArray::try_new(vec![values, null], chunk.dtype().clone())?.into_array(); - Some(value) - }; - let values = values.execute::(exec_ctx)?.into_array(); - - let mut res = Vec::new(); - let fresh = self.encoder.is_none(); - let mut encoder = match self.encoder.take() { - Some(encoder) => encoder, - None => dict_encoder_in(&values, &self.constraints, exec_ctx.allocator().clone()), - }; - let value_codes = encoder.encode(&values, exec_ctx)?; - if value_codes.len() < values.len() { - if fresh { - // The chunk's values alone do not fit, so the values the encoder took are dropped - // with it and the rows are encoded one by one. - let canonical = chunk.execute::(exec_ctx)?.into_array(); - return self.encode(labeler, canonical, exec_ctx); - } - // Close the running dictionary; values it took from this chunk stay unreferenced. - res.push(labeler.values(encoder.reset())); - res.extend(self.encode_dict(labeler, chunk, exec_ctx)?); - return Ok(res); - } - - let row_codes = codes - .into_array() - .cast(DType::Primitive(PType::U32, Nullability::Nullable))? - .execute::(exec_ctx)?; - let indices: BufferMut = row_codes - .as_slice::() - .iter() - .enumerate() - .map(|(row, &code)| match null_value { - Some(null_value) if !mask.value(row) => null_value, - _ => code, - }) - .collect(); - let ptype = encoder.codes_ptype(); - let encoded = value_codes - .into_array() - .take(indices.freeze().into_array())? - .execute::(exec_ctx)? - .into_array(); - res.push(labeler.codes(encoded, ptype)); - self.encoder = Some(encoder); - Ok(res) - } - fn drain_values(&mut self, labeler: &mut DictChunkLabeler) -> Vec { match self.encoder.as_mut() { None => Vec::new(), @@ -612,14 +527,6 @@ impl Stream for DictionaryTransformer { /// Peeks chunks until they hold `probe_len` rows, returning the stream with those chunks back in /// front and its first `probe_len` rows, or `None` if the stream is empty. -/// Whether `sample` is a dictionary, or a chunked array of dictionaries. -fn is_dict_encoded(sample: &ArrayRef) -> bool { - sample.is::() - || sample - .as_opt::() - .is_some_and(|chunked| chunked.iter_chunks().all(|chunk| chunk.is::())) -} - async fn peek_probe_sample( mut stream: BoxStream<'static, SequencedChunk>, probe_len: usize, @@ -707,11 +614,8 @@ mod tests { use vortex_array::IntoArray; use vortex_array::MaskFuture; use vortex_array::VortexSessionExecute; - use vortex_array::arrays::ChunkedArray; - use vortex_array::arrays::DictArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinArray; - use vortex_array::arrays::VarBinViewArray; use vortex_array::assert_arrays_eq; use vortex_array::builders::dict::DictConstraints; use vortex_array::builders::dict::dict_encode; @@ -843,65 +747,6 @@ mod tests { .sum() } - /// Rebuilds the rows that dictionary chunks encode: each run of codes refers to the values - /// that follow it. - fn decode_runs(chunks: Vec) -> VortexResult> { - let mut decoded = Vec::new(); - let mut pending = Vec::new(); - for chunk in chunks { - match chunk { - DictionaryChunk::Codes { codes, .. } => pending.push(codes), - DictionaryChunk::Values((_, values)) => { - for codes in pending.drain(..) { - decoded.push(DictArray::try_new(codes, values.clone())?.into_array()); - } - } - } - } - assert!(pending.is_empty(), "codes without values"); - Ok(decoded) - } - - /// Dictionary chunks, including null codes, are encoded from their values without decoding - /// their rows, and a full dictionary starts a new run. - #[rstest] - #[case::one_run(100)] - #[case::split_runs(4)] - fn encode_dictionary_chunks(#[case] max_len: usize) -> VortexResult<()> { - let mut ctx = SESSION.create_execution_ctx(); - let dict = |codes: Vec>, values: &[&str]| -> VortexResult { - Ok(DictArray::try_new( - PrimitiveArray::from_option_iter(codes).into_array(), - VarBinViewArray::from_iter_nullable_str(values.iter().map(Some)).into_array(), - )? - .into_array()) - }; - let chunks = vec![ - dict(vec![Some(0), None, Some(1), Some(0)], &["a", "b", "unused"])?, - dict(vec![Some(2), Some(0), None], &["b", "c", "d"])?, - ]; - let mut state = DictStreamState { - encoder: None, - constraints: DictConstraints { - max_bytes: 1 << 20, - max_len, - }, - }; - - let mut labeler = DictChunkLabeler::new(SequenceId::root().advance()); - let mut encoded = Vec::new(); - for chunk in &chunks { - encoded.extend(state.encode(&mut labeler, chunk.clone(), &mut ctx)?); - } - encoded.extend(state.drain_values(&mut labeler)); - - let dtype = chunks[0].dtype().clone(); - let actual = ChunkedArray::try_new(decode_runs(encoded)?, dtype.clone())?; - let expected = ChunkedArray::try_new(chunks, dtype)?; - assert_arrays_eq!(&actual, &expected, &mut ctx); - Ok(()) - } - /// A 49-byte value plus its 16-byte `BinaryView` exceeds the 64-byte budget of an empty /// dictionary. Encoding must still make progress instead of looping forever. #[test] diff --git a/vortex-layout/src/layouts/repartition.rs b/vortex-layout/src/layouts/repartition.rs index e49e81a11d5..aa960f49958 100644 --- a/vortex-layout/src/layouts/repartition.rs +++ b/vortex-layout/src/layouts/repartition.rs @@ -25,23 +25,28 @@ use vortex_array::arrays::DictArray; use vortex_array::arrays::List; use vortex_array::arrays::ListArray; use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::Primitive; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; +use vortex_array::arrays::VarBinView; use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::arrays::list::ListArraySlotsExt; use vortex_array::arrays::listview::list_from_list_view; use vortex_array::arrays::struct_::StructArrayExt; +use vortex_array::builders::dict::dict_encode; use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; use vortex_buffer::BitBufferMut; use vortex_buffer::BufferMut; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_err; use vortex_mask::AllOr; +use vortex_mask::Mask; use vortex_session::VortexSession; use crate::LayoutRef; @@ -442,11 +447,15 @@ fn concat_validity>( Ok(Validity::from_mask(bits.freeze().into(), nullability)) } -/// Concatenates dictionary chunks into one dictionary over all of their values. +/// Concatenates dictionary chunks into one compact dictionary. /// -/// Each chunk's codes are rebased onto its values. A constant chunk, as slicing a dictionary with -/// constant codes produces, contributes a single value. Returns `None` unless every chunk is a -/// dictionary or a constant and at least one is a dictionary. +/// Each chunk's codes are rebased onto its values, then the values no code references are dropped +/// and duplicate values are merged. Only the values are hashed, never the rows, so chunks that +/// each carry a whole Parquet column chunk dictionary end up sharing one small dictionary. +/// +/// A constant chunk, as slicing a dictionary with constant codes produces, contributes a single +/// value. Returns `None` unless every chunk is a dictionary or a constant and at least one is a +/// dictionary. fn concat_dicts( chunks: &[ArrayRef], dtype: &DType, @@ -526,15 +535,76 @@ fn concat_dicts( .into_array() .execute::(ctx)? .into_array(); + let validity = validity.freeze(); + let (codes, values) = compact_dict_values(codes, &validity, values, ctx)?; + let codes = PrimitiveArray::new( codes.freeze(), - Validity::from_mask(validity.freeze().into(), codes_nullability), + Validity::from_mask(validity.into(), codes_nullability), ); - // SAFETY: every code was rebased onto the values its chunk contributed. + // SAFETY: every code was rebased onto the values its chunk contributed, then remapped onto + // the compacted values. let dict = unsafe { DictArray::new_unchecked(codes.into_array(), values) }; Ok(Some(dict.into_array())) } +/// Drops the values no valid code references and merges duplicate values, remapping the codes. +/// +/// Values that cannot be hashed are only filtered, and values holding nulls are left as they are. +fn compact_dict_values( + mut codes: BufferMut, + validity: &BitBuffer, + values: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult<(BufferMut, ArrayRef)> { + let mut referenced = vec![false; values.len()]; + for (code, valid) in codes.iter().zip(validity.iter()) { + if valid { + referenced[*code as usize] = true; + } + } + // Null codes point at value 0, so it is kept to stay in bounds. + if let Some(first) = referenced.first_mut() { + *first = true; + } + if values.all_invalid(ctx)? || !values.all_valid(ctx)? { + return Ok((codes, values)); + } + + let mut remap = vec![0u32; values.len()]; + let mut kept = 0u32; + for (slot, &is_referenced) in remap.iter_mut().zip(&referenced) { + if is_referenced { + *slot = kept; + kept += 1; + } + } + let values = values + .filter(Mask::from(BitBuffer::from(referenced.clone())))? + .execute::(ctx)? + .into_array(); + + let values = if values.is::() || values.is::() { + let unique = dict_encode(&values, ctx)?; + let unique_codes = unique + .codes() + .cast(DType::Primitive(PType::U32, Nullability::NonNullable))? + .execute::(ctx)?; + let unique_codes = unique_codes.as_slice::(); + for slot in &mut remap { + *slot = unique_codes[*slot as usize]; + } + unique.values().clone() + } else { + values + }; + + for code in codes.iter_mut() { + *code = remap[*code as usize]; + } + Ok((codes, values)) +} + /// The bytes a block's rows reference, or the bytes it holds for a dtype that cannot be /// measured, such as variant. fn exact_size(chunk: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { @@ -867,7 +937,7 @@ mod tests { } #[test] - fn concat_dicts_merges_chunks() -> VortexResult<()> { + fn concat_dicts_merges_and_compacts_values() -> VortexResult<()> { let mut ctx = SESSION.create_execution_ctx(); let dtype = DType::Utf8(NonNullable); let chunks = vec![ @@ -883,7 +953,7 @@ mod tests { let dict = block .as_opt::() .vortex_expect("block stays a dictionary"); - assert_eq!(dict.values().len(), 6); + assert_eq!(dict.values().len(), 2); assert_arrays_eq!(&block, &expected, &mut ctx); Ok(()) } From 56b49031df9c4bd1a3865475b06b6b12595e3eac Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 30 Oct 2026 16:11:47 +0000 Subject: [PATCH 7/8] Export chunked dictionaries to Arrow without re-encoding rows Under the list layout, list element columns are read back as chunked arrays of per-chunk dictionaries. to_arrow_dictionary only kept a single Dict or Constant, so these were canonicalized and re-encoded row by row. Export each chunk as an Arrow dictionary and concatenate them instead, which merges only the dictionary values. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AUHM7GB9hiFe56JK863CKS Signed-off-by: Joe Isaacs --- Cargo.lock | 1 + vortex-arrow/Cargo.toml | 1 + vortex-arrow/src/executor/dictionary.rs | 53 ++++++++++++++++++++++++- 3 files changed, 54 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 0a025bc047c..254a92b20df 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11459,6 +11459,7 @@ dependencies = [ "arrow-cast 59.3.2", "arrow-data 59.3.2", "arrow-schema 59.3.2", + "arrow-select 59.3.2", "codspeed-divan-compat", "itertools 0.14.0", "mimalloc", diff --git a/vortex-arrow/Cargo.toml b/vortex-arrow/Cargo.toml index 276ba42fe3d..4ac3cb99a4b 100644 --- a/vortex-arrow/Cargo.toml +++ b/vortex-arrow/Cargo.toml @@ -24,6 +24,7 @@ arrow-buffer = { workspace = true } arrow-cast = { workspace = true } arrow-data = { workspace = true } arrow-schema = { workspace = true, features = ["canonical_extension_types"] } +arrow-select = { workspace = true } itertools = { workspace = true } num-traits = { workspace = true } simdutf8 = { workspace = true } diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index 343c08357af..1e2eb1c2e31 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -13,6 +13,8 @@ use arrow_schema::DataType; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; +use vortex_array::arrays::Chunked; +use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Dict; @@ -27,13 +29,16 @@ use crate::ArrowExporter; use crate::executor::execute_arrow_naive; /// Matches the encodings [`to_arrow_dictionary`] requires for export. +/// +/// `Chunked` is matched so that each chunk keeps its own dictionary instead of the chunks being +/// canonicalized and re-encoded row by row. struct ArrowDictExportable; impl Matcher for ArrowDictExportable { type Match<'a> = &'a ArrayRef; fn try_match(array: &ArrayRef) -> Option> { - (array.is::() || array.is::()).then_some(array) + (array.is::() || array.is::() || array.is::()).then_some(array) } } @@ -54,6 +59,18 @@ pub(super) fn to_arrow_dictionary( Ok(constant) => return constant_to_dict(constant, codes_type, values_type, exporter, ctx), Err(array) => array, }; + if let Some(chunked) = array.as_opt::() { + let chunks = chunked + .iter_chunks() + .map(|chunk| to_arrow_dictionary(chunk.clone(), codes_type, values_type, exporter, ctx)) + .collect::>>()?; + let chunks = chunks.iter().map(|chunk| chunk.as_ref()).collect::>(); + // Concatenating dictionaries merges their values without re-encoding every row. It fails + // only when the merged values overflow the codes type, which the row-wise path handles. + if let Ok(array) = arrow_select::concat::concat(&chunks) { + return Ok(array); + } + } // Otherwise, we should try and build a dictionary. // Arrow hides this functionality inside the cast module! @@ -161,6 +178,8 @@ mod tests { use arrow_array::DictionaryArray as ArrowDictArray; use arrow_array::StringArray; + use arrow_array::UInt8Array; + use arrow_array::cast::AsArray; use arrow_array::types::UInt8Type; use arrow_array::types::UInt32Type; use arrow_schema::DataType; @@ -168,9 +187,11 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ChunkedArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::VarBinViewArray; use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability::NonNullable; use vortex_array::dtype::Nullability::Nullable; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::mask::Mask; @@ -249,6 +270,36 @@ mod tests { Ok(()) } + #[test] + fn chunked_dicts_keep_their_values() -> VortexResult<()> { + let first = DictArray::try_new( + buffer![1u8, 0].into_array(), + VarBinViewArray::from_iter_str(["a", "b"]).into_array(), + )?; + let second = DictArray::try_new( + buffer![0u8, 0].into_array(), + VarBinViewArray::from_iter_str(["c"]).into_array(), + )?; + let chunked = ChunkedArray::try_new( + vec![first.into_array(), second.into_array()], + DType::Utf8(NonNullable), + )?; + + let actual = execute( + chunked.into_array(), + &dict_type(DataType::UInt8, DataType::Utf8), + )?; + + // Re-encoding the rows would order the values by first appearance: `b`, `a`, `c`. + let actual = actual.as_dictionary::(); + assert_eq!( + actual.values().as_ref(), + &StringArray::from(vec!["a", "b", "c"]) as &dyn arrow_array::Array + ); + assert_eq!(actual.keys(), &UInt8Array::from(vec![1, 0, 2, 2])); + Ok(()) + } + #[test] fn mask_wrapped_dict_exports() -> VortexResult<()> { // Dictionary behind a lazy `mask` scalar-fn — the shape a scan produces when a row From fa42c42833cf4b19a2dee43ce969ee601eb55f16 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 30 Oct 2026 16:18:17 +0000 Subject: [PATCH 8/8] Format chunked dictionary export Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AUHM7GB9hiFe56JK863CKS Signed-off-by: Joe Isaacs --- vortex-arrow/src/executor/dictionary.rs | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vortex-arrow/src/executor/dictionary.rs b/vortex-arrow/src/executor/dictionary.rs index 1e2eb1c2e31..b6bc49945a9 100644 --- a/vortex-arrow/src/executor/dictionary.rs +++ b/vortex-arrow/src/executor/dictionary.rs @@ -14,11 +14,11 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::Chunked; -use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::arrays::Constant; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::Dict; use vortex_array::arrays::DictArray; +use vortex_array::arrays::chunked::ChunkedArrayExt; use vortex_array::arrays::dict::DictArraySlotsExt; use vortex_array::matcher::Matcher; use vortex_error::VortexError; @@ -64,7 +64,10 @@ pub(super) fn to_arrow_dictionary( .iter_chunks() .map(|chunk| to_arrow_dictionary(chunk.clone(), codes_type, values_type, exporter, ctx)) .collect::>>()?; - let chunks = chunks.iter().map(|chunk| chunk.as_ref()).collect::>(); + let chunks = chunks + .iter() + .map(|chunk| chunk.as_ref()) + .collect::>(); // Concatenating dictionaries merges their values without re-encoding every row. It fails // only when the merged values overflow the codes type, which the row-wise path handles. if let Ok(array) = arrow_select::concat::concat(&chunks) {