Keep dictionary columns encoded on write, selected by field path

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 }
 base64 = { workspace = true }
 itertools = { workspace = true }
 num-traits = { workspace = true }

vortex-arrow/src/executor/dictionary.rs

@@ -13,10 +13,12 @@ use arrow_schema::DataType;
 use arrow_schema::Field;
 use vortex_array::ArrayRef;
 use vortex_array::ExecutionCtx;
+use vortex_array::arrays::Chunked;
 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;
@@ -28,13 +30,16 @@ use crate::ExecuteArrowExt;
 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<Self::Match<'_>> {
-        (array.is::<Dict>() || array.is::<Constant>()).then_some(array)
+        (array.is::<Dict>() || array.is::<Constant>() || array.is::<Chunked>()).then_some(array)
     }
 }
 
@@ -61,6 +66,21 @@ pub(super) fn to_arrow_dictionary(
         Ok(constant) => return constant_to_dict(constant, codes_type, values_field, options, ctx),
         Err(array) => array,
     };
+    if let Some(chunked) = array.as_opt::<Chunked>() {
+        let chunks = chunked
+            .iter_chunks()
+            .map(|chunk| to_arrow_dictionary(chunk.clone(), codes_type, values_field, options, ctx))
+            .collect::<VortexResult<Vec<_>>>()?;
+        let chunks = chunks
+            .iter()
+            .map(|chunk| chunk.as_ref())
+            .collect::<Vec<_>>();
+        // 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!
@@ -177,6 +197,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;
@@ -185,10 +207,12 @@ mod tests {
     use vortex_array::IntoArray;
     use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::BoolArray;
+    use vortex_array::arrays::ChunkedArray;
     use vortex_array::arrays::FixedSizeListArray;
     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;
@@ -301,6 +325,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::<UInt8Type>();
+        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

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;
@@ -31,8 +33,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;
@@ -43,6 +47,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::Nullability;
 use vortex::dtype::StructFields;
@@ -126,7 +131,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<ArrayRef> = parquet_to_vortex_stream(builder.build()?)
+        let chunks: Vec<ArrayRef> = parquet_to_vortex_stream(builder.build()?, false)
             .map(|r| r.map_err(anyhow::Error::from))
             .try_collect()
             .await?;
@@ -149,7 +154,7 @@ pub async fn parquet_to_vortex_chunks_with_batch_size(
     let dtype = SESSION.arrow().from_arrow_schema(schema.as_ref())?;
     let chunks = (0..combined.num_rows()).step_by(batch_size).map(|start| {
         let len = batch_size.min(combined.num_rows() - start);
-        record_batch_to_vortex(combined.slice(start, len))
+        record_batch_to_vortex(combined.slice(start, len), false)
     });
     let expected_nchunks = combined.num_rows().div_ceil(batch_size);
     Ok(chunks.process_results(|chunks| {
@@ -158,10 +163,16 @@ pub async fn parquet_to_vortex_chunks_with_batch_size(
     })?)
 }
 
-/// Convert one Arrow [`RecordBatch`] into a canonical Vortex array.
-fn record_batch_to_vortex(batch: RecordBatch) -> VortexResult<ArrayRef> {
+/// 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<ArrayRef> {
     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());
 
@@ -174,14 +185,16 @@ fn record_batch_to_vortex(batch: RecordBatch) -> VortexResult<ArrayRef> {
 /// 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<File>,
+    preserve_dicts: bool,
 ) -> impl futures::Stream<Item = VortexResult<ArrayRef>> {
     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))
     })
 }
 
@@ -216,6 +229,7 @@ fn parquet_row_groups_to_vortex_stream(
     parquet_path: PathBuf,
     metadata: ArrowReaderMetadata,
     concurrency: usize,
+    preserve_dicts: bool,
 ) -> impl futures::Stream<Item = VortexResult<ArrayRef>> {
     let num_row_groups = metadata.metadata().num_row_groups();
     futures::stream::iter(0..num_row_groups)
@@ -230,7 +244,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::<Vec<_>>()
                     .await
             })
@@ -261,16 +275,21 @@ 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 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(
             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)));
 
@@ -281,12 +300,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),
+        dict_fields,
+    )
+    .write(
+        &mut output_file,
+        ArrayStreamExt::boxed(ArrayStreamAdapter::new(dtype, stream)),
+    )
+    .await?;
 
     Ok(())
 }
@@ -297,14 +321,53 @@ fn is_spatialbench(path: &Path) -> bool {
         .any(|component| component.as_os_str() == "spatialbench")
 }
 
+/// The paths of the Arrow dictionary fields at any depth of structs and lists.
+fn dictionary_field_paths(fields: &Fields) -> Vec<FieldPath> {
+    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 collect_dictionary_paths(data_type: &DataType, path: FieldPath, paths: &mut Vec<FieldPath>) {
+    match data_type {
+        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) => {
+            collect_dictionary_paths(field.data_type(), path.push(Field::ElementType), paths)
+        }
+        _ => {}
+    }
+}
+
 /// 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.
+///
+/// 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,
+    dict_fields: Vec<FieldPath>,
 ) -> VortexWriteOptions {
     let binary_fields: Vec<_> = match dtype {
         DType::Struct(fields, _) if skip_binary_dict => fields
@@ -316,11 +379,14 @@ fn write_options_for(
             .collect(),
         _ => Vec::new(),
     };
-    if binary_fields.is_empty() {
+    if binary_fields.is_empty() && dict_fields.is_empty() {
         return compaction.apply_options(SESSION.write_options());
     }
 
     let mut builder = WriteStrategyBuilder::from_session(&SESSION);
+    for field in dict_fields {
+        builder = builder.with_allow_encodings(field, [Dict.id()]);
+    }
     if matches!(compaction, CompactionStrategy::Compact) {
         builder = builder.with_btrblocks_builder(
             BtrBlocksCompressorBuilder::from_session(&SESSION).with_compact(),