fix(encoding): normalize out-of-range null dictionary keys

rust/lance-encoding/src/encodings/logical/primitive.rs

@@ -5527,6 +5527,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
@@ -7656,6 +7657,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::<Int32Type>::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::<Int32Type>::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_min_file_version(LanceFileVersion::V2_1)
+            .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::{

rust/lance-encoding/src/encodings/logical/primitive/dict.rs

@@ -110,6 +110,65 @@ pub fn normalize_dict_nulls(array: Arc<dyn Array>) -> Result<Arc<dyn Array>> {
     }
 }
 
+fn clear_out_of_range_null_keys_impl<K: ArrowDictionaryKeyType>(
+    array: Arc<dyn Array>,
+) -> Result<Arc<dyn Array>> {
+    let dict_array = array.as_dictionary_opt::<K>().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::<K>::from_iter(dict_array.keys().iter());
+    let values = dict_array.values().clone();
+    Ok(Arc::new(DictionaryArray::<K>::try_new(keys, values)?) as Arc<dyn Array>)
+}
+
+/// 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<dyn Array>) -> Result<Arc<dyn Array>> {
+    match array.data_type() {
+        DataType::Dictionary(key_type, _) => match key_type.as_ref() {
+            DataType::UInt8 => clear_out_of_range_null_keys_impl::<UInt8Type>(array),
+            DataType::UInt16 => clear_out_of_range_null_keys_impl::<UInt16Type>(array),
+            DataType::UInt32 => clear_out_of_range_null_keys_impl::<UInt32Type>(array),
+            DataType::UInt64 => clear_out_of_range_null_keys_impl::<UInt64Type>(array),
+            DataType::Int8 => clear_out_of_range_null_keys_impl::<Int8Type>(array),
+            DataType::Int16 => clear_out_of_range_null_keys_impl::<Int16Type>(array),
+            DataType::Int32 => clear_out_of_range_null_keys_impl::<Int32Type>(array),
+            DataType::Int64 => clear_out_of_range_null_keys_impl::<Int64Type>(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<T>(
     variable_width_data_block: &VariableWidthBlock,
     bits_per_offset: u8,