From 2fbf202edd0ed18989ebf8e9a2bcfd5444152430 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 08:11:28 +0000 Subject: [PATCH 1/2] fix(encoding): handle empty all-null dictionaries --- .../src/encodings/logical/primitive/dict.rs | 30 ++++++++++++- .../dataset/tests/fragment_write_columns.rs | 44 +++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) diff --git a/rust/lance-encoding/src/encodings/logical/primitive/dict.rs b/rust/lance-encoding/src/encodings/logical/primitive/dict.rs index 30d79ec7255..beee69819bc 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/dict.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/dict.rs @@ -11,6 +11,7 @@ pub const DICT_INDICES_BITS_PER_VALUE: u64 = 32; use arrow_array::{ Array, DictionaryArray, PrimitiveArray, UInt64Array, cast::AsArray, + new_null_array, types::{ ArrowDictionaryKeyType, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, @@ -35,6 +36,16 @@ fn normalize_dict_nulls_impl( let dict_array = array.as_dictionary_opt::().expect_ok()?; + if dict_array.values().is_empty() && !dict_array.is_empty() { + // Valid non-empty dictionaries with no values can only contain null keys. The + // structural encoder removes key validity after recording it as rep/def levels, + // so retain an unreachable value at index zero to keep those exposed key bytes valid. + let values = new_null_array(dict_array.values().data_type(), 1); + return Ok( + Arc::new(DictionaryArray::new(dict_array.keys().clone(), values)) as Arc, + ); + } + if dict_array.values().null_count() == 0 { return Ok(array); } @@ -374,9 +385,26 @@ mod tests { buffer::LanceBuffer, data::{BlockInfo, FixedWidthDataBlock}, }; - use arrow_array::{Array, StringArray}; + use arrow_array::{Array, DictionaryArray, Int32Array, StringArray, types::Int32Type}; use std::sync::Arc; + #[test] + fn test_normalize_empty_all_null_dictionary() { + let keys = Int32Array::new_null(2); + let values = Arc::new(StringArray::from(Vec::<&str>::new())) as Arc; + let dictionary = Arc::new(DictionaryArray::::new(keys, values)); + + let normalized = normalize_dict_nulls(dictionary).unwrap(); + let data_without_key_nulls = normalized + .to_data() + .into_builder() + .nulls(None) + .build() + .unwrap(); + + assert_eq!(data_without_key_nulls.len(), 2); + } + #[test] fn test_dictionary_encode_abort_fixed_width() { // Create a u128 block with very high cardinality where dict encoding diff --git a/rust/lance/src/dataset/tests/fragment_write_columns.rs b/rust/lance/src/dataset/tests/fragment_write_columns.rs index 89e6872a866..e3a29207b7c 100644 --- a/rust/lance/src/dataset/tests/fragment_write_columns.rs +++ b/rust/lance/src/dataset/tests/fragment_write_columns.rs @@ -140,6 +140,50 @@ async fn declare_all_null(dataset: &mut Dataset, name: &str) { .unwrap(); } +#[tokio::test] +async fn test_compact_metadata_only_all_null_dictionary() { + let batch = batch_of( + vec![ArrowField::new("id", DataType::Int32, false)], + vec![ints(vec![1, 2])], + ); + let schema = batch.schema(); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://", + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_3), + max_rows_per_file: 1, + enable_stable_row_ids: false, + ..Default::default() + }), + ) + .await + .unwrap(); + + let dictionary_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + dataset + .add_columns( + NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new(vec![ArrowField::new( + "category", + dictionary_type.clone(), + true, + )]))), + None, + None, + ) + .await + .unwrap(); + + compact_files(&mut dataset, CompactionOptions::default(), None) + .await + .unwrap(); + + let batch = dataset.scan().try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 2); + assert_eq!(batch["category"].data_type(), &dictionary_type); + assert_eq!(batch["category"].null_count(), 2); +} + /// Stage `values` for an existing `column` of one fragment. async fn stage_column( dataset: &Dataset, From 9538d284bab3e3aad265f9ec33c8daa87e6f1358 Mon Sep 17 00:00:00 2001 From: Gatefixer <312823363+lance-gatefixer[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 08:40:44 +0000 Subject: [PATCH 2/2] fix(encoding): rebuild empty dictionary chunks at flush Port the page-grouping approach from #8010 so empty chunks reuse neighboring dictionary values and canonicalize hidden key payloads. Co-authored-by: Michal Matczuk Co-authored-by: Beau Hartshorne --- .../src/encodings/logical/primitive.rs | 179 +++++++++++++++++- .../src/encodings/logical/primitive/dict.rs | 30 +-- 2 files changed, 176 insertions(+), 33 deletions(-) diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 20703f15eba..3e8908ddfb8 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -24,7 +24,9 @@ use crate::{ pb21::{self, CompressiveEncoding, PageLayout, compressive_encoding::Compression}, }, }; -use arrow_array::{Array, ArrayRef, PrimitiveArray, cast::AsArray, make_array, types::UInt64Type}; +use arrow_array::{ + Array, ArrayRef, PrimitiveArray, cast::AsArray, make_array, new_null_array, types::UInt64Type, +}; use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, NullBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field as ArrowField}; use bytes::Bytes; @@ -6688,6 +6690,97 @@ impl PrimitiveStructuralEncoder { } } + /// Rebuilds dictionary chunks whose values array is empty. + /// + /// Such chunks are entirely null and retain their key validity until rep/def has been + /// recorded. At flush time their keys are zeroed and attached to a neighboring non-empty + /// dictionary. If every chunk is empty, one non-null placeholder value makes key zero valid. + /// Rep/def retains the logical nullness, so these replacement keys are never exposed. + fn rebuild_empty_dictionary_chunks(arrays: Vec) -> Result> { + if !arrays.iter().any(|array| { + array + .as_any_dictionary_opt() + .is_some_and(|dictionary| dictionary.values().is_empty()) + }) { + return Ok(arrays); + } + + let zeroed = |data_type: &DataType, len: usize| { + new_null_array(data_type, len) + .to_data() + .into_builder() + .nulls(None) + .build() + .map(make_array) + }; + let rebuild = + |keys: Vec, data_type: &DataType, values: ArrayRef| -> Result { + let keys = keys.iter().map(|keys| keys.as_ref()).collect::>(); + let keys = arrow_select::concat::concat(&keys)?; + let data = keys + .to_data() + .into_builder() + .data_type(data_type.clone()) + .child_data(vec![values.to_data()]) + .build()?; + Ok(make_array(data)) + }; + + let mut rebuilt = Vec::with_capacity(arrays.len()); + let mut pending_keys = Vec::with_capacity(arrays.len()); + let mut empty_data_type = None; + for array in arrays { + let Some(dictionary) = array.as_any_dictionary_opt() else { + return Err(Error::invalid_input_source( + "Cannot mix dictionary and non-dictionary chunks".into(), + )); + }; + if dictionary.values().is_empty() { + pending_keys.push(zeroed(dictionary.keys().data_type(), array.len())?); + empty_data_type.get_or_insert_with(|| array.data_type().clone()); + } else if pending_keys.is_empty() { + rebuilt.push(array); + } else { + pending_keys.push(make_array(dictionary.keys().to_data())); + rebuilt.push(rebuild( + std::mem::take(&mut pending_keys), + array.data_type(), + dictionary.values().clone(), + )?); + } + } + + if pending_keys.is_empty() { + return Ok(rebuilt); + } + if let Some(array) = rebuilt.pop() { + let Some(dictionary) = array.as_any_dictionary_opt() else { + return Err(Error::invalid_input_source( + "Cannot mix dictionary and non-dictionary chunks".into(), + )); + }; + let mut keys = Vec::with_capacity(pending_keys.len() + 1); + keys.push(make_array(dictionary.keys().to_data())); + keys.append(&mut pending_keys); + rebuilt.push(rebuild( + keys, + array.data_type(), + dictionary.values().clone(), + )?); + } else { + let data_type = empty_data_type.ok_or_else(|| { + Error::internal("Missing data type for an empty dictionary chunk") + })?; + let DataType::Dictionary(_, value_type) = &data_type else { + return Err(Error::internal(format!( + "Expected dictionary data type, got {data_type}" + ))); + }; + rebuilt.push(rebuild(pending_keys, &data_type, zeroed(value_type, 1)?)?); + } + Ok(rebuilt) + } + // Creates encode tasks, consuming all buffered data fn do_flush( &mut self, @@ -6696,6 +6789,7 @@ impl PrimitiveStructuralEncoder { row_number: u64, num_rows: u64, ) -> Result> { + let arrays = Self::rebuild_empty_dictionary_chunks(arrays)?; DataBlock::validate_arrays(&arrays, &self.field.name)?; let num_values = arrays.iter().map(|arr| arr.len() as u64).sum(); let is_simple_validity = repdefs.iter().all(|rd| rd.is_simple_validity()); @@ -6759,6 +6853,14 @@ impl PrimitiveStructuralEncoder { } else { repdef.add_validity_bitmap(deep_copy_nulls(Some(validity)).unwrap()); } + // Empty dictionaries have no valid key payload. Keep the validity until rep/def is + // grouped with the buffered page; `do_flush` then rebuilds the keys against either a + // neighboring values array or an all-empty-page placeholder. + if let Some(dictionary) = array.as_any_dictionary_opt() + && dictionary.values().is_empty() + { + return Ok(array); + } let data_no_nulls = array.to_data().into_builder().nulls(None).build()?; Ok(make_array(data_no_nulls)) } else { @@ -7162,10 +7264,10 @@ mod tests { use crate::testing::TestEncoding; use crate::testing::{TestCases, check_round_trip_encoding_of_data}; use arrow_array::{ - Array, ArrayRef, FixedSizeListArray, Float32Array, Int8Array, StringArray, UInt8Array, - make_array, + Array, ArrayRef, DictionaryArray, FixedSizeListArray, Float32Array, Int8Array, + PrimitiveArray, StringArray, UInt8Array, make_array, new_null_array, types::Int32Type, }; - use arrow_buffer::ScalarBuffer; + use arrow_buffer::{BooleanBuffer, NullBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field as ArrowField}; use std::collections::HashMap; use std::{collections::VecDeque, sync::Arc}; @@ -7190,6 +7292,75 @@ mod tests { assert!((!PrimitiveStructuralEncoder::is_narrow(&block))); } + fn valued_dictionary() -> ArrayRef { + Arc::new( + DictionaryArray::::try_new( + PrimitiveArray::::from(vec![0]), + Arc::new(StringArray::from(vec!["a"])), + ) + .unwrap(), + ) + } + + fn empty_dictionary() -> ArrayRef { + new_null_array( + &DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + 1, + ) + } + + fn null_valued_dictionary() -> ArrayRef { + Arc::new( + DictionaryArray::::try_new( + PrimitiveArray::::from(vec![0]), + Arc::new(StringArray::from(vec![None::<&str>])), + ) + .unwrap(), + ) + } + + fn sliced_empty_dictionary_with_hidden_key_payload() -> ArrayRef { + let keys = PrimitiveArray::::new( + ScalarBuffer::from(vec![5, 7, 9]), + Some(NullBuffer::new(BooleanBuffer::new_unset(3))), + ); + let values = Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef; + let dictionary = DictionaryArray::::try_new(keys, values).unwrap(); + Arc::new(dictionary.slice(1, 1)) + } + + #[rstest::rstest] + #[case::empty_after_value(vec![valued_dictionary(), empty_dictionary()])] + #[case::empty_before_value(vec![empty_dictionary(), valued_dictionary()])] + #[case::null_value_after_value(vec![valued_dictionary(), null_valued_dictionary()])] + #[case::hidden_sliced_after_value(vec![ + valued_dictionary(), + sliced_empty_dictionary_with_hidden_key_payload(), + ])] + #[tokio::test] + async fn test_mixed_valued_and_all_null_dictionary_chunks(#[case] dictionaries: Vec) { + check_round_trip_encoding_of_data( + dictionaries, + &TestCases::default() + .with_structural_encodings() + .with_page_sizes(vec![4096]), + HashMap::new(), + ) + .await; + } + + #[tokio::test] + async fn test_sliced_empty_dictionary_with_hidden_key_payload() { + check_round_trip_encoding_of_data( + vec![sliced_empty_dictionary_with_hidden_key_payload()], + &TestCases::default() + .with_structural_encodings() + .with_page_sizes(vec![4096]), + HashMap::new(), + ) + .await; + } + #[test] fn test_primitive_decoder_empty_page_queue_returns_error() { let field = Arc::new(ArrowField::new("vector", DataType::Float32, true)); diff --git a/rust/lance-encoding/src/encodings/logical/primitive/dict.rs b/rust/lance-encoding/src/encodings/logical/primitive/dict.rs index beee69819bc..30d79ec7255 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/dict.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/dict.rs @@ -11,7 +11,6 @@ pub const DICT_INDICES_BITS_PER_VALUE: u64 = 32; use arrow_array::{ Array, DictionaryArray, PrimitiveArray, UInt64Array, cast::AsArray, - new_null_array, types::{ ArrowDictionaryKeyType, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, @@ -36,16 +35,6 @@ fn normalize_dict_nulls_impl( let dict_array = array.as_dictionary_opt::().expect_ok()?; - if dict_array.values().is_empty() && !dict_array.is_empty() { - // Valid non-empty dictionaries with no values can only contain null keys. The - // structural encoder removes key validity after recording it as rep/def levels, - // so retain an unreachable value at index zero to keep those exposed key bytes valid. - let values = new_null_array(dict_array.values().data_type(), 1); - return Ok( - Arc::new(DictionaryArray::new(dict_array.keys().clone(), values)) as Arc, - ); - } - if dict_array.values().null_count() == 0 { return Ok(array); } @@ -385,26 +374,9 @@ mod tests { buffer::LanceBuffer, data::{BlockInfo, FixedWidthDataBlock}, }; - use arrow_array::{Array, DictionaryArray, Int32Array, StringArray, types::Int32Type}; + use arrow_array::{Array, StringArray}; use std::sync::Arc; - #[test] - fn test_normalize_empty_all_null_dictionary() { - let keys = Int32Array::new_null(2); - let values = Arc::new(StringArray::from(Vec::<&str>::new())) as Arc; - let dictionary = Arc::new(DictionaryArray::::new(keys, values)); - - let normalized = normalize_dict_nulls(dictionary).unwrap(); - let data_without_key_nulls = normalized - .to_data() - .into_builder() - .nulls(None) - .build() - .unwrap(); - - assert_eq!(data_without_key_nulls.len(), 2); - } - #[test] fn test_dictionary_encode_abort_fixed_width() { // Create a u128 block with very high cardinality where dict encoding