diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 2956d7b7f17..7e4e57a2b9e 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -6782,6 +6782,7 @@ impl PrimitiveStructuralEncoder { } DataType::Dictionary(_, _) => { array = dict::normalize_dict_nulls(array)?; + array = dict::clear_out_of_range_null_keys(array)?; Self::extract_validity_buf(array, repdef, keep_original_array) } // Extract our validity buf but NOT any child validity bufs. (they will be encoded in @@ -9760,6 +9761,43 @@ mod tests { check_round_trip_encoding_of_data(vec![arr], &test_cases, HashMap::new()).await; } + fn hand_built_dictionary_with_out_of_range_null_keys() -> ArrayRef { + use arrow_array::{DictionaryArray, Int32Array, types::Int32Type}; + use arrow_buffer::NullBuffer; + + let keys = Int32Array::new( + vec![0, 7, 7].into(), + Some(NullBuffer::from(vec![true, false, false])), + ); + let values = Arc::new(StringArray::from(vec!["a"])); + Arc::new(DictionaryArray::::try_new(keys, values).unwrap()) as ArrayRef + } + + fn concatenated_dictionary_with_out_of_range_null_keys() -> ArrayRef { + use arrow_array::{builder::StringDictionaryBuilder, new_null_array, types::Int32Type}; + + let mut builder = StringDictionaryBuilder::::new(); + builder.append_value("a"); + for _ in 0..7 { + builder.append_null(); + } + let valued = Arc::new(builder.finish()) as ArrayRef; + let all_null = new_null_array(valued.data_type(), 8); + arrow_select::concat::concat(&[valued.as_ref(), all_null.as_ref()]).unwrap() + } + + #[rstest::rstest] + #[case::hand_built(hand_built_dictionary_with_out_of_range_null_keys())] + #[case::concatenated(concatenated_dictionary_with_out_of_range_null_keys())] + #[tokio::test] + async fn test_dictionary_out_of_range_null_keys_round_trip(#[case] dictionary: ArrayRef) { + let test_cases = TestCases::default() + .with_encoding(TestEncoding::StructuralU32) + .with_page_sizes(vec![4096]); + + check_round_trip_encoding_of_data(vec![dictionary], &test_cases, HashMap::new()).await; + } + #[test] fn test_encode_decode_complex_all_null_vals_roundtrip() { use crate::compression::{DecompressionStrategy, DefaultDecompressionStrategy}; diff --git a/rust/lance-encoding/src/encodings/logical/primitive/dict.rs b/rust/lance-encoding/src/encodings/logical/primitive/dict.rs index 30d79ec7255..19582c167fb 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive/dict.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive/dict.rs @@ -110,6 +110,65 @@ pub fn normalize_dict_nulls(array: Arc) -> Result> { } } +fn clear_out_of_range_null_keys_impl( + array: Arc, +) -> Result> { + let dict_array = array.as_dictionary_opt::().expect_ok()?; + let num_values = dict_array.values().len(); + let Some(nulls) = dict_array.keys().nulls() else { + return Ok(array); + }; + + // There is no valid replacement key for an empty dictionary, so that case + // requires separate handling and must remain unchanged here. + if num_values == 0 { + return Ok(array); + } + + let has_out_of_range_null_key = dict_array + .keys() + .values() + .iter() + .zip(nulls.iter()) + .any(|(key, is_valid)| !is_valid && key.to_usize().is_none_or(|key| key >= num_values)); + if !has_out_of_range_null_key { + return Ok(array); + } + + // Building from the logical iterator writes the default physical key into + // every null slot while preserving the original validity bitmap. + let keys = PrimitiveArray::::from_iter(dict_array.keys().iter()); + let values = dict_array.values().clone(); + Ok(Arc::new(DictionaryArray::::try_new(keys, values)?) as Arc) +} + +/// Replaces out-of-range physical keys in null dictionary slots with a valid key. +/// +/// Arrow permits arbitrary keys in null slots, but the structural encoder removes +/// key validity after recording it as rep-def. The replacement keeps the array +/// valid when that null buffer is removed without changing its logical values. +pub(super) fn clear_out_of_range_null_keys(array: Arc) -> Result> { + match array.data_type() { + DataType::Dictionary(key_type, _) => match key_type.as_ref() { + DataType::UInt8 => clear_out_of_range_null_keys_impl::(array), + DataType::UInt16 => clear_out_of_range_null_keys_impl::(array), + DataType::UInt32 => clear_out_of_range_null_keys_impl::(array), + DataType::UInt64 => clear_out_of_range_null_keys_impl::(array), + DataType::Int8 => clear_out_of_range_null_keys_impl::(array), + DataType::Int16 => clear_out_of_range_null_keys_impl::(array), + DataType::Int32 => clear_out_of_range_null_keys_impl::(array), + DataType::Int64 => clear_out_of_range_null_keys_impl::(array), + _ => Err(Error::not_supported_source( + format!("Unsupported dictionary key type: {}", key_type).into(), + )), + }, + _ => Err(Error::internal(format!( + "Data type is not a dictionary: {}", + array.data_type() + ))), + } +} + fn dict_encode_variable_width( variable_width_data_block: &VariableWidthBlock, bits_per_offset: u8,